src/knowledge_base.py

# 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}")