File size: 4,270 Bytes
d44b33d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 | """ChromaDB persistence and LangChain ``Chroma`` vector store helpers.
Collections are named per ingest target; documents are stored with UUID chunk ids.
Telemetry is disabled at the client level for quieter logs in production.
"""
from datetime import datetime, timezone
from pathlib import Path
from uuid import uuid4
import chromadb
from chromadb.config import Settings
from langchain_chroma import Chroma
from langchain_core.documents import Document
from langchain_core.embeddings import Embeddings
_CHROMA_CLIENT_SETTINGS = Settings(anonymized_telemetry=False)
def _utc_now_iso() -> str:
return datetime.now(timezone.utc).replace(microsecond=0).isoformat().replace("+00:00", "Z")
def _chroma_client(persist_directory: str) -> chromadb.PersistentClient:
Path(persist_directory).mkdir(parents=True, exist_ok=True)
return chromadb.PersistentClient(path=persist_directory, settings=_CHROMA_CLIENT_SETTINGS)
def get_vector_store(
persist_directory: str,
collection_name: str,
embedding_function: Embeddings,
) -> Chroma:
"""Open or create a persisted Chroma collection wired to the given embedder."""
client = _chroma_client(persist_directory)
try:
client.get_collection(name=collection_name)
except Exception:
client.get_or_create_collection(
name=collection_name,
metadata={"created_at": _utc_now_iso()},
)
return Chroma(
collection_name=collection_name,
embedding_function=embedding_function,
persist_directory=persist_directory,
client_settings=_CHROMA_CLIENT_SETTINGS,
)
def add_documents(vector_store: Chroma, chunks: list[Document]) -> list[str]:
"""Embed and insert chunks; return the generated vector ids."""
document_ids = [str(uuid4()) for _ in chunks]
vector_store.add_documents(documents=chunks, ids=document_ids)
return document_ids
def list_collection_names(persist_directory: str) -> list[str]:
"""Sorted list of collection names in the persist directory."""
client = _chroma_client(persist_directory)
return sorted(c.name for c in client.list_collections())
def delete_collection(persist_directory: str, collection_name: str) -> int:
"""Delete a collection and return the number of documents that were removed (best effort)."""
client = _chroma_client(persist_directory)
removed = 0
try:
col = client.get_collection(name=collection_name)
removed = int(col.count())
except Exception:
removed = 0
client.delete_collection(name=collection_name)
return removed
def collection_document_count(persist_directory: str, collection_name: str) -> int:
"""Number of vectors in a collection, or 0 if the collection does not exist."""
client = _chroma_client(persist_directory)
try:
col = client.get_collection(name=collection_name)
return int(col.count())
except Exception:
return 0
def collection_created_at(persist_directory: str, collection_name: str) -> str | None:
"""Return collection metadata ``created_at`` if present (Chroma-specific)."""
client = _chroma_client(persist_directory)
try:
col = client.get_collection(name=collection_name)
meta = getattr(col, "metadata", None) or {}
if isinstance(meta, dict):
raw = meta.get("created_at") or meta.get("created")
if raw is not None:
return str(raw)
except Exception:
pass
return None
def ensure_collection_created_at(
persist_directory: str,
collection_name: str,
*,
fallback: str | None = None,
) -> str | None:
"""Persist ``created_at`` on the Chroma collection when missing; never overwrites an existing value."""
client = _chroma_client(persist_directory)
try:
col = client.get_collection(name=collection_name)
except Exception:
return None
meta = getattr(col, "metadata", None) or {}
if not isinstance(meta, dict):
meta = {}
raw = meta.get("created_at") or meta.get("created")
if raw is not None:
return str(raw)
value = fallback or _utc_now_iso()
updated = dict(meta)
updated["created_at"] = value
col.modify(metadata=updated)
return value
|