diff --git a/README.md b/README.md
index c6a715ef..b4995f86 100644
--- a/README.md
+++ b/README.md
@@ -17,6 +17,8 @@ Flo AI is a Python framework for building structured AI agents with support for
+ 
+
@@ -45,6 +47,7 @@ Flo AI is a Python framework that makes building production-ready AI agents and
- ๐ **Truly Composable**: Build complex AI systems by combining smaller, reusable components
- ๐๏ธ **Production-Ready**: Built-in best practices and optimizations for production deployments
- ๐ **YAML-First**: Define your entire agent architecture in simple YAML
+- ๐ง **LLM-Powered Routing**: Intelligent routing decisions made by LLMs, no code required
- ๐ง **Flexible**: Use pre-built components or create your own
- ๐ค **Team-Oriented**: Create and manage teams of AI agents working together
- ๐ **Langchain Compatible**: Works with all your favorite Langchain tools
@@ -82,6 +85,8 @@ Flo AI is a Python framework that makes building production-ready AI agents and
- [Memory and Context Sharing](#memory-and-context-sharing)
- [๐ Use Cases for Arium](#-use-cases-for-arium)
- [Builder Pattern Benefits](#builder-pattern-benefits)
+ - [๐ YAML-Based Arium Workflows](#-yaml-based-arium-workflows)
+ - [๐ง LLM-Powered Routers in YAML (NEW!)](#-llm-powered-routers-in-yaml-new)
- [๐ Documentation](#-documentation)
- [๐ Why Flo AI?](#-why-flo-ai)
- [๐ฏ Use Cases](#-use-cases)
@@ -419,8 +424,6 @@ asyncio.run(multi_agent_variables())
Variables work seamlessly with YAML-based agent configuration:
```yaml
-apiVersion: flo/alpha-v1
-kind: FloAgent
metadata:
name: personalized-assistant
version: 1.0.0
@@ -837,7 +840,8 @@ Arium is Flo AI's powerful workflow orchestration engine that allows you to crea
- **๐ Multi-Agent Workflows**: Orchestrate multiple agents working together
- **๐ฏ Flexible Routing**: Route between agents based on context and conditions
-- **๐ง Shared Memory**: Agents share conversation history and context
+- **๐ง LLM Routers**: Intelligent routing powered by LLMs, define routing logic in YAML
+- **๐พ Shared Memory**: Agents share conversation history and context
- **๐ Visual Workflows**: Generate flow diagrams of your agent interactions
- **โก Builder Pattern**: Fluent API for easy workflow construction
- **๐ Reusable Workflows**: Build once, run multiple times with different inputs
@@ -1133,6 +1137,465 @@ result: List[Any] = await (
- โ
Validates routing functions
- โ
Checks for unreachable nodes
+### ๐ YAML-Based Arium Workflows
+
+One of Flo AI's most powerful features is the ability to define entire multi-agent workflows using YAML configuration. This approach makes workflows reproducible, versionable, and easy to modify without changing code.
+
+#### Simple YAML Workflow
+
+```yaml
+metadata:
+ name: "content-analysis-workflow"
+ version: "1.0.0"
+ description: "Multi-agent content analysis and summarization pipeline"
+
+arium:
+ # Define agents inline
+ agents:
+ - name: "analyzer"
+ role: "Content Analyst"
+ job: "Analyze the input content and extract key insights, themes, and important information."
+ model:
+ provider: "openai"
+ name: "gpt-4o-mini"
+ settings:
+ temperature: 0.2
+ max_retries: 3
+ reasoning_pattern: "COT"
+
+ - name: "summarizer"
+ role: "Content Summarizer"
+ job: "Create a concise, actionable summary based on the analysis provided."
+ model:
+ provider: "anthropic"
+ name: "claude-3-5-sonnet-20240620"
+ settings:
+ temperature: 0.1
+ reasoning_pattern: "DIRECT"
+
+ # Define the workflow
+ workflow:
+ start: "analyzer"
+ edges:
+ - from: "analyzer"
+ to: ["summarizer"]
+ end: ["summarizer"]
+```
+
+```python
+import asyncio
+from typing import Any, List
+from flo_ai.arium import AriumBuilder
+
+async def run_yaml_workflow() -> List[Any]:
+ yaml_config = """...""" # Your YAML configuration
+
+ # Create workflow from YAML
+ result: List[Any] = await (
+ AriumBuilder()
+ .from_yaml(yaml_config)
+ .build_and_run(["Analyze this quarterly business report..."])
+ )
+
+ return result
+
+asyncio.run(run_yaml_workflow())
+```
+
+#### Advanced YAML Workflow with Tools and Routing
+
+```yaml
+metadata:
+ name: "research-workflow"
+ version: "2.0.0"
+ description: "Intelligent research workflow with conditional routing"
+
+arium:
+ # Define agents with tool references
+ agents:
+ - name: "classifier"
+ role: "Content Classifier"
+ job: "Classify input as 'research', 'calculation', or 'analysis' task."
+ model:
+ provider: "openai"
+ name: "gpt-4o-mini"
+ tools: ["web_search"] # Reference tools provided in Python
+
+ - name: "researcher"
+ role: "Research Specialist"
+ job: "Conduct thorough research on with analysis."
+ model:
+ provider: "anthropic"
+ name: "claude-3-5-sonnet-20240620"
+ tools: ["web_search"]
+ settings:
+ temperature: 0.3
+ reasoning_pattern: "REACT"
+
+ - name: "analyst"
+ role: "Data Analyst"
+ job: "Analyze numerical data and provide insights for ."
+ model:
+ provider: "openai"
+ name: "gpt-4o"
+ tools: ["calculator", "web_search"]
+ settings:
+ reasoning_pattern: "COT"
+
+ - name: "synthesizer"
+ role: "Information Synthesizer"
+ job: "Combine research and analysis into final recommendations."
+ model:
+ provider: "gemini"
+ name: "gemini-2.5-flash"
+
+ # Complex workflow with conditional routing
+ workflow:
+ start: "classifier"
+ edges:
+ # Conditional routing based on classification
+ - from: "classifier"
+ to: ["researcher", "analyst"]
+ router: "classification_router" # Router function provided in Python
+
+ # Both specialists feed into synthesizer
+ - from: "researcher"
+ to: ["synthesizer"]
+
+ - from: "analyst"
+ to: ["synthesizer"]
+
+ end: ["synthesizer"]
+```
+
+```python
+import asyncio
+from typing import Any, Dict, List, Literal
+from flo_ai.arium import AriumBuilder
+from flo_ai.tool.base_tool import Tool
+from flo_ai.arium.memory import BaseMemory
+
+# Define tools in Python (cannot be defined in YAML)
+async def web_search(query: str) -> str:
+ # Your search implementation
+ return f"Search results for: {query}"
+
+async def calculate(expression: str) -> str:
+ # Your calculation implementation
+ try:
+ result = eval(expression) # Note: Use safely in production
+ return f"Calculation result: {result}"
+ except:
+ return "Invalid expression"
+
+# Create tool objects
+tools: Dict[str, Tool] = {
+ "web_search": Tool(
+ name="web_search",
+ description="Search the web for current information",
+ function=web_search,
+ parameters={
+ "query": {
+ "type": "string",
+ "description": "Search query"
+ }
+ }
+ ),
+ "calculator": Tool(
+ name="calculator",
+ description="Perform mathematical calculations",
+ function=calculate,
+ parameters={
+ "expression": {
+ "type": "string",
+ "description": "Mathematical expression to calculate"
+ }
+ }
+ )
+}
+
+# Define router functions in Python (cannot be defined in YAML)
+def classification_router(memory: BaseMemory) -> Literal["researcher", "analyst"]:
+ """Route based on task classification"""
+ content = str(memory.get()[-1]).lower()
+ if 'research' in content or 'investigate' in content:
+ return 'researcher'
+ elif 'calculate' in content or 'analyze data' in content:
+ return 'analyst'
+ return 'researcher' # default
+
+routers: Dict[str, callable] = {
+ "classification_router": classification_router
+}
+
+async def run_workflow() -> List[Any]:
+ yaml_config = """...""" # Your YAML configuration from above
+
+ # Create workflow with tools and routers provided as Python objects
+ result: List[Any] = await (
+ AriumBuilder()
+ .from_yaml(
+ yaml_str=yaml_config,
+ tools=tools, # Tools must be provided as Python objects
+ routers=routers # Routers must be provided as Python functions
+ )
+ .build_and_run(["Research the latest trends in renewable energy"])
+ )
+
+ return result
+```
+
+#### ๐ง LLM-Powered Routers in YAML (NEW!)
+
+One of the most powerful new features is the ability to define **intelligent LLM routers directly in YAML**. No more writing router functions - just describe your routing logic and let the LLM handle the decisions!
+
+```yaml
+metadata:
+ name: "intelligent-content-workflow"
+ version: "1.0.0"
+ description: "Content creation with intelligent LLM-based routing"
+
+arium:
+ agents:
+ - name: "content_creator"
+ role: "Content Creator"
+ job: "Create initial content based on the request"
+ model:
+ provider: "openai"
+ name: "gpt-4o-mini"
+
+ - name: "technical_writer"
+ role: "Technical Writer"
+ job: "Refine content for technical accuracy and clarity"
+ model:
+ provider: "openai"
+ name: "gpt-4o-mini"
+
+ - name: "creative_writer"
+ role: "Creative Writer"
+ job: "Enhance content with creativity and storytelling"
+ model:
+ provider: "openai"
+ name: "gpt-4o-mini"
+
+ - name: "marketing_writer"
+ role: "Marketing Writer"
+ job: "Optimize content for engagement and conversion"
+ model:
+ provider: "openai"
+ name: "gpt-4o-mini"
+
+ # โจ LLM Router definitions - No code required!
+ routers:
+ - name: "content_type_router"
+ type: "smart" # Uses LLM to make intelligent routing decisions
+ routing_options:
+ technical_writer: "Technical content, documentation, tutorials, how-to guides"
+ creative_writer: "Creative writing, storytelling, fiction, brand narratives"
+ marketing_writer: "Marketing copy, sales content, landing pages, ad campaigns"
+ model:
+ provider: "openai"
+ name: "gpt-4o-mini"
+ settings:
+ temperature: 0.3
+ fallback_strategy: "first"
+
+ - name: "task_classifier"
+ type: "task_classifier" # Keyword-based classification
+ task_categories:
+ math_solver:
+ description: "Mathematical calculations and problem solving"
+ keywords: ["calculate", "solve", "equation", "math", "formula"]
+ examples: ["Calculate 2+2", "Solve x^2 + 5x + 6 = 0"]
+ code_helper:
+ description: "Programming and code assistance"
+ keywords: ["code", "program", "debug", "function", "algorithm"]
+ examples: ["Write a Python function", "Debug this code"]
+ model:
+ provider: "openai"
+ name: "gpt-4o-mini"
+
+ workflow:
+ start: "content_creator"
+ edges:
+ - from: "content_creator"
+ to: ["technical_writer", "creative_writer", "marketing_writer"]
+ router: "content_type_router" # LLM automatically routes based on content type!
+ end: ["technical_writer", "creative_writer", "marketing_writer"]
+```
+
+**๐ฏ LLM Router Types:**
+
+1. **Smart Router** (`type: smart`): General-purpose routing based on content analysis
+2. **Task Classifier** (`type: task_classifier`): Routes based on keywords and examples
+3. **Conversation Analysis** (`type: conversation_analysis`): Context-aware routing
+
+**โจ Key Benefits:**
+- ๐ซ **No Code Required**: Define routing logic purely in YAML
+- ๐ฏ **Intelligent Decisions**: LLMs understand context and make smart routing choices
+- ๐ **Easy Configuration**: Simple, declarative syntax
+- ๐ **Version Control**: Track routing changes in YAML files
+- ๐๏ธ **Model Flexibility**: Each router can use different LLM models
+
+```python
+# Using LLM routers is incredibly simple!
+async def run_intelligent_workflow():
+ # No routers dictionary needed - they're defined in YAML!
+ result = await (
+ AriumBuilder()
+ .from_yaml(yaml_str=intelligent_workflow_yaml)
+ .build_and_run(["Write a technical tutorial on Docker containers"])
+ )
+ # The LLM will automatically route to technical_writer! โจ
+ return result
+```
+
+#### YAML Workflow with Variables
+
+```yaml
+metadata:
+ name: "personalized-workflow"
+ version: "1.0.0"
+ description: "Workflow that adapts based on input variables"
+
+arium:
+ agents:
+ - name: "specialist"
+ role: ""
+ job: "You are a specializing in . Provide for ."
+ model:
+ provider: ""
+ name: ""
+ settings:
+ temperature: 0.3
+ reasoning_pattern: ""
+
+ - name: "reviewer"
+ role: "Quality Reviewer"
+ job: "Review the for and provide feedback."
+ model:
+ provider: "openai"
+ name: "gpt-4o"
+
+ workflow:
+ start: "specialist"
+ edges:
+ - from: "specialist"
+ to: ["reviewer"]
+ end: ["reviewer"]
+```
+
+```python
+import asyncio
+from typing import Any, Dict, List
+from flo_ai.arium import AriumBuilder
+
+async def run_personalized_workflow() -> List[Any]:
+ yaml_config = """...""" # Your YAML configuration with variables
+
+ # Define variables for the workflow
+ variables: Dict[str, str] = {
+ 'expert_role': 'Data Scientist',
+ 'domain': 'machine learning and predictive analytics',
+ 'output_type': 'technical analysis report',
+ 'target_audience': 'engineering team',
+ 'preferred_llm_provider': 'anthropic',
+ 'model_name': 'claude-3-5-sonnet-20240620',
+ 'reasoning_style': 'COT',
+ 'quality_criteria': 'technical accuracy and clarity'
+ }
+
+ result: List[Any] = await (
+ AriumBuilder()
+ .from_yaml(yaml_config)
+ .build_and_run(
+ ["Analyze our customer churn prediction model performance"],
+ variables=variables
+ )
+ )
+
+ return result
+```
+
+#### Using Pre-built Agents in YAML Workflows
+
+```yaml
+metadata:
+ name: "hybrid-workflow"
+ version: "1.0.0"
+ description: "Mix of inline agents and pre-built agent references"
+
+# Import existing agent configurations
+imports:
+ - "agents/content_analyzer.yaml"
+ - "agents/technical_reviewer.yaml"
+
+arium:
+ # Mix of imported and inline agents
+ agents:
+ # Reference imported agent
+ - import: "content_analyzer"
+ name: "analyzer" # Override name if needed
+
+ # Define new agent inline
+ - name: "formatter"
+ role: "Content Formatter"
+ job: "Format the analysis into a professional report structure."
+ model:
+ provider: "openai"
+ name: "gpt-4o-mini"
+
+ # Reference another imported agent
+ - import: "technical_reviewer"
+ name: "reviewer"
+
+ workflow:
+ start: "analyzer"
+ edges:
+ - from: "analyzer"
+ to: ["formatter"]
+ - from: "formatter"
+ to: ["reviewer"]
+ end: ["reviewer"]
+```
+
+#### YAML Workflow Best Practices
+
+1. **Modular Design**: Define reusable agents in YAML, create tools in Python separately
+2. **Clear Naming**: Use descriptive names for agents and workflows
+3. **Variable Usage**: Leverage variables for environment-specific configurations
+4. **Version Control**: Track workflow versions in metadata
+5. **Documentation**: Include descriptions for complex workflows
+6. **Router Functions**: Keep routing logic simple and provide as Python functions
+7. **Tool Management**: Create tools as Python objects and pass them to the builder
+
+#### What Can Be Defined in YAML vs Python
+
+**โ
YAML Configuration Supports:**
+- Agent definitions (name, role, job, model settings)
+- Workflow structure (start, edges, end nodes)
+- Agent-to-agent connections
+- Tool and router references (by name)
+- Variables and settings
+- Model configurations
+
+**โ YAML Configuration Does NOT Support:**
+- Tool function implementations (must be Python objects)
+- Router function code (must be Python functions)
+- Custom logic execution
+- Direct function definitions
+
+**๐ก Best Practice**: Use YAML for workflow structure and agent configuration, Python for executable logic (tools and routers).
+
+#### Benefits of YAML Workflows
+
+- **๐ Reproducible**: Version-controlled workflow definitions
+- **๐ Maintainable**: Easy to modify workflow structure without code changes
+- **๐งช Testable**: Different configurations for testing vs. production
+- **๐ฅ Collaborative**: Non-developers can modify workflow structure
+- **๐ Deployable**: Easy CI/CD integration with YAML configurations
+- **๐ Auditable**: Clear workflow definitions for compliance
+
> ๐ **For detailed Arium documentation and advanced patterns, see [flo_ai/flo_ai/arium/README.md](flo_ai/flo_ai/arium/README.md)**
## ๐ Documentation
diff --git a/flo_ai/docs/arium_yaml_guide.md b/flo_ai/docs/arium_yaml_guide.md
new file mode 100644
index 00000000..419444a5
--- /dev/null
+++ b/flo_ai/docs/arium_yaml_guide.md
@@ -0,0 +1,815 @@
+# Arium YAML Configuration Guide
+
+This guide explains how to define multi-agent workflows using YAML configuration files instead of programmatic builder patterns.
+
+## Overview
+
+The `AriumBuilder.from_yaml()` method allows you to define complex multi-agent workflows using declarative YAML configuration. This approach provides several benefits:
+
+- **Declarative Configuration**: Define workflows in a human-readable format
+- **Version Control**: Track workflow changes easily
+- **Reusability**: Share and reuse workflow definitions
+- **Separation of Concerns**: Keep workflow logic separate from implementation code
+- **Clean Structure**: Focus on workflow definition without implementation details like memory
+
+## YAML Structure
+
+### Basic Structure
+
+```yaml
+metadata:
+ name: workflow-name
+ version: 1.0.0
+ description: "Description of your workflow"
+
+arium:
+ agents:
+ # Method 1: Reference pre-built agents (cleanest approach)
+ - name: content_analyst # Must exist in agents parameter
+ - name: summarizer # Must exist in agents parameter
+
+ # Method 2: Direct configuration
+ - name: validator
+ role: "Quality Validator"
+ job: "Validate analysis and summary quality"
+ model:
+ provider: openai
+ name: gpt-4o-mini
+ settings:
+ temperature: 0.1
+ reasoning_pattern: DIRECT
+
+ tools:
+ - name: tool1
+ - name: tool2
+
+ # LLM Router definitions (NEW!)
+ routers:
+ - name: content_router
+ type: smart # smart, task_classifier, conversation_analysis
+ routing_options:
+ technical_writer: "Handle technical documentation tasks"
+ creative_writer: "Handle creative writing tasks"
+ editor: "Handle editing and review tasks"
+ model:
+ provider: openai
+ name: gpt-4o-mini
+ settings:
+ temperature: 0.3
+ fallback_strategy: first
+
+ workflow:
+ start: content_analyst
+ edges:
+ - from: content_analyst
+ to: [summarizer]
+ - from: summarizer
+ to: [validator, tool1]
+ router: content_router # References router defined above
+ - from: validator
+ to: [end]
+ - from: tool1
+ to: [end]
+ end: [validator, tool1]
+```
+
+### Metadata Section
+
+```yaml
+metadata:
+ name: workflow-name # Required: Unique workflow identifier
+ version: 1.0.0 # Optional: Semantic version
+ description: "Description" # Optional: Human-readable description
+ tags: ["tag1", "tag2"] # Optional: Classification tags
+```
+
+### Arium Section
+
+#### Memory Configuration
+
+Memory is **not configured in YAML**. Instead, pass it as a parameter to `from_yaml()`:
+
+```python
+from flo_ai.arium.memory import MessageMemory, CustomMemory
+
+# Use default MessageMemory
+builder = AriumBuilder.from_yaml(yaml_str=config)
+
+# Or pass custom memory
+custom_memory = MessageMemory() # or your custom memory implementation
+builder = AriumBuilder.from_yaml(yaml_str=config, memory=custom_memory)
+```
+
+**Why memory is handled as a parameter:**
+- โ
**Cleaner YAML**: Focuses on workflow structure, not implementation details
+- โ
**Runtime Flexibility**: Same workflow can use different memory implementations
+- โ
**Better Separation**: Memory is an execution concern, not a workflow definition concern
+- โ
**Easier Testing**: Mock different memory types without changing YAML
+
+#### LLM Router Configuration
+
+LLM routers can be defined directly in YAML, eliminating the need to create router functions programmatically:
+
+```yaml
+routers:
+ - name: my_smart_router
+ type: smart
+ routing_options:
+ technical_writer: "Handle technical documentation and tutorials"
+ creative_writer: "Handle creative writing and storytelling"
+ marketing_writer: "Handle marketing copy and sales content"
+ model:
+ provider: openai
+ name: gpt-4o-mini
+ settings:
+ temperature: 0.3
+ fallback_strategy: first
+
+ - name: task_classifier
+ type: task_classifier
+ task_categories:
+ math_solver:
+ description: "Mathematical calculations and problem solving"
+ keywords: ["calculate", "solve", "equation", "math"]
+ examples: ["Calculate 2+2", "Solve x^2 + 5x + 6 = 0"]
+ code_helper:
+ description: "Programming and code assistance"
+ keywords: ["code", "program", "debug", "function"]
+ examples: ["Write a Python function", "Debug this code"]
+ model:
+ provider: openai
+ name: gpt-4o-mini
+ settings:
+ temperature: 0.2
+
+ - name: conversation_router
+ type: conversation_analysis
+ routing_logic:
+ reviewer: "Route to reviewer for quality check"
+ finalizer: "Route to finalizer for completion"
+ model:
+ provider: openai
+ name: gpt-4o-mini
+ settings:
+ temperature: 0.1
+ analysis_depth: 2
+```
+
+**Router Types:**
+- **`smart`**: General-purpose routing based on content analysis
+- **`task_classifier`**: Routes based on task categorization with keywords and examples
+- **`conversation_analysis`**: Routes based on conversation context analysis
+
+**Key benefits of YAML LLM routers:**
+- โ
**Declarative Configuration**: No code needed to create routers
+- โ
**Easy Modification**: Change routing logic without code changes
+- โ
**Version Control**: Track router changes in YAML files
+- โ
**Automatic Creation**: Routers are built automatically from configuration
+- โ
**Model Flexibility**: Each router can use different LLM models/settings
+
+#### Agent Configuration
+
+You can define agents in four ways:
+
+**1. Reference Pre-built Agents (New):**
+```yaml
+agents:
+ - name: content_analyst # Must exist in agents parameter
+ - name: summarizer # Must exist in agents parameter
+```
+
+**2. Direct Configuration (Recommended):**
+```yaml
+agents:
+ - name: my_agent
+ role: Assistant # optional
+ job: "You are a helpful assistant"
+ model:
+ provider: openai
+ name: gpt-4o-mini
+ base_url: "https://api.openai.com/v1" # optional
+ settings:
+ temperature: 0.7
+ max_retries: 3
+ reasoning_pattern: DIRECT # DIRECT, REACT, COT
+ tools: ["calculator", "web_search"] # optional
+ parser: # optional for structured output
+ name: MyParser
+ fields:
+ - name: result
+ type: str
+ description: "The result"
+```
+
+**3. Inline YAML Configuration:**
+```yaml
+agents:
+ - name: my_agent
+ yaml_config: |
+ agent:
+ name: my_agent
+ role: Assistant
+ job: "You are a helpful assistant"
+ model:
+ provider: openai
+ name: gpt-4o-mini
+ settings:
+ temperature: 0.7
+ max_retries: 3
+ reasoning_pattern: DIRECT
+```
+
+**4. External File Reference:**
+```yaml
+agents:
+ - name: my_agent
+ yaml_file: "configs/my_agent.yaml"
+```
+
+#### Tool Configuration
+
+```yaml
+tools:
+ - name: calculator
+ - name: text_processor
+ - name: web_search
+```
+
+**Note**: Tools must be provided as a dictionary when calling `from_yaml()` since they require actual Python functions.
+
+#### Comparison of Agent Configuration Methods
+
+**Reference Pre-built Agents (New):**
+- โ
Maximum reusability across workflows
+- โ
Programmatic agent building with YAML workflows
+- โ
Clean separation of agent definition and workflow
+- โ
Perfect for agent libraries
+- โ ๏ธ Requires separate agent building step
+
+**Direct Configuration (Recommended):**
+- โ
Clean, flat structure
+- โ
No nested YAML-in-YAML
+- โ
IDE syntax highlighting and validation
+- โ
Easier to read and maintain
+- โ
Supports all agent features directly
+
+**Inline YAML Configuration:**
+- โ ๏ธ Requires nested YAML string
+- โ ๏ธ Limited IDE support for nested content
+- โ
Maintains existing workflow compatibility
+- โ
Good for complex parser configurations
+
+**External File Reference:**
+- โ
Best for reusable agent definitions
+- โ
Supports modular architecture
+- โ
Version control friendly
+- โ
Keeps main workflow clean
+
+#### Workflow Configuration
+
+```yaml
+workflow:
+ start: agent1 # Starting node name
+ edges:
+ - from: agent1 # Source node
+ to: [tool1, agent2] # Target nodes
+ router: decision_router # Optional custom router
+ - from: tool1
+ to: [end] # 'end' is a special keyword
+ - from: agent2
+ to: [end]
+ end: [tool1, agent2] # End nodes
+```
+
+## Usage Examples
+
+### Simple Linear Workflow
+
+```python
+import asyncio
+from flo_ai.arium.builder import AriumBuilder
+
+yaml_config = """
+metadata:
+ name: simple-workflow
+
+arium:
+ agents:
+ - name: analyzer
+ role: Content Analyst
+ job: "Analyze the input content and extract key insights"
+ model:
+ provider: openai
+ name: gpt-4o-mini
+ settings:
+ temperature: 0.3
+
+ - name: summarizer
+ role: Summary Generator
+ job: "Create a concise summary based on the analysis"
+ model:
+ provider: openai
+ name: gpt-4o-mini
+ settings:
+ temperature: 0.2
+
+ workflow:
+ start: analyzer
+ edges:
+ - from: analyzer
+ to: [summarizer]
+ - from: summarizer
+ to: [end]
+ end: [summarizer]
+"""
+
+async def main():
+ builder = AriumBuilder.from_yaml(yaml_str=yaml_config)
+ result = await builder.build_and_run(["Your input text here"])
+ print(result)
+
+asyncio.run(main())
+```
+
+**Note**: To use custom memory, pass it as a parameter:
+```python
+from flo_ai.arium.memory import MessageMemory
+
+custom_memory = MessageMemory()
+builder = AriumBuilder.from_yaml(yaml_str=yaml_config, memory=custom_memory)
+```
+
+### Complex Workflow with Tools and Routing
+
+```python
+import asyncio
+from flo_ai.arium.builder import AriumBuilder
+from flo_ai.tool.base_tool import Tool
+from flo_ai.arium.memory import BaseMemory
+
+# Custom router function
+def smart_router(memory: BaseMemory) -> str:
+ content = str(memory.get()[-1]).lower()
+ if 'calculate' in content:
+ return 'calculator'
+ elif 'search' in content:
+ return 'web_search'
+ else:
+ return 'summarizer'
+
+# Create tools
+async def calculate(x: float, y: float, op: str) -> str:
+ ops = {'add': x + y, 'sub': x - y, 'mul': x * y, 'div': x / y}
+ return f"Result: {ops.get(op, 'Invalid operation')}"
+
+tools = {
+ 'calculator': Tool(
+ name='calculator',
+ description='Mathematical calculations',
+ function=calculate,
+ parameters={
+ 'x': {'type': 'number', 'description': 'First number'},
+ 'y': {'type': 'number', 'description': 'Second number'},
+ 'op': {'type': 'string', 'description': 'Operation (add, sub, mul, div)'}
+ }
+ )
+}
+
+routers = {'smart_router': smart_router}
+
+yaml_config = """
+metadata:
+ name: complex-workflow
+
+arium:
+ agents:
+ - name: dispatcher
+ role: Request Dispatcher
+ job: "Analyze input and decide next action"
+ model:
+ provider: openai
+ name: gpt-4o-mini
+ settings:
+ temperature: 0.1
+ reasoning_pattern: REACT
+
+ - name: summarizer
+ role: Final Summarizer
+ job: "Create final summary from all previous results"
+ model:
+ provider: openai
+ name: gpt-4o-mini
+ settings:
+ temperature: 0.3
+ tools: ["calculator"] # This agent can also use tools
+
+ tools:
+ - name: calculator
+
+ workflow:
+ start: dispatcher
+ edges:
+ - from: dispatcher
+ to: [calculator, summarizer]
+ router: smart_router
+ - from: calculator
+ to: [summarizer]
+ - from: summarizer
+ to: [end]
+ end: [summarizer]
+"""
+
+async def main():
+ builder = AriumBuilder.from_yaml(
+ yaml_str=yaml_config,
+ tools=tools,
+ routers=routers
+ )
+
+ result = await builder.build_and_run([
+ "Please calculate 15 + 25 and then summarize the result"
+ ])
+ print(result)
+
+asyncio.run(main())
+```
+
+### Mixed Configuration Approaches
+
+You can also mix different configuration approaches in the same workflow:
+
+```python
+yaml_config = """
+metadata:
+ name: mixed-workflow
+
+arium:
+ agents:
+ # Direct configuration
+ - name: input_processor
+ role: Input Processor
+ job: "Process and validate input data"
+ model:
+ provider: openai
+ name: gpt-4o-mini
+ settings:
+ temperature: 0.1
+
+ # Inline YAML configuration
+ - name: data_analyzer
+ yaml_config: |
+ agent:
+ name: data_analyzer
+ role: Data Analyst
+ job: "Perform detailed data analysis"
+ model:
+ provider: anthropic
+ name: claude-3-5-sonnet-20240620
+ settings:
+ temperature: 0.3
+ reasoning_pattern: COT
+
+ # External file reference
+ - name: report_generator
+ yaml_file: "agents/report_generator.yaml"
+
+ workflow:
+ start: input_processor
+ edges:
+ - from: input_processor
+ to: [data_analyzer]
+ - from: data_analyzer
+ to: [report_generator]
+ - from: report_generator
+ to: [end]
+ end: [report_generator]
+"""
+```
+
+### Using Pre-built Agents
+
+You can build agents separately using `AgentBuilder` and then reference them in your workflow YAML:
+
+```python
+import asyncio
+from flo_ai.arium.builder import AriumBuilder
+from flo_ai.builder.agent_builder import AgentBuilder
+from flo_ai.llm import OpenAI
+
+# Build agents separately from YAML files
+async def create_agents():
+ llm = OpenAI(model="gpt-4o-mini")
+
+ # Agent 1: Built from YAML file
+ content_analyst_yaml = """
+ agent:
+ name: content_analyst
+ role: Content Analyst
+ job: >
+ Analyze content and extract key insights, themes, and important information.
+ Provide structured analysis with clear categorization.
+ model:
+ provider: openai
+ name: gpt-4o-mini
+ settings:
+ temperature: 0.3
+ reasoning_pattern: COT
+ """
+
+ content_analyst = AgentBuilder.from_yaml(yaml_str=content_analyst_yaml).build()
+
+ # Agent 2: Built programmatically
+ summarizer = (AgentBuilder()
+ .with_name("summarizer")
+ .with_role("Summary Generator")
+ .with_prompt("Create concise, actionable summaries from analysis")
+ .with_llm(llm)
+ .with_reasoning(ReasoningPattern.DIRECT)
+ .build())
+
+ # Agent 3: Built from external file
+ # reporter = AgentBuilder.from_yaml(yaml_file="agents/reporter.yaml").build()
+
+ return {
+ 'content_analyst': content_analyst,
+ 'summarizer': summarizer,
+ # 'reporter': reporter
+ }
+
+# Clean workflow YAML that only references agents
+workflow_yaml = """
+metadata:
+ name: content-analysis-workflow
+ version: 1.0.0
+ description: "Content analysis workflow using pre-built agents"
+
+arium:
+ agents:
+ # Reference pre-built agents by name only
+ - name: content_analyst
+ - name: summarizer
+
+ # You can also mix with other configuration methods
+ - name: validator
+ role: Content Validator
+ job: "Validate the analysis and summary for accuracy"
+ model:
+ provider: openai
+ name: gpt-4o-mini
+ settings:
+ temperature: 0.1
+
+ workflow:
+ start: content_analyst
+ edges:
+ - from: content_analyst
+ to: [summarizer]
+ - from: summarizer
+ to: [validator]
+ - from: validator
+ to: [end]
+ end: [validator]
+"""
+
+async def main():
+ # Build agents separately
+ agents = await create_agents()
+
+ # Create workflow with pre-built agents
+ builder = AriumBuilder.from_yaml(
+ yaml_str=workflow_yaml,
+ agents=agents # Pass pre-built agents
+ )
+
+ result = await builder.build_and_run([
+ "Artificial Intelligence is revolutionizing industries across the globe. "
+ "From healthcare diagnostics to financial trading, AI systems are providing "
+ "unprecedented capabilities. However, challenges include ethical considerations, "
+ "data privacy, and the need for human oversight in critical decisions."
+ ])
+
+ print(result)
+
+asyncio.run(main())
+```
+
+## Method Signature
+
+```python
+@classmethod
+def from_yaml(
+ cls,
+ yaml_str: Optional[str] = None,
+ yaml_file: Optional[str] = None,
+ memory: Optional[BaseMemory] = None,
+ agents: Optional[Dict[str, Agent]] = None,
+ tools: Optional[Dict[str, Tool]] = None,
+ routers: Optional[Dict[str, Callable]] = None,
+ base_llm: Optional[BaseLLM] = None,
+) -> 'AriumBuilder':
+```
+
+### Parameters
+
+- **yaml_str**: YAML configuration as a string
+- **yaml_file**: Path to YAML configuration file
+- **memory**: Memory instance to use for the workflow (defaults to MessageMemory)
+- **agents**: Dictionary mapping agent names to pre-built Agent instances
+- **tools**: Dictionary mapping tool names to Tool instances
+- **routers**: Dictionary mapping router names to router functions
+- **base_llm**: Base LLM to use for all agents if not specified individually
+
+**Note**: Exactly one of `yaml_str` or `yaml_file` must be provided.
+
+## Router Functions
+
+Router functions determine the next node in the workflow based on the current memory state. They must have the following signature:
+
+```python
+def router_function(memory: BaseMemory) -> str:
+ """
+ Args:
+ memory: Current workflow memory containing conversation history
+
+ Returns:
+ str: Name of the next node to execute
+ """
+ # Your routing logic here
+ return "next_node_name"
+```
+
+### Router Return Types
+
+Router functions must return a `Literal` type that matches the possible target nodes:
+
+```python
+from typing import Literal
+
+def my_router(memory: BaseMemory) -> Literal["node1", "node2", "node3"]:
+ # Routing logic
+ if condition1:
+ return "node1"
+ elif condition2:
+ return "node2"
+ else:
+ return "node3"
+```
+
+## Best Practices
+
+### 1. Organize Agent Configurations
+
+Keep agent configurations in separate files for better maintainability:
+
+```
+project/
+โโโ workflows/
+โ โโโ main_workflow.yaml
+โ โโโ analysis_workflow.yaml
+โโโ agents/
+โ โโโ content_analyzer.yaml
+โ โโโ summarizer.yaml
+โ โโโ researcher.yaml
+โโโ tools/
+ โโโ tool_definitions.py
+```
+
+### 2. Use Descriptive Names
+
+Choose clear, descriptive names for agents, tools, and workflows:
+
+```yaml
+# Good
+agents:
+ - name: financial_analyst
+ - name: risk_assessor
+ - name: report_generator
+
+# Avoid
+agents:
+ - name: agent1
+ - name: agent2
+ - name: agent3
+```
+
+### 3. Document Your Workflows
+
+Include comprehensive metadata and comments:
+
+```yaml
+metadata:
+ name: financial-analysis-workflow
+ version: 2.1.0
+ description: >
+ Multi-stage financial analysis workflow that processes
+ financial data through risk assessment and generates
+ comprehensive reports with recommendations.
+ tags: ["finance", "analysis", "reporting"]
+ author: "Data Science Team"
+ created: "2024-01-15"
+ last_modified: "2024-02-01"
+```
+
+### 4. Version Your Configurations
+
+Use semantic versioning for your workflow configurations and maintain changelog documentation.
+
+### 5. Validate Configurations
+
+Always test your YAML configurations thoroughly:
+
+```python
+# Test configuration loading
+try:
+ builder = AriumBuilder.from_yaml(yaml_file="workflow.yaml")
+ print("Configuration loaded successfully")
+except Exception as e:
+ print(f"Configuration error: {e}")
+```
+
+## Error Handling
+
+Common errors and their solutions:
+
+### Missing Required Sections
+```
+ValueError: YAML must contain an "arium" section
+```
+**Solution**: Ensure your YAML has the required `arium` section.
+
+### Agent Configuration Errors
+```
+ValueError: Agent {name} must have either yaml_config or yaml_file
+```
+**Solution**: Each agent must specify either inline `yaml_config` or `yaml_file`.
+
+### Tool Not Found
+```
+ValueError: Tool {name} not found in provided tools dictionary
+```
+**Solution**: Ensure all referenced tools are provided in the `tools` parameter.
+
+### Router Not Found
+```
+ValueError: Router {name} not found in provided routers dictionary
+```
+**Solution**: Ensure all referenced routers are provided in the `routers` parameter.
+
+### Invalid Workflow Structure
+```
+ValueError: Workflow must specify a start node
+ValueError: Workflow must specify end nodes
+```
+**Solution**: Ensure `workflow` section has both `start` and `end` specifications.
+
+## Migration from Programmatic Builder
+
+To convert existing programmatic builder code to YAML:
+
+**Before:**
+```python
+builder = (AriumBuilder()
+ .add_agent(agent1)
+ .add_agent(agent2)
+ .start_with(agent1)
+ .connect(agent1, agent2)
+ .end_with(agent2))
+```
+
+**After:**
+```yaml
+arium:
+ agents:
+ - name: agent1
+ yaml_config: |
+ # agent1 configuration
+ - name: agent2
+ yaml_config: |
+ # agent2 configuration
+
+ workflow:
+ start: agent1
+ edges:
+ - from: agent1
+ to: [agent2]
+ - from: agent2
+ to: [end]
+ end: [agent2]
+```
+
+## Limitations
+
+1. **Dynamic Configuration**: Router functions and tools must be provided programmatically
+2. **Complex Logic**: Very complex routing logic may be better suited for programmatic definition
+3. **Runtime Modification**: YAML configurations are static; runtime modifications require programmatic approach
+4. **Memory Types**: Memory selection is done at runtime via parameter, not in YAML
+
+## Future Enhancements
+
+Planned features for future versions:
+
+- Plugin system for router functions
+- YAML-based tool definitions using function references
+- Configuration validation schemas
+- Hot-reloading of configurations
+- Workflow debugging and visualization tools
\ No newline at end of file
diff --git a/flo_ai/flo_ai/arium/examples.py b/flo_ai/examples/arium_examples.py
similarity index 100%
rename from flo_ai/flo_ai/arium/examples.py
rename to flo_ai/examples/arium_examples.py
diff --git a/flo_ai/examples/arium_yaml_example.py b/flo_ai/examples/arium_yaml_example.py
new file mode 100644
index 00000000..f2be9e4a
--- /dev/null
+++ b/flo_ai/examples/arium_yaml_example.py
@@ -0,0 +1,552 @@
+"""
+Example demonstrating YAML-based Arium workflow construction.
+
+This example shows how to define multi-agent workflows using YAML configuration
+instead of programmatic builder patterns.
+"""
+
+import asyncio
+from typing import Dict, Literal
+from flo_ai.arium.builder import AriumBuilder
+from flo_ai.tool.base_tool import Tool
+from flo_ai.llm import OpenAI
+from flo_ai.arium.memory import BaseMemory
+
+
+# Example YAML configuration for a simple linear workflow (using direct agent definition)
+SIMPLE_WORKFLOW_YAML = """
+metadata:
+ name: simple-analysis-workflow
+ version: 1.0.0
+ description: "A simple two-agent workflow for content analysis"
+
+arium:
+ agents:
+ - name: content_analyst
+ role: Content Analyst
+ job: >
+ You are a content analyst. When you receive input, analyze it and provide:
+ 1. A brief summary of the content
+ 2. The main topics covered
+ 3. Any insights or observations
+
+ Pass your analysis to the next agent for final processing.
+ model:
+ provider: openai
+ name: gpt-4o-mini
+ settings:
+ temperature: 0.3
+ max_retries: 3
+
+ - name: summary_generator
+ role: Summary Generator
+ job: >
+ You are a summary generator. You receive analysis from the content analyst.
+ Your job is to create a concise, well-structured final summary that includes:
+ 1. Key takeaways
+ 2. Actionable insights
+ 3. A clear conclusion
+
+ Make your response clear and professional.
+ model:
+ provider: openai
+ name: gpt-4o-mini
+ settings:
+ temperature: 0.2
+
+ workflow:
+ start: content_analyst
+ edges:
+ - from: content_analyst
+ to: [summary_generator]
+ - from: summary_generator
+ to: [end]
+ end: [summary_generator]
+"""
+
+
+# Example YAML configuration with tools and conditional routing (using direct agent definition)
+COMPLEX_WORKFLOW_YAML = """
+metadata:
+ name: research-workflow
+ version: 1.0.0
+ description: "A complex workflow with tools and conditional routing"
+
+arium:
+ agents:
+ - name: researcher
+ role: Research Analyst
+ job: >
+ You are a research analyst. Analyze the input and determine if you need to:
+ 1. Use the calculator tool for mathematical analysis
+ 2. Use the text processor for text analysis
+ 3. Or proceed directly to the final summarizer
+
+ Based on your analysis, decide the appropriate next step.
+ model:
+ provider: openai
+ name: gpt-4o-mini
+ settings:
+ reasoning_pattern: REACT
+
+ - name: math_specialist
+ role: Mathematics Specialist
+ job: >
+ You are a mathematics specialist. Use the calculator tool to perform
+ mathematical calculations and provide detailed analysis of the results.
+ model:
+ provider: openai
+ name: gpt-4o-mini
+ tools: ["calculator"]
+ settings:
+ reasoning_pattern: REACT
+
+ - name: text_specialist
+ role: Text Analysis Specialist
+ job: >
+ You are a text analysis specialist. Use the text processor tool to
+ analyze text content and provide insights about the text structure and content.
+ model:
+ provider: openai
+ name: gpt-4o-mini
+ tools: ["text_processor"]
+ settings:
+ reasoning_pattern: REACT
+
+ - name: final_summarizer
+ role: Final Summarizer
+ job: >
+ You are the final summarizer. Take all previous analysis and create
+ a comprehensive final report based on the work done by previous agents.
+ model:
+ provider: openai
+ name: gpt-4o-mini
+
+ workflow:
+ start: researcher
+ edges:
+ - from: researcher
+ to: [math_specialist, text_specialist, final_summarizer]
+ router: research_router
+ - from: math_specialist
+ to: [final_summarizer]
+ - from: text_specialist
+ to: [final_summarizer]
+ - from: final_summarizer
+ to: [end]
+ end: [final_summarizer]
+"""
+
+
+# Example showing mixed configuration approaches
+MIXED_CONFIG_YAML = """
+metadata:
+ name: mixed-configuration-example
+ version: 1.0.0
+ description: "Example showing different agent configuration methods"
+
+arium:
+ agents:
+ # Method 1: Direct configuration (new approach)
+ - name: input_validator
+ role: Input Validator
+ job: "Validate and preprocess the input data"
+ model:
+ provider: openai
+ name: gpt-4o-mini
+ settings:
+ temperature: 0.1
+ max_retries: 2
+
+ # Method 2: Inline YAML configuration (existing approach)
+ - name: data_processor
+ yaml_config: |
+ agent:
+ name: data_processor
+ role: Data Processor
+ job: >
+ Process the validated data using advanced analytics
+ and prepare it for final reporting.
+ model:
+ provider: openai
+ name: gpt-4o-mini
+ settings:
+ temperature: 0.3
+ reasoning_pattern: COT
+
+ # Method 3: Direct configuration with structured output
+ - name: report_generator
+ role: Report Generator
+ job: "Generate a structured final report"
+ model:
+ provider: openai
+ name: gpt-4o-mini
+ settings:
+ temperature: 0.2
+ parser:
+ name: ReportFormat
+ fields:
+ - name: summary
+ type: str
+ description: "Executive summary"
+ - name: findings
+ type: array
+ items:
+ type: str
+ description: "Key findings list"
+ - name: recommendations
+ type: array
+ items:
+ type: str
+ description: "Recommended actions"
+
+ workflow:
+ start: input_validator
+ edges:
+ - from: input_validator
+ to: [data_processor]
+ - from: data_processor
+ to: [report_generator]
+ - from: report_generator
+ to: [end]
+ end: [report_generator]
+"""
+
+
+# Custom router function for the complex workflow
+def research_router(
+ memory: BaseMemory,
+) -> Literal['math_specialist', 'text_specialist', 'final_summarizer']:
+ """
+ Custom router that decides the next step based on memory content.
+
+ This is a simple example - in practice, this could be more sophisticated,
+ potentially using LLM-based decision making.
+ """
+ # Get the latest content from memory
+ memory_content = memory.get()
+ latest_message = memory_content[-1] if memory_content else {}
+
+ # Simple text-based routing logic
+ content_text = str(latest_message).lower()
+
+ if any(
+ keyword in content_text
+ for keyword in ['calculate', 'math', 'number', 'compute']
+ ):
+ return 'math_specialist'
+ elif any(
+ keyword in content_text for keyword in ['text', 'analyze', 'process', 'parse']
+ ):
+ return 'text_specialist'
+ else:
+ return 'final_summarizer'
+
+
+async def create_example_tools() -> Dict[str, Tool]:
+ """Create example tools for the workflow."""
+
+ # Calculator tool
+ async def calculate(operation: str, x: float, y: float) -> str:
+ operations = {
+ 'add': lambda: x + y,
+ 'subtract': lambda: x - y,
+ 'multiply': lambda: x * y,
+ 'divide': lambda: x / y if y != 0 else 'Cannot divide by zero',
+ }
+ if operation not in operations:
+ return f'Unknown operation: {operation}'
+ result = operations[operation]()
+ return f'Calculation result: {x} {operation} {y} = {result}'
+
+ calculator_tool = Tool(
+ name='calculator',
+ description='Perform basic mathematical calculations',
+ function=calculate,
+ parameters={
+ 'operation': {
+ 'type': 'string',
+ 'description': 'The operation to perform (add, subtract, multiply, divide)',
+ },
+ 'x': {'type': 'number', 'description': 'First number'},
+ 'y': {'type': 'number', 'description': 'Second number'},
+ },
+ )
+
+ # Text processor tool
+ async def process_text(text: str, operation: str = 'analyze') -> str:
+ operations = {
+ 'analyze': lambda t: f'Text analysis: {len(t)} characters, {len(t.split())} words',
+ 'uppercase': lambda t: t.upper(),
+ 'lowercase': lambda t: t.lower(),
+ 'word_count': lambda t: f'Word count: {len(t.split())}',
+ }
+
+ if operation not in operations:
+ return f'Unknown text operation: {operation}'
+
+ result = operations[operation](text)
+ return f'Text processing result: {result}'
+
+ text_processor_tool = Tool(
+ name='text_processor',
+ description='Process and analyze text content',
+ function=process_text,
+ parameters={
+ 'text': {'type': 'string', 'description': 'Text to process'},
+ 'operation': {
+ 'type': 'string',
+ 'description': 'Operation to perform (analyze, uppercase, lowercase, word_count)',
+ 'required': False,
+ },
+ },
+ )
+
+ return {
+ 'calculator': calculator_tool,
+ 'text_processor': text_processor_tool,
+ }
+
+
+async def run_simple_example():
+ """Run the simple workflow example."""
+ print('=' * 60)
+ print('RUNNING SIMPLE WORKFLOW EXAMPLE')
+ print('=' * 60)
+
+ # Create builder from YAML
+ builder = AriumBuilder.from_yaml(yaml_str=SIMPLE_WORKFLOW_YAML)
+
+ # Run the workflow
+ result = await builder.build_and_run(
+ [
+ 'Machine learning is transforming healthcare by enabling predictive analytics, '
+ 'personalized treatment recommendations, and automated medical imaging analysis. '
+ 'However, challenges include data privacy concerns, the need for regulatory approval, '
+ 'and ensuring AI systems are transparent and unbiased in their decision-making.'
+ ]
+ )
+
+ print('Result:')
+ for i, message in enumerate(result):
+ print(f'{i+1}. {message}')
+
+ return result
+
+
+async def run_complex_example():
+ """Run the complex workflow example with tools and routing."""
+ print('\n' + '=' * 60)
+ print('RUNNING COMPLEX WORKFLOW EXAMPLE')
+ print('=' * 60)
+
+ # Create tools
+ tools = await create_example_tools()
+
+ # Create routers dictionary
+ routers = {'research_router': research_router}
+
+ # Create builder from YAML
+ builder = AriumBuilder.from_yaml(
+ yaml_str=COMPLEX_WORKFLOW_YAML, tools=tools, routers=routers
+ )
+
+ # Test with mathematical content
+ print('\nTesting with mathematical content:')
+ result1 = await builder.build_and_run(
+ ['Please calculate the sum of 25 and 17, then multiply the result by 3.']
+ )
+
+ print('Result:')
+ for i, message in enumerate(result1):
+ print(f'{i+1}. {message}')
+
+ # Reset and test with text content
+ print('\nTesting with text content:')
+ builder.reset()
+ builder = AriumBuilder.from_yaml(
+ yaml_str=COMPLEX_WORKFLOW_YAML, tools=tools, routers=routers
+ )
+
+ result2 = await builder.build_and_run(
+ [
+ "Please analyze this text and process it: 'The quick brown fox jumps over the lazy dog. "
+ "This sentence contains every letter of the alphabet at least once.'"
+ ]
+ )
+
+ print('Result:')
+ for i, message in enumerate(result2):
+ print(f'{i+1}. {message}')
+
+ return result1, result2
+
+
+async def run_mixed_config_example():
+ """Run the mixed configuration example showing different agent definition methods."""
+ print('\n' + '=' * 60)
+ print('RUNNING MIXED CONFIGURATION EXAMPLE')
+ print('=' * 60)
+
+ # Create builder from YAML (no tools needed for this example)
+ builder = AriumBuilder.from_yaml(yaml_str=MIXED_CONFIG_YAML)
+
+ # Run the workflow
+ result = await builder.build_and_run(
+ [
+ 'Please analyze this business report: Our Q3 revenue increased by 15% compared to Q2, '
+ 'driven primarily by strong performance in the software division. However, hardware sales '
+ 'declined by 8%. Customer satisfaction scores improved to 4.2/5.0. We recommend focusing '
+ 'on digital transformation initiatives and reconsidering the hardware product line.'
+ ]
+ )
+
+ print('Result:')
+ for i, message in enumerate(result):
+ print(f'{i+1}. {message}')
+
+ return result
+
+
+async def run_prebuilt_agents_example():
+ """Run example using pre-built agents with YAML workflow."""
+ print('\n' + '=' * 60)
+ print('RUNNING PRE-BUILT AGENTS EXAMPLE')
+ print('=' * 60)
+
+ # Build agents separately using AgentBuilder
+ from flo_ai.builder.agent_builder import AgentBuilder
+ from flo_ai.models.base_agent import ReasoningPattern
+
+ llm = OpenAI(model='gpt-4o-mini')
+
+ # Agent 1: Built from YAML string
+ content_analyst_yaml = """
+ agent:
+ name: content_analyst
+ role: Content Analyst
+ job: >
+ You are a content analyst. Analyze the input and extract:
+ 1. Main themes and topics
+ 2. Key insights and findings
+ 3. Important facts and figures
+
+ Provide a structured analysis for the next agent.
+ model:
+ provider: openai
+ name: gpt-4o-mini
+ settings:
+ temperature: 0.3
+ reasoning_pattern: COT
+ """
+
+ content_analyst = AgentBuilder.from_yaml(yaml_str=content_analyst_yaml).build()
+
+ # Agent 2: Built programmatically
+ summarizer = (
+ AgentBuilder()
+ .with_name('summarizer')
+ .with_role('Executive Summarizer')
+ .with_prompt(
+ 'Create a concise executive summary from the content analysis. Focus on actionable insights and key recommendations.'
+ )
+ .with_llm(llm)
+ .with_reasoning(ReasoningPattern.DIRECT)
+ .build()
+ )
+
+ # Dictionary of pre-built agents
+ prebuilt_agents = {'content_analyst': content_analyst, 'summarizer': summarizer}
+
+ # Clean workflow YAML that references the pre-built agents
+ workflow_yaml = """
+ metadata:
+ name: prebuilt-agents-workflow
+ version: 1.0.0
+ description: "Workflow using pre-built agents and inline definitions"
+
+ arium:
+ agents:
+ # Reference pre-built agents (must exist in agents parameter)
+ - name: content_analyst
+ - name: summarizer
+
+ # Mix with direct configuration
+ - name: quality_checker
+ role: Quality Assurance
+ job: "Review the analysis and summary for completeness and accuracy"
+ model:
+ provider: openai
+ name: gpt-4o-mini
+ settings:
+ temperature: 0.1
+ max_retries: 2
+
+ workflow:
+ start: content_analyst
+ edges:
+ - from: content_analyst
+ to: [summarizer]
+ - from: summarizer
+ to: [quality_checker]
+ - from: quality_checker
+ to: [end]
+ end: [quality_checker]
+ """
+
+ # Create builder with pre-built agents
+ builder = AriumBuilder.from_yaml(
+ yaml_str=workflow_yaml,
+ agents=prebuilt_agents, # Pass the pre-built agents
+ )
+
+ # Run the workflow
+ result = await builder.build_and_run(
+ [
+ 'The global renewable energy market reached $1.1 trillion in 2023, representing a 15% '
+ 'increase from the previous year. Solar energy dominated with 45% market share, followed '
+ 'by wind energy at 35%. Government incentives in Europe and Asia drove significant growth, '
+ 'while corporate sustainability commitments increased private sector investment. However, '
+ 'supply chain challenges and raw material costs remain key obstacles. Industry experts '
+ 'predict continued expansion, with the market expected to reach $1.8 trillion by 2030.'
+ ]
+ )
+
+ print('Result:')
+ for i, message in enumerate(result):
+ print(f'{i+1}. {message}')
+
+ return result
+
+
+async def main():
+ """Main function to run all examples."""
+ try:
+ # Run simple example
+ await run_simple_example()
+
+ # Run complex example
+ await run_complex_example()
+
+ # Run mixed configuration example
+ await run_mixed_config_example()
+
+ # Run pre-built agents example
+ await run_prebuilt_agents_example()
+
+ print('\n' + '=' * 60)
+ print('ALL EXAMPLES COMPLETED SUCCESSFULLY!')
+ print('=' * 60)
+ print('\n๐ Examples demonstrated:')
+ print(' โข Simple linear workflow with direct agent configuration')
+ print(' โข Complex workflow with tools and conditional routing')
+ print(' โข Mixed configuration approaches')
+ print(' โข Pre-built agents with YAML workflow (NEW!)')
+
+ except Exception as e:
+ print(f'Error running examples: {e}')
+ raise
+
+
+if __name__ == '__main__':
+ asyncio.run(main())
diff --git a/flo_ai/examples/llm_router_example.py b/flo_ai/examples/llm_router_example.py
new file mode 100644
index 00000000..738b64fb
--- /dev/null
+++ b/flo_ai/examples/llm_router_example.py
@@ -0,0 +1,457 @@
+#!/usr/bin/env python3
+"""
+Example demonstrating LLM-powered routing in Arium workflows.
+
+This example shows different ways to use intelligent routing where the LLM
+makes decisions about which agent should handle the next step in the workflow.
+"""
+
+import asyncio
+import os
+from typing import Literal
+
+from flo_ai.arium import AriumBuilder, create_llm_router, llm_router
+from flo_ai.builder.agent_builder import AgentBuilder
+from flo_ai.llm import OpenAI
+from flo_ai.models.base_agent import ReasoningPattern
+from flo_ai.arium.memory import BaseMemory
+
+
+# Set up environment
+# Make sure to set your OpenAI API key
+# export OPENAI_API_KEY=your_api_key_here
+
+
+async def example_1_smart_router():
+ """
+ Example 1: Using create_llm_router for a research workflow
+ """
+ print('=' * 60)
+ print('EXAMPLE 1: Smart LLM Router')
+ print('=' * 60)
+
+ llm = OpenAI(model='gpt-4o-mini')
+
+ # Create an intelligent router using the factory function
+ intelligent_router = create_llm_router(
+ 'smart',
+ routing_options={
+ 'researcher': 'Gather information and conduct research on topics',
+ 'analyst': 'Analyze data, perform calculations, and provide data-driven insights',
+ 'summarizer': 'Create summaries, conclusions, and final reports',
+ },
+ context_description='a research and analysis workflow',
+ llm=llm,
+ )
+
+ # Build workflow with YAML-like structure
+ workflow_yaml = """
+ metadata:
+ name: intelligent-research-workflow
+ version: 1.0.0
+ description: "Research workflow with LLM-powered routing"
+
+ arium:
+ agents:
+ - name: classifier
+ role: "Task Coordinator"
+ job: "Analyze incoming tasks and coordinate the workflow by deciding what type of work is needed."
+ model:
+ provider: openai
+ name: gpt-4o-mini
+ settings:
+ temperature: 0.3
+
+ - name: researcher
+ role: "Research Specialist"
+ job: "Gather comprehensive information and conduct thorough research."
+ model:
+ provider: openai
+ name: gpt-4o-mini
+
+ - name: analyst
+ role: "Data Analyst"
+ job: "Analyze data and provide insights."
+ model:
+ provider: openai
+ name: gpt-4o-mini
+
+ - name: summarizer
+ role: "Summary Specialist"
+ job: "Create final summaries and reports."
+ model:
+ provider: openai
+ name: gpt-4o-mini
+
+ workflow:
+ start: classifier
+ edges:
+ - from: classifier
+ to: [researcher, analyst, summarizer]
+ router: intelligent_router
+ - from: researcher
+ to: [summarizer]
+ - from: analyst
+ to: [summarizer]
+ end: [summarizer]
+ """
+
+ # Create the workflow
+ result = await (
+ AriumBuilder()
+ .from_yaml(
+ yaml_str=workflow_yaml, routers={'intelligent_router': intelligent_router}
+ )
+ .build_and_run(
+ [
+ 'I need to understand the current state of renewable energy adoption globally. '
+ 'Please provide statistics, trends, and market analysis for solar and wind energy.'
+ ]
+ )
+ )
+
+ print('Result:')
+ for i, message in enumerate(result):
+ print(f'{i+1}. {message}')
+
+ return result
+
+
+async def example_2_decorator_router():
+ """
+ Example 2: Using the @llm_router decorator
+ """
+ print('\n' + '=' * 60)
+ print('EXAMPLE 2: LLM Router Decorator')
+ print('=' * 60)
+
+ llm = OpenAI(model='gpt-4o-mini')
+
+ # Define a router using the decorator with better routing logic
+ @llm_router(
+ {
+ 'editor': 'Edit and improve the content if it needs refinement (max 2 editing rounds)',
+ 'reviewer': 'Move to final review and approval when content is ready',
+ },
+ llm=llm,
+ context_description='a content creation workflow that should progress to review after 1-2 editing rounds',
+ )
+ def content_workflow_router(memory: BaseMemory) -> Literal['editor', 'reviewer']:
+ """Smart router for content creation workflow"""
+ pass # Implementation provided by decorator
+
+ # Create agents
+ content_creator = (
+ AgentBuilder()
+ .with_name('content_creator')
+ .with_prompt(
+ 'You are a creative content creator. Write engaging, original content on any topic.'
+ )
+ .with_llm(llm)
+ .build()
+ )
+
+ editor = (
+ AgentBuilder()
+ .with_name('editor')
+ .with_prompt(
+ 'You are a professional editor. Improve content clarity, flow, and structure.'
+ )
+ .with_llm(llm)
+ .build()
+ )
+
+ reviewer = (
+ AgentBuilder()
+ .with_name('reviewer')
+ .with_prompt(
+ 'You are a quality reviewer. Provide final review and approval of content.'
+ )
+ .with_llm(llm)
+ .build()
+ )
+
+ # Build workflow with progressive structure (no loops back to previous stages)
+ result = await (
+ AriumBuilder()
+ .add_agents([content_creator, editor, reviewer])
+ .start_with(content_creator)
+ .add_edge(
+ content_creator, [editor, reviewer], content_workflow_router
+ ) # Content creator can only go forward
+ .add_edge(
+ editor, [reviewer]
+ ) # Editor always goes to reviewer (no router needed for single destination)
+ .end_with(reviewer)
+ .build_and_run(
+ [
+ 'Write a blog post about the benefits of remote work for software developers. '
+ 'Make it engaging and include practical tips.'
+ ]
+ )
+ )
+
+ print('Result:')
+ for i, message in enumerate(result):
+ print(f'{i+1}. {message}')
+
+ return result
+
+
+async def example_3_task_classifier_router():
+ """
+ Example 3: Using TaskClassifierRouter for specialized routing
+ """
+ print('\n' + '=' * 60)
+ print('EXAMPLE 3: Task Classifier Router')
+ print('=' * 60)
+
+ llm = OpenAI(model='gpt-4o-mini')
+
+ # Create a task classifier router
+ task_router = create_llm_router(
+ 'task_classifier',
+ task_categories={
+ 'math_specialist': {
+ 'description': 'Handle mathematical calculations, statistics, and numerical analysis',
+ 'keywords': [
+ 'calculate',
+ 'math',
+ 'number',
+ 'statistics',
+ 'compute',
+ 'formula',
+ ],
+ 'examples': [
+ 'Calculate the average of these numbers',
+ "What's the statistical significance?",
+ 'Compute the growth rate',
+ ],
+ },
+ 'text_specialist': {
+ 'description': 'Handle text analysis, writing, and language tasks',
+ 'keywords': [
+ 'write',
+ 'text',
+ 'analyze',
+ 'grammar',
+ 'language',
+ 'content',
+ ],
+ 'examples': [
+ 'Analyze this text for sentiment',
+ 'Write a summary of this document',
+ 'Check grammar and style',
+ ],
+ },
+ 'research_specialist': {
+ 'description': 'Handle research, fact-checking, and information gathering',
+ 'keywords': [
+ 'research',
+ 'find',
+ 'investigate',
+ 'facts',
+ 'information',
+ 'search',
+ ],
+ 'examples': [
+ 'Research the latest trends in AI',
+ 'Find information about this company',
+ 'Investigate the causes of climate change',
+ ],
+ },
+ },
+ llm=llm,
+ )
+
+ # Create specialized agents
+ math_specialist = (
+ AgentBuilder()
+ .with_name('math_specialist')
+ .with_prompt(
+ 'You are a mathematics expert. Handle all calculations, statistical analysis, and numerical computations.'
+ )
+ .with_llm(llm)
+ .with_reasoning(ReasoningPattern.COT)
+ .build()
+ )
+
+ text_specialist = (
+ AgentBuilder()
+ .with_name('text_specialist')
+ .with_prompt(
+ 'You are a text analysis expert. Handle writing, editing, and language-related tasks.'
+ )
+ .with_llm(llm)
+ .build()
+ )
+
+ research_specialist = (
+ AgentBuilder()
+ .with_name('research_specialist')
+ .with_prompt(
+ 'You are a research expert. Gather information, verify facts, and provide comprehensive research.'
+ )
+ .with_llm(llm)
+ .build()
+ )
+
+ # Test different types of tasks
+ test_tasks = [
+ 'Calculate the compound annual growth rate for an investment that grew from $1000 to $1500 over 3 years.',
+ 'Research the current market leaders in electric vehicle manufacturing and their market share.',
+ "Analyze this text for sentiment and writing quality: 'The new product launch was a tremendous success, exceeding all expectations and delighting customers worldwide.'",
+ ]
+
+ for i, task in enumerate(test_tasks):
+ print(f'\nTask {i+1}: {task}')
+
+ result = await (
+ AriumBuilder()
+ .add_agents([math_specialist, text_specialist, research_specialist])
+ .start_with(math_specialist) # Start node (will be routed immediately)
+ .add_edge(
+ math_specialist,
+ [math_specialist, text_specialist, research_specialist],
+ task_router,
+ )
+ .add_edge(
+ text_specialist,
+ [math_specialist, text_specialist, research_specialist],
+ task_router,
+ )
+ .add_edge(
+ research_specialist,
+ [math_specialist, text_specialist, research_specialist],
+ task_router,
+ )
+ .end_with(math_specialist)
+ .end_with(text_specialist)
+ .end_with(research_specialist)
+ .build_and_run([task])
+ )
+
+ print(f'Result: {result[-1]}')
+
+
+async def example_4_conversation_analysis_router():
+ """
+ Example 4: Using ConversationAnalysisRouter for flow-based routing
+ """
+ print('\n' + '=' * 60)
+ print('EXAMPLE 4: Conversation Analysis Router')
+ print('=' * 60)
+
+ llm = OpenAI(model='gpt-4o-mini')
+
+ # Create a conversation analysis router
+ flow_router = create_llm_router(
+ 'conversation_analysis',
+ routing_logic={
+ 'planner': 'Route here when starting a new task or when planning is needed',
+ 'executor': "Route here when there's a clear plan and work needs to be executed",
+ 'reviewer': 'Route here when work is complete and needs review or when ready to finalize',
+ },
+ analysis_depth=3, # Analyze last 3 messages
+ llm=llm,
+ )
+
+ # Create agents
+ planner = (
+ AgentBuilder()
+ .with_name('planner')
+ .with_prompt(
+ 'You are a strategic planner. Break down complex tasks into actionable plans and next steps.'
+ )
+ .with_llm(llm)
+ .build()
+ )
+
+ executor = (
+ AgentBuilder()
+ .with_name('executor')
+ .with_prompt(
+ 'You are a task executor. Follow plans and complete specific work items with attention to detail.'
+ )
+ .with_llm(llm)
+ .build()
+ )
+
+ reviewer = (
+ AgentBuilder()
+ .with_name('reviewer')
+ .with_prompt(
+ 'You are a quality reviewer. Review completed work and provide final assessment and recommendations.'
+ )
+ .with_llm(llm)
+ .build()
+ )
+
+ # Build and run workflow
+ result = await (
+ AriumBuilder()
+ .add_agents([planner, executor, reviewer])
+ .start_with(planner)
+ .add_edge(planner, [planner, executor, reviewer], flow_router)
+ .add_edge(executor, [planner, executor, reviewer], flow_router)
+ .end_with(reviewer)
+ .build_and_run(
+ [
+ 'I need to create a comprehensive marketing strategy for launching a new mobile app. '
+ 'The app is a personal finance tracker targeting millennials.'
+ ]
+ )
+ )
+
+ print('Result:')
+ for i, message in enumerate(result):
+ print(f'{i+1}. {message}')
+
+ return result
+
+
+async def main():
+ """Run all LLM router examples"""
+ try:
+ print('๐ LLM-Powered Router Examples')
+ print(
+ 'These examples demonstrate intelligent routing using Large Language Models'
+ )
+ print('to make dynamic decisions about workflow progression.\n')
+
+ # Run examples
+ await example_1_smart_router()
+ await example_2_decorator_router()
+ await example_3_task_classifier_router()
+ await example_4_conversation_analysis_router()
+
+ print('\n' + '=' * 60)
+ print('๐ ALL LLM ROUTER EXAMPLES COMPLETED!')
+ print('=' * 60)
+ print('\n๐ Examples demonstrated:')
+ print(' โข Smart router with factory function')
+ print(' โข Decorator-based router definition')
+ print(' โข Task classifier for specialized routing')
+ print(' โข Convenience functions for common patterns')
+ print(' โข Conversation analysis for flow-based routing')
+ print('\n๐ก Key benefits:')
+ print(' โข Dynamic routing based on content analysis')
+ print(' โข Context-aware workflow progression')
+ print(' โข Reduced need for complex conditional logic')
+ print(' โข Self-adapting workflows')
+
+ except Exception as e:
+ print(f'โ Error running examples: {e}')
+ import traceback
+
+ traceback.print_exc()
+
+
+if __name__ == '__main__':
+ # Make sure you have set your OpenAI API key
+ if not os.getenv('OPENAI_API_KEY'):
+ print('โ ๏ธ Warning: OPENAI_API_KEY environment variable not set!')
+ print(' Set it with: export OPENAI_API_KEY=your_api_key_here')
+ print(' Some examples may fail without a valid API key.\n')
+
+ asyncio.run(main())
diff --git a/flo_ai/examples/usage.py b/flo_ai/examples/tool_usage.py
similarity index 100%
rename from flo_ai/examples/usage.py
rename to flo_ai/examples/tool_usage.py
diff --git a/flo_ai/examples/yaml_llm_router_example.py b/flo_ai/examples/yaml_llm_router_example.py
new file mode 100644
index 00000000..9954a2ec
--- /dev/null
+++ b/flo_ai/examples/yaml_llm_router_example.py
@@ -0,0 +1,372 @@
+"""
+Example demonstrating LLM routers in YAML-based Arium workflows.
+
+This example shows how to define LLM routers directly in YAML configuration
+without needing to create them programmatically.
+"""
+
+import asyncio
+from flo_ai.arium.builder import AriumBuilder
+from flo_ai.llm import OpenAI
+
+# Example YAML configuration with LLM routers
+CONTENT_WORKFLOW_YAML = """
+metadata:
+ name: content-workflow-with-llm-routers
+ version: 1.0.0
+ description: "Content creation workflow with intelligent LLM-based routing"
+
+arium:
+ agents:
+ - name: content_creator
+ role: Content Creator
+ job: >
+ You are a content creator. Analyze the input request and create initial content.
+ Consider the type of content requested (technical, creative, marketing, etc.)
+ and create appropriate base content that can be further processed.
+ model:
+ provider: openai
+ name: gpt-4o-mini
+ settings:
+ temperature: 0.7
+
+ - name: technical_writer
+ role: Technical Writer
+ job: >
+ You are a technical writer specializing in documentation, tutorials, and
+ technical content. Refine the content to be clear, accurate, and technically sound.
+ Focus on structure, clarity, and technical accuracy.
+ model:
+ provider: openai
+ name: gpt-4o-mini
+ settings:
+ temperature: 0.3
+
+ - name: creative_writer
+ role: Creative Writer
+ job: >
+ You are a creative writer specializing in engaging, storytelling, and
+ creative content. Enhance the content with creativity, compelling narratives,
+ and emotional engagement while maintaining the core message.
+ model:
+ provider: openai
+ name: gpt-4o-mini
+ settings:
+ temperature: 0.8
+
+ - name: marketing_writer
+ role: Marketing Writer
+ job: >
+ You are a marketing writer specializing in persuasive, conversion-focused content.
+ Optimize the content for engagement, conversion, and brand alignment.
+ Focus on clear value propositions and calls to action.
+ model:
+ provider: openai
+ name: gpt-4o-mini
+ settings:
+ temperature: 0.6
+
+ - name: editor
+ role: Editor
+ job: >
+ You are an editor responsible for final review and polishing.
+ Review the content for grammar, style, coherence, and overall quality.
+ Make final improvements and ensure the content meets high standards.
+ model:
+ provider: openai
+ name: gpt-4o-mini
+ settings:
+ temperature: 0.2
+
+ # LLM Router definitions
+ routers:
+ - name: content_type_router
+ type: smart
+ routing_options:
+ technical_writer: "Technical content, documentation, tutorials, how-to guides, API docs, technical explanations"
+ creative_writer: "Creative writing, storytelling, fiction, poetry, creative marketing, brand narratives"
+ marketing_writer: "Marketing content, sales copy, landing pages, email campaigns, product descriptions, ad copy"
+ model:
+ provider: openai
+ name: gpt-4o-mini
+ settings:
+ temperature: 0.3
+ fallback_strategy: first
+
+ - name: quality_router
+ type: conversation_analysis
+ routing_logic:
+ editor: "Content that needs final editing and quality review"
+ model:
+ provider: openai
+ name: gpt-4o-mini
+ settings:
+ temperature: 0.1
+ analysis_depth: 2
+
+ workflow:
+ start: content_creator
+ edges:
+ - from: content_creator
+ to: [technical_writer, creative_writer, marketing_writer]
+ router: content_type_router
+ - from: technical_writer
+ to: [editor]
+ router: quality_router
+ - from: creative_writer
+ to: [editor]
+ router: quality_router
+ - from: marketing_writer
+ to: [editor]
+ router: quality_router
+ - from: editor
+ to: [end]
+ end: [editor]
+"""
+
+# Example with task classifier router
+TASK_CLASSIFICATION_YAML = """
+metadata:
+ name: task-classification-workflow
+ version: 1.0.0
+ description: "Multi-purpose workflow with task classification routing"
+
+arium:
+ agents:
+ - name: task_analyzer
+ role: Task Analyzer
+ job: >
+ You are a task analyzer. Examine the incoming request and identify
+ what type of task it represents. Provide initial analysis and context.
+ model:
+ provider: openai
+ name: gpt-4o-mini
+ settings:
+ temperature: 0.3
+
+ - name: math_solver
+ role: Math Solver
+ job: >
+ You are a mathematics expert. Solve mathematical problems, equations,
+ calculations, and provide step-by-step explanations.
+ model:
+ provider: openai
+ name: gpt-4o-mini
+ settings:
+ temperature: 0.1
+
+ - name: code_helper
+ role: Code Assistant
+ job: >
+ You are a programming expert. Help with coding problems, debugging,
+ code review, and programming best practices.
+ model:
+ provider: openai
+ name: gpt-4o-mini
+ settings:
+ temperature: 0.2
+
+ - name: general_assistant
+ role: General Assistant
+ job: >
+ You are a general-purpose assistant. Handle various questions,
+ provide information, and assist with general tasks.
+ model:
+ provider: openai
+ name: gpt-4o-mini
+ settings:
+ temperature: 0.5
+
+ routers:
+ - name: task_classifier
+ type: task_classifier
+ task_categories:
+ math_solver:
+ description: "Mathematical calculations, equations, and problem solving"
+ keywords: ["calculate", "solve", "equation", "math", "formula", "arithmetic"]
+ examples: ["Calculate 2+2", "Solve x^2 + 5x + 6 = 0", "What is the derivative of x^2?"]
+ code_helper:
+ description: "Programming, code review, debugging, and software development"
+ keywords: ["code", "program", "debug", "function", "class", "algorithm", "python", "javascript"]
+ examples: ["Write a Python function", "Debug this code", "Explain this algorithm"]
+ general_assistant:
+ description: "General questions, information requests, and other tasks"
+ keywords: ["what", "how", "explain", "tell me", "help", "information"]
+ examples: ["What is the weather like?", "Explain quantum physics", "Help me plan a trip"]
+ model:
+ provider: openai
+ name: gpt-4o-mini
+ settings:
+ temperature: 0.2
+
+ workflow:
+ start: task_analyzer
+ edges:
+ - from: task_analyzer
+ to: [math_solver, code_helper, general_assistant]
+ router: task_classifier
+ - from: math_solver
+ to: [end]
+ - from: code_helper
+ to: [end]
+ - from: general_assistant
+ to: [end]
+ end: [math_solver, code_helper, general_assistant]
+"""
+
+
+async def run_content_workflow_example():
+ """Example 1: Content workflow with smart router"""
+ print('๐ EXAMPLE 1: Content Workflow with Smart LLM Router')
+ print('=' * 60)
+
+ # Create workflow from YAML - LLM routers are automatically created!
+ builder = AriumBuilder.from_yaml(
+ yaml_str=CONTENT_WORKFLOW_YAML,
+ base_llm=OpenAI(
+ model='gpt-4o-mini', api_key='dummy-key'
+ ), # Dummy key for example
+ )
+
+ # Build the workflow
+ builder.build()
+
+ # Test with different content types
+ test_inputs = [
+ 'Write a technical tutorial on how to use Docker containers',
+ 'Create a compelling story about a robot learning to love',
+ 'Write marketing copy for a new fitness app that helps people stay motivated',
+ ]
+
+ for i, test_input in enumerate(test_inputs, 1):
+ print(f'\n๐ Test {i}: {test_input}')
+ print('-' * 40)
+
+ try:
+ # In a real scenario, this would run the workflow
+ # For demo purposes, we'll just show the structure
+ print('โ
Workflow built successfully with LLM routers!')
+ print(" - Router 'content_type_router' will intelligently route to:")
+ print(' โข technical_writer (for technical content)')
+ print(' โข creative_writer (for creative content)')
+ print(' โข marketing_writer (for marketing content)')
+ print(
+ " - Router 'quality_router' will route everything to editor for final review"
+ )
+
+ except Exception as e:
+ print(f'โ Error: {e}')
+
+
+async def run_task_classification_example():
+ """Example 2: Task classification workflow"""
+ print('\n\n๐ฏ EXAMPLE 2: Task Classification with LLM Router')
+ print('=' * 60)
+
+ # Create workflow from YAML
+ builder = AriumBuilder.from_yaml(
+ yaml_str=TASK_CLASSIFICATION_YAML,
+ base_llm=OpenAI(
+ model='gpt-4o-mini', api_key='dummy-key'
+ ), # Dummy key for example
+ )
+
+ # Build the workflow
+ builder.build()
+
+ # Test with different task types
+ test_inputs = [
+ 'Calculate the area of a circle with radius 5',
+ 'Write a Python function to sort a list',
+ 'What is the capital of France?',
+ ]
+
+ for i, test_input in enumerate(test_inputs, 1):
+ print(f'\n๐ฏ Test {i}: {test_input}')
+ print('-' * 40)
+
+ try:
+ print('โ
Workflow built successfully with task classifier!')
+ print(
+ " - Router 'task_classifier' will intelligently classify and route to:"
+ )
+ print(' โข math_solver (for mathematical problems)')
+ print(' โข code_helper (for programming tasks)')
+ print(' โข general_assistant (for general questions)')
+
+ except Exception as e:
+ print(f'โ Error: {e}')
+
+
+def demonstrate_yaml_schema():
+ """Show the YAML schema for LLM routers"""
+ print('\n\n๐ LLM Router YAML Schema')
+ print('=' * 60)
+
+ schema_example = """
+# LLM Router definitions in YAML
+routers:
+ - name: my_router_name
+ type: smart # Options: smart, task_classifier, conversation_analysis
+
+ # For smart routers:
+ routing_options:
+ option1: "Description of when to route to option1"
+ option2: "Description of when to route to option2"
+
+ # For task classifier routers:
+ task_categories:
+ category1:
+ description: "What this category handles"
+ keywords: ["keyword1", "keyword2"]
+ examples: ["Example 1", "Example 2"]
+
+ # For conversation analysis routers:
+ routing_logic:
+ target1: "Logic for routing to target1"
+ target2: "Logic for routing to target2"
+
+ # Model configuration (optional - uses base_llm if not specified)
+ model:
+ provider: openai # openai, anthropic, gemini, ollama
+ name: gpt-4o-mini
+ base_url: "https://api.openai.com/v1" # optional
+
+ # Router settings (optional)
+ settings:
+ temperature: 0.3
+ fallback_strategy: first # first, last, random
+ analysis_depth: 2 # for conversation_analysis type
+"""
+
+ print(schema_example)
+
+
+async def main():
+ """Run all examples"""
+ print('๐ YAML LLM Router Examples')
+ print('=' * 80)
+ print('This example demonstrates how to define LLM routers directly in YAML!')
+ print('No need to create router functions programmatically anymore! ๐')
+
+ # Show the schema first
+ demonstrate_yaml_schema()
+
+ # Run examples
+ await run_content_workflow_example()
+ await run_task_classification_example()
+
+ print('\n\n๐ All examples completed!')
+ print('=' * 80)
+ print('Key benefits of YAML LLM routers:')
+ print('โ
Declarative configuration - no code needed')
+ print('โ
Easy to modify and version control')
+ print('โ
Clear separation of workflow logic and implementation')
+ print('โ
Automatic LLM router creation from config')
+ print(
+ 'โ
Support for all router types: smart, task_classifier, conversation_analysis'
+ )
+
+
+if __name__ == '__main__':
+ asyncio.run(main())
diff --git a/flo_ai/flo_ai/arium/README.md b/flo_ai/flo_ai/arium/README.md
index e0d4b06e..6de5fa3e 100644
--- a/flo_ai/flo_ai/arium/README.md
+++ b/flo_ai/flo_ai/arium/README.md
@@ -24,6 +24,7 @@ result = await (AriumBuilder()
- **Default Memory**: Uses `MessageMemory` by default if none provided
- **Easy Connections**: Simple `connect()` method for linear workflows
- **Flexible Routing**: Full support for custom router functions
+- **๐ง LLM-Powered Routing**: Intelligent routing using Large Language Models
- **Visualization**: Built-in graph visualization support
- **Reusable Workflows**: Build once, run multiple times
@@ -166,6 +167,43 @@ The builder includes comprehensive validation:
- Validates router function signatures
- Checks for orphaned nodes
+## ๐ง LLM-Powered Routing
+
+Arium now supports intelligent routing using Large Language Models! Instead of writing complex routing logic, let an LLM analyze the conversation and decide which agent should handle the next step.
+
+### Quick Example
+
+```python
+from flo_ai.arium import create_llm_router
+
+# Create an intelligent router
+smart_router = create_llm_router(
+ "smart",
+ routing_options={
+ "researcher": "Gather information and conduct research",
+ "analyst": "Analyze data and provide insights",
+ "writer": "Create summaries and reports"
+ }
+)
+
+# Use in your workflow
+result = await (AriumBuilder()
+ .add_agents([researcher, analyst, writer])
+ .start_with(researcher)
+ .add_edge(researcher, [analyst, writer], smart_router)
+ .end_with(writer)
+ .build_and_run(["Research AI trends and create a report"]))
+```
+
+### Router Types
+
+- **SmartRouter**: General-purpose intelligent routing
+- **TaskClassifierRouter**: Classify tasks and route to specialists
+- **ConversationAnalysisRouter**: Analyze conversation flow for routing
+
+**๐ For detailed LLM routing documentation, see [README_LLM_Router.md](README_LLM_Router.md)**
+
## Examples
-See `examples.py` for complete working examples of different workflow patterns.
\ No newline at end of file
+See `examples.py` for complete working examples of different workflow patterns.
+See `examples/llm_router_example.py` for comprehensive LLM routing examples.
\ No newline at end of file
diff --git a/flo_ai/flo_ai/arium/README_LLM_Router.md b/flo_ai/flo_ai/arium/README_LLM_Router.md
new file mode 100644
index 00000000..0cc64767
--- /dev/null
+++ b/flo_ai/flo_ai/arium/README_LLM_Router.md
@@ -0,0 +1,336 @@
+# LLM-Powered Routing for Arium Workflows
+
+This document describes the LLM-powered routing functionality that enables intelligent, context-aware routing decisions in Arium workflows.
+
+## Overview
+
+Traditional routers require pre-written logic to determine the next agent in a workflow. LLM-powered routers use Large Language Models to analyze conversation context and make intelligent routing decisions dynamically.
+
+## Key Benefits
+
+- ๐ง **Intelligent Decision Making**: Uses LLMs to understand context and intent
+- ๐ **Dynamic Routing**: No need to pre-program all routing scenarios
+- ๐ **Context Awareness**: Considers full conversation history when routing
+- ๐ ๏ธ **Easy Integration**: Drop-in replacement for traditional routers
+- ๐ฏ **Specialized Patterns**: Pre-built routers for common workflows
+
+## Router Types
+
+### 1. SmartRouter
+General-purpose router that analyzes tasks and routes to appropriate agents.
+
+```python
+from flo_ai.arium import create_llm_router
+
+# Create a smart router
+router = create_llm_router(
+ "smart",
+ routing_options={
+ "researcher": "Gather information and conduct research",
+ "analyst": "Analyze data and perform calculations",
+ "writer": "Create reports and summaries"
+ },
+ context_description="a research workflow"
+)
+```
+
+### 2. TaskClassifierRouter
+Specialized router that classifies tasks based on keywords and examples.
+
+```python
+router = create_llm_router(
+ "task_classifier",
+ task_categories={
+ "math_specialist": {
+ "description": "Handle mathematical calculations and analysis",
+ "keywords": ["calculate", "math", "statistics", "compute"],
+ "examples": ["Calculate the average", "What's the growth rate?"]
+ },
+ "text_specialist": {
+ "description": "Handle text analysis and writing tasks",
+ "keywords": ["write", "analyze", "grammar", "content"],
+ "examples": ["Write a summary", "Analyze sentiment"]
+ }
+ }
+)
+```
+
+### 3. ConversationAnalysisRouter
+Router that analyzes conversation flow and determines next steps.
+
+```python
+router = create_llm_router(
+ "conversation_analysis",
+ routing_logic={
+ "planner": "Route here when planning is needed",
+ "executor": "Route here when work needs to be executed",
+ "reviewer": "Route here when work needs review"
+ },
+ analysis_depth=3 # Analyze last 3 messages
+)
+```
+
+## Usage Patterns
+
+### Factory Function Pattern
+
+The simplest way to create LLM routers:
+
+```python
+from flo_ai.arium import create_llm_router, AriumBuilder
+
+# Create router
+intelligent_router = create_llm_router(
+ "smart",
+ routing_options={
+ "agent1": "Description of agent1's role",
+ "agent2": "Description of agent2's role"
+ }
+)
+
+# Use in workflow
+result = await (
+ AriumBuilder()
+ .add_agents([agent1, agent2])
+ .start_with(agent1)
+ .add_edge(agent1, [agent1, agent2], intelligent_router)
+ .end_with(agent2)
+ .build_and_run(inputs)
+)
+```
+
+### Decorator Pattern
+
+For clean, declarative router definitions:
+
+```python
+from flo_ai.arium import llm_router
+from typing import Literal
+
+@llm_router({
+ "researcher": "Conduct research and gather information",
+ "analyst": "Analyze data and provide insights"
+})
+def my_router(memory: BaseMemory) -> Literal["researcher", "analyst"]:
+ """Smart router for research workflow"""
+ pass # Implementation provided by decorator
+```
+
+### YAML Integration
+
+LLM routers work seamlessly with YAML workflows:
+
+```yaml
+arium:
+ workflow:
+ start: classifier
+ edges:
+ - from: classifier
+ to: [researcher, analyst, writer]
+ router: intelligent_router # LLM router provided in Python
+```
+
+```python
+# Provide router to YAML workflow
+result = await (
+ AriumBuilder()
+ .from_yaml(
+ yaml_str=workflow_yaml,
+ routers={"intelligent_router": my_llm_router}
+ )
+ .build_and_run(inputs)
+)
+```
+
+## Configuration Options
+
+### LLM Configuration
+
+```python
+from flo_ai.llm import OpenAI
+
+# Configure LLM for routing
+llm = OpenAI(model="gpt-4o", temperature=0.1)
+
+router = create_llm_router(
+ "smart",
+ routing_options=options,
+ llm=llm, # Custom LLM
+ temperature=0.1, # Lower temperature for more deterministic routing
+ max_retries=3, # Retry failed LLM calls
+ fallback_strategy="first" # Fallback when LLM fails
+)
+```
+
+### Fallback Strategies
+
+When LLM routing fails, these strategies are available:
+
+- `"first"`: Route to first option (default)
+- `"last"`: Route to last option
+- `"random"`: Route to random option
+
+## Best Practices
+
+### 1. Clear Option Descriptions
+
+Provide clear, specific descriptions for routing options:
+
+```python
+# Good
+routing_options = {
+ "data_analyst": "Analyze numerical data, create charts, and calculate statistics",
+ "text_analyst": "Analyze text content, extract insights, and summarize information"
+}
+
+# Avoid
+routing_options = {
+ "analyst1": "Handles analysis",
+ "analyst2": "Does other analysis"
+}
+```
+
+### 2. Context Descriptions
+
+Add context to help the LLM understand the workflow:
+
+```python
+router = create_llm_router(
+ "smart",
+ routing_options=options,
+ context_description="a financial analysis workflow for investment decisions"
+)
+```
+
+### 3. Task Categories with Examples
+
+For TaskClassifierRouter, include relevant keywords and examples:
+
+```python
+task_categories = {
+ "researcher": {
+ "description": "Information gathering and research tasks",
+ "keywords": ["research", "find", "investigate", "gather", "search"],
+ "examples": [
+ "Find information about market trends",
+ "Research competitor analysis",
+ "Gather data on customer preferences"
+ ]
+ }
+}
+```
+
+### 4. Error Handling
+
+Configure appropriate retry and fallback behavior:
+
+```python
+router = create_llm_router(
+ "smart",
+ routing_options=options,
+ max_retries=3, # Retry failed calls
+ fallback_strategy="first", # Clear fallback
+ temperature=0.1 # Deterministic responses
+)
+```
+
+## Advanced Usage
+
+### Custom Router Classes
+
+Create custom routers by extending `BaseLLMRouter`:
+
+```python
+from flo_ai.arium.llm_router import BaseLLMRouter
+
+class CustomRouter(BaseLLMRouter):
+ def get_routing_options(self) -> Dict[str, str]:
+ return {"agent1": "Description", "agent2": "Description"}
+
+ def get_routing_prompt(self, memory, options) -> str:
+ # Custom prompt logic
+ return "Your custom routing prompt"
+
+# Use custom router
+router_instance = CustomRouter()
+def custom_router_function(memory):
+ return asyncio.run(router_instance.route(memory))
+```
+
+### Multi-Stage Routing
+
+Combine multiple routers for complex workflows:
+
+```python
+# First stage: Task classification
+task_router = create_llm_router("task_classifier", ...)
+
+# Second stage: Specialist routing
+specialist_router = create_llm_router("smart", ...)
+
+# Use in multi-stage workflow
+workflow = (
+ AriumBuilder()
+ .add_agent(classifier)
+ .add_edge(classifier, [research_specialist, analysis_specialist], task_router)
+ .add_edge(research_specialist, [writer, reviewer], specialist_router)
+ # ... more stages
+)
+```
+
+## Troubleshooting
+
+### Common Issues
+
+1. **Router Returns Invalid Option**
+ - Check that routing options match return type annotations
+ - Ensure LLM has clear instructions
+ - Verify option descriptions are distinct
+
+2. **LLM Calls Failing**
+ - Check API key configuration
+ - Verify network connectivity
+ - Consider increasing `max_retries`
+
+3. **Inconsistent Routing**
+ - Lower `temperature` for more deterministic results
+ - Provide more specific option descriptions
+ - Add examples to task categories
+
+### Debugging
+
+Enable debug logging to see routing decisions:
+
+```python
+import logging
+logging.getLogger('flo_ai').setLevel(logging.DEBUG)
+
+# Router decisions will be logged
+router = create_llm_router(...)
+```
+
+## Examples
+
+See `examples/llm_router_example.py` for comprehensive usage examples demonstrating:
+
+- Smart router with factory functions
+- Decorator-based router definitions
+- Task classifier for specialized routing
+- Convenience functions for common patterns
+- Conversation analysis for flow-based routing
+
+## API Reference
+
+### Functions
+
+- `create_llm_router(router_type, **config)`: Factory function for creating LLM routers
+- `llm_router(routing_options, **kwargs)`: Decorator for creating routers
+
+### Classes
+
+- `BaseLLMRouter`: Abstract base class for LLM routers
+- `SmartRouter`: General-purpose intelligent router
+- `TaskClassifierRouter`: Task classification-based router
+- `ConversationAnalysisRouter`: Conversation flow analysis router
+
+For detailed API documentation, see the docstrings in `flo_ai/arium/llm_router.py`.
diff --git a/flo_ai/flo_ai/arium/__init__.py b/flo_ai/flo_ai/arium/__init__.py
index 4417d222..04f27bd5 100644
--- a/flo_ai/flo_ai/arium/__init__.py
+++ b/flo_ai/flo_ai/arium/__init__.py
@@ -3,6 +3,14 @@
from .builder import AriumBuilder, create_arium
from .memory import MessageMemory, BaseMemory
from .models import StartNode, EndNode, Edge
+from .llm_router import (
+ BaseLLMRouter,
+ SmartRouter,
+ TaskClassifierRouter,
+ ConversationAnalysisRouter,
+ create_llm_router,
+ llm_router,
+)
__all__ = [
'Arium',
@@ -14,4 +22,11 @@
'StartNode',
'EndNode',
'Edge',
+ # LLM Router functionality
+ 'BaseLLMRouter',
+ 'SmartRouter',
+ 'TaskClassifierRouter',
+ 'ConversationAnalysisRouter',
+ 'create_llm_router',
+ 'llm_router',
]
diff --git a/flo_ai/flo_ai/arium/arium.py b/flo_ai/flo_ai/arium/arium.py
index a1a20183..2e7679f8 100644
--- a/flo_ai/flo_ai/arium/arium.py
+++ b/flo_ai/flo_ai/arium/arium.py
@@ -12,6 +12,7 @@
validate_multi_agent_variables,
resolve_variables,
)
+import asyncio
class Arium(BaseArium):
@@ -53,8 +54,42 @@ async def _execute_graph(self, inputs: List[str | ImageMessage]):
current_node = self.nodes[self.start_node_name]
current_edge = self.edges[self.start_node_name]
+ # Loop prevention: track execution steps and node visits
+ max_iterations = 20 # Reasonable limit to prevent infinite loops
+ iteration_count = 0
+ node_visit_count = {} # Track how many times each node is visited
+ execution_path = [] # Track the path for debugging
+
logger.info(f'Executing graph from {current_node.name}')
- while current_node.name != self.end_node_name:
+ while current_node.name not in self.end_node_names:
+ # Check for iteration limit
+ iteration_count += 1
+ if iteration_count > max_iterations:
+ logger.error(
+ f"Maximum iterations ({max_iterations}) exceeded. Execution path: {' -> '.join(execution_path)}"
+ )
+ raise RuntimeError(
+ f'Workflow exceeded maximum iterations ({max_iterations}). Possible infinite loop detected.'
+ )
+
+ # Track node visits
+ node_visit_count[current_node.name] = (
+ node_visit_count.get(current_node.name, 0) + 1
+ )
+ execution_path.append(current_node.name)
+
+ # Check for excessive node visits (same node visited too many times)
+ if node_visit_count[current_node.name] > 3:
+ logger.error(
+ f"Node '{current_node.name}' visited {node_visit_count[current_node.name]} times. Execution path: {' -> '.join(execution_path)}"
+ )
+ raise RuntimeError(
+ f"Node '{current_node.name}' visited too many times ({node_visit_count[current_node.name]}). Possible infinite loop detected."
+ )
+
+ logger.info(
+ f'Executing node: {current_node.name} (iteration {iteration_count})'
+ )
# execute current node
result = await self._execute_node(current_node)
@@ -62,7 +97,28 @@ async def _execute_graph(self, inputs: List[str | ImageMessage]):
self._add_to_memory(result)
# find next node post current node
- next_node_name = current_edge.router_fn(memory=self.memory)
+ # Prepare execution context for router functions
+ execution_context = {
+ 'node_visit_count': node_visit_count,
+ 'execution_path': execution_path,
+ 'iteration_count': iteration_count,
+ 'current_node': current_node.name,
+ }
+
+ # Handle both sync and async router functions
+ # Try to call with execution context, fallback to memory only
+ try:
+ router_result = current_edge.router_fn(
+ memory=self.memory, execution_context=execution_context
+ )
+ except TypeError:
+ # Router function doesn't accept execution_context parameter
+ router_result = current_edge.router_fn(memory=self.memory)
+
+ if asyncio.iscoroutine(router_result):
+ next_node_name = await router_result
+ else:
+ next_node_name = router_result
# find next edge
# TODO: next_node_name might not be in self.edges if it's the end node. Handle this case
diff --git a/flo_ai/flo_ai/arium/base.py b/flo_ai/flo_ai/arium/base.py
index 5ce2c8a4..26943295 100644
--- a/flo_ai/flo_ai/arium/base.py
+++ b/flo_ai/flo_ai/arium/base.py
@@ -11,7 +11,7 @@
class BaseArium:
def __init__(self):
self.start_node_name = '__start__'
- self.end_node_name = '__end__'
+ self.end_node_names: set = set() # Support multiple end nodes
self.nodes: Dict[str, Agent | Tool | StartNode | EndNode] = dict()
self.edges: Dict[str, Edge] = dict()
@@ -28,13 +28,21 @@ def start_at(self, node: Agent | Tool | StartNode | EndNode):
)
def add_end_to(self, node: Agent | Tool | StartNode | EndNode):
+ # Create a unique end node name for this specific node
+ end_node_name = f'__end__{node.name}__'
end_node = EndNode()
- if end_node.name in self.nodes:
- raise ValueError(f'End node {end_node.name} already exists')
- self.nodes[end_node.name] = end_node
+ end_node.name = end_node_name
+
+ # Add this end node name to our set of possible end nodes
+ self.end_node_names.add(end_node_name)
+
+ # Only add the end node if it doesn't exist yet
+ if end_node_name not in self.nodes:
+ self.nodes[end_node_name] = end_node
+
self.edges[node.name] = Edge(
- router_fn=partial(default_router, to_node=end_node.name),
- to_nodes=[end_node.name],
+ router_fn=partial(default_router, to_node=end_node_name),
+ to_nodes=[end_node_name],
)
def _check_router_return_type(self, router: Callable) -> Optional[List]:
diff --git a/flo_ai/flo_ai/arium/builder.py b/flo_ai/flo_ai/arium/builder.py
index b343bcfa..fc0821fa 100644
--- a/flo_ai/flo_ai/arium/builder.py
+++ b/flo_ai/flo_ai/arium/builder.py
@@ -4,6 +4,10 @@
from flo_ai.models.agent import Agent
from flo_ai.tool.base_tool import Tool
from flo_ai.llm.base_llm import ImageMessage
+import yaml
+from flo_ai.builder.agent_builder import AgentBuilder
+from flo_ai.llm import BaseLLM
+from flo_ai.arium.llm_router import create_llm_router
class AriumBuilder:
@@ -158,6 +162,470 @@ def reset(self) -> 'AriumBuilder':
self._arium = None
return self
+ @classmethod
+ def from_yaml(
+ cls,
+ yaml_str: Optional[str] = None,
+ yaml_file: Optional[str] = None,
+ memory: Optional[BaseMemory] = None,
+ agents: Optional[Dict[str, Agent]] = None,
+ tools: Optional[Dict[str, Tool]] = None,
+ routers: Optional[Dict[str, Callable]] = None,
+ base_llm: Optional[BaseLLM] = None,
+ ) -> 'AriumBuilder':
+ """Create an AriumBuilder from a YAML configuration.
+
+ Args:
+ yaml_str: YAML string containing arium configuration
+ yaml_file: Path to YAML file containing arium configuration
+ memory: Memory instance to use for the workflow (defaults to MessageMemory)
+ agents: Dictionary mapping agent names to pre-built Agent instances
+ tools: Dictionary mapping tool names to Tool instances
+ routers: Dictionary mapping router names to router functions
+ base_llm: Base LLM to use for all agents if not specified in individual agent configs
+
+ Returns:
+ AriumBuilder: Configured builder instance
+
+ Example YAML structure:
+ metadata:
+ name: my-workflow
+ version: 1.0.0
+ description: "Example workflow"
+
+ arium:
+ agents:
+ # Method 1: Reference pre-built agents
+ - name: content_analyst # Must exist in agents parameter
+ - name: summarizer # Must exist in agents parameter
+
+ # Method 2: Direct agent definition
+ - name: validator
+ role: "Data Validator"
+ job: "You are a data validator"
+ model:
+ provider: openai
+ name: gpt-4o-mini
+ settings:
+ temperature: 0.1
+
+ # Method 3: Inline YAML configuration
+ - name: processor
+ yaml_config: |
+ agent:
+ name: processor
+ job: "You are a data processor"
+ model:
+ provider: openai
+ name: gpt-4o-mini
+
+ # Method 4: External file reference
+ - name: reporter
+ yaml_file: "path/to/reporter.yaml"
+
+ tools:
+ - name: tool1
+ - name: tool2
+
+ # LLM Router definitions (NEW)
+ routers:
+ - name: content_router
+ type: smart # smart, task_classifier, conversation_analysis
+ routing_options:
+ technical_writer: "Handle technical documentation tasks"
+ creative_writer: "Handle creative writing tasks"
+ editor: "Handle editing and review tasks"
+ model:
+ provider: openai
+ name: gpt-4o-mini
+ settings:
+ temperature: 0.3
+ fallback_strategy: first
+
+ workflow:
+ start: content_analyst
+ edges:
+ - from: content_analyst
+ to: [validator, summarizer]
+ router: content_router # References router defined above
+ - from: validator
+ to: [processor]
+ - from: summarizer
+ to: [reporter]
+ - from: processor
+ to: [end]
+ - from: reporter
+ to: [end]
+ end: [processor, reporter]
+ """
+ if yaml_str is None and yaml_file is None:
+ raise ValueError('Either yaml_str or yaml_file must be provided')
+
+ if yaml_str and yaml_file:
+ raise ValueError('Only one of yaml_str or yaml_file should be provided')
+
+ # Load YAML configuration
+ if yaml_str:
+ config = yaml.safe_load(yaml_str)
+ else:
+ with open(yaml_file, 'r') as f:
+ config = yaml.safe_load(f)
+
+ if 'arium' not in config:
+ raise ValueError('YAML must contain an "arium" section')
+
+ arium_config = config['arium']
+ builder = cls()
+
+ # Configure memory - use provided memory or default to MessageMemory
+ if memory is not None:
+ builder.with_memory(memory)
+ else:
+ builder.with_memory(MessageMemory())
+
+ # Process agents
+ agents_config = arium_config.get('agents', [])
+ agents_dict = {}
+
+ for agent_config in agents_config:
+ agent_name = agent_config['name']
+
+ # Method 1: Reference pre-built agent
+ if len(agent_config) == 1 and 'name' in agent_config:
+ # Only has name field, so it's a reference to a pre-built agent
+ if agents and agent_name in agents:
+ agent = agents[agent_name]
+ else:
+ raise ValueError(
+ f'Agent {agent_name} not found in provided agents dictionary. '
+ f'Available agents: {list(agents.keys()) if agents else []}. '
+ f'Either provide the agent in the agents parameter or add configuration fields.'
+ )
+
+ # Method 2: Direct agent definition
+ elif (
+ 'job' in agent_config
+ and 'yaml_config' not in agent_config
+ and 'yaml_file' not in agent_config
+ ):
+ agent = cls._create_agent_from_direct_config(
+ agent_config, base_llm, tools
+ )
+
+ # Method 3: Inline YAML config
+ elif 'yaml_config' in agent_config:
+ agent_builder = AgentBuilder.from_yaml(
+ yaml_str=agent_config['yaml_config'], base_llm=base_llm
+ )
+ agent = agent_builder.build()
+
+ # Method 4: External file reference
+ elif 'yaml_file' in agent_config:
+ agent_builder = AgentBuilder.from_yaml(
+ yaml_file=agent_config['yaml_file'], base_llm=base_llm
+ )
+ agent = agent_builder.build()
+
+ else:
+ raise ValueError(
+ f'Agent {agent_name} must have either:\n'
+ f' - Only a name (to reference pre-built agent),\n'
+ f' - Direct configuration (job field),\n'
+ f' - yaml_config, or\n'
+ f' - yaml_file'
+ )
+
+ agents_dict[agent_name] = agent
+ builder.add_agent(agent)
+
+ # Process tools
+ tools_config = arium_config.get('tools', [])
+ tools_dict = {}
+
+ for tool_config in tools_config:
+ tool_name = tool_config['name']
+
+ # Look up tool in provided tools dictionary
+ if tools and tool_name in tools:
+ tool = tools[tool_name]
+ tools_dict[tool_name] = tool
+ builder.add_tool(tool)
+ else:
+ raise ValueError(
+ f'Tool {tool_name} not found in provided tools dictionary. '
+ f'Available tools: {list(tools.keys()) if tools else []}'
+ )
+
+ # Process LLM routers (if defined in YAML)
+ routers_config = arium_config.get('routers', [])
+ yaml_routers = {} # Store routers created from YAML config
+
+ for router_config in routers_config:
+ router_name = router_config['name']
+ router_type = router_config.get('type', 'smart')
+
+ # Create LLM instance for router
+ router_llm = None
+ if 'model' in router_config:
+ router_llm = cls._create_llm_from_config(
+ router_config['model'], base_llm
+ )
+ else:
+ router_llm = base_llm # Use base LLM if no specific model configured
+
+ # Extract router-specific settings
+ settings = router_config.get('settings', {})
+
+ # Create router based on type
+ if router_type == 'smart':
+ routing_options = router_config.get('routing_options', {})
+ if not routing_options:
+ raise ValueError(
+ f'Smart router {router_name} must specify routing_options'
+ )
+
+ router_fn = create_llm_router(
+ router_type='smart',
+ routing_options=routing_options,
+ llm=router_llm,
+ **settings,
+ )
+
+ elif router_type == 'task_classifier':
+ task_categories = router_config.get('task_categories', {})
+ if not task_categories:
+ raise ValueError(
+ f'Task classifier router {router_name} must specify task_categories'
+ )
+
+ router_fn = create_llm_router(
+ router_type='task_classifier',
+ task_categories=task_categories,
+ llm=router_llm,
+ **settings,
+ )
+
+ elif router_type == 'conversation_analysis':
+ routing_logic = router_config.get('routing_logic', {})
+ if not routing_logic:
+ raise ValueError(
+ f'Conversation analysis router {router_name} must specify routing_logic'
+ )
+
+ router_fn = create_llm_router(
+ router_type='conversation_analysis',
+ routing_logic=routing_logic,
+ llm=router_llm,
+ **settings,
+ )
+ else:
+ raise ValueError(
+ f'Unknown router type: {router_type}. Supported types: smart, task_classifier, conversation_analysis'
+ )
+
+ yaml_routers[router_name] = router_fn
+
+ # Merge YAML routers with provided routers
+ all_routers = {}
+ if routers:
+ all_routers.update(routers)
+ all_routers.update(yaml_routers)
+
+ # Process workflow
+ workflow_config = arium_config.get('workflow', {})
+
+ # Set start node
+ start_node_name = workflow_config.get('start')
+ if not start_node_name:
+ raise ValueError('Workflow must specify a start node')
+
+ start_node = agents_dict.get(start_node_name) or tools_dict.get(start_node_name)
+ if not start_node:
+ raise ValueError(
+ f'Start node {start_node_name} not found in agents or tools'
+ )
+
+ builder.start_with(start_node)
+
+ # Process edges
+ edges_config = workflow_config.get('edges', [])
+
+ for edge_config in edges_config:
+ from_node_name = edge_config['from']
+ to_nodes_names = edge_config['to']
+ router_name = edge_config.get('router')
+
+ # Find from node
+ from_node = agents_dict.get(from_node_name) or tools_dict.get(
+ from_node_name
+ )
+ if not from_node:
+ raise ValueError(f'From node {from_node_name} not found')
+
+ # Find to nodes (handle special 'end' case)
+ to_nodes = []
+ for to_node_name in to_nodes_names:
+ if to_node_name == 'end':
+ # 'end' will be handled in end nodes processing
+ continue
+
+ to_node = agents_dict.get(to_node_name) or tools_dict.get(to_node_name)
+ if not to_node:
+ raise ValueError(f'To node {to_node_name} not found')
+ to_nodes.append(to_node)
+
+ # Find router function
+ router_fn = None
+ if router_name:
+ if all_routers and router_name in all_routers:
+ router_fn = all_routers[router_name]
+ else:
+ raise ValueError(
+ f'Router {router_name} not found. '
+ f'Available routers: {list(all_routers.keys()) if all_routers else []}'
+ )
+
+ # Add edge (only if there are actual to_nodes, not just 'end')
+ if to_nodes:
+ builder.add_edge(from_node, to_nodes, router_fn)
+
+ # Set end nodes
+ end_nodes_names = workflow_config.get('end', [])
+ if not end_nodes_names:
+ raise ValueError('Workflow must specify end nodes')
+
+ for end_node_name in end_nodes_names:
+ end_node = agents_dict.get(end_node_name) or tools_dict.get(end_node_name)
+ if not end_node:
+ raise ValueError(f'End node {end_node_name} not found')
+ builder.end_with(end_node)
+
+ return builder
+
+ @staticmethod
+ def _create_llm_from_config(
+ model_config: Dict[str, Any], base_llm: Optional[BaseLLM] = None
+ ) -> BaseLLM:
+ """Create an LLM instance from model configuration.
+
+ Args:
+ model_config: Dictionary containing model configuration
+ base_llm: Base LLM to use as fallback
+
+ Returns:
+ BaseLLM: Configured LLM instance
+ """
+ from flo_ai.llm import OpenAI, Anthropic, Gemini, OllamaLLM
+
+ provider = model_config.get('provider', 'openai').lower()
+ model_name = model_config.get('name')
+ base_url = model_config.get('base_url')
+
+ if not model_name:
+ raise ValueError('Model name must be specified in model configuration')
+
+ if provider == 'openai':
+ llm = OpenAI(model=model_name, base_url=base_url)
+ elif provider == 'anthropic':
+ llm = Anthropic(model=model_name, base_url=base_url)
+ elif provider == 'gemini':
+ llm = Gemini(model=model_name, base_url=base_url)
+ elif provider == 'ollama':
+ llm = OllamaLLM(model=model_name, base_url=base_url)
+ else:
+ raise ValueError(f'Unsupported model provider: {provider}')
+
+ return llm
+
+ @staticmethod
+ def _create_agent_from_direct_config(
+ agent_config: Dict[str, Any],
+ base_llm: Optional[BaseLLM] = None,
+ available_tools: Optional[Dict[str, Tool]] = None,
+ ) -> Agent:
+ """Create an Agent from direct YAML configuration.
+
+ Args:
+ agent_config: Dictionary containing agent configuration
+ base_llm: Base LLM to use if not specified in config
+ available_tools: Available tools dictionary for tool lookup
+
+ Returns:
+ Agent: Configured agent instance
+ """
+ from flo_ai.models.base_agent import ReasoningPattern
+ # from flo_ai.llm import OpenAI, Anthropic, Gemini, OllamaLLM
+
+ # Extract basic configuration
+ name = agent_config['name']
+ job = agent_config['job']
+ role = agent_config.get('role')
+
+ # Configure LLM
+ if 'model' in agent_config and base_llm is None:
+ llm = AriumBuilder._create_llm_from_config(agent_config['model'])
+ elif base_llm:
+ llm = base_llm
+ else:
+ raise ValueError(
+ f'Model must be specified for agent {name} or base_llm must be provided'
+ )
+
+ # Extract settings
+ settings = agent_config.get('settings', {})
+ temperature = settings.get('temperature')
+ max_retries = settings.get('max_retries', 3)
+ reasoning_pattern_str = settings.get('reasoning_pattern', 'DIRECT')
+
+ # Convert reasoning pattern string to enum
+ try:
+ reasoning_pattern = ReasoningPattern[reasoning_pattern_str.upper()]
+ except KeyError:
+ raise ValueError(f'Invalid reasoning pattern: {reasoning_pattern_str}')
+
+ # Set LLM temperature if specified
+ if temperature is not None:
+ llm.temperature = temperature
+
+ # Extract and resolve tools
+ agent_tools = []
+ tool_names = agent_config.get('tools', [])
+ if tool_names and available_tools:
+ for tool_name in tool_names:
+ if tool_name in available_tools:
+ agent_tools.append(available_tools[tool_name])
+ else:
+ raise ValueError(
+ f'Tool {tool_name} for agent {name} not found in available tools. '
+ f'Available: {list(available_tools.keys())}'
+ )
+
+ # Extract output schema if present
+ output_schema = agent_config.get('output_schema')
+
+ # Handle parser configuration if present
+ if 'parser' in agent_config:
+ from flo_ai.formatter.yaml_format_parser import FloYamlParser
+
+ # Convert agent_config to the format expected by FloYamlParser
+ parser_config = {'agent': {'parser': agent_config['parser']}}
+ parser = FloYamlParser.create(yaml_dict=parser_config)
+ output_schema = parser.get_format()
+
+ # Create and return the agent
+ agent = Agent(
+ name=name,
+ system_prompt=job,
+ llm=llm,
+ tools=agent_tools,
+ max_retries=max_retries,
+ reasoning_pattern=reasoning_pattern,
+ output_schema=output_schema,
+ role=role,
+ )
+
+ return agent
+
# Convenience function for creating a builder
def create_arium() -> AriumBuilder:
diff --git a/flo_ai/flo_ai/arium/llm_router.py b/flo_ai/flo_ai/arium/llm_router.py
new file mode 100644
index 00000000..900d2cf8
--- /dev/null
+++ b/flo_ai/flo_ai/arium/llm_router.py
@@ -0,0 +1,660 @@
+"""
+LLM-Powered Router Functions for Arium Workflows
+
+This module provides intelligent routing capabilities using Large Language Models
+to make dynamic routing decisions based on conversation context and history.
+"""
+
+from abc import ABC, abstractmethod
+from typing import Dict, Optional, Callable, Any, Union, get_args
+from functools import wraps
+from flo_ai.arium.memory import BaseMemory
+from flo_ai.llm.base_llm import BaseLLM
+from flo_ai.llm import OpenAI
+from flo_ai.utils.logger import logger
+
+
+class BaseLLMRouter(ABC):
+ """
+ Base class for LLM-powered routers that make intelligent routing decisions
+ based on conversation context and history.
+ """
+
+ def __init__(
+ self,
+ llm: Optional[BaseLLM] = None,
+ temperature: float = 0.1,
+ max_retries: int = 3,
+ fallback_strategy: str = 'first',
+ ):
+ """
+ Initialize the LLM router.
+
+ Args:
+ llm: The LLM instance to use for routing decisions. Defaults to GPT-4o-mini.
+ temperature: Temperature for LLM calls (lower = more deterministic)
+ max_retries: Maximum number of retries for LLM calls
+ fallback_strategy: Strategy when LLM fails ("first", "last", "random")
+ """
+ self.llm = llm or OpenAI(model='gpt-4o-mini', temperature=temperature)
+ self.temperature = temperature
+ self.max_retries = max_retries
+ self.fallback_strategy = fallback_strategy
+
+ @abstractmethod
+ def get_routing_options(self) -> Dict[str, str]:
+ """
+ Return a dictionary mapping route names to their descriptions.
+
+ Returns:
+ Dict[str, str]: Mapping of route names to descriptions
+ """
+ pass
+
+ @abstractmethod
+ def get_routing_prompt(
+ self,
+ memory: BaseMemory,
+ options: Dict[str, str],
+ execution_context: dict = None,
+ ) -> str:
+ """
+ Generate the prompt for the LLM to make routing decisions.
+
+ Args:
+ memory: The conversation memory
+ options: Available routing options with descriptions
+
+ Returns:
+ str: The prompt for the LLM
+ """
+ pass
+
+ def get_fallback_route(self, options: Dict[str, str]) -> str:
+ """
+ Get fallback route when LLM fails.
+
+ Args:
+ options: Available routing options
+
+ Returns:
+ str: The fallback route name
+ """
+ routes = list(options.keys())
+
+ if self.fallback_strategy == 'first':
+ return routes[0]
+ elif self.fallback_strategy == 'last':
+ return routes[-1]
+ elif self.fallback_strategy == 'random':
+ import random
+
+ return random.choice(routes)
+ else:
+ return routes[0]
+
+ async def route(self, memory: BaseMemory, execution_context: dict = None) -> str:
+ """
+ Make a routing decision using the LLM.
+
+ Args:
+ memory: The conversation memory
+
+ Returns:
+ str: The name of the route to take
+ """
+ options = self.get_routing_options()
+
+ for attempt in range(self.max_retries):
+ try:
+ prompt = self.get_routing_prompt(memory, options, execution_context)
+
+ messages = [{'role': 'user', 'content': prompt}]
+ response = await self.llm.generate(messages)
+ decision = self.llm.get_message_content(response).strip().lower()
+
+ # Find matching option (case-insensitive)
+ for option_name in options:
+ if (
+ option_name.lower() == decision
+ or option_name.lower() in decision
+ ):
+ logger.info(f'LLM router selected: {option_name}')
+ return option_name
+
+ # If no exact match, try partial matching
+ for option_name in options:
+ if (
+ decision in option_name.lower()
+ or option_name.lower() in decision
+ ):
+ logger.info(
+ f'LLM router selected (partial match): {option_name}'
+ )
+ return option_name
+
+ logger.warning(
+ f"LLM router attempt {attempt + 1}: Invalid decision '{decision}', retrying..."
+ )
+
+ except Exception as e:
+ logger.error(f'LLM router attempt {attempt + 1} failed: {e}')
+
+ # Fallback strategy
+ fallback = self.get_fallback_route(options)
+ logger.warning(f'LLM router failed, using fallback: {fallback}')
+ return fallback
+
+
+class SmartRouter(BaseLLMRouter):
+ """
+ A general-purpose smart router that can route between different types of agents
+ based on task analysis and conversation context.
+ """
+
+ def __init__(
+ self,
+ routing_options: Dict[str, str],
+ llm: Optional[BaseLLM] = None,
+ context_description: Optional[str] = None,
+ **kwargs,
+ ):
+ """
+ Initialize the smart router.
+
+ Args:
+ routing_options: Dict mapping route names to descriptions
+ llm: LLM instance for routing decisions
+ context_description: Additional context about the workflow
+ **kwargs: Additional arguments for BaseLLMRouter
+ """
+ super().__init__(llm=llm, **kwargs)
+ self.routing_options = routing_options
+ self.context_description = context_description or 'a multi-agent workflow'
+
+ def get_routing_options(self) -> Dict[str, str]:
+ return self.routing_options
+
+ def get_routing_prompt(
+ self,
+ memory: BaseMemory,
+ options: Dict[str, str],
+ execution_context: dict = None,
+ ) -> str:
+ conversation = memory.get()
+
+ # Format conversation history
+ if isinstance(conversation, list):
+ conversation_text = '\n'.join(
+ [str(msg) for msg in conversation[-5:]]
+ ) # Last 5 messages
+ else:
+ conversation_text = str(conversation)
+
+ # Format options
+ options_text = '\n'.join(
+ [f'- {name}: {desc}' for name, desc in options.items()]
+ )
+
+ # Add execution context if available
+ context_info = ''
+ if execution_context:
+ visit_counts = execution_context.get('node_visit_count', {})
+ current_node = execution_context.get('current_node', 'unknown')
+ iteration = execution_context.get('iteration_count', 0)
+
+ # Create visit count warning if any node has been visited multiple times
+ visit_warnings = []
+ for node_name, count in visit_counts.items():
+ if count >= 2:
+ visit_warnings.append(f'- {node_name}: visited {count} times')
+
+ if visit_warnings:
+ context_info = f"""
+
+โ ๏ธ EXECUTION CONTEXT (Avoid Infinite Loops):
+Current iteration: {iteration}
+Current node: {current_node}
+Node visit counts:
+{chr(10).join(visit_warnings)}
+
+IMPORTANT: To prevent infinite loops, strongly prefer moving to different nodes or completing the workflow.
+"""
+
+ prompt = f"""You are a workflow coordinator for {self.context_description}.
+
+Based on the conversation history below, decide which agent should handle the next step.
+{context_info}
+Available agents:
+{options_text}
+
+Recent conversation:
+{conversation_text}
+
+Instructions:
+1. Analyze the conversation to understand what type of work is needed next
+2. Choose the most appropriate agent from the available options
+3. If you notice nodes being visited repeatedly, prefer different options or completion
+4. Respond with ONLY the agent name (no explanations or additional text)
+
+Agent to route to:"""
+
+ return prompt
+
+
+class TaskClassifierRouter(BaseLLMRouter):
+ """
+ A router specialized for classifying tasks and routing to appropriate specialists.
+ """
+
+ def __init__(
+ self,
+ task_categories: Dict[str, Dict[str, Any]],
+ llm: Optional[BaseLLM] = None,
+ **kwargs,
+ ):
+ """
+ Initialize the task classifier router.
+
+ Args:
+ task_categories: Dict mapping category names to their config:
+ {
+ "category_name": {
+ "description": "What this category handles",
+ "keywords": ["keyword1", "keyword2"], # Optional
+ "examples": ["example task 1", "example task 2"] # Optional
+ }
+ }
+ llm: LLM instance for routing decisions
+ **kwargs: Additional arguments for BaseLLMRouter
+ """
+ super().__init__(llm=llm, **kwargs)
+ self.task_categories = task_categories
+
+ def get_routing_options(self) -> Dict[str, str]:
+ return {
+ name: config['description'] for name, config in self.task_categories.items()
+ }
+
+ def get_routing_prompt(
+ self,
+ memory: BaseMemory,
+ options: Dict[str, str],
+ execution_context: dict = None,
+ ) -> str:
+ conversation = memory.get()
+
+ # Get the latest user input or task
+ if isinstance(conversation, list) and conversation:
+ latest_task = str(conversation[-1])
+ else:
+ latest_task = str(conversation)
+
+ # Build detailed category descriptions
+ categories_detail = []
+ for name, config in self.task_categories.items():
+ detail = f"- {name}: {config['description']}"
+
+ if 'keywords' in config:
+ detail += f"\n Keywords: {', '.join(config['keywords'])}"
+
+ if 'examples' in config:
+ detail += f"\n Examples: {', '.join(config['examples'])}"
+
+ categories_detail.append(detail)
+
+ categories_text = '\n\n'.join(categories_detail)
+
+ prompt = f"""You are a task classifier that routes requests to specialized agents.
+
+Task to classify:
+{latest_task}
+
+Available categories:
+{categories_text}
+
+Instructions:
+1. Analyze the task to understand what type of work it requires
+2. Choose the most appropriate category from the available options
+3. Consider keywords and examples to make the best match
+4. Respond with ONLY the category name (no explanations)
+
+Category:"""
+
+ return prompt
+
+
+class ConversationAnalysisRouter(BaseLLMRouter):
+ """
+ A router that analyzes conversation flow and context to make routing decisions.
+ """
+
+ def __init__(
+ self,
+ routing_logic: Dict[str, str],
+ analysis_depth: int = 3,
+ llm: Optional[BaseLLM] = None,
+ **kwargs,
+ ):
+ """
+ Initialize the conversation analysis router.
+
+ Args:
+ routing_logic: Dict mapping route names to routing criteria
+ analysis_depth: Number of recent messages to analyze
+ llm: LLM instance for routing decisions
+ **kwargs: Additional arguments for BaseLLMRouter
+ """
+ super().__init__(llm=llm, **kwargs)
+ self.routing_logic = routing_logic
+ self.analysis_depth = analysis_depth
+
+ def get_routing_options(self) -> Dict[str, str]:
+ return self.routing_logic
+
+ def get_routing_prompt(
+ self,
+ memory: BaseMemory,
+ options: Dict[str, str],
+ execution_context: dict = None,
+ ) -> str:
+ conversation = memory.get()
+
+ # Analyze recent conversation
+ if isinstance(conversation, list):
+ recent_messages = conversation[-self.analysis_depth :]
+ conversation_text = '\n'.join(
+ [f'Message {i+1}: {msg}' for i, msg in enumerate(recent_messages)]
+ )
+ else:
+ conversation_text = str(conversation)
+
+ # Format routing logic
+ logic_text = '\n'.join(
+ [f'- {name}: {criteria}' for name, criteria in options.items()]
+ )
+
+ # Add execution context if available
+ context_info = ''
+ if execution_context:
+ visit_counts = execution_context.get('node_visit_count', {})
+ current_node = execution_context.get('current_node', 'unknown')
+ iteration = execution_context.get('iteration_count', 0)
+
+ # Create visit count warning if any node has been visited multiple times
+ visit_warnings = []
+ for node_name, count in visit_counts.items():
+ if count >= 2:
+ visit_warnings.append(f'- {node_name}: visited {count} times')
+
+ if visit_warnings:
+ context_info = f"""
+
+โ ๏ธ EXECUTION CONTEXT (Avoid Infinite Loops):
+Current iteration: {iteration}
+Current node: {current_node}
+Node visit counts:
+{chr(10).join(visit_warnings)}
+
+CRITICAL: Excessive node revisits detected! Consider completing the workflow or choosing different paths.
+"""
+
+ prompt = f"""You are a conversation flow analyzer that determines the next step in a workflow.
+
+Recent conversation (last {self.analysis_depth} messages):
+{conversation_text}
+{context_info}
+Routing logic:
+{logic_text}
+
+Instructions:
+1. Analyze the conversation flow and current state
+2. Consider what has been accomplished and what needs to happen next
+3. If nodes are being revisited too frequently, strongly prefer completion or different paths
+4. Choose the route that best matches the current conversation state
+5. Respond with ONLY the route name (no explanations)
+
+Next route:"""
+
+ return prompt
+
+
+def create_llm_router(router_type: str, **config) -> Callable[[BaseMemory], str]:
+ """
+ Factory function to create LLM-powered routers with different configurations.
+
+ Args:
+ router_type: Type of router ("smart", "task_classifier", "conversation_analysis")
+ **config: Configuration specific to the router type
+
+ Returns:
+ Callable router function that can be used in Arium workflows
+
+ Examples:
+ # Smart router
+ router = create_llm_router(
+ "smart",
+ routing_options={
+ "researcher": "Gather information and conduct research",
+ "analyst": "Analyze data and perform calculations",
+ "writer": "Create reports and summaries"
+ }
+ )
+
+ # Task classifier router
+ router = create_llm_router(
+ "task_classifier",
+ task_categories={
+ "research": {
+ "description": "Research and information gathering tasks",
+ "keywords": ["research", "find", "investigate", "gather"],
+ "examples": ["Find information about...", "Research the topic of..."]
+ },
+ "analysis": {
+ "description": "Data analysis and computational tasks",
+ "keywords": ["analyze", "calculate", "compute", "data"],
+ "examples": ["Analyze the data...", "Calculate the..."]
+ }
+ }
+ )
+ """
+
+ if router_type == 'smart':
+ if 'routing_options' not in config:
+ raise ValueError("SmartRouter requires 'routing_options' parameter")
+
+ router_instance = SmartRouter(**config)
+
+ elif router_type == 'task_classifier':
+ if 'task_categories' not in config:
+ raise ValueError(
+ "TaskClassifierRouter requires 'task_categories' parameter"
+ )
+
+ router_instance = TaskClassifierRouter(**config)
+
+ elif router_type == 'conversation_analysis':
+ if 'routing_logic' not in config:
+ raise ValueError(
+ "ConversationAnalysisRouter requires 'routing_logic' parameter"
+ )
+
+ router_instance = ConversationAnalysisRouter(**config)
+
+ else:
+ raise ValueError(f'Unknown router type: {router_type}')
+
+ # Get the routing options for type annotation
+ options = router_instance.get_routing_options()
+ option_names = tuple(options.keys())
+
+ # Create proper Literal type for validation
+ from typing import Literal
+
+ if len(option_names) == 1:
+ # Handle single option case
+ literal_type = Literal[option_names[0]]
+ else:
+ # Handle multiple options case
+ literal_type = Literal[option_names]
+
+ # Return a function that can be used as a router
+ async def router_function(memory: BaseMemory, execution_context: dict = None):
+ """Generated router function that uses LLM for routing decisions"""
+ return await router_instance.route(memory, execution_context)
+
+ # Add proper type annotations for validation
+ router_function.__annotations__ = {'memory': BaseMemory, 'return': literal_type}
+
+ return router_function
+
+
+def llm_router(
+ routing_options: Dict[str, str],
+ llm: Optional[BaseLLM] = None,
+ context_description: Optional[str] = None,
+ **kwargs,
+):
+ """
+ Decorator to create LLM-powered routers with a simple interface.
+
+ Args:
+ routing_options: Dict mapping route names to descriptions
+ llm: LLM instance for routing decisions
+ context_description: Additional context about the workflow
+ **kwargs: Additional router configuration
+
+ Returns:
+ Decorator function
+
+ Example:
+ @llm_router({
+ "researcher": "Gather information and conduct research",
+ "analyst": "Analyze data and perform calculations",
+ "writer": "Create reports and summaries"
+ })
+ def my_smart_router(memory: BaseMemory) -> Literal["researcher", "analyst", "writer"]:
+ pass # Implementation is provided by decorator
+ """
+
+ def decorator(func):
+ # Extract return type annotation to validate routing options
+ if hasattr(func, '__annotations__') and 'return' in func.__annotations__:
+ return_annotation = func.__annotations__['return']
+
+ # Validate that routing options match the return type
+ if (
+ hasattr(return_annotation, '__origin__')
+ and return_annotation.__origin__ is Union
+ ):
+ # Handle Literal types
+ args = get_args(return_annotation)
+ if args and hasattr(args[0], '__origin__'):
+ literal_values = get_args(args[0])
+ option_keys = set(routing_options.keys())
+ literal_set = set(literal_values)
+
+ if option_keys != literal_set:
+ logger.warning(
+ f"Routing options {option_keys} don't match return type {literal_set}"
+ )
+
+ # Create the router instance
+ router_instance = SmartRouter(
+ routing_options=routing_options,
+ llm=llm,
+ context_description=context_description,
+ **kwargs,
+ )
+
+ @wraps(func)
+ async def wrapper(memory: BaseMemory, execution_context: dict = None):
+ return await router_instance.route(memory, execution_context)
+
+ # Preserve the original function's type annotations including return type
+ wrapper.__annotations__ = func.__annotations__.copy()
+
+ # Ensure the return annotation is properly set
+ if 'return' in func.__annotations__:
+ wrapper.__annotations__['return'] = func.__annotations__['return']
+
+ return wrapper
+
+ return decorator
+
+
+# Convenience functions for common routing patterns
+
+
+def create_research_analysis_router(
+ research_agent: str = 'researcher',
+ analysis_agent: str = 'analyst',
+ summary_agent: str = 'summarizer',
+ llm: Optional[BaseLLM] = None,
+) -> Callable[[BaseMemory], str]:
+ """
+ Create a router for common research -> analysis -> summary workflows.
+
+ Args:
+ research_agent: Name of the research agent
+ analysis_agent: Name of the analysis agent
+ summary_agent: Name of the summary agent
+ llm: LLM instance for routing decisions
+
+ Returns:
+ Router function for research/analysis workflows
+ """
+ return create_llm_router(
+ 'task_classifier',
+ task_categories={
+ research_agent: {
+ 'description': 'Research and information gathering tasks',
+ 'keywords': [
+ 'research',
+ 'find',
+ 'investigate',
+ 'gather',
+ 'search',
+ 'lookup',
+ ],
+ 'examples': [
+ 'Find information about...',
+ 'Research the topic of...',
+ 'Gather data on...',
+ ],
+ },
+ analysis_agent: {
+ 'description': 'Data analysis and computational tasks',
+ 'keywords': [
+ 'analyze',
+ 'calculate',
+ 'compute',
+ 'data',
+ 'numbers',
+ 'statistics',
+ ],
+ 'examples': [
+ 'Analyze the data...',
+ 'Calculate the...',
+ 'Compute statistics for...',
+ ],
+ },
+ summary_agent: {
+ 'description': 'Summarization and conclusion tasks',
+ 'keywords': [
+ 'summarize',
+ 'conclude',
+ 'report',
+ 'final',
+ 'overview',
+ 'wrap up',
+ ],
+ 'examples': [
+ 'Summarize the findings...',
+ 'Create a final report...',
+ 'Conclude the analysis...',
+ ],
+ },
+ },
+ llm=llm,
+ )
diff --git a/flo_ai/tests/test_arium_yaml.py b/flo_ai/tests/test_arium_yaml.py
new file mode 100644
index 00000000..e1935656
--- /dev/null
+++ b/flo_ai/tests/test_arium_yaml.py
@@ -0,0 +1,1120 @@
+"""
+Tests for YAML-based Arium workflow construction.
+"""
+
+import pytest
+from unittest.mock import Mock, patch
+
+from flo_ai.arium.builder import AriumBuilder
+from flo_ai.arium.memory import MessageMemory, BaseMemory
+from flo_ai.models.agent import Agent
+from flo_ai.tool.base_tool import Tool
+from flo_ai.llm import OpenAI
+
+
+class TestAriumYamlBuilder:
+ """Test class for YAML-based Arium builder functionality."""
+
+ def test_from_yaml_validation_no_params(self):
+ """Test that from_yaml fails when no parameters are provided."""
+ with pytest.raises(
+ ValueError, match='Either yaml_str or yaml_file must be provided'
+ ):
+ AriumBuilder.from_yaml()
+
+ def test_from_yaml_validation_both_params(self):
+ """Test that from_yaml fails when both parameters are provided."""
+ with pytest.raises(
+ ValueError, match='Only one of yaml_str or yaml_file should be provided'
+ ):
+ AriumBuilder.from_yaml(yaml_str='test', yaml_file='test.yaml')
+
+ def test_from_yaml_validation_missing_arium_section(self):
+ """Test that from_yaml fails when YAML doesn't contain arium section."""
+ yaml_config = """
+ metadata:
+ name: test
+ """
+ with pytest.raises(ValueError, match='YAML must contain an "arium" section'):
+ AriumBuilder.from_yaml(yaml_str=yaml_config)
+
+ def test_from_yaml_simple_configuration(self):
+ """Test basic YAML configuration parsing."""
+ yaml_config = """
+ metadata:
+ name: test-workflow
+ version: 1.0.0
+ description: "Test workflow"
+
+ arium:
+ agents:
+ - name: test_agent
+ yaml_config: |
+ agent:
+ name: test_agent
+ job: "You are a test agent"
+ model:
+ provider: openai
+ name: gpt-4o-mini
+
+ workflow:
+ start: test_agent
+ edges:
+ - from: test_agent
+ to: [end]
+ end: [test_agent]
+ """
+
+ # Mock the AgentBuilder.from_yaml to avoid actual LLM calls
+ with patch('flo_ai.arium.builder.AgentBuilder') as mock_agent_builder:
+ mock_agent = Mock(spec=Agent)
+ mock_agent.name = 'test_agent'
+
+ mock_builder_instance = Mock()
+ mock_builder_instance.build.return_value = mock_agent
+ mock_agent_builder.from_yaml.return_value = mock_builder_instance
+
+ builder = AriumBuilder.from_yaml(yaml_str=yaml_config)
+
+ # Verify the builder was configured correctly
+ assert len(builder._agents) == 1
+ assert builder._agents[0].name == 'test_agent'
+ assert builder._start_node == mock_agent
+ assert mock_agent in builder._end_nodes
+
+ def test_from_yaml_with_custom_memory(self):
+ """Test YAML configuration with custom memory parameter."""
+ yaml_config = """
+ arium:
+ agents:
+ - name: test_agent
+ yaml_config: |
+ agent:
+ name: test_agent
+ job: "You are a test agent"
+ model:
+ provider: openai
+ name: gpt-4o-mini
+
+ workflow:
+ start: test_agent
+ edges:
+ - from: test_agent
+ to: [end]
+ end: [test_agent]
+ """
+
+ # Create custom memory
+ custom_memory = Mock(spec=MessageMemory)
+
+ with patch('flo_ai.arium.builder.AgentBuilder') as mock_agent_builder:
+ mock_agent = Mock(spec=Agent)
+ mock_agent.name = 'test_agent'
+
+ mock_builder_instance = Mock()
+ mock_builder_instance.build.return_value = mock_agent
+ mock_agent_builder.from_yaml.return_value = mock_builder_instance
+
+ builder = AriumBuilder.from_yaml(yaml_str=yaml_config, memory=custom_memory)
+
+ # Verify custom memory was used
+ assert builder._memory == custom_memory
+
+ def test_from_yaml_default_memory(self):
+ """Test that default MessageMemory is used when no memory is provided."""
+ yaml_config = """
+ arium:
+ agents:
+ - name: test_agent
+ yaml_config: |
+ agent:
+ name: test_agent
+ job: "You are a test agent"
+ model:
+ provider: openai
+ name: gpt-4o-mini
+
+ workflow:
+ start: test_agent
+ edges:
+ - from: test_agent
+ to: [end]
+ end: [test_agent]
+ """
+
+ with patch('flo_ai.arium.builder.AgentBuilder') as mock_agent_builder:
+ mock_agent = Mock(spec=Agent)
+ mock_agent.name = 'test_agent'
+
+ mock_builder_instance = Mock()
+ mock_builder_instance.build.return_value = mock_agent
+ mock_agent_builder.from_yaml.return_value = mock_builder_instance
+
+ builder = AriumBuilder.from_yaml(yaml_str=yaml_config)
+
+ # Verify default MessageMemory was created
+ assert builder._memory is not None
+ assert isinstance(builder._memory, MessageMemory)
+
+ def test_from_yaml_with_tools(self):
+ """Test YAML configuration with tools."""
+ yaml_config = """
+ arium:
+ agents:
+ - name: test_agent
+ yaml_config: |
+ agent:
+ name: test_agent
+ job: "Test agent with tools"
+ model:
+ provider: openai
+ name: gpt-4o-mini
+
+ tools:
+ - name: test_tool
+
+ workflow:
+ start: test_agent
+ edges:
+ - from: test_agent
+ to: [test_tool]
+ - from: test_tool
+ to: [end]
+ end: [test_tool]
+ """
+
+ # Create mock tool
+ mock_tool = Mock(spec=Tool)
+ mock_tool.name = 'test_tool'
+ tools = {'test_tool': mock_tool}
+
+ with patch('flo_ai.arium.builder.AgentBuilder') as mock_agent_builder:
+ mock_agent = Mock(spec=Agent)
+ mock_agent.name = 'test_agent'
+
+ mock_builder_instance = Mock()
+ mock_builder_instance.build.return_value = mock_agent
+ mock_agent_builder.from_yaml.return_value = mock_builder_instance
+
+ builder = AriumBuilder.from_yaml(yaml_str=yaml_config, tools=tools)
+
+ # Verify tools were added
+ assert len(builder._tools) == 1
+ assert builder._tools[0] == mock_tool
+
+ def test_from_yaml_with_routers(self):
+ """Test YAML configuration with custom routers."""
+ yaml_config = """
+ arium:
+ agents:
+ - name: agent1
+ yaml_config: |
+ agent:
+ name: agent1
+ job: "First agent"
+ model:
+ provider: openai
+ name: gpt-4o-mini
+ - name: agent2
+ yaml_config: |
+ agent:
+ name: agent2
+ job: "Second agent"
+ model:
+ provider: openai
+ name: gpt-4o-mini
+
+ workflow:
+ start: agent1
+ edges:
+ - from: agent1
+ to: [agent2]
+ router: test_router
+ - from: agent2
+ to: [end]
+ end: [agent2]
+ """
+
+ # Create mock router
+ def test_router(memory: BaseMemory) -> str:
+ return 'agent2'
+
+ routers = {'test_router': test_router}
+
+ with patch('flo_ai.arium.builder.AgentBuilder') as mock_agent_builder:
+ mock_agent1 = Mock(spec=Agent)
+ mock_agent1.name = 'agent1'
+ mock_agent2 = Mock(spec=Agent)
+ mock_agent2.name = 'agent2'
+
+ mock_builder_instance = Mock()
+ mock_builder_instance.build.side_effect = [mock_agent1, mock_agent2]
+ mock_agent_builder.from_yaml.return_value = mock_builder_instance
+
+ builder = AriumBuilder.from_yaml(yaml_str=yaml_config, routers=routers)
+
+ # Verify router was assigned
+ assert len(builder._edges) == 1
+ from_node, to_nodes, router = builder._edges[0]
+ assert from_node == mock_agent1
+ assert to_nodes == [mock_agent2]
+ assert router == test_router
+
+ def test_from_yaml_missing_tool_error(self):
+ """Test error when referenced tool is not provided."""
+ yaml_config = """
+ arium:
+ agents:
+ - name: test_agent
+ yaml_config: |
+ agent:
+ name: test_agent
+ job: "Test agent"
+ model:
+ provider: openai
+ name: gpt-4o-mini
+
+ tools:
+ - name: missing_tool
+
+ workflow:
+ start: test_agent
+ edges:
+ - from: test_agent
+ to: [missing_tool]
+ - from: missing_tool
+ to: [end]
+ end: [missing_tool]
+ """
+
+ with patch('flo_ai.arium.builder.AgentBuilder'):
+ with pytest.raises(
+ ValueError,
+ match='Tool missing_tool not found in provided tools dictionary',
+ ):
+ AriumBuilder.from_yaml(yaml_str=yaml_config, tools={})
+
+ def test_from_yaml_missing_router_error(self):
+ """Test error when referenced router is not provided."""
+ yaml_config = """
+ arium:
+ agents:
+ - name: agent1
+ yaml_config: |
+ agent:
+ name: agent1
+ job: "Test agent"
+ model:
+ provider: openai
+ name: gpt-4o-mini
+ - name: agent2
+ yaml_config: |
+ agent:
+ name: agent2
+ job: "Test agent"
+ model:
+ provider: openai
+ name: gpt-4o-mini
+
+ workflow:
+ start: agent1
+ edges:
+ - from: agent1
+ to: [agent2]
+ router: missing_router
+ - from: agent2
+ to: [end]
+ end: [agent2]
+ """
+
+ with patch('flo_ai.arium.builder.AgentBuilder'):
+ with pytest.raises(
+ ValueError,
+ match='Router missing_router not found',
+ ):
+ AriumBuilder.from_yaml(yaml_str=yaml_config, routers={})
+
+ def test_from_yaml_missing_start_node_error(self):
+ """Test error when workflow doesn't specify start node."""
+ yaml_config = """
+ arium:
+ agents:
+ - name: test_agent
+ yaml_config: |
+ agent:
+ name: test_agent
+ job: "Test agent"
+ model:
+ provider: openai
+ name: gpt-4o-mini
+
+ workflow:
+ edges: []
+ end: [test_agent]
+ """
+
+ with patch('flo_ai.arium.builder.AgentBuilder'):
+ with pytest.raises(ValueError, match='Workflow must specify a start node'):
+ AriumBuilder.from_yaml(yaml_str=yaml_config)
+
+ def test_from_yaml_missing_end_nodes_error(self):
+ """Test error when workflow doesn't specify end nodes."""
+ yaml_config = """
+ arium:
+ agents:
+ - name: test_agent
+ yaml_config: |
+ agent:
+ name: test_agent
+ job: "Test agent"
+ model:
+ provider: openai
+ name: gpt-4o-mini
+
+ workflow:
+ start: test_agent
+ edges:
+ - from: test_agent
+ to: [end]
+ """
+
+ with patch('flo_ai.arium.builder.AgentBuilder'):
+ with pytest.raises(ValueError, match='Workflow must specify end nodes'):
+ AriumBuilder.from_yaml(yaml_str=yaml_config)
+
+ def test_from_yaml_invalid_agent_config_error(self):
+ """Test error when agent doesn't have yaml_config or yaml_file."""
+ yaml_config = """
+ arium:
+ agents:
+ - name: invalid_agent
+ # Missing yaml_config or yaml_file
+
+ workflow:
+ start: invalid_agent
+ edges:
+ - from: invalid_agent
+ to: [end]
+ end: [invalid_agent]
+ """
+
+ with pytest.raises(
+ ValueError,
+ match='Agent invalid_agent not found in provided agents dictionary',
+ ):
+ AriumBuilder.from_yaml(yaml_str=yaml_config)
+
+ def test_from_yaml_external_file_reference(self):
+ """Test YAML configuration with external agent file reference."""
+ yaml_config = """
+ arium:
+ agents:
+ - name: external_agent
+ yaml_file: "path/to/agent.yaml"
+
+ workflow:
+ start: external_agent
+ edges:
+ - from: external_agent
+ to: [end]
+ end: [external_agent]
+ """
+
+ with patch('flo_ai.arium.builder.AgentBuilder') as mock_agent_builder:
+ mock_agent = Mock(spec=Agent)
+ mock_agent.name = 'external_agent'
+
+ mock_builder_instance = Mock()
+ mock_builder_instance.build.return_value = mock_agent
+ mock_agent_builder.from_yaml.return_value = mock_builder_instance
+
+ AriumBuilder.from_yaml(yaml_str=yaml_config)
+
+ # Verify AgentBuilder.from_yaml was called with yaml_file
+ mock_agent_builder.from_yaml.assert_called_with(
+ yaml_file='path/to/agent.yaml', base_llm=None
+ )
+
+ def test_from_yaml_with_base_llm(self):
+ """Test YAML configuration with base LLM parameter."""
+ yaml_config = """
+ arium:
+ agents:
+ - name: test_agent
+ yaml_config: |
+ agent:
+ name: test_agent
+ job: "Test agent"
+
+ workflow:
+ start: test_agent
+ edges:
+ - from: test_agent
+ to: [end]
+ end: [test_agent]
+ """
+
+ mock_llm = Mock(spec=OpenAI)
+
+ with patch('flo_ai.arium.builder.AgentBuilder') as mock_agent_builder:
+ mock_agent = Mock(spec=Agent)
+ mock_agent.name = 'test_agent'
+
+ mock_builder_instance = Mock()
+ mock_builder_instance.build.return_value = mock_agent
+ mock_agent_builder.from_yaml.return_value = mock_builder_instance
+
+ AriumBuilder.from_yaml(yaml_str=yaml_config, base_llm=mock_llm)
+
+ # Verify AgentBuilder.from_yaml was called with base_llm
+ calls = mock_agent_builder.from_yaml.call_args_list
+ assert len(calls) == 1
+ call_args, call_kwargs = calls[0]
+ assert 'yaml_str' in call_kwargs
+ assert 'base_llm' in call_kwargs
+ assert call_kwargs['base_llm'] == mock_llm
+ # Just verify it contains the expected content
+ assert 'agent:' in call_kwargs['yaml_str']
+ assert 'name: test_agent' in call_kwargs['yaml_str']
+ assert 'job: "Test agent"' in call_kwargs['yaml_str']
+
+ def test_from_yaml_complex_workflow(self):
+ """Test complex workflow with multiple agents, tools, and routers."""
+ yaml_config = """
+ metadata:
+ name: complex-workflow
+ version: 2.0.0
+
+ arium:
+ agents:
+ - name: dispatcher
+ yaml_config: |
+ agent:
+ name: dispatcher
+ role: Dispatcher
+ job: "Route requests to appropriate handlers"
+ model:
+ provider: openai
+ name: gpt-4o-mini
+ - name: processor
+ yaml_config: |
+ agent:
+ name: processor
+ role: Processor
+ job: "Process the data"
+ model:
+ provider: openai
+ name: gpt-4o-mini
+ - name: summarizer
+ yaml_config: |
+ agent:
+ name: summarizer
+ role: Summarizer
+ job: "Create final summary"
+ model:
+ provider: openai
+ name: gpt-4o-mini
+
+ tools:
+ - name: data_tool
+ - name: analysis_tool
+
+ workflow:
+ start: dispatcher
+ edges:
+ - from: dispatcher
+ to: [data_tool, analysis_tool, processor]
+ router: dispatch_router
+ - from: data_tool
+ to: [summarizer]
+ - from: analysis_tool
+ to: [summarizer]
+ - from: processor
+ to: [summarizer]
+ - from: summarizer
+ to: [end]
+ end: [summarizer]
+ """
+
+ # Create mocks
+ def dispatch_router(memory: BaseMemory) -> str:
+ return 'processor'
+
+ mock_data_tool = Mock(spec=Tool)
+ mock_data_tool.name = 'data_tool'
+ mock_analysis_tool = Mock(spec=Tool)
+ mock_analysis_tool.name = 'analysis_tool'
+
+ tools = {'data_tool': mock_data_tool, 'analysis_tool': mock_analysis_tool}
+ routers = {'dispatch_router': dispatch_router}
+
+ with patch('flo_ai.arium.builder.AgentBuilder') as mock_agent_builder:
+ mock_dispatcher = Mock(spec=Agent)
+ mock_dispatcher.name = 'dispatcher'
+ mock_processor = Mock(spec=Agent)
+ mock_processor.name = 'processor'
+ mock_summarizer = Mock(spec=Agent)
+ mock_summarizer.name = 'summarizer'
+
+ mock_builder_instance = Mock()
+ mock_builder_instance.build.side_effect = [
+ mock_dispatcher,
+ mock_processor,
+ mock_summarizer,
+ ]
+ mock_agent_builder.from_yaml.return_value = mock_builder_instance
+
+ builder = AriumBuilder.from_yaml(
+ yaml_str=yaml_config, tools=tools, routers=routers
+ )
+
+ # Verify all components were configured
+ assert len(builder._agents) == 3
+ assert len(builder._tools) == 2
+ assert len(builder._edges) == 4 # 4 edge definitions
+ assert builder._start_node == mock_dispatcher
+ assert mock_summarizer in builder._end_nodes
+
+ def test_from_yaml_end_keyword_handling(self):
+ """Test proper handling of 'end' keyword in edge definitions."""
+ yaml_config = """
+ arium:
+ agents:
+ - name: test_agent
+ yaml_config: |
+ agent:
+ name: test_agent
+ job: "Test agent"
+ model:
+ provider: openai
+ name: gpt-4o-mini
+
+ workflow:
+ start: test_agent
+ edges:
+ - from: test_agent
+ to: [end] # Should not create actual edge, handled by end nodes
+ end: [test_agent]
+ """
+
+ with patch('flo_ai.arium.builder.AgentBuilder') as mock_agent_builder:
+ mock_agent = Mock(spec=Agent)
+ mock_agent.name = 'test_agent'
+
+ mock_builder_instance = Mock()
+ mock_builder_instance.build.return_value = mock_agent
+ mock_agent_builder.from_yaml.return_value = mock_builder_instance
+
+ builder = AriumBuilder.from_yaml(yaml_str=yaml_config)
+
+ # Should not create any edges since only 'end' was in to_nodes
+ assert len(builder._edges) == 0
+ assert mock_agent in builder._end_nodes
+
+ def test_from_yaml_direct_agent_configuration(self):
+ """Test direct agent configuration without nested YAML."""
+ yaml_config = """
+ metadata:
+ name: test-workflow
+ version: 1.0.0
+
+ arium:
+ agents:
+ - name: test_agent
+ role: Test Agent
+ job: "You are a test agent for validation"
+ model:
+ provider: openai
+ name: gpt-4o-mini
+ base_url: "https://api.openai.com/v1"
+ settings:
+ temperature: 0.5
+ max_retries: 2
+ reasoning_pattern: REACT
+
+ workflow:
+ start: test_agent
+ edges:
+ - from: test_agent
+ to: [end]
+ end: [test_agent]
+ """
+
+ with patch('flo_ai.llm.OpenAI') as mock_openai:
+ mock_llm = Mock()
+ mock_openai.return_value = mock_llm
+
+ builder = AriumBuilder.from_yaml(yaml_str=yaml_config)
+
+ # Verify the builder was configured correctly
+ assert len(builder._agents) == 1
+ agent = builder._agents[0]
+ assert agent.name == 'test_agent'
+ assert agent.role == 'Test Agent'
+ assert (
+ agent.system_prompt
+ == 'You are Test Agent. You are a test agent for validation'
+ )
+ assert mock_llm.temperature == 0.5
+
+ def test_from_yaml_direct_config_with_tools(self):
+ """Test direct agent configuration with tools."""
+ yaml_config = """
+ arium:
+ agents:
+ - name: test_agent
+ job: "Test agent with tools"
+ model:
+ provider: openai
+ name: gpt-4o-mini
+ tools: ["calculator", "web_search"]
+
+ tools:
+ - name: calculator
+ - name: web_search
+
+ workflow:
+ start: test_agent
+ edges:
+ - from: test_agent
+ to: [end]
+ end: [test_agent]
+ """
+
+ # Create mock tools
+ mock_calculator = Mock(spec=Tool)
+ mock_calculator.name = 'calculator'
+ mock_web_search = Mock(spec=Tool)
+ mock_web_search.name = 'web_search'
+
+ tools = {'calculator': mock_calculator, 'web_search': mock_web_search}
+
+ with patch('flo_ai.llm.OpenAI') as mock_openai:
+ mock_llm = Mock()
+ mock_openai.return_value = mock_llm
+
+ builder = AriumBuilder.from_yaml(yaml_str=yaml_config, tools=tools)
+
+ # Verify agent was configured with tools
+ assert len(builder._agents) == 1
+ agent = builder._agents[0]
+ assert len(agent.tools) == 2
+ assert mock_calculator in agent.tools
+ assert mock_web_search in agent.tools
+
+ def test_from_yaml_direct_config_with_parser(self):
+ """Test direct agent configuration with structured output parser."""
+ yaml_config = """
+ arium:
+ agents:
+ - name: test_agent
+ job: "Test agent with structured output"
+ model:
+ provider: openai
+ name: gpt-4o-mini
+ parser:
+ name: TestParser
+ fields:
+ - name: result
+ type: str
+ description: "The result"
+ - name: confidence
+ type: float
+ description: "Confidence score"
+
+ workflow:
+ start: test_agent
+ edges:
+ - from: test_agent
+ to: [end]
+ end: [test_agent]
+ """
+
+ with patch('flo_ai.llm.OpenAI') as mock_openai:
+ with patch(
+ 'flo_ai.formatter.yaml_format_parser.FloYamlParser'
+ ) as mock_parser:
+ mock_llm = Mock()
+ mock_openai.return_value = mock_llm
+
+ mock_parser_instance = Mock()
+ mock_parser_instance.get_format.return_value = {'type': 'object'}
+ mock_parser.create.return_value = mock_parser_instance
+
+ builder = AriumBuilder.from_yaml(yaml_str=yaml_config)
+
+ # Verify parser was configured
+ assert len(builder._agents) == 1
+ agent = builder._agents[0]
+ assert agent.output_schema == {'type': 'object'}
+
+ def test_from_yaml_mixed_configuration_methods(self):
+ """Test mixing different agent configuration methods in one workflow."""
+ yaml_config = """
+ arium:
+ agents:
+ # Direct configuration
+ - name: direct_agent
+ role: Direct Agent
+ job: "Directly configured agent"
+ model:
+ provider: openai
+ name: gpt-4o-mini
+
+ # Inline YAML configuration
+ - name: yaml_agent
+ yaml_config: |
+ agent:
+ name: yaml_agent
+ role: YAML Agent
+ job: "YAML configured agent"
+ model:
+ provider: openai
+ name: gpt-4o-mini
+
+ # External file reference
+ - name: file_agent
+ yaml_file: "path/to/agent.yaml"
+
+ workflow:
+ start: direct_agent
+ edges:
+ - from: direct_agent
+ to: [yaml_agent]
+ - from: yaml_agent
+ to: [file_agent]
+ - from: file_agent
+ to: [end]
+ end: [file_agent]
+ """
+
+ with patch('flo_ai.llm.OpenAI') as mock_openai:
+ with patch('flo_ai.arium.builder.AgentBuilder') as mock_agent_builder:
+ mock_llm = Mock()
+ mock_openai.return_value = mock_llm
+
+ # Mock for inline YAML config
+ mock_yaml_agent = Mock(spec=Agent)
+ mock_yaml_agent.name = 'yaml_agent'
+
+ # Mock for external file config
+ mock_file_agent = Mock(spec=Agent)
+ mock_file_agent.name = 'file_agent'
+
+ mock_builder_instance = Mock()
+ mock_builder_instance.build.side_effect = [
+ mock_yaml_agent,
+ mock_file_agent,
+ ]
+ mock_agent_builder.from_yaml.return_value = mock_builder_instance
+
+ builder = AriumBuilder.from_yaml(yaml_str=yaml_config)
+
+ # Verify all three agents were created
+ assert len(builder._agents) == 3
+
+ # Check direct agent
+ direct_agent = next(
+ a for a in builder._agents if a.name == 'direct_agent'
+ )
+ assert direct_agent.role == 'Direct Agent'
+
+ # Check other agents were added
+ assert any(a.name == 'yaml_agent' for a in builder._agents)
+ assert any(a.name == 'file_agent' for a in builder._agents)
+
+ def test_from_yaml_direct_config_validation_errors(self):
+ """Test validation errors for direct agent configuration."""
+
+ # Test missing required field
+ yaml_config_missing_job = """
+ arium:
+ agents:
+ - name: test_agent
+ role: Test Agent
+ # missing job field
+ model:
+ provider: openai
+ name: gpt-4o-mini
+ workflow:
+ start: test_agent
+ edges: []
+ end: [test_agent]
+ """
+
+ with pytest.raises(ValueError, match='Agent test_agent must have either'):
+ AriumBuilder.from_yaml(yaml_str=yaml_config_missing_job)
+
+ # Test invalid reasoning pattern
+ yaml_config_invalid_pattern = """
+ arium:
+ agents:
+ - name: test_agent
+ job: "Test agent"
+ model:
+ provider: openai
+ name: gpt-4o-mini
+ settings:
+ reasoning_pattern: INVALID_PATTERN
+ workflow:
+ start: test_agent
+ edges: []
+ end: [test_agent]
+ """
+
+ with patch('flo_ai.llm.OpenAI'):
+ with pytest.raises(ValueError, match='Invalid reasoning pattern'):
+ AriumBuilder.from_yaml(yaml_str=yaml_config_invalid_pattern)
+
+ # Test missing model when no base_llm provided
+ yaml_config_missing_model = """
+ arium:
+ agents:
+ - name: test_agent
+ job: "Test agent"
+ # missing model field
+ workflow:
+ start: test_agent
+ edges: []
+ end: [test_agent]
+ """
+
+ with pytest.raises(ValueError, match='Model must be specified'):
+ AriumBuilder.from_yaml(yaml_str=yaml_config_missing_model)
+
+ def test_from_yaml_direct_config_with_base_llm(self):
+ """Test direct agent configuration with base LLM override."""
+ yaml_config = """
+ arium:
+ agents:
+ - name: test_agent
+ job: "Test agent without model config"
+ settings:
+ temperature: 0.7
+
+ workflow:
+ start: test_agent
+ edges: []
+ end: [test_agent]
+ """
+
+ mock_base_llm = Mock(spec=OpenAI)
+
+ builder = AriumBuilder.from_yaml(yaml_str=yaml_config, base_llm=mock_base_llm)
+
+ # Verify agent was created with base LLM
+ assert len(builder._agents) == 1
+ agent = builder._agents[0]
+ assert agent.llm == mock_base_llm
+ assert mock_base_llm.temperature == 0.7
+
+ def test_from_yaml_prebuilt_agents(self):
+ """Test using pre-built agents with YAML workflow."""
+ yaml_config = """
+ arium:
+ agents:
+ # Reference pre-built agents (only name specified)
+ - name: prebuilt_agent1
+ - name: prebuilt_agent2
+
+ workflow:
+ start: prebuilt_agent1
+ edges:
+ - from: prebuilt_agent1
+ to: [prebuilt_agent2]
+ - from: prebuilt_agent2
+ to: [end]
+ end: [prebuilt_agent2]
+ """
+
+ # Create mock pre-built agents
+ mock_agent1 = Mock(spec=Agent)
+ mock_agent1.name = 'prebuilt_agent1'
+ mock_agent2 = Mock(spec=Agent)
+ mock_agent2.name = 'prebuilt_agent2'
+
+ prebuilt_agents = {
+ 'prebuilt_agent1': mock_agent1,
+ 'prebuilt_agent2': mock_agent2,
+ }
+
+ builder = AriumBuilder.from_yaml(yaml_str=yaml_config, agents=prebuilt_agents)
+
+ # Verify pre-built agents were used
+ assert len(builder._agents) == 2
+ assert mock_agent1 in builder._agents
+ assert mock_agent2 in builder._agents
+ assert builder._start_node == mock_agent1
+ assert mock_agent2 in builder._end_nodes
+
+ def test_from_yaml_prebuilt_agents_missing_error(self):
+ """Test error when referenced pre-built agent is not provided."""
+ yaml_config = """
+ arium:
+ agents:
+ - name: missing_agent
+
+ workflow:
+ start: missing_agent
+ edges: []
+ end: [missing_agent]
+ """
+
+ with pytest.raises(
+ ValueError,
+ match='Agent missing_agent not found in provided agents dictionary',
+ ):
+ AriumBuilder.from_yaml(yaml_str=yaml_config, agents={})
+
+ def test_from_yaml_mixed_prebuilt_and_configured_agents(self):
+ """Test mixing pre-built agents with other configuration methods."""
+ yaml_config = """
+ arium:
+ agents:
+ # Pre-built agent reference
+ - name: prebuilt_agent
+
+ # Direct configuration
+ - name: direct_agent
+ role: Direct Agent
+ job: "Directly configured agent"
+ model:
+ provider: openai
+ name: gpt-4o-mini
+
+ # Inline YAML config
+ - name: yaml_agent
+ yaml_config: |
+ agent:
+ name: yaml_agent
+ job: "YAML configured agent"
+ model:
+ provider: openai
+ name: gpt-4o-mini
+
+ workflow:
+ start: prebuilt_agent
+ edges:
+ - from: prebuilt_agent
+ to: [direct_agent]
+ - from: direct_agent
+ to: [yaml_agent]
+ - from: yaml_agent
+ to: [end]
+ end: [yaml_agent]
+ """
+
+ # Create mock pre-built agent
+ mock_prebuilt_agent = Mock(spec=Agent)
+ mock_prebuilt_agent.name = 'prebuilt_agent'
+
+ prebuilt_agents = {'prebuilt_agent': mock_prebuilt_agent}
+
+ with patch('flo_ai.llm.OpenAI') as mock_openai:
+ with patch('flo_ai.arium.builder.AgentBuilder') as mock_agent_builder:
+ mock_llm = Mock()
+ mock_openai.return_value = mock_llm
+
+ # Mock for inline YAML config
+ mock_yaml_agent = Mock(spec=Agent)
+ mock_yaml_agent.name = 'yaml_agent'
+
+ mock_builder_instance = Mock()
+ mock_builder_instance.build.return_value = mock_yaml_agent
+ mock_agent_builder.from_yaml.return_value = mock_builder_instance
+
+ builder = AriumBuilder.from_yaml(
+ yaml_str=yaml_config, agents=prebuilt_agents
+ )
+
+ # Verify all agents were created/added
+ assert len(builder._agents) == 3
+
+ # Check pre-built agent
+ assert mock_prebuilt_agent in builder._agents
+
+ # Check direct agent was created
+ direct_agent = next(
+ a for a in builder._agents if a.name == 'direct_agent'
+ )
+ assert direct_agent.role == 'Direct Agent'
+
+ # Check YAML agent was added
+ assert mock_yaml_agent in builder._agents
+
+ def test_from_yaml_prebuilt_agents_parameter_validation(self):
+ """Test parameter validation for pre-built agents."""
+ yaml_config = """
+ arium:
+ agents:
+ - name: test_agent
+ # Has additional fields, so not a pure reference
+ role: "Some Role"
+
+ workflow:
+ start: test_agent
+ edges: []
+ end: [test_agent]
+ """
+
+ # This should not be treated as a pre-built agent reference
+ # because it has additional fields beyond just 'name'
+ with patch('flo_ai.llm.OpenAI'):
+ with pytest.raises(ValueError, match='Agent test_agent must have either'):
+ AriumBuilder.from_yaml(yaml_str=yaml_config)
+
+ def test_from_yaml_prebuilt_agents_with_tools_and_routers(self):
+ """Test pre-built agents working together with tools and routers."""
+ yaml_config = """
+ arium:
+ agents:
+ - name: dispatcher
+ - name: processor
+
+ tools:
+ - name: calculator
+
+ workflow:
+ start: dispatcher
+ edges:
+ - from: dispatcher
+ to: [calculator, processor]
+ router: smart_router
+ - from: calculator
+ to: [processor]
+ - from: processor
+ to: [end]
+ end: [processor]
+ """
+
+ # Create mocks
+ mock_dispatcher = Mock(spec=Agent)
+ mock_dispatcher.name = 'dispatcher'
+ mock_processor = Mock(spec=Agent)
+ mock_processor.name = 'processor'
+
+ mock_calculator = Mock(spec=Tool)
+ mock_calculator.name = 'calculator'
+
+ def smart_router(memory):
+ return 'processor'
+
+ prebuilt_agents = {'dispatcher': mock_dispatcher, 'processor': mock_processor}
+ tools = {'calculator': mock_calculator}
+ routers = {'smart_router': smart_router}
+
+ builder = AriumBuilder.from_yaml(
+ yaml_str=yaml_config, agents=prebuilt_agents, tools=tools, routers=routers
+ )
+
+ # Verify everything was configured correctly
+ assert len(builder._agents) == 2
+ assert len(builder._tools) == 1
+ assert len(builder._edges) == 2
+ assert mock_dispatcher in builder._agents
+ assert mock_processor in builder._agents
+ assert mock_calculator in builder._tools
+
+
+if __name__ == '__main__':
+ pytest.main([__file__])
diff --git a/flo_ai/tests/test_llm_router.py b/flo_ai/tests/test_llm_router.py
new file mode 100644
index 00000000..728ebdc9
--- /dev/null
+++ b/flo_ai/tests/test_llm_router.py
@@ -0,0 +1,310 @@
+"""
+Test cases for LLM-powered routers in Arium workflows.
+"""
+
+import pytest
+from unittest.mock import Mock, AsyncMock
+from typing import Literal
+
+from flo_ai.arium.llm_router import (
+ SmartRouter,
+ TaskClassifierRouter,
+ ConversationAnalysisRouter,
+ create_llm_router,
+ llm_router,
+)
+from flo_ai.arium.memory import MessageMemory
+from flo_ai.llm.base_llm import BaseLLM
+
+
+class MockLLM(BaseLLM):
+ """Mock LLM for testing"""
+
+ def __init__(self, response_text: str = 'researcher'):
+ super().__init__(model='mock')
+ self.response_text = response_text
+ self.call_count = 0
+
+ async def generate(self, messages, **kwargs):
+ self.call_count += 1
+ return {'response': self.response_text}
+
+ def get_message_content(self, response):
+ return response.get('response', 'researcher')
+
+ def format_tool_for_llm(self, tool):
+ """Mock implementation for tool formatting"""
+ return {'name': tool.name, 'description': 'mock tool'}
+
+ def format_tools_for_llm(self, tools):
+ """Mock implementation for tools formatting"""
+ return [self.format_tool_for_llm(tool) for tool in tools]
+
+ def format_image_in_message(self, image):
+ """Mock implementation for image formatting"""
+ return 'mock_image_content'
+
+
+@pytest.fixture
+def mock_memory():
+ """Create a mock memory with sample conversation"""
+ memory = MessageMemory()
+ memory.add('I need to research market trends for renewable energy')
+ memory.add('Please analyze the data and provide insights')
+ return memory
+
+
+class TestSmartRouter:
+ """Test SmartRouter functionality"""
+
+ def test_initialization(self):
+ """Test SmartRouter initialization"""
+ routing_options = {'researcher': 'Research tasks', 'analyst': 'Analysis tasks'}
+
+ mock_llm = MockLLM('researcher')
+ router = SmartRouter(routing_options, llm=mock_llm)
+
+ assert router.get_routing_options() == routing_options
+ assert router.llm == mock_llm
+
+ @pytest.mark.asyncio
+ async def test_route_decision(self, mock_memory):
+ """Test routing decision making"""
+ routing_options = {'researcher': 'Research tasks', 'analyst': 'Analysis tasks'}
+
+ mock_llm = MockLLM('researcher')
+ router = SmartRouter(routing_options, llm=mock_llm)
+
+ result = await router.route(mock_memory)
+
+ assert result == 'researcher'
+ assert mock_llm.call_count == 1
+
+ @pytest.mark.asyncio
+ async def test_fallback_on_invalid_response(self, mock_memory):
+ """Test fallback when LLM returns invalid response"""
+ routing_options = {'researcher': 'Research tasks', 'analyst': 'Analysis tasks'}
+
+ mock_llm = MockLLM('invalid_response')
+ router = SmartRouter(
+ routing_options, llm=mock_llm, max_retries=1, fallback_strategy='first'
+ )
+
+ result = await router.route(mock_memory)
+
+ assert result == 'researcher' # First option as fallback
+ assert mock_llm.call_count == 1 # Should retry once
+
+ def test_get_routing_prompt(self, mock_memory):
+ """Test routing prompt generation"""
+ routing_options = {'researcher': 'Research tasks', 'analyst': 'Analysis tasks'}
+
+ mock_llm = MockLLM('researcher')
+ router = SmartRouter(routing_options, llm=mock_llm)
+ prompt = router.get_routing_prompt(mock_memory, routing_options)
+
+ assert 'researcher' in prompt
+ assert 'analyst' in prompt
+ assert 'Research tasks' in prompt
+ assert 'Analysis tasks' in prompt
+
+
+class TestTaskClassifierRouter:
+ """Test TaskClassifierRouter functionality"""
+
+ def test_initialization(self):
+ """Test TaskClassifierRouter initialization"""
+ task_categories = {
+ 'math': {
+ 'description': 'Math tasks',
+ 'keywords': ['calculate', 'math'],
+ 'examples': ['Calculate sum'],
+ }
+ }
+
+ mock_llm = MockLLM('math')
+ router = TaskClassifierRouter(task_categories, llm=mock_llm)
+
+ assert router.task_categories == task_categories
+ assert router.get_routing_options() == {'math': 'Math tasks'}
+
+ @pytest.mark.asyncio
+ async def test_route_with_keywords(self, mock_memory):
+ """Test routing based on keywords"""
+ task_categories = {
+ 'math': {
+ 'description': 'Math tasks',
+ 'keywords': ['calculate', 'math'],
+ 'examples': ['Calculate sum'],
+ },
+ 'research': {
+ 'description': 'Research tasks',
+ 'keywords': ['research', 'find'],
+ 'examples': ['Find information'],
+ },
+ }
+
+ mock_llm = MockLLM('research')
+ router = TaskClassifierRouter(task_categories, llm=mock_llm)
+
+ result = await router.route(mock_memory)
+
+ assert result == 'research'
+
+
+class TestConversationAnalysisRouter:
+ """Test ConversationAnalysisRouter functionality"""
+
+ def test_initialization(self):
+ """Test ConversationAnalysisRouter initialization"""
+ routing_logic = {'planner': 'Plan tasks', 'executor': 'Execute tasks'}
+
+ mock_llm = MockLLM('planner')
+ router = ConversationAnalysisRouter(
+ routing_logic, analysis_depth=2, llm=mock_llm
+ )
+
+ assert router.routing_logic == routing_logic
+ assert router.analysis_depth == 2
+
+ @pytest.mark.asyncio
+ async def test_route_with_conversation_analysis(self, mock_memory):
+ """Test routing based on conversation analysis"""
+ routing_logic = {'planner': 'Plan tasks', 'executor': 'Execute tasks'}
+
+ mock_llm = MockLLM('executor')
+ router = ConversationAnalysisRouter(routing_logic, llm=mock_llm)
+
+ result = await router.route(mock_memory)
+
+ assert result == 'executor'
+
+
+class TestFactoryFunction:
+ """Test create_llm_router factory function"""
+
+ def test_create_smart_router(self):
+ """Test creating SmartRouter via factory"""
+ routing_options = {'researcher': 'Research tasks', 'analyst': 'Analysis tasks'}
+
+ mock_llm = MockLLM('researcher')
+ router_fn = create_llm_router(
+ 'smart', routing_options=routing_options, llm=mock_llm
+ )
+
+ assert callable(router_fn)
+ assert hasattr(router_fn, '__annotations__')
+
+ def test_create_task_classifier_router(self):
+ """Test creating TaskClassifierRouter via factory"""
+ task_categories = {
+ 'math': {
+ 'description': 'Math tasks',
+ 'keywords': ['calculate'],
+ 'examples': ['Calculate sum'],
+ }
+ }
+
+ mock_llm = MockLLM('math')
+ router_fn = create_llm_router(
+ 'task_classifier', task_categories=task_categories, llm=mock_llm
+ )
+
+ assert callable(router_fn)
+
+ def test_create_conversation_analysis_router(self):
+ """Test creating ConversationAnalysisRouter via factory"""
+ routing_logic = {'planner': 'Plan tasks', 'executor': 'Execute tasks'}
+
+ mock_llm = MockLLM('planner')
+ router_fn = create_llm_router(
+ 'conversation_analysis', routing_logic=routing_logic, llm=mock_llm
+ )
+
+ assert callable(router_fn)
+
+ def test_invalid_router_type(self):
+ """Test error on invalid router type"""
+ with pytest.raises(ValueError, match='Unknown router type'):
+ create_llm_router('invalid_type')
+
+
+class TestDecorator:
+ """Test @llm_router decorator"""
+
+ def test_decorator_creates_router(self):
+ """Test that decorator creates working router"""
+ routing_options = {'researcher': 'Research tasks', 'analyst': 'Analysis tasks'}
+
+ mock_llm = MockLLM('researcher')
+
+ @llm_router(routing_options, llm=mock_llm)
+ def test_router(memory) -> Literal['researcher', 'analyst']:
+ pass
+
+ assert callable(test_router)
+ # Test would require actual execution which needs async setup
+
+
+class TestErrorHandling:
+ """Test error handling and edge cases"""
+
+ def test_missing_required_config(self):
+ """Test error when required config is missing"""
+ with pytest.raises(ValueError, match="requires 'routing_options'"):
+ create_llm_router('smart')
+
+ with pytest.raises(ValueError, match="requires 'task_categories'"):
+ create_llm_router('task_classifier')
+
+ with pytest.raises(ValueError, match="requires 'routing_logic'"):
+ create_llm_router('conversation_analysis')
+
+ @pytest.mark.asyncio
+ async def test_llm_failure_fallback(self, mock_memory):
+ """Test fallback when LLM fails completely"""
+ routing_options = {'researcher': 'Research tasks', 'analyst': 'Analysis tasks'}
+
+ # Mock LLM that raises exception
+ mock_llm = Mock()
+ mock_llm.generate = AsyncMock(side_effect=Exception('LLM Error'))
+
+ router = SmartRouter(
+ routing_options, llm=mock_llm, max_retries=1, fallback_strategy='first'
+ )
+
+ result = await router.route(mock_memory)
+
+ assert result == 'researcher' # Should fallback to first option
+
+ def test_fallback_strategies(self, mock_memory):
+ """Test different fallback strategies"""
+ routing_options = {
+ 'researcher': 'Research tasks',
+ 'analyst': 'Analysis tasks',
+ 'writer': 'Writing tasks',
+ }
+
+ # Test "first" strategy
+ mock_llm = MockLLM('invalid_response') # Will trigger fallback
+ router_first = SmartRouter(
+ routing_options, fallback_strategy='first', llm=mock_llm
+ )
+ assert router_first.get_fallback_route(routing_options) == 'researcher'
+
+ # Test "last" strategy
+ router_last = SmartRouter(
+ routing_options, fallback_strategy='last', llm=mock_llm
+ )
+ assert router_last.get_fallback_route(routing_options) == 'writer'
+
+ # Test "random" strategy
+ router_random = SmartRouter(
+ routing_options, fallback_strategy='random', llm=mock_llm
+ )
+ fallback = router_random.get_fallback_route(routing_options)
+ assert fallback in routing_options.keys()
+
+
+if __name__ == '__main__':
+ pytest.main([__file__])
diff --git a/flo_ai/tests/test_router_fix.py b/flo_ai/tests/test_router_fix.py
new file mode 100644
index 00000000..291fa5f7
--- /dev/null
+++ b/flo_ai/tests/test_router_fix.py
@@ -0,0 +1,100 @@
+#!/usr/bin/env python3
+"""
+Quick test to verify the router type annotation fix works.
+"""
+
+import inspect
+from typing import get_origin, get_args, Literal
+from flo_ai.arium import create_llm_router
+from flo_ai.llm import OpenAI
+
+
+def test_router_type_annotation():
+ """Test that our router functions have proper type annotations"""
+
+ # Create a router with dummy LLM
+ llm = OpenAI(model='gpt-4o-mini', api_key='dummy-key')
+ router = create_llm_router(
+ 'smart',
+ routing_options={
+ 'researcher': 'Research tasks',
+ 'analyst': 'Analysis tasks',
+ 'writer': 'Writing tasks',
+ },
+ llm=llm,
+ )
+
+ # Check the function signature
+ sig = inspect.signature(router)
+ return_annotation = sig.return_annotation
+
+ print(f'Return annotation: {return_annotation}')
+ print(f'Return annotation type: {type(return_annotation)}')
+
+ # Check if it's a Literal type
+ origin = get_origin(return_annotation)
+ print(f'Origin: {origin}')
+ print(f'Is Literal: {origin is Literal}')
+
+ if origin is Literal:
+ literal_values = list(get_args(return_annotation))
+ print(f'Literal values: {literal_values}')
+ assert True, 'Router function has correct Literal type annotation'
+ else:
+ print('โ Not a Literal type!')
+ assert False, 'Router function should have Literal type annotation'
+
+
+def test_validation_logic():
+ """Test the exact validation logic from base.py"""
+
+ llm = OpenAI(model='gpt-4o-mini', api_key='dummy-key')
+ router = create_llm_router(
+ 'smart',
+ routing_options={'researcher': 'Research tasks', 'analyst': 'Analysis tasks'},
+ llm=llm,
+ )
+
+ try:
+ # Get the function signature (same as base.py)
+ sig = inspect.signature(router)
+ return_annotation = sig.return_annotation
+
+ # Check if there's no return annotation
+ if return_annotation == inspect.Signature.empty:
+ print('โ No return annotation')
+ return False
+
+ # Check if the return type is a Literal
+ origin = get_origin(return_annotation)
+
+ # In Python 3.8+, Literal types have get_origin() return typing.Literal
+ if origin is Literal:
+ # Extract the literal values
+ literal_values = list(get_args(return_annotation))
+ print(f'โ
Validation passed! Literal values: {literal_values}')
+ assert True, 'Validation logic works correctly'
+ else:
+ print(f'โ Validation failed! Origin is {origin}, not Literal')
+ assert False, f'Validation failed! Origin is {origin}, not Literal'
+
+ except Exception as e:
+ print(f'โ Exception during validation: {e}')
+ assert False, f'Exception during validation: {e}'
+
+
+if __name__ == '__main__':
+ print('๐งช Testing Router Type Annotation Fix')
+ print('=' * 50)
+
+ print('\n1. Testing type annotation:')
+ result1 = test_router_type_annotation()
+
+ print('\n2. Testing validation logic:')
+ result2 = test_validation_logic()
+
+ print('\n' + '=' * 50)
+ if result1 and result2:
+ print('โ
All tests passed! Router type annotations are working correctly.')
+ else:
+ print('โ Tests failed! Router type annotations need fixing.')