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

"""知识库同步定时任务: 三源(article/microbiome/dan_knowledge)幂等同步到 ChromaDB

治理改造（2026-10-19）：
- article 由「只增量追加」改为「全量重建」：旧实现永不删除 article 旧块，
  导致下架/删除文章残留在向量库、HNSW 索引无限膨胀（曾达 11.8GB）。
- 三源 metadata 键统一（id/source/post_id/title/category/tags/source_url/pub_date/chunk_index/chunk_total），
  消除旧实现 retriever.initialize 的 metadata.type=article 与同步任务 metadata.source=article 双份数据。
- 显式传入 Chroma doc id（source_itemId_chunkIdx），写入幂等，可安全重试。
- 每源独立 try/except，一处失败不影响其他源。
"""
import asyncio
import datetime
import json
import logging
import os

from langchain_chroma import Chroma
from langchain_core.documents import Document

from app.config import settings
from app.rag.embeddings import get_embeddings
from app.rag.loader import KnowledgeLoader
from app.rag.splitter import get_knowledge_splitter

logger = logging.getLogger(__name__)

SYNC_STATE_FILE = os.path.join(settings.chroma_db_path, ".sync_state")
BATCH_SIZE = 100          # 每批写入的块数
MAX_WRITE_RETRY = 3       # 单批写入失败重试次数（配合 embedding 429 退避）


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) -> tuple[list[Document], list[str]]:
    """将条目分块，返回 (chunks, doc_ids)。

    doc_id 形如 source_itemId_chunkIdx，作为 Chroma 文档 ID 显式传入，
    保证同一知识条目的块可被幂等覆盖（而非每次生成随机 UUID 无限累积）。
    兼容旧 metadata 字段映射（sourceUrl/pubDate）。
    """
    chunks: list[Document] = []
    ids: list[str] = []
    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
        # 若 content 已含 title 前缀则不重复拼接
        full_text = f"{title}\n\n{content}" if title and title not in content else content
        split_texts = splitter.split_text(full_text)
        for i, text in enumerate(split_texts):
            doc_id = f"{source}_{item_id}_{i}"
            metadata = {
                "id": doc_id,
                "source": source,
                "post_id": str(item_id),
                "title": title,
                "category": item.get("category", ""),
                "tags": item.get("tags", ""),
                "source_url": item.get("source_url", "") or item.get("sourceUrl", ""),
                "pub_date": item.get("pub_date", "") or item.get("pubDate", ""),
                "chunk_index": i,
                "chunk_total": len(split_texts),
            }
            chunks.append(Document(page_content=text, metadata=metadata))
            ids.append(doc_id)
    return chunks, ids


async def _add_batch_with_retry(vectorstore, chunks: list, ids: list, label: str, start: int):
    """单批写入，失败指数退避重试（embedding 网关限流时自动恢复）。"""
    delay = 2.0
    for attempt in range(MAX_WRITE_RETRY + 1):
        try:
            await vectorstore.aadd_documents(chunks, ids=ids)
            return
        except Exception as e:
            if attempt == MAX_WRITE_RETRY:
                raise
            logger.warning("向量化写入失败[%s] %d-%d，%.0fs 后重试 (%d/%d): %s",
                           label, start, start + len(chunks), delay, attempt + 1, MAX_WRITE_RETRY, e)
            await asyncio.sleep(delay)
            delay = min(delay * 2, 30.0)


async def _write_chunks(vectorstore, chunks: list, ids: list, label: str):
    """分批写入并输出进度日志。"""
    total = len(chunks)
    if total == 0:
        logger.info("知识库同步[%s]: 无内容可写入", label)
        return
    for i in range(0, total, BATCH_SIZE):
        batch_chunks = chunks[i:i + BATCH_SIZE]
        batch_ids = ids[i:i + BATCH_SIZE]
        await _add_batch_with_retry(vectorstore, batch_chunks, batch_ids, label, i)
        done = min(i + BATCH_SIZE, total)
        logger.info("向量化写入进度[%s]: %d/%d", label, done, total)


async def _sync_source(vectorstore, splitter, source: str, fetch_func) -> int:
    """全量重建一个 source：先删旧块，再全量拉取重建，返回写入块数。"""
    try:
        _delete_by_prefix(vectorstore, source)
        items = await fetch_func()
        chunks, ids = _build_chunks(splitter, items, source)
        await _write_chunks(vectorstore, chunks, ids, source)
        logger.info("知识库同步[%s]: %d 篇 -> %d 块", source, len(items), len(chunks))
        return len(chunks)
    except Exception as e:
        logger.error("知识库同步[%s] 失败: %s", source, e)
        return 0


def _format_article_item(a: dict) -> dict:
    """兼容 Java 文章结构，统一映射为题名/摘要/正文/标签/时间。"""
    title = a.get("title", "")
    return {
        "id": a["id"],
        "title": title,
        "content": f"{a.get('summary', '')}\n\n{a.get('content', '')}".strip(),
        "tags": a.get("tags", ""),
        "pub_date": a.get("updatedAt", ""),
    }


async def sync_knowledge_base():
    """同步三个数据源: article/microbiome/dan_knowledge 全量重建到 cfc_knowledge。

    每源独立失败隔离；全部成功后更新 .sync_state 记录 last_sync。
    """
    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()
    stats: dict = {}

    async def fetch_articles():
        articles = await loader.load_all_articles()
        return [_format_article_item(a) for a in articles if a.get("id") is not None]

    stats["article"] = await _sync_source(vectorstore, splitter, "article", fetch_articles)
    stats["microbiome"] = await _sync_source(
        vectorstore, splitter, "microbiome", loader.java.get_microbiome_articles)
    stats["dan_knowledge"] = await _sync_source(
        vectorstore, splitter, "dan_knowledge", loader.java.get_dan_knowledge_base)

    if not any(stats.values()):
        logger.warning("知识库同步: 全部数据源失败，不更新 last_sync")
        return

    state["last_sync"] = datetime.datetime.now().isoformat()
    _save_sync_state(state)
    logger.info("知识库同步完成: 各源块数 %s", stats)