AI Agent(学习四-加短期Memory)

前面几篇，分别从基础搭建到自动选择Tool，虽然进入了AI Agent的门，但是此时的ai还属于基础款，接下来，我们需要给它加记忆功能，加Memory这步，是AI Agent 从“一次性对话” → “像人一样有上下文”的关键跃迁。

加Memory后，可以做到以下的功能：

• 记住之前说过的话
• Memory不参与算数Tool（避免污染）
• 结构清晰、企业可用
• 不引入LangChain复杂Memory体系（先稳）

一、短期记忆

短期记忆：

• 只在当前程序运行期间有效
• 用于：
  • 上下文理解
  • 连续对话
  • 指代（“刚才那个”）
• Memory只给模型决策和回答用，Tool执行只用当前输入

二、实现

在项目中创建文件：main3.py，具体实现如下

import json
from langchain_groq import ChatGroq
from dotenv import load_dotenv
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS


load_dotenv()

# ========== 1️⃣ 初始化 LLM ==========
llm = ChatGroq(
    model="llama-3.1-8b-instant",
    temperature=0
)


# ========== 2️⃣ 定义 Tools ==========
def calculator_tool(expression: str) -> str:
    try:
        result = eval(expression)
        return f"计算结果：{result}"
    except Exception as e:
        return f"计算错误：{str(e)}"

def echo_tool(text: str) -> str:
    return f"Echo Tool 收到：{text}"



TOOLS = {
    "calculator": calculator_tool,
    "echo": echo_tool
}

# ========== 3️⃣ Memory（核心新增） ==========
conversation_memory = []  # List[{"role": "...", "content": "..."}]

MAX_MEMORY_LENGTH = 6  # 防止 prompt 无限变长


def add_to_memory(role: str, content: str):
    conversation_memory.append({"role": role, "content": content})
    # 只保留最近 N 条
    if len(conversation_memory) > MAX_MEMORY_LENGTH:
        del conversation_memory[0]


def render_memory() -> str:
    """
    把 Memory 渲染成 prompt 文本
    """
    if not conversation_memory:
        return ""

    memory_text = "以下是之前的对话内容（供你理解上下文）：\n"
    for msg in conversation_memory:
        memory_text += f"{msg['role']}：{msg['content']}\n"
    return memory_text


# ========== 4️⃣ Tool Router Prompt ==========
ROUTER_PROMPT = """
你是一个 AI Agent 的调度器。
你只能从以下工具中选择一个：

1. calculator：用于数学计算
2. echo：原样返回用户输入

请根据【上下文 + 用户输入】，判断是否需要使用工具。
如果需要，严格返回 JSON（不要解释）：

{
  "tool": "<工具名>",
  "input": "<传给工具的参数>"
}

如果不需要任何工具，返回：

{
  "tool": "none",
  "input": "<直接回复用户的话>"
}

"""

# ========== 5️⃣ Agent 核心 ==========
def agent_run(user_input: str) -> str:
    memory_text = render_memory()

    prompt = memory_text + ROUTER_PROMPT + "用户输入：\n" + user_input
    decision_msg = llm.invoke(prompt)
    decision_text = decision_msg.content
    #decision_text = llm.invoke(prompt)

    try:
        decision = json.loads(decision_text)
    except json.JSONDecodeError:
        # 模型异常，兜底直接回答
        reply = llm.invoke(memory_text + user_input)
        add_to_memory("用户", user_input)
        add_to_memory("Agent", reply)
        return reply

    tool_name = decision.get("tool")
    tool_input = decision.get("input", "")

    # ===== Tool 执行 =====
    if tool_name in TOOLS:
        result = TOOLS[tool_name](tool_input)
        add_to_memory("用户", user_input)
        add_to_memory("Agent", result)
        return result

    # ===== 普通对话 =====
    reply = tool_input
    add_to_memory("用户", user_input)
    add_to_memory("Agent", reply)
    return reply


# ========== 6️⃣ 启动 ==========
print("AI Agent（自动 Tool + Memory）已启动，输入 exit 退出")

while True:
    user_input = input("你：")

    if user_input in ("exit", "quit"):
        print("Agent 已退出")
        break

    result = agent_run(user_input)
    print("Agent：", result)

运行

python main3.py

测试1

你：我叫wp
Agent： 你好，wp
你：我叫什么
Agent： 你好，wp

测试2

你：计算9+8
Agent： 计算结果：17
你：这个结果乘3
Agent： 计算结果：51

三、结论

我们在其中加了一段加记忆的代码，参考3️⃣段的代码，这段代码就是在模拟“人类的短期记忆”（只记住了最近几轮的对话，用来帮助ai理解上下文，不是永久存储）

拆解这段代码

1、建立Memory数据结构

conversation_memory = []

这是一个列表，里面的每一项都是下面的结构：

{
  "role": "用户" 或 "Agent",
  "content": "说的话"
}

用上面的测试1:

[
  {"role": "用户", "content": "我叫wp"},
  {"role": "Agent", "content": "你好，wp"},
]

这和 ChatGPT 内部的 message 结构是同一思想

2、限制 Memory 长度（非常关键）

MAX_MEMORY_LENGTH = 6  # 防止 prompt 无限变长

为什么要限制？

• LLM有上下文长度限制
• Prompt太长
  • 变慢
  • 变贵
  • 容易跑偏

所以我设置了6条，差不多就是3轮对话

3、记忆写入逻辑

def add_to_memory(role: str, content: str):
    conversation_memory.append({"role": role, "content": content})
    # 只保留最近 N 条
    if len(conversation_memory) > MAX_MEMORY_LENGTH:
        del conversation_memory[0]

这个方法定义，当用户说完一句话，并且Agent给出回复后，触发这个方法，通过conversation_memory.append将2条memory加入到Memory中

后面俩句句逻辑是做记忆裁剪，当Memory中的记忆条数大于定义的MAX_MEMORY_LENGTH长度后，就删除最早的记忆：conversation_memory[0]，效果如下：

[旧 旧 旧 新 新 新]  →  [旧 新 新 新]

这个叫“滑动窗口”

4、把记忆变成 Prompt

def render_memory() -> str:
    """
    把 Memory 渲染成 prompt 文本
    """
    if not conversation_memory:
        return ""

    memory_text = "以下是之前的对话内容（供你理解上下文）：\n"
    for msg in conversation_memory:
        memory_text += f"{msg['role']}：{msg['content']}\n"
    return memory_text

作用不是存，而是给模型用

def render_memory() -> str:

空记忆直接返回,避免prompt 里出现废话

  if not conversation_memory:
        return ""

Prompt 头部说明

memory_text = "以下是之前的对话内容（供你理解上下文）：\n"

逐条拼接成自然语言

for msg in conversation_memory:
        memory_text += f"{msg['role']}：{msg['content']}\n"

结论

1. Memory不参与Tool执行
2. Memory给模型看，不是给程序用
3. Memory是可替代吃层
   1. 换成Redis
   2. 换成 DB
   3. 换成 Session
   4. 换成 LangGraph