Noah_
Noah_

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

Answers (0)

Related Questions