Port docs-mcp-template upgrades (eval, citations, prefixes) (#26)

Co-authored-by: claude <[email protected]>

rag/embeddings.py

@@ -14,8 +14,14 @@ import os
 import logging
 from typing import Any
 
-import httpx
-from chromadb import EmbeddingFunction, Documents, Embeddings
+try:
+    import httpx
+    from chromadb import Documents, EmbeddingFunction, Embeddings
+except ImportError:  # unit tests for apply_prefix / embed_texts
+    httpx = None  # type: ignore[assignment]
+    EmbeddingFunction = object  # type: ignore[misc,assignment]
+    Documents = list  # type: ignore[misc,assignment]
+    Embeddings = list  # type: ignore[misc,assignment]
 
 log = logging.getLogger(__name__)
 
@@ -23,6 +29,21 @@ OLLAMA_URLS = [u.strip() for u in os.environ.get("OLLAMA_URL",
                "http://localhost:11434").split(",") if u.strip()]
 EMBED_MODEL = os.environ.get("EMBED_MODEL", "nomic-embed-text")
 EMBED_DIM = int(os.environ.get("EMBED_DIM", "768"))
+# nomic asymmetric prefixes — embed time only. Empty string disables.
+EMBED_QUERY_PREFIX = os.environ.get("EMBED_QUERY_PREFIX", "search_query: ")
+EMBED_DOC_PREFIX = os.environ.get("EMBED_DOC_PREFIX", "search_document: ")
+
+
+def apply_prefix(text: str, prefix: str) -> str:
+    if not prefix:
+        return text
+    return prefix + text
+
+
+def embed_texts(texts: list[str], *, prefix: str, ef: Any | None = None) -> list[list[float]]:
+    """Embed texts after applying prefix. Callers store the original strings."""
+    fn = ef if ef is not None else embedding_function()
+    return list(fn([apply_prefix(t, prefix) for t in texts]))
 
 
 class OllamaEmbeddings(EmbeddingFunction):

rag/index.py

@@ -23,7 +23,7 @@ import chromadb
 from chromadb.config import Settings
 
 from .chunk import chunks_from_variety, chunks_from_trial
-from .embeddings import embedding_function
+from .embeddings import EMBED_DOC_PREFIX, embed_texts, embedding_function
 
 log = logging.getLogger(__name__)
 logging.basicConfig(level=logging.INFO, format="%(asctime)s  %(message)s")
@@ -83,9 +83,12 @@ def upsert_to_chroma(records: list[dict]) -> int:
     total = 0
     for i in range(0, len(records), BATCH):
         chunk = records[i:i + BATCH]
+        texts = [r["text"] for r in chunk]
+        vectors = embed_texts(texts, prefix=EMBED_DOC_PREFIX, ef=embedding_function())
         col.upsert(
             ids=[r["id"] for r in chunk],
-            documents=[r["text"] for r in chunk],
+            documents=texts,
+            embeddings=vectors,
             metadatas=[r["metadata"] for r in chunk],
         )
         total += len(chunk)

rag/test_embeddings.py

@@ -0,0 +1,41 @@
+"""Prefix helpers — no Ollama, no Chroma."""
+from __future__ import annotations
+
+import unittest
+
+from rag.embeddings import apply_prefix, embed_texts
+
+
+class FakeEmbedder:
+    def __init__(self) -> None:
+        self.seen: list[str] = []
+
+    def __call__(self, texts: list[str]) -> list[list[float]]:
+        self.seen.extend(texts)
+        return [[0.0, 1.0] for _ in texts]
+
+
+class PrefixTests(unittest.TestCase):
+    def test_index_sends_doc_prefix_keeps_raw_text(self) -> None:
+        fake = FakeEmbedder()
+        raw = ["hello"]
+        vecs = embed_texts(raw, prefix="search_document: ", ef=fake)
+        self.assertEqual(fake.seen, ["search_document: hello"])
+        self.assertEqual(raw, ["hello"])  # caller storage unchanged
+        self.assertEqual(len(vecs), 1)
+
+    def test_query_sends_query_prefix(self) -> None:
+        fake = FakeEmbedder()
+        embed_texts(["q"], prefix="search_query: ", ef=fake)
+        self.assertEqual(fake.seen, ["search_query: q"])
+        self.assertNotIn("q", fake.seen)
+
+    def test_empty_prefix_is_raw(self) -> None:
+        self.assertEqual(apply_prefix("hello", ""), "hello")
+        fake = FakeEmbedder()
+        embed_texts(["hello"], prefix="", ef=fake)
+        self.assertEqual(fake.seen, ["hello"])
+
+
+if __name__ == "__main__":
+    unittest.main()