chat-ai-api/app/core/qdrant_client.py

"""
Qdrant 向量数据库客户端
用于存储和检索文档切片的向量嵌入
"""
from qdrant_client import QdrantClient
from qdrant_client.models import Distance, VectorParams, PointStruct, Filter, FieldCondition, MatchValue
from typing import List, Dict, Optional
import os
import uuid


class QdrantManager:
    def __init__(self):
        # 从环境变量读取 Qdrant 配置
        self.host = os.getenv("QDRANT_HOST", "localhost")
        self.port = int(os.getenv("QDRANT_PORT", "6333"))
        self.collection_name = os.getenv("QDRANT_COLLECTION", "kb_vectors")
        self.vector_size = 1536  # doubao-embedding-vision-251215 的向量维度

        # 初始化客户端
        self.client = QdrantClient(host=self.host, port=self.port)

        # 确保集合存在
        self._ensure_collection()

    def _ensure_collection(self):
        """确保 Qdrant 集合存在，不存在则创建"""
        collections = self.client.get_collections().collections
        collection_names = [col.name for col in collections]

        if self.collection_name not in collection_names:
            self.client.create_collection(
                collection_name=self.collection_name,
                vectors_config=VectorParams(
                    size=self.vector_size,
                    distance=Distance.COSINE
                )
            )
            print(f"创建 Qdrant 集合: {self.collection_name}")

    def add_vectors(self, tenant_id: str, doc_id: str, chunks: List[Dict]) -> List[str]:
        """
        添加文档切片的向量到 Qdrant
        :param tenant_id: 租户ID
        :param doc_id: 文档ID
        :param chunks: 切片列表 [{"content": "...", "vector": [...]}, ...]
        :return: vector_id 列表
        """
        points = []
        vector_ids = []

        for idx, chunk in enumerate(chunks):
            vector_id = str(uuid.uuid4())
            vector_ids.append(vector_id)

            point = PointStruct(
                id=vector_id,
                vector=chunk["vector"],
                payload={
                    "tenant_id": tenant_id,
                    "doc_id": doc_id,
                    "chunk_index": idx,
                    "content": chunk["content"][:500],  # 存储前500字符用于预览
                }
            )
            points.append(point)

        # 批量插入
        self.client.upsert(
            collection_name=self.collection_name,
            points=points
        )

        return vector_ids

    def search_vectors(self, tenant_id: str, query_vector: List[float], top_k: int = 5) -> List[Dict]:
        """
        向量相似度搜索
        :param tenant_id: 租户ID（只搜索该租户的数据）
        :param query_vector: 查询向量
        :param top_k: 返回Top K个结果
        :return: 搜索结果列表
        """
        search_result = self.client.search(
            collection_name=self.collection_name,
            query_vector=query_vector,
            query_filter=Filter(
                must=[
                    FieldCondition(
                        key="tenant_id",
                        match=MatchValue(value=tenant_id)
                    )
                ]
            ),
            limit=top_k
        )

        results = []
        for hit in search_result:
            results.append({
                "vector_id": hit.id,
                "score": hit.score,
                "content": hit.payload.get("content", ""),
                "doc_id": hit.payload.get("doc_id", ""),
                "chunk_index": hit.payload.get("chunk_index", 0),
            })

        return results

    def delete_vectors_by_doc(self, doc_id: str):
        """
        删除文档的所有向量
        :param doc_id: 文档ID
        """
        self.client.delete(
            collection_name=self.collection_name,
            points_selector=Filter(
                must=[
                    FieldCondition(
                        key="doc_id",
                        match=MatchValue(value=doc_id)
                    )
                ]
            )
        )


# 全局单例
_qdrant_manager = None


def get_qdrant_manager() -> QdrantManager:
    """获取 Qdrant 管理器单例"""
    global _qdrant_manager
    if _qdrant_manager is None:
        _qdrant_manager = QdrantManager()
    return _qdrant_manager