Improve RAG system with advanced retrieval techniques - #11
Open
ThomasJButler wants to merge 1 commit into
Open
Conversation
Key improvements: - Add missing reranker config settings (use_reranker, reranker_model) - Add multi-query expansion settings (enable_multi_query, multi_query_count) - Integrate cross-encoder reranker into HybridRAG retrieval pipeline - Create EnhancedRetriever with multi-query expansion and RRF fusion - Add semantic chunking with markdown header awareness - Add parent-child chunking for better context retrieval - Add contextual chunking with document-level context - Add contextual compression to extract relevant excerpts - Update module exports in __init__.py files New features: - SemanticChunker: Structure-aware chunking respecting markdown headers - ParentChildChunker: Small chunks for retrieval, large parents for context - ContextualChunker: Prepends document context for better embeddings - EnhancedRetriever: Full pipeline with query rewriting, multi-query, RRF, reranking https://claude.ai/code/session_01BRVpxN2kXDmfHbPdkdUhov
|
The latest updates on your projects. Learn more about Vercel for GitHub.
|
📝 WalkthroughWalkthroughThis PR extends RAG capabilities by introducing advanced retrieval features: a new EnhancedRetriever module implementing multi-query expansion with Reciprocal Rank Fusion, reranker integration in HybridRAG, and enhanced chunking strategies (semantic, parent-child, contextual). Configuration adds reranker model settings and multi-query controls. Changes
Sequence DiagramsequenceDiagram
participant Client
participant EnhancedRetriever
participant QueryVariator
participant Embedder
participant PineconeIndex
participant RRFFusion
participant Reranker
participant Compressor
Client->>EnhancedRetriever: retrieve_enhanced(query)
EnhancedRetriever->>QueryVariator: generate_query_variations(query)
QueryVariator-->>EnhancedRetriever: [query, var1, var2, ...]
EnhancedRetriever->>Embedder: embed_queries_batch(queries)
Embedder-->>EnhancedRetriever: [[embedding1], [embedding2], ...]
loop For each query variation
EnhancedRetriever->>PineconeIndex: search_single_query(query, embedding, top_k)
PineconeIndex-->>EnhancedRetriever: [results]
end
EnhancedRetriever->>RRFFusion: reciprocal_rank_fusion(result_lists)
RRFFusion-->>EnhancedRetriever: [fused_results]
alt if use_compression
EnhancedRetriever->>Compressor: compress_context(query, contexts)
Compressor-->>EnhancedRetriever: [compressed_contexts]
end
alt if use_reranker
EnhancedRetriever->>Reranker: rerank_contexts(query, contexts)
Reranker-->>EnhancedRetriever: [reranked_contexts]
end
EnhancedRetriever-->>Client: (contexts, metrics)
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Poem
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing touches
🧪 Generate unit tests (beta)
Important Action Needed: IP Allowlist UpdateIf your organization protects your Git platform with IP whitelisting, please add the new CodeRabbit IP address to your allowlist:
Reviews will stop working after February 8, 2026 if the new IP is not added to your allowlist. 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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
backend/app/rag/hybrid.py (1)
332-379:⚠️ Potential issue | 🟠 MajorEnsure reranked results respect requested
top_k.If
rerank_top_kexceeds the requestedtop_k(or a caller passes a smallertop_k), the rerank path can return more results than requested. Cap reranking totop_kand/or slice after filtering.🩹 Proposed fix
- # 5. Take top-k after merging (get extra for reranking) - merged_results = merged_results[:top_k * 2] + # 5. Take top-k after merging (get extra for reranking) + merged_results = merged_results[:top_k * 2] + rerank_limit = min(top_k, settings.rerank_top_k) @@ - reranked_results = await self.reranker.rerank( - query, contexts_for_rerank, top_k=settings.rerank_top_k - ) + reranked_results = await self.reranker.rerank( + query, contexts_for_rerank, top_k=rerank_limit + ) @@ - filtered_results = [ + filtered_results = [ r for r in merged_results if r["score"] >= settings.min_relevance_score ] + filtered_results = filtered_results[:top_k]
🤖 Fix all issues with AI agents
In `@backend/app/rag/enhanced_retriever.py`:
- Around line 90-110: In generate_query_variations, ensure empty lists from the
rewriter don't propagate: after awaiting self.query_rewriter.expand_query(query,
num_expansions=num_variations) check if the returned variations is falsy or has
length 0 and if so return [query] as a fallback; keep the existing logging and
exception handling (logger.info/logger.error) but log when falling back due to
empty variations so callers of generate_query_variations (and downstream batch
embedding) always receive at least one query.
- Around line 196-269: The metrics.reranked flag currently copies
settings.use_reranker even when reranking was never applied or failed; update
retrieve_with_multi_query to set reranked based on the actual rerank outcome:
initialize a local boolean (e.g., did_rerank = False), only set it True after a
successful await self.reranker.rerank(...) call that returns non-empty results
(and catch exceptions around the reranker call to log the error and leave
did_rerank False), ensure that when top_results is empty you skip reranking and
leave did_rerank False, and finally assign metrics.reranked = did_rerank instead
of using settings.use_reranker.
In `@backend/app/utils/chunking.py`:
- Around line 314-320: The paragraph grouping logic incorrectly omits the
separator length when checking if adding para would exceed self.chunk_size:
update the check that uses current_length, para_length, current_group, and
chunk_size so it accounts for the separator ("\n\n") length (2) when
current_group is non-empty (e.g., calculate sep = 2 if current_group else 0 and
use current_length + para_length + sep > self.chunk_size), ensuring the same
separator length is used in both the check and when updating
current_length/current_group to prevent chunks from exceeding chunk_size.
🧹 Nitpick comments (5)
backend/app/utils/chunking.py (4)
465-476: Consider memory implications of storingparent_textin every child chunk.Each child chunk stores the full
parent_text, which means the parent content is duplicated N times in memory (once per child). For large documents, this can significantly increase memory usage.Consider storing only the
parent_idin children and looking up the parent text when needed, or making the inclusion ofparent_textoptional.♻️ Suggested approach - make parent_text inclusion optional
def __init__( self, child_chunk_size: int = 400, parent_chunk_size: Optional[int] = None, chunk_overlap: Optional[int] = None, + include_parent_text: bool = True, ): + self.include_parent_text = include_parent_text ... # In chunk_text method: child_data = { "text": child_text, "parent_id": parent_id, - "parent_text": parent_text, # Include parent for retrieval context "chunk_index": child_index, "chunk_type": "child", } + if self.include_parent_text: + child_data["parent_text"] = parent_text
518-560: Consider adding error handling for consistency with other chunkers.Unlike
SemanticChunker.chunk_textandParentChildChunker.chunk_text, this method lacks try/except error handling. For consistency and robustness, consider wrapping the chunking logic.♻️ Suggested improvement
def chunk_with_context( self, text: str, document_context: str, metadata: Optional[dict] = None, ) -> List[dict]: ... if not text or not text.strip(): return [] + try: chunks = self.text_splitter.split_text(text) chunked_docs = [] # ... rest of implementation ... logger.info(f"Contextual chunking produced {len(chunked_docs)} chunks") return chunked_docs + except Exception as e: + logger.error(f"Error in contextual chunking: {e}", exc_info=True) + # Fallback to basic chunking + chunker = DocumentChunker(self.chunk_size, self.chunk_overlap) + return chunker.chunk_text(text, metadata)
612-634: Missingchunk_overlapparameter in convenience function.The
chunk_text_parent_childfunction doesn't accept or pass achunk_overlapparameter, unlikechunk_text_semantic. SinceParentChildChunkersupports this parameter, consider exposing it for consistency.♻️ Proposed fix
def chunk_text_parent_child( text: str, child_chunk_size: int = 400, parent_chunk_size: Optional[int] = None, + chunk_overlap: Optional[int] = None, metadata: Optional[dict] = None, ) -> Tuple[List[dict], List[dict]]: """ Parent-child chunking function. Creates small chunks for retrieval that reference larger parent chunks for context. Args: text: Text to chunk child_chunk_size: Size of retrieval chunks parent_chunk_size: Size of context chunks + chunk_overlap: Overlap between parent chunks metadata: Optional metadata Returns: Tuple of (child_chunks, parent_chunks) """ - chunker = ParentChildChunker(child_chunk_size, parent_chunk_size) + chunker = ParentChildChunker(child_chunk_size, parent_chunk_size, chunk_overlap) return chunker.chunk_text(text, metadata)
229-237: Markdown detection patterns could be refined.A few minor observations on the regex patterns:
- The bold pattern
r'^\*\*.*\*\*'and italic patternr'^\*.*\*'overlap (italic will also match bold lines)- Unordered list detection only catches
-but markdown also supports*and+as list markersThis won't break functionality but may slightly affect detection accuracy.
♻️ Refined patterns (optional)
markdown_patterns = [ r'^#{1,6}\s+', # Headers - r'^\*\*.*\*\*', # Bold - r'^\*.*\*', # Italic - r'^\-\s+', # Unordered list + r'^\*\*[^*].*\*\*', # Bold (non-greedy, excludes italic) + r'^\*[^*].*\*$', # Italic (single asterisk) + r'^[-*+]\s+', # Unordered list (-, *, +) r'^\d+\.\s+', # Ordered list r'^```', # Code blocks r'^\|.*\|', # Tables ]backend/app/core/config.py (1)
78-84: Guardmulti_query_countagainst zero/negative values.A non‑positive value can lead to empty query lists and downstream embedding failures. Adding a constraint keeps the setting safe by default.
✅ Suggested constraint
- multi_query_count: int = Field( - default=3, description="Number of query variations to generate" - ) + multi_query_count: int = Field( + default=3, ge=1, description="Number of query variations to generate" + )
Comment on lines
+90
to
+110
| async def generate_query_variations( | ||
| self, query: str, num_variations: int = 3 | ||
| ) -> List[str]: | ||
| """ | ||
| Generate multiple query variations for better recall. | ||
|
|
||
| Uses LLM to create semantically similar but differently phrased queries. | ||
| """ | ||
| if not settings.enable_multi_query: | ||
| return [query] | ||
|
|
||
| try: | ||
| # Use the query rewriter's expand function | ||
| variations = await self.query_rewriter.expand_query( | ||
| query, num_expansions=num_variations | ||
| ) | ||
| logger.info(f"Generated {len(variations)} query variations") | ||
| return variations | ||
| except Exception as e: | ||
| logger.error(f"Query variation generation failed: {e}") | ||
| return [query] |
There was a problem hiding this comment.
Handle empty variation lists from the rewriter.
If the rewriter returns an empty list, the batch embedding call can fail. Add a fallback to [query] when no variations are produced.
✅ Suggested fix
- variations = await self.query_rewriter.expand_query(
- query, num_expansions=num_variations
- )
+ variations = await self.query_rewriter.expand_query(
+ query, num_expansions=num_variations
+ )
+ if not variations:
+ logger.warning("Query rewriter returned no variations; using original query.")
+ return [query]📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| async def generate_query_variations( | |
| self, query: str, num_variations: int = 3 | |
| ) -> List[str]: | |
| """ | |
| Generate multiple query variations for better recall. | |
| Uses LLM to create semantically similar but differently phrased queries. | |
| """ | |
| if not settings.enable_multi_query: | |
| return [query] | |
| try: | |
| # Use the query rewriter's expand function | |
| variations = await self.query_rewriter.expand_query( | |
| query, num_expansions=num_variations | |
| ) | |
| logger.info(f"Generated {len(variations)} query variations") | |
| return variations | |
| except Exception as e: | |
| logger.error(f"Query variation generation failed: {e}") | |
| return [query] | |
| async def generate_query_variations( | |
| self, query: str, num_variations: int = 3 | |
| ) -> List[str]: | |
| """ | |
| Generate multiple query variations for better recall. | |
| Uses LLM to create semantically similar but differently phrased queries. | |
| """ | |
| if not settings.enable_multi_query: | |
| return [query] | |
| try: | |
| # Use the query rewriter's expand function | |
| variations = await self.query_rewriter.expand_query( | |
| query, num_expansions=num_variations | |
| ) | |
| if not variations: | |
| logger.warning("Query rewriter returned no variations; using original query.") | |
| return [query] | |
| logger.info(f"Generated {len(variations)} query variations") | |
| return variations | |
| except Exception as e: | |
| logger.error(f"Query variation generation failed: {e}") | |
| return [query] |
🤖 Prompt for AI Agents
In `@backend/app/rag/enhanced_retriever.py` around lines 90 - 110, In
generate_query_variations, ensure empty lists from the rewriter don't propagate:
after awaiting self.query_rewriter.expand_query(query,
num_expansions=num_variations) check if the returned variations is falsy or has
length 0 and if so return [query] as a fallback; keep the existing logging and
exception handling (logger.info/logger.error) but log when falling back due to
empty variations so callers of generate_query_variations (and downstream batch
embedding) always receive at least one query.
Comment on lines
+196
to
+269
| async def retrieve_with_multi_query( | ||
| self, | ||
| query: str, | ||
| top_k: Optional[int] = None, | ||
| namespace: str = None, | ||
| ) -> Tuple[List[dict], RetrievalMetrics]: | ||
| """ | ||
| Enhanced retrieval using multi-query expansion and RRF. | ||
|
|
||
| Steps: | ||
| 1. Generate query variations | ||
| 2. Embed all queries in batch | ||
| 3. Search with each query in parallel | ||
| 4. Combine results with RRF | ||
| 5. Apply reranking (if enabled) | ||
| """ | ||
| start_time = time.time() | ||
| top_k = top_k or settings.top_k_results | ||
| namespace = namespace or "default" | ||
|
|
||
| try: | ||
| # 1. Generate query variations | ||
| queries = await self.generate_query_variations( | ||
| query, num_variations=settings.multi_query_count | ||
| ) | ||
|
|
||
| # 2. Batch embed all queries | ||
| embeddings = await self.embed_queries_batch(queries) | ||
|
|
||
| # 3. Search with each query in parallel | ||
| search_tasks = [ | ||
| self.search_single_query(q, emb, top_k, namespace) | ||
| for q, emb in zip(queries, embeddings) | ||
| ] | ||
| result_lists = await asyncio.gather(*search_tasks) | ||
|
|
||
| # 4. Combine results with RRF | ||
| combined_results = self.reciprocal_rank_fusion(result_lists) | ||
|
|
||
| # 5. Take top results | ||
| top_results = combined_results[:top_k * 2] # Extra for reranking | ||
|
|
||
| # 6. Apply reranking if enabled | ||
| if settings.use_reranker and top_results: | ||
| top_results = await self.reranker.rerank( | ||
| query, top_results, top_k=settings.rerank_top_k | ||
| ) | ||
|
|
||
| # 7. Filter by minimum score and limit | ||
| final_results = [ | ||
| r for r in top_results | ||
| if r.get("score", 0) >= settings.min_relevance_score | ||
| ][:top_k] | ||
|
|
||
| # Calculate metrics | ||
| query_time = (time.time() - start_time) * 1000 | ||
| metrics = RetrievalMetrics( | ||
| query_time_ms=query_time, | ||
| num_results=len(final_results), | ||
| reranked=settings.use_reranker, | ||
| top_score=final_results[0]["score"] if final_results else None, | ||
| average_score=( | ||
| sum(r["score"] for r in final_results) / len(final_results) | ||
| if final_results | ||
| else None | ||
| ), | ||
| ) | ||
|
|
||
| logger.info( | ||
| f"Multi-query retrieval: {len(queries)} queries, " | ||
| f"{len(final_results)} results in {query_time:.2f}ms" | ||
| ) | ||
|
|
||
| return final_results, metrics |
There was a problem hiding this comment.
Set metrics.reranked based on actual rerank outcome.
Right now it mirrors the setting even when reranking is skipped (e.g., no results) or fails. Tracking the outcome avoids misleading metrics.
🔧 Proposed fix
- # 6. Apply reranking if enabled
- if settings.use_reranker and top_results:
- top_results = await self.reranker.rerank(
- query, top_results, top_k=settings.rerank_top_k
- )
+ # 6. Apply reranking if enabled
+ reranked = False
+ if settings.use_reranker and top_results:
+ top_results = await self.reranker.rerank(
+ query, top_results, top_k=settings.rerank_top_k
+ )
+ reranked = any(r.get("reranked") for r in top_results)
@@
- metrics = RetrievalMetrics(
+ metrics = RetrievalMetrics(
query_time_ms=query_time,
num_results=len(final_results),
- reranked=settings.use_reranker,
+ reranked=reranked,
top_score=final_results[0]["score"] if final_results else None,🤖 Prompt for AI Agents
In `@backend/app/rag/enhanced_retriever.py` around lines 196 - 269, The
metrics.reranked flag currently copies settings.use_reranker even when reranking
was never applied or failed; update retrieve_with_multi_query to set reranked
based on the actual rerank outcome: initialize a local boolean (e.g., did_rerank
= False), only set it True after a successful await self.reranker.rerank(...)
call that returns non-empty results (and catch exceptions around the reranker
call to log the error and leave did_rerank False), ensure that when top_results
is empty you skip reranking and leave did_rerank False, and finally assign
metrics.reranked = did_rerank instead of using settings.use_reranker.
Comment on lines
+314
to
+320
| if current_length + para_length > self.chunk_size and current_group: | ||
| groups.append('\n\n'.join(current_group)) | ||
| current_group = [para] | ||
| current_length = para_length | ||
| else: | ||
| current_group.append(para) | ||
| current_length += para_length + 2 # +2 for \n\n |
There was a problem hiding this comment.
Inconsistent size accounting in paragraph grouping.
The size check at line 314 doesn't include the separator length (+2 for \n\n), but when adding to current_length at line 320, you add +2. This can cause chunks to slightly exceed chunk_size when joining paragraphs.
🔧 Proposed fix
- if current_length + para_length > self.chunk_size and current_group:
+ # Account for separator when checking size
+ separator_len = 2 if current_group else 0 # \n\n between paragraphs
+ if current_length + separator_len + para_length > self.chunk_size and current_group:🤖 Prompt for AI Agents
In `@backend/app/utils/chunking.py` around lines 314 - 320, The paragraph grouping
logic incorrectly omits the separator length when checking if adding para would
exceed self.chunk_size: update the check that uses current_length, para_length,
current_group, and chunk_size so it accounts for the separator ("\n\n") length
(2) when current_group is non-empty (e.g., calculate sep = 2 if current_group
else 0 and use current_length + para_length + sep > self.chunk_size), ensuring
the same separator length is used in both the check and when updating
current_length/current_group to prevent chunks from exceeding chunk_size.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Key improvements:
New features:
https://claude.ai/code/session_01BRVpxN2kXDmfHbPdkdUhov
Summary by CodeRabbit
New Features
Configuration