feat(retrieval): DB-native RRF via Supabase RPC, drop Python-side fusion - #4
Conversation
Replace the LlamaIndex VectorStoreIndex retriever in query.py with a direct asyncpg call to a new hybrid_search_rrf() Postgres RPC function. The RPC fuses dense cosine vector search with BM25-style full-text search using Reciprocal Rank Fusion — all inside the database, so only the final top-K rows cross the network. Key changes: - New migration 002: pg_trgm extension - New migration 003: hybrid_search_rrf() RPC function - query.py: async rewrite with asyncpg connection pool, direct embedding generation via GoogleGenAIEmbedding, and LLM call via GoogleGenAI - Removed SIMILARITY_CUTOFF (RRF scores are relative, not cosine) - Added requirements.txt with asyncpg>=0.29.0 Co-Authored-By: Claude Sonnet 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 retrieval-augmented generation (RAG) system by implementing a database-native hybrid search mechanism. By moving the complex fusion logic for vector and full-text search into a PostgreSQL RPC function, the system reduces data transfer overhead and leverages the database's efficiency. This change streamlines the query process, making it more performant and robust, while maintaining the flexibility of using LlamaIndex for embedding and LLM interactions. Highlights
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
|
There was a problem hiding this comment.
Code Review
This pull request introduces a significant architectural improvement by moving the hybrid search and fusion logic into a PostgreSQL RPC function. This reduces data transfer and leverages the database's strengths. The refactoring of query.py to use asyncpg is clean and aligns with this new approach.
My review includes two main points:
- A critical security fix to prevent potential SQL injection in
query.pyby using parameterized queries correctly for the embedding vector. - A high-impact performance suggestion for the
hybrid_search_rrfSQL function to create a GIN index, which will substantially speed up full-text search operations.
After addressing these points, this will be an excellent enhancement to the RAG system's retrieval performance and security.
| embedding_str = "[" + ",".join(str(v) for v in query_embedding) + "]" | ||
| log.debug("Calling hybrid_search_rrf (match_count=%d) ...", top_k) | ||
|
|
||
| log.info("Question: %s", question) | ||
| log.debug("Searching vector store (top_k=%d, cutoff=%.2f) ...", top_k, similarity_cutoff) | ||
| try: | ||
| rows = await pool.fetch( | ||
| "SELECT id, content, metadata, score " | ||
| "FROM hybrid_search_rrf($1, $2::vector(3072), $3)", | ||
| question, | ||
| embedding_str, | ||
| top_k, | ||
| ) |
There was a problem hiding this comment.
There is a critical SQL injection vulnerability here. The embedding vector is being converted to a string and then used in the SQL query. This is unsafe. You should pass the query_embedding list directly as a parameter to pool.fetch and let asyncpg handle the safe serialization. This prevents any possibility of SQL injection through the embedding data.
Here's how to fix it:
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,
query_embedding,
top_k,
)| fulltext AS ( | ||
| SELECT | ||
| f.id, | ||
| ROW_NUMBER() OVER ( | ||
| ORDER BY ts_rank_cd( | ||
| to_tsvector('english', f.metadata ->> 'text'), | ||
| websearch_to_tsquery('english', query_text) | ||
| ) DESC | ||
| ) AS rank_ix | ||
| FROM vecs.openclaw_docs f | ||
| WHERE | ||
| f.metadata ->> 'text' IS NOT NULL | ||
| AND to_tsvector('english', f.metadata ->> 'text') | ||
| @@ websearch_to_tsquery('english', query_text) | ||
| LIMIT (match_count * 2) |
There was a problem hiding this comment.
The full-text search portion of this query is inefficient because it computes to_tsvector on the fly for every row during every query execution. This will not scale well with a larger number of documents.
To significantly improve performance, you should create a GIN index on the tsvector expression. This will allow PostgreSQL to use an index for the full-text search operator (@@), making it much faster.
I recommend creating a new migration file (e.g., 004_fts_index.sql) with the following content:
CREATE INDEX IF NOT EXISTS openclaw_docs_fts_idx
ON vecs.openclaw_docs
USING GIN (to_tsvector('english', metadata ->> 'text'));This is a non-destructive operation that will dramatically speed up your hybrid search.
There was a problem hiding this comment.
Pull request overview
This PR shifts the RAG query path from LlamaIndex’s SupabaseVectorStore query engine to a DB-native hybrid search RPC (RRF fusion of vector + full-text), and documents the required Postgres migrations to support it.
Changes:
- Replace
rag-agent/query.pyretrieval withasyncpg+hybrid_search_rrf()RPC and keep prompt-injection guardrails + source attribution. - Add SQL migrations/docs for enabling full-text support and creating the hybrid search function.
- Update setup documentation to include running migrations and add
asyncpgto Python deps.
Reviewed changes
Copilot reviewed 8 out of 8 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
| README.md | Adds a setup step for running DB migrations after initial ingestion. |
| rag-agent/requirements.txt | Introduces an explicit Python dependency file (currently only asyncpg). |
| rag-agent/query.py | Switches query-time retrieval to asyncpg RPC-based hybrid search + builds LLM prompt manually. |
| rag-agent/migrations/README.md | Documents migration purpose/order and how to run via psql. |
| rag-agent/migrations/002_pg_trgm.sql | Enables pg_trgm extension. |
| rag-agent/migrations/003_hybrid_search_rrf.sql | Adds hybrid_search_rrf() SQL function for fused retrieval. |
| rag-agent/.env.example | Updates query config notes to reflect hybrid RPC usage. |
| CLAUDE.md | Adds repo-specific workflow/architecture guidance for Claude Code. |
Comments suppressed due to low confidence (1)
rag-agent/query.py:56
DB_CONNECTION_STRINGis read but never validated. If it is unset/empty,_get_pool()will raise whenasyncpg.create_pool()is called; add an explicit check (similar toGOOGLE_API_KEY) and fail with a clear error message.
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":
log.error("GOOGLE_API_KEY not set in .env")
return
💡 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.
| @@ -0,0 +1 @@ | |||
| asyncpg>=0.29.0 | |||
| "SELECT id, content, metadata, score " | ||
| "FROM hybrid_search_rrf($1, $2::vector(3072), $3)", |
| @@ -0,0 +1,88 @@ | |||
| -- Migration 003: Create hybrid_search_rrf() RPC function. | |||
| -- | |||
| -- Run this ONCE in the Supabase SQL editor (or via psql) AFTER migrations 001 and 002. | |||
The vecs library stores node text inside metadata->_node_content as a serialized JSON string, not directly in metadata->text. Update the SQL RPC and comments accordingly. Switch query.py to async equivalents (aget_query_embedding, acomplete) so the event loop is not blocked during embedding and LLM generation. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
- Complete requirements.txt with all direct Python dependencies - Bump package.json version to 0.1.0-beta - Gitignore .claude/ and root tools/ (personal utilities) - Remove corpus CHANGELOG.md from repo root (not a project changelog) - Update README: remove CHANGELOG ref, add pip install step, add Known Limitations Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Summary
hybrid_search_rrf()Postgres RPC that fuses dense cosine vector search with BM25-style full-text search using Reciprocal Rank Fusion — all inside the DB, only top-K rows cross the networkquery.pyto use asyncpg connection pool + direct RPC call, replacing the LlamaIndexVectorStoreIndexretriever. LlamaIndex retained only for embedding generation (GoogleGenAIEmbedding) and LLM calls (GoogleGenAI)002_pg_trgm.sql(trigram extension) and003_hybrid_search_rrf.sql(RPC function with parameterized weights, over-fetching, and metadata stripping)requirements.txtwithasyncpg>=0.29.0SIMILARITY_CUTOFF(RRF scores are relative, not cosine distances)Phase 3 of 4
This implements Mandate 1 (DB-Native Hybrid Search) from
SPRINT_PLAN.md. Phases 1 (Foundation) and 2 (Incremental Ingestion) are already merged.Test plan
psql ... -f migrations/002_pg_trgm.sql -f migrations/003_hybrid_search_rrf.sqlpython query.py "What is the gateway?"— returns answer with sources and RRF scorespython query.py "northflank"— keyword-exact query ranks relevant docs higher than Phase 2 baselineLOG_LEVEL=DEBUG python query.py "..."— shows embedding, RPC, and LLM stepspy_compilepasses on Python 3.10 + 3.11UserWarningabout missing covering index (HNSW from Phase 1 covers dense path; we bypass vecs query entirely)🤖 Generated with Claude Code