diff --git a/application/parser/chunking_strategies.py b/application/parser/chunking_strategies.py index e59cdd1b..87b40839 100644 --- a/application/parser/chunking_strategies.py +++ b/application/parser/chunking_strategies.py @@ -4,11 +4,12 @@ Each strategy honours ``max_tokens`` / ``min_tokens`` and reuses the classic ``Chunker``'s tiktoken encoding for token counting, so token budgets stay consistent across strategies. Selecting a strategy is ingest-time only; changing it requires a re-ingest (D8). Registered keys: ``recursive``, -``markdown``, ``parent_child``. ``semantic`` is intentionally deferred. +``markdown``, ``parent_child``, ``semantic``. """ from __future__ import annotations +import logging import re from typing import List @@ -17,6 +18,8 @@ from application.parser.chunking_creator import ChunkerCreator from application.parser.schema.base import Document from application.utils import get_encoding +logger = logging.getLogger(__name__) + class _BaseStrategyChunker: """Shared token helpers for strategy chunkers. @@ -64,6 +67,25 @@ class _BaseStrategyChunker: }, ) + def _merge_to_min(self, pieces: List[str], joiner: str) -> List[str]: + """Accumulate pieces up to ``max_tokens``, flushing past ``min_tokens``.""" + merged: List[str] = [] + buffer = "" + for piece in pieces: + candidate = f"{buffer}{joiner}{piece}" if buffer else piece + if self._token_count(candidate) <= self.max_tokens: + buffer = candidate + else: + if buffer: + merged.append(buffer) + buffer = piece + if buffer and self._token_count(buffer) >= self.min_tokens: + merged.append(buffer) + buffer = "" + if buffer: + merged.append(buffer) + return merged + class RecursiveChunker(_BaseStrategyChunker): """Split on a separator hierarchy, capping at ``max_tokens``. @@ -95,22 +117,7 @@ class RecursiveChunker(_BaseStrategyChunker): def _merge(self, fragments: List[str]) -> List[str]: """Merge small fragments up to ``max_tokens`` to clear ``min_tokens``.""" - merged: List[str] = [] - buffer = "" - for frag in fragments: - candidate = buffer + frag if buffer else frag - if self._token_count(candidate) <= self.max_tokens: - buffer = candidate - else: - if buffer: - merged.append(buffer) - buffer = frag - if buffer and self._token_count(buffer) >= self.min_tokens: - merged.append(buffer) - buffer = "" - if buffer: - merged.append(buffer) - return merged + return self._merge_to_min(fragments, "") def chunk(self, documents: List[Document]) -> List[Document]: processed: List[Document] = [] @@ -210,9 +217,108 @@ class ParentChildChunker(_BaseStrategyChunker): return processed +class SemanticChunker(_BaseStrategyChunker): + """Group adjacent sentences by embedding similarity into coherent chunks. + + Sentences are embedded in one batched call and split where the cosine + distance between consecutive sentences crosses a high percentile, so + chunk boundaries fall at topic shifts. Chunks are then token-capped at + ``max_tokens`` and merged up to clear ``min_tokens``. Any failure (too + few sentences, embedding error, degenerate distances) falls back to + ``RecursiveChunker`` so ingest never crashes. + """ + + _SENTENCE = re.compile(r"(?<=[.!?])\s+") + _PERCENTILE = 95.0 + + def _split_sentences(self, text: str) -> List[str]: + return [s for s in (p.strip() for p in self._SENTENCE.split(text)) if s] + + def _fallback(self, documents: List[Document]) -> List[Document]: + recursive = RecursiveChunker( + chunking_strategy=self.chunking_strategy, + max_tokens=self.max_tokens, + min_tokens=self.min_tokens, + duplicate_headers=self.duplicate_headers, + ) + return recursive.chunk(documents) + + def _breakpoints(self, embeddings) -> set: + """Indices after which a new chunk starts, from high-distance gaps.""" + import numpy as np + + matrix = np.asarray(embeddings, dtype=np.float64) + if matrix.ndim != 2 or matrix.shape[0] < 2: + return set() + norms = np.linalg.norm(matrix, axis=1) + norms[norms == 0] = 1.0 + unit = matrix / norms[:, None] + sims = np.sum(unit[:-1] * unit[1:], axis=1) + distances = 1.0 - sims + if not np.all(np.isfinite(distances)): + raise ValueError("non-finite semantic distances") + if float(distances.max()) <= float(distances.min()): + return set() + threshold = np.percentile(distances, self._PERCENTILE) + return {int(i) for i in np.where(distances >= threshold)[0]} + + def _group(self, sentences: List[str], breakpoints: set) -> List[str]: + groups: List[str] = [] + current: List[str] = [] + for idx, sentence in enumerate(sentences): + current.append(sentence) + if idx in breakpoints: + groups.append(" ".join(current)) + current = [] + if current: + groups.append(" ".join(current)) + return groups + + def _enforce_tokens(self, groups: List[str]) -> List[str]: + """Hard-split groups over ``max_tokens`` then merge to clear min.""" + capped: List[str] = [] + for group in groups: + if self._token_count(group) <= self.max_tokens: + capped.append(group) + else: + capped.extend(p for p in self._split_by_tokens(group) if p.strip()) + return self._merge_to_min(capped, " ") + + def _chunk_text(self, text: str) -> List[str]: + sentences = self._split_sentences(text) + if len(sentences) < 2: + raise ValueError("too few sentences for semantic chunking") + from application.vectorstore.base import EmbeddingsSingleton + from application.core.settings import settings + + embeddings = EmbeddingsSingleton.get_instance( + settings.EMBEDDINGS_NAME + ).embed_documents(sentences) + breakpoints = self._breakpoints(embeddings) + groups = self._group(sentences, breakpoints) + return [g for g in self._enforce_tokens(groups) if g.strip()] + + def chunk(self, documents: List[Document]) -> List[Document]: + processed: List[Document] = [] + for doc in documents: + try: + texts = self._chunk_text(doc.text) + except Exception as exc: + logger.warning( + "Semantic chunking failed (%s); falling back to recursive.", + exc, + ) + processed.extend(self._fallback([doc])) + continue + for idx, text in enumerate(texts): + processed.append(self._emit(doc, idx, text)) + return processed + + ChunkerCreator.register("recursive", RecursiveChunker) ChunkerCreator.register("markdown", MarkdownChunker) ChunkerCreator.register("parent_child", ParentChildChunker) +ChunkerCreator.register("semantic", SemanticChunker) # Reuse the classic Chunker reference so this module can be the single import # that pulls every strategy into the registry. @@ -220,5 +326,6 @@ __all__ = [ "RecursiveChunker", "MarkdownChunker", "ParentChildChunker", + "SemanticChunker", "Chunker", ] diff --git a/application/requirements.txt b/application/requirements.txt index 5414df03..822259c6 100644 --- a/application/requirements.txt +++ b/application/requirements.txt @@ -69,6 +69,7 @@ pandas==3.0.2 openpyxl==3.1.5 pathable==0.5.0 pdf2image>=1.17.0 +pgvector>=0.2,<1 pillow portalocker>=2.7.0,<4.0.0 prompt-toolkit==3.0.52 diff --git a/application/retriever/classic_rag.py b/application/retriever/classic_rag.py index a95e7750..154fc7b0 100644 --- a/application/retriever/classic_rag.py +++ b/application/retriever/classic_rag.py @@ -144,6 +144,21 @@ class ClassicRAG(BaseRetriever): logging.error(f"Error rephrasing query: {e}", exc_info=True) return self.original_question + def _fetch_candidates(self, docsearch, question, src_k, score_threshold): + """Fetch candidate hits for one vector store (vector search). + + Subclasses override this to change candidate sourcing (e.g. RRF fusion) + while inheriting the surrounding per-source resolution and budgeting. + """ + # ``score_threshold`` is honoured by pgvector/mongodb and safely ignored + # by stores whose ``search`` swallows kwargs. The candidate count is + # clamped to a ceiling to bound memory/latency. + k = min(max(src_k * 2, 20), 500) + search_kwargs = {"k": k} + if score_threshold is not None: + search_kwargs["score_threshold"] = score_threshold + return docsearch.search(question, **search_kwargs) + def _get_data(self): if self.chunks == 0 or not self.vectorstores: logging.info( @@ -193,12 +208,9 @@ class ClassicRAG(BaseRetriever): docsearch = VectorCreator.create_vectorstore( settings.VECTOR_STORE, vectorstore_id, settings.EMBEDDINGS_KEY ) - # ``score_threshold`` is honoured by pgvector/mongodb and - # safely ignored by stores whose ``search`` swallows kwargs. - search_kwargs = {"k": max(src_k * 2, 20)} - if score_threshold is not None: - search_kwargs["score_threshold"] = score_threshold - docs_temp = docsearch.search(question, **search_kwargs) + docs_temp = self._fetch_candidates( + docsearch, question, src_k, score_threshold + ) for doc in docs_temp: if cumulative_tokens >= token_budget: diff --git a/application/retriever/dispatcher.py b/application/retriever/dispatcher.py index f995aaa1..38ebd325 100644 --- a/application/retriever/dispatcher.py +++ b/application/retriever/dispatcher.py @@ -32,6 +32,11 @@ logger = logging.getLogger(__name__) # for the all-classic case. _CLASSIC_KEYS = frozenset({"classic", "default"}) +# Retriever keys that subclass ClassicRAG and accept ``defer_rephrase``, so the +# eager ctor rephrase can be skipped and computed lazily per-source. Unknown +# keys are excluded: a future non-ClassicRAG retriever may not accept the kwarg. +_DEFERRABLE_KEYS = _CLASSIC_KEYS | {"hybrid"} + # Fields that, when left at their defaults, mean the source did not opt into any # per-source retrieval override — so it can flow through the global ClassicRAG # path unchanged (byte-identical parity). @@ -220,7 +225,7 @@ class Dispatcher(BaseRetriever): kwargs["chunks"] = max(self.chunks, candidate_k or 0) # With per-source configs the rephrase decision is per-source, so defer # the eager rephrase to let a rephrase_query=False source skip the call. - if group["retrievals"] and retriever_key in _CLASSIC_KEYS: + if group["retrievals"] and retriever_key in _DEFERRABLE_KEYS: kwargs["defer_rephrase"] = True retriever = RetrieverCreator.create_retriever(retriever_key, **kwargs) # Hand the per-source retrieval configs to the classic retriever so it diff --git a/application/retriever/hybrid_rag.py b/application/retriever/hybrid_rag.py new file mode 100644 index 00000000..b789ddc5 --- /dev/null +++ b/application/retriever/hybrid_rag.py @@ -0,0 +1,62 @@ +"""Hybrid retriever fusing vector and keyword search via RRF. + +Subclasses :class:`ClassicRAG` so the Dispatcher builds it with identical +ctor kwargs and it inherits rephrase + token-budgeting. Only the per-source +fetch is overridden: for each vector store it pulls vector hits and keyword +hits, then fuses them with Reciprocal Rank Fusion. Stores without keyword +support (``keyword_search`` returns ``[]``) reduce to exact vector-only +behaviour. +""" + +from application.retriever.classic_rag import ClassicRAG + +RRF_K = 60 + + +def _doc_key(doc): + """Stable identity for a hit so the same chunk fuses across both lists.""" + if hasattr(doc, "page_content") and hasattr(doc, "metadata"): + content = doc.page_content + metadata = doc.metadata or {} + else: + content = doc.get("text", doc.get("page_content", "")) + metadata = doc.get("metadata") or {} + source = metadata.get("source", "") + return (source, content) + + +def reciprocal_rank_fusion(vector_hits, keyword_hits, k=RRF_K): + """Fuse two ranked hit lists into one by Reciprocal Rank Fusion. + + Each list contributes ``1 / (k + rank)`` per document (rank 0-based); + documents are returned ordered by summed score, highest first. A document + present in only one list is ranked solely on that list's contribution, so + an empty ``keyword_hits`` yields exactly the vector ordering. + """ + scores = {} + docs = {} + for hits in (vector_hits, keyword_hits): + for rank, doc in enumerate(hits): + key = _doc_key(doc) + scores[key] = scores.get(key, 0.0) + 1.0 / (k + rank) + if key not in docs: + docs[key] = doc + ordered = sorted(docs.keys(), key=lambda key: scores[key], reverse=True) + return [docs[key] for key in ordered] + + +class HybridRetriever(ClassicRAG): + """ClassicRAG variant that fuses vector + keyword search with RRF.""" + + def _fetch_candidates(self, docsearch, question, src_k, score_threshold): + """Return RRF-fused vector+keyword hits for one vector store. + + Inherits the per-source resolution and budgeting from + :meth:`ClassicRAG._get_data`; only candidate sourcing differs. + RRF scores are not cosine similarities, so ``score_threshold`` is + intentionally not applied to the fused list. + """ + candidate_k = min(max(src_k * 2, 20), 500) + vector_hits = docsearch.search(question, k=candidate_k) + keyword_hits = docsearch.keyword_search(question, k=candidate_k) + return reciprocal_rank_fusion(vector_hits, keyword_hits) diff --git a/application/retriever/retriever_creator.py b/application/retriever/retriever_creator.py index 2b507bf3..43d75035 100644 --- a/application/retriever/retriever_creator.py +++ b/application/retriever/retriever_creator.py @@ -1,10 +1,12 @@ from application.retriever.classic_rag import ClassicRAG +from application.retriever.hybrid_rag import HybridRetriever class RetrieverCreator: retrievers = { "classic": ClassicRAG, "default": ClassicRAG, + "hybrid": HybridRetriever, } @classmethod diff --git a/application/storage/db/source_config.py b/application/storage/db/source_config.py index fdf648dd..540073e7 100644 --- a/application/storage/db/source_config.py +++ b/application/storage/db/source_config.py @@ -73,6 +73,15 @@ class RetrievalConfig(BaseModel): reranker: Optional[dict] = None # reserved: future cross-encoder/LLM reorder prescreen: Optional[dict] = None # None = off; else PreScreenConfig dict (D12) + @field_validator("chunks") + @classmethod + def _bounded_chunks(cls, value: int) -> int: + if value < 1: + raise ValueError("must be >= 1") + if value > 500: + raise ValueError("must be <= 500") + return value + @model_validator(mode="after") def _validate_prescreen(self) -> "RetrievalConfig": """Validate ``prescreen`` through ``PreScreenConfig`` when present. diff --git a/application/vectorstore/base.py b/application/vectorstore/base.py index b3cf4da4..d61177a8 100644 --- a/application/vectorstore/base.py +++ b/application/vectorstore/base.py @@ -136,6 +136,14 @@ class BaseVectorStore(ABC): """Search for similar documents/chunks in the vectorstore""" pass + def keyword_search(self, question, k=10): + """Keyword/full-text search. + + Default returns no results so hybrid retrieval degrades to vector-only + on stores without keyword support. Override in stores that support it. + """ + return [] + @abstractmethod def add_texts(self, texts, metadatas=None, *args, **kwargs): """Add texts with their embeddings to the vectorstore""" diff --git a/application/vectorstore/pgvector.py b/application/vectorstore/pgvector.py index 43b1585f..78925cbb 100644 --- a/application/vectorstore/pgvector.py +++ b/application/vectorstore/pgvector.py @@ -102,11 +102,18 @@ class PGVectorStore(BaseVectorStore): # Create index for source_id filtering source_index_query = f""" - CREATE INDEX IF NOT EXISTS {self._table_name}_source_id_idx + CREATE INDEX IF NOT EXISTS {self._table_name}_source_id_idx ON {self._table_name} (source_id); """ cursor.execute(source_index_query) - + + # Functional GIN index backing keyword_search full-text queries. + fts_index_query = f""" + CREATE INDEX IF NOT EXISTS {self._table_name}_text_fts_idx + ON {self._table_name} USING gin(to_tsvector('english', {self._text_column})); + """ + cursor.execute(fts_index_query) + conn.commit() except Exception as e: conn.rollback() @@ -170,6 +177,48 @@ class PGVectorStore(BaseVectorStore): finally: cursor.close() + def keyword_search(self, question: str, k: int = 10) -> List[Document]: + """Full-text keyword search using Postgres ``websearch_to_tsquery``. + + Returns the same ``Document`` shape as :meth:`search`. The question is + bound as a query parameter (never interpolated) to prevent injection. + """ + conn = self._get_connection() + cursor = conn.cursor() + + try: + keyword_query = f""" + SELECT {self._text_column}, {self._metadata_column}, + ts_rank( + to_tsvector('english', {self._text_column}), + websearch_to_tsquery('english', %s) + ) AS rank + FROM {self._table_name} + WHERE source_id = %s + AND to_tsvector('english', {self._text_column}) + @@ websearch_to_tsquery('english', %s) + ORDER BY rank DESC + LIMIT %s; + """ + + cursor.execute( + keyword_query, (question, self._source_id, question, k) + ) + results = cursor.fetchall() + + documents = [] + for text, metadata, _rank in results: + metadata = metadata or {} + documents.append(Document(page_content=text, metadata=metadata)) + + return documents + + except Exception as e: + logging.error(f"Error in keyword search: {e}", exc_info=True) + return [] + finally: + cursor.close() + def add_texts( self, texts: List[str], diff --git a/frontend/src/locale/en.json b/frontend/src/locale/en.json index 3bd8400a..4f5330e7 100644 --- a/frontend/src/locale/en.json +++ b/frontend/src/locale/en.json @@ -145,8 +145,10 @@ "rephraseQuery": "Rephrase query before search", "retriever": "Retriever", "retrievers": { - "classic": "Classic" + "classic": "Classic", + "hybrid": "Hybrid" }, + "retrieverHint": "Hybrid (keyword + vector) search is supported on pgvector; other vector stores fall back to vector-only.", "exposure": "Search exposure", "exposures": { "prefetch": "Pre-fetch into context", @@ -172,7 +174,8 @@ "classic_chunk": "Classic", "recursive": "Recursive", "markdown": "Markdown", - "parent_child": "Parent / child" + "parent_child": "Parent / child", + "semantic": "Semantic" }, "maxTokens": "Max tokens per chunk", "minTokens": "Min tokens per chunk", diff --git a/frontend/src/models/misc.ts b/frontend/src/models/misc.ts index 004e397a..7b821f2b 100644 --- a/frontend/src/models/misc.ts +++ b/frontend/src/models/misc.ts @@ -11,7 +11,8 @@ export type ChunkingStrategy = | 'classic_chunk' | 'recursive' | 'markdown' - | 'parent_child'; + | 'parent_child' + | 'semantic'; export type RetrievalExposure = 'prefetch' | 'agentic_tool'; diff --git a/frontend/src/settings/components/RetrievalOptions.tsx b/frontend/src/settings/components/RetrievalOptions.tsx index 9552fff1..fb2a82a6 100644 --- a/frontend/src/settings/components/RetrievalOptions.tsx +++ b/frontend/src/settings/components/RetrievalOptions.tsx @@ -83,8 +83,11 @@ const CHUNKING_STRATEGIES: ChunkingStrategy[] = [ 'recursive', 'markdown', 'parent_child', + 'semantic', ]; +const RETRIEVERS = ['classic', 'hybrid']; + /** * Group eyebrow: an uppercase, tracked title with a normal-weight muted ` · tag` * suffix. Reads as more prominent than the individual field labels below it. @@ -302,6 +305,34 @@ export default function RetrievalOptions({
+ + + + zero distances -> no split. + text = "Same one. Same two. Same three. Same four." + vectors = [[1.0, 0.0]] * 4 + chunker = SemanticChunker(max_tokens=2000, min_tokens=0) + with patch(_EMB_TARGET, return_value=_FakeEmbeddings(vectors)): + out = chunker.chunk([Document(text=text, doc_id="d")]) + assert len(out) == 1 + assert out[0].extra_info["token_count"] == _tok(out[0].text) + + def test_max_tokens_enforced(self): + # A single semantic group larger than max_tokens is hard-split. + long_sentence = "word " * 300 + "." + text = f"{long_sentence} {long_sentence}" + vectors = [[1.0, 0.0], [1.0, 0.0]] + chunker = SemanticChunker(max_tokens=40, min_tokens=0) + with patch(_EMB_TARGET, return_value=_FakeEmbeddings(vectors)): + out = chunker.chunk([Document(text=text, doc_id="d")]) + assert len(out) > 1 + for c in out: + assert _tok(c.text) <= 40 + + def test_min_tokens_merges_neighbours(self): + # Non-uniform distances yield several breakpoints and tiny groups, + # which must merge until they clear min_tokens. + text = "A. B. C. D. E. F." + vectors = [ + [1.0, 0.0], + [0.0, 1.0], + [0.0, 1.0], + [1.0, 0.0], + [0.0, 1.0], + [0.0, 1.0], + ] + chunker = SemanticChunker(max_tokens=2000, min_tokens=8) + with patch(_EMB_TARGET, return_value=_FakeEmbeddings(vectors)): + out = chunker.chunk([Document(text=text, doc_id="d")]) + assert len(out) < 6 + assert _tok(out[0].text) >= 8 + + def test_embeddings_error_falls_back_to_recursive(self): + text = "First sentence here. Second sentence here. Third one." + + def _boom(*args, **kwargs): + raise RuntimeError("model unavailable") + + chunker = SemanticChunker(max_tokens=2000, min_tokens=0) + with patch(_EMB_TARGET, side_effect=_boom): + out = chunker.chunk([Document(text=text, doc_id="d")]) + recursive = RecursiveChunker(max_tokens=2000, min_tokens=0) + expected = recursive.chunk([Document(text=text, doc_id="d")]) + assert [c.text for c in out] == [c.text for c in expected] + + def test_too_few_sentences_falls_back(self): + # A single sentence cannot be semantically split. + chunker = SemanticChunker(max_tokens=2000, min_tokens=0) + with patch(_EMB_TARGET, side_effect=AssertionError("must not embed")): + out = chunker.chunk([Document(text="just one sentence", doc_id="d")]) + assert len(out) == 1 + assert out[0].text.strip() == "just one sentence" + + def test_source_and_extra_info_preserved(self): + text = "Alpha one. Alpha two. Beta one. Beta two." + vectors = [[1.0, 0.0], [1.0, 0.0], [0.0, 1.0], [0.0, 1.0]] + doc = Document( + text=text, + doc_id="d", + extra_info={"source": "file.md", "title": "T"}, + ) + chunker = SemanticChunker(max_tokens=2000, min_tokens=0) + with patch(_EMB_TARGET, return_value=_FakeEmbeddings(vectors)): + out = chunker.chunk([doc]) + assert len(out) == 2 + for c in out: + assert c.extra_info["source"] == "file.md" + assert c.extra_info["title"] == "T" + assert c.extra_info["token_count"] == _tok(c.text) + assert c.doc_id.startswith("d-") + + @pytest.mark.unit class TestClassicByteIdentical: def test_classic_chunk_unchanged(self): diff --git a/tests/retriever/test_hybrid.py b/tests/retriever/test_hybrid.py new file mode 100644 index 00000000..7f5efa5f --- /dev/null +++ b/tests/retriever/test_hybrid.py @@ -0,0 +1,188 @@ +"""Tests for the hybrid retriever (vector + keyword RRF fusion).""" + +from unittest.mock import MagicMock, Mock, patch + +import pytest + +from application.retriever.hybrid_rag import HybridRetriever, reciprocal_rank_fusion +from application.retriever.retriever_creator import RetrieverCreator + + +@pytest.fixture +def _patch_llm_creator(mock_llm, monkeypatch): + monkeypatch.setattr( + "application.retriever.classic_rag.LLMCreator.create_llm", + Mock(return_value=mock_llm), + ) + return mock_llm + + +def _make_doc(page_content, source="s", title="t"): + doc = Mock() + doc.page_content = page_content + doc.metadata = {"title": title, "source": source} + return doc + + +def _make_hybrid(source=None, **overrides): + defaults = dict( + source=source or {"question": "q", "active_docs": ["vs1"]}, + chat_history=None, + prompt="", + chunks=2, + doc_token_limit=50000, + model_id="test-model", + llm_name="openai", + api_key="fake", + decoded_token={"sub": "user1"}, + ) + defaults.update(overrides) + return HybridRetriever(**defaults) + + +# ── RRF fusion ────────────────────────────────────────────────────────────── + + +@pytest.mark.unit +class TestReciprocalRankFusion: + def test_doc_in_both_lists_outranks_singletons(self): + shared = _make_doc("shared", source="x") + only_vec = _make_doc("vec_only", source="v") + only_kw = _make_doc("kw_only", source="k") + # "shared" is rank-1 in vector and rank-0 in keyword → highest summed score. + vector_hits = [only_vec, shared] + keyword_hits = [shared, only_kw] + + fused = reciprocal_rank_fusion(vector_hits, keyword_hits) + + assert fused[0].page_content == "shared" + assert {d.page_content for d in fused} == {"shared", "vec_only", "kw_only"} + + def test_empty_keyword_is_vector_only_order(self): + vector_hits = [_make_doc("a", source="a"), _make_doc("b", source="b")] + fused = reciprocal_rank_fusion(vector_hits, []) + assert [d.page_content for d in fused] == ["a", "b"] + + def test_dedupes_same_doc(self): + d_vec = _make_doc("same", source="same") + d_kw = _make_doc("same", source="same") + fused = reciprocal_rank_fusion([d_vec], [d_kw]) + assert len(fused) == 1 + + def test_higher_keyword_rank_can_promote(self): + # Vector top is "v0"; keyword strongly favours "kw" (rank 0 vs v0's rank 1). + v0 = _make_doc("v0", source="v0") + kw = _make_doc("kw", source="kw") + fused = reciprocal_rank_fusion([v0, kw], [kw]) + assert fused[0].page_content == "kw" + + +# ── HybridRetriever._get_data ──────────────────────────────────────────────── + + +@pytest.mark.unit +class TestHybridGetData: + @patch("application.retriever.classic_rag.VectorCreator") + @patch("application.retriever.classic_rag.num_tokens_from_string", return_value=10) + def test_fuses_vector_and_keyword(self, _tok, mock_vc, _patch_llm_creator): + docsearch = MagicMock() + docsearch.search.return_value = [_make_doc("vec", source="vec")] + docsearch.keyword_search.return_value = [_make_doc("kw", source="kw")] + mock_vc.create_vectorstore.return_value = docsearch + + rag = _make_hybrid() + docs = rag._get_data() + + docsearch.search.assert_called_once() + docsearch.keyword_search.assert_called_once() + assert {d["text"] for d in docs} == {"vec", "kw"} + + @patch("application.retriever.classic_rag.VectorCreator") + @patch("application.retriever.classic_rag.num_tokens_from_string", return_value=10) + def test_keyword_empty_equals_vector_only(self, _tok, mock_vc, _patch_llm_creator): + vec_docs = [_make_doc("a", source="a"), _make_doc("b", source="b")] + + # Hybrid with empty keyword results. + ds_hybrid = MagicMock() + ds_hybrid.search.return_value = list(vec_docs) + ds_hybrid.keyword_search.return_value = [] + mock_vc.create_vectorstore.return_value = ds_hybrid + hybrid_out = _make_hybrid().search("query") + + # Vector-only baseline: same vector hits, no keyword call. + from application.retriever.classic_rag import ClassicRAG + + with patch( + "application.retriever.classic_rag.VectorCreator" + ) as mock_vc_classic, patch( + "application.retriever.classic_rag.num_tokens_from_string", return_value=10 + ): + ds_classic = MagicMock() + ds_classic.search.return_value = list(vec_docs) + mock_vc_classic.create_vectorstore.return_value = ds_classic + classic_out = ClassicRAG( + source={"question": "q", "active_docs": ["vs1"]}, + chat_history=None, + chunks=2, + doc_token_limit=50000, + model_id="test-model", + llm_name="openai", + api_key="fake", + decoded_token={"sub": "user1"}, + ).search("query") + + assert hybrid_out == classic_out + + @patch("application.retriever.classic_rag.VectorCreator") + @patch("application.retriever.classic_rag.num_tokens_from_string", return_value=10) + def test_score_threshold_not_applied_to_fused(self, _tok, mock_vc, _patch_llm_creator): + from application.storage.db.source_config import RetrievalConfig + + docsearch = MagicMock() + docsearch.search.return_value = [_make_doc("a", source="a")] + docsearch.keyword_search.return_value = [] + mock_vc.create_vectorstore.return_value = docsearch + + rag = _make_hybrid() + rag.per_source_retrieval = {"vs1": RetrievalConfig(score_threshold=0.9)} + rag._get_data() + + # RRF scores are not cosine — score_threshold must not reach the store. + assert "score_threshold" not in docsearch.search.call_args.kwargs + assert "score_threshold" not in docsearch.keyword_search.call_args.kwargs + + @patch("application.retriever.classic_rag.VectorCreator") + @patch("application.retriever.classic_rag.num_tokens_from_string", return_value=10) + def test_chunks_zero_returns_empty(self, _tok, mock_vc, _patch_llm_creator): + rag = _make_hybrid(chunks=0) + assert rag._get_data() == [] + mock_vc.create_vectorstore.assert_not_called() + + @patch("application.retriever.classic_rag.VectorCreator") + @patch("application.retriever.classic_rag.num_tokens_from_string", return_value=10) + def test_store_error_continues(self, _tok, mock_vc, _patch_llm_creator): + mock_vc.create_vectorstore.side_effect = RuntimeError("boom") + rag = _make_hybrid() + assert rag._get_data() == [] + + +# ── Registry resolution ────────────────────────────────────────────────────── + + +@pytest.mark.unit +class TestHybridRegistration: + def test_hybrid_resolves_via_creator(self): + assert RetrieverCreator.retrievers["hybrid"] is HybridRetriever + + def test_create_retriever_builds_hybrid(self, _patch_llm_creator): + retriever = RetrieverCreator.create_retriever( + "hybrid", + source={"question": "q", "active_docs": ["vs1"]}, + chunks=2, + doc_token_limit=50000, + model_id="m", + llm_name="openai", + api_key="fake", + decoded_token={"sub": "u"}, + ) + assert isinstance(retriever, HybridRetriever) diff --git a/tests/storage/db/test_source_config.py b/tests/storage/db/test_source_config.py index b0f86cc5..f1883697 100644 --- a/tests/storage/db/test_source_config.py +++ b/tests/storage/db/test_source_config.py @@ -68,6 +68,18 @@ class TestStrictWrite: with pytest.raises(ValidationError): SourceConfig.model_validate({"unexpected": True}) + def test_chunks_upper_bound_rejected(self): + with pytest.raises(ValidationError): + RetrievalConfig(chunks=501) + + def test_chunks_lower_bound_rejected(self): + with pytest.raises(ValidationError): + RetrievalConfig(chunks=0) + + def test_chunks_accepts_small_and_ceiling_values(self): + assert RetrievalConfig(chunks=2).chunks == 2 + assert RetrievalConfig(chunks=500).chunks == 500 + def test_model_validate_accepts_full_valid_config(self): cfg = SourceConfig.model_validate( { diff --git a/tests/vectorstore/test_pgvector.py b/tests/vectorstore/test_pgvector.py index 1771ea5b..5d87c09b 100644 --- a/tests/vectorstore/test_pgvector.py +++ b/tests/vectorstore/test_pgvector.py @@ -131,6 +131,61 @@ class TestPGVectorStoreSearch: assert len(results) == 2 +@pytest.mark.unit +class TestPGVectorStoreKeywordSearch: + def test_keyword_search_returns_documents(self): + store, _, mock_cursor, mock_emb = _make_store(source_id="src1") + mock_cursor.fetchall.return_value = [ + ("hello world", {"source": "a.txt"}, 0.9), + ("foo bar", {"source": "b.txt"}, 0.3), + ] + + results = store.keyword_search("hello", k=5) + + # Keyword search must not embed the query. + mock_emb.embed_query.assert_not_called() + assert len(results) == 2 + assert results[0].page_content == "hello world" + assert results[0].metadata == {"source": "a.txt"} + + def test_keyword_search_is_parameterized(self): + store, _, mock_cursor, _ = _make_store(source_id="src1") + mock_cursor.fetchall.return_value = [] + + store.keyword_search("DROP TABLE documents; --", k=7) + + sql, params = mock_cursor.execute.call_args[0] + # The raw question must never be interpolated into the SQL text. + assert "DROP TABLE documents" not in sql + assert "websearch_to_tsquery('english', %s)" in sql + # Question is bound twice (rank + WHERE), then source_id and k. + assert params == ("DROP TABLE documents; --", "src1", "DROP TABLE documents; --", 7) + + def test_keyword_search_handles_null_metadata(self): + store, _, mock_cursor, _ = _make_store() + mock_cursor.fetchall.return_value = [("text", None, 0.5)] + + results = store.keyword_search("query") + assert len(results) == 1 + assert results[0].metadata == {} + + def test_keyword_search_returns_empty_on_error(self): + store, _, mock_cursor, _ = _make_store() + mock_cursor.execute.side_effect = Exception("fts failed") + + assert store.keyword_search("query") == [] + + def test_ensure_table_exists_creates_fts_index(self): + store, mock_conn, mock_cursor, _ = _make_store() + store._ensure_table_exists() + + executed = " ".join( + str(call.args[0]) for call in mock_cursor.execute.call_args_list + ) + assert "documents_text_fts_idx" in executed + assert "gin(to_tsvector('english'" in executed + + @pytest.mark.unit class TestPGVectorStoreAddTexts: def test_add_texts_inserts_and_returns_ids(self):