Reputation: 11
Langgraph State not being passed properly with Ollama
I'm trying to create this simple workflow using Ollama and Langgraph which is a basic travel agent to a user. I am not having any errors, but the output is not what I expect. It seems like there are some issues with the way the state of each agent is being passed to other agents, but I cannot figure out the issue.
I am new to ollama and langgraph and referenced this blog post for manually using ollama schemas: https://ollama.com/blog/structured-outputs, but I think I have them implemented properly.
Ollama is properly installed, and all other libraires are properly installed.
Here's the code:
from langgraph.graph import StateGraph, MessagesState, START, END
from langchain_core.messages import HumanMessage
from ollama import chat
from pydantic import BaseModel
# --------------------------
# Define Schemas
# --------------------------
class WeatherData(BaseModel):
day: int
condition: str
high: int
low: int
class ItineraryItem(BaseModel):
day: int
morning: str
afternoon: str
weather: str
class Itinerary(BaseModel):
itinerary: list[ItineraryItem]
class Budget(BaseModel):
total_cost: float
breakdown: dict
# --------------------------
# Ollama Chat Model
# --------------------------
def ollama_chat(model, messages, schema=None):
"""Wrapper for the Ollama chat call with an optional JSON schema."""
return chat(model=model, messages=messages, format=schema)
# --------------------------
# Weather Agent
# --------------------------
def weather_agent(state: dict) -> dict:
print("Weather Agent: Starting...")
print("DEBUG - state at start of weather_agent:", state)
weather_data = {
1: {"condition": "Sunny", "high": 25, "low": 15},
2: {"condition": "Cloudy", "high": 22, "low": 14},
3: {"condition": "Rainy", "high": 20, "low": 12},
}
print(f"Weather Data: {weather_data}")
new_state = {
**state,
"weather_data": weather_data,
"goto": "itinerary_node"
}
print("DEBUG - state at end of weather_agent:", new_state)
return new_state
# --------------------------
# Itinerary Agent
# --------------------------
def itinerary_agent(state: dict) -> dict:
print("Itinerary Agent: Starting...")
print("DEBUG - state at start of itinerary_agent:", state)
weather_data = state.get("weather_data")
if not weather_data:
print("Error: No weather data found in state")
return {**state, "goto": END}
schema = Itinerary.model_json_schema()
print(f"Calling Ollama with weather data: {weather_data}")
try:
response = ollama_chat(
model="llama3.2",
messages=[
{"role": "system", "content": "You are an itinerary planning agent."},
{"role": "user", "content": f"Plan a 3-day trip to Paris with this weather data: {weather_data}"},
],
schema=schema
)
itinerary = Itinerary.model_validate_json(response.message.content)
print(f"Generated Itinerary: {itinerary}")
new_state = {
**state,
"itinerary": itinerary,
"goto": "budget_node"
}
print("DEBUG - state at end of itinerary_agent:", new_state)
return new_state
except Exception as e:
print(f"Error in itinerary generation: {e}")
return {**state, "goto": END}
# --------------------------
# Budget Agent
# --------------------------
def budget_agent(state: dict) -> dict:
print("Budget Agent: Starting...")
print("DEBUG - state at start of budget_agent:", state)
itinerary = state.get("itinerary")
if not itinerary:
print("Error: No itinerary found in state")
return {**state, "goto": END}
schema = Budget.model_json_schema()
print(f"Calling Ollama with itinerary: {itinerary}")
try:
response = ollama_chat(
model="llama3.2",
messages=[
{"role": "system", "content": "You are a budget estimation agent."},
{"role": "user", "content": f"Estimate costs for this itinerary: {itinerary}"},
],
schema=schema
)
budget = Budget.model_validate_json(response.message.content)
print(f"Generated Budget: {budget}")
new_state = {
**state,
"budget": budget,
"goto": END
}
print("DEBUG - state at end of budget_agent:", new_state)
return new_state
except Exception as e:
print(f"Error in budget estimation: {e}")
return {**state, "goto": END}
# --------------------------
# Supervisor Node
# --------------------------
def supervisor_node(state: dict) -> dict:
print("Supervisor Node: Starting...")
print("DEBUG - state at start of supervisor_node:", state)
messages = state.get("messages", [])
print(f"Messages received: {messages}")
# Check the content of the first message
user_request = messages[0].content if messages else ""
print(f"User request: {user_request}")
if "trip" in user_request.lower():
print("Routing to weather agent...")
return {
**state,
"goto": "weather_node"
}
print("No trip request detected, ending...")
return {
**state,
"goto": END
}
# --------------------------
# Build the Graph
# --------------------------
graph = StateGraph(MessagesState)
graph.add_node("supervisor_node", supervisor_node)
graph.add_node("weather_node", weather_agent)
graph.add_node("itinerary_node", itinerary_agent)
graph.add_node("budget_node", budget_agent)
graph.add_edge(START, "supervisor_node")
graph.add_edge("supervisor_node","weather_node")
graph.add_edge("weather_node", "itinerary_node")
graph.add_edge("itinerary_node", "budget_node")
graph.add_edge("budget_node", END)
# Compile the graph
graph = graph.compile()
# --------------------------
# Run the Graph
# --------------------------
if __name__ == "__main__":
print("\nStarting graph execution...")
initial_state = {
"messages": [HumanMessage(content="Plan a 3-day trip to Paris.")]
}
result = graph.invoke(initial_state)
print("\nFinal Result:")
print("Weather:", result.get("weather_data"))
print("Itinerary:", result.get("itinerary"))
print("Budget:", result.get("budget"))`
and my output:
Starting graph execution... Supervisor Node: Starting... DEBUG - state at start of supervisor_node: {'messages': [HumanMessage(content='Plan a 3-day trip to Paris.', additional_kwargs={}, response_metadata={}, id='b74e1a47-e126-4c5d-831f-e9e4b14497a1')]} Messages received: [HumanMessage(content='Plan a 3-day trip to Paris.', additional_kwargs={}, response_metadata={}, id='b74e1a47-e126-4c5d-831f-e9e4b14497a1')] User request: Plan a 3-day trip to Paris. Routing to weather agent... Weather Agent: Starting... DEBUG - state at start of weather_agent: {'messages': [HumanMessage(content='Plan a 3-day trip to Paris.', additional_kwargs={}, response_metadata={}, id='b74e1a47-e126-4c5d-831f-e9e4b14497a1')]} Weather Data: {1: {'condition': 'Sunny', 'high': 25, 'low': 15}, 2: {'condition': 'Cloudy', 'high': 22, 'low': 14}, 3: {'condition': 'Rainy', 'high': 20, 'low': 12}} DEBUG - state at end of weather_agent: {'messages': [HumanMessage(content='Plan a 3-day trip to Paris.', additional_kwargs={}, response_metadata={}, id='b74e1a47-e126-4c5d-831f-e9e4b14497a1')], 'weather_data': {1: {'condition': 'Sunny', 'high': 25, 'low': 15}, 2: {'condition': 'Cloudy', 'high': 22, 'low': 14}, 3: {'condition': 'Rainy', 'high': 20, 'low': 12}}, 'goto': 'itinerary_node'} Itinerary Agent: Starting... DEBUG - state at start of itinerary_agent: {'messages': [HumanMessage(content='Plan a 3-day trip to Paris.', additional_kwargs={}, response_metadata={}, id='b74e1a47-e126-4c5d-831f-e9e4b14497a1')]} Error: No weather data found in state Budget Agent: Starting... DEBUG - state at start of budget_agent: {'messages': [HumanMessage(content='Plan a 3-day trip to Paris.', additional_kwargs={}, response_metadata={}, id='b74e1a47-e126-4c5d-831f-e9e4b14497a1')]} Error: No itinerary found in state
Final Result: Weather: None Itinerary: None Budget: None
Thanks for the help, I've been trying to use ollama for open-source projects and keep running into this issue.
Upvotes: 1
Views: 168