catonline_ai/vw-agentic-rag/docs/topics/LANGGRAPH_IMPROVEMENTS.md

LangGraph Implementation Analysis and Improvements

Official Example vs Current Implementation

Key Differences Found

1. Graph Structure

Official Example:

workflow = StateGraph(AgentState)
workflow.add_node("agent", call_model)
workflow.add_node("tools", run_tools) 
workflow.set_entry_point("agent")
workflow.add_conditional_edges("agent", should_continue, ["tools", END])
workflow.add_edge("tools", "agent")
graph = workflow.compile()

Current Implementation:

class AgentWorkflow:
    def __init__(self):
        self.agent_node = AgentNode()
        self.post_process_node = PostProcessNode()
    
    async def astream(self, state, stream_callback):
        state = await self.agent_node(state, stream_callback)
        state = await self.post_process_node(state, stream_callback)

2. State Management

Official Example:

class AgentState(TypedDict):
    messages: Annotated[list, add_messages]

Current Implementation:

class TurnState(BaseModel):
    session_id: str
    messages: List[Message] = Field(default_factory=list)
    tool_results: List[ToolResult] = Field(default_factory=list)
    citations: List[Citation] = Field(default_factory=list)
    # ... many more fields

3. Tool Handling

Official Example:

@tool
def get_stock_price(stock_symbol: str):
    return mock_stock_data[stock_symbol]

tools = [get_stock_price]
tool_node = ToolNode(tools)

Current Implementation:

async def _execute_tool_call(self, tool_call, state, stream_callback):
    async with RetrievalTools() as retrieval:
        if tool_name == "retrieve_standard_regulation":
            result = await retrieval.retrieve_standard_regulation(**tool_args)
        # Manual tool execution logic

Recommendations for Improvement

1. Use Standard LangGraph Patterns

• Adopt StateGraph with add_node() and add_edge()
• Use @tool decorators for cleaner tool definitions
• Leverage ToolNode for automatic tool execution

2. Simplify State Management

• Reduce state complexity where possible
• Use LangGraph's add_messages helper for message handling
• Keep only essential fields in the main state

3. Improve Code Organization

• Separate concerns: graph definition, tool definitions, state
• Use factory functions for graph creation
• Follow LangGraph's recommended patterns

4. Better Tool Integration

• Use @tool decorators for automatic schema generation
• Leverage LangGraph's built-in tool execution
• Reduce manual tool call handling

Implementation Plan

Phase 1: Create Simplified Graph (✅ Done)

• service/graph/simplified_graph.py - follows LangGraph patterns
• Uses @tool decorators
• Cleaner state management
• Reduced complexity

Phase 2: Update Main Implementation

• Refactor existing graph.py to use LangGraph patterns
• Keep existing functionality but improve structure
• Maintain backward compatibility

Phase 3: Testing and Migration

• Test simplified implementation
• Gradual migration of features
• Performance comparison

Code Comparison

Tool Definition

Before:

async def _execute_tool_call(self, tool_call, state, stream_callback):
    tool_name = tool_call["name"]
    tool_args = tool_call["args"]
    async with RetrievalTools() as retrieval:
        if tool_name == "retrieve_standard_regulation":
            result = await retrieval.retrieve_standard_regulation(**tool_args)
        # 20+ lines of manual handling

After:

@tool
async def retrieve_standard_regulation(query: str, conversation_history: str = "") -> str:
    async with RetrievalTools() as retrieval:
        result = await retrieval.retrieve_standard_regulation(query=query, conversation_history=conversation_history)
        return f"Found {len(result.results)} results"

Graph Creation

Before:

class AgentWorkflow:
    def __init__(self):
        self.agent_node = AgentNode()
        self.post_process_node = PostProcessNode()

After:

def create_agent_graph():
    workflow = StateGraph(AgentState)
    workflow.add_node("agent", call_model)
    workflow.add_node("tools", run_tools)
    workflow.set_entry_point("agent")
    workflow.add_conditional_edges("agent", should_continue, ["tools", END])
    return workflow.compile()

Benefits of LangGraph Patterns

1. Declarative: Graph structure is explicit and easy to understand
2. Modular: Nodes and edges can be easily modified
3. Testable: Individual nodes can be tested in isolation
4. Standard: Follows LangGraph community conventions
5. Maintainable: Less custom logic, more framework features
6. Debuggable: LangGraph provides built-in debugging tools