Spaces:
Sleeping
Sleeping
File size: 1,654 Bytes
847db01 | 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 | from functools import lru_cache
import chromadb
from chromadb.api.models.Collection import Collection
from src.config import settings
@lru_cache(maxsize=1)
def get_chroma_client() -> chromadb.ClientAPI:
return chromadb.PersistentClient(path=str(settings.chroma_db_path))
def get_collection() -> Collection:
client = get_chroma_client()
return client.get_or_create_collection(
name=settings.chroma_collection_name,
metadata={"hnsw:space": "cosine"},
)
def add_chunks(
ids: list[str],
documents: list[str],
embeddings: list[list[float]],
metadatas: list[dict],
) -> None:
collection = get_collection()
# ChromaDB has a batch limit; process in batches of 5000
batch_size = 5000
for i in range(0, len(ids), batch_size):
collection.add(
ids=ids[i : i + batch_size],
documents=documents[i : i + batch_size],
embeddings=embeddings[i : i + batch_size],
metadatas=metadatas[i : i + batch_size],
)
def query_chunks(
query_embedding: list[float],
n_results: int = 5,
where: dict | None = None,
) -> dict:
collection = get_collection()
kwargs = {
"query_embeddings": [query_embedding],
"n_results": n_results,
"include": ["documents", "metadatas", "distances"],
}
if where:
kwargs["where"] = where
return collection.query(**kwargs)
def delete_by_book_id(book_id: str) -> None:
collection = get_collection()
collection.delete(where={"book_id": book_id})
def get_collection_count() -> int:
collection = get_collection()
return collection.count()
|