"""知识库同步定时任务: 三源(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", ""), } # 同步互斥锁：防止定时任务与手动触发并发写同一 Chroma collection。 # gunicorn 单 worker 单事件循环下 asyncio.Lock 足够；两任务并发曾导致 # microbiome 块被相互删除（缺 844 块），加锁后后到者直接跳过。 _sync_lock = asyncio.Lock() async def sync_knowledge_base(): """同步三个数据源: article/microbiome/dan_knowledge 全量重建到 cfc_knowledge。 加并发锁（后到者跳过）；每源独立失败隔离；全部成功后更新 .sync_state。 """ if _sync_lock.locked(): logger.info("已有知识库同步任务进行中，本次触发跳过") return async with _sync_lock: await _do_sync() async def _do_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)