from langchain.agents import AgentState from typing import NotRequired
classCustomState(AgentState): """Extended agent state with custom fields.""" model_call_count: NotRequired[int] # How many times the model was called session_start: NotRequired[str] # When this session started
在工具中访问 State
LangChain 在工具中提供了一个保留的 runtime 参数。通过它可以读取 state：
1 2 3 4 5 6 7
from langchain.tools import tool, ToolRuntime
@tool defmy_tool(runtime: ToolRuntime): """A tool that reads agent state.""" count = runtime.state.get("model_call_count", 0) returnf"Model has been called {count} times so far."
# Record session start time on the first interaction if message_count <= 2: updates["session_start"] = datetime.now()
return Command(update=updates)
注册自定义 State
创建智能体时，指向你的自定义 state schema：
1 2 3 4 5 6 7 8 9 10
from langchain.agents import create_agent from langgraph.checkpoint.memory import InMemorySaver
agent = create_agent( "deepseek-chat", tools=[update_state], state_schema=CustomState, # Tell the agent to use our state checkpointer=InMemorySaver(), system_prompt="You are a helpful assistant. Always call update_state to track your invocation." )
检查 State
调用结束后，你可以获取完整的 state 快照：
1 2 3 4 5 6 7 8 9
config = {"configurable": {"thread_id": "1"}} response = agent.invoke( {"messages": [HumanMessage(content="Hi, my name is Alex")]}, config )
# View the full state snapshot snapshot = agent.get_state(config) print(snapshot.values)
Store — Long-term Memory
State 是短暂的，而 Store 跨会话持久化。它用来存放用户偏好、领域知识、故障日志以及一切需要超越单次对话而存在的数据。
from langgraph_cli.schemas import IndexConfig from langchain_community.embeddings import DashScopeEmbeddings import os
# Initialize embedding model embedding_model = DashScopeEmbeddings( model="text-embedding-v4", dashscope_api_key=os.getenv("DASHSCOPE_API_KEY") )
# Create vector-backed store memory_store = InMemoryStore(index=IndexConfig( embed=embedding_model, dims=1024 ))
使用向量存储后，可以进行语义搜索：
1 2 3 4 5 6 7
results = memory_store.search( ("users",), query="technical staff", # semantic search limit=5 ) for item in results: print(f"Score: {item.score:.3f} | Data: {item.value}")
score 越高表示相似度越高。
在工具中访问 Store
和 State 一样，Store 也可以通过 runtime 访问：
1 2 3 4 5 6 7 8 9 10 11
@tool defget_user_info(user_id: str, runtime: ToolRuntime) -> str: """Look up user info from the store.""" if runtime.store isNone: return"Store not available"
user_info = runtime.store.get(("users",), user_id) if user_info isNone: return"User not found"
returnf"User info: {user_info.value}"
将 Store 注册到智能体
1 2 3 4 5
agent = create_agent( model="deepseek-chat", tools=[get_user_info], store=memory_store # Attach the store )
@tool defget_users(runtime: ToolRuntime[UserContext]): """Query all users — requires clearance level 3+.""" store = runtime.store if store isNone: return"Store not available"
# Who is making the request? user_id = runtime.context.user_id if user_id isNoneor user_id == "": return"User not logged in."
# Check permissions user = store.get(("users",), user_id) if user isNone: return"User not found."
user_info = dict(user.value) if user_info.get("clearance_level", 0) < 3: return"Insufficient permissions!"
results = store.search(("users",)) return [item.value for item in results]
@tool defget_user_preferences(runtime: ToolRuntime[UserContext]): """Get preferences for the current user.""" user_id = runtime.context.user_id pref = runtime.store.get(("preferences",), user_id) if pref isNone: return"No preferences found." return pref.value
在调用时传入 Context
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17
agent = create_agent( model="deepseek-chat", tools=[get_users, get_user_preferences], store=memory_store, context_schema=UserContext, system_prompt=""" # identity You are a helpful assistant that can look up user info and preferences. # instruction Always format results according to the user's preference style. """ )
response = agent.invoke( {"messages": [HumanMessage("Hello, show me all users")]}, context=UserContext(user_id="user_001") # Pass context here )
agent = create_agent( model="deepseek-chat", tools=[transfer_money], middleware=[hitl], checkpointer=InMemorySaver(), system_prompt="You are an account management assistant. Help users transfer money." )
from langgraph.runtime import Runtime from langchain.agents import AgentState from typing import NotRequired, Any from langchain.agents.middleware import after_model
# 1. Define custom state with a counter classCustomAgentState(AgentState): """Extended state with a model call counter.""" model_call_count: NotRequired[int]
# 2. Define the middleware @after_model(state_schema=CustomAgentState) defincrement_counter(state: CustomAgentState, runtime: Runtime) -> dict[str, Any]: """After each model call, bump the counter.""" current = state.get("model_call_count", 0) return {"model_call_count": current + 1}
from langchain.agents.middleware import ( wrap_model_call, ModelRequest, ModelResponse, ) from typing importCallable
@wrap_model_call defretry_model( request: ModelRequest, handler: Callable[[ModelRequest], ModelResponse], ) -> ModelResponse: """Retry model calls up to 3 times on failure.""" for attempt inrange(3): try: return handler(request) except Exception as e: print(f"Retry {attempt + 1}/3 after error: {e}") if attempt == 2: raise returnNone# unreachable, but satisfies the type checker
from langchain.agents.middleware import AgentMiddleware from langchain.agents.middleware.types import ModelCallResult, ToolCallRequest from langgraph.types import Command from typing importCallable
for chunk, metadata in agent.stream( {"messages": [HumanMessage("How's the weather in Hangzhou?")]}, stream_mode="messages" ): if chunk and chunk.content: print(chunk.content, end="", flush=True)
输出：
1 2 3 4 5 6 7
=== Calling model with 1 messages === Let me check the weather in Hangzhou for you. === Calling tool: get_weather === Args: {'location': 'Hangzhou'}
=== Tool call succeeded === The weather in Hangzhou is sunny, 25°C — perfect for going out! 🌞
from langchain.agents.middleware import wrap_model_call from collections.abc importCallable from pydantic.dataclasses import dataclass from langchain.chat_models import init_chat_model
# Initialize two models reasoning_model = init_chat_model(model="deepseek-reasoner") chat_model = init_chat_model(model="deepseek-chat")
# Context to control which model to use @dataclass classUserContext: reasoning: bool = False
@wrap_model_call defdynamic_model_selector( request: ModelRequest, handler: Callable[[ModelRequest], ModelResponse], ) -> ModelResponse: """Pick a model based on the runtime context.""" is_reasoning = getattr(request.runtime.context, 'reasoning', False) selected = reasoning_model if is_reasoning else chat_model print(f"Selected model: {selected.model_name} (reasoning={is_reasoning})")
# Use the reasoning model for this invocation response = agent.invoke( {"messages": [HumanMessage("Solve this math problem: ...")]}, context=UserContext(reasoning=True) )
from langgraph.runtime import Runtime from typing import NotRequired, Any from langchain.agents.middleware import AgentMiddleware, hook_config, AgentState
classCustomAgentState(AgentState): """State with a call counter.""" model_call_count: NotRequired[int]