File size: 4,343 Bytes
38830c1 | 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 | from __future__ import annotations
from pathlib import Path
from rag.schemas import DocumentChunk, RetrievalHit
import chromadb
class ChromaDocumentStore:
def __init__(self, persist_path: str, collection_name: str = "uploaded_documents") -> None:
self.persist_path = Path(persist_path)
self.collection_name = collection_name
self.persist_path = Path("/tmp/chroma_db") # force safe path
self.persist_path.mkdir(parents=True, exist_ok=True)
self.persist_path.mkdir(parents=True, exist_ok=True)
self._client = None
self._collection = None
self._available = None
@property
def is_available(self) -> bool:
if self._available is None:
try:
self._available = self._ensure_collection() is not None
except Exception:
self._available = False
return bool(self._available)
@property
def is_available(self) -> bool:
if self._available is None:
try:
self._available = self._ensure_collection() is not None
except Exception:
self._available = False
return bool(self._available)
def upsert_chunks(self, chunks: list[DocumentChunk], embeddings: list[list[float]]) -> None:
collection = self._ensure_collection()
if collection is None:
raise RuntimeError("ChromaDB is not installed. Add `chromadb` to your environment.")
collection.upsert(
ids=[chunk.chunk_id for chunk in chunks],
documents=[chunk.text for chunk in chunks],
metadatas=[self._sanitize_metadata(chunk.metadata | {"document_id": chunk.document_id}) for chunk in chunks],
embeddings=embeddings,
)
def similarity_search(self, query_embedding: list[float], top_k: int = 5) -> list[RetrievalHit]:
collection = self._ensure_collection()
if collection is None:
return []
result = collection.query(query_embeddings=[query_embedding], n_results=top_k)
ids = result.get("ids", [[]])[0]
docs = result.get("documents", [[]])[0]
metadatas = result.get("metadatas", [[]])[0]
distances = result.get("distances", [[]])[0]
hits: list[RetrievalHit] = []
for chunk_id, text, metadata, distance in zip(ids, docs, metadatas, distances):
hits.append(
RetrievalHit(
chunk_id=chunk_id,
text=text,
score=1.0 / (1.0 + float(distance)),
metadata=metadata or {},
)
)
return hits
def list_documents(self) -> list[dict]:
collection = self._ensure_collection()
if collection is None:
return []
payload = collection.get(include=["metadatas"])
seen: dict[str, dict] = {}
for metadata in payload.get("metadatas", []):
if not metadata:
continue
document_id = str(metadata.get("document_id", ""))
if document_id and document_id not in seen:
seen[document_id] = {
"document_id": document_id,
"file_name": metadata.get("file_name", ""),
}
return sorted(seen.values(), key=lambda item: item["file_name"])
def has_documents(self) -> bool:
return bool(self.list_documents())
def _ensure_collection(self):
if self._collection is not None:
return self._collection
try:
import chromadb
except Exception:
self._available = False
return None
##self._client = chromadb.PersistentClient(path=str(self.persist_path))
self._client = chromadb.Client() # in-memory (no SQLite issues)
self._collection = self._client.get_or_create_collection(name=self.collection_name)
self._available = True
return self._collection
@staticmethod
def _sanitize_metadata(metadata: dict) -> dict:
sanitized = {}
for key, value in metadata.items():
if isinstance(value, (str, int, float, bool)) or value is None:
sanitized[key] = value
else:
sanitized[key] = str(value)
return sanitized
|