diff --git a/alembic/versions/031_rag_chunks_fulltext.py b/alembic/versions/031_rag_chunks_fulltext.py new file mode 100644 index 00000000..d219ae2f --- /dev/null +++ b/alembic/versions/031_rag_chunks_fulltext.py @@ -0,0 +1,94 @@ +"""Add a full-text (tsvector) column + GIN index to every RAG chunk table. + +Hybrid retrieval fuses pgvector cosine similarity with Postgres native +full-text search (keyword/BM25-style) so recall no longer depends on a +per-query HyDE LLM call. Each ``chunks_`` table gets a generated +``tsv`` column (``to_tsvector('english', content)``, auto-maintained on +insert/update) and a GIN index over it. + +The work runs inside a plpgsql ``DO`` block guarding on ``information_schema``, +so it is idempotent, offline-renderable (``alembic --sql``), and safe whether a +table is present, already has ``tsv``, or is absent (e.g. ``chunks_code``). +Depends on 030 having reshaped the column to ``content``. + +Revision ID: 031_rag_chunks_fulltext +Revises: 030_rag_chunks_content_schema +Create Date: 2026-06-15 +""" + +from __future__ import annotations + +from alembic import op + +revision = "031_rag_chunks_fulltext" +down_revision = "030_rag_chunks_content_schema" +branch_labels = None +depends_on = None + +CHUNK_TABLES = ( + "chunks_code", + "chunks_documentation", + "chunks_conversations", + "chunks_journals", + "chunks_errors", + "chunks_standards", + "chunks_decisions", + "chunks_reviews", + "chunks_learnings", +) + +_TABLES_SQL = ", ".join(f"'{name}'" for name in CHUNK_TABLES) + + +def upgrade() -> None: + op.execute( + f""" + DO $$ + DECLARE + t text; + BEGIN + FOREACH t IN ARRAY ARRAY[{_TABLES_SQL}] LOOP + IF EXISTS ( + SELECT 1 FROM information_schema.columns + WHERE table_schema = 'public' + AND table_name = t AND column_name = 'content' + ) THEN + IF NOT EXISTS ( + SELECT 1 FROM information_schema.columns + WHERE table_schema = 'public' + AND table_name = t AND column_name = 'tsv' + ) THEN + EXECUTE format( + 'ALTER TABLE %I ADD COLUMN tsv tsvector ' + 'GENERATED ALWAYS AS ' + '(to_tsvector(''english'', content)) STORED', + t + ); + END IF; + EXECUTE format( + 'CREATE INDEX IF NOT EXISTS %I ON %I USING gin (tsv)', + t || '_tsv_idx', t + ); + END IF; + END LOOP; + END $$; + """ + ) + + +def downgrade() -> None: + op.execute( + f""" + DO $$ + DECLARE + t text; + BEGIN + FOREACH t IN ARRAY ARRAY[{_TABLES_SQL}] LOOP + EXECUTE format('DROP INDEX IF EXISTS %I', t || '_tsv_idx'); + EXECUTE format( + 'ALTER TABLE IF EXISTS %I DROP COLUMN IF EXISTS tsv', t + ); + END LOOP; + END $$; + """ + ) diff --git a/roboco/api/schemas/optimal.py b/roboco/api/schemas/optimal.py index 0f8deab5..c8ba56fe 100644 --- a/roboco/api/schemas/optimal.py +++ b/roboco/api/schemas/optimal.py @@ -102,7 +102,7 @@ class RAGHealthResponse(BaseModel): healthy: bool embedding_status: str = Field(..., description="Embedding model status") - llm_status: str = Field(..., description="LLM (HyDE) status") + llm_status: str = Field(..., description="LLM (answer synthesis) status") vector_store_status: str = Field(..., description="Vector store status") details: dict[str, Any] = Field(default_factory=dict) diff --git a/roboco/config.py b/roboco/config.py index 9a0903a7..dba27b52 100644 --- a/roboco/config.py +++ b/roboco/config.py @@ -131,11 +131,6 @@ class Settings(BaseSettings): default=1024, ge=100, description="Chunk size for journals/reflections" ) rag_chunk_overlap: int = Field(default=128, ge=0) - rag_use_hyde: bool = Field( - default=True, - description="Use HyDE (hypothetical document embeddings). " - "Makes one LLM call per query for better semantic matching.", - ) rag_auto_update_enabled: bool = Field(default=True) rag_auto_update_interval: int = Field( default=300, ge=60, description="Seconds between auto-updates" @@ -165,10 +160,11 @@ class Settings(BaseSettings): description="Embedding dimensions (1024 for qwen3-embedding)", ) - # Local LLM for RAG (HyDE, reranking, etc.) + # Local LLM for RAG answer synthesis local_llm_model: str = Field( default="glm-5:cloud", - description="Local LLM for HyDE/RAG (non-thinking models are faster)", + description="Local LLM for RAG answer synthesis " + "(non-thinking models are faster)", ) local_llm_base_url: str = Field( default="http://roboco-ollama:11434/v1", diff --git a/roboco/services/optimal.py b/roboco/services/optimal.py index 2311662a..9b064210 100644 --- a/roboco/services/optimal.py +++ b/roboco/services/optimal.py @@ -1085,8 +1085,8 @@ class OptimalService: if not plugins: return [] - # Embed the query ONCE (HyDE + embed), then run every index's vector - # search concurrently — instead of each index re-embedding in series. + # Embed the query ONCE, then run every index's hybrid search + # concurrently — instead of each index re-embedding in series. try: query_embedding = await plugins[0][1].compute_query_embedding(query) except Exception as e: @@ -1095,7 +1095,7 @@ class OptimalService: outcomes = await asyncio.gather( *( - plugin.search_with_embedding(query_embedding, top_k=top_k) + plugin.search_with_embedding(query_embedding, query, top_k=top_k) for _, plugin in plugins ), return_exceptions=True, @@ -1107,18 +1107,21 @@ class OptimalService: async def _search_single_index( self, - index_type: IndexType, - plugin: BaseIndexPlugin, + entry: tuple[IndexType, BaseIndexPlugin], query_embedding: list[float], + query_text: str, top_k: int, buf: _QueryAggregationBuffer, ) -> None: """Search one index with a pre-computed embedding; update buf in place.""" + index_type, plugin = entry if await plugin.count() == 0: logger.debug("Skipping empty index", index_type=index_type.value) return - outcome = await plugin.search_with_embedding(query_embedding, top_k=top_k) + outcome = await plugin.search_with_embedding( + query_embedding, query_text, top_k=top_k + ) if outcome.success: buf.stats[index_type.value] = len(outcome.results) buf.citations.extend(outcome.results) @@ -1148,7 +1151,9 @@ class OptimalService: await asyncio.gather( *( - self._search_single_index(it, plugin, query_embedding, top_k, buf) + self._search_single_index( + (it, plugin), query_embedding, query, top_k, buf + ) for it, plugin in plugins ), return_exceptions=True, diff --git a/roboco/services/optimal_brain/indexes/base.py b/roboco/services/optimal_brain/indexes/base.py index 7ca1726f..72b4840d 100644 --- a/roboco/services/optimal_brain/indexes/base.py +++ b/roboco/services/optimal_brain/indexes/base.py @@ -67,7 +67,6 @@ class IndexConfig: chunk_strategy: str = "fixed" chunk_size: int = 512 chunk_overlap: int = 50 - use_hyde: bool = True embedding_model: str = "qwen3-embedding:0.6b" llm_model: str = "glm-5:cloud" llm_base_url: str = "http://roboco-ollama:11434/v1" @@ -88,7 +87,6 @@ class IndexConfig: chunk_strategy=settings.rag_chunk_strategy, chunk_size=chunk_size, chunk_overlap=settings.rag_chunk_overlap, - use_hyde=settings.rag_use_hyde, embedding_model=settings.default_embedding_model, llm_model=settings.local_llm_model, llm_base_url=settings.local_llm_base_url, @@ -721,70 +719,12 @@ class BaseIndexPlugin(ABC): return preprocessed - async def _generate_hyde_passage(self, query: str) -> str: - """Generate a hypothetical passage for HyDE query expansion. - - Calls the configured Ollama LLM to produce a short passage that would - plausibly answer *query*. Returns an empty string on any failure so - the caller can fall back to raw query embedding. - - Args: - query: The (preprocessed) search query. - - Returns: - A short hypothetical passage, or ``""`` on LLM error. - """ - import httpx - - prompt = ( - "Write a concise technical passage (2-4 sentences) that directly " - "answers the following question. Include specific technical details " - "and terminology. Do not use tags.\n\n" - f"Question: {query}\n\n" - "Answer:" - ) - - try: - async with httpx.AsyncClient(timeout=30.0) as client: - resp = await client.post( - f"{self.config.llm_base_url}/chat/completions", - json={ - "model": self.config.llm_model, - "messages": [{"role": "user", "content": prompt}], - "max_tokens": 200, - }, - ) - if resp.is_success: - data = resp.json() - passage = self._extract_from_think_tags( - data["choices"][0]["message"]["content"] - ) - logger.debug( - "HyDE passage generated", - index_type=self.index_type.value, - passage_len=len(passage), - ) - return passage - logger.debug( - "HyDE LLM call returned non-success status", - status=resp.status_code, - index_type=self.index_type.value, - ) - except Exception as e: - logger.debug( - "HyDE LLM call failed, will use raw query embedding", - index_type=self.index_type.value, - error=str(e), - ) - return "" - async def _compute_query_embedding(self, query: str) -> list[float]: """Preprocess the query and embed it using the configured embedder. - When ``config.use_hyde`` is ``True``, first calls the Ollama LLM to - generate a hypothetical answer passage (HyDE) and embeds that instead - of the raw query. Falls back to raw query embedding if the LLM call - fails or returns an empty passage. + Hybrid retrieval (vector + full-text) covers the question/document + vocabulary gap HyDE used to paper over, so the query is embedded + directly — no per-query LLM round-trip in the search hot path. """ import asyncio @@ -799,29 +739,24 @@ class BaseIndexPlugin(ABC): index_type=self.index_type.value, ) - # HyDE: try to embed a hypothetical passage instead of the raw query - text_to_embed = processed_query - if self.config.use_hyde: - passage = await self._generate_hyde_passage(processed_query) - if passage: - text_to_embed = passage - if hasattr(embedder, "aembed_query"): - result: list[float] = await embedder.aembed_query(text_to_embed) + result: list[float] = await embedder.aembed_query(processed_query) return result - return await asyncio.to_thread(embedder.embed_query, text_to_embed) + return await asyncio.to_thread(embedder.embed_query, processed_query) async def _fetch_citations( self, query_embedding: list[float], + query_text: str, top_k: int, has_filters: bool, ) -> list[Citation]: - """Fetch citations from the vector store.""" + """Fetch citations via hybrid (vector + full-text) search.""" store = self._require_store fetch_k = top_k * 3 if has_filters else top_k - return await store.search( + return await store.hybrid_search( query_embedding, + query_text, top_k=fetch_k, min_chunk_length=100, ) @@ -855,30 +790,32 @@ class BaseIndexPlugin(ABC): return results async def compute_query_embedding(self, query: str) -> list[float]: - """Embed a query (incl. HyDE) once. + """Embed a query once. Lets the service compute one embedding and fan it out across every - index, instead of each index re-running HyDE + embed sequentially. + index, instead of each index re-embedding. """ return await self._compute_query_embedding(query) async def search_with_embedding( self, query_embedding: list[float], + query_text: str, top_k: int = 5, filters: dict[str, Any] | None = None, ) -> SearchOutcome: - """Search using a pre-computed query embedding (skips HyDE/embed). + """Search using a pre-computed query embedding (skips embed). Splitting embed from fetch lets the service embed once and run every - index's vector search concurrently. + index's hybrid search concurrently. ``query_text`` drives the full-text + half of the hybrid query. """ import time start_time = time.time() try: chunks = await self._fetch_citations( - query_embedding, top_k, has_filters=bool(filters) + query_embedding, query_text, top_k, has_filters=bool(filters) ) results = self._citations_to_results(chunks, top_k, filters) elapsed_ms = (time.time() - start_time) * 1000 @@ -932,7 +869,7 @@ class BaseIndexPlugin(ABC): search_time_ms=elapsed_ms, ) return await self.search_with_embedding( - query_embedding, top_k=top_k, filters=filters + query_embedding, query, top_k=top_k, filters=filters ) def _fallback_answer(self, _search_results: list[SearchResult]) -> str: diff --git a/roboco/services/optimal_brain/vector_store.py b/roboco/services/optimal_brain/vector_store.py index 90a09cb4..2d66011a 100644 --- a/roboco/services/optimal_brain/vector_store.py +++ b/roboco/services/optimal_brain/vector_store.py @@ -134,7 +134,9 @@ class VectorStore: source TEXT NOT NULL, embedding vector({self._vector_dimension}), metadata JSONB NOT NULL DEFAULT '{{}}', - created_at TIMESTAMPTZ NOT NULL DEFAULT NOW() + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + tsv tsvector GENERATED ALWAYS AS + (to_tsvector('english', content)) STORED ) """ ) @@ -147,6 +149,14 @@ class VectorStore: WITH (lists = 100) """ ) + # GIN index for the full-text (keyword) half of hybrid search. + await conn.execute( + f""" + CREATE INDEX IF NOT EXISTS {self._table_name}_tsv_idx + ON {self._table_name} + USING gin (tsv) + """ + ) logger.info( "VectorStore initialised", @@ -280,6 +290,96 @@ class VectorStore: for row in rows ] + async def hybrid_search( + self, + embedding: list[float], + query_text: str, + top_k: int = 5, + min_chunk_length: int = 100, + candidate_pool: int = 50, + ) -> list[Citation]: + """Hybrid retrieval: fuse pgvector cosine with full-text keyword search. + + Per chunk the score is + ``min(1, cosine + 0.3 * normalized_ts_rank)``: a vector-only match keeps + its cosine score (so downstream thresholds — decisions/reviewer — are + unchanged), a keyword match adds a bounded boost (the recall win), and a + keyword-only match stays low (<= 0.3). Empty/garbage ``query_text`` + degrades gracefully to pure vector search. + + Args: + embedding: Query embedding vector. + query_text: Raw query for the full-text half. + top_k: Maximum number of fused results to return. + min_chunk_length: Minimum character length of returned chunks. + candidate_pool: Per-side candidate count before fusion. + """ + pool = self._require_pool() + emb_str = _vec_to_str(embedding) + cand = max(candidate_pool, top_k) + + async with pool.acquire() as conn: + rows = await conn.fetch( + self._q( + """ + WITH q AS ( + SELECT websearch_to_tsquery('english', $2) AS tsq + ), + vec AS ( + SELECT id, content, source, metadata, + 1 - (embedding <=> $1::vector) AS cos + FROM {table} + WHERE length(content) >= $3 AND embedding IS NOT NULL + ORDER BY embedding <=> $1::vector + LIMIT $4 + ), + kw AS ( + SELECT c.id, c.content, c.source, c.metadata, + ts_rank(c.tsv, q.tsq) AS kw_rank + FROM {table} c, q + WHERE c.tsv @@ q.tsq AND length(c.content) >= $3 + ORDER BY kw_rank DESC + LIMIT $4 + ), + maxk AS (SELECT NULLIF(MAX(kw_rank), 0) AS m FROM kw), + fused AS ( + SELECT COALESCE(vec.id, kw.id) AS id, + COALESCE(vec.content, kw.content) AS content, + COALESCE(vec.source, kw.source) AS source, + COALESCE(vec.metadata, kw.metadata) AS metadata, + COALESCE(vec.cos, 0) AS cos, + COALESCE(kw.kw_rank, 0) AS kw_rank + FROM vec FULL OUTER JOIN kw ON vec.id = kw.id + ) + SELECT content, source, metadata, + LEAST( + 1.0, + cos + 0.3 * COALESCE( + kw_rank / (SELECT m FROM maxk), 0 + ) + ) AS score + FROM fused + ORDER BY score DESC + LIMIT $5 + """ + ), + emb_str, + query_text, + min_chunk_length, + cand, + top_k, + ) + + return [ + Citation( + chunk=row["content"], + source=row["source"], + score=float(row["score"]), + metadata=_as_dict(row["metadata"]), + ) + for row in rows + ] + async def count(self) -> int: """Return the total number of chunk rows in the table.""" pool = self._require_pool() diff --git a/tests/unit/services/optimal_brain/test_hybrid_search_wiring.py b/tests/unit/services/optimal_brain/test_hybrid_search_wiring.py new file mode 100644 index 00000000..83528735 --- /dev/null +++ b/tests/unit/services/optimal_brain/test_hybrid_search_wiring.py @@ -0,0 +1,36 @@ +"""The plugin search path routes to hybrid_search with the query text. + +The full-text half of hybrid retrieval only works if the raw query text is +threaded all the way to ``VectorStore.hybrid_search``. Mocks resolve any +attribute, so this guards against the wiring silently regressing to pure-vector +``store.search`` or dropping ``query_text``. +""" + +from __future__ import annotations + +from unittest.mock import AsyncMock + +import pytest +from roboco.services.optimal_brain.indexes.standards import StandardsIndexPlugin + + +@pytest.mark.asyncio +async def test_search_with_embedding_calls_hybrid_with_query_text( + monkeypatch: pytest.MonkeyPatch, +) -> None: + plugin = StandardsIndexPlugin() + store = AsyncMock() + store.hybrid_search = AsyncMock(return_value=[]) + monkeypatch.setattr(plugin, "_store", store) + monkeypatch.setattr(plugin, "_initialized", True) + + outcome = await plugin.search_with_embedding( + [0.1, 0.2, 0.3], "claim a task", top_k=4 + ) + + assert outcome.success + store.hybrid_search.assert_awaited_once() + call = store.hybrid_search.call_args + assert call.args[0] == [0.1, 0.2, 0.3] # embedding + assert call.args[1] == "claim a task" # raw query text drives full-text half + store.search.assert_not_awaited() # not the pure-vector path