""" Librarian Agent — KnowHub 的知识管理 Agent 同一个 Agent 处理多种任务模式（查询 / 上传等），调用方通过 skills 参数选择策略： - skills=["ask_strategy"] → 知识库查询整合 - skills=["upload_strategy"] → 知识上传图谱编排 通过 HTTP API 被 FastAPI server 调用，每次请求一次 AgentRunner.run()。 续跑由 caller 显式传入 continue_from 指定。 """ import json import logging import sys from pathlib import Path from typing import Optional, Dict, Any, List # 确保项目路径可用 sys.path.insert(0, str(Path(__file__).parent.parent.parent)) from agent.core.runner import AgentRunner, RunConfig from agent.trace import FileSystemTraceStore from agent.llm import create_qwen_llm_call from agent.llm.prompts import SimplePrompt from agent.tools.builtin.knowledge import KnowledgeConfig logger = logging.getLogger("agents.librarian") # ===== 配置 ===== ENABLE_DATABASE_COMMIT = False # Librarian 允许调用方注入的 skill 白名单。不在列表中的 skill 会被过滤掉（只打 warning，不报错）。 ALLOWED_SKILLS = ["ask_strategy", "upload_strategy"] DEFAULT_SKILLS = ["ask_strategy"] # 调用方不传时的默认值 def get_librarian_config(enable_db_commit: bool = ENABLE_DATABASE_COMMIT) -> RunConfig: """获取 Librarian Agent 配置""" tools = [ "knowledge_search", "tool_search", "capability_search", "requirement_search", "relation_search", "read_file", "write_file", "list_cache_status", "match_tree_nodes", "sync_atomic_capabilities", "skill", ] if enable_db_commit: tools.extend(["commit_to_database", "organize_cached_data", "cache_research_data"]) else: tools.extend(["organize_cached_data", "cache_research_data"]) return RunConfig( model="qwen3.5-plus", temperature=0.2, max_iterations=30, agent_type="default", name="Librarian Agent", goal_compression="on_complete", skills=[], # 不注入通用 skills（planning/research/browser），使用指定注入 knowledge=KnowledgeConfig( enable_extraction=False, enable_completion_extraction=False, enable_injection=False, ), tools=tools, tool_groups=[], # 精确指定工具，不从分组加载 ) def _register_internal_tools(): """注册内部工具""" try: sys.path.insert(0, str(Path(__file__).parent.parent)) from internal_tools.cache_manager import ( cache_research_data, organize_cached_data, commit_to_database, list_cache_status, ) from internal_tools.tree_matcher import match_tree_nodes from internal_tools.capability_extractor import sync_atomic_capabilities from agent.tools import get_tool_registry registry = get_tool_registry() registry.register(cache_research_data) registry.register(organize_cached_data) registry.register(commit_to_database) registry.register(list_cache_status) registry.register(match_tree_nodes) registry.register(sync_atomic_capabilities) logger.info("✓ 已注册 Librarian 内部工具") except Exception as e: logger.error(f"✗ 注册内部工具失败: {e}") # ===== 单例 Runner ===== _runner: Optional[AgentRunner] = None _prompt_messages = None _initialized = False def _ensure_initialized(): """延迟初始化 Runner 和 Prompt（首次调用时执行）""" global _runner, _prompt_messages, _initialized if _initialized: return _initialized = True _register_internal_tools() _runner = AgentRunner( trace_store=FileSystemTraceStore(base_path=".trace"), llm_call=create_qwen_llm_call(model="qwen3.5-plus"), skills_dir=str(Path(__file__).parent / "skills"), debug=True, logger_name="agents.librarian", ) prompt_path = Path(__file__).parent / "librarian_agent.prompt" if prompt_path.exists(): prompt = SimplePrompt(prompt_path) _prompt_messages = prompt.build_messages() else: _prompt_messages = [] logger.warning(f"Librarian prompt 文件不存在: {prompt_path}") logger.info("✓ Librarian Agent 已初始化") # ===== 核心方法 ===== async def run_librarian( query: str, continue_from: Optional[str] = None, skills: Optional[List[str]] = None, ) -> Dict[str, Any]: """ 同步运行 Librarian Agent。由 skills 参数决定当前是什么模式： - skills=["ask_strategy"] → 查询模式（默认） - skills=["upload_strategy"] → 上传模式（query 应为 JSON 字符串） Args: query: 任务内容。查询模式为自然语言问题；上传模式为 JSON 字符串 {knowledge, tools, resources} continue_from: 已有 sub_trace_id，传入则续跑该 trace skills: 调用方指定的 skill 列表，会被 ALLOWED_SKILLS 过滤 Returns: {"status", "sub_trace_id", "summary", "stats", "error"?} """ _ensure_initialized() # Skill 过滤 skills = _filter_skills(skills) # 判断模式：upload_strategy 走上传路径（会写 buffer + 解析 JSON） is_upload = "upload_strategy" in skills buffer_file = None if is_upload: try: data = json.loads(query) if isinstance(query, str) else query except json.JSONDecodeError as e: return _fail(f"upload 模式下 query 不是合法 JSON: {e}") if not isinstance(data, dict): return _fail("upload 模式下 query 应为 JSON 对象") if not (data.get("knowledge") or data.get("tools") or data.get("resources")): return _fail("upload 模式下 data 中无有效条目") buffer_file = _write_upload_buffer(data) content = f"[UPLOAD:BATCH] 收到上传请求，请处理：\n{json.dumps(data, ensure_ascii=False)}" else: content = f"[ASK] {query}" # 运行 Librarian config = get_librarian_config() config.trace_id = continue_from if continue_from is None: messages = _prompt_messages + [{"role": "user", "content": content}] else: messages = [{"role": "user", "content": content}] try: result = await _runner.run_result( messages=messages, config=config, inject_skills=skills, ) actual_trace_id = result.get("trace_id") if buffer_file: _update_buffer_status(buffer_file, "completed", trace_id=actual_trace_id) return { "status": result.get("status", "completed"), "sub_trace_id": actual_trace_id, "summary": result.get("summary", ""), "stats": result.get("stats", {}), "error": result.get("error"), } except Exception as e: if buffer_file: _update_buffer_status(buffer_file, "failed", error=str(e)) logger.error(f"[Librarian] run 失败: {e}") return _fail(str(e)) def _filter_skills(skills: Optional[List[str]]) -> List[str]: """把调用方传的 skills 按 ALLOWED_SKILLS 过滤；空列表退化到 DEFAULT_SKILLS。""" if not skills: return list(DEFAULT_SKILLS) allowed = [s for s in skills if s in ALLOWED_SKILLS] dropped = [s for s in skills if s not in ALLOWED_SKILLS] if dropped: logger.warning(f"[Librarian] 忽略不在白名单的 skills: {dropped}（允许: {ALLOWED_SKILLS}）") return allowed or list(DEFAULT_SKILLS) def _fail(error: str) -> Dict[str, Any]: return {"status": "failed", "sub_trace_id": None, "summary": "", "stats": {}, "error": error} def _write_upload_buffer(data: Dict[str, Any]) -> Optional[str]: """把 upload 数据写到 buffer 目录，便于审计和失败重跑。""" try: from datetime import datetime as dt buffer_dir = Path(".cache/.knowledge/buffer") buffer_dir.mkdir(parents=True, exist_ok=True) timestamp = dt.now().strftime("%Y%m%d_%H%M%S") path = buffer_dir / f"upload_{timestamp}.json" path.write_text(json.dumps({ "data": data, "received_at": dt.now().isoformat(), }, ensure_ascii=False, indent=2), encoding="utf-8") return str(path) except Exception as e: logger.warning(f"写 buffer 失败: {e}") return None def _update_buffer_status(buffer_file: Optional[str], status: str, trace_id: str = None, error: str = None): """更新 buffer 文件中的处理状态""" if not buffer_file: return try: from datetime import datetime as dt path = Path(buffer_file) if not path.exists(): return data = json.loads(path.read_text(encoding="utf-8")) data["status"] = status data["processed_at"] = dt.now().isoformat() if trace_id: data["librarian_trace_id"] = trace_id if error: data["error"] = error path.write_text(json.dumps(data, ensure_ascii=False, indent=2), encoding="utf-8") except Exception as e: logger.warning(f"更新 buffer 状态失败: {e}") def list_pending_uploads() -> list: """列出所有未处理或失败的 upload buffer 文件""" buffer_dir = Path(".cache/.knowledge/buffer") if not buffer_dir.exists(): return [] pending = [] for f in sorted(buffer_dir.glob("upload_*.json")): try: data = json.loads(f.read_text(encoding="utf-8")) status = data.get("status", "pending") if status in ("pending", "failed"): pending.append({ "file": str(f), "status": status, "received_at": data.get("received_at", ""), "error": data.get("error", ""), "trace_id": data.get("trace_id", ""), }) except Exception: pass return pending