Ai/rag/vector_store.py

"""
Qdrant Vector Store — ذخیره و جستجوی وکتورها
"""
from qdrant_client import QdrantClient
from qdrant_client.http import models as qmodels

from .client import get_qdrant_client
from .config import load_rag_config, RAGConfig


class QdrantVectorStore:
    """
    ذخیره و جستجوی documents در Qdrant.
    """

    def __init__(self, config: RAGConfig | None = None):
        self.config = config or load_rag_config()
        self.qdrant = self.config.qdrant
        self._client: QdrantClient | None = None

    @property
    def client(self) -> QdrantClient:
        if self._client is None:
            self._client = get_qdrant_client(self.qdrant)
        return self._client

    def ensure_collection(self, recreate: bool = False) -> None:
        """
        اطمینان از وجود collection با نام و اندازه مناسب.
        """
        name = self.qdrant.collection_name
        size = self.qdrant.vector_size

        try:
            self.client.get_collection(name)
            if recreate:
                self.client.delete_collection(name)
                self.client.create_collection(
                    collection_name=name,
                    vectors_config=qmodels.VectorParams(
                        size=size,
                        distance=qmodels.Distance.COSINE,
                    ),
                )
        except Exception:
            self.client.create_collection(
                collection_name=name,
                vectors_config=qmodels.VectorParams(
                    size=size,
                    distance=qmodels.Distance.COSINE,
                ),
            )

    def add_documents(
        self,
        ids: list[str],
        embeddings: list[list[float]],
        documents: list[str],
        metadatas: list[dict] | None = None,
    ) -> int:
        """
        افزودن documents به collection.
        metadata فقط str, int, float, bool پشتیبانی می‌شود.
        """
        self.ensure_collection()
        metas = metadatas or [{}] * len(ids)

        def _serialize(m: dict) -> dict:
            out = {}
            for k, v in m.items():
                if v is None:
                    continue
                if isinstance(v, (str, int, float, bool)):
                    out[k] = v
                else:
                    out[k] = str(v)
            return out

        payloads = [
            {"text": doc, "doc_id": sid, **_serialize(m)}
            for doc, m, sid in zip(documents, metas, ids)
        ]

        self.client.upsert(
            collection_name=self.qdrant.collection_name,
            points=[
                qmodels.PointStruct(id=pid, vector=emb, payload=pl)
                for pid, emb, pl in zip(ids, embeddings, payloads)
            ],
        )
        return len(ids)

    def search(
        self,
        query_vector: list[float],
        limit: int = 5,
        score_threshold: float | None = None,
        sensor_uuid: str | None = None,
        kb_name: str | None = None,
        sensor_uuids: list[str] | None = None,
        kb_names: list[str] | None = None,
    ) -> list[dict]:
        """
        جستجوی شباهت بر اساس query vector.
        روی نسخه‌های جدید از query_points و روی نسخه‌های قدیمی‌تر از search استفاده می‌کند.
        sensor_uuid: اجباری — فقط chunks مربوط به این سنسور یا __global__ برگردانده می‌شود.
        kb_name: اختیاری — فیلتر بر اساس پایگاه دانش (chat/irrigation/fertilization).
                 اگر مشخص شود، فقط chunks همان KB و __all__ برگردانده می‌شود.
        """
        must_conditions = []

        sensor_values = [value for value in (sensor_uuids or ([sensor_uuid] if sensor_uuid else [])) if value]
        if sensor_values:
            must_conditions.append(
                qmodels.Filter(
                    should=[
                        qmodels.FieldCondition(
                            key="sensor_uuid",
                            match=qmodels.MatchValue(value=value),
                        )
                        for value in sensor_values
                    ]
                )
            )

        kb_values = [value for value in (kb_names or ([kb_name] if kb_name else [])) if value]
        if kb_values:
            must_conditions.append(
                qmodels.Filter(
                    should=[
                        qmodels.FieldCondition(
                            key="kb_name",
                            match=qmodels.MatchValue(value=value),
                        )
                        for value in kb_values
                    ]
                )
            )

        query_filter = None
        if must_conditions:
            query_filter = qmodels.Filter(must=must_conditions)

        if hasattr(self.client, "query_points"):
            response = self.client.query_points(
                collection_name=self.qdrant.collection_name,
                query=query_vector,
                limit=limit,
                score_threshold=score_threshold,
                query_filter=query_filter,
            )
            points = getattr(response, "points", []) or []
        else:
            points = self.client.search(
                collection_name=self.qdrant.collection_name,
                query_vector=query_vector,
                limit=limit,
                score_threshold=score_threshold,
                query_filter=query_filter,
            )

        return [
            {
                "id": str(r.id),
                "score": float(r.score) if r.score is not None else 0.0,
                "text": (r.payload or {}).get("text", ""),
                "metadata": {
                    k: v for k, v in (r.payload or {}).items() if k != "text"
                },
            }
            for r in points
        ]