From c5e600b0a2dc151c2dcf5cd17ee0fba73868c8fb Mon Sep 17 00:00:00 2001 From: vizsatiz Date: Wed, 6 Aug 2025 17:32:09 +0530 Subject: [PATCH 1/8] 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/8] 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 3c3a75f9176d9e296d9fb6199bf5e52a6f866852 Mon Sep 17 00:00:00 2001 From: vizsatiz Date: Fri, 15 Aug 2025 18:12:43 +0530 Subject: [PATCH 3/8] Studio first commit --- .gitignore | 1 + flo_studio/README.md | 354 ++++++++++++++++++ flo_studio/package.json | 54 +++ flo_studio/postcss.config.js | 6 + flo_studio/src/App.css | 45 +++ flo_studio/src/App.tsx | 97 +++++ .../src/components/editors/AgentEditor.tsx | 281 ++++++++++++++ flo_studio/src/components/flow/AgentNode.tsx | 106 ++++++ flo_studio/src/components/flow/FlowCanvas.tsx | 118 ++++++ flo_studio/src/components/flow/ToolNode.tsx | 71 ++++ flo_studio/src/components/sidebar/Sidebar.tsx | 170 +++++++++ flo_studio/src/components/ui/button.tsx | 52 +++ flo_studio/src/components/ui/dialog.tsx | 119 ++++++ flo_studio/src/components/ui/input.tsx | 24 ++ flo_studio/src/components/ui/label.tsx | 23 ++ flo_studio/src/components/ui/select.tsx | 155 ++++++++ flo_studio/src/components/ui/textarea.tsx | 23 ++ flo_studio/src/index.css | 37 ++ flo_studio/src/lib/utils.ts | 6 + flo_studio/src/main.tsx | 10 + flo_studio/src/store/designerStore.ts | 251 +++++++++++++ flo_studio/src/types/agent.ts | 92 +++++ flo_studio/src/types/reactflow.ts | 31 ++ flo_studio/src/utils/yamlExport.ts | 214 +++++++++++ flo_studio/tailwind.config.js | 77 ++++ flo_studio/tsconfig.json | 31 ++ flo_studio/tsconfig.node.json | 10 + flo_studio/vite.config.ts | 13 + index.html | 225 ----------- 29 files changed, 2471 insertions(+), 225 deletions(-) create mode 100644 flo_studio/README.md create mode 100644 flo_studio/package.json create mode 100644 flo_studio/postcss.config.js create mode 100644 flo_studio/src/App.css create mode 100644 flo_studio/src/App.tsx create mode 100644 flo_studio/src/components/editors/AgentEditor.tsx create mode 100644 flo_studio/src/components/flow/AgentNode.tsx create mode 100644 flo_studio/src/components/flow/FlowCanvas.tsx create mode 100644 flo_studio/src/components/flow/ToolNode.tsx create mode 100644 flo_studio/src/components/sidebar/Sidebar.tsx create mode 100644 flo_studio/src/components/ui/button.tsx create mode 100644 flo_studio/src/components/ui/dialog.tsx create mode 100644 flo_studio/src/components/ui/input.tsx create mode 100644 flo_studio/src/components/ui/label.tsx create mode 100644 flo_studio/src/components/ui/select.tsx create mode 100644 flo_studio/src/components/ui/textarea.tsx create mode 100644 flo_studio/src/index.css create mode 100644 flo_studio/src/lib/utils.ts create mode 100644 flo_studio/src/main.tsx create mode 100644 flo_studio/src/store/designerStore.ts create mode 100644 flo_studio/src/types/agent.ts create mode 100644 flo_studio/src/types/reactflow.ts create mode 100644 flo_studio/src/utils/yamlExport.ts create mode 100644 flo_studio/tailwind.config.js create mode 100644 flo_studio/tsconfig.json create mode 100644 flo_studio/tsconfig.node.json create mode 100644 flo_studio/vite.config.ts delete mode 100644 index.html diff --git a/.gitignore b/.gitignore index 2fb783ef..2e0165e9 100644 --- a/.gitignore +++ b/.gitignore @@ -21,3 +21,4 @@ scratch_pad.py *.html usecases/ compare_gemini_outputs_v1.py +node_modules/ diff --git a/flo_studio/README.md b/flo_studio/README.md new file mode 100644 index 00000000..69ee426c --- /dev/null +++ b/flo_studio/README.md @@ -0,0 +1,354 @@ +# Flo AI Studio + +A powerful visual designer for creating YAML-based AI agent workflows. Build complex multi-agent workflows with an intuitive drag-and-drop interface, configure agents with comprehensive forms, set up routing logic, and export everything as production-ready YAML. + +## 🌟 Overview + +Flo AI Studio is a React-based visual editor that makes it easy to design and configure AI workflows for the Flo AI framework. It provides a user-friendly interface for creating complex agent orchestrations without writing code. + +## ✨ Features + +- **🎨 Visual Workflow Design**: Drag-and-drop interface using React Flow +- **🤖 Agent Management**: Create and edit agents with comprehensive configuration forms +- **🔧 Tool Integration**: Add and configure tools for your agents +- **🔀 Router Configuration**: Define custom routing logic between workflow nodes +- **📄 YAML Export**: Generate production-ready YAML configurations +- **📋 Template System**: Quick agent templates for common use cases +- **⚙️ Configuration Management**: Manage available tools, LLMs, and routers +- **💾 State Management**: Robust state management with Zustand +- **🎯 TypeScript**: Fully typed for better development experience + +## 🚀 Quick Start + +### Installation + +```bash +# Navigate to flo_studio directory +cd flo_studio + +# Install dependencies +pnpm install + +# Start development server +pnpm dev +``` + +### Building for Production + +```bash +# Build the application +pnpm build + +# Preview the build +pnpm preview +``` + +## 🏗️ Architecture + +### Project Structure + +``` +flo_studio/ +├── src/ +│ ├── components/ # React components +│ │ ├── editors/ # Modal editors (Agent, Edge, Config) +│ │ ├── flow/ # React Flow components +│ │ ├── sidebar/ # Sidebar components +│ │ ├── toolbar/ # Toolbar components +│ │ └── ui/ # Reusable UI components +│ ├── store/ # Zustand store +│ ├── types/ # TypeScript definitions +│ ├── utils/ # Utility functions +│ └── lib/ # Shared utilities +├── public/ # Static assets +└── dist/ # Build output +``` + +### Key Technologies + +- **React 18** - Modern React with hooks and concurrent features +- **TypeScript** - Type safety and better developer experience +- **React Flow** - Graph visualization and interaction +- **Zustand** - Lightweight state management +- **Tailwind CSS** - Utility-first CSS framework +- **Radix UI** - Accessible component primitives +- **React Hook Form** - Form handling with validation +- **js-yaml** - YAML parsing and generation +- **Vite** - Fast build tool and dev server + +## 🎯 Usage Guide + +### Creating Your First Workflow + +1. **Start the Application** + ```bash + pnpm dev + ``` + +2. **Create an Agent** + - Click the "Agent" button in the toolbar + - Fill in the agent configuration form: + - Name and role + - Job description + - LLM model selection + - Tools (optional) + - Output parser (optional) + +3. **Build the Workflow** + - Drag nodes from the sidebar onto the canvas + - Connect nodes by dragging from output handles to input handles + - Configure edge routers by clicking on connections + +4. **Export Configuration** + - Click "Export" in the toolbar + - Review the generated YAML + - Download or copy the configuration + +### Agent Configuration + +Agents can be configured with: + +- **Basic Information**: Name, role, job description +- **Model Settings**: Provider (OpenAI, Anthropic, etc.), model name, temperature +- **Tools**: Select from available tools or add custom ones +- **Output Parser**: Define structured output schemas +- **Reasoning Pattern**: DIRECT, COT (Chain of Thought), or REACT + +### Workflow Features + +- **Visual Connections**: Drag to connect agents and tools +- **Router Functions**: Configure conditional routing between nodes +- **Start/End Nodes**: Automatically detected based on connections +- **Validation**: Real-time validation of workflow structure + +## 🔧 Configuration + +### Available LLMs + +The studio comes pre-configured with popular LLM providers: + +- **OpenAI**: GPT-4o, GPT-4o-mini +- **Anthropic**: Claude-3.5-Sonnet, Claude-3.5-Haiku +- **Google**: Gemini-2.5-Flash, Gemini-2.5-Pro +- **Ollama**: Llama2, Llama3 (local models) + +### Available Tools + +Default tools include: + +- **calculator** - Mathematical calculations +- **web_search** - Web search functionality +- **file_reader** - File reading and analysis +- **email_sender** - Email sending capabilities +- **text_processor** - Text processing and analysis +- **image_analyzer** - Image analysis and processing + +### Router Functions + +Pre-configured router functions: + +- **default_router** - Simple pass-through routing +- **content_router** - Routes based on content analysis +- **classification_router** - Routes based on classification results +- **sentiment_router** - Routes based on sentiment analysis + +## 📊 YAML Export Format + +The studio generates YAML compatible with the Flo AI framework: + +```yaml +metadata: + name: "My Workflow" + version: "1.0.0" + description: "Generated with Flo AI Studio" + tags: ["flo-ai", "studio-generated"] + +arium: + agents: + - name: "content_analyzer" + role: "Content Analyst" + job: "Analyze content and extract insights" + model: + provider: "openai" + name: "gpt-4o-mini" + settings: + temperature: 0.3 + reasoning_pattern: "COT" + + workflow: + start: "content_analyzer" + edges: + - from: "content_analyzer" + to: ["summarizer"] + end: ["summarizer"] +``` + +## 🔌 Integration with Flo AI + +Use the exported YAML with the Flo AI framework: + +```python +from flo_ai.arium.builder import AriumBuilder + +# Load your exported workflow +builder = AriumBuilder.from_yaml(yaml_file="my-workflow.yaml") + +# Run the workflow +result = await builder.build_and_run(["Your input here"]) +``` + +## 🛠️ Development + +### Adding New Components + +1. **Create Component** + ```typescript + // src/components/MyComponent.tsx + import React from 'react'; + + export const MyComponent: React.FC = () => { + return
My Component
; + }; + ``` + +2. **Add to Store** (if needed) + ```typescript + // src/store/designerStore.ts + interface DesignerState { + // Add new state properties + myNewFeature: boolean; + setMyNewFeature: (value: boolean) => void; + } + ``` + +### Adding New Tool Templates + +Edit the store configuration: + +```typescript +// src/store/designerStore.ts +const defaultConfig: DesignerConfig = { + availableTools: [ + // Add new tools + { name: 'my_tool', description: 'My custom tool' }, + ], + // ... +}; +``` + +### Customizing Themes + +Update CSS variables in `src/index.css`: + +```css +:root { + --primary: 222.2 47.4% 11.2%; + --primary-foreground: 210 40% 98%; + /* Add custom colors */ +} +``` + +## 🐛 Troubleshooting + +### Common Issues + +1. **Build Errors** + - Ensure all dependencies are installed: `pnpm install` + - Clear node_modules and reinstall if needed + +2. **TypeScript Errors** + - Check type definitions in `src/types/` + - Ensure proper imports and exports + +3. **React Flow Issues** + - Verify React Flow version compatibility + - Check node and edge data structures + +### Performance Optimization + +- Use React.memo for heavy components +- Implement virtual scrolling for large workflows +- Optimize store subscriptions with selectors + +## 🚀 Deployment + +### Building for Production + +```bash +# Build the application +pnpm build + +# The dist/ folder contains the built application +``` + +### Deployment Options + +- **Static Hosting**: Deploy `dist/` to Netlify, Vercel, or GitHub Pages +- **Docker**: Create a Docker container with nginx +- **CDN**: Upload to S3 + CloudFront or similar + +### Environment Variables + +Create `.env` files for different environments: + +```bash +# .env.development +VITE_API_URL=http://localhost:3000 + +# .env.production +VITE_API_URL=https://api.myapp.com +``` + +## 📈 Roadmap + +### Phase 1 (Current) +- ✅ Basic visual editor +- ✅ Agent configuration +- ✅ YAML export +- ✅ TypeScript support + +### Phase 2 (Planned) +- [ ] YAML import functionality +- [ ] Workflow validation +- [ ] Advanced routing configuration +- [ ] Template library + +### Phase 3 (Future) +- [ ] Real-time collaboration +- [ ] Workflow simulation +- [ ] Plugin system +- [ ] Cloud deployment + +## 🤝 Contributing + +1. Fork the repository +2. Create a feature branch: `git checkout -b feature/my-feature` +3. Make changes and add tests +4. Commit: `git commit -am 'Add my feature'` +5. Push: `git push origin feature/my-feature` +6. Create a Pull Request + +### Development Guidelines + +- Use TypeScript for all new code +- Follow the existing component structure +- Add proper error handling +- Write meaningful commit messages +- Update documentation as needed + +## 📄 License + +This project is part of the Flo AI framework and follows the same licensing terms. + +## 🙏 Acknowledgments + +- Built for the [Flo AI framework](../flo_ai/) +- Powered by React Flow for graph visualization +- UI components from Radix UI +- Icons from Lucide React + +--- + +**Happy Building! 🚀** + +For more information about the Flo AI framework, check out the [main documentation](../flo_ai/README.md). diff --git a/flo_studio/package.json b/flo_studio/package.json new file mode 100644 index 00000000..c72e4a1a --- /dev/null +++ b/flo_studio/package.json @@ -0,0 +1,54 @@ +{ + "name": "flo-ai-studio", + "version": "1.0.0", + "description": "Visual designer for Flo AI workflows", + "private": true, + "scripts": { + "dev": "vite", + "build": "tsc && vite build", + "lint": "eslint . --ext ts,tsx --report-unused-disable-directives --max-warnings 0", + "preview": "vite preview" + }, + "dependencies": { + "@hookform/resolvers": "^3.3.2", + "@radix-ui/react-dialog": "^1.0.5", + "@radix-ui/react-dropdown-menu": "^2.0.6", + "@radix-ui/react-form": "^0.0.3", + "@radix-ui/react-label": "^2.0.2", + "@radix-ui/react-select": "^2.0.0", + "@radix-ui/react-separator": "^1.0.3", + "@radix-ui/react-slot": "^1.0.2", + "@radix-ui/react-switch": "^1.0.3", + "@radix-ui/react-tabs": "^1.0.4", + "@radix-ui/react-tooltip": "^1.0.7", + "class-variance-authority": "^0.7.0", + "clsx": "^2.0.0", + "js-yaml": "^4.1.0", + "lucide-react": "^0.294.0", + "react": "^18.2.0", + "react-dom": "^18.2.0", + "react-hook-form": "^7.48.2", + "reactflow": "^11.10.4", + "tailwind-merge": "^2.0.0", + "tailwindcss-animate": "^1.0.7", + "zod": "^3.22.4", + "zustand": "^4.4.7" + }, + "devDependencies": { + "@types/js-yaml": "^4.0.9", + "@types/node": "^20.10.4", + "@types/react": "^18.2.43", + "@types/react-dom": "^18.2.17", + "@typescript-eslint/eslint-plugin": "^6.14.0", + "@typescript-eslint/parser": "^6.14.0", + "@vitejs/plugin-react": "^4.2.1", + "autoprefixer": "^10.4.16", + "eslint": "^8.55.0", + "eslint-plugin-react-hooks": "^4.6.0", + "eslint-plugin-react-refresh": "^0.4.5", + "postcss": "^8.4.32", + "tailwindcss": "^3.3.6", + "typescript": "^5.2.2", + "vite": "^5.0.8" + } +} diff --git a/flo_studio/postcss.config.js b/flo_studio/postcss.config.js new file mode 100644 index 00000000..2e7af2b7 --- /dev/null +++ b/flo_studio/postcss.config.js @@ -0,0 +1,6 @@ +export default { + plugins: { + tailwindcss: {}, + autoprefixer: {}, + }, +} diff --git a/flo_studio/src/App.css b/flo_studio/src/App.css new file mode 100644 index 00000000..bdd2e3e9 --- /dev/null +++ b/flo_studio/src/App.css @@ -0,0 +1,45 @@ +@import 'tailwindcss/base'; +@import 'tailwindcss/components'; +@import 'tailwindcss/utilities'; + +/* React Flow Customizations */ +.react-flow__node { + font-family: inherit; +} + +.react-flow__edge { + stroke-width: 2px; +} + +.react-flow__edge.selected { + stroke: #3b82f6; + stroke-width: 3px; +} + +.react-flow__connection-line { + stroke: #3b82f6; + stroke-width: 2px; +} + +.react-flow__handle { + width: 10px; + height: 10px; + border: 2px solid white; +} + +/* Focus styles for better accessibility */ +.react-flow__node:focus, +.react-flow__node:focus-visible { + outline: 2px solid #3b82f6; + outline-offset: 2px; +} + +/* Animation for node hover */ +.react-flow__node { + transition: transform 0.1s ease, box-shadow 0.1s ease; +} + +.react-flow__node:hover { + transform: scale(1.02); + box-shadow: 0 10px 25px -5px rgba(0, 0, 0, 0.1); +} diff --git a/flo_studio/src/App.tsx b/flo_studio/src/App.tsx new file mode 100644 index 00000000..c206acf8 --- /dev/null +++ b/flo_studio/src/App.tsx @@ -0,0 +1,97 @@ +import React, { useState } from 'react'; +import { ReactFlowProvider } from 'reactflow'; +import { useDesignerStore } from '@/store/designerStore'; +import { Button } from '@/components/ui/button'; +import { Plus, Download, Settings, FileText } from 'lucide-react'; +import FlowCanvas from '@/components/flow/FlowCanvas'; +import Sidebar from '@/components/sidebar/Sidebar'; +import AgentEditor from '@/components/editors/AgentEditor'; +import { generateAriumYAML, downloadYAML } from '@/utils/yamlExport'; +import './App.css'; + +// Simplified Config Editor Modal +const ConfigEditorModal: React.FC<{ isOpen: boolean; onClose: () => void }> = ({ isOpen, onClose }) => { + if (!isOpen) return null; + + return ( +
+
+

Configuration

+

+ Configuration editor coming soon! For now, tools and LLMs are pre-configured. +

+ +
+
+ ); +}; + +const ToolbarComponent: React.FC = () => { + const { + openAgentEditor, + nodes, + edges, + workflowName, + workflowDescription, + workflowVersion + } = useDesignerStore(); + + const [isConfigOpen, setIsConfigOpen] = useState(false); + + const handleExport = () => { + const yaml = generateAriumYAML({ + nodes, + edges, + workflowName, + workflowDescription, + workflowVersion, + }); + + downloadYAML(yaml, `${workflowName.replace(/\s+/g, '-').toLowerCase()}.yaml`); + }; + + return ( + <> +
+
+

Flo AI Studio

+
Visual Workflow Designer
+
+
+ + + +
+
+ setIsConfigOpen(false)} /> + + ); +}; + +function App() { + return ( +
+ +
+ + + + +
+ + {/* Modals */} + +
+ ); +} + +export default App; \ No newline at end of file diff --git a/flo_studio/src/components/editors/AgentEditor.tsx b/flo_studio/src/components/editors/AgentEditor.tsx new file mode 100644 index 00000000..7b1dfcae --- /dev/null +++ b/flo_studio/src/components/editors/AgentEditor.tsx @@ -0,0 +1,281 @@ +import React, { useEffect, useState } from 'react'; +import { + Dialog, + DialogContent, + DialogHeader, + DialogTitle, + DialogFooter, +} from '@/components/ui/dialog'; +import { Button } from '@/components/ui/button'; +import { Input } from '@/components/ui/input'; +import { Label } from '@/components/ui/label'; +import { Textarea } from '@/components/ui/textarea'; +import { + Select, + SelectContent, + SelectItem, + SelectTrigger, + SelectValue, +} from '@/components/ui/select'; +import { useDesignerStore } from '@/store/designerStore'; +import { Agent } from '@/types/agent'; + +const AgentEditor: React.FC = () => { + const { + isAgentEditorOpen, + closeAgentEditor, + selectedNode, + updateAgent, + addAgent, + config, + } = useDesignerStore(); + + const [formData, setFormData] = useState({ + name: '', + role: '', + job: '', + provider: 'openai' as const, + modelName: 'gpt-4o-mini', + temperature: 0.7, + maxRetries: 3, + reasoningPattern: 'DIRECT' as const, + tools: [] as string[], + }); + + const [isNewAgent, setIsNewAgent] = useState(false); + + useEffect(() => { + if (selectedNode && selectedNode.type === 'agent') { + const agentData = selectedNode.data as any; + const agent = agentData.agent; + setIsNewAgent(false); + + setFormData({ + name: agent.name || '', + role: agent.role || '', + job: agent.job || '', + provider: agent.model.provider || 'openai', + modelName: agent.model.name || 'gpt-4o-mini', + temperature: agent.settings?.temperature || 0.7, + maxRetries: agent.settings?.max_retries || 3, + reasoningPattern: agent.settings?.reasoning_pattern || 'DIRECT', + tools: agent.tools || [], + }); + } else if (isAgentEditorOpen) { + setIsNewAgent(true); + setFormData({ + name: '', + role: '', + job: '', + provider: 'openai', + modelName: 'gpt-4o-mini', + temperature: 0.7, + maxRetries: 3, + reasoningPattern: 'DIRECT', + tools: [], + }); + } + }, [selectedNode, isAgentEditorOpen]); + + const handleSubmit = (e: React.FormEvent) => { + e.preventDefault(); + + const agent: Agent = { + id: isNewAgent ? `agent_${Date.now()}` : selectedNode!.id, + name: formData.name, + role: formData.role || undefined, + job: formData.job, + model: { + provider: formData.provider, + name: formData.modelName, + }, + settings: { + temperature: formData.temperature, + max_retries: formData.maxRetries, + reasoning_pattern: formData.reasoningPattern, + }, + tools: formData.tools.length > 0 ? formData.tools : undefined, + }; + + if (isNewAgent) { + addAgent(agent, { x: 100, y: 100 }); + } else { + updateAgent(agent.id, agent); + } + + closeAgentEditor(); + }; + + const availableModels = config.availableLLMs.filter( + (llm) => llm.provider === formData.provider + ); + + const handleToolToggle = (toolName: string) => { + setFormData(prev => ({ + ...prev, + tools: prev.tools.includes(toolName) + ? prev.tools.filter(t => t !== toolName) + : [...prev.tools, toolName] + })); + }; + + return ( + + + + + {isNewAgent ? 'Create New Agent' : 'Edit Agent'} + + + +
+
+
+ + setFormData(prev => ({ ...prev, name: e.target.value }))} + placeholder="e.g., Content Analyzer" + required + /> +
+
+ + setFormData(prev => ({ ...prev, role: e.target.value }))} + placeholder="e.g., Data Analyst" + /> +
+
+ +
+ +