Skip to content

Mode Based Agent Architecture

Jump to bottom
Joe Curlee (w4ffl35) edited this page Nov 5, 2025 · 1 revision

Mode-Based Agent Architecture Analysis & Implementation Plan

Executive Summary

Based on the LangSmith Complex Agent Evaluation Guide, we need to implement a mode-based hierarchical agent architecture with specialized subgraphs for different contexts (author mode, coding mode, research mode, etc.). This document outlines the current state, gaps, and implementation plan.

Current Architecture Analysis

✅ What We Have

  1. LangGraph Integration: /src/airunner/components/llm/managers/workflow_manager.py

    • Single StateGraph(WorkflowState) with model → tools → model loop
    • Checkpoint persistence via DatabaseCheckpointSaver
    • Tool redundancy detection and forced response handling
    • Supports ReAct-style tool calling

  2. Tool Registry System: /src/airunner/components/llm/core/tool_registry.py

    • 37 registered tools across multiple categories
    • Decorator-based registration: @tool(category=ToolCategory.X)
    • Categories: CHAT, IMAGE, SYSTEM, FILE, SEARCH, CONVERSATION, RAG, MOOD, ANALYSIS, WORKFLOW, MATH

  3. Visual Workflow Builder: /src/airunner/components/nodegraph/

    • GUI node graph editor for LangGraph workflows
    • Generates Python code at runtime
    • Database-backed workflow persistence
    • NOT currently integrated with main LLM workflow

❌ Critical Gaps

  1. No Mode-Based Routing

    • Missing parent graph that classifies user intent
    • No automatic mode switching based on conversation context
    • All tools active at once → overwhelming for LLM

  2. No Specialized Subgraphs

    • Single monolithic workflow for all tasks
    • No author_mode, code_mode, research_mode, etc.
    • No tool isolation per mode

  3. Incomplete Eval Testing

    • Current eval tests only check tool triggering + response content
    • Missing trajectory evaluation (path through nodes/tools)
    • Missing single-step evaluation (intent classification)
    • No mode-switching evaluation

  4. Tool Category Misalignment

    • Current categories don't map to modes (SYSTEM, FILE, SEARCH, etc.)
    • Need mode-specific categories (AUTHOR, CODE, RESEARCH, etc.)

Recommended Architecture (Based on LangSmith Pattern)

Parent Graph (Intent Classifier)

┌─────────────┐
│   START     │
└──────┬──────┘
       │
       ▼
┌─────────────────────────┐
│  Intent Classifier      │  ← Classifies: author, code, research, qa, refund, etc.
│  (GPT-4o-mini)          │
└────┬────────────────────┘
     │
     ├──────► author_agent ──────► compile_followup ──► END
     ├──────► code_agent   ──────► compile_followup ──► END
     ├──────► research_agent ────► compile_followup ──► END
     ├──────► qa_agent     ──────► compile_followup ──► END
     └──────► general_agent ────► compile_followup ──► END

Specialized Subgraphs

Author Mode (Writing Assistant)

Tools: Writing analysis, style checking, grammar, thesaurus, research notes

author_graph = StateGraph(WorkflowState)
author_graph.add_node("gather_writing_context", gather_context)
author_graph.add_node("analyze_style", analyze_writing)
author_graph.add_node("generate_content", generate_writing)
author_graph.add_node("tools", ToolNode([
    improve_writing,
    check_grammar,
    find_synonyms,
    research_topic,
]))

Code Mode (Programming Assistant)

Tools: Code execution, debugging, file operations, git operations

code_graph = StateGraph(WorkflowState)
code_graph.add_node("analyze_code_request", analyze_request)
code_graph.add_node("execute_code", execute_code_safely)
code_graph.add_node("tools", ToolNode([
    run_python_code,
    read_file,
    write_file,
    git_operations,
    search_codebase,
]))

Research Mode (Information Gathering)

Tools: Web search, RAG search, knowledge extraction, summarization

research_graph = StateGraph(WorkflowState)
research_graph.add_node("plan_research", plan_research_strategy)
research_graph.add_node("gather_sources", gather_information)
research_graph.add_node("synthesize", synthesize_findings)
research_graph.add_node("tools", ToolNode([
    search_web,
    rag_search,
    save_to_knowledge_base,
    record_knowledge,
]))

QA Mode (Question Answering)

Tools: RAG retrieval, knowledge recall, reasoning

qa_graph = StateGraph(WorkflowState)
qa_graph.add_node("retrieve_context", retrieve_relevant_docs)
qa_graph.add_node("answer", generate_answer)
qa_graph.add_node("tools", ToolNode([
    rag_search,
    recall_knowledge,
    recall_by_category,
]))

Implementation Plan

Phase 1: Reorganize Tool Categories (Week 1)

  1. Define Mode-Based Categories

    class ToolCategory(Enum):
    # Mode-specific categories
    AUTHOR = "author"           # Writing tools
    CODE = "code"               # Programming tools  
    RESEARCH = "research"       # Research/search tools
    QA = "qa"                   # Q&A/RAG tools
    SYSTEM = "system"           # System operations
    
    # Cross-mode categories
    CONVERSATION = "conversation"
    IMAGE = "image"
    MATH = "math"

  2. Reclassify Existing Tools

    • Move web_tools → RESEARCH
    • Move rag_tools → QA + RESEARCH
    • Move knowledge_tools → QA
    • Move file tools → CODE
    • Create writing_tools.py → AUTHOR

  3. Create Mode-Specific Tool Modules

    src/airunner/components/llm/tools/
├── author_tools.py      # Writing assistant tools
├── code_tools.py        # Programming tools (NEW)
├── research_tools.py    # Research/search tools
└── qa_tools.py          # Question answering tools

Phase 2: Implement Parent Routing Graph (Week 2)

  1. Create Intent Classifier

    # src/airunner/components/llm/managers/mode_router.py

class UserIntent(TypedDict):
    intent: Literal["author", "code", "research", "qa", "general"]

router_llm = ChatModel().with_structured_output(UserIntent)

async def intent_classifier(state: WorkflowState):
    response = router_llm.invoke([
        {"role": "system", "content": ROUTING_INSTRUCTIONS},
        *state["messages"]
    ])
    return Command(goto=response["intent"] + "_agent")

  2. Build Parent Graph

    parent_graph = StateGraph(WorkflowState)
parent_graph.add_node("intent_classifier", intent_classifier)
parent_graph.add_node("author_agent", author_graph)
parent_graph.add_node("code_agent", code_graph)
parent_graph.add_node("research_agent", research_graph)
parent_graph.add_node("qa_agent", qa_graph)
parent_graph.add_node("general_agent", general_graph)
parent_graph.add_node("compile_followup", compile_followup)

parent_graph.set_entry_point("intent_classifier")
# All agents route to compile_followup → END

  3. Integrate with WorkflowManager

    • Replace single graph with parent graph
    • Pass mode-specific tools to each subgraph
    • Maintain backward compatibility

Phase 3: Implement Specialized Subgraphs (Week 3-4)

  1. Author Agent (src/airunner/components/llm/agents/author_agent.py)

    • Tools: Writing analysis, grammar checking, style improvement
    • State: Writing context, style preferences
    • Workflow: Analyze → Generate → Refine

  2. Code Agent (src/airunner/components/llm/agents/code_agent.py)

    • Tools: Code execution, file operations, debugging
    • State: Code context, execution results
    • Workflow: Analyze → Execute → Validate

  3. Research Agent (src/airunner/components/llm/agents/research_agent.py)

    • Tools: Web search, RAG search, knowledge storage
    • State: Research context, gathered sources
    • Workflow: Plan → Gather → Synthesize

  4. QA Agent (src/airunner/components/llm/agents/qa_agent.py)

    • Tools: RAG retrieval, knowledge recall
    • State: Retrieved context, confidence scores
    • Workflow: Retrieve → Reason → Answer

Phase 4: Enhanced Evaluation Framework (Week 5)

  1. Trajectory Evaluation

    def trajectory_eval(outputs: dict, reference_outputs: dict) -> float:
    """Evaluate agent path through nodes/tools."""
    expected_trajectory = reference_outputs['trajectory']
    actual_trajectory = outputs['trajectory']
    
    # Check subsequence match (order matters)
    matches = sum(1 for exp in expected_trajectory 
                  if exp in actual_trajectory)
    return matches / len(expected_trajectory)

  2. Single-Step Evaluation

    # Test intent classification in isolation
intent_dataset = [
    {
        "inputs": {"messages": [{"role": "user", "content": "Help me write a story"}]},
        "outputs": {"route": "author_agent"}
    },
    {
        "inputs": {"messages": [{"role": "user", "content": "Debug this Python code"}]},
        "outputs": {"route": "code_agent"}
    },
]

def test_intent_classifier():
    command = parent_graph.nodes['intent_classifier'].invoke(inputs)
    return {"route": command.goto}

  3. Mode-Specific Eval Tests

    • test_author_mode_eval.py - Writing task quality
    • test_code_mode_eval.py - Code execution correctness
    • test_research_mode_eval.py - Source quality + synthesis
    • test_qa_mode_eval.py - Answer accuracy + grounding

  4. Update Existing Eval Tests

    • Add trajectory tracking to all tests
    • Add expected_trajectory to datasets
    • Implement _track_trajectory() helper
    • Test mode switching mid-conversation

Phase 5: Integration & Testing (Week 6)

  1. GUI Integration

    • Update chat interface to display current mode
    • Add mode override controls (force specific mode)
    • Visualize trajectory in debug mode

  2. Performance Optimization

    • Cache intent classification for follow-ups
    • Implement mode memory (stick to mode if context unchanged)
    • Optimize tool loading per mode

  3. Documentation

    • Mode-specific user guides
    • Tool catalog organized by mode
    • Trajectory visualization examples

Evaluation Metrics (LangSmith Pattern)

Final Response Evaluation

  • Accuracy: Is the final answer correct/useful?
  • Grounding: Is answer based on retrieved context?
  • Completeness: Does answer address all aspects?

Trajectory Evaluation

  • Efficiency: Did agent take shortest path?
  • Correctness: Did agent use right tools in right order?
  • Recovery: Did agent recover from errors?

Single-Step Evaluation

  • Intent Classification: Did router choose correct mode? (Target: >95%)
  • Tool Selection: Did agent pick right tools? (Per-mode target: >90%)
  • Context Gathering: Did agent retrieve relevant info?

Success Criteria

  1. Intent Classification Accuracy: >95% on test dataset
  2. Mode-Specific Tool Usage: <10 tools active per mode (vs 37 global)
  3. Trajectory Correctness: >85% match expected path
  4. Response Quality: Maintain/improve current quality scores
  5. User Experience: Clear mode indicators, faster responses

Technical Considerations

Backward Compatibility

  • Keep general_agent as fallback for ambiguous intents
  • Maintain current API for headless usage
  • Support manual mode override

Scalability

  • Easy to add new modes (gaming_mode, teaching_mode, etc.)
  • Tool registration remains unchanged
  • Each mode can have its own LLM model/config

Testing

  • Unit tests for each subgraph
  • Integration tests for mode switching
  • Eval tests with trajectory validation
  • Benchmark tests for performance

File Structure (Post-Implementation)

src/airunner/components/llm/
├── agents/
│   ├── __init__.py
│   ├── author_agent.py       # Writing assistant
│   ├── code_agent.py         # Programming assistant
│   ├── research_agent.py     # Research assistant
│   ├── qa_agent.py           # Q&A assistant
│   └── general_agent.py      # Fallback agent
├── managers/
│   ├── workflow_manager.py   # Updated with parent graph
│   └── mode_router.py        # Intent classification
├── tools/
│   ├── author_tools.py       # Writing-specific tools
│   ├── code_tools.py         # Code-specific tools
│   ├── research_tools.py     # Research-specific tools
│   ├── qa_tools.py           # Q&A-specific tools
│   └── ...
└── core/
    └── tool_registry.py      # Updated categories

Eval Test Structure (Post-Implementation)

src/airunner/components/eval/tests/
├── test_intent_classifier_eval.py    # Single-step: routing
├── test_author_mode_eval.py          # Author mode: writing quality
├── test_code_mode_eval.py            # Code mode: execution correctness
├── test_research_mode_eval.py        # Research mode: source quality
├── test_qa_mode_eval.py              # QA mode: answer accuracy
├── test_mode_switching_eval.py       # Multi-turn mode transitions
└── test_trajectory_eval.py           # Path correctness across modes

Next Steps

  1. ✅ Review this plan - Validate approach with team
  2. Phase 1: Reorganize tool categories (1 week)
  3. Phase 2: Implement parent routing graph (1 week)
  4. Phase 3: Build specialized subgraphs (2 weeks)
  5. Phase 4: Enhanced evaluation framework (1 week)
  6. Phase 5: Integration & testing (1 week)

Total Timeline: 6 weeks

References

Clone this wiki locally