diff --git a/search/schemas.py b/search/schemas.py index d4246f3..73f63f5 100644 --- a/search/schemas.py +++ b/search/schemas.py @@ -8,6 +8,7 @@ """ from __future__ import annotations +import re from typing import Any, Dict, List, Literal, Optional, Union from pydantic import BaseModel, Field, model_validator @@ -129,17 +130,59 @@ class IndexFilter(BaseModel): # ----------------------------------------------------------------------------- +_FTS_FIELD_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*(\.[A-Za-z_][A-Za-z0-9_]*)*$") + + class IndexSpec(BaseModel): id: str = Field(..., description="Caller-supplied label; echoed in results") store: StoreConfig = Field(..., discriminator="type") + mode: Literal["vector", "fts"] = Field( + default="vector", + description=( + "Search mode. 'vector' (default) performs ANN over a vector column. " + "'fts' performs Cosmos DB full-text search via FullTextScore and requires " + "the container to have a full-text indexing policy on `fts_field`." + ), + ) vector: VectorConfig = Field(default_factory=VectorConfig) - embedding: EmbeddingPolicy = Field(..., discriminator="policy") + embedding: Optional[EmbeddingPolicy] = Field( + default=None, + discriminator="policy", + description="Required when mode='vector'; ignored when mode='fts'.", + ) + fts_field: Optional[str] = Field( + default=None, + description=( + "Document path to full-text search (e.g. 'content' or 'body.text'). " + "When omitted, defaults to the first entry in `content_fields`. " + "Only used when mode='fts'." + ), + ) content_fields: List[str] = Field(default_factory=lambda: ["content"]) return_fields: List[str] = Field(default_factory=list) filter: Optional[IndexFilter] = None top_k: Optional[int] = Field(default=None, ge=1, le=500) pipeline_id: Optional[str] = Field(default=None, description="Source pipeline id; surfaced into result metadata for blob preview") + @model_validator(mode="after") + def _validate_mode(self): + if self.mode == "vector": + if self.embedding is None: + raise ValueError("embedding is required when mode='vector'") + else: # fts + if not isinstance(self.store, CosmosStore): + raise ValueError("mode='fts' is only supported for cosmosdb store in v1") + field = self.fts_field or (self.content_fields[0] if self.content_fields else None) + if not field: + raise ValueError( + "mode='fts' requires `fts_field` or at least one entry in `content_fields`" + ) + if not _FTS_FIELD_RE.match(field): + raise ValueError( + f"invalid fts_field {field!r}: must match [A-Za-z_][A-Za-z0-9_]*(\\.[A-Za-z_][A-Za-z0-9_]*)*" + ) + return self + # ----------------------------------------------------------------------------- # Merge + include diff --git a/search/searcher.py b/search/searcher.py index 5c7f337..ccffdb1 100644 --- a/search/searcher.py +++ b/search/searcher.py @@ -278,6 +278,126 @@ def _run(): return await asyncio.to_thread(_run) +# ============================================================================= +# Cosmos full-text search +# ============================================================================= + + +_FTS_FIELD_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*(\.[A-Za-z_][A-Za-z0-9_]*)*$") +_FTS_TOKEN_RE = re.compile(r"[A-Za-z0-9_'\-]+") +_FTS_MAX_TERMS = int(os.getenv("SEARCH_FTS_MAX_TERMS", "16")) + + +def _tokenize_fts_query(query: str) -> List[str]: + """Conservative tokenizer for Cosmos FullTextScore. + + Splits on non-alphanumeric (keeping hyphens, apostrophes, underscores), + lowercases, drops empties + duplicates while preserving order, and caps + at SEARCH_FTS_MAX_TERMS to bound SQL size + RU cost. + """ + if not query: + return [] + seen: Dict[str, None] = {} + for tok in _FTS_TOKEN_RE.findall(query.lower()): + if tok and tok not in seen: + seen[tok] = None + if len(seen) >= _FTS_MAX_TERMS: + break + return list(seen.keys()) + + +async def search_cosmos_fts( + store: CosmosStore, + fts_field: str, + query_terms: List[str], + top_k: int, + content_fields: List[str], + return_fields: List[str], + index_filter: Optional[IndexFilter], + include_vector: bool, +) -> List[dict]: + """Cosmos DB full-text search via `FullTextScore` (sync SDK wrapped). + + Requires the target container to have a full-text indexing policy on + `fts_field`. FullTextScore can only appear in `ORDER BY RANK`, so the + native BM25 score is not returned — hits get `score=None` and merge via + rank-derived RRF. + """ + from azure.cosmos import CosmosClient + from azure.identity import DefaultAzureCredential + + if not _FTS_FIELD_RE.match(fts_field): + raise ValueError(f"invalid fts_field {fts_field!r}") + if not query_terms: + return [] + + if getattr(store.auth, "mode", "managed_identity") == "key": + key = await resolve_secret_ref(store.auth.secret_ref) # type: ignore[attr-defined] + client_args = {"credential": key} + else: + client_args = {"credential": DefaultAzureCredential()} + + def _run(): + client = CosmosClient(store.endpoint, **client_args) + database = client.get_database_client(store.database) + container = database.get_container_client(store.container) + + ff = fts_field.lstrip("/").replace("/", ".") + + params: List[dict] = [{"name": "@top_k", "value": int(top_k)}] + term_placeholders: List[str] = [] + for i, term in enumerate(query_terms): + ph = f"@term{i}" + term_placeholders.append(ph) + params.append({"name": ph, "value": term}) + + where = "" + if index_filter and index_filter.where: + where = f"WHERE {index_filter.where}" + if isinstance(index_filter.params, dict): + for k, v in index_filter.params.items(): + name = k if k.startswith("@") else f"@{k}" + params.append({"name": name, "value": v}) + + term_args = ", ".join(term_placeholders) + query = ( + f"SELECT TOP @top_k c FROM c {where} " + f"ORDER BY RANK FullTextScore(c.{ff}, {term_args})" + ) + + hits: List[dict] = [] + for item in container.query_items( + query=query, parameters=params, enable_cross_partition_query=True + ): + if isinstance(item, dict) and "c" in item: + doc = item.get("c") or {} + else: + # Cosmos may return the doc inlined when only one column is projected. + doc = item if isinstance(item, dict) else {} + text, text_parts = _extract_text(doc, content_fields) + + metadata: Dict[str, Any] = dict(doc.get("metadata", {}) or {}) + for f in return_fields: + if f in doc and f not in metadata: + metadata[f] = doc[f] + + hit: Dict[str, Any] = { + "id": doc.get("id"), + "score": None, # FullTextScore not available in projection + "text": text, + "text_parts": text_parts, + "metadata": metadata, + "source": doc.get("source"), + "source_ref": doc.get("source_ref") or doc.get("title") or doc.get("url"), + } + if include_vector: + hit["vector"] = None + hits.append(hit) + return hits + + return await asyncio.to_thread(_run) + + async def search_pgvector( store: PgVectorStore, vec_field: str, @@ -401,6 +521,50 @@ async def _search_one_index( info = PerIndexInfo(index_id=idx.id) per_top_k = idx.top_k or default_per_index_top_k + # ------------------------------------------------------------------------- + # FTS branch — no embedding call; Cosmos-only in v1. + # ------------------------------------------------------------------------- + if idx.mode == "fts": + info.embedding_model = "fts" + if not query: + info.error = "fts mode requires text query" + return info, [] + if not isinstance(idx.store, CosmosStore): + info.error = "fts mode only supports cosmosdb store in v1" + return info, [] + terms = _tokenize_fts_query(query) + if not terms: + info.error = "fts mode: query produced no searchable tokens" + return info, [] + fts_field = idx.fts_field or (idx.content_fields[0] if idx.content_fields else "content") + t1 = time.time() + try: + hits = await asyncio.wait_for( + search_cosmos_fts( + idx.store, fts_field, terms, per_top_k, + idx.content_fields, idx.return_fields, idx.filter, include_vector, + ), + timeout=PER_INDEX_TIMEOUT_S, + ) + except asyncio.TimeoutError: + info.error = f"per-index timeout > {PER_INDEX_TIMEOUT_S}s" + info.search_ms = int((time.time() - t1) * 1000) + return info, [] + except Exception as e: + info.error = f"search: {str(e)[:180]}" + info.search_ms = int((time.time() - t1) * 1000) + return info, [] + info.search_ms = int((time.time() - t1) * 1000) + info.result_count = len(hits) + return info, hits + + # ------------------------------------------------------------------------- + # Vector branch — embed then search. + # ------------------------------------------------------------------------- + if idx.embedding is None: + info.error = "embedding policy required for vector mode" + return info, [] + # Embed t0 = time.time() try: @@ -573,6 +737,10 @@ async def run_search(http: httpx.AsyncClient, req: SearchRequest) -> SearchRespo warnings.append( "mixed embedding models with merge.strategy=score; scores are not directly comparable — consider 'rrf'" ) + if req.merge.strategy == "score" and any(ix.mode == "fts" for ix in req.indexes): + warnings.append( + "merge.strategy=score with fts indexes will rank fts hits as None — use 'rrf' for hybrid" + ) if req.strict and errors: raise RuntimeError( @@ -587,7 +755,10 @@ async def run_search(http: httpx.AsyncClient, req: SearchRequest) -> SearchRespo idx_modality: Dict[str, str] = {} idx_pipeline: Dict[str, str] = {} for ix in req.indexes: - mod = getattr(ix.embedding, "input_modality", "text") + if ix.mode == "fts": + mod = "text" + else: + mod = getattr(ix.embedding, "input_modality", "text") if ix.embedding else "text" idx_modality[ix.id] = mod pid = getattr(ix, "pipeline_id", None) if pid: @@ -646,15 +817,33 @@ async def explain_search(req: SearchRequest) -> Dict[str, Any]: """Dry-run: return the resolved plan without executing embedding/search.""" plan = [] for idx in req.indexes: + if idx.mode == "fts": + fts_field = idx.fts_field or (idx.content_fields[0] if idx.content_fields else None) + plan.append({ + "index_id": idx.id, + "mode": "fts", + "store_type": idx.store.type, + "fts_field": fts_field, + "fts_terms": _tokenize_fts_query(req.query or ""), + "embedding": None, + "content_fields": idx.content_fields, + "return_fields": idx.return_fields, + "has_filter": idx.filter is not None, + "top_k": idx.top_k or req.merge.per_index_top_k, + }) + continue emb = idx.embedding if isinstance(emb, ModelEmbedding): emb_desc = {"policy": "model", "model_id": emb.model_id, "normalize": emb.normalize} elif isinstance(emb, PipelineEmbedding): emb_desc = {"policy": "pipeline", "pipeline": emb.pipeline, "normalize": emb.normalize} - else: + elif isinstance(emb, PrecomputedEmbedding): emb_desc = {"policy": "precomputed", "dims": len(emb.vector)} + else: + emb_desc = None plan.append({ "index_id": idx.id, + "mode": "vector", "store_type": idx.store.type, "vector_field": idx.vector.field, "metric": idx.vector.metric, diff --git a/tests/unit/__snapshots__/test_openapi_snapshots.ambr b/tests/unit/__snapshots__/test_openapi_snapshots.ambr index f9b5c6e..9d463a6 100644 --- a/tests/unit/__snapshots__/test_openapi_snapshots.ambr +++ b/tests/unit/__snapshots__/test_openapi_snapshots.ambr @@ -9097,25 +9097,33 @@ 'type': 'array', }), 'embedding': dict({ - 'discriminator': dict({ - 'mapping': dict({ - 'model': '#/components/schemas/ModelEmbedding', - 'pipeline': '#/components/schemas/PipelineEmbedding', - 'precomputed': '#/components/schemas/PrecomputedEmbedding', - }), - 'propertyName': 'policy', - }), - 'oneOf': list([ - dict({ - '$ref': '#/components/schemas/ModelEmbedding', - }), + 'anyOf': list([ dict({ - '$ref': '#/components/schemas/PipelineEmbedding', + 'discriminator': dict({ + 'mapping': dict({ + 'model': '#/components/schemas/ModelEmbedding', + 'pipeline': '#/components/schemas/PipelineEmbedding', + 'precomputed': '#/components/schemas/PrecomputedEmbedding', + }), + 'propertyName': 'policy', + }), + 'oneOf': list([ + dict({ + '$ref': '#/components/schemas/ModelEmbedding', + }), + dict({ + '$ref': '#/components/schemas/PipelineEmbedding', + }), + dict({ + '$ref': '#/components/schemas/PrecomputedEmbedding', + }), + ]), }), dict({ - '$ref': '#/components/schemas/PrecomputedEmbedding', + 'type': 'null', }), ]), + 'description': "Required when mode='vector'; ignored when mode='fts'.", 'title': 'Embedding', }), 'filter': dict({ @@ -9128,11 +9136,33 @@ }), ]), }), + 'fts_field': dict({ + 'anyOf': list([ + dict({ + 'type': 'string', + }), + dict({ + 'type': 'null', + }), + ]), + 'description': "Document path to full-text search (e.g. 'content' or 'body.text'). When omitted, defaults to the first entry in `content_fields`. Only used when mode='fts'.", + 'title': 'Fts Field', + }), 'id': dict({ 'description': 'Caller-supplied label; echoed in results', 'title': 'Id', 'type': 'string', }), + 'mode': dict({ + 'default': 'vector', + 'description': "Search mode. 'vector' (default) performs ANN over a vector column. 'fts' performs Cosmos DB full-text search via FullTextScore and requires the container to have a full-text indexing policy on `fts_field`.", + 'enum': list([ + 'vector', + 'fts', + ]), + 'title': 'Mode', + 'type': 'string', + }), 'pipeline_id': dict({ 'anyOf': list([ dict({ @@ -9190,7 +9220,6 @@ 'required': list([ 'id', 'store', - 'embedding', ]), 'title': 'IndexSpec', 'type': 'object', diff --git a/tests/unit/test_search_fts.py b/tests/unit/test_search_fts.py new file mode 100644 index 0000000..59a9760 --- /dev/null +++ b/tests/unit/test_search_fts.py @@ -0,0 +1,299 @@ +"""Unit tests for Cosmos full-text search mode (search/searcher.py). + +Covers schema validation, tokenization, SQL shape, per-index branching, +and RRF merging between fts + vector indexes. +""" +from __future__ import annotations + +import asyncio +import importlib +import sys + +import pytest + + +# --------------------------------------------------------------------------- +# Module loader — reuse the same sys.path trick the search_app fixture uses +# but without instantiating the FastAPI app (we only need searcher + schemas). +# --------------------------------------------------------------------------- +@pytest.fixture +def search_mods(monkeypatch): + import pathlib + REPO_ROOT = pathlib.Path(__file__).resolve().parents[2] + SEARCH_DIR = REPO_ROOT / "search" + for n in ["main", "auth", "schemas", "searcher"]: + sys.modules.pop(n, None) + monkeypatch.syspath_prepend(str(SEARCH_DIR)) + schemas = importlib.import_module("schemas") + searcher = importlib.import_module("searcher") + return schemas, searcher + + +# --------------------------------------------------------------------------- +# Schema validation +# --------------------------------------------------------------------------- +def test_vector_mode_default_requires_embedding(search_mods): + schemas, _ = search_mods + with pytest.raises(Exception): # ValidationError + schemas.IndexSpec( + id="i1", + store=schemas.CosmosStore(endpoint="https://x", database="d", container="c"), + ) + + +def test_fts_mode_does_not_require_embedding(search_mods): + schemas, _ = search_mods + idx = schemas.IndexSpec( + id="i1", + mode="fts", + store=schemas.CosmosStore(endpoint="https://x", database="d", container="c"), + ) + assert idx.embedding is None + assert idx.fts_field is None # defaults to first content_field at search time + + +def test_fts_mode_rejects_pgvector_store(search_mods): + schemas, _ = search_mods + with pytest.raises(Exception): + schemas.IndexSpec( + id="i1", + mode="fts", + store=schemas.PgVectorStore(dsn="postgres://x", table="t"), + ) + + +def test_fts_mode_rejects_invalid_field_path(search_mods): + schemas, _ = search_mods + with pytest.raises(Exception): + schemas.IndexSpec( + id="i1", + mode="fts", + fts_field="c.content; DROP TABLE", + store=schemas.CosmosStore(endpoint="https://x", database="d", container="c"), + ) + + +def test_fts_mode_accepts_nested_path(search_mods): + schemas, _ = search_mods + idx = schemas.IndexSpec( + id="i1", + mode="fts", + fts_field="body.text", + store=schemas.CosmosStore(endpoint="https://x", database="d", container="c"), + ) + assert idx.fts_field == "body.text" + + +def test_fts_mode_requires_field_when_content_fields_empty(search_mods): + schemas, _ = search_mods + with pytest.raises(Exception): + schemas.IndexSpec( + id="i1", + mode="fts", + content_fields=[], + store=schemas.CosmosStore(endpoint="https://x", database="d", container="c"), + ) + + +# --------------------------------------------------------------------------- +# Tokenization +# --------------------------------------------------------------------------- +def test_tokenize_dedupes_and_lowercases(search_mods): + _, searcher = search_mods + assert searcher._tokenize_fts_query("Hello hello WORLD world") == ["hello", "world"] + + +def test_tokenize_handles_punctuation(search_mods): + _, searcher = search_mods + toks = searcher._tokenize_fts_query("don't break, it's fine.") + assert "don't" in toks and "it's" in toks and "fine" in toks + + +def test_tokenize_caps_terms(search_mods, monkeypatch): + _, searcher = search_mods + monkeypatch.setattr(searcher, "_FTS_MAX_TERMS", 3) + q = "alpha beta gamma delta epsilon" + assert searcher._tokenize_fts_query(q) == ["alpha", "beta", "gamma"] + + +def test_tokenize_empty(search_mods): + _, searcher = search_mods + assert searcher._tokenize_fts_query("") == [] + assert searcher._tokenize_fts_query(" ,, ") == [] + + +# --------------------------------------------------------------------------- +# search_cosmos_fts — verify SQL shape via a fake container. +# --------------------------------------------------------------------------- +class _FakeContainer: + def __init__(self, docs): + self._docs = docs + self.last_query = None + self.last_params = None + + def query_items(self, query, parameters, enable_cross_partition_query): + self.last_query = query + self.last_params = parameters + return [{"c": d} for d in self._docs] + + +class _FakeDatabase: + def __init__(self, container): + self._container = container + + def get_container_client(self, _): + return self._container + + +class _FakeClient: + def __init__(self, container): + self._db = _FakeDatabase(container) + + def get_database_client(self, _): + return self._db + + +def test_search_cosmos_fts_builds_expected_sql(search_mods, monkeypatch): + schemas, searcher = search_mods + container = _FakeContainer([ + {"id": "a", "content": "hello world", "source": "s1"}, + {"id": "b", "content": "world peace", "source": "s2"}, + ]) + monkeypatch.setattr(searcher, "CosmosClient", None, raising=False) + + # Patch the local import inside search_cosmos_fts. + import azure.cosmos + monkeypatch.setattr( + azure.cosmos, "CosmosClient", lambda *a, **kw: _FakeClient(container) + ) + import azure.identity + monkeypatch.setattr( + azure.identity, "DefaultAzureCredential", lambda *a, **kw: object() + ) + + store = schemas.CosmosStore(endpoint="https://x", database="d", container="c") + hits = asyncio.run(searcher.search_cosmos_fts( + store=store, + fts_field="content", + query_terms=["hello", "world"], + top_k=5, + content_fields=["content"], + return_fields=[], + index_filter=None, + include_vector=False, + )) + + assert len(hits) == 2 + assert hits[0]["score"] is None # FullTextScore not projectable + assert hits[0]["text"] == "hello world" + sql = container.last_query + assert "ORDER BY RANK FullTextScore(c.content, @term0, @term1)" in sql + assert "SELECT TOP @top_k c FROM c " in sql + params = {p["name"]: p["value"] for p in container.last_params} + assert params == {"@top_k": 5, "@term0": "hello", "@term1": "world"} + + +def test_search_cosmos_fts_rejects_invalid_field(search_mods): + schemas, searcher = search_mods + store = schemas.CosmosStore(endpoint="https://x", database="d", container="c") + with pytest.raises(ValueError): + asyncio.run(searcher.search_cosmos_fts( + store=store, fts_field="c.x; DROP", query_terms=["a"], top_k=5, + content_fields=["content"], return_fields=[], + index_filter=None, include_vector=False, + )) + + +def test_search_cosmos_fts_empty_terms_returns_empty(search_mods): + schemas, searcher = search_mods + store = schemas.CosmosStore(endpoint="https://x", database="d", container="c") + hits = asyncio.run(searcher.search_cosmos_fts( + store=store, fts_field="content", query_terms=[], top_k=5, + content_fields=["content"], return_fields=[], + index_filter=None, include_vector=False, + )) + assert hits == [] + + +# --------------------------------------------------------------------------- +# _search_one_index branching +# --------------------------------------------------------------------------- +def test_search_one_index_fts_skips_embed(search_mods, monkeypatch): + schemas, searcher = search_mods + + called_embed = {"n": 0} + + async def _fake_embed(*a, **kw): + called_embed["n"] += 1 + raise AssertionError("embed_query must not be called for fts mode") + + monkeypatch.setattr(searcher, "embed_query", _fake_embed) + + async def _fake_fts(*a, **kw): + return [{"id": "x", "score": None, "text": "hi", "metadata": {}, + "source": None, "source_ref": None, "text_parts": None}] + + monkeypatch.setattr(searcher, "search_cosmos_fts", _fake_fts) + + idx = schemas.IndexSpec( + id="i1", mode="fts", + store=schemas.CosmosStore(endpoint="https://x", database="d", container="c"), + ) + info, hits = asyncio.run(searcher._search_one_index( + http=None, idx=idx, default_per_index_top_k=10, + query="hello world", request_id="rid", include_vector=False, + )) + assert called_embed["n"] == 0 + assert info.error is None + assert info.embedding_model == "fts" + assert info.result_count == 1 + assert hits[0]["id"] == "x" + + +def test_search_one_index_fts_errors_on_empty_query(search_mods): + schemas, searcher = search_mods + idx = schemas.IndexSpec( + id="i1", mode="fts", + store=schemas.CosmosStore(endpoint="https://x", database="d", container="c"), + ) + info, hits = asyncio.run(searcher._search_one_index( + http=None, idx=idx, default_per_index_top_k=10, + query="", request_id="rid", include_vector=False, + )) + assert hits == [] + assert "requires text query" in (info.error or "") + + +def test_search_one_index_fts_errors_on_pure_punctuation(search_mods): + schemas, searcher = search_mods + idx = schemas.IndexSpec( + id="i1", mode="fts", + store=schemas.CosmosStore(endpoint="https://x", database="d", container="c"), + ) + info, hits = asyncio.run(searcher._search_one_index( + http=None, idx=idx, default_per_index_top_k=10, + query=",,, !!", request_id="rid", include_vector=False, + )) + assert hits == [] + assert "no searchable tokens" in (info.error or "") + + +# --------------------------------------------------------------------------- +# explain_search handles FTS without embedding +# --------------------------------------------------------------------------- +def test_explain_search_with_fts_index(search_mods): + schemas, searcher = search_mods + req = schemas.SearchRequest( + query="hello world", + indexes=[schemas.IndexSpec( + id="fts1", mode="fts", + store=schemas.CosmosStore(endpoint="https://x", database="d", container="c"), + )], + ) + plan = asyncio.run(searcher.explain_search(req)) + assert len(plan["indexes"]) == 1 + p = plan["indexes"][0] + assert p["mode"] == "fts" + assert p["embedding"] is None + assert p["fts_field"] == "content" + assert p["fts_terms"] == ["hello", "world"]