MCP RAG

"""Vector database management for MCP-RAG.""" import logging from abc import ABC, abstractmethod from typing import List, Dict, Any, Optional from dataclasses import dataclass import chromadb from chromadb.config import Settings as ChromaSettings from .config import settings logger = logging.getLogger(__name__) @dataclass class Document: """Document data structure.""" id: str content: str metadata: Dict[str, Any] embedding: Optional[List[float]] = None @dataclass class SearchResult: """Search result data structure.""" document: Document score: float class VectorDatabase(ABC): """Abstract base class for vector databases.""" @abstractmethod async def initialize(self) -> None: """Initialize the vector database.""" pass @abstractmethod async def add_document(self, content: str, collection_name: str = "default", metadata: Dict[str, Any] = None) -> None: """Add a single document to the collection.""" pass @abstractmethod async def search( self, query_embedding: List[float], collection_name: str = "default", limit: int = 5, threshold: float = 0.7 ) -> List[SearchResult]: """Search for similar documents.""" pass @abstractmethod async def delete_collection(self, collection_name: str = "default") -> None: """Delete a collection.""" pass @abstractmethod async def list_collections(self) -> List[str]: """List all collections.""" pass @abstractmethod async def list_documents(self, collection_name: str = "default", limit: int = 100, offset: int = 0, filename: Optional[str] = None) -> Dict[str, Any]: """List documents in a collection.""" pass @abstractmethod async def delete_document(self, document_id: str, collection_name: str = "default") -> bool: """Delete a document from a collection.""" pass @abstractmethod async def list_files(self, collection_name: str = "default") -> List[Dict[str, Any]]: """List all files in a collection.""" pass @abstractmethod async def delete_file(self, filename: str, collection_name: str = "default") -> bool: """Delete all documents associated with a file from a Chroma collection.""" if not self.client: raise RuntimeError("Database not initialized") try: collection = self.client.get_collection(name=collection_name) # Get all documents result = collection.get(include=["metadatas"]) # Find IDs to delete ids_to_delete = [] if result["ids"] and result["metadatas"]: for idx, metadata in enumerate(result["metadatas"]): doc_id = result["ids"][idx] # Check if metadata has filename if metadata.get("filename") == filename: ids_to_delete.append(doc_id) # For legacy chunks without filename metadata, check ID elif "_chunk_" in doc_id: base_id = doc_id.rsplit("_chunk_", 1)[0] if base_id == filename: ids_to_delete.append(doc_id) elif doc_id == filename: ids_to_delete.append(doc_id) if ids_to_delete: collection.delete(ids=ids_to_delete) logger.info(f"Deleted {len(ids_to_delete)} chunks for file '{filename}' from '{collection_name}'") else: logger.warning(f"No chunks found for file '{filename}' in '{collection_name}'") return True except Exception as e: logger.error(f"Failed to delete file '{filename}' from '{collection_name}': {e}") return False class ChromaDatabase(VectorDatabase): """Chroma vector database implementation.""" def __init__(self, embedding_function=None): self.client: Optional[chromadb.Client] = None self.collections: Dict[str, chromadb.Collection] = {} self.embedding_function = embedding_function async def initialize(self) -> None: """Initialize Chroma client.""" try: chroma_settings = ChromaSettings( persist_directory=settings.chroma_persist_directory, is_persistent=True ) self.client = chromadb.PersistentClient(path=settings.chroma_persist_directory) logger.info(f"Chroma database initialized at {settings.chroma_persist_directory}") # Ensure default collection exists await self._ensure_default_collection() except Exception as e: logger.error(f"Failed to initialize Chroma database: {e}") raise async def _ensure_default_collection(self) -> None: """Ensure the default collection exists with correct configuration.""" if not self.client: return try: # Try to get the default collection try: collection = self.client.get_collection(name="default") # Check if collection uses the correct distance metric current_space = collection.metadata.get("hnsw:space") if collection.metadata else None if current_space != "cosine": logger.warning(f"Default collection uses distance metric '{current_space}'. Recreating with cosine.") self.client.delete_collection(name="default") collection = self.client.create_collection( name="default", metadata={"hnsw:space": "cosine"} ) logger.info("Recreated default collection with cosine similarity") except Exception: # Collection doesn't exist, create it with cosine similarity collection = self.client.create_collection( name="default", metadata={"hnsw:space": "cosine"} ) logger.info("Created default collection with cosine similarity") except Exception as e: logger.error(f"Failed to ensure default collection: {e}") # Don't raise here as this is not critical for initialization async def add_document(self, content: str, collection_name: str = "default", metadata: Dict[str, Any] = None) -> None: """Add a single document to Chroma collection.""" if metadata is None: metadata = {} # Auto-generate filename if not provided if "filename" not in metadata: if "title" in metadata and metadata["title"]: # Use title as filename (sanitized) filename = metadata["title"].strip()[:50] # Limit length # Replace invalid filename characters for char in ['/', '\\', ':', '*', '?', '"', '<', '>', '|']: filename = filename.replace(char, '_') else: # Generate default filename with timestamp from datetime import datetime timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") filename = f"手动输入_{timestamp}" metadata["filename"] = filename document = Document( id=f"doc_{len(content)}_{hash(content)}", # Simple ID generation content=content, metadata=metadata ) await self.add_documents([document], collection_name) async def add_documents(self, documents: List[Document], collection_name: str = "default") -> None: """Add multiple documents to Chroma collection.""" if not self.client: raise RuntimeError("Database not initialized") try: # Get or create collection try: collection = self.client.get_collection(name=collection_name) # Check if collection uses the correct distance metric current_space = collection.metadata.get("hnsw:space") if collection.metadata else None if current_space != "cosine": logger.warning(f"Collection '{collection_name}' uses distance metric '{current_space}'. Recreating with cosine.") self.client.delete_collection(name=collection_name) collection = self.client.create_collection( name=collection_name, metadata={"hnsw:space": "cosine"} ) logger.info(f"Recreated collection '{collection_name}' with cosine similarity") except Exception: # Collection doesn't exist, create it with cosine similarity collection = self.client.create_collection( name=collection_name, metadata={"hnsw:space": "cosine"} ) logger.info(f"Created collection '{collection_name}' with cosine similarity") # Prepare data for Chroma ids = [] contents = [] metadatas = [] from .text_splitter import split_text for doc in documents: # Split text into chunks chunks = split_text(doc.content) for i, chunk in enumerate(chunks): # Create unique ID for chunk chunk_id = f"{doc.id}_chunk_{i}" # Create metadata for chunk chunk_metadata = doc.metadata.copy() chunk_metadata.update({ "chunk_index": i, "total_chunks": len(chunks), "original_id": doc.id }) ids.append(chunk_id) contents.append(chunk) metadatas.append(chunk_metadata) if not ids: logger.warning(f"No content to add to collection '{collection_name}'") return # Calculate embeddings for documents from .embedding import get_embedding_model embedding_model = await get_embedding_model() # Process in batches to avoid API limits batch_size = 10 for i in range(0, len(contents), batch_size): batch_ids = ids[i:i + batch_size] batch_contents = contents[i:i + batch_size] batch_metadatas = metadatas[i:i + batch_size] batch_embeddings = await embedding_model.encode(batch_contents) # Add documents to collection with embeddings collection.add( documents=batch_contents, metadatas=batch_metadatas, ids=batch_ids, embeddings=batch_embeddings ) logger.info(f"Added {len(documents)} documents ({len(contents)} chunks) to collection '{collection_name}'") except Exception as e: logger.error(f"Failed to add documents to collection '{collection_name}': {e}") raise async def search( self, query_embedding: List[float], collection_name: str = "default", limit: int = 5, threshold: float = 0.7 ) -> List[SearchResult]: """Search Chroma collection using built-in vector search.""" if not self.client: raise RuntimeError("Database not initialized") try: collection = self.client.get_collection(name=collection_name) if not collection: logger.warning(f"Collection '{collection_name}' not found") return [] # Use ChromaDB's built-in vector search results = collection.query( query_embeddings=[query_embedding], n_results=limit, include=["documents", "metadatas", "distances"] ) # Create search results search_results = [] if results["distances"] and results["documents"] and len(results["distances"]) > 0: distances = results["distances"][0] documents = results["documents"][0] ids = results["ids"][0] if results["ids"] else [] metadatas = results["metadatas"][0] if results["metadatas"] else [] for i, distance in enumerate(distances): # For cosine distance, similarity = 1 - distance similarity = 1 - distance if similarity >= threshold: document = Document( id=ids[i] if i < len(ids) else f"result_{i}", content=documents[i], metadata=metadatas[i] if i < len(metadatas) else {} ) search_results.append(SearchResult(document=document, score=float(similarity))) logger.info(f"Found {len(search_results)} results above threshold {threshold}") return search_results except Exception as e: logger.error(f"Failed to search collection '{collection_name}': {e}") raise async def delete_collection(self, collection_name: str = "default") -> None: """Delete Chroma collection.""" if not self.client: raise RuntimeError("Database not initialized") try: self.client.delete_collection(name=collection_name) if collection_name in self.collections: del self.collections[collection_name] logger.info(f"Deleted collection '{collection_name}'") except Exception as e: logger.error(f"Failed to delete collection '{collection_name}': {e}") raise async def list_collections(self) -> List[str]: """List all Chroma collections.""" if not self.client: raise RuntimeError("Database not initialized") try: collections = self.client.list_collections() return [col.name for col in collections] except Exception as e: logger.error(f"Failed to list collections: {e}") raise async def list_documents(self, collection_name: str = "default", limit: int = 100, offset: int = 0, filename: Optional[str] = None) -> Dict[str, Any]: """List documents in a Chroma collection.""" if not self.client: raise RuntimeError("Database not initialized") try: collection = self.client.get_collection(name=collection_name) # Get all documents with metadata where_clause = {"filename": filename} if filename else None result = collection.get( where=where_clause, limit=limit, offset=offset, include=["metadatas", "documents"] ) documents = [] if result["ids"]: for i, doc_id in enumerate(result["ids"]): documents.append({ "id": doc_id, "content": result["documents"][i] if result["documents"] else "", "metadata": result["metadatas"][i] if result["metadatas"] else {} }) # Get total count (approximation as Chroma doesn't have a direct count method efficiently exposed in all versions) # For now, we'll just return what we have. In a real app, we might want a separate count query. total = collection.count() return { "total": total, "documents": documents, "limit": limit, "offset": offset } except Exception as e: logger.error(f"Failed to list documents in '{collection_name}': {e}") # Return empty result if collection doesn't exist or other error return {"total": 0, "documents": [], "limit": limit, "offset": offset} async def delete_document(self, document_id: str, collection_name: str = "default") -> bool: """Delete a document from a Chroma collection.""" if not self.client: raise RuntimeError("Database not initialized") try: collection = self.client.get_collection(name=collection_name) collection.delete(ids=[document_id]) logger.info(f"Deleted document '{document_id}' from '{collection_name}'") return True except Exception as e: logger.error(f"Failed to delete document '{document_id}' from '{collection_name}': {e}") return False async def list_files(self, collection_name: str = "default") -> List[Dict[str, Any]]: """List all files in a Chroma collection.""" if not self.client: raise RuntimeError("Database not initialized") try: collection = self.client.get_collection(name=collection_name) # Get all metadatas and IDs result = collection.get(include=["metadatas"]) files = {} if result["metadatas"] and result["ids"]: for idx, metadata in enumerate(result["metadatas"]): doc_id = result["ids"][idx] # Try to get filename from metadata first filename = metadata.get("filename") # If no filename in metadata, try to extract from ID # For legacy chunks like "doc_150000_-9093687922038414764_chunk_34" # or new chunks like "some_id_chunk_0" if not filename: # Check if this is a chunk (has _chunk_ in ID) if "_chunk_" in doc_id: # Remove the chunk suffix to get the base ID filename = doc_id.rsplit("_chunk_", 1)[0] else: # Not a chunk, use the ID as filename filename = doc_id if filename not in files: files[filename] = { "filename": filename, "chunk_count": 0, "total_size": 0, "file_type": metadata.get("file_type", "unknown"), "upload_time": metadata.get("timestamp", "") } files[filename]["chunk_count"] += 1 # Approximate size if available, otherwise just count chunks files[filename]["total_size"] += metadata.get("size", 0) return list(files.values()) except Exception as e: logger.error(f"Failed to list files in '{collection_name}': {e}") return [] async def delete_file(self, filename: str, collection_name: str = "default") -> bool: """Delete all documents associated with a file from a Chroma collection.""" if not self.client: raise RuntimeError("Database not initialized") try: collection = self.client.get_collection(name=collection_name) # Delete where metadata['filename'] == filename collection.delete(where={"filename": filename}) logger.info(f"Deleted file '{filename}' from '{collection_name}'") return True except Exception as e: logger.error(f"Failed to delete file '{filename}' from '{collection_name}': {e}") return False # Global database instance vector_db: Optional[VectorDatabase] = None async def get_vector_database() -> VectorDatabase: """Get the global vector database instance.""" global vector_db if vector_db is None: if settings.vector_db_type == "chroma": vector_db = ChromaDatabase() else: raise ValueError(f"Unsupported vector database type: {settings.vector_db_type}") await vector_db.initialize() return vector_db