基于 openai sdk 的简单 AI Agent

主要特性:

• 支持 OpenAI, DeepSeek 等公司的大语言模型
• 内置多个工具函数(获取时间、网页搜索、执行 Python 代码等)
• 支持流式输出
• 可扩展的工具系统
• 支持 MCP 服务

安装:

pip install ez-agent

使用方法

导入并创建 Agent:

from ez_agent import Agent

agent = Agent(
    model="gpt-4o",
    api_key="你的API密钥",
    base_url="https://api.openai.com/v1",
)

支持高度自定义

agent = Agent(
    model="gpt-4o",
    api_key="你的API密钥",
    base_url="https://api.openai.com/v1",
    tools=[python_script_tool, get_time_tool, open_website_tool, browse_web_tool],
    instructions="你是一个助手，可以使用自定义工具。"
    temperature=0.9,
    max_tokens=1024,
    top_p=1,
    frequency_penalty=0,
    presence_penalty=0,
)

添加响应处理

agent.add_response_handler(lambda response: print(response.content))
agent.add_tool_call_handler(lambda tool_call: print(f"工具调用: {tool_call.function.name}"))

发送请求:

response = agent.run("解释量子力学的基本原理")

流式输出

agent.add_stream_chunk_handler(lambda chunk: print(chunk))
agent.run("解释量子力学的基本原理", stream=True)

完整示例

# 通用助手
from ez_agent import Agent

# 初始化通用助手
agent = Agent(
    model="gpt-4o",
    api_key="sk-your-api-key",
    base_url="https://api.openai.com/v1",
)

agent.add_stream_chunk_handler(lambda chunk: print(chunk, end="", flush=True))
agent.add_tool_call_handler(lambda tool_call: print(f"工具调用: {tool_call["function"]["name"]}"))

while True:
    user_input = input(">>> ")
    agent.run(user_input, stream=True)
    print()

预设的工具:

• get_time_tool 获取当前时间
• python_script_tool 执行 Python 脚本
• open_website_tool 调用浏览器打开网站
• browse_web_tool 浏览网页并获取内容

使用预设工具:

from ez_agent.prefab import python_script_tool

自定义函数工具:

使用 @FunctionTool 装饰器创建自定义工具，会自动根据类型注解和函数注释生成工具。该装饰器不影响原函数，原函数仍可被直接调用。

from ez_agent import FunctionTool, Agent

# 创建自定义工具
@FunctionTool
def my_tool(arg1: str, arg2: int, arg3: bool = False, ...) -> str:
    """
    工具描述
    """
    ...

# 使用自定义工具创建Agent
custom_agent = Agent(
    model="gpt-4o",
    api_key="你的API密钥",
    base_url="https://api.openai.com/v1",
    tools=[my_tool],
    instructions="你是一个助手，可以使用自定义工具。"
)

这样，模型便可以调用你创建的自定义工具。

MCP 工具

该框架允许用户用简单轻松的方式连接MCP服务

创建配置文件

{
    "mcpServers": {
        "<name>": { // SSE 模式
            "url": "<url>",
            "env": {
                ...
            }
        },
        "<name>": { // stdio 模式
            "command": "<command>",
            "args": ["<arg1>", "<arg2>", ...],
            "env": {
                ...
            }
        }
    }
}

将上述文件保存为mcp_config.json

使用 MCP 工具

from ez_agent import AsyncAgent
agent = AsyncAgent(...) # 创建异步Agent，方法与同步Agent相同
agent.load_mcp_config("mcp_config.json")

注意：MCP工具需手动释放资源

await agent.cleanup() # 须在异步环境中使用

设置消息过期时间

可在创建 Agent 时设置

agent = Agent(
    model="gpt-4o",
    api_key="你的API密钥",
    base_url="https://api.openai.com/v1",
    message_expiration_time=60 * 60, # 单位为秒
)

这样，1 小时前产生的消息将会被自动清理

也可手动清理

agent.clear_msg_by_time(60 * 60) # 单位为秒

并行支持

如果有多个请求同时到达，但又想让它们共享同一个Agent（例如想共享 messages），可以使用Agent.safe_modify()来确保线程安全：

async with agent.safe_modify():
    # 在此处填写代码逻辑，例如 agent.run()，不会影响原对象
    ...

注意，safe_modify()的作用域结束后会自动合并新消息，如果要完全独立地修改Agent，请使用Agent.safe_modify(merge_messages=False)

async with agent.safe_modify(merge_messages=False):
    ...

附录

Agent.run 流程图

graph TD
    A[Agent.run 开始] --> B[准备用户消息]
    B --> C[调用 LLM 获取响应]

    C --> D{检测工具调用}
    D -->|没有| E[返回最终回答]
    D -->|有| F[解析工具调用请求]

    F --> G[执行工具调用]
    G --> H[获取工具执行结果]

    H --> I[将工具结果添加到消息历史]
    I --> J[再次调用 LLM 获取响应]

    J --> K{检测新工具调用}
    K -->|有| F
    K -->|没有| E

    E --> L[Agent.run 结束]

Loading

工具继承关系图

graph TD
Tool[Tool - ABC] --> BaseFunctionTool[BaseFunctionTool]
Tool --> MCPTool[MCPTool]

    BaseFunctionTool --> FunctionTool[FunctionTool]
    BaseFunctionTool --> AsyncFunctionTool[AsyncFunctionTool]

    FunctionTool --> FoldableFunctionTool[FoldableFunctionTool]
    AsyncFunctionTool --> FoldableAsyncFunctionTool[FoldableAsyncFunctionTool]
    MCPTool --> FoldableMCPTool[FoldableMCPTool]

Loading