""" Advanced Vector Database for Document Storage and Retrieval Handles document embeddings, similarity search, and metadata management """ import os import json import logging import hashlib import uuid from datetime import datetime from typing import List, Dict, Any, Optional, Tuple from dataclasses import dataclass, asdict from pathlib import Path import numpy as np # Vector database and embedding libraries import chromadb from chromadb.config import Settings from sentence_transformers import SentenceTransformer # Optional LangChain imports with error handling try: from langchain_huggingface import HuggingFaceEmbeddings from langchain_chroma import Chroma from langchain.retrievers import ContextualCompressionRetriever from langchain.retrievers.document_compressors import LLMChainExtractor LANGCHAIN_AVAILABLE = True except ImportError as e: print(f"⚠️ LangChain components not available: {e}") print(" Basic functionality will work without LangChain features") LANGCHAIN_AVAILABLE = False # Disable ChromaDB telemetry os.environ["ANONYMIZED_TELEMETRY"] = "False" # Configure logging logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) @dataclass class SearchResult: """Represents a search result with metadata""" chunk_id: str content: str source_file: str similarity_score: float metadata: Dict[str, Any] section_type: Optional[str] = None table_data: Optional[Dict[str, Any]] = None class VectorDatabase: """Advanced vector database with GPU optimization and hybrid search""" def __init__(self, embedding_model: str = "all-MiniLM-L6-v2", collection_name: str = "documents", persist_directory: str = "./vector_db", use_gpu: bool = True): self.embedding_model = embedding_model self.collection_name = collection_name self.persist_directory = persist_directory self.use_gpu = use_gpu # Initialize embedding model self._initialize_embeddings() # Initialize ChromaDB self._initialize_chromadb() # Initialize LangChain components self._initialize_langchain() logger.info(f"Vector database initialized with model: {embedding_model}") def _initialize_embeddings(self): """Initialize the embedding model""" try: # Use GPU if available and requested device = "cuda" if self.use_gpu and self._check_gpu_availability() else "cpu" self.embedder = SentenceTransformer(self.embedding_model, device=device) # Initialize LangChain embeddings if available if LANGCHAIN_AVAILABLE: self.langchain_embeddings = HuggingFaceEmbeddings( model_name=self.embedding_model, model_kwargs={'device': device} ) else: self.langchain_embeddings = None logger.info(f"Embedding model loaded on device: {device}") except Exception as e: logger.error(f"Error initializing embedding model: {e}") raise def _check_gpu_availability(self) -> bool: """Check if GPU is available""" try: import torch return torch.cuda.is_available() except ImportError: return False def _initialize_chromadb(self): """Initialize ChromaDB client and collection""" try: # Create persist directory os.makedirs(self.persist_directory, exist_ok=True) # Initialize ChromaDB client self.chroma_client = chromadb.PersistentClient( path=self.persist_directory, settings=Settings( anonymized_telemetry=False, allow_reset=True ) ) # Get or create collection self.collection = self.chroma_client.get_or_create_collection( name=self.collection_name, metadata={"hnsw:space": "cosine"} ) logger.info(f"ChromaDB collection '{self.collection_name}' initialized") except Exception as e: logger.error(f"Error initializing ChromaDB: {e}") raise def _initialize_langchain(self): """Initialize LangChain components for advanced retrieval""" if not LANGCHAIN_AVAILABLE: logger.warning("LangChain components not available - advanced features disabled") self.langchain_chroma = None self.contextual_retriever = None return try: # Initialize LangChain Chroma self.langchain_chroma = Chroma( client=self.chroma_client, collection_name=self.collection_name, embedding_function=self.langchain_embeddings ) # Initialize contextual compression retriever self.contextual_retriever = ContextualCompressionRetriever( base_retriever=self.langchain_chroma.as_retriever( search_type="similarity", search_kwargs={"k": 10} ), base_compressor=LLMChainExtractor.from_llm( llm=None, # Will be set later prompt_template="Extract the most relevant information from the following text: {text}" ) ) logger.info("LangChain components initialized") except Exception as e: logger.error(f"Error initializing LangChain components: {e}") # Continue without LangChain components if they fail self.langchain_chroma = None self.contextual_retriever = None def add_documents(self, chunks: List[Any]) -> bool: """Add document chunks to the vector database""" try: if not chunks: logger.warning("No chunks to add") return False # Prepare data for ChromaDB ids = [] texts = [] metadatas = [] embeddings = [] for chunk in chunks: # Generate unique ID chunk_id = chunk.chunk_id if hasattr(chunk, 'chunk_id') else str(uuid.uuid4()) # Get content content = chunk.content if hasattr(chunk, 'content') else str(chunk) # Create metadata (filter out None values) metadata = { 'source_file': getattr(chunk, 'source_file', 'unknown'), 'file_type': getattr(chunk, 'file_type', 'unknown'), 'section_type': getattr(chunk, 'section_type', 'text'), 'chunk_index': getattr(chunk, 'chunk_index', 0), 'confidence_score': getattr(chunk, 'confidence_score', 1.0), 'timestamp': datetime.now().isoformat() } # Add page_number only if it's not None page_number = getattr(chunk, 'page_number', None) if page_number is not None: metadata['page_number'] = page_number # Add table data if present if hasattr(chunk, 'table_data') and chunk.table_data: metadata['table_data'] = json.dumps(chunk.table_data) # Generate embedding embedding = self.embedder.encode(content, convert_to_tensor=False) ids.append(chunk_id) texts.append(content) metadatas.append(metadata) embeddings.append(embedding.tolist()) # Add to ChromaDB self.collection.add( ids=ids, documents=texts, metadatas=metadatas, embeddings=embeddings ) logger.info(f"Successfully added {len(chunks)} chunks to vector database") return True except Exception as e: logger.error(f"Error adding documents to vector database: {e}") return False def add_document(self, content: str, metadata: Dict[str, Any]) -> bool: """Add a single document to the vector database""" try: # Generate unique ID chunk_id = str(uuid.uuid4()) # Create metadata with defaults doc_metadata = { 'source_file': metadata.get('source_file', 'unknown'), 'file_type': metadata.get('file_type', 'unknown'), 'section_type': metadata.get('section_type', 'text'), 'chunk_index': metadata.get('chunk_index', 0), 'confidence_score': metadata.get('confidence_score', 1.0), 'timestamp': datetime.now().isoformat() } # Add additional metadata for key, value in metadata.items(): if key not in doc_metadata and value is not None: doc_metadata[key] = value # Generate embedding embedding = self.embedder.encode(content, convert_to_tensor=False) # Add to ChromaDB self.collection.add( ids=[chunk_id], documents=[content], metadatas=[doc_metadata], embeddings=[embedding.tolist()] ) logger.info(f"Successfully added document to vector database") return True except Exception as e: logger.error(f"Error adding document to vector database: {e}") return False def search_documents(self, query: str, n_results: int = 5, similarity_threshold: float = 0.7) -> List[Dict[str, Any]]: """Search for documents and return as dictionary format for compatibility""" try: search_results = self.search_similar(query, n_results, similarity_threshold) # Convert to dictionary format results = [] for result in search_results: results.append({ 'content': result.content, 'source_file': result.source_file, 'similarity_score': result.similarity_score, 'metadata': result.metadata, 'section_type': result.section_type, 'table_data': result.table_data }) return results except Exception as e: logger.error(f"Error searching documents: {e}") return [] def search_similar(self, query: str, n_results: int = 5, similarity_threshold: float = 0.7, filter_metadata: Optional[Dict[str, Any]] = None) -> List[SearchResult]: """Search for similar documents using semantic similarity""" try: # Generate query embedding query_embedding = self.embedder.encode(query, convert_to_tensor=False) # Prepare where clause for filtering where_clause = None if filter_metadata: where_clause = filter_metadata # Search in ChromaDB results = self.collection.query( query_embeddings=[query_embedding.tolist()], n_results=n_results, where=where_clause, include=['documents', 'metadatas', 'distances'] ) # Process results search_results = [] for i in range(len(results['ids'][0])): chunk_id = results['ids'][0][i] content = results['documents'][0][i] metadata = results['metadatas'][0][i] distance = results['distances'][0][i] # Convert distance to similarity score similarity_score = 1 - distance # Filter by similarity threshold if similarity_score >= similarity_threshold: # Parse table data if present table_data = None if 'table_data' in metadata and metadata['table_data']: try: table_data = json.loads(metadata['table_data']) except: pass result = SearchResult( chunk_id=chunk_id, content=content, source_file=metadata.get('source_file', 'unknown'), similarity_score=similarity_score, metadata=metadata, section_type=metadata.get('section_type', 'text'), table_data=table_data ) search_results.append(result) # Sort by similarity score search_results.sort(key=lambda x: x.similarity_score, reverse=True) logger.info(f"Found {len(search_results)} similar documents for query") return search_results except Exception as e: logger.error(f"Error searching vector database: {e}") return [] def hybrid_search(self, query: str, n_results: int = 5, semantic_weight: float = 0.7, keyword_weight: float = 0.3) -> List[SearchResult]: """Perform hybrid search combining semantic and keyword matching""" try: # Semantic search semantic_results = self.search_similar(query, n_results=n_results*2) # Keyword search (simple implementation) keyword_results = self._keyword_search(query, n_results=n_results*2) # Combine and rank results combined_results = self._combine_search_results( semantic_results, keyword_results, semantic_weight, keyword_weight ) # Return top results return combined_results[:n_results] except Exception as e: logger.error(f"Error in hybrid search: {e}") return self.search_similar(query, n_results) def _keyword_search(self, query: str, n_results: int = 5) -> List[SearchResult]: """Simple keyword-based search""" try: # Get all documents all_results = self.collection.get() keyword_results = [] query_terms = query.lower().split() for i, content in enumerate(all_results['documents']): content_lower = content.lower() # Calculate keyword match score matches = sum(1 for term in query_terms if term in content_lower) if matches > 0: score = matches / len(query_terms) result = SearchResult( chunk_id=all_results['ids'][i], content=content, source_file=all_results['metadatas'][i].get('source_file', 'unknown'), similarity_score=score, metadata=all_results['metadatas'][i], section_type=all_results['metadatas'][i].get('section_type', 'text') ) keyword_results.append(result) # Sort by score keyword_results.sort(key=lambda x: x.similarity_score, reverse=True) return keyword_results[:n_results] except Exception as e: logger.error(f"Error in keyword search: {e}") return [] def _combine_search_results(self, semantic_results: List[SearchResult], keyword_results: List[SearchResult], semantic_weight: float, keyword_weight: float) -> List[SearchResult]: """Combine semantic and keyword search results""" try: # Create a dictionary to store combined scores combined_scores = {} # Add semantic results for result in semantic_results: combined_scores[result.chunk_id] = { 'result': result, 'semantic_score': result.similarity_score, 'keyword_score': 0.0 } # Add keyword results for result in keyword_results: if result.chunk_id in combined_scores: combined_scores[result.chunk_id]['keyword_score'] = result.similarity_score else: combined_scores[result.chunk_id] = { 'result': result, 'semantic_score': 0.0, 'keyword_score': result.similarity_score } # Calculate combined scores combined_results = [] for chunk_id, scores in combined_scores.items(): combined_score = (scores['semantic_score'] * semantic_weight + scores['keyword_score'] * keyword_weight) # Update the result with combined score result = scores['result'] result.similarity_score = combined_score combined_results.append(result) # Sort by combined score combined_results.sort(key=lambda x: x.similarity_score, reverse=True) return combined_results except Exception as e: logger.error(f"Error combining search results: {e}") return semantic_results def get_document_statistics(self) -> Dict[str, Any]: """Get statistics about the vector database""" try: # Get collection info count = self.collection.count() # Get unique sources all_results = self.collection.get() sources = set() file_types = set() section_types = set() for metadata in all_results['metadatas']: sources.add(metadata.get('source_file', 'unknown')) file_types.add(metadata.get('file_type', 'unknown')) section_types.add(metadata.get('section_type', 'text')) stats = { 'total_chunks': count, 'unique_sources': len(sources), 'file_types': list(file_types), 'section_types': list(section_types), 'sources': list(sources), 'embedding_model': self.embedding_model, 'collection_name': self.collection_name } return stats except Exception as e: logger.error(f"Error getting document statistics: {e}") return {} def delete_documents(self, source_file: str) -> bool: """Delete all documents from a specific source file""" try: # Get documents from the source results = self.collection.get( where={"source_file": source_file} ) if results['ids']: # Delete the documents self.collection.delete(ids=results['ids']) logger.info(f"Deleted {len(results['ids'])} chunks from {source_file}") return True else: logger.warning(f"No documents found for source: {source_file}") return False except Exception as e: logger.error(f"Error deleting documents: {e}") return False def clear_database(self) -> bool: """Clear all documents from the database""" try: self.collection.delete(where={}) logger.info("Cleared all documents from vector database") return True except Exception as e: logger.error(f"Error clearing database: {e}") return False def export_database(self, export_path: str) -> bool: """Export database statistics and metadata""" try: stats = self.get_document_statistics() with open(export_path, 'w') as f: json.dump(stats, f, indent=2) logger.info(f"Database exported to: {export_path}") return True except Exception as e: logger.error(f"Error exporting database: {e}") return False # Example usage if __name__ == "__main__": # Initialize vector database vector_db = VectorDatabase(use_gpu=True) # Test search results = vector_db.search_similar("insurance policy coverage", n_results=3) print(f"Found {len(results)} results:") for i, result in enumerate(results): print(f"\nResult {i+1}:") print(f"Score: {result.similarity_score:.3f}") print(f"Source: {result.source_file}") print(f"Content: {result.content[:100]}...") # Get statistics stats = vector_db.get_document_statistics() print(f"\nDatabase statistics: {stats}")