从零手写AI Agent：LangChain工具调用与记忆循环实现细节

用户输入 → Prompt组装 → LLM推理 → 解析输出
                ↑                      ↓
            记忆管理 ← 循环控制 ← 工具执行
                                      ↓
                                  错误处理

from langchain.tools import BaseTool
from pydantic import BaseModel, Field
import sqlite3

class SalesQueryInput(BaseModel):
    region: str = Field(description="区域名称，如'华东'")
    month: str = Field(description="月份，格式YYYY-MM")

class SalesQueryTool(BaseTool):
    name = "query_sales_db"
    description = "查询指定区域和月份的销售额数据，返回JSON字符串"
    args_schema = SalesQueryInput

    def _run(self, region: str, month: str) -> str:
        try:
            conn = sqlite3.connect("sales.db")
            cursor = conn.cursor()
            cursor.execute(
                "SELECT product, amount FROM sales WHERE region=? AND month=? ORDER BY amount DESC LIMIT 10",
                (region, month)
            )
            rows = cursor.fetchall()
            conn.close()
            if not rows:
                return f"ERROR: 未找到{region}在{month}的销售数据"
            return str([{"product": r[0], "amount": r[1]} for r in rows])
        except Exception as e:
            return f"ERROR: 数据库查询失败 - {str(e)}"

    async def _arun(self, *args, **kwargs):
        raise NotImplementedError("不支持异步")

from langchain.memory import ConversationSummaryBufferMemory
from langchain_openai import ChatOpenAI

llm = ChatOpenAI(model="gpt-4-0125-preview", temperature=0, max_tokens=1024)

memory = ConversationSummaryBufferMemory(
    llm=llm,
    max_token_limit=500,  # 超过500 token触发摘要
    return_messages=True,
    memory_key="chat_history"
)

import re
from langchain.prompts import PromptTemplate

REACT_PROMPT = """你是一个数据分析Agent。可用工具：
{tools}

使用格式：
Question: 用户问题
Thought: 思考下一步
Action: 工具名，必须是[{tool_names}]之一
Action Input: 工具参数
Observation: 工具返回结果
...（重复Thought/Action/Observation）
Thought: 我现在知道最终答案了
Final Answer: 最终回答

历史对话：{chat_history}
Question: {input}
{agent_scratchpad}"""

def run_agent(user_input: str, max_iterations: int = 6) -> str:
    tools = [SalesQueryTool()]
    tool_names = ", ".join([t.name for t in tools])
    tool_desc = "\n".join([f"{t.name}: {t.description}" for t in tools])

    scratchpad = ""
    for i in range(max_iterations):
        prompt = REACT_PROMPT.format(
            tools=tool_desc,
            tool_names=tool_names,
            chat_history=memory.load_memory_variables({})["chat_history"],
            input=user_input,
            agent_scratchpad=scratchpad
        )
        response = llm.invoke(prompt).content

        if "Final Answer:" in response:
            final = response.split("Final Answer:")[-1].strip()
            memory.save_context({"input": user_input}, {"output": final})
            return final

        # 解析Action和Action Input
        action_match = re.search(r"Action:\s*(\w+)", response)
        input_match = re.search(r"Action Input:\s*(.+)", response)

        if not action_match or not input_match:
            scratchpad += f"\n{response}\nObservation: 解析失败，请按格式输出Action和Action Input"
            continue

        tool_name = action_match.group(1)
        tool_input = input_match.group(1).strip()

        # 执行工具，错误处理内嵌
        tool = next((t for t in tools if t.name == tool_name), None)
        if not tool:
            observation = f"ERROR: 工具'{tool_name}'不存在，可用工具：{tool_names}"
        else:
            try:
                # 这里简化参数解析，实际用args_schema
                observation = tool._run(**eval(f"dict({tool_input})"))
            except Exception as e:
                observation = f"ERROR: 工具执行异常 - {str(e)}"

        scratchpad += f"\n{response}\nObservation: {observation}\n"

    # 超过最大迭代，强制返回
    return "达到最大迭代次数，当前结果：" + scratchpad[-200:]