labor-arbitration-system/backend/app/services/vector_store.py

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
        ]