core/vectorstore.py

"""
Qdrant 기반 메뉴 벡터 저장소 (MenuVectorStore).

ChromaDB → Qdrant 교체. 공개 인터페이스는 동일하게 유지하여
search_engine / 빌드 스크립트 변경을 최소화한다.

Qdrant point ID 규칙:
    Qdrant는 point ID로 int 또는 UUID만 허용. menu_id("SCR_1300")는
    uuid5(NAMESPACE, menu_id)로 결정적 UUID 변환하고, 원본 menu_id는
    payload에 저장한다.

Distance.COSINE:
    정규화된 임베딩(normalize_embeddings=True) 기준 score = 코사인 유사도.
    ChromaDB의 (1 - distance) 변환 불필요.

연결 대상 (QDRANT_MODE):
    server (기본): QDRANT_URL HTTP 연결 (docker-compose qdrant 컨테이너 / 사내 서버)
    local        : QDRANT_PATH 폴더 기반 임베디드 실행 (서버 불필요, HF Spaces 단일 컨테이너)
"""
import uuid

from qdrant_client.models import (
    Distance, VectorParams, PointStruct,
    Filter, FieldCondition, MatchValue, PayloadSchemaType,
)

from config import QDRANT_COLLECTION_NAME, EMBEDDING_DIM
from core.qdrant_conn import get_qdrant_client

_UUID_NS = uuid.UUID("6f9619ff-8b86-d011-b42d-00c04fc964ff")  # 고정 네임스페이스


def _point_id(menu_id: str) -> str:
    """menu_id → 결정적 UUID (동일 menu_id는 항상 동일 point ID → upsert 멱등)"""
    return str(uuid.uuid5(_UUID_NS, menu_id))


class MenuVectorStore:
    """
    Qdrant 기반 메뉴 벡터 저장소.

    재시작해도 컬렉션 영속(볼륨). 컬렉션이 없으면 자동 생성하고
    category payload 인덱스를 만들어 카테고리 필터를 pre-filter로 처리한다.
    """

    def __init__(self):
        # 공유 클라이언트 사용 (local 모드 단일 인스턴스 제약 + 쿼리 인덱스와 공유)
        self.client = get_qdrant_client()
        self.collection_name = QDRANT_COLLECTION_NAME
        self._ensure_collection()

    def _ensure_collection(self):
        """컬렉션 없으면 생성 + category 인덱스 부여"""
        existing = {c.name for c in self.client.get_collections().collections}
        if self.collection_name not in existing:
            self.client.create_collection(
                collection_name=self.collection_name,
                vectors_config=VectorParams(size=EMBEDDING_DIM, distance=Distance.COSINE),
            )
            # 카테고리 pre-filter용 payload 인덱스
            self.client.create_payload_index(
                collection_name=self.collection_name,
                field_name="category",
                field_schema=PayloadSchemaType.KEYWORD,
            )
            print(f"[VectorStore] Qdrant 컬렉션 생성: {self.collection_name} (dim={EMBEDDING_DIM}, cosine)")

    # -------------------------------------------------------------------------
    # 쓰기
    # -------------------------------------------------------------------------

    def upsert_menus(self, menu_ids: list, embeddings: list, documents: list, metadatas: list):
        """메뉴 벡터 일괄 업서트 (menu_id·document를 payload에 동봉)"""
        points = []
        for mid, emb, doc, meta in zip(menu_ids, embeddings, documents, metadatas):
            payload = dict(meta) if meta else {}
            payload["menu_id"] = mid
            payload["document"] = doc
            points.append(PointStruct(id=_point_id(mid), vector=emb, payload=payload))

        # 대용량 대비 배치 업서트
        BATCH = 256
        for i in range(0, len(points), BATCH):
            self.client.upsert(
                collection_name=self.collection_name,
                points=points[i:i + BATCH],
            )
        print(f"[완료] {len(points)}개 메뉴 저장 완료")

    # -------------------------------------------------------------------------
    # 검색
    # -------------------------------------------------------------------------

    def search(self, query_embedding: list, top_n: int = 5, where: dict = None) -> list:
        """
        유사도 검색.

        Args:
            query_embedding: 쿼리 벡터 (1024dim, 정규화됨)
            top_n: 반환 개수
            where: {"category": "국내주식"} 형태. Qdrant pre-filter로 적용.

        Returns:
            [{"menu_id", "similarity", "document", "metadata"}, ...]
        """
        if self.count() == 0:
            return []

        qfilter = None
        if where and where.get("category"):
            qfilter = Filter(must=[
                FieldCondition(key="category", match=MatchValue(value=where["category"]))
            ])

        # qdrant-client 1.12+ : search() 제거 → query_points() 사용
        resp = self.client.query_points(
            collection_name=self.collection_name,
            query=query_embedding,
            query_filter=qfilter,
            limit=top_n,
            with_payload=True,
        )

        output = []
        for h in resp.points:
            payload = dict(h.payload or {})
            menu_id = payload.pop("menu_id", None)
            document = payload.pop("document", "")
            output.append({
                "menu_id": menu_id,
                "similarity": round(float(h.score), 4),  # COSINE → 유사도 직접
                "document": document,
                "metadata": payload,  # menu_id·document 제외한 나머지 = 기존 metadata
            })
        return output

    # -------------------------------------------------------------------------
    # 유틸
    # -------------------------------------------------------------------------

    def count(self) -> int:
        """총 포인트 수 (search_engine에서 dense_top_k 산정에 사용)"""
        try:
            return self.client.count(self.collection_name, exact=True).count
        except Exception:
            return 0

    def get_stats(self) -> dict:
        return {
            "total_menus": self.count(),
            "collection_name": self.collection_name,
        }

    def reset(self):
        """컬렉션 초기화 (재구축 시 사용)"""
        try:
            self.client.delete_collection(self.collection_name)
        except Exception:
            pass
        self._ensure_collection()
        print("컬렉션 초기화 완료")