ConversationBufferMemory 迁移到 v1：thread_id 与历史裁剪实测

"""What the old memory classes do (langchain 0.3.x): one object, one string, no session key."""
import warnings
 
from langchain.memory import ConversationBufferMemory, ConversationBufferWindowMemory
 
warnings.filterwarnings("ignore")  # the deprecation notice is printed once per class otherwise
 
 
def exchange(memory, question, answer):
    memory.save_context({"input": question}, {"output": answer})
 
 
print("A. buffer memory returns one string")
buffer = ConversationBufferMemory()
exchange(buffer, "hi, my name is bob", "hello bob")
exchange(buffer, "what is 2+2", "4")
print(buffer.load_memory_variables({}))
 
print()
print("B. return_messages=True returns message objects")
as_messages = ConversationBufferMemory(return_messages=True)
exchange(as_messages, "hi", "hello")
print([(m.type, m.content) for m in as_messages.load_memory_variables({})["history"]])
 
print()
print("C. window memory keeps the last k exchanges")
window = ConversationBufferWindowMemory(k=1)
for i in range(3):
    exchange(window, f"q{i}", f"a{i}")
print(window.load_memory_variables({}))
 
print()
print("D. one shared instance has no notion of a user")
shared = ConversationBufferMemory()
exchange(shared, "I am alice, my card ends in 4242", "noted")
print("bob asks for his history:", shared.load_memory_variables({}))
 
print()
print("E. the buffer grows without bound")
growing = ConversationBufferMemory()
for turn in range(1, 101):
    exchange(growing, f"question number {turn}", f"answer number {turn}")
    if turn in (10, 50, 100):
        print(f"after {turn} exchanges: {len(growing.load_memory_variables({})['history'])} characters")
 
print()
print("F. clear() empties the one shared history")
shared.clear()
print(shared.load_memory_variables({}))

A. buffer memory returns one string
{'history': 'Human: hi, my name is bob\nAI: hello bob\nHuman: what is 2+2\nAI: 4'}

B. return_messages=True returns message objects
[('human', 'hi'), ('ai', 'hello')]

C. window memory keeps the last k exchanges
{'history': 'Human: q2\nAI: a2'}

D. one shared instance has no notion of a user
bob asks for his history: {'history': 'Human: I am alice, my card ends in 4242\nAI: noted'}

E. the buffer grows without bound
after 10 exchanges: 451 characters
after 50 exchanges: 2331 characters
after 100 exchanges: 4683 characters

F. clear() empties the one shared history
{'history': ''}

"""A chat model that answers with a counter and remembers what it was shown."""
from typing import Any, List, Optional, Sequence
 
from langchain_core.language_models.chat_models import BaseChatModel
from langchain_core.messages import AIMessage, BaseMessage
from langchain_core.outputs import ChatGeneration, ChatResult
 
 
class RecordingModel(BaseChatModel):
    calls: int = 0
    seen: List[List[BaseMessage]] = []
 
    @property
    def _llm_type(self) -> str:
        return "recording"
 
    def bind_tools(self, tools: Sequence[Any], **kwargs: Any):
        return self
 
    def _generate(
        self,
        messages: List[BaseMessage],
        stop: Optional[List[str]] = None,
        run_manager: Any = None,
        **kwargs: Any,
    ) -> ChatResult:
        self.calls += 1
        self.seen.append(list(messages))
        return ChatResult(generations=[ChatGeneration(message=AIMessage(content=f"answer {self.calls}"))])

"""The v1 replacement: a checkpointer plus a thread_id decides who remembers what."""
from langchain.agents import create_agent
from langchain_core.messages import HumanMessage
from langgraph.checkpoint.memory import InMemorySaver
from recording_model import RecordingModel
 
 
def turns(agent, model, config, count, label):
    seen_sizes = []
    for turn in range(count):
        before = len(model.seen)
        agent.invoke({"messages": [HumanMessage(f"turn {turn + 1}")]}, config)
        seen_sizes.append(len(model.seen[before]))
    print(f"{label:<44} messages the model saw per turn: {seen_sizes}")
 
 
model = RecordingModel()
agent = create_agent(model, tools=[], checkpointer=InMemorySaver())
turns(agent, model, {"configurable": {"thread_id": "alice"}}, 3, "checkpointer, thread alice")
turns(agent, model, {"configurable": {"thread_id": "bob"}}, 2, "same agent, thread bob")
turns(agent, model, {"configurable": {"thread_id": "alice"}}, 1, "thread alice again")
stored = agent.get_state({"configurable": {"thread_id": "alice"}}).values["messages"]
print(f"stored messages for alice: {len(stored)} ({[m.type for m in stored][:4]} ...)")
 
print()
model = RecordingModel()
agent = create_agent(model, tools=[])
turns(agent, model, {"configurable": {"thread_id": "alice"}}, 3, "no checkpointer, thread_id passed anyway")
 
print()
model = RecordingModel()
agent = create_agent(model, tools=[], checkpointer=InMemorySaver())
try:
    agent.invoke({"messages": [HumanMessage("hello")]})
except Exception as exc:
    print(f"checkpointer but no thread_id: {type(exc).__name__}: {str(exc).splitlines()[0]}")

checkpointer, thread alice                   messages the model saw per turn: [1, 3, 5]
same agent, thread bob                       messages the model saw per turn: [1, 3]
thread alice again                           messages the model saw per turn: [7]
stored messages for alice: 8 (['human', 'ai', 'human', 'ai'] ...)

no checkpointer, thread_id passed anyway     messages the model saw per turn: [1, 1, 1]

checkpointer but no thread_id: ValueError: Checkpointer requires one or more of the following 'configurable' keys: thread_id, checkpoint_ns, checkpoint_id

"""Bounding the history the model sees: delete from state, or trim only the model's view."""
from typing import Any
 
from langchain.agents import AgentState, create_agent
from langchain.agents.middleware import before_model, wrap_model_call
from langchain.messages import RemoveMessage
from langchain_core.messages import HumanMessage, trim_messages
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.graph.message import REMOVE_ALL_MESSAGES
from langgraph.runtime import Runtime
from recording_model import RecordingModel
 
TURNS = 8
 
 
@before_model
def trim_state(state: AgentState, runtime: Runtime) -> dict[str, Any] | None:
    """Keep the first message and the last few, as in the official short-term memory page."""
    messages = state["messages"]
    if len(messages) <= 3:
        return None
    first = messages[0]
    recent = messages[-3:] if len(messages) % 2 == 0 else messages[-4:]
    return {"messages": [RemoveMessage(id=REMOVE_ALL_MESSAGES), first, *recent]}
 
 
@wrap_model_call
def trim_view(request, handler):
    """Trim only what this model call sees; the stored history stays complete."""
    kept = trim_messages(request.messages, strategy="last", token_counter=len, max_tokens=4, start_on="human")
    return handler(request.override(messages=kept))
 
 
def run(label, middleware):
    model = RecordingModel()
    agent = create_agent(model, tools=[], middleware=middleware, checkpointer=InMemorySaver())
    config = {"configurable": {"thread_id": "t1"}}
    for turn in range(1, TURNS + 1):
        agent.invoke({"messages": [HumanMessage(f"turn {turn}")]}, config)
    sizes = [len(call) for call in model.seen]
    last_view = [m.content for m in model.seen[-1] if m.type == "human"]
    stored = agent.get_state(config).values["messages"]
    print(f"{label}")
    print(f"  messages seen per turn : {sizes}")
    print(f"  human turns in the last model call: {last_view}")
    print(f"  messages stored after {TURNS} turns: {len(stored)}")
    print()
 
 
run("no middleware", [])
run("@before_model trim_state (deletes from state)", [trim_state])
run("@wrap_model_call trim_view (trims the view only)", [trim_view])

no middleware
  messages seen per turn : [1, 3, 5, 7, 9, 11, 13, 15]
  human turns in the last model call: ['turn 1', 'turn 2', 'turn 3', 'turn 4', 'turn 5', 'turn 6', 'turn 7', 'turn 8']
  messages stored after 8 turns: 16

@before_model trim_state (deletes from state)
  messages seen per turn : [1, 3, 5, 5, 5, 5, 5, 5]
  human turns in the last model call: ['turn 1', 'turn 7', 'turn 8']
  messages stored after 8 turns: 6

@wrap_model_call trim_view (trims the view only)
  messages seen per turn : [1, 3, 3, 3, 3, 3, 3, 3]
  human turns in the last model call: ['turn 7', 'turn 8']
  messages stored after 8 turns: 16

"""Does the conversation survive a process restart? InMemorySaver versus SqliteSaver."""
import sqlite3
import tempfile
from pathlib import Path
 
from langchain.agents import create_agent
from langchain_core.messages import HumanMessage
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.checkpoint.sqlite import SqliteSaver
from recording_model import RecordingModel
 
CONFIG = {"configurable": {"thread_id": "alice"}}
 
 
def ask(saver, text):
    """Build a fresh agent around the saver, like a newly started process, and ask once."""
    model = RecordingModel()
    agent = create_agent(model, tools=[], checkpointer=saver)
    agent.invoke({"messages": [HumanMessage(text)]}, CONFIG)
    return len(model.seen[0]), len(agent.get_state(CONFIG).values["messages"])
 
 
print("InMemorySaver")
first = ask(InMemorySaver(), "turn 1")
restarted = ask(InMemorySaver(), "turn 2")  # a new saver object is a new process
print(f"  before restart: model saw {first[0]}, stored {first[1]}")
print(f"  after restart : model saw {restarted[0]}, stored {restarted[1]}")
 
print()
print("SqliteSaver on a file")
with tempfile.TemporaryDirectory() as tmp:
    db = str(Path(tmp) / "memory.sqlite")
 
    def open_saver():
        return SqliteSaver(sqlite3.connect(db, check_same_thread=False))
 
    saver = open_saver()
    first = ask(saver, "turn 1")
    saver.conn.close()
 
    saver = open_saver()
    restarted = ask(saver, "turn 2")
    print(f"  before restart: model saw {first[0]}, stored {first[1]}")
    print(f"  after restart : model saw {restarted[0]}, stored {restarted[1]}")
 
    saver.delete_thread("alice")
    deleted = ask(saver, "turn 3")
    print(f"  after delete_thread: model saw {deleted[0]}, stored {deleted[1]}")
    saver.conn.close()

InMemorySaver
  before restart: model saw 1, stored 2
  after restart : model saw 1, stored 2

SqliteSaver on a file
  before restart: model saw 1, stored 2
  after restart : model saw 3, stored 4
  after delete_thread: model saw 1, stored 2