RAGGER/vectorstore.py

from qdrant_client import QdrantClient
from qdrant_client.models import Distance, VectorParams, PointStruct, Filter, FieldCondition, MatchValue
from pathlib import Path
from typing import Optional
import uuid

from config import settings
from ollama_client import ollama
from chunker import Chunk
from logger import get_logger

log = get_logger("vectorstore")


class VectorStore:
    def __init__(self):
        self.path = Path(settings.QDRANT_PATH)
        self.path.mkdir(parents=True, exist_ok=True)
        self.collection = settings.QDRANT_COLLECTION
        log.debug(f"Initializing Qdrant at: {self.path}")
        self.client = QdrantClient(path=str(self.path), force_disable_check_same_thread=True)
        self._ensure_collection()
    
    def _ensure_collection(self):
        collections = self.client.get_collections().collections
        names = [c.name for c in collections]
        
        if self.collection not in names:
            log.info(f"Creating collection: {self.collection}")
            self.client.create_collection(
                collection_name=self.collection,
                vectors_config=VectorParams(
                    size=768,
                    distance=Distance.COSINE
                )
            )
        else:
            log.debug(f"Collection exists: {self.collection}")
    
    def add_chunk(self, chunk: Chunk) -> str:
        log.debug(f"Adding chunk: source={chunk.source}, page={chunk.page}, index={chunk.index}")
        log.debug(f"Chunk text length: {len(chunk.text)} chars")
        
        embedding = ollama.embed(chunk.text)
        log.debug(f"Embedding generated: {len(embedding)} dims")
        
        point_id = str(uuid.uuid4())
        
        self.client.upsert(
            collection_name=self.collection,
            points=[
                PointStruct(
                    id=point_id,
                    vector=embedding,
                    payload={
                        "text": chunk.text,
                        "source": chunk.source,
                        "page": chunk.page,
                        "sheet": chunk.sheet,
                        "index": chunk.index,
                        "metadata": chunk.metadata or {}
                    }
                )
            ]
        )
        
        log.info(f"Chunk added: {point_id[:8]}... from {chunk.source}")
        return point_id
    
    def add_chunks(self, chunks: list[Chunk]) -> list[str]:
        log.info(f"Adding {len(chunks)} chunks...")
        ids = []
        for i, c in enumerate(chunks):
            log.debug(f"Processing chunk {i+1}/{len(chunks)}")
            ids.append(self.add_chunk(c))
        log.info(f"Added {len(ids)} chunks total")
        return ids
    
    def search(self, query: str, limit: int = 5) -> list[dict]:
        log.debug(f"Searching: '{query[:50]}...' (limit={limit})")
        
        query_embedding = ollama.embed(query)
        log.debug(f"Query embedding: {len(query_embedding)} dims")
        
        results = self.client.query_points(
            collection_name=self.collection,
            query=query_embedding,
            limit=limit
        )
        
        log.debug(f"Found {len(results.points)} results")
        
        output = []
        for r in results.points:
            payload = r.payload or {}
            output.append({
                "text": payload.get("text", ""),
                "source": payload.get("source", "unknown"),
                "page": payload.get("page"),
                "sheet": payload.get("sheet"),
                "score": r.score
            })
            if payload.get("sheet"):
                log.debug(f"  Result: {payload.get('source')} / {payload.get('sheet')} (score={r.score:.4f})")
            else:
                log.debug(f"  Result: {payload.get('source')} (score={r.score:.4f})")
        
        return output
    
    def delete_by_source(self, source: str) -> int:
        log.info(f"Deleting chunks from: {source}")
        result = self.client.delete(
            collection_name=self.collection,
            points_selector=Filter(
                must=[
                    FieldCondition(
                        key="source",
                        match=MatchValue(value=source)
                    )
                ]
            )
        )
        log.info(f"Delete result: {result}")
        return result.operation_id if hasattr(result, 'operation_id') else 0
    
    def count(self) -> int:
        info = self.client.get_collection(self.collection)
        count = info.points_count or 0
        log.debug(f"Collection count: {count}")
        return count
    
    def clear(self) -> bool:
        log.warning(f"Clearing collection: {self.collection}")
        try:
            self.client.delete_collection(self.collection)
            log.info("Collection deleted")
            self._ensure_collection()
            log.info("Collection recreated (empty)")
            return True
        except Exception as e:
            log.error(f"Error clearing collection: {e}")
            return False
    
    def reset(self) -> bool:
        log.warning("RESETTING VectorStore - deleting all data!")
        try:
            if self.collection in [c.name for c in self.client.get_collections().collections]:
                self.client.delete_collection(self.collection)
                log.info("Collection deleted")
            self._ensure_collection()
            log.info("VectorStore reset complete")
            return True
        except Exception as e:
            log.error(f"Error resetting VectorStore: {e}")
            return False


vectorstore = VectorStore()