Example of LangGraph with Ollama for conditional logic

Here, we will use the LangGraph library to create a simple AI agent that can decide whether to answer a user’s question directly or use a “search” tool to find the answer first.

The system starts by receiving user input in the form of text or voice commands. This input is then processed through several stages, each of which performs specific functions to analyze and understand the input, and to generate an appropriate response.

The first stage of processing is natural language understanding, which uses machine learning algorithms to identify the meaning and intent behind the user’s input. This allows the system to determine whether the user is asking a question or making a statement, and what specific information they are seeking.

Once the system has identified the intent behind the user’s input, it can begin to generate a response that is relevant to their query. This response may be generated using pre-defined rules, which provide the system with a set of options for how to respond to different types of queries. Alternatively, the system may use machine learning algorithms to generate a customized response based on the user’s input.

The second stage of processing is natural language generation, which involves taking the information that the system has gathered from the user and producing an appropriate output in the form of text or voice commands. This output is then returned to the user as the system’s response.

---

1. Defining the Tool 🛠️

Python

@tool
def search(query: str):
    """Call to surf the web to answer questions."""
    print(f"---SEARCHING THE WEB FOR: '{query}'---")
    if "langgraph" in query.lower():
        return "LangGraph is a library for building stateful, multi-agent applications with LLMs."
    return "I don't know the answer to that."

tools = [search]

• @tool: This is a decorator from LangChain that registers the search function as a tool the AI model can use.
• search(query: str): This is a simple, “mock” search function. It doesn’t actually search the web. It’s programmed to give a specific answer if the query contains “langgraph” and a generic “I don’t know” for everything else. In a real application, this function would connect to a search engine API.

---

2. Defining the State 🧠

Python

class AgentState(TypedDict):
    messages: Annotated[Sequence[BaseMessage], operator.add]
    tool_call: dict | None

This is the memory or “state” of our agent. The graph will update this state as it runs.

• AgentState(TypedDict): This defines the structure of our state object, which is like a Python dictionary with defined types.
• messages: This is a list that will store the entire conversation history (e.g., human questions, AI answers). The Annotated[..., operator.add] part is special for LangGraph; it tells the graph that whenever this state is updated, new messages should be added to the existing list rather than replacing it.
• tool_call: This is a placeholder to store the details of a tool call if the model decides to use one. It will hold a dictionary like {'tool_name': 'search', 'arguments': {...}} or be None if no tool is being called.

---

3. Defining the Nodes and Edges 🗺️

Nodes are the steps in our workflow (like functions to be executed), and edges are the connections between them.

The Nodes (The Steps)

• call_model(state: AgentState): This is the “brain” of the agent.
  1. It creates a system prompt that instructs the model (Ollama llama3) on how to behave and, crucially, how to format its output as JSON when it wants to use a tool.
  2. It uses format="json" with the ChatOllama model to strongly encourage it to output valid JSON.
  3. It tries to parse the model’s output. If the output is a valid JSON representing a tool call, it updates the state with the tool_call details.
  4. If the model outputs something that isn’t a valid JSON tool call (or isn’t JSON at all), the try...except block gracefully handles this by treating the output as a final, direct text answer.
• call_tool_node(state: AgentState): This node is executed only if the call_model node decided to use a tool.
  1. It reads the tool_call information from the state.
  2. It executes the corresponding tool (in this case, our search function).
  3. It takes the result from the tool and formats it as a new AI message.
  4. It updates the state by adding this new message to the messages list and clearing the tool_call so the process can end.

The Edge Logic (The Decision)

• should_continue(state: AgentState): This function acts as a router. After the call_model node runs, this function checks the state.
  • If the tool_call field in the state has been filled, it returns "call_tool", directing the graph to run the call_tool_node next.
  • If tool_call is empty (None), it returns "end", telling the graph the process is finished.

---

4. Building the Graph 🏗️

Python

workflow = StateGraph(AgentState)

# Add nodes
workflow.add_node("agent", call_model)
workflow.add_node("call_tool", call_tool_node)

# Set entry point
workflow.set_entry_point("agent")

# Add edges
workflow.add_conditional_edges(
    "agent",
    should_continue,
    {"call_tool": "call_tool", "end": END},
)
workflow.add_edge("call_tool", END)

app = workflow.compile()

This section assembles the nodes and edges into a complete, runnable application.

1. StateGraph(AgentState): Initializes a new graph defined by our AgentState memory structure.
2. add_node(...): Adds our functions as nodes named “agent” and “call_tool”.
3. set_entry_point("agent"): Specifies that any run should start with the “agent” (call_model) node.
4. add_conditional_edges(...): This is the core logic. It says: “After running the ‘agent’ node, call the should_continue function. Based on its return value, go to the ‘call_tool’ node or go to the special END state.”
5. add_edge("call_tool", END): This is a simple, direct edge. It says: “After running the ‘call_tool’ node, the process is finished (END).”
6. app = workflow.compile(): Compiles these definitions into an executable object.

---

5. Running the Graph 🚀

Python

inputs = {"messages": [HumanMessage(content="What is langgraph?")], "tool_call": None}
final_state = app.invoke(inputs)

final_answer = final_state['messages'][-1]
print(final_answer.content)

This is the execution phase.

1. An inputs dictionary is created that matches the AgentState structure, containing the user’s initial question.
2. app.invoke(inputs) runs the entire graph from start to finish.
3. Based on the input “What is langgraph?”, the execution flow is:
   • Start -> agent node: The model is called. It sees the word “langgraph” and, following its system prompt, correctly outputs a JSON object requesting to use the search tool. The state is updated with this tool_call.
   • Conditional Edge: The should_continue function sees the tool_call in the state and directs the flow to the call_tool node.
   • call_tool node: The search function is executed with the query “what is LangGraph?”. It returns its pre-programmed answer. This answer is added to the messages list, and tool_call is cleared.
   • Edge -> END: The graph finishes.
4. Finally, the script accesses the messages list from the final_state and prints the content of the very last message, which is the result from the tool.

Whole code:

import json
from typing import TypedDict, Annotated, Sequence
import operator

from langchain_core.messages import BaseMessage, HumanMessage, AIMessage
from langchain_core.output_parsers import JsonOutputParser
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.tools import tool
from langchain_ollama import ChatOllama
from langgraph.graph import StateGraph, END

# --- 1. Define the tools ---
@tool
def search(query: str):
    """Call to surf the web to answer questions."""
    print(f"---SEARCHING THE WEB FOR: '{query}'---")
    if "langgraph" in query.lower():
        return "LangGraph is a library for building stateful, multi-agent applications with LLMs."
    return "I don't know the answer to that."

tools = [search]

# --- 2. Define the State ---
# We add a place to store the parsed tool call from the model
class AgentState(TypedDict):
    messages: Annotated[Sequence[BaseMessage], operator.add]
    tool_call: dict | None  # To store the parsed JSON for the tool call

# --- 3. Define the Nodes and Edges ---
def should_continue(state: AgentState):
    """Decide whether to call a tool or end the conversation."""
    if state.get("tool_call"):
        return "call_tool"
    else:
        return "end"

def call_model(state: AgentState):
    """
    Call the Ollama model. This is now driven by a system prompt
    that instructs the model to output JSON for tool calls.
    """
    print("---CALLING THE MODEL---")
    
    # The system prompt describes the tools and how to call them
    system_prompt = (
        "You are a helpful assistant. You have access to a search tool. "
        "If you need to search for information, respond with a JSON object "
        "containing 'tool_name' and 'arguments'. 'arguments' should be a "
        "dictionary with a single key: 'query'.\n"
        "For example: {{\"tool_name\": \"search\", \"arguments\": {{\"query\": \"what is LangGraph?\"}}}}\n"
        "If you don't need to use the tool, just respond to the user directly."
    )
    
    prompt = ChatPromptTemplate.from_messages([
        ("system", system_prompt),
        MessagesPlaceholder(variable_name="messages"),
    ])
    
    model = ChatOllama(model="llama3", format="json") # format="json" helps guide the model
    
    # We chain the prompt, model, and a JSON parser
    chain = prompt | model | JsonOutputParser()
    
    try:
        # Invoke the chain with the current conversation history
        response = chain.invoke({"messages": state['messages']})
        # If we get valid JSON with a tool_name, we have a tool call
        if "tool_name" in response and "arguments" in response:
            print(f"AI: Requested to use tool: {response['tool_name']}")
            return {"tool_call": response}
        else:
            # The model returned JSON, but it's not a tool call
            # This can happen if it gets confused. We'll treat it as a final message.
            return {"messages": [AIMessage(content=json.dumps(response))], "tool_call": None}
    except Exception:
        # If the model's output is not valid JSON, we assume it's a direct answer.
        # We re-run the model without the JSON parser.
        print("---MODEL FAILED TO RETURN JSON, TREATING AS TEXT---")
        text_model = ChatOllama(model="llama3")
        text_chain = prompt | text_model
        response_text = text_chain.invoke({"messages": state['messages']})
        return {"messages": [response_text], "tool_call": None}


def call_tool_node(state: AgentState):
    """Execute the tool that the model requested."""
    tool_call = state["tool_call"]
    tool_name = tool_call["tool_name"]
    arguments = tool_call["arguments"]
    
    print(f"---CALLING TOOL: '{tool_name}'---")
    
    if tool_name == "search":
        result = search.invoke(arguments)
    else:
        result = f"Error: Unknown tool '{tool_name}'"
        
    # We create a new AIMessage to represent the tool's output
    # This is more semantic than a ToolMessage for this prompt-based approach
    message = AIMessage(content=f"The search results are: {result}")
    
    # We clear the tool_call state and append the result message
    return {"messages": [message], "tool_call": None}

# --- 4. Build the Graph ---
workflow = StateGraph(AgentState)

workflow.add_node("agent", call_model)
workflow.add_node("call_tool", call_tool_node)

workflow.set_entry_point("agent")

workflow.add_conditional_edges(
    "agent",
    should_continue,
    {"call_tool": "call_tool", "end": END},
)

# After the tool is called, we don't need to loop back. The agent's job is done.
# The `call_tool_node` now produces the final answer.
workflow.add_edge("call_tool", END)

app = workflow.compile()

# --- 5. Run the Graph ---
inputs = {"messages": [HumanMessage(content="What is langgraph?")], "tool_call": None}
final_state = app.invoke(inputs)

# The final answer is the last message in the state
final_answer = final_state['messages'][-1]
print("\n------\n🏁 Final Answer:")
print(final_answer.content)