feat: CRAG pipeline hardening — hybrid retrieval, citations, config unification#8
feat: CRAG pipeline hardening — hybrid retrieval, citations, config unification#8dtype2100 wants to merge 2 commits into
Conversation
…nification - Add ChunkHit helpers and fix dict/string type mismatches in policies/evaluators - Wire hybrid_retrieve to hybrid_retriever + BM25 corpus with multi-query RRF - Enable parent-child ingest, corpus registry, and metadata-aware expansion - Return structured SourceCitation objects from /query API - Consolidate env settings (VECTOR_BACKEND, RERANKER, MULTI_QUERY, API_KEY) - Add optional API key auth, LLM health probe, and SSE /query/stream endpoint - Expand tests (72 passing), Makefile PATH fix, README and .env.example update Co-authored-by: jwlee <dtype2100@users.noreply.github.com>
|
Important Review skippedDraft detected. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request upgrades the RAG pipeline to a self-corrective RAG (CRAG) architecture, introducing hybrid retrieval (vector + BM25), parent-child chunking, context expansion, and an SSE streaming endpoint, alongside optional API key authentication. Feedback on these changes highlights a critical process isolation issue where the in-memory registry fails to sync across async ARQ workers or multi-worker API setups, suggesting Redis as a shared store. Additionally, the review identifies a performance bottleneck in the BM25 corpus registry, an event loop blocking issue in the async health endpoint due to synchronous HTTP calls, and opportunities for defensive programming when handling missing metadata or non-dictionary retrieval results.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| docstore = {p["metadata"]["chunk_id"]: p["text"] for p in parents} | ||
| register_docstore(docstore) |
There was a problem hiding this comment.
⚠️ Process Isolation Issue with In-Memory Registry
Calling register_docstore and register_documents updates an in-memory registry (_docstore and _corpus_docs) inside the current process. This introduces a major architectural issue in production:
- ARQ Worker Isolation: When asynchronous ingestion is enabled (
INGEST_QUEUE_ASYNC=true), the ingestion pipeline runs in the background ARQ worker process (make worker). The worker will update its own in-memory registry, but the API worker process (handling/api/v1/query) will remain empty. As a result, BM25 hybrid retrieval and parent-child expansion will fail on the query path. - Multi-Worker API: Even with synchronous ingestion, if the FastAPI app is run with multiple Uvicorn workers (e.g.,
--workers 4), the in-memory state will not be shared or synchronized across the processes.
Suggested Resolution
To make this production-ready, you should persist and share this state across processes. Since Redis is already required for the async queue, you can use Redis as the shared store:
- Store the parent docstore in Redis using a Hash (
HSET/HGET). - For BM25, consider persisting the serialized BM25 index/corpus to Redis or a shared volume, or leverage a database/search backend that natively supports hybrid search (like Qdrant's hybrid search or pgvector) instead of an in-memory BM25 registry.
| _corpus_docs: list[dict[str, Any]] = [] | ||
| _docstore: dict[str, str] = {} | ||
|
|
||
|
|
||
| def register_documents(docs: list[dict[str, Any]]) -> None: | ||
| """Append indexed chunks to the BM25 corpus, deduplicating by text.""" | ||
| seen = {d["text"] for d in _corpus_docs} | ||
| for doc in docs: | ||
| text = doc.get("text", "") | ||
| if text and text not in seen: | ||
| seen.add(text) | ||
| _corpus_docs.append( | ||
| { | ||
| "text": text, | ||
| "score": 0.0, | ||
| "metadata": dict(doc.get("metadata") or {}), | ||
| } | ||
| ) | ||
|
|
||
|
|
||
| def register_docstore(entries: dict[str, str]) -> None: | ||
| """Merge parent-id → parent-text mappings into the docstore.""" | ||
| _docstore.update(entries) | ||
|
|
||
|
|
||
| def get_corpus_docs() -> list[dict[str, Any]]: | ||
| """Return a snapshot of the in-memory BM25 corpus.""" | ||
| return list(_corpus_docs) | ||
|
|
||
|
|
||
| def get_docstore() -> dict[str, str]: | ||
| """Return a snapshot of the parent docstore.""" | ||
| return dict(_docstore) | ||
|
|
||
|
|
||
| def clear() -> None: | ||
| """Reset corpus and docstore (primarily for tests).""" | ||
| _corpus_docs.clear() | ||
| _docstore.clear() |
There was a problem hiding this comment.
⚡ Performance Bottleneck: $O(N^2)$ Batch Ingestion
Currently, register_documents builds a new set of all existing document texts (seen = {d["text"] for d in _corpus_docs}) on every single call. As the corpus grows, this set creation becomes increasingly expensive, leading to
We can optimize this to _seen_texts to track already registered document texts.
_corpus_docs: list[dict[str, Any]] = []
_docstore: dict[str, str] = {}
_seen_texts: set[str] = set()
def register_documents(docs: list[dict[str, Any]]) -> None:
"""Append indexed chunks to the BM25 corpus, deduplicating by text."""
for doc in docs:
text = doc.get("text", "")
if text and text not in _seen_texts:
_seen_texts.add(text)
_corpus_docs.append(
{
"text": text,
"score": 0.0,
"metadata": dict(doc.get("metadata") or {}),
}
)
def register_docstore(entries: dict[str, str]) -> None:
"""Merge parent-id → parent-text mappings into the docstore."""
_docstore.update(entries)
def get_corpus_docs() -> list[dict[str, Any]]:
"""Return a snapshot of the in-memory BM25 corpus."""
return list(_corpus_docs)
def get_docstore() -> dict[str, str]:
"""Return a snapshot of the parent docstore."""
return dict(_docstore)
def clear() -> None:
"""Reset corpus and docstore (primarily for tests)."""
_corpus_docs.clear()
_docstore.clear()
_seen_texts.clear()| def probe_llm() -> str: | ||
| """Return a short status string for the configured LLM backend.""" | ||
| if settings.llm_backend == "openai": | ||
| if not settings.openai_api_key: | ||
| return "not_configured" | ||
| return "configured" | ||
|
|
||
| try: | ||
| url = f"{settings.vllm_base_url.rstrip('/')}/models" | ||
| with httpx.Client(timeout=2.0) as client: | ||
| response = client.get(url) | ||
| if response.status_code == 200: | ||
| return "ready" | ||
| return f"error: HTTP {response.status_code}" | ||
| except Exception as exc: | ||
| logger.debug("LLM health probe failed: %s", exc) | ||
| return f"unreachable: {exc}" |
There was a problem hiding this comment.
🔄 Event Loop Blocking in Async Endpoint
The probe_llm function performs a synchronous HTTP request using httpx.Client with a 2.0-second timeout. Since this function is called directly inside the asynchronous /health endpoint (async def health()), it will block the ASGI event loop for up to 2 seconds if the LLM backend is slow or unreachable. This can degrade the performance of the entire application.
We should make probe_llm asynchronous using httpx.AsyncClient so it can be awaited without blocking the event loop.
| def probe_llm() -> str: | |
| """Return a short status string for the configured LLM backend.""" | |
| if settings.llm_backend == "openai": | |
| if not settings.openai_api_key: | |
| return "not_configured" | |
| return "configured" | |
| try: | |
| url = f"{settings.vllm_base_url.rstrip('/')}/models" | |
| with httpx.Client(timeout=2.0) as client: | |
| response = client.get(url) | |
| if response.status_code == 200: | |
| return "ready" | |
| return f"error: HTTP {response.status_code}" | |
| except Exception as exc: | |
| logger.debug("LLM health probe failed: %s", exc) | |
| return f"unreachable: {exc}" | |
| async def probe_llm() -> str: | |
| """Return a short status string for the configured LLM backend.""" | |
| if settings.llm_backend == "openai": | |
| if not settings.openai_api_key: | |
| return "not_configured" | |
| return "configured" | |
| try: | |
| url = f"{settings.vllm_base_url.rstrip('/')}/models" | |
| async with httpx.AsyncClient(timeout=2.0) as client: | |
| response = await client.get(url) | |
| if response.status_code == 200: | |
| return "ready" | |
| return f"error: HTTP {response.status_code}" | |
| except Exception as exc: | |
| logger.debug("LLM health probe failed: %s", exc) | |
| return f"unreachable: {exc}" |
| status="ok", | ||
| llm_backend=settings.llm_backend, | ||
| llm_model=settings.llm_model, | ||
| llm=probe_llm(), |
There was a problem hiding this comment.
Await the asynchronous probe_llm() function to prevent blocking the event loop.
| llm=probe_llm(), | |
| llm=await probe_llm(), |
| for child in children: | ||
| parent_id = child["metadata"].get("parent_id") | ||
| if parent_id and parent_id in docstore: | ||
| child["metadata"]["parent_text"] = docstore[parent_id] |
There was a problem hiding this comment.
🛡️ Defensive Programming: Safe Metadata Access
If child does not contain a "metadata" key, or if "metadata" is None, accessing child["metadata"].get("parent_id") will raise a KeyError or AttributeError. We should safely access the metadata dictionary.
| for child in children: | |
| parent_id = child["metadata"].get("parent_id") | |
| if parent_id and parent_id in docstore: | |
| child["metadata"]["parent_text"] = docstore[parent_id] | |
| for child in children: | |
| meta = child.get("metadata") | |
| if meta is None: | |
| meta = {} | |
| child["metadata"] = meta | |
| parent_id = meta.get("parent_id") | |
| if parent_id and parent_id in docstore: | |
| meta["parent_text"] = docstore[parent_id] |
| def build_source_citations(results: list[dict[str, Any]]) -> list[SourceCitation]: | ||
| """Convert retrieval dicts into structured ``SourceCitation`` models. | ||
|
|
||
| Args: | ||
| results: List of ``{text, score, metadata}`` retrieval hits. | ||
|
|
||
| Returns: | ||
| Ordered list of citations for the API response. | ||
| """ | ||
| citations: list[SourceCitation] = [] | ||
| for result in results: | ||
| meta = result.get("metadata") or {} | ||
| citations.append( | ||
| SourceCitation( | ||
| text=str(result.get("text", "")), | ||
| source=str(meta.get("source", "unknown")), | ||
| page=str(meta.get("page", "")), | ||
| score=float(result.get("score", 0.0)), | ||
| ) | ||
| ) | ||
| return citations |
There was a problem hiding this comment.
🛡️ Defensive Programming: Handle Non-Dict Results
To prevent potential runtime AttributeErrors (e.g., if a retriever or mock returns plain strings instead of dictionaries), we should defensively handle cases where a result is a string rather than a dictionary.
def build_source_citations(results: list[str | dict[str, Any]]) -> list[SourceCitation]:
"""Convert retrieval dicts into structured ``SourceCitation`` models.
Args:
results: List of ``{text, score, metadata}`` retrieval hits.
Returns:
Ordered list of citations for the API response.
"""
citations: list[SourceCitation] = []
for result in results:
if isinstance(result, str):
citations.append(SourceCitation(text=result))
continue
meta = result.get("metadata") or {}
citations.append(
SourceCitation(
text=str(result.get("text", "")),
source=str(meta.get("source", "unknown")),
page=str(meta.get("page", "")),
score=float(result.get("score", 0.0)),
)
)
return citations- Route low grounding scores through run_judge before retry (hybrid loop) - Add mark_rejected node and wire judge outcomes in CRAG graph - Redis-backed chat history with in-memory fallback; session_id multi-turn - Include chat_history in generate_answer prompt context - Add React Web UI (Chat + Studio), CORS, and /api/v1/studio endpoints - Add eval smoke tests (tests/evals/) and make test-evals target - 79 tests passing Co-authored-by: jwlee <dtype2100@users.noreply.github.com>
2 participants
Summary
Implements the recommended improvement roadmap for the Advanced-RAG CRAG pipeline, closing the gap between designed architecture and the live query path.
P0 — Runtime fixes
ChunkHithelpers (app/rag/types.py) to unify dict/string chunk handlingexpansion_policyandgrounding_evaluatortype mismatchesSourceCitationobjects from/api/v1/queryinstead of raw dictsP1 — Retrieval & ingest integration
hybrid_retrievenode tohybrid_retriever(vector + BM25 RRF) with optional multi-query fusioncorpus_registrysynced during ingest for BM25 corpus + parent docstoretop_kfrom API request into graph stateP1 — Configuration consolidation
VECTOR_BACKEND,RERANKER_BACKEND,MULTI_QUERY,USE_PARENT_CHILD_CHUNKING,API_KEY, and judge LLM overrides intoSettings.env.exampleP2 — Production readiness
X-API-Key)/api/v1/health(llmfield)POST /api/v1/query/streampytest/ruffin~/.local/binP3 — Documentation
pyproject.toml(bm25,reranker,loaders)Test plan
make lintpassesmake test— 72 tests passingpip install -e ".[bm25]")POST /api/v1/query/stream