Skip to content

Instantly share code, notes, and snippets.

View kennethleungty's full-sized avatar
🎯
Focusing

Kenneth Leung kennethleungty

🎯
Focusing
View GitHub Profile
def wrap_model_call(
self,
request: ModelRequest,
handler: Callable[[ModelRequest], ModelResponse],
) -> ModelCallResult:
"""Update the system message to include the todo system prompt."""
if request.system_message is not None:
new_system_content = [
*request.system_message.content_blocks,
{"type": "text", "text": f"\n\n{self.system_prompt}"},
user_input = "Book business trip to Singapore between 20 and 24 Sep 2025 for an AI conference"
response = agent.invoke({"messages": [HumanMessage(user_input)]})
# Show the to-do list generated by the planning agent
print(response["todos"])
WRITE_TODOS_SYSTEM_PROMPT = """## `write_todos`
You have access to the `write_todos` tool to help you manage and plan complex objectives.
Use this tool for complex objectives to ensure that you are tracking each necessary step and giving the user visibility into your progress.
This tool is very helpful for planning complex objectives, and for breaking down these larger complex objectives into smaller steps.
It is critical that you mark todos as completed as soon as you are done with a step. Do not batch up multiple steps before marking them as completed.
For simple objectives that only require a few steps, it is better to just complete the objective directly and NOT use this tool.
Writing todos takes time and tokens, use it when it is helpful for managing complex many-step problems! But not for simple few-step requests.
WRITE_TODOS_TOOL_DESCRIPTION = """Use this tool to create and manage a structured task list for your current work session. This helps you track progress, organize complex tasks, and demonstrate thoroughness to the user.
Only use this tool if you think it will be helpful in staying organized. If the user's request is trivial and takes less than 3 steps, it is better to NOT use this tool and just do the task directly.
## When to Use This Tool
Use this tool in these scenarios:
1. Complex multi-step tasks - When a task requires 3 or more distinct steps or actions
2. Non-trivial and complex tasks - Tasks that require careful planning or multiple operations
3. User explicitly requests todo list - When the user directly asks you to use the todo list
@tool(description=WRITE_TODOS_TOOL_DESCRIPTION)
def write_todos(todos: list[Todo], tool_call_id: Annotated[str, InjectedToolCallId]) -> Command:
"""Create and manage a structured task list for your current work session."""
return Command(
update={
"todos": todos,
"messages": [ToolMessage(f"Updated todo list to {todos}", tool_call_id=tool_call_id)],
}
)
class PlanningState(AgentState):
"""State schema for the todo middleware."""
todos: Annotated[NotRequired[list[Todo]], OmitFromInput]
"""List of todo items for tracking task progress."""
class Todo(TypedDict):
"""A single todo item with content and status."""
content: str
"""The content/description of the todo item."""
status: Literal["pending", "in_progress", "completed"]
"""The current status of the todo item."""
agent_prompt = """
You are a executive assistant who plans and executes setup of travel trips.
If any details are missing, no need to ask user, just make realistic assumptions directly.
After entire plan is completed, succinctly explain rationale for sequence of actions of the plan.
"""
agent = create_agent(
system_prompt=agent_prompt,
model="openai:gpt-5.1",
tools=[buy_travel_insurance, book_flight, book_hotel],
def supervisor_conditional_node(state: SupervisorState) -> dict:
messages = state["messages"]
# Call supervisor with full conversation history
response = supervisor_llm.invoke([SystemMessage(content=SUPERVISOR_PROMPT)] + messages)
return {"messages": [AIMessage(content=response.content, name="supervisor")]}
def supervisor_command_node(
state: SupervisorState,
) -> Command[Literal["transaction_history_agent", "property_profile_agent"]]:
"""Supervisor node using Command for routing + state updates.
Uses structured output to:
1. Determine which agent to route to
2. Extract property_name from the query
3. Update state via Command.update