src/vector_store.py

from typing import List, Dict, Any, Optional, Set
from langchain_qdrant import QdrantVectorStore
from langchain_core.documents import Document
from qdrant_client import QdrantClient
from qdrant_client.http.models import Distance, VectorParams, Filter, FieldCondition, MatchValue

from config import (
    QDRANT_PATH,
    COLLECTION_NAME,
    EMBEDDING_DIMENSION,
    USE_MEMORY_MODE
)
from embeddings import get_embedder
from qdrant_client_manager import get_qdrant_client


class VectorStore:
    """Qdrant vector store wrapper."""
    
    def __init__(
        self,
        path: str = QDRANT_PATH,
        collection_name: str = COLLECTION_NAME,
        use_memory: bool = USE_MEMORY_MODE,
        embedder=None
    ):
        self.collection_name = collection_name
        self.use_memory = use_memory
        self.path = path
        
        # Use shared Qdrant client to prevent multiple instance conflicts
        self._client = get_qdrant_client()
        
        self._ensure_collection_exists()
        self._embedder = embedder or get_embedder()
        
        self._vector_store = QdrantVectorStore(
            client=self._client,
            collection_name=self.collection_name,
            embedding=self._embedder
        )
    
    def _ensure_collection_exists(self):
        """Create collection if it doesn't exist."""
        collections = self._client.get_collections().collections
        names = [c.name for c in collections]
        
        if self.collection_name not in names:
            self._client.create_collection(
                collection_name=self.collection_name,
                vectors_config=VectorParams(
                    size=EMBEDDING_DIMENSION,
                    distance=Distance.COSINE
                )
            )
    
    def add_documents(
        self,
        texts: List[str],
        metadatas: Optional[List[Dict[str, Any]]] = None
    ) -> List[str]:
        """Add documents to the vector store."""
        if not texts:
            return []
        
        if metadatas is None:
            metadatas = [{} for _ in texts]
        
        documents = [
            Document(page_content=text, metadata=meta)
            for text, meta in zip(texts, metadatas)
        ]
        
        ids = self._vector_store.add_documents(documents)
        return ids
    
    def search(self, query: str, top_k: int = 20) -> List[Dict[str, Any]]:
        """Search for similar documents."""
        results = self._vector_store.similarity_search_with_score(
            query=query,
            k=top_k
        )
        
        formatted = []
        for doc, score in results:
            formatted.append({
                "id": doc.metadata.get("_id", ""),
                "score": score,
                "text": doc.page_content,
                "source": doc.metadata.get("source", "Unknown"),
                "chunk_index": doc.metadata.get("chunk_index", -1),
                "page_number": doc.metadata.get("page_number", -1),
                "metadata": doc.metadata
            })
        
        return formatted
    
    def get_collection_stats(self) -> Dict[str, Any]:
        """Get collection statistics."""
        try:
            info = self._client.get_collection(self.collection_name)
            count = info.points_count or 0
            return {
                "name": self.collection_name,
                "vectors_count": count,
                "points_count": count,
                "status": str(info.status)
            }
        except Exception:
            return {
                "name": self.collection_name,
                "vectors_count": 0,
                "points_count": 0,
                "status": "error"
            }
    
    def clear_collection(self):
        """Delete and recreate the collection."""
        self._client.delete_collection(self.collection_name)
        self._ensure_collection_exists()
        
        self._vector_store = QdrantVectorStore(
            client=self._client,
            collection_name=self.collection_name,
            embedding=self._embedder
        )
    
    def collection_exists(self) -> bool:
        """Check if collection has documents."""
        stats = self.get_collection_stats()
        return stats["points_count"] > 0

    def document_exists(self, source: str) -> bool:
        """
        Check if a document with the given source name exists in the collection.
        
        Args:
            source: The source filename to check for
            
        Returns:
            True if document exists, False otherwise
        """
        try:
            result = self._client.scroll(
                collection_name=self.collection_name,
                scroll_filter=Filter(
                    must=[
                        FieldCondition(
                            key="metadata.source",
                            match=MatchValue(value=source)
                        )
                    ]
                ),
                limit=1,
                with_payload=False,
                with_vectors=False
            )
            points, _ = result
            return len(points) > 0
        except Exception:
            return False

    def get_loaded_sources(self) -> Set[str]:
        """
        Get set of all unique source names in the collection.
        
        Returns:
            Set of source filenames
        """
        sources = set()
        try:
            offset = None
            while True:
                result = self._client.scroll(
                    collection_name=self.collection_name,
                    limit=100,
                    offset=offset,
                    with_payload=True,
                    with_vectors=False
                )
                points, offset = result
                
                for point in points:
                    if point.payload:
                        # Check both possible metadata structures
                        source = None
                        if "metadata" in point.payload and isinstance(point.payload["metadata"], dict):
                            source = point.payload["metadata"].get("source")
                        if not source:
                            source = point.payload.get("source")
                        if source:
                            sources.add(source)
                
                if offset is None:
                    break
            
            return sources
        except Exception:
            return sources


_vector_store_instance = None


def get_vector_store() -> VectorStore:
    """Return singleton vector store instance."""
    global _vector_store_instance
    if _vector_store_instance is None:
        _vector_store_instance = VectorStore()
    return _vector_store_instance


def reset_vector_store():
    """Reset singleton for testing."""
    global _vector_store_instance
    if _vector_store_instance is not None:
        try:
            _vector_store_instance.clear_collection()
        except Exception:
            pass
    _vector_store_instance = None


if __name__ == "__main__":
    store = VectorStore(use_memory=True)
    
    texts = ["Atlas ERP sistemi.", "Finans modülü özellikleri."]
    metadatas = [
        {"source": "test.pdf", "chunk_index": 0},
        {"source": "test.pdf", "chunk_index": 1}
    ]
    
    store.add_documents(texts, metadatas)
    print(f"Stats: {store.get_collection_stats()}")
    
    results = store.search("ERP nedir?", top_k=2)
    for r in results:
        print(f"Score: {r['score']:.4f} - {r['text'][:50]}")