cfc/cfc-langgraph/app/tasks/knowledge_sync.py

"""知识库同步定时任务: 定期从 Java 拉取文章, 更新 ChromaDB"""
from app.rag.loader import KnowledgeLoader
from app.rag.splitter import get_knowledge_splitter
from app.rag.embeddings import get_embeddings
from app.config import settings
from langchain_chroma import Chroma
from langchain_core.documents import Document
import logging
import os
import json

logger = logging.getLogger(__name__)

SYNC_STATE_FILE = os.path.join(settings.chroma_db_path, ".sync_state")


def _load_sync_state() -> dict:
    try:
        if os.path.exists(SYNC_STATE_FILE):
            with open(SYNC_STATE_FILE) as f:
                return json.load(f)
    except Exception:
        pass
    return {"last_sync": "2000-01-01T00:00:00"}


def _save_sync_state(state: dict):
    os.makedirs(os.path.dirname(SYNC_STATE_FILE), exist_ok=True)
    with open(SYNC_STATE_FILE, "w") as f:
        json.dump(state, f)


def _delete_by_prefix(vectorstore, prefix: str):
    """删除指定 source 的所有旧块。

    修复: 原实现 vectorstore.get() 会拉取整个 collection (数百万条 embedding)
    到内存, 导致 gunicorn worker OOM 反复重启。改为按 metadata where 过滤
    只取目标 source 的 ids 再删除, 内存占用与数据总量无关。
    """
    try:
        collection = vectorstore._collection
        existing = collection.get(where={"source": prefix}, include=[])
        ids = existing.get("ids", []) or []
        if ids:
            collection.delete(ids=ids)
            logger.info("已删除 source=%s 的 %d 个旧块", prefix, len(ids))
        else:
            logger.info("无 source=%s 旧块可删除", prefix)
    except Exception as e:
        logger.warning("删除 source=%s 旧块失败: %s", prefix, e)


def _build_chunks(splitter, items: list[dict], source: str) -> list[Document]:
    chunks = []
    for item in items:
        item_id = item.get("id")
        title = item.get("title", "")
        content = item.get("content", "")
        if item_id is None or not content:
            continue
        full_text = f"{title}\n\n{content}"
        split_texts = splitter.split_text(full_text)
        for i, text in enumerate(split_texts):
            metadata = {
                "id": f"{source}_{item_id}_{i}",
                "source": source,
                "post_id": str(item_id),
                "title": title,
                "category": item.get("category", ""),
                "tags": item.get("tags", ""),
                "source_url": item.get("sourceUrl", ""),
                "pub_date": item.get("pubDate", ""),
                "chunk_index": i,
                "chunk_total": len(split_texts),
            }
            chunks.append(Document(page_content=text, metadata=metadata))
    return chunks


async def sync_knowledge_base():
    """同步三个数据源: articles 增量, microbiome/dan_knowledge 全量重建到 cfc_knowledge"""
    loader = KnowledgeLoader()
    splitter = get_knowledge_splitter()
    embeddings = get_embeddings()

    vectorstore = Chroma(
        collection_name="cfc_knowledge",
        embedding_function=embeddings,
        persist_directory=settings.chroma_db_path,
    )

    state = _load_sync_state()
    BATCH_SIZE = 100
    stats = {}

    articles = await loader.load_updated_since(state["last_sync"])
    article_chunks = []
    if articles:
        docs = loader.format_for_indexing(articles)
        for doc in docs:
            if not doc.get("content"):
                continue
            split_texts = splitter.split_text(doc["content"])
            for i, text in enumerate(split_texts):
                metadata = dict(doc["metadata"])
                metadata["chunk_index"] = i
                article_chunks.append(Document(page_content=text, metadata=metadata))
    stats["article"] = len(article_chunks)
    logger.info("文章增量: %d 篇 -> %d 块", len(articles), len(article_chunks))

    microbiome = await loader.java.get_microbiome_articles()
    _delete_by_prefix(vectorstore, "microbiome")
    microbiome_chunks = _build_chunks(splitter, microbiome, "microbiome")
    stats["microbiome"] = len(microbiome_chunks)
    logger.info("菌群知识库: %d 篇 -> %d 块", len(microbiome), len(microbiome_chunks))

    dan_kb = await loader.java.get_dan_knowledge_base()
    _delete_by_prefix(vectorstore, "dan_knowledge")
    dan_chunks = _build_chunks(splitter, dan_kb, "dan_knowledge")
    stats["dan_knowledge"] = len(dan_chunks)
    logger.info("知识库(dan): %d 篇 -> %d 块", len(dan_kb), len(dan_chunks))

    all_chunks = article_chunks + microbiome_chunks + dan_chunks
    if not all_chunks:
        logger.info("知识库同步: 无新增内容")
        return

    for i in range(0, len(all_chunks), BATCH_SIZE):
        batch = all_chunks[i:i + BATCH_SIZE]
        await vectorstore.aadd_documents(batch)
        logger.debug("知识库同步进度: %d/%d", i + len(batch), len(all_chunks))

    import datetime
    state["last_sync"] = datetime.datetime.now().isoformat()
    _save_sync_state(state)

    logger.info("知识库同步完成: 各源块数 %s", stats)