diff --git a/flo_ai/docs/plan_execution_guide.md b/flo_ai/docs/plan_execution_guide.md new file mode 100644 index 00000000..f37130a9 --- /dev/null +++ b/flo_ai/docs/plan_execution_guide.md @@ -0,0 +1,123 @@ +# Plan Execution Framework Guide + +The Flo AI framework provides built-in support for plan-and-execute workflows, making it easy to create multi-step, coordinated agent workflows. + +## Quick Start + +### 1. Basic Software Development Workflow + +```python +import asyncio +from flo_ai.llm import OpenAI +from flo_ai.arium.memory import PlanAwareMemory +from flo_ai.arium.llm_router import create_plan_execute_router +from flo_ai.arium import AriumBuilder +from flo_ai.models.plan_agents import create_software_development_agents + +async def main(): + # Setup (3 lines!) + llm = OpenAI(model='gpt-4o', api_key='your-key') + memory = PlanAwareMemory() + agents = create_software_development_agents(memory, llm) + + # Create router + router = create_plan_execute_router( + planner_agent='planner', + executor_agent='developer', + reviewer_agent='reviewer', + additional_agents={'tester': 'Tests implementations'}, + llm=llm, + ) + + # Build workflow + agent_list = list(agents.values()) + arium = ( + AriumBuilder() + .with_memory(memory) + .add_agents(agent_list) + .start_with(agents['planner']) + .add_edge(agents['planner'], agent_list, router) + .add_edge(agents['developer'], agent_list, router) + .add_edge(agents['tester'], agent_list, router) + .add_edge(agents['reviewer'], agent_list, router) + .end_with(agents['reviewer']) + .build() + ) + + # Execute + result = await arium.run(['Create a user authentication API']) + +asyncio.run(main()) +``` + +### 2. Custom Plan Workflow + +```python +from flo_ai.models.plan_agents import PlannerAgent, ExecutorAgent + +# Create custom agents +planner = PlannerAgent(memory, llm, name='planner') +researcher = ExecutorAgent(memory, llm, name='researcher') +analyst = ExecutorAgent(memory, llm, name='analyst') +writer = ExecutorAgent(memory, llm, name='writer') + +# Use the same pattern as above with your custom agents +``` + +## Key Components + +### Plan Agents + +- **`PlannerAgent`**: Creates execution plans automatically +- **`ExecutorAgent`**: Executes plan steps and tracks progress +- **`create_software_development_agents()`**: Pre-configured dev team + +### Plan Tools + +- **`PlanTool`**: Parses and stores execution plans (from `flo_ai.tool.plan_tool`) +- **`StepTool`**: Marks steps as completed (from `flo_ai.tool.plan_tool`) +- **`PlanStatusTool`**: Checks plan progress (from `flo_ai.tool.plan_tool`) + +### Memory + +- **`PlanAwareMemory`**: Stores both conversations and execution plans + +### Router + +- **`create_plan_execute_router()`**: Intelligent routing for plan workflows + +## How It Works + +1. **Planning Phase**: Router sends task to planner agent +2. **Plan Storage**: Planner creates and stores ExecutionPlan in memory +3. **Execution Phase**: Router routes to appropriate agents based on plan steps +4. **Progress Tracking**: Agents mark steps as completed using tools +5. **Completion**: Router detects when all steps are done + +## Plan Format + +Plans are created in this standard format: + +``` +EXECUTION PLAN: [Title] +DESCRIPTION: [Description] + +STEPS: +1. step_1: [Task description] β†’ agent_name +2. step_2: [Task description] β†’ agent_name (depends on: step_1) +3. step_3: [Task description] β†’ agent_name (depends on: step_1, step_2) +``` + +## Benefits + +- **Minimal Code**: Pre-built components handle all the complexity +- **Automatic Plan Management**: Plans are created, stored, and tracked automatically +- **Flexible**: Create custom agents for any domain +- **Robust**: Built-in error handling and progress tracking +- **Reusable**: Tools and agents work across different workflows + +## Examples + +See the `examples/` directory for: +- `fixed_plan_execute_demo.py` - Basic software development workflow +- `custom_plan_execute_demo.py` - Custom research workflow diff --git a/flo_ai/examples/custom_plan_execute_demo.py b/flo_ai/examples/custom_plan_execute_demo.py new file mode 100644 index 00000000..e7499358 --- /dev/null +++ b/flo_ai/examples/custom_plan_execute_demo.py @@ -0,0 +1,127 @@ +""" +Custom Plan-Execute Demo - Creating Your Own Plan Workflows + +This demo shows how to create custom plan-execute workflows +using the framework's plan execution components. +""" + +import asyncio +import os +from flo_ai.llm import OpenAI +from flo_ai.arium.memory import PlanAwareMemory +from flo_ai.arium.llm_router import create_plan_execute_router +from flo_ai.arium import AriumBuilder +from flo_ai.models.plan_agents import PlannerAgent, ExecutorAgent + + +async def main(): + """Custom plan-execute workflow example""" + print('🎯 Custom Plan-Execute Demo') + print('=' * 35) + + # Check API key + api_key = os.getenv('OPENAI_API_KEY') + if not api_key: + print('❌ OPENAI_API_KEY environment variable not set') + return + + # Setup + llm = OpenAI(model='gpt-4o', api_key=api_key) + memory = PlanAwareMemory() + + # Create custom agents for research workflow + planner = PlannerAgent( + memory=memory, + llm=llm, + name='planner', + system_prompt="""You are a research project planner. Create plans for research tasks. + +EXECUTION PLAN: [Title] +DESCRIPTION: [Description] + +STEPS: +1. step_1: [Research task] β†’ researcher +2. step_2: [Analysis task] β†’ analyst (depends on: step_1) +3. step_3: [Writing task] β†’ writer (depends on: step_2) + +Use agents: researcher, analyst, writer +IMPORTANT: After generating the plan, use store_execution_plan to save it.""", + ) + + researcher = ExecutorAgent( + memory=memory, + llm=llm, + name='researcher', + system_prompt="""You are a researcher who gathers information and data. +Check plan status first, then execute research steps thoroughly.""", + ) + + analyst = ExecutorAgent( + memory=memory, + llm=llm, + name='analyst', + system_prompt="""You are an analyst who processes and analyzes research data. +Check plan status first, then execute analysis steps thoroughly.""", + ) + + writer = ExecutorAgent( + memory=memory, + llm=llm, + name='writer', + system_prompt="""You are a writer who creates reports and summaries. +Check plan status first, then execute writing steps thoroughly.""", + ) + + agents = [planner, researcher, analyst, writer] + + # Create router + router = create_plan_execute_router( + planner_agent='planner', + executor_agent='researcher', + reviewer_agent='writer', + additional_agents={'analyst': 'Analyzes research data and findings'}, + llm=llm, + ) + + # Build workflow + arium = ( + AriumBuilder() + .with_memory(memory) + .add_agents(agents) + .start_with(planner) + .add_edge(planner, agents, router) + .add_edge(researcher, agents, router) + .add_edge(analyst, agents, router) + .add_edge(writer, agents, router) + .end_with(writer) + .build() + ) + + # Execute task + task = 'Research the impact of AI on software development productivity' + print(f'πŸ“‹ Task: {task}') + print('πŸ”„ Executing custom research workflow...\n') + + try: + result = await arium.run([task]) + + print('\n' + '=' * 40) + print('πŸŽ‰ CUSTOM WORKFLOW COMPLETED!') + print('=' * 40) + + if result: + final_result = result[-1] if isinstance(result, list) else result + print(f'\nπŸ“„ Final Result:\n{final_result}') + + # Show plan status + current_plan = memory.get_current_plan() + if current_plan: + print(f'\nπŸ“Š Plan: {current_plan.title}') + print(f'βœ… Completed: {current_plan.is_completed()}') + + except Exception as e: + print(f'❌ Error: {e}') + + +if __name__ == '__main__': + asyncio.run(main()) diff --git a/flo_ai/examples/fixed_plan_execute_demo.py b/flo_ai/examples/fixed_plan_execute_demo.py deleted file mode 100644 index a5976c8b..00000000 --- a/flo_ai/examples/fixed_plan_execute_demo.py +++ /dev/null @@ -1,401 +0,0 @@ -""" -Fixed PlanExecuteRouter Demo - Resolves the planner loop issue - -The previous demo had an issue where the planner kept looping because the plan -wasn't being properly stored in PlanAwareMemory. This version fixes that by: - -1. Creating a special PlannerAgent that stores ExecutionPlan objects -2. Ensuring the router can detect when plans are created -3. Proper integration between plan creation and execution phases -""" - -import asyncio -import os -import uuid -import re -from flo_ai.models.agent import Agent -from flo_ai.llm import OpenAI -from flo_ai.arium.memory import PlanAwareMemory, ExecutionPlan, PlanStep, StepStatus -from flo_ai.arium.llm_router import create_plan_execute_router -from flo_ai.arium import AriumBuilder - - -class PlannerAgent(Agent): - """Special agent that creates ExecutionPlan objects and stores them in memory""" - - def __init__(self, memory: PlanAwareMemory, **kwargs): - super().__init__(**kwargs) - self.memory = memory - - async def run(self, input_data, **kwargs): - """Override run to create and store ExecutionPlan after generating plan text""" - - # First, generate the plan text using the normal agent behavior - plan_text = await super().run(input_data, **kwargs) - - # Parse the plan text and create ExecutionPlan object - execution_plan = self._parse_plan_text(plan_text) - - # Store the plan in memory - if execution_plan: - self.memory.add_plan(execution_plan) - print(f'βœ… Plan stored in memory: {execution_plan.title}') - print(f'πŸ“Š Steps: {len(execution_plan.steps)}') - - # Show the plan - for i, step in enumerate(execution_plan.steps, 1): - deps = ( - f" (depends: {', '.join(step.dependencies)})" - if step.dependencies - else '' - ) - print(f' {i}. {step.id}: {step.description} β†’ {step.agent}{deps}') - - return plan_text - - def _parse_plan_text(self, plan_text: str) -> ExecutionPlan: - """Parse LLM-generated plan text into ExecutionPlan object""" - - # Extract title - title_match = re.search(r'EXECUTION PLAN:\s*(.+)', plan_text) - title = title_match.group(1).strip() if title_match else 'Generated Plan' - - # Extract description - desc_match = re.search(r'DESCRIPTION:\s*(.+)', plan_text) - description = desc_match.group(1).strip() if desc_match else 'Execution plan' - - # Extract steps using regex - steps = [] - step_pattern = ( - r'(\d+)\.\s*(\w+):\s*(.+?)\s*β†’\s*(\w+)(?:\s*\(depends on:\s*([^)]+)\))?' - ) - - for match in re.finditer(step_pattern, plan_text, re.MULTILINE): - step_num, step_id, step_desc, agent, deps_str = match.groups() - - dependencies = [] - if deps_str: - dependencies = [dep.strip() for dep in deps_str.split(',')] - - step = PlanStep( - id=step_id, - description=step_desc.strip(), - agent=agent, - dependencies=dependencies, - status=StepStatus.PENDING, - ) - steps.append(step) - - return ExecutionPlan( - id=str(uuid.uuid4()), - title=title, - description=description, - steps=steps, - created_by=self.name, - ) - - -class ExecutorAgent(Agent): - """Special agent that marks steps as completed when executing them""" - - def __init__(self, memory: PlanAwareMemory, **kwargs): - super().__init__(**kwargs) - self.memory = memory - - async def run(self, input_data, **kwargs): - """Override run to update step status after execution""" - - # Get current plan and next steps - current_plan = self.memory.get_current_plan() - if current_plan: - next_steps = current_plan.get_next_steps() - - # Find steps assigned to this agent - my_steps = [step for step in next_steps if step.agent == self.name] - - if my_steps: - step = my_steps[0] # Execute first available step - print(f'⏳ Executing step: {step.id} - {step.description}') - - # Mark step as in progress - step.status = StepStatus.IN_PROGRESS - self.memory.update_plan(current_plan) - - # Execute the step using normal agent behavior - result = await super().run( - f'Execute this step: {step.description}. Context: {input_data}', - **kwargs, - ) - - # Mark step as completed - step.status = StepStatus.COMPLETED - step.result = result - self.memory.update_plan(current_plan) - - print(f'βœ… Completed step: {step.id}') - return result - - # If no steps to execute, just run normally - return await super().run(input_data, **kwargs) - - -async def run_fixed_demo(): - """Run the fixed plan-execute demo""" - print('πŸ”§ Fixed PlanExecuteRouter Demo') - print('=' * 40) - - # Check API key - api_key = os.getenv('OPENAI_API_KEY') - if not api_key: - print('❌ OPENAI_API_KEY environment variable not set') - print(' Set it with: export OPENAI_API_KEY=your_key_here') - return - - print('βœ… OpenAI API key found') - - # Create LLM - llm = OpenAI(model='gpt-4o', api_key=api_key) - - # Create plan-aware memory - memory = PlanAwareMemory() - - # Create special planner that stores ExecutionPlan objects - planner = PlannerAgent( - memory=memory, - name='planner', - system_prompt="""You are a project planner. Create execution plans in this EXACT format: - -EXECUTION PLAN: [Clear Title] -DESCRIPTION: [Brief description] - -STEPS: -1. step_1: [Description] β†’ developer -2. step_2: [Description] β†’ developer (depends on: step_1) -3. step_3: [Description] β†’ tester (depends on: step_2) -4. step_4: [Description] β†’ reviewer (depends on: step_3) - -Use only these agents: developer, tester, reviewer -Keep steps clear and actionable. -Always include dependencies where steps must be done in sequence.""", - llm=llm, - ) - - # Create executor agents that update step status - developer = ExecutorAgent( - memory=memory, - name='developer', - system_prompt="""You are a developer executing implementation steps. -Provide detailed implementation for the given step.""", - llm=llm, - ) - - tester = ExecutorAgent( - memory=memory, - name='tester', - system_prompt="""You are a tester validating implementations. -Create test scenarios and validate the implementation.""", - llm=llm, - ) - - reviewer = Agent( - name='reviewer', - system_prompt="""You are a reviewer providing final validation. -Review all completed work and give final approval.""", - llm=llm, - ) - - print('βœ… Created special agents with plan integration') - - # Create plan-execute router - plan_router = create_plan_execute_router( - planner_agent='planner', - executor_agent='developer', - reviewer_agent='reviewer', - additional_agents={'tester': 'Tests implementations'}, - llm=llm, - ) - - print('βœ… Created PlanExecuteRouter') - - # Create workflow - arium = ( - AriumBuilder() - .with_memory(memory) - .add_agents([planner, developer, tester, reviewer]) - .start_with(planner) - .add_edge(planner, [developer, tester, reviewer, planner], plan_router) - .add_edge(developer, [tester, reviewer, developer, planner], plan_router) - .add_edge(tester, [developer, reviewer, tester, planner], plan_router) - .add_edge(reviewer, [developer, tester, reviewer, planner], plan_router) - .end_with(reviewer) - .build() - ) - - print('βœ… Built Arium workflow') - - # Task to execute - task = 'Create a simple user login API endpoint' - - print(f'\nπŸ“‹ Task: {task}') - print('\nπŸ”„ Starting fixed plan-execute workflow...') - print(' This should now properly:') - print(' 1. Create and store execution plan') - print(' 2. Execute steps sequentially') - print(' 3. Track progress through completion') - - try: - # Execute workflow - result = await arium.run([task]) - - print('\n' + '=' * 50) - print('πŸŽ‰ WORKFLOW COMPLETED!') - print('=' * 50) - - # Show results - if result: - final_result = result[-1] if isinstance(result, list) else result - print('\nπŸ“„ Final Output:') - print('-' * 30) - print(final_result) - - # Show execution plan status - current_plan = memory.get_current_plan() - if current_plan: - print(f'\nπŸ“Š Final Plan Status: {current_plan.title}') - print(f'Completed: {current_plan.is_completed()}') - - print('\nπŸ“‹ Step Status:') - for step in current_plan.steps: - status_icon = { - StepStatus.PENDING: 'β—‹', - StepStatus.IN_PROGRESS: '⏳', - StepStatus.COMPLETED: 'βœ…', - StepStatus.FAILED: '❌', - }.get(step.status, 'β—‹') - print(f' {status_icon} {step.id}: {step.description} β†’ {step.agent}') - if step.result: - print(f' Result: {step.result[:100]}...') - - print('\nπŸ’‘ What was fixed:') - print(' β€’ PlannerAgent now creates and stores ExecutionPlan objects') - print(' β€’ ExecutorAgent updates step status during execution') - print(' β€’ Router can detect when plans exist in memory') - print(' β€’ No more infinite loops in the planner!') - - except Exception as e: - print(f'\n❌ Error: {e}') - print('Check the logs above for more details.') - - -async def run_simple_test(): - """Run a simple test to verify the fix works""" - print('\nπŸ“‹ Simple Plan Creation Test') - print('=' * 35) - - api_key = os.getenv('OPENAI_API_KEY') - if not api_key: - print('❌ OPENAI_API_KEY not set') - return - - # Create memory and planner - memory = PlanAwareMemory() - llm = OpenAI(model='gpt-4o', api_key=api_key) - - planner = PlannerAgent( - memory=memory, - name='planner', - system_prompt="""Create a plan in this format: - -EXECUTION PLAN: [Title] -DESCRIPTION: [Description] - -STEPS: -1. step_1: [Task] β†’ developer -2. step_2: [Task] β†’ tester (depends on: step_1) -3. step_3: [Task] β†’ reviewer (depends on: step_2)""", - llm=llm, - ) - - print('πŸ”„ Testing plan creation and storage...') - - try: - # Create a plan - await planner.run('Create a plan for building a simple calculator API') - - # Check if plan was stored - current_plan = memory.get_current_plan() - if current_plan: - print('βœ… Plan successfully created and stored!') - print(f'βœ… Title: {current_plan.title}') - print(f'βœ… Steps: {len(current_plan.steps)}') - - # Test next steps - next_steps = current_plan.get_next_steps() - print(f'βœ… Ready to execute: {len(next_steps)} steps') - - for step in next_steps: - print(f' β†’ {step.id}: {step.agent}') - else: - print('❌ Plan was not stored in memory') - - except Exception as e: - print(f'❌ Error: {e}') - - -def show_fix_explanation(): - """Explain what was fixed""" - print('\nπŸ”§ What Was Fixed') - print('=' * 20) - - explanation = """ -❌ Problem: Planner Loop -The original demo had the planner stuck in a loop because: -β€’ Planner generated plan text but didn't store ExecutionPlan objects -β€’ Router couldn't detect that a plan existed in memory -β€’ Router kept routing back to planner infinitely - -βœ… Solution: Special Agent Classes -Created specialized agents that integrate with PlanAwareMemory: -β€’ PlannerAgent: Parses plan text and stores ExecutionPlan objects -β€’ ExecutorAgent: Updates step status during execution -β€’ Proper integration between plan creation and execution - -🎯 Key Changes: -1. PlannerAgent.run() now creates and stores ExecutionPlan objects -2. ExecutorAgent.run() marks steps as in-progress/completed -3. Router can detect when plans exist and route accordingly -4. Proper step-by-step execution with status tracking - -πŸš€ Result: -β€’ No more infinite loops -β€’ Proper plan-execute workflow -β€’ Real-time step progress tracking -β€’ Seamless integration with PlanAwareMemory -""" - - print(explanation) - - -async def main(): - """Main demo function""" - print('🎯 Fixed PlanExecuteRouter Demo with Real OpenAI LLM') - print('This version fixes the planner loop issue!\n') - - if not os.getenv('OPENAI_API_KEY'): - print('❌ OPENAI_API_KEY not set') - print(' Set it with: export OPENAI_API_KEY=your_key_here') - return - - # Show what was fixed - show_fix_explanation() - - # Run tests - await run_simple_test() - await run_fixed_demo() - - print('\n\nπŸŽ‰ Fixed Demo Complete!') - print('The PlanExecuteRouter now works properly without infinite loops! πŸš€') - - -if __name__ == '__main__': - asyncio.run(main()) diff --git a/flo_ai/examples/live_plan_execute_demo.py b/flo_ai/examples/live_plan_execute_demo.py deleted file mode 100644 index bf56267f..00000000 --- a/flo_ai/examples/live_plan_execute_demo.py +++ /dev/null @@ -1,307 +0,0 @@ -""" -Live PlanExecuteRouter Demo - Real OpenAI LLM Integration - -A focused demo showing PlanExecuteRouter with actual LLM calls. -Set OPENAI_API_KEY environment variable to run. - -This demonstrates: -1. Automatic task breakdown into execution plans -2. Step-by-step execution with progress tracking -3. Intelligent routing between planning and execution phases -""" - -import asyncio -import os -from flo_ai.models.agent import Agent -from flo_ai.llm import OpenAI -from flo_ai.arium.memory import PlanAwareMemory, ExecutionPlan, PlanStep, StepStatus -from flo_ai.arium.llm_router import create_plan_execute_router -from flo_ai.arium import AriumBuilder - - -class PlanParser: - """Helper to parse LLM-generated plans into ExecutionPlan objects""" - - @staticmethod - def parse_plan_from_text( - plan_text: str, planner_agent: str = 'planner' - ) -> ExecutionPlan: - """Parse structured plan text into ExecutionPlan object""" - import uuid - import re - - # Extract title and description - title_match = re.search(r'EXECUTION PLAN:\s*(.+)', plan_text) - title = title_match.group(1).strip() if title_match else 'Generated Plan' - - desc_match = re.search(r'DESCRIPTION:\s*(.+)', plan_text) - description = desc_match.group(1).strip() if desc_match else 'Execution plan' - - # Extract steps - steps = [] - step_pattern = ( - r'(\d+)\.\s*(\w+):\s*(.+?)\s*β†’\s*(\w+)(?:\s*\(depends on:\s*([^)]+)\))?' - ) - - for match in re.finditer(step_pattern, plan_text): - step_num, step_id, step_desc, agent, deps_str = match.groups() - - dependencies = [] - if deps_str: - dependencies = [dep.strip() for dep in deps_str.split(',')] - - step = PlanStep( - id=step_id, - description=step_desc.strip(), - agent=agent, - dependencies=dependencies, - status=StepStatus.PENDING, - ) - steps.append(step) - - return ExecutionPlan( - id=str(uuid.uuid4()), - title=title, - description=description, - steps=steps, - created_by=planner_agent, - ) - - -async def run_live_demo(): - """Run a live demo with real OpenAI API calls""" - print('πŸš€ Live PlanExecuteRouter Demo') - print('=' * 40) - - # Check API key - api_key = os.getenv('OPENAI_API_KEY') - if not api_key: - print('❌ OPENAI_API_KEY environment variable not set') - print(' Set it with: export OPENAI_API_KEY=your_key_here') - return - - print('βœ… OpenAI API key found') - - # Create LLM - llm = OpenAI(model='gpt-4o-mini', api_key=api_key) - - # Create agents with specific roles - planner = Agent( - name='planner', - system_prompt="""You are a project planner. Create execution plans in this exact format: - -EXECUTION PLAN: [Title] -DESCRIPTION: [Brief description] - -STEPS: -1. step_1: [Description] β†’ [agent] -2. step_2: [Description] β†’ [agent] (depends on: step_1) -3. step_3: [Description] β†’ [agent] (depends on: step_1, step_2) - -Use agents: developer, tester, reviewer -Keep steps clear and actionable.""", - llm=llm, - ) - - developer = Agent( - name='developer', - system_prompt="""You are a developer executing specific implementation steps. -Acknowledge the step you're working on and provide detailed implementation.""", - llm=llm, - ) - - tester = Agent( - name='tester', - system_prompt="""You are a tester validating implementations. -Create test scenarios and report validation results.""", - llm=llm, - ) - - reviewer = Agent( - name='reviewer', - system_prompt="""You are a reviewer providing final validation. -Review completed work and give final approval.""", - llm=llm, - ) - - print('βœ… Created agents: planner, developer, tester, reviewer') - - # Create memory and router - memory = PlanAwareMemory() - - plan_router = create_plan_execute_router( - planner_agent='planner', - executor_agent='developer', - reviewer_agent='reviewer', - additional_agents={'tester': 'Tests implementations'}, - llm=llm, - ) - - print('βœ… Created PlanExecuteRouter with PlanAwareMemory') - - # Create workflow - arium = ( - AriumBuilder() - .with_memory(memory) - .add_agents([planner, developer, tester, reviewer]) - .start_with(planner) - .add_edge(planner, [developer, tester, reviewer, planner], plan_router) - .add_edge(developer, [tester, reviewer, developer, planner], plan_router) - .add_edge(tester, [developer, reviewer, tester, planner], plan_router) - .add_edge(reviewer, [developer, tester, reviewer, planner], plan_router) - .end_with(reviewer) - .build() - ) - - print('βœ… Built Arium workflow') - - # Task to execute - task = 'Create a simple user registration API with email validation' - - print(f'\nπŸ“‹ Task: {task}') - print('\nπŸ”„ Executing plan-and-execute workflow...') - print( - ' Watch as the router coordinates planning β†’ development β†’ testing β†’ review' - ) - - try: - # Execute workflow - result = await arium.run([task]) - - print('\n' + '=' * 50) - print('πŸŽ‰ WORKFLOW COMPLETED!') - print('=' * 50) - - # Show results - if result: - final_result = result[-1] if isinstance(result, list) else result - print('\nπŸ“„ Final Output:') - print('-' * 30) - print(final_result) - - # Show execution plan if created - current_plan = memory.get_current_plan() - if current_plan: - print(f'\nπŸ“Š Execution Plan: {current_plan.title}') - print(f'Description: {current_plan.description}') - print(f'Steps: {len(current_plan.steps)}') - print(f'Completed: {current_plan.is_completed()}') - - print('\nπŸ“‹ Step Progress:') - for step in current_plan.steps: - status_icon = { - StepStatus.PENDING: 'β—‹', - StepStatus.IN_PROGRESS: '⏳', - StepStatus.COMPLETED: 'βœ…', - StepStatus.FAILED: '❌', - }.get(step.status, 'β—‹') - deps = ( - f" (deps: {', '.join(step.dependencies)})" - if step.dependencies - else '' - ) - print( - f' {status_icon} {step.id}: {step.description} β†’ {step.agent}{deps}' - ) - - print('\nπŸ’‘ What happened:') - print(' β€’ Router started by routing to planner') - print(' β€’ Planner created detailed execution plan') - print(' β€’ Router routed to developer for implementation steps') - print(' β€’ Router routed to tester for validation') - print(' β€’ Router routed to reviewer for final approval') - print(' β€’ Plan state tracked in PlanAwareMemory throughout') - - except Exception as e: - print(f'\n❌ Error: {e}') - print('This could be due to API rate limits or network issues.') - - -async def run_simple_plan_creation(): - """Demo just the plan creation part""" - print('\n\nπŸ“‹ Plan Creation Demo') - print('=' * 30) - - api_key = os.getenv('OPENAI_API_KEY') - if not api_key: - print('❌ OPENAI_API_KEY not set') - return - - llm = OpenAI(model='gpt-4o-mini', api_key=api_key) - - planner = Agent( - name='planner', - system_prompt="""Create an execution plan in this format: - -EXECUTION PLAN: [Title] -DESCRIPTION: [Description] - -STEPS: -1. step_1: [Task description] β†’ developer -2. step_2: [Task description] β†’ developer (depends on: step_1) -3. step_3: [Task description] β†’ tester (depends on: step_2) -4. step_4: [Task description] β†’ reviewer (depends on: step_3)""", - llm=llm, - ) - - task = 'Build a simple blog API with posts and comments' - print(f'Task: {task}') - print('\nπŸ”„ Generating execution plan...') - - try: - plan_text = await planner.run(f'Create a detailed execution plan for: {task}') - print('\nπŸ“‹ Generated Plan:') - print('-' * 30) - print(plan_text) - - # Parse into ExecutionPlan object - execution_plan = PlanParser.parse_plan_from_text(plan_text) - - print('\nβœ… Parsed into ExecutionPlan:') - print(f'Title: {execution_plan.title}') - print(f'Steps: {len(execution_plan.steps)}') - - for step in execution_plan.steps: - deps = ( - f" (depends: {', '.join(step.dependencies)})" - if step.dependencies - else '' - ) - print(f' β€’ {step.id}: {step.description} β†’ {step.agent}{deps}') - - except Exception as e: - print(f'❌ Error: {e}') - - -def show_setup_instructions(): - """Show how to set up and run the demo""" - print('πŸ“– Setup Instructions') - print('=' * 25) - print('1. Set your OpenAI API key:') - print(' export OPENAI_API_KEY=your_api_key_here') - print('\n2. Run the demo:') - print(' python examples/live_plan_execute_demo.py') - print('\n3. Watch the PlanExecuteRouter in action! πŸš€') - - -async def main(): - """Main demo function""" - print('🎯 Live PlanExecuteRouter Demo with Real OpenAI LLM') - print('This demo shows the full plan-and-execute workflow with actual API calls!\n') - - if not os.getenv('OPENAI_API_KEY'): - show_setup_instructions() - return - - # Run demos - await run_simple_plan_creation() - await run_live_demo() - - print('\n\nπŸŽ‰ Demo Complete!') - print( - 'You just saw PlanExecuteRouter break down a task and execute it step by step! πŸš€' - ) - - -if __name__ == '__main__': - asyncio.run(main()) diff --git a/flo_ai/flo_ai/models/agent.py b/flo_ai/flo_ai/models/agent.py index 00682d71..22f8b526 100644 --- a/flo_ai/flo_ai/models/agent.py +++ b/flo_ai/flo_ai/models/agent.py @@ -190,7 +190,6 @@ async def _run_with_tools( while tool_call_count < max_tool_calls: formatted_tools = self.llm.format_tools_for_llm(self.tools) - print(messages) response = await self.llm.generate( messages, functions=formatted_tools, diff --git a/flo_ai/flo_ai/models/plan_agents.py b/flo_ai/flo_ai/models/plan_agents.py new file mode 100644 index 00000000..d49df782 --- /dev/null +++ b/flo_ai/flo_ai/models/plan_agents.py @@ -0,0 +1,220 @@ +""" +Plan Execution Agents for Flo AI Framework + +This module provides standard agent classes for plan-based execution patterns, +making it easy to create plan-and-execute workflows. +""" + +from typing import List, Optional +from flo_ai.models.agent import Agent +from flo_ai.llm.base_llm import BaseLLM +from flo_ai.arium.memory import PlanAwareMemory +from flo_ai.tool.plan_tool import PlanTool, StepTool, PlanStatusTool + + +class PlannerAgent(Agent): + """ + Agent specialized for creating execution plans. + + Automatically equipped with tools to store plans in PlanAwareMemory. + """ + + def __init__( + self, + memory: PlanAwareMemory, + llm: BaseLLM, + name: str = 'planner', + system_prompt: Optional[str] = None, + **kwargs, + ): + """ + Initialize the planner agent. + + Args: + memory: PlanAwareMemory instance to store plans in + llm: LLM instance for the agent + name: Agent name (default: "planner") + system_prompt: Custom system prompt, or uses default if None + **kwargs: Additional arguments for Agent + """ + + # Default system prompt for planners + if system_prompt is None: + system_prompt = """You are a project planner. Create execution plans in this EXACT format: + +EXECUTION PLAN: [Clear Title] +DESCRIPTION: [Brief description] + +STEPS: +1. step_1: [Description] β†’ [agent_name] +2. step_2: [Description] β†’ [agent_name] (depends on: step_1) +3. step_3: [Description] β†’ [agent_name] (depends on: step_1, step_2) + +Rules: +- Use clear, actionable step descriptions +- Assign steps to appropriate agent names +- Include dependencies where steps must be done in sequence +- Keep step IDs simple (step_1, step_2, etc.) + +IMPORTANT: After generating the plan text, use the store_execution_plan tool to save it.""" + + # Create plan tool + plan_tool = PlanTool(memory) + plan_status_tool = PlanStatusTool(memory) + + super().__init__( + name=name, + system_prompt=system_prompt, + llm=llm, + tools=[plan_tool, plan_status_tool], + **kwargs, + ) + + +class ExecutorAgent(Agent): + """ + Agent specialized for executing plan steps. + + Automatically equipped with tools to mark steps as completed in PlanAwareMemory. + """ + + def __init__( + self, + memory: PlanAwareMemory, + llm: BaseLLM, + name: str, + system_prompt: Optional[str] = None, + **kwargs, + ): + """ + Initialize the executor agent. + + Args: + memory: PlanAwareMemory instance to update plans in + llm: LLM instance for the agent + name: Agent name (must match agent names used in plans) + system_prompt: Custom system prompt, or uses default if None + **kwargs: Additional arguments for Agent + """ + + # Default system prompt for executors + if system_prompt is None: + system_prompt = f"""You are a {name} executing specific steps from an execution plan. + +Process: +1. Check the current execution plan status using check_plan_status +2. Look for steps assigned to "{name}" that are ready to execute +3. Execute the step and provide detailed results +4. Use complete_step tool to mark the step as completed with your results + +Focus on providing high-quality, detailed work for each step you execute.""" + + # Create step tools + step_tool = StepTool(memory, name) + plan_status_tool = PlanStatusTool(memory) + + super().__init__( + name=name, + system_prompt=system_prompt, + llm=llm, + tools=[step_tool, plan_status_tool], + **kwargs, + ) + + +def create_plan_execution_agents( + memory: PlanAwareMemory, + llm: BaseLLM, + executor_agents: List[str], + planner_name: str = 'planner', +) -> dict: + """ + Factory function to create a complete set of plan execution agents. + + Args: + memory: PlanAwareMemory instance + llm: LLM instance for all agents + executor_agents: List of executor agent names (e.g., ["developer", "tester", "reviewer"]) + planner_name: Name for the planner agent + + Returns: + Dict mapping agent names to agent instances + """ + agents = {} + + # Create planner + agents[planner_name] = PlannerAgent(memory=memory, llm=llm, name=planner_name) + + # Create executors + for agent_name in executor_agents: + agents[agent_name] = ExecutorAgent(memory=memory, llm=llm, name=agent_name) + + return agents + + +def create_software_development_agents(memory: PlanAwareMemory, llm: BaseLLM) -> dict: + """ + Create a standard set of agents for software development workflows. + + Args: + memory: PlanAwareMemory instance + llm: LLM instance for all agents + + Returns: + Dict with planner, developer, tester, and reviewer agents + """ + agents = {} + + # Planner with software development focus + agents['planner'] = PlannerAgent( + memory=memory, + llm=llm, + name='planner', + system_prompt="""You are a software development project planner. Create execution plans for development tasks. + +EXECUTION PLAN: [Clear Title] +DESCRIPTION: [Brief description] + +STEPS: +1. step_1: [Development task] β†’ developer +2. step_2: [Development task] β†’ developer (depends on: step_1) +3. step_3: [Testing task] β†’ tester (depends on: step_2) +4. step_4: [Review task] β†’ reviewer (depends on: step_3) + +Use these agents: developer, tester, reviewer +Focus on breaking down development tasks into logical, sequential steps. + +IMPORTANT: After generating the plan, use store_execution_plan to save it.""", + ) + + # Developer + agents['developer'] = ExecutorAgent( + memory=memory, + llm=llm, + name='developer', + system_prompt="""You are a software developer executing implementation steps. +Provide detailed code implementations, technical designs, and development work. +Always check the plan status first, then execute your assigned steps thoroughly.""", + ) + + # Tester + agents['tester'] = ExecutorAgent( + memory=memory, + llm=llm, + name='tester', + system_prompt="""You are a QA tester validating implementations. +Create comprehensive test scenarios, validate functionality, and report results. +Always check the plan status first, then execute your assigned testing steps.""", + ) + + # Reviewer + agents['reviewer'] = ExecutorAgent( + memory=memory, + llm=llm, + name='reviewer', + system_prompt="""You are a code reviewer providing final validation. +Review completed work, check quality, and provide final approval or feedback. +Always check the plan status first, then execute your assigned review steps.""", + ) + + return agents diff --git a/flo_ai/flo_ai/tool/plan_tool.py b/flo_ai/flo_ai/tool/plan_tool.py new file mode 100644 index 00000000..73770b00 --- /dev/null +++ b/flo_ai/flo_ai/tool/plan_tool.py @@ -0,0 +1,226 @@ +""" +Plan Execution Tools for Flo AI Framework + +This module provides reusable tools for plan-based execution patterns, +enabling agents to create, store, and manage execution plans automatically. +""" + +import uuid +import re +from flo_ai.tool.base_tool import Tool +from flo_ai.arium.memory import PlanAwareMemory, ExecutionPlan, PlanStep, StepStatus + + +class PlanTool(Tool): + """Tool for creating and storing execution plans in PlanAwareMemory""" + + def __init__(self, memory: PlanAwareMemory): + """ + Initialize the plan tool. + + Args: + memory: PlanAwareMemory instance to store plans in + """ + self.memory = memory + super().__init__( + name='store_execution_plan', + description='Create and store an execution plan from plan text. Use this after generating a plan.', + function=self._execute_plan_storage, + parameters={ + 'plan_text': { + 'type': 'string', + 'description': 'The generated plan text in the required format', + } + }, + ) + + async def _execute_plan_storage(self, plan_text: str) -> str: + """Parse plan text and store ExecutionPlan object in memory""" + try: + execution_plan = self._parse_plan_text(plan_text) + + if execution_plan: + self.memory.add_plan(execution_plan) + + plan_summary = f'βœ… Plan stored: {execution_plan.title}\n' + plan_summary += f'πŸ“Š Steps: {len(execution_plan.steps)}\n' + + for i, step in enumerate(execution_plan.steps, 1): + deps = ( + f" (depends: {', '.join(step.dependencies)})" + if step.dependencies + else '' + ) + plan_summary += ( + f' {i}. {step.id}: {step.description} β†’ {step.agent}{deps}\n' + ) + + return plan_summary + else: + return '❌ Failed to parse plan text into ExecutionPlan' + + except Exception as e: + return f'❌ Error storing plan: {str(e)}' + + def _parse_plan_text(self, plan_text: str) -> ExecutionPlan: + """Parse LLM-generated plan text into ExecutionPlan object""" + + # Extract title + title_match = re.search(r'EXECUTION PLAN:\s*(.+)', plan_text) + title = title_match.group(1).strip() if title_match else 'Generated Plan' + + # Extract description + desc_match = re.search(r'DESCRIPTION:\s*(.+)', plan_text) + description = desc_match.group(1).strip() if desc_match else 'Execution plan' + + # Extract steps using regex + steps = [] + step_pattern = ( + r'(\d+)\.\s*(\w+):\s*(.+?)\s*β†’\s*(\w+)(?:\s*\(depends on:\s*([^)]+)\))?' + ) + + for match in re.finditer(step_pattern, plan_text, re.MULTILINE): + step_num, step_id, step_desc, agent, deps_str = match.groups() + + dependencies = [] + if deps_str: + dependencies = [dep.strip() for dep in deps_str.split(',')] + + step = PlanStep( + id=step_id, + description=step_desc.strip(), + agent=agent, + dependencies=dependencies, + status=StepStatus.PENDING, + ) + steps.append(step) + + return ExecutionPlan( + id=str(uuid.uuid4()), + title=title, + description=description, + steps=steps, + created_by='planner', + ) + + +class StepTool(Tool): + """Tool for marking execution steps as completed""" + + def __init__(self, memory: PlanAwareMemory, agent_name: str): + """ + Initialize the step tool. + + Args: + memory: PlanAwareMemory instance to update plans in + agent_name: Name of the agent this tool belongs to + """ + self.memory = memory + self.agent_name = agent_name + super().__init__( + name='complete_step', + description='Mark a plan step as completed after executing it', + function=self._execute_step_completion, + parameters={ + 'step_id': { + 'type': 'string', + 'description': 'The ID of the step that was completed', + }, + 'result': { + 'type': 'string', + 'description': 'The result or output of completing the step', + }, + }, + ) + + async def _execute_step_completion(self, step_id: str, result: str) -> str: + """Mark a step as completed and store the result""" + try: + current_plan = self.memory.get_current_plan() + if not current_plan: + return '❌ No current plan found' + + step = current_plan.get_step(step_id) + if not step: + return f'❌ Step {step_id} not found in current plan' + + if step.agent != self.agent_name: + return f'❌ Step {step_id} is assigned to {step.agent}, not {self.agent_name}' + + # Mark step as completed + step.status = StepStatus.COMPLETED + step.result = result + self.memory.update_plan(current_plan) + + return f'βœ… Step {step_id} marked as completed' + + except Exception as e: + return f'❌ Error completing step: {str(e)}' + + +class PlanStatusTool(Tool): + """Tool for checking the current plan status and next steps""" + + def __init__(self, memory: PlanAwareMemory): + """ + Initialize the plan status tool. + + Args: + memory: PlanAwareMemory instance to check + """ + self.memory = memory + super().__init__( + name='check_plan_status', + description='Check the current execution plan status and get next available steps', + function=self._execute_status_check, + parameters={}, + ) + + async def _execute_status_check(self) -> str: + """Get current plan status and next steps""" + try: + current_plan = self.memory.get_current_plan() + if not current_plan: + return '❌ No current execution plan found' + + # Get plan overview + status_info = f'πŸ“‹ Current Plan: {current_plan.title}\n' + status_info += f'πŸ“ Description: {current_plan.description}\n' + status_info += f'βœ… Completed: {current_plan.is_completed()}\n' + + # Get next steps + next_steps = current_plan.get_next_steps() + if next_steps: + status_info += f'\n🎯 Next Steps ({len(next_steps)} available):\n' + for step in next_steps: + deps = ( + f" (depends: {', '.join(step.dependencies)})" + if step.dependencies + else '' + ) + status_info += ( + f' β€’ {step.id}: {step.description} β†’ {step.agent}{deps}\n' + ) + else: + if current_plan.is_completed(): + status_info += '\nπŸŽ‰ All steps completed!' + else: + status_info += '\n⏳ Waiting for dependencies to complete' + + # Show all steps with status + status_info += '\nπŸ“Š All Steps:\n' + for step in current_plan.steps: + status_icon = { + StepStatus.PENDING: 'β—‹', + StepStatus.IN_PROGRESS: '⏳', + StepStatus.COMPLETED: 'βœ…', + StepStatus.FAILED: '❌', + }.get(step.status, 'β—‹') + status_info += ( + f' {status_icon} {step.id}: {step.description} β†’ {step.agent}\n' + ) + + return status_info + + except Exception as e: + return f'❌ Error checking plan status: {str(e)}' diff --git a/studio/src/components/drawer/YamlPreviewDrawer.tsx b/studio/src/components/drawer/YamlPreviewDrawer.tsx index 0f4bf3fc..0f8dc251 100644 --- a/studio/src/components/drawer/YamlPreviewDrawer.tsx +++ b/studio/src/components/drawer/YamlPreviewDrawer.tsx @@ -9,7 +9,7 @@ const YamlPreviewDrawer: React.FC = () => { const [yamlContent, setYamlContent] = useState(''); const [copied, setCopied] = useState(false); - const { nodes, edges, workflowName, workflowDescription, workflowVersion } = useDesignerStore(); + const { nodes, edges, workflowName, workflowDescription, workflowVersion, startNodeId, endNodeIds } = useDesignerStore(); // Update YAML content when workflow changes useEffect(() => { @@ -21,6 +21,8 @@ const YamlPreviewDrawer: React.FC = () => { workflowName, workflowDescription, workflowVersion, + startNodeId, + endNodeIds, }); setYamlContent(yaml); } catch (error) { @@ -29,7 +31,7 @@ const YamlPreviewDrawer: React.FC = () => { } else { setYamlContent('# Create agents and connect them to see YAML preview\n# Drag agents from the sidebar to get started'); } - }, [nodes, edges, workflowName, workflowDescription, workflowVersion]); + }, [nodes, edges, workflowName, workflowDescription, workflowVersion, startNodeId, endNodeIds]); const handleExport = () => { if (nodes.length > 0) { @@ -39,6 +41,8 @@ const YamlPreviewDrawer: React.FC = () => { workflowName, workflowDescription, workflowVersion, + startNodeId, + endNodeIds, }); downloadYAML(yaml, `${workflowName.replace(/\s+/g, '-').toLowerCase()}.yaml`); } diff --git a/studio/src/components/editors/AgentEditor.tsx b/studio/src/components/editors/AgentEditor.tsx index 7b1dfcae..41f55c5f 100644 --- a/studio/src/components/editors/AgentEditor.tsx +++ b/studio/src/components/editors/AgentEditor.tsx @@ -17,8 +17,10 @@ import { SelectTrigger, SelectValue, } from '@/components/ui/select'; +import { Play, Square, Flag } from 'lucide-react'; import { useDesignerStore } from '@/store/designerStore'; import { Agent } from '@/types/agent'; +import { cn } from '@/lib/utils'; const AgentEditor: React.FC = () => { const { @@ -28,6 +30,10 @@ const AgentEditor: React.FC = () => { updateAgent, addAgent, config, + startNodeId, + endNodeIds, + setStartNode, + toggleEndNode, } = useDesignerStore(); const [formData, setFormData] = useState({ @@ -119,6 +125,22 @@ const AgentEditor: React.FC = () => { })); }; + const currentNodeId = selectedNode?.id; + const isCurrentStart = currentNodeId === startNodeId; + const isCurrentEnd = currentNodeId ? endNodeIds.includes(currentNodeId) : false; + + const handleSetStart = () => { + if (currentNodeId && !isNewAgent) { + setStartNode(currentNodeId); + } + }; + + const handleToggleEnd = () => { + if (currentNodeId && !isNewAgent) { + toggleEndNode(currentNodeId); + } + }; + return ( @@ -264,6 +286,61 @@ const AgentEditor: React.FC = () => { + {/* Workflow Start/End Controls */} + {!isNewAgent && ( +
+ +
+
+ + +

+ Only one start node allowed per workflow +

+
+
+ + +

+ Multiple end nodes allowed +

+
+
+
+ )} + - - - + + )} ); }; diff --git a/studio/src/components/flow/SmoothEdge.tsx b/studio/src/components/flow/SmoothEdge.tsx new file mode 100644 index 00000000..e4ebeb21 --- /dev/null +++ b/studio/src/components/flow/SmoothEdge.tsx @@ -0,0 +1,107 @@ +import React from 'react'; +import { + BaseEdge, + EdgeLabelRenderer, + EdgeProps, + getSmoothStepPath, +} from 'reactflow'; +import { Route } from 'lucide-react'; + +import { CustomEdgeData } from '@/types/reactflow'; + +const SmoothEdge: React.FC> = ({ + id, + sourceX, + sourceY, + targetX, + targetY, + sourcePosition, + targetPosition, + data, + selected, +}) => { + const [edgePath, labelX, labelY] = getSmoothStepPath({ + sourceX, + sourceY, + sourcePosition, + targetX, + targetY, + targetPosition, + borderRadius: 20, + offset: 20, + }); + + // Generate unique marker ID for this edge + const markerId = `arrow-smooth-${id}`; + + + + // Determine edge color based on router type + const getEdgeColor = () => { + if (!data?.router) return selected ? '#3b82f6' : '#64748b'; + + if (data.router.includes('reflection')) return selected ? '#9333ea' : '#a855f7'; + if (data.router.includes('plan_execute')) return selected ? '#ea580c' : '#f97316'; + return selected ? '#059669' : '#10b981'; + }; + + const edgeColor = getEdgeColor(); + + return ( + <> + {/* Arrow marker definition */} + + + + + + + + + {/* Only show label when there's a router */} + {data?.router && ( + +
+
+ + + {data.router.replace(/_/g, ' ')} + +
+
+ + )} + + ); +}; + +export default SmoothEdge; diff --git a/studio/src/components/panels/ValidationPanel.tsx b/studio/src/components/panels/ValidationPanel.tsx index ebc68a25..78f6a71a 100644 --- a/studio/src/components/panels/ValidationPanel.tsx +++ b/studio/src/components/panels/ValidationPanel.tsx @@ -5,18 +5,18 @@ import { validateWorkflow, ValidationResult, ValidationIssue, getValidationSumma import { cn } from '@/lib/utils'; const ValidationPanel: React.FC = () => { - const { nodes, edges, setSelectedNode, setSelectedEdge } = useDesignerStore(); + const { nodes, edges, setSelectedNode, setSelectedEdge, startNodeId, endNodeIds } = useDesignerStore(); const [validationResult, setValidationResult] = useState(null); const [expandedSections, setExpandedSections] = useState>(new Set(['errors'])); useEffect(() => { if (nodes.length > 0 || edges.length > 0) { - const result = validateWorkflow(nodes, edges); + const result = validateWorkflow(nodes, edges, startNodeId, endNodeIds); setValidationResult(result); } else { setValidationResult(null); } - }, [nodes, edges]); + }, [nodes, edges, startNodeId, endNodeIds]); const toggleSection = (section: string) => { const newExpanded = new Set(expandedSections); diff --git a/studio/src/components/sidebar/Sidebar.tsx b/studio/src/components/sidebar/Sidebar.tsx index 7dc16d12..c7e98f11 100644 --- a/studio/src/components/sidebar/Sidebar.tsx +++ b/studio/src/components/sidebar/Sidebar.tsx @@ -6,7 +6,17 @@ import { useDesignerStore } from '@/store/designerStore'; import { Agent, Tool, Router } from '@/types/agent'; const Sidebar: React.FC = () => { - const { config, addAgent, addTool, addRouter, openAgentEditor, openRouterEditor } = useDesignerStore(); + const { + config, + addAgent, + addTool, + addRouter, + openAgentEditor, + openRouterEditor, + onConnect, + setStartNode, + addEndNode + } = useDesignerStore(); const [searchTerm, setSearchTerm] = useState(''); const filteredTools = config.availableTools.filter((tool) => @@ -57,6 +67,34 @@ const Sidebar: React.FC = () => { model: { provider: 'openai' as const, name: 'gpt-4o' }, tools: ['web_search'], }, + planner: { + name: 'Planner Agent', + role: 'Task Planner', + job: 'You are an expert project planner who breaks down complex tasks into detailed, sequential steps. Create comprehensive execution plans with clear dependencies and assigned agents. When given a task, analyze it thoroughly and create a structured plan that can be executed step by step. Focus on logical sequencing, resource allocation, and clear deliverables for each step.', + model: { provider: 'openai' as const, name: 'gpt-4o-mini' }, + settings: { temperature: 0.3, reasoning_pattern: 'COT' as const }, + }, + executor: { + name: 'Executor Agent', + role: 'Task Executor', + job: 'You are a task executor who implements plans step by step. Execute development tasks, research activities, or any assigned work with precision and attention to detail. Report on progress, identify blockers, and ensure each step is completed thoroughly before moving to the next. Focus on delivering high-quality results that meet the specified requirements.', + model: { provider: 'openai' as const, name: 'gpt-4o-mini' }, + settings: { temperature: 0.5, reasoning_pattern: 'DIRECT' as const }, + }, + critic: { + name: 'Critic Agent', + role: 'Quality Critic', + job: 'You are a constructive critic who reviews work and provides actionable feedback. Analyze outputs for accuracy, completeness, clarity, and areas for improvement. Provide specific, detailed feedback that helps improve the quality of work. Focus on: logical consistency, thoroughness, clarity of communication, and alignment with requirements. Be constructive and specific in your criticism.', + model: { provider: 'openai' as const, name: 'gpt-4o-mini' }, + settings: { temperature: 0.3, reasoning_pattern: 'COT' as const }, + }, + finalizer: { + name: 'Finalizer Agent', + role: 'Quality Finalizer', + job: 'You are the final agent responsible for polishing and finalizing work. Take refined outputs and ensure they meet high quality standards. Format professionally, add final touches, ensure consistency, and prepare the final deliverable. Focus on presentation, completeness, and professional polish while maintaining the core substance of the work.', + model: { provider: 'openai' as const, name: 'gpt-4o-mini' }, + settings: { temperature: 0.5, reasoning_pattern: 'DIRECT' as const }, + }, }; const template_config = templates[template as keyof typeof templates]; @@ -141,6 +179,195 @@ const Sidebar: React.FC = () => { } }; + const createPlanExecuteWorkflow = () => { + const timestamp = Date.now(); + + // Create planner agent + const planner: Agent = { + id: `planner_${timestamp}`, + name: 'Planner', + role: 'Task Planner', + job: 'You are an expert project planner who breaks down complex tasks into detailed, sequential steps. Create comprehensive execution plans with clear dependencies and assigned agents. When given a task, analyze it thoroughly and create a structured plan that can be executed step by step.', + model: { provider: 'openai', name: 'gpt-4o-mini' }, + settings: { temperature: 0.3, reasoning_pattern: 'COT' }, + }; + + // Create executor agent + const executor: Agent = { + id: `executor_${timestamp}`, + name: 'Executor', + role: 'Task Executor', + job: 'You are a task executor who implements plans step by step. Execute development tasks, research activities, or any assigned work with precision and attention to detail. Report on progress and ensure each step is completed thoroughly.', + model: { provider: 'openai', name: 'gpt-4o-mini' }, + settings: { temperature: 0.5, reasoning_pattern: 'DIRECT' }, + }; + + // Create reviewer agent + const reviewer: Agent = { + id: `reviewer_${timestamp}`, + name: 'Reviewer', + role: 'Quality Reviewer', + job: 'You are a quality reviewer who validates completed work. Review implementations for correctness, completeness, and quality. Provide final approval or request improvements. Ensure all requirements are met.', + model: { provider: 'openai', name: 'gpt-4o-mini' }, + settings: { temperature: 0.2, reasoning_pattern: 'COT' }, + }; + + // Create plan-execute router + const planExecuteRouter: Router = { + id: `plan_execute_router_${timestamp}`, + name: 'Plan Execute Router', + description: 'Coordinates plan-and-execute workflow pattern', + type: 'plan_execute', + model: { provider: 'openai', name: 'gpt-4o-mini' }, + settings: { temperature: 0.2, fallback_strategy: 'first' }, + }; + + // Add agents and router to the workflow + addAgent(planner, { x: 100, y: 100 }); + addAgent(executor, { x: 400, y: 100 }); + addAgent(reviewer, { x: 700, y: 100 }); + addRouter(planExecuteRouter, { x: 400, y: 250 }); + + // Create edges for plan-execute pattern (agents connect through router) + setTimeout(() => { + // Agents connect through the router - router coordinates the workflow + + // Planner → Router + onConnect({ + source: `planner_${timestamp}`, + target: `plan_execute_router_${timestamp}`, + sourceHandle: null, + targetHandle: null, + }); + + // Router → Executor + onConnect({ + source: `plan_execute_router_${timestamp}`, + target: `executor_${timestamp}`, + sourceHandle: null, + targetHandle: null, + }); + + // Router → Reviewer + onConnect({ + source: `plan_execute_router_${timestamp}`, + target: `reviewer_${timestamp}`, + sourceHandle: null, + targetHandle: null, + }); + + // Executor → Router (for coordination) + onConnect({ + source: `executor_${timestamp}`, + target: `plan_execute_router_${timestamp}`, + sourceHandle: null, + targetHandle: null, + }); + + // Set initial start and end nodes for plan-execute workflow + setStartNode(`planner_${timestamp}`); + addEndNode(`reviewer_${timestamp}`); + + // The router connections will be converted to proper YAML workflow edges + // with router fields by the YAML export logic + }, 100); + }; + + const createReflectionWorkflow = () => { + const timestamp = Date.now(); + + // Create main agent + const mainAgent: Agent = { + id: `main_agent_${timestamp}`, + name: 'Main Agent', + role: 'Main Agent', + job: 'You are the main agent responsible for analyzing tasks and creating initial solutions. When you receive input, analyze it thoroughly and provide an initial response. If you receive feedback from the critic, incorporate it to improve your work.', + model: { provider: 'openai', name: 'gpt-4o-mini' }, + settings: { temperature: 0.7, reasoning_pattern: 'COT' }, + }; + + // Create critic agent + const critic: Agent = { + id: `critic_${timestamp}`, + name: 'Critic', + role: 'Quality Critic', + job: 'You are a constructive critic who reviews work and provides actionable feedback. Analyze outputs for accuracy, completeness, clarity, and areas for improvement. Provide specific, detailed feedback that helps improve the quality of work.', + model: { provider: 'openai', name: 'gpt-4o-mini' }, + settings: { temperature: 0.3, reasoning_pattern: 'COT' }, + }; + + // Create finalizer agent + const finalizer: Agent = { + id: `finalizer_${timestamp}`, + name: 'Finalizer', + role: 'Quality Finalizer', + job: 'You are the final agent responsible for polishing and finalizing work. Take refined outputs and ensure they meet high quality standards. Format professionally and prepare the final deliverable.', + model: { provider: 'openai', name: 'gpt-4o-mini' }, + settings: { temperature: 0.5, reasoning_pattern: 'DIRECT' }, + }; + + // Create reflection router + const reflectionRouter: Router = { + id: `reflection_router_${timestamp}`, + name: 'Reflection Router', + description: 'Coordinates A→B→A→C reflection workflow pattern', + type: 'reflection', + model: { provider: 'openai', name: 'gpt-4o-mini' }, + settings: { temperature: 0.2, fallback_strategy: 'first', allow_early_exit: false }, + flow_pattern: ['main_agent', 'critic', 'main_agent', 'finalizer'], + }; + + // Add agents and router to the workflow + addAgent(mainAgent, { x: 100, y: 100 }); + addAgent(critic, { x: 400, y: 100 }); + addAgent(finalizer, { x: 700, y: 100 }); + addRouter(reflectionRouter, { x: 400, y: 250 }); + + // Create edges for reflection pattern (agents connect through router) + setTimeout(() => { + // Agents connect through the router - router is the coordination hub + + // Main Agent → Router + onConnect({ + source: `main_agent_${timestamp}`, + target: `reflection_router_${timestamp}`, + sourceHandle: null, + targetHandle: null, + }); + + // Router → Critic + onConnect({ + source: `reflection_router_${timestamp}`, + target: `critic_${timestamp}`, + sourceHandle: null, + targetHandle: null, + }); + + // Router → Finalizer + onConnect({ + source: `reflection_router_${timestamp}`, + target: `finalizer_${timestamp}`, + sourceHandle: null, + targetHandle: null, + }); + + // Critic → Router (for reflection back) + onConnect({ + source: `critic_${timestamp}`, + target: `reflection_router_${timestamp}`, + sourceHandle: null, + targetHandle: null, + }); + + // Set initial start and end nodes for reflection workflow + setStartNode(`main_agent_${timestamp}`); + addEndNode(`finalizer_${timestamp}`); + + // The router connections will be converted to proper YAML workflow edges + // with router fields by the YAML export logic + }, 100); + }; + return (
@@ -173,25 +400,55 @@ const Sidebar: React.FC = () => { {/* Quick Templates */}

Agent Templates

-
- {[ - { key: 'analyzer', name: 'Content Analyzer', desc: 'Analyze and extract insights' }, - { key: 'summarizer', name: 'Summarizer', desc: 'Create concise summaries' }, - { key: 'classifier', name: 'Classifier', desc: 'Classify and route content' }, - { key: 'researcher', name: 'Researcher', desc: 'Research with tools' }, - ].map((template) => ( -
createQuickAgent(template.key)} - > -
- - {template.name} + + {/* Core Agents */} +
+

Core Agents

+
+ {[ + { key: 'analyzer', name: 'Content Analyzer', desc: 'Analyze and extract insights' }, + { key: 'summarizer', name: 'Summarizer', desc: 'Create concise summaries' }, + { key: 'classifier', name: 'Classifier', desc: 'Classify and route content' }, + { key: 'researcher', name: 'Researcher', desc: 'Research with tools' }, + ].map((template) => ( +
createQuickAgent(template.key)} + > +
+ + {template.name} +
+

{template.desc}

-

{template.desc}

-
- ))} + ))} +
+
+ + {/* Advanced Pattern Agents */} +
+

Advanced Pattern Agents

+
+ {[ + { key: 'planner', name: 'Planner Agent', desc: 'Break down tasks into execution plans' }, + { key: 'executor', name: 'Executor Agent', desc: 'Execute plans step by step' }, + { key: 'critic', name: 'Critic Agent', desc: 'Provide constructive feedback' }, + { key: 'finalizer', name: 'Finalizer Agent', desc: 'Polish and finalize outputs' }, + ].map((template) => ( +
createQuickAgent(template.key)} + > +
+ + {template.name} +
+

{template.desc}

+
+ ))} +
@@ -200,11 +457,11 @@ const Sidebar: React.FC = () => {

Router Templates

{[ - { key: 'smart', name: 'Smart Router', desc: 'LLM-powered intelligent routing' }, - { key: 'classifier', name: 'Task Classifier', desc: 'Route based on task categories' }, - { key: 'conversation', name: 'Conversation Router', desc: 'Route based on conversation flow' }, - { key: 'reflection', name: 'Reflection Router', desc: 'A→B→A reflection patterns' }, - { key: 'plan_execute', name: 'Plan Execute Router', desc: 'Cursor-style plan-and-execute workflows' }, + { key: 'smart', name: 'Smart Router', desc: 'LLM-powered intelligent routing based on content' }, + { key: 'classifier', name: 'Task Classifier', desc: 'Route based on task categories and keywords' }, + { key: 'conversation', name: 'Conversation Router', desc: 'Route based on conversation context analysis' }, + { key: 'reflection', name: 'Reflection Router', desc: 'A→B→A→C reflection patterns (needs main+critic agents)' }, + { key: 'plan_execute', name: 'Plan Execute Router', desc: 'Plan-and-execute workflows (needs planner+executor agents)' }, ].map((template) => (
{ )}
+ {/* Quick Workflow Templates */} +
+

Quick Workflows

+
+ + + +
+
+ {/* Getting Started */}

Getting Started

    -
  1. 1. Create agents using templates or custom forms
    2. -
  2. 2. Add tools to enhance agent capabilities
    3. +
  3. 1. Try quick workflows above for advanced patterns
    4. +
  4. 2. Or create agents using templates below
  5. 3. Connect agents by dragging from output to input
    6. -
  6. 4. Export as YAML when ready
    7. +
  7. 4. Add routers for intelligent routing
    8. +
  8. 5. Export as YAML when ready
+
+ πŸ’‘ Smart Routing: Reflection & Plan-Execute routers automatically track execution context to prevent infinite loops and manage complex flow patterns intelligently. +
diff --git a/studio/src/store/designerStore.ts b/studio/src/store/designerStore.ts index e684d19f..6435a600 100644 --- a/studio/src/store/designerStore.ts +++ b/studio/src/store/designerStore.ts @@ -1,5 +1,5 @@ import { create } from 'zustand'; -import { Node, Edge, addEdge, Connection, applyNodeChanges, applyEdgeChanges, NodeChange, EdgeChange } from 'reactflow'; +import { addEdge, Connection, applyNodeChanges, applyEdgeChanges, NodeChange, EdgeChange } from 'reactflow'; import { Agent, Tool, Router, DesignerConfig } from '@/types/agent'; import { CustomNode, CustomEdge, AgentNodeData, ToolNodeData, RouterNodeData } from '@/types/reactflow'; @@ -24,6 +24,10 @@ interface DesignerState { workflowDescription: string; workflowVersion: string; + // Start/End node management + startNodeId?: string; + endNodeIds: string[]; + // Actions setConfig: (config: DesignerConfig) => void; @@ -59,6 +63,12 @@ interface DesignerState { // Workflow metadata actions setWorkflowMetadata: (metadata: { name: string; description: string; version: string }) => void; + // Start/End node actions + setStartNode: (nodeId: string) => void; + addEndNode: (nodeId: string) => void; + removeEndNode: (nodeId: string) => void; + toggleEndNode: (nodeId: string) => void; + // Utility actions clearWorkflow: () => void; loadWorkflow: (workflow: any) => void; @@ -131,7 +141,7 @@ const defaultConfig: DesignerConfig = { ], }; -export const useDesignerStore = create((set, get) => ({ +export const useDesignerStore = create((set) => ({ // Initial state config: defaultConfig, nodes: [], @@ -145,6 +155,8 @@ export const useDesignerStore = create((set, get) => ({ workflowName: 'New Workflow', workflowDescription: '', workflowVersion: '1.0.0', + startNodeId: undefined, + endNodeIds: [], // Configuration actions setConfig: (config) => set({ config }), @@ -316,6 +328,34 @@ export const useDesignerStore = create((set, get) => ({ }); }, + // Start/End node actions + setStartNode: (nodeId) => { + set({ startNodeId: nodeId }); + }, + + addEndNode: (nodeId) => { + set((state) => ({ + endNodeIds: [...state.endNodeIds.filter(id => id !== nodeId), nodeId], + })); + }, + + removeEndNode: (nodeId) => { + set((state) => ({ + endNodeIds: state.endNodeIds.filter(id => id !== nodeId), + })); + }, + + toggleEndNode: (nodeId) => { + set((state) => { + const isCurrentlyEnd = state.endNodeIds.includes(nodeId); + return { + endNodeIds: isCurrentlyEnd + ? state.endNodeIds.filter(id => id !== nodeId) + : [...state.endNodeIds, nodeId], + }; + }); + }, + // Utility actions clearWorkflow: () => { set({ @@ -326,6 +366,8 @@ export const useDesignerStore = create((set, get) => ({ workflowName: 'New Workflow', workflowDescription: '', workflowVersion: '1.0.0', + startNodeId: undefined, + endNodeIds: [], }); }, diff --git a/studio/src/utils/workflowValidation.ts b/studio/src/utils/workflowValidation.ts index feeda6b2..5682d91a 100644 --- a/studio/src/utils/workflowValidation.ts +++ b/studio/src/utils/workflowValidation.ts @@ -16,7 +16,7 @@ export interface ValidationResult { suggestions: ValidationIssue[]; } -export function validateWorkflow(nodes: CustomNode[], edges: CustomEdge[]): ValidationResult { +export function validateWorkflow(nodes: CustomNode[], edges: CustomEdge[], startNodeId?: string, endNodeIds?: string[]): ValidationResult { const issues: ValidationIssue[] = []; const warnings: ValidationIssue[] = []; const suggestions: ValidationIssue[] = []; @@ -112,14 +112,46 @@ export function validateWorkflow(nodes: CustomNode[], edges: CustomEdge[]): Vali } // Validate router-specific configurations - if (router.type === 'reflection' && (!router.flow_pattern || router.flow_pattern.length < 2)) { - issues.push({ - id: `reflection_no_pattern_${node.id}`, - type: 'error', - message: 'Reflection router must have a flow pattern with at least 2 agents', - nodeId: node.id, - category: 'configuration', - }); + if (router.type === 'reflection') { + if (!router.flow_pattern || router.flow_pattern.length < 3) { + issues.push({ + id: `reflection_no_pattern_${node.id}`, + type: 'error', + message: 'Reflection router must have a flow pattern with at least 3 agents (main→critic→main→finalizer)', + nodeId: node.id, + category: 'configuration', + }); + } + + // Check if the workflow has the necessary agents for reflection pattern + const hasMainAgent = nodes.some(n => n.type === 'agent' && (n.data as any).agent.name.toLowerCase().includes('main')); + const hasCriticAgent = nodes.some(n => n.type === 'agent' && (n.data as any).agent.name.toLowerCase().includes('critic')); + + if (!hasMainAgent || !hasCriticAgent) { + warnings.push({ + id: `reflection_missing_agents_${node.id}`, + type: 'warning', + message: 'Reflection router works best with "Main Agent" and "Critic Agent" types. Use the quick workflow templates.', + nodeId: node.id, + category: 'best_practice', + }); + } + } + + if (router.type === 'plan_execute') { + // Check if the workflow has the necessary agents for plan-execute pattern + const hasPlannerAgent = nodes.some(n => n.type === 'agent' && (n.data as any).agent.name.toLowerCase().includes('planner')); + const hasExecutorAgent = nodes.some(n => n.type === 'agent' && (n.data as any).agent.name.toLowerCase().includes('executor')); + + if (!hasPlannerAgent || !hasExecutorAgent) { + warnings.push({ + id: `plan_execute_missing_agents_${node.id}`, + type: 'warning', + message: 'Plan-execute router works best with "Planner Agent" and "Executor Agent" types. Use the quick workflow templates.', + nodeId: node.id, + category: 'best_practice', + }); + } } if (router.type === 'task_classifier' && (!router.task_categories || Object.keys(router.task_categories).length === 0)) { @@ -134,6 +166,63 @@ export function validateWorkflow(nodes: CustomNode[], edges: CustomEdge[]): Vali } } + // Validate start/end node configuration + if (startNodeId) { + const startNode = nodes.find(n => n.id === startNodeId); + if (!startNode) { + issues.push({ + id: 'invalid_start_node', + type: 'error', + message: 'Start node no longer exists in the workflow', + category: 'structure', + }); + } else if (startNode.type === 'router') { + issues.push({ + id: 'router_start_node', + type: 'error', + message: 'Start node cannot be a router - must be an agent or tool', + nodeId: startNodeId, + category: 'structure', + }); + } + } else { + warnings.push({ + id: 'no_start_node_set', + type: 'warning', + message: 'No start node defined. Click "Set Start" on an agent to define the workflow entry point.', + category: 'structure', + }); + } + + if (endNodeIds && endNodeIds.length > 0) { + endNodeIds.forEach((endNodeId) => { + const endNode = nodes.find(n => n.id === endNodeId); + if (!endNode) { + issues.push({ + id: `invalid_end_node_${endNodeId}`, + type: 'error', + message: 'End node no longer exists in the workflow', + category: 'structure', + }); + } else if (endNode.type === 'router') { + issues.push({ + id: `router_end_node_${endNodeId}`, + type: 'error', + message: 'End node cannot be a router - must be an agent or tool', + nodeId: endNodeId, + category: 'structure', + }); + } + }); + } else { + suggestions.push({ + id: 'no_end_nodes_set', + type: 'info', + message: 'No end nodes defined. Click "Set End" on agents to define workflow completion points.', + category: 'structure', + }); + } + // Validate workflow structure const { startNodes, endNodes, isolatedNodes } = analyzeWorkflowStructure(nodes, edges); diff --git a/studio/src/utils/yamlExport.ts b/studio/src/utils/yamlExport.ts index 650af4a9..5a04d8cb 100644 --- a/studio/src/utils/yamlExport.ts +++ b/studio/src/utils/yamlExport.ts @@ -8,6 +8,8 @@ export interface ExportData { workflowName: string; workflowDescription: string; workflowVersion: string; + startNodeId?: string; + endNodeIds: string[]; } // Utility function to convert names to snake_case @@ -22,7 +24,7 @@ function toSnakeCase(str: string): string { } export function generateAriumYAML(data: ExportData): string { - const { nodes, edges, workflowName, workflowDescription, workflowVersion } = data; + const { nodes, edges, workflowName, workflowDescription, workflowVersion, startNodeId, endNodeIds } = data; // Create mapping from node IDs to snake_case names const nodeIdToName = new Map(); @@ -59,38 +61,50 @@ export function generateAriumYAML(data: ExportData): string { } }); - // Determine start and end nodes - const nodeConnections = new Map(); + // Determine start and end nodes (use user-defined or auto-detect) + let finalStartNodeId = startNodeId; + let finalEndNodeIds = [...endNodeIds]; - // Initialize connections map - nodes.forEach((node) => { - nodeConnections.set(node.id, { incoming: [], outgoing: [] }); - }); - - // Build connections - edges.forEach((edge) => { - const sourceConnections = nodeConnections.get(edge.source); - const targetConnections = nodeConnections.get(edge.target); + // If no user-defined start/end nodes, auto-detect from connections + if (!finalStartNodeId || finalEndNodeIds.length === 0) { + const nodeConnections = new Map(); - if (sourceConnections) { - sourceConnections.outgoing.push(edge.target); - } - if (targetConnections) { - targetConnections.incoming.push(edge.source); - } - }); + // Initialize connections map + nodes.forEach((node) => { + nodeConnections.set(node.id, { incoming: [], outgoing: [] }); + }); - // Find start nodes (nodes with no incoming connections) - const startNodes = nodes.filter((node) => { - const connections = nodeConnections.get(node.id); - return connections && connections.incoming.length === 0; - }); + // Build connections + edges.forEach((edge) => { + const sourceConnections = nodeConnections.get(edge.source); + const targetConnections = nodeConnections.get(edge.target); + + if (sourceConnections) { + sourceConnections.outgoing.push(edge.target); + } + if (targetConnections) { + targetConnections.incoming.push(edge.source); + } + }); - // Find end nodes (nodes with no outgoing connections) - const endNodes = nodes.filter((node) => { - const connections = nodeConnections.get(node.id); - return connections && connections.outgoing.length === 0; - }); + // Auto-detect start node if not set + if (!finalStartNodeId) { + const startNodes = nodes.filter((node) => { + const connections = nodeConnections.get(node.id); + return connections && connections.incoming.length === 0 && node.type !== 'router'; + }); + finalStartNodeId = startNodes[0]?.id; + } + + // Auto-detect end nodes if none set + if (finalEndNodeIds.length === 0) { + const endNodes = nodes.filter((node) => { + const connections = nodeConnections.get(node.id); + return connections && connections.outgoing.length === 0 && node.type !== 'router'; + }); + finalEndNodeIds = endNodes.map(node => node.id); + } + } // Build workflow edges using snake_case names const workflowEdges: Array<{ @@ -113,50 +127,70 @@ export function generateAriumYAML(data: ExportData): string { const sourceName = nodeIdToName.get(sourceId); if (!sourceName) return; // Skip if source node not found - const targets = sourceEdges.map((edge) => edge.target); - const targetNames = targets - .map(targetId => nodeIdToName.get(targetId)) - .filter(name => name !== undefined) as string[]; + const sourceNode = nodes.find(n => n.id === sourceId); - // Check if any target is a router node - const routerTarget = targets.find(targetId => { - const targetNode = nodes.find(n => n.id === targetId); - return targetNode?.type === 'router'; - }); - - if (routerTarget) { - // If connecting to a router, find what the router connects to - const routerName = nodeIdToName.get(routerTarget); - const routerEdges = edgesBySource.get(routerTarget) || []; - const routerTargetNames = routerEdges - .map(edge => nodeIdToName.get(edge.target)) - .filter(name => name !== undefined) as string[]; + // Handle agent nodes that connect to routers + if (sourceNode?.type === 'agent') { + // Check if this agent connects to a router + const routerTargets = sourceEdges + .map(edge => edge.target) + .map(targetId => nodes.find(n => n.id === targetId)) + .filter(node => node?.type === 'router'); - if (routerTargetNames.length > 0) { - workflowEdges.push({ - from: sourceName, - to: routerTargetNames, // Use router's targets as destinations - router: routerName, // Specify router for decision making - }); - } - } else { - // Direct connection without router - if (targetNames.length > 0) { - workflowEdges.push({ - from: sourceName, - to: targetNames, + if (routerTargets.length > 0) { + // This agent connects through a router + const routerNode = routerTargets[0]; + const routerName = nodeIdToName.get(routerNode!.id); + + // Find what the router connects to + const routerEdges = edgesBySource.get(routerNode!.id) || []; + const routerTargetNames = routerEdges + .map(edge => nodeIdToName.get(edge.target)) + .filter(name => name !== undefined) as string[]; + + // Create workflow edge with router + if (routerTargetNames.length > 0) { + workflowEdges.push({ + from: sourceName, + to: routerTargetNames, + router: routerName, + }); + } + } else { + // Direct agent-to-agent connections + const targets = sourceEdges.map((edge) => edge.target); + const targetNames = targets + .map(targetId => nodeIdToName.get(targetId)) + .filter(name => name !== undefined) as string[]; + + // Skip router nodes as targets + const agentTargetNames = targetNames.filter(name => { + const targetNode = nodes.find(n => nodeIdToName.get(n.id) === name); + return targetNode?.type !== 'router'; }); + + if (agentTargetNames.length > 0) { + const workflowEdge: any = { + from: sourceName, + to: agentTargetNames, + }; + + // Check if edges have router data + const routerNames = sourceEdges + .map(edge => edge.data?.router) + .filter(router => router !== undefined); + + if (routerNames.length > 0) { + workflowEdge.router = routerNames[0]; + } + + workflowEdges.push(workflowEdge); + } } } }); - // Remove router-only edges (edges that start from routers) - const filteredWorkflowEdges = workflowEdges.filter(edge => { - const isRouterSource = nodes.some(node => - nodeIdToName.get(node.id) === edge.from && node.type === 'router' - ); - return !isRouterSource; - }); + // Note: workflowEdges already filters out router-only edges in the main logic above // Convert agents to YAML format const yamlAgents = agents.map((agent) => { @@ -277,9 +311,9 @@ export function generateAriumYAML(data: ExportData): string { tools: tools.length > 0 ? tools.map(name => ({ name })) : undefined, routers: yamlRouters.length > 0 ? yamlRouters : undefined, workflow: { - start: startNodes.length > 0 ? (nodeIdToName.get(startNodes[0].id) || startNodes[0].id) : (agents[0]?.name || ''), - edges: filteredWorkflowEdges.filter(edge => edge.to.length > 0), - end: endNodes.map((node) => nodeIdToName.get(node.id) || node.id), + start: finalStartNodeId ? (nodeIdToName.get(finalStartNodeId) || finalStartNodeId) : (agents[0]?.name || ''), + edges: workflowEdges.filter(edge => edge.to.length > 0), + end: finalEndNodeIds.map((nodeId) => nodeIdToName.get(nodeId) || nodeId).filter(name => name), }, }, };