quant-llm-wiki 0.2.0__py3-none-any.whl

quant_llm_wiki/agent/__init__.py

@@ -0,0 +1,3 @@
+from quant_llm_wiki.agent.graph import create_agent
+
+__all__ = ["create_agent"]

quant_llm_wiki/agent/cli.py

@@ -0,0 +1,133 @@
+#!/usr/bin/env python3
+"""CLI entry point for the knowledge base agent.
+
+Usage:
+    # Interactive multi-turn conversation
+    python3 agent_cli.py
+
+    # Single command
+    python3 agent_cli.py --query "list all articles"
+"""
+from __future__ import annotations
+
+import argparse
+import sys
+from pathlib import Path
+from typing import Any
+
+from langchain_core.messages import HumanMessage
+
+from quant_llm_wiki.shared import _sanitize_lone_surrogates
+
+ROOT = Path(__file__).resolve().parent.parent.parent
+if str(ROOT) not in sys.path:
+    sys.path.insert(0, str(ROOT))
+
+
+def _safe_text(value: Any) -> str:
+    """Coerce to str and strip lone surrogates so print/encode never crashes.
+
+    Why: defense in depth at the output boundary — the agent's `post_model_hook`
+    and `post_llm_json` already clean upstream, but nothing stops a stray
+    surrogate from arriving here via a tool message, exception repr, etc.
+    """
+    text = value if isinstance(value, str) else str(value)
+    return _sanitize_lone_surrogates(text)
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Knowledge base agent CLI")
+    parser.add_argument("--query", help="Single query to run (non-interactive mode)")
+    return parser.parse_args()
+
+
+def _extract_last_ai_content(messages) -> str:
+    """Extract the last AI message content that isn't a pure tool call."""
+    for msg in reversed(messages):
+        if hasattr(msg, "type") and msg.type == "ai" and msg.content and not getattr(msg, "tool_calls", None):
+            return msg.content
+        if hasattr(msg, "type") and msg.type == "ai" and msg.content:
+            return msg.content
+    return str(messages[-1].content) if messages else "No response."
+
+
+def run_query(agent, query: str) -> str:
+    """Run a single query through the agent, streaming intermediate output."""
+    messages = []
+    for state in agent.stream(
+        {"messages": [{"role": "user", "content": query}]},
+        stream_mode="values",
+    ):
+        messages = state.get("messages", messages)
+    return _extract_last_ai_content(messages)
+
+
+def interactive_loop(agent) -> None:
+    """Run an interactive multi-turn conversation."""
+    print("Knowledge Base Agent (type 'quit' or 'exit' to stop)")
+    print("-" * 50)
+    messages: list = []
+    while True:
+        try:
+            user_input = input("\nYou: ").strip()
+        except (EOFError, KeyboardInterrupt):
+            print("\nBye!")
+            break
+        if not user_input:
+            continue
+        if user_input.lower() in ("quit", "exit", "q"):
+            print("Bye!")
+            break
+
+        messages.append(HumanMessage(content=user_input))
+        try:
+            for state in agent.stream(
+                {"messages": messages},
+                stream_mode="values",
+            ):
+                messages = state.get("messages", messages)
+            print(f"\nAgent: {_safe_text(_extract_last_ai_content(messages))}")
+        except KeyboardInterrupt:
+            print("\nInterrupted.")
+            break
+        except Exception as exc:
+            print(_safe_text(f"\nError ({type(exc).__name__}): {exc}"))
+
+
+def register(parser: argparse.ArgumentParser) -> None:
+    """Attach this module's CLI flags to `parser`. Called by quant_llm_wiki.cli."""
+    parser.add_argument("--query", help="Single query to run (non-interactive mode)")
+    parser.set_defaults(func=_run)
+
+
+def _run(args) -> int:
+    """The module's command body. Receives parsed args from the dispatcher."""
+    from quant_llm_wiki.agent import create_agent
+
+    agent = create_agent()
+
+    if args.query:
+        try:
+            print(_safe_text(run_query(agent, args.query)))
+        except Exception as exc:
+            print(
+                _safe_text(f"Error ({type(exc).__name__}): {exc}"),
+                file=sys.stderr,
+            )
+            return 1
+        return 0
+
+    interactive_loop(agent)
+    return 0
+
+
+def main() -> int:
+    """Standalone entry: python -m quant_llm_wiki.agent.cli ..."""
+    parser = argparse.ArgumentParser(description="Knowledge base agent CLI")
+    register(parser)
+    args = parser.parse_args()
+    return args.func(args)
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

quant_llm_wiki/agent/graph.py

@@ -0,0 +1,92 @@
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+ROOT = Path(__file__).resolve().parent.parent.parent
+if str(ROOT) not in sys.path:
+    sys.path.insert(0, str(ROOT))
+
+from typing import Any
+
+from langchain_core.messages import AIMessage
+from langchain_openai import ChatOpenAI
+from langgraph.prebuilt import create_react_agent
+
+from quant_llm_wiki.agent.prompts import SYSTEM_PROMPT
+from quant_llm_wiki.agent.tools import ALL_TOOLS
+from quant_llm_wiki.shared import _sanitize_response_strings, get_llm_config
+
+# AIMessage fields that may carry provider-emitted strings and therefore can
+# transport lone UTF-16 surrogates. We deliberately skip `response_metadata`
+# (model name / finish reason / token counts — low-risk, rarely user-visible)
+# and `usage_metadata` (numeric).
+_SANITIZED_FIELDS = ("content", "tool_calls", "invalid_tool_calls", "additional_kwargs")
+
+
+def _sanitize_agent_message(state: dict[str, Any]) -> dict[str, Any]:
+    """Strip lone UTF-16 surrogates from the most recent AIMessage.
+
+    Why: ChatOpenAI bypasses our `post_llm_json` sanitizer, so any surrogate
+    halves emitted by the provider would reach LangGraph's stream/print path
+    and crash on UTF-8 encode. Returning a copy of the same message (same id)
+    causes the `add_messages` reducer to replace it in place rather than
+    appending a duplicate.
+
+    Scope: only AIMessage. Tool outputs (ToolMessage) come from our own tool
+    code — if a tool ever returned surrogate-bearing text, the output-boundary
+    `_safe_text` in `agent/cli.py` catches it on print. Sanitizing ToolMessage
+    in the hook would shift framework-internal state for a contingency we
+    don't actually have a producer for.
+    """
+    messages = state.get("messages") or []
+    if not messages:
+        return {}
+    last = messages[-1]
+    if not isinstance(last, AIMessage):
+        return {}
+
+    update: dict[str, Any] = {}
+    for field in _SANITIZED_FIELDS:
+        original = getattr(last, field, None)
+        if not original:
+            continue
+        cleaned = _sanitize_response_strings(original)
+        if cleaned != original:
+            update[field] = cleaned
+
+    if not update:
+        return {}
+
+    return {"messages": [last.model_copy(update=update)]}
+
+
+def create_agent():
+    """Create and return a compiled LangGraph ReAct agent.
+
+    Uses the OpenAI-compatible API configured via LLM_* or ZHIPU_* env vars.
+    Works with any provider: Zhipu GLM, DeepSeek, Moonshot, Qwen, OpenAI, etc.
+    """
+    try:
+        api_key, base_url, model = get_llm_config()
+    except RuntimeError as exc:
+        raise RuntimeError(
+            f"Failed to initialize agent LLM: {exc}\n"
+            f"Configure via .env file or environment variables. "
+            f"See llm_config.example.env for examples."
+        ) from exc
+
+    llm = ChatOpenAI(
+        model=model,
+        api_key=api_key,
+        base_url=base_url,
+        temperature=0.1,
+    )
+
+    agent = create_react_agent(
+        model=llm,
+        tools=ALL_TOOLS,
+        prompt=SYSTEM_PROMPT,
+        post_model_hook=_sanitize_agent_message,
+    )
+    return agent

quant_llm_wiki/agent/prompts.py

@@ -0,0 +1,48 @@
+SYSTEM_PROMPT = """你是量化投研知识库管理助手。你管理一个完整的知识库流水线，包括文章抓取、LLM结构化增强、状态审核、向量索引、Wiki概念合成和RAG问答/脑暴。
+
+你可以使用以下工具：
+
+1. **ingest_article** — 抓取文章并保存到 articles/raw/。支持多种输入：
+   - url: 单个URL（自动识别 WeChat / 通用网页 / PDF）
+   - urls: 多个URL（换行/逗号分隔）
+   - url_list_file: URL 列表文件
+   - html_file: 本地 HTML 文件（WeChat 风格）
+   - pdf_file: 本地 PDF 文件
+   - pdf_url: 远程 PDF URL
+2. **enrich_articles** — 对原始文章进行 LLM 结构化增强（生成 idea_blocks 等字段）
+3. **list_articles** — 列出各阶段文章
+4. **review_articles** — 展示待审核文章
+5. **set_article_status** — 批量更新文章状态
+6. **embed_knowledge** — 构建/更新 ChromaDB 向量索引（同时索引 wiki/）
+7. **query_knowledge_base** — 问答(ask) / 脑暴(brainstorm)。Wiki 概念优先；向量库仅作为补充/兜底
+8. **compile_wiki** — 由文章合成 wiki 概念文章和 source 摘要。模式: incremental（默认）/ rebuild
+9. **list_concepts** — 列出 wiki 概念，按状态筛选（stable / proposed / deprecated）
+10. **set_concept_status** — 批准 / 弃用 / 删除概念（stable / deprecated / deleted）
+11. **read_wiki** — 读取 INDEX、概念文章或 source 摘要
+
+## Wiki 层使用指南
+
+- **"解释 X" / "梳理 Y" / "总结知识库对 Z 怎么说"** → 优先 read_wiki，target 用概念 slug。如概念不存在，才退回 query_knowledge_base
+- **"脑暴" / "组合想法" / "新策略"** → query_knowledge_base(mode='brainstorm')。它会自动优先检索 wiki 概念，再用复杂检索找互补文章
+- **"找包含 X 的文章" / "做新颖度检查"** → query_knowledge_base(mode='ask')
+
+## 典型工作流
+
+### 完整入库流程
+ingest_article → enrich_articles → review_articles → set_article_status → compile_wiki → embed_knowledge
+
+注意：所有文章统一存放在 raw/ 下，frontmatter 的 status 字段决定其阶段（reviewed / high_value / rejected）。compile_wiki 读取所有非 raw 状态的文章。embed_knowledge 在 compile_wiki 之后运行，使新合成的 wiki 内容也进入向量索引。
+
+### 概念审核流程
+当 compile_wiki 报告有 N 个 proposed 概念时：
+1. 调用 list_concepts(status='proposed') 展示
+2. 等待用户决定哪些批准、哪些拒绝
+3. 调用 set_concept_status 批量处理
+4. 如有批准，建议再次运行 compile_wiki 以让批准的概念被纳入合成
+
+## 规则
+- 用用户使用的语言回复（中文或英文）
+- 报告结果时清晰简洁，不要编造
+- 链式操作时，每步完成后报告结果再继续下一步
+- 只执行用户明确要求的操作，不要自动链式执行未请求的步骤
+"""