feat: citations, eval sidecar/pvalue/trace, nomic prefixes
mcp 2.x already on main. Does not replace the variety chunker (one chunk per variety is the anti-hallucination contract). - Numbered [1] citations on search_docs / search_trials - Eval JSONL sidecar, k-curve, eval.pvalue, eval.trace - Nomic prefixes at embed time only; stored text unprefixed Closes #25
This commit is contained in:
+23
-2
@@ -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):
|
||||
|
||||
+5
-2
@@ -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)
|
||||
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user