使用 Python 从零构建一个轻量级 AI Agent:支持工具调用和思维链
背景介绍
过去一年里,AI Agent 成为大语言模型应用领域最热门的话题之一。从 AutoGPT 到 Claude Agent,从单智能体到多智能体协作,各种框架层出不穷。但对于大多数开发者来说,这些框架要么过于复杂,要么依赖特定的云服务,学习成本不低。
我最近在项目中需要构建一个简单的 AI Agent,用来处理用户的技术咨询。它需要能够理解用户意图、调用内部 API 查询信息、进行多轮对话、生成结构化的回复。现有的框架要么功能太强大导致难以定制,要么太简单不够用。干脆自己动手,从零开始。
问题描述
构建一个最小可用的 AI Agent,需要解决这几个核心问题:
LLM 的调用方式。 需要一种标准化的方式来调用大语言模型,支持流式输出、函数调用、上下文管理。
工具系统的设计。 Agent 需要能够使用工具来扩展自己的能力。如何定义工具、注册工具、调用工具、解析结果,这套机制要设计好。
思维链的实现。 让 Agent 能够分步思考,在每个步骤中决定下一步做什么。
对话上下文的维护。 多轮对话需要维护历史消息,同时要控制上下文长度,避免超出模型的 token 限制。
详细步骤
环境准备
首先安装必要的依赖。我选择使用 OpenAI 的 API 作为 LLM 后端,因为它提供了完善的函数调用功能。
bash
pip install openai python-dotenv
创建项目的目录结构:
agent_project/
├── main.py
├── agent/
│ ├── __init__.py
│ ├── llm.py
│ ├── tools.py
│ ├── memory.py
│ └── agent.py
└── .env
LLM 客户端封装
在 agent/llm.py 中封装一个 LLM 客户端,提供统一的调用接口:
“`python
import os
from openai import OpenAI
from typing import Generator, Optional, List, Dict, Any
class LLMClient:
“””统一的大语言模型客户端”””
def __init__(self, model: str = "gpt-4o"):
api_key = os.getenv("OPENAI_API_KEY")
if not api_key:
raise ValueError("请在 .env 文件中设置 OPENAI_API_KEY")
self.client = OpenAI(api_key=api_key)
self.model = model
def chat(
self,
messages: List[Dict[str, str]],
tools: Optional[List[Dict]] = None,
temperature: float = 0.7
) -> Dict[str, Any]:
"""发送聊天请求"""
params = {
"model": self.model,
"messages": messages,
"temperature": temperature
}
if tools:
params["tools"] = tools
response = self.client.chat.completions.create(**params)
return response.choices[0].message
def chat_stream(
self,
messages: List[Dict[str, str]],
tools: Optional[List[Dict]] = None,
temperature: float = 0.7
) -> Generator[str, None, None]:
"""流式输出"""
params = {
"model": self.model,
"messages": messages,
"temperature": temperature,
"stream": True
}
if tools:
params["tools"] = tools
response = self.client.chat.completions.create(**params)
for chunk in response:
if chunk.choices[0].delta.content:
yield chunk.choices[0].delta.content
“`
工具系统设计
在 agent/tools.py 中实现工具注册和调用的核心逻辑:
“`python
from typing import Callable, Any, Dict, List
from functools import wraps
class Tool:
“””工具定义类”””
def __init__(self, name: str, description: str, func: Callable):
self.name = name
self.description = description
self.func = func
self.parameters = self._extract_parameters(func)
def _extract_parameters(self, func: Callable) -> Dict:
"""从函数签名中提取参数定义"""
import inspect
sig = inspect.signature(func)
params = {
"type": "object",
"properties": {},
"required": []
}
for param_name, param in sig.parameters.items():
if param_name == "self":
continue
param_type = "string"
if param.annotation == int:
param_type = "integer"
elif param.annotation == float:
param_type = "number"
elif param.annotation == bool:
param_type = "boolean"
params["properties"][param_name] = {
"type": param_type,
"description": f"参数 {param_name}"
}
if param.default == inspect.Parameter.empty:
params["required"].append(param_name)
return params
def to_openai_format(self) -> Dict:
"""转换为 OpenAI 函数调用格式"""
return {
"type": "function",
"function": {
"name": self.name,
"description": self.description,
"parameters": self.parameters
}
}
def execute(self, arguments: Dict[str, Any]) -> Any:
"""执行工具函数"""
return self.func(**arguments)
class ToolRegistry:
“””工具注册中心”””
def __init__(self):
self.tools: Dict[str, Tool] = {}
def register(self, name: str, description: str) -> Callable:
"""装饰器:注册工具"""
def decorator(func: Callable) -> Callable:
tool = Tool(name, description, func)
self.tools[name] = tool
return func
return decorator
def get_tool(self, name: str) -> Tool:
"""获取工具"""
if name not in self.tools:
raise ValueError(f"工具 {name} 不存在")
return self.tools[name]
def list_tools(self) -> List[Dict]:
"""列出所有工具"""
return [tool.to_openai_format() for tool in self.tools.values()]
def execute_tool(self, name: str, arguments: Dict[str, Any]) -> Any:
"""执行工具"""
tool = self.get_tool(name)
return tool.execute(arguments)
全局工具注册表
registry = ToolRegistry()
“`
定义一些具体工具:
“`python
定义示例工具
@registry.register(name=”search_knowledge”, description=”搜索知识库中的相关信息”)
def search_knowledge(query: str) -> str:
“””搜索知识库”””
# 实际项目中,这里会连接向量数据库
knowledge_db = {
“python”: “Python 是一种高级编程语言,由 Guido van Rossum 于 1991 年创建。”,
“openai”: “OpenAI 是一家 AI 研究公司,开发了 GPT 系列模型。”,
“agent”: “AI Agent 是指能够自主执行任务的人工智能系统。”
}
for key, value in knowledge_db.items():
if key in query.lower():
return value
return "未找到相关信息"
@registry.register(name=”calculate”, description=”执行数学计算”)
def calculate(expression: str) -> str:
“””数学计算”””
try:
result = eval(expression)
return str(result)
except Exception as e:
return f”计算错误: {str(e)}”
@registry.register(name=”get_time”, description=”获取当前时间”)
def get_time() -> str:
“””获取当前时间”””
from datetime import datetime
return datetime.now().strftime(“%Y-%m-%d %H:%M:%S”)
“`
对话上下文管理
在 agent/memory.py 中实现对话上下文的管理:
“`python
from typing import List, Dict
from collections import deque
class ConversationMemory:
“””对话记忆管理”””
def __init__(self, max_tokens: int = 8000):
self.max_tokens = max_tokens
self.messages: deque = deque(maxlen=100)
def add_message(self, role: str, content: str):
"""添加消息"""
self.messages.append({
"role": role,
"content": content
})
def get_messages(self) -> List[Dict[str, str]]:
"""获取所有消息"""
return list(self.messages)
def clear(self):
"""清空记忆"""
self.messages.clear()
def summarize_if_needed(self, llm_client) -> str:
"""如果消息过长,进行摘要"""
if len(self.messages) > 20:
system_msg = None
if self.messages[0]["role"] == "system":
system_msg = self.messages[0]
self.messages.popleft()
recent = list(self.messages)[-10:]
self.messages.clear()
if system_msg:
self.messages.append(system_msg)
self.messages.extend(recent)
return "历史消息已压缩"
return ""
“`
Agent 核心实现
在 agent/agent.py 中实现 Agent 的核心逻辑:
“`python
from typing import Optional, Callable
from .llm import LLMClient
from .tools import registry, ToolRegistry
from .memory import ConversationMemory
class Agent:
“””AI Agent 主类”””
def __init__(
self,
system_prompt: str,
model: str = "gpt-4o",
tools: Optional[ToolRegistry] = None
):
self.llm = LLMClient(model=model)
self.memory = ConversationMemory()
self.tools = tools or registry
self.system_prompt = system_prompt
self.memory.add_message("system", system_prompt)
self.max_iterations = 10
self.verbose = True
def think(self, user_input: str) -> str:
"""核心思考循环:实现思维链"""
self.memory.add_message("user", user_input)
iteration = 0
final_response = None
while iteration < self.max_iterations:
iteration += 1
messages = self.memory.get_messages()
response = self.llm.chat(
messages=messages,
tools=self.tools.list_tools() if self.tools.tools else None
)
if response.tool_calls:
for tool_call in response.tool_calls:
tool_name = tool_call.function.name
tool_args = tool_call.function.arguments
if self.verbose:
print(f"[思考 {iteration}] 调用工具: {tool_name}")
print(f"[参数] {tool_args}")
try:
result = self.tools.execute_tool(tool_name, tool_args)
if self.verbose:
print(f"[结果] {result}")
self.memory.add_message(
"assistant",
f"调用了工具 {tool_name},结果: {result}"
)
self.memory.add_message(
"tool",
str(result)
)
except Exception as e:
error_msg = f"工具执行错误: {str(e)}"
if self.verbose:
print(f"[错误] {error_msg}")
self.memory.add_message("tool", error_msg)
continue
else:
final_response = response.content
self.memory.add_message("assistant", final_response)
break
if final_response is None:
final_response = "抱歉,我需要更多时间来处理这个请求。"
return final_response
def run(self, user_input: str) -> str:
"""运行 Agent"""
return self.think(user_input)
def reset(self):
"""重置 Agent 状态"""
self.memory.clear()
self.memory.add_message("system", self.system_prompt)
“`
主程序入口
在 main.py 中创建和运行 Agent:
“`python
import os
from dotenv import load_dotenv
from agent.agent import Agent
from agent.tools import registry
load_dotenv()
SYSTEM_PROMPT = “””你是一个智能技术助手。
你可以使用以下工具来帮助用户:
– search_knowledge: 搜索知识库
– calculate: 进行数学计算
– get_time: 获取当前时间
请根据用户的问题,选择合适的工具来回答。”””
def main():
agent = Agent(
system_prompt=SYSTEM_PROMPT,
model=”gpt-4o”,
tools=registry
)
print("=" * 50)
print("智能助手已启动!输入 quit 退出")
print("=" * 50)
while True:
user_input = input("\n你: ").strip()
if user_input.lower() in ["quit", "exit", "q"]:
print("再见!")
break
if not user_input:
continue
response = agent.run(user_input)
print(f"\n助手: {response}")
if name == “main“:
main()
“`
配置文件
创建 .env 文件:
OPENAI_API_KEY=your-api-key-here
运行结果
交互示例
运行程序后:
“`
智能助手已启动!输入 quit 退出
你: 什么是 Python?
助手: Python 是一种高级编程语言,由 Guido van Rossum 于 1991 年创建。它以简洁的语法和强大的功能著称,广泛应用于 Web 开发、数据分析、人工智能等领域。
你: 现在几点了?
助手: 当前时间是 2026-04-21 09:00:00。
你: 计算 123 * 456 = ?
助手: 123 * 456 的结果是 56088。
你: 我想学习 AI Agent,请推荐一些知识
助手: 根据知识库,AI Agent 是指能够自主执行任务的人工智能系统。它可以感知环境、进行推理、做出决策并执行动作。目前业界主流的 AI Agent 框架包括 LangChain、AutoGen、Claude Agent 等。
“`
调试模式输出
开启 verbose 模式可以看到 Agent 的思考过程:
“`
你: 帮我计算 100 + 200 并搜索一下 OpenAI 的信息
[思考 1] 调用工具: calculate
[参数] {“expression”: “100 + 200”}
[结果] 300
[思考 2] 调用工具: search_knowledge
[参数] {“query”: “OpenAI”}
[结果] OpenAI 是一家 AI 研究公司,开发了 GPT 系列模型。
助手: 计算结果是 300。OpenAI 是一家 AI 研究公司,由 Elon Musk 和 Sam Altman 等人于 2015 年创立,总部位于美国旧金山。他们开发了 GPT 系列语言模型,包括 GPT-3、GPT-4 等,在 AI 领域产生了深远影响。
“`
总结
通过以上步骤,我们从零构建了一个轻量级的 AI Agent 框架,具备以下核心能力:
统一的 LLM 接口 – 封装 OpenAI API,轻松切换不同模型,支持流式输出。
灵活的工具系统 – 使用装饰器模式,添加新工具只需一个装饰器。
思维链实现 – Agent 在循环中分步思考,根据上一步结果决定下一步行动。
上下文管理 – ConversationMemory 自动管理对话历史,避免上下文溢出。
可扩展性 – 框架简洁,方便添加更多功能,比如多 Agent 协作、长期记忆、反思机制等。
代码量只有几百行,已经具备 Agent 的核心要素。继续扩展可以考虑:添加向量数据库实现 RAG、加入 ReAct 推理模式、实现反思机制、或者添加多 Agent 协作能力。
AI Agent 领域发展很快,更多有趣的玩法等待你去探索。