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)


async def sync_knowledge_base():
    """执行知识库同步 (全量+增量)"""
    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()

    articles = await loader.load_updated_since(state["last_sync"])
    if not articles:
        logger.info("知识库同步: 无更新内容")
        return

    docs = loader.format_for_indexing(articles)

    chunks = []
    for doc in docs:
        split_texts = splitter.split_text(doc["content"])
        for i, text in enumerate(split_texts):
            metadata = dict(doc["metadata"])
            metadata["chunk_index"] = i
            chunks.append(Document(page_content=text, metadata=metadata))

    if not chunks:
        logger.info("知识库同步: 无新增块")
        return

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

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

    logger.info("知识库同步完成: 新增 %d 篇文章, %d 个块", len(articles), len(chunks))