🎯 What You'll Learn Today
LangGraph Tutorial: Intelligent Tool Selection System - Unit 2.2 Exercise 4
This tutorial is also available in Google Colab here or for download here
Joint Initiative: This tutorial is part of a collaboration between AI Product Engineer and the Nebius Academy.
This tutorial demonstrates how to build a robust tool selection system using LangGraph, implementing proper message typing and state management.
Key Concepts Covered
- Message Type Hierarchy
- State Management with TypedDict
- Tool Selection Logic
- Conversation Flow Control
import uuid
from typing import Annotated, TypedDict
!pip install langchain-core
!pip install langgraph
from langchain_core.messages import (
AIMessage,
BaseMessage,
HumanMessage,
SystemMessage,
ToolMessage,
)
from langgraph.graph.message import add_messagesStep 1: State Definition
We define our state structure to manage the conversation and tool selection.
Why This Matters
Proper state structure is crucial because
- Enables type-safe message handling
- Maintains tool availability tracking
- Provides clear tool selection state
Debug Tips
-
State Structure Issues:
- Verify TypedDict fields are properly annotated
- Check message list initialization
- Ensure tool_name can be None
class State(TypedDict):
"""State container for tool selection system.
Attributes:
messages: Conversation history with proper message typing
available_tools: List of tools the agent can access
tool_name: Currently selected tool (or None if no selection)
"""
messages: Annotated[list[BaseMessage], add_messages]
available_tools: list[str]
tool_name: str | NoneStep 2: System Initialization
Set up the initial system state with proper context and instructions.
Why This Matters
System initialization is essential because
- Sets up consistent starting state
- Provides clear tool instructions
- Establishes conversation context
Debug Tips
-
Initialization Issues:
- Verify SystemMessage content
- Check AIMessage formatting
- Ensure message list structure
def initialize_system() -> list[BaseMessage]:
"""Initialize the system with instructions and context.
Returns:
list[BaseMessage]: Initial system messages
"""
return [
SystemMessage(
content="""
I am an AI assistant that can help you with various tasks.
Available tools:
- Calculator: For mathematical calculations
- Weather: For checking weather information
- Search: For finding general information
"""
),
AIMessage(content="Hello! I'm ready to help. What would you like to know?"),
]Step 3: Message Type Management
Implement message type handling and retrieval.
Why This Matters
Message type management is crucial because
- Ensures proper conversation flow
- Maintains message type safety
- Enables targeted message processing
Debug Tips
-
Message Retrieval Issues:
- Add logging for message type checks
- Verify message order handling
- Check type comparison logic
def get_last_message_by_type(state: State, message_type) -> BaseMessage | None:
"""Retrieve the most recent message of a specific type.
Args:
state: Current conversation state
message_type: Type of message to find
Returns:
BaseMessage | None: Latest matching message or None
"""
for message in reversed(state["messages"]):
if isinstance(message, message_type):
return message
return NoneStep 4: Tool Selection Logic
Implement intelligent tool selection based on message content.
Why This Matters
Tool selection logic is important because
- Determines appropriate tool for user needs
- Maintains conversation coherence
- Provides clear selection reasoning
Debug Tips
-
Selection Logic Issues:
- Monitor keyword matching
- Verify tool availability
- Check message generation
def tool_selector(state: State) -> State:
"""Select appropriate tool based on message content.
Args:
state: Current conversation state
Returns:
State: Updated state with tool selection
"""
if not state.get("messages"):
return {**state, "messages": initialize_system()}
last_human = get_last_message_by_type(state, HumanMessage)
if not last_human:
return state
message = last_human.content.lower()
tool_name = None
reasoning = ""
if "weather" in message:
tool_name = "check_weather"
reasoning = "I'll use the weather tool to check weather conditions."
elif any(word in message for word in ["calculate", "compute", "+", "-", "*", "/"]):
tool_name = "calculator"
reasoning = "I'll use the calculator tool to help with this calculation."
else:
tool_name = "search"
reasoning = "I'll use the search tool to find this information."
tool_message = ToolMessage(
content=f"Selected tool: {tool_name}",
tool_call_id=str(uuid.uuid4()),
name="tool_selector",
)
ai_message = AIMessage(content=reasoning)
return {
**state,
"tool_name": tool_name,
"messages": state["messages"] + [tool_message, ai_message],
}Step 5: System Demonstration
Demonstrate the tool selection system in action.
Why This Matters
System demonstration is valuable because
- Verifies system functionality
- Shows expected behavior
- Provides usage patterns
Debug Tips
-
Demonstration Issues:
- Monitor state transitions
- Verify message sequence
- Check output formatting
def demonstrate_tool_selection():
"""Demonstrate tool selection with various queries."""
state = {
"messages": initialize_system(),
"available_tools": ["calculator", "check_weather", "search"],
"tool_name": None,
}
queries = [
"What's the weather in Paris?",
"Calculate 2 + 2",
"Who won the World Cup?",
]
print("Tool Selection Demonstration")
print("==========================")
for query in queries:
state["messages"].append(HumanMessage(content=query))
result = tool_selector(state)
print(f"\nUser Query: {query}")
tool_msg = get_last_message_by_type(result, ToolMessage)
ai_msg = get_last_message_by_type(result, AIMessage)
if tool_msg:
print(f"Tool Selection: {tool_msg.content}")
if ai_msg:
print(f"Assistant: {ai_msg.content}")
print(f"Selected Tool: {result['tool_name']}")
print("-" * 40)Common Pitfalls
- Not handling missing messages properly
- Incorrect message type comparison
- Mutable state modifications
- Incomplete tool selection logic
Key Takeaways
- Message type safety is crucial
- Immutable state updates prevent bugs
- Clear tool selection logic improves reliability
- Proper demonstration validates functionality
Next Steps
- Add tool selection confidence scoring
- Implement multi-tool selection
- Add conversation context tracking
- Implement tool usage analytics
Expected Output
Tool Selection Demonstration
User Query: What's the weather in Paris?
Tool Selection: Selected tool: check_weather
Assistant: I'll use the weather tool to check weather conditions.
## Selected Tool: check_weather
User Query: Calculate 2 + 2
Tool Selection: Selected tool: calculator
Assistant: I'll use the calculator tool to help with this calculation.
## Selected Tool: calculator
User Query: Who won the World Cup?
Tool Selection: Selected tool: search
Assistant: I'll use the search tool to find this information.
## Selected Tool: search
if __name__ == "__main__":
demonstrate_tool_selection()
🇬🇧 Chapter