1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93
| from dataclasses import dataclass
from langchain.agents import create_agent from langchain.chat_models import init_chat_model from langchain.tools import tool, ToolRuntime from langgraph.checkpoint.memory import InMemorySaver
SYSTEM_PROMPT = """You are an expert weather forecaster, who speaks in puns.
You have access to two tools:
- get_weather_for_location: use this to get the weather for a specific location - get_user_location: use this to get the user's location
If a user asks you for the weather, make sure you know the location. If you can tell from the question that they mean wherever they are, use the get_user_location tool to find their location."""
@dataclass class Context: """Custom runtime context schema.""" user_id: str
@tool def get_weather_for_location(city: str) -> str: """Get weather for a given city.""" return f"It's always sunny in {city}!"
@tool def get_user_location(runtime: ToolRuntime[Context]) -> str: """Retrieve user information based on user ID.""" user_id = runtime.context.user_id return "Florida" if user_id == "1" else "SF"
model = init_chat_model( "claude-sonnet-4-5-20250929", temperature=0 )
@dataclass class ResponseFormat: """Response schema for the agent.""" punny_response: str weather_conditions: str | None = None
checkpointer = InMemorySaver()
agent = create_agent( model=model, system_prompt=SYSTEM_PROMPT, tools=[get_user_location, get_weather_for_location], context_schema=Context, response_format=ResponseFormat, checkpointer=checkpointer )
config = {"configurable": {"thread_id": "1"}}
response = agent.invoke( {"messages": [{"role": "user", "content": "外面天气怎么样?"}]}, config=config, context=Context(user_id="1") )
print(response['structured_response'])
response = agent.invoke( {"messages": [{"role": "user", "content": "thank you!"}]}, config=config, context=Context(user_id="1") )
print(response['structured_response'])
|