Port docs-mcp-template upgrades (eval, citations, prefixes) (#26)
Image rebuild (skip scrape) / build (push) Successful in 3m58s
Image rebuild (skip scrape) / build (push) Successful in 3m58s
Co-authored-by: claude <[email protected]>
This commit was merged in pull request #26.
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):
|
||||
|
||||
Reference in New Issue
Block a user