feat(phase4): FlashRank reranking, MCP server, mimalloc docs - #5
feat(phase4): FlashRank reranking, MCP server, mimalloc docs#5johansabent wants to merge 2 commits into
Conversation
Refactor query.py: extract retrieve_and_answer() returning QueryResult dataclass so both CLI and MCP server can consume structured results. Move logging.basicConfig into __main__ guard to prevent stdout corruption when imported by MCP server (stdio transport). Add optional FlashRank cross-encoder reranking (RERANKER_ENABLED env var) with lazy model loading and dual-score source attribution. Add mcp_server.py exposing query_docs tool over MCP stdio transport for Claude Desktop / VS Code integration. Update README with reranking, MCP server, and WSL2 mimalloc deployment documentation. Update CLAUDE.md roadmap and config tables. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
|
You have reached your Codex usage limits for code reviews. You can see your limits in the Codex usage dashboard. |
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request significantly enhances the RAG pipeline by integrating advanced reranking capabilities and enabling external tool integration. It provides a more robust and extensible query system, allowing for improved answer precision and seamless interaction with development environments, while also updating the project's documentation to reflect these new features. Highlights
Ignored Files
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here. Footnotes
|
![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 introduces several significant features, including FlashRank reranking and an MCP server, alongside a valuable refactoring of query.py. The code is well-structured, particularly the extraction of retrieve_and_answer and the introduction of the QueryResult dataclass. My feedback focuses on a point of code duplication in formatting the query results, which can be centralized to improve maintainability.
| @dataclass | ||
| class QueryResult: | ||
| answer: str | ||
| sources: list[dict] = field(default_factory=list) # [{file_name, score, rerank_score?}, ...] | ||
| error: str | None = None |
There was a problem hiding this comment.
The logic for formatting the sources list into display strings is duplicated in mcp_server.py (lines 43-49) and the query function in this file (lines 188-194). To improve maintainability and adhere to the DRY (Don't Repeat Yourself) principle, you can centralize this logic by adding a method to the QueryResult dataclass that returns the formatted source lines. This makes the QueryResult object responsible for its own representation.
@dataclass
class QueryResult:
answer: str
sources: list[dict] = field(default_factory=list) # [{file_name, score, rerank_score?}, ...]
error: str | None = None
def get_formatted_source_lines(self) -> list[str]:
"""Returns a list of formatted strings for each source."""
lines = []
for src in self.sources:
score = f"{src['score']:.4f}"
if "rerank_score" in src:
rerank = f"{src['rerank_score']:.4f}"
lines.append(f" [RRF {score} | rerank {rerank}] {src['file_name']}")
else:
lines.append(f" [{score}] {src['file_name']}")
return lines| for src in result.sources: | ||
| score = f"{src['score']:.4f}" | ||
| if "rerank_score" in src: | ||
| rerank = f"{src['rerank_score']:.4f}" | ||
| print(f" [RRF {score} | rerank {rerank}] {src['file_name']}") | ||
| else: | ||
| print(f" [{score}] {src['file_name']}") |
There was a problem hiding this comment.
Following the suggestion to add a get_formatted_source_lines method to QueryResult, this loop can be simplified to iterate over the pre-formatted lines. This removes the duplicated formatting logic and makes the intent clearer.
for line in result.get_formatted_source_lines():
print(line)| parts.append("\nSources:") | ||
| for src in result.sources: | ||
| score = f"{src['score']:.4f}" | ||
| if "rerank_score" in src: | ||
| rerank = f"{src['rerank_score']:.4f}" | ||
| parts.append(f" [RRF {score} | rerank {rerank}] {src['file_name']}") | ||
| else: | ||
| parts.append(f" [{score}] {src['file_name']}") |
There was a problem hiding this comment.
Following the suggestion to add a get_formatted_source_lines method to QueryResult in query.py, this block can be simplified to remove the duplicated formatting logic.
parts.append("\nSources:")
parts.extend(result.get_formatted_source_lines())There was a problem hiding this comment.
Pull request overview
Adds Phase 4 capabilities to the RAG agent by introducing optional FlashRank reranking and an MCP stdio server, refactoring the query pipeline to return structured results usable by both CLI and MCP clients, and updating docs/CI accordingly.
Changes:
- Refactor
query.pyto exposeretrieve_and_answer()returning aQueryResult(CLI prints the structured result). - Add optional FlashRank cross-encoder reranking controlled via env vars.
- Add
mcp_server.pyand extend docs/CI to support MCP usage and compilation checks.
Reviewed changes
Copilot reviewed 7 out of 7 changed files in this pull request and generated 1 comment.
Show a summary per file
| File | Description |
|---|---|
| rag-agent/requirements.txt | Adds flashrank and mcp dependencies for reranking and MCP server. |
| rag-agent/query.py | Refactors query flow into retrieve_and_answer() + optional reranking and structured results. |
| rag-agent/mcp_server.py | New MCP stdio server exposing a query_docs tool. |
| rag-agent/.env.example | Documents reranking configuration env vars. |
| README.md | Adds reranking/MCP usage docs and WSL2 mimalloc deployment notes. |
| CLAUDE.md | Updates pipeline description, key files, config table, and phase status. |
| .github/workflows/ci.yml | Extends syntax check to include mcp_server.py. |
Comments suppressed due to low confidence (2)
rag-agent/query.py:94
DB_CONNECTION_STRINGis read but never validated. If it's missing/empty,_get_pool(db_connection)will raise and crash the CLI/MCP server instead of returning a structuredQueryResulterror. Add an explicit check similar toGOOGLE_API_KEYand returnQueryResult(..., error=...)(or raise a clearSystemExitfor CLI) before calling_get_pool().
api_key = os.getenv("GOOGLE_API_KEY")
db_connection = os.getenv("DB_CONNECTION_STRING")
if not api_key or api_key == "YOUR_API_KEY_HERE":
return QueryResult(answer="", error="GOOGLE_API_KEY not set in .env")
embed_model_name = os.getenv("EMBED_MODEL", "models/gemini-embedding-2-preview")
llm_model = os.getenv("LLM_MODEL", "models/gemini-3.1-flash-lite-preview")
dimensions = int(os.getenv("EMBED_DIMENSIONS", "3072"))
top_k = int(os.getenv("SIMILARITY_TOP_K", "5"))
# --- Embedding ---
embed_model = GoogleGenAIEmbedding(
model_name=embed_model_name,
api_key=api_key,
output_dimensionality=dimensions,
)
log.info("Question: %s", question)
log.debug("Generating query embedding ...")
query_embedding = await embed_model.aget_query_embedding(question)
# --- Hybrid retrieval via DB-native RPC ---
pool = await _get_pool(db_connection)
rag-agent/query.py:104
dimensionsis configurable viaEMBED_DIMENSIONS, but the query embedding is cast asvector(3072)in the RPC call. IfEMBED_DIMENSIONSis changed (and ingestion uses it), queries will fail due to a dimension mismatch. Either enforcedimensions == 3072with a clear error, or update the SQL/RPC and this cast so the dimension is consistent end-to-end.
embed_model_name = os.getenv("EMBED_MODEL", "models/gemini-embedding-2-preview")
llm_model = os.getenv("LLM_MODEL", "models/gemini-3.1-flash-lite-preview")
dimensions = int(os.getenv("EMBED_DIMENSIONS", "3072"))
top_k = int(os.getenv("SIMILARITY_TOP_K", "5"))
# --- Embedding ---
embed_model = GoogleGenAIEmbedding(
model_name=embed_model_name,
api_key=api_key,
output_dimensionality=dimensions,
)
log.info("Question: %s", question)
log.debug("Generating query embedding ...")
query_embedding = await embed_model.aget_query_embedding(question)
# --- Hybrid retrieval via DB-native RPC ---
pool = await _get_pool(db_connection)
embedding_str = "[" + ",".join(str(v) for v in query_embedding) + "]"
log.debug("Calling hybrid_search_rrf (match_count=%d) ...", top_k)
try:
rows = await pool.fetch(
"SELECT id, content, metadata, score "
"FROM hybrid_search_rrf($1, $2::vector(3072), $3)",
question,
embedding_str,
top_k,
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
You can also share your feedback on Copilot code review. Take the survey.
| for i, row in enumerate(rows) if row["content"] | ||
| ] | ||
| rerank_request = RerankRequest(query=question, passages=passages) | ||
| reranked = ranker.rerank(rerank_request)[:rerank_top_n] |
- Add format_sources() method to QueryResult dataclass to DRY up source formatting duplicated across query() and mcp_server.py - Wrap synchronous ranker.rerank() in asyncio.to_thread() to avoid blocking the event loop under MCP server Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Summary
query.py: Extractretrieve_and_answer()returning aQueryResultdataclass so both CLI and MCP server consume structured results. Movelogging.basicConfiginto__main__guard to prevent stdout corruption under MCP stdio transport.RERANKER_ENABLED=true) with lazy model loading, configurable top-N and model name, dual-score source attribution (RRF + rerank).mcp_server.pyexposingquery_docstool over stdio transport for Claude Desktop / VS Code integration. Logs routed to stderr.mcp_server.py.Test plan
python query.py "What channels does OpenClaw support?"— output identical to Phase 3 (regression)RERANKER_ENABLED=true SIMILARITY_TOP_K=20 python query.py "What channels does OpenClaw support?"— shows dual RRF/rerank scorespython mcp_server.py— starts without error on stdiopython -m py_compile ingest.py query.py mcp_server.py— passes (mirrors CI)🤖 Generated with Claude Code