enterprise-kb-rag/2.2_LangChain基础知识/18_memory_agent_short_term.py

from pathlib import Path
import sys
import time

sys.path.append(str(Path(__file__).resolve().parents[1]))

from _chapter2_config import get_chat_model, print_dependency_hint


try:
    from langchain.agents import create_agent
    from langchain_core.messages import HumanMessage
    try:
        from langgraph.checkpoint.memory import InMemorySaver
    except ImportError:
        from langgraph.checkpoint.memory import MemorySaver as InMemorySaver
except ModuleNotFoundError as exc:
    print_dependency_hint(exc)
    raise SystemExit(1)


def main() -> None:
    try:
        model = get_chat_model(temperature=0)
    except ModuleNotFoundError as exc:
        print_dependency_hint(exc)
        raise SystemExit(1)

    checkpointer = InMemorySaver()
    timestamp = int(time.time())
    config = {"configurable": {"thread_id": f"{timestamp}"}}

    agent = create_agent(
        model=model,
        tools=[],
        checkpointer=checkpointer,
        system_prompt="你是一个翻译官，擅长中英互译。",
    )

    input_1 = {"messages": [HumanMessage(content="你好，我叫小呆。")]}
    response_a = agent.invoke(input=input_1, config=config)

    input_2 = {"messages": [HumanMessage(content="你好，我叫什么？")]}
    config_2 = {"configurable": {"thread_id": f"{timestamp}-new"}}

    response_b = agent.invoke(input=input_2, config=config_2)
    response_c = agent.invoke(input=input_2, config=config)

    print("Response A (第一次对话):")
    print(response_a)
    print("\nResponse B (新线程，无记忆):")
    print(response_b)
    print("\nResponse C (回到原线程，有记忆):")
    print(response_c)


if __name__ == "__main__":
    main()