第 8 章：记忆管理

## 示例：使用 InMemorySessionService
## 这适用于不需要跨应用程序重启的数据持久性的本地开发和测试。
from google.adk.sessions import InMemorySessionService

session_service = InMemorySessionService()

## 示例：使用 DatabaseSessionService
## 这适用于需要持久存储的生产或开发。
## 您需要配置数据库 URL（例如，用于 SQLite、PostgreSQL 等）。
## 需要：pip install google-adk[sqlalchemy] 和数据库驱动程序（例如，PostgreSQL 的 psycopg2）
from google.adk.sessions import DatabaseSessionService

## 示例使用本地 SQLite 文件：
db_url = "sqlite:///./my_agent_data.db"
session_service = DatabaseSessionService(db_url=db_url)

## 示例：使用 VertexAiSessionService
## 这适用于 Google Cloud Platform 上的可扩展生产，利用
## Vertex AI 基础设施进行会话管理。
## 需要：pip install google-adk[vertexai] 和 GCP 设置/身份验证
from google.adk.sessions import VertexAiSessionService

PROJECT_ID = "your-gcp-project-id"  # 替换为您的 GCP 项目 ID
LOCATION = "us-central1"  # 替换为您想要的 GCP 位置

## 与此服务一起使用的 app_name 应对应于 Reasoning Engine ID 或名称
REASONING_ENGINE_APP_NAME = "projects/your-gcp-project-id/locations/us-central1/reasoningEngines/your-engine-id"  # 替换为您的 Reasoning Engine 资源名称

session_service = VertexAiSessionService(project=PROJECT_ID, location=LOCATION)

## 使用此服务时，将 REASONING_ENGINE_APP_NAME 传递给服务方法：
## session_service.create_session(app_name=REASONING_ENGINE_APP_NAME, ...)
## session_service.get_session(app_name=REASONING_ENGINE_APP_NAME, ...)
## session_service.append_event(session, event, app_name=REASONING_ENGINE_APP_NAME)
## session_service.delete_session(app_name=REASONING_ENGINE_APP_NAME, ...)

## 从 Google 智能体开发工具包 (ADK) 导入必要的类
from google.adk.agents import LlmAgent
from google.adk.sessions import InMemorySessionService, Session
from google.adk.runners import Runner
from google.genai.types import Content, Part

## 定义带有 output_key 的 LlmAgent。
greeting_agent = LlmAgent(
    name="Greeter",
    model="gemini-2.0-flash",
    instruction="生成一个简短、友好的问候语。",
    output_key="last_greeting"
)

## --- 设置 Runner 和 Session ---
app_name, user_id, session_id = "state_app", "user1", "session1"
session_service = InMemorySessionService()
runner = Runner(
    agent=greeting_agent,
    app_name=app_name,
    session_service=session_service
)

session = session_service.create_session(
    app_name=app_name,
    user_id=user_id,
    session_id=session_id
)

print(f"初始状态：{session.state}")

## --- 运行智能体 ---
user_message = Content(parts=[Part(text="你好")])
print("\n--- 运行 agent ---")
for event in runner.run(
    user_id=user_id,
    session_id=session_id,
    new_message=user_message
):
    if event.is_final_response():
        print("Agent 已响应。")

## --- 检查更新的状态 ---
## 在 runner 完成处理所有事件*之后*正确检查状态。
updated_session = session_service.get_session(app_name, user_id, session_id)
print(f"\nAgent 运行后的状态：{updated_session.state}")

import time
from google.adk.tools.tool_context import ToolContext
from google.adk.sessions import InMemorySessionService

## --- 定义推荐的基于工具的方法 ---
def log_user_login(tool_context: ToolContext) -> dict:
    """
    在用户登录事件时更新会话状态。
    此工具封装了与用户登录相关的所有状态更改。
    参数：
        tool_context：由 ADK 自动提供，提供对会话状态的访问。
    返回：
        确认操作成功的字典。
    """
    # 通过提供的上下文直接访问状态。
    state = tool_context.state

    # 获取当前值或默认值，然后更新状态。
    # 这更加清晰并将逻辑共置。
    login_count = state.get("user:login_count", 0) + 1
    state["user:login_count"] = login_count
    state["task_status"] = "active"
    state["user:last_login_ts"] = time.time()
    state["temp:validation_needed"] = True

    print("从 `log_user_login` 工具内部更新了状态。")

    return {
        "status": "success",
        "message": f"已跟踪用户登录。总登录次数：{login_count}。"
    }

## --- 使用演示 ---
## 在真实应用程序中，LLM智能体将调用此工具。
## 在这里，我们模拟直接调用以进行演示。

## 1. 设置
session_service = InMemorySessionService()
app_name, user_id, session_id = "state_app_tool", "user3", "session3"
session = session_service.create_session(
    app_name=app_name,
    user_id=user_id,
    session_id=session_id,
    state={"user:login_count": 0, "task_status": "idle"}
)

print(f"初始状态：{session.state}")

## 2. 模拟工具调用（在真实应用中，ADK Runner 会执行此操作）
## 我们仅为此独立示例手动创建 ToolContext。
from google.adk.tools.tool_context import InvocationContext
mock_context = ToolContext(
    invocation_context=InvocationContext(
        app_name=app_name, user_id=user_id, session_id=session_id,
        session=session, session_service=session_service
    )
)

## 3. 执行工具
log_user_login(mock_context)

## 4. 检查更新的状态
updated_session = session_service.get_session(app_name, user_id, session_id)
print(f"工具执行后的状态：{updated_session.state}")

## 预期输出将显示与"之前"情况相同的状态更改，
## 但代码组织明显更清晰
## 且更健壮。

## 示例：使用 InMemoryMemoryService
## 这适用于不需要跨应用程序重启数据持久性的本地开发和测试
## 应用停止时内存内容会丢失
from google.adk.memory import InMemoryMemoryService

memory_service = InMemoryMemoryService()

## 示例：使用 VertexAiRagMemoryService
## 这适用于 GCP 上的可扩展生产，利用
## Vertex AI RAG（检索增强生成）实现持久的、
## 可搜索的记忆。
## 需要：pip install google-adk[vertexai]，GCP
## 设置/身份验证和 Vertex AI RAG Corpus。
from google.adk.memory import VertexAiRagMemoryService

## 您的 Vertex AI RAG Corpus 的资源名称
RAG_CORPUS_RESOURCE_NAME = "projects/your-gcp-project-id/locations/us-central1/ragCorpora/your-corpus-id"  # 替换为您的 Corpus 资源名称

## 检索行为的可选配置
SIMILARITY_TOP_K = 5  # 要检索的顶部结果数
VECTOR_DISTANCE_THRESHOLD = 0.7  # 向量相似度阈值

memory_service = VertexAiRagMemoryService(
    rag_corpus=RAG_CORPUS_RESOURCE_NAME,
    similarity_top_k=SIMILARITY_TOP_K,
    vector_distance_threshold=VECTOR_DISTANCE_THRESHOLD
)

## 使用此服务时，像 add_session_to_memory
## 和 search_memory 这样的方法将与指定的 Vertex AI
## RAG Corpus 交互。

from langchain.memory import ChatMessageHistory

## 初始化历史对象
history = ChatMessageHistory()

## 添加用户和 AI 消息
history.add_user_message("我下周要去纽约。")
history.add_ai_message("太好了！这是一个很棒的城市。")

## 访问消息列表
print(history.messages)

from langchain.memory import ConversationBufferMemory

## 初始化记忆
memory = ConversationBufferMemory()

## 保存对话轮次
memory.save_context({"input": "天气怎么样？"}, {"output": "今天是晴天。"})

## 将记忆作为字符串加载
print(memory.load_memory_variables({}))

from langchain_openai import OpenAI
from langchain.chains import LLMChain
from langchain.prompts import PromptTemplate
from langchain.memory import ConversationBufferMemory

## 1. 定义 LLM 和提示词
llm = OpenAI(temperature=0)
template = """你是一个有帮助的旅行代理。
之前的对话：
{history}
新问题：{question}
响应："""
prompt = PromptTemplate.from_template(template)

## 2. 配置记忆
## memory_key "history" 与提示词中的变量匹配
memory = ConversationBufferMemory(memory_key="history")

## 3. 构建链
conversation = LLMChain(llm=llm, prompt=prompt, memory=memory)

## 4. 运行对话
response = conversation.predict(question="我想预订航班。")
print(response)
response = conversation.predict(question="顺便说一下，我叫 Sam。")
print(response)
response = conversation.predict(question="我的名字是什么？")
print(response)

from langchain_openai import ChatOpenAI
from langchain.chains import LLMChain
from langchain.memory import ConversationBufferMemory
from langchain_core.prompts import (
    ChatPromptTemplate,
    MessagesPlaceholder,
    SystemMessagePromptTemplate,
    HumanMessagePromptTemplate,
)

## 1. 定义聊天模型和提示词
llm = ChatOpenAI()
prompt = ChatPromptTemplate(
    messages=[
        SystemMessagePromptTemplate.from_template("你是一个友好的助手。"),
        MessagesPlaceholder(variable_name="chat_history"),
        HumanMessagePromptTemplate.from_template("{question}")
    ]
)

## 2. 配置记忆
## return_messages=True 对于聊天模型是必不可少的
memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)

## 3. 构建链
conversation = LLMChain(llm=llm, prompt=prompt, memory=memory)

## 4. 运行对话
response = conversation.predict(question="嗨，我是 Jane。")
print(response)
response = conversation.predict(question="你记得我的名字吗？")
print(response)

## 更新 agent 指令的节点
def update_instructions(state: State, store: BaseStore):
    namespace = ("instructions",)
    # 从存储中获取当前指令
    current_instructions = store.search(namespace)[0]

    # 创建提示词要求 LLM 反思对话
    # 并生成新的、改进的指令
    prompt = prompt_template.format(
        instructions=current_instructions.value["instructions"],
        conversation=state["messages"]
    )

    # 从 LLM 获取新指令
    output = llm.invoke(prompt)
    new_instructions = output['new_instructions']

    # 将更新的指令保存回存储
    store.put(("agent_instructions",), "agent_a", {"instructions": new_instructions})

## 使用指令生成响应的节点
def call_model(state: State, store: BaseStore):
    namespace = ("agent_instructions", )
    # 从存储中检索最新指令
    instructions = store.get(namespace, key="agent_a")[0]

    # 使用检索到的指令格式化提示词
    prompt = prompt_template.format(instructions=instructions.value["instructions"])
    # ... 应用程序逻辑继续

from langgraph.store.memory import InMemoryStore

## 真实嵌入函数的占位符
def embed(texts: list[str]) -> list[list[float]]:
    # 在真实应用程序中，使用适当的嵌入模型
    return [[1.0, 2.0] for _ in texts]

## 初始化记忆存储。对于生产，使用数据库支持的存储。
store = InMemoryStore(index={"embed": embed, "dims": 2})

## 为特定用户和应用程序上下文定义命名空间
user_id = "my-user"
application_context = "chitchat"
namespace = (user_id, application_context)

## 1. 将记忆放入存储
store.put(
    namespace,
    "a-memory",  # 此记忆的键
    {
        "rules": [
            "用户喜欢简短、直接的语言",
            "用户只说英语和 python",
        ],
        "my-key": "my-value",
    },
)

## 2. 通过其命名空间和键获取记忆
item = store.get(namespace, "a-memory")
print("检索的项目：", item)

## 3. 在命名空间内搜索记忆，按内容过滤
## 并按与查询的向量相似度排序。
items = store.search(
    namespace,
    filter={"my-key": "my-value"},
    query="语言偏好"
)
print("搜索结果：", items)

from google.adk.memory import VertexAiMemoryBankService

agent_engine_id = agent_engine.api_resource.name.split("/")[-1]
memory_service = VertexAiMemoryBankService(
    project="PROJECT_ID",
    location="LOCATION",
    agent_engine_id=agent_engine_id
)

session = await session_service.get_session(
    app_name=app_name,
    user_id="USER_ID",
    session_id=session.id
)

await memory_service.add_session_to_memory(session)