cfc/cfc-langgraph/app/rag/retriever.py

from langchain_chroma import Chroma
from langchain_openai import ChatOpenAI
from app.config import settings
from app.rag.embeddings import get_embeddings
from app.tools.java_client import JavaClient
from typing import Optional
import logging

logger = logging.getLogger(__name__)


class RagRetriever:
    """简单的向量检索器 - 基于 ChromaDB"""

    def __init__(self, collection_name: str = "cfc_knowledge"):
        embeddings = get_embeddings()
        self.vectorstore = Chroma(
            collection_name=collection_name,
            embedding_function=embeddings,
            persist_directory=settings.chroma_db_path,
        )
        self.java_client = JavaClient()

    async def initialize(self):
        """初始化向量库连接（不写入任何内容）。

        内容统一由 app.tasks.knowledge_sync.sync_knowledge_base 维护（article/microbiome/dan_knowledge 三源）。

        旧实现在此额外写入一版粗粒度文章索引（metadata.type=article，标题+摘要+标签），
        与同步任务写入的 metadata.source=article 分块索引（标题+摘要+正文）构成同一批文章的
        双份向量，且每次容器重启都会追加，导致 HNSW 索引无限膨胀。治理后统一移除。
        """
        try:
            count = self.vectorstore._collection.count()
            logger.info("知识库初始化完成：collection 当前 %d 个文档（内容同步由 knowledge_sync 维护）", count)
        except Exception as e:
            logger.warning("知识库初始化失败：%s", e)

    async def retrieve(
        self,
        query: str,
        filters: Optional[dict] = None,
        k: int = 5,
    ) -> list[dict]:
        """向量相似度检索"""
        results = []
        
        # 向量检索
        filter_query = {"user_id": filters["user_id"]} if filters and "user_id" in filters else None
        
        docs = self.vectorstore.similarity_search(
            query,
            k=k * 2,  # 多取一些
            filter=filter_query,
        )
        
        # 格式化结果并去重
        seen = set()
        for doc in docs:
            content_hash = hash(doc.page_content[:100])
            if content_hash in seen:
                continue
            seen.add(content_hash)
            
            results.append({
                "content": doc.page_content,
                "metadata": doc.metadata,
                "score": doc.metadata.get("score", 0) if hasattr(doc, "metadata") else 0,
            })
            
            if len(results) >= k:
                break
        
        return results[:k]