Tutorial: Advanced Water Management Agent¶

This tutorial demonstrates how to build an advanced LangGraph/DLLMForge agent with hydrological related examples. The agent showcases core LangGraph concepts including nodes, edges, and conditional edges using practical water management examples.

Learning Objectives¶

After completing this tutorial, you will understand:

How to create specialized nodes for different types of water analysis
How to implement conditional edges that route queries based on water problem types
How to group tools by domain (calculations vs. information retrieval)
How to build a workflow that intelligently handles hydrological questions

Core Concepts Demonstrated¶

Nodes: Different types of processing units

Agent node: Makes routing decisions
Calculation node: Handles water calculations
Info node: Retrieves aquifer and watershed information
Unified node: Handles complex queries requiring both calculations and information
Concise summary node: Generates a brief, numeric-faithful summary
Extended summary node: Generates a single-paragraph, detailed summary
Steps taken node: Prints routing decisions, tool calls, results (and unified toolset when applicable)

Edges: Connections between nodes, this control is what makes it more advanced than the simple agent, as it allows explicit routing rather than relying on the prompt only.

Conditional edges (agentic): Agent decides which tool node to use → calculation_node | info_node | unified_node
Regular edges (deterministic): - calculation_node → concise_summary → steps_taken → END - info_node → concise_summary → steps_taken → END - unified_node → extended_summary → steps_taken → END

Tools: Grouped by domain for clarity

Water calculation tools: Flow rate, groundwater storage, water balance, Darcy velocity
Information tools: Aquifer and watershed data (with RAG system placeholders)

Workflow Overview¶

The water management agent uses intelligent routing to determine the appropriate processing path (copy the code into mermaid.live to visualize the workflow):

Water Calculation Tools¶

The agent includes four essential water calculation tools based on fundamental hydrological equations:

Flow Rate Calculation

@tool
def calculate_flow_rate(area: float, velocity: float) -> float:
    """Calculate flow rate using Q = A × V (discharge = area × velocity)."""
    return area * velocity

Inputs/Units:

area: channel cross-sectional area (e.g., m²)
velocity: mean flow velocity (e.g., m/s)

Returns:

Discharge (e.g., m³/s)

Example:

Q = 15 m² × 1.5 m/s = 22.5 m³/s

Groundwater Storage Calculation

@tool
def calculate_groundwater_storage(aquifer_area: float, thickness: float, porosity: float) -> str:
    """Calculate groundwater storage volume using V = A × h × n and return a readable result."""
    volume = aquifer_area * thickness * porosity
    return f"Groundwater storage V = A×h×n = {aquifer_area} × {thickness} × {porosity} = {volume}"

Inputs/Units:

aquifer_area (A): plan-view area (e.g., m² or km² converted to m²)
thickness (h): saturated thickness (e.g., m)
porosity (n): effective porosity (0–1)

Returns:

Readable string containing the computed groundwater storage volume (e.g., m³)

Example:

V = 2,000 × 30 × 0.25 = 15,000 (units depend on input units)

Water Balance Calculation

@tool
def calculate_water_balance(precipitation: float, evapotranspiration: float, runoff: float) -> float:
    """Calculate water balance: P - ET - R = ΔS (change in storage)."""
    return precipitation - evapotranspiration - runoff

Inputs/Units:

precipitation (P): total precipitation (e.g., mm)
evapotranspiration (ET): actual ET (e.g., mm)
runoff (R): surface runoff (e.g., mm)

Returns:

ΔS (change in storage), in same units as inputs

Example:

ΔS = 1000 - 600 - 200 = 200 (mm)

Darcy Velocity Calculation

@tool
def calculate_darcy_velocity(hydraulic_conductivity: float, hydraulic_gradient: float) -> float:
    """Calculate Darcy velocity using v = K × i."""
    return hydraulic_conductivity * hydraulic_gradient

Inputs/Units:

hydraulic_conductivity (K): e.g., m/s
hydraulic_gradient (i): dimensionless

Returns:

Darcy velocity (specific discharge), e.g., m/s

Example:

v = 0.01 × 0.05 = 5×10⁻⁴ m/s

Information Retrieval Tools¶

The agent includes information tools with placeholders for connecting to RAG (Retrieval Augmented Generation) systems:

Aquifer Information Tool

@tool  # placeholder for connecting a Retrieval Augmented Generation (RAG) system
def get_aquifer_info(aquifer_name: str) -> str:
    """Get information about a specific aquifer."""
    # This would connect to your RAG system for real-time aquifer data
    aquifer_data = {
        "ogallala": "The Ogallala Aquifer is a major water source...",
        "floridan": "The Floridan Aquifer System is one of the world's most productive aquifers...",
        "edwards": "The Edwards Aquifer is a karst limestone system supplying water to San Antonio; known for rapid recharge and springflows.",
        "high plains": "The High Plains (Ogallala) Aquifer underlies parts of 8 U.S. states and supports extensive irrigation."
        # Additional aquifer data
    }
    return aquifer_data.get(aquifer_name.lower(), f"Sorry, I don't have specific information about the {aquifer_name} aquifer.")

Watershed Information Tool

@tool  # placeholder for connecting a Retrieval Augmented Generation (RAG) system
def get_watershed_info(watershed_name: str) -> str:
    """Get information about a specific watershed."""
    # This would connect to your RAG system for real-time watershed data
    watershed_data = {
        "mississippi": "The Mississippi River watershed drains 41% of the continental US...",
        "colorado": "The Colorado River watershed spans 7 US states...",
        "amazon": "The Amazon watershed is the largest on Earth, with vast biodiversity and significant discharge to the Atlantic.",
        "nile": "The Nile watershed spans 11 countries, sustaining agriculture and livelihoods along the river corridor."
        # Additional watershed data
    }
    return watershed_data.get(watershed_name.lower(), f"Sorry, I don't have specific information about the {watershed_name} watershed.")

Conditional Edge Implementation¶

The conditional edge function determines which node to route to based on the tools the agent wants to call:

def route_to_node(state):
    """Conditional edge: AI decides which node to use based on the water problem."""
    messages = state["messages"]
    last_message = messages[-1]

    # Check if the agent wants to call tools
    if hasattr(last_message, 'tool_calls') and last_message.tool_calls:
        tool_calls = last_message.tool_calls
        tool_names = [tool_call['name'] for tool_call in tool_calls]

        # Check which type of tools are needed
        calculation_tools_needed = any(tool in ['calculate_flow_rate', 'calculate_groundwater_storage', 'calculate_water_balance', 'calculate_darcy_velocity'] for tool in tool_names)
        info_tools_needed = any(tool in ['get_aquifer_info', 'get_watershed_info'] for tool in tool_names)

        # Route based on tool types
        if calculation_tools_needed and info_tools_needed:
            return "unified_node"  # Both calculation and info needed
        elif calculation_tools_needed:
            return "calculation_node"  # Only calculations needed
        elif info_tools_needed:
            return "info_node"  # Only info needed
        else:
            return "unified_node"  # Default fallback
    else:
        # No tools needed, route to concise summary
        return "concise_summary"

Workflow Assembly¶

The workflow is assembled by adding nodes and edges:

# Group tools by domain
calculation_tools = [calculate_flow_rate, calculate_groundwater_storage, calculate_water_balance, calculate_darcy_velocity]
info_tools = [get_aquifer_info, get_watershed_info]
all_tools = calculation_tools + info_tools

# Create specialized tool nodes
calculation_node = ToolNode(calculation_tools)
info_node = ToolNode(info_tools)
unified_node = ToolNode(all_tools)

# Add nodes to the workflow
agent.add_node("calculation_node", calculation_node)
agent.add_node("info_node", info_node)
agent.add_node("unified_node", unified_node)
agent.add_node("concise_summary", concise_summary)
agent.add_node("extended_summary", extended_summary)
agent.add_node("steps_taken", steps_taken)

# Add conditional edge
agent.add_conditional_edge("agent", route_to_node)

# Add regular edges
agent.add_edge(START, "agent")
agent.add_edge("calculation_node", "concise_summary")
agent.add_edge("info_node", "concise_summary")
agent.add_edge("unified_node", "extended_summary")
agent.add_edge("concise_summary", "steps_taken")
agent.add_edge("extended_summary", "steps_taken")
agent.add_edge("steps_taken", END)

Testing the Workflow¶

The tutorial includes test cases that demonstrate different routing scenarios:

Test Case 1: Flow Rate Calculation

Query: “What’s the flow rate for a channel with area 15 m² and velocity 1.5 m/s?”
Route: Agent → Calculation Node → Concise Summary → Steps Taken → End
Tools Used: calculate_flow_rate

Test Case 2: Aquifer Information

Query: “Tell me about the Ogallala aquifer”
Route: Agent → Info Node → Concise Summary → Steps Taken → End
Tools Used: get_aquifer_info

Test Case 3: Complex Analysis

Query: “Calculate groundwater storage for area 2000 km², thickness 30 m, porosity 0.25 and tell me about the Floridan aquifer”
Route: Agent → Unified Node → Extended Summary → Steps Taken → End
Tools Used: calculate_groundwater_storage + get_aquifer_info

Additional Examples¶

Test Case 4: Water Balance

Query: “Calculate the water balance for precipitation 1200 mm, evapotranspiration 700 mm, and runoff 350 mm”
Route: Agent → Calculation Node → Concise Summary → Steps Taken → End
Tools Used: calculate_water_balance

Test Case 5: Darcy Velocity

Query: “What’s the Darcy velocity for hydraulic conductivity 0.008 m/s and hydraulic gradient 0.04?”
Route: Agent → Calculation Node → Concise Summary → Steps Taken → End
Tools Used: calculate_darcy_velocity

Test Case 6: Watershed Context + Calculation

Query: “Give me the Colorado watershed context and compute Darcy velocity for K=0.01 m/s, i=0.03”
Route: Agent → Unified Node → Extended Summary → Steps Taken → End
Tools Used: get_watershed_info + calculate_darcy_velocity

Test Case 7: No Tools Needed (Direct Summary)

Query: “Summarize your capabilities in one sentence”
Route: Agent → Concise Summary → Steps Taken → End
Tools Used: none

Running the Tutorial¶

To run the complete tutorial:

python tutorial_advanced_agent.py

This will: 1. Display the workflow structure 2. Run the test cases 3. Show the routing decisions for each query type

Testing with Custom Queries¶

Once the agent is compiled, you can test it with your own water management queries:

# Test flow rate calculation
agent.process_query("What's the flow rate for a channel with area 15 m² and velocity 1.5 m/s?", stream=True)

# Test aquifer information
agent.process_query("Tell me about the Ogallala aquifer", stream=True)

# Test complex analysis requiring both calculations and information
agent.process_query("Calculate groundwater storage for area 2000 km², thickness 30 m, porosity 0.25 and tell me about the Floridan aquifer", stream=True)

# Test water balance calculation
agent.process_query("Calculate the water balance for precipitation 1000 mm, evapotranspiration 600 mm, and runoff 200 mm", stream=True)

# Test Darcy velocity
agent.process_query("What's the Darcy velocity for hydraulic conductivity 0.01 m/s and hydraulic gradient 0.05?", stream=True)

The agent will automatically: 1. Analyze the query to determine which tools are needed 2. Route to the appropriate node (calculation, info, or unified) 3. Execute the required tools for calculations or information retrieval 4. Generate a comprehensive summary of the results

Key Benefits for Water Professionals¶

Domain-Specific Tools: All tools are designed for hydrological calculations and information retrieval

Intelligent Routing: The agent automatically determines whether calculations, information lookup, or both are needed

RAG Integration Ready: Information tools include placeholders for connecting to real-time databases

Scalable Architecture: Easy to add new calculation tools or information sources

Professional Context: Uses proper hydrological terminology and real-world examples

Next Steps¶

Connect the information tools to your RAG system for real-time aquifer and watershed data
Add more specialized calculation tools for your specific hydrological needs
Implement additional nodes for data visualization or report generation
Extend the conditional routing logic for more complex decision trees

This tutorial provides a solid foundation for building sophisticated water management agents using LangGraph’s powerful workflow capabilities.

Tutorial: Advanced Water Management Agent¶

Learning Objectives¶

Core Concepts Demonstrated¶

Workflow Overview¶

Water Calculation Tools¶

Information Retrieval Tools¶

Conditional Edge Implementation¶

Workflow Assembly¶

Testing the Workflow¶

Additional Examples¶

Running the Tutorial¶

Testing with Custom Queries¶

Key Benefits for Water Professionals¶

Next Steps¶

dllmforge

Navigation

Related Topics