Spaces:
Sleeping
Sleeping
File size: 6,750 Bytes
db4d559 3786a3f db4d559 3786a3f db4d559 3786a3f db4d559 3786a3f feca495 db4d559 3786a3f db4d559 3786a3f db4d559 3786a3f db4d559 feca495 db4d559 feca495 db4d559 3786a3f db4d559 feca495 db4d559 0527a95 db4d559 | 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 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 | import os
import uuid
import logging
from typing import List
from django.conf import settings
import chromadb
from chromadb.utils.embedding_functions import ONNXMiniLM_L6_V2
logger = logging.getLogger(__name__)
class VectorRetriever:
def __init__(self):
logger.info("Initializing VectorRetriever service.")
# 1. Initialize Persistent ChromaDB Client
self.persist_directory = getattr(settings, 'CHROMADB_DIR', os.path.join(settings.BASE_DIR, "chroma_db"))
os.makedirs(self.persist_directory, exist_ok=True)
self.chroma_client = chromadb.PersistentClient(path=self.persist_directory)
# 2. Load the Embedding Model via ChromaDB's built-in ONNX path
# Uses ONNX runtime only — no torch/CUDA needed, no sentence-transformers package
self.embedding_fn = ONNXMiniLM_L6_V2(preferred_providers=["CPUExecutionProvider"])
# 3. Setup text splitter for document chunking
def chunk_text(text: str, chunk_size: int = 800, chunk_overlap: int = 100) -> List[str]:
words = text.split()
chunks = []
start = 0
while start < len(words):
end = min(start + chunk_size, len(words))
chunks.append(' '.join(words[start:end]))
if end == len(words):
break
start = end - chunk_overlap
return chunks
self.text_splitter = chunk_text
def _get_user_collection(self, user_id):
"""
Enforce multi-tenancy by returning a collection isolated for each user.
"""
collection_name = f"user_collection_{str(user_id).replace('-', '_')}"
return self.chroma_client.get_or_create_collection(
name=collection_name,
metadata={"hnsw:space": "cosine"} # Use cosine similarity
)
def index_document(self, text_content: str, doc_name: str, user_id: str):
"""
Splits document text into chunks, generates embeddings, and saves them to ChromaDB.
"""
if not text_content or not text_content.strip():
logger.warning("Empty text content provided for vector indexing.")
return
logger.info("Starting vector indexing for document '%s' (User: %s)", doc_name, user_id)
try:
# Split text into chunks
chunks = self.text_splitter(text_content)
logger.info("Split document into %d vector chunks.", len(chunks))
collection = self._get_user_collection(user_id)
# Prepare inputs for ChromaDB
doc_id = str(uuid.uuid4())[:8]
ids = [f"{doc_id}_{doc_name}_chunk_{i}" for i in range(len(chunks))]
# Generate vector representations using ONNX-based embedding
embeddings = self.embedding_fn(chunks)
metadatas = [{"source_doc": doc_name, "page": i + 1, "chunk_index": i} for i in range(len(chunks))]
# Insert or update in ChromaDB
collection.upsert(
ids=ids,
embeddings=embeddings,
documents=chunks,
metadatas=metadatas
)
logger.info("Successfully indexed %d chunks in ChromaDB for document: %s", len(chunks), doc_name)
except Exception as e:
logger.error("Failed to index document in ChromaDB. Error: %s", str(e), exc_info=True)
raise e
def retrieve_relevant_chunks(self, query: str, user_id: str, limit: int = 5, doc_names: List[str] = None) -> List[dict]:
"""
Queries ChromaDB to retrieve the most semantically relevant text passages.
"""
logger.info("Searching ChromaDB for query: '%s' (Limit: %d, User: %s, Docs: %s)", query, limit, user_id, doc_names)
try:
collection = self._get_user_collection(user_id)
query_vector = self.embedding_fn([query])[0]
query_params = {
"query_embeddings": [query_vector],
"n_results": limit
}
if doc_names:
query_params["where"] = {"source_doc": {"$in": [str(d) for d in doc_names]}}
results = collection.query(**query_params)
retrieved = []
if results and results["documents"]:
documents = results["documents"][0]
metadatas = results["metadatas"][0]
distances = results["distances"][0] if "distances" in results else [0.0] * len(documents)
for doc, meta, dist in zip(documents, metadatas, distances):
# Cosine distance (0.0 is exact match, 1.0 is opposite)
# Convert distance to a similarity score (1.0 - distance)
similarity = round(1.0 - dist, 4)
retrieved.append({
"text": doc,
"source_doc": meta.get("source_doc", "unknown"),
"page": meta.get("page", 1),
"similarity_score": similarity
})
logger.info("Retrieved %d relevant text chunks from ChromaDB.", len(retrieved))
return retrieved
except Exception as e:
logger.error("Error retrieving from ChromaDB: %s", str(e), exc_info=True)
return []
def get_chunk_by_page(self, doc_name: str, page: int, user_id: str) -> str:
"""
Retrieves the verbatim text content of a specific page/chunk from ChromaDB.
"""
try:
collection = self._get_user_collection(user_id)
results = collection.get(
where={
"$and": [
{"source_doc": str(doc_name)},
{"page": int(page)}
]
},
limit=1
)
if results and results["documents"]:
return results["documents"][0]
except Exception as e:
logger.error("Error fetching chunk from ChromaDB for %s page %d: %s", doc_name, page, str(e))
return ""
def delete_document_vectors(self, doc_name: str, user_id: str):
"""
Removes all vectors belonging to a deleted document.
"""
logger.info("Deleting vectors for document '%s' from ChromaDB (User: %s)", doc_name, user_id)
try:
collection = self._get_user_collection(user_id)
collection.delete(where={"source_doc": doc_name})
logger.info("Successfully deleted all vectors for document '%s' from ChromaDB.", doc_name)
except Exception as e:
logger.error("Failed to delete document vectors from ChromaDB: %s", str(e), exc_info=True)
|