Spaces:
Sleeping
Sleeping
| """ | |
| ChromaDB vector store operations. | |
| Per-user collections for data isolation. | |
| """ | |
| import logging | |
| from typing import List, Dict, Any, Optional | |
| from app.config import get_settings | |
| logger = logging.getLogger(__name__) | |
| settings = get_settings() | |
| # ── Singleton ChromaDB client ──────────────────────── | |
| _chroma_client = None | |
| def get_chroma_client(): | |
| """Get or create persistent ChromaDB client.""" | |
| global _chroma_client | |
| if _chroma_client is None: | |
| import os | |
| import chromadb | |
| from chromadb.config import Settings as ChromaSettings | |
| os.makedirs(settings.CHROMA_PERSIST_DIR, exist_ok=True) | |
| _chroma_client = chromadb.PersistentClient( | |
| path=settings.CHROMA_PERSIST_DIR, | |
| settings=ChromaSettings(anonymized_telemetry=False), | |
| ) | |
| logger.info(f"ChromaDB initialized at {settings.CHROMA_PERSIST_DIR}") | |
| return _chroma_client | |
| def get_collection_name(user_id: str) -> str: | |
| """Generate a valid collection name for a user.""" | |
| # ChromaDB collection names must be 3-63 chars, alphanumeric + underscores | |
| clean_id = user_id.replace("-", "_") | |
| name = f"user_{clean_id}" | |
| # Truncate if too long | |
| return name[:63] | |
| def store_chunks( | |
| chunks: List[Dict[str, Any]], | |
| document_id: str, | |
| filename: str, | |
| user_id: str, | |
| ) -> int: | |
| """ | |
| Embed and store document chunks in ChromaDB, and build a local BM25 index. | |
| Returns the number of chunks stored. | |
| """ | |
| if not chunks: | |
| return 0 | |
| # Delete existing chunks for this document and user before inserting new ones | |
| # to avoid stale/orphaned chunks in ChromaDB and BM25. | |
| delete_document_chunks(document_id, user_id) | |
| # Build and store BM25 index | |
| from app.rag.bm25 import store_bm25_index | |
| try: | |
| store_bm25_index(chunks, document_id, filename, user_id) | |
| except Exception as e: | |
| logger.error(f"Could not build BM25 index: {e}") | |
| # Generate captions for any extracted images before embedding | |
| try: | |
| from app.rag.vision import generate_captions_for_chunks | |
| generate_captions_for_chunks(chunks) | |
| except Exception as e: | |
| logger.warning(f"Could not generate image captions: {e}") | |
| client = get_chroma_client() | |
| from app.rag.embeddings import get_embedding_model | |
| embedding_model = get_embedding_model() | |
| collection_name = get_collection_name(user_id) | |
| collection = client.get_or_create_collection( | |
| name=collection_name, | |
| metadata={"hnsw:space": "cosine"}, | |
| ) | |
| # ── Prepare batch data ─────────────────────────── | |
| texts = [chunk["text"] for chunk in chunks] | |
| ids = [f"{document_id}_{chunk['chunk_index']}" for chunk in chunks] | |
| metadatas = [ | |
| { | |
| "text": chunk["text"], | |
| "filename": filename, | |
| "document_id": document_id, | |
| "page": chunk["page"], | |
| "chunk_index": chunk["chunk_index"], | |
| "chunk_type": chunk.get("chunk_type", "text"), | |
| **({"bbox": chunk.get("bbox", "")} if chunk.get("bbox") else {}), | |
| **({"table_index": chunk.get("table_index", 0)} if chunk.get("chunk_type") == "table" else {}), | |
| # Indicate whether this chunk was originally an image and include a short caption | |
| **({"is_image": True, "image_caption": chunk.get("image_caption", "")} | |
| if chunk.get("is_image") else {}), | |
| } | |
| for chunk in chunks | |
| ] | |
| # ── Embed and upsert in batches ────────────────── | |
| batch_size = 50 | |
| total_stored = 0 | |
| for i in range(0, len(texts), batch_size): | |
| batch_texts = texts[i:i + batch_size] | |
| batch_ids = ids[i:i + batch_size] | |
| batch_metadatas = metadatas[i:i + batch_size] | |
| # Generate embeddings | |
| embeddings = embedding_model.embed_documents(batch_texts) | |
| collection.add( | |
| ids=batch_ids, | |
| embeddings=embeddings, | |
| metadatas=batch_metadatas, | |
| documents=batch_texts, | |
| ) | |
| total_stored += len(batch_texts) | |
| logger.info(f"Stored {total_stored} chunks for document {document_id}") | |
| return total_stored | |
| def query_chunks( | |
| query_embedding: List[float], | |
| user_id: str, | |
| document_id: Optional[str] = None, | |
| document_ids: Optional[List[str]] = None, | |
| top_k: int = 10, | |
| ) -> List[Dict[str, Any]]: | |
| """ | |
| Query ChromaDB for relevant chunks. | |
| Returns list of dicts with text, metadata, and distance. | |
| """ | |
| client = get_chroma_client() | |
| collection_name = get_collection_name(user_id) | |
| try: | |
| collection = client.get_collection(name=collection_name) | |
| except Exception: | |
| logger.warning(f"Collection {collection_name} not found") | |
| return [] | |
| # ── Build filter ───────────────────────────────── | |
| where_filter = None | |
| if document_id: | |
| where_filter = {"document_id": {"$eq": document_id}} | |
| elif document_ids: | |
| where_filter = {"document_id": {"$in": document_ids}} | |
| # ── Query ──────────────────────────────────────── | |
| results = collection.query( | |
| query_embeddings=[query_embedding], | |
| n_results=top_k, | |
| where=where_filter, | |
| include=["documents", "metadatas", "distances", "ids"], | |
| ) | |
| # ── Format results ─────────────────────────────── | |
| chunks = [] | |
| if results and results["documents"] and results["documents"][0]: | |
| for i, doc in enumerate(results["documents"][0]): | |
| metadata = results["metadatas"][0][i] if results["metadatas"] else {} | |
| distance = results["distances"][0][i] if results["distances"] else 0 | |
| chunk_id = results["ids"][0][i] if results.get("ids") and len(results["ids"]) > 0 else None | |
| # Convert cosine distance to similarity score (0-1) | |
| similarity = 1 - distance | |
| chunks.append({ | |
| "id": chunk_id, | |
| "text": doc, | |
| "filename": metadata.get("filename", ""), | |
| "document_id": metadata.get("document_id", ""), | |
| "page": metadata.get("page", 1), | |
| "chunk_type": metadata.get("chunk_type", "text"), | |
| "bbox": metadata.get("bbox", ""), | |
| "score": round(similarity, 4), | |
| }) | |
| return chunks | |
| def delete_document_chunks(document_id: str, user_id: str): | |
| """Delete all chunks for a specific document.""" | |
| client = get_chroma_client() | |
| collection_name = get_collection_name(user_id) | |
| try: | |
| from app.rag.bm25 import delete_bm25_index | |
| delete_bm25_index(document_id, user_id) | |
| except Exception as e: | |
| logger.warning(f"Error deleting BM25 index: {e}") | |
| try: | |
| collection = client.get_collection(name=collection_name) | |
| # Get all IDs for this document | |
| results = collection.get( | |
| where={"document_id": {"$eq": document_id}}, | |
| include=[], | |
| ) | |
| if results["ids"]: | |
| collection.delete(ids=results["ids"]) | |
| logger.info(f"Deleted {len(results['ids'])} chunks for document {document_id}") | |
| except Exception as e: | |
| logger.warning(f"Error deleting chunks: {e}") | |
| def delete_user_collection(user_id: str): | |
| """Delete entire collection for a user.""" | |
| client = get_chroma_client() | |
| collection_name = get_collection_name(user_id) | |
| try: | |
| from app.rag.bm25 import delete_user_bm25_indexes | |
| delete_user_bm25_indexes(user_id) | |
| except Exception as e: | |
| logger.warning(f"Error deleting user BM25 indexes: {e}") | |
| try: | |
| client.delete_collection(name=collection_name) | |
| logger.info(f"Deleted collection {collection_name}") | |
| except Exception as e: | |
| logger.warning(f"Error deleting collection: {e}") | |