python-library-ai-agent 0.1.0__py3-none-any.whl

ai_agent/loop.py

@@ -0,0 +1,124 @@
+from __future__ import annotations
+
+from typing import AsyncIterator
+
+from collections.abc import Sequence
+
+from ai_agent.context import AgentContext, RunContext, RunStatus, ToolInvocation
+from ai_agent.listener import (
+    AgentListener,
+    notify_output_delta,
+    notify_run_end,
+    notify_run_start,
+    notify_thinking_delta,
+    notify_tool_end,
+    notify_tool_start,
+)
+from ai_agent.llm import StreamChunk, StreamKind
+from ai_agent.react_tool_turn import deferred_tool_reply, should_defer_tool_in_batch
+
+
+class ReactLoop:
+    """多步 ReAct 运行；更新 RunContext 中的思考、工具实例与最终 output。"""
+
+    def __init__(self, context: AgentContext) -> None:
+        self._context = context
+
+    async def run(self, run: RunContext) -> AsyncIterator[RunContext]:
+        """
+        驱动单轮上下文直至完成、失败或超步数；每有更新 yield 同一对象。
+        """
+        run.status = RunStatus.RUNNING
+        run.thinking = ""
+        run.output = ""
+        await notify_run_start(self._context.listeners, run)
+        yield run
+
+        try:
+            for _ in range(self._context.max_steps):
+                run.thinking = ""
+                pending_calls: list[ToolInvocation] = []
+                api_tools = self._context.tools.api_tools() or None
+
+                async for chunk in self._context.llm.stream(run, tools=api_tools):
+                    self._apply_chunk(run, chunk, pending_calls)
+                    await self._notify_chunk(self._context.listeners, chunk, run)
+                    yield run
+                    if chunk.kind == StreamKind.DONE:
+                        break
+
+                if pending_calls:
+                    turn: list[ToolInvocation] = []
+                    for inv in pending_calls:
+                        run.tool_invocations.append(inv)
+                        turn.append(inv)
+                        yield run
+                        await notify_tool_start(self._context.listeners, inv, run)
+                        if should_defer_tool_in_batch(inv, pending_calls):
+                            inv.answer = deferred_tool_reply()
+                            inv.ok = True
+                        else:
+                            await self._context.tools.execute(inv)
+                        await notify_tool_end(self._context.listeners, inv, run)
+                        yield run
+                    run.tool_turns.append(turn)
+                    run.output = ""
+                    continue
+                if run.output:
+                    run.status = RunStatus.COMPLETED
+                    yield run
+                    return
+
+                if run.thinking and not run.output:
+                    run.output = run.thinking
+                    run.status = RunStatus.COMPLETED
+                    yield run
+                    return
+
+            run.status = RunStatus.MAX_STEPS
+            yield run
+
+        except Exception:  # noqa: BLE001 — 运行边界
+            run.status = RunStatus.FAILED
+            yield run
+            raise
+        finally:
+            await notify_run_end(self._context.listeners, run)
+
+    async def _notify_chunk(
+        self,
+        listeners: Sequence[AgentListener],
+        chunk: StreamChunk,
+        run: RunContext,
+    ) -> None:
+        if not chunk.delta:
+            return
+        if chunk.kind == StreamKind.REASONING:
+            await notify_thinking_delta(listeners, chunk.delta, run)
+        elif chunk.kind == StreamKind.TEXT:
+            await notify_output_delta(listeners, chunk.delta, run)
+
+    def _apply_chunk(
+        self,
+        run: RunContext,
+        chunk: StreamChunk,
+        pending_calls: list[ToolInvocation],
+    ) -> None:
+        if chunk.kind == StreamKind.REASONING:
+            if chunk.delta:
+                run.thinking += chunk.delta
+                if pending_calls:
+                    pending_calls[-1].thinking += chunk.delta
+        elif chunk.kind == StreamKind.TEXT:
+            if chunk.delta:
+                run.output += chunk.delta
+        elif chunk.kind == StreamKind.TOOL_CALL:
+            if chunk.tool_call_id and chunk.tool_name is not None:
+                inv = ToolInvocation(
+                    call_id=chunk.tool_call_id,
+                    tool_name=chunk.tool_name,
+                    arguments=dict(chunk.tool_arguments or {}),
+                    thinking=run.thinking,
+                )
+                pending_calls.append(inv)
+                run.thinking = ""

ai_agent/mcp_config.py

@@ -0,0 +1,54 @@
+from __future__ import annotations
+
+from collections.abc import Mapping
+from typing import Any
+
+from pydantic import BaseModel, ConfigDict, Field, field_validator
+
+
+class McpStdioServerConfig(BaseModel):
+    """单个 MCP stdio 子进程启动项。"""
+
+    model_config = ConfigDict(extra="forbid")
+
+    command: str = Field(description="启动 MCP server 的可执行文件路径或命令名")
+    args: list[str] = Field(default_factory=list, description="传给 command 的参数列表")
+    env: dict[str, str] | None = Field(
+        default=None,
+        description="传入子进程的环境变量；与 MCP SDK 默认项合并，不继承宿主进程 os.environ",
+    )
+    cwd: str | None = Field(default=None, description="子进程工作目录；省略则使用当前工作目录")
+
+    @field_validator("command")
+    @classmethod
+    def _command_non_empty(cls, value: str) -> str:
+        stripped = value.strip()
+        if not stripped:
+            raise ValueError("command 不能为空")
+        return stripped
+
+
+class McpConfig(BaseModel):
+    """MCP 工具加载配置，与常见 ``mcpServers`` JSON 结构一致。"""
+
+    model_config = ConfigDict(extra="forbid", populate_by_name=True)
+
+    mcp_servers: dict[str, McpStdioServerConfig] = Field(
+        alias="mcpServers",
+        description="按名称索引的 MCP server 启动配置",
+    )
+
+
+def parse_mcp_config(data: McpConfig | Mapping[str, Any]) -> McpConfig:
+    """
+    将映射或已有模型规范为 ``McpConfig``。
+
+    Args:
+        data: 含 ``mcpServers`` 键的映射，或已校验的 ``McpConfig``
+
+    Returns:
+        校验后的配置模型
+    """
+    if isinstance(data, McpConfig):
+        return data
+    return McpConfig.model_validate(dict(data))

ai_agent/mcp_loader.py

@@ -0,0 +1,110 @@
+from __future__ import annotations
+
+from collections.abc import Mapping
+from contextlib import AsyncExitStack
+from typing import Any
+
+from mcp import ClientSession, StdioServerParameters
+from mcp.client.stdio import get_default_environment, stdio_client
+
+from ai_agent.mcp_config import McpConfig, McpStdioServerConfig, parse_mcp_config
+from ai_agent.tools import Tool
+
+
+def _stdio_env(server_env: dict[str, str] | None) -> dict[str, str]:
+    """
+    组装 MCP 子进程环境变量。
+
+    仅合并 MCP SDK 默认项与 ``mcp.json`` 中该 server 的 ``env``，
+    不继承当前进程的 ``os.environ``，避免与示例 ``.env`` 中的 LLM 等变量串扰。
+    """
+    merged = dict(get_default_environment())
+    if server_env:
+        merged.update(server_env)
+    return merged
+
+
+class MCPToolLoader:
+    """
+    按已校验的 MCP 配置启动 stdio 服务并保持会话，将远程工具包装为库内 ``Tool``。
+
+    会话在 ``close()`` 之前须保持存活，以便后续 ReAct 循环通过同一连接调用工具。
+    配置文件读取与多格式解析由调用方负责，本类只接收 ``McpConfig`` 或等价映射。
+    """
+
+    def __init__(self) -> None:
+        self._stack = AsyncExitStack()
+
+    async def load(self, config: McpConfig | Mapping[str, Any]) -> list[Tool]:
+        """
+        按配置启动全部 MCP server 并收集工具。
+
+        Args:
+            config: ``McpConfig`` 或含 ``mcpServers`` 的映射（经 Pydantic 校验）
+
+        Returns:
+            各 server 暴露的工具列表
+        """
+        parsed = parse_mcp_config(config)
+        tools: list[Tool] = []
+        for server_name, server_config in parsed.mcp_servers.items():
+            server_tools = await self._load_server(server_name, server_config)
+            tools.extend(server_tools)
+        return tools
+
+    async def _load_server(
+        self,
+        server_name: str,
+        server_config: McpStdioServerConfig,
+    ) -> list[Tool]:
+        params = StdioServerParameters(
+            command=server_config.command,
+            args=server_config.args,
+            env=_stdio_env(server_config.env),
+            cwd=server_config.cwd,
+        )
+
+        read, write = await self._stack.enter_async_context(stdio_client(params))
+        session = await self._stack.enter_async_context(ClientSession(read, write))
+        await session.initialize()
+
+        result = await session.list_tools()
+        tools: list[Tool] = []
+
+        for mcp_tool in result.tools:
+            ai_tool_name = f"{server_name}__{mcp_tool.name}"
+            handler = _make_tool_handler(session, mcp_tool.name)
+            tools.append(
+                Tool(
+                    name=ai_tool_name,
+                    description=mcp_tool.description or "",
+                    parameters=mcp_tool.inputSchema,
+                    handler=handler,
+                )
+            )
+
+        return tools
+
+    async def close(self) -> None:
+        """关闭已启动的 MCP 会话与 stdio 子进程。"""
+        await self._stack.aclose()
+
+    @staticmethod
+    def _result_to_text(result: Any) -> str:
+        parts: list[str] = []
+        for item in getattr(result, "content", []):
+            text = getattr(item, "text", None)
+            if text is not None:
+                parts.append(text)
+        return "\n".join(parts)
+
+
+def _make_tool_handler(
+    session: ClientSession,
+    mcp_tool_name: str,
+) -> Any:
+    async def handler(**arguments: Any) -> str:
+        call_result = await session.call_tool(mcp_tool_name, arguments)
+        return MCPToolLoader._result_to_text(call_result)
+
+    return handler

ai_agent/memory/__init__.py

@@ -0,0 +1,9 @@
+from ai_agent.memory.config import MemoryConfig
+from ai_agent.memory.context_builder import BuiltMemoryContext
+from ai_agent.memory.memory_system import MemorySystem
+
+__all__ = [
+    "BuiltMemoryContext",
+    "MemoryConfig",
+    "MemorySystem",
+]

ai_agent/memory/compression_work.py

@@ -0,0 +1,71 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Union
+
+from ai_agent.memory.models import (
+    DateMemoryEntry,
+    DateMemoryDay,
+    ImportantMemoryEntry,
+    LongTermChunk,
+    MemoryMessage,
+)
+
+
+@dataclass(frozen=True)
+class ShortToDateWork:
+    batch: tuple[MemoryMessage, ...]
+    batch_size: int
+
+
+@dataclass(frozen=True)
+class ShortToDateResult:
+    entries: tuple[DateMemoryEntry, ...]
+    important_texts: tuple[str, ...]
+    day_label: str
+
+
+@dataclass(frozen=True)
+class DateToLongWork:
+    day_label: str
+    entries: tuple[DateMemoryEntry, ...]
+
+
+@dataclass(frozen=True)
+class DateToLongResult:
+    chunk: LongTermChunk
+
+
+@dataclass(frozen=True)
+class CompressLongWork:
+    chunks: tuple[LongTermChunk, ...]
+
+
+@dataclass(frozen=True)
+class CompressLongResult:
+    chunks: tuple[LongTermChunk, ...]
+
+
+@dataclass(frozen=True)
+class CompressImportantWork:
+    entries: tuple[ImportantMemoryEntry, ...]
+
+
+@dataclass(frozen=True)
+class CompressImportantResult:
+    entries: tuple[ImportantMemoryEntry, ...]
+
+
+CompressionWork = Union[
+    ShortToDateWork,
+    DateToLongWork,
+    CompressLongWork,
+    CompressImportantWork,
+]
+
+CompressionResult = Union[
+    ShortToDateResult,
+    DateToLongResult,
+    CompressLongResult,
+    CompressImportantResult,
+]