Skip to content

GitTanish/K-STATE

Repository files navigation

🧠 K-State

Corrective RAG that doesn't trust its own retriever.

Python LangGraph FAISS Groq Redis FastAPI


📋 Table of Contents


Why K-State?

Most RAG systems are a straight line: retrieve → generate → ship. That works until it doesn't — and in production, it doesn't often enough to matter.

K-State replaces that line with an intelligent loop:

Problem Traditional RAG K-State
Semantic drift Hallucinates from wrong chunks LLM-as-judge scores each chunk
Stale knowledge Stuck in indexed documents Web search fallback on low scores
Context poisoning LLM sees all retrieved text Sentence-level filtering
Repeated queries Re-processes every time Multi-tier Redis cache
Complex questions Single-shot retrieval Multi-turn agentic search

K-State is skeptical of its own retrievers. Every chunk is evaluated before reaching the LLM. Bad scores trigger web fallbacks. Everything gets filtered at the sentence level. The system validates itself at every step.


Architecture

graph TB
    Q[Query Input] --> Cache{Redis Cache?}
    Cache -->|Hit| Out[Return Response]
    Cache -->|Miss| Hybrid[Hybrid Retrieval]

    Hybrid --> Dense[FAISS Dense]
    Hybrid --> Sparse[BM25 Sparse]
    Dense --> RRF[RRF Fusion]
    Sparse --> RRF

    RRF --> Eval[CRAG Evaluator]
    Eval --> Score{Score Analysis}

    Score -->|any > 0.7| Correct[CORRECT]
    Score -->|some 0.3-0.7| Amb[AMBIGUOUS]
    Score -->|all < 0.3| Inc[INCORRECT]

    Correct --> Refine[Sentence Refinement]
    Amb --> Rewrite[Rewrite Query]
    Inc --> Rewrite

    Rewrite --> Web[Web Search]
    Web --> Refine

    Refine --> Filter[LLM Sentence Filter]
    Filter --> Gen[Generate Response]
    Gen --> Store[Store in Cache]
    Store --> Out
Loading

Data Flow

Query
  │
  ▼
Redis Cache Check ────► (HIT) ──► Return cached response
  │
  ▼ (MISS)
Hybrid Retrieval (FAISS + BM25 + RRF)
  │
  ▼
CRAG Evaluator (LLM scores 0.0-1.0 per chunk)
  │
  ├── any score > 0.7 ──────────────────────────► CORRECT  → local docs only
  │
  ├── scores between 0.3–0.7 ──► rewrite → web ──► AMBIGUOUS → merge local + web
  │
  └── all scores < 0.3 ────────► rewrite → web ──► INCORRECT → web only
  │
  ▼
Sentence Refinement (decompose → LLM filter → recompose)
  │
  ▼
Generation (LLM synthesis)
  │
  ▼
Redis Cache Store → Response

Core Components

1. Hybrid Retrieval (FAISS + BM25 + RRF)

# Combines dense and sparse retrieval via Reciprocal Rank Fusion
dense_results  = vector_store.similarity_search_with_score(query, k=60)
bm25_scores    = bm25.get_scores(tokenized_query)
fused          = _rrf_fusion(dense_results, sparse_results)  # top_k=4 returned

Dense retrieval finds semantic meaning; BM25 finds exact keyword matches. RRF fusion ranks both without manually tuning weights.

2. CRAG Evaluator (LLM as Judge)

Each retrieved chunk is scored independently with a structured-output LLM call:

Score Verdict Action
any > 0.7 CORRECT Use local chunks only
some 0.3–0.7 AMBIGUOUS Merge local + web results
all < 0.3 INCORRECT Abandon local, use web only

This prevents garbage-in-garbage-out. The generator never sees chunks that scored poorly.

3. Query Rewriting + Web Search

When verdict is AMBIGUOUS or INCORRECT, the original question is rewritten into a terse web-search query (6–14 words), then searched via DuckDuckGo with a two-layer fallback (API → DDGS scraper).

4. Sentence-Level Refinement

raw context (all source docs)
    ↓ decompose_to_sentences()    [min length: 25 chars]
["Sentence 1.", "Sentence 2.", ...]
    ↓ LLM filter (keep? yes/no per sentence)
["Sentence 1.", "Sentence 3."]
    ↓ recompose
refined_context  (passed to generator)

If the filter removes >90% of sentences, the full original context is used as fallback.

5. Multi-Turn Agentic Search

For questions longer than 5 words, K-State enters a multi-turn loop before the CRAG evaluation:

Turn 1/3: Evaluate context → "PARTIAL" (confidence: 0.32)
         → Decompose into sub-questions
         → Search each sub-question (FAISS + web)
         → Self-edit accumulated context

Turn 2/3: Evaluate → "COMPLETE" (confidence: 0.89) → stop early

Max turns is configurable (max_agentic_turns, default 3).

6. Redis Multi-Tier Cache

Tier TTL Use Case
HOT 1 hour Frequent exact matches
WARM 24 hours Less frequent queries
COLD 7 days Historical lookups
SESSION 30 min Current session (hardcoded)

Cache is optional — if Redis is unavailable, the system degrades gracefully and logs a warning.


Quick Start

Prerequisites

python --version  # 3.10+
# Redis is optional but recommended for caching

Installation

git clone https://github.com/GitTanish/K-STATE.git
cd K-STATE

# Use the project's virtual environment, or create a fresh one
pip install -r requirements.txt

Environment Setup

Create a .env file in the project root:

GROQ_API_KEY=gsk_your_key_here

Get your key from console.groq.com.

Add Your Documents

# Place PDF files in the books/ directory
cp /path/to/your/pdfs/*.pdf books/

Run

# Terminal 1 (optional): Redis for caching
redis-server

# Terminal 2: FastAPI server
python api.py
# Server starts at http://localhost:8000

# Terminal 3: CLI (connects to local server via direct graph invocation)
python main.py --mode cli

First-run note: If faiss_index/ doesn't exist, the startup will build it from your PDFs (2–5 minutes for large corpora). The index is saved and reused on subsequent runs.

Modes

Mode Command Description
cli (default) python main.py Interactive CLI, calls graph directly
api python main.py --mode api Interactive CLI, calls running API server
direct python main.py --mode direct --question "..." Single question, graph direct, JSON output

API Documentation

Base URL: http://localhost:8000

GET /

Health check.

curl http://localhost:8000/
{
  "service": "CRAG API",
  "status": "healthy",
  "cache_enabled": false,
  "version": "2.0.0"
}

GET /health

curl http://localhost:8000/health
{
  "status": "healthy",
  "redis": false,
  "graph": "loaded"
}

POST /query

Submit a question.

curl -X POST http://localhost:8000/query \
  -H "Content-Type: application/json" \
  -d '{"question": "What is batch normalization?", "use_cache": true}'
{
  "question": "What is batch normalization?",
  "answer": "Batch Normalization (BatchNorm) is a layer inserted into deep neural networks...",
  "verdict": "WEB_SEARCH",
  "reason": "All retrieved chunks scored < 0.3.",
  "agentic_turns": 0,
  "cached": false,
  "latency_ms": 7537.8,
  "tokens_used": null
}

Verdict values:

Value Meaning
CORRECT Local chunks scored well; answer from local context
AMBIGUOUS Mixed scores; local + web merged
INCORRECT All chunks scored poorly; web only attempted
WEB_SEARCH Web search was executed (covers AMBIGUOUS + INCORRECT paths)

POST /query/stream

Stream response as Server-Sent Events.

import requests

response = requests.post(
    "http://localhost:8000/query/stream",
    json={"question": "What is attention?"},
    stream=True
)
for line in response.iter_lines():
    if line:
        print(line.decode())

Events: {"type": "start"}{"type": "token", "content": "..."} × N → {"type": "end"}

GET /cache/stats

curl http://localhost:8000/cache/stats

Returns {"enabled": false} when Redis is not running; otherwise returns per-tier entry counts and sample keys.

DELETE /cache/{tier}

curl -X DELETE http://localhost:8000/cache/hot

Valid tiers: hot, warm, cold, session.


Configuration

All settings live in config.py as a dataclass:

@dataclass
class Config:
    # CRAG thresholds
    upper_th: float = 0.7        # Score above this → CORRECT
    lower_th: float = 0.3        # Score below this → INCORRECT

    # Retrieval
    top_k: int = 4               # Final chunks returned to evaluator
    hybrid_candidates: int = 60  # Candidates fetched per retriever

    # Redis
    redis_host: str = "localhost"
    redis_port: int = 6379
    redis_db: int = 0
    cache_ttl_hot: int = 3600    # 1 hour
    cache_ttl_warm: int = 86400  # 24 hours
    cache_ttl_cold: int = 604800 # 7 days

    # Agentic search
    max_agentic_turns: int = 3

    # Paths
    faiss_persist_dir: str = "./faiss_index"
    books_dir: str = "books"

    # Models
    embedding_model: str = "BAAI/bge-small-en-v1.5"
    llm_model: str = "openai/gpt-oss-120b"
    llm_temperature: float = 0.0

Tuning Recommendations

Use Case upper_th lower_th max_agentic_turns
Factual Q&A 0.8 0.4 1
Research/Exploration 0.6 0.2 3
News/Current Events 0.5 0.1 2
Technical Documentation 0.7 0.3 2

Performance

Observed Latency (tested)

Operation Observed
Server startup ~2 min (FAISS load + BM25 index build on 3 large PDFs)
CORRECT path ~3–5s
WEB_SEARCH path ~7–10s
Cache hit ~50ms (when Redis running)

Note: Latency depends heavily on Groq's inference speed and network conditions. The agentic path makes multiple LLM calls per turn.

Resource Usage

Resource Notes
Memory ~350MB for FAISS index (3 large PDFs)
Disk ~6MB for faiss_index/ (index.faiss + index.pkl)
CPU Moderate — embedding runs on CPU
GPU Not required

Benchmarking & Evaluation

  • Iterations: ran the tests/golden_set.py benchmark after incremental improvements (adaptive thresholds, evaluator calibration, retriever routing).
  • Latest golden set: 7 questions — 5 passed, 2 failed (pass rate ~71%).
  • Metrics and logs:
    • Query logs: metrics/queries.jsonl
    • Evaluator calibration traces: metrics/calibration.jsonl

Reproduce the golden-set benchmark locally:

# activate venv (Windows PowerShell)
.venv\Scripts\Activate.ps1
setx GROQ_MODEL_POOL "llama-3.1-8b-instant,openai/gpt-oss-120b"
python -u -c "from tests.golden_set import GoldenSet; from evaluation.experience_store import metrics_logger; print(GoldenSet().run(verbose=False, delay_seconds=1.0)); print(metrics_logger.get_health_check()); print(metrics_logger.analyze_thresholds())"

Notes:

  • If you change embeddings or retrievers, rebuild the index files: faiss_index/, bm25_index.pkl, chunks.pkl.
  • To experiment with thresholds, edit config.py (upper_th, lower_th) or rely on the adaptive router in evaluation/router_optimizer.py.
  • Benchmark corpus (documents used in the golden-set runs):
    • books/book1.pdf — "Deep Learning" by Ian Goodfellow et al.
    • books/book2.pdf — "Pattern Recognition and Machine Learning" by Christopher Bishop
    • books/book3.pdf — "Hands-On Machine Learning with Scikit-Learn and TensorFlow" by Aurélien Géron
    • books/sample.txt — auxiliary sample content

Results Summary

Metric Value
Golden set size 7 questions
Pass rate 71% (5/7)
CORRECT path (local only) 4
AMBIGUOUS path (merged local+web) 1
INCORRECT / WEB_FALLBACK 2
Avg latency (CORRECT path) ~3–5s
Avg latency (INCORRECT / WEB path) ~7–10s

Why the 2 failures occurred

  • Case 1 — What is layer normalization?: routed to web search in one run because the retrieved local chunks had low evaluator scores; investigation suggests incomplete chunk coverage for this topic in the indexed corpus (index coverage issue).
  • Case 2 — Batch norm vs layer norm (comparison): classified as CORRECT rather than the expected AMBIGUOUS because the evaluator favored a synthesized comparative answer when the benchmark expected a conservative ambiguous verdict (evaluator calibration / decision policy mismatch).

Notes: these failure explanations are intentionally brief — they point to root causes (index quality vs evaluator policy) we can act on: rebuild/index more targeted sections for layer normalization, or adjust evaluator few‑shot examples / decision rules for comparison questions.

Project Structure

K-STATE/
├── api.py            # FastAPI server (REST + SSE streaming)
├── caching.py        # Redis multi-tier cache manager
├── config.py         # Centralized configuration dataclass
├── graph.py          # LangGraph pipeline (nodes + edges)
├── main.py           # CLI entry point (Rich terminal UI)
├── models.py         # Pydantic models + LangGraph State TypedDict
├── retrieval.py      # HybridRetriever, AgenticSearchLoop, load_vector_store
├── requirements.txt  # Python dependencies
├── .env              # API keys (not committed)
├── books/            # PDF source documents
│   ├── book1.pdf
│   ├── book2.pdf
│   └── book3.pdf
└── faiss_index/      # Persisted FAISS vector store (auto-created)
    ├── index.faiss
    └── index.pkl

Stack

Layer Technology Role
Orchestration LangGraph Stateful DAG with conditional routing
Embeddings BAAI/bge-small-en-v1.5 Local, CPU-friendly, 384-dim
Dense Retrieval FAISS Persisted vector index
Sparse Retrieval BM25 (rank_bm25) Keyword matching + RRF fusion
LLM openai/gpt-oss-120b via Groq Evaluator, generator, rewriter, filter
Web Search DuckDuckGo + DDGS fallback No API key required
Caching Redis Optional multi-tier TTL cache
API FastAPI + Uvicorn REST + SSE, auto-reload in dev
CLI Rich Terminal table + interactive prompt

Known Limitations

Latency

Not built for sub-second responses. Each LLM call adds 500–800ms on Groq. The AMBIGUOUS/INCORRECT paths make 3–4 LLM calls total. Pre-warm common queries with a cache to reduce perceived latency.

Sentence Filter Aggression

Dense technical content with cross-references can get over-filtered. The filter falls back to the full sentence list if >90% is removed, but you may want to adjust the filter prompt in graph.py → refine() for your domain.

Web Search Reliability

DuckDuckGo has no formal SLA. The two-layer fallback (API wrapper → DDGS scraper) handles most rate limits. For production, consider swapping in Tavily:

from langchain_community.tools.tavily_search import TavilySearchResults

Redis is Optional, Not Integrated by Default

If Redis isn't running, the system works without cache. Startup will log a warning and cache_enabled will be false.

tokens_used Not Populated

The tokens_used field in QueryResponse is always null currently — Groq usage tracking is not implemented.


Troubleshooting

Server takes 2+ minutes to start

This is normal on first run (or after model cache is cleared). The startup sequence:

  1. Download/load BAAI/bge-small-en-v1.5 embeddings (~35MB)
  2. Load FAISS index from disk
  3. Build BM25 index over all chunks in memory
  4. Connect to Redis (graceful fail if unavailable)
  5. Compile LangGraph DAG

Subsequent startups are faster once HuggingFace weights are cached locally.

FAISS index not found

RuntimeError: Books directory not found: books
# or
RuntimeError: No PDF files found in books

Ensure your PDFs are in the books/ directory relative to where you run the server from.

Groq auth error

401 Unauthorized

Check .env contains a valid key: GROQ_API_KEY=gsk_...

Redis warning on startup

Warning: Redis unavailable - caching disabled

This is non-fatal. Start redis-server to enable caching, or ignore it.

DuckDuckGo rate limiting

The system automatically falls back from the DuckDuckGo API wrapper to the DDGS scraper. If both fail, the web_docs list will be empty and the response will fall through to whatever context was already gathered.


License

MIT License. See LICENSE for details.


Citation

@software{kstate2026,
  title   = {K-State: Corrective Retrieval-Augmented Generation with Agentic Search},
  author  = {Tanish},
  year    = {2026},
  url     = {https://github.com/GitTanish/K-STATE}
}

Acknowledgments


Built with skepticism. Runs on skepticism.

Report Bug · Request Feature · Star on GitHub


About

Stateful Corrective RAG (CRAG) built with LangGraph. Features hybrid retrieval, an LLM judge, and multi-turn web search fallbacks.

Resources

Stars

Watchers

Forks

Releases

Packages

Contributors

Languages