From c5e600b0a2dc151c2dcf5cd17ee0fba73868c8fb Mon Sep 17 00:00:00 2001 From: vizsatiz Date: Wed, 6 Aug 2025 17:32:09 +0530 Subject: [PATCH 1/7] Arium yaml chanegs --- flo_ai/docs/arium_yaml_guide.md | 739 ++++++++++++++++ flo_ai/examples/arium_yaml_example.py | 552 ++++++++++++ flo_ai/flo_ai/arium/builder.py | 359 ++++++++ flo_ai/tests/test_arium_yaml.py | 1134 +++++++++++++++++++++++++ 4 files changed, 2784 insertions(+) create mode 100644 flo_ai/docs/arium_yaml_guide.md create mode 100644 flo_ai/examples/arium_yaml_example.py create mode 100644 flo_ai/tests/test_arium_yaml.py diff --git a/flo_ai/docs/arium_yaml_guide.md b/flo_ai/docs/arium_yaml_guide.md new file mode 100644 index 00000000..989c6c8e --- /dev/null +++ b/flo_ai/docs/arium_yaml_guide.md @@ -0,0 +1,739 @@ +# 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 + + workflow: + start: content_analyst + edges: + - from: content_analyst + to: [summarizer] + - from: summarizer + to: [validator, tool1] + router: quality_router # optional + - 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 + +#### 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/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/flo_ai/arium/builder.py b/flo_ai/flo_ai/arium/builder.py index b343bcfa..dad79f37 100644 --- a/flo_ai/flo_ai/arium/builder.py +++ b/flo_ai/flo_ai/arium/builder.py @@ -4,6 +4,9 @@ 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 class AriumBuilder: @@ -158,6 +161,362 @@ 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 + + workflow: + start: content_analyst + edges: + - from: content_analyst + to: [validator, summarizer] + router: my_router # optional + - 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 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 routers and router_name in routers: + router_fn = routers[router_name] + else: + raise ValueError( + f'Router {router_name} not found in provided routers dictionary. ' + f'Available routers: {list(routers.keys()) if 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_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: + model_config = agent_config['model'] + 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(f'Model name must be specified for agent {name}') + + 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}') + 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') + if 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/tests/test_arium_yaml.py b/flo_ai/tests/test_arium_yaml.py new file mode 100644 index 00000000..6303e533 --- /dev/null +++ b/flo_ai/tests/test_arium_yaml.py @@ -0,0 +1,1134 @@ +""" +Tests for YAML-based Arium workflow construction. +""" + +import pytest +from unittest.mock import Mock + +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 pytest.mock.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 pytest.mock.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 pytest.mock.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 pytest.mock.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 pytest.mock.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 pytest.mock.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 pytest.mock.patch('flo_ai.arium.builder.AgentBuilder'): + with pytest.raises( + ValueError, + match='Router missing_router not found in provided routers dictionary', + ): + 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 pytest.mock.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 pytest.mock.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 must have either yaml_config or yaml_file', + ): + 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 pytest.mock.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 pytest.mock.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 + mock_agent_builder.from_yaml.assert_called_with( + yaml_str=yaml_config.split('yaml_config: |')[1].strip(), + base_llm=mock_llm, + ) + + 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 pytest.mock.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 pytest.mock.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 pytest.mock.patch('flo_ai.arium.builder.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 pytest.mock.patch('flo_ai.arium.builder.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 pytest.mock.patch('flo_ai.arium.builder.OpenAI') as mock_openai: + with pytest.mock.patch('flo_ai.arium.builder.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 pytest.mock.patch('flo_ai.arium.builder.OpenAI') as mock_openai: + with pytest.mock.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(KeyError): + 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 pytest.mock.patch('flo_ai.arium.builder.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 pytest.mock.patch('flo_ai.arium.builder.OpenAI') as mock_openai: + with pytest.mock.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 pytest.mock.patch('flo_ai.arium.builder.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__]) From 3c78390670c3a2fbd7e13b825617da77a8c73b2b Mon Sep 17 00:00:00 2001 From: vizsatiz Date: Fri, 15 Aug 2025 12:07:40 +0530 Subject: [PATCH 2/7] fix readme for format --- README.md | 2 -- 1 file changed, 2 deletions(-) diff --git a/README.md b/README.md index 24dc995d..c83506b6 100644 --- a/README.md +++ b/README.md @@ -418,8 +418,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 From cfe16bcec0a3eab0c36c3a3b1de0e3db97861304 Mon Sep 17 00:00:00 2001 From: vizsatiz Date: Fri, 15 Aug 2025 21:00:06 +0530 Subject: [PATCH 3/7] Update README and fix tests --- README.md | 288 ++++++++++++++++++++++++++++++++ flo_ai/flo_ai/arium/builder.py | 18 +- flo_ai/tests/test_arium_yaml.py | 90 +++++----- 3 files changed, 335 insertions(+), 61 deletions(-) diff --git a/README.md b/README.md index c83506b6..648f7f30 100644 --- a/README.md +++ b/README.md @@ -81,6 +81,7 @@ 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) - [📖 Documentation](#-documentation) - [🌟 Why Flo AI?](#-why-flo-ai) - [🎯 Use Cases](#-use-cases) @@ -1111,6 +1112,293 @@ 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" + +# Define reusable tools +tools: + - name: "web_search" + description: "Search the web for current information" + parameters: + query: + type: "string" + description: "Search query" + + - name: "calculator" + description: "Perform mathematical calculations" + parameters: + expression: + type: "string" + description: "Mathematical expression to calculate" + +arium: + # Reference external agent configurations + 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 defined above + + - 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" + + # Both specialists feed into synthesizer + - from: "researcher" + to: ["synthesizer"] + + - from: "analyst" + to: ["synthesizer"] + + end: ["synthesizer"] + +# Define router functions +routers: + classification_router: + description: "Route based on task classification" + code: | + def route(memory: BaseMemory) -> str: + 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 +``` + +#### 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 and tools separately +2. **Clear Naming**: Use descriptive names for agents, tools, 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 well-documented + +#### Benefits of YAML Workflows + +- **🔄 Reproducible**: Version-controlled workflow definitions +- **📝 Maintainable**: Easy to modify without code changes +- **🧪 Testable**: Different configurations for testing vs. production +- **👥 Collaborative**: Non-developers can modify workflow logic +- **🚀 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/flo_ai/arium/builder.py b/flo_ai/flo_ai/arium/builder.py index dad79f37..78ad8a85 100644 --- a/flo_ai/flo_ai/arium/builder.py +++ b/flo_ai/flo_ai/arium/builder.py @@ -493,15 +493,15 @@ def _create_agent_from_direct_config( # Extract output schema if present output_schema = agent_config.get('output_schema') - if 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() + + # 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( diff --git a/flo_ai/tests/test_arium_yaml.py b/flo_ai/tests/test_arium_yaml.py index 6303e533..542615ac 100644 --- a/flo_ai/tests/test_arium_yaml.py +++ b/flo_ai/tests/test_arium_yaml.py @@ -3,7 +3,7 @@ """ import pytest -from unittest.mock import Mock +from unittest.mock import Mock, patch from flo_ai.arium.builder import AriumBuilder from flo_ai.arium.memory import MessageMemory, BaseMemory @@ -66,9 +66,7 @@ def test_from_yaml_simple_configuration(self): """ # Mock the AgentBuilder.from_yaml to avoid actual LLM calls - with pytest.mock.patch( - 'flo_ai.arium.builder.AgentBuilder' - ) as mock_agent_builder: + with patch('flo_ai.arium.builder.AgentBuilder') as mock_agent_builder: mock_agent = Mock(spec=Agent) mock_agent.name = 'test_agent' @@ -109,9 +107,7 @@ def test_from_yaml_with_custom_memory(self): # Create custom memory custom_memory = Mock(spec=MessageMemory) - with pytest.mock.patch( - 'flo_ai.arium.builder.AgentBuilder' - ) as mock_agent_builder: + with patch('flo_ai.arium.builder.AgentBuilder') as mock_agent_builder: mock_agent = Mock(spec=Agent) mock_agent.name = 'test_agent' @@ -146,9 +142,7 @@ def test_from_yaml_default_memory(self): end: [test_agent] """ - with pytest.mock.patch( - 'flo_ai.arium.builder.AgentBuilder' - ) as mock_agent_builder: + with patch('flo_ai.arium.builder.AgentBuilder') as mock_agent_builder: mock_agent = Mock(spec=Agent) mock_agent.name = 'test_agent' @@ -194,9 +188,7 @@ def test_from_yaml_with_tools(self): mock_tool.name = 'test_tool' tools = {'test_tool': mock_tool} - with pytest.mock.patch( - 'flo_ai.arium.builder.AgentBuilder' - ) as mock_agent_builder: + with patch('flo_ai.arium.builder.AgentBuilder') as mock_agent_builder: mock_agent = Mock(spec=Agent) mock_agent.name = 'test_agent' @@ -249,9 +241,7 @@ def test_router(memory: BaseMemory) -> str: routers = {'test_router': test_router} - with pytest.mock.patch( - 'flo_ai.arium.builder.AgentBuilder' - ) as mock_agent_builder: + 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) @@ -297,7 +287,7 @@ def test_from_yaml_missing_tool_error(self): end: [missing_tool] """ - with pytest.mock.patch('flo_ai.arium.builder.AgentBuilder'): + with patch('flo_ai.arium.builder.AgentBuilder'): with pytest.raises( ValueError, match='Tool missing_tool not found in provided tools dictionary', @@ -337,7 +327,7 @@ def test_from_yaml_missing_router_error(self): end: [agent2] """ - with pytest.mock.patch('flo_ai.arium.builder.AgentBuilder'): + with patch('flo_ai.arium.builder.AgentBuilder'): with pytest.raises( ValueError, match='Router missing_router not found in provided routers dictionary', @@ -363,7 +353,7 @@ def test_from_yaml_missing_start_node_error(self): end: [test_agent] """ - with pytest.mock.patch('flo_ai.arium.builder.AgentBuilder'): + 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) @@ -388,7 +378,7 @@ def test_from_yaml_missing_end_nodes_error(self): to: [end] """ - with pytest.mock.patch('flo_ai.arium.builder.AgentBuilder'): + with patch('flo_ai.arium.builder.AgentBuilder'): with pytest.raises(ValueError, match='Workflow must specify end nodes'): AriumBuilder.from_yaml(yaml_str=yaml_config) @@ -410,7 +400,7 @@ def test_from_yaml_invalid_agent_config_error(self): with pytest.raises( ValueError, - match='Agent invalid_agent must have either yaml_config or yaml_file', + match='Agent invalid_agent not found in provided agents dictionary', ): AriumBuilder.from_yaml(yaml_str=yaml_config) @@ -430,9 +420,7 @@ def test_from_yaml_external_file_reference(self): end: [external_agent] """ - with pytest.mock.patch( - 'flo_ai.arium.builder.AgentBuilder' - ) as mock_agent_builder: + with patch('flo_ai.arium.builder.AgentBuilder') as mock_agent_builder: mock_agent = Mock(spec=Agent) mock_agent.name = 'external_agent' @@ -468,9 +456,7 @@ def test_from_yaml_with_base_llm(self): mock_llm = Mock(spec=OpenAI) - with pytest.mock.patch( - 'flo_ai.arium.builder.AgentBuilder' - ) as mock_agent_builder: + with patch('flo_ai.arium.builder.AgentBuilder') as mock_agent_builder: mock_agent = Mock(spec=Agent) mock_agent.name = 'test_agent' @@ -481,10 +467,16 @@ def test_from_yaml_with_base_llm(self): AriumBuilder.from_yaml(yaml_str=yaml_config, base_llm=mock_llm) # Verify AgentBuilder.from_yaml was called with base_llm - mock_agent_builder.from_yaml.assert_called_with( - yaml_str=yaml_config.split('yaml_config: |')[1].strip(), - base_llm=mock_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.""" @@ -556,9 +548,7 @@ def dispatch_router(memory: BaseMemory) -> str: tools = {'data_tool': mock_data_tool, 'analysis_tool': mock_analysis_tool} routers = {'dispatch_router': dispatch_router} - with pytest.mock.patch( - 'flo_ai.arium.builder.AgentBuilder' - ) as mock_agent_builder: + 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) @@ -607,9 +597,7 @@ def test_from_yaml_end_keyword_handling(self): end: [test_agent] """ - with pytest.mock.patch( - 'flo_ai.arium.builder.AgentBuilder' - ) as mock_agent_builder: + with patch('flo_ai.arium.builder.AgentBuilder') as mock_agent_builder: mock_agent = Mock(spec=Agent) mock_agent.name = 'test_agent' @@ -652,7 +640,7 @@ def test_from_yaml_direct_agent_configuration(self): end: [test_agent] """ - with pytest.mock.patch('flo_ai.arium.builder.OpenAI') as mock_openai: + with patch('flo_ai.llm.OpenAI') as mock_openai: mock_llm = Mock() mock_openai.return_value = mock_llm @@ -701,7 +689,7 @@ def test_from_yaml_direct_config_with_tools(self): tools = {'calculator': mock_calculator, 'web_search': mock_web_search} - with pytest.mock.patch('flo_ai.arium.builder.OpenAI') as mock_openai: + with patch('flo_ai.llm.OpenAI') as mock_openai: mock_llm = Mock() mock_openai.return_value = mock_llm @@ -742,8 +730,10 @@ def test_from_yaml_direct_config_with_parser(self): end: [test_agent] """ - with pytest.mock.patch('flo_ai.arium.builder.OpenAI') as mock_openai: - with pytest.mock.patch('flo_ai.arium.builder.FloYamlParser') as mock_parser: + 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 @@ -798,10 +788,8 @@ def test_from_yaml_mixed_configuration_methods(self): end: [file_agent] """ - with pytest.mock.patch('flo_ai.arium.builder.OpenAI') as mock_openai: - with pytest.mock.patch( - 'flo_ai.arium.builder.AgentBuilder' - ) as mock_agent_builder: + 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 @@ -854,7 +842,7 @@ def test_from_yaml_direct_config_validation_errors(self): end: [test_agent] """ - with pytest.raises(KeyError): + with pytest.raises(ValueError, match='Agent test_agent must have either'): AriumBuilder.from_yaml(yaml_str=yaml_config_missing_job) # Test invalid reasoning pattern @@ -874,7 +862,7 @@ def test_from_yaml_direct_config_validation_errors(self): end: [test_agent] """ - with pytest.mock.patch('flo_ai.arium.builder.OpenAI'): + with patch('flo_ai.llm.OpenAI'): with pytest.raises(ValueError, match='Invalid reasoning pattern'): AriumBuilder.from_yaml(yaml_str=yaml_config_invalid_pattern) @@ -1022,10 +1010,8 @@ def test_from_yaml_mixed_prebuilt_and_configured_agents(self): prebuilt_agents = {'prebuilt_agent': mock_prebuilt_agent} - with pytest.mock.patch('flo_ai.arium.builder.OpenAI') as mock_openai: - with pytest.mock.patch( - 'flo_ai.arium.builder.AgentBuilder' - ) as mock_agent_builder: + 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 @@ -1073,7 +1059,7 @@ def test_from_yaml_prebuilt_agents_parameter_validation(self): # This should not be treated as a pre-built agent reference # because it has additional fields beyond just 'name' - with pytest.mock.patch('flo_ai.arium.builder.OpenAI'): + with patch('flo_ai.llm.OpenAI'): with pytest.raises(ValueError, match='Agent test_agent must have either'): AriumBuilder.from_yaml(yaml_str=yaml_config) From fb44ea7df68eaf805a2544d710c0ff87b9b678e5 Mon Sep 17 00:00:00 2001 From: vizsatiz Date: Sun, 17 Aug 2025 12:49:06 +0530 Subject: [PATCH 4/7] Feature to have llm routers --- README.md | 139 +++- .../arium_examples.py} | 0 flo_ai/examples/llm_router_example.py | 457 ++++++++++++ flo_ai/flo_ai/arium/README.md | 40 +- flo_ai/flo_ai/arium/README_LLM_Router.md | 336 +++++++++ flo_ai/flo_ai/arium/__init__.py | 15 + flo_ai/flo_ai/arium/arium.py | 60 +- flo_ai/flo_ai/arium/base.py | 20 +- flo_ai/flo_ai/arium/llm_router.py | 660 ++++++++++++++++++ flo_ai/test_router_fix.py | 95 +++ flo_ai/tests/test_llm_router.py | 279 ++++++++ 11 files changed, 2056 insertions(+), 45 deletions(-) rename flo_ai/{flo_ai/arium/examples.py => examples/arium_examples.py} (100%) create mode 100644 flo_ai/examples/llm_router_example.py create mode 100644 flo_ai/flo_ai/arium/README_LLM_Router.md create mode 100644 flo_ai/flo_ai/arium/llm_router.py create mode 100644 flo_ai/test_router_fix.py create mode 100644 flo_ai/tests/test_llm_router.py diff --git a/README.md b/README.md index a3d0f751..a6bc8be1 100644 --- a/README.md +++ b/README.md @@ -1205,24 +1205,8 @@ metadata: version: "2.0.0" description: "Intelligent research workflow with conditional routing" -# Define reusable tools -tools: - - name: "web_search" - description: "Search the web for current information" - parameters: - query: - type: "string" - description: "Search query" - - - name: "calculator" - description: "Perform mathematical calculations" - parameters: - expression: - type: "string" - description: "Mathematical expression to calculate" - arium: - # Reference external agent configurations + # Define agents with tool references agents: - name: "classifier" role: "Content Classifier" @@ -1230,7 +1214,7 @@ arium: model: provider: "openai" name: "gpt-4o-mini" - tools: ["web_search"] # Reference tools defined above + tools: ["web_search"] # Reference tools provided in Python - name: "researcher" role: "Research Specialist" @@ -1267,7 +1251,7 @@ arium: # Conditional routing based on classification - from: "classifier" to: ["researcher", "analyst"] - router: "classification_router" + router: "classification_router" # Router function provided in Python # Both specialists feed into synthesizer - from: "researcher" @@ -1277,19 +1261,83 @@ arium: 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 router functions -routers: - classification_router: - description: "Route based on task classification" - code: | - def route(memory: BaseMemory) -> str: - 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 +# 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 ``` #### YAML Workflow with Variables @@ -1403,19 +1451,38 @@ arium: #### YAML Workflow Best Practices -1. **Modular Design**: Define reusable agents and tools separately -2. **Clear Naming**: Use descriptive names for agents, tools, and workflows +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 well-documented +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 without code changes +- **📝 Maintainable**: Easy to modify workflow structure without code changes - **🧪 Testable**: Different configurations for testing vs. production -- **👥 Collaborative**: Non-developers can modify workflow logic +- **👥 Collaborative**: Non-developers can modify workflow structure - **🚀 Deployable**: Easy CI/CD integration with YAML configurations - **🔍 Auditable**: Clear workflow definitions for compliance 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/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/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/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/test_router_fix.py b/flo_ai/test_router_fix.py new file mode 100644 index 00000000..237b6157 --- /dev/null +++ b/flo_ai/test_router_fix.py @@ -0,0 +1,95 @@ +#!/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 + + +def test_router_type_annotation(): + """Test that our router functions have proper type annotations""" + + # Create a router + router = create_llm_router( + 'smart', + routing_options={ + 'researcher': 'Research tasks', + 'analyst': 'Analysis tasks', + 'writer': 'Writing tasks', + }, + ) + + # 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}') + return True + else: + print('❌ Not a Literal type!') + return False + + +def test_validation_logic(): + """Test the exact validation logic from base.py""" + + router = create_llm_router( + 'smart', + routing_options={'researcher': 'Research tasks', 'analyst': 'Analysis tasks'}, + ) + + 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}') + return True + else: + print(f'❌ Validation failed! Origin is {origin}, not Literal') + return False + + except Exception as e: + print(f'❌ Exception during validation: {e}') + return False + + +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.') diff --git a/flo_ai/tests/test_llm_router.py b/flo_ai/tests/test_llm_router.py new file mode 100644 index 00000000..dad952f2 --- /dev/null +++ b/flo_ai/tests/test_llm_router.py @@ -0,0 +1,279 @@ +""" +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') + + +@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'} + + router = SmartRouter(routing_options) + 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'], + } + } + + router = TaskClassifierRouter(task_categories) + + 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'} + + router = ConversationAnalysisRouter(routing_logic, analysis_depth=2) + + 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'} + + router_fn = create_llm_router('smart', routing_options=routing_options) + + 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'], + } + } + + router_fn = create_llm_router( + 'task_classifier', task_categories=task_categories + ) + + assert callable(router_fn) + + def test_create_conversation_analysis_router(self): + """Test creating ConversationAnalysisRouter via factory""" + routing_logic = {'planner': 'Plan tasks', 'executor': 'Execute tasks'} + + router_fn = create_llm_router( + 'conversation_analysis', routing_logic=routing_logic + ) + + 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'} + + @llm_router(routing_options) + 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 + router_first = SmartRouter(routing_options, fallback_strategy='first') + assert router_first.get_fallback_route(routing_options) == 'researcher' + + # Test "last" strategy + router_last = SmartRouter(routing_options, fallback_strategy='last') + assert router_last.get_fallback_route(routing_options) == 'writer' + + # Test "random" strategy + router_random = SmartRouter(routing_options, fallback_strategy='random') + fallback = router_random.get_fallback_route(routing_options) + assert fallback in routing_options.keys() + + +if __name__ == '__main__': + pytest.main([__file__]) From cf79d0a184a41759fc3e1182a58e61968a52d77b Mon Sep 17 00:00:00 2001 From: vizsatiz Date: Sun, 17 Aug 2025 12:57:49 +0530 Subject: [PATCH 5/7] Fixing failing tests --- README.md | 2 ++ flo_ai/test_router_fix.py | 10 +++++----- flo_ai/tests/test_llm_router.py | 12 ++++++++++++ 3 files changed, 19 insertions(+), 5 deletions(-) diff --git a/README.md b/README.md index a6bc8be1..e201ccce 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 License + Tests Passing +

diff --git a/flo_ai/test_router_fix.py b/flo_ai/test_router_fix.py index 237b6157..fa1adb31 100644 --- a/flo_ai/test_router_fix.py +++ b/flo_ai/test_router_fix.py @@ -36,10 +36,10 @@ def test_router_type_annotation(): if origin is Literal: literal_values = list(get_args(return_annotation)) print(f'Literal values: {literal_values}') - return True + assert True, 'Router function has correct Literal type annotation' else: print('❌ Not a Literal type!') - return False + assert False, 'Router function should have Literal type annotation' def test_validation_logic(): @@ -68,14 +68,14 @@ def test_validation_logic(): # Extract the literal values literal_values = list(get_args(return_annotation)) print(f'✅ Validation passed! Literal values: {literal_values}') - return True + assert True, 'Validation logic works correctly' else: print(f'❌ Validation failed! Origin is {origin}, not Literal') - return False + assert False, f'Validation failed! Origin is {origin}, not Literal' except Exception as e: print(f'❌ Exception during validation: {e}') - return False + assert False, f'Exception during validation: {e}' if __name__ == '__main__': diff --git a/flo_ai/tests/test_llm_router.py b/flo_ai/tests/test_llm_router.py index dad952f2..1943efab 100644 --- a/flo_ai/tests/test_llm_router.py +++ b/flo_ai/tests/test_llm_router.py @@ -32,6 +32,18 @@ async def generate(self, messages, **kwargs): 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(): From 68454a3a3c29e37c16b4a2c0078c0f8208a81f09 Mon Sep 17 00:00:00 2001 From: vizsatiz Date: Sun, 17 Aug 2025 13:05:54 +0530 Subject: [PATCH 6/7] fix failing tests --- flo_ai/tests/test_llm_router.py | 39 ++++++++++++++++++++------- flo_ai/{ => tests}/test_router_fix.py | 7 ++++- 2 files changed, 35 insertions(+), 11 deletions(-) rename flo_ai/{ => tests}/test_router_fix.py (93%) diff --git a/flo_ai/tests/test_llm_router.py b/flo_ai/tests/test_llm_router.py index 1943efab..728ebdc9 100644 --- a/flo_ai/tests/test_llm_router.py +++ b/flo_ai/tests/test_llm_router.py @@ -99,7 +99,8 @@ def test_get_routing_prompt(self, mock_memory): """Test routing prompt generation""" routing_options = {'researcher': 'Research tasks', 'analyst': 'Analysis tasks'} - router = SmartRouter(routing_options) + mock_llm = MockLLM('researcher') + router = SmartRouter(routing_options, llm=mock_llm) prompt = router.get_routing_prompt(mock_memory, routing_options) assert 'researcher' in prompt @@ -121,7 +122,8 @@ def test_initialization(self): } } - router = TaskClassifierRouter(task_categories) + 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'} @@ -157,7 +159,10 @@ def test_initialization(self): """Test ConversationAnalysisRouter initialization""" routing_logic = {'planner': 'Plan tasks', 'executor': 'Execute tasks'} - router = ConversationAnalysisRouter(routing_logic, analysis_depth=2) + 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 @@ -182,7 +187,10 @@ def test_create_smart_router(self): """Test creating SmartRouter via factory""" routing_options = {'researcher': 'Research tasks', 'analyst': 'Analysis tasks'} - router_fn = create_llm_router('smart', routing_options=routing_options) + 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__') @@ -197,8 +205,9 @@ def test_create_task_classifier_router(self): } } + mock_llm = MockLLM('math') router_fn = create_llm_router( - 'task_classifier', task_categories=task_categories + 'task_classifier', task_categories=task_categories, llm=mock_llm ) assert callable(router_fn) @@ -207,8 +216,9 @@ 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 + 'conversation_analysis', routing_logic=routing_logic, llm=mock_llm ) assert callable(router_fn) @@ -226,7 +236,9 @@ def test_decorator_creates_router(self): """Test that decorator creates working router""" routing_options = {'researcher': 'Research tasks', 'analyst': 'Analysis tasks'} - @llm_router(routing_options) + mock_llm = MockLLM('researcher') + + @llm_router(routing_options, llm=mock_llm) def test_router(memory) -> Literal['researcher', 'analyst']: pass @@ -274,15 +286,22 @@ def test_fallback_strategies(self, mock_memory): } # Test "first" strategy - router_first = SmartRouter(routing_options, fallback_strategy='first') + 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') + 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') + 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() diff --git a/flo_ai/test_router_fix.py b/flo_ai/tests/test_router_fix.py similarity index 93% rename from flo_ai/test_router_fix.py rename to flo_ai/tests/test_router_fix.py index fa1adb31..291fa5f7 100644 --- a/flo_ai/test_router_fix.py +++ b/flo_ai/tests/test_router_fix.py @@ -6,12 +6,14 @@ 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 + # Create a router with dummy LLM + llm = OpenAI(model='gpt-4o-mini', api_key='dummy-key') router = create_llm_router( 'smart', routing_options={ @@ -19,6 +21,7 @@ def test_router_type_annotation(): 'analyst': 'Analysis tasks', 'writer': 'Writing tasks', }, + llm=llm, ) # Check the function signature @@ -45,9 +48,11 @@ def test_router_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: From 34b957b3296abd95f0bccf9565d01777e3f9b9e9 Mon Sep 17 00:00:00 2001 From: vizsatiz Date: Sun, 17 Aug 2025 21:11:07 +0530 Subject: [PATCH 7/7] support for custom yaml based routers --- README.md | 110 +++++- flo_ai/docs/arium_yaml_guide.md | 78 +++- flo_ai/examples/{usage.py => tool_usage.py} | 0 flo_ai/examples/yaml_llm_router_example.py | 372 ++++++++++++++++++++ flo_ai/flo_ai/arium/builder.py | 157 +++++++-- flo_ai/tests/test_arium_yaml.py | 2 +- 6 files changed, 692 insertions(+), 27 deletions(-) rename flo_ai/examples/{usage.py => tool_usage.py} (100%) create mode 100644 flo_ai/examples/yaml_llm_router_example.py diff --git a/README.md b/README.md index e201ccce..b4995f86 100644 --- a/README.md +++ b/README.md @@ -47,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 @@ -85,6 +86,7 @@ Flo AI is a Python framework that makes building production-ready AI agents and - [📊 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) @@ -838,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 @@ -1342,6 +1345,111 @@ async def run_workflow() -> List[Any]: 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 diff --git a/flo_ai/docs/arium_yaml_guide.md b/flo_ai/docs/arium_yaml_guide.md index 989c6c8e..419444a5 100644 --- a/flo_ai/docs/arium_yaml_guide.md +++ b/flo_ai/docs/arium_yaml_guide.md @@ -43,6 +43,21 @@ arium: - 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: @@ -50,7 +65,7 @@ arium: to: [summarizer] - from: summarizer to: [validator, tool1] - router: quality_router # optional + router: content_router # References router defined above - from: validator to: [end] - from: tool1 @@ -91,6 +106,67 @@ builder = AriumBuilder.from_yaml(yaml_str=config, memory=custom_memory) - ✅ **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: 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/builder.py b/flo_ai/flo_ai/arium/builder.py index 78ad8a85..fc0821fa 100644 --- a/flo_ai/flo_ai/arium/builder.py +++ b/flo_ai/flo_ai/arium/builder.py @@ -7,6 +7,7 @@ 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: @@ -226,12 +227,27 @@ def from_yaml( - 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: my_router # optional + router: content_router # References router defined above - from: validator to: [processor] - from: summarizer @@ -340,6 +356,81 @@ def from_yaml( 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', {}) @@ -386,12 +477,12 @@ def from_yaml( # Find router function router_fn = None if router_name: - if routers and router_name in routers: - router_fn = routers[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 in provided routers dictionary. ' - f'Available routers: {list(routers.keys()) if routers else []}' + 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') @@ -411,6 +502,41 @@ def from_yaml( 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], @@ -428,7 +554,7 @@ def _create_agent_from_direct_config( Agent: Configured agent instance """ from flo_ai.models.base_agent import ReasoningPattern - from flo_ai.llm import OpenAI, Anthropic, Gemini, OllamaLLM + # from flo_ai.llm import OpenAI, Anthropic, Gemini, OllamaLLM # Extract basic configuration name = agent_config['name'] @@ -437,24 +563,7 @@ def _create_agent_from_direct_config( # Configure LLM if 'model' in agent_config and base_llm is None: - model_config = agent_config['model'] - 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(f'Model name must be specified for agent {name}') - - 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}') + llm = AriumBuilder._create_llm_from_config(agent_config['model']) elif base_llm: llm = base_llm else: diff --git a/flo_ai/tests/test_arium_yaml.py b/flo_ai/tests/test_arium_yaml.py index 542615ac..e1935656 100644 --- a/flo_ai/tests/test_arium_yaml.py +++ b/flo_ai/tests/test_arium_yaml.py @@ -330,7 +330,7 @@ def test_from_yaml_missing_router_error(self): with patch('flo_ai.arium.builder.AgentBuilder'): with pytest.raises( ValueError, - match='Router missing_router not found in provided routers dictionary', + match='Router missing_router not found', ): AriumBuilder.from_yaml(yaml_str=yaml_config, routers={})