chore: import upstream snapshot with attribution
This commit is contained in:
@@ -0,0 +1,3 @@
|
||||
recursive-include yuxi/config/static *.txt *.yaml
|
||||
recursive-include yuxi/agents/skills *.md
|
||||
recursive-include yuxi/agents/skills/buildin/mysql-reporter/scripts *.py
|
||||
@@ -0,0 +1,5 @@
|
||||
# yuxi
|
||||
|
||||
Yuxi 核心后端 Python 包。
|
||||
|
||||
该目录用于将 `yuxi` 作为本地可构建、可安装的 Python package 提供给 `backend` 项目使用。
|
||||
@@ -0,0 +1,118 @@
|
||||
[build-system]
|
||||
requires = ["setuptools>=80", "wheel"]
|
||||
build-backend = "setuptools.build_meta"
|
||||
|
||||
[project]
|
||||
name = "yuxi"
|
||||
version = "0.7.1.beta1"
|
||||
description = "Yuxi 智能知识库与知识图谱平台核心后端包"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.12,<3.14"
|
||||
dependencies = [
|
||||
"agent-sandbox>=0.0.26",
|
||||
"aioboto3>=13.0.0",
|
||||
"aiofiles>=24.1.0",
|
||||
"aiohttp>=3.14.1",
|
||||
"aiosqlite>=0.20.0",
|
||||
"argon2-cffi>=25.1.0",
|
||||
"asyncpg>=0.30.0",
|
||||
"chardet>=5.0.0",
|
||||
"colorlog>=6.9.0",
|
||||
"dashscope>=1.23.2",
|
||||
"deepagents>=0.6.7",
|
||||
"docling>=2.68.0",
|
||||
"docx2txt>=0.9",
|
||||
"httpx>=0.27.0",
|
||||
"igraph>=0.11.8",
|
||||
"json-repair>=0.54.0",
|
||||
"langchain>=1.3.9",
|
||||
"langchain-community>=0.4",
|
||||
"langchain-core>=1.2.5",
|
||||
"langchain-deepseek>=1.0",
|
||||
"langchain-mcp-adapters>=0.1.9",
|
||||
"langchain-openai>=1.0.2",
|
||||
"langchain-tavily>=0.2.13",
|
||||
"langchain-text-splitters>=1.0",
|
||||
"langfuse>=4.0.0",
|
||||
"langgraph>=1.0.1",
|
||||
"langgraph-checkpoint-postgres>=2.0.0",
|
||||
"langgraph-checkpoint-sqlite>=3.0",
|
||||
"langgraph-cli[inmem]>=0.4",
|
||||
"langsmith>=0.4",
|
||||
"llama-index>=0.14",
|
||||
"llama-index-readers-file>=0.4.7",
|
||||
"loguru>=0.7.3",
|
||||
"markdownify>=1.1.0",
|
||||
"mcp>=1.20",
|
||||
"minio>=7.2.7",
|
||||
"neo4j>=5.28.1",
|
||||
"networkx>=3.5",
|
||||
"openai>=1.109",
|
||||
"opencv-python-headless>=4.11.0.86",
|
||||
"onnxruntime>=1.20.0",
|
||||
"Pillow>=10.5.0",
|
||||
"psycopg[binary,pool]>=3.3.3",
|
||||
"pyjwt>=2.13.0",
|
||||
"pymilvus>=2.5.8",
|
||||
"pymupdf>=1.25.5",
|
||||
"pymysql>=1.1.0",
|
||||
"pypinyin>=0.55.0",
|
||||
"python-dotenv>=1.1.0",
|
||||
"python-multipart>=0.0.31",
|
||||
"pyyaml>=6.0.2",
|
||||
"rapidocr>=3.0.0",
|
||||
"readability-lxml>=0.8.1",
|
||||
"redis>=5.2.0",
|
||||
"requests>=2.34.0",
|
||||
"rich>=13.7.1",
|
||||
"sqlalchemy[asyncio]>=2.0.0",
|
||||
"tabulate>=0.9.0",
|
||||
"tavily-python>=0.7.0",
|
||||
"tenacity>=8.0.0",
|
||||
"tomli",
|
||||
"tomli-w",
|
||||
"torch>=2.12.1",
|
||||
"torchvision==0.27.1",
|
||||
"tqdm>=4.66.4",
|
||||
"typer>=0.16.0",
|
||||
"unstructured>=0.17.2",
|
||||
"wcmatch>=8.0.0",
|
||||
# Markdown 解析语义切分相关依赖
|
||||
"markdown-it-py>=3.0.0",
|
||||
"mdit-py-plugins>=0.4.0",
|
||||
"scikit-learn>=1.3.0",
|
||||
"nltk>=3.8.1",
|
||||
"beautifulsoup4>=4.12.0",
|
||||
]
|
||||
|
||||
[tool.setuptools]
|
||||
include-package-data = true
|
||||
|
||||
[tool.setuptools.packages.find]
|
||||
where = ["."]
|
||||
include = ["yuxi*"]
|
||||
|
||||
|
||||
[tool.uv]
|
||||
constraint-dependencies = [
|
||||
"cryptography>=48.0.1",
|
||||
"langchain-anthropic>=1.4.6",
|
||||
"pypdf>=6.13.3",
|
||||
"starlette>=1.3.1",
|
||||
]
|
||||
|
||||
[[tool.uv.index]]
|
||||
name = "tsinghua"
|
||||
url = "https://pypi.tuna.tsinghua.edu.cn/simple"
|
||||
default = true
|
||||
|
||||
|
||||
[[tool.uv.index]]
|
||||
name = "pytorch-cpu"
|
||||
url = "https://download.pytorch.org/whl/cpu"
|
||||
explicit = true
|
||||
|
||||
[tool.uv.sources]
|
||||
yuxi = { path = "package", editable = true }
|
||||
torch = { index = "pytorch-cpu" }
|
||||
torchvision = { index = "pytorch-cpu" }
|
||||
Generated
+5472
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,33 @@
|
||||
from dotenv import load_dotenv
|
||||
|
||||
load_dotenv(".env", override=True)
|
||||
|
||||
from concurrent.futures import ThreadPoolExecutor # noqa: E402
|
||||
from importlib import import_module # noqa: E402
|
||||
|
||||
from yuxi.config import config as config # noqa: E402
|
||||
|
||||
try:
|
||||
from importlib.metadata import version
|
||||
|
||||
__version__ = version("yuxi")
|
||||
except Exception:
|
||||
__version__ = "unknown"
|
||||
|
||||
executor = ThreadPoolExecutor() # noqa: E402
|
||||
|
||||
|
||||
def get_version():
|
||||
"""Return the Yuxi version."""
|
||||
return __version__
|
||||
|
||||
|
||||
def __getattr__(name: str):
|
||||
if name == "knowledge_base":
|
||||
knowledge = import_module("yuxi.knowledge")
|
||||
return getattr(knowledge, name)
|
||||
raise AttributeError(f"module 'yuxi' has no attribute {name!r}")
|
||||
|
||||
|
||||
def __dir__():
|
||||
return sorted(set(globals()) | {"knowledge_base"})
|
||||
@@ -0,0 +1,29 @@
|
||||
# Base classes - 核心基类
|
||||
from yuxi.agents.base import BaseAgent
|
||||
|
||||
# 从 buildin 模块导入 agent_manager
|
||||
from yuxi.agents.context import BaseContext
|
||||
|
||||
# MCP - Agent 层统一入口(自动过滤 disabled_tools)
|
||||
from yuxi.agents.mcp.service import get_enabled_mcp_tools
|
||||
|
||||
# Model utilities - 模型加载
|
||||
from yuxi.agents.models import load_chat_model, resolve_chat_model_spec
|
||||
from yuxi.agents.state import BaseState
|
||||
|
||||
# Tools - 核心工具函数
|
||||
from yuxi.agents.toolkits.utils import get_tool_info
|
||||
|
||||
__all__ = [
|
||||
# Base classes
|
||||
"BaseAgent",
|
||||
"BaseContext",
|
||||
"BaseState",
|
||||
# Model utilities
|
||||
"load_chat_model",
|
||||
"resolve_chat_model_spec",
|
||||
# Core tools
|
||||
"get_tool_info",
|
||||
# Core MCP
|
||||
"get_enabled_mcp_tools",
|
||||
]
|
||||
@@ -0,0 +1,49 @@
|
||||
from deepagents.backends import CompositeBackend, StateBackend
|
||||
|
||||
from .composite import create_agent_composite_backend, create_agent_filesystem_middleware
|
||||
from .knowledge_base_backend import resolve_visible_knowledge_bases_for_context
|
||||
from .sandbox import (
|
||||
SKILLS_PATH,
|
||||
USER_DATA_PATH,
|
||||
VIRTUAL_PATH_PREFIX,
|
||||
ProvisionerSandboxBackend,
|
||||
ProvisionerSandboxProvider,
|
||||
SandboxConnection,
|
||||
get_sandbox_provider,
|
||||
init_sandbox_provider,
|
||||
resolve_virtual_path,
|
||||
sandbox_id_for_thread,
|
||||
sandbox_outputs_dir,
|
||||
sandbox_uploads_dir,
|
||||
sandbox_user_data_dir,
|
||||
sandbox_workspace_dir,
|
||||
shutdown_sandbox_provider,
|
||||
virtual_path_for_thread_file,
|
||||
)
|
||||
from .skills_backend import SelectedSkillsReadonlyBackend
|
||||
|
||||
__all__ = [
|
||||
"CompositeBackend",
|
||||
"StateBackend",
|
||||
"SelectedSkillsReadonlyBackend",
|
||||
"create_agent_composite_backend",
|
||||
"create_agent_filesystem_middleware",
|
||||
"ProvisionerSandboxBackend",
|
||||
"ProvisionerSandboxProvider",
|
||||
"SandboxConnection",
|
||||
"get_sandbox_provider",
|
||||
"init_sandbox_provider",
|
||||
"shutdown_sandbox_provider",
|
||||
"resolve_virtual_path",
|
||||
"resolve_visible_knowledge_bases_for_context",
|
||||
"virtual_path_for_thread_file",
|
||||
"sandbox_id_for_thread",
|
||||
"sandbox_user_data_dir",
|
||||
"sandbox_workspace_dir",
|
||||
"sandbox_uploads_dir",
|
||||
"sandbox_outputs_dir",
|
||||
# Config paths
|
||||
"VIRTUAL_PATH_PREFIX",
|
||||
"USER_DATA_PATH",
|
||||
"SKILLS_PATH",
|
||||
]
|
||||
@@ -0,0 +1,201 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
from deepagents.backends.composite import (
|
||||
CompositeBackend,
|
||||
_remap_file_info_path,
|
||||
_route_for_path,
|
||||
_strip_route_from_pattern,
|
||||
)
|
||||
from deepagents.backends.protocol import FileInfo, GlobResult
|
||||
from deepagents.middleware.filesystem import FilesystemMiddleware
|
||||
|
||||
from yuxi.agents.skills.service import normalize_string_list
|
||||
from yuxi.utils.paths import VIRTUAL_PATH_CONVERSATION_HISTORY, VIRTUAL_PATH_LARGE_TOOL_RESULTS, VIRTUAL_PATH_OUTPUTS
|
||||
|
||||
from .sandbox import ProvisionerSandboxBackend
|
||||
from .skills_backend import SelectedSkillsReadonlyBackend
|
||||
|
||||
_TOOL_RESULT_EVICTION_EXEMPT_TOOLS = frozenset({"read_file", "open_kb_document"})
|
||||
|
||||
|
||||
def _coerce_glob_result(result) -> GlobResult:
|
||||
if isinstance(result, GlobResult):
|
||||
return result
|
||||
return GlobResult(matches=result or [])
|
||||
|
||||
|
||||
class CustomCompositeBackend(CompositeBackend):
|
||||
"""修复 glob 路由逻辑的 CompositeBackend。"""
|
||||
|
||||
def glob(self, pattern: str, path: str = "/") -> GlobResult:
|
||||
backend, backend_path, route_prefix = _route_for_path(
|
||||
default=self.default,
|
||||
sorted_routes=self.sorted_routes,
|
||||
path=path,
|
||||
)
|
||||
if route_prefix is not None:
|
||||
result = _coerce_glob_result(backend.glob(pattern, backend_path))
|
||||
if result.error:
|
||||
return result
|
||||
return GlobResult(matches=[_remap_file_info_path(fi, route_prefix) for fi in (result.matches or [])])
|
||||
|
||||
if path is None or path == "/":
|
||||
results: list[FileInfo] = []
|
||||
default_result = _coerce_glob_result(self.default.glob(pattern, path))
|
||||
if default_result.error:
|
||||
return default_result
|
||||
results.extend(default_result.matches or [])
|
||||
for route_prefix, backend in self.routes.items():
|
||||
route_pattern = _strip_route_from_pattern(pattern, route_prefix)
|
||||
result = _coerce_glob_result(backend.glob(route_pattern, "/"))
|
||||
if result.error:
|
||||
return result
|
||||
results.extend(_remap_file_info_path(fi, route_prefix) for fi in (result.matches or []))
|
||||
results.sort(key=lambda x: x.get("path", ""))
|
||||
return GlobResult(matches=results)
|
||||
|
||||
return _coerce_glob_result(self.default.glob(pattern, path))
|
||||
|
||||
async def aglob(self, pattern: str, path: str = "/") -> GlobResult:
|
||||
backend, backend_path, route_prefix = _route_for_path(
|
||||
default=self.default,
|
||||
sorted_routes=self.sorted_routes,
|
||||
path=path,
|
||||
)
|
||||
if route_prefix is not None:
|
||||
result = _coerce_glob_result(await backend.aglob(pattern, backend_path))
|
||||
if result.error:
|
||||
return result
|
||||
return GlobResult(matches=[_remap_file_info_path(fi, route_prefix) for fi in (result.matches or [])])
|
||||
|
||||
if path is None or path == "/":
|
||||
results: list[FileInfo] = []
|
||||
default_result = _coerce_glob_result(await self.default.aglob(pattern, path))
|
||||
if default_result.error:
|
||||
return default_result
|
||||
results.extend(default_result.matches or [])
|
||||
for route_prefix, backend in self.routes.items():
|
||||
route_pattern = _strip_route_from_pattern(pattern, route_prefix)
|
||||
result = _coerce_glob_result(await backend.aglob(route_pattern, "/"))
|
||||
if result.error:
|
||||
return result
|
||||
results.extend(_remap_file_info_path(fi, route_prefix) for fi in (result.matches or []))
|
||||
results.sort(key=lambda x: x.get("path", ""))
|
||||
return GlobResult(matches=results)
|
||||
|
||||
return _coerce_glob_result(await self.default.aglob(pattern, path))
|
||||
|
||||
|
||||
class YuxiFilesystemMiddleware(FilesystemMiddleware):
|
||||
"""Filesystem middleware that budgets large tool outputs before they hit model context."""
|
||||
|
||||
def wrap_tool_call(self, request, handler):
|
||||
tool_result = handler(request)
|
||||
|
||||
if self._tool_token_limit_before_evict is None:
|
||||
return tool_result
|
||||
if request.tool_call["name"] in _TOOL_RESULT_EVICTION_EXEMPT_TOOLS:
|
||||
return tool_result
|
||||
|
||||
return self._intercept_large_tool_result(tool_result, request.runtime)
|
||||
|
||||
async def awrap_tool_call(self, request, handler):
|
||||
tool_result = await handler(request)
|
||||
|
||||
if self._tool_token_limit_before_evict is None:
|
||||
return tool_result
|
||||
if request.tool_call["name"] in _TOOL_RESULT_EVICTION_EXEMPT_TOOLS:
|
||||
return tool_result
|
||||
|
||||
return await self._aintercept_large_tool_result(tool_result, request.runtime)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _BackendScope:
|
||||
thread_id: str
|
||||
uid: str
|
||||
readable_skills: list[str]
|
||||
file_thread_id: str
|
||||
skills_thread_id: str
|
||||
|
||||
@classmethod
|
||||
def from_runtime(cls, runtime) -> _BackendScope:
|
||||
config = getattr(runtime, "config", None)
|
||||
configurable = config.get("configurable", {}) if isinstance(config, dict) else {}
|
||||
context = getattr(runtime, "context", None)
|
||||
state = getattr(runtime, "state", None)
|
||||
return cls.from_sources(
|
||||
configurable if isinstance(configurable, dict) else {},
|
||||
context,
|
||||
state if isinstance(state, dict) else {},
|
||||
readable_skills_source=context,
|
||||
error_context="runtime configurable context",
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_sources(cls, *sources, readable_skills_source, error_context: str) -> _BackendScope:
|
||||
def string_value(key: str) -> str | None:
|
||||
for source in sources:
|
||||
value = source.get(key) if isinstance(source, dict) else getattr(source, key, None)
|
||||
if isinstance(value, str) and value.strip():
|
||||
return value.strip()
|
||||
return None
|
||||
|
||||
thread_id = string_value("thread_id")
|
||||
if not thread_id:
|
||||
raise ValueError(f"thread_id is required in {error_context}")
|
||||
|
||||
uid = string_value("uid")
|
||||
if not uid:
|
||||
raise ValueError(f"uid is required in {error_context}")
|
||||
|
||||
selected = getattr(readable_skills_source, "_readable_skills", [])
|
||||
return cls(
|
||||
thread_id=thread_id,
|
||||
uid=uid,
|
||||
readable_skills=normalize_string_list(selected if isinstance(selected, list) else []),
|
||||
file_thread_id=string_value("file_thread_id") or thread_id,
|
||||
skills_thread_id=string_value("skills_thread_id") or thread_id,
|
||||
)
|
||||
|
||||
def create_backend(self) -> CompositeBackend:
|
||||
return CustomCompositeBackend(
|
||||
default=ProvisionerSandboxBackend(
|
||||
thread_id=self.thread_id,
|
||||
uid=self.uid,
|
||||
readable_skills=self.readable_skills,
|
||||
file_thread_id=self.file_thread_id,
|
||||
skills_thread_id=self.skills_thread_id,
|
||||
),
|
||||
routes={
|
||||
"/skills/": SelectedSkillsReadonlyBackend(selected_slugs=self.readable_skills),
|
||||
},
|
||||
artifacts_root=VIRTUAL_PATH_OUTPUTS,
|
||||
)
|
||||
|
||||
|
||||
def create_agent_composite_backend(runtime) -> CompositeBackend:
|
||||
return _BackendScope.from_runtime(runtime).create_backend()
|
||||
|
||||
|
||||
def create_agent_filesystem_middleware(
|
||||
tool_token_limit_before_evict: int | None = None,
|
||||
*,
|
||||
context=None,
|
||||
) -> FilesystemMiddleware:
|
||||
backend = create_agent_composite_backend
|
||||
if context is not None:
|
||||
backend = _BackendScope.from_sources(
|
||||
context,
|
||||
readable_skills_source=context,
|
||||
error_context="runtime context",
|
||||
).create_backend()
|
||||
middleware = YuxiFilesystemMiddleware(
|
||||
backend=backend,
|
||||
tool_token_limit_before_evict=tool_token_limit_before_evict,
|
||||
)
|
||||
middleware._large_tool_results_prefix = VIRTUAL_PATH_LARGE_TOOL_RESULTS
|
||||
middleware._conversation_history_prefix = VIRTUAL_PATH_CONVERSATION_HISTORY
|
||||
return middleware
|
||||
@@ -0,0 +1,22 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
|
||||
async def resolve_visible_knowledge_bases_for_context(context) -> list[dict[str, Any]]:
|
||||
from yuxi import knowledge_base
|
||||
|
||||
uid = getattr(context, "uid", None)
|
||||
if not uid:
|
||||
setattr(context, "_visible_knowledge_bases", [])
|
||||
return []
|
||||
|
||||
result = await knowledge_base.get_databases_by_uid(str(uid))
|
||||
databases = result.get("databases") or []
|
||||
enabled_knowledges = getattr(context, "knowledges", None)
|
||||
if enabled_knowledges is not None:
|
||||
enabled_ids = {str(value).strip() for value in enabled_knowledges if str(value).strip()}
|
||||
databases = [db for db in databases if str(db.get("kb_id") or "").strip() in enabled_ids]
|
||||
|
||||
setattr(context, "_visible_knowledge_bases", databases)
|
||||
return databases
|
||||
@@ -0,0 +1,47 @@
|
||||
from .backend import ProvisionerSandboxBackend
|
||||
from .paths import (
|
||||
VIRTUAL_PATH_PREFIX,
|
||||
ensure_thread_dirs,
|
||||
ensure_workspace_default_files,
|
||||
resolve_virtual_path,
|
||||
sandbox_outputs_dir,
|
||||
sandbox_uploads_dir,
|
||||
sandbox_user_data_dir,
|
||||
sandbox_workspace_agents_prompt_file,
|
||||
sandbox_workspace_dir,
|
||||
virtual_path_for_thread_file,
|
||||
)
|
||||
from .provider import (
|
||||
ProvisionerSandboxProvider,
|
||||
SandboxConnection,
|
||||
get_sandbox_provider,
|
||||
init_sandbox_provider,
|
||||
sandbox_id_for_thread,
|
||||
shutdown_sandbox_provider,
|
||||
)
|
||||
|
||||
# Sandbox-visible paths for viewer/filesystem services.
|
||||
USER_DATA_PATH = VIRTUAL_PATH_PREFIX
|
||||
SKILLS_PATH = "/home/gem/skills"
|
||||
|
||||
__all__ = [
|
||||
"ProvisionerSandboxBackend",
|
||||
"ProvisionerSandboxProvider",
|
||||
"SandboxConnection",
|
||||
"VIRTUAL_PATH_PREFIX",
|
||||
"ensure_thread_dirs",
|
||||
"ensure_workspace_default_files",
|
||||
"get_sandbox_provider",
|
||||
"init_sandbox_provider",
|
||||
"resolve_virtual_path",
|
||||
"sandbox_id_for_thread",
|
||||
"sandbox_outputs_dir",
|
||||
"sandbox_uploads_dir",
|
||||
"sandbox_user_data_dir",
|
||||
"sandbox_workspace_agents_prompt_file",
|
||||
"sandbox_workspace_dir",
|
||||
"shutdown_sandbox_provider",
|
||||
"virtual_path_for_thread_file",
|
||||
"USER_DATA_PATH",
|
||||
"SKILLS_PATH",
|
||||
]
|
||||
@@ -0,0 +1,629 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import uuid
|
||||
from contextlib import suppress
|
||||
from datetime import datetime
|
||||
from pathlib import PurePosixPath
|
||||
from typing import Any
|
||||
|
||||
from deepagents.backends.protocol import (
|
||||
EditResult,
|
||||
ExecuteResponse,
|
||||
FileDownloadResponse,
|
||||
FileInfo,
|
||||
FileUploadResponse,
|
||||
GlobResult,
|
||||
GrepMatch,
|
||||
GrepResult,
|
||||
LsResult,
|
||||
ReadResult,
|
||||
WriteResult,
|
||||
)
|
||||
from deepagents.backends.sandbox import MAX_BINARY_BYTES, BaseSandbox
|
||||
from deepagents.backends.utils import _get_file_type
|
||||
|
||||
from yuxi import config as conf
|
||||
from yuxi.agents.skills.service import sync_thread_readable_skills
|
||||
from yuxi.utils.logging_config import logger
|
||||
from yuxi.utils.paths import (
|
||||
OUTPUTS_DIR_NAME,
|
||||
UPLOADS_DIR_NAME,
|
||||
VIRTUAL_PATH_PREFIX,
|
||||
VIRTUAL_SKILLS_PATH,
|
||||
WORKSPACE_DIR_NAME,
|
||||
)
|
||||
|
||||
from .provider import get_sandbox_provider, sandbox_id_for_thread
|
||||
|
||||
_USER_DATA_ROOT = "/" + VIRTUAL_PATH_PREFIX.strip("/")
|
||||
_WORKSPACE_ROOT = f"{_USER_DATA_ROOT}/{WORKSPACE_DIR_NAME}"
|
||||
_UPLOADS_ROOT = f"{_USER_DATA_ROOT}/{UPLOADS_DIR_NAME}"
|
||||
_OUTPUTS_ROOT = f"{_USER_DATA_ROOT}/{OUTPUTS_DIR_NAME}"
|
||||
_SKILLS_ROOT = "/" + VIRTUAL_SKILLS_PATH.strip("/")
|
||||
_READABLE_ROOTS = (_USER_DATA_ROOT, _SKILLS_ROOT)
|
||||
_WRITABLE_ROOTS = (_WORKSPACE_ROOT, _OUTPUTS_ROOT)
|
||||
_BINARY_PREVIEW_TOO_LARGE_ERROR = f"Binary file exceeds maximum preview size of {MAX_BINARY_BYTES} bytes"
|
||||
|
||||
|
||||
def _normalize_path(path: str) -> str:
|
||||
raw = str(path or "").strip()
|
||||
if not raw:
|
||||
raise ValueError("path is required")
|
||||
if not raw.startswith("/"):
|
||||
raise ValueError("path must start with /")
|
||||
pure = PurePosixPath(raw)
|
||||
if ".." in pure.parts:
|
||||
raise ValueError("path traversal is not allowed")
|
||||
return str(pure)
|
||||
|
||||
|
||||
def _is_same_or_child(path: str, root: str) -> bool:
|
||||
root = root.rstrip("/") or "/"
|
||||
if root == "/":
|
||||
return path == "/" or path.startswith("/")
|
||||
return path == root or path.startswith(f"{root}/")
|
||||
|
||||
|
||||
def _path_overlaps_root(path: str, root: str) -> bool:
|
||||
return _is_same_or_child(path, root) or _is_same_or_child(root, path)
|
||||
|
||||
|
||||
def _can_read_path(path: str) -> bool:
|
||||
return any(_is_same_or_child(path, root) for root in _READABLE_ROOTS)
|
||||
|
||||
|
||||
def _can_list_path(path: str) -> bool:
|
||||
return any(_path_overlaps_root(path, root) for root in _READABLE_ROOTS)
|
||||
|
||||
|
||||
def _can_write_path(path: str) -> bool:
|
||||
return any(_is_same_or_child(path, root) for root in _WRITABLE_ROOTS)
|
||||
|
||||
|
||||
def _readable_search_paths(path: str) -> list[str]:
|
||||
if _can_read_path(path):
|
||||
return [path]
|
||||
return [root for root in _READABLE_ROOTS if _is_same_or_child(root, path)]
|
||||
|
||||
|
||||
def _glob_for_search_root(pattern: str, root: str) -> str:
|
||||
bare_pattern = str(pattern or "*").lstrip("/")
|
||||
bare_root = root.strip("/")
|
||||
if bare_pattern == bare_root:
|
||||
return "*"
|
||||
root_prefix = f"{bare_root}/"
|
||||
if bare_pattern.startswith(root_prefix):
|
||||
return bare_pattern[len(root_prefix) :] or "*"
|
||||
return pattern
|
||||
|
||||
|
||||
def _filter_readable_infos(infos: list[FileInfo]) -> list[FileInfo]:
|
||||
result: list[FileInfo] = []
|
||||
for info in infos:
|
||||
try:
|
||||
path = _normalize_path(info.get("path", ""))
|
||||
except ValueError:
|
||||
continue
|
||||
if _can_list_path(path):
|
||||
result.append(info)
|
||||
return result
|
||||
|
||||
|
||||
def _filter_readable_matches(matches: list[GrepMatch]) -> list[GrepMatch]:
|
||||
result: list[GrepMatch] = []
|
||||
for match in matches:
|
||||
try:
|
||||
path = _normalize_path(match.get("path", ""))
|
||||
except ValueError:
|
||||
continue
|
||||
if _can_read_path(path):
|
||||
result.append(match)
|
||||
return result
|
||||
|
||||
|
||||
def _permission_error(operation: str, path: str) -> str:
|
||||
return f"permission denied for {operation} on '{path}'"
|
||||
|
||||
|
||||
def _describe_read_error(file_path: str, exc: Exception) -> str:
|
||||
if isinstance(exc, FileNotFoundError):
|
||||
return f"Error: File '{file_path}' not found"
|
||||
if isinstance(exc, IsADirectoryError):
|
||||
return f"Error: Path '{file_path}' is a directory"
|
||||
if isinstance(exc, PermissionError):
|
||||
return f"Error: Access denied for '{file_path}'"
|
||||
if isinstance(exc, ValueError):
|
||||
return f"Error: Invalid path '{file_path}': {exc}"
|
||||
detail = str(exc).strip()
|
||||
if detail:
|
||||
return f"Error: Failed to read '{file_path}': {detail}"
|
||||
return f"Error: Failed to read '{file_path}'"
|
||||
|
||||
|
||||
def _is_missing_file_error(exc: Exception) -> bool:
|
||||
if isinstance(exc, FileNotFoundError):
|
||||
return True
|
||||
|
||||
status_code = getattr(exc, "status_code", None)
|
||||
response = getattr(exc, "response", None)
|
||||
if status_code == 404 or getattr(response, "status_code", None) == 404:
|
||||
return True
|
||||
|
||||
detail = str(exc).lower()
|
||||
return "status_code: 404" in detail or "file does not exist" in detail
|
||||
|
||||
|
||||
def _looks_like_binary(content: bytes) -> bool:
|
||||
if not content:
|
||||
return False
|
||||
if b"\x00" in content:
|
||||
return True
|
||||
try:
|
||||
content.decode("utf-8")
|
||||
return False
|
||||
except UnicodeDecodeError:
|
||||
return True
|
||||
|
||||
|
||||
def _is_utf8_decode_failure(exc: Exception) -> bool:
|
||||
detail = str(exc).lower()
|
||||
return "utf-8" in detail and "can't decode" in detail
|
||||
|
||||
|
||||
class ProvisionerSandboxBackend(BaseSandbox):
|
||||
def __init__(
|
||||
self,
|
||||
thread_id: str,
|
||||
*,
|
||||
uid: str,
|
||||
readable_skills: list[str] | None = None,
|
||||
file_thread_id: str | None = None,
|
||||
skills_thread_id: str | None = None,
|
||||
):
|
||||
self._thread_id = str(thread_id or "").strip()
|
||||
if not self._thread_id:
|
||||
raise ValueError("thread_id is required for ProvisionerSandboxBackend")
|
||||
self._file_thread_id = str(file_thread_id or self._thread_id).strip()
|
||||
if not self._file_thread_id:
|
||||
raise ValueError("file_thread_id is required for ProvisionerSandboxBackend")
|
||||
self._skills_thread_id = str(skills_thread_id or self._thread_id).strip()
|
||||
if not self._skills_thread_id:
|
||||
raise ValueError("skills_thread_id is required for ProvisionerSandboxBackend")
|
||||
self._uid = str(uid or "").strip()
|
||||
if not self._uid:
|
||||
raise ValueError("uid is required for ProvisionerSandboxBackend")
|
||||
|
||||
self._readable_skills = list(readable_skills or [])
|
||||
self._provider = get_sandbox_provider()
|
||||
self._id = sandbox_id_for_thread(self._file_thread_id, self._skills_thread_id, uid=self._uid)
|
||||
self._client: Any | None = None
|
||||
self._client_url: str | None = None
|
||||
self._command_timeout_seconds = int(getattr(conf, "sandbox_exec_timeout_seconds", 180))
|
||||
self._max_output_bytes = int(getattr(conf, "sandbox_max_output_bytes", 262_144))
|
||||
|
||||
@property
|
||||
def id(self) -> str:
|
||||
return self._id
|
||||
|
||||
def _build_client(self, sandbox_url: str):
|
||||
try:
|
||||
from agent_sandbox import Sandbox as AgentSandboxClient
|
||||
except Exception as exc: # noqa: BLE001
|
||||
raise RuntimeError(
|
||||
"agent-sandbox is required. Install dependency `agent-sandbox` in the docker image."
|
||||
) from exc
|
||||
|
||||
return AgentSandboxClient(base_url=sandbox_url, timeout=self._command_timeout_seconds)
|
||||
|
||||
def _get_client(self) -> Any:
|
||||
sync_thread_readable_skills(self._skills_thread_id, self._readable_skills)
|
||||
connection = self._provider.get(
|
||||
self._thread_id,
|
||||
uid=self._uid,
|
||||
create_if_missing=True,
|
||||
file_thread_id=self._file_thread_id,
|
||||
skills_thread_id=self._skills_thread_id,
|
||||
)
|
||||
if connection is None:
|
||||
raise RuntimeError(f"sandbox is unavailable for thread {self._thread_id}")
|
||||
|
||||
if self._client is None or self._client_url != connection.sandbox_url:
|
||||
self._client = self._build_client(connection.sandbox_url)
|
||||
self._client_url = connection.sandbox_url
|
||||
|
||||
return self._client
|
||||
|
||||
def _read_binary(self, path: str, offset: int = 0, limit: int | None = None) -> bytes:
|
||||
"""Read file content from the sandbox file API and normalize it to bytes.
|
||||
|
||||
The underlying API returns plain text by default and may include an
|
||||
explicit `encoding="base64"` marker for binary payloads. This helper is
|
||||
the single normalization point used by read(), edit(), and download_files().
|
||||
"""
|
||||
start_line = max(0, int(offset))
|
||||
end_line = start_line + int(limit) if limit is not None else None
|
||||
|
||||
result = self._get_client().file.read_file(
|
||||
file=path,
|
||||
start_line=start_line,
|
||||
end_line=end_line,
|
||||
)
|
||||
|
||||
content = result.data.content
|
||||
if content is None:
|
||||
return b""
|
||||
if isinstance(content, bytes):
|
||||
return content
|
||||
if not isinstance(content, str):
|
||||
return str(content).encode("utf-8")
|
||||
|
||||
encoding = getattr(result.data, "encoding", None)
|
||||
if isinstance(encoding, str) and encoding.lower() == "base64":
|
||||
return base64.b64decode(content, validate=True)
|
||||
return content.encode("utf-8")
|
||||
|
||||
def _file_size_bytes(self, path: str) -> int:
|
||||
path_b64 = base64.b64encode(path.encode("utf-8")).decode("ascii")
|
||||
command = (
|
||||
'python3 -c "'
|
||||
"import base64, os, stat; "
|
||||
f"path = base64.b64decode('{path_b64}').decode('utf-8'); "
|
||||
"st = os.stat(path); "
|
||||
"print(st.st_size if stat.S_ISREG(st.st_mode) else -1)"
|
||||
'"'
|
||||
)
|
||||
result = self.execute(command)
|
||||
if result.exit_code not in (0, None):
|
||||
detail = (result.output or "").strip()
|
||||
raise RuntimeError(detail or f"failed to stat '{path}'")
|
||||
|
||||
output = (result.output or "").strip().splitlines()
|
||||
try:
|
||||
size = int(output[-1])
|
||||
except (IndexError, ValueError) as exc:
|
||||
raise RuntimeError(f"failed to stat '{path}'") from exc
|
||||
if size < 0:
|
||||
raise IsADirectoryError(path)
|
||||
return size
|
||||
|
||||
def _read_file_base64(self, path: str) -> str:
|
||||
path_b64 = base64.b64encode(path.encode("utf-8")).decode("ascii")
|
||||
output_path = f"/tmp/yuxi-read-file-{uuid.uuid4().hex}.b64"
|
||||
output_path_b64 = base64.b64encode(output_path.encode("utf-8")).decode("ascii")
|
||||
command = (
|
||||
'python3 -c "'
|
||||
"import base64; "
|
||||
f"path = base64.b64decode('{path_b64}').decode('utf-8'); "
|
||||
f"output_path = base64.b64decode('{output_path_b64}').decode('utf-8'); "
|
||||
"open(output_path, 'w').write(base64.b64encode(open(path, 'rb').read()).decode('ascii'))"
|
||||
'"'
|
||||
)
|
||||
client = self._get_client()
|
||||
try:
|
||||
result = client.shell.exec_command(
|
||||
command=command,
|
||||
timeout=self._command_timeout_seconds,
|
||||
truncate=False,
|
||||
)
|
||||
output = result.data.output or ""
|
||||
if result.data.exit_code not in (0, None):
|
||||
raise RuntimeError(output.strip() or f"failed to read '{path}'")
|
||||
|
||||
content = self._read_binary(output_path).decode("ascii").strip()
|
||||
base64.b64decode(content, validate=True)
|
||||
return content
|
||||
finally:
|
||||
with suppress(Exception):
|
||||
client.shell.exec_command(command=f"rm -f {output_path}", timeout=10)
|
||||
|
||||
def _read_base64_file(self, path: str) -> ReadResult:
|
||||
if self._file_size_bytes(path) > MAX_BINARY_BYTES:
|
||||
return ReadResult(error=_BINARY_PREVIEW_TOO_LARGE_ERROR)
|
||||
return ReadResult(file_data={"content": self._read_file_base64(path), "encoding": "base64"})
|
||||
|
||||
def read(
|
||||
self,
|
||||
file_path: str,
|
||||
offset: int = 0,
|
||||
limit: int = 2000,
|
||||
) -> ReadResult:
|
||||
"""Read allowed file content via the sandbox file API."""
|
||||
try:
|
||||
normalized_path = _normalize_path(file_path)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
return ReadResult(error=f"Invalid path '{file_path}': {exc}")
|
||||
if not _can_read_path(normalized_path):
|
||||
return ReadResult(error=_permission_error("read", normalized_path))
|
||||
|
||||
try:
|
||||
if _get_file_type(normalized_path) != "text":
|
||||
return self._read_base64_file(normalized_path)
|
||||
|
||||
try:
|
||||
content = self._read_binary(normalized_path, offset=offset, limit=limit)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
if not _is_utf8_decode_failure(exc):
|
||||
raise
|
||||
return self._read_base64_file(normalized_path)
|
||||
|
||||
if not _looks_like_binary(content):
|
||||
return ReadResult(file_data={"content": content.decode("utf-8"), "encoding": "utf-8"})
|
||||
|
||||
return self._read_base64_file(normalized_path)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
error = _describe_read_error(file_path, exc)
|
||||
return ReadResult(error=error.removeprefix("Error: "))
|
||||
|
||||
def execute(self, command: str, *, timeout: int | None = None) -> ExecuteResponse:
|
||||
"""Execute a shell command in the sandbox.
|
||||
|
||||
Output is normalized to text and truncated to the configured maximum
|
||||
payload size before being returned.
|
||||
"""
|
||||
try:
|
||||
kwargs: dict[str, Any] = {"command": command}
|
||||
if timeout is not None:
|
||||
kwargs["timeout"] = timeout
|
||||
result = self._get_client().shell.exec_command(**kwargs)
|
||||
|
||||
output = result.data.output or ""
|
||||
exit_code = result.data.exit_code
|
||||
|
||||
truncated = False
|
||||
encoded = output.encode("utf-8", errors="ignore")
|
||||
if len(encoded) > self._max_output_bytes:
|
||||
output = encoded[: self._max_output_bytes].decode("utf-8", errors="ignore")
|
||||
truncated = True
|
||||
|
||||
return ExecuteResponse(
|
||||
output=output,
|
||||
exit_code=exit_code if isinstance(exit_code, int) else None,
|
||||
truncated=truncated,
|
||||
)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
logger.error(f"Sandbox execute failed for thread {self._thread_id}: {exc}")
|
||||
return ExecuteResponse(output=f"Error: {exc}", exit_code=1, truncated=False)
|
||||
|
||||
def ls(self, path: str) -> LsResult:
|
||||
"""List direct children of an allowed sandbox path with lightweight metadata."""
|
||||
try:
|
||||
normalized_path = _normalize_path(path)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
return LsResult(error=f"Invalid path '{path}': {exc}")
|
||||
if not _can_list_path(normalized_path):
|
||||
return LsResult(error=_permission_error("read", normalized_path))
|
||||
|
||||
try:
|
||||
result = self._get_client().file.list_path(path=normalized_path, recursive=False, include_size=True)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
return LsResult(error=str(exc) or f"Failed to list '{path}'")
|
||||
|
||||
entries = result.data.files or []
|
||||
infos: list[FileInfo] = []
|
||||
for entry in entries:
|
||||
info: FileInfo = {"path": entry.path, "is_dir": entry.is_directory}
|
||||
size = entry.size
|
||||
if isinstance(size, int):
|
||||
info["size"] = size
|
||||
modified_time = entry.modified_time
|
||||
if modified_time:
|
||||
if isinstance(modified_time, str) and modified_time.isdigit():
|
||||
info["modified_at"] = datetime.fromtimestamp(int(modified_time)).isoformat()
|
||||
elif isinstance(modified_time, str):
|
||||
try:
|
||||
info["modified_at"] = datetime.fromisoformat(modified_time).isoformat()
|
||||
except ValueError:
|
||||
info["modified_at"] = modified_time
|
||||
elif isinstance(modified_time, (int, float)):
|
||||
info["modified_at"] = datetime.fromtimestamp(modified_time).isoformat()
|
||||
infos.append(info)
|
||||
return LsResult(entries=_filter_readable_infos(infos))
|
||||
|
||||
def write(self, file_path: str, content: str) -> WriteResult:
|
||||
"""Write a new text file.
|
||||
|
||||
This method is intentionally text-only. Binary payloads should go through
|
||||
upload_files(), which uses base64 encoding for the sandbox file API.
|
||||
"""
|
||||
try:
|
||||
normalized_path = _normalize_path(file_path)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
return WriteResult(error=f"Error: Invalid path '{file_path}': {exc}")
|
||||
if not _can_write_path(normalized_path):
|
||||
return WriteResult(error=f"Error: {_permission_error('write', normalized_path)}")
|
||||
if not isinstance(content, str):
|
||||
return WriteResult(error="Error: write() only supports text content; use upload_files() for binary data")
|
||||
try:
|
||||
self._read_binary(normalized_path)
|
||||
except Exception: # noqa: BLE001
|
||||
pass
|
||||
else:
|
||||
return WriteResult(error=f"Error: File '{file_path}' already exists")
|
||||
|
||||
try:
|
||||
result = self._get_client().file.write_file(file=normalized_path, content=content)
|
||||
if not result.success:
|
||||
return WriteResult(error=result.message or f"Failed to write file '{file_path}'")
|
||||
except Exception as exc: # noqa: BLE001
|
||||
return WriteResult(error=str(exc) or f"Failed to write file '{file_path}'")
|
||||
|
||||
return WriteResult(path=normalized_path)
|
||||
|
||||
def edit(
|
||||
self,
|
||||
file_path: str,
|
||||
old_string: str,
|
||||
new_string: str,
|
||||
replace_all: bool = False, # noqa: FBT001, FBT002
|
||||
) -> EditResult:
|
||||
"""Edit an existing text file by replacing string content.
|
||||
|
||||
This method operates on UTF-8-decoded text content only. Binary files
|
||||
are not supported here and should be handled via download/upload flows.
|
||||
"""
|
||||
try:
|
||||
normalized_path = _normalize_path(file_path)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
return EditResult(error=f"Error: Invalid path '{file_path}': {exc}")
|
||||
if not _can_write_path(normalized_path):
|
||||
return EditResult(error=f"Error: {_permission_error('write', normalized_path)}")
|
||||
|
||||
# Check if old_string exists
|
||||
try:
|
||||
text = self._read_binary(normalized_path).decode("utf-8", errors="replace")
|
||||
except Exception: # noqa: BLE001
|
||||
return EditResult(error=f"Error: File '{file_path}' not found")
|
||||
|
||||
count = text.count(old_string)
|
||||
if count == 0:
|
||||
return EditResult(error=f"Error: String not found in file: '{old_string}'")
|
||||
if count > 1 and not replace_all:
|
||||
return EditResult(
|
||||
error=(
|
||||
f"Error: String '{old_string}' appears multiple times. "
|
||||
"Use replace_all=True to replace all occurrences."
|
||||
)
|
||||
)
|
||||
|
||||
# Use str_replace_editor API
|
||||
replace_mode = "ALL" if replace_all else "FIRST"
|
||||
try:
|
||||
result = self._get_client().file.str_replace_editor(
|
||||
command="str_replace",
|
||||
path=normalized_path,
|
||||
old_str=old_string,
|
||||
new_str=new_string,
|
||||
replace_mode=replace_mode,
|
||||
)
|
||||
if not result.success:
|
||||
return EditResult(error=result.message or f"Error editing file '{file_path}'")
|
||||
except Exception as exc: # noqa: BLE001
|
||||
return EditResult(error=f"Error editing file: {exc}")
|
||||
|
||||
return EditResult(path=normalized_path, occurrences=count if replace_all else 1)
|
||||
|
||||
def grep(
|
||||
self,
|
||||
pattern: str,
|
||||
path: str | None = None,
|
||||
glob: str | None = None,
|
||||
) -> GrepResult:
|
||||
"""Search allowed sandbox paths for literal text."""
|
||||
try:
|
||||
normalized_path = _normalize_path(path or "/")
|
||||
except Exception as exc: # noqa: BLE001
|
||||
return GrepResult(error=f"Invalid path '{path or '/'}': {exc}")
|
||||
|
||||
search_paths = _readable_search_paths(normalized_path)
|
||||
if not search_paths:
|
||||
return GrepResult(error=_permission_error("read", normalized_path))
|
||||
|
||||
matches: list[GrepMatch] = []
|
||||
for search_path in search_paths:
|
||||
result = super().grep(pattern=pattern, path=search_path, glob=glob)
|
||||
if result.error:
|
||||
return result
|
||||
matches.extend(result.matches or [])
|
||||
return GrepResult(matches=_filter_readable_matches(matches))
|
||||
|
||||
def glob(self, pattern: str, path: str = "/") -> GlobResult:
|
||||
"""Return files matching a glob pattern under allowed sandbox paths."""
|
||||
try:
|
||||
normalized_path = _normalize_path(path)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
return GlobResult(error=f"Invalid path '{path}': {exc}")
|
||||
if ".." in PurePosixPath(str(pattern or "")).parts:
|
||||
return GlobResult(error="Invalid glob pattern: path traversal is not allowed")
|
||||
|
||||
search_paths = _readable_search_paths(normalized_path)
|
||||
if not search_paths:
|
||||
return GlobResult(error=_permission_error("read", normalized_path))
|
||||
|
||||
infos: list[FileInfo] = []
|
||||
for search_path in search_paths:
|
||||
try:
|
||||
result = self._get_client().file.find_files(
|
||||
path=search_path,
|
||||
glob=_glob_for_search_root(pattern, search_path),
|
||||
)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
return GlobResult(error=str(exc) or f"Failed to glob '{path}'")
|
||||
for file_path in result.data.files or []:
|
||||
infos.append({"path": file_path})
|
||||
infos = _filter_readable_infos(infos)
|
||||
infos.sort(key=lambda item: item.get("path", ""))
|
||||
return GlobResult(matches=infos)
|
||||
|
||||
def upload_files(self, files: list[tuple[str, bytes]]) -> list[FileUploadResponse]:
|
||||
"""Upload binary or text file payloads via the sandbox file API.
|
||||
|
||||
Contents are base64-encoded before calling the remote write_file API so
|
||||
arbitrary bytes can be transferred safely.
|
||||
"""
|
||||
responses: list[FileUploadResponse] = []
|
||||
for path, content in files:
|
||||
try:
|
||||
normalized_path = _normalize_path(path)
|
||||
if not _can_write_path(normalized_path):
|
||||
responses.append(FileUploadResponse(path=normalized_path, error="permission_denied"))
|
||||
continue
|
||||
result = self._get_client().file.write_file(
|
||||
file=normalized_path,
|
||||
content=base64.b64encode(content).decode("ascii"),
|
||||
encoding="base64",
|
||||
)
|
||||
if not result.success:
|
||||
raise Exception(result.message or "Upload failed")
|
||||
responses.append(FileUploadResponse(path=normalized_path, error=None))
|
||||
except PermissionError:
|
||||
normalized_path = str(path)
|
||||
responses.append(FileUploadResponse(path=normalized_path, error="permission_denied"))
|
||||
except IsADirectoryError:
|
||||
normalized_path = str(path)
|
||||
responses.append(FileUploadResponse(path=normalized_path, error="is_directory"))
|
||||
except FileNotFoundError:
|
||||
normalized_path = str(path)
|
||||
responses.append(FileUploadResponse(path=normalized_path, error="file_not_found"))
|
||||
except Exception as exc: # noqa: BLE001
|
||||
normalized_path = str(path)
|
||||
logger.warning(f"Upload to sandbox failed for {normalized_path}: {exc}")
|
||||
responses.append(FileUploadResponse(path=normalized_path, error="invalid_path"))
|
||||
return responses
|
||||
|
||||
def download_files(self, paths: list[str]) -> list[FileDownloadResponse]:
|
||||
"""Download file payloads as raw bytes from the sandbox file API.
|
||||
|
||||
_read_binary() normalizes the sandbox file API response to bytes.
|
||||
"""
|
||||
responses: list[FileDownloadResponse] = []
|
||||
for path in paths:
|
||||
try:
|
||||
normalized_path = _normalize_path(path)
|
||||
if not _can_read_path(normalized_path):
|
||||
responses.append(
|
||||
FileDownloadResponse(path=normalized_path, content=None, error="permission_denied")
|
||||
)
|
||||
continue
|
||||
content = self._read_binary(normalized_path)
|
||||
responses.append(FileDownloadResponse(path=normalized_path, content=content, error=None))
|
||||
except PermissionError:
|
||||
normalized_path = str(path)
|
||||
responses.append(FileDownloadResponse(path=normalized_path, content=None, error="permission_denied"))
|
||||
except IsADirectoryError:
|
||||
normalized_path = str(path)
|
||||
responses.append(FileDownloadResponse(path=normalized_path, content=None, error="is_directory"))
|
||||
except FileNotFoundError:
|
||||
normalized_path = str(path)
|
||||
responses.append(FileDownloadResponse(path=normalized_path, content=None, error="file_not_found"))
|
||||
except ValueError:
|
||||
normalized_path = str(path)
|
||||
responses.append(FileDownloadResponse(path=normalized_path, content=None, error="invalid_path"))
|
||||
except Exception as exc: # noqa: BLE001
|
||||
normalized_path = str(path)
|
||||
if _is_missing_file_error(exc):
|
||||
responses.append(FileDownloadResponse(path=normalized_path, content=None, error="file_not_found"))
|
||||
continue
|
||||
logger.warning(f"Download from sandbox failed for {normalized_path}: {exc}")
|
||||
responses.append(FileDownloadResponse(path=normalized_path, content=None, error=f"read_failed: {exc}"))
|
||||
return responses
|
||||
@@ -0,0 +1,196 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from pathlib import Path
|
||||
|
||||
from yuxi import config as conf
|
||||
from yuxi.utils.logging_config import logger
|
||||
from yuxi.utils.paths import (
|
||||
OUTPUTS_DIR_NAME,
|
||||
UPLOADS_DIR_NAME,
|
||||
VIRTUAL_PATH_PREFIX,
|
||||
WORKSPACE_AGENTS_DIR_NAME,
|
||||
WORKSPACE_AGENTS_PROMPT_FILE_NAME,
|
||||
WORKSPACE_DIR_NAME,
|
||||
)
|
||||
|
||||
_SAFE_ID_RE = re.compile(r"^[A-Za-z0-9_-]+$")
|
||||
|
||||
|
||||
def get_virtual_path_prefix() -> str:
|
||||
return "/" + VIRTUAL_PATH_PREFIX.strip("/")
|
||||
|
||||
|
||||
def validate_thread_id(thread_id: str) -> str:
|
||||
value = str(thread_id or "").strip()
|
||||
if not value:
|
||||
raise ValueError("thread_id is required")
|
||||
if not _SAFE_ID_RE.match(value):
|
||||
raise ValueError("thread_id contains invalid characters")
|
||||
return value
|
||||
|
||||
|
||||
def _thread_root_dir(thread_id: str) -> Path:
|
||||
safe_thread_id = validate_thread_id(thread_id)
|
||||
return Path(conf.save_dir) / "threads" / safe_thread_id / "user-data"
|
||||
|
||||
|
||||
def _validate_uid(uid: str) -> str:
|
||||
value = str(uid or "").strip()
|
||||
if not value:
|
||||
raise ValueError("uid is required")
|
||||
if not _SAFE_ID_RE.match(value):
|
||||
raise ValueError("uid contains invalid characters")
|
||||
return value
|
||||
|
||||
|
||||
def _global_user_data_dir(uid: str) -> Path:
|
||||
"""Return the shared host-side directory used for one user's workspace files."""
|
||||
safe_uid = _validate_uid(uid)
|
||||
return Path(conf.save_dir) / "threads" / "shared" / safe_uid
|
||||
|
||||
|
||||
def sandbox_user_data_dir(thread_id: str) -> Path:
|
||||
return _thread_root_dir(thread_id)
|
||||
|
||||
|
||||
def sandbox_workspace_dir(thread_id: str, uid: str) -> Path:
|
||||
validate_thread_id(thread_id)
|
||||
return _global_user_data_dir(uid) / WORKSPACE_DIR_NAME
|
||||
|
||||
|
||||
def sandbox_workspace_agents_prompt_file(thread_id: str, uid: str) -> Path:
|
||||
return sandbox_workspace_dir(thread_id, uid) / WORKSPACE_AGENTS_DIR_NAME / WORKSPACE_AGENTS_PROMPT_FILE_NAME
|
||||
|
||||
|
||||
def _threads_root_dir() -> Path:
|
||||
return (Path(conf.save_dir) / "threads").resolve(strict=False)
|
||||
|
||||
|
||||
def _resolve_threads_child_path(path: Path) -> Path:
|
||||
root = _threads_root_dir()
|
||||
resolved = path.resolve(strict=False)
|
||||
if not resolved.is_relative_to(root):
|
||||
raise ValueError("path resolved outside threads root")
|
||||
return resolved
|
||||
|
||||
|
||||
def _chmod_writable(path: Path, *, dir: bool = False) -> None:
|
||||
safe_path = _resolve_threads_child_path(path)
|
||||
mode = 0o777 if dir else 0o666
|
||||
try:
|
||||
safe_path.chmod(mode)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
|
||||
def ensure_workspace_default_files(workspace_dir: Path) -> None:
|
||||
workspace_dir = _resolve_threads_child_path(workspace_dir)
|
||||
agents_dir = workspace_dir / WORKSPACE_AGENTS_DIR_NAME
|
||||
agents_file = agents_dir / WORKSPACE_AGENTS_PROMPT_FILE_NAME
|
||||
|
||||
try:
|
||||
agents_dir.mkdir(parents=True, exist_ok=True)
|
||||
_chmod_writable(agents_dir, dir=True)
|
||||
except FileExistsError:
|
||||
logger.warning("工作区默认 Agents 目录创建失败:路径已被文件占用")
|
||||
return
|
||||
except OSError as exc:
|
||||
logger.warning(f"工作区默认 Agents 目录初始化失败: {exc}")
|
||||
return
|
||||
|
||||
try:
|
||||
with agents_file.open("xb"):
|
||||
pass
|
||||
_chmod_writable(agents_file)
|
||||
except FileExistsError:
|
||||
if agents_file.is_dir():
|
||||
logger.warning("工作区默认 AGENTS.md 创建失败:路径已被目录占用")
|
||||
except OSError as exc:
|
||||
logger.warning(f"工作区默认 Agents 文件初始化失败: {exc}")
|
||||
|
||||
|
||||
def sandbox_uploads_dir(thread_id: str) -> Path:
|
||||
return _thread_root_dir(thread_id) / UPLOADS_DIR_NAME
|
||||
|
||||
|
||||
def sandbox_outputs_dir(thread_id: str) -> Path:
|
||||
return _thread_root_dir(thread_id) / OUTPUTS_DIR_NAME
|
||||
|
||||
|
||||
def ensure_thread_dirs(thread_id: str, uid: str) -> None:
|
||||
_resolve_threads_child_path(_global_user_data_dir(uid)).mkdir(parents=True, exist_ok=True)
|
||||
workspace_dir = _resolve_threads_child_path(sandbox_workspace_dir(thread_id, uid))
|
||||
workspace_dir.mkdir(parents=True, exist_ok=True)
|
||||
ensure_workspace_default_files(workspace_dir)
|
||||
_resolve_threads_child_path(sandbox_uploads_dir(thread_id)).mkdir(parents=True, exist_ok=True)
|
||||
_resolve_threads_child_path(sandbox_outputs_dir(thread_id)).mkdir(parents=True, exist_ok=True)
|
||||
|
||||
|
||||
def _resolve_user_data_base_dir(thread_id: str, uid: str, relative_path: str) -> tuple[Path, Path]:
|
||||
"""Map a virtual user-data path to the correct host-side base directory."""
|
||||
parts = Path(relative_path).parts
|
||||
if not parts:
|
||||
base_dir = sandbox_user_data_dir(thread_id)
|
||||
return base_dir.resolve(), base_dir.resolve()
|
||||
|
||||
namespace = parts[0]
|
||||
if namespace == WORKSPACE_DIR_NAME:
|
||||
# Workspace is shared across one user's threads, so it lives outside the per-thread root.
|
||||
base_dir = sandbox_workspace_dir(thread_id, uid)
|
||||
target_path = base_dir.joinpath(*parts[1:]) if len(parts) > 1 else base_dir
|
||||
return base_dir.resolve(), target_path.resolve()
|
||||
if namespace == UPLOADS_DIR_NAME:
|
||||
base_dir = sandbox_uploads_dir(thread_id)
|
||||
target_path = base_dir.joinpath(*parts[1:]) if len(parts) > 1 else base_dir
|
||||
return base_dir.resolve(), target_path.resolve()
|
||||
if namespace == OUTPUTS_DIR_NAME:
|
||||
base_dir = sandbox_outputs_dir(thread_id)
|
||||
target_path = base_dir.joinpath(*parts[1:]) if len(parts) > 1 else base_dir
|
||||
return base_dir.resolve(), target_path.resolve()
|
||||
|
||||
base_dir = sandbox_user_data_dir(thread_id)
|
||||
return base_dir.resolve(), (base_dir / relative_path).resolve()
|
||||
|
||||
|
||||
def resolve_virtual_path(thread_id: str, virtual_path: str, *, uid: str) -> Path:
|
||||
clean_virtual_path = "/" + str(virtual_path or "").strip().lstrip("/")
|
||||
virtual_prefix = get_virtual_path_prefix()
|
||||
|
||||
if clean_virtual_path != virtual_prefix and not clean_virtual_path.startswith(f"{virtual_prefix}/"):
|
||||
raise ValueError(f"path must start with {virtual_prefix}")
|
||||
|
||||
relative_path = clean_virtual_path[len(virtual_prefix) :].lstrip("/")
|
||||
base_dir, target_path = _resolve_user_data_base_dir(thread_id, uid, relative_path)
|
||||
|
||||
try:
|
||||
target_path.relative_to(base_dir)
|
||||
except ValueError as exc:
|
||||
raise ValueError("path traversal detected") from exc
|
||||
|
||||
return target_path
|
||||
|
||||
|
||||
def virtual_path_for_thread_file(thread_id: str, path: str | Path, *, uid: str) -> str:
|
||||
target_path = Path(path).resolve()
|
||||
thread_root = sandbox_user_data_dir(thread_id).resolve()
|
||||
global_workspace_root = sandbox_workspace_dir(thread_id, uid).resolve()
|
||||
|
||||
try:
|
||||
relative_path = target_path.relative_to(global_workspace_root)
|
||||
except ValueError:
|
||||
try:
|
||||
relative_path = target_path.relative_to(thread_root)
|
||||
except ValueError as exc:
|
||||
raise ValueError("file is outside allowed user-data directories") from exc
|
||||
relative_path_str = relative_path.as_posix()
|
||||
else:
|
||||
workspace_relative = relative_path.as_posix()
|
||||
relative_path_str = (
|
||||
WORKSPACE_DIR_NAME if workspace_relative in {"", "."} else f"{WORKSPACE_DIR_NAME}/{workspace_relative}"
|
||||
)
|
||||
|
||||
prefix = get_virtual_path_prefix().rstrip("/")
|
||||
if not relative_path_str:
|
||||
return prefix
|
||||
return f"{prefix}/{relative_path_str}"
|
||||
@@ -0,0 +1,278 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import threading
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
|
||||
from yuxi import config as conf
|
||||
from yuxi.utils.logging_config import logger
|
||||
|
||||
from .provisioner_client import ProvisionerClient, SandboxRecord
|
||||
|
||||
|
||||
def sandbox_id_for_thread(thread_id: str, skills_thread_id: str | None = None, *, uid: str | None = None) -> str:
|
||||
file_thread_id = str(thread_id or "").strip()
|
||||
skills_id = str(skills_thread_id or file_thread_id).strip()
|
||||
uid_id = str(uid or "").strip()
|
||||
scope = file_thread_id if skills_id == file_thread_id else f"{file_thread_id}:{skills_id}"
|
||||
identity = f"{uid_id}:{scope}" if uid_id else scope
|
||||
digest = hashlib.sha256(identity.encode("utf-8")).hexdigest()
|
||||
return digest[:12]
|
||||
|
||||
|
||||
def _sandbox_key(uid: str, file_thread_id: str, skills_thread_id: str) -> str:
|
||||
return f"{uid}::{file_thread_id}::{skills_thread_id}"
|
||||
|
||||
|
||||
def normalize_env(env: dict | None) -> dict[str, str]:
|
||||
if not isinstance(env, dict):
|
||||
return {}
|
||||
return {str(key): "" if value is None else str(value) for key, value in env.items() if str(key)}
|
||||
|
||||
|
||||
def postgres_conninfo() -> str:
|
||||
db_url = os.getenv("POSTGRES_URL", "").strip()
|
||||
return db_url.replace("+asyncpg", "").replace("+psycopg", "")
|
||||
|
||||
|
||||
def load_user_agent_env(uid: str) -> dict[str, str]:
|
||||
conninfo = postgres_conninfo()
|
||||
if not conninfo:
|
||||
return {}
|
||||
|
||||
try:
|
||||
import psycopg
|
||||
|
||||
with psycopg.connect(conninfo, connect_timeout=3) as conn:
|
||||
with conn.cursor() as cursor:
|
||||
cursor.execute("SELECT env FROM agent_envs WHERE uid = %s", (uid,))
|
||||
row = cursor.fetchone()
|
||||
except Exception as exc:
|
||||
raise RuntimeError(f"failed to load agent env for uid {uid}: {exc}") from exc
|
||||
|
||||
if not row:
|
||||
return {}
|
||||
|
||||
value = row[0]
|
||||
if isinstance(value, str):
|
||||
try:
|
||||
value = json.loads(value)
|
||||
except json.JSONDecodeError as exc:
|
||||
raise RuntimeError(f"stored agent env for uid {uid} is not valid JSON") from exc
|
||||
return normalize_env(value)
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class SandboxConnection:
|
||||
cache_key: str
|
||||
thread_id: str
|
||||
file_thread_id: str
|
||||
skills_thread_id: str
|
||||
uid: str
|
||||
sandbox_id: str
|
||||
sandbox_url: str
|
||||
|
||||
|
||||
class ProvisionerSandboxProvider:
|
||||
def __init__(self):
|
||||
provider_name = str(getattr(conf, "sandbox_provider", "provisioner")).strip().lower()
|
||||
if provider_name != "provisioner":
|
||||
raise RuntimeError("only sandbox_provider=provisioner is supported")
|
||||
|
||||
provisioner_url = str(getattr(conf, "sandbox_provisioner_url", "") or "").strip()
|
||||
if not provisioner_url:
|
||||
raise RuntimeError("sandbox_provisioner_url is required")
|
||||
|
||||
self._client = ProvisionerClient(provisioner_url)
|
||||
self._lock = threading.Lock()
|
||||
self._thread_locks: dict[str, threading.Lock] = {}
|
||||
self._connections: dict[str, SandboxConnection] = {}
|
||||
self._last_touch_at: dict[str, float] = {}
|
||||
self._touch_interval_seconds = int(getattr(conf, "sandbox_keepalive_interval_seconds", 30))
|
||||
|
||||
def _thread_lock(self, cache_key: str) -> threading.Lock:
|
||||
with self._lock:
|
||||
lock = self._thread_locks.get(cache_key)
|
||||
if lock is None:
|
||||
lock = threading.Lock()
|
||||
self._thread_locks[cache_key] = lock
|
||||
return lock
|
||||
|
||||
def _record_to_connection(
|
||||
self,
|
||||
*,
|
||||
cache_key: str,
|
||||
thread_id: str,
|
||||
file_thread_id: str,
|
||||
skills_thread_id: str,
|
||||
uid: str,
|
||||
record: SandboxRecord,
|
||||
) -> SandboxConnection:
|
||||
connection = SandboxConnection(
|
||||
cache_key=cache_key,
|
||||
thread_id=thread_id,
|
||||
file_thread_id=file_thread_id,
|
||||
skills_thread_id=skills_thread_id,
|
||||
uid=uid,
|
||||
sandbox_id=record.sandbox_id,
|
||||
sandbox_url=record.sandbox_url,
|
||||
)
|
||||
self._connections[cache_key] = connection
|
||||
self._last_touch_at[cache_key] = time.time()
|
||||
return connection
|
||||
|
||||
def _should_touch(self, cache_key: str) -> bool:
|
||||
if self._touch_interval_seconds <= 0:
|
||||
return False
|
||||
last_touch = self._last_touch_at.get(cache_key)
|
||||
if last_touch is None:
|
||||
return True
|
||||
return (time.time() - last_touch) >= self._touch_interval_seconds
|
||||
|
||||
def _touch_if_needed(self, connection: SandboxConnection) -> bool:
|
||||
if not self._should_touch(connection.cache_key):
|
||||
return True
|
||||
is_alive = self._client.touch(connection.sandbox_id)
|
||||
self._last_touch_at[connection.cache_key] = time.time()
|
||||
return is_alive
|
||||
|
||||
def acquire(
|
||||
self,
|
||||
thread_id: str,
|
||||
*,
|
||||
uid: str,
|
||||
file_thread_id: str | None = None,
|
||||
skills_thread_id: str | None = None,
|
||||
) -> str:
|
||||
file_id = str(file_thread_id or thread_id).strip()
|
||||
skills_id = str(skills_thread_id or thread_id).strip()
|
||||
cache_key = _sandbox_key(uid, file_id, skills_id)
|
||||
lock = self._thread_lock(cache_key)
|
||||
with lock:
|
||||
current = self._connections.get(cache_key)
|
||||
if current:
|
||||
if current.uid != uid:
|
||||
raise RuntimeError(f"sandbox scope {cache_key} belongs to uid {current.uid}, not {uid}")
|
||||
try:
|
||||
if self._touch_if_needed(current):
|
||||
return current.sandbox_id
|
||||
self._connections.pop(cache_key, None)
|
||||
self._last_touch_at.pop(cache_key, None)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
logger.warning(f"Failed to touch sandbox {current.sandbox_id} for {cache_key}: {exc}")
|
||||
return current.sandbox_id
|
||||
|
||||
sandbox_id = sandbox_id_for_thread(file_id, skills_id, uid=uid)
|
||||
logger.info(f"Ensuring sandbox {sandbox_id} for file thread {file_id} and skills thread {skills_id}")
|
||||
record = self._client.create(
|
||||
sandbox_id,
|
||||
thread_id,
|
||||
uid,
|
||||
load_user_agent_env(uid),
|
||||
file_thread_id=file_id,
|
||||
skills_thread_id=skills_id,
|
||||
)
|
||||
|
||||
connection = self._record_to_connection(
|
||||
cache_key=cache_key,
|
||||
thread_id=thread_id,
|
||||
file_thread_id=file_id,
|
||||
skills_thread_id=skills_id,
|
||||
uid=uid,
|
||||
record=record,
|
||||
)
|
||||
return connection.sandbox_id
|
||||
|
||||
def get(
|
||||
self,
|
||||
thread_id: str,
|
||||
*,
|
||||
uid: str,
|
||||
create_if_missing: bool = False,
|
||||
file_thread_id: str | None = None,
|
||||
skills_thread_id: str | None = None,
|
||||
) -> SandboxConnection | None:
|
||||
file_id = str(file_thread_id or thread_id).strip()
|
||||
skills_id = str(skills_thread_id or thread_id).strip()
|
||||
cache_key = _sandbox_key(uid, file_id, skills_id)
|
||||
lock = self._thread_lock(cache_key)
|
||||
with lock:
|
||||
current = self._connections.get(cache_key)
|
||||
if current:
|
||||
if current.uid != uid:
|
||||
raise RuntimeError(f"sandbox scope {cache_key} belongs to uid {current.uid}, not {uid}")
|
||||
try:
|
||||
if self._touch_if_needed(current):
|
||||
return current
|
||||
self._connections.pop(cache_key, None)
|
||||
self._last_touch_at.pop(cache_key, None)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
logger.warning(f"Failed to touch sandbox {current.sandbox_id} for {cache_key}: {exc}")
|
||||
return current
|
||||
|
||||
sandbox_id = sandbox_id_for_thread(file_id, skills_id, uid=uid)
|
||||
if create_if_missing:
|
||||
record = self._client.create(
|
||||
sandbox_id,
|
||||
thread_id,
|
||||
uid,
|
||||
load_user_agent_env(uid),
|
||||
file_thread_id=file_id,
|
||||
skills_thread_id=skills_id,
|
||||
)
|
||||
else:
|
||||
record = self._client.discover(sandbox_id)
|
||||
if record is None:
|
||||
return None
|
||||
|
||||
return self._record_to_connection(
|
||||
cache_key=cache_key,
|
||||
thread_id=thread_id,
|
||||
file_thread_id=file_id,
|
||||
skills_thread_id=skills_id,
|
||||
uid=uid,
|
||||
record=record,
|
||||
)
|
||||
|
||||
def shutdown(self) -> None:
|
||||
with self._lock:
|
||||
connections = list(self._connections.values())
|
||||
self._connections.clear()
|
||||
self._last_touch_at.clear()
|
||||
|
||||
for connection in connections:
|
||||
try:
|
||||
self._client.delete(connection.sandbox_id)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
logger.warning(f"Failed to release sandbox {connection.sandbox_id} for {connection.cache_key}: {exc}")
|
||||
|
||||
|
||||
_sandbox_provider: ProvisionerSandboxProvider | None = None
|
||||
_sandbox_provider_lock = threading.Lock()
|
||||
|
||||
|
||||
def init_sandbox_provider() -> ProvisionerSandboxProvider:
|
||||
global _sandbox_provider
|
||||
with _sandbox_provider_lock:
|
||||
if _sandbox_provider is None:
|
||||
_sandbox_provider = ProvisionerSandboxProvider()
|
||||
return _sandbox_provider
|
||||
|
||||
|
||||
def get_sandbox_provider() -> ProvisionerSandboxProvider:
|
||||
provider = _sandbox_provider
|
||||
if provider is not None:
|
||||
return provider
|
||||
return init_sandbox_provider()
|
||||
|
||||
|
||||
def shutdown_sandbox_provider() -> None:
|
||||
global _sandbox_provider
|
||||
with _sandbox_provider_lock:
|
||||
provider = _sandbox_provider
|
||||
_sandbox_provider = None
|
||||
if provider is not None:
|
||||
provider.shutdown()
|
||||
@@ -0,0 +1,88 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
import httpx
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class SandboxRecord:
|
||||
sandbox_id: str
|
||||
sandbox_url: str
|
||||
status: str | None = None
|
||||
|
||||
|
||||
class ProvisionerClient:
|
||||
def __init__(self, base_url: str, *, timeout_seconds: int = 20):
|
||||
self._base_url = base_url.rstrip("/")
|
||||
self._timeout = httpx.Timeout(timeout_seconds)
|
||||
|
||||
def _request(self, method: str, path: str, **kwargs) -> httpx.Response:
|
||||
return httpx.request(
|
||||
method=method,
|
||||
url=f"{self._base_url}{path}",
|
||||
timeout=self._timeout,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
def health(self) -> bool:
|
||||
response = self._request("GET", "/health")
|
||||
return response.status_code == 200
|
||||
|
||||
def create(
|
||||
self,
|
||||
sandbox_id: str,
|
||||
thread_id: str,
|
||||
uid: str,
|
||||
env: dict[str, str] | None = None,
|
||||
*,
|
||||
file_thread_id: str | None = None,
|
||||
skills_thread_id: str | None = None,
|
||||
) -> SandboxRecord:
|
||||
response = self._request(
|
||||
"POST",
|
||||
"/api/sandboxes",
|
||||
json={
|
||||
"sandbox_id": sandbox_id,
|
||||
"thread_id": thread_id,
|
||||
"file_thread_id": file_thread_id or thread_id,
|
||||
"skills_thread_id": skills_thread_id or thread_id,
|
||||
"uid": uid,
|
||||
"env": env or {},
|
||||
},
|
||||
)
|
||||
if response.status_code >= 400:
|
||||
raise RuntimeError(f"failed to create sandbox {sandbox_id}: {response.status_code} {response.text}")
|
||||
payload = response.json()
|
||||
return SandboxRecord(
|
||||
sandbox_id=payload["sandbox_id"],
|
||||
sandbox_url=payload["sandbox_url"],
|
||||
status=payload.get("status"),
|
||||
)
|
||||
|
||||
def discover(self, sandbox_id: str) -> SandboxRecord | None:
|
||||
response = self._request("GET", f"/api/sandboxes/{sandbox_id}")
|
||||
if response.status_code == 404:
|
||||
return None
|
||||
if response.status_code >= 400:
|
||||
raise RuntimeError(f"failed to discover sandbox {sandbox_id}: {response.status_code} {response.text}")
|
||||
payload = response.json()
|
||||
return SandboxRecord(
|
||||
sandbox_id=payload["sandbox_id"],
|
||||
sandbox_url=payload["sandbox_url"],
|
||||
status=payload.get("status"),
|
||||
)
|
||||
|
||||
def touch(self, sandbox_id: str) -> bool:
|
||||
response = self._request("POST", f"/api/sandboxes/{sandbox_id}/touch")
|
||||
if response.status_code == 404:
|
||||
return False
|
||||
if response.status_code >= 400:
|
||||
raise RuntimeError(f"failed to touch sandbox {sandbox_id}: {response.status_code} {response.text}")
|
||||
return True
|
||||
|
||||
def delete(self, sandbox_id: str) -> None:
|
||||
response = self._request("DELETE", f"/api/sandboxes/{sandbox_id}")
|
||||
if response.status_code in {200, 404}:
|
||||
return
|
||||
raise RuntimeError(f"failed to delete sandbox {sandbox_id}: {response.status_code} {response.text}")
|
||||
@@ -0,0 +1,133 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import PurePosixPath
|
||||
|
||||
from deepagents.backends import FilesystemBackend
|
||||
from deepagents.backends.protocol import (
|
||||
EditResult,
|
||||
FileDownloadResponse,
|
||||
FileInfo,
|
||||
FileUploadResponse,
|
||||
GlobResult,
|
||||
GrepMatch,
|
||||
GrepResult,
|
||||
LsResult,
|
||||
ReadResult,
|
||||
WriteResult,
|
||||
)
|
||||
|
||||
from yuxi.agents.skills.service import get_skills_root_dir, is_valid_skill_slug
|
||||
|
||||
|
||||
class SelectedSkillsReadonlyBackend(FilesystemBackend):
|
||||
"""只读 skills backend,仅暴露选中的技能目录。"""
|
||||
|
||||
def __init__(self, *, selected_slugs: list[str] | None):
|
||||
super().__init__(root_dir=get_skills_root_dir(), virtual_mode=True)
|
||||
self._selected_slugs = {
|
||||
str(slug).strip()
|
||||
for slug in (selected_slugs or [])
|
||||
if isinstance(slug, str) and is_valid_skill_slug(str(slug))
|
||||
}
|
||||
|
||||
def _extract_slug(self, path: str | None) -> str | None:
|
||||
if not path:
|
||||
return None
|
||||
normalized = (path or "").strip()
|
||||
if not normalized or normalized == "/":
|
||||
return None
|
||||
pure = PurePosixPath(normalized if normalized.startswith("/") else f"/{normalized}")
|
||||
parts = [p for p in pure.parts if p not in ("/", "")]
|
||||
return parts[0] if parts else None
|
||||
|
||||
def _is_allowed_path(self, path: str | None) -> bool:
|
||||
slug = self._extract_slug(path)
|
||||
if slug is None:
|
||||
return True
|
||||
return slug in self._selected_slugs
|
||||
|
||||
def _is_allowed_file(self, file_path: str) -> bool:
|
||||
slug = self._extract_slug(file_path)
|
||||
return slug is not None and slug in self._selected_slugs
|
||||
|
||||
def _filter_infos(self, infos: list[FileInfo]) -> list[FileInfo]:
|
||||
return [item for item in infos if self._extract_slug(item.get("path", "")) in self._selected_slugs]
|
||||
|
||||
def _filter_matches(self, matches: list[GrepMatch]) -> list[GrepMatch]:
|
||||
return [item for item in matches if self._extract_slug(item.get("path", "")) in self._selected_slugs]
|
||||
|
||||
def ls(self, path: str) -> LsResult:
|
||||
if not self._selected_slugs:
|
||||
return LsResult(entries=[])
|
||||
|
||||
normalized = (path or "/").strip() or "/"
|
||||
if not self._is_allowed_path(normalized):
|
||||
return LsResult(error="Access denied: path is outside selected skills.")
|
||||
|
||||
result = super().ls(normalized)
|
||||
if result.error:
|
||||
return result
|
||||
infos = result.entries or []
|
||||
if normalized == "/":
|
||||
infos = self._filter_infos(infos)
|
||||
return LsResult(entries=infos)
|
||||
|
||||
def read(self, file_path: str, offset: int = 0, limit: int = 2000) -> ReadResult:
|
||||
if not self._is_allowed_file(file_path):
|
||||
return ReadResult(error="Access denied: file is outside selected skills.")
|
||||
return super().read(file_path, offset=offset, limit=limit)
|
||||
|
||||
def grep(self, pattern: str, path: str | None = None, glob: str | None = None) -> GrepResult:
|
||||
if not self._selected_slugs:
|
||||
return GrepResult(matches=[])
|
||||
|
||||
if path is not None:
|
||||
if not self._is_allowed_path(path):
|
||||
return GrepResult(error="Access denied: path is outside selected skills.")
|
||||
result = super().grep(pattern=pattern, path=path, glob=glob)
|
||||
if result.error:
|
||||
return result
|
||||
return GrepResult(matches=self._filter_matches(result.matches or []))
|
||||
|
||||
matches: list[GrepMatch] = []
|
||||
for slug in sorted(self._selected_slugs):
|
||||
result = super().grep(pattern=pattern, path=f"/{slug}", glob=glob)
|
||||
if result.error:
|
||||
continue
|
||||
matches.extend(result.matches or [])
|
||||
return GrepResult(matches=matches)
|
||||
|
||||
def glob(self, pattern: str, path: str = "/") -> GlobResult:
|
||||
if not self._selected_slugs:
|
||||
return GlobResult(matches=[])
|
||||
if not self._is_allowed_path(path):
|
||||
return GlobResult(error="Access denied: path is outside selected skills.")
|
||||
result = super().glob(pattern=pattern, path=path)
|
||||
if result.error:
|
||||
return result
|
||||
return GlobResult(matches=self._filter_infos(result.matches or []))
|
||||
|
||||
def write(self, file_path: str, content: str) -> WriteResult:
|
||||
return WriteResult(error="Skills path is read-only.")
|
||||
|
||||
def edit(self, file_path: str, old_string: str, new_string: str, replace_all: bool = False) -> EditResult:
|
||||
return EditResult(error="Skills path is read-only.")
|
||||
|
||||
def upload_files(self, files: list[tuple[str, bytes]]) -> list[FileUploadResponse]:
|
||||
return [FileUploadResponse(path=p, error="permission_denied") for p, _ in files]
|
||||
|
||||
def download_files(self, paths: list[str]) -> list[FileDownloadResponse]:
|
||||
responses: list[FileDownloadResponse] = []
|
||||
for path in paths:
|
||||
if not self._is_allowed_file(path):
|
||||
responses.append(FileDownloadResponse(path=path, content=None, error="invalid_path"))
|
||||
continue
|
||||
target = self._resolve_path(path)
|
||||
if not target.exists():
|
||||
responses.append(FileDownloadResponse(path=path, content=None, error="file_not_found"))
|
||||
continue
|
||||
if target.is_dir():
|
||||
responses.append(FileDownloadResponse(path=path, content=None, error="is_directory"))
|
||||
continue
|
||||
responses.append(FileDownloadResponse(path=path, content=target.read_bytes(), error=None))
|
||||
return responses
|
||||
@@ -0,0 +1,421 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
from abc import abstractmethod
|
||||
from contextlib import suppress
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from langchain_core.messages import ToolMessage
|
||||
from langgraph.checkpoint.memory import InMemorySaver
|
||||
from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver, aiosqlite
|
||||
from langgraph.graph.state import CompiledStateGraph
|
||||
from langgraph.stream.transformers import CustomTransformer
|
||||
from langgraph.types import Command
|
||||
|
||||
from yuxi import config as sys_config
|
||||
from yuxi.agents.context import DEFAULT_MAX_EXECUTION_STEPS, BaseContext, resolve_agent_resource_options
|
||||
from yuxi.storage.postgres.manager import pg_manager
|
||||
from yuxi.utils import logger
|
||||
from yuxi.utils.hash_utils import subagent_child_thread_id
|
||||
from yuxi.utils.thread_utils import extract_thread_id as _metadata_thread_id
|
||||
|
||||
|
||||
def _json_safe(value: Any) -> Any:
|
||||
if value is None or isinstance(value, str | int | float | bool):
|
||||
return value
|
||||
if isinstance(value, dict):
|
||||
return {str(key): _json_safe(child) for key, child in value.items()}
|
||||
if isinstance(value, list | tuple):
|
||||
return [_json_safe(child) for child in value]
|
||||
if hasattr(value, "model_dump"):
|
||||
return _json_safe(value.model_dump())
|
||||
return str(value)
|
||||
|
||||
|
||||
def _normalize_tool_event_data(data: Any) -> Any:
|
||||
"""规整 tools 流事件:write_todos / task 等返回 Command 的工具,其 tool-finished
|
||||
output 是 Command 对象,_json_safe 只能退化成 repr 字符串,前端无法关联结果。
|
||||
这里从 Command.update["messages"] 取出真正的 ToolMessage,使其与普通工具一致。"""
|
||||
if not isinstance(data, dict) or data.get("event") != "tool-finished":
|
||||
return data
|
||||
output = data.get("output")
|
||||
if not isinstance(output, Command):
|
||||
return data
|
||||
update = output.update if isinstance(output.update, dict) else {}
|
||||
messages = update.get("messages")
|
||||
if not isinstance(messages, list):
|
||||
return data
|
||||
tool_call_id = data.get("tool_call_id")
|
||||
tool_message = next(
|
||||
(m for m in messages if isinstance(m, ToolMessage) and m.tool_call_id == tool_call_id),
|
||||
next((m for m in messages if isinstance(m, ToolMessage)), None),
|
||||
)
|
||||
if tool_message is None:
|
||||
return data
|
||||
return {**data, "output": tool_message}
|
||||
|
||||
|
||||
def _subagent_route_for_namespace(
|
||||
routes: dict[tuple[str, ...], dict[str, str]], namespace: list[str]
|
||||
) -> dict[str, str] | None:
|
||||
ns = tuple(namespace)
|
||||
for path, route in sorted(routes.items(), key=lambda item: len(item[0]), reverse=True):
|
||||
if ns[: len(path)] == path:
|
||||
return route
|
||||
return None
|
||||
|
||||
|
||||
async def _collect_subagent_routes(run, parent_thread_id: str, routes: dict[tuple[str, ...], dict[str, str]]) -> None:
|
||||
subagents = getattr(run, "yuxi_subagents", None)
|
||||
if subagents is None:
|
||||
subagents = getattr(run, "subagents", None)
|
||||
if subagents is None:
|
||||
return
|
||||
|
||||
try:
|
||||
async for subagent in subagents:
|
||||
path = tuple(getattr(subagent, "path", ()) or ())
|
||||
subagent_slug = getattr(subagent, "name", None) or getattr(subagent, "graph_name", None)
|
||||
cause = getattr(subagent, "cause", None)
|
||||
tool_call_id = (
|
||||
cause.get("tool_call_id") if isinstance(cause, dict) else getattr(subagent, "trigger_call_id", None)
|
||||
)
|
||||
state = getattr(subagent, "state", None)
|
||||
metadata = getattr(subagent, "metadata", None)
|
||||
thread_id = _metadata_thread_id(metadata) or _metadata_thread_id(state)
|
||||
if not thread_id and isinstance(subagent_slug, str) and isinstance(tool_call_id, str) and tool_call_id:
|
||||
thread_id = subagent_child_thread_id(parent_thread_id, subagent_slug, tool_call_id)
|
||||
if path and isinstance(subagent_slug, str) and isinstance(tool_call_id, str) and tool_call_id and thread_id:
|
||||
routes[path] = {
|
||||
"thread_id": thread_id,
|
||||
"parent_thread_id": parent_thread_id,
|
||||
"subagent_slug": subagent_slug,
|
||||
"tool_call_id": tool_call_id,
|
||||
}
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except Exception as exc:
|
||||
logger.debug(f"collect subagent stream routes failed: {exc}")
|
||||
|
||||
|
||||
def _recursion_limit_from_context(context: BaseContext, default: int) -> int:
|
||||
value = getattr(context, "max_execution_steps", default)
|
||||
return int(value) if isinstance(value, int) and value > 0 else default
|
||||
|
||||
|
||||
class BaseAgent:
|
||||
"""
|
||||
定义一个基础 Agent 供 各类 graph 继承
|
||||
"""
|
||||
|
||||
name = "base_agent"
|
||||
description = "base_agent"
|
||||
capabilities: list[str] = [] # 智能体能力列表,如 ["file_upload", "web_search"] 等
|
||||
context_schema: type[BaseContext] = BaseContext # 智能体上下文 schema
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
self.graph = None # will be covered by get_graph
|
||||
self.checkpointer = None
|
||||
self._async_conn = None
|
||||
self.workdir = Path(sys_config.save_dir) / "agents" / self.module_name
|
||||
self.workdir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
@property
|
||||
def module_name(self) -> str:
|
||||
"""Get the module name of the agent class."""
|
||||
return self.__class__.__module__.split(".")[-2]
|
||||
|
||||
@property
|
||||
def id(self) -> str:
|
||||
"""Get the agent's class name."""
|
||||
return self.__class__.__name__
|
||||
|
||||
async def get_info(
|
||||
self,
|
||||
include_configurable_items: bool = True,
|
||||
user_role: str | None = None,
|
||||
db=None,
|
||||
user=None,
|
||||
):
|
||||
# metadata 固定在代码中,由各 Agent 的类属性提供
|
||||
metadata = self.load_metadata()
|
||||
configurable_items = {}
|
||||
if include_configurable_items:
|
||||
configurable_items = self.context_schema.get_configurable_items(user_role=user_role)
|
||||
if db is not None and user is not None:
|
||||
resource_fields = {
|
||||
item["kind"]
|
||||
for item in configurable_items.values()
|
||||
if item.get("kind") in {"tools", "knowledges", "mcps", "skills", "subagents"}
|
||||
}
|
||||
resource_options = await resolve_agent_resource_options(resource_fields, db=db, user=user)
|
||||
for item in configurable_items.values():
|
||||
if item.get("kind") in resource_options:
|
||||
item["options"] = resource_options[item["kind"]]
|
||||
|
||||
# Merge metadata with class attributes, metadata takes precedence
|
||||
return {
|
||||
"id": self.id,
|
||||
"name": getattr(self, "name", "Unknown"),
|
||||
"description": getattr(self, "description", "Unknown"),
|
||||
"metadata": metadata,
|
||||
"configurable_items": configurable_items,
|
||||
"capabilities": getattr(self, "capabilities", []), # 智能体能力列表
|
||||
}
|
||||
|
||||
async def get_config(self):
|
||||
return self.context_schema()
|
||||
|
||||
async def stream_values(self, messages: list[str], input_context=None, **kwargs):
|
||||
context = self.context_schema()
|
||||
context.update_from_dict(input_context or {})
|
||||
graph = await self.get_graph(context=context)
|
||||
for event in graph.astream({"messages": messages}, stream_mode="values", context=context):
|
||||
yield event["messages"]
|
||||
|
||||
async def stream_messages(self, messages: list[str], input_context=None, **kwargs):
|
||||
context = self.context_schema()
|
||||
context.update_from_dict(input_context or {})
|
||||
graph = await self.get_graph(context=context)
|
||||
logger.debug(f"stream_messages: {context=}")
|
||||
|
||||
# 构建配置:LangGraph 会自动从 checkpointer 恢复 state
|
||||
input_config = {
|
||||
"configurable": {"thread_id": context.thread_id, "uid": context.uid},
|
||||
"recursion_limit": _recursion_limit_from_context(context, DEFAULT_MAX_EXECUTION_STEPS),
|
||||
}
|
||||
|
||||
# langfuse metadata and callbacks integration
|
||||
if callbacks := kwargs.get("callbacks"):
|
||||
input_config["callbacks"] = list(callbacks)
|
||||
if metadata := kwargs.get("metadata"):
|
||||
input_config["metadata"] = dict(metadata)
|
||||
if tags := kwargs.get("tags"):
|
||||
input_config["tags"] = list(tags)
|
||||
|
||||
async for msg, metadata in graph.astream(
|
||||
{"messages": messages},
|
||||
stream_mode="messages",
|
||||
context=context,
|
||||
config=input_config,
|
||||
):
|
||||
yield msg, metadata
|
||||
|
||||
async def _stream_input_with_state(self, graph_input, input_context=None, **kwargs):
|
||||
context = self.context_schema()
|
||||
context.update_from_dict(input_context or {})
|
||||
graph = await self.get_graph(context=context)
|
||||
logger.debug(f"stream_with_state: {context=}")
|
||||
|
||||
input_config = {
|
||||
"configurable": {"thread_id": context.thread_id, "uid": context.uid},
|
||||
"recursion_limit": _recursion_limit_from_context(context, DEFAULT_MAX_EXECUTION_STEPS),
|
||||
}
|
||||
|
||||
if callbacks := kwargs.get("callbacks"):
|
||||
input_config["callbacks"] = list(callbacks)
|
||||
if metadata := kwargs.get("metadata"):
|
||||
input_config["metadata"] = dict(metadata)
|
||||
if tags := kwargs.get("tags"):
|
||||
input_config["tags"] = list(tags)
|
||||
|
||||
run = await graph.astream_events(
|
||||
graph_input,
|
||||
context=context,
|
||||
config=input_config,
|
||||
version="v3",
|
||||
transformers=[CustomTransformer],
|
||||
)
|
||||
subagent_routes: dict[tuple[str, ...], dict[str, str]] = {}
|
||||
route_task = asyncio.create_task(_collect_subagent_routes(run, context.thread_id, subagent_routes))
|
||||
try:
|
||||
async for event in run:
|
||||
params = event.get("params") or {}
|
||||
namespace = list(params.get("namespace") or [])
|
||||
method = event.get("method")
|
||||
data = params.get("data")
|
||||
subagent_route = _subagent_route_for_namespace(subagent_routes, namespace)
|
||||
|
||||
if method == "custom":
|
||||
yield "custom", data
|
||||
continue
|
||||
if method == "messages":
|
||||
msg, metadata = data
|
||||
metadata = dict(metadata or {})
|
||||
actual_thread_id = (subagent_route or {}).get("thread_id") or _metadata_thread_id(metadata)
|
||||
metadata["namespace"] = namespace
|
||||
metadata["stream_event"] = {"method": method, "namespace": namespace}
|
||||
if subagent_route:
|
||||
metadata.update(subagent_route)
|
||||
if actual_thread_id:
|
||||
metadata["thread_id"] = actual_thread_id
|
||||
yield "messages", (msg, metadata)
|
||||
elif method == "values" and not namespace:
|
||||
yield "values", data
|
||||
elif method in {"tasks", "tools", "lifecycle"}:
|
||||
if method == "tools":
|
||||
data = _normalize_tool_event_data(data)
|
||||
event_payload = {
|
||||
"method": method,
|
||||
"namespace": namespace,
|
||||
"data": _json_safe(data),
|
||||
}
|
||||
actual_thread_id = (subagent_route or {}).get("thread_id") or _metadata_thread_id(params)
|
||||
if subagent_route:
|
||||
event_payload.update(subagent_route)
|
||||
if actual_thread_id:
|
||||
event_payload["thread_id"] = actual_thread_id
|
||||
yield "stream_event", event_payload
|
||||
finally:
|
||||
route_task.cancel()
|
||||
with suppress(asyncio.CancelledError):
|
||||
await route_task
|
||||
|
||||
async def stream_messages_with_state(self, messages: list[str], input_context=None, **kwargs):
|
||||
async for event in self._stream_input_with_state({"messages": messages}, input_context, **kwargs):
|
||||
yield event
|
||||
|
||||
async def stream_resume_with_state(self, resume_input, input_context=None, **kwargs):
|
||||
async for event in self._stream_input_with_state(resume_input, input_context, **kwargs):
|
||||
yield event
|
||||
|
||||
async def invoke_messages(self, messages: list[str], input_context=None, **kwargs):
|
||||
context = self.context_schema()
|
||||
context.update_from_dict(input_context or {})
|
||||
graph = await self.get_graph(context=context)
|
||||
logger.debug(f"invoke_messages: {context}")
|
||||
|
||||
# 构建配置
|
||||
input_config = {
|
||||
"configurable": {"thread_id": context.thread_id, "uid": context.uid},
|
||||
"recursion_limit": _recursion_limit_from_context(context, DEFAULT_MAX_EXECUTION_STEPS),
|
||||
}
|
||||
|
||||
# langfuse metadata and callbacks integration
|
||||
if callbacks := kwargs.get("callbacks"):
|
||||
input_config["callbacks"] = list(callbacks)
|
||||
if metadata := kwargs.get("metadata"):
|
||||
input_config["metadata"] = dict(metadata)
|
||||
if tags := kwargs.get("tags"):
|
||||
input_config["tags"] = list(tags)
|
||||
|
||||
msg = await graph.ainvoke(
|
||||
{"messages": messages},
|
||||
context=context,
|
||||
config=input_config,
|
||||
)
|
||||
return msg
|
||||
|
||||
async def check_checkpointer(self):
|
||||
app = await self.get_graph()
|
||||
if not hasattr(app, "checkpointer") or app.checkpointer is None:
|
||||
return False
|
||||
return True
|
||||
|
||||
async def get_history(self, uid, thread_id) -> list[dict]:
|
||||
"""获取历史消息"""
|
||||
try:
|
||||
app = await self.get_graph()
|
||||
|
||||
if not await self.check_checkpointer():
|
||||
return []
|
||||
|
||||
config = {"configurable": {"thread_id": thread_id, "uid": uid}}
|
||||
state = await app.aget_state(config)
|
||||
|
||||
result = []
|
||||
if state:
|
||||
messages = state.values.get("messages", [])
|
||||
for msg in messages:
|
||||
if hasattr(msg, "model_dump"):
|
||||
msg_dict = msg.model_dump() # 转换成字典
|
||||
else:
|
||||
msg_dict = dict(msg) if hasattr(msg, "__dict__") else {"content": str(msg)}
|
||||
result.append(msg_dict)
|
||||
|
||||
return result
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"获取智能体 {self.name} 历史消息出错: {e}")
|
||||
return []
|
||||
|
||||
def reload_graph(self):
|
||||
"""重置 graph 缓存,强制下次调用 get_graph 时重新构建"""
|
||||
self.graph = None
|
||||
logger.info(f"{self.name} graph 缓存已清空,将在下次调用时重新构建")
|
||||
|
||||
@abstractmethod
|
||||
async def get_graph(self, **kwargs) -> CompiledStateGraph:
|
||||
"""
|
||||
获取并编译对话图实例。
|
||||
必须确保在编译时设置 checkpointer,否则将无法获取历史记录。
|
||||
例如: graph = workflow.compile(checkpointer=sqlite_checkpointer)
|
||||
"""
|
||||
pass
|
||||
|
||||
async def _get_checkpointer(self):
|
||||
if self.checkpointer is not None:
|
||||
return self.checkpointer
|
||||
|
||||
checkpointer = None
|
||||
backend = os.getenv("LANGGRAPH_CHECKPOINTER_BACKEND", "sqlite").strip().lower()
|
||||
|
||||
if backend == "postgres":
|
||||
checkpointer = await self._create_postgres_checkpointer()
|
||||
|
||||
if checkpointer is None:
|
||||
try:
|
||||
checkpointer = AsyncSqliteSaver(await self.get_async_conn())
|
||||
except Exception as e:
|
||||
logger.error(f"构建 sqlite checkpointer 失败: {e}, 尝试使用内存存储")
|
||||
checkpointer = InMemorySaver()
|
||||
|
||||
self.checkpointer = checkpointer
|
||||
return self.checkpointer
|
||||
|
||||
async def _create_postgres_checkpointer(self):
|
||||
postgres_url = os.getenv("POSTGRES_URL")
|
||||
if not postgres_url:
|
||||
logger.warning("POSTGRES_URL 未配置,无法启用 postgres checkpointer,回退 sqlite")
|
||||
return None
|
||||
|
||||
try:
|
||||
from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver # type: ignore
|
||||
except Exception as e:
|
||||
logger.warning(f"langgraph postgres checkpointer 不可用,回退 sqlite: {e}")
|
||||
return None
|
||||
|
||||
try:
|
||||
saver = AsyncPostgresSaver(pg_manager.langgraph_pool)
|
||||
|
||||
logger.info(f"{self.name} 使用 postgres checkpointer")
|
||||
return saver
|
||||
except Exception as e:
|
||||
logger.warning(f"初始化 postgres checkpointer 失败,回退 sqlite: {e}")
|
||||
return None
|
||||
|
||||
async def get_async_conn(self) -> aiosqlite.Connection:
|
||||
"""获取异步数据库连接"""
|
||||
if self._async_conn is not None:
|
||||
return self._async_conn
|
||||
|
||||
conn = await aiosqlite.connect(os.path.join(self.workdir, "aio_history.db"))
|
||||
# Patch: langgraph's AsyncSqliteSaver expects is_alive() method which aiosqlite may not have
|
||||
if not hasattr(conn, "is_alive"):
|
||||
conn.is_alive = lambda: True
|
||||
self._async_conn = conn
|
||||
return self._async_conn
|
||||
|
||||
async def get_aio_memory(self) -> AsyncSqliteSaver:
|
||||
"""获取异步存储实例"""
|
||||
return AsyncSqliteSaver(await self.get_async_conn())
|
||||
|
||||
def load_metadata(self) -> dict:
|
||||
"""Load metadata from agent class attribute."""
|
||||
metadata = getattr(self, "metadata", {})
|
||||
if isinstance(metadata, dict):
|
||||
return metadata
|
||||
logger.warning(f"Agent {self.module_name} metadata is not a dict, fallback to empty metadata")
|
||||
return {}
|
||||
@@ -0,0 +1,99 @@
|
||||
import asyncio
|
||||
import importlib
|
||||
import inspect
|
||||
from pathlib import Path
|
||||
|
||||
from yuxi.agents.base import BaseAgent
|
||||
from yuxi.utils import logger
|
||||
from yuxi.utils.singleton import SingletonMeta
|
||||
|
||||
|
||||
class AgentManager(metaclass=SingletonMeta):
|
||||
def __init__(self):
|
||||
self._classes = {}
|
||||
self._instances = {} # 存储已创建的 agent 实例
|
||||
|
||||
def register_agent(self, agent_class):
|
||||
self._classes[agent_class.__name__] = agent_class
|
||||
|
||||
def init_all_agents(self):
|
||||
for agent_id in self._classes.keys():
|
||||
self.get_agent(agent_id)
|
||||
|
||||
def get_agent(self, agent_id, reload=False, reload_graph=False, **kwargs):
|
||||
# 检查是否已经创建了该 agent 的实例
|
||||
if reload or agent_id not in self._instances:
|
||||
agent_class = self._classes[agent_id]
|
||||
self._instances[agent_id] = agent_class()
|
||||
|
||||
# 如果仅需要重新加载 graph,则清空 graph 缓存
|
||||
if reload_graph and agent_id in self._instances:
|
||||
self._instances[agent_id].reload_graph()
|
||||
|
||||
return self._instances[agent_id]
|
||||
|
||||
def get_agents(self):
|
||||
return list(self._instances.values())
|
||||
|
||||
async def reload_all(self):
|
||||
for agent_id in self._classes.keys():
|
||||
self.get_agent(agent_id, reload=True)
|
||||
|
||||
async def get_agents_info(self, include_configurable_items: bool = True):
|
||||
agents = self.get_agents()
|
||||
return await asyncio.gather(
|
||||
*[a.get_info(include_configurable_items=include_configurable_items) for a in agents]
|
||||
)
|
||||
|
||||
def auto_discover_agents(self):
|
||||
"""自动发现并注册 yuxi/agents/buildin/ 下的所有智能体。
|
||||
|
||||
遍历 yuxi/agents/buildin/ 目录下的所有子文件夹,如果子文件夹包含 __init__.py,
|
||||
则尝试从中导入 BaseAgent 的子类并注册。(使用自动导入的方式,支持私有agent)
|
||||
"""
|
||||
# 获取 agents 目录的路径
|
||||
agents_dir = Path(__file__).parent
|
||||
|
||||
# 遍历所有子目录
|
||||
for item in agents_dir.iterdir():
|
||||
# logger.info(f"尝试导入模块:{item}")
|
||||
# 跳过非目录、common 目录、__pycache__ 等
|
||||
if not item.is_dir() or item.name.startswith("_"):
|
||||
continue
|
||||
|
||||
# 检查是否有 __init__.py 文件
|
||||
init_file = item / "__init__.py"
|
||||
if not init_file.exists():
|
||||
logger.warning(f"{item} 不是一个有效的模块")
|
||||
continue
|
||||
|
||||
# 尝试导入模块
|
||||
try:
|
||||
module_name = f"yuxi.agents.buildin.{item.name}"
|
||||
module = importlib.import_module(module_name)
|
||||
|
||||
# 查找模块中所有 BaseAgent 的子类
|
||||
for name, obj in inspect.getmembers(module):
|
||||
if (
|
||||
inspect.isclass(obj)
|
||||
and issubclass(obj, BaseAgent)
|
||||
and obj is not BaseAgent
|
||||
and obj.__module__.startswith(module_name)
|
||||
):
|
||||
logger.info(f"自动发现智能体: {obj.__name__} 来自 {item.name}")
|
||||
self.register_agent(obj)
|
||||
|
||||
except Exception as e:
|
||||
logger.warning(f"无法从 {item.name} 加载智能体: {e}")
|
||||
|
||||
|
||||
agent_manager = AgentManager()
|
||||
# 自动发现并注册所有智能体
|
||||
agent_manager.auto_discover_agents()
|
||||
agent_manager.init_all_agents()
|
||||
|
||||
__all__ = ["agent_manager"]
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pass
|
||||
@@ -0,0 +1,5 @@
|
||||
from .context import ChatBotContext
|
||||
from .graph import ChatbotAgent
|
||||
from .state import ChatBotState, SubAgentRunState, merge_subagent_runs
|
||||
|
||||
__all__ = ["ChatBotContext", "ChatBotState", "ChatbotAgent", "SubAgentRunState", "merge_subagent_runs"]
|
||||
@@ -0,0 +1,17 @@
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
from yuxi.agents.context import BaseContext
|
||||
|
||||
|
||||
@dataclass(kw_only=True)
|
||||
class ChatBotContext(BaseContext):
|
||||
subagents: list[str] | None = field(
|
||||
default=None,
|
||||
metadata={
|
||||
"name": "子智能体",
|
||||
"options": [],
|
||||
"description": "可选子智能体列表,为空表示启用当前用户可见的全部子智能体。",
|
||||
"type": "list",
|
||||
"kind": "subagents",
|
||||
},
|
||||
)
|
||||
@@ -0,0 +1,113 @@
|
||||
from deepagents.middleware.patch_tool_calls import PatchToolCallsMiddleware
|
||||
from langchain.agents import create_agent
|
||||
from langchain.agents.middleware import ModelRetryMiddleware, TodoListMiddleware
|
||||
|
||||
from yuxi.agents import BaseAgent, load_chat_model, resolve_chat_model_spec
|
||||
from yuxi.agents.backends import create_agent_filesystem_middleware
|
||||
from yuxi.agents.context import (
|
||||
DEFAULT_SUMMARY_KEEP_MESSAGES,
|
||||
DEFAULT_SUMMARY_L2_TRIGGER_RATIO,
|
||||
DEFAULT_SUMMARY_THRESHOLD_K,
|
||||
DEFAULT_SUMMARY_TOOL_RESULT_TOKEN_LIMIT,
|
||||
DEFAULT_TOOL_RESULT_EVICTION_K_TOKENS,
|
||||
DEFAULT_YUXI_SUMMARY_PROMPT,
|
||||
prepare_agent_runtime_context,
|
||||
)
|
||||
from yuxi.agents.middlewares import (
|
||||
TokenUsageMiddleware,
|
||||
create_summary_middleware,
|
||||
save_attachments_to_fs,
|
||||
)
|
||||
from yuxi.agents.middlewares.skills import SkillsMiddleware
|
||||
from yuxi.agents.middlewares.subagent_task import create_subagent_task_middleware
|
||||
from yuxi.agents.toolkits.service import resolve_configured_runtime_tools
|
||||
|
||||
from .context import ChatBotContext
|
||||
from .prompt import TODO_MID_PROMPT, build_prompt_with_context
|
||||
from .state import ChatBotState
|
||||
|
||||
|
||||
async def _build_middlewares(context):
|
||||
"""构建中间件列表"""
|
||||
# summary middleware
|
||||
# 主 Agent 上下文优化:默认 100k tokens 触发压缩,压缩后保留最近 10 条消息
|
||||
summary_trigger_tokens = getattr(context, "summary_threshold", DEFAULT_SUMMARY_THRESHOLD_K) * 1024
|
||||
summary_keep_messages = getattr(context, "summary_keep_messages", DEFAULT_SUMMARY_KEEP_MESSAGES)
|
||||
summary_prompt = getattr(context, "summary_prompt", None) or DEFAULT_YUXI_SUMMARY_PROMPT
|
||||
summary_tool_result_token_limit = getattr(
|
||||
context,
|
||||
"summary_tool_result_token_limit",
|
||||
DEFAULT_SUMMARY_TOOL_RESULT_TOKEN_LIMIT,
|
||||
)
|
||||
summary_l2_trigger_ratio = getattr(context, "summary_l2_trigger_ratio", DEFAULT_SUMMARY_L2_TRIGGER_RATIO)
|
||||
model_spec = resolve_chat_model_spec(context.model)
|
||||
summary_middleware = create_summary_middleware(
|
||||
model=load_chat_model(fully_specified_name=model_spec),
|
||||
trigger=("tokens", summary_trigger_tokens),
|
||||
keep=("messages", summary_keep_messages),
|
||||
summary_prompt=summary_prompt,
|
||||
trim_tokens_to_summarize=summary_trigger_tokens,
|
||||
tool_result_offload_token_limit=summary_tool_result_token_limit,
|
||||
l1_l2_trigger_ratio=summary_l2_trigger_ratio,
|
||||
)
|
||||
|
||||
middlewares = [
|
||||
create_agent_filesystem_middleware(
|
||||
getattr(context, "tool_token_limit", DEFAULT_TOOL_RESULT_EVICTION_K_TOKENS) * 1024,
|
||||
context=context,
|
||||
),
|
||||
save_attachments_to_fs,
|
||||
SkillsMiddleware(),
|
||||
]
|
||||
subagent_middleware = await create_subagent_task_middleware(context)
|
||||
if subagent_middleware:
|
||||
middlewares.append(subagent_middleware)
|
||||
middlewares.extend(
|
||||
[
|
||||
summary_middleware,
|
||||
TodoListMiddleware(system_prompt=TODO_MID_PROMPT),
|
||||
PatchToolCallsMiddleware(),
|
||||
ModelRetryMiddleware(max_retries=getattr(context, "model_retry_times", 2)),
|
||||
TokenUsageMiddleware(),
|
||||
]
|
||||
)
|
||||
return middlewares
|
||||
|
||||
|
||||
class ChatbotAgent(BaseAgent):
|
||||
name = "智能助手"
|
||||
description = "基础的对话机器人,可以回答问题,可在配置中启用需要的工具。"
|
||||
capabilities = ["file_upload", "files"] # 支持文件上传功能
|
||||
context_schema = ChatBotContext
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
|
||||
async def get_graph(self, context=None, **kwargs):
|
||||
|
||||
context = await prepare_agent_runtime_context(
|
||||
context or self.context_schema(),
|
||||
context_schema=self.context_schema,
|
||||
)
|
||||
|
||||
# 使用 create_agent 创建智能体
|
||||
model_spec = resolve_chat_model_spec(context.model)
|
||||
graph = create_agent(
|
||||
model=load_chat_model(fully_specified_name=model_spec),
|
||||
tools=await resolve_configured_runtime_tools(context),
|
||||
system_prompt=build_prompt_with_context(context),
|
||||
middleware=await _build_middlewares(context),
|
||||
state_schema=ChatBotState,
|
||||
checkpointer=await self._get_checkpointer(),
|
||||
)
|
||||
|
||||
return graph
|
||||
|
||||
|
||||
def main():
|
||||
pass
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
# asyncio.run(main())
|
||||
@@ -0,0 +1,91 @@
|
||||
from yuxi.utils.datetime_utils import shanghai_now
|
||||
from yuxi.utils.paths import (
|
||||
VIRTUAL_PATH_OUTPUTS,
|
||||
VIRTUAL_PATH_PREFIX,
|
||||
VIRTUAL_PATH_UPLOADS,
|
||||
VIRTUAL_PATH_WORKSPACE,
|
||||
)
|
||||
|
||||
PROMPT = f"""
|
||||
你是一个交互式智能体“语析“。
|
||||
|
||||
专门用来回答用户的问题。请根据用户提供的信息,尽可能详细地回答问题。
|
||||
如果你不确定答案,可以说你不知道,但请尽量提供相关的信息或建议。请保持礼貌和专业。
|
||||
|
||||
<| 内部执行约束:重要 |>
|
||||
以下内容仅用于指导你的内部执行过程,不属于面向用户的基本设定。除非用户明确询问系统如何工作,
|
||||
否则不要主动向用户说明工作区、文件系统、知识库路径、工具调用方式等内部实现细节。
|
||||
|
||||
<| 文件系统约束 |>
|
||||
系统主要工作路径为 {VIRTUAL_PATH_PREFIX},但必须遵守规范:
|
||||
- {VIRTUAL_PATH_OUTPUTS}:用于写入的文件夹
|
||||
- {VIRTUAL_PATH_OUTPUTS}/tmp/:用于存放中间结果或备份内容
|
||||
- {VIRTUAL_PATH_UPLOADS}:用于存放用户上传的附件(只读,除非用户要求,否则不得写入)
|
||||
- {VIRTUAL_PATH_WORKSPACE}:用于存放用户文件(用户私人目录,除非用户要求,否则不得写入)
|
||||
- 其他路径:非必要不写入其他路径
|
||||
|
||||
<| 风格规范 |>
|
||||
保持专业严谨,减少使用 Emoji
|
||||
|
||||
<| 可视化 HTML 辅助组件规范 |>
|
||||
回答的主要表达载体始终是 Markdown。只有当普通 Markdown 难以清晰表达数值对比、层级关系、流程结构、
|
||||
时间线、关键指标或布局示意时,才可以额外使用 Markdown 围栏代码块语言标记 `html:preview`
|
||||
输出一个轻量静态 HTML 辅助组件:
|
||||
```html:preview
|
||||
自包含的静态 HTML/CSS 内容
|
||||
```
|
||||
使用要求:
|
||||
- `html:preview` 只用于补齐 Markdown 的短板,不能替代正文回答;核心解释、推理、背景、风险、
|
||||
结论展开和完整明细必须放在普通 Markdown 中。
|
||||
- 如果 Markdown 的标题、列表、表格、引用或代码块已经足够清楚,不要使用 `html:preview`。
|
||||
- 预览内容应优先使用静态 HTML/CSS;可以引用方便访问、稳定、无需登录鉴权的 HTTPS 外链资源
|
||||
(如公开图片或字体),但必须保证没有外链时核心信息仍可读,不要依赖跨域受限、内网、
|
||||
临时链接或不稳定资源,不要编写 JavaScript。
|
||||
- 这是嵌入在回答中的辅助可视化组件,不是完整网页、不是正文容器、不是自带外壳的信息卡片;
|
||||
不要设计导航栏、页脚、登录态、表单、复杂按钮、营销页 Hero 或多屏网页结构。
|
||||
- 外层预览容器已经提供 12px 圆角、边框和裁切;HTML 内容本身不要再套卡片壳、面板壳或页面壳,
|
||||
不要给最外层内容添加大圆角、阴影、厚边框、额外外边距或整页背景。
|
||||
- 内容组织必须以“快速看懂”为中心:优先呈现少量关键指标、对比关系、趋势/阶段、状态和极短备注,
|
||||
避免为了视觉效果牺牲可读性。
|
||||
- 默认按 800px * 360px 的展示尺寸设计;前端最大可能支持到 700px 高度,真实宽高也会随容器变化,因此布局必须响应式。
|
||||
- HTML 内部不要写死整体画布高度;优先使用 `max-width: 100%`、`box-sizing: border-box`、
|
||||
弹性网格、换行和适度压缩间距来适配不同宽高。
|
||||
- 必须保证核心内容在 800px * 360px 内可读且不依赖滚动;如果预计放不下,必须减少内容,而不是缩小到难以阅读或继续堆叠。
|
||||
- 可视化组件最多呈现 1 个短标题、3-5 个关键指标或一组简短对比;不要在组件里放完整明细、长表格、长列表或多段说明。
|
||||
- 当数据超过 6 项时,不要逐项做卡片网格;应汇总为趋势、最大/最小值、异常点、Top 3、分布或区间。
|
||||
完整列表、明细表或逐日解释放在 `html:preview` 之后用普通 Markdown 展示。
|
||||
- 可视化组件内禁止放成段文字、长句解释、新闻正文、报告段落、多行预警说明或叙事性文案;
|
||||
组件内文字应以短标签、短结论、数字、单位、状态词和极短备注为主。
|
||||
- 单个说明文本建议不超过 20 个中文字符;超过一句话的解释、背景、风险说明、数据来源详情必须放在
|
||||
`html:preview` 后面的普通 Markdown 中。
|
||||
- 设计应克制、清晰、信息密度适中;优先使用紧凑指标组、摘要表、对比条、状态标签、时间轴和简单关系图,
|
||||
不要做复杂装饰、大图标、密集网格或过重视觉效果。
|
||||
- 如果用户是在询问 HTML 源码、教程示例或需要复制代码,必须使用普通 `html` 代码块,不要使用 `html:preview`。
|
||||
"""
|
||||
|
||||
# 效果不好,暂时不启用
|
||||
SOURCE_CITE_PROMPT = """
|
||||
|
||||
<| 引用来源 |>
|
||||
当你提供的信息来自于用户上传的文件或者知识库中的内容时,请务必在回答中注明信息来源,以增加答案的可信度和透明度。
|
||||
|
||||
对于论断内容,需要添加参考文献信息,将对应段落的末尾添加 cite 信息。使用
|
||||
<cite source="$SOURCE" type="$TYPE">$INDEX</cite>
|
||||
|
||||
- $SOURCE:信息来源,可以是文件名,可以是url
|
||||
- $TYPE:引用类型,可以是 "file"、"url",对于网络搜索应该使用 "url",对于用户上传的文件或者知识库中的内容应该使用 "file"
|
||||
- $INDEX:引用索引,应该从 1 开始
|
||||
|
||||
比如 <cite source="食品工艺学.pdf" type="file">1</cite>
|
||||
"""
|
||||
|
||||
TODO_MID_PROMPT = """
|
||||
你需要根据任务的复杂程度来使用 write_todos 来记录规划和待办事项,确保任务的每个步骤都被记录和跟踪。
|
||||
每个待办任务名称必须简短,控制在 20 个中文汉字以内。
|
||||
"""
|
||||
|
||||
|
||||
def build_prompt_with_context(context):
|
||||
current_date = f"当前日期:{shanghai_now().strftime('%Y-%m-%d')}"
|
||||
system_prompt = f"{current_date}\n\n{PROMPT.strip()}\n\n{context.system_prompt or ''}"
|
||||
return system_prompt.strip()
|
||||
@@ -0,0 +1,59 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Annotated, Literal, TypedDict
|
||||
|
||||
from yuxi.agents.state import BaseState
|
||||
|
||||
|
||||
class SubAgentRunState(TypedDict, total=False):
|
||||
id: str
|
||||
run_id: str
|
||||
subagent_slug: str
|
||||
subagent_name: str
|
||||
child_thread_id: str
|
||||
description: str
|
||||
status: Literal["pending", "running", "completed", "failed", "cancel_requested", "cancelled", "interrupted"]
|
||||
created_at: str
|
||||
completed_at: str
|
||||
error: str | None
|
||||
artifacts: list[str]
|
||||
events_url: str
|
||||
result_url: str
|
||||
|
||||
|
||||
def merge_subagent_runs(
|
||||
existing: list[SubAgentRunState] | None,
|
||||
new: list[SubAgentRunState] | None,
|
||||
) -> list[SubAgentRunState]:
|
||||
"""LangGraph state reducer:增量合并父 Agent 记录的子智能体运行摘要。
|
||||
|
||||
`run_id` 是一次真实子智能体执行的身份。只有相同 `run_id` 才会更新同一条记录;
|
||||
没有 `run_id` 的增量记录直接追加,不用工具调用 ID 或子线程 ID 做旧状态兼容匹配。
|
||||
"""
|
||||
if existing is None:
|
||||
return list(new or [])
|
||||
if new is None:
|
||||
return existing
|
||||
|
||||
merged = [dict(item) for item in existing]
|
||||
run_id_index = {item.get("run_id"): position for position, item in enumerate(merged) if item.get("run_id")}
|
||||
for item in new:
|
||||
run = dict(item)
|
||||
run_id = run.get("run_id")
|
||||
position = None
|
||||
if run_id and run_id in run_id_index:
|
||||
position = run_id_index[run_id]
|
||||
|
||||
if position is None:
|
||||
position = len(merged)
|
||||
merged.append(run)
|
||||
else:
|
||||
merged[position] = {**merged[position], **run}
|
||||
|
||||
if run_id:
|
||||
run_id_index[run_id] = position
|
||||
return merged
|
||||
|
||||
|
||||
class ChatBotState(BaseState):
|
||||
subagent_runs: Annotated[list[SubAgentRunState], merge_subagent_runs]
|
||||
@@ -0,0 +1,4 @@
|
||||
from .context import SubAgentContext
|
||||
from .graph import SubAgentBackend
|
||||
|
||||
__all__ = ["SubAgentBackend", "SubAgentContext"]
|
||||
@@ -0,0 +1,23 @@
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
from yuxi.agents.context import BaseContext
|
||||
|
||||
|
||||
@dataclass(kw_only=True)
|
||||
class SubAgentContext(BaseContext):
|
||||
parent_thread_id: str | None = field(
|
||||
default=None,
|
||||
metadata={"name": "父线程ID", "configurable": False, "hide": True},
|
||||
)
|
||||
file_thread_id: str | None = field(
|
||||
default=None,
|
||||
metadata={"name": "文件线程ID", "configurable": False, "hide": True},
|
||||
)
|
||||
skills_thread_id: str | None = field(
|
||||
default=None,
|
||||
metadata={"name": "Skills线程ID", "configurable": False, "hide": True},
|
||||
)
|
||||
is_subagent_runtime: bool = field(
|
||||
default=False,
|
||||
metadata={"name": "子智能体运行态", "configurable": False, "hide": True},
|
||||
)
|
||||
@@ -0,0 +1,127 @@
|
||||
from typing import Any
|
||||
|
||||
from deepagents.middleware.patch_tool_calls import PatchToolCallsMiddleware
|
||||
from langchain.agents import create_agent
|
||||
from langchain.agents.middleware import ModelRetryMiddleware, TodoListMiddleware
|
||||
from langchain.agents.middleware.types import AgentMiddleware
|
||||
|
||||
from yuxi.agents import BaseAgent, BaseState, load_chat_model, resolve_chat_model_spec
|
||||
from yuxi.agents.backends import create_agent_filesystem_middleware
|
||||
from yuxi.agents.buildin.chatbot.prompt import TODO_MID_PROMPT, build_prompt_with_context
|
||||
from yuxi.agents.buildin.subagent.context import SubAgentContext
|
||||
from yuxi.agents.context import (
|
||||
DEFAULT_SUMMARY_KEEP_MESSAGES,
|
||||
DEFAULT_SUMMARY_L2_TRIGGER_RATIO,
|
||||
DEFAULT_SUMMARY_THRESHOLD_K,
|
||||
DEFAULT_SUMMARY_TOOL_RESULT_TOKEN_LIMIT,
|
||||
DEFAULT_TOOL_RESULT_EVICTION_K_TOKENS,
|
||||
DEFAULT_YUXI_SUMMARY_PROMPT,
|
||||
prepare_agent_runtime_context,
|
||||
)
|
||||
from yuxi.agents.middlewares import TokenUsageMiddleware, create_summary_middleware, save_attachments_to_fs
|
||||
from yuxi.agents.middlewares.skills import SkillsMiddleware
|
||||
from yuxi.agents.toolkits.service import resolve_configured_runtime_tools
|
||||
|
||||
_SUBAGENT_DISABLED_TOOLS = frozenset({"present_artifacts", "ask_user_question", "install_skill"})
|
||||
|
||||
|
||||
def _tool_name(tool) -> str | None:
|
||||
if isinstance(tool, dict):
|
||||
name = tool.get("name")
|
||||
else:
|
||||
name = getattr(tool, "name", None)
|
||||
return name if isinstance(name, str) else None
|
||||
|
||||
|
||||
def _filter_disabled_tools(tools):
|
||||
return [tool for tool in tools if _tool_name(tool) not in _SUBAGENT_DISABLED_TOOLS]
|
||||
|
||||
|
||||
class _SubAgentToolFilterMiddleware(AgentMiddleware[Any, Any, Any]):
|
||||
def wrap_model_call(self, request, handler):
|
||||
return handler(request.override(tools=_filter_disabled_tools(request.tools or [])))
|
||||
|
||||
async def awrap_model_call(self, request, handler):
|
||||
return await handler(request.override(tools=_filter_disabled_tools(request.tools or [])))
|
||||
|
||||
|
||||
async def _build_middlewares(context):
|
||||
summary_trigger_tokens = getattr(context, "summary_threshold", DEFAULT_SUMMARY_THRESHOLD_K) * 1024
|
||||
summary_keep_messages = getattr(context, "summary_keep_messages", DEFAULT_SUMMARY_KEEP_MESSAGES)
|
||||
summary_prompt = getattr(context, "summary_prompt", None) or DEFAULT_YUXI_SUMMARY_PROMPT
|
||||
summary_tool_result_token_limit = getattr(
|
||||
context,
|
||||
"summary_tool_result_token_limit",
|
||||
DEFAULT_SUMMARY_TOOL_RESULT_TOKEN_LIMIT,
|
||||
)
|
||||
summary_l2_trigger_ratio = getattr(context, "summary_l2_trigger_ratio", DEFAULT_SUMMARY_L2_TRIGGER_RATIO)
|
||||
model_spec = resolve_chat_model_spec(context.model)
|
||||
summary_middleware = create_summary_middleware(
|
||||
model=load_chat_model(fully_specified_name=model_spec),
|
||||
trigger=("tokens", summary_trigger_tokens),
|
||||
keep=("messages", summary_keep_messages),
|
||||
summary_prompt=summary_prompt,
|
||||
trim_tokens_to_summarize=summary_trigger_tokens,
|
||||
tool_result_offload_token_limit=summary_tool_result_token_limit,
|
||||
l1_l2_trigger_ratio=summary_l2_trigger_ratio,
|
||||
)
|
||||
|
||||
return [
|
||||
create_agent_filesystem_middleware(
|
||||
getattr(context, "tool_token_limit", DEFAULT_TOOL_RESULT_EVICTION_K_TOKENS) * 1024,
|
||||
context=context,
|
||||
),
|
||||
save_attachments_to_fs,
|
||||
SkillsMiddleware(),
|
||||
summary_middleware,
|
||||
TodoListMiddleware(system_prompt=TODO_MID_PROMPT),
|
||||
PatchToolCallsMiddleware(),
|
||||
_SubAgentToolFilterMiddleware(),
|
||||
ModelRetryMiddleware(),
|
||||
TokenUsageMiddleware(),
|
||||
]
|
||||
|
||||
|
||||
class SubAgentBackend(BaseAgent):
|
||||
name = "子智能体"
|
||||
description = "用于被主智能体通过 task 工具调用的专用智能体后端。"
|
||||
capabilities = ["file_upload", "files"]
|
||||
context_schema = SubAgentContext
|
||||
|
||||
async def get_info(
|
||||
self,
|
||||
include_configurable_items: bool = True,
|
||||
user_role: str | None = None,
|
||||
db=None,
|
||||
user=None,
|
||||
):
|
||||
info = await super().get_info(
|
||||
include_configurable_items=include_configurable_items,
|
||||
user_role=user_role,
|
||||
db=db,
|
||||
user=user,
|
||||
)
|
||||
tools_item = (info.get("configurable_items") or {}).get("tools")
|
||||
if isinstance(tools_item, dict):
|
||||
tools_item["options"] = [
|
||||
option
|
||||
for option in tools_item.get("options") or []
|
||||
if option.get("key") not in _SUBAGENT_DISABLED_TOOLS
|
||||
]
|
||||
return info
|
||||
|
||||
async def get_graph(self, context=None, **kwargs):
|
||||
context = await prepare_agent_runtime_context(
|
||||
context or self.context_schema(),
|
||||
context_schema=self.context_schema,
|
||||
)
|
||||
model_spec = resolve_chat_model_spec(context.model)
|
||||
|
||||
return create_agent(
|
||||
model=load_chat_model(fully_specified_name=model_spec),
|
||||
tools=_filter_disabled_tools(await resolve_configured_runtime_tools(context)),
|
||||
system_prompt=build_prompt_with_context(context),
|
||||
middleware=await _build_middlewares(context),
|
||||
state_schema=BaseState,
|
||||
checkpointer=await self._get_checkpointer(),
|
||||
)
|
||||
@@ -0,0 +1,558 @@
|
||||
"""Define the configurable parameters for the agent."""
|
||||
|
||||
import asyncio
|
||||
import uuid
|
||||
from dataclasses import MISSING, dataclass, field, fields
|
||||
from typing import Any, get_origin
|
||||
|
||||
from yuxi.agents.backends.sandbox.paths import sandbox_workspace_agents_prompt_file
|
||||
from yuxi.utils.logging_config import logger
|
||||
|
||||
WORKSPACE_AGENTS_PROMPT_MAX_BYTES = 64 * 1024
|
||||
DEFAULT_SUMMARY_THRESHOLD_K = 100 # 100K tokens
|
||||
DEFAULT_SUMMARY_KEEP_MESSAGES = 10
|
||||
DEFAULT_SUMMARY_TOOL_RESULT_TOKEN_LIMIT = 300
|
||||
DEFAULT_SUMMARY_L2_TRIGGER_RATIO = 0.4
|
||||
DEFAULT_MAX_EXECUTION_STEPS = 300
|
||||
DEFAULT_TOOL_RESULT_EVICTION_K_TOKENS = 3
|
||||
DEFAULT_YUXI_SUMMARY_PROMPT = """你是对话上下文压缩助手。
|
||||
你的任务是把下面的对话历史压缩成后续智能体继续工作所需的高价值上下文。
|
||||
|
||||
请特别保留并清晰记录:
|
||||
|
||||
## SESSION INTENT
|
||||
用户当前的主要目标、任务范围和最终交付物。
|
||||
|
||||
## USER REQUIREMENTS AND PREFERENCES
|
||||
用户明确提出的要求、偏好、禁忌、输出格式、语言风格、技术约束、验收标准,以及对实现方式的取舍意见。只记录仍然可能影响后续回答或执行的内容。
|
||||
|
||||
## PROGRESS AND DECISIONS
|
||||
已经完成的步骤、关键结论、已确认的方案、被否定的方案及原因。
|
||||
|
||||
## ARTIFACTS AND REFERENCES
|
||||
已经创建、修改、读取或需要继续关注的文件、路径、工具输出路径、线程或运行标识。保留具体路径和关键标识符。
|
||||
|
||||
## NEXT STEPS
|
||||
为了完成用户目标,后续最应该继续做的具体步骤。没有待办时写 None。
|
||||
|
||||
要求:
|
||||
- 不要逐字复述冗长工具输出;保留结论、路径和必要证据。
|
||||
- 不要编造没有出现在对话中的事实。
|
||||
- 如果存在未解决的问题或风险,明确记录。
|
||||
- 使用与用户主要对话一致的语言。
|
||||
|
||||
<messages>
|
||||
{messages}
|
||||
</messages>
|
||||
|
||||
只输出压缩后的上下文,不要添加额外说明。"""
|
||||
|
||||
|
||||
def _role_can_access(auth: str | None, role: str | None) -> bool:
|
||||
if not auth:
|
||||
return True
|
||||
if auth == "admin":
|
||||
return role in {"admin", "superadmin"}
|
||||
if auth == "superadmin":
|
||||
return role == "superadmin"
|
||||
return False
|
||||
|
||||
|
||||
def _load_workspace_agents_prompt(thread_id: str, uid: str) -> str:
|
||||
prompt_file = sandbox_workspace_agents_prompt_file(thread_id, uid)
|
||||
try:
|
||||
with prompt_file.open("rb") as buffer:
|
||||
content = buffer.read(WORKSPACE_AGENTS_PROMPT_MAX_BYTES + 1)
|
||||
except FileNotFoundError:
|
||||
return ""
|
||||
except IsADirectoryError:
|
||||
logger.warning("读取工作区 AGENTS.md 失败: 路径是目录")
|
||||
return ""
|
||||
except OSError as exc:
|
||||
logger.warning(f"读取工作区 AGENTS.md 失败: {exc}")
|
||||
return ""
|
||||
|
||||
prompt = content[:WORKSPACE_AGENTS_PROMPT_MAX_BYTES].decode("utf-8", errors="replace").strip()
|
||||
if not prompt:
|
||||
return ""
|
||||
if len(content) > WORKSPACE_AGENTS_PROMPT_MAX_BYTES:
|
||||
return f"{prompt}\n\n[AGENTS.md 内容已截断]"
|
||||
return prompt
|
||||
|
||||
|
||||
async def build_agent_input_context(
|
||||
agent_config: dict | None,
|
||||
*,
|
||||
thread_id: str,
|
||||
uid: str,
|
||||
run_id: str | None = None,
|
||||
request_id: str | None = None,
|
||||
) -> dict:
|
||||
input_context = dict(agent_config or {})
|
||||
agents_prompt = await asyncio.to_thread(_load_workspace_agents_prompt, thread_id, uid)
|
||||
|
||||
if agents_prompt:
|
||||
agents_section = f"用户工作区 agents/AGENTS.md 内容:\n{agents_prompt}"
|
||||
base_prompt = str(input_context.get("system_prompt") or "").rstrip()
|
||||
input_context["system_prompt"] = f"{base_prompt}\n\n{agents_section}" if base_prompt else agents_section
|
||||
|
||||
input_context.update({"uid": uid, "thread_id": thread_id, "run_id": run_id, "request_id": request_id})
|
||||
return input_context
|
||||
|
||||
|
||||
def filter_config_by_role(
|
||||
config_json: dict,
|
||||
role: str | None,
|
||||
context_schema: type["BaseContext"] | None = None,
|
||||
) -> dict:
|
||||
"""按 Context 字段 metadata.auth 过滤 config_json.context。"""
|
||||
if not isinstance(config_json, dict):
|
||||
return {}
|
||||
|
||||
schema = context_schema or BaseContext
|
||||
restricted_fields = {
|
||||
f.name
|
||||
for f in fields(schema)
|
||||
if f.metadata.get("auth") and not _role_can_access(str(f.metadata.get("auth")), role)
|
||||
}
|
||||
if not restricted_fields:
|
||||
return dict(config_json)
|
||||
|
||||
filtered = dict(config_json)
|
||||
context = filtered.get("context")
|
||||
if isinstance(context, dict):
|
||||
filtered["context"] = {key: value for key, value in context.items() if key not in restricted_fields}
|
||||
return filtered
|
||||
|
||||
|
||||
@dataclass(kw_only=True)
|
||||
class BaseContext:
|
||||
"""
|
||||
定义一个基础 Context 供 各类 graph 继承
|
||||
|
||||
配置优先级:
|
||||
1. 运行时配置(RunnableConfig):最高优先级,直接从函数参数传入
|
||||
2. 类默认配置:最低优先级,类中定义的默认值
|
||||
"""
|
||||
|
||||
def update(self, data: dict):
|
||||
"""更新配置字段"""
|
||||
for key, value in data.items():
|
||||
if hasattr(self, key):
|
||||
setattr(self, key, value)
|
||||
|
||||
thread_id: str = field(
|
||||
default_factory=lambda: str(uuid.uuid4()),
|
||||
metadata={"name": "线程ID", "configurable": False, "description": "用来唯一标识一个对话线程"},
|
||||
)
|
||||
|
||||
uid: str = field(
|
||||
default_factory=lambda: str(uuid.uuid4()),
|
||||
metadata={"name": "UID", "configurable": False, "description": "用来唯一标识一个用户"},
|
||||
)
|
||||
|
||||
run_id: str | None = field(
|
||||
default=None,
|
||||
metadata={"name": "运行 ID", "configurable": False, "hide": True},
|
||||
)
|
||||
|
||||
request_id: str | None = field(
|
||||
default=None,
|
||||
metadata={"name": "请求 ID", "configurable": False, "hide": True},
|
||||
)
|
||||
|
||||
system_prompt: str = field(
|
||||
default="You are a helpful assistant.",
|
||||
metadata={"name": "系统提示词", "description": "用来描述智能体的角色和行为", "kind": "prompt"},
|
||||
)
|
||||
|
||||
model: str = field(
|
||||
default="",
|
||||
metadata={
|
||||
"name": "智能体模型",
|
||||
"options": [],
|
||||
"description": "智能体的驱动模型,留空时使用系统默认模型。",
|
||||
"kind": "llm",
|
||||
},
|
||||
)
|
||||
|
||||
tools: list[str] | None = field(
|
||||
default=None,
|
||||
metadata={
|
||||
"name": "工具",
|
||||
"description": "内置的工具。默认选择当前用户可用的全部工具。",
|
||||
"type": "list",
|
||||
"kind": "tools",
|
||||
},
|
||||
)
|
||||
|
||||
knowledges: list[str] | None = field(
|
||||
default=None,
|
||||
metadata={
|
||||
"name": "知识库",
|
||||
"description": "知识库列表,可以在左侧知识库页面中创建知识库。默认选择当前用户可访问的全部知识库。",
|
||||
"type": "list",
|
||||
"kind": "knowledges",
|
||||
},
|
||||
)
|
||||
|
||||
mcps: list[str] | None = field(
|
||||
default=None,
|
||||
metadata={
|
||||
"name": "MCP服务器",
|
||||
"options": [],
|
||||
"description": (
|
||||
"MCP服务器列表,默认选择当前用户可用的全部 MCP 服务器。建议使用支持 SSE 的 MCP 服务器,"
|
||||
"如果需要使用 uvx 或 npx 运行的服务器,也请在项目外部启动 MCP 服务器,并在项目中配置 MCP 服务器。"
|
||||
),
|
||||
"type": "list",
|
||||
"kind": "mcps",
|
||||
},
|
||||
)
|
||||
|
||||
skills: list[str] | None = field(
|
||||
default=None,
|
||||
metadata={
|
||||
"name": "Skills",
|
||||
"options": [],
|
||||
"description": "可选 Skill 拓展列表,默认选择当前用户可用的全部 Skill 拓展。"
|
||||
"Skill 拓展依赖的工具和 MCP 服务器也会被自动挂载。",
|
||||
"type": "list",
|
||||
"kind": "skills",
|
||||
},
|
||||
)
|
||||
|
||||
summary_threshold: int = field(
|
||||
default=DEFAULT_SUMMARY_THRESHOLD_K,
|
||||
metadata={
|
||||
"name": "上下文摘要触发阈值 (K)",
|
||||
"description": (
|
||||
f"当上下文大小超过该值时,启用摘要功能以优化上下文使用。单位为 K,默认值为 "
|
||||
f"{DEFAULT_SUMMARY_THRESHOLD_K}K。"
|
||||
),
|
||||
"type": "number",
|
||||
"auth": "admin",
|
||||
},
|
||||
)
|
||||
|
||||
summary_keep_messages: int = field(
|
||||
default=DEFAULT_SUMMARY_KEEP_MESSAGES,
|
||||
metadata={
|
||||
"name": "摘要后保留消息数",
|
||||
"description": (
|
||||
f"上下文摘要触发后,除摘要消息外保留最近的消息数量,默认 {DEFAULT_SUMMARY_KEEP_MESSAGES} 条。"
|
||||
),
|
||||
"type": "number",
|
||||
"auth": "admin",
|
||||
},
|
||||
)
|
||||
|
||||
summary_prompt: str = field(
|
||||
default=DEFAULT_YUXI_SUMMARY_PROMPT,
|
||||
metadata={
|
||||
"name": "上下文摘要提示词",
|
||||
"description": "触发上下文摘要时使用的提示词,必须能接收 {messages} 作为待摘要消息占位符。",
|
||||
"type": "string",
|
||||
"kind": "prompt",
|
||||
"auth": "admin",
|
||||
},
|
||||
)
|
||||
|
||||
summary_tool_result_token_limit: int = field(
|
||||
default=DEFAULT_SUMMARY_TOOL_RESULT_TOKEN_LIMIT,
|
||||
metadata={
|
||||
"name": "摘要工具结果 token 上限",
|
||||
"description": (
|
||||
"上下文摘要 L1 清洗历史工具结果时,超过该 token 数的 ToolMessage 会写入 outputs,"
|
||||
"并在上下文中保留不超过该 token 数的预览;未超过则保持原样。默认 "
|
||||
f"{DEFAULT_SUMMARY_TOOL_RESULT_TOKEN_LIMIT}。"
|
||||
),
|
||||
"type": "number",
|
||||
"auth": "admin",
|
||||
},
|
||||
)
|
||||
|
||||
summary_l2_trigger_ratio: float = field(
|
||||
default=DEFAULT_SUMMARY_L2_TRIGGER_RATIO,
|
||||
metadata={
|
||||
"name": "L2 摘要触发比例",
|
||||
"description": (
|
||||
"L1 结构精简后,剩余上下文超过 摘要触发阈值 * 该比例 时才进入 L2 summary。"
|
||||
"建议范围 0.1 到 1.0,值越小越容易触发 L2,默认 "
|
||||
f"{DEFAULT_SUMMARY_L2_TRIGGER_RATIO}。"
|
||||
),
|
||||
"type": "number",
|
||||
"auth": "admin",
|
||||
},
|
||||
)
|
||||
|
||||
max_execution_steps: int = field(
|
||||
default=DEFAULT_MAX_EXECUTION_STEPS,
|
||||
metadata={
|
||||
"name": "最大执行步数",
|
||||
"description": (
|
||||
"单次 Agent 运行允许的最大 LangGraph 执行步数,对应 recursion_limit,默认 "
|
||||
f"{DEFAULT_MAX_EXECUTION_STEPS}。"
|
||||
),
|
||||
"type": "number",
|
||||
"auth": "admin",
|
||||
},
|
||||
)
|
||||
|
||||
model_retry_times: int = field(
|
||||
default=2,
|
||||
metadata={
|
||||
"name": "模型重试次数",
|
||||
"description": "模型调用失败时的最大重试次数,默认值为 2。",
|
||||
"type": "number",
|
||||
"auth": "admin",
|
||||
},
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def get_configurable_items(cls, user_role: str | None = None):
|
||||
"""实现一个可配置的参数列表,在 UI 上配置时使用"""
|
||||
configurable_items = {}
|
||||
for f in fields(cls):
|
||||
if f.init and not f.metadata.get("hide", False):
|
||||
if user_role is not None and not _role_can_access(f.metadata.get("auth"), user_role):
|
||||
continue
|
||||
if f.metadata.get("configurable", True):
|
||||
type_name = cls._get_type_name(f.type)
|
||||
|
||||
options = f.metadata.get("options", [])
|
||||
if callable(options):
|
||||
options = options()
|
||||
|
||||
configurable_items[f.name] = {
|
||||
"type": f.metadata.get("type", type_name),
|
||||
"name": f.metadata.get("name", f.name),
|
||||
"options": options,
|
||||
"default": f.default
|
||||
if f.default is not MISSING
|
||||
else f.default_factory()
|
||||
if f.default_factory is not MISSING
|
||||
else None,
|
||||
"description": f.metadata.get("description", ""),
|
||||
"kind": f.metadata.get("kind", ""),
|
||||
}
|
||||
|
||||
return configurable_items
|
||||
|
||||
@classmethod
|
||||
def _get_type_name(cls, field_type) -> str:
|
||||
"""获取类型名称"""
|
||||
origin = get_origin(field_type)
|
||||
if origin is not None:
|
||||
if hasattr(origin, "__name__"):
|
||||
return origin.__name__
|
||||
return str(origin)
|
||||
elif hasattr(field_type, "__name__"):
|
||||
return field_type.__name__
|
||||
else:
|
||||
return str(field_type)
|
||||
|
||||
def update_from_dict(self, data: dict):
|
||||
"""从字典更新配置字段"""
|
||||
for key, value in data.items():
|
||||
if hasattr(self, key):
|
||||
setattr(self, key, value)
|
||||
|
||||
|
||||
_DEFAULT_ALL_CONTEXT_FIELDS = frozenset({"tools", "knowledges", "mcps", "skills"})
|
||||
_EMPTY_ALL_CONTEXT_FIELDS = frozenset({"subagents"})
|
||||
_AGENT_RESOURCE_FIELDS = _DEFAULT_ALL_CONTEXT_FIELDS | _EMPTY_ALL_CONTEXT_FIELDS
|
||||
|
||||
|
||||
def _normalize_selected_resource_keys(value: Any, available: list[str]) -> list[str]:
|
||||
if not isinstance(value, list):
|
||||
return []
|
||||
|
||||
allowed = set(available)
|
||||
normalized: list[str] = []
|
||||
seen: set[str] = set()
|
||||
for item in value:
|
||||
if not isinstance(item, str):
|
||||
continue
|
||||
key = item.strip()
|
||||
if not key or key in seen or key not in allowed:
|
||||
continue
|
||||
seen.add(key)
|
||||
normalized.append(key)
|
||||
return normalized
|
||||
|
||||
|
||||
def _resource_fields_requiring_available_keys(normalized: dict, resource_fields: set[str]) -> set[str]:
|
||||
fields_to_load: set[str] = set()
|
||||
for field_name in resource_fields:
|
||||
current = normalized.get(field_name)
|
||||
if current is None:
|
||||
if field_name in _DEFAULT_ALL_CONTEXT_FIELDS | _EMPTY_ALL_CONTEXT_FIELDS:
|
||||
fields_to_load.add(field_name)
|
||||
else:
|
||||
normalized[field_name] = []
|
||||
elif field_name in _EMPTY_ALL_CONTEXT_FIELDS and current == []:
|
||||
normalized[field_name] = None
|
||||
fields_to_load.add(field_name)
|
||||
elif isinstance(current, list) and current:
|
||||
fields_to_load.add(field_name)
|
||||
else:
|
||||
normalized[field_name] = []
|
||||
return fields_to_load
|
||||
|
||||
|
||||
def _resource_option(key: Any, name: Any = None, description: Any = None) -> dict[str, str]:
|
||||
key_value = str(key)
|
||||
return {
|
||||
"key": key_value,
|
||||
"name": str(name or key_value),
|
||||
"description": str(description or ""),
|
||||
}
|
||||
|
||||
|
||||
async def resolve_agent_resource_options(
|
||||
resource_fields: set[str] | None = None,
|
||||
*,
|
||||
db,
|
||||
user,
|
||||
) -> dict[str, list[dict[str, str]]]:
|
||||
fields_to_load = _AGENT_RESOURCE_FIELDS if resource_fields is None else resource_fields
|
||||
if not fields_to_load:
|
||||
return {}
|
||||
|
||||
options: dict[str, list[dict[str, str]]] = {}
|
||||
|
||||
if "tools" in fields_to_load:
|
||||
from yuxi.agents.toolkits.service import get_tool_metadata
|
||||
|
||||
options["tools"] = [
|
||||
_resource_option(tool["slug"], tool.get("name"), tool.get("description"))
|
||||
for tool in get_tool_metadata(category="buildin")
|
||||
if tool.get("slug")
|
||||
]
|
||||
if "knowledges" in fields_to_load:
|
||||
from yuxi.knowledge import knowledge_base
|
||||
|
||||
databases = (await knowledge_base.get_databases_by_user(user)).get("databases", [])
|
||||
options["knowledges"] = [
|
||||
_resource_option(item.get("kb_id"), item.get("name"), item.get("description"))
|
||||
for item in databases
|
||||
if isinstance(item, dict) and item.get("kb_id")
|
||||
]
|
||||
if "mcps" in fields_to_load:
|
||||
from yuxi.agents.mcp.service import get_all_mcp_servers
|
||||
|
||||
servers = await get_all_mcp_servers(db)
|
||||
options["mcps"] = [
|
||||
_resource_option(server.slug, server.name, server.description)
|
||||
for server in servers
|
||||
if server.enabled and server.slug
|
||||
]
|
||||
if "skills" in fields_to_load:
|
||||
from yuxi.agents.skills.service import list_accessible_skills
|
||||
|
||||
skills = await list_accessible_skills(db, user)
|
||||
options["skills"] = [
|
||||
_resource_option(skill.slug, skill.name, skill.description) for skill in skills if skill.slug
|
||||
]
|
||||
if "subagents" in fields_to_load:
|
||||
from yuxi.repositories.agent_repository import AgentRepository
|
||||
|
||||
subagents = await AgentRepository(db).list_visible_subagents(user=user)
|
||||
options["subagents"] = [
|
||||
_resource_option(agent.slug, agent.name, agent.description) for agent in subagents if agent.slug
|
||||
]
|
||||
|
||||
return options
|
||||
|
||||
|
||||
async def normalize_agent_context_config(
|
||||
context: dict | None,
|
||||
*,
|
||||
db,
|
||||
user,
|
||||
context_schema: type[BaseContext] | None = None,
|
||||
) -> dict:
|
||||
schema = context_schema or BaseContext
|
||||
raw_context = dict(context) if isinstance(context, dict) else {}
|
||||
filtered = filter_config_by_role({"context": raw_context}, getattr(user, "role", None), schema)
|
||||
normalized = dict(filtered.get("context") or {})
|
||||
field_names = {item.name for item in fields(schema)}
|
||||
resource_fields = _AGENT_RESOURCE_FIELDS & field_names
|
||||
if not resource_fields:
|
||||
return normalized
|
||||
|
||||
fields_to_load = _resource_fields_requiring_available_keys(normalized, resource_fields)
|
||||
if not fields_to_load:
|
||||
return normalized
|
||||
|
||||
resource_options = await resolve_agent_resource_options(fields_to_load, db=db, user=user)
|
||||
available = {
|
||||
field_name: [option["key"] for option in field_options]
|
||||
for field_name, field_options in resource_options.items()
|
||||
}
|
||||
|
||||
for field_name, available_keys in available.items():
|
||||
current = normalized.get(field_name)
|
||||
if current is None:
|
||||
normalized[field_name] = available_keys
|
||||
else:
|
||||
normalized[field_name] = _normalize_selected_resource_keys(current, available_keys)
|
||||
|
||||
return normalized
|
||||
|
||||
|
||||
async def prepare_agent_runtime_context(
|
||||
context: BaseContext,
|
||||
*,
|
||||
context_schema: type[BaseContext] | None = None,
|
||||
) -> BaseContext:
|
||||
"""准备 Agent 运行时上下文,主要是根据 context 中的 uid 加载用户可访问的资源列表,并进行规范化处理。"""
|
||||
schema = context_schema or type(context)
|
||||
uid = str(getattr(context, "uid", "") or "").strip()
|
||||
if not uid:
|
||||
return context
|
||||
|
||||
from yuxi.agents.backends.knowledge_base_backend import resolve_visible_knowledge_bases_for_context
|
||||
from yuxi.agents.middlewares.skills import resolve_runtime_skills_for_context
|
||||
from yuxi.repositories.user_repository import UserRepository
|
||||
from yuxi.storage.postgres.manager import pg_manager
|
||||
|
||||
resource_fields = _AGENT_RESOURCE_FIELDS
|
||||
async with pg_manager.get_async_session_context() as db:
|
||||
user = await UserRepository().get_by_uid_with_db(db, uid)
|
||||
if user is None:
|
||||
for field_name in resource_fields:
|
||||
if hasattr(context, field_name):
|
||||
setattr(context, field_name, [])
|
||||
setattr(context, "_visible_knowledge_bases", [])
|
||||
setattr(context, "_prompt_skills", [])
|
||||
setattr(context, "_readable_skills", [])
|
||||
setattr(context, "_runtime_skill_metadata", {})
|
||||
setattr(context, "_runtime_skill_dependency_map", {})
|
||||
return context
|
||||
|
||||
raw_resources = {
|
||||
field_name: getattr(context, field_name, None)
|
||||
for field_name in resource_fields
|
||||
if hasattr(context, field_name)
|
||||
}
|
||||
normalized = await normalize_agent_context_config(
|
||||
raw_resources,
|
||||
db=db,
|
||||
user=user,
|
||||
context_schema=schema,
|
||||
)
|
||||
for field_name in resource_fields:
|
||||
if hasattr(context, field_name):
|
||||
setattr(context, field_name, normalized.get(field_name, []))
|
||||
|
||||
await resolve_visible_knowledge_bases_for_context(context)
|
||||
skill_scope = await resolve_runtime_skills_for_context(context, db=db, user=user)
|
||||
context.skills = skill_scope["context_skills"]
|
||||
setattr(context, "_prompt_skills", skill_scope["prompt_skills"])
|
||||
setattr(context, "_readable_skills", skill_scope["readable_skills"])
|
||||
setattr(context, "_runtime_skill_metadata", skill_scope["runtime_skill_metadata"])
|
||||
setattr(context, "_runtime_skill_dependency_map", skill_scope["runtime_skill_dependency_map"])
|
||||
|
||||
return context
|
||||
@@ -0,0 +1,83 @@
|
||||
"""MCP 服务器数据访问层 - Repository"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from sqlalchemy import select
|
||||
|
||||
from yuxi.storage.postgres.manager import pg_manager
|
||||
from yuxi.storage.postgres.models_business import MCPServer
|
||||
|
||||
|
||||
class MCPServerRepository:
|
||||
"""MCP 服务器数据访问层"""
|
||||
|
||||
async def get_by_slug(self, slug: str) -> MCPServer | None:
|
||||
"""根据名称获取 MCP 服务器"""
|
||||
async with pg_manager.get_async_session_context() as session:
|
||||
result = await session.execute(select(MCPServer).where(MCPServer.slug == slug))
|
||||
return result.scalar_one_or_none()
|
||||
|
||||
async def list(self) -> list[MCPServer]:
|
||||
"""获取所有 MCP 服务器"""
|
||||
async with pg_manager.get_async_session_context() as session:
|
||||
result = await session.execute(select(MCPServer))
|
||||
return list(result.scalars().all())
|
||||
|
||||
async def list_enabled(self) -> list[MCPServer]:
|
||||
"""获取所有启用的 MCP 服务器"""
|
||||
async with pg_manager.get_async_session_context() as session:
|
||||
result = await session.execute(select(MCPServer).where(MCPServer.enabled == 1))
|
||||
return list(result.scalars().all())
|
||||
|
||||
async def create(self, data: dict[str, Any]) -> MCPServer:
|
||||
"""创建 MCP 服务器"""
|
||||
async with pg_manager.get_async_session_context() as session:
|
||||
server = MCPServer(**data)
|
||||
session.add(server)
|
||||
return server
|
||||
|
||||
async def update(self, slug: str, data: dict[str, Any]) -> MCPServer | None:
|
||||
"""更新 MCP 服务器"""
|
||||
async with pg_manager.get_async_session_context() as session:
|
||||
result = await session.execute(select(MCPServer).where(MCPServer.slug == slug))
|
||||
server = result.scalar_one_or_none()
|
||||
if server is None:
|
||||
return None
|
||||
for key, value in data.items():
|
||||
if key != "slug":
|
||||
setattr(server, key, value)
|
||||
return server
|
||||
|
||||
async def delete(self, slug: str) -> bool:
|
||||
"""删除 MCP 服务器"""
|
||||
async with pg_manager.get_async_session_context() as session:
|
||||
result = await session.execute(select(MCPServer).where(MCPServer.slug == slug))
|
||||
server = result.scalar_one_or_none()
|
||||
if server is None:
|
||||
return False
|
||||
await session.delete(server)
|
||||
return True
|
||||
|
||||
async def upsert(self, data: dict[str, Any]) -> MCPServer:
|
||||
"""插入或更新 MCP 服务器"""
|
||||
slug = data.get("slug")
|
||||
async with pg_manager.get_async_session_context() as session:
|
||||
result = await session.execute(select(MCPServer).where(MCPServer.slug == slug))
|
||||
existing = result.scalar_one_or_none()
|
||||
if existing is None:
|
||||
server = MCPServer(**data)
|
||||
session.add(server)
|
||||
else:
|
||||
for key, value in data.items():
|
||||
if key != "slug":
|
||||
setattr(existing, key, value)
|
||||
server = existing
|
||||
return server
|
||||
|
||||
async def exists_by_slug(self, slug: str) -> bool:
|
||||
"""检查 MCP 服务器是否存在"""
|
||||
async with pg_manager.get_async_session_context() as session:
|
||||
result = await session.execute(select(MCPServer.id).where(MCPServer.slug == slug))
|
||||
return result.scalar_one_or_none() is not None
|
||||
@@ -0,0 +1,609 @@
|
||||
"""MCP Service - Unified business logic and state management for MCP.
|
||||
|
||||
Responsibilities:
|
||||
- Server configuration CRUD operations
|
||||
- Built-in configuration synchronization (Code <-> Database)
|
||||
- Unified entry point for Agent tool retrieval (auto-filtering disabled_tools)
|
||||
- MCP Client and Tools management (formerly in agents/common/mcp.py)
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import hashlib
|
||||
import json
|
||||
import re
|
||||
from collections.abc import Callable
|
||||
from typing import Any, cast
|
||||
|
||||
from langchain_mcp_adapters.client import MultiServerMCPClient
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from yuxi.storage.postgres.models_business import MCPServer
|
||||
from yuxi.utils import logger
|
||||
|
||||
# =============================================================================
|
||||
# === Global Cache & State ===
|
||||
# =============================================================================
|
||||
|
||||
# Global Lock for MCP state
|
||||
_mcp_lock = asyncio.Lock()
|
||||
|
||||
# 本地仅缓存工具对象。配置始终以数据库为准,每次按 server_slug 现查。
|
||||
# cache key 使用 server_slug:config_hash,当配置变化时会自然失效。
|
||||
_mcp_tools_cache: dict[str, list[Callable[..., Any]]] = {}
|
||||
|
||||
# MCP tools statistics (for reporting enabled/disabled counts)
|
||||
_mcp_tools_stats: dict[str, dict[str, int]] = {}
|
||||
_UNSET = object()
|
||||
|
||||
# Default MCP Server configurations (Imported to DB on first run)
|
||||
_DEFAULT_MCP_SERVERS = {
|
||||
"mcp-server-chart": {
|
||||
"command": "npx",
|
||||
"args": ["-y", "@antv/mcp-server-chart"],
|
||||
"transport": "stdio",
|
||||
"description": "图表生成工具,支持生成各类图表(柱状图、折线图、饼图等)",
|
||||
"icon": "📊",
|
||||
"tags": ["内置", "图表"],
|
||||
},
|
||||
}
|
||||
|
||||
_RETIRED_BUILTIN_MCP_SERVER_SLUGS = ("sequentialthinking",)
|
||||
|
||||
_SYNCED_MCP_FIELDS = (
|
||||
"description",
|
||||
"transport",
|
||||
"url",
|
||||
"command",
|
||||
"args",
|
||||
"env",
|
||||
"headers",
|
||||
"timeout",
|
||||
"sse_read_timeout",
|
||||
"tags",
|
||||
"icon",
|
||||
)
|
||||
|
||||
# =============================================================================
|
||||
# === Core Logic (Moved from agents/common/mcp.py) ===
|
||||
# =============================================================================
|
||||
|
||||
|
||||
async def ensure_builtin_mcp_servers_in_db() -> None:
|
||||
"""Ensure built-in MCP server definitions exist in the database."""
|
||||
from yuxi.storage.postgres.manager import pg_manager
|
||||
|
||||
try:
|
||||
async with pg_manager.get_async_session_context() as session:
|
||||
any_changed = False
|
||||
for slug in _RETIRED_BUILTIN_MCP_SERVER_SLUGS:
|
||||
result = await session.execute(
|
||||
select(MCPServer).filter(MCPServer.slug == slug, MCPServer.created_by == "system")
|
||||
)
|
||||
retired = result.scalar_one_or_none()
|
||||
if retired:
|
||||
await session.delete(retired)
|
||||
clear_mcp_server_tools_cache(slug)
|
||||
any_changed = True
|
||||
logger.info(f"Removed retired built-in MCP server '{slug}' from database")
|
||||
|
||||
for slug, config in _DEFAULT_MCP_SERVERS.items():
|
||||
result = await session.execute(select(MCPServer).filter(MCPServer.slug == slug))
|
||||
existing = result.scalar_one_or_none()
|
||||
if not existing:
|
||||
session.add(
|
||||
MCPServer(
|
||||
slug=slug,
|
||||
name=config.get("name", slug),
|
||||
description=config.get("description"),
|
||||
transport=config["transport"],
|
||||
url=config.get("url"),
|
||||
command=config.get("command"),
|
||||
args=config.get("args"),
|
||||
env=config.get("env"),
|
||||
headers=config.get("headers"),
|
||||
timeout=config.get("timeout"),
|
||||
sse_read_timeout=config.get("sse_read_timeout"),
|
||||
tags=config.get("tags"),
|
||||
icon=config.get("icon"),
|
||||
enabled=0,
|
||||
created_by="system",
|
||||
updated_by="system",
|
||||
)
|
||||
)
|
||||
any_changed = True
|
||||
logger.info(f"Added built-in MCP server '{slug}' to database")
|
||||
continue
|
||||
|
||||
server_changed = False
|
||||
for field in _SYNCED_MCP_FIELDS:
|
||||
next_value = config.get(field)
|
||||
if getattr(existing, field) != next_value:
|
||||
setattr(existing, field, next_value)
|
||||
server_changed = True
|
||||
if server_changed:
|
||||
existing.updated_by = "system"
|
||||
any_changed = True
|
||||
|
||||
if any_changed:
|
||||
await session.commit()
|
||||
|
||||
except Exception as e:
|
||||
logger.exception(f"Failed to ensure builtin MCP servers in database: {e}")
|
||||
|
||||
|
||||
async def get_mcp_client(
|
||||
server_configs: dict[str, Any] | None = None,
|
||||
) -> MultiServerMCPClient | None:
|
||||
"""Initializes an MCP client with the given server configurations."""
|
||||
try:
|
||||
client = MultiServerMCPClient(server_configs) # pyright: ignore[reportArgumentType]
|
||||
logger.info(f"Initialized MCP client with servers: {list(server_configs.keys())}")
|
||||
return client
|
||||
except Exception as e:
|
||||
logger.error("Failed to initialize MCP client: {}", e)
|
||||
return None
|
||||
|
||||
|
||||
def to_camel_case(s: str) -> str:
|
||||
"""Convert string to lowerCamelCase."""
|
||||
|
||||
# Handle - and _
|
||||
s = re.sub(r"[-_]+(.)", lambda m: m.group(1).upper(), s)
|
||||
# Lowercase first letter
|
||||
if len(s) > 0:
|
||||
s = s[0].lower() + s[1:]
|
||||
return s
|
||||
|
||||
|
||||
async def _load_enabled_mcp_server_configs(
|
||||
*,
|
||||
names: list[str] | None = None,
|
||||
db: AsyncSession | None = None,
|
||||
) -> dict[str, dict[str, Any]]:
|
||||
"""Load enabled MCP server configs directly from the database."""
|
||||
if db is not None:
|
||||
stmt = select(MCPServer).where(MCPServer.enabled == 1)
|
||||
if names:
|
||||
stmt = stmt.where(MCPServer.slug.in_(names))
|
||||
result = await db.execute(stmt)
|
||||
servers = result.scalars().all()
|
||||
return {server.slug: server.to_mcp_config() for server in servers}
|
||||
|
||||
from yuxi.storage.postgres.manager import pg_manager
|
||||
|
||||
async with pg_manager.get_async_session_context() as session:
|
||||
return await _load_enabled_mcp_server_configs(names=names, db=session)
|
||||
|
||||
|
||||
async def get_enabled_mcp_server_config(server_slug: str, *, db: AsyncSession | None = None) -> dict[str, Any] | None:
|
||||
"""Get the latest enabled MCP server config from the database."""
|
||||
configs = await _load_enabled_mcp_server_configs(names=[server_slug], db=db)
|
||||
return configs.get(server_slug)
|
||||
|
||||
|
||||
async def get_enabled_mcp_server_slugs(*, db: AsyncSession | None = None) -> list[str]:
|
||||
"""Get enabled MCP server slugs from the database."""
|
||||
if db is not None:
|
||||
result = await db.execute(select(MCPServer.slug).where(MCPServer.enabled == 1))
|
||||
return [name for name in result.scalars().all() if isinstance(name, str)]
|
||||
|
||||
from yuxi.storage.postgres.manager import pg_manager
|
||||
|
||||
async with pg_manager.get_async_session_context() as session:
|
||||
return await get_enabled_mcp_server_slugs(db=session)
|
||||
|
||||
|
||||
async def get_mcp_tools(
|
||||
server_slug: str,
|
||||
additional_servers: dict[str, dict[str, Any]] | None = None,
|
||||
disabled_tools: list[str] = None,
|
||||
cache: bool = True,
|
||||
force_refresh: bool = False,
|
||||
) -> list[Callable[..., Any]]:
|
||||
"""Get MCP tools for a specific server.
|
||||
|
||||
Architecture:
|
||||
1. Fetching: Connects to MCP server to get ALL tools.
|
||||
2. Caching: Stores the FULL, UNFILTERED list of tools in `_mcp_tools_cache`.
|
||||
3. Filtering: Filters the return value based on `disabled_tools` argument.
|
||||
|
||||
Args:
|
||||
server_slug: Server slug
|
||||
additional_servers: Additional server configurations
|
||||
disabled_tools: List of tool names to filter out from the RETURN value (does not affect cache)
|
||||
cache: Whether to use/update the cache (default: True)
|
||||
force_refresh: Whether to force a refresh from the server (default: False)
|
||||
"""
|
||||
if additional_servers and server_slug in additional_servers:
|
||||
server_config = additional_servers[server_slug]
|
||||
else:
|
||||
server_config = await get_enabled_mcp_server_config(server_slug)
|
||||
|
||||
if server_config is None:
|
||||
logger.warning(f"MCP server '{server_slug}' not found in database or disabled")
|
||||
return []
|
||||
|
||||
# 配置 hash 直接基于完整配置生成。只要数据库中的配置发生变化,
|
||||
# 本地工具缓存 key 就会变化,从而自然触发重建。
|
||||
config_payload = json.dumps(server_config, sort_keys=True, ensure_ascii=True, separators=(",", ":"))
|
||||
config_hash = hashlib.sha256(config_payload.encode("utf-8")).hexdigest()[:16]
|
||||
cache_key = f"{server_slug}:{config_hash}"
|
||||
|
||||
all_processed_tools: list[Callable[..., Any]] = []
|
||||
|
||||
async with _mcp_lock:
|
||||
if not force_refresh and cache and cache_key in _mcp_tools_cache:
|
||||
all_processed_tools = _mcp_tools_cache[cache_key]
|
||||
|
||||
if not all_processed_tools:
|
||||
try:
|
||||
# disabled_tools 只影响返回值过滤,不参与 MCP client 建连参数。
|
||||
client_config = {k: v for k, v in server_config.items() if k not in ("disabled_tools",)}
|
||||
|
||||
client = await get_mcp_client({server_slug: client_config})
|
||||
if client is None:
|
||||
return []
|
||||
|
||||
raw_tools = cast(list[Any], await client.get_tools())
|
||||
|
||||
server_cc = to_camel_case(server_slug)
|
||||
for tool in raw_tools:
|
||||
original_name = tool.name
|
||||
tool_cc = to_camel_case(original_name)
|
||||
unique_id = f"mcp__{server_cc}__{tool_cc}"
|
||||
|
||||
if tool.metadata is None:
|
||||
tool.metadata = {}
|
||||
tool.metadata["id"] = unique_id
|
||||
# 开启错误处理,防止工具调用抛出 ToolException 时击穿服务
|
||||
tool.handle_tool_error = True
|
||||
all_processed_tools.append(tool)
|
||||
|
||||
if cache:
|
||||
async with _mcp_lock:
|
||||
stale_keys = [
|
||||
key for key in _mcp_tools_cache if key.startswith(f"{server_slug}:") and key != cache_key
|
||||
]
|
||||
for stale_key in stale_keys:
|
||||
_mcp_tools_cache.pop(stale_key, None)
|
||||
_mcp_tools_cache[cache_key] = all_processed_tools
|
||||
|
||||
global_config_disabled = server_config.get("disabled_tools") or []
|
||||
enabled_count = len([t for t in all_processed_tools if t.name not in global_config_disabled])
|
||||
_mcp_tools_stats[server_slug] = {
|
||||
"total": len(all_processed_tools),
|
||||
"enabled": enabled_count,
|
||||
"disabled": len(all_processed_tools) - enabled_count,
|
||||
}
|
||||
|
||||
logger.info(
|
||||
f"Refreshed MCP tools cache for '{server_slug}' with key '{cache_key}': "
|
||||
f"{len(all_processed_tools)} tools loaded."
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.exception(f"Failed to load tools from MCP server '{server_slug}': {e}")
|
||||
return []
|
||||
|
||||
# 3. Filtering (Apply to Return Value Only)
|
||||
if disabled_tools:
|
||||
filtered_tools = [t for t in all_processed_tools if t.name not in disabled_tools]
|
||||
logger.debug(
|
||||
f"Returning {len(filtered_tools)}/{len(all_processed_tools)} tools for '{server_slug}' "
|
||||
f"(filtered {len(disabled_tools)} by argument)"
|
||||
)
|
||||
return filtered_tools
|
||||
|
||||
return all_processed_tools
|
||||
|
||||
|
||||
async def get_tools_from_all_servers() -> list[Callable[..., Any]]:
|
||||
"""Get all tools from all configured MCP servers."""
|
||||
server_configs = await _load_enabled_mcp_server_configs()
|
||||
all_tools = []
|
||||
for server_slug in server_configs:
|
||||
tools = await get_mcp_tools(server_slug, additional_servers=server_configs)
|
||||
all_tools.extend(tools)
|
||||
return all_tools
|
||||
|
||||
|
||||
def clear_mcp_cache() -> None:
|
||||
"""Clear the MCP tools cache (useful for testing)."""
|
||||
global _mcp_tools_cache, _mcp_tools_stats
|
||||
_mcp_tools_cache = {}
|
||||
_mcp_tools_stats = {}
|
||||
|
||||
|
||||
def clear_mcp_server_tools_cache(server_slug: str) -> None:
|
||||
"""Clear the tools cache for a specific MCP server."""
|
||||
global _mcp_tools_cache, _mcp_tools_stats
|
||||
server_prefix = f"{server_slug}:"
|
||||
stale_keys = [key for key in _mcp_tools_cache if key.startswith(server_prefix)]
|
||||
for stale_key in stale_keys:
|
||||
_mcp_tools_cache.pop(stale_key, None)
|
||||
_mcp_tools_stats.pop(server_slug, None)
|
||||
logger.info(f"Cleared tools cache for MCP server '{server_slug}'")
|
||||
|
||||
|
||||
def get_mcp_tools_stats(server_slug: str) -> dict[str, int] | None:
|
||||
"""Get tools statistics for a MCP server.
|
||||
|
||||
Returns:
|
||||
dict with 'total', 'enabled', 'disabled' counts, or None if not available
|
||||
"""
|
||||
return _mcp_tools_stats.get(server_slug)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# === Server Config CRUD ===
|
||||
# =============================================================================
|
||||
|
||||
|
||||
async def get_mcp_server(db: AsyncSession, slug: str) -> MCPServer | None:
|
||||
"""Get single server configuration by slug."""
|
||||
result = await db.execute(select(MCPServer).filter(MCPServer.slug == slug))
|
||||
return result.scalar_one_or_none()
|
||||
|
||||
|
||||
async def get_all_mcp_servers(db: AsyncSession) -> list[MCPServer]:
|
||||
"""Get all server configurations."""
|
||||
result = await db.execute(select(MCPServer))
|
||||
return list(result.scalars().all())
|
||||
|
||||
|
||||
async def create_mcp_server(
|
||||
db: AsyncSession,
|
||||
slug: str,
|
||||
name: str,
|
||||
transport: str,
|
||||
url: str = None,
|
||||
command: str = None,
|
||||
args: list = None,
|
||||
env: dict = None,
|
||||
description: str = None,
|
||||
headers: dict = None,
|
||||
timeout: int = None,
|
||||
sse_read_timeout: int = None,
|
||||
tags: list = None,
|
||||
icon: str = None,
|
||||
created_by: str = None,
|
||||
) -> MCPServer:
|
||||
"""Create server."""
|
||||
existing = await get_mcp_server(db, slug)
|
||||
if existing:
|
||||
raise ValueError(f"Server slug '{slug}' already exists")
|
||||
|
||||
server = MCPServer(
|
||||
slug=slug,
|
||||
name=name,
|
||||
description=description,
|
||||
transport=transport,
|
||||
url=url,
|
||||
command=command,
|
||||
args=args,
|
||||
env=env,
|
||||
headers=headers,
|
||||
timeout=timeout,
|
||||
sse_read_timeout=sse_read_timeout,
|
||||
tags=tags,
|
||||
icon=icon,
|
||||
enabled=1,
|
||||
created_by=created_by,
|
||||
updated_by=created_by,
|
||||
)
|
||||
db.add(server)
|
||||
await db.commit()
|
||||
await db.refresh(server)
|
||||
|
||||
clear_mcp_server_tools_cache(slug)
|
||||
|
||||
logger.info(f"Created MCP server '{slug}'")
|
||||
return server
|
||||
|
||||
|
||||
async def update_mcp_server(
|
||||
db: AsyncSession,
|
||||
slug: str,
|
||||
name: str = None,
|
||||
description: str = None,
|
||||
transport: str = None,
|
||||
url: str = None,
|
||||
command: str = None,
|
||||
args: list = None,
|
||||
env: Any = _UNSET,
|
||||
headers: dict = None,
|
||||
timeout: int = None,
|
||||
sse_read_timeout: int = None,
|
||||
tags: list = None,
|
||||
icon: str = None,
|
||||
updated_by: str = None,
|
||||
) -> MCPServer:
|
||||
"""Update server configuration."""
|
||||
server = await get_mcp_server(db, slug)
|
||||
if not server:
|
||||
raise ValueError(f"Server '{slug}' does not exist")
|
||||
|
||||
if name is not None:
|
||||
server.name = name
|
||||
if description is not None:
|
||||
server.description = description
|
||||
if transport is not None:
|
||||
server.transport = transport
|
||||
if url is not None:
|
||||
server.url = url
|
||||
if command is not None:
|
||||
server.command = command
|
||||
if args is not None:
|
||||
server.args = args
|
||||
if env is not _UNSET:
|
||||
server.env = env
|
||||
if headers is not None:
|
||||
server.headers = headers
|
||||
if timeout is not None:
|
||||
server.timeout = timeout
|
||||
if sse_read_timeout is not None:
|
||||
server.sse_read_timeout = sse_read_timeout
|
||||
if tags is not None:
|
||||
server.tags = tags
|
||||
if icon is not None:
|
||||
server.icon = icon
|
||||
if updated_by is not None:
|
||||
server.updated_by = updated_by
|
||||
|
||||
await db.commit()
|
||||
await db.refresh(server)
|
||||
|
||||
clear_mcp_server_tools_cache(slug)
|
||||
|
||||
logger.info(f"Updated MCP server '{slug}'")
|
||||
return server
|
||||
|
||||
|
||||
async def delete_mcp_server(db: AsyncSession, slug: str) -> bool:
|
||||
"""Delete server."""
|
||||
server = await get_mcp_server(db, slug)
|
||||
if not server:
|
||||
return False
|
||||
|
||||
await db.delete(server)
|
||||
await db.commit()
|
||||
|
||||
clear_mcp_server_tools_cache(slug)
|
||||
|
||||
logger.info(f"Deleted MCP server '{slug}'")
|
||||
return True
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# === Tool Management ===
|
||||
# =============================================================================
|
||||
|
||||
|
||||
async def set_server_enabled(
|
||||
db: AsyncSession, slug: str, enabled: bool, updated_by: str = None
|
||||
) -> tuple[bool, MCPServer]:
|
||||
"""Set server enabled status."""
|
||||
server = await get_mcp_server(db, slug)
|
||||
if not server:
|
||||
raise ValueError(f"Server '{slug}' does not exist")
|
||||
|
||||
server.enabled = 1 if enabled else 0
|
||||
if updated_by is not None:
|
||||
server.updated_by = updated_by
|
||||
await db.commit()
|
||||
|
||||
is_enabled = bool(server.enabled)
|
||||
clear_mcp_server_tools_cache(slug)
|
||||
|
||||
logger.info(f"Set MCP server '{slug}' enabled={is_enabled}")
|
||||
return is_enabled, server
|
||||
|
||||
|
||||
async def toggle_tool_enabled(
|
||||
db: AsyncSession,
|
||||
server_slug: str,
|
||||
tool_name: str,
|
||||
updated_by: str = None,
|
||||
) -> tuple[bool, MCPServer]:
|
||||
"""Toggle single tool enabled status.
|
||||
|
||||
Args:
|
||||
db: Database session
|
||||
server_slug: Server slug
|
||||
tool_name: Tool name
|
||||
updated_by: Updater
|
||||
|
||||
Returns:
|
||||
(enabled, server): Tool enabled status and updated server object
|
||||
"""
|
||||
server = await get_mcp_server(db, server_slug)
|
||||
if not server:
|
||||
raise ValueError(f"Server '{server_slug}' does not exist")
|
||||
|
||||
disabled_tools = list(server.disabled_tools or [])
|
||||
|
||||
if tool_name in disabled_tools:
|
||||
disabled_tools.remove(tool_name)
|
||||
enabled = True
|
||||
else:
|
||||
disabled_tools.append(tool_name)
|
||||
enabled = False
|
||||
|
||||
server.disabled_tools = disabled_tools
|
||||
if updated_by is not None:
|
||||
server.updated_by = updated_by
|
||||
await db.commit()
|
||||
|
||||
# Clear tool cache (re-filtered on next fetch)
|
||||
clear_mcp_server_tools_cache(server_slug)
|
||||
|
||||
logger.info(f"Toggled tool '{tool_name}' for server '{server_slug}' enabled={enabled}")
|
||||
return enabled, server
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# === Unified Entry Points (Wrappers) ===
|
||||
# =============================================================================
|
||||
|
||||
|
||||
async def get_enabled_mcp_tools(server_slug: str) -> list:
|
||||
"""Get MCP server tools (auto-filtering disabled_tools).
|
||||
|
||||
Unified entry point for Agents, automatically:
|
||||
1. Gets the latest server config from database
|
||||
2. Gets all tools
|
||||
3. Filters out disabled_tools
|
||||
|
||||
Args:
|
||||
server_slug: Server slug
|
||||
|
||||
Returns:
|
||||
List of enabled tools
|
||||
"""
|
||||
config = await get_enabled_mcp_server_config(server_slug)
|
||||
if config is None:
|
||||
logger.warning(f"MCP server '{server_slug}' not found in database or disabled")
|
||||
return []
|
||||
|
||||
disabled_tools = config.get("disabled_tools") or []
|
||||
return await get_mcp_tools(server_slug, additional_servers={server_slug: config}, disabled_tools=disabled_tools)
|
||||
|
||||
|
||||
async def get_servers_config(names: list[str]) -> dict[str, dict[str, Any]]:
|
||||
"""Batch get server configurations.
|
||||
|
||||
Args:
|
||||
names: List of server names
|
||||
|
||||
Returns:
|
||||
{name: config} dictionary, containing only found servers
|
||||
"""
|
||||
return await _load_enabled_mcp_server_configs(names=names)
|
||||
|
||||
|
||||
async def get_all_mcp_tools(server_slug: str) -> list:
|
||||
"""Get all tools of an MCP server (no filtering).
|
||||
|
||||
For management UI to display tool list, supports viewing all tools and their enabled status.
|
||||
Does NOT update the global tools cache to avoid polluting agent's filtered view.
|
||||
|
||||
Args:
|
||||
server_slug: Server slug
|
||||
|
||||
Returns:
|
||||
List of all tools (unfiltered)
|
||||
"""
|
||||
config = await get_enabled_mcp_server_config(server_slug)
|
||||
if config is None:
|
||||
logger.warning(f"MCP server '{server_slug}' not found in database or disabled")
|
||||
return []
|
||||
|
||||
# Get all tools (no filtering, force refresh, no cache update)
|
||||
return await get_mcp_tools(
|
||||
server_slug,
|
||||
additional_servers={server_slug: config},
|
||||
disabled_tools=[],
|
||||
cache=False,
|
||||
force_refresh=True,
|
||||
)
|
||||
@@ -0,0 +1,15 @@
|
||||
from .attachment import inject_attachment_context, save_attachments_to_fs
|
||||
from .context import context_aware_prompt, context_based_model
|
||||
from .dynamic_tool import DynamicToolMiddleware
|
||||
from .summary import create_summary_middleware
|
||||
from .token_usage import TokenUsageMiddleware
|
||||
|
||||
__all__ = [
|
||||
"DynamicToolMiddleware",
|
||||
"TokenUsageMiddleware",
|
||||
"context_aware_prompt",
|
||||
"context_based_model",
|
||||
"create_summary_middleware",
|
||||
"inject_attachment_context", # 已废弃,使用 save_attachments_to_fs
|
||||
"save_attachments_to_fs",
|
||||
]
|
||||
@@ -0,0 +1,89 @@
|
||||
"""Attachment prompt injection middleware.
|
||||
|
||||
Read uploaded file metadata from LangGraph state and inject readable paths
|
||||
into the system prompt so the model can use `read_file` on demand.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Callable, Sequence
|
||||
from typing import NotRequired
|
||||
|
||||
from langchain.agents import AgentState
|
||||
from langchain.agents.middleware import AgentMiddleware, ModelRequest, ModelResponse
|
||||
from langchain_core.messages import SystemMessage
|
||||
|
||||
from yuxi.utils import logger
|
||||
|
||||
ATTACHMENT_PROMPT_MARKER = "<!-- attachment_context -->"
|
||||
|
||||
|
||||
class AttachmentState(AgentState):
|
||||
"""Extended state schema with uploaded files."""
|
||||
|
||||
uploads: NotRequired[list[dict]]
|
||||
|
||||
|
||||
def _build_attachment_prompt(uploads: Sequence[dict]) -> str | None:
|
||||
"""Render uploads into a concise prompt block."""
|
||||
if not uploads:
|
||||
return None
|
||||
|
||||
valid_uploads: list[tuple[str, str]] = []
|
||||
for upload in uploads:
|
||||
path = upload.get("path")
|
||||
if not isinstance(path, str) or not path.strip():
|
||||
continue
|
||||
file_name = upload.get("file_name", "未知文件")
|
||||
valid_uploads.append((str(file_name), path))
|
||||
|
||||
if not valid_uploads:
|
||||
return None
|
||||
|
||||
upload_infos = [f"- {file_name}: {path}" for file_name, path in valid_uploads]
|
||||
lines = [
|
||||
"用户上传了以下文件:",
|
||||
"",
|
||||
*upload_infos,
|
||||
"",
|
||||
"请优先使用 `read_file` 工具读取这些路径中的文件内容,再回答用户问题。",
|
||||
]
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
class AttachmentMiddleware(AgentMiddleware[AttachmentState]):
|
||||
"""Inject upload context from state.uploads into system prompt."""
|
||||
|
||||
state_schema = AttachmentState
|
||||
|
||||
async def awrap_model_call(
|
||||
self, request: ModelRequest, handler: Callable[[ModelRequest], ModelResponse]
|
||||
) -> ModelResponse:
|
||||
uploads = request.state.get("uploads", [])
|
||||
logger.info(f"AttachmentMiddleware: found {len(uploads)} uploads in state")
|
||||
|
||||
if uploads:
|
||||
attachment_prompt = _build_attachment_prompt(uploads)
|
||||
if attachment_prompt:
|
||||
logger.info("AttachmentMiddleware: injecting attachment prompt")
|
||||
existing_blocks = list(request.system_message.content_blocks) if request.system_message else []
|
||||
existing_text = "\n".join(
|
||||
block.get("text", "")
|
||||
for block in existing_blocks
|
||||
if isinstance(block, dict) and block.get("type") == "text"
|
||||
)
|
||||
|
||||
if ATTACHMENT_PROMPT_MARKER in existing_text:
|
||||
logger.info("AttachmentMiddleware: attachment prompt already injected, skip")
|
||||
return await handler(request)
|
||||
|
||||
merged_blocks = existing_blocks + [
|
||||
{"type": "text", "text": f"{ATTACHMENT_PROMPT_MARKER}\n{attachment_prompt}"}
|
||||
]
|
||||
request = request.override(system_message=SystemMessage(content=merged_blocks))
|
||||
|
||||
return await handler(request)
|
||||
|
||||
|
||||
save_attachments_to_fs = AttachmentMiddleware()
|
||||
inject_attachment_context = save_attachments_to_fs
|
||||
@@ -0,0 +1,25 @@
|
||||
"""通用的 Context 相关中间件"""
|
||||
|
||||
from collections.abc import Callable
|
||||
|
||||
from langchain.agents.middleware import ModelRequest, ModelResponse, dynamic_prompt, wrap_model_call
|
||||
|
||||
from yuxi.agents import load_chat_model, resolve_chat_model_spec
|
||||
from yuxi.utils import logger
|
||||
|
||||
|
||||
@dynamic_prompt
|
||||
def context_aware_prompt(request: ModelRequest) -> str:
|
||||
"""从 runtime context 动态生成系统提示词"""
|
||||
return request.runtime.context.system_prompt
|
||||
|
||||
|
||||
@wrap_model_call
|
||||
async def context_based_model(request: ModelRequest, handler: Callable[[ModelRequest], ModelResponse]) -> ModelResponse:
|
||||
"""从 runtime context 动态选择模型"""
|
||||
model_spec = resolve_chat_model_spec(request.runtime.context.model)
|
||||
model = load_chat_model(model_spec)
|
||||
|
||||
request = request.override(model=model)
|
||||
logger.debug(f"Using model {model_spec} for request {request.messages[-1].content[:200]}")
|
||||
return await handler(request)
|
||||
@@ -0,0 +1,69 @@
|
||||
from collections.abc import Callable
|
||||
from typing import Any
|
||||
|
||||
from langchain.agents.middleware import AgentMiddleware, ModelRequest, ModelResponse
|
||||
|
||||
from yuxi.agents.mcp.service import get_mcp_tools
|
||||
from yuxi.utils import logger
|
||||
|
||||
|
||||
class DynamicToolMiddleware(AgentMiddleware):
|
||||
"""动态工具选择中间件 - 支持 MCP 工具的动态加载和注册
|
||||
|
||||
注意:所有可能用到的 MCP 工具必须在初始化时预加载并注册到 self.tools
|
||||
运行时只是根据配置筛选工具,不能动态添加新工具
|
||||
"""
|
||||
|
||||
def __init__(self, base_tools: list[Any], mcp_servers: list[str] | None = None):
|
||||
"""初始化中间件
|
||||
|
||||
Args:
|
||||
base_tools: 基础工具列表
|
||||
mcp_servers: 需要预加载的 MCP 服务器列表(可选)
|
||||
"""
|
||||
super().__init__()
|
||||
self.tools: list[Any] = base_tools
|
||||
self._all_mcp_tools: dict[str, list[Any]] = {} # 所有已加载的 MCP 工具
|
||||
self._mcp_servers = mcp_servers or []
|
||||
|
||||
async def initialize_mcp_tools(self) -> None:
|
||||
"""异步初始化:预加载所有可能用到的 MCP 工具"""
|
||||
for mcp_name in self._mcp_servers:
|
||||
if mcp_name not in self._all_mcp_tools:
|
||||
logger.info(f"Pre-loading MCP tools from: {mcp_name}")
|
||||
mcp_tools = await get_mcp_tools(mcp_name)
|
||||
self._all_mcp_tools[mcp_name] = mcp_tools
|
||||
# 将 MCP 工具注册到 middleware.tools
|
||||
self.tools.extend(mcp_tools)
|
||||
logger.info(f"Registered {len(mcp_tools)} tools from {mcp_name}")
|
||||
|
||||
async def awrap_model_call(
|
||||
self, request: ModelRequest, handler: Callable[[ModelRequest], ModelResponse]
|
||||
) -> ModelResponse:
|
||||
"""根据配置动态选择工具(从已注册的工具中筛选)"""
|
||||
# 从 runtime context 获取配置
|
||||
selected_tools = request.runtime.context.tools
|
||||
selected_mcps = request.runtime.context.mcps
|
||||
|
||||
enabled_tools = []
|
||||
|
||||
# 根据配置筛选基础工具
|
||||
if selected_tools and isinstance(selected_tools, list) and len(selected_tools) > 0:
|
||||
enabled_tools = [tool for tool in self.tools if tool.name in selected_tools]
|
||||
|
||||
# 根据配置筛选 MCP 工具(从已注册的工具中选择)
|
||||
if selected_mcps and isinstance(selected_mcps, list) and len(selected_mcps) > 0:
|
||||
for mcp in selected_mcps:
|
||||
if mcp in self._all_mcp_tools:
|
||||
enabled_tools.extend(self._all_mcp_tools[mcp])
|
||||
else:
|
||||
logger.warning(f"MCP server '{mcp}' not pre-loaded. Please add it to mcp_servers list.")
|
||||
|
||||
logger.info(
|
||||
f"Dynamic tool selection: {len(enabled_tools)} tools enabled: {[tool.name for tool in enabled_tools]}, "
|
||||
f"selected_tools: {selected_tools}, selected_mcps: {selected_mcps}"
|
||||
)
|
||||
|
||||
# 更新 request 中的工具列表
|
||||
request = request.override(tools=enabled_tools)
|
||||
return await handler(request)
|
||||
@@ -0,0 +1,491 @@
|
||||
"""Skills 中间件 - 处理 skills 提示词注入、依赖展开、动态激活"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Callable
|
||||
from pathlib import PurePosixPath
|
||||
from typing import Annotated, Any, NotRequired, TypedDict
|
||||
|
||||
from deepagents.middleware._utils import append_to_system_message
|
||||
from deepagents.middleware.skills import SKILLS_SYSTEM_PROMPT
|
||||
from langchain.agents import AgentState
|
||||
from langchain.agents.middleware import AgentMiddleware, ModelRequest, ModelResponse
|
||||
from langchain.tools.tool_node import ToolCallRequest
|
||||
from langgraph.types import Command
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from yuxi.agents.mcp.service import get_enabled_mcp_tools
|
||||
from yuxi.agents.skills.repository import SkillRepository
|
||||
from yuxi.agents.skills.service import is_valid_skill_slug, list_accessible_skills, normalize_string_list
|
||||
from yuxi.agents.toolkits import get_all_tool_instances
|
||||
from yuxi.storage.postgres.manager import pg_manager
|
||||
from yuxi.utils.logging_config import logger
|
||||
|
||||
# =============================================================================
|
||||
# 类型定义
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class SkillPromptMetadata(TypedDict):
|
||||
name: str
|
||||
description: str
|
||||
path: str
|
||||
|
||||
|
||||
class SkillDependencyNode(TypedDict):
|
||||
tools: list[str]
|
||||
mcps: list[str]
|
||||
skills: list[str]
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# 运行时数据加载函数
|
||||
# =============================================================================
|
||||
|
||||
|
||||
async def _list_skills_from_db(db: AsyncSession | None = None, user=None) -> list:
|
||||
"""从数据库加载 skills 列表"""
|
||||
if db is not None:
|
||||
if user is not None:
|
||||
return await list_accessible_skills(db, user)
|
||||
repo = SkillRepository(db)
|
||||
return await repo.list_enabled()
|
||||
|
||||
async with pg_manager.get_async_session_context() as session:
|
||||
if user is not None:
|
||||
return await list_accessible_skills(session, user)
|
||||
repo = SkillRepository(session)
|
||||
return await repo.list_enabled()
|
||||
|
||||
|
||||
def build_prompt_metadata(skills: list) -> dict[str, SkillPromptMetadata]:
|
||||
return {
|
||||
item.slug: {
|
||||
"name": item.name,
|
||||
"description": item.description,
|
||||
"path": f"/home/gem/skills/{item.slug}/SKILL.md",
|
||||
}
|
||||
for item in skills
|
||||
if item.slug
|
||||
}
|
||||
|
||||
|
||||
def build_dependency_map(skills: list) -> dict[str, SkillDependencyNode]:
|
||||
result: dict[str, SkillDependencyNode] = {}
|
||||
for item in skills:
|
||||
if not item.slug:
|
||||
continue
|
||||
result[item.slug] = {
|
||||
"tools": normalize_string_list(item.tool_dependencies or []),
|
||||
"mcps": normalize_string_list(item.mcp_dependencies or []),
|
||||
"skills": normalize_string_list(item.skill_dependencies or []),
|
||||
}
|
||||
return result
|
||||
|
||||
|
||||
async def get_prompt_metadata(db: AsyncSession | None = None, user=None) -> dict[str, SkillPromptMetadata]:
|
||||
"""获取提示词元数据(直接从数据库加载)"""
|
||||
return build_prompt_metadata(await _list_skills_from_db(db, user))
|
||||
|
||||
|
||||
async def get_dependency_map(db: AsyncSession | None = None, user=None) -> dict[str, SkillDependencyNode]:
|
||||
"""获取依赖关系映射(直接从数据库加载)"""
|
||||
return build_dependency_map(await _list_skills_from_db(db, user))
|
||||
|
||||
|
||||
def expand_skill_closure(
|
||||
slugs: list[str] | None,
|
||||
dependency_map: dict[str, SkillDependencyNode],
|
||||
) -> list[str]:
|
||||
"""展开 skills 依赖闭包,返回包含所有依赖的列表"""
|
||||
ordered_roots = normalize_string_list(slugs)
|
||||
if not ordered_roots:
|
||||
return []
|
||||
|
||||
result: list[str] = []
|
||||
seen: set[str] = set()
|
||||
|
||||
def dfs(slug: str, stack: set[str]) -> None:
|
||||
if slug in stack:
|
||||
logger.warning(f"Cycle detected in skill dependencies, skip: {' -> '.join([*stack, slug])}")
|
||||
return
|
||||
if slug in seen:
|
||||
return
|
||||
|
||||
node = dependency_map.get(slug)
|
||||
if not node:
|
||||
logger.warning(f"Skill dependency target not found in DB, skip: {slug}")
|
||||
return
|
||||
|
||||
seen.add(slug)
|
||||
result.append(slug)
|
||||
next_stack = set(stack)
|
||||
next_stack.add(slug)
|
||||
for dep in node.get("skills", []):
|
||||
dfs(dep, next_stack)
|
||||
|
||||
for root in ordered_roots:
|
||||
dfs(root, set())
|
||||
return result
|
||||
|
||||
|
||||
async def resolve_runtime_skills_for_context(context, *, db: AsyncSession | None = None, user=None) -> dict[str, Any]:
|
||||
skill_items = await _list_skills_from_db(db, user)
|
||||
dependency_map = build_dependency_map(skill_items)
|
||||
prompt_metadata = build_prompt_metadata(skill_items)
|
||||
available = set(dependency_map)
|
||||
selected = normalize_string_list(getattr(context, "skills", None))
|
||||
context_skills = [slug for slug in selected if slug in available]
|
||||
prompt_skills = expand_skill_closure(context_skills, dependency_map)
|
||||
return {
|
||||
"context_skills": context_skills,
|
||||
"prompt_skills": prompt_skills,
|
||||
"readable_skills": prompt_skills,
|
||||
"runtime_skill_metadata": prompt_metadata,
|
||||
"runtime_skill_dependency_map": dependency_map,
|
||||
}
|
||||
|
||||
|
||||
def resolve_skill_gated_tools(context) -> list:
|
||||
"""解析当前 Agent 所有可见 Skill 依赖的本地工具实例。
|
||||
|
||||
这些工具默认不会绑定给模型(由 SkillsMiddleware 在对应 Skill 激活后才放出),
|
||||
但必须在构建期注册进 create_agent 的 ToolNode,否则即便模型发起调用,执行器也会
|
||||
判定为 "not a valid tool"。调用方应自行排除已在基础工具集中的工具,避免重复门控。
|
||||
"""
|
||||
dependency_map = getattr(context, "_runtime_skill_dependency_map", {}) or {}
|
||||
readable_skills = getattr(context, "_readable_skills", []) or []
|
||||
tool_names: set[str] = set()
|
||||
for slug in readable_skills:
|
||||
node = dependency_map.get(slug) or {}
|
||||
tool_names.update(node.get("tools", []))
|
||||
if not tool_names:
|
||||
return []
|
||||
return [tool for tool in get_all_tool_instances() if tool.name in tool_names]
|
||||
|
||||
|
||||
def _activated_skills_reducer(left: list[str] | None, right: list[str] | None) -> list[str]:
|
||||
"""合并 activated_skills 列表"""
|
||||
merged: list[str] = []
|
||||
seen: set[str] = set()
|
||||
for group in (left or [], right or []):
|
||||
for value in group:
|
||||
if not isinstance(value, str):
|
||||
continue
|
||||
slug = value.strip()
|
||||
if not slug or slug in seen:
|
||||
continue
|
||||
seen.add(slug)
|
||||
merged.append(slug)
|
||||
return merged
|
||||
|
||||
|
||||
class SkillsState(AgentState):
|
||||
"""Skills 状态定义"""
|
||||
|
||||
activated_skills: NotRequired[Annotated[list[str], _activated_skills_reducer]]
|
||||
|
||||
|
||||
class SkillsMiddleware(AgentMiddleware):
|
||||
"""Skills 中间件 - 处理 skills 提示词注入、依赖展开、动态激活
|
||||
|
||||
职责:
|
||||
- Skills 提示词注入(直接从数据库加载)
|
||||
- 依赖展开(用户配置 + 动态激活)
|
||||
- 工具/MCP 动态加载
|
||||
"""
|
||||
|
||||
state_schema = SkillsState
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
skills_context_name: str = "skills",
|
||||
enable_skills_prompt: bool = True,
|
||||
skills_sources_for_prompt: list[str] | None = None,
|
||||
):
|
||||
"""初始化中间件
|
||||
|
||||
Args:
|
||||
skills_context_name: 上下文中的 skills 列表字段名称(默认 "skills")
|
||||
enable_skills_prompt: 是否启用 skills 提示段注入(默认 True)
|
||||
skills_sources_for_prompt: skills 来源路径(用于提示词展示,默认 ["/home/gem/skills/"])
|
||||
"""
|
||||
super().__init__()
|
||||
self.skills_context_name = skills_context_name
|
||||
self.enable_skills_prompt = enable_skills_prompt
|
||||
self.skills_sources_for_prompt = skills_sources_for_prompt or ["/home/gem/skills/"]
|
||||
|
||||
async def awrap_model_call(
|
||||
self, request: ModelRequest, handler: Callable[[ModelRequest], ModelResponse]
|
||||
) -> ModelResponse:
|
||||
"""包装模型调用,处理 skills 提示词注入、动态激活和依赖展开"""
|
||||
runtime_context = request.runtime.context
|
||||
|
||||
if self.enable_skills_prompt:
|
||||
prompt_skills = getattr(runtime_context, "_prompt_skills", None)
|
||||
if isinstance(prompt_skills, list):
|
||||
prompt_skills = normalize_string_list(prompt_skills)
|
||||
if prompt_skills:
|
||||
skills_meta = self._collect_prompt_metadata(prompt_skills, runtime_context)
|
||||
skills_section = self._build_skills_section(skills_meta)
|
||||
system_message = append_to_system_message(getattr(request, "system_message", None), skills_section)
|
||||
request = request.override(system_message=system_message)
|
||||
|
||||
state = request.state if isinstance(request.state, dict) else {}
|
||||
activated = state.get("activated_skills", []) or []
|
||||
if not isinstance(activated, list):
|
||||
activated = []
|
||||
|
||||
readable_skills = self._get_readable_skills(runtime_context)
|
||||
activated = [slug for slug in normalize_string_list(activated) if slug in readable_skills]
|
||||
|
||||
deps_bundle = self._build_dependency_bundle(activated, runtime_context)
|
||||
activated_tool_names = set(deps_bundle["tools"])
|
||||
|
||||
# 门控:未激活 Skill 的依赖工具对模型不可见(保持按需加载)。
|
||||
# 这些工具已在构建期由 resolve_configured_runtime_tools 注册进 ToolNode,剔除只影响模型可见性、不影响可执行性。
|
||||
# 排除基础工具集中的工具(如 present_artifacts),它们始终可见、不受 Skill 激活影响。
|
||||
gated_tool_names = self._resolve_gated_tool_names(runtime_context) - activated_tool_names
|
||||
model_tools = list(request.tools or [])
|
||||
if gated_tool_names:
|
||||
model_tools = [t for t in model_tools if t.name not in gated_tool_names]
|
||||
|
||||
# 追加已激活 Skill 的依赖工具:本地工具确保绑定给模型,MCP 工具按需加载
|
||||
enabled_tools = []
|
||||
if activated_tool_names:
|
||||
enabled_tools = [t for t in get_all_tool_instances() if t.name in activated_tool_names]
|
||||
if deps_bundle["mcps"]:
|
||||
enabled_tools.extend(
|
||||
await self._get_mcp_tools_from_context(runtime_context, extra_mcps=deps_bundle["mcps"])
|
||||
)
|
||||
|
||||
existing_tool_names = {t.name for t in model_tools}
|
||||
for t in enabled_tools:
|
||||
if t.name not in existing_tool_names:
|
||||
model_tools.append(t)
|
||||
existing_tool_names.add(t.name)
|
||||
|
||||
if gated_tool_names or enabled_tools:
|
||||
request = request.override(tools=model_tools)
|
||||
|
||||
return await handler(request)
|
||||
|
||||
def _resolve_gated_tool_names(self, runtime_context) -> set[str]:
|
||||
"""所有可见 Skill 依赖、且不属于基础工具集的工具名集合(即「仅经 Skill 激活才放出」的工具)。"""
|
||||
dependency_map = self._get_runtime_dependency_map(runtime_context)
|
||||
readable_skills = self._get_readable_skills(runtime_context)
|
||||
base_tool_names = set(normalize_string_list(getattr(runtime_context, "tools", None)))
|
||||
gated: set[str] = set()
|
||||
for slug in readable_skills:
|
||||
gated.update(dependency_map.get(slug, {}).get("tools", []))
|
||||
return gated - base_tool_names
|
||||
|
||||
def _build_dependency_bundle(self, activated_skills: list[str], runtime_context) -> dict[str, list[str]]:
|
||||
"""根据直接激活的 skills 构建依赖包(不包含闭包展开的依赖)"""
|
||||
dependency_map = self._get_runtime_dependency_map(runtime_context)
|
||||
|
||||
tools: list[str] = []
|
||||
mcps: list[str] = []
|
||||
seen_tools: set[str] = set()
|
||||
seen_mcps: set[str] = set()
|
||||
|
||||
for slug in activated_skills:
|
||||
dep = dependency_map.get(slug, {})
|
||||
for tool_name in dep.get("tools", []):
|
||||
if tool_name in seen_tools:
|
||||
continue
|
||||
seen_tools.add(tool_name)
|
||||
tools.append(tool_name)
|
||||
for mcp_name in dep.get("mcps", []):
|
||||
if mcp_name in seen_mcps:
|
||||
continue
|
||||
seen_mcps.add(mcp_name)
|
||||
mcps.append(mcp_name)
|
||||
|
||||
return {"tools": tools, "mcps": mcps, "skills": activated_skills}
|
||||
|
||||
def _collect_prompt_metadata(self, slugs: list[str], runtime_context) -> list[SkillPromptMetadata]:
|
||||
"""收集指定 slugs 的提示词元数据"""
|
||||
prompt_metadata = self._get_runtime_prompt_metadata(runtime_context)
|
||||
|
||||
result: list[SkillPromptMetadata] = []
|
||||
seen: set[str] = set()
|
||||
|
||||
for slug in slugs:
|
||||
if not isinstance(slug, str):
|
||||
continue
|
||||
normalized = slug.strip()
|
||||
if not normalized or normalized in seen:
|
||||
continue
|
||||
seen.add(normalized)
|
||||
|
||||
item = prompt_metadata.get(normalized)
|
||||
if not item:
|
||||
logger.debug(f"Skill slug not found in prompt metadata, skip: {normalized}")
|
||||
continue
|
||||
result.append(dict(item))
|
||||
|
||||
return result
|
||||
|
||||
async def _get_mcp_tools_from_context(
|
||||
self,
|
||||
context,
|
||||
*,
|
||||
extra_mcps: list[str] | None = None,
|
||||
) -> list:
|
||||
"""从上下文配置中获取 MCP 工具列表"""
|
||||
import asyncio
|
||||
|
||||
# MCP 工具(并行加载)
|
||||
mcps = getattr(context, "mcps", None) or []
|
||||
all_mcp_names: list[str] = []
|
||||
for server_name in mcps:
|
||||
if isinstance(server_name, str):
|
||||
all_mcp_names.append(server_name)
|
||||
for server_name in extra_mcps or []:
|
||||
if isinstance(server_name, str):
|
||||
all_mcp_names.append(server_name)
|
||||
|
||||
# 去重
|
||||
unique_mcp_names = list(dict.fromkeys(all_mcp_names))
|
||||
|
||||
async def load_mcp_tools(server_name: str) -> list:
|
||||
"""加载单个 MCP 服务器的工具"""
|
||||
try:
|
||||
mcp_tools = await get_enabled_mcp_tools(server_name)
|
||||
if not mcp_tools:
|
||||
logger.warning(f"SkillsMiddleware: mcp dependency unavailable, skip: {server_name}")
|
||||
return mcp_tools
|
||||
except Exception as e:
|
||||
logger.warning(f"SkillsMiddleware: failed to load mcp dependency '{server_name}': {e}")
|
||||
return []
|
||||
|
||||
# 并行加载所有 MCP 工具
|
||||
results = await asyncio.gather(*[load_mcp_tools(name) for name in unique_mcp_names])
|
||||
selected_tools = []
|
||||
for tools in results:
|
||||
selected_tools.extend(tools)
|
||||
|
||||
return selected_tools
|
||||
|
||||
def _process_tool_call_result(self, result: Any, request: ToolCallRequest) -> Any:
|
||||
"""处理工具调用结果,检查并处理 skill 动态激活"""
|
||||
if request.tool_call.get("name") != "read_file":
|
||||
return result
|
||||
|
||||
args = request.tool_call.get("args") or {}
|
||||
file_path = args.get("file_path") if isinstance(args, dict) else None
|
||||
slug = self._extract_skill_slug_from_skill_md_path(file_path)
|
||||
|
||||
if not slug:
|
||||
return result
|
||||
|
||||
if not self._is_visible_skill_slug(request, slug):
|
||||
logger.warning(f"SkillsMiddleware: deny skill activation for invisible slug: {slug}")
|
||||
return result
|
||||
|
||||
logger.debug(f"SkillsMiddleware: activated skill by read_file: {slug}")
|
||||
return self._merge_activated_skill_update(result, slug)
|
||||
|
||||
async def awrap_tool_call(
|
||||
self,
|
||||
request: ToolCallRequest,
|
||||
handler: Callable[[ToolCallRequest], Any],
|
||||
):
|
||||
"""包装工具调用,处理 skill 动态激活"""
|
||||
result = await handler(request)
|
||||
return self._process_tool_call_result(result, request)
|
||||
|
||||
def wrap_tool_call(
|
||||
self,
|
||||
request: ToolCallRequest,
|
||||
handler: Callable[[ToolCallRequest], Any],
|
||||
):
|
||||
"""同步版本的工具调用包装"""
|
||||
result = handler(request)
|
||||
return self._process_tool_call_result(result, request)
|
||||
|
||||
def _extract_skill_slug_from_skill_md_path(self, file_path: Any) -> str | None:
|
||||
"""从文件路径中提取 skill slug"""
|
||||
if not isinstance(file_path, str):
|
||||
return None
|
||||
raw = file_path.strip()
|
||||
if not raw:
|
||||
return None
|
||||
pure = PurePosixPath(raw if raw.startswith("/") else f"/{raw}")
|
||||
parts = [p for p in pure.parts if p not in ("/", "")]
|
||||
slug: str | None = None
|
||||
if (
|
||||
len(parts) == 5
|
||||
and parts[0] == "home"
|
||||
and parts[1] == "gem"
|
||||
and parts[2] == "skills"
|
||||
and parts[4] == "SKILL.md"
|
||||
):
|
||||
slug = parts[3]
|
||||
|
||||
if not is_valid_skill_slug(slug):
|
||||
return None
|
||||
return slug
|
||||
|
||||
def _get_readable_skills(self, runtime_context) -> set[str]:
|
||||
selected = getattr(runtime_context, "_readable_skills", [])
|
||||
return set(normalize_string_list(selected if isinstance(selected, list) else []))
|
||||
|
||||
def _get_runtime_prompt_metadata(self, runtime_context) -> dict[str, SkillPromptMetadata]:
|
||||
metadata = getattr(runtime_context, "_runtime_skill_metadata", {})
|
||||
return metadata if isinstance(metadata, dict) else {}
|
||||
|
||||
def _get_runtime_dependency_map(self, runtime_context) -> dict[str, SkillDependencyNode]:
|
||||
dependency_map = getattr(runtime_context, "_runtime_skill_dependency_map", {})
|
||||
return dependency_map if isinstance(dependency_map, dict) else {}
|
||||
|
||||
def _is_visible_skill_slug(self, request: ToolCallRequest, slug: str) -> bool:
|
||||
"""检查 slug 是否可见"""
|
||||
return slug in self._get_readable_skills(request.runtime.context)
|
||||
|
||||
def _merge_activated_skill_update(self, result: Any, slug: str):
|
||||
"""合并动态激活的 skill 更新"""
|
||||
from langchain_core.messages import ToolMessage
|
||||
|
||||
if isinstance(result, Command):
|
||||
update = dict(result.update or {})
|
||||
current = update.get("activated_skills") or []
|
||||
update["activated_skills"] = _activated_skills_reducer(current, [slug])
|
||||
return Command(graph=result.graph, update=update, resume=result.resume, goto=result.goto)
|
||||
|
||||
if isinstance(result, ToolMessage):
|
||||
return Command(update={"messages": [result], "activated_skills": [slug]})
|
||||
|
||||
return result
|
||||
|
||||
def _format_skills_locations(self, sources: list[str]) -> str:
|
||||
"""格式化 skills 位置信息"""
|
||||
locations = []
|
||||
for i, source_path in enumerate(sources):
|
||||
name = PurePosixPath(source_path.rstrip("/")).name.capitalize()
|
||||
suffix = " (higher priority)" if i == len(sources) - 1 else ""
|
||||
locations.append(f"**{name} Skills**: `{source_path}`{suffix}")
|
||||
return "\n".join(locations)
|
||||
|
||||
def _format_skills_list(self, skills_meta: list[dict[str, str]]) -> str:
|
||||
"""格式化 skills 列表"""
|
||||
if not skills_meta:
|
||||
return f"(No skills available yet. You can create skills in {' or '.join(self.skills_sources_for_prompt)})"
|
||||
|
||||
lines = []
|
||||
for skill in skills_meta:
|
||||
lines.append(f"- **{skill['name']}**: {skill['description']}")
|
||||
lines.append(f" -> Read `{skill['path']}` for full instructions")
|
||||
return "\n".join(lines)
|
||||
|
||||
def _build_skills_section(self, skills_meta: list[dict[str, str]]) -> str:
|
||||
"""构建 skills 提示段"""
|
||||
skills_locations = self._format_skills_locations(self.skills_sources_for_prompt)
|
||||
skills_list = self._format_skills_list(skills_meta)
|
||||
return SKILLS_SYSTEM_PROMPT.format(
|
||||
skills_locations=skills_locations,
|
||||
skills_load_warnings="",
|
||||
skills_list=skills_list,
|
||||
)
|
||||
@@ -0,0 +1,586 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from collections.abc import Awaitable, Callable
|
||||
from dataclasses import dataclass
|
||||
from typing import Annotated, Any
|
||||
|
||||
from deepagents import SubagentTransformer as DeepAgentsSubagentTransformer
|
||||
from deepagents.middleware._utils import append_to_system_message
|
||||
from langchain.agents.middleware.types import AgentMiddleware, ContextT, ModelRequest, ModelResponse, ResponseT
|
||||
from langchain_core.messages import ToolMessage
|
||||
from langchain_core.tools import StructuredTool
|
||||
from langgraph.prebuilt.tool_node import ToolRuntime
|
||||
from langgraph.types import Command
|
||||
|
||||
from yuxi.repositories.agent_repository import AgentRepository
|
||||
from yuxi.repositories.agent_run_repository import TERMINAL_RUN_STATUSES
|
||||
from yuxi.repositories.user_repository import UserRepository
|
||||
from yuxi.services.input_message_service import build_chat_input_message
|
||||
from yuxi.storage.postgres.manager import pg_manager
|
||||
from yuxi.storage.postgres.models_business import Agent
|
||||
|
||||
YUXI_SUBAGENTS_STREAM_KEY = "yuxi_subagents"
|
||||
|
||||
|
||||
def _subagent_run_service_module():
|
||||
from yuxi.services import subagent_run_service
|
||||
|
||||
return subagent_run_service
|
||||
|
||||
|
||||
def _async_only_tool(*, name: str, coroutine: Callable[..., Awaitable[Any]], description: str) -> StructuredTool:
|
||||
"""后台子智能体工具只在异步链路执行;仅声明 coroutine,同步调用由 LangChain 直接报错。"""
|
||||
return StructuredTool.from_function(name=name, coroutine=coroutine, description=description, infer_schema=True)
|
||||
|
||||
|
||||
class YuxiSubagentTransformer(DeepAgentsSubagentTransformer):
|
||||
def init(self) -> dict[str, Any]:
|
||||
return {YUXI_SUBAGENTS_STREAM_KEY: self._log}
|
||||
|
||||
|
||||
TASK_SYSTEM_PROMPT = """## `task`(子智能体任务工具)
|
||||
|
||||
你可以使用 `task` 工具把复杂、独立的子任务交给已配置的子智能体处理。子智能体只返回最终结果,你看不到它的中间步骤。
|
||||
工具结果会包含子智能体线程 ID,后续需要继续同一个子任务时,把该 ID 作为 `thread_id` 传回 `task`。
|
||||
|
||||
使用原则:
|
||||
- 任务足够复杂、可以独立完成、或需要隔离上下文时使用。
|
||||
- 多个互不依赖的子任务可以并行调用多个 `task`。
|
||||
- 继续既有子智能体任务时传入之前结果中的 `thread_id`;新任务不要填写 `thread_id`。
|
||||
- 不要并行调用同一个 `thread_id`,避免多个续跑请求同时写入同一子线程。
|
||||
- 简单问题或少量直接工具调用不要委派。
|
||||
- 调用时必须选择下方可用的 `subagent_slug`,并在 `description` 中写清目标、上下文和期望输出。
|
||||
- 不要通过 shell、curl、HTTP API 或命令行间接调用子智能体;需要子智能体时必须使用 `task` 工具。
|
||||
|
||||
后台子智能体:
|
||||
- 长任务或多个可并行任务优先使用 `subagent_start`,它会立即返回 `run_id` 和 `thread_id`,父智能体可以继续工作。
|
||||
- 后续用 `subagent_status` 查询状态,`subagent_events` 读取增量事件,
|
||||
`subagent_cancel` 取消,`subagent_await` 在明确需要结果时等待。
|
||||
- `thread_id` 是子智能体长期上下文 ID;同一个 `thread_id` 完成后可以继续创建新的 run。
|
||||
若同线程已有运行中 run,会返回 busy,不会隐藏排队。
|
||||
- 短任务且父智能体必须立刻依赖结果时继续使用 `task`。
|
||||
|
||||
Available subagent slugs:
|
||||
|
||||
{available_agents}"""
|
||||
|
||||
TASK_TOOL_DESCRIPTION = """Launch a configured Yuxi subagent to handle an isolated task.
|
||||
|
||||
Available subagent slugs:
|
||||
{available_agents}
|
||||
|
||||
Use `subagent_slug` to select one available subagent and put the full task brief in `description`.
|
||||
Omit `thread_id` for a new task. To continue a previous subagent task, pass the child thread ID returned by
|
||||
that prior task result as `thread_id`.
|
||||
Do not call subagents through shell, curl, HTTP APIs, or command-line indirection."""
|
||||
|
||||
SUBAGENT_START_DESCRIPTION = """Start a configured Yuxi subagent asynchronously.
|
||||
|
||||
Returns a child thread ID for future continuation and a run ID for status/events/cancel/result checks.
|
||||
Use this for long-running or parallelizable subagent work. If `thread_id` is provided, it continues that subagent
|
||||
thread when no active run is currently writing to it."""
|
||||
|
||||
SUBAGENT_STATUS_DESCRIPTION = """Check a subagent run status by run_id.
|
||||
|
||||
Returns the current run status, a compact progress summary with the latest 3 readable messages, and the final result
|
||||
when the run has reached a terminal status."""
|
||||
|
||||
SUBAGENT_EVENTS_DESCRIPTION = """Read recent events for a subagent run by run_id and Redis stream cursor."""
|
||||
|
||||
SUBAGENT_CANCEL_DESCRIPTION = """Cancel a running subagent run by run_id."""
|
||||
|
||||
SUBAGENT_AWAIT_DESCRIPTION = """Wait for a subagent run to finish and return its final result."""
|
||||
|
||||
TASK_DESCRIPTION_ARG = "需要子智能体独立完成的任务描述,包含必要上下文和期望输出。"
|
||||
SUBAGENT_SLUG_ARG = "要调用的子智能体 slug,必须是工具描述中列出的可用项之一。"
|
||||
TASK_THREAD_ID_ARG = "可选。要继续的既有子智能体线程 ID,通常来自之前 task 工具结果;新任务不要填写。"
|
||||
ASYNC_THREAD_ID_ARG = "可选。要继续的后台子智能体线程 ID,来自之前 subagent_start 返回的 thread_id;新任务不要填写。"
|
||||
SUBAGENT_RUN_ID_ARG = "子智能体运行 ID,由 subagent_start 返回。"
|
||||
SUBAGENT_AFTER_SEQ_ARG = "可选。事件流游标,首次读取传 0-0;后续传上次返回的 last_seq。"
|
||||
SUBAGENT_EVENT_LIMIT_ARG = "可选。读取事件数量,范围 1-50。"
|
||||
|
||||
|
||||
async def create_subagent_task_middleware(parent_context) -> YuxiSubAgentMiddleware | None:
|
||||
"""根据父智能体上下文加载可用子智能体,并在存在可调用项时创建 task 中间件。"""
|
||||
selected_slugs = [
|
||||
str(slug).strip() for slug in (getattr(parent_context, "subagents", None) or []) if str(slug).strip()
|
||||
]
|
||||
uid = str(getattr(parent_context, "uid", "") or "").strip()
|
||||
if not uid:
|
||||
return None
|
||||
|
||||
async with pg_manager.get_async_session_context() as db:
|
||||
user = await UserRepository().get_by_uid_with_db(db, uid)
|
||||
if user is None:
|
||||
return None
|
||||
repo = AgentRepository(db)
|
||||
if selected_slugs:
|
||||
subagents: list[Agent] = []
|
||||
seen: set[str] = set()
|
||||
for slug in selected_slugs:
|
||||
if slug in seen:
|
||||
continue
|
||||
seen.add(slug)
|
||||
agent = await repo.get_visible_by_slug(slug=slug, user=user, kind="subagent")
|
||||
if agent:
|
||||
subagents.append(agent)
|
||||
else:
|
||||
subagents = await repo.list_visible_subagents(user=user)
|
||||
|
||||
if not subagents:
|
||||
return None
|
||||
return YuxiSubAgentMiddleware(parent_context=parent_context, subagents=subagents)
|
||||
|
||||
|
||||
class YuxiSubAgentMiddleware(AgentMiddleware[Any, ContextT, ResponseT]):
|
||||
def __init__(self, *, parent_context, subagents: list[Agent]) -> None:
|
||||
super().__init__()
|
||||
self.parent_context = parent_context
|
||||
self.subagents = {agent.slug: agent for agent in subagents}
|
||||
available_agents = "\n".join(f"- {agent.slug}: {agent.description or agent.name}" for agent in subagents)
|
||||
self.system_prompt = TASK_SYSTEM_PROMPT.format(available_agents=available_agents)
|
||||
self.tools = [self._build_task_tool(available_agents), *self._build_async_subagent_tools(available_agents)]
|
||||
self.subagent_names = frozenset(self.subagents)
|
||||
self.transformers = [lambda scope: YuxiSubagentTransformer(scope, subagent_names=self.subagent_names)]
|
||||
|
||||
def wrap_model_call(
|
||||
self,
|
||||
request: ModelRequest[ContextT],
|
||||
handler: Callable[[ModelRequest[ContextT]], ModelResponse[ResponseT]],
|
||||
) -> ModelResponse[ResponseT]:
|
||||
return handler(
|
||||
request.override(system_message=append_to_system_message(request.system_message, self.system_prompt))
|
||||
)
|
||||
|
||||
async def awrap_model_call(
|
||||
self,
|
||||
request: ModelRequest[ContextT],
|
||||
handler: Callable[[ModelRequest[ContextT]], Awaitable[ModelResponse[ResponseT]]],
|
||||
) -> ModelResponse[ResponseT]:
|
||||
return await handler(
|
||||
request.override(system_message=append_to_system_message(request.system_message, self.system_prompt))
|
||||
)
|
||||
|
||||
def _build_task_tool(self, available_agents: str) -> StructuredTool:
|
||||
"""构建 task 工具:启动子智能体后阻塞等待其最终结果。"""
|
||||
|
||||
async def atask(
|
||||
description: Annotated[str, TASK_DESCRIPTION_ARG],
|
||||
subagent_slug: Annotated[str, SUBAGENT_SLUG_ARG],
|
||||
runtime: ToolRuntime,
|
||||
thread_id: Annotated[str | None, TASK_THREAD_ID_ARG] = None,
|
||||
) -> str | Command:
|
||||
started, error = await self._start_subagent(
|
||||
description=description,
|
||||
subagent_slug=subagent_slug,
|
||||
runtime=runtime,
|
||||
thread_id=thread_id,
|
||||
error_prefix="无法调用子智能体",
|
||||
)
|
||||
if error is not None:
|
||||
return error
|
||||
|
||||
# 阻塞父智能体运行,直到子 run 终结再读取最终结果;运行失败时 result 含 error 信息
|
||||
parent_runtime = started.parent_runtime
|
||||
subagent_service = _subagent_run_service_module()
|
||||
try:
|
||||
from yuxi.services.agent_run_service import AgentRunWaitTimeout, await_agent_run_result
|
||||
|
||||
result = await await_agent_run_result(run_id=started.result.run.id, current_uid=parent_runtime.uid)
|
||||
run = await self._get_verified_subagent_run(
|
||||
run_id=started.result.run.id,
|
||||
uid=parent_runtime.uid,
|
||||
created_by_run_id=parent_runtime.created_by_run_id,
|
||||
)
|
||||
except AgentRunWaitTimeout as exc:
|
||||
try:
|
||||
run = await self._get_verified_subagent_run(
|
||||
run_id=started.result.run.id,
|
||||
uid=parent_runtime.uid,
|
||||
created_by_run_id=parent_runtime.created_by_run_id,
|
||||
)
|
||||
except ValueError as verify_exc:
|
||||
return str(verify_exc)
|
||||
subagent_run = subagent_service.serialize_subagent_run_state(run)
|
||||
return _task_wait_timeout_response(exc.result, runtime.tool_call_id, subagent_run)
|
||||
except ValueError as exc:
|
||||
return str(exc)
|
||||
|
||||
subagent_run = subagent_service.serialize_subagent_run_state(run)
|
||||
return _task_result_response(result, runtime.tool_call_id, subagent_run)
|
||||
|
||||
return _async_only_tool(
|
||||
name="task",
|
||||
coroutine=atask,
|
||||
description=TASK_TOOL_DESCRIPTION.format(available_agents=available_agents),
|
||||
)
|
||||
|
||||
def _build_async_subagent_tools(self, available_agents: str) -> list[StructuredTool]:
|
||||
"""构建后台子智能体生命周期工具:start/status/events/cancel/await。"""
|
||||
|
||||
async def asubagent_start(
|
||||
description: Annotated[str, TASK_DESCRIPTION_ARG],
|
||||
subagent_slug: Annotated[str, SUBAGENT_SLUG_ARG],
|
||||
runtime: ToolRuntime,
|
||||
thread_id: Annotated[str | None, ASYNC_THREAD_ID_ARG] = None,
|
||||
) -> str | Command:
|
||||
started, error = await self._start_subagent(
|
||||
description=description,
|
||||
subagent_slug=subagent_slug,
|
||||
runtime=runtime,
|
||||
thread_id=thread_id,
|
||||
error_prefix="无法启动子智能体",
|
||||
)
|
||||
if error is not None:
|
||||
return error
|
||||
|
||||
result, agent_item = started.result, started.agent_item
|
||||
subagent_service = _subagent_run_service_module()
|
||||
payload = {
|
||||
"status": "started" if result.created else "existing",
|
||||
"run_id": result.run.id,
|
||||
"thread_id": result.relation.child_thread_id,
|
||||
"subagent_slug": subagent_slug,
|
||||
"subagent_name": agent_item.name,
|
||||
"created_by_run_id": result.run.created_by_run_id,
|
||||
"run_status": result.run.status,
|
||||
"continuing": result.continuing,
|
||||
"subagent_thread_relation_id": result.relation.id,
|
||||
**subagent_service.subagent_run_urls(result.run.id),
|
||||
}
|
||||
subagent_run = subagent_service.serialize_subagent_run_state(result.run)
|
||||
return _json_tool_command(payload, runtime.tool_call_id, subagent_run=subagent_run)
|
||||
|
||||
async def asubagent_status(
|
||||
run_id: Annotated[str, SUBAGENT_RUN_ID_ARG],
|
||||
runtime: ToolRuntime,
|
||||
) -> str | Command:
|
||||
from yuxi.services.agent_run_service import get_agent_run_progress, get_agent_run_result
|
||||
|
||||
parent_runtime, runtime_error = self._require_async_parent_runtime("无法查询子智能体")
|
||||
if runtime_error:
|
||||
return runtime_error
|
||||
try:
|
||||
run = await self._get_verified_subagent_run(
|
||||
uid=parent_runtime.uid,
|
||||
created_by_run_id=parent_runtime.created_by_run_id,
|
||||
run_id=run_id,
|
||||
)
|
||||
|
||||
# 如果 run 已经终结,则尝试读取最终结果;否则 result 保持 None
|
||||
result = None
|
||||
if run.status in TERMINAL_RUN_STATUSES:
|
||||
async with pg_manager.get_async_session_context() as db:
|
||||
result = await get_agent_run_result(run_id=run.id, current_uid=parent_runtime.uid, db=db)
|
||||
|
||||
except ValueError as exc:
|
||||
return str(exc)
|
||||
|
||||
subagent_service = _subagent_run_service_module()
|
||||
payload = {
|
||||
"status": run.status,
|
||||
"run_id": run.id,
|
||||
"thread_id": run.conversation_thread_id,
|
||||
"subagent_slug": run.agent_slug,
|
||||
"error": run.error_message,
|
||||
"progress": await get_agent_run_progress(run.id),
|
||||
**subagent_service.subagent_run_urls(run.id),
|
||||
}
|
||||
if result:
|
||||
payload["result"] = result
|
||||
subagent_run = subagent_service.serialize_subagent_run_state(run)
|
||||
return _json_tool_command(payload, runtime.tool_call_id, subagent_run=subagent_run)
|
||||
|
||||
async def asubagent_events(
|
||||
run_id: Annotated[str, SUBAGENT_RUN_ID_ARG],
|
||||
runtime: ToolRuntime,
|
||||
after_seq: Annotated[str, SUBAGENT_AFTER_SEQ_ARG] = "0-0",
|
||||
limit: Annotated[int, SUBAGENT_EVENT_LIMIT_ARG] = 20,
|
||||
) -> str | Command:
|
||||
from yuxi.services.run_queue_service import list_run_stream_events, normalize_after_seq
|
||||
|
||||
parent_runtime, runtime_error = self._require_async_parent_runtime("无法读取子智能体事件")
|
||||
if runtime_error:
|
||||
return runtime_error
|
||||
try:
|
||||
await self._get_verified_subagent_run(
|
||||
run_id=run_id,
|
||||
uid=parent_runtime.uid,
|
||||
created_by_run_id=parent_runtime.created_by_run_id,
|
||||
) # 校验子智能体归属
|
||||
except ValueError as exc:
|
||||
return str(exc)
|
||||
|
||||
normalized_after_seq = normalize_after_seq(after_seq)
|
||||
event_limit = min(50, max(1, int(limit or 20)))
|
||||
events = await list_run_stream_events(run_id, after_seq=normalized_after_seq, limit=event_limit)
|
||||
payload = {
|
||||
"status": "ok",
|
||||
"run_id": run_id,
|
||||
"after_seq": normalized_after_seq,
|
||||
"last_seq": str(events[-1]["seq"]) if events else normalized_after_seq,
|
||||
"events": events,
|
||||
}
|
||||
return _json_tool_command(payload, runtime.tool_call_id)
|
||||
|
||||
async def asubagent_cancel(
|
||||
run_id: Annotated[str, SUBAGENT_RUN_ID_ARG],
|
||||
runtime: ToolRuntime,
|
||||
) -> str | Command:
|
||||
from yuxi.services.agent_run_service import request_cancel_agent_run
|
||||
|
||||
parent_runtime, runtime_error = self._require_async_parent_runtime("无法取消子智能体")
|
||||
if runtime_error:
|
||||
return runtime_error
|
||||
try:
|
||||
await self._get_verified_subagent_run(
|
||||
run_id=run_id,
|
||||
uid=parent_runtime.uid,
|
||||
created_by_run_id=parent_runtime.created_by_run_id,
|
||||
) # 校验子智能体归属
|
||||
|
||||
# 取消子智能体运行,返回最新 run 状态
|
||||
async with pg_manager.get_async_session_context() as db:
|
||||
run = await request_cancel_agent_run(run_id=run_id, current_uid=parent_runtime.uid, db=db)
|
||||
|
||||
except ValueError as exc:
|
||||
return str(exc)
|
||||
|
||||
subagent_service = _subagent_run_service_module()
|
||||
payload = {
|
||||
"status": run.status,
|
||||
"run_id": run.id,
|
||||
"thread_id": run.conversation_thread_id,
|
||||
**subagent_service.subagent_run_urls(run.id),
|
||||
}
|
||||
subagent_run = subagent_service.serialize_subagent_run_state(run)
|
||||
return _json_tool_command(payload, runtime.tool_call_id, subagent_run=subagent_run)
|
||||
|
||||
async def asubagent_await(
|
||||
run_id: Annotated[str, SUBAGENT_RUN_ID_ARG],
|
||||
runtime: ToolRuntime,
|
||||
) -> str | Command:
|
||||
from yuxi.services.agent_run_service import AgentRunWaitTimeout, await_agent_run_result
|
||||
|
||||
parent_runtime, runtime_error = self._require_async_parent_runtime("无法等待子智能体")
|
||||
if runtime_error:
|
||||
return runtime_error
|
||||
wait_timed_out = False
|
||||
try:
|
||||
# 等待前校验 run 归属,避免越权等待其它子任务
|
||||
await self._get_verified_subagent_run(
|
||||
run_id=run_id,
|
||||
uid=parent_runtime.uid,
|
||||
created_by_run_id=parent_runtime.created_by_run_id,
|
||||
)
|
||||
# 等待结束后重新读取已验证的最新 run 状态
|
||||
result = await await_agent_run_result(run_id=run_id, current_uid=parent_runtime.uid)
|
||||
run = await self._get_verified_subagent_run(
|
||||
run_id=run_id,
|
||||
uid=parent_runtime.uid,
|
||||
created_by_run_id=parent_runtime.created_by_run_id,
|
||||
)
|
||||
except AgentRunWaitTimeout as exc:
|
||||
wait_timed_out = True
|
||||
result = exc.result
|
||||
try:
|
||||
run = await self._get_verified_subagent_run(
|
||||
run_id=run_id,
|
||||
uid=parent_runtime.uid,
|
||||
created_by_run_id=parent_runtime.created_by_run_id,
|
||||
)
|
||||
except ValueError as verify_exc:
|
||||
return str(verify_exc)
|
||||
except ValueError as exc:
|
||||
return str(exc)
|
||||
|
||||
subagent_service = _subagent_run_service_module()
|
||||
payload = {
|
||||
"status": run.status,
|
||||
"run_id": run.id,
|
||||
"thread_id": run.conversation_thread_id,
|
||||
"result": result,
|
||||
}
|
||||
if wait_timed_out:
|
||||
payload["wait_timed_out"] = True
|
||||
payload["message"] = "子智能体仍在运行,等待最终结果超时;请稍后继续查询。"
|
||||
subagent_run = subagent_service.serialize_subagent_run_state(run)
|
||||
return _json_tool_command(payload, runtime.tool_call_id, subagent_run=subagent_run)
|
||||
|
||||
return [
|
||||
_async_only_tool(
|
||||
name="subagent_start",
|
||||
coroutine=asubagent_start,
|
||||
description=SUBAGENT_START_DESCRIPTION + "\n\nAvailable subagent slugs:\n" + available_agents,
|
||||
),
|
||||
_async_only_tool(
|
||||
name="subagent_status",
|
||||
coroutine=asubagent_status,
|
||||
description=SUBAGENT_STATUS_DESCRIPTION,
|
||||
),
|
||||
_async_only_tool(
|
||||
name="subagent_events",
|
||||
coroutine=asubagent_events,
|
||||
description=SUBAGENT_EVENTS_DESCRIPTION,
|
||||
),
|
||||
_async_only_tool(
|
||||
name="subagent_cancel",
|
||||
coroutine=asubagent_cancel,
|
||||
description=SUBAGENT_CANCEL_DESCRIPTION,
|
||||
),
|
||||
_async_only_tool(
|
||||
name="subagent_await",
|
||||
coroutine=asubagent_await,
|
||||
description=SUBAGENT_AWAIT_DESCRIPTION,
|
||||
),
|
||||
]
|
||||
|
||||
def _parent_runtime(self) -> _ParentRuntime:
|
||||
"""从父智能体 context 中抽取子智能体运行所需的最小父运行信息。"""
|
||||
parent_thread_id = str(getattr(self.parent_context, "parent_thread_id", None) or self.parent_context.thread_id)
|
||||
file_thread_id = str(getattr(self.parent_context, "file_thread_id", None) or parent_thread_id)
|
||||
uid = str(getattr(self.parent_context, "uid", "") or "").strip()
|
||||
created_by_run_id = str(getattr(self.parent_context, "run_id", "") or "").strip()
|
||||
return _ParentRuntime(
|
||||
file_thread_id=file_thread_id,
|
||||
uid=uid,
|
||||
created_by_run_id=created_by_run_id,
|
||||
)
|
||||
|
||||
def _require_async_parent_runtime(self, error_prefix: str) -> tuple[_ParentRuntime, str | None]:
|
||||
"""校验后台子智能体工具必须依赖的父运行上下文。"""
|
||||
parent_runtime = self._parent_runtime()
|
||||
if not parent_runtime.uid:
|
||||
return parent_runtime, f"{error_prefix}:当前运行时缺少 uid"
|
||||
if not parent_runtime.created_by_run_id:
|
||||
return parent_runtime, f"{error_prefix}:当前运行时缺少父运行 ID"
|
||||
return parent_runtime, None
|
||||
|
||||
async def _start_subagent(
|
||||
self,
|
||||
*,
|
||||
description: str,
|
||||
subagent_slug: str,
|
||||
runtime: ToolRuntime,
|
||||
thread_id: str | None,
|
||||
error_prefix: str,
|
||||
) -> tuple[_StartedSubagent | None, str | Command | None]:
|
||||
"""校验并启动(或继续)后台子智能体 run;成功返回启动结果,失败返回可直接回传的错误响应。"""
|
||||
if subagent_slug not in self.subagents:
|
||||
allowed = ", ".join(f"`{slug}`" for slug in self.subagents)
|
||||
return None, f"无法调用子智能体 {subagent_slug},可用子智能体只有:{allowed}"
|
||||
if not runtime.tool_call_id:
|
||||
raise ValueError("Tool call ID is required for subagent invocation")
|
||||
|
||||
parent_runtime, runtime_error = self._require_async_parent_runtime(error_prefix)
|
||||
if runtime_error:
|
||||
return None, runtime_error
|
||||
|
||||
agent_item = self.subagents[subagent_slug]
|
||||
input_message = build_chat_input_message(description)
|
||||
subagent_service = _subagent_run_service_module()
|
||||
try:
|
||||
async with pg_manager.get_async_session_context() as db:
|
||||
result = await subagent_service.SubagentRunService(db).start(
|
||||
uid=parent_runtime.uid,
|
||||
created_by_run_id=parent_runtime.created_by_run_id,
|
||||
agent_item=agent_item,
|
||||
input_message=input_message,
|
||||
tool_call_id=runtime.tool_call_id,
|
||||
requested_thread_id=thread_id,
|
||||
file_thread_id=parent_runtime.file_thread_id,
|
||||
model_spec=self._subagent_model_override(agent_item),
|
||||
)
|
||||
except subagent_service.SubagentRunBusy as exc:
|
||||
return None, _json_tool_command(exc.to_payload(), runtime.tool_call_id)
|
||||
except ValueError as exc:
|
||||
return None, str(exc)
|
||||
return _StartedSubagent(result=result, parent_runtime=parent_runtime, agent_item=agent_item), None
|
||||
|
||||
def _subagent_model_override(self, agent_item: Agent) -> str | None:
|
||||
"""当子智能体未显式配置模型时,沿用父智能体当前模型。"""
|
||||
config_context = (
|
||||
(agent_item.config_json or {}).get("context") if isinstance(agent_item.config_json, dict) else None
|
||||
)
|
||||
configured_model = ""
|
||||
if isinstance(config_context, dict):
|
||||
configured_model = str(config_context.get("model") or "").strip()
|
||||
if configured_model:
|
||||
return None
|
||||
return str(getattr(self.parent_context, "model", None) or "").strip() or None
|
||||
|
||||
async def _get_verified_subagent_run(self, *, run_id: str, uid: str, created_by_run_id: str):
|
||||
"""在工具调用前按父 run 作用域校验子 run 归属。"""
|
||||
subagent_service = _subagent_run_service_module()
|
||||
async with pg_manager.get_async_session_context() as db:
|
||||
return await subagent_service.SubagentRunService(db).get_run_for_creator(
|
||||
uid=uid,
|
||||
created_by_run_id=created_by_run_id,
|
||||
run_id=run_id,
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _ParentRuntime:
|
||||
file_thread_id: str
|
||||
uid: str
|
||||
created_by_run_id: str
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _StartedSubagent:
|
||||
"""``_start_subagent`` 的结果:子 run 启动产物及其依赖的父运行上下文。"""
|
||||
|
||||
result: Any # SubagentStartResult
|
||||
parent_runtime: _ParentRuntime
|
||||
agent_item: Agent
|
||||
|
||||
|
||||
def _task_result_response(result: dict[str, Any], tool_call_id: str, subagent_run: dict[str, Any]) -> Command:
|
||||
"""把后台子智能体 run 的最终结果转换为同步 task 工具结果。"""
|
||||
output = str(result.get("output") or "").strip()
|
||||
error = result.get("error") if isinstance(result.get("error"), dict) else None
|
||||
if not output and error:
|
||||
output = str(error.get("message") or "子智能体运行失败")
|
||||
if not output:
|
||||
output = "子智能体已完成任务,但没有返回文本结果。"
|
||||
|
||||
tool_result = _tool_result_with_thread_id(subagent_run["child_thread_id"], output)
|
||||
return Command(
|
||||
update={"messages": [ToolMessage(tool_result, tool_call_id=tool_call_id)], "subagent_runs": [subagent_run]}
|
||||
)
|
||||
|
||||
|
||||
def _task_wait_timeout_response(result: dict[str, Any], tool_call_id: str, subagent_run: dict[str, Any]) -> Command:
|
||||
"""同步 task 等待到达上限时,明确告诉父智能体子 run 仍未终结。"""
|
||||
status = str(result.get("status") or subagent_run.get("status") or "running")
|
||||
run_id = str(result.get("agent_run_id") or subagent_run["run_id"])
|
||||
output = (
|
||||
f"子智能体仍在运行(status: {status}),尚未返回最终文本结果。\n"
|
||||
f"run_id: {run_id}\n"
|
||||
"请稍后使用 subagent_status 或 subagent_await 查询结果;不要把当前结果视为任务已完成。"
|
||||
)
|
||||
tool_result = _tool_result_with_thread_id(subagent_run["child_thread_id"], output)
|
||||
return Command(
|
||||
update={"messages": [ToolMessage(tool_result, tool_call_id=tool_call_id)], "subagent_runs": [subagent_run]}
|
||||
)
|
||||
|
||||
|
||||
def _json_tool_command(
|
||||
payload: dict[str, Any],
|
||||
tool_call_id: str,
|
||||
*,
|
||||
subagent_run: dict[str, Any] | None = None,
|
||||
) -> Command:
|
||||
"""把后台子智能体工具的结构化结果包装成 ToolMessage。"""
|
||||
content = json.dumps(payload, ensure_ascii=False, indent=2)
|
||||
update: dict[str, Any] = {"messages": [ToolMessage(content, tool_call_id=tool_call_id)]}
|
||||
if subagent_run is not None:
|
||||
update["subagent_runs"] = [subagent_run]
|
||||
return Command(update=update)
|
||||
|
||||
|
||||
def _tool_result_with_thread_id(child_thread_id: str, content: str) -> str:
|
||||
"""把子线程 ID 放进工具结果,方便后续继续同一子任务。"""
|
||||
return f"> 子智能体线程 ID: {child_thread_id}\n\n---\n\n{content}"
|
||||
@@ -0,0 +1,777 @@
|
||||
"""Yuxi adapter for DeepAgents conversation summarization middleware."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import hashlib
|
||||
import re
|
||||
import warnings
|
||||
from collections.abc import Awaitable, Callable, Iterable
|
||||
from contextvars import ContextVar
|
||||
from typing import Any
|
||||
|
||||
from deepagents.middleware.summarization import (
|
||||
Command,
|
||||
ContextOverflowError,
|
||||
SummarizationMiddleware,
|
||||
_aclip_overflow_tail,
|
||||
_clip_overflow_tail,
|
||||
)
|
||||
from langchain.agents.middleware.summarization import ContextSize
|
||||
from langchain.agents.middleware.types import ExtendedModelResponse, ModelRequest, ModelResponse
|
||||
from langchain.chat_models import BaseChatModel
|
||||
from langchain_core.messages import AIMessage, AnyMessage, ToolMessage, get_buffer_string
|
||||
from langchain_core.messages.utils import count_tokens_approximately
|
||||
from langgraph.config import get_stream_writer
|
||||
from langgraph.constants import TAG_NOSTREAM
|
||||
|
||||
from yuxi.agents.backends.composite import create_agent_composite_backend
|
||||
from yuxi.utils.logging_config import logger
|
||||
from yuxi.utils.paths import VIRTUAL_PATH_CONVERSATION_HISTORY, VIRTUAL_PATH_LARGE_TOOL_RESULTS
|
||||
|
||||
_APPROX_CHARS_PER_TOKEN = 4
|
||||
_DEFAULT_SUMMARY_TOOL_RESULT_LIMIT_TOKENS = 300
|
||||
_DEFAULT_L1_L2_TRIGGER_RATIO = 0.4
|
||||
_DEFAULT_TOOL_ARG_MAX_LENGTH = 2000
|
||||
_TRUNCATED_TOOL_ARG_TEXT = "...(argument truncated for context view)"
|
||||
_TOOL_RESULT_SAVED_MARKER = "yuxi_tool_result_saved"
|
||||
_SUMMARY_BACKEND: ContextVar[Any | None] = ContextVar("yuxi_summary_backend", default=None)
|
||||
_SUMMARY_SANITIZED_MESSAGES: ContextVar[dict[tuple[int, ...], list[AnyMessage]] | None] = ContextVar(
|
||||
"yuxi_summary_sanitized_messages",
|
||||
default=None,
|
||||
)
|
||||
_SUMMARY_COMPRESSION_STATE: ContextVar[dict[str, bool] | None] = ContextVar(
|
||||
"yuxi_summary_compression_state",
|
||||
default=None,
|
||||
)
|
||||
|
||||
|
||||
def _emit_compression(status: str, **extra: Any) -> None:
|
||||
try:
|
||||
writer = get_stream_writer()
|
||||
except RuntimeError:
|
||||
return
|
||||
writer({"type": "yuxi.context_compression", "status": status, **extra})
|
||||
|
||||
|
||||
def _emit_compression_started_once() -> None:
|
||||
state = _SUMMARY_COMPRESSION_STATE.get()
|
||||
if state is not None and state.get("started"):
|
||||
return
|
||||
if state is not None:
|
||||
state["started"] = True
|
||||
_emit_compression("started")
|
||||
|
||||
|
||||
def _count_tokens_for_summary_trigger(messages: Iterable[Any], **kwargs: Any) -> int:
|
||||
kwargs.pop("use_usage_metadata_scaling", None)
|
||||
return count_tokens_approximately(messages, use_usage_metadata_scaling=False, **kwargs)
|
||||
|
||||
|
||||
def _extract_text_content(content: Any) -> str:
|
||||
if isinstance(content, str):
|
||||
return content
|
||||
if isinstance(content, list):
|
||||
parts: list[str] = []
|
||||
for item in content:
|
||||
if isinstance(item, str):
|
||||
parts.append(item)
|
||||
elif isinstance(item, dict) and isinstance(item.get("text"), str):
|
||||
parts.append(item["text"])
|
||||
return "\n".join(part for part in parts if part)
|
||||
return "" if content is None else str(content)
|
||||
|
||||
|
||||
def _tool_result_path(tool_name: str | None, content: str, large_tool_results_prefix: str) -> str:
|
||||
safe_tool_name = re.sub(r"[^A-Za-z0-9_.-]+", "-", (tool_name or "").strip()).strip(".-") or "tool-result"
|
||||
digest = hashlib.sha256(content.encode("utf-8")).hexdigest()[:16]
|
||||
return f"{large_tool_results_prefix}/{safe_tool_name}-{digest}.txt"
|
||||
|
||||
|
||||
def _preview_tool_result(content: str, token_limit: int | None) -> tuple[str, int]:
|
||||
text = content.strip()
|
||||
if token_limit is None:
|
||||
return text, 0
|
||||
if token_limit <= 0:
|
||||
return "", len(text)
|
||||
|
||||
max_chars = token_limit * _APPROX_CHARS_PER_TOKEN
|
||||
if len(text) <= max_chars:
|
||||
return text, 0
|
||||
|
||||
preview = text[:max_chars].rstrip()
|
||||
return preview, len(text) - len(preview)
|
||||
|
||||
|
||||
def _write_tool_result(backend, path: str, content: str) -> str | None:
|
||||
if backend is None:
|
||||
return None
|
||||
|
||||
result = backend.write(path, content)
|
||||
error = getattr(result, "error", None)
|
||||
if not error:
|
||||
return path
|
||||
if "already exists" in str(error).lower():
|
||||
return path
|
||||
raise RuntimeError(f"Failed to write tool result to {path}: {error}")
|
||||
|
||||
|
||||
def _estimated_content_tokens(content: str) -> int:
|
||||
return max((len(content) + _APPROX_CHARS_PER_TOKEN - 1) // _APPROX_CHARS_PER_TOKEN, 1)
|
||||
|
||||
|
||||
def _tool_result_replacement_content(
|
||||
message: ToolMessage,
|
||||
*,
|
||||
backend,
|
||||
tool_result_offload_token_limit: int | None,
|
||||
large_tool_results_prefix: str,
|
||||
) -> str:
|
||||
content = _extract_text_content(message.content)
|
||||
approx_tokens = max((len(content) + _APPROX_CHARS_PER_TOKEN - 1) // _APPROX_CHARS_PER_TOKEN, 1)
|
||||
tool_name = message.name if isinstance(message.name, str) and message.name else None
|
||||
path = _write_tool_result(
|
||||
backend,
|
||||
_tool_result_path(tool_name, content, large_tool_results_prefix),
|
||||
content,
|
||||
)
|
||||
preview, omitted_chars = _preview_tool_result(content, tool_result_offload_token_limit)
|
||||
|
||||
lines = [
|
||||
"[Tool result saved]",
|
||||
f"Tool: {tool_name or 'unknown'}",
|
||||
f"Approx tokens: {approx_tokens}",
|
||||
]
|
||||
if path:
|
||||
lines.append(f"Full output path: {path}")
|
||||
if preview:
|
||||
lines.extend(["", "Output preview:", preview])
|
||||
if omitted_chars:
|
||||
lines.append(f"\n[Truncated {omitted_chars} chars. Read the full output from the saved file.]")
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
def _should_offload_tool_message(message: ToolMessage, trigger_tokens: int | None) -> bool:
|
||||
if trigger_tokens is None:
|
||||
return True
|
||||
if trigger_tokens <= 0:
|
||||
return True
|
||||
content = _extract_text_content(message.content)
|
||||
return _estimated_content_tokens(content) > trigger_tokens
|
||||
|
||||
|
||||
def _replace_tool_message_content(
|
||||
message: ToolMessage,
|
||||
*,
|
||||
backend,
|
||||
tool_result_offload_token_limit: int | None,
|
||||
large_tool_results_prefix: str,
|
||||
) -> ToolMessage:
|
||||
additional_kwargs = dict(getattr(message, "additional_kwargs", {}) or {})
|
||||
additional_kwargs[_TOOL_RESULT_SAVED_MARKER] = True
|
||||
return message.model_copy(
|
||||
update={
|
||||
"content": _tool_result_replacement_content(
|
||||
message,
|
||||
backend=backend,
|
||||
tool_result_offload_token_limit=tool_result_offload_token_limit,
|
||||
large_tool_results_prefix=large_tool_results_prefix,
|
||||
),
|
||||
"additional_kwargs": additional_kwargs,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _truncate_string_arg(value: str, max_length: int) -> str:
|
||||
if len(value) <= max_length:
|
||||
return value
|
||||
return f"{value[:20]}{_TRUNCATED_TOOL_ARG_TEXT}"
|
||||
|
||||
|
||||
def _truncate_tool_call_args(tool_call: dict[str, Any], max_length: int) -> tuple[dict[str, Any], bool]:
|
||||
if tool_call.get("name") not in {"write_file", "edit_file"}:
|
||||
return tool_call, False
|
||||
args = tool_call.get("args")
|
||||
if not isinstance(args, dict):
|
||||
return tool_call, False
|
||||
|
||||
truncated_args: dict[str, Any] = {}
|
||||
modified = False
|
||||
for key, value in args.items():
|
||||
if isinstance(value, str) and len(value) > max_length:
|
||||
truncated_args[key] = _truncate_string_arg(value, max_length)
|
||||
modified = True
|
||||
else:
|
||||
truncated_args[key] = value
|
||||
if not modified:
|
||||
return tool_call, False
|
||||
return {**tool_call, "args": truncated_args}, True
|
||||
|
||||
|
||||
def _truncate_provider_tool_calls(additional_kwargs: dict[str, Any], max_length: int) -> tuple[dict[str, Any], bool]:
|
||||
raw_tool_calls = additional_kwargs.get("tool_calls")
|
||||
if not isinstance(raw_tool_calls, list):
|
||||
return additional_kwargs, False
|
||||
|
||||
updated_tool_calls: list[Any] = []
|
||||
modified = False
|
||||
for raw_call in raw_tool_calls:
|
||||
if not isinstance(raw_call, dict):
|
||||
updated_tool_calls.append(raw_call)
|
||||
continue
|
||||
function = raw_call.get("function")
|
||||
if not isinstance(function, dict):
|
||||
updated_tool_calls.append(raw_call)
|
||||
continue
|
||||
if function.get("name") not in {"write_file", "edit_file"}:
|
||||
updated_tool_calls.append(raw_call)
|
||||
continue
|
||||
arguments = function.get("arguments")
|
||||
if not isinstance(arguments, str) or len(arguments) <= max_length:
|
||||
updated_tool_calls.append(raw_call)
|
||||
continue
|
||||
|
||||
updated_function = {**function, "arguments": _truncate_string_arg(arguments, max_length)}
|
||||
updated_tool_calls.append({**raw_call, "function": updated_function})
|
||||
modified = True
|
||||
|
||||
if not modified:
|
||||
return additional_kwargs, False
|
||||
return {**additional_kwargs, "tool_calls": updated_tool_calls}, True
|
||||
|
||||
|
||||
def _truncate_ai_tool_call_args(message: AIMessage, *, max_length: int) -> AIMessage:
|
||||
if not message.tool_calls and not getattr(message, "additional_kwargs", None):
|
||||
return message
|
||||
|
||||
updated_tool_calls: list[Any] = []
|
||||
tool_calls_modified = False
|
||||
for tool_call in message.tool_calls or []:
|
||||
if not isinstance(tool_call, dict):
|
||||
updated_tool_calls.append(tool_call)
|
||||
continue
|
||||
updated, modified = _truncate_tool_call_args(tool_call, max_length)
|
||||
updated_tool_calls.append(updated)
|
||||
tool_calls_modified = tool_calls_modified or modified
|
||||
|
||||
additional_kwargs = dict(getattr(message, "additional_kwargs", {}) or {})
|
||||
additional_kwargs, additional_kwargs_modified = _truncate_provider_tool_calls(additional_kwargs, max_length)
|
||||
|
||||
if not tool_calls_modified and not additional_kwargs_modified:
|
||||
return message
|
||||
|
||||
updated_message = message.model_copy(update={"additional_kwargs": additional_kwargs})
|
||||
if tool_calls_modified:
|
||||
updated_message.tool_calls = updated_tool_calls
|
||||
return updated_message
|
||||
|
||||
|
||||
def sanitize_messages_for_summary(
|
||||
messages: list[AnyMessage],
|
||||
*,
|
||||
backend=None,
|
||||
tool_result_offload_token_limit: int | None = _DEFAULT_SUMMARY_TOOL_RESULT_LIMIT_TOKENS,
|
||||
large_tool_results_prefix: str = VIRTUAL_PATH_LARGE_TOOL_RESULTS,
|
||||
) -> list[AnyMessage]:
|
||||
"""Build a compact summary/offload view by replacing only ToolMessage content."""
|
||||
sanitized: list[AnyMessage] = []
|
||||
for message in messages:
|
||||
if isinstance(message, ToolMessage):
|
||||
if getattr(message, "additional_kwargs", {}).get(_TOOL_RESULT_SAVED_MARKER) is True:
|
||||
sanitized.append(message)
|
||||
continue
|
||||
sanitized.append(
|
||||
_replace_tool_message_content(
|
||||
message,
|
||||
backend=backend,
|
||||
tool_result_offload_token_limit=tool_result_offload_token_limit,
|
||||
large_tool_results_prefix=large_tool_results_prefix,
|
||||
)
|
||||
)
|
||||
continue
|
||||
sanitized.append(message)
|
||||
return sanitized
|
||||
|
||||
|
||||
class YuxiSummarizationMiddleware(SummarizationMiddleware):
|
||||
"""DeepAgents summarization middleware with Yuxi-specific tool-call sanitization."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*args,
|
||||
tool_result_offload_token_limit: int | None = _DEFAULT_SUMMARY_TOOL_RESULT_LIMIT_TOKENS,
|
||||
l1_l2_trigger_ratio: float = _DEFAULT_L1_L2_TRIGGER_RATIO,
|
||||
tool_arg_max_length: int = _DEFAULT_TOOL_ARG_MAX_LENGTH,
|
||||
**kwargs,
|
||||
) -> None:
|
||||
super().__init__(*args, **kwargs)
|
||||
self.tool_result_offload_token_limit = tool_result_offload_token_limit
|
||||
self.l1_l2_trigger_ratio = l1_l2_trigger_ratio
|
||||
self.tool_arg_max_length = tool_arg_max_length
|
||||
|
||||
def _should_summarize(self, messages: list[AnyMessage], total_tokens: int) -> bool:
|
||||
if not self._lc_helper._trigger_clauses:
|
||||
return False
|
||||
|
||||
for clause in self._lc_helper._trigger_clauses:
|
||||
clause_met = True
|
||||
for kind, value in clause.items():
|
||||
if kind == "messages":
|
||||
if len(messages) < value:
|
||||
clause_met = False
|
||||
break
|
||||
elif kind == "tokens":
|
||||
if total_tokens < value:
|
||||
clause_met = False
|
||||
break
|
||||
elif kind == "fraction":
|
||||
max_input_tokens = self._get_profile_limits()
|
||||
if max_input_tokens is None:
|
||||
clause_met = False
|
||||
break
|
||||
threshold = int(max_input_tokens * value)
|
||||
if threshold <= 0:
|
||||
threshold = 1
|
||||
if total_tokens < threshold:
|
||||
clause_met = False
|
||||
break
|
||||
if clause_met:
|
||||
return True
|
||||
return False
|
||||
|
||||
def _sanitize_messages_for_summary(
|
||||
self,
|
||||
messages: list[AnyMessage],
|
||||
*,
|
||||
backend,
|
||||
) -> list[AnyMessage]:
|
||||
_ = backend
|
||||
cache = _SUMMARY_SANITIZED_MESSAGES.get()
|
||||
cache_key = tuple(id(message) for message in messages)
|
||||
if cache is not None and cache_key in cache:
|
||||
return cache[cache_key]
|
||||
|
||||
sanitized = messages
|
||||
if cache is not None:
|
||||
cache[cache_key] = sanitized
|
||||
return sanitized
|
||||
|
||||
def _sanitize_messages_for_l1(
|
||||
self,
|
||||
messages: list[AnyMessage],
|
||||
*,
|
||||
backend,
|
||||
) -> list[AnyMessage]:
|
||||
compacted: list[AnyMessage] = []
|
||||
modified = False
|
||||
for message in messages:
|
||||
if isinstance(message, AIMessage):
|
||||
updated = _truncate_ai_tool_call_args(message, max_length=self.tool_arg_max_length)
|
||||
compacted.append(updated)
|
||||
modified = modified or updated is not message
|
||||
continue
|
||||
if isinstance(message, ToolMessage) and _should_offload_tool_message(
|
||||
message,
|
||||
self.tool_result_offload_token_limit,
|
||||
):
|
||||
updated = _replace_tool_message_content(
|
||||
message,
|
||||
backend=backend,
|
||||
tool_result_offload_token_limit=self.tool_result_offload_token_limit,
|
||||
large_tool_results_prefix=self._large_tool_results_prefix,
|
||||
)
|
||||
compacted.append(updated)
|
||||
modified = modified or updated is not message
|
||||
continue
|
||||
compacted.append(message)
|
||||
return compacted if modified else messages
|
||||
|
||||
def _count_request_tokens(
|
||||
self,
|
||||
messages: list[AnyMessage],
|
||||
*,
|
||||
system_message,
|
||||
tools,
|
||||
) -> int:
|
||||
counted_messages = [system_message, *messages] if system_message is not None else messages
|
||||
try:
|
||||
return self.token_counter(counted_messages, tools=tools) # ty: ignore[unknown-argument]
|
||||
except TypeError:
|
||||
return self.token_counter(counted_messages)
|
||||
|
||||
def _entry_trigger_tokens(self) -> int | None:
|
||||
token_thresholds: list[int] = []
|
||||
for clause in getattr(self._lc_helper, "_trigger_clauses", []) or []:
|
||||
value = clause.get("tokens")
|
||||
if isinstance(value, int) and value > 0:
|
||||
token_thresholds.append(value)
|
||||
fraction = clause.get("fraction")
|
||||
if isinstance(fraction, int | float):
|
||||
max_input_tokens = self._get_profile_limits()
|
||||
if max_input_tokens is not None:
|
||||
threshold = int(max_input_tokens * fraction)
|
||||
token_thresholds.append(max(threshold, 1))
|
||||
return min(token_thresholds) if token_thresholds else None
|
||||
|
||||
def _should_run_l1(self, messages: list[AnyMessage], total_tokens: int) -> bool:
|
||||
return self._should_summarize(messages, total_tokens)
|
||||
|
||||
def _should_run_l2(self, compacted_total_tokens: int, entry_threshold_tokens: int | None) -> bool:
|
||||
if entry_threshold_tokens is None:
|
||||
return True
|
||||
threshold = max(int(entry_threshold_tokens * self.l1_l2_trigger_ratio), 1)
|
||||
return compacted_total_tokens > threshold
|
||||
|
||||
def _backend_for_request(self, request: ModelRequest):
|
||||
try:
|
||||
return self._get_backend(request.state, request.runtime)
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _summarization_event_from_result(result: Any) -> dict | None:
|
||||
if not isinstance(result, ExtendedModelResponse):
|
||||
return None
|
||||
command = getattr(result, "command", None)
|
||||
update = getattr(command, "update", None) if command is not None else None
|
||||
if not isinstance(update, dict):
|
||||
return None
|
||||
event = update.get("_summarization_event")
|
||||
return event if isinstance(event, dict) else None
|
||||
|
||||
def _emit_completed(self, result: Any) -> None:
|
||||
event = self._summarization_event_from_result(result)
|
||||
if event is not None:
|
||||
_emit_compression(
|
||||
"completed",
|
||||
cutoff_index=event.get("cutoff_index"),
|
||||
file_path=event.get("file_path"),
|
||||
)
|
||||
|
||||
# 重写 _create_summary/_acreate_summary 以在摘要 LLM 调用上挂 TAG_NOSTREAM:父类
|
||||
# 的 model.invoke 带 lc_source 元数据但无 nostream 标记,其 token 流会被 LangGraph
|
||||
# messages stream 捕获并广播到前端,形成 phantom 摘要消息。带 TAG_NOSTREAM 后流式
|
||||
# 层在源头跳过该调用,无需 chat_service 下游过滤,主 messages 流天然只含用户可见回复。
|
||||
# 父类硬编码 invoke config 且无 tags 钩子(self.model 为中间件实例共享属性,并发下不能
|
||||
# 临时换绑 bind(tags=...)),故只能重写;trim/format 是纯同步逻辑,抽到 _build_summary_prompt
|
||||
# 供 sync/async 两条路径共用,避免逐字重复。
|
||||
_SUMMARY_INVOKE_CONFIG = {"metadata": {"lc_source": "summarization"}, "tags": [TAG_NOSTREAM]}
|
||||
|
||||
def _build_summary_prompt(self, sanitized: list[AnyMessage]) -> str | None:
|
||||
trimmed = self._lc_helper._trim_messages_for_summary(sanitized)
|
||||
if not trimmed:
|
||||
return None
|
||||
return self._lc_helper.summary_prompt.format(messages=get_buffer_string(trimmed, format="xml")).rstrip()
|
||||
|
||||
def _create_summary(self, messages_to_summarize: list[AnyMessage]) -> str:
|
||||
sanitized = self._sanitize_messages_for_summary(
|
||||
messages_to_summarize,
|
||||
backend=_SUMMARY_BACKEND.get(),
|
||||
)
|
||||
if not sanitized:
|
||||
return "No previous conversation history."
|
||||
prompt = self._build_summary_prompt(sanitized)
|
||||
if prompt is None:
|
||||
return "Previous conversation was too long to summarize."
|
||||
try:
|
||||
return self.model.invoke(prompt, config=self._SUMMARY_INVOKE_CONFIG).text.strip()
|
||||
except Exception as e:
|
||||
return f"Error generating summary: {e!s}"
|
||||
|
||||
async def _acreate_summary(self, messages_to_summarize: list[AnyMessage]) -> str:
|
||||
sanitized = self._sanitize_messages_for_summary(
|
||||
messages_to_summarize,
|
||||
backend=_SUMMARY_BACKEND.get(),
|
||||
)
|
||||
if not sanitized:
|
||||
return "No previous conversation history."
|
||||
prompt = self._build_summary_prompt(sanitized)
|
||||
if prompt is None:
|
||||
return "Previous conversation was too long to summarize."
|
||||
try:
|
||||
response = await self.model.ainvoke(prompt, config=self._SUMMARY_INVOKE_CONFIG)
|
||||
return response.text.strip()
|
||||
except Exception as e:
|
||||
return f"Error generating summary: {e!s}"
|
||||
|
||||
def _offload_to_backend(self, backend, messages: list[AnyMessage]) -> str | None:
|
||||
_emit_compression_started_once()
|
||||
return super()._offload_to_backend(
|
||||
backend,
|
||||
self._sanitize_messages_for_summary(
|
||||
messages,
|
||||
backend=backend,
|
||||
),
|
||||
)
|
||||
|
||||
async def _aoffload_to_backend(self, backend, messages: list[AnyMessage]) -> str | None:
|
||||
_emit_compression_started_once()
|
||||
return await super()._aoffload_to_backend(
|
||||
backend,
|
||||
self._sanitize_messages_for_summary(
|
||||
messages,
|
||||
backend=backend,
|
||||
),
|
||||
)
|
||||
|
||||
def wrap_model_call(
|
||||
self,
|
||||
request: ModelRequest,
|
||||
handler: Callable[[ModelRequest], ModelResponse],
|
||||
) -> ModelResponse:
|
||||
backend_token = _SUMMARY_BACKEND.set(self._backend_for_request(request))
|
||||
sanitized_token = _SUMMARY_SANITIZED_MESSAGES.set({})
|
||||
compression_state: dict[str, bool] = {"started": False}
|
||||
compression_token = _SUMMARY_COMPRESSION_STATE.set(compression_state)
|
||||
try:
|
||||
try:
|
||||
result = self._wrap_model_call_with_l1(request, handler)
|
||||
except Exception as exc:
|
||||
if compression_state.get("started"):
|
||||
_emit_compression("failed", error=repr(exc))
|
||||
raise
|
||||
self._emit_completed(result)
|
||||
return result
|
||||
finally:
|
||||
_SUMMARY_COMPRESSION_STATE.reset(compression_token)
|
||||
_SUMMARY_SANITIZED_MESSAGES.reset(sanitized_token)
|
||||
_SUMMARY_BACKEND.reset(backend_token)
|
||||
|
||||
async def awrap_model_call(
|
||||
self,
|
||||
request: ModelRequest,
|
||||
handler: Callable[[ModelRequest], Awaitable[ModelResponse]],
|
||||
) -> ModelResponse:
|
||||
backend_token = _SUMMARY_BACKEND.set(self._backend_for_request(request))
|
||||
sanitized_token = _SUMMARY_SANITIZED_MESSAGES.set({})
|
||||
compression_state: dict[str, bool] = {"started": False}
|
||||
compression_token = _SUMMARY_COMPRESSION_STATE.set(compression_state)
|
||||
try:
|
||||
try:
|
||||
result = await self._awrap_model_call_with_l1(request, handler)
|
||||
except Exception as exc:
|
||||
if compression_state.get("started"):
|
||||
_emit_compression("failed", error=repr(exc))
|
||||
raise
|
||||
self._emit_completed(result)
|
||||
return result
|
||||
finally:
|
||||
_SUMMARY_COMPRESSION_STATE.reset(compression_token)
|
||||
_SUMMARY_SANITIZED_MESSAGES.reset(sanitized_token)
|
||||
_SUMMARY_BACKEND.reset(backend_token)
|
||||
|
||||
def _wrap_model_call_with_l1(
|
||||
self,
|
||||
request: ModelRequest,
|
||||
handler: Callable[[ModelRequest], ModelResponse],
|
||||
) -> ModelResponse | ExtendedModelResponse:
|
||||
effective_messages = self._get_effective_messages(request)
|
||||
truncated_messages, _ = self._truncate_args(
|
||||
effective_messages,
|
||||
request.system_message,
|
||||
request.tools,
|
||||
)
|
||||
total_tokens = self._count_request_tokens(
|
||||
truncated_messages,
|
||||
system_message=request.system_message,
|
||||
tools=request.tools,
|
||||
)
|
||||
should_run_l1 = self._should_run_l1(truncated_messages, total_tokens)
|
||||
|
||||
overflow_triggered = False
|
||||
if not should_run_l1:
|
||||
try:
|
||||
return handler(request.override(messages=truncated_messages))
|
||||
except ContextOverflowError:
|
||||
overflow_triggered = True
|
||||
|
||||
backend = self._get_backend(request.state, request.runtime)
|
||||
l1_messages = self._sanitize_messages_for_l1(truncated_messages, backend=backend)
|
||||
if should_run_l1:
|
||||
_emit_compression_started_once()
|
||||
|
||||
l1_total_tokens = self._count_request_tokens(
|
||||
l1_messages,
|
||||
system_message=request.system_message,
|
||||
tools=request.tools,
|
||||
)
|
||||
should_run_l2 = overflow_triggered or self._should_run_l2(l1_total_tokens, self._entry_trigger_tokens())
|
||||
|
||||
if not should_run_l2:
|
||||
try:
|
||||
response = handler(request.override(messages=l1_messages))
|
||||
except ContextOverflowError:
|
||||
overflow_triggered = True
|
||||
else:
|
||||
if should_run_l1:
|
||||
_emit_compression("completed")
|
||||
return response
|
||||
|
||||
cutoff_index = self._determine_cutoff_index(l1_messages)
|
||||
if cutoff_index <= 0:
|
||||
response = handler(request.override(messages=l1_messages))
|
||||
if should_run_l1:
|
||||
_emit_compression("completed")
|
||||
return response
|
||||
|
||||
messages_to_summarize, preserved_messages = self._partition_messages(l1_messages, cutoff_index)
|
||||
new_state_tail: list[AnyMessage] = []
|
||||
if overflow_triggered:
|
||||
preserved_messages, new_state_tail = _clip_overflow_tail(
|
||||
preserved_messages,
|
||||
backend,
|
||||
keep=self._lc_helper.keep,
|
||||
max_input_tokens=self._get_profile_limits(),
|
||||
token_counter=self.token_counter,
|
||||
large_tool_results_prefix=self._large_tool_results_prefix,
|
||||
)
|
||||
|
||||
file_path = self._offload_to_backend(backend, messages_to_summarize)
|
||||
if file_path is None:
|
||||
msg = (
|
||||
"Offloading conversation history to backend failed during summarization. "
|
||||
"Older messages will not be recoverable."
|
||||
)
|
||||
logger.error(msg)
|
||||
warnings.warn(msg, stacklevel=2)
|
||||
|
||||
summary = self._create_summary(messages_to_summarize)
|
||||
new_messages = self._build_new_messages_with_path(summary, file_path)
|
||||
previous_event = request.state.get("_summarization_event")
|
||||
state_cutoff_index = self._compute_state_cutoff(previous_event, cutoff_index)
|
||||
new_event = {
|
||||
"cutoff_index": state_cutoff_index,
|
||||
"summary_message": new_messages[0],
|
||||
"file_path": file_path,
|
||||
}
|
||||
|
||||
response = handler(request.override(messages=[*new_messages, *preserved_messages]))
|
||||
update: dict[str, Any] = {"_summarization_event": new_event}
|
||||
if new_state_tail:
|
||||
update["messages"] = list(new_state_tail)
|
||||
return ExtendedModelResponse(model_response=response, command=Command(update=update))
|
||||
|
||||
async def _awrap_model_call_with_l1(
|
||||
self,
|
||||
request: ModelRequest,
|
||||
handler: Callable[[ModelRequest], Awaitable[ModelResponse]],
|
||||
) -> ModelResponse | ExtendedModelResponse:
|
||||
effective_messages = self._get_effective_messages(request)
|
||||
truncated_messages, _ = self._truncate_args(
|
||||
effective_messages,
|
||||
request.system_message,
|
||||
request.tools,
|
||||
)
|
||||
total_tokens = self._count_request_tokens(
|
||||
truncated_messages,
|
||||
system_message=request.system_message,
|
||||
tools=request.tools,
|
||||
)
|
||||
should_run_l1 = self._should_run_l1(truncated_messages, total_tokens)
|
||||
|
||||
overflow_triggered = False
|
||||
if not should_run_l1:
|
||||
try:
|
||||
return await handler(request.override(messages=truncated_messages))
|
||||
except ContextOverflowError:
|
||||
overflow_triggered = True
|
||||
|
||||
backend = self._get_backend(request.state, request.runtime)
|
||||
l1_messages = self._sanitize_messages_for_l1(truncated_messages, backend=backend)
|
||||
if should_run_l1:
|
||||
_emit_compression_started_once()
|
||||
|
||||
l1_total_tokens = self._count_request_tokens(
|
||||
l1_messages,
|
||||
system_message=request.system_message,
|
||||
tools=request.tools,
|
||||
)
|
||||
should_run_l2 = overflow_triggered or self._should_run_l2(l1_total_tokens, self._entry_trigger_tokens())
|
||||
|
||||
if not should_run_l2:
|
||||
try:
|
||||
response = await handler(request.override(messages=l1_messages))
|
||||
except ContextOverflowError:
|
||||
overflow_triggered = True
|
||||
else:
|
||||
if should_run_l1:
|
||||
_emit_compression("completed")
|
||||
return response
|
||||
|
||||
cutoff_index = self._determine_cutoff_index(l1_messages)
|
||||
if cutoff_index <= 0:
|
||||
response = await handler(request.override(messages=l1_messages))
|
||||
if should_run_l1:
|
||||
_emit_compression("completed")
|
||||
return response
|
||||
|
||||
messages_to_summarize, preserved_messages = self._partition_messages(l1_messages, cutoff_index)
|
||||
new_state_tail: list[AnyMessage] = []
|
||||
if overflow_triggered:
|
||||
preserved_messages, new_state_tail = await _aclip_overflow_tail(
|
||||
preserved_messages,
|
||||
backend,
|
||||
keep=self._lc_helper.keep,
|
||||
max_input_tokens=self._get_profile_limits(),
|
||||
token_counter=self.token_counter,
|
||||
large_tool_results_prefix=self._large_tool_results_prefix,
|
||||
)
|
||||
|
||||
# Offload 与 summary 互相独立,并发执行以避免串行等待一次文件 I/O + 一次
|
||||
# LLM 调用;_SUMMARY_SANITIZED_MESSAGES 的 id 缓存保证两路 sanitize 不会重复
|
||||
# 写入工具结果文件,offload 失败返回 None 时 summary 仍可独立完成。
|
||||
file_path, summary = await asyncio.gather(
|
||||
self._aoffload_to_backend(backend, messages_to_summarize),
|
||||
self._acreate_summary(messages_to_summarize),
|
||||
)
|
||||
if file_path is None:
|
||||
msg = (
|
||||
"Offloading conversation history to backend failed during summarization. "
|
||||
"Older messages will not be recoverable."
|
||||
)
|
||||
logger.error(msg)
|
||||
warnings.warn(msg, stacklevel=2)
|
||||
|
||||
new_messages = self._build_new_messages_with_path(summary, file_path)
|
||||
previous_event = request.state.get("_summarization_event")
|
||||
state_cutoff_index = self._compute_state_cutoff(previous_event, cutoff_index)
|
||||
new_event = {
|
||||
"cutoff_index": state_cutoff_index,
|
||||
"summary_message": new_messages[0],
|
||||
"file_path": file_path,
|
||||
}
|
||||
|
||||
response = await handler(request.override(messages=[*new_messages, *preserved_messages]))
|
||||
update: dict[str, Any] = {"_summarization_event": new_event}
|
||||
if new_state_tail:
|
||||
update["messages"] = list(new_state_tail)
|
||||
return ExtendedModelResponse(model_response=response, command=Command(update=update))
|
||||
|
||||
|
||||
def create_summary_middleware(
|
||||
model: str | BaseChatModel,
|
||||
*,
|
||||
trigger: ContextSize | list[ContextSize] | None,
|
||||
keep: ContextSize | list[ContextSize] | None,
|
||||
summary_prompt: str | None = None,
|
||||
trim_tokens_to_summarize: int | None = None,
|
||||
tool_result_offload_token_limit: int | None = _DEFAULT_SUMMARY_TOOL_RESULT_LIMIT_TOKENS,
|
||||
l1_l2_trigger_ratio: float = _DEFAULT_L1_L2_TRIGGER_RATIO,
|
||||
) -> SummarizationMiddleware:
|
||||
"""Create DeepAgents summarization middleware using Yuxi's virtual outputs root."""
|
||||
middleware_kwargs = {
|
||||
"model": model,
|
||||
"backend": create_agent_composite_backend,
|
||||
"trigger": trigger,
|
||||
"keep": keep,
|
||||
"token_counter": _count_tokens_for_summary_trigger,
|
||||
"trim_tokens_to_summarize": trim_tokens_to_summarize,
|
||||
"tool_result_offload_token_limit": tool_result_offload_token_limit,
|
||||
"l1_l2_trigger_ratio": l1_l2_trigger_ratio,
|
||||
}
|
||||
if summary_prompt and summary_prompt.strip():
|
||||
middleware_kwargs["summary_prompt"] = summary_prompt
|
||||
middleware = YuxiSummarizationMiddleware(**middleware_kwargs)
|
||||
middleware._history_path_prefix = VIRTUAL_PATH_CONVERSATION_HISTORY
|
||||
middleware._large_tool_results_prefix = VIRTUAL_PATH_LARGE_TOOL_RESULTS
|
||||
return middleware
|
||||
@@ -0,0 +1,191 @@
|
||||
"""Token usage observation middleware for Yuxi agents."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Awaitable, Callable, Iterable, Mapping
|
||||
from datetime import UTC, datetime
|
||||
from typing import Any, NotRequired, TypedDict
|
||||
|
||||
from langchain.agents.middleware.types import (
|
||||
AgentMiddleware,
|
||||
AgentState,
|
||||
ExtendedModelResponse,
|
||||
ModelRequest,
|
||||
ModelResponse,
|
||||
)
|
||||
from langchain_core.messages import AIMessage, AnyMessage
|
||||
from langchain_core.messages.utils import count_tokens_approximately
|
||||
from langgraph.types import Command
|
||||
|
||||
|
||||
class TokenUsagePayload(TypedDict, total=False):
|
||||
"""Serializable token usage snapshot stored in LangGraph state."""
|
||||
|
||||
state_message_count: int
|
||||
state_message_count_before_call: int
|
||||
state_messages_tokens: int
|
||||
state_messages_tokens_before_call: int
|
||||
llm_message_count: int
|
||||
llm_messages_tokens: int
|
||||
llm_content_message_count: int
|
||||
llm_content_message_tokens: int
|
||||
llm_tool_message_count: int
|
||||
llm_tool_message_tokens: int
|
||||
llm_input_tokens: int
|
||||
system_tokens: int
|
||||
tools_tokens: int
|
||||
tool_count: int
|
||||
context_window: int | None
|
||||
context_usage_ratio: float | None
|
||||
remaining_context_tokens: int | None
|
||||
summary_active: bool
|
||||
summary_message_tokens: int
|
||||
summary_trigger_tokens: int | None
|
||||
model_usage: dict[str, int]
|
||||
counter: str
|
||||
estimate: bool
|
||||
measured_at: str
|
||||
|
||||
|
||||
class TokenUsageState(AgentState):
|
||||
"""Agent state extension with the latest token usage snapshot."""
|
||||
|
||||
token_usage: NotRequired[TokenUsagePayload]
|
||||
|
||||
|
||||
def _safe_int(value: Any) -> int | None:
|
||||
if isinstance(value, bool):
|
||||
return None
|
||||
if isinstance(value, int):
|
||||
return value
|
||||
if isinstance(value, float) and value.is_integer():
|
||||
return int(value)
|
||||
return None
|
||||
|
||||
|
||||
def _model_context_window(model: Any) -> int | None:
|
||||
profile = getattr(model, "profile", None)
|
||||
if not isinstance(profile, Mapping):
|
||||
return None
|
||||
max_input_tokens = profile.get("max_input_tokens")
|
||||
return max_input_tokens if isinstance(max_input_tokens, int) and max_input_tokens > 0 else None
|
||||
|
||||
|
||||
def _summary_trigger_tokens(runtime_context: Any) -> int | None:
|
||||
threshold = _safe_int(getattr(runtime_context, "summary_threshold", None))
|
||||
if threshold is None or threshold <= 0:
|
||||
return None
|
||||
return threshold * 1024
|
||||
|
||||
|
||||
def _is_summary_message(message: AnyMessage) -> bool:
|
||||
return getattr(message, "additional_kwargs", {}).get("lc_source") == "summarization"
|
||||
|
||||
|
||||
def _is_tool_message(message: AnyMessage) -> bool:
|
||||
return getattr(message, "type", None) == "tool" or getattr(message, "role", None) == "tool"
|
||||
|
||||
|
||||
def _model_usage_from_response(response: ModelResponse) -> dict[str, int]:
|
||||
for message in reversed(response.result):
|
||||
if not isinstance(message, AIMessage):
|
||||
continue
|
||||
usage = getattr(message, "usage_metadata", None)
|
||||
if not isinstance(usage, Mapping):
|
||||
continue
|
||||
return {str(key): value for key, value in usage.items() if isinstance(value, int)}
|
||||
return {}
|
||||
|
||||
|
||||
class TokenUsageMiddleware(AgentMiddleware[TokenUsageState]):
|
||||
"""Record approximate context token usage for the current model request."""
|
||||
|
||||
state_schema = TokenUsageState
|
||||
|
||||
def __init__(self, token_counter=count_tokens_approximately) -> None:
|
||||
super().__init__()
|
||||
self.token_counter = token_counter
|
||||
|
||||
def _count_tokens(self, messages: Iterable[Any], *, tools: list[Any] | None = None) -> int:
|
||||
message_list = list(messages)
|
||||
if tools is not None:
|
||||
return int(self.token_counter(message_list, tools=tools))
|
||||
return int(self.token_counter(message_list))
|
||||
|
||||
def _build_snapshot(self, request: ModelRequest, response: ModelResponse) -> TokenUsagePayload:
|
||||
state_messages = list(request.state.get("messages") or [])
|
||||
llm_messages = list(request.messages or [])
|
||||
system_messages = [request.system_message] if request.system_message is not None else []
|
||||
tools = list(request.tools or [])
|
||||
response_messages = list(response.result or [])
|
||||
|
||||
state_tokens_before_call = self._count_tokens(state_messages)
|
||||
next_state_messages = [*state_messages, *response_messages]
|
||||
state_messages_tokens = self._count_tokens(next_state_messages)
|
||||
llm_messages_tokens = self._count_tokens(llm_messages)
|
||||
system_tokens = self._count_tokens(system_messages)
|
||||
tools_tokens = self._count_tokens([], tools=tools) if tools else 0
|
||||
llm_input_tokens = self._count_tokens([*system_messages, *llm_messages], tools=tools)
|
||||
|
||||
context_window = _model_context_window(request.model)
|
||||
context_usage_ratio = None
|
||||
remaining_context_tokens = None
|
||||
if context_window:
|
||||
context_usage_ratio = min(1.0, round(llm_input_tokens / context_window, 4))
|
||||
remaining_context_tokens = max(context_window - llm_input_tokens, 0)
|
||||
|
||||
summary_message = llm_messages[0] if llm_messages and _is_summary_message(llm_messages[0]) else None
|
||||
llm_tool_messages = [message for message in llm_messages if _is_tool_message(message)]
|
||||
llm_content_messages = [
|
||||
message for message in llm_messages if not _is_tool_message(message) and not _is_summary_message(message)
|
||||
]
|
||||
summary_trigger_tokens = _summary_trigger_tokens(getattr(request.runtime, "context", None))
|
||||
|
||||
return {
|
||||
"state_message_count": len(next_state_messages),
|
||||
"state_message_count_before_call": len(state_messages),
|
||||
"state_messages_tokens": state_messages_tokens,
|
||||
"state_messages_tokens_before_call": state_tokens_before_call,
|
||||
"llm_message_count": len(llm_messages),
|
||||
"llm_messages_tokens": llm_messages_tokens,
|
||||
"llm_content_message_count": len(llm_content_messages),
|
||||
"llm_content_message_tokens": self._count_tokens(llm_content_messages),
|
||||
"llm_tool_message_count": len(llm_tool_messages),
|
||||
"llm_tool_message_tokens": self._count_tokens(llm_tool_messages),
|
||||
"llm_input_tokens": llm_input_tokens,
|
||||
"system_tokens": system_tokens,
|
||||
"tools_tokens": tools_tokens,
|
||||
"tool_count": len(tools),
|
||||
"context_window": context_window,
|
||||
"context_usage_ratio": context_usage_ratio,
|
||||
"remaining_context_tokens": remaining_context_tokens,
|
||||
"summary_active": summary_message is not None,
|
||||
"summary_message_tokens": self._count_tokens([summary_message]) if summary_message else 0,
|
||||
"summary_trigger_tokens": summary_trigger_tokens,
|
||||
"model_usage": _model_usage_from_response(response),
|
||||
"counter": "langchain.count_tokens_approximately",
|
||||
"estimate": True,
|
||||
"measured_at": datetime.now(UTC).isoformat(),
|
||||
}
|
||||
|
||||
def wrap_model_call(
|
||||
self,
|
||||
request: ModelRequest,
|
||||
handler: Callable[[ModelRequest], ModelResponse],
|
||||
) -> ExtendedModelResponse:
|
||||
response = handler(request)
|
||||
return ExtendedModelResponse(
|
||||
model_response=response,
|
||||
command=Command(update={"token_usage": self._build_snapshot(request, response)}),
|
||||
)
|
||||
|
||||
async def awrap_model_call(
|
||||
self,
|
||||
request: ModelRequest,
|
||||
handler: Callable[[ModelRequest], Awaitable[ModelResponse]],
|
||||
) -> ExtendedModelResponse:
|
||||
response = await handler(request)
|
||||
return ExtendedModelResponse(
|
||||
model_response=response,
|
||||
command=Command(update={"token_usage": self._build_snapshot(request, response)}),
|
||||
)
|
||||
@@ -0,0 +1,193 @@
|
||||
from typing import Any
|
||||
|
||||
from langchain.chat_models import BaseChatModel
|
||||
from langchain_openai import ChatOpenAI
|
||||
from pydantic import SecretStr
|
||||
|
||||
from yuxi import config as sys_config
|
||||
from yuxi.models.providers.cache import model_cache
|
||||
from yuxi.utils import get_docker_safe_url
|
||||
from yuxi.utils.logging_config import logger
|
||||
|
||||
_TOOL_IMAGE_USER_TEXT = "Images returned by read_file are attached below. Inspect them when answering."
|
||||
|
||||
|
||||
def resolve_chat_model_spec(model_spec: str | None, *, fallback: str | None = None) -> str:
|
||||
"""解析空模型配置,不吞掉已经配置但无效的模型值。
|
||||
|
||||
这里仅处理模型为空时的优先级:请求或配置值、调用方 fallback、系统默认模型;
|
||||
具体模型是否存在、是否为聊天模型仍由 model_cache 校验。
|
||||
"""
|
||||
for candidate in (model_spec, fallback, sys_config.default_model):
|
||||
if isinstance(candidate, str) and candidate.strip():
|
||||
return candidate.strip()
|
||||
raise ValueError("model spec 不能为空")
|
||||
|
||||
|
||||
def load_chat_model(fully_specified_name: str | None, **kwargs) -> BaseChatModel:
|
||||
fully_specified_name = resolve_chat_model_spec(fully_specified_name)
|
||||
|
||||
info = model_cache.get_model_info(fully_specified_name)
|
||||
if not info:
|
||||
available_specs = model_cache.get_all_specs("chat")
|
||||
available_ids = [item.spec for item in available_specs[:10]]
|
||||
raise ValueError(
|
||||
f"Unknown model spec: '{fully_specified_name}'. "
|
||||
f"Available chat models ({len(available_specs)}): {available_ids}"
|
||||
)
|
||||
|
||||
if info.model_type != "chat":
|
||||
raise ValueError(f"Model {fully_specified_name} is not a chat model (type={info.model_type})")
|
||||
|
||||
api_key = info.api_key
|
||||
base_url = get_docker_safe_url(info.base_url)
|
||||
|
||||
logger.debug(f"Loading model {fully_specified_name} with provider_type={info.provider_type}")
|
||||
|
||||
if info.provider_type == "anthropic":
|
||||
from langchain_anthropic import ChatAnthropic
|
||||
|
||||
return ChatAnthropic(
|
||||
model=info.model_id,
|
||||
api_key=SecretStr(api_key),
|
||||
base_url=base_url,
|
||||
**kwargs,
|
||||
)
|
||||
if info.provider_type == "gemini":
|
||||
from langchain_google_genai import ChatGoogleGenerativeAI
|
||||
|
||||
return ChatGoogleGenerativeAI(
|
||||
model=info.model_id,
|
||||
google_api_key=SecretStr(api_key),
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
return _ToolCallChunkFixChatOpenAI(
|
||||
model=info.model_id,
|
||||
api_key=SecretStr(api_key),
|
||||
base_url=base_url,
|
||||
stream_usage=True,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
|
||||
class _ToolCallChunkFixChatOpenAI(ChatOpenAI):
|
||||
"""归一化流式 tool_call 续片中的空串 name/id,规避 v3 流式累积缺陷。"""
|
||||
|
||||
def _get_request_payload(self, input_, *, stop=None, **kwargs):
|
||||
"""Override to bridge tool image blocks to user messages."""
|
||||
payload = super()._get_request_payload(input_, stop=stop, **kwargs)
|
||||
return _bridge_tool_images_to_user_messages(payload)
|
||||
|
||||
async def _astream(self, *args, **kwargs):
|
||||
async for chunk in super()._astream(*args, **kwargs):
|
||||
_normalize_tool_call_chunks(chunk.message)
|
||||
yield chunk
|
||||
|
||||
def _stream(self, *args, **kwargs):
|
||||
for chunk in super()._stream(*args, **kwargs):
|
||||
_normalize_tool_call_chunks(chunk.message)
|
||||
yield chunk
|
||||
|
||||
|
||||
def _bridge_tool_images_to_user_messages(payload: dict[str, Any]) -> dict[str, Any]:
|
||||
"""将工具调用返回的 image_url 块桥接到用户消息中,避免工具消息中包含图片导致的渲染问题。"""
|
||||
messages = payload.get("messages")
|
||||
if not isinstance(messages, list):
|
||||
return payload
|
||||
if not any(isinstance(m, dict) and m.get("role") == "tool" and _tool_image_blocks(m) for m in messages):
|
||||
return payload
|
||||
|
||||
bridged_messages: list[dict[str, Any]] = []
|
||||
pending_images: list[dict[str, Any]] = []
|
||||
|
||||
def flush_pending_images() -> None:
|
||||
nonlocal pending_images
|
||||
if not pending_images:
|
||||
return
|
||||
|
||||
bridged_messages.append(
|
||||
{
|
||||
"role": "user",
|
||||
"content": [{"type": "text", "text": _TOOL_IMAGE_USER_TEXT}, *pending_images],
|
||||
}
|
||||
)
|
||||
pending_images = []
|
||||
|
||||
for message in messages:
|
||||
if not isinstance(message, dict):
|
||||
flush_pending_images()
|
||||
bridged_messages.append(message)
|
||||
continue
|
||||
|
||||
role = message.get("role")
|
||||
if role != "tool":
|
||||
flush_pending_images()
|
||||
|
||||
image_blocks = _tool_image_blocks(message) if role == "tool" else []
|
||||
if image_blocks:
|
||||
pending_images.extend(image_blocks)
|
||||
|
||||
content = _text_without_images(message.get("content"), image_blocks)
|
||||
if not content:
|
||||
content = (
|
||||
f"read_file returned {len(image_blocks)} image(s). "
|
||||
"The image content is attached in the following user message for visual inspection."
|
||||
)
|
||||
message = {**message, "content": content}
|
||||
|
||||
bridged_messages.append(message)
|
||||
|
||||
flush_pending_images()
|
||||
|
||||
return {**payload, "messages": bridged_messages}
|
||||
|
||||
|
||||
def _normalize_tool_call_chunks(message) -> None:
|
||||
"""把工具调用续片里空字符串的 name/id 归一化为 None。
|
||||
|
||||
LangGraph v3 流式累积对 tool_call 字段是“后值覆盖”:部分 OpenAI 兼容提供商
|
||||
(siliconflow、阿里云百炼等)在续片里把 name/id 下发为空字符串 "",会覆盖首片
|
||||
的真实值(siliconflow 丢 name、百炼丢 id),导致工具结果无法按 tool_call_id
|
||||
关联、工具状态停留在“进行中”。OpenAI 官方在续片里发 None 不会触发覆盖,这里
|
||||
把空串归一化为 None 对齐该行为。待上游修复 v3 协议后可移除。
|
||||
"""
|
||||
for chunk in message.tool_call_chunks:
|
||||
if chunk.get("name") == "":
|
||||
chunk["name"] = None
|
||||
if chunk.get("id") == "":
|
||||
chunk["id"] = None
|
||||
|
||||
|
||||
def _tool_image_blocks(message: dict[str, Any]) -> list[dict[str, Any]]:
|
||||
content = message.get("content")
|
||||
if not isinstance(content, list):
|
||||
return []
|
||||
return [
|
||||
block
|
||||
for block in content
|
||||
if isinstance(block, dict)
|
||||
and block.get("type") == "image_url"
|
||||
and isinstance(block.get("image_url"), dict)
|
||||
and isinstance(block["image_url"].get("url"), str)
|
||||
]
|
||||
|
||||
|
||||
def _text_without_images(content: Any, image_blocks: list[dict[str, Any]]) -> str:
|
||||
if isinstance(content, str):
|
||||
return content
|
||||
if not isinstance(content, list):
|
||||
return ""
|
||||
|
||||
image_ids = {id(block) for block in image_blocks}
|
||||
parts: list[str] = []
|
||||
for block in content:
|
||||
if id(block) in image_ids:
|
||||
continue
|
||||
if isinstance(block, str):
|
||||
parts.append(block)
|
||||
elif isinstance(block, dict) and block.get("type") in {"text", "input_text"}:
|
||||
text = block.get("text")
|
||||
if isinstance(text, str):
|
||||
parts.append(text)
|
||||
return "\n".join(part for part in parts if part)
|
||||
@@ -0,0 +1,56 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class BuiltinSkillSpec:
|
||||
slug: str
|
||||
source_dir: Path
|
||||
description: str = ""
|
||||
version: str = "1.0.0"
|
||||
tool_dependencies: tuple[str, ...] = ()
|
||||
mcp_dependencies: tuple[str, ...] = ()
|
||||
skill_dependencies: tuple[str, ...] = ()
|
||||
|
||||
|
||||
_SKILLS_ROOT = Path(__file__).resolve().parent
|
||||
|
||||
BUILTIN_SKILLS: list[BuiltinSkillSpec] = [
|
||||
BuiltinSkillSpec(
|
||||
slug="image-gen",
|
||||
source_dir=_SKILLS_ROOT / "image-gen",
|
||||
description="在 Agent 沙盒中生成图片并保存到 outputs,默认支持 Qwen-Image,也可接入其它图片生成接口。",
|
||||
version="2026.06.02",
|
||||
tool_dependencies=("present_artifacts",),
|
||||
),
|
||||
BuiltinSkillSpec(
|
||||
slug="deep-research",
|
||||
source_dir=_SKILLS_ROOT / "deep-research",
|
||||
description="深度研究编排方法论:澄清范围、拆解规划、并行调度子智能体调研、对抗式核验、综合成带引用的结构化报告。",
|
||||
version="2026.06.05",
|
||||
tool_dependencies=("tavily_search",),
|
||||
),
|
||||
BuiltinSkillSpec(
|
||||
slug="knowledge-base",
|
||||
source_dir=_SKILLS_ROOT / "knowledge-base",
|
||||
description="使用 Yuxi 知识库进行检索、打开文档、文档内定位和查看思维导图。",
|
||||
version="2026.06.24",
|
||||
tool_dependencies=(
|
||||
"list_kbs",
|
||||
"query_kb",
|
||||
"find_kb_document",
|
||||
"open_kb_document",
|
||||
"get_mindmap",
|
||||
"search_file",
|
||||
),
|
||||
),
|
||||
BuiltinSkillSpec(
|
||||
slug="mysql-reporter",
|
||||
source_dir=_SKILLS_ROOT / "mysql-reporter",
|
||||
description="基于 MySQL 数据库生成查询报表和可视化图表,适合分析业务指标、统计趋势,并用 Charts MCP 展示结果。",
|
||||
version="2026.06.05",
|
||||
mcp_dependencies=("mcp-server-chart",),
|
||||
),
|
||||
]
|
||||
@@ -0,0 +1,50 @@
|
||||
---
|
||||
name: deep-research
|
||||
description: "深度研究编排方法论:澄清范围、拆解规划、并行调度子智能体调研、对抗式核验、综合成带引用的结构化报告。当任务需要多来源、可追溯、需事实核查的深度研究时使用此技能。"
|
||||
---
|
||||
|
||||
# 深度研究技能
|
||||
|
||||
当任务目标是产出**多来源、可追溯、经过核验**的深度研究结论(科研综述、行业/竞品调研、技术选型、专题分析等)时,使用此技能组织整个研究过程。本技能的核心是**编排**:你负责整体把控与子智能体调度,把繁重的检索与核验工作派发出去,自己专注规划与综合。
|
||||
|
||||
## 可用子智能体
|
||||
|
||||
通过 `task` 工具调度(可并行多开,互不依赖的子任务同时派发):
|
||||
|
||||
- `research-explorer`(调研探索员):围绕一个明确子问题做多轮网页/知识库检索,返回按要点组织、带 `<cite>` 引用的结构化发现。**这是主力,按子问题并行多开。**
|
||||
- `fact-verifier`(事实核查员):对给定的关键论断做对抗式核验,逐条给出 支持 / 存疑 / 反驳 + 依据来源 + 置信度,并标注冲突。
|
||||
|
||||
## 编排流程
|
||||
|
||||
### 1. 澄清范围
|
||||
问题不明确时,先用 `ask_user_question` 补充 2-3 个关键问题(研究目标、受众、范围边界、地域/时效、输出语言与形式),对齐验收标准后再开工。已经清晰的任务不要反复追问。
|
||||
|
||||
### 2. 规划拆解
|
||||
用 `write_todos` 把研究目标拆成**可独立调研**的子问题,每个子问题写明产出标准(要回答什么、需要哪类证据)。子问题应正交、覆盖完整,避免重叠或遗漏关键角度。
|
||||
|
||||
### 3. 并行派发调研
|
||||
- 把互不依赖的子问题用**多个 `task` 调用并行**派发给 `research-explorer`。
|
||||
- 每次派发在 `description` 中写清:子问题目标、已知上下文、期望输出格式(要点 + `<cite source="$URL" type="url">$INDEX</cite>` 引用 + 参考来源列表)。
|
||||
- 何时派发 vs 自己直检:子问题复杂、需多轮检索、可隔离上下文、可并行时一律派发子智能体;仅在澄清范围、补一两个零散事实、或快速校正方向时才自己少量直接检索。
|
||||
- 子问题之间有依赖时,先派发前置子问题,拿到结果后再派发后续。
|
||||
|
||||
### 4. 核验关键结论
|
||||
对**影响最终结论的关键论断**、数字、以及子智能体之间相互冲突的发现,派发 `fact-verifier` 做对抗式核验。要求其默认倾向「证据不足即标注存疑」。核验未通过的结论不要写进正文,或必须明确降级标注。
|
||||
|
||||
### 5. 综合成稿
|
||||
证据充分后,由你统一综合为结构化报告,**不要**简单拼接子智能体返回的原文。组织顺序:问题定义 → 证据整理 → 分析比较 → 结论与建议 → 来源。围绕「论证」而非「资料堆砌」,每个结论都要有证据支撑。
|
||||
|
||||
### 6. 停止准则
|
||||
信息饱和、或确认无法获取更多有效信息即停。明确标注证据缺口与不确定性,不臆断、不编造来源。
|
||||
|
||||
## 引用规范
|
||||
|
||||
- 报告中关键结论、数据、观点必须绑定来源。
|
||||
- 沿用 `<cite source="$URL" type="url">$INDEX</cite>` 标注,$INDEX 从 1 起递增,引用紧跟结论后、不单独成行。
|
||||
- 文末单列「来源」章节,逐条列出标题与 URL;引用用户附件/知识库时标明文件名或路径。
|
||||
|
||||
## 输出约束
|
||||
|
||||
- 最终交付的是一份可直接使用的报告,而不是「我打算怎么研究」。
|
||||
- 不要外泄中间推理过程、原始检索日志,也不要把待办清单原样输出成正文。
|
||||
- 报告语言与用户提问语言一致,使用正式、克制、可复核的书面表达。
|
||||
@@ -0,0 +1,82 @@
|
||||
---
|
||||
name: image-gen
|
||||
description: "在 Agent 沙盒中生成图片并保存到 outputs。当用户要求生成图片、海报、插画、文生图,或指定 Qwen-Image、其它兼容图片生成接口时使用此技能。"
|
||||
---
|
||||
|
||||
# 图片生成技能
|
||||
|
||||
当用户要求生成图片、海报、插画、文生图,或明确提到 Qwen-Image 时,使用此技能组织图片生成流程。
|
||||
|
||||
## 默认生成接口
|
||||
|
||||
默认使用 SiliconFlow 的 Qwen-Image 接口:
|
||||
|
||||
- Endpoint: `POST https://api.siliconflow.cn/v1/images/generations`
|
||||
- Model: `Qwen/Qwen-Image`
|
||||
- 默认参数:
|
||||
- `negative_prompt`: `""`
|
||||
- `num_inference_steps`: `20`
|
||||
- `guidance_scale`: `7.5`
|
||||
|
||||
调用外部接口时,必须在 Agent 沙盒执行环境中读取 `SILICONFLOW_API_KEY`。不要依赖后端进程环境变量。
|
||||
|
||||
## 操作流程
|
||||
|
||||
1. 明确用户要生成的图片内容、风格、尺寸或约束;信息不足但不影响生成时,使用合理默认值,不要反复追问。
|
||||
2. 使用可用的执行工具在沙盒中运行脚本,调用图片生成接口,传入用户需求整理后的 `prompt`,并按需传入 `negative_prompt`、`num_inference_steps`、`guidance_scale`。
|
||||
3. 从生成接口响应中读取图片地址,默认路径为 `images[0].url`。
|
||||
4. 在同一个沙盒脚本中用 `Authorization: Bearer $SILICONFLOW_API_KEY` 下载该图片地址;如果接口直接返回 base64,则直接解码保存。
|
||||
5. 将最终图片保存到 `/home/gem/user-data/outputs/` 下,例如 `/home/gem/user-data/outputs/generated-image.png`。
|
||||
6. 调用 `present_artifacts`,传入保存后的 outputs 虚拟路径,让前端展示图片产物。
|
||||
7. 最终回复简要说明图片已生成,不要把外部临时 URL 当作最终结果展示。
|
||||
|
||||
## 脚本示例
|
||||
|
||||
可根据用户需求调整 prompt 和输出文件名:
|
||||
|
||||
```python
|
||||
import os
|
||||
import requests
|
||||
from pathlib import Path
|
||||
|
||||
api_key = os.environ["SILICONFLOW_API_KEY"]
|
||||
prompt = "根据用户需求整理后的图片提示词"
|
||||
|
||||
response = requests.post(
|
||||
"https://api.siliconflow.cn/v1/images/generations",
|
||||
headers={"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"},
|
||||
json={
|
||||
"model": "Qwen/Qwen-Image",
|
||||
"prompt": prompt,
|
||||
"negative_prompt": "",
|
||||
"num_inference_steps": 20,
|
||||
"guidance_scale": 7.5,
|
||||
},
|
||||
timeout=120,
|
||||
)
|
||||
response.raise_for_status()
|
||||
image_url = response.json()["images"][0]["url"]
|
||||
|
||||
image_response = requests.get(
|
||||
image_url,
|
||||
headers={"Authorization": f"Bearer {api_key}"},
|
||||
timeout=120,
|
||||
)
|
||||
image_response.raise_for_status()
|
||||
|
||||
output_path = Path("/home/gem/user-data/outputs/generated-image.png")
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
output_path.write_bytes(image_response.content)
|
||||
print(output_path.as_posix())
|
||||
```
|
||||
|
||||
## 多模型扩展
|
||||
|
||||
如果用户指定其它图片生成模型或兼容接口,可以按该接口的协议先生成图片。只要最终拿到图片 bytes 或 base64,就保存到 `/home/gem/user-data/outputs/`,再调用 `present_artifacts` 展示。
|
||||
|
||||
## 关键约束
|
||||
|
||||
- 不要把外部生成接口返回的临时 URL 当作最终结果直接展示给用户。
|
||||
- 不要调用后端 MinIO 上传工具;图片生成和下载都应在沙盒内完成。
|
||||
- 如果 `SILICONFLOW_API_KEY` 缺失,应明确提示用户需要在 Agent 沙盒环境变量中配置。
|
||||
- 保存到 outputs 后必须调用 `present_artifacts`,否则前端不会自动展示生成图片。
|
||||
@@ -0,0 +1,33 @@
|
||||
---
|
||||
name: knowledge-base
|
||||
slug: knowledge-base
|
||||
description: "使用 Yuxi 知识库进行检索、打开文档、文档内定位和查看思维导图。当用户需要基于已配置知识库回答问题、核验资料或引用文档内容时使用此技能。"
|
||||
---
|
||||
|
||||
# 知识库技能
|
||||
|
||||
当用户要求基于项目知识库、内部资料、上传入库文档或知识图谱相关内容回答问题时,使用此技能。
|
||||
|
||||
## 可用工具
|
||||
|
||||
- `list_kbs`:列出当前会话可访问且已启用的知识库。
|
||||
- `query_kb`:按 `kb_id` 在指定知识库中检索内容,返回 `file_id` 和相关片段。
|
||||
- `open_kb_document`:按 `kb_id` 和 `file_id` 打开文档原文窗口,适合查看更完整上下文。
|
||||
- `find_kb_document`:在已知文档内用关键词或正则定位段落。
|
||||
- `get_mindmap`:查看知识库思维导图结构。
|
||||
- `search_file`:按文件名关键词搜索知识库中的文件,支持指定知识库或跨知识库,返回文件列表与分页信息。
|
||||
|
||||
## 操作流程
|
||||
|
||||
1. 需要先确认当前会话有哪些知识库可用;不确定时调用 `list_kbs`。
|
||||
2. 针对用户问题选择最相关的知识库,使用 `query_kb` 检索。
|
||||
3. 如果检索片段不足以回答,使用返回的 `file_id` 调用 `open_kb_document` 查看上下文。
|
||||
4. 如果用户要求定位术语、指标、章节或原文证据,使用 `find_kb_document` 在候选文档内查找。
|
||||
5. 当用户关心知识库结构、文件分类或知识框架时,使用 `get_mindmap`。
|
||||
|
||||
## 关键约束
|
||||
|
||||
- 只能访问当前会话配置和用户权限允许的知识库。
|
||||
- 不要编造 `kb_id` 或 `file_id`;优先从 `list_kbs` 和 `query_kb` 的返回结果中获取。
|
||||
- 回答需要可追溯时,应说明依据来自哪个知识库、文件或检索片段。
|
||||
- Dify 等外部只读知识库可能只支持检索,不一定支持打开全文或文档内查找;遇到工具返回限制说明时,应如实告知用户。
|
||||
@@ -0,0 +1,49 @@
|
||||
---
|
||||
name: mysql reporter
|
||||
slug: mysql-reporter
|
||||
description: "生成 MySQL 查询报表并生成可视化图表。当用户需要查询 MySQL 数据库并以报表形式展示结果时使用此技能,包括:统计销售数据、分析用户行为、生成业务报表、查询业务指标等。"
|
||||
---
|
||||
|
||||
# MySQL 报表技能
|
||||
|
||||
根据用户的指令,通过终端脚本访问 MySQL 数据库,并结合图表绘制工具构建 SQL 查询报告。
|
||||
|
||||
## 操作流程
|
||||
|
||||
1. 理解用户的指令,明确报表的需求和目标
|
||||
2. 通过 terminal 进入技能目录:`cd /home/gem/skills/mysql-reporter`
|
||||
3. 使用 `uv run scripts/list_tables.py` 查看可用表;如果脚本提示缺少 MySQL 配置,按“环境变量缺失处理”回复用户
|
||||
4. 必要时用 `uv run scripts/describe_table.py --table 表名` 查看表结构
|
||||
5. 生成正确且高效的只读 SQL,通过 `uv run scripts/query.py --sql "SQL语句" --timeout 60` 执行查询并获取结果
|
||||
6. 使用 Charts MCP 生成图表
|
||||
7. 将图表以 markdown 图片格式嵌入报表
|
||||
|
||||
## 环境变量缺失处理
|
||||
|
||||
脚本只读取 Agent 沙盒中的环境变量,不读取后端 `.env` 或 Docker Compose 变量。必填变量包括:
|
||||
|
||||
- `MYSQL_HOST`
|
||||
- `MYSQL_USER`
|
||||
- `MYSQL_PASSWORD`
|
||||
- `MYSQL_DATABASE`
|
||||
|
||||
可选变量包括:
|
||||
|
||||
- `MYSQL_PORT`:默认 `3306`
|
||||
- `MYSQL_DATABASE_DESCRIPTION`:数据库业务说明,用于辅助理解表和指标含义
|
||||
|
||||
如果执行脚本时出现 `MySQL configuration missing required key`,不要继续猜测连接信息或编造报表。应明确告诉用户:需要在个人设置中的「沙盒环境变量」里配置缺失的 `MYSQL_*` 变量;保存后仅对新建沙盒生效,需要重新发起任务或新建会话后再执行。
|
||||
|
||||
## 关键约束
|
||||
|
||||
- 生成的 SQL 查询必须正确且高效,避免全表扫描
|
||||
- MySQL 操作必须通过本技能 `scripts/` 下的 CLI 脚本执行,不要调用平台内置 MySQL tools
|
||||
- 不要在报表或错误说明中输出 `MYSQL_PASSWORD` 等敏感环境变量的值,只能说明缺少哪些变量名
|
||||
- 图表生成工具的返回结果不会默认渲染,必须在最终报表中以 `` 格式嵌入
|
||||
- 只返回报表相关的结论,不要返回原始 SQL 查询语句
|
||||
|
||||
## 允许的工具
|
||||
|
||||
- terminal:执行 `scripts/list_tables.py`、`scripts/describe_table.py`、`scripts/query.py`
|
||||
- Charts MCP:生成可视化图表
|
||||
- 网络检索工具:必要时补充背景信息
|
||||
@@ -0,0 +1,177 @@
|
||||
# /// script
|
||||
# dependencies = [
|
||||
# "pymysql>=1.1.0",
|
||||
# ]
|
||||
# ///
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import re
|
||||
import sys
|
||||
import time
|
||||
from typing import Any
|
||||
|
||||
import pymysql
|
||||
from pymysql import MySQLError
|
||||
from pymysql.cursors import DictCursor
|
||||
|
||||
|
||||
class MySQLConnectionError(Exception):
|
||||
"""MySQL 连接异常"""
|
||||
|
||||
|
||||
class MySQLSecurityChecker:
|
||||
"""MySQL 安全检查器"""
|
||||
|
||||
@classmethod
|
||||
def validate_table_name(cls, table_name: str) -> bool:
|
||||
"""验证表名的安全性"""
|
||||
if not table_name:
|
||||
return False
|
||||
|
||||
return bool(re.match(r"^[a-zA-Z_][a-zA-Z0-9_]*$", table_name))
|
||||
|
||||
|
||||
def load_mysql_config() -> dict[str, Any]:
|
||||
config: dict[str, Any] = {
|
||||
"host": os.getenv("MYSQL_HOST"),
|
||||
"user": os.getenv("MYSQL_USER"),
|
||||
"password": os.getenv("MYSQL_PASSWORD"),
|
||||
"database": os.getenv("MYSQL_DATABASE"),
|
||||
"port": int(os.getenv("MYSQL_PORT") or "3306"),
|
||||
"charset": "utf8mb4",
|
||||
"description": os.getenv("MYSQL_DATABASE_DESCRIPTION") or "默认 MySQL 数据库",
|
||||
}
|
||||
|
||||
required_keys = ["host", "user", "password", "database"]
|
||||
for key in required_keys:
|
||||
if not config[key]:
|
||||
raise MySQLConnectionError(
|
||||
f"MySQL configuration missing required key: {key}, please check your environment variables."
|
||||
)
|
||||
|
||||
return config
|
||||
|
||||
|
||||
def create_connection(config: dict[str, Any]) -> pymysql.Connection:
|
||||
max_retries = 3
|
||||
for attempt in range(max_retries):
|
||||
try:
|
||||
return pymysql.connect(
|
||||
host=config["host"],
|
||||
user=config["user"],
|
||||
password=config["password"],
|
||||
database=config["database"],
|
||||
port=config["port"],
|
||||
charset=config.get("charset", "utf8mb4"),
|
||||
cursorclass=DictCursor,
|
||||
connect_timeout=10,
|
||||
read_timeout=60,
|
||||
write_timeout=30,
|
||||
autocommit=True,
|
||||
)
|
||||
except MySQLError as exc:
|
||||
if attempt < max_retries - 1:
|
||||
time.sleep(2**attempt)
|
||||
continue
|
||||
raise ConnectionError(f"MySQL connection failed: {exc}") from exc
|
||||
|
||||
raise ConnectionError("MySQL connection failed")
|
||||
|
||||
|
||||
def describe_table(table_name: str) -> str:
|
||||
if not MySQLSecurityChecker.validate_table_name(table_name):
|
||||
raise ValueError("表名包含非法字符,请检查表名")
|
||||
|
||||
config = load_mysql_config()
|
||||
connection = create_connection(config)
|
||||
try:
|
||||
with connection.cursor() as cursor:
|
||||
cursor.execute(f"DESCRIBE `{table_name}`")
|
||||
columns = cursor.fetchall()
|
||||
|
||||
if not columns:
|
||||
return f"表 {table_name} 不存在或没有字段"
|
||||
|
||||
column_comments: dict[str, str] = {}
|
||||
try:
|
||||
cursor.execute(
|
||||
"""
|
||||
SELECT COLUMN_NAME, COLUMN_COMMENT
|
||||
FROM information_schema.COLUMNS
|
||||
WHERE TABLE_NAME = %s AND TABLE_SCHEMA = %s
|
||||
""",
|
||||
(table_name, config["database"]),
|
||||
)
|
||||
comment_rows = cursor.fetchall()
|
||||
for row in comment_rows:
|
||||
column_name = row.get("COLUMN_NAME")
|
||||
if column_name:
|
||||
column_comments[column_name] = row.get("COLUMN_COMMENT") or ""
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
result = f"表 `{table_name}` 的结构:\n\n"
|
||||
result += "字段名\t\t类型\t\tNULL\t键\t默认值\t\t额外\t备注\n"
|
||||
result += "-" * 80 + "\n"
|
||||
|
||||
for col in columns:
|
||||
field = col["Field"] or ""
|
||||
type_str = col["Type"] or ""
|
||||
null_str = col["Null"] or ""
|
||||
key_str = col["Key"] or ""
|
||||
default_str = col.get("Default") or ""
|
||||
extra_str = col.get("Extra") or ""
|
||||
comment_str = column_comments.get(field, "")
|
||||
|
||||
result += (
|
||||
f"{field:<16}\t{type_str:<16}\t{null_str:<8}\t{key_str:<4}\t"
|
||||
f"{default_str:<16}\t{extra_str:<16}\t{comment_str}\n"
|
||||
)
|
||||
|
||||
try:
|
||||
cursor.execute(f"SHOW INDEX FROM `{table_name}`")
|
||||
indexes = cursor.fetchall()
|
||||
|
||||
if indexes:
|
||||
result += "\n索引信息:\n"
|
||||
index_dict: dict[str, list[str]] = {}
|
||||
for idx in indexes:
|
||||
key_name = idx["Key_name"]
|
||||
if key_name not in index_dict:
|
||||
index_dict[key_name] = []
|
||||
index_dict[key_name].append(idx["Column_name"])
|
||||
|
||||
for key_name, index_columns in index_dict.items():
|
||||
result += f"- {key_name}: {', '.join(index_columns)}\n"
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
return result
|
||||
finally:
|
||||
connection.close()
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description="描述 MySQL 表结构")
|
||||
parser.add_argument("--table", required=True, help="要查询结构的表名")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main() -> int:
|
||||
args = parse_args()
|
||||
try:
|
||||
print(describe_table(args.table))
|
||||
return 0
|
||||
except ValueError as exc:
|
||||
print(str(exc), file=sys.stderr)
|
||||
return 1
|
||||
except Exception as exc:
|
||||
print(f"获取表 {args.table} 结构失败: {exc}", file=sys.stderr)
|
||||
return 1
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -0,0 +1,112 @@
|
||||
# /// script
|
||||
# dependencies = [
|
||||
# "pymysql>=1.1.0",
|
||||
# ]
|
||||
# ///
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
from typing import Any
|
||||
|
||||
import pymysql
|
||||
from pymysql import MySQLError
|
||||
from pymysql.cursors import DictCursor
|
||||
|
||||
|
||||
class MySQLConnectionError(Exception):
|
||||
"""MySQL 连接异常"""
|
||||
|
||||
|
||||
def load_mysql_config() -> dict[str, Any]:
|
||||
config: dict[str, Any] = {
|
||||
"host": os.getenv("MYSQL_HOST"),
|
||||
"user": os.getenv("MYSQL_USER"),
|
||||
"password": os.getenv("MYSQL_PASSWORD"),
|
||||
"database": os.getenv("MYSQL_DATABASE"),
|
||||
"port": int(os.getenv("MYSQL_PORT") or "3306"),
|
||||
"charset": "utf8mb4",
|
||||
"description": os.getenv("MYSQL_DATABASE_DESCRIPTION") or "默认 MySQL 数据库",
|
||||
}
|
||||
|
||||
required_keys = ["host", "user", "password", "database"]
|
||||
for key in required_keys:
|
||||
if not config[key]:
|
||||
raise MySQLConnectionError(
|
||||
f"MySQL configuration missing required key: {key}, please check your environment variables."
|
||||
)
|
||||
|
||||
return config
|
||||
|
||||
|
||||
def create_connection(config: dict[str, Any]) -> pymysql.Connection:
|
||||
max_retries = 3
|
||||
for attempt in range(max_retries):
|
||||
try:
|
||||
return pymysql.connect(
|
||||
host=config["host"],
|
||||
user=config["user"],
|
||||
password=config["password"],
|
||||
database=config["database"],
|
||||
port=config["port"],
|
||||
charset=config.get("charset", "utf8mb4"),
|
||||
cursorclass=DictCursor,
|
||||
connect_timeout=10,
|
||||
read_timeout=60,
|
||||
write_timeout=30,
|
||||
autocommit=True,
|
||||
)
|
||||
except MySQLError as exc:
|
||||
if attempt < max_retries - 1:
|
||||
time.sleep(2**attempt)
|
||||
continue
|
||||
raise ConnectionError(f"MySQL connection failed: {exc}") from exc
|
||||
|
||||
raise ConnectionError("MySQL connection failed")
|
||||
|
||||
|
||||
def list_tables() -> str:
|
||||
config = load_mysql_config()
|
||||
connection = create_connection(config)
|
||||
try:
|
||||
with connection.cursor() as cursor:
|
||||
cursor.execute("SHOW TABLES")
|
||||
tables = cursor.fetchall()
|
||||
|
||||
if not tables:
|
||||
return "数据库中没有找到任何表"
|
||||
|
||||
table_names = []
|
||||
for table in tables:
|
||||
table_name = list(table.values())[0]
|
||||
table_names.append(table_name)
|
||||
|
||||
all_table_names = "\n".join(table_names)
|
||||
result = f"数据库中的表:\n{all_table_names}"
|
||||
if db_note := config.get("description"):
|
||||
result = f"数据库说明: {db_note}\n\n" + result
|
||||
return result
|
||||
finally:
|
||||
connection.close()
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description="列出当前 MySQL 数据库中的所有表")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main() -> int:
|
||||
parse_args()
|
||||
try:
|
||||
print(list_tables())
|
||||
return 0
|
||||
except Exception as exc:
|
||||
print(f"获取表名失败: {exc}", file=sys.stderr)
|
||||
return 1
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -0,0 +1,299 @@
|
||||
# /// script
|
||||
# dependencies = [
|
||||
# "pymysql>=1.1.0",
|
||||
# ]
|
||||
# ///
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import concurrent.futures
|
||||
import os
|
||||
import re
|
||||
import sys
|
||||
import time
|
||||
from typing import Any
|
||||
|
||||
import pymysql
|
||||
from pymysql import MySQLError
|
||||
from pymysql.cursors import DictCursor
|
||||
|
||||
|
||||
class MySQLConnectionError(Exception):
|
||||
"""MySQL 连接异常"""
|
||||
|
||||
|
||||
class QueryTimeoutError(Exception):
|
||||
"""查询超时异常"""
|
||||
|
||||
|
||||
class MySQLSecurityChecker:
|
||||
"""MySQL 安全检查器"""
|
||||
|
||||
ALLOWED_OPERATIONS = {"SELECT", "SHOW", "DESCRIBE", "EXPLAIN"}
|
||||
|
||||
DANGEROUS_KEYWORDS = {
|
||||
"DROP",
|
||||
"DELETE",
|
||||
"UPDATE",
|
||||
"INSERT",
|
||||
"CREATE",
|
||||
"ALTER",
|
||||
"TRUNCATE",
|
||||
"REPLACE",
|
||||
"LOAD",
|
||||
"GRANT",
|
||||
"REVOKE",
|
||||
"SET",
|
||||
"COMMIT",
|
||||
"ROLLBACK",
|
||||
"UNLOCK",
|
||||
"KILL",
|
||||
"SHUTDOWN",
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def validate_sql(cls, sql: str) -> bool:
|
||||
"""验证SQL语句的安全性"""
|
||||
if not sql:
|
||||
return False
|
||||
|
||||
sql_clean = re.sub(r"--.*$", "", sql, flags=re.MULTILINE)
|
||||
sql_clean = re.sub(r"/\*.*?\*/", "", sql_clean, flags=re.DOTALL)
|
||||
sql_upper = sql_clean.strip().upper()
|
||||
sql_without_trailing_semicolon = sql_upper.rstrip()
|
||||
if sql_without_trailing_semicolon.endswith(";"):
|
||||
sql_without_trailing_semicolon = sql_without_trailing_semicolon[:-1].rstrip()
|
||||
if ";" in sql_without_trailing_semicolon:
|
||||
return False
|
||||
|
||||
if not any(sql_upper.startswith(op) for op in cls.ALLOWED_OPERATIONS):
|
||||
return False
|
||||
|
||||
first_word_match = re.match(r"^\s*(\w+)", sql_upper)
|
||||
first_word = first_word_match.group(1) if first_word_match else ""
|
||||
if first_word in cls.DANGEROUS_KEYWORDS:
|
||||
return False
|
||||
|
||||
sql_injection_patterns = [
|
||||
r"\bor\s+1\s*=\s*1\b",
|
||||
r"\bunion\s+select\b",
|
||||
r"\bexec\s*\(",
|
||||
r"\bxp_cmdshell\b",
|
||||
r"\bsleep\s*\(",
|
||||
r"\bbenchmark\s*\(",
|
||||
r"\bwaitfor\s+delay\b",
|
||||
]
|
||||
sql_injection_patterns.extend(rf"\b;\s*{re.escape(keyword)}\b" for keyword in cls.DANGEROUS_KEYWORDS)
|
||||
|
||||
for pattern in sql_injection_patterns:
|
||||
if re.search(pattern, sql_upper, re.IGNORECASE):
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
@classmethod
|
||||
def validate_timeout(cls, timeout: int) -> bool:
|
||||
"""验证timeout参数"""
|
||||
return isinstance(timeout, int) and 1 <= timeout <= 600
|
||||
|
||||
|
||||
def load_mysql_config() -> dict[str, Any]:
|
||||
config: dict[str, Any] = {
|
||||
"host": os.getenv("MYSQL_HOST"),
|
||||
"user": os.getenv("MYSQL_USER"),
|
||||
"password": os.getenv("MYSQL_PASSWORD"),
|
||||
"database": os.getenv("MYSQL_DATABASE"),
|
||||
"port": int(os.getenv("MYSQL_PORT") or "3306"),
|
||||
"charset": "utf8mb4",
|
||||
"description": os.getenv("MYSQL_DATABASE_DESCRIPTION") or "默认 MySQL 数据库",
|
||||
}
|
||||
|
||||
required_keys = ["host", "user", "password", "database"]
|
||||
for key in required_keys:
|
||||
if not config[key]:
|
||||
raise MySQLConnectionError(
|
||||
f"MySQL configuration missing required key: {key}, please check your environment variables."
|
||||
)
|
||||
|
||||
return config
|
||||
|
||||
|
||||
def create_connection(config: dict[str, Any]) -> pymysql.Connection:
|
||||
max_retries = 3
|
||||
for attempt in range(max_retries):
|
||||
try:
|
||||
return pymysql.connect(
|
||||
host=config["host"],
|
||||
user=config["user"],
|
||||
password=config["password"],
|
||||
database=config["database"],
|
||||
port=config["port"],
|
||||
charset=config.get("charset", "utf8mb4"),
|
||||
cursorclass=DictCursor,
|
||||
connect_timeout=10,
|
||||
read_timeout=60,
|
||||
write_timeout=30,
|
||||
autocommit=True,
|
||||
)
|
||||
except MySQLError as exc:
|
||||
if attempt < max_retries - 1:
|
||||
time.sleep(2**attempt)
|
||||
continue
|
||||
raise ConnectionError(f"MySQL connection failed: {exc}") from exc
|
||||
|
||||
raise ConnectionError("MySQL connection failed")
|
||||
|
||||
|
||||
def execute_query_with_timeout(
|
||||
connection: pymysql.Connection, sql: str, params: tuple | None = None, timeout: int = 10
|
||||
):
|
||||
"""使用线程池实现超时控制,避免信号导致的生成器问题"""
|
||||
|
||||
def query_worker():
|
||||
cursor = connection.cursor(DictCursor)
|
||||
try:
|
||||
if params is None:
|
||||
cursor.execute(sql)
|
||||
else:
|
||||
cursor.execute(sql, params)
|
||||
return cursor.fetchall()
|
||||
finally:
|
||||
cursor.close()
|
||||
|
||||
executor = concurrent.futures.ThreadPoolExecutor(max_workers=1)
|
||||
future = executor.submit(query_worker)
|
||||
try:
|
||||
return future.result(timeout=timeout)
|
||||
except concurrent.futures.TimeoutError as exc:
|
||||
future.cancel()
|
||||
try:
|
||||
connection.close()
|
||||
except Exception:
|
||||
pass
|
||||
raise QueryTimeoutError(f"Query timeout after {timeout} seconds") from exc
|
||||
finally:
|
||||
executor.shutdown(wait=False, cancel_futures=True)
|
||||
|
||||
|
||||
def limit_result_size(result: list, max_chars: int = 10000) -> list:
|
||||
"""限制结果大小"""
|
||||
if not result:
|
||||
return result
|
||||
|
||||
result_str = str(result)
|
||||
if len(result_str) > max_chars:
|
||||
limited_result = []
|
||||
current_chars = 0
|
||||
for row in result:
|
||||
row_str = str(row)
|
||||
if current_chars + len(row_str) > max_chars:
|
||||
break
|
||||
limited_result.append(row)
|
||||
current_chars += len(row_str)
|
||||
return limited_result
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def format_query_result(result: list[dict[str, Any]]) -> str:
|
||||
if not result:
|
||||
return "查询执行成功,但没有返回任何结果"
|
||||
|
||||
limited_result = limit_result_size(result, max_chars=10000)
|
||||
|
||||
if len(limited_result) < len(result):
|
||||
warning = f"\n\n⚠️ 警告: 查询结果过大,只显示了前 {len(limited_result)} 行(共 {len(result)} 行)。\n"
|
||||
warning += "建议使用更精确的查询条件或使用LIMIT子句来减少返回的数据量。"
|
||||
else:
|
||||
warning = ""
|
||||
|
||||
if limited_result:
|
||||
columns = list(limited_result[0].keys())
|
||||
|
||||
col_widths = {}
|
||||
for col in columns:
|
||||
col_widths[col] = max(len(str(col)), max(len(str(row.get(col, ""))) for row in limited_result))
|
||||
col_widths[col] = min(col_widths[col], 50)
|
||||
|
||||
header = "| " + " | ".join(f"{col:<{col_widths[col]}}" for col in columns) + " |"
|
||||
separator = "|" + "|".join("-" * (col_widths[col] + 2) for col in columns) + "|"
|
||||
|
||||
rows = []
|
||||
for row in limited_result:
|
||||
row_str = "| " + " | ".join(f"{str(row.get(col, '')):<{col_widths[col]}}" for col in columns) + " |"
|
||||
rows.append(row_str)
|
||||
|
||||
result_str = f"查询结果(共 {len(limited_result)} 行):\n\n"
|
||||
result_str += header + "\n" + separator + "\n"
|
||||
result_str += "\n".join(rows[:50])
|
||||
|
||||
if len(rows) > 50:
|
||||
result_str += f"\n\n... 还有 {len(rows) - 50} 行未显示 ..."
|
||||
|
||||
result_str += warning
|
||||
return result_str
|
||||
|
||||
return "查询执行成功,但返回数据为空"
|
||||
|
||||
|
||||
def run_query(sql: str, timeout: int) -> str:
|
||||
if not MySQLSecurityChecker.validate_sql(sql):
|
||||
raise ValueError("SQL语句包含不安全的操作或可能的注入攻击,请检查SQL语句")
|
||||
|
||||
if not MySQLSecurityChecker.validate_timeout(timeout):
|
||||
raise ValueError("timeout参数必须在1-600之间")
|
||||
|
||||
config = load_mysql_config()
|
||||
connection = create_connection(config)
|
||||
try:
|
||||
result = execute_query_with_timeout(connection, sql, timeout=timeout or 60)
|
||||
return format_query_result(result)
|
||||
finally:
|
||||
if connection.open:
|
||||
connection.close()
|
||||
|
||||
|
||||
def build_query_error(exc: Exception, sql: str) -> str:
|
||||
error_msg = f"SQL查询执行失败: {exc}\n\n{sql}"
|
||||
|
||||
if "timeout" in str(exc).lower():
|
||||
error_msg += "\n\n💡 建议:查询超时了,请尝试以下方法:\n"
|
||||
error_msg += "1. 减少查询的数据量(使用WHERE条件过滤)\n"
|
||||
error_msg += "2. 使用LIMIT子句限制返回行数\n"
|
||||
error_msg += "3. 增加timeout参数值(最大600秒)"
|
||||
elif "table" in str(exc).lower() and "doesn't exist" in str(exc).lower():
|
||||
error_msg += "\n\n💡 建议:表不存在,请使用 scripts/list_tables.py 查看可用的表名"
|
||||
elif "column" in str(exc).lower() and "doesn't exist" in str(exc).lower():
|
||||
error_msg += "\n\n💡 建议:列不存在,请使用 scripts/describe_table.py 查看表结构"
|
||||
elif "not enough arguments for format string" in str(exc).lower():
|
||||
error_msg += (
|
||||
"\n\n💡 建议:SQL 中的百分号 (%) 被当作参数占位符使用。"
|
||||
" 如需匹配包含百分号的文本,请将百分号写成双百分号 (%%) 或使用参数化查询。"
|
||||
)
|
||||
|
||||
return error_msg
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description="执行只读 MySQL SQL 查询")
|
||||
parser.add_argument("--sql", required=True, help="要执行的SQL查询语句")
|
||||
parser.add_argument("--timeout", type=int, default=60, help="查询超时时间(秒),默认60秒,最大600秒")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main() -> int:
|
||||
args = parse_args()
|
||||
try:
|
||||
print(run_query(args.sql, args.timeout))
|
||||
return 0
|
||||
except ValueError as exc:
|
||||
print(str(exc), file=sys.stderr)
|
||||
return 1
|
||||
except Exception as exc:
|
||||
print(build_query_error(exc, args.sql), file=sys.stderr)
|
||||
return 1
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -0,0 +1,380 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import re
|
||||
import shutil
|
||||
import tempfile
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from yuxi.agents.skills.service import import_skill_dir, is_valid_skill_slug
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from yuxi.storage.postgres.models_business import Skill
|
||||
|
||||
ANSI_ESCAPE_RE = re.compile(r"\x1B\[[0-?]*[ -/]*[@-~]")
|
||||
CONTROL_SEQUENCE_RE = re.compile(r"\x1B\][^\x07]*(?:\x07|\x1B\\)|\x1B[\(\)][A-Za-z0-9]")
|
||||
CLI_TIMEOUT_SECONDS = 300
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class RemoteSkillsBatchPreparation:
|
||||
temp_home: str | None
|
||||
results: list[dict]
|
||||
|
||||
async def cleanup(self) -> None:
|
||||
if self.temp_home:
|
||||
await asyncio.to_thread(shutil.rmtree, self.temp_home, ignore_errors=True)
|
||||
|
||||
|
||||
def _normalize_source(source: str) -> str:
|
||||
value = str(source or "").strip()
|
||||
if not value:
|
||||
raise ValueError("source 不能为空")
|
||||
if any(ch in value for ch in ("\n", "\r", "\x00")):
|
||||
raise ValueError("source 包含非法字符")
|
||||
return value
|
||||
|
||||
|
||||
def _normalize_skill_name(skill: str) -> str:
|
||||
value = str(skill or "").strip()
|
||||
if not is_valid_skill_slug(value):
|
||||
raise ValueError("skill 名称不合法")
|
||||
return value
|
||||
|
||||
|
||||
def _clean_cli_output(output: str) -> list[str]:
|
||||
cleaned = ANSI_ESCAPE_RE.sub("", output or "")
|
||||
cleaned = CONTROL_SEQUENCE_RE.sub("", cleaned)
|
||||
cleaned = cleaned.replace("\r", "\n")
|
||||
normalized_lines: list[str] = []
|
||||
for line in cleaned.splitlines():
|
||||
stripped = line.strip()
|
||||
stripped = re.sub(r"^[│┌└◇◒◐◓◑■●]+\s*", "", stripped)
|
||||
normalized_lines.append(stripped.strip())
|
||||
return normalized_lines
|
||||
|
||||
|
||||
def _parse_available_skills(output: str) -> list[dict[str, str]]:
|
||||
lines = _clean_cli_output(output)
|
||||
items: list[dict[str, str]] = []
|
||||
seen: set[str] = set()
|
||||
collecting = False
|
||||
|
||||
for idx, line in enumerate(lines):
|
||||
if not collecting:
|
||||
if "Available Skills" in line:
|
||||
collecting = True
|
||||
continue
|
||||
|
||||
if not line:
|
||||
continue
|
||||
if "Use --skill " in line:
|
||||
break
|
||||
if not is_valid_skill_slug(line):
|
||||
continue
|
||||
if line in seen:
|
||||
continue
|
||||
|
||||
description = ""
|
||||
next_index = idx + 1
|
||||
while next_index < len(lines):
|
||||
next_line = lines[next_index]
|
||||
next_index += 1
|
||||
if not next_line:
|
||||
continue
|
||||
if "Use --skill " in next_line:
|
||||
break
|
||||
if is_valid_skill_slug(next_line):
|
||||
break
|
||||
if next_line and next_line[0].isalpha():
|
||||
description = next_line
|
||||
else:
|
||||
continue
|
||||
break
|
||||
|
||||
seen.add(line)
|
||||
items.append({"name": line, "description": description})
|
||||
|
||||
return items
|
||||
|
||||
|
||||
async def _run_skills_cli(
|
||||
args: list[str],
|
||||
*,
|
||||
env: dict[str, str],
|
||||
cwd: str,
|
||||
) -> str:
|
||||
process = await asyncio.create_subprocess_exec(
|
||||
*args,
|
||||
cwd=cwd,
|
||||
env=env,
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
stderr=asyncio.subprocess.PIPE,
|
||||
)
|
||||
try:
|
||||
stdout, stderr = await asyncio.wait_for(process.communicate(), timeout=CLI_TIMEOUT_SECONDS)
|
||||
except TimeoutError:
|
||||
process.kill()
|
||||
await process.communicate()
|
||||
raise ValueError("skills CLI 执行超时") from None
|
||||
|
||||
output = (stdout or b"").decode("utf-8", errors="replace")
|
||||
error_output = (stderr or b"").decode("utf-8", errors="replace")
|
||||
combined = "\n".join(part for part in [output.strip(), error_output.strip()] if part)
|
||||
if process.returncode != 0:
|
||||
cleaned_lines = _clean_cli_output(combined)
|
||||
error_msg = "\n".join(line for line in cleaned_lines if line)[:500]
|
||||
raise ValueError(error_msg or "skills CLI 执行失败")
|
||||
return combined
|
||||
|
||||
|
||||
def _create_isolated_workdir() -> tuple[str, dict[str, str], str]:
|
||||
temp_home = tempfile.mkdtemp(prefix=".remote-skills-")
|
||||
env = os.environ.copy()
|
||||
env["HOME"] = temp_home
|
||||
workdir = str(Path(temp_home) / "workspace")
|
||||
Path(workdir).mkdir(parents=True, exist_ok=True)
|
||||
return temp_home, env, workdir
|
||||
|
||||
|
||||
async def list_remote_skills(source: str) -> list[dict[str, str]]:
|
||||
normalized_source = _normalize_source(source)
|
||||
|
||||
temp_home, env, workdir = _create_isolated_workdir()
|
||||
try:
|
||||
output = await _run_skills_cli(
|
||||
["npx", "-y", "skills", "add", normalized_source, "--list"],
|
||||
env=env,
|
||||
cwd=workdir,
|
||||
)
|
||||
finally:
|
||||
await asyncio.to_thread(shutil.rmtree, temp_home, ignore_errors=True)
|
||||
|
||||
skills = _parse_available_skills(output)
|
||||
if not skills:
|
||||
raise ValueError("未发现可安装的 skills")
|
||||
return skills
|
||||
|
||||
|
||||
async def install_remote_skill(
|
||||
db: AsyncSession,
|
||||
*,
|
||||
source: str,
|
||||
skill: str,
|
||||
created_by: str | None,
|
||||
) -> Skill:
|
||||
normalized_source = _normalize_source(source)
|
||||
normalized_skill = _normalize_skill_name(skill)
|
||||
|
||||
temp_home, env, workdir = _create_isolated_workdir()
|
||||
try:
|
||||
available_skills = _parse_available_skills(
|
||||
await _run_skills_cli(
|
||||
["npx", "-y", "skills", "add", normalized_source, "--list"],
|
||||
env=env,
|
||||
cwd=workdir,
|
||||
)
|
||||
)
|
||||
available_names = {item["name"] for item in available_skills}
|
||||
if normalized_skill not in available_names:
|
||||
raise ValueError(f"远程仓库中不存在 skill: {normalized_skill}")
|
||||
|
||||
await _run_skills_cli(
|
||||
[
|
||||
"npx",
|
||||
"-y",
|
||||
"skills",
|
||||
"add",
|
||||
normalized_source,
|
||||
"--skill",
|
||||
normalized_skill,
|
||||
"-g",
|
||||
"-y",
|
||||
"--copy",
|
||||
],
|
||||
env=env,
|
||||
cwd=workdir,
|
||||
)
|
||||
|
||||
skills_dir = Path(temp_home).resolve() / ".agents" / "skills"
|
||||
installed_dir = _find_skill_dir(skills_dir, normalized_skill)
|
||||
if installed_dir is None:
|
||||
raise ValueError("skills CLI 未生成预期的技能目录")
|
||||
|
||||
return await import_skill_dir(
|
||||
db,
|
||||
source_dir=installed_dir,
|
||||
created_by=created_by,
|
||||
)
|
||||
finally:
|
||||
await asyncio.to_thread(shutil.rmtree, temp_home, ignore_errors=True)
|
||||
|
||||
|
||||
async def install_remote_skills_batch(
|
||||
db: AsyncSession,
|
||||
*,
|
||||
source: str,
|
||||
skills: list[str],
|
||||
created_by: str | None,
|
||||
) -> list[dict]:
|
||||
"""批量从同一个远程仓库安装多个 skills(仅一次克隆)。
|
||||
|
||||
Args:
|
||||
db: 数据库会话。
|
||||
source: 远程仓库来源,如 ``owner/repo`` 或 GitHub URL。
|
||||
skills: 需要安装的 skill 名称列表。
|
||||
created_by: 操作者标识。
|
||||
|
||||
Returns:
|
||||
每个 skill 的安装结果列表,顺序与请求一致: ``[{slug, success, error?}, ...]``
|
||||
"""
|
||||
preparation = await prepare_remote_skills_batch(source=source, skills=skills)
|
||||
try:
|
||||
results = preparation.results
|
||||
for index, result in enumerate(results):
|
||||
if not result.get("success"):
|
||||
continue
|
||||
|
||||
source_dir = result.get("source_dir")
|
||||
try:
|
||||
item = await import_skill_dir(
|
||||
db,
|
||||
source_dir=source_dir,
|
||||
created_by=created_by,
|
||||
)
|
||||
results[index] = {"slug": item.slug, "success": True}
|
||||
except Exception as e:
|
||||
if hasattr(db, "rollback"):
|
||||
await db.rollback()
|
||||
results[index] = {"slug": result["slug"], "success": False, "error": str(e)}
|
||||
|
||||
return results
|
||||
finally:
|
||||
await preparation.cleanup()
|
||||
|
||||
|
||||
async def prepare_remote_skills_batch(
|
||||
*,
|
||||
source: str,
|
||||
skills: list[str],
|
||||
) -> RemoteSkillsBatchPreparation:
|
||||
"""批量从远程仓库拉取 skill 目录,但不写数据库。"""
|
||||
normalized_source = _normalize_source(source)
|
||||
if not skills:
|
||||
raise ValueError("skills 列表不能为空")
|
||||
|
||||
# 预分配结果数组(按请求顺序),校验非法名并记录失败
|
||||
results: list[dict] = [{"slug": "", "success": False, "error": "unset"} for _ in range(len(skills))]
|
||||
normalized_skills: list[str] = []
|
||||
valid_indices: list[int] = []
|
||||
for i, skill in enumerate(skills):
|
||||
try:
|
||||
normalized_skills.append(_normalize_skill_name(skill))
|
||||
valid_indices.append(i)
|
||||
except ValueError as e:
|
||||
results[i] = {"slug": skill, "success": False, "error": str(e)}
|
||||
|
||||
if not normalized_skills:
|
||||
return RemoteSkillsBatchPreparation(temp_home=None, results=results)
|
||||
|
||||
temp_home, env, workdir = _create_isolated_workdir()
|
||||
try:
|
||||
skill_args: list[str] = []
|
||||
for name in normalized_skills:
|
||||
skill_args.extend(["--skill", name])
|
||||
|
||||
cli_failed = False
|
||||
try:
|
||||
await _run_skills_cli(
|
||||
[
|
||||
"npx",
|
||||
"-y",
|
||||
"skills",
|
||||
"add",
|
||||
normalized_source,
|
||||
*skill_args,
|
||||
"-g",
|
||||
"-y",
|
||||
"--copy",
|
||||
],
|
||||
env=env,
|
||||
cwd=workdir,
|
||||
)
|
||||
except ValueError:
|
||||
# CLI 对不匹配的 skill 会退出码非零,但已安装的目录仍在
|
||||
cli_failed = True
|
||||
|
||||
skills_dir = Path(temp_home).resolve() / ".agents" / "skills"
|
||||
for original_index, name in zip(valid_indices, normalized_skills):
|
||||
installed_dir = _find_skill_dir(skills_dir, name)
|
||||
if installed_dir is None:
|
||||
error_msg = "CLI 安装失败" if cli_failed else "skills CLI 未生成预期的技能目录"
|
||||
results[original_index] = {"slug": name, "success": False, "error": error_msg}
|
||||
else:
|
||||
results[original_index] = {"slug": name, "success": True, "source_dir": installed_dir}
|
||||
|
||||
return RemoteSkillsBatchPreparation(temp_home=temp_home, results=results)
|
||||
except Exception:
|
||||
await asyncio.to_thread(shutil.rmtree, temp_home, ignore_errors=True)
|
||||
raise
|
||||
|
||||
|
||||
def _find_skill_dir(skills_dir: Path, name: str) -> Path | None:
|
||||
"""在 skills 安装目录下按名称查找 skill 子目录。"""
|
||||
if not skills_dir.is_dir():
|
||||
return None
|
||||
for candidate in skills_dir.iterdir():
|
||||
if candidate.name == name and candidate.is_dir():
|
||||
return candidate
|
||||
return None
|
||||
|
||||
|
||||
def _parse_search_skills(output: str) -> list[dict[str, str]]:
|
||||
"""解析 npx skills find 命令的输出。"""
|
||||
lines = _clean_cli_output(output)
|
||||
results: list[dict[str, str]] = []
|
||||
# 匹配形如 "owner/repo@skill-name [installs]"
|
||||
# 例如:vercel-labs/agent-skills@web-design-guidelines 339.3K installs
|
||||
pattern = re.compile(r"^([a-zA-Z0-9_\-\.]+/[a-zA-Z0-9_\-\.]+)\@([a-zA-Z0-9_\-\.]+)(?:\s+(.*))?$")
|
||||
for line in lines:
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
match = pattern.match(line)
|
||||
if match:
|
||||
source, name, extra = match.groups()
|
||||
installs = extra.strip() if extra else ""
|
||||
results.append(
|
||||
{
|
||||
"source": source,
|
||||
"name": name,
|
||||
"installs": installs,
|
||||
}
|
||||
)
|
||||
return results
|
||||
|
||||
|
||||
async def search_remote_skills(query: str) -> list[dict[str, str]]:
|
||||
"""使用 npx skills find <query> 搜索远程 skills。"""
|
||||
query_val = str(query or "").strip()
|
||||
if not query_val:
|
||||
return []
|
||||
if any(ch in query_val for ch in ("\n", "\r", "\x00")):
|
||||
raise ValueError("搜索关键字包含非法字符")
|
||||
|
||||
temp_home, env, workdir = _create_isolated_workdir()
|
||||
try:
|
||||
output = await _run_skills_cli(
|
||||
["npx", "-y", "skills", "find", query_val],
|
||||
env=env,
|
||||
cwd=workdir,
|
||||
)
|
||||
finally:
|
||||
await asyncio.to_thread(shutil.rmtree, temp_home, ignore_errors=True)
|
||||
|
||||
return _parse_search_skills(output)
|
||||
@@ -0,0 +1,154 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from yuxi.storage.postgres.models_business import Skill
|
||||
from yuxi.utils.datetime_utils import utc_now_naive
|
||||
|
||||
|
||||
class SkillRepository:
|
||||
def __init__(self, db_session: AsyncSession):
|
||||
self.db = db_session
|
||||
|
||||
async def list_all(self) -> list[Skill]:
|
||||
result = await self.db.execute(select(Skill).order_by(Skill.updated_at.desc(), Skill.id.desc()))
|
||||
return list(result.scalars().all())
|
||||
|
||||
async def list_enabled(self) -> list[Skill]:
|
||||
result = await self.db.execute(
|
||||
select(Skill).where(Skill.enabled.is_(True)).order_by(Skill.updated_at.desc(), Skill.id.desc())
|
||||
)
|
||||
return list(result.scalars().all())
|
||||
|
||||
async def list_by_slugs(self, slugs: list[str]) -> list[Skill]:
|
||||
normalized = [slug for slug in dict.fromkeys(slugs) if isinstance(slug, str) and slug]
|
||||
if not normalized:
|
||||
return []
|
||||
result = await self.db.execute(select(Skill).where(Skill.slug.in_(normalized)))
|
||||
items = list(result.scalars().all())
|
||||
item_map = {item.slug: item for item in items}
|
||||
return [item_map[slug] for slug in normalized if slug in item_map]
|
||||
|
||||
async def get_by_slug(self, slug: str, *, for_update: bool = False) -> Skill | None:
|
||||
stmt = select(Skill).where(Skill.slug == slug)
|
||||
if for_update:
|
||||
stmt = stmt.with_for_update()
|
||||
result = await self.db.execute(stmt)
|
||||
return result.scalar_one_or_none()
|
||||
|
||||
async def exists_slug(self, slug: str) -> bool:
|
||||
return (await self.get_by_slug(slug)) is not None
|
||||
|
||||
async def create(
|
||||
self,
|
||||
*,
|
||||
slug: str,
|
||||
name: str,
|
||||
description: str,
|
||||
source_type: str,
|
||||
tool_dependencies: list[str] | None,
|
||||
mcp_dependencies: list[str] | None,
|
||||
skill_dependencies: list[str] | None,
|
||||
dir_path: str,
|
||||
share_config: dict,
|
||||
enabled: bool = True,
|
||||
version: str | None = None,
|
||||
content_hash: str | None = None,
|
||||
created_by: str | None,
|
||||
) -> Skill:
|
||||
now = utc_now_naive()
|
||||
item = Skill(
|
||||
slug=slug,
|
||||
name=name,
|
||||
description=description,
|
||||
source_type=source_type,
|
||||
tool_dependencies=tool_dependencies or [],
|
||||
mcp_dependencies=mcp_dependencies or [],
|
||||
skill_dependencies=skill_dependencies or [],
|
||||
dir_path=dir_path,
|
||||
version=version,
|
||||
content_hash=content_hash,
|
||||
share_config=share_config,
|
||||
enabled=enabled,
|
||||
created_by=created_by,
|
||||
updated_by=created_by,
|
||||
created_at=now,
|
||||
updated_at=now,
|
||||
)
|
||||
self.db.add(item)
|
||||
await self.db.commit()
|
||||
await self.db.refresh(item)
|
||||
return item
|
||||
|
||||
async def update_builtin_install(
|
||||
self,
|
||||
item: Skill,
|
||||
*,
|
||||
version: str,
|
||||
content_hash: str,
|
||||
updated_by: str | None,
|
||||
) -> Skill:
|
||||
item.version = version
|
||||
item.content_hash = content_hash
|
||||
item.source_type = "builtin"
|
||||
item.share_config = {"access_level": "global", "department_ids": [], "user_uids": []}
|
||||
item.updated_by = updated_by
|
||||
item.updated_at = utc_now_naive()
|
||||
await self.db.commit()
|
||||
await self.db.refresh(item)
|
||||
return item
|
||||
|
||||
async def update_dependencies(
|
||||
self,
|
||||
item: Skill,
|
||||
*,
|
||||
tool_dependencies: list[str],
|
||||
mcp_dependencies: list[str],
|
||||
skill_dependencies: list[str],
|
||||
updated_by: str | None,
|
||||
) -> Skill:
|
||||
item.tool_dependencies = tool_dependencies
|
||||
item.mcp_dependencies = mcp_dependencies
|
||||
item.skill_dependencies = skill_dependencies
|
||||
item.updated_by = updated_by
|
||||
item.updated_at = utc_now_naive()
|
||||
await self.db.commit()
|
||||
await self.db.refresh(item)
|
||||
return item
|
||||
|
||||
async def update_metadata(
|
||||
self,
|
||||
item: Skill,
|
||||
*,
|
||||
name: str,
|
||||
description: str,
|
||||
updated_by: str | None,
|
||||
) -> Skill:
|
||||
item.name = name
|
||||
item.description = description
|
||||
item.updated_by = updated_by
|
||||
item.updated_at = utc_now_naive()
|
||||
await self.db.commit()
|
||||
await self.db.refresh(item)
|
||||
return item
|
||||
|
||||
async def update_share_config(self, item: Skill, *, share_config: dict, updated_by: str | None) -> Skill:
|
||||
item.share_config = share_config
|
||||
item.updated_by = updated_by
|
||||
item.updated_at = utc_now_naive()
|
||||
await self.db.commit()
|
||||
await self.db.refresh(item)
|
||||
return item
|
||||
|
||||
async def update_enabled(self, item: Skill, *, enabled: bool, updated_by: str | None) -> Skill:
|
||||
item.enabled = enabled
|
||||
item.updated_by = updated_by
|
||||
item.updated_at = utc_now_naive()
|
||||
await self.db.commit()
|
||||
await self.db.refresh(item)
|
||||
return item
|
||||
|
||||
async def delete(self, item: Skill) -> None:
|
||||
await self.db.delete(item)
|
||||
await self.db.commit()
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,32 @@
|
||||
"""Define the state structures for the agent."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Annotated, TypedDict
|
||||
|
||||
from langchain.agents import AgentState
|
||||
|
||||
|
||||
def merge_artifacts(existing: list[str] | None, new: list[str] | None) -> list[str]:
|
||||
"""Merge artifact file paths while preserving order and removing duplicates."""
|
||||
if existing is None:
|
||||
return new or []
|
||||
if new is None:
|
||||
return existing
|
||||
return list(dict.fromkeys(existing + new))
|
||||
|
||||
|
||||
class BaseState(AgentState):
|
||||
"""Shared state fields for Yuxi agents."""
|
||||
|
||||
artifacts: Annotated[list[str], merge_artifacts]
|
||||
|
||||
|
||||
class AgentStatePayload(TypedDict):
|
||||
"""Serialized agent state payload consumed by the frontend."""
|
||||
|
||||
todos: list
|
||||
files: dict
|
||||
artifacts: list[str]
|
||||
subagent_runs: list[dict]
|
||||
token_usage: dict | None
|
||||
@@ -0,0 +1,25 @@
|
||||
# toolkits 包
|
||||
# 触发各模块的 @tool 装饰器执行,自动注册工具
|
||||
from . import buildin, debug
|
||||
|
||||
# 工具获取函数
|
||||
from .kbs import get_common_kb_tools
|
||||
from .registry import (
|
||||
ToolExtraMetadata,
|
||||
get_all_extra_metadata,
|
||||
get_all_tool_instances,
|
||||
get_extra_metadata,
|
||||
tool,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"get_extra_metadata",
|
||||
"get_all_extra_metadata",
|
||||
"get_all_tool_instances",
|
||||
"ToolExtraMetadata",
|
||||
"tool",
|
||||
"get_common_kb_tools",
|
||||
# 触发各模块的 @tool 装饰器执行,自动注册工具
|
||||
"buildin",
|
||||
"debug",
|
||||
]
|
||||
@@ -0,0 +1,10 @@
|
||||
# buildin 工具包
|
||||
from .install_skill import install_skill
|
||||
from .tools import ask_user_question, ocr_parse_file, present_artifacts
|
||||
|
||||
__all__ = [
|
||||
"ask_user_question",
|
||||
"install_skill",
|
||||
"ocr_parse_file",
|
||||
"present_artifacts",
|
||||
]
|
||||
@@ -0,0 +1,302 @@
|
||||
import shutil
|
||||
import tempfile
|
||||
from pathlib import Path, PurePosixPath
|
||||
from typing import Annotated
|
||||
|
||||
from langchain.tools import InjectedToolCallId
|
||||
from langchain_core.messages import ToolMessage
|
||||
from langgraph.prebuilt.tool_node import ToolRuntime
|
||||
from langgraph.types import Command
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from yuxi.agents.toolkits.registry import tool
|
||||
from yuxi.repositories.agent_repository import AgentRepository
|
||||
from yuxi.repositories.conversation_repository import ConversationRepository
|
||||
from yuxi.storage.postgres.manager import pg_manager
|
||||
from yuxi.utils.logging_config import logger
|
||||
|
||||
SANDBOX_PATH_HINT = (
|
||||
"请使用 /home/gem/user-data/workspace/...、/home/gem/user-data/uploads/... 或 /home/gem/user-data/outputs/..."
|
||||
)
|
||||
MAX_SANDBOX_SKILL_FILES = 1000
|
||||
|
||||
|
||||
class InstallSkillInput(BaseModel):
|
||||
source: str = Field(
|
||||
description="Skill 来源,支持两种格式:\n"
|
||||
"1. Sandbox 路径: /home/gem/user-data/workspace/...、"
|
||||
"/home/gem/user-data/uploads/... 或 /home/gem/user-data/outputs/...(/ 开头)\n"
|
||||
"2. Git 仓库: owner/repo 或完整 GitHub URL"
|
||||
)
|
||||
skill_names: list[str] | None = Field(
|
||||
default=None, description="Git 安装时指定要安装的 skill slug 列表(至少一个)。Sandbox 路径安装时忽略此参数。"
|
||||
)
|
||||
|
||||
|
||||
def _personal_skill_share_config(uid: str) -> dict:
|
||||
return {"access_level": "user", "department_ids": [], "user_uids": [str(uid)]}
|
||||
|
||||
|
||||
def _collect_sandbox_file_paths(backend, remote_dir: str, file_paths: list[str] | None = None) -> list[str]:
|
||||
file_paths = file_paths if file_paths is not None else []
|
||||
result = backend.ls(remote_dir)
|
||||
if result.error:
|
||||
raise ValueError(result.error)
|
||||
entries = result.entries or []
|
||||
for entry in entries:
|
||||
path = entry["path"]
|
||||
if entry.get("is_dir"):
|
||||
_collect_sandbox_file_paths(backend, path, file_paths)
|
||||
else:
|
||||
if len(file_paths) >= MAX_SANDBOX_SKILL_FILES:
|
||||
raise ValueError(f"Skill 目录文件数超过限制(最多 {MAX_SANDBOX_SKILL_FILES} 个文件)")
|
||||
file_paths.append(path)
|
||||
return file_paths
|
||||
|
||||
|
||||
def _download_skill_dir(backend, remote_dir: str, local_dir: Path) -> None:
|
||||
"""递归通过沙盒 API 下载 skill 目录到本地。"""
|
||||
remote_root = PurePosixPath(remote_dir.rstrip("/"))
|
||||
file_paths = _collect_sandbox_file_paths(backend, remote_dir)
|
||||
if not file_paths:
|
||||
raise ValueError(f"沙盒路径 {remote_dir} 中未发现可下载文件")
|
||||
|
||||
responses = backend.download_files(file_paths)
|
||||
if len(responses) != len(file_paths):
|
||||
raise ValueError("沙盒文件下载结果数量不匹配")
|
||||
|
||||
for expected_path, response in zip(file_paths, responses):
|
||||
error = getattr(response, "error", None)
|
||||
content = getattr(response, "content", None)
|
||||
if error or content is None:
|
||||
raise ValueError(f"下载沙盒文件失败: {expected_path} ({error or 'empty_content'})")
|
||||
|
||||
pure_path = PurePosixPath(expected_path)
|
||||
try:
|
||||
relative_path = pure_path.relative_to(remote_root)
|
||||
except ValueError:
|
||||
relative_path = PurePosixPath(pure_path.name)
|
||||
|
||||
target_path = local_dir / Path(relative_path.as_posix())
|
||||
target_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
target_path.write_bytes(content)
|
||||
|
||||
|
||||
def _prepare_skill_from_sandbox(sandbox_path: str, thread_id: str, uid: str, staging_root: Path) -> tuple[Path, str]:
|
||||
"""从 Sandbox 路径准备 skill 目录。返回 (本地目录, 原始 skill name)。"""
|
||||
from yuxi.agents.backends.sandbox import ProvisionerSandboxBackend, resolve_virtual_path
|
||||
from yuxi.agents.skills.service import (
|
||||
_parse_skill_markdown,
|
||||
is_valid_skill_slug,
|
||||
)
|
||||
|
||||
slug = PurePosixPath(sandbox_path.rstrip("/")).name
|
||||
if not is_valid_skill_slug(slug):
|
||||
raise ValueError(f"slug '{slug}' 不合法(仅允许小写字母、数字和连字符)")
|
||||
|
||||
if not sandbox_path.startswith("/home/gem/user-data/"):
|
||||
raise ValueError(f"不支持的沙盒路径: {sandbox_path}。{SANDBOX_PATH_HINT}")
|
||||
|
||||
staging = staging_root / slug
|
||||
|
||||
# 优先尝试共享卷路径(性能更好,无需走沙盒 API)。
|
||||
try:
|
||||
local_path = resolve_virtual_path(thread_id, sandbox_path, uid=uid)
|
||||
if (local_path / "SKILL.md").exists():
|
||||
shutil.copytree(local_path, staging)
|
||||
else:
|
||||
raise FileNotFoundError(f"{local_path} 中未找到 SKILL.md")
|
||||
except FileNotFoundError:
|
||||
staging.mkdir(parents=True, exist_ok=True)
|
||||
backend = ProvisionerSandboxBackend(thread_id=thread_id, uid=uid)
|
||||
_download_skill_dir(backend, sandbox_path, staging)
|
||||
if not (staging / "SKILL.md").exists():
|
||||
raise ValueError(f"沙盒路径 {sandbox_path} 中未找到 SKILL.md")
|
||||
|
||||
content = (staging / "SKILL.md").read_text(encoding="utf-8")
|
||||
parsed_name, _, _ = _parse_skill_markdown(content)
|
||||
return staging, parsed_name
|
||||
|
||||
|
||||
async def _enable_skills_in_current_config(db, thread_id: str, uid: str, skill_slugs: list[str]) -> bool:
|
||||
"""在当前会话绑定且当前用户拥有的 Agent 配置中启用新安装的 skill。"""
|
||||
conv_repo = ConversationRepository(db)
|
||||
conv = await conv_repo.get_conversation_by_thread_id(thread_id)
|
||||
if not conv or str(conv.uid) != str(uid):
|
||||
return False
|
||||
|
||||
agent_repo = AgentRepository(db)
|
||||
agent = await agent_repo.get_by_slug(conv.agent_id)
|
||||
if not agent or agent.created_by != str(uid):
|
||||
return False
|
||||
|
||||
config_json = dict(agent.config_json or {})
|
||||
context = dict(config_json.get("context") or {})
|
||||
skills = [item for item in context.get("skills") or [] if isinstance(item, str) and item.strip()]
|
||||
seen = set(skills)
|
||||
for slug in skill_slugs:
|
||||
if slug not in seen:
|
||||
skills.append(slug)
|
||||
seen.add(slug)
|
||||
context["skills"] = skills
|
||||
config_json["context"] = context
|
||||
await agent_repo.update(agent, config_json=config_json, updated_by=str(uid))
|
||||
return True
|
||||
|
||||
|
||||
async def _run_install_task(
|
||||
source: str,
|
||||
runtime: ToolRuntime,
|
||||
tool_call_id: str,
|
||||
skill_names: list[str] | None = None,
|
||||
) -> Command:
|
||||
"""执行异步安装任务的核心逻辑。"""
|
||||
runtime_context = getattr(runtime, "context", None)
|
||||
if getattr(runtime_context, "is_subagent_runtime", False):
|
||||
return Command(
|
||||
update={
|
||||
"messages": [
|
||||
ToolMessage(
|
||||
content="错误:install_skill 只能在主智能体中使用,子智能体无法安装 Skill",
|
||||
tool_call_id=tool_call_id,
|
||||
)
|
||||
]
|
||||
}
|
||||
)
|
||||
|
||||
source = str(source or "").strip()
|
||||
uid = getattr(runtime_context, "uid", None)
|
||||
thread_id = getattr(runtime_context, "thread_id", None)
|
||||
|
||||
logger.info(f"install_skill called with uid={uid}, thread_id={thread_id}, source={source}")
|
||||
|
||||
if not uid or not thread_id:
|
||||
return Command(
|
||||
update={"messages": [ToolMessage(content="错误:无法获取当前会话信息", tool_call_id=tool_call_id)]}
|
||||
)
|
||||
if not source:
|
||||
return Command(
|
||||
update={"messages": [ToolMessage(content="错误:Skill 来源不能为空", tool_call_id=tool_call_id)]}
|
||||
)
|
||||
|
||||
personal_share_config = _personal_skill_share_config(uid)
|
||||
|
||||
try:
|
||||
from yuxi.agents.skills.service import (
|
||||
import_skill_dir,
|
||||
normalize_string_list,
|
||||
sync_thread_readable_skills,
|
||||
)
|
||||
|
||||
installed_slugs: list[str] = []
|
||||
failed_items: list[dict] = []
|
||||
slug_warnings: list[str] = []
|
||||
config_success = True
|
||||
|
||||
if source.startswith("/"):
|
||||
with tempfile.TemporaryDirectory(prefix=".skill-install-") as tmp:
|
||||
source_dir, parsed_name = _prepare_skill_from_sandbox(source, thread_id, uid, Path(tmp))
|
||||
async with pg_manager.get_async_session_context() as db:
|
||||
item = await import_skill_dir(
|
||||
db,
|
||||
source_dir=source_dir,
|
||||
created_by=uid,
|
||||
share_config=personal_share_config,
|
||||
)
|
||||
installed_slugs = [item.slug]
|
||||
if item.slug != parsed_name:
|
||||
slug_warnings.append(f"⚠️ 技能 slug '{item.slug}' 已存在,已自动重命名安装")
|
||||
config_success = await _enable_skills_in_current_config(db, thread_id, uid, installed_slugs)
|
||||
else:
|
||||
_skill_names = skill_names or []
|
||||
if not _skill_names:
|
||||
return Command(
|
||||
update={
|
||||
"messages": [
|
||||
ToolMessage(
|
||||
content="❌ 错误: 从 Git 安装时必须通过 skill_names 指定技能名称",
|
||||
tool_call_id=tool_call_id,
|
||||
)
|
||||
]
|
||||
}
|
||||
)
|
||||
|
||||
from yuxi.agents.skills.remote_install import prepare_remote_skills_batch
|
||||
|
||||
preparation = await prepare_remote_skills_batch(source=source, skills=_skill_names)
|
||||
try:
|
||||
async with pg_manager.get_async_session_context() as db:
|
||||
for result in preparation.results:
|
||||
if not result.get("success"):
|
||||
failed_items.append(result)
|
||||
continue
|
||||
|
||||
try:
|
||||
item = await import_skill_dir(
|
||||
db,
|
||||
source_dir=result["source_dir"],
|
||||
created_by=uid,
|
||||
share_config=personal_share_config,
|
||||
)
|
||||
installed_slugs.append(item.slug)
|
||||
except Exception as e:
|
||||
await db.rollback()
|
||||
failed_items.append({"slug": result["slug"], "success": False, "error": str(e)})
|
||||
|
||||
config_success = True
|
||||
if installed_slugs:
|
||||
config_success = await _enable_skills_in_current_config(db, thread_id, uid, installed_slugs)
|
||||
finally:
|
||||
preparation.cleanup()
|
||||
|
||||
current_skills = normalize_string_list(getattr(runtime_context, "skills", None))
|
||||
sync_thread_readable_skills(thread_id, normalize_string_list(current_skills + installed_slugs))
|
||||
|
||||
lines = []
|
||||
if installed_slugs:
|
||||
lines.append(f"✅ 成功安装并激活技能: {', '.join(installed_slugs)}")
|
||||
for warning in slug_warnings:
|
||||
lines.append(warning)
|
||||
if failed_items:
|
||||
for item in failed_items:
|
||||
lines.append(f"❌ 安装失败 ({item['slug']}): {item.get('error', '未知错误')}")
|
||||
if not config_success:
|
||||
lines.append("⚠️ 已添加这个 Skill 拓展(仅当前会话生效,未写入当前 Agent 配置)")
|
||||
if not installed_slugs and not failed_items:
|
||||
lines.append("ℹ️ 未发现需要安装的技能")
|
||||
|
||||
return Command(
|
||||
update={
|
||||
"activated_skills": installed_slugs,
|
||||
"messages": [ToolMessage(content="\n".join(lines), tool_call_id=tool_call_id)],
|
||||
}
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.exception("install_skill 异常")
|
||||
return Command(
|
||||
update={
|
||||
"messages": [
|
||||
ToolMessage(
|
||||
content=f"❌ 安装异常: {str(e)}",
|
||||
tool_call_id=tool_call_id,
|
||||
)
|
||||
]
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@tool(
|
||||
category="buildin",
|
||||
tags=["skill", "安装"],
|
||||
display_name="安装技能",
|
||||
args_schema=InstallSkillInput,
|
||||
)
|
||||
async def install_skill(
|
||||
source: str,
|
||||
skill_names: list[str] | None = None,
|
||||
runtime: ToolRuntime = None,
|
||||
tool_call_id: Annotated[str, InjectedToolCallId] = "",
|
||||
) -> Command:
|
||||
"""安装新的 Skill 到当前用户的私有空间,并在当前主智能体会话中激活。"""
|
||||
return await _run_install_task(source, runtime, tool_call_id, skill_names)
|
||||
@@ -0,0 +1,388 @@
|
||||
import os
|
||||
import re
|
||||
from pathlib import Path
|
||||
from typing import Annotated
|
||||
|
||||
from langchain.tools import InjectedToolCallId
|
||||
from langchain_core.messages import ToolMessage
|
||||
from langgraph.prebuilt.tool_node import ToolRuntime
|
||||
from langgraph.types import Command, interrupt
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from yuxi.agents.toolkits.registry import ToolExtraMetadata, _all_tool_instances, _extra_registry, tool
|
||||
from yuxi.utils import logger
|
||||
from yuxi.utils.paths import (
|
||||
CONVERSATION_HISTORY_DIR_NAME,
|
||||
LARGE_TOOL_RESULTS_DIR_NAME,
|
||||
OUTPUTS_DIR_NAME,
|
||||
UPLOADS_DIR_NAME,
|
||||
VIRTUAL_PATH_OUTPUTS,
|
||||
WORKSPACE_DIR_NAME,
|
||||
)
|
||||
from yuxi.utils.question_utils import normalize_questions
|
||||
|
||||
# Lazy initialization for TavilySearch (only when API key is available)
|
||||
_tavily_search_instance = None
|
||||
|
||||
_PRESENT_ARTIFACTS_INTERNAL_DIR_NAMES = frozenset(
|
||||
{CONVERSATION_HISTORY_DIR_NAME, LARGE_TOOL_RESULTS_DIR_NAME, "large_tool_history"}
|
||||
)
|
||||
_OCR_PARSE_ALLOWED_DIRS = frozenset({WORKSPACE_DIR_NAME, UPLOADS_DIR_NAME, OUTPUTS_DIR_NAME})
|
||||
_OCR_OUTPUT_DIR_NAME = "ocr"
|
||||
_OCR_PREVIEW_LIMIT = 1200
|
||||
_SAFE_OUTPUT_STEM_RE = re.compile(r"[^A-Za-z0-9._\-\u4e00-\u9fff]+")
|
||||
|
||||
|
||||
def _create_tavily_search():
|
||||
"""Create and register TavilySearch tool with metadata."""
|
||||
global _tavily_search_instance
|
||||
if _tavily_search_instance is None:
|
||||
from langchain_tavily import TavilySearch
|
||||
|
||||
_tavily_search_instance = TavilySearch()
|
||||
|
||||
return _tavily_search_instance
|
||||
|
||||
|
||||
# 注册 TavilySearch 工具(延迟初始化)
|
||||
def _register_tavily_tool():
|
||||
"""Register TavilySearch tool with extra metadata."""
|
||||
tavily_instance = _create_tavily_search()
|
||||
# 手动注册到全局注册表
|
||||
_extra_registry["tavily_search"] = ToolExtraMetadata(
|
||||
category="buildin",
|
||||
tags=["搜索"],
|
||||
display_name="Tavily 网页搜索",
|
||||
)
|
||||
# 添加到工具实例列表
|
||||
_all_tool_instances.append(tavily_instance)
|
||||
|
||||
|
||||
# 模块加载时注册
|
||||
if os.getenv("TAVILY_API_KEY"):
|
||||
try:
|
||||
_register_tavily_tool()
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to register TavilySearch tool: {e}")
|
||||
|
||||
|
||||
class PresentArtifactsInput(BaseModel):
|
||||
"""Expose artifact files to the frontend after the agent finishes."""
|
||||
|
||||
filepaths: list[str] = Field(
|
||||
description=f"需要展示给用户的文件绝对路径列表,只允许位于 {VIRTUAL_PATH_OUTPUTS} 下,且不能是内部运行文件"
|
||||
)
|
||||
|
||||
|
||||
def _normalize_presented_artifact_path(filepath: str, runtime: ToolRuntime) -> str:
|
||||
from yuxi.agents.backends.sandbox.paths import (
|
||||
VIRTUAL_PATH_PREFIX,
|
||||
ensure_thread_dirs,
|
||||
resolve_virtual_path,
|
||||
sandbox_outputs_dir,
|
||||
)
|
||||
|
||||
outputs_virtual_prefix = f"{VIRTUAL_PATH_PREFIX}/outputs"
|
||||
runtime_context = runtime.context
|
||||
thread_id = getattr(runtime_context, "file_thread_id", None) or getattr(runtime_context, "thread_id", None)
|
||||
if not thread_id:
|
||||
raise ValueError("当前运行时缺少 thread_id")
|
||||
uid = getattr(runtime_context, "uid", None)
|
||||
if not uid:
|
||||
raise ValueError("当前运行时缺少 uid")
|
||||
|
||||
ensure_thread_dirs(thread_id, str(uid))
|
||||
outputs_dir = sandbox_outputs_dir(thread_id).resolve()
|
||||
normalized_input = str(filepath or "").strip()
|
||||
if not normalized_input:
|
||||
raise ValueError("文件路径不能为空")
|
||||
|
||||
stripped = normalized_input.lstrip("/")
|
||||
virtual_prefix = VIRTUAL_PATH_PREFIX.lstrip("/")
|
||||
if stripped == virtual_prefix or stripped.startswith(f"{virtual_prefix}/"):
|
||||
actual_path = resolve_virtual_path(thread_id, normalized_input, uid=str(uid))
|
||||
else:
|
||||
actual_path = Path(normalized_input).expanduser().resolve()
|
||||
|
||||
if not actual_path.exists() or not actual_path.is_file():
|
||||
raise ValueError(f"文件不存在或不是普通文件: {normalized_input}")
|
||||
|
||||
try:
|
||||
relative_path = actual_path.relative_to(outputs_dir)
|
||||
except ValueError as exc:
|
||||
raise ValueError(f"只允许展示 {outputs_virtual_prefix}/ 下的文件: {normalized_input}") from exc
|
||||
|
||||
if relative_path.parts and relative_path.parts[0] in _PRESENT_ARTIFACTS_INTERNAL_DIR_NAMES:
|
||||
raise ValueError(f"不允许展示工具调用阶段文件: {outputs_virtual_prefix}/{relative_path.as_posix()}")
|
||||
|
||||
return f"{outputs_virtual_prefix}/{relative_path.as_posix()}"
|
||||
|
||||
|
||||
PRESENT_ARTIFACTS_DESCRIPTION = f"""
|
||||
将已经生成好的结果文件展示给用户。
|
||||
|
||||
使用场景:
|
||||
1. 你已经在 `{VIRTUAL_PATH_OUTPUTS}` 下写好了最终结果文件
|
||||
2. 你希望前端在对话结束后显示这些结果文件卡片
|
||||
3. 这些文件需要支持下载或预览
|
||||
|
||||
注意事项:
|
||||
1. 只能传入 `{VIRTUAL_PATH_OUTPUTS}` 下的文件
|
||||
2. 不要传入中间过程文件,只有真正需要给用户看的结果文件才调用
|
||||
3. 不要传入工具调用阶段文件,例如:
|
||||
- `{VIRTUAL_PATH_OUTPUTS}/{LARGE_TOOL_RESULTS_DIR_NAME}`
|
||||
- `{VIRTUAL_PATH_OUTPUTS}/{CONVERSATION_HISTORY_DIR_NAME}`
|
||||
4. 可以一次传多个文件
|
||||
"""
|
||||
|
||||
|
||||
@tool(
|
||||
category="buildin",
|
||||
tags=["文件", "交付物"],
|
||||
display_name="展示交付物",
|
||||
description=PRESENT_ARTIFACTS_DESCRIPTION,
|
||||
args_schema=PresentArtifactsInput,
|
||||
)
|
||||
def present_artifacts(
|
||||
filepaths: list[str],
|
||||
runtime: ToolRuntime,
|
||||
tool_call_id: Annotated[str, InjectedToolCallId],
|
||||
) -> Command:
|
||||
"""登记当前线程 outputs 目录下的交付物文件,使前端在对话结束后展示给用户。"""
|
||||
try:
|
||||
normalized_paths = [_normalize_presented_artifact_path(filepath, runtime) for filepath in filepaths]
|
||||
except ValueError as exc:
|
||||
return Command(update={"messages": [ToolMessage(content=f"Error: {exc}", tool_call_id=tool_call_id)]})
|
||||
|
||||
return Command(
|
||||
update={
|
||||
"artifacts": normalized_paths,
|
||||
"messages": [ToolMessage(content="已将交付物展示给用户", tool_call_id=tool_call_id)],
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
class OcrParseFileInput(BaseModel):
|
||||
"""Parse a sandbox file with OCR and save the Markdown result."""
|
||||
|
||||
file_path: str = Field(description="需要 OCR 解析的沙盒虚拟路径,必须位于 /home/gem/user-data 下")
|
||||
ocr_engine: str | None = Field(default=None, description="可选 OCR 引擎;省略时使用系统默认 OCR 引擎")
|
||||
|
||||
|
||||
OCR_PARSE_FILE_DESCRIPTION = f"""
|
||||
将沙盒中的 PDF 或图片文件解析为 Markdown 文本,并把结果保存为文件。
|
||||
|
||||
使用场景:
|
||||
1. 用户上传了 PDF/图片附件,需要提取其中的文字内容
|
||||
2. 工作区、uploads 或 outputs 下已有文件,需要转成可读取的 Markdown
|
||||
3. 解析结果较长,后续应使用 read_file 读取保存后的 Markdown 文件
|
||||
|
||||
注意事项:
|
||||
1. file_path 必须是 /home/gem/user-data 下的虚拟路径
|
||||
2. 只允许读取 workspace、uploads、outputs 下的普通文件
|
||||
3. 解析结果会写入 {VIRTUAL_PATH_OUTPUTS}/{_OCR_OUTPUT_DIR_NAME}/
|
||||
4. 工具只返回结果文件路径和短预览,不直接返回完整 OCR 文本
|
||||
5. 如需在前端展示结果文件,请再调用 present_artifacts
|
||||
"""
|
||||
|
||||
|
||||
@tool(
|
||||
category="buildin",
|
||||
tags=["文件", "OCR"],
|
||||
display_name="OCR 解析文件",
|
||||
description=OCR_PARSE_FILE_DESCRIPTION,
|
||||
args_schema=OcrParseFileInput,
|
||||
)
|
||||
async def ocr_parse_file(file_path: str, runtime: ToolRuntime, ocr_engine: str | None = None) -> dict:
|
||||
"""Parse a sandbox file with OCR, persist Markdown output, and return only a short result summary."""
|
||||
from yuxi.agents.backends.sandbox.paths import virtual_path_for_thread_file
|
||||
from yuxi.knowledge.parser import Parser
|
||||
|
||||
file_thread_id, uid, actual_path = _resolve_ocr_source_path(file_path, runtime)
|
||||
engine = _resolve_ocr_engine(ocr_engine)
|
||||
markdown = await Parser.aparse(str(actual_path), params={"ocr_engine": engine})
|
||||
|
||||
output_path = _next_ocr_output_path(file_thread_id, actual_path)
|
||||
output_path.write_text(markdown, encoding="utf-8")
|
||||
parsed_path = virtual_path_for_thread_file(file_thread_id, output_path, uid=uid)
|
||||
source_virtual_path = virtual_path_for_thread_file(file_thread_id, actual_path, uid=uid)
|
||||
preview, truncated = _ocr_preview(markdown)
|
||||
|
||||
return {
|
||||
"source_path": source_virtual_path,
|
||||
"parsed_path": parsed_path,
|
||||
"ocr_engine": engine,
|
||||
"char_count": len(markdown),
|
||||
"preview": preview,
|
||||
"truncated": truncated,
|
||||
}
|
||||
|
||||
|
||||
def _resolve_ocr_source_path(file_path: str, runtime: ToolRuntime) -> tuple[str, str, Path]:
|
||||
"""Resolve a sandbox virtual path to a host file inside the Agent-visible user-data roots."""
|
||||
from yuxi.agents.backends.sandbox.paths import get_virtual_path_prefix, resolve_virtual_path
|
||||
|
||||
file_thread_id, uid = _resolve_runtime_file_scope(runtime)
|
||||
|
||||
normalized_input = str(file_path or "").strip()
|
||||
if not normalized_input:
|
||||
raise ValueError("文件路径不能为空")
|
||||
|
||||
virtual_prefix = get_virtual_path_prefix().rstrip("/")
|
||||
clean_virtual_path = "/" + normalized_input.lstrip("/")
|
||||
if clean_virtual_path != virtual_prefix and not clean_virtual_path.startswith(f"{virtual_prefix}/"):
|
||||
raise ValueError(f"只允许解析 {virtual_prefix} 下的沙盒虚拟路径")
|
||||
|
||||
relative_path = clean_virtual_path[len(virtual_prefix) :].lstrip("/")
|
||||
namespace = Path(relative_path).parts[0] if relative_path else ""
|
||||
if namespace not in _OCR_PARSE_ALLOWED_DIRS:
|
||||
allowed = ", ".join(f"{virtual_prefix}/{item}" for item in sorted(_OCR_PARSE_ALLOWED_DIRS))
|
||||
raise ValueError(f"只允许解析 {allowed} 下的文件")
|
||||
|
||||
try:
|
||||
actual_path = resolve_virtual_path(file_thread_id, clean_virtual_path, uid=uid)
|
||||
except ValueError as exc:
|
||||
raise ValueError(f"只允许解析 {virtual_prefix} 下的沙盒虚拟路径") from exc
|
||||
if not actual_path.exists():
|
||||
raise ValueError(f"文件不存在: {clean_virtual_path}")
|
||||
if not actual_path.is_file():
|
||||
raise ValueError(f"路径不是普通文件: {clean_virtual_path}")
|
||||
|
||||
return file_thread_id, uid, actual_path
|
||||
|
||||
|
||||
def _resolve_runtime_file_scope(runtime: ToolRuntime) -> tuple[str, str]:
|
||||
"""Read the thread and user scope needed for sandbox path mapping from ToolRuntime."""
|
||||
thread_id = _runtime_scope_value(runtime, "file_thread_id") or _runtime_scope_value(runtime, "thread_id")
|
||||
uid = _runtime_scope_value(runtime, "uid")
|
||||
if not thread_id:
|
||||
raise ValueError("当前运行时缺少 thread_id")
|
||||
if not uid:
|
||||
raise ValueError("当前运行时缺少 uid")
|
||||
return thread_id, uid
|
||||
|
||||
|
||||
def _runtime_scope_value(runtime: ToolRuntime, key: str) -> str | None:
|
||||
"""Look up a runtime scope value from LangGraph config, context, or state."""
|
||||
config = getattr(runtime, "config", None)
|
||||
configurable = config.get("configurable", {}) if isinstance(config, dict) else {}
|
||||
sources = (
|
||||
configurable if isinstance(configurable, dict) else {},
|
||||
getattr(runtime, "context", None),
|
||||
getattr(runtime, "state", None) if isinstance(getattr(runtime, "state", None), dict) else {},
|
||||
)
|
||||
for source in sources:
|
||||
value = source.get(key) if isinstance(source, dict) else getattr(source, key, None)
|
||||
if isinstance(value, str) and value.strip():
|
||||
return value.strip()
|
||||
return None
|
||||
|
||||
|
||||
def _resolve_ocr_engine(ocr_engine: str | None) -> str:
|
||||
"""Validate the requested OCR engine, falling back to the system default when omitted."""
|
||||
from yuxi import config
|
||||
from yuxi.knowledge.parser.factory import DocumentProcessorFactory
|
||||
|
||||
engine = str(ocr_engine or config.default_ocr_engine).strip() or config.default_ocr_engine
|
||||
allowed = {"disable", *DocumentProcessorFactory.get_available_processors()}
|
||||
if engine not in allowed:
|
||||
raise ValueError(f"不支持的 OCR 引擎: {engine}")
|
||||
return engine
|
||||
|
||||
|
||||
def _next_ocr_output_path(thread_id: str, source_path: Path) -> Path:
|
||||
"""Choose a non-conflicting Markdown output path under the thread outputs/ocr directory."""
|
||||
from yuxi.agents.backends.sandbox.paths import sandbox_outputs_dir
|
||||
|
||||
output_dir = sandbox_outputs_dir(thread_id) / _OCR_OUTPUT_DIR_NAME
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
base_name = _safe_ocr_output_stem(source_path)
|
||||
candidate = output_dir / f"{base_name}.md"
|
||||
index = 1
|
||||
while candidate.exists():
|
||||
candidate = output_dir / f"{base_name}-{index}.md"
|
||||
index += 1
|
||||
return candidate
|
||||
|
||||
|
||||
def _safe_ocr_output_stem(source_path: Path) -> str:
|
||||
"""Build a filesystem-friendly output filename stem from the source file name."""
|
||||
stem = source_path.stem.strip() or "ocr_result"
|
||||
safe_stem = _SAFE_OUTPUT_STEM_RE.sub("_", stem).strip("._-")
|
||||
return safe_stem or "ocr_result"
|
||||
|
||||
|
||||
def _ocr_preview(markdown: str) -> tuple[str, bool]:
|
||||
"""Return the short preview included in the tool result and whether it was truncated."""
|
||||
if len(markdown) <= _OCR_PREVIEW_LIMIT:
|
||||
return markdown, False
|
||||
return markdown[:_OCR_PREVIEW_LIMIT].rstrip(), True
|
||||
|
||||
|
||||
ASK_USER_QUESTION_DESCRIPTION = """
|
||||
在执行过程中,当你需要用户做决定或补充需求时,使用这个工具向用户提问。
|
||||
|
||||
适用场景:
|
||||
1. 收集用户偏好或需求(例如风格、范围、优先级)
|
||||
2. 澄清模糊指令(存在多种合理解释时)
|
||||
3. 在实现过程中让用户选择方案方向
|
||||
4. 在有明显权衡时让用户做取舍
|
||||
|
||||
使用规范:
|
||||
1. questions 提供 1-5 个问题,每项包含:question、options、multi_select、allow_other
|
||||
2. 每个问题的 options 提供 2-5 个有区分度的选项,每项包含 label 和 value
|
||||
3. 若有推荐选项:把推荐项放在第一位,并在 label 末尾加 "(Recommended)"
|
||||
4. 若需要多选:将该问题的 multi_select 设为 true
|
||||
5. allow_other 通常保持 true,用户可通过 Other 输入自定义答案
|
||||
|
||||
注意事项:
|
||||
1. 不要用这个工具询问“是否继续执行”“计划是否准备好”这类流程控制问题
|
||||
2. 不要在信息已充分、无需用户决策时滥用该工具
|
||||
3. 先基于现有上下文自行决策,只有关键不确定性时才提问
|
||||
|
||||
返回结果:
|
||||
answer 为 object,格式为 {question_id: answer}。
|
||||
其中 answer 可能是 string(单选)、list(多选)或 object(Other 文本)。
|
||||
"""
|
||||
|
||||
|
||||
@tool(
|
||||
category="buildin",
|
||||
tags=["交互"],
|
||||
display_name="向用户提问",
|
||||
description=ASK_USER_QUESTION_DESCRIPTION,
|
||||
)
|
||||
def ask_user_question(
|
||||
questions: Annotated[
|
||||
list[dict] | str | None,
|
||||
"问题列表,每项格式 {question, options, multi_select, allow_other, question_id(optional)}",
|
||||
] = None,
|
||||
) -> dict:
|
||||
"""向用户发起问题并等待回答。"""
|
||||
# 解析 questions 参数:如果是字符串,尝试解析为 JSON
|
||||
if isinstance(questions, str):
|
||||
try:
|
||||
import json
|
||||
|
||||
questions = json.loads(questions)
|
||||
logger.debug(f"Parsed string questions to list: {questions}")
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to parse questions string: {e}, using None")
|
||||
questions = None
|
||||
|
||||
normalized_questions = normalize_questions(questions or [])
|
||||
|
||||
if not normalized_questions:
|
||||
raise ValueError("questions 至少需要包含一个有效问题")
|
||||
|
||||
interrupt_payload = {
|
||||
"questions": normalized_questions,
|
||||
"source": "ask_user_question",
|
||||
}
|
||||
answer = interrupt(interrupt_payload)
|
||||
|
||||
return {
|
||||
"questions": normalized_questions,
|
||||
"answer": answer,
|
||||
}
|
||||
@@ -0,0 +1,3 @@
|
||||
# debug 工具包
|
||||
|
||||
__all__ = []
|
||||
@@ -0,0 +1,7 @@
|
||||
from .tools import (
|
||||
find_kb_document,
|
||||
get_common_kb_tools,
|
||||
open_kb_document,
|
||||
)
|
||||
|
||||
__all__ = ["find_kb_document", "get_common_kb_tools", "open_kb_document"]
|
||||
@@ -0,0 +1,473 @@
|
||||
"""知识库工具模块"""
|
||||
|
||||
import inspect
|
||||
from typing import Any
|
||||
|
||||
from langgraph.prebuilt.tool_node import ToolRuntime
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from yuxi.agents.toolkits.registry import tool
|
||||
from yuxi.knowledge.base import KnowledgeBase
|
||||
from yuxi.knowledge.schemas import (
|
||||
FindInputSchema,
|
||||
FindOutputSchema,
|
||||
OpenInputSchema,
|
||||
OpenOutputSchema,
|
||||
SearchInputSchema,
|
||||
SearchOutputSchema,
|
||||
)
|
||||
from yuxi.utils import logger
|
||||
|
||||
# ========== 通用知识库工具函数 ==========
|
||||
|
||||
|
||||
def _get_knowledge_base():
|
||||
from yuxi import knowledge_base
|
||||
|
||||
return knowledge_base
|
||||
|
||||
|
||||
class ListKBsInput(BaseModel):
|
||||
"""列出用户可访问的知识库输入模型"""
|
||||
|
||||
# Langchain 的 runtime 注入机制要求必须有参数
|
||||
dummy: str = Field(default="", description="Dummy parameter - ignore") # Add this
|
||||
|
||||
|
||||
@tool(category="knowledge", tags=["知识库"], args_schema=ListKBsInput)
|
||||
async def list_kbs(dummy: str, runtime: ToolRuntime) -> str: # Now has 2 params
|
||||
"""列出当前用户可访问的知识库列表
|
||||
|
||||
返回用户基于权限可访问的知识库名称列表。这个列表是根据用户的角色和部门信息过滤后的结果,
|
||||
但不包括用户在当前对话中未启用的知识库。
|
||||
|
||||
Returns:
|
||||
用户可访问的知识库名称列表(字符串格式)
|
||||
"""
|
||||
# 从 runtime.context 获取用户信息
|
||||
runtime_context = runtime.context
|
||||
uid = getattr(runtime_context, "uid", None)
|
||||
if not uid:
|
||||
return "无法获取用户信息"
|
||||
|
||||
# 打印 runtime—context 中的所有信息以进行调试
|
||||
logger.debug(f"Runtime context: {runtime_context.__dict__}")
|
||||
|
||||
enabled_kb_names = getattr(runtime_context, "knowledges", None)
|
||||
|
||||
try:
|
||||
from yuxi.agents.backends.knowledge_base_backend import resolve_visible_knowledge_bases_for_context
|
||||
|
||||
available_kbs = await resolve_visible_knowledge_bases_for_context(runtime_context)
|
||||
except Exception as e:
|
||||
logger.error(f"获取用户知识库列表失败: {e}")
|
||||
return f"获取知识库列表失败: {str(e)}"
|
||||
|
||||
all_kb_names = [kb["name"] for kb in available_kbs]
|
||||
|
||||
logger.debug(f"用户 {uid} 可访问的知识库列表: {all_kb_names}")
|
||||
logger.debug(f"用户 {uid} 当前对话启用的知识库列表: {enabled_kb_names}")
|
||||
|
||||
if not available_kbs:
|
||||
return "当前没有可访问的知识库"
|
||||
|
||||
# 格式化输出(包含名称和描述)
|
||||
kb_list = []
|
||||
for kb in available_kbs:
|
||||
name = kb.get("name", "")
|
||||
desc = kb.get("description") or "无描述"
|
||||
kb_list.append({"kb_id": kb.get("kb_id"), "name": name, "description": desc})
|
||||
|
||||
return kb_list
|
||||
|
||||
|
||||
class GetMindmapInput(BaseModel):
|
||||
"""获取思维导图输入模型"""
|
||||
|
||||
kb_name: str = Field(description="知识库名称,用于指定要获取思维导图的知识库")
|
||||
|
||||
|
||||
@tool(category="knowledge", tags=["知识库"], args_schema=GetMindmapInput)
|
||||
async def get_mindmap(kb_name: str, runtime: ToolRuntime) -> str:
|
||||
"""获取指定知识库的思维导图结构
|
||||
|
||||
当用户想要了解知识库的整体结构、文件分类、知识架构时使用此工具。
|
||||
返回知识库的思维导图层级结构。
|
||||
|
||||
Args:
|
||||
kb_name: 知识库名称
|
||||
|
||||
Returns:
|
||||
知识库的思维导图结构(文本格式)
|
||||
"""
|
||||
if not kb_name:
|
||||
return "请提供知识库名称"
|
||||
|
||||
# 获取所有检索器
|
||||
knowledge_base = _get_knowledge_base()
|
||||
retrievers = knowledge_base.get_retrievers()
|
||||
|
||||
# 查找对应的知识库
|
||||
target_kb_id = None
|
||||
target_info = None
|
||||
for kb_id, info in retrievers.items():
|
||||
if info["name"] == kb_name:
|
||||
target_kb_id = kb_id
|
||||
target_info = info
|
||||
break
|
||||
|
||||
if not target_kb_id:
|
||||
return f"知识库 '{kb_name}' 不存在"
|
||||
|
||||
try:
|
||||
from yuxi.repositories.knowledge_base_repository import KnowledgeBaseRepository
|
||||
|
||||
kb_repo = KnowledgeBaseRepository()
|
||||
kb = await kb_repo.get_by_kb_id(target_kb_id)
|
||||
|
||||
if kb is None:
|
||||
return f"知识库 {target_info['name']} 不存在"
|
||||
|
||||
mindmap_data = kb.mindmap
|
||||
|
||||
if not mindmap_data:
|
||||
return f"知识库 {target_info['name']} 还没有生成思维导图。"
|
||||
|
||||
# 将思维导图数据转换为文本格式
|
||||
def mindmap_to_text(node, level=0):
|
||||
"""递归将思维导图JSON转换为层级文本"""
|
||||
indent = " " * level
|
||||
text = f"{indent}- {node.get('content', '')}\n"
|
||||
for child in node.get("children", []):
|
||||
text += mindmap_to_text(child, level + 1)
|
||||
return text
|
||||
|
||||
mindmap_text = f"知识库 {target_info['name']} 的思维导图结构:\n\n"
|
||||
mindmap_text += mindmap_to_text(mindmap_data)
|
||||
|
||||
return mindmap_text
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"获取思维导图失败: {e}")
|
||||
return f"获取思维导图失败: {str(e)}"
|
||||
|
||||
|
||||
QueryKBInput = SearchInputSchema
|
||||
OpenKBDocumentInput = OpenInputSchema
|
||||
FindKBDocumentInput = FindInputSchema
|
||||
|
||||
|
||||
async def _resolve_visible_knowledge_bases_for_query(runtime: ToolRuntime | None) -> list[dict[str, Any]]:
|
||||
if runtime is None:
|
||||
return []
|
||||
|
||||
context = getattr(runtime, "context", None)
|
||||
if context is None:
|
||||
return []
|
||||
|
||||
visible_kbs = getattr(context, "_visible_knowledge_bases", None)
|
||||
if isinstance(visible_kbs, list):
|
||||
return visible_kbs
|
||||
|
||||
try:
|
||||
from yuxi.agents.backends.knowledge_base_backend import resolve_visible_knowledge_bases_for_context
|
||||
|
||||
return await resolve_visible_knowledge_bases_for_context(context)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
logger.warning(f"解析会话可见知识库失败: {exc}")
|
||||
return []
|
||||
|
||||
|
||||
def _find_query_target(
|
||||
*,
|
||||
kb_id: str,
|
||||
retrievers: dict[str, Any],
|
||||
visible_kbs: list[dict[str, Any]],
|
||||
) -> tuple[dict[str, Any] | None, str | None, str | None]:
|
||||
if not visible_kbs:
|
||||
return None, None, "无法获取当前会话可访问的知识库"
|
||||
|
||||
normalized_kb_id = str(kb_id or "").strip()
|
||||
visible_kb_ids = {str(kb.get("kb_id") or "").strip() for kb in visible_kbs}
|
||||
if normalized_kb_id not in visible_kb_ids:
|
||||
return None, None, f"知识库资源 '{normalized_kb_id}' 不存在或当前会话未启用"
|
||||
|
||||
target_info = retrievers.get(normalized_kb_id)
|
||||
if target_info is None:
|
||||
return None, None, f"知识库资源 '{normalized_kb_id}' 不存在"
|
||||
return target_info, normalized_kb_id, None
|
||||
|
||||
|
||||
async def _build_query_output(target_kb_id: str, result: Any) -> Any:
|
||||
if isinstance(result, dict) and result.get("kb_id") == target_kb_id and isinstance(result.get("results"), list):
|
||||
return SearchOutputSchema(**result).model_dump()
|
||||
return KnowledgeBase.build_search_output(target_kb_id, result)
|
||||
|
||||
|
||||
@tool(category="knowledge", tags=["知识库"], args_schema=QueryKBInput)
|
||||
async def query_kb(kb_id: str, query_text: str, file_name: str | None = None, runtime: ToolRuntime = None) -> Any:
|
||||
"""在指定知识库中检索内容
|
||||
|
||||
当用户需要查询具体内容时使用此工具。kb_id 是知识库资源 ID,也就是 kb_id;返回结果中的
|
||||
file_id 可继续用于 find_kb_document 或 open_kb_document。
|
||||
"""
|
||||
if not kb_id:
|
||||
return "请提供 kb_id"
|
||||
if not query_text:
|
||||
return "请提供查询内容"
|
||||
|
||||
knowledge_base = _get_knowledge_base()
|
||||
retrievers = knowledge_base.get_retrievers()
|
||||
visible_kbs = await _resolve_visible_knowledge_bases_for_query(runtime)
|
||||
target_info, target_kb_id, target_error = _find_query_target(
|
||||
kb_id=kb_id,
|
||||
retrievers=retrievers,
|
||||
visible_kbs=visible_kbs,
|
||||
)
|
||||
if target_error:
|
||||
return target_error
|
||||
|
||||
try:
|
||||
retriever = target_info["retriever"]
|
||||
kwargs = {}
|
||||
if file_name:
|
||||
kwargs["file_name"] = file_name
|
||||
|
||||
if inspect.iscoroutinefunction(retriever):
|
||||
result = await retriever(query_text, **kwargs)
|
||||
else:
|
||||
result = retriever(query_text, **kwargs)
|
||||
|
||||
return await _build_query_output(target_kb_id, result)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"检索失败: {e}")
|
||||
return f"检索失败: {str(e)}"
|
||||
|
||||
|
||||
@tool(category="knowledge", tags=["知识库"], args_schema=OpenKBDocumentInput)
|
||||
async def open_kb_document(
|
||||
kb_id: str,
|
||||
file_id: str,
|
||||
line: int | None = None,
|
||||
offset: int | None = None,
|
||||
window_size: int = 1800,
|
||||
runtime: ToolRuntime = None,
|
||||
) -> dict[str, Any] | str:
|
||||
"""按行窗口打开知识库文档原文
|
||||
|
||||
当 query_kb 返回的片段不足以回答问题,或需要查看某个文档的上下文时使用。
|
||||
kb_id 是知识库资源 ID,也就是 kb_id;file_id 是知识库文件 ID。
|
||||
"""
|
||||
normalized_kb_id = str(kb_id or "").strip()
|
||||
normalized_file_id = str(file_id or "").strip()
|
||||
if not normalized_kb_id:
|
||||
return "请提供 kb_id"
|
||||
if not normalized_file_id:
|
||||
return "请提供 file_id"
|
||||
|
||||
visible_kbs = await _resolve_visible_knowledge_bases_for_query(runtime)
|
||||
if not visible_kbs:
|
||||
return "无法获取当前会话可访问的知识库"
|
||||
|
||||
visible_kb_ids = {str(kb.get("kb_id") or "").strip() for kb in visible_kbs}
|
||||
if normalized_kb_id not in visible_kb_ids:
|
||||
return f"知识库资源 '{normalized_kb_id}' 不存在或当前会话未启用"
|
||||
|
||||
knowledge_base = _get_knowledge_base()
|
||||
retrievers = knowledge_base.get_retrievers()
|
||||
target_info = retrievers.get(normalized_kb_id)
|
||||
if target_info is None:
|
||||
return f"知识库资源 '{normalized_kb_id}' 不存在"
|
||||
|
||||
metadata = target_info.get("metadata") if isinstance(target_info, dict) else None
|
||||
kb_type = str((metadata or {}).get("kb_type") or "").strip().lower()
|
||||
if kb_type == "dify":
|
||||
return "Dify 知识库为外部只读检索源,当前不支持通过 Open 打开全文"
|
||||
|
||||
try:
|
||||
start_offset = int(line) - 1 if line is not None else int(offset or 0)
|
||||
window = await knowledge_base.open_file_content(
|
||||
normalized_kb_id,
|
||||
normalized_file_id,
|
||||
offset=start_offset,
|
||||
limit=window_size,
|
||||
)
|
||||
return OpenOutputSchema(kb_id=normalized_kb_id, file_id=normalized_file_id, **window).model_dump()
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"打开知识库文档失败: {e}")
|
||||
return f"打开知识库文档失败: {str(e)}"
|
||||
|
||||
|
||||
@tool(category="knowledge", tags=["知识库"], args_schema=FindKBDocumentInput)
|
||||
async def find_kb_document(
|
||||
kb_id: str,
|
||||
file_id: str,
|
||||
patterns: list[str],
|
||||
use_regex: bool = False,
|
||||
case_sensitive: bool = False,
|
||||
max_windows: int = 5,
|
||||
window_size: int = 80,
|
||||
runtime: ToolRuntime = None,
|
||||
) -> dict[str, Any] | str:
|
||||
"""在已知知识库文件内做关键词或正则定位。
|
||||
|
||||
当 query_kb 已找到候选文件,但需要在该文件内定位术语、指标、章节或实体时使用。
|
||||
"""
|
||||
normalized_kb_id = str(kb_id or "").strip()
|
||||
normalized_file_id = str(file_id or "").strip()
|
||||
if not normalized_kb_id:
|
||||
return "请提供 kb_id"
|
||||
if not normalized_file_id:
|
||||
return "请提供 file_id"
|
||||
if not patterns:
|
||||
return "请提供 patterns"
|
||||
|
||||
visible_kbs = await _resolve_visible_knowledge_bases_for_query(runtime)
|
||||
if not visible_kbs:
|
||||
return "无法获取当前会话可访问的知识库"
|
||||
|
||||
visible_kb_ids = {str(kb.get("kb_id") or "").strip() for kb in visible_kbs}
|
||||
if normalized_kb_id not in visible_kb_ids:
|
||||
return f"知识库资源 '{normalized_kb_id}' 不存在或当前会话未启用"
|
||||
|
||||
knowledge_base = _get_knowledge_base()
|
||||
retrievers = knowledge_base.get_retrievers()
|
||||
target_info = retrievers.get(normalized_kb_id)
|
||||
if target_info is None:
|
||||
return f"知识库资源 '{normalized_kb_id}' 不存在"
|
||||
|
||||
metadata = target_info.get("metadata") if isinstance(target_info, dict) else None
|
||||
kb_type = str((metadata or {}).get("kb_type") or "").strip().lower()
|
||||
if kb_type == "dify":
|
||||
return "Dify 知识库为外部只读检索源,当前不支持通过 Find 检索全文"
|
||||
|
||||
try:
|
||||
result = await knowledge_base.find_file_content(
|
||||
normalized_kb_id,
|
||||
normalized_file_id,
|
||||
patterns,
|
||||
use_regex=use_regex,
|
||||
case_sensitive=case_sensitive,
|
||||
max_windows=max_windows,
|
||||
window_size=window_size,
|
||||
)
|
||||
return FindOutputSchema(kb_id=normalized_kb_id, file_id=normalized_file_id, **result).model_dump()
|
||||
except Exception as e:
|
||||
logger.error(f"知识库文档内检索失败: {e}")
|
||||
return f"知识库文档内检索失败: {str(e)}"
|
||||
|
||||
|
||||
# 单个知识库一次最多扫描的文件数(与仓储层 list_by_kb_id_after 的硬上限保持一致),
|
||||
# 用于在内存中做文件名过滤与准确计数,避免按 limit+offset 截断导致 total/has_more 失真。
|
||||
_KB_FILE_SCAN_LIMIT = 5000
|
||||
|
||||
|
||||
class SearchFileInput(BaseModel):
|
||||
"""搜索文件输入模型"""
|
||||
|
||||
kb_name: str | None = Field(default=None, description="知识库名称,为空时搜索所有知识库")
|
||||
query: str | None = Field(default=None, description="搜索关键词,为空时返回所有文件")
|
||||
offset: int = Field(default=0, ge=0, description="偏移量,从 0 开始")
|
||||
limit: int = Field(default=300, ge=1, le=5000, description="返回数量限制,默认 300")
|
||||
|
||||
|
||||
@tool(category="knowledge", tags=["知识库"], args_schema=SearchFileInput)
|
||||
async def search_file(
|
||||
kb_name: str | None = None,
|
||||
query: str | None = None,
|
||||
offset: int = 0,
|
||||
limit: int = 300,
|
||||
runtime: ToolRuntime = None,
|
||||
) -> dict[str, Any] | str:
|
||||
"""搜索知识库中的文件
|
||||
|
||||
当用户需要查找特定文件时使用此工具。可以指定知识库名称和搜索关键词。
|
||||
如果不指定知识库,将搜索所有可访问的知识库。
|
||||
如果不指定搜索关键词,将返回所有文件。
|
||||
|
||||
Args:
|
||||
kb_name: 知识库名称,为空时搜索所有知识库
|
||||
query: 搜索关键词,为空时返回所有文件
|
||||
offset: 偏移量,从 0 开始
|
||||
limit: 返回数量限制,默认 300
|
||||
|
||||
Returns:
|
||||
匹配的文件列表和分页信息
|
||||
"""
|
||||
if not kb_name and not query:
|
||||
return "请提供知识库名称或搜索关键词,不能同时为空"
|
||||
|
||||
visible_kbs = await _resolve_visible_knowledge_bases_for_query(runtime)
|
||||
if not visible_kbs:
|
||||
return "无法获取当前会话可访问的知识库"
|
||||
|
||||
if kb_name:
|
||||
target_kbs = [kb for kb in visible_kbs if kb.get("name") == kb_name]
|
||||
if not target_kbs:
|
||||
return f"知识库 '{kb_name}' 不存在或当前会话未启用"
|
||||
else:
|
||||
target_kbs = visible_kbs
|
||||
|
||||
from yuxi.repositories.knowledge_file_repository import KnowledgeFileRepository
|
||||
|
||||
repo = KnowledgeFileRepository()
|
||||
|
||||
all_files = []
|
||||
for kb in target_kbs:
|
||||
kb_id = kb.get("kb_id")
|
||||
if not kb_id:
|
||||
continue
|
||||
|
||||
files = await repo.list_by_kb_id_after(
|
||||
kb_id=kb_id,
|
||||
limit=_KB_FILE_SCAN_LIMIT,
|
||||
files_only=True,
|
||||
)
|
||||
|
||||
if query:
|
||||
query_lower = query.lower()
|
||||
files = [f for f in files if query_lower in f.filename.lower()]
|
||||
|
||||
for f in files:
|
||||
all_files.append(
|
||||
{
|
||||
"kb_id": kb_id,
|
||||
"kb_name": kb.get("name"),
|
||||
"file_id": f.file_id,
|
||||
"filename": f.filename,
|
||||
"file_type": f.file_type,
|
||||
"status": f.status,
|
||||
"created_at": str(f.created_at) if f.created_at else None,
|
||||
"updated_at": str(f.updated_at) if f.updated_at else None,
|
||||
"file_size": f.file_size,
|
||||
}
|
||||
)
|
||||
|
||||
all_files.sort(key=lambda x: x.get("updated_at") or "", reverse=True)
|
||||
|
||||
total = len(all_files)
|
||||
paginated_files = all_files[offset : offset + limit]
|
||||
|
||||
return {
|
||||
"files": paginated_files,
|
||||
"total": total,
|
||||
"offset": offset,
|
||||
"limit": limit,
|
||||
"has_more": offset + limit < total,
|
||||
}
|
||||
|
||||
|
||||
def get_common_kb_tools() -> list:
|
||||
"""获取通用知识库工具列表
|
||||
|
||||
返回 6 个通用工具:
|
||||
- list_kbs: 列出用户可访问的知识库
|
||||
- get_mindmap: 获取指定知识库的思维导图
|
||||
- query_kb: 在指定知识库中检索
|
||||
- find_kb_document: 在指定文件内定位关键词或正则模式
|
||||
- open_kb_document: 按 file_id 分段打开知识库文档
|
||||
- search_file: 搜索知识库中的文件
|
||||
"""
|
||||
return [list_kbs, get_mindmap, query_kb, find_kb_document, open_kb_document, search_file]
|
||||
@@ -0,0 +1,89 @@
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
|
||||
@dataclass
|
||||
class ToolExtraMetadata:
|
||||
"""附加元数据(用装饰器注册)"""
|
||||
|
||||
category: str = "" # 分类: buildin, knowledge, mysql, subagents, debug
|
||||
tags: list[str] = field(default_factory=list)
|
||||
display_name: str = "" # 显示名称(给人看的名字)
|
||||
icon: str = ""
|
||||
config_guide: str = "" # 配置说明(给人看的使用前配置提示)
|
||||
|
||||
|
||||
# 全局注册表: tool_name -> ToolExtraMetadata
|
||||
_extra_registry: dict[str, ToolExtraMetadata] = {}
|
||||
|
||||
# 全局工具实例列表(由 @tool 装饰器自动收集)
|
||||
_all_tool_instances: list = []
|
||||
|
||||
|
||||
def get_extra_metadata(tool_name: str) -> ToolExtraMetadata | None:
|
||||
"""获取工具附加元数据"""
|
||||
return _extra_registry.get(tool_name)
|
||||
|
||||
|
||||
def get_all_extra_metadata() -> dict[str, ToolExtraMetadata]:
|
||||
"""获取所有附加元数据"""
|
||||
return _extra_registry.copy()
|
||||
|
||||
|
||||
def get_all_tool_instances() -> list:
|
||||
"""获取所有工具实例(由 @tool 装饰器自动收集)"""
|
||||
return _all_tool_instances
|
||||
|
||||
|
||||
# 基于 langchain.tool 的拓展装饰器
|
||||
def tool(
|
||||
category: str = "",
|
||||
tags: list[str] = None,
|
||||
display_name: str = "",
|
||||
icon: str = "",
|
||||
config_guide: str = "",
|
||||
name_or_callable: str | Callable | None = None,
|
||||
description: str | None = None,
|
||||
args_schema: type | None = None,
|
||||
return_direct: bool = False,
|
||||
):
|
||||
"""基于 langchain.tool 的拓展装饰器,同时注册元数据
|
||||
|
||||
使用方式:
|
||||
@tool(category="buildin", tags=["计算"], display_name="计算器")
|
||||
def calculator(a: float, b: float, operation: str) -> float:
|
||||
...
|
||||
|
||||
或者保留原有的 name_or_callable 和 description 参数来自定义工具名称和说明。
|
||||
"""
|
||||
from langchain.tools import tool as langchain_tool
|
||||
|
||||
# 先应用 langchain tool 装饰器
|
||||
langchain_decorator = langchain_tool(
|
||||
name_or_callable=name_or_callable,
|
||||
description=description,
|
||||
args_schema=args_schema,
|
||||
return_direct=return_direct,
|
||||
)
|
||||
|
||||
def decorator(func: Callable) -> Callable:
|
||||
# 应用 langchain 装饰器
|
||||
tool_obj = langchain_decorator(func)
|
||||
|
||||
# 注册附加元数据
|
||||
tool_name = tool_obj.name
|
||||
_extra_registry[tool_name] = ToolExtraMetadata(
|
||||
category=category,
|
||||
tags=tags or [],
|
||||
display_name=display_name,
|
||||
icon=icon,
|
||||
config_guide=config_guide,
|
||||
)
|
||||
|
||||
# 自动收集工具实例
|
||||
tool_obj.handle_tool_error = True
|
||||
_all_tool_instances.append(tool_obj)
|
||||
|
||||
return tool_obj
|
||||
|
||||
return decorator
|
||||
@@ -0,0 +1,144 @@
|
||||
from typing import Any
|
||||
|
||||
from yuxi.utils import logger
|
||||
|
||||
# 工具元数据缓存
|
||||
_metadata_cache: list[dict] = []
|
||||
|
||||
|
||||
def _extract_tool_info(tool_obj) -> dict:
|
||||
"""从 tool_obj 提取基础信息"""
|
||||
metadata = getattr(tool_obj, "metadata", {}) or {}
|
||||
info = {
|
||||
"slug": tool_obj.name,
|
||||
"name": metadata.get("name", tool_obj.name), # 显示名称优先从 metadata 获取
|
||||
"description": tool_obj.description,
|
||||
"metadata": metadata,
|
||||
"args": [],
|
||||
}
|
||||
|
||||
if hasattr(tool_obj, "args_schema") and tool_obj.args_schema:
|
||||
schema = tool_obj.args_schema
|
||||
if hasattr(schema, "schema"):
|
||||
schema = schema.schema()
|
||||
for arg_name, arg_info in schema.get("properties", {}).items():
|
||||
info["args"].append(
|
||||
{
|
||||
"name": arg_name,
|
||||
"type": arg_info.get("type", ""),
|
||||
"description": arg_info.get("description", ""),
|
||||
}
|
||||
)
|
||||
return info
|
||||
|
||||
|
||||
def _ensure_metadata_loaded():
|
||||
"""延迟加载工具元数据(首次调用时自动触发)"""
|
||||
global _metadata_cache
|
||||
|
||||
if _metadata_cache: # 已加载
|
||||
return
|
||||
|
||||
from yuxi.agents.toolkits.registry import (
|
||||
get_all_extra_metadata,
|
||||
get_all_tool_instances,
|
||||
)
|
||||
|
||||
# 获取所有工具实例
|
||||
all_tools = get_all_tool_instances()
|
||||
extra_meta = get_all_extra_metadata()
|
||||
|
||||
for tool in all_tools:
|
||||
tool_name = tool.name
|
||||
runtime_info = _extract_tool_info(tool)
|
||||
|
||||
# 合并附加元数据
|
||||
if tool_name in extra_meta:
|
||||
extra = extra_meta[tool_name]
|
||||
runtime_info["category"] = extra.category
|
||||
runtime_info["tags"] = extra.tags
|
||||
runtime_info["config_guide"] = extra.config_guide
|
||||
# display_name 优先级高于 tool.name
|
||||
if extra.display_name:
|
||||
runtime_info["name"] = extra.display_name
|
||||
else:
|
||||
# 未注册,设为默认分类
|
||||
runtime_info["category"] = "buildin"
|
||||
runtime_info["tags"] = []
|
||||
runtime_info["config_guide"] = ""
|
||||
|
||||
_metadata_cache.append(runtime_info)
|
||||
|
||||
logger.info(f"Tool service loaded {len(_metadata_cache)} tools (lazy load)")
|
||||
|
||||
|
||||
def get_tool_metadata(category: str = None) -> list[dict]:
|
||||
"""获取工具元数据列表(延迟加载)"""
|
||||
_ensure_metadata_loaded()
|
||||
|
||||
if category:
|
||||
return [t for t in _metadata_cache if t.get("category") == category]
|
||||
return _metadata_cache
|
||||
|
||||
|
||||
def get_tool_instances_by_category(category: str) -> list[Any]:
|
||||
from yuxi.agents.toolkits.registry import get_all_extra_metadata, get_all_tool_instances
|
||||
|
||||
extra_meta = get_all_extra_metadata()
|
||||
tools = []
|
||||
for tool in get_all_tool_instances():
|
||||
tool_meta = extra_meta.get(tool.name)
|
||||
tool_category = tool_meta.category if tool_meta else "buildin"
|
||||
if tool_category == category:
|
||||
tools.append(tool)
|
||||
return tools
|
||||
|
||||
|
||||
async def resolve_configured_runtime_tools(context) -> list[Any]:
|
||||
from yuxi.agents.mcp.service import get_enabled_mcp_tools
|
||||
|
||||
selected_tools = []
|
||||
selected_tool_names: set[str] = set()
|
||||
buildin_tools = {tool.name: tool for tool in get_tool_instances_by_category("buildin")}
|
||||
|
||||
for tool_name in getattr(context, "tools", None) or []:
|
||||
if not isinstance(tool_name, str) or tool_name in selected_tool_names:
|
||||
continue
|
||||
tool = buildin_tools.get(tool_name)
|
||||
if tool is None:
|
||||
logger.warning(f"Configured buildin tool not found, skip: {tool_name}")
|
||||
continue
|
||||
selected_tools.append(tool)
|
||||
selected_tool_names.add(tool_name)
|
||||
|
||||
selected_mcp_servers: set[str] = set()
|
||||
for server_name in getattr(context, "mcps", None) or []:
|
||||
if not isinstance(server_name, str) or server_name in selected_mcp_servers:
|
||||
continue
|
||||
selected_mcp_servers.add(server_name)
|
||||
try:
|
||||
mcp_tools = await get_enabled_mcp_tools(server_name)
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to load configured MCP tools '{server_name}': {e}")
|
||||
continue
|
||||
if not mcp_tools:
|
||||
logger.warning(f"Configured MCP unavailable, skip: {server_name}")
|
||||
continue
|
||||
for tool in mcp_tools:
|
||||
if tool.name in selected_tool_names:
|
||||
continue
|
||||
selected_tools.append(tool)
|
||||
selected_tool_names.add(tool.name)
|
||||
|
||||
# Skill 依赖的本地工具:必须随基础工具一起注册进 create_agent 的 ToolNode 才可执行,
|
||||
# 否则 Skill 激活后模型虽能发起调用,执行器仍报 "not a valid tool"。
|
||||
# 默认绑定给模型的可见性由 SkillsMiddleware 按 Skill 激活状态门控(保持按需加载)。
|
||||
from yuxi.agents.middlewares.skills import resolve_skill_gated_tools
|
||||
|
||||
for tool in resolve_skill_gated_tools(context):
|
||||
if tool.name in selected_tool_names:
|
||||
continue
|
||||
selected_tools.append(tool)
|
||||
selected_tool_names.add(tool.name)
|
||||
|
||||
return selected_tools
|
||||
@@ -0,0 +1,55 @@
|
||||
import traceback
|
||||
from typing import Any
|
||||
|
||||
from yuxi.utils import logger
|
||||
|
||||
|
||||
def get_tool_info(tools) -> list[dict[str, Any]]:
|
||||
"""获取所有工具的信息(用于前端展示)"""
|
||||
tools_info = []
|
||||
|
||||
try:
|
||||
# 获取注册的工具信息
|
||||
for tool_obj in tools:
|
||||
try:
|
||||
metadata = getattr(tool_obj, "metadata", {}) or {}
|
||||
info = {
|
||||
"id": tool_obj.name,
|
||||
"name": metadata.get("name", tool_obj.name),
|
||||
"description": tool_obj.description,
|
||||
"metadata": metadata,
|
||||
"args": [],
|
||||
# "is_async": is_async # Include async information
|
||||
}
|
||||
|
||||
if hasattr(tool_obj, "args_schema") and tool_obj.args_schema:
|
||||
if isinstance(tool_obj.args_schema, dict):
|
||||
schema = tool_obj.args_schema
|
||||
else:
|
||||
schema = tool_obj.args_schema.schema()
|
||||
|
||||
for arg_name, arg_info in schema.get("properties", {}).items():
|
||||
info["args"].append(
|
||||
{
|
||||
"name": arg_name,
|
||||
"type": arg_info.get("type", ""),
|
||||
"description": arg_info.get("description", ""),
|
||||
}
|
||||
)
|
||||
|
||||
tools_info.append(info)
|
||||
# logger.debug(f"Successfully processed tool info for {tool_obj.name}")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"Failed to process tool {getattr(tool_obj, 'name', 'unknown')}: {e}\n{traceback.format_exc()}. "
|
||||
f"Details: {dict(tool_obj.__dict__)}"
|
||||
)
|
||||
continue
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to get tools info: {e}\n{traceback.format_exc()}")
|
||||
return []
|
||||
|
||||
logger.info(f"Successfully extracted info for {len(tools_info)} tools")
|
||||
return tools_info
|
||||
@@ -0,0 +1,4 @@
|
||||
from .app import config
|
||||
from .user import UserConfig, UserConfigSchema
|
||||
|
||||
__all__ = ["UserConfig", "UserConfigSchema", "config"]
|
||||
@@ -0,0 +1,211 @@
|
||||
"""应用配置模块。"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import tomli
|
||||
import tomli_w
|
||||
from pydantic import BaseModel, Field, PrivateAttr
|
||||
|
||||
from yuxi.config import cache as runtime_cache
|
||||
from yuxi.utils.logging_config import logger
|
||||
|
||||
READONLY_CONFIG_FIELDS = frozenset({"save_dir"})
|
||||
DEFAULT_OCR_ENGINE = "rapid_ocr"
|
||||
|
||||
|
||||
def _get_available_ocr_engines() -> set[str]:
|
||||
from yuxi.knowledge.parser.factory import DocumentProcessorFactory
|
||||
|
||||
return {"disable", *DocumentProcessorFactory.get_available_processors()}
|
||||
|
||||
|
||||
def _normalize_default_ocr_engine(value: Any) -> str:
|
||||
engine = str(value or "").strip() or DEFAULT_OCR_ENGINE
|
||||
if engine not in _get_available_ocr_engines():
|
||||
raise ValueError(f"不支持的默认 OCR 引擎: {engine}")
|
||||
return engine
|
||||
|
||||
|
||||
class Config(BaseModel):
|
||||
"""应用配置类。
|
||||
|
||||
`save_dir` 只在启动时决定配置文件位置,运行时不可修改。管理员保存配置时先写
|
||||
`base.toml`,再把可运行时同步的字段写入 Redis 快照(`yuxi:runtime_config`)。
|
||||
其他进程通过 `start_runtime_sync()` 启动的后台线程周期性拉取该快照刷新内存值。
|
||||
"""
|
||||
|
||||
save_dir: str = Field(default="saves", description="保存目录", exclude=True)
|
||||
enable_content_guard: bool = Field(default=False, description="是否启用内容审查")
|
||||
enable_content_guard_llm: bool = Field(default=False, description="是否启用LLM内容审查")
|
||||
default_model: str = Field(
|
||||
default="siliconflow-cn:Pro/MiniMaxAI/MiniMax-M2.5",
|
||||
description="默认对话模型",
|
||||
)
|
||||
fast_model: str = Field(
|
||||
default="siliconflow-cn:Pro/MiniMaxAI/MiniMax-M2.5",
|
||||
description="快速响应模型",
|
||||
)
|
||||
embed_model: str = Field(
|
||||
default="siliconflow-cn:Pro/BAAI/bge-m3",
|
||||
description="默认 Embedding 模型",
|
||||
)
|
||||
reranker: str = Field(
|
||||
default="siliconflow-cn:Pro/BAAI/bge-reranker-v2-m3",
|
||||
description="默认 Re-Ranker 模型",
|
||||
)
|
||||
content_guard_llm_model: str = Field(
|
||||
default="siliconflow-cn:Pro/MiniMaxAI/MiniMax-M2.5",
|
||||
description="内容审查LLM模型",
|
||||
)
|
||||
default_ocr_engine: str = Field(default=DEFAULT_OCR_ENGINE, description="默认 OCR 解析引擎")
|
||||
|
||||
sandbox_provider: str = Field(default="provisioner", description="沙箱提供者")
|
||||
sandbox_provisioner_url: str = Field(default="http://sandbox-provisioner:8002", description="沙箱服务地址")
|
||||
sandbox_virtual_path_prefix: str = Field(default="/home/gem/user-data", description="沙箱用户目录前缀")
|
||||
sandbox_exec_timeout_seconds: int = Field(default=180, description="沙箱执行超时时间(秒)")
|
||||
sandbox_max_output_bytes: int = Field(default=262144, description="沙箱最大输出字节数")
|
||||
sandbox_keepalive_interval_seconds: int = Field(default=30, description="沙箱保活间隔")
|
||||
|
||||
_config_file: Path | None = PrivateAttr(default=None)
|
||||
_runtime_sync_thread: Any = PrivateAttr(default=None)
|
||||
|
||||
model_config = {"arbitrary_types_allowed": True, "extra": "allow"}
|
||||
|
||||
def __init__(self, **data):
|
||||
super().__init__(**data)
|
||||
self._setup_paths()
|
||||
self._load_user_config()
|
||||
self._handle_environment()
|
||||
|
||||
def _setup_paths(self) -> None:
|
||||
self._config_file = Path(self.save_dir) / "config" / "base.toml"
|
||||
self._config_file.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
def _load_user_config(self) -> None:
|
||||
if not self._config_file or not self._config_file.exists():
|
||||
logger.info(f"Config file not found, using defaults: {self._config_file}")
|
||||
return
|
||||
|
||||
logger.info(f"Loading config from {self._config_file}")
|
||||
try:
|
||||
with open(self._config_file, "rb") as f:
|
||||
user_config = tomli.load(f)
|
||||
|
||||
for key, value in user_config.items():
|
||||
if key in READONLY_CONFIG_FIELDS:
|
||||
logger.warning(f"Readonly config key ignored: {key}")
|
||||
elif key in type(self).model_fields:
|
||||
try:
|
||||
setattr(self, key, self._normalize_config_value(key, value))
|
||||
except ValueError as exc:
|
||||
logger.warning(f"Invalid config key ignored: {key} ({exc})")
|
||||
else:
|
||||
logger.warning(f"Unknown config key: {key}")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to load config from {self._config_file}: {e}")
|
||||
|
||||
def _handle_environment(self) -> None:
|
||||
self.sandbox_provider = (os.getenv("SANDBOX_PROVIDER") or self.sandbox_provider or "provisioner").strip()
|
||||
self.sandbox_provisioner_url = (
|
||||
os.getenv("SANDBOX_PROVISIONER_URL") or self.sandbox_provisioner_url or "http://sandbox-provisioner:8002"
|
||||
).strip()
|
||||
self.sandbox_virtual_path_prefix = (
|
||||
os.getenv("SANDBOX_VIRTUAL_PATH_PREFIX") or self.sandbox_virtual_path_prefix or "/home/gem/user-data"
|
||||
).strip()
|
||||
self.sandbox_exec_timeout_seconds = int(
|
||||
os.getenv("SANDBOX_EXEC_TIMEOUT_SECONDS") or self.sandbox_exec_timeout_seconds or 180
|
||||
)
|
||||
self.sandbox_max_output_bytes = int(
|
||||
os.getenv("SANDBOX_MAX_OUTPUT_BYTES") or self.sandbox_max_output_bytes or 262144
|
||||
)
|
||||
self.sandbox_keepalive_interval_seconds = int(
|
||||
os.getenv("SANDBOX_KEEPALIVE_INTERVAL_SECONDS") or self.sandbox_keepalive_interval_seconds or 30
|
||||
)
|
||||
|
||||
if self.sandbox_provider.lower() != "provisioner":
|
||||
raise ValueError("Only sandbox_provider=provisioner is supported.")
|
||||
if not self.sandbox_provisioner_url:
|
||||
raise ValueError("SANDBOX_PROVISIONER_URL is required when sandbox provider is provisioner.")
|
||||
if not self.sandbox_virtual_path_prefix.startswith("/"):
|
||||
self.sandbox_virtual_path_prefix = f"/{self.sandbox_virtual_path_prefix}"
|
||||
|
||||
def start_runtime_sync(self, interval: float = runtime_cache.RUNTIME_CONFIG_SYNC_INTERVAL_SECONDS) -> None:
|
||||
"""启动后台线程周期性从 Redis 同步运行时配置。多次调用仅启动一次。"""
|
||||
self._runtime_sync_thread = runtime_cache.start_runtime_sync(
|
||||
self,
|
||||
self._runtime_sync_thread,
|
||||
interval=interval,
|
||||
)
|
||||
|
||||
def refresh(self) -> None:
|
||||
"""从 Redis 快照刷新公开配置字段到内存;Redis 不可用或无快照时保持当前值。"""
|
||||
runtime_cache.refresh_runtime_config(self)
|
||||
|
||||
def save(self) -> None:
|
||||
if not self._config_file:
|
||||
logger.warning("Config file path not set")
|
||||
return
|
||||
|
||||
logger.info(f"Saving config to {self._config_file}")
|
||||
user_modified = {}
|
||||
for field_name, field_info in type(self).model_fields.items():
|
||||
if field_info.exclude:
|
||||
continue
|
||||
current_value = getattr(self, field_name)
|
||||
if current_value != field_info.default:
|
||||
user_modified[field_name] = current_value
|
||||
|
||||
try:
|
||||
with open(self._config_file, "wb") as f:
|
||||
tomli_w.dump(user_modified, f)
|
||||
logger.info(f"Config saved to {self._config_file}")
|
||||
runtime_cache.save_runtime_config(self)
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to save config to {self._config_file}: {e}")
|
||||
|
||||
def dump_config(self) -> dict[str, Any]:
|
||||
config_dict = self.model_dump()
|
||||
fields_info = {}
|
||||
for field_name, field_info in Config.model_fields.items():
|
||||
if field_info.exclude:
|
||||
continue
|
||||
fields_info[field_name] = {
|
||||
"des": field_info.description,
|
||||
"default": field_info.default,
|
||||
"type": field_info.annotation.__name__
|
||||
if hasattr(field_info.annotation, "__name__")
|
||||
else str(field_info.annotation),
|
||||
"exclude": field_info.exclude if hasattr(field_info, "exclude") else False,
|
||||
}
|
||||
config_dict["_config_items"] = fields_info
|
||||
return config_dict
|
||||
|
||||
def update(self, other: dict[str, Any]) -> None:
|
||||
for key, value in other.items():
|
||||
if self.can_update(key):
|
||||
self.set_value(key, value)
|
||||
elif key in READONLY_CONFIG_FIELDS:
|
||||
logger.warning(f"Readonly config key ignored: {key}")
|
||||
else:
|
||||
logger.warning(f"Unknown config key: {key}")
|
||||
|
||||
def can_update(self, key: object) -> bool:
|
||||
return isinstance(key, str) and key in type(self).model_fields and key not in READONLY_CONFIG_FIELDS
|
||||
|
||||
def set_value(self, key: str, value: Any) -> None:
|
||||
if not self.can_update(key):
|
||||
raise ValueError(f"配置项不可修改: {key}")
|
||||
setattr(self, key, self._normalize_config_value(key, value))
|
||||
|
||||
def _normalize_config_value(self, key: str, value: Any) -> Any:
|
||||
if key == "default_ocr_engine":
|
||||
return _normalize_default_ocr_engine(value)
|
||||
return value
|
||||
|
||||
|
||||
config = Config()
|
||||
@@ -0,0 +1,99 @@
|
||||
"""运行时配置 Redis 快照同步。"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import threading
|
||||
import time
|
||||
from collections.abc import Iterator
|
||||
from typing import Any
|
||||
|
||||
from yuxi.storage.redis import RedisConfig, sync_redis_client
|
||||
from yuxi.utils.logging_config import logger
|
||||
|
||||
RUNTIME_CONFIG_REDIS_KEY = "yuxi:runtime_config"
|
||||
RUNTIME_CONFIG_SYNC_INTERVAL_SECONDS = 5.0
|
||||
_RUNTIME_CONFIG_REDIS_TIMEOUT_SECONDS = 0.2
|
||||
|
||||
|
||||
def _runtime_config_redis_config() -> RedisConfig:
|
||||
return RedisConfig.from_env(
|
||||
decode_responses=True,
|
||||
socket_timeout=_RUNTIME_CONFIG_REDIS_TIMEOUT_SECONDS,
|
||||
socket_connect_timeout=_RUNTIME_CONFIG_REDIS_TIMEOUT_SECONDS,
|
||||
)
|
||||
|
||||
|
||||
def _runtime_fields(config: Any) -> Iterator[str]:
|
||||
for field_name, field_info in type(config).model_fields.items():
|
||||
if field_info.exclude:
|
||||
continue
|
||||
yield field_name
|
||||
|
||||
|
||||
def _runtime_snapshot(config: Any) -> dict[str, Any]:
|
||||
return {field_name: getattr(config, field_name) for field_name in _runtime_fields(config)}
|
||||
|
||||
|
||||
def _load_snapshot() -> dict[str, Any] | None:
|
||||
try:
|
||||
with sync_redis_client(_runtime_config_redis_config()) as redis_client:
|
||||
raw = redis_client.get(RUNTIME_CONFIG_REDIS_KEY)
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to load runtime config from Redis: {e}")
|
||||
return None
|
||||
|
||||
if not raw:
|
||||
return None
|
||||
|
||||
try:
|
||||
snapshot = json.loads(raw)
|
||||
except json.JSONDecodeError as e:
|
||||
logger.warning(f"Failed to decode runtime config snapshot: {e}")
|
||||
return None
|
||||
|
||||
return snapshot if isinstance(snapshot, dict) else None
|
||||
|
||||
|
||||
def save_runtime_config(config: Any) -> None:
|
||||
try:
|
||||
with sync_redis_client(_runtime_config_redis_config()) as redis_client:
|
||||
redis_client.set(
|
||||
RUNTIME_CONFIG_REDIS_KEY,
|
||||
json.dumps(_runtime_snapshot(config), ensure_ascii=False),
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to save runtime config to Redis: {e}")
|
||||
|
||||
|
||||
def refresh_runtime_config(config: Any) -> None:
|
||||
snapshot = _load_snapshot()
|
||||
if snapshot is None:
|
||||
return
|
||||
|
||||
for field_name in _runtime_fields(config):
|
||||
if field_name in snapshot:
|
||||
setattr(config, field_name, snapshot[field_name])
|
||||
|
||||
|
||||
def start_runtime_sync(
|
||||
config: Any,
|
||||
current_thread: threading.Thread | None,
|
||||
*,
|
||||
interval: float = RUNTIME_CONFIG_SYNC_INTERVAL_SECONDS,
|
||||
) -> threading.Thread:
|
||||
if current_thread is not None:
|
||||
return current_thread
|
||||
|
||||
def runtime_sync_loop() -> None:
|
||||
while True:
|
||||
time.sleep(interval)
|
||||
try:
|
||||
refresh_runtime_config(config)
|
||||
except Exception as e:
|
||||
logger.warning(f"Runtime config sync iteration failed: {e}")
|
||||
|
||||
thread = threading.Thread(target=runtime_sync_loop, name="config-runtime-sync", daemon=True)
|
||||
thread.start()
|
||||
logger.info(f"Runtime config sync thread started (interval={interval}s)")
|
||||
return thread
|
||||
@@ -0,0 +1,129 @@
|
||||
# ==========================================================
|
||||
# 有害内容过滤关键词列表
|
||||
# 说明:
|
||||
# 1. 本列表仅为基础示例,请根据您的具体应用场景进行修改和扩展。
|
||||
# 2. 很多词语具有多义性,请务必结合语境进行判断,以防误伤。
|
||||
# 3. 网络用语和暗语变化迅速,建议定期更新此列表。
|
||||
# ==========================================================
|
||||
|
||||
# ----------------------------------------
|
||||
# 1. 仇恨言论 (Hate Speech)
|
||||
# ----------------------------------------
|
||||
|
||||
# ** 种族与民族歧视 **
|
||||
尼哥
|
||||
黑鬼
|
||||
白皮猪
|
||||
棒子
|
||||
阿三
|
||||
|
||||
# ** 宗教歧视 **
|
||||
异教徒
|
||||
邪教
|
||||
|
||||
# ** 性别与性取向歧视 **
|
||||
死娘炮
|
||||
男人婆
|
||||
同性恋不得好死
|
||||
恐同
|
||||
|
||||
|
||||
# ----------------------------------------
|
||||
# 2. 暴力与极端内容 (Violence & Extremism)
|
||||
# ----------------------------------------
|
||||
|
||||
# ** 煽动暴力 **
|
||||
杀全家
|
||||
弄死他
|
||||
打断腿
|
||||
血洗
|
||||
人肉他
|
||||
|
||||
# ** 恐怖主义与极端思想 **
|
||||
圣战
|
||||
ISIS
|
||||
东突
|
||||
|
||||
# ** 血腥与残忍 **
|
||||
虐杀
|
||||
活剐
|
||||
分尸
|
||||
电锯惊魂
|
||||
|
||||
# ** 武器与爆炸物 **
|
||||
土制炸弹
|
||||
买枪
|
||||
军火
|
||||
C4炸药
|
||||
炸药爆
|
||||
制作炸药
|
||||
制作武器
|
||||
炸弹
|
||||
|
||||
# ----------------------------------------
|
||||
# 3. 非法活动与违禁品 (Illegal Activities)
|
||||
# ----------------------------------------
|
||||
|
||||
# ** 毒品 **
|
||||
冰毒
|
||||
K粉
|
||||
海洛因
|
||||
吸毒
|
||||
贩毒
|
||||
摇头丸
|
||||
毒品走私
|
||||
|
||||
|
||||
# ** 赌博 **
|
||||
线上赌场
|
||||
澳门首家
|
||||
性感荷官
|
||||
百家乐
|
||||
六合彩
|
||||
|
||||
# ** 诈骗 **
|
||||
杀猪盘
|
||||
刷单兼职
|
||||
网贷陷阱
|
||||
冒充公检法
|
||||
|
||||
# ** 违禁品交易 **
|
||||
出售个人信息
|
||||
办假证
|
||||
针孔摄像头
|
||||
窃听器
|
||||
|
||||
|
||||
# ----------------------------------------
|
||||
# 4. 色情与成人内容 (Pornography)
|
||||
# ----------------------------------------
|
||||
AV女优
|
||||
草榴
|
||||
91大神
|
||||
黄片
|
||||
色情网站
|
||||
约炮
|
||||
裸聊
|
||||
福利姬
|
||||
|
||||
# ----------------------------------------
|
||||
# 5. 自残与危险行为 (Self-harm)
|
||||
# ----------------------------------------
|
||||
自杀教程
|
||||
割腕
|
||||
烧炭
|
||||
无痛自杀
|
||||
一起死
|
||||
|
||||
# ----------------------------------------
|
||||
# 6. 网络欺凌与骚扰 (Cyberbullying)
|
||||
# ----------------------------------------
|
||||
biss
|
||||
nmsl
|
||||
开盒
|
||||
挂人
|
||||
孤儿
|
||||
|
||||
# ==========================================================
|
||||
# 列表结束
|
||||
# ==========================================================
|
||||
@@ -0,0 +1,27 @@
|
||||
# 单位信息配置文件
|
||||
# 用于配置网站的基本信息、品牌信息等
|
||||
|
||||
# 组织信息
|
||||
organization:
|
||||
name: "江南语析" # 完整组织名称
|
||||
logo: "/favicon.svg" # Logo文件路径(放在 web/public 目录下)
|
||||
avatar: "/avatar.jpg" # 头像文件路径(放在 web/public 目录下)
|
||||
login_bg: "/login-bg.jpg" # 登录背景图片路径(放在 web/public 目录下)
|
||||
|
||||
# 项目信息
|
||||
branding:
|
||||
name: "Yuxi"
|
||||
title: "让智能体可构建、可编排、可落地" # 系统标题
|
||||
subtitle: "开源智能体平台套件,融合 RAG 与知识图谱" # 副标题(subtitles 为空时使用)
|
||||
subtitles:
|
||||
- "开源智能体平台套件,融合 RAG 与知识图谱"
|
||||
- "统一编排 Agent、知识库、图谱与工具链"
|
||||
- "让智能体可构建,让知识可连接,让决策可验证"
|
||||
- "让智能体可落地,让流程可编排,让协作可扩展"
|
||||
- "让数据可沉淀,让能力可复用,让系统可进化"
|
||||
|
||||
# 页脚信息
|
||||
footer:
|
||||
copyright: "© 江南语析 2026 v{{YUXI_VERSION}}"
|
||||
user_agreement_url: "/protocols/user-agreement.template.html"
|
||||
privacy_policy_url: "/protocols/privacy-policy.template.html"
|
||||
@@ -0,0 +1,80 @@
|
||||
"""用户级配置模块。"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import datetime
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field
|
||||
from sqlalchemy import select, update
|
||||
from sqlalchemy.exc import IntegrityError
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from yuxi.storage.postgres.models_business import UserConfig as UserConfigRecord
|
||||
from yuxi.utils.datetime_utils import format_utc_datetime, utc_now_naive
|
||||
|
||||
|
||||
class UserConfigSchema(BaseModel):
|
||||
"""用户专属配置 schema。"""
|
||||
|
||||
enable_memory: bool = Field(default=False, description="是否启用 Memory")
|
||||
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
|
||||
class UserConfig:
|
||||
"""用户级配置访问器。每次加载都从 PostgreSQL 查询,不做进程缓存。"""
|
||||
|
||||
def __init__(self, uid: str, schema: UserConfigSchema | None = None, updated_at: datetime | None = None):
|
||||
self.uid = uid
|
||||
self.schema = schema or UserConfigSchema()
|
||||
self.updated_at = updated_at
|
||||
|
||||
@classmethod
|
||||
async def load(cls, db: AsyncSession, uid: str) -> UserConfig:
|
||||
result = await db.execute(select(UserConfigRecord).filter(UserConfigRecord.uid == uid))
|
||||
record = result.scalar_one_or_none()
|
||||
if record is None:
|
||||
return cls(uid=uid)
|
||||
return cls(
|
||||
uid=uid,
|
||||
schema=UserConfigSchema(enable_memory=bool(record.enable_memory)),
|
||||
updated_at=record.updated_at,
|
||||
)
|
||||
|
||||
async def save(self, db: AsyncSession) -> UserConfig:
|
||||
now = utc_now_naive()
|
||||
result = await db.execute(
|
||||
update(UserConfigRecord)
|
||||
.where(UserConfigRecord.uid == self.uid)
|
||||
.values(enable_memory=self.schema.enable_memory, updated_at=now)
|
||||
)
|
||||
if result.rowcount == 0:
|
||||
db.add(
|
||||
UserConfigRecord(
|
||||
uid=self.uid,
|
||||
enable_memory=self.schema.enable_memory,
|
||||
created_at=now,
|
||||
updated_at=now,
|
||||
)
|
||||
)
|
||||
try:
|
||||
await db.commit()
|
||||
except IntegrityError:
|
||||
await db.rollback()
|
||||
now = utc_now_naive()
|
||||
retry_result = await db.execute(
|
||||
update(UserConfigRecord)
|
||||
.where(UserConfigRecord.uid == self.uid)
|
||||
.values(enable_memory=self.schema.enable_memory, updated_at=now)
|
||||
)
|
||||
if retry_result.rowcount == 0:
|
||||
raise
|
||||
await db.commit()
|
||||
return type(self)(uid=self.uid, schema=self.schema, updated_at=now)
|
||||
|
||||
def dump_config(self) -> dict[str, str | bool | None]:
|
||||
return {
|
||||
"uid": self.uid,
|
||||
"enable_memory": self.schema.enable_memory,
|
||||
"updated_at": format_utc_datetime(self.updated_at),
|
||||
}
|
||||
@@ -0,0 +1,24 @@
|
||||
import os
|
||||
|
||||
from ..config import config
|
||||
from .factory import KnowledgeBaseFactory
|
||||
from .implementations.dify import DifyKB
|
||||
from .implementations.milvus import MilvusKB
|
||||
from .implementations.notion import NotionKB
|
||||
from .manager import KnowledgeBaseManager
|
||||
|
||||
_LITE_MODE = os.environ.get("LITE_MODE", "").lower() in ("true", "1")
|
||||
_SKIP_APP_INIT = os.environ.get("YUXI_SKIP_APP_INIT") == "1"
|
||||
|
||||
if not _LITE_MODE:
|
||||
# 注册知识库类型
|
||||
KnowledgeBaseFactory.register(MilvusKB)
|
||||
|
||||
KnowledgeBaseFactory.register(DifyKB)
|
||||
KnowledgeBaseFactory.register(NotionKB)
|
||||
|
||||
# 创建知识库管理器
|
||||
work_dir = os.path.join(config.save_dir, "knowledge_base_data")
|
||||
knowledge_base = KnowledgeBaseManager(work_dir)
|
||||
|
||||
__all__ = ["knowledge_base"]
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1 @@
|
||||
__all__ = []
|
||||
@@ -0,0 +1,3 @@
|
||||
from yuxi.knowledge.chunking.ragflow_like.dispatcher import chunk_file, chunk_markdown
|
||||
|
||||
__all__ = ["chunk_file", "chunk_markdown"]
|
||||
@@ -0,0 +1,85 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from yuxi.knowledge.chunking.ragflow_like.parsers import book, general, laws, qa, semantic, separator
|
||||
from yuxi.knowledge.chunking.ragflow_like.presets import map_to_internal_parser_id, normalize_chunk_preset_id
|
||||
|
||||
|
||||
def _build_chunk_records(
|
||||
text_chunks: list[str], file_id: str, filename: str, source_text: str | None = None
|
||||
) -> list[dict[str, Any]]:
|
||||
records: list[dict[str, Any]] = []
|
||||
search_from = 0
|
||||
|
||||
for idx, chunk_content in enumerate(text_chunks):
|
||||
text = (chunk_content or "").strip()
|
||||
if not text:
|
||||
continue
|
||||
|
||||
start_char_pos = None
|
||||
end_char_pos = None
|
||||
if source_text:
|
||||
found_at = source_text.find(text, search_from)
|
||||
if found_at >= 0:
|
||||
start_char_pos = found_at
|
||||
end_char_pos = found_at + len(text)
|
||||
search_from = end_char_pos
|
||||
|
||||
records.append(
|
||||
{
|
||||
"id": f"{file_id}_chunk_{idx}",
|
||||
"content": text,
|
||||
"file_id": file_id,
|
||||
"filename": filename,
|
||||
"chunk_index": idx,
|
||||
"source": filename,
|
||||
"chunk_id": f"{file_id}_chunk_{idx}",
|
||||
"start_char_pos": start_char_pos,
|
||||
"end_char_pos": end_char_pos,
|
||||
"start_token_pos": None,
|
||||
"end_token_pos": None,
|
||||
"extraction_result": None,
|
||||
}
|
||||
)
|
||||
|
||||
return records
|
||||
|
||||
|
||||
def _dispatch_markdown_parser(
|
||||
preset_id: str, filename: str, markdown_content: str, parser_config: dict[str, Any]
|
||||
) -> list[str]:
|
||||
parser_id = map_to_internal_parser_id(preset_id)
|
||||
|
||||
if parser_id == "naive":
|
||||
return general.chunk_markdown(markdown_content, parser_config)
|
||||
if parser_id == "qa":
|
||||
return qa.chunk_markdown(filename, markdown_content, parser_config)
|
||||
if parser_id == "book":
|
||||
return book.chunk_markdown(markdown_content, parser_config)
|
||||
if parser_id == "laws":
|
||||
return laws.chunk_markdown(filename, markdown_content, parser_config)
|
||||
if parser_id == "semantic":
|
||||
return semantic.chunk_markdown(markdown_content, parser_config)
|
||||
if parser_id == "separator":
|
||||
return separator.chunk_markdown(markdown_content, parser_config)
|
||||
|
||||
return general.chunk_markdown(markdown_content, parser_config)
|
||||
|
||||
|
||||
def chunk_markdown(
|
||||
markdown_content: str, file_id: str, filename: str, processing_params: dict[str, Any]
|
||||
) -> list[dict[str, Any]]:
|
||||
params = dict(processing_params or {})
|
||||
preset_id = normalize_chunk_preset_id(params.get("chunk_preset_id"))
|
||||
parser_config = params.get("chunk_parser_config") if isinstance(params.get("chunk_parser_config"), dict) else {}
|
||||
|
||||
text_chunks = _dispatch_markdown_parser(preset_id, filename, markdown_content, parser_config)
|
||||
return _build_chunk_records(text_chunks, file_id, filename, markdown_content)
|
||||
|
||||
|
||||
def chunk_file(
|
||||
file_content: str, file_id: str, filename: str, processing_params: dict[str, Any]
|
||||
) -> list[dict[str, Any]]:
|
||||
# 当前链路中入库前均已转换为 markdown,因此与 chunk_markdown 保持同实现。
|
||||
return chunk_markdown(file_content, file_id, filename, processing_params)
|
||||
@@ -0,0 +1,583 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import random
|
||||
import re
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
BULLET_PATTERN = [
|
||||
[
|
||||
r"第[零一二三四五六七八九十百0-9]+(分?编|部分)",
|
||||
r"第[零一二三四五六七八九十百0-9]+章",
|
||||
r"第[零一二三四五六七八九十百0-9]+节",
|
||||
r"第[零一二三四五六七八九十百0-9]+条",
|
||||
r"[\((][零一二三四五六七八九十百]+[\))]",
|
||||
],
|
||||
[
|
||||
r"第[0-9]+章",
|
||||
r"第[0-9]+节",
|
||||
r"[0-9]{,2}[\. 、]",
|
||||
r"[0-9]{,2}\.[0-9]{,2}[^a-zA-Z/%~-]",
|
||||
r"[0-9]{,2}\.[0-9]{,2}\.[0-9]{,2}",
|
||||
r"[0-9]{,2}\.[0-9]{,2}\.[0-9]{,2}\.[0-9]{,2}",
|
||||
],
|
||||
[
|
||||
r"第[零一二三四五六七八九十百0-9]+章",
|
||||
r"第[零一二三四五六七八九十百0-9]+节",
|
||||
r"[零一二三四五六七八九十百]+[ 、]",
|
||||
r"[\((][零一二三四五六七八九十百]+[\))]",
|
||||
r"[\((][0-9]{,2}[\))]",
|
||||
],
|
||||
[
|
||||
r"PART (ONE|TWO|THREE|FOUR|FIVE|SIX|SEVEN|EIGHT|NINE|TEN)",
|
||||
r"Chapter (I+V?|VI*|XI|IX|X)",
|
||||
r"Section [0-9]+",
|
||||
r"Article [0-9]+",
|
||||
],
|
||||
[
|
||||
r"^#[^#]",
|
||||
r"^##[^#]",
|
||||
r"^###.*",
|
||||
r"^####.*",
|
||||
r"^#####.*",
|
||||
r"^######.*",
|
||||
],
|
||||
]
|
||||
|
||||
MARKDOWN_BULLET_GROUP_INDEX = 4
|
||||
|
||||
|
||||
def count_tokens(text: str) -> int:
|
||||
"""近似 token 计数,避免引入额外依赖。"""
|
||||
if not text:
|
||||
return 0
|
||||
# 英文单词 + 数字 + CJK 单字
|
||||
parts = re.findall(r"[A-Za-z0-9_]+|[\u4e00-\u9fff]", text)
|
||||
return max(1, len(parts)) if text.strip() else 0
|
||||
|
||||
|
||||
def hard_split_by_token_limit(text: str, chunk_token_num: int, hard_limit_token_num: int | None = None) -> list[str]:
|
||||
"""将文本按 token 上限硬切,用于 naive_merge 之后的兜底保护。
|
||||
|
||||
hard_limit_token_num 只在调用方显式传入时生效,用于允许略超目标长度的块
|
||||
保持完整;默认保持严格不超过 chunk_token_num 的历史行为。
|
||||
"""
|
||||
token_iter = list(re.finditer(r"[A-Za-z0-9_]+|[一-鿿]", text or ""))
|
||||
if not token_iter:
|
||||
cleaned = (text or "").strip()
|
||||
return [cleaned] if cleaned else []
|
||||
|
||||
max_tokens = max(int(chunk_token_num or 0), 1)
|
||||
hard_limit = None
|
||||
if hard_limit_token_num is not None:
|
||||
hard_limit = max(int(hard_limit_token_num or 0), max_tokens)
|
||||
if len(token_iter) <= hard_limit:
|
||||
cleaned = (text or "").strip()
|
||||
return [cleaned] if cleaned else []
|
||||
|
||||
spans: list[tuple[int, int]] = []
|
||||
start = 0
|
||||
index = 0
|
||||
|
||||
while index < len(token_iter):
|
||||
next_index = min(index + max_tokens, len(token_iter))
|
||||
if next_index < len(token_iter):
|
||||
end = token_iter[next_index].start()
|
||||
else:
|
||||
end = len(text)
|
||||
if text[start:end].strip():
|
||||
spans.append((start, end))
|
||||
start = end
|
||||
index = next_index
|
||||
|
||||
tail = text[start:].strip()
|
||||
if tail:
|
||||
spans.append((start, len(text)))
|
||||
|
||||
if hard_limit is not None and len(spans) >= 2:
|
||||
prev_start, _ = spans[-2]
|
||||
_, tail_end = spans[-1]
|
||||
candidate = text[prev_start:tail_end].strip()
|
||||
if count_tokens(candidate) <= hard_limit:
|
||||
spans[-2] = (prev_start, tail_end)
|
||||
spans.pop()
|
||||
|
||||
return [text[start:end].strip() for start, end in spans if text[start:end].strip()]
|
||||
|
||||
|
||||
def random_choices(arr: list[str], k: int) -> list[str]:
|
||||
if not arr:
|
||||
return []
|
||||
return random.choices(arr, k=min(len(arr), k))
|
||||
|
||||
|
||||
def is_english(texts: str | list[str]) -> bool:
|
||||
if not texts:
|
||||
return False
|
||||
|
||||
patt = re.compile(r"[`a-zA-Z0-9\s.,':;/\"?<>!\(\)\-]+")
|
||||
if isinstance(texts, str):
|
||||
seq = [texts]
|
||||
else:
|
||||
seq = [t for t in texts if isinstance(t, str) and t.strip()]
|
||||
|
||||
if not seq:
|
||||
return False
|
||||
|
||||
hits = sum(1 for t in seq if patt.fullmatch(t.strip()))
|
||||
return (hits / len(seq)) > 0.8
|
||||
|
||||
|
||||
def not_bullet(line: str) -> bool:
|
||||
patt = [
|
||||
r"0",
|
||||
r"[0-9]+ +[0-9~个只-]",
|
||||
r"[0-9]+\.{2,}",
|
||||
]
|
||||
return any(re.match(p, line) for p in patt)
|
||||
|
||||
|
||||
def is_probable_heading_line(line: str) -> bool:
|
||||
text = (line or "").strip()
|
||||
if not text:
|
||||
return False
|
||||
|
||||
if re.match(r"^#{1,6}\s+\S", text):
|
||||
return True
|
||||
|
||||
# 表格/HTML 残留通常不是标题。
|
||||
if re.search(r"</?(table|tr|td|th|caption|tbody|thead)[^>]*>", text, flags=re.IGNORECASE):
|
||||
return False
|
||||
|
||||
# 超长行基本是正文或条款,不是章节标题。
|
||||
if len(text) > 96:
|
||||
return False
|
||||
if count_tokens(text) > 72:
|
||||
return False
|
||||
|
||||
# 标题前段通常不会出现明显句号/逗号;出现则大概率是正文。
|
||||
if re.search(r"[,。;!?!?::]", text[:24]):
|
||||
return False
|
||||
|
||||
if text.endswith(("。", ";", "!", "!", "?", "?")) and len(text) > 20:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
|
||||
def _is_mid_sentence_bullet(line: str) -> bool:
|
||||
text = (line or "").strip()
|
||||
if not text:
|
||||
return False
|
||||
|
||||
if re.match(r"^#{1,6}\s+\S", text):
|
||||
return False
|
||||
|
||||
marker = re.search(
|
||||
r"([一二三四五六七八九十百]+、|[\((][一二三四五六七八九十百]+[\))]|[0-9]{1,2}[\.、])",
|
||||
text,
|
||||
)
|
||||
if not marker:
|
||||
return False
|
||||
|
||||
if marker.start() == 0:
|
||||
return False
|
||||
|
||||
prev = text[marker.start() - 1]
|
||||
return prev not in {"#", "\n"}
|
||||
|
||||
|
||||
def bullets_category(sections: list[str]) -> int:
|
||||
hits: list[float] = [0.0] * len(BULLET_PATTERN)
|
||||
|
||||
def bullet_weight(group_idx: int, line: str) -> float:
|
||||
# 对 markdown 标题候选增加权重,避免“正文里的 一、/(一)”压过真正的 # 标题层级。
|
||||
if group_idx != MARKDOWN_BULLET_GROUP_INDEX:
|
||||
return 1.0
|
||||
|
||||
heading = line.strip()
|
||||
if not re.match(r"^#{1,6}\s+\S", heading):
|
||||
return 1.0
|
||||
|
||||
level = len(heading) - len(heading.lstrip("#"))
|
||||
if level <= 2:
|
||||
return 4.0
|
||||
if level <= 4:
|
||||
return 3.0
|
||||
return 2.0
|
||||
|
||||
for i, pro in enumerate(BULLET_PATTERN):
|
||||
for sec in sections:
|
||||
sec = sec.strip()
|
||||
for p in pro:
|
||||
if re.match(p, sec) and not not_bullet(sec):
|
||||
w = bullet_weight(i, sec)
|
||||
if _is_mid_sentence_bullet(sec):
|
||||
w *= 0.1
|
||||
if i != MARKDOWN_BULLET_GROUP_INDEX and not is_probable_heading_line(sec):
|
||||
w *= 0.2
|
||||
hits[i] += w
|
||||
break
|
||||
maximum = 0
|
||||
res = -1
|
||||
for i, hit in enumerate(hits):
|
||||
if hit <= maximum:
|
||||
continue
|
||||
res = i
|
||||
maximum = hit
|
||||
return res
|
||||
|
||||
|
||||
def _get_text(section: str | tuple[str, str]) -> str:
|
||||
if isinstance(section, str):
|
||||
return section.strip()
|
||||
return (section[0] or "").strip()
|
||||
|
||||
|
||||
def remove_contents_table(sections: list[str] | list[tuple[str, str]], eng: bool = False) -> None:
|
||||
i = 0
|
||||
while i < len(sections):
|
||||
line = re.sub(r"( | |\u3000)+", "", _get_text(sections[i]).split("@@")[0], flags=re.IGNORECASE)
|
||||
if not re.match(r"(contents|目录|目次|tableofcontents|致谢|acknowledge)$", line, flags=re.IGNORECASE):
|
||||
i += 1
|
||||
continue
|
||||
|
||||
sections.pop(i)
|
||||
if i >= len(sections):
|
||||
break
|
||||
|
||||
prefix = _get_text(sections[i])[:3] if not eng else " ".join(_get_text(sections[i]).split()[:2])
|
||||
while not prefix and i < len(sections):
|
||||
sections.pop(i)
|
||||
if i >= len(sections):
|
||||
break
|
||||
prefix = _get_text(sections[i])[:3] if not eng else " ".join(_get_text(sections[i]).split()[:2])
|
||||
|
||||
if i >= len(sections) or not prefix:
|
||||
break
|
||||
|
||||
sections.pop(i)
|
||||
if i >= len(sections):
|
||||
break
|
||||
|
||||
for j in range(i, min(i + 128, len(sections))):
|
||||
if not re.match(re.escape(prefix), _get_text(sections[j])):
|
||||
continue
|
||||
for _ in range(i, j):
|
||||
sections.pop(i)
|
||||
break
|
||||
|
||||
|
||||
def make_colon_as_title(sections: list[str] | list[tuple[str, str]]) -> list[str] | list[tuple[str, str]]:
|
||||
if not sections:
|
||||
return sections
|
||||
if isinstance(sections[0], str):
|
||||
return sections
|
||||
|
||||
i = 0
|
||||
while i < len(sections):
|
||||
text, layout = sections[i]
|
||||
i += 1
|
||||
text = text.split("@")[0].strip()
|
||||
if not text or text[-1] not in "::":
|
||||
continue
|
||||
|
||||
rev = text[::-1]
|
||||
arr = re.split(r"([。?!!?;;]| \.)", rev)
|
||||
if len(arr) < 2 or len(arr[1]) < 32:
|
||||
continue
|
||||
|
||||
sections.insert(i - 1, (arr[0][::-1], "title"))
|
||||
i += 1
|
||||
|
||||
return sections
|
||||
|
||||
|
||||
def not_title(text: str) -> bool:
|
||||
if re.match(r"第[零一二三四五六七八九十百0-9]+条", text):
|
||||
return False
|
||||
if len(text.split()) > 12 or (" " not in text and len(text) >= 32):
|
||||
return True
|
||||
return bool(re.search(r"[,;,。;!!]", text))
|
||||
|
||||
|
||||
def tree_merge(bull: int, sections: list[str] | list[tuple[str, str]], depth: int) -> list[str]:
|
||||
if not sections or bull < 0:
|
||||
return [s if isinstance(s, str) else s[0] for s in sections]
|
||||
|
||||
if isinstance(sections[0], str):
|
||||
typed_sections: list[tuple[str, str]] = [(s, "") for s in sections]
|
||||
else:
|
||||
typed_sections = sections # type: ignore[assignment]
|
||||
|
||||
typed_sections = [
|
||||
(t, o)
|
||||
for t, o in typed_sections
|
||||
if t and len(t.split("@")[0].strip()) > 1 and not re.match(r"[0-9]+$", t.split("@")[0].strip())
|
||||
]
|
||||
|
||||
def get_level(section: tuple[str, str]) -> tuple[int, str]:
|
||||
text, layout = section
|
||||
text = re.sub(r"\u3000", " ", text).strip()
|
||||
|
||||
for i, patt in enumerate(BULLET_PATTERN[bull]):
|
||||
if re.match(patt, text) and is_probable_heading_line(text):
|
||||
return i + 1, text
|
||||
|
||||
if re.search(r"(title|head)", layout) and not not_title(text):
|
||||
return len(BULLET_PATTERN[bull]) + 1, text
|
||||
|
||||
return len(BULLET_PATTERN[bull]) + 2, text
|
||||
|
||||
lines: list[tuple[int, str]] = []
|
||||
level_set: set[int] = set()
|
||||
for section in typed_sections:
|
||||
level, text = get_level(section)
|
||||
if not text.strip("\n"):
|
||||
continue
|
||||
lines.append((level, text))
|
||||
level_set.add(level)
|
||||
|
||||
if not lines:
|
||||
return []
|
||||
|
||||
sorted_levels = sorted(level_set)
|
||||
target_level = sorted_levels[depth - 1] if depth <= len(sorted_levels) else sorted_levels[-1]
|
||||
|
||||
max_body_level = len(BULLET_PATTERN[bull]) + 2
|
||||
if target_level == max_body_level:
|
||||
target_level = sorted_levels[-2] if len(sorted_levels) > 1 else sorted_levels[0]
|
||||
|
||||
root = Node(level=0, depth=target_level, texts=[])
|
||||
root.build_tree(lines)
|
||||
return [item for item in root.get_tree() if item]
|
||||
|
||||
|
||||
def hierarchical_merge(bull: int, sections: list[str] | list[tuple[str, str]], depth: int) -> list[list[str]]:
|
||||
if not sections or bull < 0:
|
||||
return []
|
||||
|
||||
if isinstance(sections[0], str):
|
||||
typed_sections: list[tuple[str, str]] = [(s, "") for s in sections]
|
||||
else:
|
||||
typed_sections = sections # type: ignore[assignment]
|
||||
|
||||
typed_sections = [
|
||||
(t, o)
|
||||
for t, o in typed_sections
|
||||
if t and len(t.split("@")[0].strip()) > 1 and not re.match(r"[0-9]+$", t.split("@")[0].strip())
|
||||
]
|
||||
|
||||
bullets_size = len(BULLET_PATTERN[bull])
|
||||
levels: list[list[int]] = [[] for _ in range(bullets_size + 2)]
|
||||
|
||||
for i, (text, layout) in enumerate(typed_sections):
|
||||
for j, patt in enumerate(BULLET_PATTERN[bull]):
|
||||
if re.match(patt, text.strip()) and is_probable_heading_line(text):
|
||||
levels[j].append(i)
|
||||
break
|
||||
else:
|
||||
if re.search(r"(title|head)", layout) and not not_title(text):
|
||||
levels[bullets_size].append(i)
|
||||
else:
|
||||
levels[bullets_size + 1].append(i)
|
||||
|
||||
pure_sections = [t for t, _ in typed_sections]
|
||||
|
||||
def binary_search(arr: list[int], target: int) -> int:
|
||||
if not arr:
|
||||
return -1
|
||||
if target > arr[-1]:
|
||||
return len(arr) - 1
|
||||
if target < arr[0]:
|
||||
return -1
|
||||
|
||||
s, e = 0, len(arr)
|
||||
while e - s > 1:
|
||||
mid = (e + s) // 2
|
||||
if target > arr[mid]:
|
||||
s = mid
|
||||
elif target < arr[mid]:
|
||||
e = mid
|
||||
else:
|
||||
return mid
|
||||
return s
|
||||
|
||||
cks: list[list[int]] = []
|
||||
readed = [False] * len(pure_sections)
|
||||
levels = list(reversed(levels))
|
||||
for i, arr in enumerate(levels[:depth]):
|
||||
for j in arr:
|
||||
if readed[j]:
|
||||
continue
|
||||
readed[j] = True
|
||||
cks.append([j])
|
||||
if i + 1 == len(levels) - 1:
|
||||
continue
|
||||
|
||||
for ii in range(i + 1, len(levels)):
|
||||
jj = binary_search(levels[ii], j)
|
||||
if jj < 0:
|
||||
continue
|
||||
if levels[ii][jj] > cks[-1][-1]:
|
||||
cks[-1].pop(-1)
|
||||
cks[-1].append(levels[ii][jj])
|
||||
|
||||
for ii in cks[-1]:
|
||||
readed[ii] = True
|
||||
|
||||
if not cks:
|
||||
return []
|
||||
|
||||
for i in range(len(cks)):
|
||||
cks[i] = [pure_sections[j] for j in reversed(cks[i])]
|
||||
|
||||
res: list[list[str]] = [[]]
|
||||
num = [0]
|
||||
for ck in cks:
|
||||
if len(ck) == 1:
|
||||
n = count_tokens(re.sub(r"@@[0-9]+.*", "", ck[0]))
|
||||
if n + num[-1] < 218:
|
||||
res[-1].append(ck[0])
|
||||
num[-1] += n
|
||||
continue
|
||||
res.append(ck)
|
||||
num.append(n)
|
||||
continue
|
||||
res.append(ck)
|
||||
num.append(218)
|
||||
|
||||
return [chunk for chunk in res if chunk]
|
||||
|
||||
|
||||
def _remove_pdf_tags(text: str) -> str:
|
||||
return re.sub(r"@@[0-9-]+\t[0-9.\t]+##", "", text or "")
|
||||
|
||||
|
||||
def _extract_custom_delimiters(delimiter: str) -> list[str]:
|
||||
return [m.group(1) for m in re.finditer(r"`([^`]+)`", delimiter or "")]
|
||||
|
||||
|
||||
def naive_merge(
|
||||
sections: str | list[str] | list[tuple[str, str]],
|
||||
chunk_token_num: int = 128,
|
||||
delimiter: str = "\n。;!?",
|
||||
overlapped_percent: int = 0,
|
||||
) -> list[str]:
|
||||
if not sections:
|
||||
return []
|
||||
|
||||
if isinstance(sections, str):
|
||||
typed_sections: list[tuple[str, str]] = [(sections, "")]
|
||||
elif isinstance(sections[0], str):
|
||||
typed_sections = [(s, "") for s in sections] # type: ignore[index]
|
||||
else:
|
||||
typed_sections = sections # type: ignore[assignment]
|
||||
|
||||
chunk_token_num = max(int(chunk_token_num or 0), 0)
|
||||
overlap = max(0, min(int(overlapped_percent or 0), 99))
|
||||
|
||||
custom_delimiters = _extract_custom_delimiters(delimiter)
|
||||
if custom_delimiters:
|
||||
pattern = "|".join(re.escape(t) for t in sorted(set(custom_delimiters), key=len, reverse=True))
|
||||
chunks: list[str] = []
|
||||
for sec, pos in typed_sections:
|
||||
split_sec = re.split(rf"({pattern})", sec, flags=re.DOTALL)
|
||||
for sub in split_sec:
|
||||
if re.fullmatch(pattern, sub or ""):
|
||||
continue
|
||||
text = "\n" + sub
|
||||
local_pos = pos if count_tokens(text) >= 8 else ""
|
||||
if local_pos and local_pos not in text:
|
||||
text += local_pos
|
||||
if text.strip():
|
||||
chunks.append(text)
|
||||
return chunks
|
||||
|
||||
if chunk_token_num <= 0:
|
||||
merged = "\n".join(sec for sec, _ in typed_sections if sec and sec.strip())
|
||||
return [merged] if merged.strip() else []
|
||||
|
||||
chunks = [""]
|
||||
token_nums = [0]
|
||||
|
||||
def add_chunk(text: str, pos: str) -> None:
|
||||
tnum = count_tokens(text)
|
||||
local_pos = pos or ""
|
||||
if tnum < 8:
|
||||
local_pos = ""
|
||||
|
||||
threshold = chunk_token_num * (100 - overlap) / 100.0
|
||||
if chunks[-1] == "" or token_nums[-1] > threshold:
|
||||
if chunks:
|
||||
prev = _remove_pdf_tags(chunks[-1])
|
||||
start = int(len(prev) * (100 - overlap) / 100.0)
|
||||
text = prev[start:] + text
|
||||
if local_pos and local_pos not in text:
|
||||
text += local_pos
|
||||
chunks.append(text)
|
||||
token_nums.append(tnum)
|
||||
else:
|
||||
if local_pos and local_pos not in chunks[-1]:
|
||||
text += local_pos
|
||||
chunks[-1] += text
|
||||
token_nums[-1] += tnum
|
||||
|
||||
for sec, pos in typed_sections:
|
||||
if not sec:
|
||||
continue
|
||||
add_chunk("\n" + sec, pos)
|
||||
|
||||
return [chunk for chunk in chunks if chunk.strip()]
|
||||
|
||||
|
||||
@dataclass
|
||||
class Node:
|
||||
level: int
|
||||
depth: int = -1
|
||||
texts: list[str] = field(default_factory=list)
|
||||
children: list[Node] = field(default_factory=list)
|
||||
|
||||
def add_child(self, child_node: Node) -> None:
|
||||
self.children.append(child_node)
|
||||
|
||||
def add_text(self, text: str) -> None:
|
||||
self.texts.append(text)
|
||||
|
||||
def build_tree(self, lines: list[tuple[int, str]]) -> Node:
|
||||
stack: list[Node] = [self]
|
||||
for level, text in lines:
|
||||
if self.depth != -1 and level > self.depth:
|
||||
stack[-1].add_text(text)
|
||||
continue
|
||||
|
||||
while len(stack) > 1 and level <= stack[-1].level:
|
||||
stack.pop()
|
||||
|
||||
node = Node(level=level, texts=[text])
|
||||
stack[-1].add_child(node)
|
||||
stack.append(node)
|
||||
|
||||
return self
|
||||
|
||||
def get_tree(self) -> list[str]:
|
||||
tree_list: list[str] = []
|
||||
self._dfs(self, tree_list, [])
|
||||
return tree_list
|
||||
|
||||
def _dfs(self, node: Node, tree_list: list[str], titles: list[str]) -> None:
|
||||
level = node.level
|
||||
texts = node.texts
|
||||
child = node.children
|
||||
|
||||
if level == 0 and texts:
|
||||
tree_list.append("\n".join(titles + texts))
|
||||
|
||||
path_titles = titles + texts if 1 <= level <= self.depth else titles
|
||||
|
||||
if level > self.depth and texts:
|
||||
tree_list.append("\n".join(path_titles + texts))
|
||||
elif not child and (1 <= level <= self.depth):
|
||||
tree_list.append("\n".join(path_titles))
|
||||
|
||||
for c in child:
|
||||
self._dfs(c, tree_list, path_titles)
|
||||
@@ -0,0 +1,3 @@
|
||||
from yuxi.knowledge.chunking.ragflow_like.parsers import book, general, laws, qa, separator
|
||||
|
||||
__all__ = ["general", "qa", "book", "laws", "separator"]
|
||||
@@ -0,0 +1,79 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from yuxi.knowledge.chunking.ragflow_like import nlp
|
||||
|
||||
|
||||
def _unescape_delimiter(delimiter: str) -> str:
|
||||
return delimiter.replace("\\n", "\n").replace("\\r", "\r").replace("\\t", "\t").replace("\\\\", "\\")
|
||||
|
||||
|
||||
def _iter_sections(markdown_content: str) -> list[tuple[str, str]]:
|
||||
sections: list[tuple[str, str]] = []
|
||||
for line in (markdown_content or "").splitlines():
|
||||
text = line.strip()
|
||||
if not text:
|
||||
continue
|
||||
sections.append((text, ""))
|
||||
|
||||
if not sections and markdown_content and markdown_content.strip():
|
||||
sections.append((markdown_content.strip(), ""))
|
||||
|
||||
return sections
|
||||
|
||||
|
||||
def _ensure_chunk_token_limit(chunks: list[str], chunk_token_num: int) -> list[str]:
|
||||
max_tokens = int(chunk_token_num or 0)
|
||||
normalized = [chunk.strip() for chunk in chunks if chunk and chunk.strip()]
|
||||
if max_tokens <= 0:
|
||||
return normalized
|
||||
|
||||
protected: list[str] = []
|
||||
for chunk in normalized:
|
||||
if nlp.count_tokens(chunk) <= max_tokens:
|
||||
protected.append(chunk)
|
||||
else:
|
||||
protected.extend(nlp.hard_split_by_token_limit(chunk, max_tokens))
|
||||
return protected
|
||||
|
||||
|
||||
def chunk_markdown(markdown_content: str, parser_config: dict[str, Any] | None = None) -> list[str]:
|
||||
parser_config = parser_config or {}
|
||||
|
||||
delimiter = _unescape_delimiter(str(parser_config.get("delimiter", "\n") or "\n"))
|
||||
chunk_token_num = int(parser_config.get("chunk_token_num", 512) or 512)
|
||||
overlapped_percent = int(parser_config.get("overlapped_percent", 0) or 0)
|
||||
|
||||
sections = _iter_sections(markdown_content)
|
||||
if not sections:
|
||||
return []
|
||||
|
||||
section_texts = [text for text, _ in sections]
|
||||
nlp.remove_contents_table(sections, eng=nlp.is_english(nlp.random_choices(section_texts, k=200)))
|
||||
nlp.make_colon_as_title(sections)
|
||||
|
||||
bull = nlp.bullets_category([t for t in nlp.random_choices([t for t, _ in sections], k=100)])
|
||||
|
||||
if bull >= 0:
|
||||
chunks = ["\n".join(ck) for ck in nlp.hierarchical_merge(bull, sections, depth=5)]
|
||||
else:
|
||||
chunks = nlp.naive_merge(
|
||||
sections,
|
||||
chunk_token_num=chunk_token_num,
|
||||
delimiter=delimiter,
|
||||
overlapped_percent=overlapped_percent,
|
||||
)
|
||||
|
||||
if chunks:
|
||||
return _ensure_chunk_token_limit(chunks, chunk_token_num)
|
||||
|
||||
return _ensure_chunk_token_limit(
|
||||
nlp.naive_merge(
|
||||
sections,
|
||||
chunk_token_num=chunk_token_num,
|
||||
delimiter=delimiter,
|
||||
overlapped_percent=overlapped_percent,
|
||||
),
|
||||
chunk_token_num,
|
||||
)
|
||||
@@ -0,0 +1,68 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from yuxi.knowledge.chunking.ragflow_like import nlp
|
||||
|
||||
GENERAL_HARD_LIMIT_RATIO = 1.5
|
||||
|
||||
|
||||
def _unescape_delimiter(delimiter: str) -> str:
|
||||
return delimiter.replace("\\n", "\n").replace("\\r", "\r").replace("\\t", "\t").replace("\\\\", "\\")
|
||||
|
||||
|
||||
def _iter_sections(markdown_content: str, delimiter: str) -> list[tuple[str, str]]:
|
||||
sections: list[tuple[str, str]] = []
|
||||
text = markdown_content or ""
|
||||
if delimiter and delimiter not in {"\n", "\r\n"} and "`" not in delimiter:
|
||||
for part in text.split(delimiter):
|
||||
block = part.strip()
|
||||
if block:
|
||||
sections.append((block, ""))
|
||||
else:
|
||||
for line in text.splitlines():
|
||||
block = line.strip()
|
||||
if not block:
|
||||
continue
|
||||
sections.append((block, ""))
|
||||
|
||||
if not sections and text.strip():
|
||||
sections.append((text.strip(), ""))
|
||||
|
||||
return sections
|
||||
|
||||
|
||||
def _ensure_chunk_token_limit(chunks: list[str], chunk_token_num: int) -> list[str]:
|
||||
"""对输出 chunk 做 token 上限保护:默认 512 token 时允许到 768 再硬切。"""
|
||||
max_tokens = int(chunk_token_num or 0)
|
||||
if max_tokens <= 0:
|
||||
return [c.strip() for c in chunks if c and c.strip()]
|
||||
|
||||
hard_limit = max(max_tokens, int(max_tokens * GENERAL_HARD_LIMIT_RATIO))
|
||||
protected: list[str] = []
|
||||
for chunk in chunks:
|
||||
cleaned = (chunk or "").strip()
|
||||
if not cleaned:
|
||||
continue
|
||||
if nlp.count_tokens(cleaned) <= max_tokens:
|
||||
protected.append(cleaned)
|
||||
else:
|
||||
protected.extend(nlp.hard_split_by_token_limit(cleaned, max_tokens, hard_limit_token_num=hard_limit))
|
||||
return protected
|
||||
|
||||
|
||||
def chunk_markdown(markdown_content: str, parser_config: dict[str, Any] | None = None) -> list[str]:
|
||||
parser_config = parser_config or {}
|
||||
|
||||
delimiter = _unescape_delimiter(str(parser_config.get("delimiter", "\n") or "\n"))
|
||||
chunk_token_num = int(parser_config.get("chunk_token_num", 512) or 512)
|
||||
overlapped_percent = int(parser_config.get("overlapped_percent", 0) or 0)
|
||||
|
||||
sections = _iter_sections(markdown_content, delimiter)
|
||||
chunks = nlp.naive_merge(
|
||||
sections,
|
||||
chunk_token_num=chunk_token_num,
|
||||
delimiter=delimiter,
|
||||
overlapped_percent=overlapped_percent,
|
||||
)
|
||||
return _ensure_chunk_token_limit(chunks, chunk_token_num)
|
||||
@@ -0,0 +1,183 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from typing import Any
|
||||
|
||||
from yuxi.knowledge.chunking.ragflow_like import nlp
|
||||
|
||||
_ARTICLE_PATTERN = re.compile(r"^(第[零一二三四五六七八九十百千万0-9]+条)[\s ::]*(.*)$")
|
||||
|
||||
|
||||
def _unescape_delimiter(delimiter: str) -> str:
|
||||
return delimiter.replace("\\n", "\n").replace("\\r", "\r").replace("\\t", "\t").replace("\\\\", "\\")
|
||||
|
||||
|
||||
def _iter_lines(markdown_content: str) -> list[str]:
|
||||
return [line.strip() for line in (markdown_content or "").splitlines() if line.strip()]
|
||||
|
||||
|
||||
def _normalize_law_line(line: str) -> str:
|
||||
# 法规 markdown 常见的 #、-、** 装饰会干扰层级识别,这里先做轻量归一化。
|
||||
text = (line or "").strip()
|
||||
text = re.sub(r"^#{1,6}\s+", "", text)
|
||||
text = re.sub(r"^[-*+]\s+", "", text)
|
||||
text = text.replace("**", "").replace("__", "").replace("`", "")
|
||||
text = re.sub(r"[ \t]+", " ", text)
|
||||
return text.strip()
|
||||
|
||||
|
||||
def _expand_article_line(line: str) -> list[str]:
|
||||
normalized = _normalize_law_line(line)
|
||||
if not normalized:
|
||||
return []
|
||||
|
||||
matched = _ARTICLE_PATTERN.match(normalized)
|
||||
if not matched:
|
||||
return [normalized]
|
||||
|
||||
article = matched.group(1).strip()
|
||||
body = matched.group(2).strip()
|
||||
if not body:
|
||||
return [article]
|
||||
return [article, body]
|
||||
|
||||
|
||||
def _iter_law_sections(markdown_content: str) -> list[str]:
|
||||
sections: list[str] = []
|
||||
for line in _iter_lines(markdown_content):
|
||||
sections.extend(_expand_article_line(line))
|
||||
return [section for section in sections if section]
|
||||
|
||||
|
||||
def _docx_heading_tree(markdown_content: str) -> list[str]:
|
||||
lines: list[tuple[int, str]] = []
|
||||
level_set: set[int] = set()
|
||||
|
||||
for raw in (markdown_content or "").splitlines():
|
||||
text = raw.strip()
|
||||
if not text:
|
||||
continue
|
||||
|
||||
heading_match = re.match(r"^(#{1,6})\s+(.*)$", text)
|
||||
if heading_match:
|
||||
level = len(heading_match.group(1))
|
||||
value = heading_match.group(2).strip()
|
||||
else:
|
||||
level = 99
|
||||
value = text
|
||||
|
||||
if not value:
|
||||
continue
|
||||
|
||||
lines.append((level, value))
|
||||
level_set.add(level)
|
||||
|
||||
if not lines:
|
||||
return []
|
||||
|
||||
sorted_levels = sorted(level_set)
|
||||
h2_level = sorted_levels[1] if len(sorted_levels) > 1 else 1
|
||||
h2_level = sorted_levels[-2] if h2_level == sorted_levels[-1] and len(sorted_levels) > 2 else h2_level
|
||||
|
||||
root = nlp.Node(level=0, depth=h2_level, texts=[])
|
||||
root.build_tree(lines)
|
||||
return [element for element in root.get_tree() if element]
|
||||
|
||||
|
||||
def _ensure_chunk_token_limit(
|
||||
chunks: list[str], chunk_token_num: int, delimiter: str, overlapped_percent: int
|
||||
) -> list[str]:
|
||||
"""
|
||||
对输出 chunk 做 token 上限保护:
|
||||
1) 先尝试按行用 naive_merge 再切一次;
|
||||
2) 仍超长时才硬切,避免 embeddings 413。
|
||||
"""
|
||||
max_tokens = int(chunk_token_num or 0)
|
||||
normalized = [chunk.strip() for chunk in chunks if chunk and chunk.strip()]
|
||||
if max_tokens <= 0:
|
||||
return normalized
|
||||
|
||||
protected: list[str] = []
|
||||
for chunk in normalized:
|
||||
if nlp.count_tokens(chunk) <= max_tokens:
|
||||
protected.append(chunk)
|
||||
continue
|
||||
|
||||
refined = nlp.naive_merge(
|
||||
[(_line, "") for _line in _iter_lines(chunk)],
|
||||
chunk_token_num=max_tokens,
|
||||
delimiter=delimiter,
|
||||
overlapped_percent=overlapped_percent,
|
||||
)
|
||||
if not refined:
|
||||
refined = [chunk]
|
||||
|
||||
for item in refined:
|
||||
cleaned = (item or "").strip()
|
||||
if not cleaned:
|
||||
continue
|
||||
if nlp.count_tokens(cleaned) <= max_tokens:
|
||||
protected.append(cleaned)
|
||||
else:
|
||||
sentence_refined = nlp.naive_merge(
|
||||
[(_sentence, "") for _sentence in re.split(r"(?<=[。!?;;!?])", cleaned) if _sentence.strip()],
|
||||
chunk_token_num=max_tokens,
|
||||
delimiter=delimiter,
|
||||
overlapped_percent=overlapped_percent,
|
||||
)
|
||||
if not sentence_refined:
|
||||
sentence_refined = [cleaned]
|
||||
|
||||
for sentence_chunk in sentence_refined:
|
||||
text = sentence_chunk.strip()
|
||||
if not text:
|
||||
continue
|
||||
if nlp.count_tokens(text) <= max_tokens:
|
||||
protected.append(text)
|
||||
else:
|
||||
protected.extend(nlp.hard_split_by_token_limit(text, max_tokens))
|
||||
|
||||
return [chunk for chunk in protected if chunk.strip()]
|
||||
|
||||
|
||||
def chunk_markdown(filename: str, markdown_content: str, parser_config: dict[str, Any] | None = None) -> list[str]:
|
||||
"""
|
||||
法规分块主流程(简化版):
|
||||
- docx 优先尝试标题树;
|
||||
- 其余格式先做法规文本归一化,再按章/节/条树形切分(depth=3);
|
||||
- 最后统一执行超长保护。
|
||||
"""
|
||||
parser_config = parser_config or {}
|
||||
|
||||
delimiter = _unescape_delimiter(str(parser_config.get("delimiter", "\n") or "\n"))
|
||||
chunk_token_num = int(parser_config.get("chunk_token_num", 512) or 512)
|
||||
overlapped_percent = int(parser_config.get("overlapped_percent", 0) or 0)
|
||||
|
||||
if re.search(r"\.docx$", filename or "", re.IGNORECASE):
|
||||
docx_chunks = _docx_heading_tree(markdown_content)
|
||||
if len(docx_chunks) > 1:
|
||||
return _ensure_chunk_token_limit(docx_chunks, chunk_token_num, delimiter, overlapped_percent)
|
||||
|
||||
sections = _iter_law_sections(markdown_content)
|
||||
if not sections:
|
||||
return []
|
||||
|
||||
nlp.remove_contents_table(sections, eng=nlp.is_english(sections))
|
||||
typed_sections = [(s, "") for s in sections]
|
||||
nlp.make_colon_as_title(typed_sections)
|
||||
|
||||
bull = nlp.bullets_category([s for s, _ in typed_sections])
|
||||
if bull == nlp.MARKDOWN_BULLET_GROUP_INDEX:
|
||||
# 归一化后仍命中 markdown 组时,回退到中文法规层级组,避免退化成“按章大块”。
|
||||
bull = 0
|
||||
|
||||
merged = nlp.tree_merge(bull, typed_sections, depth=3)
|
||||
if not merged:
|
||||
merged = nlp.naive_merge(
|
||||
typed_sections,
|
||||
chunk_token_num=chunk_token_num,
|
||||
delimiter=delimiter,
|
||||
overlapped_percent=overlapped_percent,
|
||||
)
|
||||
|
||||
return _ensure_chunk_token_limit(merged, chunk_token_num, delimiter, overlapped_percent)
|
||||
@@ -0,0 +1,260 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import csv
|
||||
import re
|
||||
from typing import Any
|
||||
|
||||
|
||||
def _rm_prefix(text: str) -> str:
|
||||
return re.sub(
|
||||
r"^(问题|答案|回答|user|assistant|Q|A|Question|Answer|问|答)[\t:: ]+",
|
||||
"",
|
||||
(text or "").strip(),
|
||||
flags=re.IGNORECASE,
|
||||
)
|
||||
|
||||
|
||||
def _to_qa_chunk(question: str, answer: str, eng: bool = False) -> str:
|
||||
qprefix = "Question: " if eng else "问题:"
|
||||
aprefix = "Answer: " if eng else "回答:"
|
||||
return "\t".join([qprefix + _rm_prefix(question), aprefix + _rm_prefix(answer)])
|
||||
|
||||
|
||||
def _guess_delimiter(lines: list[str]) -> str:
|
||||
comma = 0
|
||||
tab = 0
|
||||
for line in lines:
|
||||
if len(line.split(",")) == 2:
|
||||
comma += 1
|
||||
if len(line.split("\t")) == 2:
|
||||
tab += 1
|
||||
return "\t" if tab >= comma else ","
|
||||
|
||||
|
||||
def _extract_pairs_with_delimiter(lines: list[str], delimiter: str) -> list[tuple[str, str]]:
|
||||
pairs: list[tuple[str, str]] = []
|
||||
question = ""
|
||||
answer = ""
|
||||
|
||||
for line in lines:
|
||||
arr = line.split(delimiter)
|
||||
if len(arr) != 2:
|
||||
if question:
|
||||
answer += "\n" + line
|
||||
continue
|
||||
|
||||
if question and answer:
|
||||
pairs.append((question, answer))
|
||||
question, answer = arr
|
||||
|
||||
if question:
|
||||
pairs.append((question, answer))
|
||||
|
||||
return [(q.strip(), a.strip()) for q, a in pairs if q.strip()]
|
||||
|
||||
|
||||
def _extract_pairs_from_csv(lines: list[str], delimiter: str) -> list[tuple[str, str]]:
|
||||
pairs: list[tuple[str, str]] = []
|
||||
question = ""
|
||||
answer = ""
|
||||
|
||||
reader = csv.reader(lines, delimiter=delimiter)
|
||||
for row, raw_line in zip(reader, lines, strict=False):
|
||||
if len(row) != 2:
|
||||
if question:
|
||||
answer += "\n" + raw_line
|
||||
continue
|
||||
|
||||
if question and answer:
|
||||
pairs.append((question, answer))
|
||||
question, answer = row
|
||||
|
||||
if question:
|
||||
pairs.append((question, answer))
|
||||
|
||||
return [(q.strip(), a.strip()) for q, a in pairs if q.strip()]
|
||||
|
||||
|
||||
def _parse_markdown_table_row(line: str) -> list[str] | None:
|
||||
if "|" not in line:
|
||||
return None
|
||||
|
||||
text = line.strip()
|
||||
if not text:
|
||||
return None
|
||||
|
||||
if text.startswith("|"):
|
||||
text = text[1:]
|
||||
if text.endswith("|"):
|
||||
text = text[:-1]
|
||||
|
||||
cells = [cell.strip() for cell in text.split("|")]
|
||||
if not cells:
|
||||
return None
|
||||
|
||||
if all(re.fullmatch(r":?-{3,}:?", c.replace(" ", "")) for c in cells if c):
|
||||
return None
|
||||
|
||||
return cells
|
||||
|
||||
|
||||
def _extract_pairs_from_markdown_tables(markdown_content: str) -> list[tuple[str, str]]:
|
||||
pairs: list[tuple[str, str]] = []
|
||||
|
||||
for line in (markdown_content or "").splitlines():
|
||||
cells = _parse_markdown_table_row(line)
|
||||
if not cells or len(cells) < 2:
|
||||
continue
|
||||
|
||||
question = cells[0]
|
||||
answer = cells[1]
|
||||
if question and answer:
|
||||
pairs.append((question, answer))
|
||||
|
||||
return pairs
|
||||
|
||||
|
||||
def _md_question_level(line: str) -> tuple[int, str]:
|
||||
match = re.match(r"^#*", line)
|
||||
if not match:
|
||||
return 0, line
|
||||
return len(match.group(0)), line.lstrip("#").lstrip()
|
||||
|
||||
|
||||
def _extract_pairs_from_markdown_headings(markdown_content: str) -> list[tuple[str, str]]:
|
||||
lines = (markdown_content or "").splitlines()
|
||||
if not lines:
|
||||
return []
|
||||
|
||||
pairs: list[tuple[str, str]] = []
|
||||
last_answer = ""
|
||||
question_stack: list[str] = []
|
||||
level_stack: list[int] = []
|
||||
code_block = False
|
||||
|
||||
for line in lines:
|
||||
if line.strip().startswith("```"):
|
||||
code_block = not code_block
|
||||
|
||||
question_level = 0
|
||||
question = ""
|
||||
if not code_block:
|
||||
question_level, question = _md_question_level(line)
|
||||
|
||||
if not question_level or question_level > 6:
|
||||
last_answer = f"{last_answer}\n{line}"
|
||||
continue
|
||||
|
||||
if last_answer.strip():
|
||||
sum_question = "\n".join(question_stack)
|
||||
if sum_question:
|
||||
pairs.append((sum_question, last_answer.strip()))
|
||||
last_answer = ""
|
||||
|
||||
while question_stack and question_level <= level_stack[-1]:
|
||||
question_stack.pop()
|
||||
level_stack.pop()
|
||||
|
||||
question_stack.append(question)
|
||||
level_stack.append(question_level)
|
||||
|
||||
if last_answer.strip():
|
||||
sum_question = "\n".join(question_stack)
|
||||
if sum_question:
|
||||
pairs.append((sum_question, last_answer.strip()))
|
||||
|
||||
return pairs
|
||||
|
||||
|
||||
def _extract_pairs_by_prefix(lines: list[str]) -> list[tuple[str, str]]:
|
||||
pairs: list[tuple[str, str]] = []
|
||||
question = ""
|
||||
answer_lines: list[str] = []
|
||||
|
||||
for line in lines:
|
||||
if re.match(r"^(Q|Question|问|问题)\s*[::]", line, flags=re.IGNORECASE):
|
||||
if question:
|
||||
pairs.append((question, "\n".join(answer_lines)))
|
||||
question = re.sub(r"^(Q|Question|问|问题)\s*[::]", "", line, flags=re.IGNORECASE).strip()
|
||||
answer_lines = []
|
||||
continue
|
||||
|
||||
if re.match(r"^(A|Answer|答|回答)\s*[::]", line, flags=re.IGNORECASE):
|
||||
answer_lines.append(re.sub(r"^(A|Answer|答|回答)\s*[::]", "", line, flags=re.IGNORECASE).strip())
|
||||
continue
|
||||
|
||||
if question:
|
||||
answer_lines.append(line)
|
||||
|
||||
if question:
|
||||
pairs.append((question, "\n".join(answer_lines)))
|
||||
|
||||
return [(q.strip(), a.strip()) for q, a in pairs if q.strip() and a.strip()]
|
||||
|
||||
|
||||
def _dedupe_pairs(pairs: list[tuple[str, str]]) -> list[tuple[str, str]]:
|
||||
res: list[tuple[str, str]] = []
|
||||
seen: set[tuple[str, str]] = set()
|
||||
|
||||
for question, answer in pairs:
|
||||
q = question.strip()
|
||||
a = answer.strip()
|
||||
if not q or not a:
|
||||
continue
|
||||
key = (q, a)
|
||||
if key in seen:
|
||||
continue
|
||||
seen.add(key)
|
||||
res.append((q, a))
|
||||
|
||||
return res
|
||||
|
||||
|
||||
def chunk_markdown(filename: str, markdown_content: str, parser_config: dict[str, Any] | None = None) -> list[str]:
|
||||
parser_config = parser_config or {}
|
||||
eng = str(parser_config.get("language", "Chinese")).lower() == "english"
|
||||
|
||||
suffix = ""
|
||||
if filename and "." in filename:
|
||||
suffix = "." + filename.lower().split(".")[-1]
|
||||
|
||||
lines = [line for line in (markdown_content or "").splitlines() if line.strip()]
|
||||
pairs: list[tuple[str, str]] = []
|
||||
|
||||
if suffix in {".xlsx", ".xls"}:
|
||||
pairs.extend(_extract_pairs_from_markdown_tables(markdown_content))
|
||||
if not pairs:
|
||||
delimiter = _guess_delimiter(lines)
|
||||
pairs.extend(_extract_pairs_with_delimiter(lines, delimiter))
|
||||
elif suffix == ".csv":
|
||||
pairs.extend(_extract_pairs_from_markdown_tables(markdown_content))
|
||||
delimiter = "\t" if any("\t" in line for line in lines) else ","
|
||||
pairs.extend(_extract_pairs_from_csv(lines, delimiter))
|
||||
elif suffix == ".txt":
|
||||
delimiter = _guess_delimiter(lines)
|
||||
pairs.extend(_extract_pairs_with_delimiter(lines, delimiter))
|
||||
elif suffix in {".md", ".markdown", ".mdx"}:
|
||||
pairs.extend(_extract_pairs_from_markdown_headings(markdown_content))
|
||||
pairs.extend(_extract_pairs_from_markdown_tables(markdown_content))
|
||||
elif suffix == ".docx":
|
||||
pairs.extend(_extract_pairs_from_markdown_headings(markdown_content))
|
||||
pairs.extend(_extract_pairs_from_markdown_tables(markdown_content))
|
||||
else:
|
||||
pairs.extend(_extract_pairs_from_markdown_headings(markdown_content))
|
||||
pairs.extend(_extract_pairs_from_markdown_tables(markdown_content))
|
||||
pairs.extend(_extract_pairs_by_prefix(lines))
|
||||
if not pairs:
|
||||
delimiter = _guess_delimiter(lines)
|
||||
pairs.extend(_extract_pairs_with_delimiter(lines, delimiter))
|
||||
|
||||
pairs = _dedupe_pairs(pairs)
|
||||
|
||||
if not pairs and lines:
|
||||
# 最后兜底:把内容按 2 行一组构成问答
|
||||
for i in range(0, len(lines), 2):
|
||||
q = lines[i]
|
||||
a = lines[i + 1] if i + 1 < len(lines) else ""
|
||||
if q.strip() and a.strip():
|
||||
pairs.append((q, a))
|
||||
|
||||
return [_to_qa_chunk(q, a, eng=eng) for q, a in pairs]
|
||||
@@ -0,0 +1,298 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from typing import Any
|
||||
|
||||
from markdown_it import MarkdownIt
|
||||
from mdit_py_plugins.dollarmath import dollarmath_plugin
|
||||
|
||||
from yuxi.knowledge.chunking.ragflow_like.nlp import count_tokens
|
||||
from yuxi.utils.logging_config import logger
|
||||
|
||||
from ..utils.md_parser_utils import (
|
||||
extract_table_block,
|
||||
get_title_path,
|
||||
split_text_by_length_and_newline,
|
||||
)
|
||||
from ..utils.table_utils import html_table_to_key_value
|
||||
|
||||
|
||||
def _flush_content(
|
||||
result: list,
|
||||
current_content: list,
|
||||
title_stack: list,
|
||||
max_length: int,
|
||||
embed_fn: Any,
|
||||
special_element: str = None,
|
||||
allow_split: bool = False,
|
||||
) -> None:
|
||||
if not current_content:
|
||||
return
|
||||
|
||||
content = "\n".join(current_content).strip()
|
||||
if not content:
|
||||
current_content.clear()
|
||||
return
|
||||
|
||||
level = next((i + 1 for i in range(5, -1, -1) if title_stack[i]), 1)
|
||||
title_path = get_title_path(title_stack)
|
||||
|
||||
if special_element and not allow_split:
|
||||
header = f"{'#' * level} {title_path}|{special_element}" if title_path else f"{'#' * level} {special_element}"
|
||||
result.extend([header, content, "-" * 10])
|
||||
else:
|
||||
if count_tokens(content) > max_length:
|
||||
chunks = split_text_by_length_and_newline(
|
||||
content, max_length, embed_fn=embed_fn, token_count_fn=count_tokens
|
||||
)
|
||||
for idx, chunk in enumerate(chunks, 1):
|
||||
base_header = f"{'#' * level} {title_path}" if title_path else f"{'#' * level}"
|
||||
if special_element:
|
||||
header = f"{base_header}|{special_element}|Part {idx}"
|
||||
else:
|
||||
header = f"{base_header}|Part {idx}"
|
||||
result.extend([header, chunk, "-" * 10])
|
||||
else:
|
||||
base_header = f"{'#' * level} {title_path}" if title_path else f"{'#' * level}"
|
||||
if special_element:
|
||||
header = f"{base_header}|{special_element}"
|
||||
else:
|
||||
header = base_header
|
||||
|
||||
if header:
|
||||
result.append(header)
|
||||
result.append("")
|
||||
result.extend([content, "-" * 10])
|
||||
|
||||
current_content.clear()
|
||||
|
||||
|
||||
def _handle_image_caption(tokens, i, result, current_content, title_stack, max_length, embed_fn):
|
||||
token = tokens[i]
|
||||
if token.type != "paragraph_open":
|
||||
return False, i
|
||||
|
||||
inline_token = tokens[i + 1]
|
||||
if inline_token.type != "inline":
|
||||
return False, i
|
||||
|
||||
content = inline_token.content.strip()
|
||||
image_pattern = r"^!\[.*?\]\(.*?\)\s*$"
|
||||
caption_pattern = r"^(?:Figure|图|Fig\.|表|Table)\s*[\d\w\.]+"
|
||||
|
||||
img_match = re.search(r"^(!\[.*?\]\(.*?\))", content)
|
||||
if img_match:
|
||||
rest = content[img_match.end() :].strip()
|
||||
if rest and re.match(caption_pattern, rest, re.IGNORECASE):
|
||||
_flush_content(result, current_content, title_stack, max_length, embed_fn)
|
||||
current_content.append(content)
|
||||
caption_title = rest.split("\n")[0].strip()
|
||||
_flush_content(result, current_content, title_stack, max_length, embed_fn, special_element=caption_title)
|
||||
return True, i + 3
|
||||
|
||||
if re.match(image_pattern, content):
|
||||
next_p_idx = i + 3
|
||||
if next_p_idx + 1 < len(tokens) and tokens[next_p_idx].type == "paragraph_open":
|
||||
next_inline = tokens[next_p_idx + 1]
|
||||
if next_inline.type == "inline":
|
||||
next_content = next_inline.content.strip()
|
||||
if re.match(caption_pattern, next_content, re.IGNORECASE):
|
||||
_flush_content(result, current_content, title_stack, max_length, embed_fn)
|
||||
current_content.append(content)
|
||||
current_content.append(next_content)
|
||||
_flush_content(
|
||||
result, current_content, title_stack, max_length, embed_fn, special_element=next_content
|
||||
)
|
||||
return True, i + 6
|
||||
|
||||
if current_content and re.match(caption_pattern, content, re.IGNORECASE):
|
||||
last_item = current_content[-1].strip()
|
||||
if re.match(image_pattern, last_item):
|
||||
image_tag = current_content.pop()
|
||||
_flush_content(result, current_content, title_stack, max_length, embed_fn)
|
||||
current_content.append(image_tag)
|
||||
current_content.append(content)
|
||||
_flush_content(result, current_content, title_stack, max_length, embed_fn, special_element=content)
|
||||
return True, i + 3
|
||||
|
||||
return False, i
|
||||
|
||||
|
||||
def chunk_markdown(
|
||||
markdown_content: str, parser_config: dict[str, Any] | None = None, embed_fn: Any | None = None
|
||||
) -> list[str]:
|
||||
"""
|
||||
语义化切分 Markdown 内容。
|
||||
|
||||
Args:
|
||||
markdown_content: 待切分的 Markdown 文本
|
||||
parser_config: 切分参数,如 chunk_token_num
|
||||
embed_fn: 可选。传入用于生成向量的函数。如果不传,将从系统配置中加载模型。
|
||||
通过注入此参数可以避免在单元测试中加载重型资源。
|
||||
"""
|
||||
parser_config = parser_config or {}
|
||||
max_length = int(parser_config.get("chunk_token_num", 512))
|
||||
logger.info(f"语义切分开始: max_length={max_length}, content_length={len(markdown_content)}")
|
||||
|
||||
# 延迟加载重型资源,仅在没有注入 embed_fn 时触发
|
||||
if embed_fn is None:
|
||||
try:
|
||||
from yuxi.config.app import config
|
||||
from yuxi.models.embed import select_embedding_model
|
||||
|
||||
embed_model_id = parser_config.get("embed_model_id") or config.embed_model
|
||||
logger.info(f"语义切分加载Embedding模型: {embed_model_id}")
|
||||
embed_model = select_embedding_model(embed_model_id)
|
||||
embed_fn = embed_model.encode
|
||||
except Exception as e:
|
||||
logger.error(f"加载 Embedding 模型失败: {e}。将退化为简单切分。")
|
||||
embed_fn = None
|
||||
|
||||
md = MarkdownIt("commonmark").enable("table")
|
||||
md.use(dollarmath_plugin, allow_space=True, allow_digits=True)
|
||||
|
||||
tokens: list = md.parse(markdown_content)
|
||||
original_lines: list = markdown_content.split("\n")
|
||||
|
||||
result: list = []
|
||||
current_content: list = []
|
||||
title_stack: list = [""] * 6
|
||||
|
||||
i = 0
|
||||
while i < len(tokens):
|
||||
token = tokens[i]
|
||||
if token.type == "heading_open":
|
||||
_flush_content(result, current_content, title_stack, max_length, embed_fn)
|
||||
level = int(token.tag[1:]) if token.tag and len(token.tag) > 1 else 1
|
||||
inline_token = tokens[i + 1]
|
||||
if inline_token.type == "inline":
|
||||
full_title = inline_token.content.strip()
|
||||
title_stack[level - 1] = full_title
|
||||
for j in range(level, 6):
|
||||
title_stack[j] = ""
|
||||
i += 3
|
||||
continue
|
||||
elif token.type == "table_open":
|
||||
_flush_content(result, current_content, title_stack, max_length, embed_fn)
|
||||
j, table_content = extract_table_block(tokens, i, original_lines)
|
||||
current_content.append(table_content)
|
||||
_flush_content(result, current_content, title_stack, max_length, embed_fn, special_element="Table")
|
||||
i = j + 1 if j < len(tokens) else len(tokens)
|
||||
continue
|
||||
elif token.type == "paragraph_open":
|
||||
handled, new_i = _handle_image_caption(
|
||||
tokens, i, result, current_content, title_stack, max_length, embed_fn
|
||||
)
|
||||
if handled:
|
||||
i = new_i
|
||||
continue
|
||||
inline_token = tokens[i + 1]
|
||||
if inline_token.type == "inline":
|
||||
current_content.append(inline_token.content.strip())
|
||||
i += 3
|
||||
continue
|
||||
elif token.type == "fence":
|
||||
current_content.append(f"```\n{token.content}\n```")
|
||||
i += 1
|
||||
continue
|
||||
elif token.type == "ordered_list_open":
|
||||
# _flush_content(result, current_content, title_stack, max_length, embed_fn)
|
||||
list_content = []
|
||||
j = i + 1
|
||||
list_item_counter = 1
|
||||
while j < len(tokens) and tokens[j].type != "ordered_list_close":
|
||||
if tokens[j].type == "list_item_open":
|
||||
k = j + 1
|
||||
while k < len(tokens) and tokens[k].type != "list_item_close":
|
||||
if (
|
||||
tokens[k].type == "paragraph_open"
|
||||
and k + 1 < len(tokens)
|
||||
and tokens[k + 1].type == "inline"
|
||||
):
|
||||
list_content.append(f"{list_item_counter}. {tokens[k + 1].content.strip()}")
|
||||
list_item_counter += 1
|
||||
k += 1
|
||||
j += 1
|
||||
if list_content:
|
||||
current_content.extend(list_content)
|
||||
_flush_content(result, current_content, title_stack, max_length, embed_fn, special_element=token.type)
|
||||
i = j + 1
|
||||
continue
|
||||
elif token.type == "bullet_list_open":
|
||||
# _flush_content(result, current_content, title_stack, max_length, embed_fn)
|
||||
list_content = []
|
||||
j = i + 1
|
||||
while j < len(tokens) and tokens[j].type != "bullet_list_close":
|
||||
if tokens[j].type == "list_item_open":
|
||||
k = j + 1
|
||||
while k < len(tokens) and tokens[k].type != "list_item_close":
|
||||
if (
|
||||
tokens[k].type == "paragraph_open"
|
||||
and k + 1 < len(tokens)
|
||||
and tokens[k + 1].type == "inline"
|
||||
):
|
||||
list_content.append(f"- {tokens[k + 1].content.strip()}")
|
||||
k += 1
|
||||
j += 1
|
||||
if list_content:
|
||||
current_content.extend(list_content)
|
||||
_flush_content(result, current_content, title_stack, max_length, embed_fn, special_element=token.type)
|
||||
i = j + 1
|
||||
continue
|
||||
elif token.type == "html_block":
|
||||
_flush_content(result, current_content, title_stack, max_length, embed_fn)
|
||||
content = token.content.strip()
|
||||
is_converted_table = False
|
||||
if "<table" in content.lower():
|
||||
try:
|
||||
kv_list = html_table_to_key_value(content)
|
||||
if kv_list:
|
||||
content = "\n".join([f"- {item}" for item in kv_list])
|
||||
is_converted_table = True
|
||||
except Exception as e:
|
||||
logger.warning(f"HTML表格转KV失败: {e}")
|
||||
|
||||
current_content.append(content)
|
||||
if is_converted_table:
|
||||
_flush_content(
|
||||
result,
|
||||
current_content,
|
||||
title_stack,
|
||||
max_length,
|
||||
embed_fn,
|
||||
special_element="Table KV",
|
||||
allow_split=True,
|
||||
)
|
||||
else:
|
||||
_flush_content(result, current_content, title_stack, max_length, embed_fn, special_element=token.type)
|
||||
i += 1
|
||||
continue
|
||||
elif token.type in ["list_item_close", "ordered_list_close", "bullet_list_close", "list_item_open"]:
|
||||
i += 1
|
||||
continue
|
||||
elif token.type == "math_block":
|
||||
# _flush_content(result, current_content, title_stack, max_length, embed_fn)
|
||||
current_content.append(f"$ {token.content} $")
|
||||
_flush_content(result, current_content, title_stack, max_length, embed_fn, special_element="Math Block")
|
||||
i += 1
|
||||
continue
|
||||
else:
|
||||
i += 1
|
||||
|
||||
_flush_content(result, current_content, title_stack, max_length, embed_fn)
|
||||
|
||||
chunks = []
|
||||
current_chunk_parts = []
|
||||
for item in result:
|
||||
if item == "-" * 10:
|
||||
if current_chunk_parts:
|
||||
chunks.append("\n".join(current_chunk_parts).strip())
|
||||
current_chunk_parts = []
|
||||
else:
|
||||
current_chunk_parts.append(item)
|
||||
|
||||
if current_chunk_parts:
|
||||
chunks.append("\n".join(current_chunk_parts).strip())
|
||||
|
||||
logger.info(f"语义切分完成: chunks={len(chunks)}")
|
||||
return chunks
|
||||
@@ -0,0 +1,86 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from yuxi.knowledge.chunking.ragflow_like import nlp
|
||||
from yuxi.knowledge.chunking.ragflow_like.parsers.general import _iter_sections, _unescape_delimiter
|
||||
|
||||
|
||||
def _slice_text_by_tokens(text: str, max_tokens: int, overlap_tokens: int) -> list[str]:
|
||||
if max_tokens <= 0:
|
||||
return [text] if text.strip() else []
|
||||
|
||||
units = [part for part in text]
|
||||
chunks: list[str] = []
|
||||
start = 0
|
||||
|
||||
while start < len(units):
|
||||
current = ""
|
||||
current_tokens = 0
|
||||
end = start
|
||||
|
||||
while end < len(units):
|
||||
next_text = current + units[end]
|
||||
next_tokens = nlp.count_tokens(next_text)
|
||||
if current and next_tokens > max_tokens:
|
||||
break
|
||||
current = next_text
|
||||
current_tokens = next_tokens
|
||||
end += 1
|
||||
if current_tokens >= max_tokens:
|
||||
break
|
||||
|
||||
chunk = current.strip()
|
||||
if chunk:
|
||||
chunks.append(chunk)
|
||||
|
||||
if end >= len(units):
|
||||
break
|
||||
|
||||
if overlap_tokens <= 0:
|
||||
start = end
|
||||
continue
|
||||
|
||||
backtrack = end
|
||||
overlap_text = ""
|
||||
while backtrack > start:
|
||||
candidate = units[backtrack - 1] + overlap_text
|
||||
if nlp.count_tokens(candidate) > overlap_tokens:
|
||||
break
|
||||
overlap_text = candidate
|
||||
backtrack -= 1
|
||||
|
||||
start = backtrack if backtrack < end else end
|
||||
|
||||
return chunks
|
||||
|
||||
|
||||
def _split_section_with_overlap(section: str, chunk_token_num: int, overlapped_percent: int) -> list[str]:
|
||||
overlap_tokens = 0
|
||||
if chunk_token_num > 0 and overlapped_percent > 0:
|
||||
overlap_tokens = int(chunk_token_num * max(0, min(overlapped_percent, 99)) / 100)
|
||||
return _slice_text_by_tokens(section, chunk_token_num, overlap_tokens)
|
||||
|
||||
|
||||
def chunk_markdown(markdown_content: str, parser_config: dict[str, Any] | None = None) -> list[str]:
|
||||
parser_config = parser_config or {}
|
||||
|
||||
delimiter = _unescape_delimiter(str(parser_config.get("delimiter", "\n") or "\n"))
|
||||
chunk_token_num = int(parser_config.get("chunk_token_num", 512) or 512)
|
||||
overlapped_percent = int(parser_config.get("overlapped_percent", 0) or 0)
|
||||
|
||||
sections = _iter_sections(markdown_content, delimiter)
|
||||
chunks: list[str] = []
|
||||
|
||||
for section, _ in sections:
|
||||
text = (section or "").strip()
|
||||
if not text:
|
||||
continue
|
||||
|
||||
if chunk_token_num > 0 and nlp.count_tokens(text) > chunk_token_num:
|
||||
chunks.extend(_split_section_with_overlap(text, chunk_token_num, overlapped_percent))
|
||||
continue
|
||||
|
||||
chunks.append(text)
|
||||
|
||||
return chunks
|
||||
@@ -0,0 +1,129 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from copy import deepcopy
|
||||
from typing import Any
|
||||
|
||||
from yuxi.utils import logger
|
||||
|
||||
DEFAULT_CHUNK_PRESET_ID = "general"
|
||||
|
||||
CHUNK_PRESETS: dict[str, dict[str, str]] = {
|
||||
"general": {
|
||||
"label": "General",
|
||||
"description": "通用分块:按分隔符和长度切分,适合大多数普通文档。",
|
||||
},
|
||||
"qa": {
|
||||
"label": "QA",
|
||||
"description": "问答分块:优先抽取问题-回答结构,适合 FAQ、题库、问答手册。",
|
||||
},
|
||||
"book": {
|
||||
"label": "Book",
|
||||
"description": "书籍分块:强化章节标题识别并做层级合并,适合教材、手册、长章节文档。",
|
||||
},
|
||||
"laws": {
|
||||
"label": "Laws",
|
||||
"description": "法规分块:按法条层级组织与合并,适合法律法规、制度规范类文本。",
|
||||
},
|
||||
"semantic": {
|
||||
"label": "Semantic",
|
||||
"description": "语义分块:利用嵌入和聚类算法进行语义切分,并自动增强标题上下文。",
|
||||
},
|
||||
"separator": {
|
||||
"label": "Separator",
|
||||
"description": "严格分隔:命中分隔符即切分,仅超长片段内部继续按长度切分。",
|
||||
},
|
||||
}
|
||||
|
||||
CHUNK_PRESET_IDS = set(CHUNK_PRESETS)
|
||||
|
||||
CHUNK_ENGINE_VERSION = "ragflow_like_v1"
|
||||
GENERAL_INTERNAL_PARSER_ID = "naive"
|
||||
|
||||
|
||||
def deep_merge(base: dict[str, Any], override: dict[str, Any]) -> dict[str, Any]:
|
||||
result = deepcopy(base)
|
||||
for key, value in (override or {}).items():
|
||||
if isinstance(value, dict) and isinstance(result.get(key), dict):
|
||||
result[key] = deep_merge(result[key], value)
|
||||
else:
|
||||
result[key] = value
|
||||
return result
|
||||
|
||||
|
||||
def normalize_chunk_preset_id(value: str | None) -> str:
|
||||
if not value:
|
||||
return DEFAULT_CHUNK_PRESET_ID
|
||||
|
||||
normalized = str(value).strip().lower()
|
||||
if normalized == GENERAL_INTERNAL_PARSER_ID:
|
||||
return DEFAULT_CHUNK_PRESET_ID
|
||||
|
||||
if normalized in CHUNK_PRESET_IDS:
|
||||
return normalized
|
||||
|
||||
logger.warning(f"Unknown chunk preset id '{value}', fallback to general")
|
||||
return DEFAULT_CHUNK_PRESET_ID
|
||||
|
||||
|
||||
def map_to_internal_parser_id(preset_id: str) -> str:
|
||||
normalized = normalize_chunk_preset_id(preset_id)
|
||||
if normalized == DEFAULT_CHUNK_PRESET_ID:
|
||||
return GENERAL_INTERNAL_PARSER_ID
|
||||
return normalized
|
||||
|
||||
|
||||
def get_default_chunk_parser_config(preset_id: str) -> dict[str, Any]:
|
||||
normalize_chunk_preset_id(preset_id)
|
||||
return {}
|
||||
|
||||
|
||||
def ensure_chunk_defaults_in_additional_params(additional_params: dict[str, Any] | None) -> dict[str, Any]:
|
||||
params = dict(additional_params or {})
|
||||
params["chunk_preset_id"] = normalize_chunk_preset_id(params.get("chunk_preset_id"))
|
||||
|
||||
if "chunk_parser_config" in params and not isinstance(params.get("chunk_parser_config"), dict):
|
||||
logger.warning("Invalid chunk_parser_config in additional_params, fallback to empty dict")
|
||||
params["chunk_parser_config"] = {}
|
||||
|
||||
return params
|
||||
|
||||
|
||||
def resolve_chunk_processing_params(
|
||||
kb_additional_params: dict[str, Any] | None,
|
||||
file_processing_params: dict[str, Any] | None,
|
||||
request_params: dict[str, Any] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
kb_additional = ensure_chunk_defaults_in_additional_params(kb_additional_params)
|
||||
file_params = dict(file_processing_params or {})
|
||||
request = dict(request_params or {})
|
||||
|
||||
preset_id = normalize_chunk_preset_id(
|
||||
request.get("chunk_preset_id") or file_params.get("chunk_preset_id") or kb_additional.get("chunk_preset_id")
|
||||
)
|
||||
|
||||
parser_config = get_default_chunk_parser_config(preset_id)
|
||||
|
||||
kb_parser_config = kb_additional.get("chunk_parser_config")
|
||||
if isinstance(kb_parser_config, dict):
|
||||
parser_config = deep_merge(parser_config, kb_parser_config)
|
||||
|
||||
file_parser_config = file_params.get("chunk_parser_config")
|
||||
if isinstance(file_parser_config, dict):
|
||||
parser_config = deep_merge(parser_config, file_parser_config)
|
||||
|
||||
req_parser_config = request.get("chunk_parser_config")
|
||||
if isinstance(req_parser_config, dict):
|
||||
parser_config = deep_merge(parser_config, req_parser_config)
|
||||
|
||||
return {
|
||||
"chunk_preset_id": preset_id,
|
||||
"chunk_parser_config": parser_config,
|
||||
"chunk_engine_version": CHUNK_ENGINE_VERSION,
|
||||
}
|
||||
|
||||
|
||||
def get_chunk_preset_options() -> list[dict[str, str]]:
|
||||
return [
|
||||
{"value": preset_id, "label": preset["label"], "description": preset["description"]}
|
||||
for preset_id, preset in CHUNK_PRESETS.items()
|
||||
]
|
||||
@@ -0,0 +1,150 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from collections.abc import Callable
|
||||
from typing import Any
|
||||
|
||||
from .semantic_utils import semantic_chunking_with_auto_clusters
|
||||
|
||||
|
||||
def infer_heading_level(title: str) -> int:
|
||||
"""
|
||||
根据标题文本推断其层级级别(1-6级)。
|
||||
|
||||
逻辑说明:
|
||||
1. 数字序号推断:
|
||||
- 匹配如 "1.", "1.1", "1.2.3" 等格式。
|
||||
- 根据点号分隔的数量确定层级,例如 "1.1" 为 2 级,"1.2.3" 为 3 级。
|
||||
- 层级限制在 1-6 之间。
|
||||
2. 中文序号推断:
|
||||
- 匹配如 "一、", "二." 等中文数字序号。
|
||||
- 统一归类为 1 级标题。
|
||||
3. 默认处理:
|
||||
- 若不匹配以上规则,默认返回 1 级。
|
||||
"""
|
||||
m = re.match(r"^\s*(\d+(?:\.\d+)*)[.)、]?\s*", title)
|
||||
if m:
|
||||
return max(1, min(len(m.group(1).split(".")), 6))
|
||||
m_zh = re.match(r"^\s*[一二三四五六七八九十百千]+[、.]\s*", title)
|
||||
if m_zh:
|
||||
return 1
|
||||
return 1
|
||||
|
||||
|
||||
def get_title_path(stack: list[str]) -> str:
|
||||
"""
|
||||
根据标题栈生成标题路径,用"|"分隔。
|
||||
"""
|
||||
return "|".join([t for t in stack if t])
|
||||
|
||||
|
||||
def extract_table_block(tokens: list[Any], i: int, original_lines: list[str]) -> tuple[int, str]:
|
||||
"""
|
||||
从token流和原始文本中提取完整的表格块。
|
||||
|
||||
逻辑说明:
|
||||
1. 定位起始:通过当前 token (i) 的 `map` 属性获取表格在原始行中的起始行号 `table_start`。
|
||||
2. 查找结束 token:遍历后续 tokens 直到找到 `table_close`。
|
||||
3. 确定结束行号 (`table_end`):
|
||||
- 优先使用 `table_close` token 的 `map` 属性。
|
||||
- 若不存在,则尝试查找下一个带有 `map` 信息的 token 的起始行作为当前表格的结束。
|
||||
- 若上述均失败(如文件末尾或解析异常),则回退到基于文本内容的启发式扫描:
|
||||
从 `table_start` 开始向下扫描,直到遇到不符合 Markdown 表格特征(不以 '|' 开头且不含 '|')的行为止。
|
||||
4. 返回结果:返回 `table_close` 的索引 `j` 以及拼接后的表格原始字符串。
|
||||
"""
|
||||
token = tokens[i]
|
||||
table_start = token.map[0] if token.map else 0
|
||||
j = i + 1
|
||||
while j < len(tokens) and tokens[j].type != "table_close":
|
||||
j += 1
|
||||
if j < len(tokens):
|
||||
end_token = tokens[j]
|
||||
if end_token.map and end_token.map[1] is not None:
|
||||
table_end = end_token.map[1]
|
||||
else:
|
||||
table_end = None
|
||||
for k in range(j + 1, len(tokens)):
|
||||
if tokens[k].map and tokens[k].map[0] is not None:
|
||||
table_end = tokens[k].map[0]
|
||||
break
|
||||
if table_end is None:
|
||||
table_end = table_start + 1
|
||||
for line_idx in range(table_start, len(original_lines)):
|
||||
line = original_lines[line_idx].strip()
|
||||
if not line or not (line.startswith("|") or "|" in line):
|
||||
table_end = line_idx
|
||||
break
|
||||
else:
|
||||
table_end = table_start + 1
|
||||
for line_idx in range(table_start, len(original_lines)):
|
||||
line = original_lines[line_idx].strip()
|
||||
if not line or not (line.startswith("|") or "|" in line):
|
||||
table_end = line_idx
|
||||
break
|
||||
return j, "\n".join(original_lines[table_start:table_end])
|
||||
|
||||
|
||||
def split_text_by_length_and_newline(
|
||||
text: str, max_length: int, embed_fn: Callable[[list[str]], Any] | None, token_count_fn: Callable[[str], int]
|
||||
) -> list[str]:
|
||||
"""
|
||||
层次化文本切分策略。
|
||||
"""
|
||||
chunks = []
|
||||
|
||||
paragraphs = text.split("\n\n")
|
||||
|
||||
for paragraph in paragraphs:
|
||||
paragraph = paragraph.strip()
|
||||
if not paragraph:
|
||||
continue
|
||||
|
||||
paragraph_token_count = token_count_fn(paragraph)
|
||||
|
||||
# 如果当前段落长度未超过最大 Token 数量,直接作为独立分块放入chunks
|
||||
# 否则继续尝试按行切分
|
||||
if paragraph_token_count <= max_length:
|
||||
chunks.append(paragraph)
|
||||
continue
|
||||
|
||||
# 把段落进一步使用换行符进行切分为行
|
||||
lines = paragraph.split("\n")
|
||||
current_chunk_lines = []
|
||||
current_chunk_tokens = 0
|
||||
|
||||
for line in lines:
|
||||
line = line.strip()
|
||||
if not line: # 跳过空行
|
||||
continue
|
||||
|
||||
line_token_count = token_count_fn(line) # 计算当前行的 Token 数量
|
||||
# 为了考虑行之间的空格,需要在计算 Token 数量时加 1(如果当前行不是第一行,需要添加一个换行符的Token数量)
|
||||
added_tokens = line_token_count + (1 if current_chunk_lines else 0)
|
||||
# 如果当前行的 Token 数量超过最大 Token 数量,直接作为独立分块放入chunks
|
||||
if line_token_count > max_length:
|
||||
if current_chunk_lines:
|
||||
chunks.append("\n".join(current_chunk_lines))
|
||||
current_chunk_lines = []
|
||||
current_chunk_tokens = 0
|
||||
|
||||
sub_chunks = semantic_chunking_with_auto_clusters(
|
||||
line, embed_fn=embed_fn, token_count_fn=token_count_fn, max_chunk_size=max_length
|
||||
)
|
||||
chunks.extend(sub_chunks)
|
||||
# 如果当前行的 Token 数量与当前分块的 Token 数量合并后超过最大 Token 数量,直接作为独立分块放入chunks
|
||||
elif current_chunk_tokens + added_tokens > max_length:
|
||||
# 把之前的分块内容放入chunks
|
||||
chunks.append("\n".join(current_chunk_lines))
|
||||
# 重置当前分块为当前行的内容
|
||||
current_chunk_lines = [line]
|
||||
# 更新当前分块的 Token 数量
|
||||
current_chunk_tokens = line_token_count
|
||||
# 如果当前行的内容加入当前分块后不会超过最大 Token 数量,直接加入当前分块
|
||||
else:
|
||||
current_chunk_lines.append(line)
|
||||
current_chunk_tokens += added_tokens # 更新当前分块的 Token 数量
|
||||
# 最后的收尾,把最后一行内容放入chunks
|
||||
if current_chunk_lines:
|
||||
chunks.append("\n".join(current_chunk_lines))
|
||||
|
||||
return chunks
|
||||
@@ -0,0 +1,193 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from collections.abc import Callable
|
||||
from typing import Any
|
||||
|
||||
import nltk
|
||||
from nltk.tokenize import sent_tokenize
|
||||
from sklearn.cluster import AgglomerativeClustering
|
||||
from sklearn.metrics import silhouette_score
|
||||
|
||||
_punkt_checked = False
|
||||
|
||||
|
||||
def _ensure_punkt_tab() -> None:
|
||||
"""首次使用分句能力时检查 NLTK punkt_tab 资源,缺失时给出可操作错误。"""
|
||||
global _punkt_checked
|
||||
if _punkt_checked:
|
||||
return
|
||||
|
||||
try:
|
||||
nltk.data.find("tokenizers/punkt_tab")
|
||||
except LookupError as e:
|
||||
raise RuntimeError("缺少 NLTK 资源 punkt_tab。请先执行: python -m nltk.downloader punkt_tab") from e
|
||||
|
||||
_punkt_checked = True
|
||||
|
||||
|
||||
def split_sentences_chinese(text: str) -> list[str]:
|
||||
"""
|
||||
使用正则表达式将中文文本分割成句子。
|
||||
|
||||
逻辑:
|
||||
- 匹配中文句号、感叹号、问号(。!?)作为分隔点。
|
||||
- 使用正向/反向预查处理引号:确保如果标点后面紧跟引号(”’"),该引号会被保留在当前句子末尾,而不是被切分到下一句。
|
||||
- 返回去除两端空格且非空的句子列表。
|
||||
"""
|
||||
pattern = r'(?<=[。!?][”’"])|(?<=[。!?])(?![”’"])'
|
||||
sentences = re.split(pattern, text)
|
||||
return [s.strip() for s in sentences if s.strip()]
|
||||
|
||||
|
||||
def split_mixed_sentences(text: str) -> list[str]:
|
||||
"""
|
||||
处理中英文混合文本的分句逻辑,支持按物理段落分发不同的分句策略。
|
||||
|
||||
该函数采用“分而治之”的策略来处理复杂的混合文本:
|
||||
1. **物理分块**:首先按换行符 (`\\n+`) 将原始文本切分为多个物理段落(chunks),确保物理结构不被破坏。
|
||||
2. **语言检测与分发**:
|
||||
- **英文/混合路径**:若段落中包含英文字母 (`[A-Za-z]`),则视为英文或混合文本,
|
||||
调用 NLTK 的 `sent_tokenize` 进行处理。NLTK 能更好地处理英文缩写、句点等复杂情况。
|
||||
- **中文路径**:若段落不含字母,则视为纯中文文本,调用 `split_sentences_chinese`。
|
||||
该方法通过正则精准匹配中文标点及后续引号。
|
||||
- **兜底方案**:若上述方法未产生结果,则使用简单的正则表达式按中文标点强制分割。
|
||||
3. **清洗与过滤**:汇总所有子句,去除两端空白字符,并过滤掉空字符串。
|
||||
|
||||
Args:
|
||||
text: 待分句的原始字符串。
|
||||
|
||||
Returns:
|
||||
List[str]: 分割后的句子列表。
|
||||
"""
|
||||
_ensure_punkt_tab()
|
||||
|
||||
chunks = re.split(r"(\n+)", text)
|
||||
sentences = []
|
||||
|
||||
for ch in chunks:
|
||||
if not ch.strip():
|
||||
continue
|
||||
if re.search(r"[A-Za-z]", ch):
|
||||
parts = sent_tokenize(ch)
|
||||
sentences.extend([p.strip() for p in parts if p.strip()])
|
||||
else:
|
||||
sents = split_sentences_chinese(ch)
|
||||
if sents:
|
||||
sentences.extend([s.strip() for s in sents if s.strip()])
|
||||
else:
|
||||
parts = re.split(r"(?<=[。!?])", ch)
|
||||
sentences.extend([p.strip() for p in parts if p.strip()])
|
||||
return sentences
|
||||
|
||||
|
||||
def find_best_num_clusters(embeddings: Any, min_clusters: int = 2, max_clusters: int = 10) -> int:
|
||||
"""
|
||||
使用轮廓系数选择最佳聚类数量。让每个分段语义集中,且分段之间界限分明
|
||||
|
||||
逻辑:
|
||||
- 遍历可能的聚类数量(从 min_clusters 到 max_clusters)。
|
||||
- 对每个聚类数量,使用 `AgglomerativeClustering` 进行聚类。
|
||||
- 计算轮廓系数(Silhouette Score)。
|
||||
- 选择轮廓系数最高的聚类数量作为最佳聚类数量。
|
||||
- 如果聚类数量为 1 或更少,直接返回 1。
|
||||
|
||||
Args:
|
||||
embeddings: 待聚类的向量数据(所有句子的嵌入向量列表)。
|
||||
min_clusters: 搜索的最佳聚类数量下限,默认为 2。
|
||||
max_clusters: 搜索的最佳聚类数量上限,默认为 10。
|
||||
|
||||
Returns:
|
||||
int: 轮廓系数表现最好的聚类数量。
|
||||
"""
|
||||
if len(embeddings) <= min_clusters:
|
||||
return len(embeddings)
|
||||
|
||||
best_score = -1
|
||||
best_k = min_clusters
|
||||
|
||||
limit_k = min(max_clusters, len(embeddings))
|
||||
for k in range(min_clusters, limit_k + 1):
|
||||
labels = AgglomerativeClustering(n_clusters=k, metric="cosine", linkage="average").fit_predict(embeddings)
|
||||
if len(set(labels)) <= 1:
|
||||
continue
|
||||
score = silhouette_score(embeddings, labels, metric="cosine")
|
||||
if score > best_score:
|
||||
best_score = score
|
||||
best_k = k
|
||||
|
||||
return best_k
|
||||
|
||||
|
||||
def semantic_chunking_with_auto_clusters(
|
||||
text: str,
|
||||
embed_fn: Callable[[list[str]], Any] | None,
|
||||
token_count_fn: Callable[[str], int],
|
||||
max_chunk_size: int = 512,
|
||||
) -> list[str]:
|
||||
"""
|
||||
对传入的文本进行语义切分,过程中会自动选择最佳的聚集数量。
|
||||
|
||||
逻辑:
|
||||
- 先将文本中的句子按语言进行分发,英文/混合文本使用NLTK的sent_tokenize,中文文本使用split_sentences_chinese。
|
||||
- 对每个句子进行嵌入向量化。
|
||||
- 确定最佳的聚类数量(根据轮廓系数)。
|
||||
- 对句子进行聚类后,按原文顺序遍历:当聚类标签变化或达到长度上限时切分,形成连续分块。
|
||||
- 如果嵌入模型缺失,则退化为原始切分方式
|
||||
"""
|
||||
sentences = split_mixed_sentences(text)
|
||||
if len(sentences) < 2:
|
||||
return [text.strip()]
|
||||
|
||||
# 计算每个句子的token数量
|
||||
sentence_token_counts = [token_count_fn(s) for s in sentences]
|
||||
total_tokens = sum(sentence_token_counts)
|
||||
|
||||
# 如果没有提供向量化函数,或者整体未超长,则直接进行简单合并/返回
|
||||
if embed_fn is None or total_tokens <= max_chunk_size:
|
||||
chunks = []
|
||||
current_chunk = ""
|
||||
current_chunk_tokens = 0
|
||||
for s, cnt in zip(sentences, sentence_token_counts):
|
||||
if current_chunk_tokens + cnt > max_chunk_size and current_chunk:
|
||||
chunks.append(current_chunk.strip())
|
||||
current_chunk = s
|
||||
current_chunk_tokens = cnt
|
||||
else:
|
||||
current_chunk += s
|
||||
current_chunk_tokens += cnt
|
||||
if current_chunk:
|
||||
chunks.append(current_chunk.strip())
|
||||
return chunks
|
||||
|
||||
# 向量化每个句子, 得到他们的嵌入向量
|
||||
embeddings = embed_fn(sentences)
|
||||
|
||||
# 决定合适的聚集数量:超长时按上限向上取整,避免整除时多切一块
|
||||
best_k = (total_tokens + max_chunk_size - 1) // max_chunk_size
|
||||
best_k = min(best_k, len(sentences))
|
||||
|
||||
# 根据指定的聚集数量、相似度判断方式、联动方式,对句子进行聚类
|
||||
# labels 是每个句子的聚类标签列表(如 [0,0,1,2,2]),后续会按原文顺序在标签变化处切分连续分块
|
||||
labels = AgglomerativeClustering(n_clusters=best_k, metric="cosine", linkage="average").fit_predict(embeddings)
|
||||
|
||||
chunks = []
|
||||
current_chunk = ""
|
||||
current_chunk_tokens = 0
|
||||
current_label = labels[0]
|
||||
|
||||
for sentence, label, token_count in zip(sentences, labels, sentence_token_counts):
|
||||
if label != current_label or current_chunk_tokens + token_count > max_chunk_size:
|
||||
if current_chunk.strip():
|
||||
chunks.append(current_chunk.strip())
|
||||
current_chunk = sentence
|
||||
current_chunk_tokens = token_count
|
||||
current_label = label
|
||||
else:
|
||||
current_chunk += sentence
|
||||
current_chunk_tokens += token_count
|
||||
|
||||
if current_chunk.strip():
|
||||
chunks.append(current_chunk.strip())
|
||||
|
||||
return chunks
|
||||
@@ -0,0 +1,147 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from bs4 import BeautifulSoup
|
||||
|
||||
|
||||
def html_table_to_markdown(html: str) -> str:
|
||||
"""
|
||||
将HTML表格转换为Markdown格式的具体实现
|
||||
"""
|
||||
soup = BeautifulSoup(html, "html.parser")
|
||||
table = soup.find("table")
|
||||
if table is None:
|
||||
return ""
|
||||
|
||||
rows = table.find_all("tr")
|
||||
if not rows:
|
||||
return ""
|
||||
|
||||
grid = []
|
||||
|
||||
for r_idx, row in enumerate(rows):
|
||||
while len(grid) <= r_idx:
|
||||
grid.append([])
|
||||
|
||||
cells = row.find_all(["td", "th"])
|
||||
c_idx = 0
|
||||
|
||||
for cell in cells:
|
||||
while c_idx < len(grid[r_idx]) and grid[r_idx][c_idx] is not None:
|
||||
c_idx += 1
|
||||
|
||||
text = cell.get_text(strip=True)
|
||||
text = text.replace("\n", " ")
|
||||
|
||||
rowspan = int(cell.get("rowspan", 1))
|
||||
colspan = int(cell.get("colspan", 1))
|
||||
|
||||
for r in range(rowspan):
|
||||
target_r = r_idx + r
|
||||
while len(grid) <= target_r:
|
||||
grid.append([])
|
||||
|
||||
for c in range(colspan):
|
||||
target_c = c_idx + c
|
||||
while len(grid[target_r]) <= target_c:
|
||||
grid[target_r].append(None)
|
||||
|
||||
grid[target_r][target_c] = text
|
||||
|
||||
c_idx += colspan
|
||||
|
||||
if not grid:
|
||||
return ""
|
||||
|
||||
markdown_lines = []
|
||||
max_cols = max(len(r) for r in grid)
|
||||
for r in grid:
|
||||
while len(r) < max_cols:
|
||||
r.append("")
|
||||
|
||||
header = grid[0]
|
||||
header = [h if h is not None else "" for h in header]
|
||||
markdown_lines.append("| " + " | ".join(header) + " |")
|
||||
markdown_lines.append("|" + "|".join([" --- " for _ in range(max_cols)]) + "|")
|
||||
|
||||
for row in grid[1:]:
|
||||
row_clean = [cell if cell is not None else "" for cell in row]
|
||||
line = "| " + " | ".join(row_clean) + " |"
|
||||
markdown_lines.append(line)
|
||||
|
||||
return "\n".join(markdown_lines)
|
||||
|
||||
|
||||
def html_table_to_key_value(html: str) -> list[str]:
|
||||
"""
|
||||
将HTML表格转换为键值对格式的列表,为了应对过长的表格的切分问题。
|
||||
|
||||
处理逻辑:
|
||||
1. **网格重建**:由于 HTML 表格可能包含 `rowspan` 和 `colspan`(合并单元格),
|
||||
函数首先构建一个完整的二维网格(grid)。
|
||||
2. **单元格展开**:遍历 HTML 行和列,遇到合并单元格时,将其内容填充到网格中受影响的所有坐标点。
|
||||
这确保了原本被合并的区域在逻辑网格中每个点都有对应的值。
|
||||
3. **键值对转换**:
|
||||
- 将网格的第一行视为表头(Key)。
|
||||
- 从第二行开始,将每一行与表头对应,生成 "键:值" 形式的字符串。
|
||||
|
||||
例如:
|
||||
- 输入:HTML表格,包含姓名、年龄、性别三列。
|
||||
- 输出:['姓名:张三;年龄:25;性别:男', '姓名:李四;年龄:30;性别:女']
|
||||
"""
|
||||
soup = BeautifulSoup(html, "html.parser")
|
||||
table = soup.find("table")
|
||||
if table is None:
|
||||
return []
|
||||
|
||||
rows = table.find_all("tr")
|
||||
if not rows:
|
||||
return []
|
||||
|
||||
grid = []
|
||||
|
||||
for r_idx, row in enumerate(rows):
|
||||
while len(grid) <= r_idx:
|
||||
grid.append([])
|
||||
|
||||
cells = row.find_all(["td", "th"])
|
||||
c_idx = 0
|
||||
|
||||
for cell in cells:
|
||||
while c_idx < len(grid[r_idx]) and grid[r_idx][c_idx] is not None:
|
||||
c_idx += 1
|
||||
|
||||
text = cell.get_text(strip=True)
|
||||
rowspan = int(cell.get("rowspan", 1))
|
||||
colspan = int(cell.get("colspan", 1))
|
||||
|
||||
for r in range(rowspan):
|
||||
target_r = r_idx + r
|
||||
while len(grid) <= target_r:
|
||||
grid.append([])
|
||||
|
||||
for c in range(colspan):
|
||||
target_c = c_idx + c
|
||||
while len(grid[target_r]) <= target_c:
|
||||
grid[target_r].append(None)
|
||||
grid[target_r][target_c] = text
|
||||
c_idx += colspan
|
||||
|
||||
if not grid:
|
||||
return []
|
||||
|
||||
headers = grid[0]
|
||||
headers = [h if h is not None else "" for h in headers]
|
||||
|
||||
kv_lines = []
|
||||
for row_values in grid[1:]:
|
||||
min_len = min(len(headers), len(row_values))
|
||||
row_parts = []
|
||||
for i in range(min_len):
|
||||
key = headers[i]
|
||||
val = row_values[i] if row_values[i] is not None else ""
|
||||
if key:
|
||||
row_parts.append(f"{key}:{val}")
|
||||
if row_parts:
|
||||
kv_lines.append(";".join(row_parts) + ";")
|
||||
|
||||
return kv_lines
|
||||
@@ -0,0 +1 @@
|
||||
"""知识库评估核心能力。"""
|
||||
@@ -0,0 +1,352 @@
|
||||
import asyncio
|
||||
import json
|
||||
import random
|
||||
from collections.abc import AsyncIterator, Callable
|
||||
from typing import Any
|
||||
|
||||
import json_repair
|
||||
|
||||
from yuxi.models import select_model
|
||||
from yuxi.utils import logger
|
||||
|
||||
DEFAULT_BENCHMARK_GENERATION_CONCURRENCY = 10
|
||||
MAX_BENCHMARK_GENERATION_CONCURRENCY = 20
|
||||
DEFAULT_GRAPH_EXPAND_TOP_K = 1
|
||||
MAX_GRAPH_EXPAND_TOP_K = 3
|
||||
GRAPH_SEED_DECAY = 0.9
|
||||
GRAPH_PPR_DAMPING = 0.85
|
||||
GRAPH_PPR_MAX_NODES = 10000
|
||||
|
||||
|
||||
async def collect_kb_chunks(kb_instance: Any, kb_id: str) -> list[dict[str, Any]]:
|
||||
del kb_instance
|
||||
|
||||
from yuxi.repositories.knowledge_chunk_repository import KnowledgeChunkRepository
|
||||
|
||||
return [
|
||||
{
|
||||
"id": chunk.chunk_id,
|
||||
"content": chunk.content or "",
|
||||
"file_id": chunk.file_id,
|
||||
"chunk_index": chunk.chunk_index,
|
||||
"graph_indexed": bool(chunk.graph_indexed),
|
||||
"ent_ids": chunk.ent_ids or [],
|
||||
"tags": chunk.tags or [],
|
||||
"extraction_result": chunk.extraction_result,
|
||||
}
|
||||
for chunk in await KnowledgeChunkRepository().list_by_kb_id(kb_id)
|
||||
]
|
||||
|
||||
|
||||
def clamp_neighbors_count(neighbors_count: int) -> int:
|
||||
return min(max(neighbors_count, 0), 10)
|
||||
|
||||
|
||||
def normalize_generation_concurrency_count(value: Any) -> int:
|
||||
if value in (None, ""):
|
||||
return DEFAULT_BENCHMARK_GENERATION_CONCURRENCY
|
||||
return min(max(1, int(value)), MAX_BENCHMARK_GENERATION_CONCURRENCY)
|
||||
|
||||
|
||||
def normalize_graph_expand_top_k(value: Any) -> int:
|
||||
if value in (None, ""):
|
||||
return DEFAULT_GRAPH_EXPAND_TOP_K
|
||||
return min(max(1, int(value)), MAX_GRAPH_EXPAND_TOP_K)
|
||||
|
||||
|
||||
def _chunk_entity_ids(chunk: dict[str, Any]) -> list[str]:
|
||||
return [str(entity_id) for entity_id in chunk.get("ent_ids") or [] if entity_id]
|
||||
|
||||
|
||||
def _is_anchor_chunk(candidate: dict[str, Any], anchor_chunk: dict[str, Any]) -> bool:
|
||||
metadata = candidate.get("metadata") or {}
|
||||
candidate_id = metadata.get("chunk_id")
|
||||
if candidate_id is not None and str(candidate_id) == str(anchor_chunk.get("id")):
|
||||
return True
|
||||
|
||||
candidate_file_id = metadata.get("file_id")
|
||||
candidate_chunk_index = metadata.get("chunk_index")
|
||||
return candidate_file_id == anchor_chunk.get("file_id") and candidate_chunk_index == anchor_chunk.get("chunk_index")
|
||||
|
||||
|
||||
async def select_neighbor_chunks_by_kb_query(
|
||||
*, kb_instance: Any, kb_id: str, anchor_chunk: dict[str, Any], neighbors_count: int
|
||||
) -> list[dict[str, Any]]:
|
||||
if neighbors_count <= 0:
|
||||
return []
|
||||
|
||||
anchor_content = anchor_chunk.get("content", "")
|
||||
if not anchor_content:
|
||||
return []
|
||||
|
||||
candidates = await kb_instance.aquery(
|
||||
anchor_content,
|
||||
kb_id,
|
||||
search_mode="vector",
|
||||
final_top_k=neighbors_count + 3,
|
||||
use_reranker=False,
|
||||
similarity_threshold=0.0,
|
||||
)
|
||||
|
||||
chunks = []
|
||||
for candidate in candidates:
|
||||
if _is_anchor_chunk(candidate, anchor_chunk):
|
||||
continue
|
||||
|
||||
metadata = candidate.get("metadata") or {}
|
||||
chunk_id = metadata.get("chunk_id")
|
||||
content = candidate.get("content", "")
|
||||
if not chunk_id or not content:
|
||||
continue
|
||||
|
||||
chunks.append(
|
||||
{
|
||||
"id": str(chunk_id),
|
||||
"content": content,
|
||||
"file_id": metadata.get("file_id"),
|
||||
"chunk_index": metadata.get("chunk_index"),
|
||||
}
|
||||
)
|
||||
if len(chunks) >= neighbors_count:
|
||||
break
|
||||
|
||||
return chunks
|
||||
|
||||
|
||||
async def select_graph_enhanced_chunks(
|
||||
*,
|
||||
kb_id: str,
|
||||
anchor_chunk: dict[str, Any],
|
||||
chunks_by_id: dict[str, dict[str, Any]],
|
||||
context_count: int,
|
||||
graph_expand_top_k: int,
|
||||
) -> list[dict[str, Any]] | None:
|
||||
if context_count <= 1:
|
||||
return [anchor_chunk]
|
||||
|
||||
from yuxi.knowledge.graphs.milvus_graph_service import MilvusGraphService
|
||||
|
||||
anchor_entity_ids = _chunk_entity_ids(anchor_chunk)
|
||||
if not anchor_entity_ids:
|
||||
return None
|
||||
|
||||
graph_service = MilvusGraphService()
|
||||
selected = [anchor_chunk]
|
||||
selected_ids = {str(anchor_chunk.get("id"))}
|
||||
seed_weights = {entity_id: 1.0 for entity_id in anchor_entity_ids}
|
||||
round_index = 1
|
||||
|
||||
while len(selected) < context_count:
|
||||
for entity_id in anchor_entity_ids:
|
||||
seed_weights[entity_id] = 1.0
|
||||
|
||||
ranked_chunks = await graph_service.query_and_rank_chunks_by_ppr(
|
||||
kb_id,
|
||||
seed_weights,
|
||||
max_nodes=GRAPH_PPR_MAX_NODES,
|
||||
top_k=max(context_count * 5, 20),
|
||||
damping=GRAPH_PPR_DAMPING,
|
||||
)
|
||||
if not ranked_chunks:
|
||||
return None
|
||||
|
||||
new_chunks = []
|
||||
for chunk_id, _ in ranked_chunks:
|
||||
chunk_id = str(chunk_id)
|
||||
if chunk_id in selected_ids:
|
||||
continue
|
||||
chunk = chunks_by_id.get(chunk_id)
|
||||
if chunk is None:
|
||||
continue
|
||||
new_chunks.append(chunk)
|
||||
if len(new_chunks) >= min(graph_expand_top_k, context_count - len(selected)):
|
||||
break
|
||||
|
||||
if not new_chunks:
|
||||
return None
|
||||
|
||||
new_weight = GRAPH_SEED_DECAY**round_index
|
||||
for chunk in new_chunks:
|
||||
selected.append(chunk)
|
||||
selected_ids.add(str(chunk.get("id")))
|
||||
for entity_id in _chunk_entity_ids(chunk):
|
||||
seed_weights[entity_id] = max(seed_weights.get(entity_id, 0.0), new_weight)
|
||||
round_index += 1
|
||||
|
||||
return selected
|
||||
|
||||
|
||||
def build_benchmark_generation_prompt(ctx_items: list[tuple[str, str]]) -> str:
|
||||
context_text = "\n\n".join([f"片段ID={cid}\n{content}" for cid, content in ctx_items])
|
||||
return (
|
||||
"你将基于以下上下文生成一个可由上下文准确回答的问题与标准答案。"
|
||||
"仅返回一个JSON对象,不要包含其他文字。"
|
||||
"键为 query、gold_answer、gold_chunk_ids。gold_chunk_ids 必须是上述上下文片段的ID子集。\n\n"
|
||||
"上下文:\n" + context_text + "\n"
|
||||
)
|
||||
|
||||
|
||||
async def _generate_benchmark_item_once(
|
||||
*,
|
||||
kb_instance: Any,
|
||||
kb_id: str,
|
||||
all_chunks: list[dict[str, Any]],
|
||||
llm: Any,
|
||||
context_count: int,
|
||||
generation_mode: str,
|
||||
graph_expand_top_k: int,
|
||||
chunks_by_id: dict[str, dict[str, Any]],
|
||||
) -> dict[str, Any] | None:
|
||||
if generation_mode == "graph_enhanced":
|
||||
graph_anchor_chunks = [
|
||||
chunk for chunk in all_chunks if chunk.get("graph_indexed") is True and _chunk_entity_ids(chunk)
|
||||
]
|
||||
if not graph_anchor_chunks:
|
||||
raise ValueError("No graph indexed chunks with entities found in knowledge base")
|
||||
anchor_chunk = graph_anchor_chunks[random.randrange(len(graph_anchor_chunks))]
|
||||
ctx_chunks = await select_graph_enhanced_chunks(
|
||||
kb_id=kb_id,
|
||||
anchor_chunk=anchor_chunk,
|
||||
chunks_by_id=chunks_by_id,
|
||||
context_count=context_count,
|
||||
graph_expand_top_k=graph_expand_top_k,
|
||||
)
|
||||
if ctx_chunks is None:
|
||||
return None
|
||||
else:
|
||||
anchor_chunk = all_chunks[random.randrange(len(all_chunks))]
|
||||
neighbor_chunks = await select_neighbor_chunks_by_kb_query(
|
||||
kb_instance=kb_instance,
|
||||
kb_id=kb_id,
|
||||
anchor_chunk=anchor_chunk,
|
||||
neighbors_count=context_count - 1,
|
||||
)
|
||||
ctx_chunks = [anchor_chunk] + neighbor_chunks
|
||||
ctx_items = [(chunk["id"], chunk["content"]) for chunk in ctx_chunks]
|
||||
allowed_ids = {cid for cid, _ in ctx_items}
|
||||
|
||||
try:
|
||||
resp = await llm.call(build_benchmark_generation_prompt(ctx_items), False)
|
||||
obj = json_repair.loads(resp.content if resp else "")
|
||||
query = obj.get("query")
|
||||
answer = obj.get("gold_answer")
|
||||
gold_ids = obj.get("gold_chunk_ids")
|
||||
if not query or not answer or not isinstance(gold_ids, list):
|
||||
logger.warning(f"Generated JSON missing fields or invalid format: {obj}")
|
||||
return None
|
||||
|
||||
gold_ids = [str(item) for item in gold_ids if str(item) in allowed_ids]
|
||||
if not gold_ids:
|
||||
logger.warning("Generated gold_chunk_ids not found in allowed context")
|
||||
return None
|
||||
|
||||
return {"query": query, "gold_chunk_ids": gold_ids, "gold_answer": answer}
|
||||
except Exception as e:
|
||||
logger.warning(f"Benchmark generation failed for one item: {e}")
|
||||
return None
|
||||
|
||||
|
||||
async def iter_generated_benchmark_items(
|
||||
*,
|
||||
kb_instance: Any,
|
||||
kb_id: str,
|
||||
count: int,
|
||||
neighbors_count: int,
|
||||
llm_model_spec: str | None,
|
||||
concurrency_count: int = DEFAULT_BENCHMARK_GENERATION_CONCURRENCY,
|
||||
generation_mode: str = "vector",
|
||||
graph_expand_top_k: int = DEFAULT_GRAPH_EXPAND_TOP_K,
|
||||
progress_cb: Callable[[int, str], Any] | None = None,
|
||||
cancel_cb: Callable[[], Any] | None = None,
|
||||
) -> AsyncIterator[dict[str, Any]]:
|
||||
if progress_cb:
|
||||
await progress_cb(5, "加载chunks")
|
||||
|
||||
all_chunks = await collect_kb_chunks(kb_instance, kb_id)
|
||||
if not all_chunks:
|
||||
raise ValueError("No chunks found in knowledge base")
|
||||
chunks_by_id = {str(chunk["id"]): chunk for chunk in all_chunks if chunk.get("id") is not None}
|
||||
|
||||
if generation_mode not in {"vector", "graph_enhanced"}:
|
||||
raise ValueError("Unsupported benchmark generation mode")
|
||||
graph_expand_top_k = normalize_graph_expand_top_k(graph_expand_top_k)
|
||||
|
||||
if progress_cb:
|
||||
await progress_cb(15, "准备生成样本")
|
||||
|
||||
if not llm_model_spec:
|
||||
raise ValueError("llm_model_spec 不能为空")
|
||||
|
||||
llm = select_model(model_spec=llm_model_spec)
|
||||
context_count = max(clamp_neighbors_count(neighbors_count), 1)
|
||||
max_attempts = max(count * 5, 50)
|
||||
worker_count = normalize_generation_concurrency_count(concurrency_count)
|
||||
actual_worker_count = min(worker_count, max(count, 1), max_attempts)
|
||||
generated = 0
|
||||
results: list[tuple[int, dict[str, Any]]] = []
|
||||
state_lock = asyncio.Lock()
|
||||
queue: asyncio.Queue[int] = asyncio.Queue()
|
||||
|
||||
for attempt_no in range(max_attempts):
|
||||
queue.put_nowait(attempt_no)
|
||||
|
||||
async def worker() -> None:
|
||||
nonlocal generated
|
||||
while True:
|
||||
if cancel_cb:
|
||||
await cancel_cb()
|
||||
async with state_lock:
|
||||
if generated >= count:
|
||||
return
|
||||
try:
|
||||
attempt_no = queue.get_nowait()
|
||||
except asyncio.QueueEmpty:
|
||||
return
|
||||
try:
|
||||
item = await _generate_benchmark_item_once(
|
||||
kb_instance=kb_instance,
|
||||
kb_id=kb_id,
|
||||
all_chunks=all_chunks,
|
||||
llm=llm,
|
||||
context_count=context_count,
|
||||
generation_mode=generation_mode,
|
||||
graph_expand_top_k=graph_expand_top_k,
|
||||
chunks_by_id=chunks_by_id,
|
||||
)
|
||||
if item is None:
|
||||
continue
|
||||
progress = None
|
||||
message = None
|
||||
async with state_lock:
|
||||
if generated >= count:
|
||||
continue
|
||||
generated += 1
|
||||
results.append((attempt_no, item))
|
||||
if progress_cb:
|
||||
progress = int(99 * generated / max(count, 1))
|
||||
message = f"已生成 {generated}/{count}"
|
||||
if progress_cb:
|
||||
await progress_cb(progress, message)
|
||||
finally:
|
||||
queue.task_done()
|
||||
|
||||
workers = [asyncio.create_task(worker()) for _ in range(actual_worker_count)]
|
||||
try:
|
||||
await asyncio.gather(*workers)
|
||||
except asyncio.CancelledError:
|
||||
for task in workers:
|
||||
task.cancel()
|
||||
await asyncio.gather(*workers, return_exceptions=True)
|
||||
raise
|
||||
except Exception:
|
||||
for task in workers:
|
||||
task.cancel()
|
||||
await asyncio.gather(*workers, return_exceptions=True)
|
||||
raise
|
||||
|
||||
for _, item in sorted(results, key=lambda pair: pair[0]):
|
||||
yield item
|
||||
|
||||
|
||||
def dump_benchmark_item(item: dict[str, Any]) -> str:
|
||||
return json.dumps(item, ensure_ascii=False, separators=(",", ":")) + "\n"
|
||||
@@ -0,0 +1,137 @@
|
||||
from collections.abc import Callable
|
||||
from typing import Any
|
||||
|
||||
from yuxi.knowledge.eval.metrics import EvaluationMetricsCalculator
|
||||
from yuxi.utils import logger
|
||||
|
||||
|
||||
def normalize_query_result(query_result: Any) -> tuple[str, list[dict[str, Any]]]:
|
||||
if isinstance(query_result, dict):
|
||||
return query_result.get("answer", ""), query_result.get("retrieved_chunks", [])
|
||||
if isinstance(query_result, list):
|
||||
return "", query_result
|
||||
return "", []
|
||||
|
||||
|
||||
def build_answer_prompt(query: str, retrieved_chunks: list[dict[str, Any]], max_docs: int = 5) -> str:
|
||||
context_docs = []
|
||||
for idx, chunk in enumerate(retrieved_chunks[:max_docs]):
|
||||
content = chunk.get("content", "")
|
||||
if content:
|
||||
context_docs.append(f"文档 {idx + 1}:\n{content}")
|
||||
|
||||
context_text = "\\n\\n".join(context_docs)
|
||||
return (
|
||||
f"基于以下上下文信息,请回答用户的问题。\n\n"
|
||||
f"上下文信息:{context_text}\n\n"
|
||||
f"用户问题:{query}\n\n"
|
||||
"请根据上下文信息准确回答问题。\n\n"
|
||||
"如果上下文中缺少相关信息,请回答“信息不足,无法回答”。\n\n"
|
||||
)
|
||||
|
||||
|
||||
async def generate_answer_if_needed(
|
||||
*,
|
||||
query: str,
|
||||
generated_answer: str,
|
||||
retrieved_chunks: list[dict[str, Any]],
|
||||
retrieval_config: dict[str, Any],
|
||||
select_model_fn: Callable[..., Any],
|
||||
) -> str:
|
||||
if generated_answer:
|
||||
return generated_answer
|
||||
if not retrieved_chunks or not retrieval_config.get("answer_llm"):
|
||||
return ""
|
||||
|
||||
logger.debug(f"使用 LLM {retrieval_config.get('answer_llm')} 生成答案...")
|
||||
try:
|
||||
llm = select_model_fn(model_spec=retrieval_config["answer_llm"])
|
||||
response = await llm.call(build_answer_prompt(query, retrieved_chunks), stream=False)
|
||||
generated_answer = response.content if response else ""
|
||||
logger.debug(f"LLM 生成的答案长度: {len(generated_answer) if generated_answer else 0}")
|
||||
return generated_answer
|
||||
except Exception as e:
|
||||
logger.error(f"LLM 生成答案失败: {e}")
|
||||
return ""
|
||||
|
||||
|
||||
async def evaluate_question(
|
||||
*,
|
||||
kb_instance: Any,
|
||||
kb_id: str,
|
||||
question_data: dict[str, Any],
|
||||
retrieval_config: dict[str, Any],
|
||||
has_gold_chunks: bool,
|
||||
has_gold_answers: bool,
|
||||
judge_llm: Any | None,
|
||||
select_model_fn: Callable[..., Any],
|
||||
) -> dict[str, Any]:
|
||||
query = question_data["query"]
|
||||
query_result = await kb_instance.aquery(query, kb_id, **retrieval_config)
|
||||
generated_answer, retrieved_chunks = normalize_query_result(query_result)
|
||||
generated_answer = await generate_answer_if_needed(
|
||||
query=query,
|
||||
generated_answer=generated_answer,
|
||||
retrieved_chunks=retrieved_chunks,
|
||||
retrieval_config=retrieval_config,
|
||||
select_model_fn=select_model_fn,
|
||||
)
|
||||
|
||||
current_metrics = {}
|
||||
retrieval_scores = {}
|
||||
answer_scores = {}
|
||||
|
||||
if has_gold_chunks and question_data.get("gold_chunk_ids"):
|
||||
retrieval_scores = EvaluationMetricsCalculator.calculate_retrieval_metrics(
|
||||
retrieved_chunks, question_data["gold_chunk_ids"]
|
||||
)
|
||||
current_metrics.update(retrieval_scores)
|
||||
|
||||
if has_gold_answers and question_data.get("gold_answer"):
|
||||
if judge_llm:
|
||||
answer_scores = await EvaluationMetricsCalculator.calculate_answer_metrics(
|
||||
query=query,
|
||||
generated_answer=generated_answer,
|
||||
gold_answer=question_data["gold_answer"],
|
||||
judge_llm=judge_llm,
|
||||
)
|
||||
current_metrics.update(answer_scores)
|
||||
else:
|
||||
logger.warning("需要计算答案指标但未配置 Judge LLM")
|
||||
|
||||
return {
|
||||
"detail": {
|
||||
"query_text": query,
|
||||
"gold_chunk_ids": question_data.get("gold_chunk_ids"),
|
||||
"gold_answer": question_data.get("gold_answer"),
|
||||
"generated_answer": generated_answer,
|
||||
"retrieved_chunks": retrieved_chunks,
|
||||
"metrics": current_metrics,
|
||||
},
|
||||
"retrieval_scores": retrieval_scores,
|
||||
"answer_scores": answer_scores,
|
||||
}
|
||||
|
||||
|
||||
def aggregate_metrics(
|
||||
retrieval_metrics_list: list[dict[str, float]],
|
||||
answer_metrics_list: list[dict[str, Any]],
|
||||
*,
|
||||
include_overall_score: bool = False,
|
||||
) -> tuple[dict[str, Any], float | None]:
|
||||
overall_metrics = {}
|
||||
|
||||
if retrieval_metrics_list:
|
||||
keys = retrieval_metrics_list[0].keys()
|
||||
for key in keys:
|
||||
overall_metrics[key] = sum(m.get(key, 0) for m in retrieval_metrics_list) / len(retrieval_metrics_list)
|
||||
|
||||
if answer_metrics_list:
|
||||
scores = [m.get("score", 0) for m in answer_metrics_list]
|
||||
overall_metrics["answer_correctness"] = sum(scores) / len(scores) if scores else 0.0
|
||||
|
||||
overall_score = EvaluationMetricsCalculator.calculate_overall_score(retrieval_metrics_list, answer_metrics_list)
|
||||
if include_overall_score:
|
||||
overall_metrics["overall_score"] = overall_score
|
||||
|
||||
return overall_metrics, overall_score
|
||||
@@ -0,0 +1,141 @@
|
||||
"""
|
||||
RAG评估指标计算工具
|
||||
简化版:只保留Recall/F1(检索)和 LLM Judge(答案准确性)
|
||||
"""
|
||||
|
||||
import textwrap
|
||||
from typing import Any
|
||||
|
||||
import json_repair
|
||||
|
||||
from yuxi.utils import logger
|
||||
|
||||
|
||||
class RetrievalMetrics:
|
||||
"""检索评估指标计算"""
|
||||
|
||||
@staticmethod
|
||||
def precision_at_k(retrieved_ids: list[str], relevant_ids: list[str], k: int) -> float:
|
||||
"""计算Precision@K"""
|
||||
if not retrieved_ids[:k]:
|
||||
return 0.0
|
||||
retrieved_set = set(retrieved_ids[:k])
|
||||
relevant_set = set(relevant_ids)
|
||||
return len(retrieved_set & relevant_set) / k
|
||||
|
||||
@staticmethod
|
||||
def recall_at_k(retrieved_ids: list[str], relevant_ids: list[str], k: int) -> float:
|
||||
"""计算Recall@K"""
|
||||
if not relevant_ids:
|
||||
return 0.0
|
||||
retrieved_set = set(retrieved_ids[:k])
|
||||
relevant_set = set(relevant_ids)
|
||||
return len(retrieved_set & relevant_set) / len(relevant_set)
|
||||
|
||||
@staticmethod
|
||||
def f1_score_at_k(retrieved_ids: list[str], relevant_ids: list[str], k: int) -> float:
|
||||
"""计算F1@K"""
|
||||
precision = RetrievalMetrics.precision_at_k(retrieved_ids, relevant_ids, k)
|
||||
recall = RetrievalMetrics.recall_at_k(retrieved_ids, relevant_ids, k)
|
||||
if precision + recall == 0:
|
||||
return 0.0
|
||||
return 2 * precision * recall / (precision + recall)
|
||||
|
||||
|
||||
class AnswerMetrics:
|
||||
"""答案评估指标计算"""
|
||||
|
||||
@staticmethod
|
||||
async def judge_correctness(query: str, generated_answer: str, gold_answer: str, judge_llm: Any) -> dict[str, Any]:
|
||||
"""
|
||||
使用LLM判断生成的答案是否正确
|
||||
"""
|
||||
if not generated_answer:
|
||||
return {"score": 0.0, "reasoning": "未生成答案"}
|
||||
if not gold_answer:
|
||||
return {"score": 0.0, "reasoning": "无参考答案"}
|
||||
|
||||
prompt = textwrap.dedent(f"""你是一个公正的评判者,请评估AI生成的答案相对于标准答案的准确性。
|
||||
|
||||
问题:{query}
|
||||
|
||||
标准答案:
|
||||
{gold_answer}
|
||||
|
||||
AI生成的答案:
|
||||
{generated_answer}
|
||||
|
||||
请判断AI生成的答案是否在事实层面与标准答案一致。
|
||||
忽略措辞、标点符号或格式上的细微差异。
|
||||
只关注核心事实是否准确包含。
|
||||
|
||||
请返回以下JSON格式的结果(不要包含其他文本、Markdown 或注释):
|
||||
{{
|
||||
"score": 1.0,
|
||||
"reasoning": "简要说明判定理由"
|
||||
}}
|
||||
score 只能是 1.0 或 0.0。
|
||||
""")
|
||||
try:
|
||||
response = await judge_llm.call(prompt, stream=False)
|
||||
content = response.content.strip()
|
||||
|
||||
# 尝试清理可能的 markdown 代码块
|
||||
if content.startswith("```json"):
|
||||
content = content[7:]
|
||||
if content.endswith("```"):
|
||||
content = content[:-3]
|
||||
content = content.strip()
|
||||
|
||||
result = json_repair.loads(content)
|
||||
return {"score": float(result.get("score", 0.0)), "reasoning": result.get("reasoning", "")}
|
||||
except Exception as e:
|
||||
logger.error(f"LLM 评判失败: {e}")
|
||||
return {"score": 0.0, "reasoning": f"评判出错: {str(e)}"}
|
||||
|
||||
|
||||
class EvaluationMetricsCalculator:
|
||||
"""综合评估指标计算器"""
|
||||
|
||||
@staticmethod
|
||||
def calculate_retrieval_metrics(
|
||||
retrieved_chunks: list[dict[str, Any]], gold_chunk_ids: list[str], k_values: list[int] = [1, 3, 5, 10]
|
||||
) -> dict[str, float]:
|
||||
"""计算检索指标 (Recall, F1)"""
|
||||
if not retrieved_chunks or not gold_chunk_ids:
|
||||
return {}
|
||||
|
||||
# 提取 ID
|
||||
retrieved_ids = []
|
||||
for chunk in retrieved_chunks:
|
||||
chunk_id = chunk.get("chunk_id") or chunk.get("metadata", {}).get("chunk_id")
|
||||
retrieved_ids.append(str(chunk_id) if chunk_id else "")
|
||||
|
||||
metrics = {}
|
||||
for k in k_values:
|
||||
metrics[f"recall@{k}"] = RetrievalMetrics.recall_at_k(retrieved_ids, gold_chunk_ids, k)
|
||||
metrics[f"f1@{k}"] = RetrievalMetrics.f1_score_at_k(retrieved_ids, gold_chunk_ids, k)
|
||||
|
||||
return metrics
|
||||
|
||||
@staticmethod
|
||||
async def calculate_answer_metrics(
|
||||
query: str, generated_answer: str, gold_answer: str, judge_llm: Any = None
|
||||
) -> dict[str, Any]:
|
||||
"""计算答案指标 (LLM Judge)"""
|
||||
if not judge_llm:
|
||||
return {}
|
||||
|
||||
return await AnswerMetrics.judge_correctness(query, generated_answer, gold_answer, judge_llm)
|
||||
|
||||
@staticmethod
|
||||
def calculate_overall_score(
|
||||
retrieval_metrics_list: list[dict[str, float]], answer_metrics_list: list[dict[str, Any]]
|
||||
) -> float | None:
|
||||
"""综合得分:有答案准确率则用准确率,否则用 recall@10。"""
|
||||
if answer_metrics_list:
|
||||
scores = [m.get("score", 0.0) for m in answer_metrics_list]
|
||||
return sum(scores) / len(scores) if scores else None
|
||||
|
||||
recalls = [m["recall@10"] for m in retrieval_metrics_list if m and "recall@10" in m]
|
||||
return sum(recalls) / len(recalls) if recalls else None
|
||||
@@ -0,0 +1,722 @@
|
||||
import json
|
||||
import re
|
||||
import uuid
|
||||
from typing import Any
|
||||
|
||||
from yuxi.knowledge import knowledge_base
|
||||
from yuxi.knowledge.eval.benchmark_generation import (
|
||||
dump_benchmark_item,
|
||||
iter_generated_benchmark_items,
|
||||
normalize_generation_concurrency_count,
|
||||
)
|
||||
from yuxi.knowledge.eval.evaluator import aggregate_metrics, evaluate_question
|
||||
from yuxi.models import select_model
|
||||
from yuxi.repositories.evaluation_repository import EvaluationRepository
|
||||
from yuxi.repositories.knowledge_base_repository import KnowledgeBaseRepository
|
||||
from yuxi.repositories.knowledge_chunk_repository import KnowledgeChunkRepository
|
||||
from yuxi.repositories.task_repository import TaskRepository
|
||||
from yuxi.services.task_service import TaskContext, tasker
|
||||
from yuxi.utils import logger
|
||||
from yuxi.utils.datetime_utils import format_utc_datetime, utc_now_naive
|
||||
|
||||
|
||||
def build_evaluation_run_name(started_at=None, hash_value: str | None = None) -> str:
|
||||
date_part = (started_at or utc_now_naive()).strftime("%Y%m%d")
|
||||
hash_part = re.sub(r"[^a-fA-F0-9]", "", hash_value or uuid.uuid4().hex).lower()[:6]
|
||||
if len(hash_part) < 6:
|
||||
hash_part = (hash_part + uuid.uuid4().hex)[:6]
|
||||
return f"eval-{date_part}-{hash_part}"
|
||||
|
||||
|
||||
class EvaluationService:
|
||||
"""RAG评估服务"""
|
||||
|
||||
def __init__(self):
|
||||
self.eval_repo = EvaluationRepository()
|
||||
self.kb_repo = KnowledgeBaseRepository()
|
||||
self.chunk_repo = KnowledgeChunkRepository()
|
||||
self.task_repo = TaskRepository()
|
||||
|
||||
def _dataset_to_dict(self, row) -> dict[str, Any]:
|
||||
return {
|
||||
"id": row.dataset_id,
|
||||
"dataset_id": row.dataset_id,
|
||||
"name": row.name,
|
||||
"description": row.description,
|
||||
"kb_id": row.kb_id,
|
||||
"item_count": row.item_count,
|
||||
"has_gold_chunks": row.has_gold_chunks,
|
||||
"has_gold_answers": row.has_gold_answers,
|
||||
"build_metadata": row.build_metadata or {},
|
||||
"created_by": row.created_by,
|
||||
"created_at": format_utc_datetime(row.created_at),
|
||||
"updated_at": format_utc_datetime(row.updated_at),
|
||||
}
|
||||
|
||||
def _dataset_item_to_dict(self, item) -> dict[str, Any]:
|
||||
return {
|
||||
"item_id": item.item_id,
|
||||
"item_index": item.item_index,
|
||||
"query": item.query_text,
|
||||
"gold_chunk_ids": item.gold_chunk_ids or [],
|
||||
"gold_answer": item.gold_answer,
|
||||
}
|
||||
|
||||
def _run_item_to_dict(self, item) -> dict[str, Any]:
|
||||
return {
|
||||
"query": item.query_text,
|
||||
"gold_chunk_ids": item.gold_chunk_ids,
|
||||
"gold_answer": item.gold_answer,
|
||||
"generated_answer": item.generated_answer,
|
||||
"retrieved_chunks": item.retrieved_chunks,
|
||||
"metrics": item.metrics or {},
|
||||
}
|
||||
|
||||
def _is_error_run_item(self, item) -> bool:
|
||||
metrics = item.metrics or {}
|
||||
return metrics.get("score", 1.0) <= 0.5 or any(
|
||||
metrics.get(key, 1.0) < 0.3 for key in metrics if key.startswith("recall@")
|
||||
)
|
||||
|
||||
def _normalize_run_name(self, name: str | None, run_id: str) -> str:
|
||||
run_name = (name or "").strip()
|
||||
if run_name:
|
||||
return run_name
|
||||
return build_evaluation_run_name(hash_value=run_id.removeprefix("run_"))
|
||||
|
||||
def _run_name_from_row(self, row) -> str:
|
||||
name = (getattr(row, "name", None) or "").strip()
|
||||
if name:
|
||||
return name
|
||||
return build_evaluation_run_name(row.started_at, hash_value=row.run_id.removeprefix("run_"))
|
||||
|
||||
async def _sync_dataset_build_metadata(self, row) -> None:
|
||||
metadata = dict(row.build_metadata or {})
|
||||
if metadata.get("source") != "generated" or metadata.get("status") not in {"pending", "running"}:
|
||||
return
|
||||
|
||||
task_id = metadata.get("task_id")
|
||||
task = await self.task_repo.get_by_id(task_id) if task_id else None
|
||||
if task is None:
|
||||
metadata.pop("progress", None)
|
||||
metadata.update(status="failed", message="生成任务不存在")
|
||||
elif task.status == "success":
|
||||
metadata.update(status="completed", progress=100, message=task.message or "完成")
|
||||
elif task.status in {"failed", "cancelled"}:
|
||||
metadata.pop("progress", None)
|
||||
metadata.update(status="failed", message=task.error or task.message or "生成任务失败")
|
||||
else:
|
||||
metadata.update(status=task.status, progress=task.progress, message=task.message)
|
||||
|
||||
if metadata != (row.build_metadata or {}):
|
||||
await self.eval_repo.update_dataset(row.dataset_id, {"build_metadata": metadata})
|
||||
row.build_metadata = metadata
|
||||
|
||||
def _build_dataset_items(
|
||||
self, dataset_id: str, kb_id: str, questions: list[dict[str, Any]]
|
||||
) -> list[dict[str, Any]]:
|
||||
return [
|
||||
{
|
||||
"item_id": f"dataset_item_{uuid.uuid4().hex[:12]}",
|
||||
"dataset_id": dataset_id,
|
||||
"kb_id": kb_id,
|
||||
"item_index": index,
|
||||
"query_text": item["query"],
|
||||
"gold_chunk_ids": item.get("gold_chunk_ids") or [],
|
||||
"gold_answer": item.get("gold_answer"),
|
||||
}
|
||||
for index, item in enumerate(questions)
|
||||
]
|
||||
|
||||
def _build_jsonl_content(self, items: list[Any]) -> str:
|
||||
lines = []
|
||||
for item in items:
|
||||
payload = {"query": item.query_text}
|
||||
if item.gold_chunk_ids:
|
||||
payload["gold_chunk_ids"] = item.gold_chunk_ids
|
||||
if item.gold_answer:
|
||||
payload["gold_answer"] = item.gold_answer
|
||||
lines.append(dump_benchmark_item(payload).rstrip("\n"))
|
||||
return "\n".join(lines) + ("\n" if lines else "")
|
||||
|
||||
def _safe_jsonl_filename(self, name: str | None, fallback: str) -> str:
|
||||
filename = (name or "").strip() or fallback
|
||||
filename = re.sub(r"[\\/:*?\"<>|]+", "_", filename).strip()
|
||||
if not filename or filename in {".", ".."}:
|
||||
filename = fallback
|
||||
return filename if filename.endswith(".jsonl") else f"{filename}.jsonl"
|
||||
|
||||
def _parse_jsonl_questions(self, file_content: bytes) -> tuple[list[dict[str, Any]], bool, bool]:
|
||||
questions = []
|
||||
has_gold_chunks = False
|
||||
has_gold_answers = False
|
||||
content = file_content.decode("utf-8")
|
||||
|
||||
for line_num, line in enumerate(content.strip().split("\n"), 1):
|
||||
if not line.strip():
|
||||
continue
|
||||
try:
|
||||
item = json.loads(line)
|
||||
except json.JSONDecodeError as e:
|
||||
raise ValueError(f"第{line_num}行JSON格式错误: {str(e)}")
|
||||
if "query" not in item:
|
||||
raise ValueError(f"第{line_num}行缺少必需的'query'字段")
|
||||
if item.get("gold_chunk_ids"):
|
||||
has_gold_chunks = True
|
||||
if item.get("gold_answer"):
|
||||
has_gold_answers = True
|
||||
questions.append(item)
|
||||
|
||||
if not questions:
|
||||
raise ValueError("文件中没有有效的问题数据")
|
||||
return questions, has_gold_chunks, has_gold_answers
|
||||
|
||||
async def upload_dataset(
|
||||
self, kb_id: str, file_content: bytes, filename: str, name: str, description: str, created_by: str
|
||||
) -> dict[str, Any]:
|
||||
try:
|
||||
questions, has_gold_chunks, has_gold_answers = self._parse_jsonl_questions(file_content)
|
||||
dataset_id = f"dataset_{uuid.uuid4().hex[:8]}"
|
||||
dataset_name = name.strip() or filename or dataset_id
|
||||
|
||||
row = await self.eval_repo.create_dataset_with_items(
|
||||
{
|
||||
"dataset_id": dataset_id,
|
||||
"kb_id": kb_id,
|
||||
"name": dataset_name,
|
||||
"description": description,
|
||||
"item_count": len(questions),
|
||||
"has_gold_chunks": has_gold_chunks,
|
||||
"has_gold_answers": has_gold_answers,
|
||||
"build_metadata": {
|
||||
"source": "upload",
|
||||
"status": "completed",
|
||||
"progress": 100,
|
||||
"filename": filename,
|
||||
},
|
||||
"created_by": created_by,
|
||||
},
|
||||
self._build_dataset_items(dataset_id, kb_id, questions),
|
||||
)
|
||||
return self._dataset_to_dict(row)
|
||||
except Exception as e:
|
||||
logger.error(f"上传评估数据集失败: {e}")
|
||||
raise
|
||||
|
||||
async def list_datasets(self, kb_id: str) -> list[dict[str, Any]]:
|
||||
try:
|
||||
rows = await self.eval_repo.list_datasets(kb_id)
|
||||
for row in rows:
|
||||
await self._sync_dataset_build_metadata(row)
|
||||
return [self._dataset_to_dict(row) for row in rows]
|
||||
except Exception as e:
|
||||
logger.error(f"获取评估数据集列表失败: {e}")
|
||||
raise
|
||||
|
||||
async def get_dataset_detail(
|
||||
self, kb_id: str, dataset_id: str, page: int = 1, page_size: int = 10
|
||||
) -> dict[str, Any]:
|
||||
try:
|
||||
row = await self.eval_repo.get_dataset(dataset_id)
|
||||
if row is None or row.kb_id != kb_id:
|
||||
raise ValueError("Dataset not found")
|
||||
if (row.build_metadata or {}).get("status", "completed") != "completed":
|
||||
raise ValueError("Dataset is not ready")
|
||||
|
||||
total_items = await self.eval_repo.count_dataset_items(dataset_id)
|
||||
items = await self.eval_repo.list_dataset_items(dataset_id, (page - 1) * page_size, page_size)
|
||||
total_pages = (total_items + page_size - 1) // page_size
|
||||
data = self._dataset_to_dict(row)
|
||||
data.update(
|
||||
{
|
||||
"items": [self._dataset_item_to_dict(item) for item in items],
|
||||
"pagination": {
|
||||
"current_page": page,
|
||||
"page_size": page_size,
|
||||
"total_items": total_items,
|
||||
"total_pages": total_pages,
|
||||
"has_next": page < total_pages,
|
||||
"has_prev": page > 1,
|
||||
},
|
||||
}
|
||||
)
|
||||
return data
|
||||
except Exception as e:
|
||||
logger.error(f"获取评估数据集详情失败: {e}")
|
||||
raise
|
||||
|
||||
async def export_dataset_jsonl(self, dataset_id: str) -> dict[str, str]:
|
||||
row = await self.eval_repo.get_dataset(dataset_id)
|
||||
if row is None:
|
||||
raise ValueError("Dataset not found")
|
||||
if (row.build_metadata or {}).get("status", "completed") != "completed":
|
||||
raise ValueError("Dataset is not ready")
|
||||
items = await self.eval_repo.list_all_dataset_items(dataset_id)
|
||||
return {
|
||||
"filename": self._safe_jsonl_filename(row.name, row.dataset_id),
|
||||
"content": self._build_jsonl_content(items),
|
||||
}
|
||||
|
||||
async def delete_dataset(self, dataset_id: str) -> None:
|
||||
try:
|
||||
row = await self.eval_repo.get_dataset(dataset_id)
|
||||
if row is None:
|
||||
raise ValueError("Dataset not found")
|
||||
await self.eval_repo.delete_dataset(dataset_id)
|
||||
logger.info(f"成功删除评估数据集: {dataset_id}")
|
||||
except Exception as e:
|
||||
logger.error(f"删除评估数据集失败: {e}")
|
||||
raise
|
||||
|
||||
async def generate_dataset(
|
||||
self,
|
||||
kb_id: str,
|
||||
name: str,
|
||||
description: str,
|
||||
count: int,
|
||||
neighbors_count: int,
|
||||
concurrency_count: int,
|
||||
llm_model_spec: str,
|
||||
generation_mode: str = "vector",
|
||||
graph_expand_top_k: int = 1,
|
||||
created_by: str = "system",
|
||||
) -> dict[str, Any]:
|
||||
dataset_id = f"dataset_{uuid.uuid4().hex[:8]}"
|
||||
count = int(count)
|
||||
neighbors_count = int(neighbors_count)
|
||||
concurrency_count = normalize_generation_concurrency_count(concurrency_count)
|
||||
graph_expand_top_k = min(max(1, int(graph_expand_top_k)), 3)
|
||||
if generation_mode not in {"vector", "graph_enhanced"}:
|
||||
raise ValueError("不支持的评估基准生成方式")
|
||||
if generation_mode == "graph_enhanced":
|
||||
indexed_count = await self.chunk_repo.count_graph_indexed_by_kb_id(kb_id)
|
||||
if indexed_count <= 0:
|
||||
raise ValueError("当前知识库尚未完成图索引,无法使用图增强构建")
|
||||
build_metadata = {
|
||||
"source": "generated",
|
||||
"status": "pending",
|
||||
"progress": 0,
|
||||
"params": {
|
||||
"count": count,
|
||||
"neighbors_count": neighbors_count,
|
||||
"concurrency_count": concurrency_count,
|
||||
"llm_model_spec": llm_model_spec,
|
||||
"generation_mode": generation_mode,
|
||||
"graph_expand_top_k": graph_expand_top_k,
|
||||
},
|
||||
}
|
||||
await self.eval_repo.create_dataset(
|
||||
{
|
||||
"dataset_id": dataset_id,
|
||||
"kb_id": kb_id,
|
||||
"name": name,
|
||||
"description": description,
|
||||
"item_count": 0,
|
||||
"has_gold_chunks": True,
|
||||
"has_gold_answers": True,
|
||||
"build_metadata": build_metadata,
|
||||
"created_by": created_by,
|
||||
}
|
||||
)
|
||||
task = await tasker.enqueue(
|
||||
name="生成评估数据集",
|
||||
task_type="dataset_generation",
|
||||
payload={
|
||||
"dataset_id": dataset_id,
|
||||
"kb_id": kb_id,
|
||||
"created_by": created_by,
|
||||
"name": name,
|
||||
"description": description,
|
||||
"count": count,
|
||||
"neighbors_count": neighbors_count,
|
||||
"concurrency_count": concurrency_count,
|
||||
"llm_model_spec": llm_model_spec,
|
||||
"generation_mode": generation_mode,
|
||||
"graph_expand_top_k": graph_expand_top_k,
|
||||
},
|
||||
coroutine=self._generate_dataset_task,
|
||||
)
|
||||
build_metadata["task_id"] = task.id
|
||||
await self.eval_repo.update_dataset(dataset_id, {"build_metadata": build_metadata})
|
||||
return {"dataset_id": dataset_id, "task_id": task.id, "message": "评估数据集生成任务已提交"}
|
||||
|
||||
async def _update_dataset_build_metadata(
|
||||
self, dataset_id: str, metadata: dict[str, Any], **updates
|
||||
) -> dict[str, Any]:
|
||||
metadata.update(updates)
|
||||
await self.eval_repo.update_dataset(dataset_id, {"build_metadata": metadata})
|
||||
return metadata
|
||||
|
||||
async def _generate_dataset_task(self, context: TaskContext):
|
||||
await context.set_progress(0, "初始化")
|
||||
payload = context.payload
|
||||
|
||||
dataset_id = payload.get("dataset_id")
|
||||
kb_id = payload.get("kb_id")
|
||||
count = int(payload.get("count", 10))
|
||||
neighbors_count = int(payload.get("neighbors_count", 1))
|
||||
concurrency_count = normalize_generation_concurrency_count(payload.get("concurrency_count"))
|
||||
llm_model_spec = payload.get("llm_model_spec")
|
||||
generation_mode = payload.get("generation_mode") or "vector"
|
||||
graph_expand_top_k = min(max(1, int(payload.get("graph_expand_top_k", 1))), 3)
|
||||
build_metadata = {
|
||||
"source": "generated",
|
||||
"status": "running",
|
||||
"progress": 0,
|
||||
"task_id": context.task_id,
|
||||
"params": {
|
||||
"count": count,
|
||||
"neighbors_count": neighbors_count,
|
||||
"concurrency_count": concurrency_count,
|
||||
"llm_model_spec": llm_model_spec,
|
||||
"generation_mode": generation_mode,
|
||||
"graph_expand_top_k": graph_expand_top_k,
|
||||
},
|
||||
}
|
||||
await self._update_dataset_build_metadata(dataset_id, build_metadata)
|
||||
|
||||
async def report_progress(progress: float, message: str | None = None) -> None:
|
||||
await context.set_progress(progress, message)
|
||||
await self._update_dataset_build_metadata(
|
||||
dataset_id,
|
||||
build_metadata,
|
||||
progress=max(0, min(round(progress), 100)),
|
||||
message=message or build_metadata.get("message", ""),
|
||||
)
|
||||
|
||||
try:
|
||||
kb_instance = await knowledge_base.aget_kb(kb_id)
|
||||
if not kb_instance:
|
||||
await report_progress(100, "知识库不存在")
|
||||
raise ValueError("Knowledge Base not found")
|
||||
if kb_instance.kb_type != "milvus":
|
||||
await report_progress(100, "仅支持 commonrag/Milvus 类型知识库生成评估数据集")
|
||||
raise ValueError("Unsupported KB type for dataset generation")
|
||||
|
||||
questions = []
|
||||
try:
|
||||
async for item in iter_generated_benchmark_items(
|
||||
kb_instance=kb_instance,
|
||||
kb_id=kb_id,
|
||||
count=count,
|
||||
neighbors_count=neighbors_count,
|
||||
llm_model_spec=llm_model_spec,
|
||||
concurrency_count=concurrency_count,
|
||||
generation_mode=generation_mode,
|
||||
graph_expand_top_k=graph_expand_top_k,
|
||||
progress_cb=report_progress,
|
||||
cancel_cb=context.raise_if_cancelled,
|
||||
):
|
||||
questions.append(item)
|
||||
except ValueError as e:
|
||||
if str(e) == "No chunks found in knowledge base":
|
||||
await report_progress(100, "知识库为空或未解析到chunks")
|
||||
raise
|
||||
|
||||
if not questions:
|
||||
raise ValueError("未生成有效评估题目")
|
||||
|
||||
await self.eval_repo.add_dataset_items(self._build_dataset_items(dataset_id, kb_id, questions))
|
||||
await self.eval_repo.update_dataset(dataset_id, {"item_count": len(questions)})
|
||||
await self._update_dataset_build_metadata(
|
||||
dataset_id,
|
||||
build_metadata,
|
||||
status="completed",
|
||||
progress=100,
|
||||
message="完成",
|
||||
)
|
||||
await context.set_progress(100, "完成")
|
||||
except Exception as e:
|
||||
await self._update_dataset_build_metadata(
|
||||
dataset_id,
|
||||
build_metadata,
|
||||
status="failed",
|
||||
progress=100,
|
||||
error_message=str(e),
|
||||
message=str(e),
|
||||
)
|
||||
raise
|
||||
|
||||
async def run_evaluation(
|
||||
self,
|
||||
kb_id: str,
|
||||
dataset_id: str,
|
||||
name: str | None = None,
|
||||
model_config: dict[str, Any] = None,
|
||||
created_by: str = "system",
|
||||
) -> str:
|
||||
try:
|
||||
run_id = f"run_{uuid.uuid4().hex[:8]}"
|
||||
run_name = self._normalize_run_name(name, run_id)
|
||||
dataset_row = await self.eval_repo.get_dataset(dataset_id)
|
||||
if dataset_row is None or dataset_row.kb_id != kb_id:
|
||||
raise ValueError("Dataset not found")
|
||||
if (dataset_row.build_metadata or {}).get("status", "completed") != "completed":
|
||||
raise ValueError("Dataset is not ready")
|
||||
|
||||
retrieval_config = {}
|
||||
try:
|
||||
kb_row = await self.kb_repo.get_by_kb_id(kb_id)
|
||||
query_params = (kb_row.query_params if kb_row else None) or {}
|
||||
retrieval_config = query_params.get("options", {}) if isinstance(query_params, dict) else {}
|
||||
if not retrieval_config:
|
||||
kb_instance = await knowledge_base.aget_kb(kb_id)
|
||||
if kb_instance:
|
||||
retrieval_config = kb_instance._get_default_query_params(kb_id).get("options", {})
|
||||
logger.info(f"从知识库 {kb_id} 加载检索配置: {list(retrieval_config.keys())}")
|
||||
except Exception as e:
|
||||
logger.error(f"获取知识库检索配置失败: {e}")
|
||||
|
||||
if model_config:
|
||||
retrieval_config.update(model_config)
|
||||
|
||||
await self.eval_repo.create_run(
|
||||
{
|
||||
"run_id": run_id,
|
||||
"name": run_name,
|
||||
"kb_id": kb_id,
|
||||
"dataset_id": dataset_id,
|
||||
"status": "running",
|
||||
"retrieval_config": retrieval_config,
|
||||
"metrics": {},
|
||||
"overall_score": None,
|
||||
"total_items": dataset_row.item_count or 0,
|
||||
"completed_items": 0,
|
||||
"started_at": utc_now_naive(),
|
||||
"completed_at": None,
|
||||
"created_by": created_by,
|
||||
}
|
||||
)
|
||||
|
||||
await tasker.enqueue(
|
||||
name=f"RAG评估({run_name})",
|
||||
task_type="rag_evaluation",
|
||||
payload={
|
||||
"run_id": run_id,
|
||||
"name": run_name,
|
||||
"kb_id": kb_id,
|
||||
"dataset_id": dataset_id,
|
||||
"retrieval_config": retrieval_config,
|
||||
"created_by": created_by,
|
||||
},
|
||||
coroutine=self._run_evaluation_task,
|
||||
)
|
||||
return run_id
|
||||
except Exception as e:
|
||||
logger.error(f"启动评估失败: {e}")
|
||||
raise
|
||||
|
||||
async def _run_evaluation_task(self, context: TaskContext):
|
||||
try:
|
||||
payload = context.payload
|
||||
|
||||
run_id = payload["run_id"]
|
||||
kb_id = payload["kb_id"]
|
||||
dataset_id = payload["dataset_id"]
|
||||
retrieval_config = payload["retrieval_config"]
|
||||
|
||||
await context.set_progress(5, "加载评估数据集")
|
||||
dataset_row = await self.eval_repo.get_dataset(dataset_id)
|
||||
if dataset_row is None or dataset_row.kb_id != kb_id:
|
||||
raise ValueError("Dataset not found")
|
||||
dataset_items = await self.eval_repo.list_all_dataset_items(dataset_id)
|
||||
if not dataset_items:
|
||||
raise ValueError("Dataset has no items")
|
||||
|
||||
kb_instance = await knowledge_base.aget_kb(kb_id)
|
||||
if not kb_instance:
|
||||
raise ValueError(f"Knowledge Base {kb_id} not found")
|
||||
|
||||
judge_llm = None
|
||||
if dataset_row.has_gold_answers:
|
||||
judge_model_spec = retrieval_config.get("judge_llm") or retrieval_config.get("answer_llm")
|
||||
if judge_model_spec:
|
||||
try:
|
||||
logger.debug(f"Initializing Judge LLM: {judge_model_spec}")
|
||||
judge_llm = select_model(model_spec=judge_model_spec)
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to load judge LLM: {e}")
|
||||
|
||||
all_retrieval_metrics = []
|
||||
all_answer_metrics = []
|
||||
total_items = len(dataset_items)
|
||||
|
||||
async def update_run_db(status=None, completed=None, metrics=None, final_score=None):
|
||||
data = {}
|
||||
if status is not None:
|
||||
data["status"] = status
|
||||
if status in ["completed", "failed"]:
|
||||
data["completed_at"] = utc_now_naive()
|
||||
if completed is not None:
|
||||
data["completed_items"] = completed
|
||||
if metrics is not None:
|
||||
data["metrics"] = metrics
|
||||
if final_score is not None:
|
||||
data["overall_score"] = final_score
|
||||
if data:
|
||||
await self.eval_repo.update_run(run_id, data)
|
||||
|
||||
for index, item in enumerate(dataset_items):
|
||||
await context.raise_if_cancelled()
|
||||
progress = 10 + (index / total_items) * 80
|
||||
await context.set_progress(progress, f"评估 {index + 1}/{total_items}")
|
||||
|
||||
question_data = {
|
||||
"query": item.query_text,
|
||||
"gold_chunk_ids": item.gold_chunk_ids or [],
|
||||
"gold_answer": item.gold_answer,
|
||||
}
|
||||
question_result = await evaluate_question(
|
||||
kb_instance=kb_instance,
|
||||
kb_id=kb_id,
|
||||
question_data=question_data,
|
||||
retrieval_config=retrieval_config,
|
||||
has_gold_chunks=dataset_row.has_gold_chunks,
|
||||
has_gold_answers=dataset_row.has_gold_answers,
|
||||
judge_llm=judge_llm,
|
||||
select_model_fn=select_model,
|
||||
)
|
||||
|
||||
if dataset_row.has_gold_chunks and question_data.get("gold_chunk_ids"):
|
||||
all_retrieval_metrics.append(question_result["retrieval_scores"])
|
||||
if dataset_row.has_gold_answers and question_data.get("gold_answer") and judge_llm:
|
||||
all_answer_metrics.append(question_result["answer_scores"])
|
||||
|
||||
await self.eval_repo.upsert_run_item(
|
||||
run_id=run_id,
|
||||
item_index=index,
|
||||
data={"dataset_item_id": item.item_id, **question_result["detail"]},
|
||||
)
|
||||
|
||||
if (index + 1) % 5 == 0 or (index + 1) == total_items:
|
||||
current_metrics, _ = aggregate_metrics(all_retrieval_metrics, all_answer_metrics)
|
||||
await context.set_result(
|
||||
{"current_metrics": current_metrics, "completed_items": index + 1, "total_items": total_items}
|
||||
)
|
||||
await update_run_db(completed=index + 1)
|
||||
|
||||
await context.set_progress(95, "计算最终指标")
|
||||
overall_metrics, overall_score = aggregate_metrics(
|
||||
all_retrieval_metrics, all_answer_metrics, include_overall_score=True
|
||||
)
|
||||
await update_run_db(
|
||||
status="completed",
|
||||
completed=total_items,
|
||||
metrics=overall_metrics,
|
||||
final_score=overall_score,
|
||||
)
|
||||
await context.set_progress(100, "完成")
|
||||
except Exception as e:
|
||||
logger.error(f"Task failed: {e}")
|
||||
try:
|
||||
if "payload" in locals():
|
||||
await self.eval_repo.update_run(
|
||||
payload["run_id"],
|
||||
{"status": "failed", "metrics": {"error": str(e)}, "completed_at": utc_now_naive()},
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.error(f"Error updating run record: {exc}")
|
||||
await context.set_message(f"Error: {str(e)}")
|
||||
raise
|
||||
|
||||
async def list_runs(self, kb_id: str) -> list[dict[str, Any]]:
|
||||
try:
|
||||
rows = await self.eval_repo.list_runs(kb_id)
|
||||
running_run_ids = {row.run_id for row in rows if row.status == "running"}
|
||||
task_by_run_id = {}
|
||||
if running_run_ids:
|
||||
tasks = await self.task_repo.list_all()
|
||||
task_by_run_id = {
|
||||
(task.payload or {}).get("run_id"): task
|
||||
for task in tasks
|
||||
if task.type == "rag_evaluation"
|
||||
and task.status in {"pending", "running"}
|
||||
and (task.payload or {}).get("run_id") in running_run_ids
|
||||
}
|
||||
|
||||
runs = []
|
||||
for row in rows:
|
||||
run = {
|
||||
"run_id": row.run_id,
|
||||
"name": self._run_name_from_row(row),
|
||||
"dataset_id": row.dataset_id,
|
||||
"status": row.status,
|
||||
"started_at": format_utc_datetime(row.started_at),
|
||||
"completed_at": format_utc_datetime(row.completed_at),
|
||||
"total_items": row.total_items,
|
||||
"completed_items": row.completed_items,
|
||||
"overall_score": row.overall_score,
|
||||
"retrieval_config": row.retrieval_config or {},
|
||||
"metrics": row.metrics or {},
|
||||
}
|
||||
if row.status == "running":
|
||||
task = task_by_run_id.get(row.run_id)
|
||||
if task:
|
||||
run.update(progress=task.progress, message=task.message)
|
||||
runs.append(run)
|
||||
return runs
|
||||
except Exception as e:
|
||||
logger.error(f"获取评估运行历史失败: {e}")
|
||||
raise
|
||||
|
||||
async def get_run_results(
|
||||
self, kb_id: str, run_id: str, page: int = 1, page_size: int = 20, error_only: bool = False
|
||||
) -> dict[str, Any]:
|
||||
if not re.match(r"^run_[a-f0-9]{8}$", run_id):
|
||||
raise ValueError("Invalid run_id format")
|
||||
row = await self.eval_repo.get_run(run_id)
|
||||
if row is None or row.kb_id != kb_id:
|
||||
task = await tasker.get_task(run_id)
|
||||
if task:
|
||||
return {"run_id": run_id, "status": task.status, "progress": task.progress, "message": task.message}
|
||||
raise ValueError(f"Run not found for {run_id}")
|
||||
|
||||
start_idx = (page - 1) * page_size
|
||||
if error_only:
|
||||
total = 0
|
||||
paged_items = []
|
||||
offset = 0
|
||||
batch_size = 200
|
||||
while True:
|
||||
batch = await self.eval_repo.list_run_items(run_id, offset, batch_size)
|
||||
if not batch:
|
||||
break
|
||||
for item in batch:
|
||||
if not self._is_error_run_item(item):
|
||||
continue
|
||||
if start_idx <= total < start_idx + page_size:
|
||||
paged_items.append(self._run_item_to_dict(item))
|
||||
total += 1
|
||||
offset += batch_size
|
||||
else:
|
||||
total = await self.eval_repo.count_run_items(run_id)
|
||||
details = await self.eval_repo.list_run_items(run_id, start_idx, page_size)
|
||||
paged_items = [self._run_item_to_dict(item) for item in details]
|
||||
return {
|
||||
"run_id": row.run_id,
|
||||
"name": self._run_name_from_row(row),
|
||||
"status": row.status,
|
||||
"started_at": format_utc_datetime(row.started_at),
|
||||
"completed_at": format_utc_datetime(row.completed_at),
|
||||
"total_items": row.total_items or 0,
|
||||
"completed_items": row.completed_items or 0,
|
||||
"overall_score": row.overall_score,
|
||||
"retrieval_config": row.retrieval_config or {},
|
||||
"items": paged_items,
|
||||
"pagination": {
|
||||
"current_page": page,
|
||||
"page_size": page_size,
|
||||
"total": total,
|
||||
"total_pages": (total + page_size - 1) // page_size,
|
||||
"error_only": error_only,
|
||||
},
|
||||
}
|
||||
|
||||
async def delete_run(self, kb_id: str, run_id: str) -> None:
|
||||
if not re.match(r"^run_[a-f0-9]{8}$", run_id):
|
||||
raise ValueError("Invalid run_id format")
|
||||
row = await self.eval_repo.get_run(run_id)
|
||||
if row is None or row.kb_id != kb_id:
|
||||
raise ValueError("Run not found")
|
||||
await self.eval_repo.delete_run(run_id)
|
||||
logger.info(f"成功删除评估运行: {run_id}")
|
||||
@@ -0,0 +1,104 @@
|
||||
from yuxi.knowledge.base import KBNotFoundError, KnowledgeBase
|
||||
from yuxi.utils import logger
|
||||
|
||||
|
||||
class KnowledgeBaseFactory:
|
||||
"""知识库工厂类,负责创建不同类型的知识库实例"""
|
||||
|
||||
# 注册的知识库类型映射 {kb_type: kb_class}
|
||||
_kb_types: dict[str, type[KnowledgeBase]] = {}
|
||||
|
||||
@classmethod
|
||||
def register(cls, kb_class: type[KnowledgeBase]):
|
||||
"""
|
||||
注册知识库类型
|
||||
|
||||
Args:
|
||||
kb_class: 知识库类
|
||||
"""
|
||||
if not issubclass(kb_class, KnowledgeBase):
|
||||
raise ValueError("Knowledge base class must inherit from KnowledgeBase")
|
||||
if not kb_class.kb_type:
|
||||
raise ValueError("Knowledge base class must define kb_type")
|
||||
|
||||
cls._kb_types[kb_class.kb_type] = kb_class
|
||||
# logger.info(f"Registered knowledge base type: {kb_class.kb_type}")
|
||||
|
||||
@classmethod
|
||||
def create(cls, kb_type: str, work_dir: str, **kwargs) -> KnowledgeBase:
|
||||
"""
|
||||
创建知识库实例
|
||||
|
||||
Args:
|
||||
kb_type: 知识库类型
|
||||
work_dir: 工作目录
|
||||
**kwargs: 其他初始化参数
|
||||
|
||||
Returns:
|
||||
知识库实例
|
||||
|
||||
Raises:
|
||||
KBNotFoundError: 未知的知识库类型
|
||||
"""
|
||||
if kb_type not in cls._kb_types:
|
||||
available_types = list(cls._kb_types.keys())
|
||||
raise KBNotFoundError(f"Unknown knowledge base type: {kb_type}. Available types: {available_types}")
|
||||
|
||||
kb_class = cls._kb_types[kb_type]
|
||||
|
||||
try:
|
||||
# 创建实例
|
||||
instance = kb_class(work_dir, **kwargs)
|
||||
logger.info(f"Created {kb_type} knowledge base instance at {work_dir}")
|
||||
return instance
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to create {kb_type} knowledge base: {e}")
|
||||
raise
|
||||
|
||||
@classmethod
|
||||
def get_available_types(cls) -> dict[str, dict]:
|
||||
"""
|
||||
获取所有可用的知识库类型
|
||||
|
||||
Returns:
|
||||
知识库类型信息字典
|
||||
"""
|
||||
result = {}
|
||||
for kb_type, kb_class in cls._kb_types.items():
|
||||
result[kb_type] = {
|
||||
"name": kb_class.name,
|
||||
"description": kb_class.description,
|
||||
"requires_embedding_model": kb_class.requires_embedding_model,
|
||||
"supports_documents": kb_class.supports_documents,
|
||||
"create_params": kb_class.get_create_params_config(),
|
||||
}
|
||||
return result
|
||||
|
||||
@classmethod
|
||||
def get_kb_class(cls, kb_type: str) -> type[KnowledgeBase]:
|
||||
"""
|
||||
获取指定类型的知识库类。
|
||||
|
||||
Args:
|
||||
kb_type: 知识库类型
|
||||
|
||||
Returns:
|
||||
知识库类
|
||||
"""
|
||||
if kb_type not in cls._kb_types:
|
||||
available_types = list(cls._kb_types.keys())
|
||||
raise KBNotFoundError(f"Unknown knowledge base type: {kb_type}. Available types: {available_types}")
|
||||
return cls._kb_types[kb_type]
|
||||
|
||||
@classmethod
|
||||
def is_type_supported(cls, kb_type: str) -> bool:
|
||||
"""
|
||||
检查是否支持指定的知识库类型
|
||||
|
||||
Args:
|
||||
kb_type: 知识库类型
|
||||
|
||||
Returns:
|
||||
是否支持
|
||||
"""
|
||||
return kb_type in cls._kb_types
|
||||
@@ -0,0 +1,3 @@
|
||||
from .milvus_graph_service import MilvusGraphService
|
||||
|
||||
__all__ = ["MilvusGraphService"]
|
||||
@@ -0,0 +1,10 @@
|
||||
from .base import GraphExtractor, normalize_extraction_result
|
||||
from .factory import GraphExtractorFactory
|
||||
from .llm import LLMGraphExtractor
|
||||
|
||||
__all__ = [
|
||||
"GraphExtractor",
|
||||
"GraphExtractorFactory",
|
||||
"LLMGraphExtractor",
|
||||
"normalize_extraction_result",
|
||||
]
|
||||
@@ -0,0 +1,165 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from collections.abc import Callable
|
||||
from typing import Any
|
||||
|
||||
from yuxi.knowledge.graphs.graph_utils import normalize_entity_name
|
||||
|
||||
|
||||
class GraphExtractor(ABC):
|
||||
extractor_type: str
|
||||
|
||||
def __init__(self, options: dict[str, Any] | None = None):
|
||||
self.options = options or {}
|
||||
|
||||
@abstractmethod
|
||||
async def extract(self, text: str, *, chunk_metadata: dict[str, Any] | None = None) -> dict[str, Any]:
|
||||
pass
|
||||
|
||||
def validate_options(self) -> None:
|
||||
return None
|
||||
|
||||
|
||||
def normalize_extraction_result(result: dict[str, Any], extractor_type: str) -> dict[str, Any]:
|
||||
if not isinstance(result, dict):
|
||||
raise ValueError("extraction_result 必须是对象")
|
||||
|
||||
entities = result.get("entities") or []
|
||||
relations = result.get("relations") or []
|
||||
if not isinstance(entities, list) or not isinstance(relations, list):
|
||||
raise ValueError("extraction_result.entities 和 relations 必须是数组")
|
||||
|
||||
normalized_entities_by_key: dict[tuple[str, str], dict[str, Any]] = {}
|
||||
entity_refs: dict[str, dict[str, Any]] = {}
|
||||
|
||||
def add_entity(entity: Any, path: str) -> dict[str, Any]:
|
||||
normalized_entity = _normalize_entity(entity, path)
|
||||
key = _entity_key(normalized_entity)
|
||||
existing = normalized_entities_by_key.get(key)
|
||||
if existing is None:
|
||||
normalized_entities_by_key[key] = normalized_entity
|
||||
existing = normalized_entity
|
||||
else:
|
||||
_merge_attributes(existing, normalized_entity)
|
||||
|
||||
for ref in _entity_refs(entity, existing):
|
||||
entity_refs[ref] = existing
|
||||
return existing
|
||||
|
||||
for index, entity in enumerate(entities):
|
||||
add_entity(entity, f"entities[{index}]")
|
||||
|
||||
normalized_relations = []
|
||||
for index, relation in enumerate(relations):
|
||||
if not isinstance(relation, dict):
|
||||
raise ValueError("relations 元素必须是对象")
|
||||
source = _normalize_relation_endpoint(
|
||||
relation.get("source"),
|
||||
entity_refs,
|
||||
add_entity,
|
||||
result,
|
||||
f"relations[{index}].source",
|
||||
)
|
||||
target = _normalize_relation_endpoint(
|
||||
relation.get("target"),
|
||||
entity_refs,
|
||||
add_entity,
|
||||
result,
|
||||
f"relations[{index}].target",
|
||||
)
|
||||
text = str(relation.get("text") or "").strip()
|
||||
if not text:
|
||||
raise ValueError("relations[].text 不能为空")
|
||||
normalized_relations.append(
|
||||
{
|
||||
"source": source,
|
||||
"target": target,
|
||||
"text": text,
|
||||
"label": str(relation.get("label") or "RELATED_TO").strip() or "RELATED_TO",
|
||||
}
|
||||
)
|
||||
|
||||
metadata = dict(result.get("metadata") or {})
|
||||
metadata.setdefault("extractor_type", extractor_type)
|
||||
metadata.setdefault("schema_version", 1)
|
||||
return {
|
||||
"entities": list(normalized_entities_by_key.values()),
|
||||
"relations": normalized_relations,
|
||||
"metadata": metadata,
|
||||
}
|
||||
|
||||
|
||||
def _normalize_relation_endpoint(
|
||||
endpoint: Any,
|
||||
entity_refs: dict[str, dict[str, Any]],
|
||||
add_entity: Callable[[Any, str], dict[str, Any]],
|
||||
result: dict[str, Any],
|
||||
path: str,
|
||||
) -> dict[str, Any]:
|
||||
if isinstance(endpoint, dict):
|
||||
return add_entity(endpoint, path)
|
||||
|
||||
endpoint_ref = str(endpoint or "").strip()
|
||||
entity = entity_refs.get(endpoint_ref)
|
||||
if entity is None:
|
||||
raise ValueError(
|
||||
f"relations[].source/target 必须是实体对象,或引用 entities[].text/id,"
|
||||
f"未找到: {path}={endpoint_ref}, Result: {result}"
|
||||
)
|
||||
return entity
|
||||
|
||||
|
||||
def _normalize_entity(entity: Any, path: str) -> dict[str, Any]:
|
||||
if not isinstance(entity, dict):
|
||||
raise ValueError(f"{path} 必须是对象")
|
||||
|
||||
text = str(entity.get("text") or "").strip()
|
||||
if not text:
|
||||
raise ValueError(f"{path}.text 不能为空")
|
||||
|
||||
attributes = entity.get("attributes") or []
|
||||
if not isinstance(attributes, list):
|
||||
raise ValueError(f"{path}.attributes 必须是数组")
|
||||
|
||||
normalized_attributes = []
|
||||
for attribute in attributes:
|
||||
if not isinstance(attribute, dict):
|
||||
raise ValueError(f"{path}.attributes 元素必须是对象")
|
||||
attr_text = str(attribute.get("text") or "").strip()
|
||||
if not attr_text:
|
||||
continue
|
||||
normalized_attributes.append(
|
||||
{
|
||||
"text": attr_text,
|
||||
"label": str(attribute.get("label") or "Attribute").strip() or "Attribute",
|
||||
}
|
||||
)
|
||||
|
||||
return {
|
||||
"text": text,
|
||||
"label": str(entity.get("label") or "Entity").strip() or "Entity",
|
||||
"attributes": normalized_attributes,
|
||||
}
|
||||
|
||||
|
||||
def _entity_key(entity: dict[str, Any]) -> tuple[str, str]:
|
||||
return (normalize_entity_name(entity["text"]), entity["label"])
|
||||
|
||||
|
||||
def _entity_refs(raw_entity: Any, entity: dict[str, Any]) -> list[str]:
|
||||
refs = [entity["text"]]
|
||||
if isinstance(raw_entity, dict):
|
||||
entity_id = str(raw_entity.get("id") or "").strip()
|
||||
if entity_id:
|
||||
refs.append(entity_id)
|
||||
return refs
|
||||
|
||||
|
||||
def _merge_attributes(target: dict[str, Any], source: dict[str, Any]) -> None:
|
||||
known_attributes = {(attr["text"], attr["label"]) for attr in target.get("attributes") or []}
|
||||
for attribute in source.get("attributes") or []:
|
||||
attribute_key = (attribute["text"], attribute["label"])
|
||||
if attribute_key not in known_attributes:
|
||||
target.setdefault("attributes", []).append(attribute)
|
||||
known_attributes.add(attribute_key)
|
||||
@@ -0,0 +1,26 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from .base import GraphExtractor
|
||||
from .llm import LLMGraphExtractor
|
||||
|
||||
|
||||
class GraphExtractorFactory:
|
||||
_registry: dict[str, type[GraphExtractor]] = {
|
||||
"llm": LLMGraphExtractor,
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def create(cls, extractor_type: str | None, options: dict[str, Any] | None = None) -> GraphExtractor:
|
||||
normalized_type = (extractor_type or "").lower()
|
||||
extractor_class = cls._registry.get(normalized_type)
|
||||
if not extractor_class:
|
||||
raise ValueError(f"不支持的图谱抽取器类型: {extractor_type}")
|
||||
extractor = extractor_class(options or {})
|
||||
extractor.validate_options()
|
||||
return extractor
|
||||
|
||||
@classmethod
|
||||
def supported_types(cls) -> list[str]:
|
||||
return list(cls._registry.keys())
|
||||
@@ -0,0 +1,65 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
import json_repair
|
||||
|
||||
from yuxi.models.chat import select_model
|
||||
|
||||
from .base import GraphExtractor
|
||||
|
||||
DEFAULT_TRIPLE_EXTRACTION_PROMPT = """请从下面文本中抽取实体和实体关系,返回严格 JSON,不要输出解释。
|
||||
JSON 格式:
|
||||
{
|
||||
"relations": [
|
||||
{
|
||||
"source": {"text": "实体文本", "label": "实体类型", "attributes": [{"text": "属性值", "label": "属性名称"}]},
|
||||
"target": {"text": "实体文本", "label": "实体类型", "attributes": [{"text": "属性值", "label": "属性名称"}]},
|
||||
"text": "关系显示文本",
|
||||
"label": "关系类型"
|
||||
}
|
||||
]
|
||||
}
|
||||
"""
|
||||
|
||||
SCHEMA_INSTRUCTION = """抽取 Schema 约束:
|
||||
{schema}
|
||||
"""
|
||||
|
||||
|
||||
class LLMGraphExtractor(GraphExtractor):
|
||||
extractor_type = "llm"
|
||||
|
||||
def validate_options(self) -> None:
|
||||
if not self.options.get("model_spec"):
|
||||
raise ValueError("LLM 抽取器需要 model_spec")
|
||||
if self.options.get("prompt"):
|
||||
raise ValueError("LLM 图谱抽取器不支持自定义完整 Prompt,请使用 schema 配置抽取约束")
|
||||
concurrency_count = self.options.get("concurrency_count", 1)
|
||||
try:
|
||||
concurrency_count = int(concurrency_count)
|
||||
except (TypeError, ValueError) as exc:
|
||||
raise ValueError("LLM 抽取器 concurrency_count 必须是整数") from exc
|
||||
if concurrency_count < 1 or concurrency_count > 1000:
|
||||
raise ValueError("LLM 抽取器 concurrency_count 必须在 1 到 1000 之间")
|
||||
if self.options.get("model_params") is not None and not isinstance(self.options["model_params"], dict):
|
||||
raise ValueError("LLM 抽取器 model_params 必须是对象")
|
||||
|
||||
async def extract(self, text: str, *, chunk_metadata: dict[str, Any] | None = None) -> dict[str, Any]:
|
||||
self.validate_options()
|
||||
model = select_model(
|
||||
model_spec=self.options["model_spec"],
|
||||
timeout=60.0,
|
||||
model_params=self.options.get("model_params") or {},
|
||||
)
|
||||
prompt = self._build_prompt(text)
|
||||
response = await model.call(prompt, stream=False)
|
||||
parsed = json_repair.loads(response.content if response else "")
|
||||
return parsed
|
||||
|
||||
def _build_prompt(self, text: str) -> str:
|
||||
extraction_prompt = DEFAULT_TRIPLE_EXTRACTION_PROMPT
|
||||
schema = str(self.options.get("schema") or "").strip()
|
||||
if schema:
|
||||
extraction_prompt = f"{extraction_prompt}\n{SCHEMA_INSTRUCTION.format(schema=schema)}"
|
||||
return f"{extraction_prompt}\n\n文本:\n{text}"
|
||||
@@ -0,0 +1,150 @@
|
||||
"""图谱构建相关的纯函数工具集。
|
||||
|
||||
将数据变换逻辑从 MilvusGraphService 中抽离,
|
||||
使 service 类专注于 I/O 和业务编排。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from yuxi.utils import hashstr
|
||||
|
||||
|
||||
def normalize_entity_name(text: str) -> str:
|
||||
"""统一实体名称:去首尾空白、小写化、压缩内部连续空白。"""
|
||||
return " ".join(text.strip().lower().split())
|
||||
|
||||
|
||||
def compute_entity_id(kb_id: str, normalized_name: str, label: str) -> str:
|
||||
return hashstr(f"{kb_id}:{normalized_name}:{label}", length=32)
|
||||
|
||||
|
||||
def compute_triple_id(
|
||||
kb_id: str,
|
||||
source_normalized_name: str,
|
||||
source_label: str,
|
||||
relation_type: str,
|
||||
target_normalized_name: str,
|
||||
target_label: str,
|
||||
) -> str:
|
||||
return hashstr(
|
||||
f"{kb_id}:{source_normalized_name}:{source_label}:{relation_type}:{target_normalized_name}:{target_label}",
|
||||
length=32,
|
||||
)
|
||||
|
||||
|
||||
def graph_entity_collection_name(kb_id: str) -> str:
|
||||
return f"{kb_id}_entity"
|
||||
|
||||
|
||||
def graph_triple_collection_name(kb_id: str) -> str:
|
||||
return f"{kb_id}_triple"
|
||||
|
||||
|
||||
def build_graph_payload(normalized_result: dict[str, Any]) -> dict[str, Any]:
|
||||
"""将抽取器产出的标准化结果转换为 Neo4j 写入所需的图结构。
|
||||
|
||||
返回的 entities 已完成去重合并:同名同 label 的实体只保留一份,
|
||||
属性(attributes)取并集。
|
||||
"""
|
||||
entities: list[dict[str, Any]] = []
|
||||
entity_by_key: dict[tuple[str, str], dict[str, Any]] = {}
|
||||
|
||||
def add_entity(entity: dict[str, Any]) -> str:
|
||||
key = (normalize_entity_name(entity["text"]), entity.get("label") or "Entity")
|
||||
existing = entity_by_key.get(key)
|
||||
if existing is not None:
|
||||
known_attributes = {(attr["text"], attr["label"]) for attr in existing.get("attributes") or []}
|
||||
for attribute in entity.get("attributes") or []:
|
||||
attribute_key = (attribute["text"], attribute["label"])
|
||||
if attribute_key not in known_attributes:
|
||||
existing.setdefault("attributes", []).append(attribute)
|
||||
known_attributes.add(attribute_key)
|
||||
return existing["id"]
|
||||
|
||||
graph_entity = {
|
||||
"id": f"e{len(entities) + 1}",
|
||||
"text": entity["text"],
|
||||
"label": entity.get("label") or "Entity",
|
||||
"attributes": list(entity.get("attributes") or []),
|
||||
}
|
||||
entities.append(graph_entity)
|
||||
entity_by_key[key] = graph_entity
|
||||
return graph_entity["id"]
|
||||
|
||||
for entity in normalized_result["entities"]:
|
||||
add_entity(entity)
|
||||
|
||||
relations = []
|
||||
for relation in normalized_result["relations"]:
|
||||
relations.append(
|
||||
{
|
||||
"source": add_entity(relation["source"]),
|
||||
"target": add_entity(relation["target"]),
|
||||
"text": relation["text"],
|
||||
"label": relation.get("label") or "RELATED_TO",
|
||||
}
|
||||
)
|
||||
|
||||
return {"entities": entities, "relations": relations, "metadata": normalized_result["metadata"]}
|
||||
|
||||
|
||||
# ─── Cypher 模板 ────────────────────────────────────────────────
|
||||
# 将大段 Cypher 字符串集中管理,提升 write_chunk_graph 的可读性。
|
||||
|
||||
|
||||
def cypher_merge_chunk(db_label: str) -> str:
|
||||
"""MERGE Chunk 节点并写入元数据。"""
|
||||
return f"""
|
||||
MERGE (c:Chunk:MilvusKB:`{db_label}` {{chunk_id: $chunk_id}})
|
||||
SET c.file_id = $file_id,
|
||||
c.kb_id = $kb_id,
|
||||
c.chunk_index = $chunk_index,
|
||||
c.content_preview = $content_preview,
|
||||
c.start_char_pos = $start_char_pos,
|
||||
c.end_char_pos = $end_char_pos
|
||||
"""
|
||||
|
||||
|
||||
def cypher_merge_entity_mention(db_label: str) -> str:
|
||||
"""MERGE Entity 节点并创建 Chunk → Entity 的 MENTIONS 关系。"""
|
||||
return f"""
|
||||
MATCH (c:Chunk:MilvusKB:`{db_label}` {{chunk_id: $chunk_id}})
|
||||
MERGE (e:Entity:MilvusKB:`{db_label}` {{
|
||||
kb_id: $kb_id,
|
||||
normalized_name: $normalized_name,
|
||||
label: $entity_label
|
||||
}})
|
||||
SET e.entity_id = $entity_id,
|
||||
e.name = $name,
|
||||
e.attributes = $attributes
|
||||
MERGE (c)-[m:MENTIONS {{chunk_id: $chunk_id, file_id: $file_id, kb_id: $kb_id}}]->(e)
|
||||
"""
|
||||
|
||||
|
||||
def cypher_merge_relation(db_label: str) -> str:
|
||||
"""MERGE 两个 Entity 之间的 RELATION 边。"""
|
||||
return f"""
|
||||
MATCH (source:Entity:MilvusKB:`{db_label}` {{
|
||||
kb_id: $kb_id,
|
||||
normalized_name: $source_name,
|
||||
label: $source_label
|
||||
}})
|
||||
MATCH (target:Entity:MilvusKB:`{db_label}` {{
|
||||
kb_id: $kb_id,
|
||||
normalized_name: $target_name,
|
||||
label: $target_label
|
||||
}})
|
||||
MERGE (source)-[r:RELATION {{
|
||||
kb_id: $kb_id,
|
||||
chunk_id: $chunk_id,
|
||||
source_name: $source_name,
|
||||
target_name: $target_name,
|
||||
type: $relation_type
|
||||
}}]->(target)
|
||||
SET r.triple_id = $triple_id,
|
||||
r.text = $text,
|
||||
r.file_id = $file_id,
|
||||
r.extractor_type = $extractor_type
|
||||
"""
|
||||
@@ -0,0 +1,950 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import weakref
|
||||
from typing import Any
|
||||
|
||||
from yuxi.knowledge.graphs.extractors import GraphExtractor, GraphExtractorFactory, normalize_extraction_result
|
||||
from yuxi.knowledge.graphs.graph_utils import (
|
||||
build_graph_payload,
|
||||
compute_entity_id,
|
||||
compute_triple_id,
|
||||
cypher_merge_chunk,
|
||||
cypher_merge_entity_mention,
|
||||
cypher_merge_relation,
|
||||
normalize_entity_name,
|
||||
)
|
||||
from yuxi.knowledge.graphs.milvus_graph_vector_store import MilvusGraphVectorStore
|
||||
from yuxi.repositories.knowledge_base_repository import KnowledgeBaseRepository
|
||||
from yuxi.repositories.knowledge_chunk_repository import KnowledgeChunkRepository
|
||||
from yuxi.repositories.knowledge_graph_repository import KnowledgeGraphRepository
|
||||
from yuxi.storage.neo4j import (
|
||||
Neo4jConnectionManager,
|
||||
get_shared_neo4j_connection,
|
||||
neo4j_read,
|
||||
neo4j_write,
|
||||
safe_neo4j_label,
|
||||
)
|
||||
from yuxi.utils import logger
|
||||
from yuxi.utils.datetime_utils import utc_isoformat
|
||||
|
||||
GRAPH_CONFIG_KEY = "graph_build_config"
|
||||
GRAPH_TASK_TYPE = "knowledge_graph_index"
|
||||
NEO4J_QUERY_OFFLOAD_LIMIT = 8
|
||||
_neo4j_query_offload_semaphore_refs: dict[
|
||||
int,
|
||||
tuple[weakref.ReferenceType[asyncio.AbstractEventLoop], weakref.ReferenceType[asyncio.Semaphore]],
|
||||
] = {}
|
||||
|
||||
|
||||
def _get_neo4j_query_offload_semaphore() -> asyncio.Semaphore:
|
||||
loop = asyncio.get_running_loop()
|
||||
loop_id = id(loop)
|
||||
entry = _neo4j_query_offload_semaphore_refs.get(loop_id)
|
||||
if entry is not None:
|
||||
loop_ref, semaphore_ref = entry
|
||||
semaphore = semaphore_ref()
|
||||
if loop_ref() is loop and semaphore is not None:
|
||||
return semaphore
|
||||
|
||||
semaphore = asyncio.Semaphore(NEO4J_QUERY_OFFLOAD_LIMIT)
|
||||
|
||||
def cleanup(ref, stale_loop_id=loop_id):
|
||||
current_entry = _neo4j_query_offload_semaphore_refs.get(stale_loop_id)
|
||||
if current_entry is not None and current_entry[1] is ref:
|
||||
_neo4j_query_offload_semaphore_refs.pop(stale_loop_id, None)
|
||||
|
||||
_neo4j_query_offload_semaphore_refs[loop_id] = (weakref.ref(loop), weakref.ref(semaphore, cleanup))
|
||||
return semaphore
|
||||
|
||||
|
||||
async def _run_neo4j_query_io(func, /, *args, **kwargs):
|
||||
semaphore = _get_neo4j_query_offload_semaphore()
|
||||
await semaphore.acquire()
|
||||
task = asyncio.create_task(asyncio.to_thread(func, *args, **kwargs))
|
||||
|
||||
def release_capacity(completed_task: asyncio.Task):
|
||||
semaphore.release()
|
||||
if completed_task.cancelled():
|
||||
return
|
||||
completed_task.exception()
|
||||
|
||||
task.add_done_callback(release_capacity)
|
||||
return await asyncio.shield(task)
|
||||
|
||||
|
||||
class MilvusGraphService:
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
kb_id: str | None = None,
|
||||
kb_repo: KnowledgeBaseRepository | None = None,
|
||||
chunk_repo: KnowledgeChunkRepository | None = None,
|
||||
graph_repo: KnowledgeGraphRepository | None = None,
|
||||
graph_vector_store: MilvusGraphVectorStore | None = None,
|
||||
neo4j_connection: Neo4jConnectionManager | None = None,
|
||||
):
|
||||
self.kb_id = kb_id
|
||||
self.kb_repo = kb_repo or KnowledgeBaseRepository()
|
||||
self.chunk_repo = chunk_repo or KnowledgeChunkRepository()
|
||||
self.graph_repo = graph_repo or KnowledgeGraphRepository()
|
||||
self._graph_vector_store = graph_vector_store
|
||||
self._connection = neo4j_connection
|
||||
|
||||
@property
|
||||
def connection(self) -> Neo4jConnectionManager:
|
||||
if self._connection is None:
|
||||
self._connection = get_shared_neo4j_connection()
|
||||
return self._connection
|
||||
|
||||
@property
|
||||
def graph_vector_store(self) -> MilvusGraphVectorStore:
|
||||
if self._graph_vector_store is None:
|
||||
self._graph_vector_store = MilvusGraphVectorStore()
|
||||
return self._graph_vector_store
|
||||
|
||||
@property
|
||||
def driver(self):
|
||||
return self.connection.driver
|
||||
|
||||
async def get_status(self, kb_id: str, *, tasker: Any = None) -> dict[str, Any]:
|
||||
kb = await self._get_milvus_kb(kb_id)
|
||||
params = dict(kb.additional_params or {})
|
||||
config = params.get(GRAPH_CONFIG_KEY) or {}
|
||||
total_chunks, pending_chunks, indexed_chunks, graph_counts = await asyncio.gather(
|
||||
self.chunk_repo.count_by_kb_id(kb_id),
|
||||
self.chunk_repo.count_graph_pending_by_kb_id(kb_id),
|
||||
self.chunk_repo.count_graph_indexed_by_kb_id(kb_id),
|
||||
self.graph_repo.count_by_kb_id(kb_id),
|
||||
)
|
||||
entity_count, relationship_count = graph_counts
|
||||
|
||||
build_task_status = None
|
||||
build_task_progress = 0
|
||||
if tasker is not None:
|
||||
active_task = await tasker.find_task_by_payload(
|
||||
task_type=GRAPH_TASK_TYPE,
|
||||
payload_match={"kb_id": kb_id},
|
||||
statuses={"pending", "running"},
|
||||
)
|
||||
if active_task:
|
||||
build_task_status = active_task.status
|
||||
build_task_progress = round(active_task.progress)
|
||||
else:
|
||||
failed_task = await tasker.find_task_by_payload(
|
||||
task_type=GRAPH_TASK_TYPE,
|
||||
payload_match={"kb_id": kb_id},
|
||||
statuses={"failed", "cancelled"},
|
||||
)
|
||||
if failed_task:
|
||||
build_task_status = "failed"
|
||||
build_task_progress = 0
|
||||
|
||||
return {
|
||||
"kb_id": kb_id,
|
||||
"kb_type": kb.kb_type,
|
||||
"configured": bool(config),
|
||||
"locked": bool(config.get("locked")),
|
||||
"config": self._public_config(config),
|
||||
"total_chunks": total_chunks,
|
||||
"pending_chunks": pending_chunks,
|
||||
"indexed_chunks": indexed_chunks,
|
||||
"entity_count": entity_count,
|
||||
"relationship_count": relationship_count,
|
||||
"build_task_status": build_task_status,
|
||||
"build_task_progress": build_task_progress,
|
||||
}
|
||||
|
||||
async def configure(
|
||||
self,
|
||||
kb_id: str,
|
||||
extractor_type: str,
|
||||
extractor_options: dict[str, Any],
|
||||
created_by: str,
|
||||
) -> dict:
|
||||
kb = await self._get_milvus_kb(kb_id)
|
||||
additional_params = dict(kb.additional_params or {})
|
||||
existing_config = additional_params.get(GRAPH_CONFIG_KEY) or {}
|
||||
normalized_extractor_type = (extractor_type or "").lower()
|
||||
if existing_config.get("locked"):
|
||||
existing_extractor_type = (existing_config.get("extractor_type") or "").lower()
|
||||
if normalized_extractor_type != existing_extractor_type:
|
||||
raise ValueError("图谱抽取器类型已锁定,只能修改模型、Schema 等抽取参数")
|
||||
|
||||
extractor_options = extractor_options or {}
|
||||
if normalized_extractor_type == "llm" and extractor_options.get("prompt"):
|
||||
raise ValueError("LLM 图谱抽取器不支持自定义完整 Prompt,请使用 schema 配置抽取约束")
|
||||
GraphExtractorFactory.create(normalized_extractor_type, extractor_options)
|
||||
config = {
|
||||
"locked": True,
|
||||
"extractor_type": normalized_extractor_type,
|
||||
"extractor_options": extractor_options or {},
|
||||
"created_at": existing_config.get("created_at") or utc_isoformat(),
|
||||
"created_by": existing_config.get("created_by") or created_by,
|
||||
}
|
||||
if existing_config.get("locked"):
|
||||
config["updated_at"] = utc_isoformat()
|
||||
config["updated_by"] = created_by
|
||||
additional_params[GRAPH_CONFIG_KEY] = config
|
||||
await self.kb_repo.update(kb_id, {"additional_params": additional_params})
|
||||
return config
|
||||
|
||||
async def build_pending_chunks(self, kb_id: str, *, batch_size: int, context=None) -> dict[str, Any]:
|
||||
kb = await self._get_milvus_kb(kb_id)
|
||||
config = self._get_locked_config(kb.additional_params or {})
|
||||
extractor_options = self._runtime_extractor_options(config)
|
||||
extractor = GraphExtractorFactory.create(config["extractor_type"], extractor_options)
|
||||
worker_count = self._get_worker_count(config)
|
||||
total_pending = await self.chunk_repo.count_graph_pending_by_kb_id(kb_id)
|
||||
processed = 0
|
||||
failed = 0
|
||||
failed_chunk_ids: set[str] = set()
|
||||
write_lock = asyncio.Lock()
|
||||
|
||||
while True:
|
||||
if context is not None:
|
||||
await context.raise_if_cancelled()
|
||||
chunks = await self.chunk_repo.list_graph_pending_by_kb_id(kb_id, batch_size)
|
||||
unprocessed = [c for c in chunks if c.chunk_id not in failed_chunk_ids]
|
||||
if not unprocessed:
|
||||
break
|
||||
|
||||
queue: asyncio.Queue[Any] = asyncio.Queue()
|
||||
for chunk in unprocessed:
|
||||
queue.put_nowait(chunk)
|
||||
|
||||
async def worker() -> None:
|
||||
nonlocal processed, failed
|
||||
while True:
|
||||
if context is not None:
|
||||
await context.raise_if_cancelled()
|
||||
try:
|
||||
chunk = queue.get_nowait()
|
||||
except asyncio.QueueEmpty:
|
||||
return
|
||||
try:
|
||||
extraction_result = await self._get_chunk_extraction_result(kb_id, chunk, extractor)
|
||||
async with write_lock:
|
||||
entities, triples = await asyncio.to_thread(
|
||||
self.write_chunk_graph,
|
||||
kb_id,
|
||||
chunk,
|
||||
extraction_result,
|
||||
)
|
||||
await self.graph_repo.upsert_chunk_graph(
|
||||
kb_id=kb_id,
|
||||
file_id=chunk.file_id,
|
||||
chunk_id=chunk.chunk_id,
|
||||
entities=entities,
|
||||
triples=triples,
|
||||
)
|
||||
await self.graph_vector_store.insert_missing_graph_records(
|
||||
kb_id=kb_id,
|
||||
embedding_model_spec=kb.embedding_model_spec,
|
||||
entities=entities,
|
||||
triples=triples,
|
||||
)
|
||||
await self.chunk_repo.mark_graph_indexed(
|
||||
chunk.chunk_id,
|
||||
ent_ids=[entity["entity_id"] for entity in entities],
|
||||
)
|
||||
processed += 1
|
||||
except Exception as exc:
|
||||
logger.error(f"Chunk 图谱构建失败 chunk_id={chunk.chunk_id}: {exc}")
|
||||
failed_chunk_ids.add(chunk.chunk_id)
|
||||
failed += 1
|
||||
finally:
|
||||
queue.task_done()
|
||||
|
||||
if context is not None:
|
||||
completed = processed + failed
|
||||
progress = 5.0 + min(90.0, completed / max(total_pending, 1) * 90.0)
|
||||
await context.set_progress(progress, f"图谱构建 {completed}/{total_pending},失败 {failed}")
|
||||
|
||||
workers = [asyncio.create_task(worker()) for _ in range(min(worker_count, len(unprocessed)))]
|
||||
try:
|
||||
await asyncio.gather(*workers)
|
||||
except Exception:
|
||||
for task in workers:
|
||||
task.cancel()
|
||||
await asyncio.gather(*workers, return_exceptions=True)
|
||||
raise
|
||||
|
||||
remaining = await self.chunk_repo.count_graph_pending_by_kb_id(kb_id)
|
||||
return {"kb_id": kb_id, "success": processed, "failed": failed, "remaining": remaining}
|
||||
|
||||
@staticmethod
|
||||
def _get_worker_count(config: dict[str, Any]) -> int:
|
||||
if (config.get("extractor_type") or "").lower() != "llm":
|
||||
return 1
|
||||
try:
|
||||
worker_count = int((config.get("extractor_options") or {}).get("concurrency_count") or 1)
|
||||
except (TypeError, ValueError):
|
||||
return 1
|
||||
return max(1, min(worker_count, 1000))
|
||||
|
||||
@staticmethod
|
||||
def _runtime_extractor_options(config: dict[str, Any]) -> dict[str, Any]:
|
||||
options = dict(config.get("extractor_options") or {})
|
||||
options.pop("prompt", None)
|
||||
return options
|
||||
|
||||
async def _get_chunk_extraction_result(self, kb_id: str, chunk, extractor: GraphExtractor) -> dict[str, Any]:
|
||||
extractor_type = extractor.extractor_type
|
||||
if chunk.extraction_result:
|
||||
return normalize_extraction_result(chunk.extraction_result, extractor_type)
|
||||
|
||||
extraction_result = await extractor.extract(
|
||||
chunk.content,
|
||||
chunk_metadata={
|
||||
"kb_id": kb_id,
|
||||
"chunk_id": chunk.chunk_id,
|
||||
"file_id": chunk.file_id,
|
||||
"chunk_index": chunk.chunk_index,
|
||||
},
|
||||
)
|
||||
normalized_result = normalize_extraction_result(extraction_result, extractor_type)
|
||||
await self.chunk_repo.update_extraction_result(chunk.chunk_id, normalized_result)
|
||||
return normalized_result
|
||||
|
||||
def write_chunk_graph(
|
||||
self,
|
||||
kb_id: str,
|
||||
chunk,
|
||||
normalized_result: dict[str, Any],
|
||||
) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
|
||||
"""将单个 chunk 的抽取结果写入 Neo4j。"""
|
||||
label = safe_neo4j_label(kb_id)
|
||||
graph_payload = build_graph_payload(normalized_result)
|
||||
relation_extractor_type = graph_payload["metadata"].get("extractor_type", "unknown")
|
||||
entities = graph_payload["entities"]
|
||||
relations = graph_payload["relations"]
|
||||
entity_by_id = {entity["id"]: entity for entity in entities}
|
||||
entity_records = self._build_entity_records(kb_id, entities)
|
||||
entity_record_by_local_id = {
|
||||
entity["id"]: record for entity, record in zip(entities, entity_records, strict=True)
|
||||
}
|
||||
triple_records = self._build_triple_records(kb_id, relations, entity_record_by_local_id, graph_payload)
|
||||
content_preview = (chunk.content or "")[:300]
|
||||
|
||||
# 预构建 Cypher 模板(同一 chunk 内复用)
|
||||
merge_chunk_cypher = cypher_merge_chunk(label)
|
||||
merge_entity_cypher = cypher_merge_entity_mention(label)
|
||||
merge_relation_cypher = cypher_merge_relation(label)
|
||||
|
||||
def query(tx):
|
||||
# 1. MERGE Chunk 节点
|
||||
tx.run(
|
||||
merge_chunk_cypher,
|
||||
chunk_id=chunk.chunk_id,
|
||||
file_id=chunk.file_id,
|
||||
kb_id=kb_id,
|
||||
chunk_index=chunk.chunk_index,
|
||||
content_preview=content_preview,
|
||||
start_char_pos=chunk.start_char_pos,
|
||||
end_char_pos=chunk.end_char_pos,
|
||||
)
|
||||
|
||||
# 2. MERGE Entity 节点 + Chunk→Entity (MENTIONS)
|
||||
for entity in entities:
|
||||
entity_record = entity_record_by_local_id[entity["id"]]
|
||||
tx.run(
|
||||
merge_entity_cypher,
|
||||
chunk_id=chunk.chunk_id,
|
||||
file_id=chunk.file_id,
|
||||
kb_id=kb_id,
|
||||
entity_id=entity_record["entity_id"],
|
||||
normalized_name=normalize_entity_name(entity["text"]),
|
||||
entity_label=entity.get("label") or "Entity",
|
||||
name=entity["text"],
|
||||
attributes=json.dumps(entity.get("attributes") or [], ensure_ascii=False),
|
||||
)
|
||||
|
||||
# 3. MERGE Entity→Entity (RELATION) 边
|
||||
for relation in relations:
|
||||
source = entity_by_id[relation["source"]]
|
||||
target = entity_by_id[relation["target"]]
|
||||
source_record = entity_record_by_local_id[relation["source"]]
|
||||
target_record = entity_record_by_local_id[relation["target"]]
|
||||
relation_type = relation.get("label") or "RELATED_TO"
|
||||
triple_id = compute_triple_id(
|
||||
kb_id,
|
||||
source_record["normalized_name"],
|
||||
source_record["label"],
|
||||
relation_type,
|
||||
target_record["normalized_name"],
|
||||
target_record["label"],
|
||||
)
|
||||
tx.run(
|
||||
merge_relation_cypher,
|
||||
kb_id=kb_id,
|
||||
chunk_id=chunk.chunk_id,
|
||||
file_id=chunk.file_id,
|
||||
source_name=normalize_entity_name(source["text"]),
|
||||
source_label=source.get("label") or "Entity",
|
||||
target_name=normalize_entity_name(target["text"]),
|
||||
target_label=target.get("label") or "Entity",
|
||||
relation_type=relation_type,
|
||||
triple_id=triple_id,
|
||||
text=relation["text"],
|
||||
extractor_type=relation_extractor_type,
|
||||
)
|
||||
|
||||
neo4j_write(self.driver, query)
|
||||
return entity_records, triple_records
|
||||
|
||||
def _build_entity_records(self, kb_id: str, entities: list[dict[str, Any]]) -> list[dict[str, Any]]:
|
||||
records = []
|
||||
for entity in entities:
|
||||
label = entity.get("label") or "Entity"
|
||||
normalized_name = normalize_entity_name(entity["text"])
|
||||
entity_id = compute_entity_id(kb_id, normalized_name, label)
|
||||
records.append(
|
||||
{
|
||||
"entity_id": entity_id,
|
||||
"kb_id": kb_id,
|
||||
"normalized_name": normalized_name,
|
||||
"label": label,
|
||||
"name": entity["text"],
|
||||
"attributes": entity.get("attributes") or [],
|
||||
"content": normalized_name,
|
||||
}
|
||||
)
|
||||
return records
|
||||
|
||||
def _build_triple_records(
|
||||
self,
|
||||
kb_id: str,
|
||||
relations: list[dict[str, Any]],
|
||||
entity_record_by_local_id: dict[str, dict[str, Any]],
|
||||
graph_payload: dict[str, Any],
|
||||
) -> list[dict[str, Any]]:
|
||||
records = []
|
||||
seen_triple_ids: set[str] = set()
|
||||
extractor_type = graph_payload["metadata"].get("extractor_type", "unknown")
|
||||
for relation in relations:
|
||||
source_record = entity_record_by_local_id[relation["source"]]
|
||||
target_record = entity_record_by_local_id[relation["target"]]
|
||||
relation_type = relation.get("label") or "RELATED_TO"
|
||||
triple_id = compute_triple_id(
|
||||
kb_id,
|
||||
source_record["normalized_name"],
|
||||
source_record["label"],
|
||||
relation_type,
|
||||
target_record["normalized_name"],
|
||||
target_record["label"],
|
||||
)
|
||||
if triple_id in seen_triple_ids:
|
||||
continue
|
||||
seen_triple_ids.add(triple_id)
|
||||
content = f"{source_record['normalized_name']} → {relation_type} → {target_record['normalized_name']}"
|
||||
records.append(
|
||||
{
|
||||
"triple_id": triple_id,
|
||||
"kb_id": kb_id,
|
||||
"source_entity_id": source_record["entity_id"],
|
||||
"target_entity_id": target_record["entity_id"],
|
||||
"relation_type": relation_type,
|
||||
"content": content,
|
||||
"text": relation["text"],
|
||||
"extractor_type": extractor_type,
|
||||
}
|
||||
)
|
||||
return records
|
||||
|
||||
async def reset(self, kb_id: str, *, clear_extraction_result: bool, clear_config: bool) -> dict[str, Any]:
|
||||
kb = await self._get_milvus_kb(kb_id)
|
||||
await asyncio.to_thread(self.delete_graph, kb_id)
|
||||
await self.graph_repo.delete_by_kb_id(kb_id)
|
||||
reset_chunks = await self.chunk_repo.reset_graph_state_by_kb_id(kb_id, clear_extraction_result)
|
||||
if clear_config:
|
||||
additional_params = dict(kb.additional_params or {})
|
||||
additional_params.pop(GRAPH_CONFIG_KEY, None)
|
||||
await self.kb_repo.update(kb_id, {"additional_params": additional_params})
|
||||
return {
|
||||
"message": "图谱构建状态已重置",
|
||||
"status": "success",
|
||||
"reset_chunks": reset_chunks,
|
||||
"clear_extraction_result": clear_extraction_result,
|
||||
"clear_config": clear_config,
|
||||
}
|
||||
|
||||
def delete_graph(self, kb_id: str) -> None:
|
||||
label = safe_neo4j_label(kb_id)
|
||||
|
||||
def query(tx):
|
||||
tx.run(f"MATCH (n:MilvusKB:`{label}`) DETACH DELETE n")
|
||||
|
||||
neo4j_write(self.driver, query)
|
||||
self.graph_vector_store.drop_graph_collections(kb_id)
|
||||
|
||||
async def delete_file_graph(self, kb_id: str, file_id: str) -> None:
|
||||
orphan_entity_ids, orphan_triple_ids = await self.graph_repo.delete_file_references(file_id)
|
||||
await self.graph_vector_store.delete_graph_records(
|
||||
kb_id,
|
||||
entity_ids=orphan_entity_ids,
|
||||
triple_ids=orphan_triple_ids,
|
||||
)
|
||||
await asyncio.to_thread(self._delete_file_graph_from_neo4j, kb_id, file_id)
|
||||
|
||||
def _delete_file_graph_from_neo4j(self, kb_id: str, file_id: str) -> None:
|
||||
label = safe_neo4j_label(kb_id)
|
||||
|
||||
def query(tx):
|
||||
tx.run(
|
||||
f"""
|
||||
MATCH (:Chunk:MilvusKB:`{label}`)-[m:MENTIONS {{kb_id: $kb_id, file_id: $file_id}}]->
|
||||
(:Entity:MilvusKB:`{label}`)
|
||||
DELETE m
|
||||
""",
|
||||
kb_id=kb_id,
|
||||
file_id=file_id,
|
||||
)
|
||||
tx.run(
|
||||
f"""
|
||||
MATCH (:Entity:MilvusKB:`{label}`)-[r:RELATION {{kb_id: $kb_id, file_id: $file_id}}]->
|
||||
(:Entity:MilvusKB:`{label}`)
|
||||
DELETE r
|
||||
""",
|
||||
kb_id=kb_id,
|
||||
file_id=file_id,
|
||||
)
|
||||
tx.run(
|
||||
f"""
|
||||
MATCH (c:Chunk:MilvusKB:`{label}` {{kb_id: $kb_id, file_id: $file_id}})
|
||||
DETACH DELETE c
|
||||
""",
|
||||
kb_id=kb_id,
|
||||
file_id=file_id,
|
||||
)
|
||||
tx.run(
|
||||
f"""
|
||||
MATCH (e:Entity:MilvusKB:`{label}` {{kb_id: $kb_id}})
|
||||
WHERE NOT ()-[:MENTIONS]->(e)
|
||||
DETACH DELETE e
|
||||
""",
|
||||
kb_id=kb_id,
|
||||
)
|
||||
|
||||
neo4j_write(self.driver, query)
|
||||
|
||||
async def query_nodes(
|
||||
self,
|
||||
kb_id: str | None = None,
|
||||
*,
|
||||
keyword: str = "",
|
||||
max_depth: int = 1,
|
||||
max_nodes: int = 50,
|
||||
exclude_chunk: bool = False,
|
||||
) -> dict[str, Any]:
|
||||
effective_kb_id = kb_id or self.kb_id
|
||||
if not effective_kb_id:
|
||||
return {"nodes": [], "edges": []}
|
||||
|
||||
label = safe_neo4j_label(effective_kb_id)
|
||||
limit = max_nodes
|
||||
try:
|
||||
return await _run_neo4j_query_io(
|
||||
self._query_nodes_sync,
|
||||
effective_kb_id,
|
||||
label,
|
||||
keyword,
|
||||
limit,
|
||||
max_depth,
|
||||
exclude_chunk,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"Milvus graph query failed: {e}")
|
||||
return {"nodes": [], "edges": []}
|
||||
|
||||
def _query_nodes_sync(
|
||||
self,
|
||||
kb_id: str,
|
||||
label: str,
|
||||
keyword: str,
|
||||
limit: int,
|
||||
max_depth: int,
|
||||
exclude_chunk: bool,
|
||||
) -> dict[str, Any]:
|
||||
with self.driver.session() as session:
|
||||
result = session.run(
|
||||
self._build_query(label, keyword, limit, max_depth, exclude_chunk),
|
||||
keyword=keyword,
|
||||
limit=limit,
|
||||
edge_limit=limit * 10,
|
||||
)
|
||||
return self._process_query_result(result, limit, kb_id, exclude_chunk)
|
||||
|
||||
async def query_seed_subgraph(
|
||||
self,
|
||||
kb_id: str,
|
||||
*,
|
||||
entity_ids: list[str],
|
||||
max_nodes: int,
|
||||
) -> dict[str, Any]:
|
||||
if not entity_ids:
|
||||
return {"nodes": [], "edges": []}
|
||||
seed_entity_ids = list(dict.fromkeys(entity_ids))
|
||||
label = safe_neo4j_label(kb_id)
|
||||
cypher = f"""
|
||||
MATCH (seed:Entity:MilvusKB:`{label}`)
|
||||
WHERE seed.entity_id IN $entity_ids
|
||||
MATCH p = (seed)-[*1..2]-(n:MilvusKB:`{label}`)
|
||||
WITH p LIMIT $path_limit
|
||||
WITH collect(p) AS paths
|
||||
UNWIND paths AS node_path
|
||||
UNWIND nodes(node_path) AS node
|
||||
WITH paths, collect(DISTINCT node) AS graph_nodes
|
||||
UNWIND paths AS rel_path
|
||||
UNWIND relationships(rel_path) AS rel
|
||||
RETURN graph_nodes AS nodes, collect(DISTINCT rel) AS edges
|
||||
"""
|
||||
try:
|
||||
return await _run_neo4j_query_io(
|
||||
self._query_seed_subgraph_sync,
|
||||
kb_id,
|
||||
cypher,
|
||||
seed_entity_ids,
|
||||
max_nodes,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"Milvus seed subgraph query failed: {e}")
|
||||
return {"nodes": [], "edges": []}
|
||||
|
||||
def _query_seed_subgraph_sync(
|
||||
self,
|
||||
kb_id: str,
|
||||
cypher: str,
|
||||
entity_ids: list[str],
|
||||
max_nodes: int,
|
||||
) -> dict[str, Any]:
|
||||
with self.driver.session() as session:
|
||||
record = session.run(
|
||||
cypher,
|
||||
entity_ids=entity_ids,
|
||||
path_limit=max(max_nodes, 1) * 4,
|
||||
).single()
|
||||
if not record:
|
||||
return {"nodes": [], "edges": []}
|
||||
return self._process_subgraph_record(record, max_nodes, kb_id)
|
||||
|
||||
async def query_and_rank_chunks_by_ppr(
|
||||
self,
|
||||
kb_id: str,
|
||||
seed_weights: dict[str, float],
|
||||
*,
|
||||
max_nodes: int,
|
||||
top_k: int,
|
||||
damping: float,
|
||||
) -> list[tuple[str, float]]:
|
||||
if not seed_weights:
|
||||
return []
|
||||
subgraph = await self.query_seed_subgraph(
|
||||
kb_id,
|
||||
entity_ids=list(seed_weights.keys()),
|
||||
max_nodes=max_nodes,
|
||||
)
|
||||
return self.rank_chunks_by_ppr(subgraph, seed_weights, top_k=top_k, damping=damping)
|
||||
|
||||
@staticmethod
|
||||
def rank_chunks_by_ppr(
|
||||
subgraph: dict[str, Any],
|
||||
seed_weights: dict[str, float],
|
||||
*,
|
||||
top_k: int,
|
||||
damping: float,
|
||||
) -> list[tuple[str, float]]:
|
||||
nodes = subgraph.get("nodes") or []
|
||||
edges = subgraph.get("edges") or []
|
||||
if not nodes:
|
||||
return []
|
||||
|
||||
try:
|
||||
import igraph as ig
|
||||
except ImportError:
|
||||
logger.error("Graph retrieval requires python-igraph. Please install igraph.")
|
||||
return []
|
||||
|
||||
node_ids = [node["id"] for node in nodes]
|
||||
index_by_id = {node_id: index for index, node_id in enumerate(node_ids)}
|
||||
edge_indices = [
|
||||
(index_by_id[edge["source_id"]], index_by_id[edge["target_id"]])
|
||||
for edge in edges
|
||||
if edge.get("source_id") in index_by_id and edge.get("target_id") in index_by_id
|
||||
]
|
||||
if not edge_indices:
|
||||
return []
|
||||
|
||||
graph = ig.Graph(n=len(nodes), edges=edge_indices, directed=False)
|
||||
reset = [0.0] * len(nodes)
|
||||
chunk_node_indexes: list[tuple[int, str]] = []
|
||||
for index, node in enumerate(nodes):
|
||||
properties = node.get("properties") or {}
|
||||
if node.get("type") == "Chunk" and properties.get("chunk_id"):
|
||||
chunk_node_indexes.append((index, properties["chunk_id"]))
|
||||
continue
|
||||
entity_id = properties.get("entity_id")
|
||||
if entity_id in seed_weights:
|
||||
reset[index] = seed_weights[entity_id]
|
||||
|
||||
reset_total = sum(reset)
|
||||
if reset_total <= 0 or not chunk_node_indexes:
|
||||
return []
|
||||
reset = [value / reset_total for value in reset]
|
||||
scores = graph.personalized_pagerank(damping=min(max(damping, 0.1), 0.99), reset=reset)
|
||||
ranked = sorted(
|
||||
((chunk_id, float(scores[index])) for index, chunk_id in chunk_node_indexes),
|
||||
key=lambda item: item[1],
|
||||
reverse=True,
|
||||
)
|
||||
return ranked[:top_k]
|
||||
|
||||
async def get_labels(self, kb_id: str | None = None) -> list[str]:
|
||||
effective_kb_id = kb_id or self.kb_id
|
||||
if not effective_kb_id:
|
||||
return []
|
||||
label = safe_neo4j_label(effective_kb_id)
|
||||
|
||||
cypher = f"""
|
||||
MATCH (n:MilvusKB:`{label}`)
|
||||
UNWIND labels(n) AS node_label
|
||||
WITH DISTINCT node_label
|
||||
WHERE node_label <> 'MilvusKB' AND node_label <> $kb_id
|
||||
RETURN node_label
|
||||
ORDER BY node_label
|
||||
"""
|
||||
try:
|
||||
records = await _run_neo4j_query_io(self._get_labels_sync, cypher, effective_kb_id)
|
||||
return [record["node_label"] for record in records]
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to get Milvus graph labels: {e}")
|
||||
return []
|
||||
|
||||
def _get_labels_sync(self, cypher: str, kb_id: str) -> list[Any]:
|
||||
return neo4j_read(self.driver, cypher, kb_id=kb_id)
|
||||
|
||||
async def get_stats(self, kb_id: str | None = None) -> dict[str, Any]:
|
||||
effective_kb_id = kb_id or self.kb_id
|
||||
if not effective_kb_id:
|
||||
return {"total_nodes": 0, "total_edges": 0, "entity_types": []}
|
||||
label = safe_neo4j_label(effective_kb_id)
|
||||
|
||||
stats_cypher = f"""
|
||||
MATCH (n:MilvusKB:`{label}`)
|
||||
WITH count(n) AS node_count
|
||||
OPTIONAL MATCH (:MilvusKB:`{label}`)-[r]->(:MilvusKB:`{label}`)
|
||||
RETURN node_count, count(r) AS edge_count
|
||||
"""
|
||||
label_cypher = f"""
|
||||
MATCH (n:Entity:MilvusKB:`{label}`)
|
||||
WITH n.label AS entity_label, count(*) AS count
|
||||
RETURN entity_label, count
|
||||
ORDER BY count DESC
|
||||
"""
|
||||
try:
|
||||
return await _run_neo4j_query_io(self._get_stats_sync, stats_cypher, label_cypher)
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to get Milvus graph stats: {e}")
|
||||
return {"total_nodes": 0, "total_edges": 0, "entity_types": []}
|
||||
|
||||
def _get_stats_sync(self, stats_cypher: str, label_cypher: str) -> dict[str, Any]:
|
||||
with self.driver.session() as session:
|
||||
stats = session.run(stats_cypher).single()
|
||||
label_stats = session.run(label_cypher)
|
||||
return {
|
||||
"total_nodes": stats["node_count"] if stats else 0,
|
||||
"total_edges": stats["edge_count"] if stats else 0,
|
||||
"entity_types": [{"type": row["entity_label"], "count": row["count"]} for row in label_stats],
|
||||
}
|
||||
|
||||
async def _get_milvus_kb(self, kb_id: str):
|
||||
kb = await self.kb_repo.get_by_kb_id(kb_id)
|
||||
if kb is None:
|
||||
raise ValueError(f"知识库 {kb_id} 不存在")
|
||||
if (kb.kb_type or "").lower() != "milvus":
|
||||
raise ValueError("仅 Milvus 知识库支持独立图谱构建")
|
||||
return kb
|
||||
|
||||
def _get_locked_config(self, additional_params: dict[str, Any]) -> dict[str, Any]:
|
||||
config = additional_params.get(GRAPH_CONFIG_KEY) or {}
|
||||
if not config.get("locked"):
|
||||
raise ValueError("请先确认并锁定图谱抽取配置")
|
||||
if not config.get("extractor_type"):
|
||||
raise ValueError("图谱抽取配置缺少 extractor_type")
|
||||
return config
|
||||
|
||||
def _public_config(self, config: dict[str, Any]) -> dict[str, Any] | None:
|
||||
if not config:
|
||||
return None
|
||||
return {
|
||||
"locked": bool(config.get("locked")),
|
||||
"extractor_type": config.get("extractor_type"),
|
||||
"extractor_options": self._runtime_extractor_options(config),
|
||||
"created_at": config.get("created_at"),
|
||||
"created_by": config.get("created_by"),
|
||||
"updated_at": config.get("updated_at"),
|
||||
"updated_by": config.get("updated_by"),
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def _build_where(exclude_chunk: bool, keyword: str) -> str:
|
||||
clauses = []
|
||||
if exclude_chunk:
|
||||
clauses.append("NOT n:Chunk")
|
||||
if keyword and keyword != "*":
|
||||
clauses.append(
|
||||
"(toLower(coalesce(n.name, '')) CONTAINS toLower($keyword)"
|
||||
" OR toLower(coalesce(n.content_preview, '')) CONTAINS toLower($keyword)"
|
||||
" OR toLower(coalesce(n.chunk_id, '')) CONTAINS toLower($keyword))"
|
||||
)
|
||||
return "WHERE " + " AND ".join(clauses) if clauses else ""
|
||||
|
||||
def _build_query(self, label: str, keyword: str, limit: int, max_depth: int, exclude_chunk: bool = False) -> str:
|
||||
where = self._build_where(exclude_chunk, keyword)
|
||||
m_exclude = " WHERE NOT m:Chunk" if exclude_chunk else ""
|
||||
|
||||
if max_depth <= 0:
|
||||
return f"""
|
||||
MATCH (n:MilvusKB:`{label}`)
|
||||
{where}
|
||||
RETURN n AS h, null AS r, null AS t
|
||||
LIMIT $limit
|
||||
"""
|
||||
|
||||
return f"""
|
||||
MATCH (n:MilvusKB:`{label}`)
|
||||
{where}
|
||||
WITH n LIMIT $limit
|
||||
OPTIONAL MATCH (n)-[r]-(m:MilvusKB:`{label}`){m_exclude}
|
||||
RETURN n AS h, r AS r, m AS t
|
||||
LIMIT $edge_limit
|
||||
"""
|
||||
|
||||
def _process_query_result(self, result, limit: int, kb_id: str, exclude_chunk: bool = False) -> dict[str, Any]:
|
||||
nodes = []
|
||||
edges = []
|
||||
node_ids = set()
|
||||
edge_ids = set()
|
||||
|
||||
for record in result:
|
||||
for key in ("h", "t"):
|
||||
raw_node = record.get(key)
|
||||
if raw_node is None:
|
||||
continue
|
||||
node = self._normalize_node(raw_node, kb_id)
|
||||
if not node or node["id"] in node_ids:
|
||||
continue
|
||||
if exclude_chunk and node.get("type") == "Chunk":
|
||||
continue
|
||||
nodes.append(node)
|
||||
node_ids.add(node["id"])
|
||||
raw_edge = record.get("r")
|
||||
if raw_edge is not None:
|
||||
edge = self._normalize_edge(raw_edge)
|
||||
if edge and edge["id"] not in edge_ids:
|
||||
edges.append(edge)
|
||||
edge_ids.add(edge["id"])
|
||||
if len(nodes) >= limit:
|
||||
break
|
||||
|
||||
return self._finalize_subgraph_result(nodes, edges, limit)
|
||||
|
||||
def _process_subgraph_record(self, record: Any, limit: int, kb_id: str) -> dict[str, Any]:
|
||||
nodes = []
|
||||
edges = []
|
||||
node_ids = set()
|
||||
edge_ids = set()
|
||||
|
||||
for raw_node in record.get("nodes") or []:
|
||||
node = self._normalize_node(raw_node, kb_id)
|
||||
if not node or node["id"] in node_ids:
|
||||
continue
|
||||
nodes.append(node)
|
||||
node_ids.add(node["id"])
|
||||
if len(nodes) >= limit:
|
||||
break
|
||||
|
||||
for raw_edge in record.get("edges") or []:
|
||||
edge = self._normalize_edge(raw_edge)
|
||||
if not edge or edge["id"] in edge_ids:
|
||||
continue
|
||||
if edge["source_id"] not in node_ids or edge["target_id"] not in node_ids:
|
||||
continue
|
||||
edges.append(edge)
|
||||
edge_ids.add(edge["id"])
|
||||
|
||||
return self._finalize_subgraph_result(nodes, edges, limit)
|
||||
|
||||
@staticmethod
|
||||
def _finalize_subgraph_result(
|
||||
nodes: list[dict[str, Any]], edges: list[dict[str, Any]], limit: int
|
||||
) -> dict[str, Any]:
|
||||
limit = max(0, limit)
|
||||
final_nodes = nodes[:limit]
|
||||
node_ids = {node["id"] for node in final_nodes}
|
||||
final_edges = [
|
||||
edge for edge in edges if edge.get("source_id") in node_ids and edge.get("target_id") in node_ids
|
||||
]
|
||||
return {"nodes": final_nodes, "edges": final_edges[: limit * 2]}
|
||||
|
||||
def _normalize_node(self, raw_node: Any, kb_id: str | None = None) -> dict[str, Any]:
|
||||
if hasattr(raw_node, "element_id"):
|
||||
node_id = raw_node.element_id
|
||||
labels = list(raw_node.labels)
|
||||
properties = dict(raw_node.items())
|
||||
elif isinstance(raw_node, dict):
|
||||
node_id = raw_node.get("id") or raw_node.get("element_id")
|
||||
labels = raw_node.get("labels", [])
|
||||
properties = raw_node.get("properties") or {k: v for k, v in raw_node.items() if k not in {"id", "labels"}}
|
||||
else:
|
||||
return {}
|
||||
|
||||
effective_kb_id = kb_id or self.kb_id
|
||||
db_label = properties.get("kb_id") or effective_kb_id
|
||||
filtered_labels = [label for label in labels if label not in {"MilvusKB", db_label}]
|
||||
entity_type = "Chunk" if "Chunk" in labels else properties.get("label", "Entity")
|
||||
name = properties.get("name") or properties.get("content_preview") or properties.get("chunk_id") or "Unknown"
|
||||
return {
|
||||
"id": node_id,
|
||||
"name": name,
|
||||
"original_id": node_id,
|
||||
"type": entity_type,
|
||||
"labels": filtered_labels,
|
||||
"properties": properties,
|
||||
"normalized": {
|
||||
"name": name,
|
||||
"type": entity_type,
|
||||
"source": "milvus",
|
||||
},
|
||||
"graph_type": "milvus",
|
||||
}
|
||||
|
||||
def _normalize_edge(self, raw_edge: Any) -> dict[str, Any]:
|
||||
if hasattr(raw_edge, "element_id"):
|
||||
edge_id = raw_edge.element_id
|
||||
edge_type = raw_edge.type
|
||||
source_id = raw_edge.start_node.element_id
|
||||
target_id = raw_edge.end_node.element_id
|
||||
properties = dict(raw_edge.items())
|
||||
edge_type = properties.get("type") or edge_type
|
||||
elif isinstance(raw_edge, dict):
|
||||
edge_id = raw_edge.get("id")
|
||||
edge_type = raw_edge.get("type")
|
||||
source_id = raw_edge.get("source_id")
|
||||
target_id = raw_edge.get("target_id")
|
||||
properties = raw_edge.get("properties", {})
|
||||
else:
|
||||
return {}
|
||||
|
||||
return {
|
||||
"id": edge_id,
|
||||
"source_id": source_id,
|
||||
"target_id": target_id,
|
||||
"type": edge_type,
|
||||
"properties": properties,
|
||||
"normalized": {
|
||||
"type": edge_type,
|
||||
"direction": "directed",
|
||||
},
|
||||
}
|
||||
@@ -0,0 +1,322 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
from functools import partial
|
||||
from typing import Any
|
||||
|
||||
from pymilvus import (
|
||||
Collection,
|
||||
CollectionSchema,
|
||||
DataType,
|
||||
FieldSchema,
|
||||
Function,
|
||||
FunctionType,
|
||||
connections,
|
||||
db,
|
||||
utility,
|
||||
)
|
||||
|
||||
from yuxi.knowledge.graphs.graph_utils import graph_entity_collection_name, graph_triple_collection_name
|
||||
from yuxi.knowledge.implementations.milvus import (
|
||||
CONTENT_ANALYZER_PARAMS,
|
||||
CONTENT_SPARSE_FIELD,
|
||||
VECTOR_METRIC_TYPE,
|
||||
_run_milvus_query_io,
|
||||
)
|
||||
from yuxi.models.embed import select_embedding_model
|
||||
from yuxi.models.providers.cache import model_cache
|
||||
from yuxi.utils import hashstr, logger
|
||||
|
||||
|
||||
class MilvusGraphVectorStore:
|
||||
def __init__(self):
|
||||
self.milvus_token = os.getenv("MILVUS_TOKEN") or ""
|
||||
self.milvus_uri = os.getenv("MILVUS_URI") or "http://localhost:19530"
|
||||
self.milvus_db = os.getenv("MILVUS_DB") or "yuxi"
|
||||
self.connection_alias = f"milvus_graph_{hashstr(self.milvus_uri, 6)}"
|
||||
self._init_connection()
|
||||
|
||||
def _init_connection(self) -> None:
|
||||
if not connections.has_connection(self.connection_alias):
|
||||
connections.connect(alias=self.connection_alias, uri=self.milvus_uri, token=self.milvus_token)
|
||||
try:
|
||||
if self.milvus_db not in db.list_database():
|
||||
db.create_database(self.milvus_db)
|
||||
db.using_database(self.milvus_db)
|
||||
except Exception as exc:
|
||||
logger.warning(f"Milvus graph database operation failed, using default: {exc}")
|
||||
|
||||
async def insert_missing_graph_records(
|
||||
self,
|
||||
*,
|
||||
kb_id: str,
|
||||
embedding_model_spec: str,
|
||||
entities: list[dict[str, Any]],
|
||||
triples: list[dict[str, Any]],
|
||||
) -> None:
|
||||
if not entities and not triples:
|
||||
return
|
||||
|
||||
embedding_info = model_cache.get_model_info(embedding_model_spec)
|
||||
if not embedding_info or embedding_info.model_type != "embedding":
|
||||
raise ValueError(f"Unsupported embedding model: {embedding_model_spec}")
|
||||
|
||||
entity_collection = self._get_or_create_entity_collection(kb_id, embedding_info)
|
||||
triple_collection = self._get_or_create_triple_collection(kb_id, embedding_info)
|
||||
|
||||
entity_ids = [entity["entity_id"] for entity in entities]
|
||||
triple_ids = [triple["triple_id"] for triple in triples]
|
||||
existing_entity_ids, existing_triple_ids = await asyncio.gather(
|
||||
asyncio.to_thread(self._query_existing_ids, entity_collection, entity_ids),
|
||||
asyncio.to_thread(self._query_existing_ids, triple_collection, triple_ids),
|
||||
)
|
||||
|
||||
missing_entities = [entity for entity in entities if entity["entity_id"] not in existing_entity_ids]
|
||||
missing_triples = [triple for triple in triples if triple["triple_id"] not in existing_triple_ids]
|
||||
if not missing_entities and not missing_triples:
|
||||
return
|
||||
|
||||
embed = self._get_embedding_function(embedding_model_spec)
|
||||
entity_embeddings, triple_embeddings = await asyncio.gather(
|
||||
embed([entity["content"] for entity in missing_entities]) if missing_entities else self._empty_embeddings(),
|
||||
embed([triple["content"] for triple in missing_triples]) if missing_triples else self._empty_embeddings(),
|
||||
)
|
||||
|
||||
if missing_entities:
|
||||
await asyncio.to_thread(self._insert_entities, entity_collection, missing_entities, entity_embeddings)
|
||||
if missing_triples:
|
||||
await asyncio.to_thread(self._insert_triples, triple_collection, missing_triples, triple_embeddings)
|
||||
|
||||
async def delete_graph_records(self, kb_id: str, *, entity_ids: list[str], triple_ids: list[str]) -> None:
|
||||
tasks = []
|
||||
if entity_ids:
|
||||
tasks.append(asyncio.to_thread(self._delete_ids, graph_entity_collection_name(kb_id), entity_ids))
|
||||
if triple_ids:
|
||||
tasks.append(asyncio.to_thread(self._delete_ids, graph_triple_collection_name(kb_id), triple_ids))
|
||||
if tasks:
|
||||
await asyncio.gather(*tasks)
|
||||
|
||||
async def search_entities(
|
||||
self,
|
||||
*,
|
||||
kb_id: str,
|
||||
query_text: str,
|
||||
embedding_model_spec: str,
|
||||
top_k: int,
|
||||
) -> list[dict[str, Any]]:
|
||||
collection_name = graph_entity_collection_name(kb_id)
|
||||
has_collection = await _run_milvus_query_io(
|
||||
utility.has_collection, collection_name, using=self.connection_alias
|
||||
)
|
||||
if not has_collection:
|
||||
return []
|
||||
return await self._search_graph_collection(
|
||||
collection_name=collection_name,
|
||||
query_text=query_text,
|
||||
embedding_model_spec=embedding_model_spec,
|
||||
top_k=top_k,
|
||||
output_fields=["id", "content"],
|
||||
)
|
||||
|
||||
async def search_triples(
|
||||
self,
|
||||
*,
|
||||
kb_id: str,
|
||||
query_text: str,
|
||||
embedding_model_spec: str,
|
||||
top_k: int,
|
||||
) -> list[dict[str, Any]]:
|
||||
collection_name = graph_triple_collection_name(kb_id)
|
||||
has_collection = await _run_milvus_query_io(
|
||||
utility.has_collection, collection_name, using=self.connection_alias
|
||||
)
|
||||
if not has_collection:
|
||||
return []
|
||||
return await self._search_graph_collection(
|
||||
collection_name=collection_name,
|
||||
query_text=query_text,
|
||||
embedding_model_spec=embedding_model_spec,
|
||||
top_k=top_k,
|
||||
output_fields=["id", "content", "source_id", "target_id"],
|
||||
)
|
||||
|
||||
def drop_graph_collections(self, kb_id: str) -> None:
|
||||
for collection_name in [graph_entity_collection_name(kb_id), graph_triple_collection_name(kb_id)]:
|
||||
try:
|
||||
if utility.has_collection(collection_name, using=self.connection_alias):
|
||||
utility.drop_collection(collection_name, using=self.connection_alias)
|
||||
logger.info(f"Dropped Milvus graph collection {collection_name}")
|
||||
except Exception as exc:
|
||||
logger.error(f"Failed to drop Milvus graph collection {collection_name}: {exc}")
|
||||
|
||||
async def _empty_embeddings(self) -> list:
|
||||
return []
|
||||
|
||||
def _get_embedding_function(self, embedding_model_spec: str):
|
||||
model = select_embedding_model(embedding_model_spec)
|
||||
batch_size = int(getattr(model, "batch_size", 40) or 40)
|
||||
return partial(model.abatch_encode, batch_size=batch_size)
|
||||
|
||||
async def _search_graph_collection(
|
||||
self,
|
||||
*,
|
||||
collection_name: str,
|
||||
query_text: str,
|
||||
embedding_model_spec: str,
|
||||
top_k: int,
|
||||
output_fields: list[str],
|
||||
) -> list[dict[str, Any]]:
|
||||
if top_k <= 0:
|
||||
return []
|
||||
embed = self._get_embedding_function(embedding_model_spec)
|
||||
query_embedding = await embed([query_text])
|
||||
return await _run_milvus_query_io(
|
||||
self._search_graph_collection_sync,
|
||||
collection_name,
|
||||
query_embedding,
|
||||
max(top_k, 1),
|
||||
output_fields,
|
||||
)
|
||||
|
||||
def _search_graph_collection_sync(
|
||||
self,
|
||||
collection_name: str,
|
||||
query_embedding: list,
|
||||
top_k: int,
|
||||
output_fields: list[str],
|
||||
) -> list[dict[str, Any]]:
|
||||
collection = Collection(name=collection_name, using=self.connection_alias)
|
||||
collection.load()
|
||||
return self._search_loaded_collection(collection, query_embedding, top_k, output_fields)
|
||||
|
||||
def _search_loaded_collection(
|
||||
self,
|
||||
collection: Collection,
|
||||
query_embedding: list,
|
||||
top_k: int,
|
||||
output_fields: list[str],
|
||||
) -> list[dict[str, Any]]:
|
||||
results = collection.search(
|
||||
data=query_embedding,
|
||||
anns_field="embedding",
|
||||
param={"metric_type": VECTOR_METRIC_TYPE, "params": {"nprobe": 10}},
|
||||
limit=top_k,
|
||||
output_fields=output_fields,
|
||||
)
|
||||
if not results or not results[0]:
|
||||
return []
|
||||
|
||||
records = []
|
||||
for hit in results[0]:
|
||||
entity = hit.entity
|
||||
record = {field: entity.get(field) for field in output_fields}
|
||||
record["score"] = float(hit.distance or 0.0)
|
||||
records.append(record)
|
||||
return records
|
||||
|
||||
def _get_or_create_entity_collection(self, kb_id: str, embedding_info: Any) -> Collection:
|
||||
collection_name = graph_entity_collection_name(kb_id)
|
||||
fields = [
|
||||
FieldSchema(name="id", dtype=DataType.VARCHAR, max_length=100, is_primary=True),
|
||||
FieldSchema(
|
||||
name="content",
|
||||
dtype=DataType.VARCHAR,
|
||||
max_length=65535,
|
||||
enable_analyzer=True,
|
||||
analyzer_params=CONTENT_ANALYZER_PARAMS,
|
||||
),
|
||||
FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=embedding_info.dimension or 1024),
|
||||
FieldSchema(name=CONTENT_SPARSE_FIELD, dtype=DataType.SPARSE_FLOAT_VECTOR),
|
||||
]
|
||||
return self._get_or_create_collection(collection_name, fields, embedding_info)
|
||||
|
||||
def _get_or_create_triple_collection(self, kb_id: str, embedding_info: Any) -> Collection:
|
||||
collection_name = graph_triple_collection_name(kb_id)
|
||||
fields = [
|
||||
FieldSchema(name="id", dtype=DataType.VARCHAR, max_length=100, is_primary=True),
|
||||
FieldSchema(
|
||||
name="content",
|
||||
dtype=DataType.VARCHAR,
|
||||
max_length=65535,
|
||||
enable_analyzer=True,
|
||||
analyzer_params=CONTENT_ANALYZER_PARAMS,
|
||||
),
|
||||
FieldSchema(name="source_id", dtype=DataType.VARCHAR, max_length=100),
|
||||
FieldSchema(name="target_id", dtype=DataType.VARCHAR, max_length=100),
|
||||
FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=embedding_info.dimension or 1024),
|
||||
FieldSchema(name=CONTENT_SPARSE_FIELD, dtype=DataType.SPARSE_FLOAT_VECTOR),
|
||||
]
|
||||
return self._get_or_create_collection(collection_name, fields, embedding_info)
|
||||
|
||||
def _get_or_create_collection(
|
||||
self, collection_name: str, fields: list[FieldSchema], embedding_info: Any
|
||||
) -> Collection:
|
||||
if utility.has_collection(collection_name, using=self.connection_alias):
|
||||
return Collection(name=collection_name, using=self.connection_alias)
|
||||
|
||||
bm25_function = Function(
|
||||
name="content_bm25",
|
||||
input_field_names=["content"],
|
||||
output_field_names=[CONTENT_SPARSE_FIELD],
|
||||
function_type=FunctionType.BM25,
|
||||
)
|
||||
schema = CollectionSchema(
|
||||
fields=fields,
|
||||
description=f"Knowledge graph collection {collection_name} using {embedding_info.model_id}",
|
||||
functions=[bm25_function],
|
||||
)
|
||||
collection = Collection(name=collection_name, schema=schema, using=self.connection_alias)
|
||||
collection.create_index(
|
||||
"embedding", {"metric_type": VECTOR_METRIC_TYPE, "index_type": "IVF_FLAT", "params": {"nlist": 1024}}
|
||||
)
|
||||
collection.create_index(
|
||||
CONTENT_SPARSE_FIELD,
|
||||
{
|
||||
"metric_type": "BM25",
|
||||
"index_type": "SPARSE_INVERTED_INDEX",
|
||||
"params": {"inverted_index_algo": "DAAT_MAXSCORE"},
|
||||
},
|
||||
)
|
||||
return collection
|
||||
|
||||
def _query_existing_ids(self, collection: Collection, ids: list[str]) -> set[str]:
|
||||
if not ids:
|
||||
return set()
|
||||
collection.load()
|
||||
existing_ids: set[str] = set()
|
||||
for start in range(0, len(ids), 1000):
|
||||
batch = ids[start : start + 1000]
|
||||
quoted_ids = ", ".join(f'"{item}"' for item in batch)
|
||||
rows = collection.query(expr=f"id in [{quoted_ids}]", output_fields=["id"])
|
||||
existing_ids.update(row["id"] for row in rows)
|
||||
return existing_ids
|
||||
|
||||
def _insert_entities(self, collection: Collection, entities: list[dict[str, Any]], embeddings: list) -> None:
|
||||
collection.insert(
|
||||
[
|
||||
[entity["entity_id"] for entity in entities],
|
||||
[entity["content"] for entity in entities],
|
||||
embeddings,
|
||||
]
|
||||
)
|
||||
|
||||
def _insert_triples(self, collection: Collection, triples: list[dict[str, Any]], embeddings: list) -> None:
|
||||
collection.insert(
|
||||
[
|
||||
[triple["triple_id"] for triple in triples],
|
||||
[triple["content"] for triple in triples],
|
||||
[triple["source_entity_id"] for triple in triples],
|
||||
[triple["target_entity_id"] for triple in triples],
|
||||
embeddings,
|
||||
]
|
||||
)
|
||||
|
||||
def _delete_ids(self, collection_name: str, ids: list[str]) -> None:
|
||||
if not utility.has_collection(collection_name, using=self.connection_alias):
|
||||
return
|
||||
collection = Collection(name=collection_name, using=self.connection_alias)
|
||||
for start in range(0, len(ids), 1000):
|
||||
batch = ids[start : start + 1000]
|
||||
quoted_ids = ", ".join(f'"{item}"' for item in batch)
|
||||
collection.delete(expr=f"id in [{quoted_ids}]")
|
||||
@@ -0,0 +1,14 @@
|
||||
"""知识库具体实现模块
|
||||
|
||||
包含各种知识库的具体实现:
|
||||
- MilvusKB: 基于 Milvus 的向量知识库
|
||||
- DifyKB: 基于 Dify 检索 API 的只读知识库
|
||||
- NotionKB: 基于 Notion Data Source 的只读知识库
|
||||
"""
|
||||
|
||||
from .dify import DifyKB
|
||||
from .milvus import MilvusKB
|
||||
from .notion import NotionKB
|
||||
from .read_only_connectors import ReadOnlyConnectors
|
||||
|
||||
__all__ = ["MilvusKB", "DifyKB", "NotionKB", "ReadOnlyConnectors"]
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user