from typing import Any from uuid import uuid4 from qdrant_client import QdrantClient from qdrant_client.http import models as qmodels from sentence_transformers import SentenceTransformer from app.config import settings class VectorStore: def __init__(self) -> None: self.client = QdrantClient(url=settings.qdrant_url) self.embedder = SentenceTransformer(settings.embedding_model_name) self.collection = settings.qdrant_collection self._ensure_collection() def _ensure_collection(self) -> None: collections = self.client.get_collections().collections exists = any(c.name == self.collection for c in collections) if not exists: self.client.create_collection( collection_name=self.collection, vectors_config=qmodels.VectorParams(size=384, distance=qmodels.Distance.COSINE), ) def add_case_vector(self, case_id: int, text: str, metadata: dict[str, Any]) -> None: vector = self.embedder.encode(text).tolist() self.client.upsert( collection_name=self.collection, points=[ qmodels.PointStruct( id=str(uuid4()), vector=vector, payload={"case_id": case_id, **metadata}, ) ], ) def search_similar(self, text: str, limit: int = 5) -> list[dict[str, Any]]: query_vector = self.embedder.encode(text).tolist() hits = self.client.search( collection_name=self.collection, query_vector=query_vector, limit=limit, with_payload=True, ) return [ { "score": hit.score, "case_id": hit.payload.get("case_id") if hit.payload else None, "case_name": hit.payload.get("case_name") if hit.payload else None, } for hit in hits ]