mirror of
https://github.com/rennf93/roboco.git
synced 2026-08-03 07:23:24 +02:00
feat(rag): hybrid retrieval (vector + full-text), retire HyDE
Recall no longer depends on a per-query HyDE LLM call — it comes from the index. Each chunks_<type> table gets a generated `tsv` column + GIN index (migration 031; the engine CREATE TABLE matches so fresh tables get it too). VectorStore.hybrid_search fuses pgvector cosine with Postgres full-text in one query: score = min(1, cosine + 0.3 * normalized_ts_rank). A vector-only match keeps its cosine score (so decisions/reviewer thresholds are unchanged), a keyword match adds a bounded boost (the recall win), and a keyword-only match stays low. Empty/garbage query text degrades to pure vector. HyDE is removed from the search hot path: _compute_query_embedding now embeds the query directly, and _generate_hyde_passage / rag_use_hyde / IndexConfig.use_hyde are deleted. So a search is one local embed + one indexed SQL — no LLM round-trip. The raw query text is threaded through the embed-once + concurrent fan-out (search_with_embedding(embedding, query_text)). Verified live via a real pgvector round-trip: vector ranking + keyword boost + [0,1] scores + empty-query fallback all correct. Adds wiring + fan-out unit tests; the fusion SQL itself is verified live (needs pgvector, not gated in CI).
This commit is contained in:
Renn F
@@ -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_<index_type>`` 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 $$;
|
||||
"""
|
||||
)
|
||||
@@ -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)
|
||||
|
||||
|
||||
+3
-7
@@ -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",
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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 <think> 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:
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user