# knowledge_base.py import torch from langchain_community.vectorstores import Chroma from langchain_community.embeddings import HuggingFaceEmbeddings from langchain_core.documents import Document from typing import List, Dict, Optional import os class KnowledgeBase: """ Vector database wrapper for RAG retrieval. Handles embedding generation and similarity search. """ def __init__( self, persist_directory: str = "../chroma_db", embedding_model_name: str = "abhinand/MedEmbed-large-v0.1", collection_name: str = "kidney_stone_kb_v2", device: Optional[str] = None # None → otomatik seç ): self.persist_directory = persist_directory self.collection_name = collection_name self.embedding_model_name = embedding_model_name # Device seçimi: GPU varsa kullan, yoksa CPU'ya düş if device is None: self.device = "cuda" if torch.cuda.is_available() else "cpu" else: self.device = device print(f"Loading embedding model: {embedding_model_name}") print(f"Device: {self.device}") if self.device == "cuda": print(f"GPU: {torch.cuda.get_device_name(0)}") # Embedding model'i yükle self.embeddings = HuggingFaceEmbeddings( model_name=embedding_model_name, model_kwargs={'device': self.device}, encode_kwargs={ 'normalize_embeddings': True, 'batch_size': 16 if self.device == "cuda" else 8 # GPU'da daha büyük batch } ) self.vectorstore = self._initialize_vectorstore() print(f"Knowledge base ready at: {persist_directory}") def _initialize_vectorstore(self) -> Chroma: """Vector store'u başlat - var olan varsa yükle, yoksa yenisini yarat""" return Chroma( persist_directory=self.persist_directory, embedding_function=self.embeddings, collection_name=self.collection_name ) def add_documents(self, documents: List[Document]) -> None: """Vector store'a yeni dokümanlar ekle""" if not documents: print("No documents to add") return print(f"Adding {len(documents)} chunks to knowledge base...") self.vectorstore.add_documents(documents) print("Documents added successfully") def search( self, query: str, k: int = 5, filter_dict: Optional[Dict] = None ) -> List[Dict]: """ Similarity search yap. Args: query: Natural language query k: Top-k results to return filter_dict: Metadata-based filtering (optional) Returns: List of retrieved chunks with metadata """ results = self.vectorstore.similarity_search_with_score( query=query, k=k, filter=filter_dict ) formatted_results = [] for doc, score in results: formatted_results.append({ "content": doc.page_content, "source": doc.metadata.get("source", "unknown"), "page": doc.metadata.get("page", "N/A"), "similarity_score": float(score) }) return formatted_results def search_with_chroma_ids( self, query: str, k: int = 5, filter_dict: Optional[Dict] = None, ) -> List[Dict]: """search() gibi ama her sonuçta ChromaDB UUID de döner.""" query_embedding = self.embeddings.embed_query(query) collection = self.vectorstore._collection kwargs: Dict = { "query_embeddings": [query_embedding], "n_results": k, "include": ["documents", "metadatas", "distances"], } if filter_dict: kwargs["where"] = filter_dict results = collection.query(**kwargs) formatted = [] for cid, doc, meta, dist in zip( results["ids"][0], results["documents"][0], results["metadatas"][0], results["distances"][0], ): formatted.append({ "chroma_id": cid, "content": doc, "source": meta.get("source", "unknown"), "page": meta.get("page", "N/A"), "similarity_score": float(1.0 - dist), "metadata": meta, }) return formatted def get_chunk_by_id(self, chroma_id: str) -> Optional[Dict]: """Tek bir chunk'ı ChromaDB UUID ile getir.""" collection = self.vectorstore._collection result = collection.get( ids=[chroma_id], include=["documents", "metadatas"] ) if not result["ids"]: return None return { "chroma_id": chroma_id, "content": result["documents"][0], "metadata": result["metadatas"][0], } def get_chunks_by_metadata(self, where: Dict) -> List[Dict]: """Metadata filtresine uyan tüm chunk'ları getir.""" collection = self.vectorstore._collection result = collection.get( where=where, include=["documents", "metadatas"] ) return [ {"chroma_id": cid, "content": doc, "metadata": meta} for cid, doc, meta in zip( result["ids"], result["documents"], result["metadatas"] ) ] def update_chunk( self, chroma_id: str, new_content: str, new_metadata: Optional[Dict] = None, ) -> None: """Chunk içeriğini yeniden embed ederek güncelle.""" new_embedding = self.embeddings.embed_query(new_content) collection = self.vectorstore._collection kwargs: Dict = { "ids": [chroma_id], "documents": [new_content], "embeddings": [new_embedding], } if new_metadata is not None: kwargs["metadatas"] = [new_metadata] collection.update(**kwargs) def delete_chunk(self, chroma_id: str) -> None: """Chunk'ı koleksiyondan kalıcı olarak sil.""" self.vectorstore._collection.delete(ids=[chroma_id]) def add_chunk(self, content: str, metadata: Dict) -> str: """Yeni chunk ekle; yeni ChromaDB UUID döner.""" import uuid as _uuid new_id = str(_uuid.uuid4()) embedding = self.embeddings.embed_query(content) clean_meta = {k: (v if v is not None else "") for k, v in metadata.items()} self.vectorstore._collection.add( ids=[new_id], documents=[content], embeddings=[embedding], metadatas=[clean_meta], ) return new_id def get_stats(self) -> Dict: """Knowledge base istatistikleri""" collection = self.vectorstore._collection return { "collection_name": self.collection_name, "document_count": collection.count(), "embedding_model": self.embedding_model_name, "persist_directory": self.persist_directory } def clear(self) -> None: """Knowledge base'i tamamen temizle (dikkatli kullan!)""" self.vectorstore.delete_collection() self.vectorstore = self._initialize_vectorstore() print("Knowledge base cleared") # Test if __name__ == "__main__": kb = KnowledgeBase() stats = kb.get_stats() print(f"\nKnowledge Base Stats:") for key, value in stats.items(): print(f" {key}: {value}")