zhkj_llm_function 0.2.0__py3-none-any.whl

zhkj_llm_function/__init__.py

@@ -0,0 +1,78 @@
+"""zhkj_llm_function：基于 @agent_def 装饰器 + Pydantic 类型注解反射的 Agent 插件框架。
+
+执行引擎：tinyagent-py（ReAct 循环、工具调用、上下文压缩、token 统计）。
+框架层：@agent_def 声明、Input/Output 模型校验、外部工具注册、统一返回结构。
+
+快速开始：
+
+    from zhkj_llm_function import BaseAgentGroup, agent_def, ToolDef
+
+    class MyPlugin(BaseAgentGroup):
+        @agent_def(ttl=86400, compact_config=None)
+        async def my_agent(self, payload: MyInput, tools: list[ToolDef] | None = None) -> MyOutput:
+            \"\"\"
+            角色：xxx
+            业务规则：
+            1. xxx
+            \"\"\"
+            ...
+
+    plugin = MyPlugin(model="gpt-5", api_key="...")
+    result = await plugin.my_agent({"...": "..."}, tools=[...])   # 返回 MyOutput 实例
+"""
+
+from .agent import AgentSpec, BaseAgentGroup, DEFAULT_TTL, agent_def
+from .errors import (
+    AgentNotRegisteredError,
+    FrameworkError,
+    InputSchemaError,
+    OutputSchemaError,
+    SessionError,
+    ToolExecutionError,
+    ToolNotFoundError,
+)
+from .protocol import (
+    TOOL_CALL_END,
+    TOOL_CALL_START,
+    ToolCall,
+    extract_json_payload,
+    parse_tool_call,
+    render_tool_call,
+)
+from .result import build_result, envelope_to_result
+from .runtime import DEFAULT_COMPACT_CONFIG, AgentRuntime
+from .session import Session, SessionStore
+from .tools import ToolDef, ToolSpec, execute_tool, normalize_tools, register_tools
+
+__all__ = [
+    "AgentSpec",
+    "AgentRuntime",
+    "BaseAgentGroup",
+    "DEFAULT_TTL",
+    "DEFAULT_COMPACT_CONFIG",
+    "agent_def",
+    "FrameworkError",
+    "AgentNotRegisteredError",
+    "InputSchemaError",
+    "OutputSchemaError",
+    "ToolExecutionError",
+    "ToolNotFoundError",
+    "SessionError",
+    "TOOL_CALL_START",
+    "TOOL_CALL_END",
+    "ToolCall",
+    "parse_tool_call",
+    "render_tool_call",
+    "extract_json_payload",
+    "build_result",
+    "envelope_to_result",
+    "Session",
+    "SessionStore",
+    "ToolDef",
+    "ToolSpec",
+    "normalize_tools",
+    "register_tools",
+    "execute_tool",
+]
+
+__version__ = "0.2.0"

zhkj_llm_function/agent.py

@@ -0,0 +1,285 @@
+"""Agent 声明层：`agent_def` 装饰器与 `BaseAgentGroup`。
+
+框架核心模式：基于 `@agent_def` 装饰器 + Pydantic 类型注解反射；
+执行引擎使用 tinyagent-py。
+"""
+
+import functools
+import inspect
+from typing import Any, Callable, Dict, List, Optional
+
+from pydantic import BaseModel
+
+from .errors import AgentNotRegisteredError, FrameworkError
+from .schema import describe_docstring, model_schema, resolve_model_hints
+
+DEFAULT_TTL = 86400
+
+
+class AgentSpec(object):
+    """一个 agent 的声明信息（由 @agent_def 生成）。"""
+
+    __slots__ = (
+        "name",
+        "ttl",
+        "compact_config",
+        "input_model",
+        "output_model",
+        "input_schema",
+        "output_schema",
+        "business_instruction",
+        "func",
+        "group_cls",
+    )
+
+    def __init__(
+        self,
+        name: str,
+        ttl: int,
+        compact_config: Optional[Dict[str, Any]],
+        input_model: type,
+        output_model: type,
+        business_instruction: str,
+        func: Callable[..., Any],
+        group_cls: Optional[type] = None,
+    ) -> None:
+        self.name = name
+        self.ttl = ttl
+        self.compact_config = compact_config
+        self.input_model = input_model
+        self.output_model = output_model
+        self.input_schema = model_schema(input_model)
+        self.output_schema = model_schema(output_model)
+        self.business_instruction = business_instruction
+        self.func = func
+        self.group_cls = group_cls
+
+    def __repr__(self) -> str:  # pragma: no cover - 调试用
+        return "AgentSpec(name={0!r}, input={1}, output={2})".format(
+            self.name, self.input_model.__name__, self.output_model.__name__
+        )
+
+
+def _require_model(annotation: Any, kind: str, agent_name: str) -> type:
+    """校验注解必须是 Pydantic BaseModel 子类。"""
+    if annotation is None or not isinstance(annotation, type):
+        raise FrameworkError(
+            "agent [{0}] 的 {1} 必须使用 Pydantic BaseModel 进行类型注解".format(agent_name, kind)
+        )
+    if not issubclass(annotation, BaseModel):
+        raise FrameworkError(
+            "agent [{0}] 的 {1} 必须是 Pydantic BaseModel 子类，当前为 {2}".format(
+                agent_name, kind, annotation.__name__
+            )
+        )
+    return annotation
+
+
+def agent_def(
+    ttl: int = DEFAULT_TTL, compact_config: Optional[Dict[str, Any]] = None
+) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
+    """标记一个 async 方法为 agent。
+
+    参数：
+    - ttl：会话有效期（秒），常用 86400（一天）。
+    - compact_config：None=使用全局默认前缀压缩策略，也可传入自定义压缩配置 dict。
+
+    方法签名固定为：
+        async def xxx_agent(self, payload: XXXInput, tools: list[ToolDef] | None = None) -> XXXOutput
+    """
+
+    def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
+        if not inspect.iscoroutinefunction(func):
+            raise FrameworkError(
+                "@agent_def 只能装饰 async 方法: {0}".format(getattr(func, "__name__", func))
+            )
+        if not isinstance(ttl, int) or ttl <= 0:
+            raise FrameworkError("@agent_def 的 ttl 必须是正整数")
+        if compact_config is not None and not isinstance(compact_config, dict):
+            raise FrameworkError("@agent_def 的 compact_config 必须是 dict 或 None")
+
+        signature = inspect.signature(func)
+        parameters = list(signature.parameters.values())
+        if len(parameters) < 2 or parameters[0].name != "self":
+            raise FrameworkError("@agent_def 方法必须保留 self 作为第一个参数")
+        if parameters[1].name != "payload":
+            raise FrameworkError("@agent_def 方法的第二个参数必须命名为 payload")
+        if len(parameters) < 3 or parameters[2].name != "tools":
+            raise FrameworkError("@agent_def 方法的第三个参数必须命名为 tools")
+
+        hints = resolve_model_hints(func)
+        input_model = _require_model(hints.get("payload"), "入参(payload)", func.__name__)
+        output_model = _require_model(hints.get("return"), "出参(return)", func.__name__)
+
+        spec = AgentSpec(
+            name=func.__name__,
+            ttl=ttl,
+            compact_config=compact_config,
+            input_model=input_model,
+            output_model=output_model,
+            business_instruction=describe_docstring(func),
+            func=func,
+        )
+
+        @functools.wraps(func)
+        async def wrapper(
+            self: Any,
+            payload: Any = None,
+            tools: Optional[List[Any]] = None,
+            session_id: Optional[str] = None,
+        ) -> Any:
+            """以普通方法方式调用：await plugin.xxx_agent(payload, tools)。
+
+            返回值即声明的出参模型实例（外层可按出参签名直接接收）；
+            扩展信息（ok / finish_reason / meta）通过代理属性提供。
+            执行失败时抛出对应异常（InputSchemaError / OutputSchemaError 等）。
+            """
+            runtime = getattr(self, "_runtime", None)
+            if runtime is not None:
+                from .result import envelope_to_result
+
+                envelope = await runtime.run(
+                    spec.name, payload, tools=tools, session_id=session_id
+                )
+                return envelope_to_result(envelope)
+            # 未接入引擎时退化为声明式空实现（插件桩）
+            return await func(self, payload, tools)
+
+        wrapper.__agent_spec__ = spec  # type: ignore[attr-defined]
+        wrapper.__signature__ = inspect.Signature(
+            [
+                inspect.Parameter("self", inspect.Parameter.POSITIONAL_OR_KEYWORD),
+                inspect.Parameter("payload", inspect.Parameter.POSITIONAL_OR_KEYWORD, default=None),
+                inspect.Parameter("tools", inspect.Parameter.POSITIONAL_OR_KEYWORD, default=None),
+                inspect.Parameter("session_id", inspect.Parameter.KEYWORD_ONLY, default=None),
+            ]
+        )
+        wrapper.__annotations__ = dict(getattr(func, "__annotations__", {}))
+        return wrapper
+
+    return decorator
+
+
+class BaseAgentGroup(object):
+    """插件基类：所有 agent 能力通过 @agent_def 注册在本类中。
+
+    执行引擎为 tinyagent-py；模型、鉴权、压缩策略等通过构造参数注入。
+    """
+
+    def __init__(
+        self,
+        model: str = "gpt-5",
+        api_key: Optional[str] = None,
+        model_kwargs: Optional[Dict[str, Any]] = None,
+        temperature: float = 0.0,
+        session_store: Any = None,
+        compact_config: Optional[Dict[str, Any]] = None,
+        global_compact_config: Optional[Dict[str, Any]] = None,
+        max_tool_iterations: int = 8,
+        max_turns: int = 10,
+        enable_todo_write: bool = False,
+        storage: Any = None,
+        logger: Any = None,
+    ) -> None:
+        from .runtime import AgentRuntime
+
+        effective_compact_config = (
+            compact_config if compact_config is not None else global_compact_config
+        )
+        self._global_compact_config = effective_compact_config
+        self._runtime = AgentRuntime(
+            group=self,
+            model=model,
+            api_key=api_key,
+            model_kwargs=model_kwargs,
+            temperature=temperature,
+            session_store=session_store,
+            global_compact_config=effective_compact_config,
+            max_tool_iterations=max_tool_iterations,
+            max_turns=max_turns,
+            enable_todo_write=enable_todo_write,
+            storage=storage,
+            logger=logger,
+        )
+
+    def __init_subclass__(cls, **kwargs: Any) -> None:
+        super().__init_subclass__(**kwargs)
+        specs: Dict[str, AgentSpec] = {}
+        for base in reversed(cls.__mro__):
+            for attr_name, attr in vars(base).items():
+                spec = getattr(attr, "__agent_spec__", None)
+                if isinstance(spec, AgentSpec):
+                    spec.group_cls = cls
+                    specs[attr_name] = spec
+        cls.__agent_specs__ = specs  # type: ignore[attr-defined]
+
+    # ---------- 声明查询 ----------
+    @classmethod
+    def agent_specs(cls) -> Dict[str, AgentSpec]:
+        """本插件注册的全部 agent 声明。"""
+        return dict(getattr(cls, "__agent_specs__", {}))
+
+    @classmethod
+    def agent_names(cls) -> List[str]:
+        """已注册的 agent 名称列表。"""
+        return sorted(cls.agent_specs().keys())
+
+    @classmethod
+    def get_agent_spec(cls, agent_name: str) -> AgentSpec:
+        """按名称取 agent 声明；不存在则抛异常。"""
+        specs = cls.agent_specs()
+        if agent_name not in specs:
+            raise AgentNotRegisteredError(
+                "未注册的 agent: {0}，可用: {1}".format(agent_name, ", ".join(sorted(specs)))
+            )
+        return specs[agent_name]
+
+    # ---------- 执行 ----------
+    @property
+    def runtime(self) -> Any:
+        """运行时对象。"""
+        return self._runtime
+
+    @property
+    def session_store(self) -> Any:
+        """会话存储（内存存储，进程内有效）。"""
+        return self._runtime.session_store
+
+    def session_ids(self) -> List[str]:
+        """当前进程内已知的 session_id 列表。"""
+        return self._runtime.session_store.list_session_ids()
+
+    def set_llm(self, client: Any) -> None:
+        """切换模型 / 鉴权信息。"""
+        if isinstance(client, str):
+            self._runtime.model = client
+        elif isinstance(client, dict):
+            self._runtime.model = client.get("model", self._runtime.model)
+            self._runtime.api_key = client.get("api_key", self._runtime.api_key)
+        else:
+            raise FrameworkError("set_llm 仅支持模型名或配置 dict")
+
+    async def run(
+        self,
+        agent_name: str,
+        payload: Any,
+        tools: Optional[List[Any]] = None,
+        session_id: Optional[str] = None,
+    ) -> Any:
+        """执行指定 agent，返回声明的出参模型实例（扩展信息见代理属性）。"""
+        from .result import envelope_to_result
+
+        envelope = await self._runtime.run(
+            agent_name, payload, tools=tools, session_id=session_id
+        )
+        return envelope_to_result(envelope)
+
+    async def run_envelope(
+        self,
+        agent_name: str,
+        payload: Any,
+        tools: Optional[List[Any]] = None,
+        session_id: Optional[str] = None,
+    ) -> Dict[str, Any]:
+        """执行指定 agent，返回框架原始封装结构 {ok, data, finish_reason, meta}。"""
+        return await self._runtime.run(agent_name, payload, tools=tools, session_id=session_id)

zhkj_llm_function/compat.py

@@ -0,0 +1,63 @@
+"""与 tinyagent / litellm 消息对象的兼容转换。
+
+问题背景：
+- tinyagent 的 `_format_conversation_for_summary()` 用 `tool_call.get("function", {})` 读取工具调用，
+  但运行期 assistant 消息里的 `tool_calls` 是 litellm/OpenAI 对象（属性访问），压缩时会抛
+  `AttributeError: 'X' object has no attribute 'get'`。
+
+因此统一提供：对象 -> dict 的转换。
+"""
+
+from typing import Any, Dict, Iterable, List
+
+
+def _get(source: Any, key: str, default: Any = None) -> Any:
+    """同时兼容 dict 与对象属性两种取值方式。"""
+    if source is None:
+        return default
+    if isinstance(source, dict):
+        return source.get(key, default)
+    return getattr(source, key, default)
+
+
+def tool_call_to_dict(call: Any) -> Dict[str, Any]:
+    """把一个工具调用（dict 或对象）转成可序列化的 dict。"""
+    if isinstance(call, dict):
+        function = call.get("function") or {}
+        return {
+            "id": call.get("id", ""),
+            "type": call.get("type", "function"),
+            "function": {
+                "name": _get(function, "name", ""),
+                "arguments": _get(function, "arguments", "{}"),
+            },
+        }
+
+    function = _get(call, "function")
+    return {
+        "id": _get(call, "id", "") or "",
+        "type": _get(call, "type", "function") or "function",
+        "function": {
+            "name": _get(function, "name", "") or "",
+            "arguments": _get(function, "arguments", "{}") or "{}",
+        },
+    }
+
+
+def message_to_plain_dict(message: Any) -> Dict[str, Any]:
+    """把一条消息转成纯 dict（tool_calls 一并转换）。"""
+    if not isinstance(message, dict):
+        message = {
+            "role": _get(message, "role", "unknown"),
+            "content": _get(message, "content", ""),
+        }
+    plain = dict(message)
+    calls = plain.get("tool_calls")
+    if isinstance(calls, Iterable) and not isinstance(calls, (str, bytes, dict)):
+        plain["tool_calls"] = [tool_call_to_dict(call) for call in calls]
+    return plain
+
+
+def messages_for_summary(messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+    """生成用于摘要的副本：tool_calls 统一为 dict，避免引擎侧属性访问报错。"""
+    return [message_to_plain_dict(message) for message in messages or []]

zhkj_llm_function/errors.py

@@ -0,0 +1,29 @@
+"""框架内部异常类型。"""
+
+
+class FrameworkError(Exception):
+    """所有框架异常的基类。"""
+
+
+class AgentNotRegisteredError(FrameworkError):
+    """请求的 agent 未在 BaseAgentGroup 子类中注册。"""
+
+
+class InputSchemaError(FrameworkError):
+    """入参不符合 agent 声明的 Input 模型。"""
+
+
+class OutputSchemaError(FrameworkError):
+    """模型最终输出不符合 agent 声明的 Output 模型。"""
+
+
+class ToolExecutionError(FrameworkError):
+    """工具执行过程中抛出异常。"""
+
+
+class ToolNotFoundError(FrameworkError):
+    """模型请求了未传入的工具。"""
+
+
+class SessionError(FrameworkError):
+    """会话相关的错误。"""

zhkj_llm_function/protocol.py

@@ -0,0 +1,109 @@
+"""工具调用协议解析。
+
+框架统一使用 <|tool_call|> 标记协议：
+
+    <|tool_call|>
+    {"name": "工具名称", "arguments": {}}
+    <|end_tool_call|>
+
+一轮只支持调用单个工具。
+"""
+
+import json
+import re
+from typing import Any, Dict, Optional
+
+TOOL_CALL_START = "<|tool_call|>"
+TOOL_CALL_END = "<|end_tool_call|>"
+
+_FENCE_RE = re.compile(r"```(?:json)?\s*(.*?)```", re.DOTALL)
+
+
+class ToolCall(object):
+    """一次工具调用请求。"""
+
+    __slots__ = ("name", "arguments", "raw")
+
+    def __init__(self, name: str, arguments: Dict[str, Any], raw: str) -> None:
+        self.name = name
+        self.arguments = arguments
+        self.raw = raw
+
+    def __repr__(self) -> str:  # pragma: no cover - 调试用
+        return "ToolCall(name={0!r}, arguments={1!r})".format(self.name, self.arguments)
+
+
+def strip_code_fence(text: str) -> str:
+    """去掉 markdown 代码围栏。"""
+    match = _FENCE_RE.search(text or "")
+    if match:
+        return match.group(1).strip()
+    return (text or "").strip()
+
+
+def extract_json_payload(text: str) -> Optional[Any]:
+    """从模型输出里抽取 JSON 对象/数组。"""
+    if text is None:
+        return None
+    candidate = strip_code_fence(text)
+    if not candidate:
+        return None
+
+    try:
+        return json.loads(candidate)
+    except Exception:
+        pass
+
+    decoder = json.JSONDecoder()
+    for index, char in enumerate(candidate):
+        if char not in "{[":
+            continue
+        try:
+            value, _ = decoder.raw_decode(candidate[index:])
+        except Exception:
+            continue
+        return value
+    return None
+
+
+def parse_tool_call(text: str) -> Optional[ToolCall]:
+    """解析文本中的工具调用标记；没有则返回 None。"""
+    if not text or TOOL_CALL_START not in text:
+        return None
+
+    segment = text.split(TOOL_CALL_START, 1)[1]
+    if TOOL_CALL_END in segment:
+        segment = segment.split(TOOL_CALL_END, 1)[0]
+    segment = strip_code_fence(segment.strip())
+    if not segment:
+        return None
+
+    payload = extract_json_payload(segment)
+    if not isinstance(payload, dict):
+        return None
+
+    name = payload.get("name")
+    if not isinstance(name, str) or not name:
+        return None
+
+    arguments = payload.get("arguments", {})
+    if arguments is None:
+        arguments = {}
+    if isinstance(arguments, str):
+        try:
+            arguments = json.loads(arguments)
+        except Exception:
+            arguments = {}
+    if not isinstance(arguments, dict):
+        arguments = {}
+
+    return ToolCall(name=name, arguments=arguments, raw=segment)
+
+
+def render_tool_call(name: str, arguments: Dict[str, Any]) -> str:
+    """按框架协议渲染工具调用文本（供 mock / 测试使用）。"""
+    return "{0}\n{1}\n{2}".format(
+        TOOL_CALL_START,
+        json.dumps({"name": name, "arguments": arguments}, ensure_ascii=False),
+        TOOL_CALL_END,
+    )

zhkj_llm_function/result.py

@@ -0,0 +1,82 @@
+"""结果对象：按 agent 声明的出参签名返回，扩展信息通过代理属性暴露。
+
+- 方法返回值就是声明的出参模型实例（`isinstance(result, XXXOutput) == True`），
+  外层可以直接按出参签名接收，`model_dump()` 只包含声明字段。
+- 框架扩展信息（ok / finish_reason / meta）通过动态子类的代理属性提供，
+  不污染出参结构，也不会进入 `model_dump()`。
+"""
+
+from typing import Any, Dict, Optional
+
+from pydantic import PrivateAttr
+
+from .errors import (
+    AgentNotRegisteredError,
+    FrameworkError,
+    InputSchemaError,
+    OutputSchemaError,
+    ToolExecutionError,
+)
+
+_FINISH_ERROR_MAP = {
+    "not_registered": AgentNotRegisteredError,
+    "input_schema_invalid": InputSchemaError,
+    "output_schema_invalid": OutputSchemaError,
+    "tool_error": ToolExecutionError,
+}
+
+
+def build_result(model: Any, meta: Optional[Dict[str, Any]] = None, finish_reason: str = "stop"):
+    """把出参模型包装成带代理属性的结果对象。
+
+    返回对象仍然是出参模型（动态子类）的实例，字段与序列化结果完全不变。
+    """
+    meta = meta or {}
+    base = type(model)
+    meta_cls = type(base)
+
+    namespace = {
+        "_framework_meta": PrivateAttr(default=None),
+        "_framework_finish_reason": PrivateAttr(default=""),
+        "_framework_ok": PrivateAttr(default=True),
+        "ok": property(lambda self: self._framework_ok),
+        "finish_reason": property(lambda self: self._framework_finish_reason),
+        "meta": property(lambda self: self._framework_meta),
+        "error": property(lambda self: (self._framework_meta or {}).get("error", "")),
+    }
+    result_cls = meta_cls(
+        "{0}Result".format(base.__name__),
+        (base,),
+        namespace,
+    )
+
+    obj = model
+    try:
+        obj.__class__ = result_cls
+        object.__setattr__(
+            obj,
+            "__pydantic_private__",
+            {
+                "_framework_meta": meta,
+                "_framework_finish_reason": finish_reason,
+                "_framework_ok": True,
+            },
+        )
+    except Exception:  # pragma: no cover - 极端情况下退化为原模型
+        return model
+    return obj
+
+
+def envelope_to_result(envelope: Dict[str, Any]):
+    """把运行时 envelope 转成出参对象；失败时抛出对应异常。"""
+    if envelope.get("ok") and envelope.get("data") is not None:
+        return build_result(
+            envelope["data"],
+            envelope.get("meta") or {},
+            envelope.get("finish_reason", "stop"),
+        )
+
+    finish_reason = envelope.get("finish_reason", "error")
+    message = envelope.get("error") or "agent 执行失败: {0}".format(finish_reason)
+    error_cls = _FINISH_ERROR_MAP.get(finish_reason, FrameworkError)
+    raise error_cls(message)