"""知识库同步定时任务: 定期从 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): try: existing = vectorstore.get() ids = [i for i in existing.get("ids", []) if i.startswith(prefix)] if ids: vectorstore.delete(ids) logger.info("已删除前缀 %s 的 %d 个旧块", prefix, len(ids)) except Exception as e: logger.warning("删除前缀 %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)