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/README.md b/README.md index b4995f86..8f1b5ed4 100644 --- a/README.md +++ b/README.md @@ -42,6 +42,51 @@ Flo AI is a Python framework for building structured AI agents with support for Flo AI is a Python framework that makes building production-ready AI agents and teams as easy as writing YAML. Think "Kubernetes for AI Agents" - compose complex AI architectures using pre-built components while maintaining the flexibility to create your own. +## 🎨 Flo AI Studio - Visual Workflow Designer + +**Create AI workflows visually with our powerful React-based studio!** + +

+ Flo AI Studio - Visual Workflow Designer +

+ +Flo AI Studio is a modern, intuitive visual editor that allows you to design complex multi-agent workflows through a drag-and-drop interface. Build sophisticated AI systems without writing code, then export them as production-ready YAML configurations. + +### πŸš€ Studio Features + +- **🎯 Visual Design**: Drag-and-drop interface for creating agent workflows +- **πŸ€– Agent Management**: Configure AI agents with different roles, models, and tools +- **πŸ”€ Smart Routing**: Visual router configuration for intelligent workflow decisions +- **πŸ“€ YAML Export**: Export workflows as Flo AI-compatible YAML configurations +- **πŸ“₯ YAML Import**: Import existing workflows for further editing +- **βœ… Workflow Validation**: Real-time validation and error checking +- **πŸ”§ Tool Integration**: Connect agents to external tools and APIs +- **πŸ“‹ Template System**: Quick start with pre-built agent and router templates + +### πŸƒβ€β™‚οΈ Quick Start with Studio + +1. **Start the Studio**: + ```bash + cd studio + pnpm install + pnpm dev + ``` + +2. **Design Your Workflow**: + - Add agents, routers, and tools to the canvas + - Configure their properties and connections + - Test with the built-in validation + +3. **Export & Run**: + ```bash + # Export YAML from the studio, then run with Flo AI + python -c " + from flo_ai.arium import AriumBuilder + builder = AriumBuilder.from_yaml(yaml_file='your_workflow.yaml') + result = await builder.build_and_run(['Your input here']) + " + ``` + ## ✨ Features - πŸ”Œ **Truly Composable**: Build complex AI systems by combining smaller, reusable components @@ -87,6 +132,8 @@ Flo AI is a Python framework that makes building production-ready AI agents and - [Builder Pattern Benefits](#builder-pattern-benefits) - [πŸ“„ YAML-Based Arium Workflows](#-yaml-based-arium-workflows) - [🧠 LLM-Powered Routers in YAML (NEW!)](#-llm-powered-routers-in-yaml-new) + - [πŸ”„ ReflectionRouter: Structured Reflection Workflows (NEW!)](#-reflectionrouter-structured-reflection-workflows-new) + - [πŸ”„ PlanExecuteRouter: Cursor-Style Plan-and-Execute Workflows (NEW!)](#-planexecuterouter-cursor-style-plan-and-execute-workflows-new) - [πŸ“– Documentation](#-documentation) - [🌟 Why Flo AI?](#-why-flo-ai) - [🎯 Use Cases](#-use-cases) @@ -1429,6 +1476,8 @@ arium: 1. **Smart Router** (`type: smart`): General-purpose routing based on content analysis 2. **Task Classifier** (`type: task_classifier`): Routes based on keywords and examples 3. **Conversation Analysis** (`type: conversation_analysis`): Context-aware routing +4. **Reflection Router** (`type: reflection`): Structured Aβ†’Bβ†’Aβ†’C patterns for reflection workflows +5. **PlanExecute Router** (`type: plan_execute`): Cursor-style plan-and-execute workflows with step tracking **✨ Key Benefits:** - 🚫 **No Code Required**: Define routing logic purely in YAML @@ -1450,6 +1499,560 @@ async def run_intelligent_workflow(): return result ``` +##### πŸ”„ ReflectionRouter: Structured Reflection Workflows (NEW!) + +The **ReflectionRouter** is designed specifically for reflection-based workflows that follow Aβ†’Bβ†’Aβ†’C patterns, commonly used for mainβ†’criticβ†’mainβ†’final agent sequences. This pattern is perfect for iterative improvement workflows where a critic agent provides feedback before final processing. + +**πŸ“‹ Key Features:** +- 🎯 **Pattern Tracking**: Automatically tracks progress through defined reflection sequences +- πŸ”„ **Self-Reference Support**: Allows routing back to the same agent (Aβ†’Bβ†’A patterns) +- πŸ“Š **Visual Progress**: Shows current position with β—‹ pending, βœ“ completed indicators +- πŸ›‘οΈ **Loop Prevention**: Built-in safety mechanisms to prevent infinite loops +- πŸŽ›οΈ **Flexible Patterns**: Supports both 2-agent (Aβ†’Bβ†’A) and 3-agent (Aβ†’Bβ†’Aβ†’C) flows + +**🎯 Supported Patterns:** + +1. **A β†’ B β†’ A** (2 agents): Main β†’ Critic β†’ Main β†’ End +2. **A β†’ B β†’ A β†’ C** (3 agents): Main β†’ Critic β†’ Main β†’ Final + +```yaml +# Simple A β†’ B β†’ A reflection pattern +metadata: + name: "content-reflection-workflow" + version: "1.0.0" + description: "Content creation with critic feedback loop" + +arium: + agents: + - name: "writer" + role: "Content Writer" + job: "Create and improve content based on feedback from critics." + model: + provider: "openai" + name: "gpt-4o-mini" + settings: + temperature: 0.7 + + - name: "critic" + role: "Content Critic" + job: "Review content and provide constructive feedback for improvement." + model: + provider: "openai" + name: "gpt-4o-mini" + settings: + temperature: 0.3 + + # ✨ ReflectionRouter definition + routers: + - name: "reflection_router" + type: "reflection" # Specialized for reflection patterns + flow_pattern: [writer, critic, writer] # A β†’ B β†’ A pattern + model: + provider: "openai" + name: "gpt-4o-mini" + settings: + temperature: 0.2 + allow_early_exit: false # Strict adherence to pattern + + workflow: + start: "writer" + edges: + - from: "writer" + to: [critic, writer] # Can go to critic or self-reference + router: "reflection_router" + - from: "critic" + to: [writer] # Always returns to writer + router: "reflection_router" + end: [writer] # Writer produces final output +``` + +```yaml +# Advanced A β†’ B β†’ A β†’ C reflection pattern +metadata: + name: "advanced-reflection-workflow" + version: "1.0.0" + description: "Full reflection cycle with dedicated final agent" + +arium: + agents: + - name: "researcher" + role: "Research Agent" + job: "Conduct research and gather information on topics." + model: + provider: "openai" + name: "gpt-4o-mini" + + - name: "reviewer" + role: "Research Reviewer" + job: "Review research quality and suggest improvements." + model: + provider: "anthropic" + name: "claude-3-5-sonnet-20240620" + + - name: "synthesizer" + role: "Information Synthesizer" + job: "Create final synthesis and conclusions from research." + model: + provider: "openai" + name: "gpt-4o" + + routers: + - name: "research_reflection_router" + type: "reflection" + flow_pattern: [researcher, reviewer, researcher, synthesizer] # A β†’ B β†’ A β†’ C + settings: + allow_early_exit: true # Allow smart early completion + + workflow: + start: "researcher" + edges: + - from: "researcher" + to: [reviewer, researcher, synthesizer] # All possible destinations + router: "research_reflection_router" + - from: "reviewer" + to: [researcher, reviewer, synthesizer] + router: "research_reflection_router" + - from: "synthesizer" + to: [end] + end: [synthesizer] +``` + +**πŸ”§ ReflectionRouter Configuration Options:** + +```yaml +routers: + - name: "my_reflection_router" + type: "reflection" + flow_pattern: [main_agent, critic, main_agent, final_agent] # Define your pattern + model: # Optional: LLM for routing decisions + provider: "openai" + name: "gpt-4o-mini" + settings: # Optional settings + temperature: 0.2 # Router temperature (lower = more deterministic) + allow_early_exit: false # Allow early completion if LLM determines pattern is done + fallback_strategy: "first" # first, last, random - fallback when LLM fails +``` + +**πŸ—οΈ Programmatic Usage:** + +```python +import asyncio +from flo_ai.arium import AriumBuilder +from flo_ai.models.agent import Agent +from flo_ai.llm import OpenAI +from flo_ai.arium.llm_router import create_main_critic_reflection_router + +async def reflection_workflow_example(): + llm = OpenAI(model='gpt-4o-mini', api_key='your-api-key') + + # Create agents + main_agent = Agent( + name='main_agent', + system_prompt='Create solutions and improve them based on feedback.', + llm=llm + ) + + critic = Agent( + name='critic', + system_prompt='Provide constructive feedback for improvement.', + llm=llm + ) + + final_agent = Agent( + name='final_agent', + system_prompt='Polish and finalize the work.', + llm=llm + ) + + # Create reflection router - A β†’ B β†’ A β†’ C pattern + reflection_router = create_main_critic_reflection_router( + main_agent='main_agent', + critic_agent='critic', + final_agent='final_agent', + allow_early_exit=False, # Strict pattern adherence + llm=llm + ) + + # Build workflow + result = await ( + AriumBuilder() + .add_agents([main_agent, critic, final_agent]) + .start_with(main_agent) + .add_edge(main_agent, [critic, final_agent], reflection_router) + .add_edge(critic, [main_agent, final_agent], reflection_router) + .end_with(final_agent) + .build_and_run(["Create a comprehensive project proposal"]) + ) + + return result + +# Alternative: Direct factory usage +from flo_ai.arium.llm_router import create_llm_router + +reflection_router = create_llm_router( + 'reflection', + flow_pattern=['writer', 'editor', 'writer'], # A β†’ B β†’ A + allow_early_exit=False, + llm=llm +) +``` + +**πŸ’‘ ReflectionRouter Intelligence:** + +The ReflectionRouter automatically: +- **Tracks Progress**: Knows which step in the pattern should execute next +- **Prevents Loops**: Uses execution context to avoid infinite cycles +- **Provides Guidance**: Shows LLM the suggested next step and current progress +- **Handles Self-Reference**: Properly validates flows that return to the same agent +- **Visual Feedback**: Displays pattern progress: `β—‹ writer β†’ βœ“ critic β†’ β—‹ writer` + +**🎯 Perfect Use Cases:** +- πŸ“ **Content Creation**: Writer β†’ Editor β†’ Writer β†’ Publisher +- πŸ”¬ **Research Workflows**: Researcher β†’ Reviewer β†’ Researcher β†’ Synthesizer +- πŸ’Ό **Business Analysis**: Analyst β†’ Critic β†’ Analyst β†’ Decision Maker +- 🎨 **Creative Processes**: Creator β†’ Critic β†’ Creator β†’ Finalizer +- πŸ§ͺ **Iterative Refinement**: Any process requiring feedback and improvement cycles + +**⚑ Quick Start Example:** + +```python +# Minimal A β†’ B β†’ A pattern +yaml_config = """ +arium: + agents: + - name: main_agent + job: "Main work agent" + model: {provider: openai, name: gpt-4o-mini} + - name: critic + job: "Feedback agent" + model: {provider: openai, name: gpt-4o-mini} + + routers: + - name: reflection_router + type: reflection + flow_pattern: [main_agent, critic, main_agent] + + workflow: + start: main_agent + edges: + - from: main_agent + to: [critic, main_agent] + router: reflection_router + - from: critic + to: [main_agent] + router: reflection_router + end: [main_agent] +""" + +result = await AriumBuilder().from_yaml(yaml_str=yaml_config).build_and_run(["Your task"]) +``` + +The ReflectionRouter makes implementing sophisticated feedback loops and iterative improvement workflows incredibly simple, whether you need a 2-agent or 3-agent pattern! πŸš€ + +##### πŸ”„ PlanExecuteRouter: Cursor-Style Plan-and-Execute Workflows (NEW!) + +The **PlanExecuteRouter** implements sophisticated plan-and-execute patterns similar to how Cursor works. It automatically breaks down complex tasks into detailed execution plans and coordinates step-by-step execution with intelligent progress tracking. + +**πŸ“‹ Key Features:** +- 🎯 **Automatic Task Breakdown**: Creates detailed execution plans from high-level tasks +- πŸ“Š **Step Tracking**: Real-time progress monitoring with visual indicators (β—‹ ⏳ βœ… ❌) +- πŸ”„ **Phase Coordination**: Intelligent routing between planning, execution, and review phases +- πŸ›‘οΈ **Dependency Management**: Handles step dependencies and execution order automatically +- πŸ’Ύ **Plan Persistence**: Uses PlanAwareMemory for stateful plan storage and updates +- πŸ”§ **Error Recovery**: Built-in retry logic for failed steps + +**🎯 Perfect for Cursor-Style Workflows:** +- πŸ’» **Software Development**: Requirements β†’ Design β†’ Implementation β†’ Testing β†’ Review +- πŸ“ **Content Creation**: Planning β†’ Writing β†’ Editing β†’ Review β†’ Publishing +- πŸ”¬ **Research Projects**: Plan β†’ Investigate β†’ Analyze β†’ Synthesize β†’ Report +- πŸ“Š **Business Processes**: Any multi-step workflow with dependencies + +**πŸ“„ YAML Configuration:** + +```yaml +# Complete Plan-Execute Workflow +metadata: + name: "development-plan-execute" + version: "1.0.0" + description: "Cursor-style development workflow" + +arium: + agents: + - name: planner + role: Project Planner + job: > + Break down complex development tasks into detailed, sequential execution plans. + Create clear steps with dependencies and agent assignments. + model: + provider: openai + name: gpt-4o-mini + settings: + temperature: 0.3 + + - name: developer + role: Software Developer + job: > + Implement features step by step according to execution plans. + Provide detailed implementation and update step status. + model: + provider: openai + name: gpt-4o-mini + settings: + temperature: 0.5 + + - name: tester + role: QA Engineer + job: > + Test implementations thoroughly and validate functionality. + Create comprehensive test scenarios and report results. + model: + provider: openai + name: gpt-4o-mini + settings: + temperature: 0.2 + + - name: reviewer + role: Senior Reviewer + job: > + Provide final quality assessment and approval. + Review completed work for best practices and requirements. + model: + provider: openai + name: gpt-4o-mini + + # PlanExecuteRouter configuration + routers: + - name: dev_plan_router + type: plan_execute # Router type for plan-execute workflows + agents: # Available agents and their capabilities + planner: "Creates detailed execution plans by breaking down tasks" + developer: "Implements features and code according to plan specifications" + tester: "Tests implementations and validates functionality" + reviewer: "Reviews and approves completed work" + model: # Optional: LLM for routing decisions + provider: openai + name: gpt-4o-mini + settings: # Optional configuration + temperature: 0.2 # Router decision temperature + planner_agent: planner # Agent responsible for creating plans + executor_agent: developer # Default agent for executing steps + reviewer_agent: reviewer # Optional agent for final review + max_retries: 3 # Maximum retries for failed steps + + workflow: + start: planner + edges: + # All agents can route to all others based on plan state + - from: planner + to: [developer, tester, reviewer, planner] + router: dev_plan_router + - from: developer + to: [developer, tester, reviewer, planner] + router: dev_plan_router + - from: tester + to: [developer, tester, reviewer, planner] + router: dev_plan_router + - from: reviewer + to: [end] + end: [reviewer] +``` + +**πŸ—οΈ Programmatic Usage:** + +```python +import asyncio +from flo_ai.arium import AriumBuilder +from flo_ai.arium.memory import PlanAwareMemory +from flo_ai.models.agent import Agent +from flo_ai.llm import OpenAI +from flo_ai.arium.llm_router import create_plan_execute_router + +async def cursor_style_workflow(): + llm = OpenAI(model='gpt-4o-mini', api_key='your-api-key') + + # Create specialized agents + planner = Agent( + name='planner', + system_prompt='Create detailed execution plans by breaking down tasks into sequential steps.', + llm=llm + ) + + developer = Agent( + name='developer', + system_prompt='Implement features step by step according to execution plans.', + llm=llm + ) + + tester = Agent( + name='tester', + system_prompt='Test implementations and validate functionality thoroughly.', + llm=llm + ) + + reviewer = Agent( + name='reviewer', + system_prompt='Review completed work and provide final approval.', + llm=llm + ) + + # Create plan-execute router + plan_router = create_plan_execute_router( + planner_agent='planner', + executor_agent='developer', + reviewer_agent='reviewer', + additional_agents={'tester': 'Tests implementations and validates quality'}, + llm=llm + ) + + # Use PlanAwareMemory for plan state persistence + memory = PlanAwareMemory() + + # Build and run workflow + result = await ( + 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, [developer, tester, reviewer, planner], plan_router) + .add_edge(tester, [developer, tester, reviewer, planner], plan_router) + .add_edge(reviewer, [developer, tester, reviewer, planner], plan_router) + .end_with(reviewer) + .build_and_run(["Create a REST API for user authentication with JWT tokens"]) + ) + + return result + +# Alternative: Factory function +from flo_ai.arium.llm_router import create_plan_execute_router + +plan_router = create_plan_execute_router( + planner_agent='planner', + executor_agent='developer', + reviewer_agent='reviewer', + llm=llm +) +``` + +**πŸ’‘ How PlanExecuteRouter Works:** + +The router intelligently coordinates workflow phases: + +1. **Planning Phase**: + - Detects when no execution plan exists + - Routes to planner agent to create detailed plan + - Plan stored as ExecutionPlan object in PlanAwareMemory + +2. **Execution Phase**: + - Analyzes plan state and step dependencies + - Routes to appropriate agents for next ready steps + - Updates step status (pending β†’ in-progress β†’ completed) + - Handles parallel execution of independent steps + +3. **Review Phase**: + - Detects when all steps are completed + - Routes to reviewer agent for final validation + - Manages error recovery for failed steps + +**πŸ“Š Plan Progress Visualization:** + +``` +πŸ“‹ EXECUTION PLAN: User Authentication API +πŸ“Š CURRENT PROGRESS: +βœ… design_schema: Design user database schema β†’ developer +βœ… implement_registration: Create registration endpoint β†’ developer +⏳ implement_login: Add login with JWT β†’ developer (depends: design_schema, implement_registration) +β—‹ add_middleware: Authentication middleware β†’ developer (depends: implement_login) +β—‹ write_tests: Comprehensive testing β†’ tester (depends: add_middleware) +β—‹ final_review: Security and code review β†’ reviewer (depends: write_tests) + +🎯 NEXT ACTION: Execute step 'implement_login' +🎯 SUGGESTED AGENT: developer +``` + +**πŸ”§ Advanced Configuration Options:** + +```yaml +routers: + - name: advanced_plan_router + type: plan_execute + agents: + planner: "Creates execution plans" + frontend_dev: "Frontend implementation" + backend_dev: "Backend implementation" + devops: "Deployment and infrastructure" + qa_tester: "Quality assurance testing" + security_reviewer: "Security review" + product_owner: "Product validation" + model: + provider: openai + name: gpt-4o + settings: + temperature: 0.1 # Lower for more deterministic routing + planner_agent: planner # Plan creation agent + executor_agent: backend_dev # Default execution agent + reviewer_agent: product_owner # Final review agent + max_retries: 5 # Retry attempts for failed steps + allow_parallel_execution: true # Enable parallel step execution + plan_validation: strict # Validate plan completeness +``` + +**⚑ Quick Start Example:** + +```python +# Minimal plan-execute workflow +yaml_config = """ +arium: + agents: + - name: planner + job: "Create execution plans" + model: {provider: openai, name: gpt-4o-mini} + - name: executor + job: "Execute plan steps" + model: {provider: openai, name: gpt-4o-mini} + - name: reviewer + job: "Review final results" + model: {provider: openai, name: gpt-4o-mini} + + routers: + - name: simple_plan_router + type: plan_execute + agents: + planner: "Creates plans" + executor: "Executes steps" + reviewer: "Reviews results" + settings: + planner_agent: planner + executor_agent: executor + reviewer_agent: reviewer + + workflow: + start: planner + edges: + - from: planner + to: [executor, reviewer, planner] + router: simple_plan_router + - from: executor + to: [executor, reviewer, planner] + router: simple_plan_router + - from: reviewer + to: [end] + end: [reviewer] +""" + +result = await AriumBuilder().from_yaml(yaml_str=yaml_config).build_and_run(["Your complex task"]) +``` + +**🎯 Use Cases and Examples:** + +- πŸ“± **App Development**: "Build a todo app with React and Node.js" +- πŸ›’ **E-commerce**: "Create a shopping cart system with payment processing" +- πŸ“Š **Data Pipeline**: "Build ETL pipeline for customer analytics" +- πŸ” **Security**: "Implement OAuth2 authentication system" +- πŸ“ˆ **Analytics**: "Create real-time dashboard with user metrics" + +The PlanExecuteRouter brings Cursor-style intelligent task automation to Flo AI, making it incredibly easy to build sophisticated multi-step workflows that adapt and execute complex tasks automatically! πŸš€ + #### YAML Workflow with Variables ```yaml 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/concept_demo.py b/flo_ai/examples/concept_demo.py new file mode 100644 index 00000000..ab573783 --- /dev/null +++ b/flo_ai/examples/concept_demo.py @@ -0,0 +1,284 @@ +""" +PlanExecuteRouter Concept Demo + +This demonstrates the plan-execute concept without requiring API calls. +Shows the architecture and explains how to fix the planner loop issue. +""" + +import asyncio +from flo_ai.arium.memory import PlanAwareMemory, ExecutionPlan, PlanStep, StepStatus +import uuid + + +def demonstrate_plan_aware_memory(): + """Show how PlanAwareMemory works with ExecutionPlan objects""" + print('πŸ“‹ PlanAwareMemory Concept Demo') + print('=' * 35) + + # Create memory + memory = PlanAwareMemory() + + # Create a sample execution plan + plan = ExecutionPlan( + id=str(uuid.uuid4()), + title='Build User Authentication API', + description='Create complete user auth system with registration and login', + steps=[ + PlanStep( + id='design_schema', + description='Design user database schema', + agent='developer', + status=StepStatus.PENDING, + ), + PlanStep( + id='implement_registration', + description='Implement user registration endpoint', + agent='developer', + dependencies=['design_schema'], + status=StepStatus.PENDING, + ), + PlanStep( + id='implement_login', + description='Implement user login with JWT tokens', + agent='developer', + dependencies=['design_schema', 'implement_registration'], + status=StepStatus.PENDING, + ), + PlanStep( + id='test_endpoints', + description='Test all authentication endpoints', + agent='tester', + dependencies=['implement_login'], + status=StepStatus.PENDING, + ), + PlanStep( + id='security_review', + description='Review security implementation', + agent='reviewer', + dependencies=['test_endpoints'], + status=StepStatus.PENDING, + ), + ], + ) + + # Store plan in memory + memory.add_plan(plan) + print(f'βœ… Plan stored: {plan.title}') + print(f'πŸ“Š Total steps: {len(plan.steps)}') + + # Show initial state + def show_plan_status(): + current = memory.get_current_plan() + print(f'\nπŸ“‹ Plan Status: {current.title}') + for step in current.steps: + status_icon = { + StepStatus.PENDING: 'β—‹', + StepStatus.IN_PROGRESS: '⏳', + StepStatus.COMPLETED: 'βœ…', + StepStatus.FAILED: '❌', + }.get(step.status, 'β—‹') + deps = ( + f" (depends: {', '.join(step.dependencies)})" + if step.dependencies + else '' + ) + print(f' {status_icon} {step.id}: {step.description} β†’ {step.agent}{deps}') + + show_plan_status() + + # Simulate step execution + print('\nπŸ”„ Simulating step-by-step execution...') + + current_plan = memory.get_current_plan() + + # Execute step 1 + next_steps = current_plan.get_next_steps() + print(f'\n🎯 Next steps ready: {len(next_steps)}') + step1 = next_steps[0] + print(f'⏳ Executing: {step1.description}') + step1.status = StepStatus.COMPLETED + step1.result = 'User table created with id, email, password_hash fields' + memory.update_plan(current_plan) + + # Execute step 2 + next_steps = current_plan.get_next_steps() + print(f'\n🎯 Next steps ready: {len(next_steps)}') + step2 = next_steps[0] + print(f'⏳ Executing: {step2.description}') + step2.status = StepStatus.COMPLETED + step2.result = 'POST /register endpoint with validation implemented' + memory.update_plan(current_plan) + + show_plan_status() + + print(f'\nπŸ“ˆ Plan completion status: {current_plan.is_completed()}') + print(f'πŸ“Š Remaining steps: {len(current_plan.get_next_steps())}') + + +def explain_planner_loop_issue(): + """Explain why the planner got stuck in a loop and how to fix it""" + print('\n\nπŸ”„ Understanding the Planner Loop Issue') + print('=' * 45) + + explanation = """ +❌ THE PROBLEM: +In the original demo, the planner kept getting called in an infinite loop because: + +1. Router asks: "Is there a plan in memory?" +2. Memory says: "No ExecutionPlan objects found" +3. Router decides: "Route to planner to create plan" +4. Planner generates plan TEXT but doesn't store ExecutionPlan OBJECTS +5. Router asks again: "Is there a plan in memory?" +6. Memory still says: "No ExecutionPlan objects found" +7. INFINITE LOOP! πŸ”„ + +βœ… THE SOLUTION: +We need to bridge the gap between plan TEXT and plan OBJECTS: + +APPROACH 1: Specialized PlannerAgent +β€’ Custom agent that parses plan text and stores ExecutionPlan objects +β€’ Router can detect when ExecutionPlan exists in memory +β€’ Automatically switches from planning to execution phase + +APPROACH 2: Content-Based Routing +β€’ Router analyzes message content instead of relying on ExecutionPlan objects +β€’ If message contains "PLAN:", switch to execution mode +β€’ Simpler but less sophisticated + +APPROACH 3: State Management +β€’ Explicitly track workflow state (planning/executing/reviewing) +β€’ Router uses state instead of trying to detect plan completion +β€’ Most reliable but requires more setup +""" + + print(explanation) + + +def show_router_intelligence(): + """Show how the PlanExecuteRouter makes intelligent decisions""" + print('\n\n🧠 PlanExecuteRouter Intelligence') + print('=' * 35) + + intelligence_demo = """ +The PlanExecuteRouter is designed to coordinate complex workflows by: + +🎯 PHASE DETECTION: +β€’ Planning Phase: No plan exists β†’ route to planner +β€’ Execution Phase: Plan exists with pending steps β†’ route to executor +β€’ Review Phase: All steps complete β†’ route to reviewer +β€’ Error Recovery: Failed steps exist β†’ route for retry + +πŸ“Š STEP MANAGEMENT: +β€’ Analyzes step dependencies automatically +β€’ Only executes steps when dependencies are met +β€’ Tracks progress with visual indicators (β—‹ ⏳ βœ… ❌) +β€’ Handles parallel execution of independent steps + +πŸ”„ INTELLIGENT ROUTING: +β€’ Context-aware prompts with plan progress +β€’ Suggests next agent based on plan state +β€’ Prevents infinite loops with completion detection +β€’ Adapts to different workflow patterns + +πŸ’‘ EXAMPLE ROUTING DECISIONS: + +Scenario 1: No plan exists +β†’ Router: "Route to planner to create execution plan" + +Scenario 2: Plan exists, step_1 ready +β†’ Router: "Route to developer to execute step_1" + +Scenario 3: Development complete, testing needed +β†’ Router: "Route to tester to validate implementation" + +Scenario 4: All steps complete +β†’ Router: "Route to reviewer for final approval" + +Scenario 5: Step failed +β†’ Router: "Route to developer to retry failed step" +""" + + print(intelligence_demo) + + +def show_working_implementation(): + """Show the key components of a working implementation""" + print('\n\nπŸ—οΈ Working Implementation Components') + print('=' * 40) + + implementation = """ +For a working PlanExecuteRouter implementation, you need: + +1. πŸ“‹ PLAN STORAGE: + ```python + memory = PlanAwareMemory() # Stores ExecutionPlan objects + plan = ExecutionPlan(title="...", steps=[...]) + memory.add_plan(plan) + ``` + +2. πŸ€– SPECIALIZED AGENTS: + ```python + class PlannerAgent(Agent): + def run(self, input_data): + plan_text = await super().run(input_data) + execution_plan = self.parse_plan(plan_text) + self.memory.add_plan(execution_plan) # KEY! + return plan_text + ``` + +3. 🎯 SMART ROUTING: + ```python + router = create_plan_execute_router( + planner_agent='planner', + executor_agent='developer', + reviewer_agent='reviewer' + ) + ``` + +4. πŸ”„ WORKFLOW COORDINATION: + ```python + arium = AriumBuilder() + .with_memory(memory) # PlanAwareMemory + .add_agents([planner, developer, tester, reviewer]) + .add_edge(planner, [...], router) + .build() + ``` + +🎯 CRITICAL SUCCESS FACTORS: +β€’ PlannerAgent MUST store ExecutionPlan objects +β€’ Use PlanAwareMemory for plan state persistence +β€’ Router needs to detect plan existence reliably +β€’ Agents should update step status during execution +""" + + print(implementation) + + +async def main(): + """Main concept demo""" + print('🎯 PlanExecuteRouter Concept Demo') + print('Understanding the architecture and fixing the planner loop\n') + + # Demonstrate core concepts + demonstrate_plan_aware_memory() + explain_planner_loop_issue() + show_router_intelligence() + show_working_implementation() + + print('\n\nπŸŽ‰ Concept Demo Complete!') + print('=' * 30) + print('Key Takeaways:') + print('βœ… PlanExecuteRouter enables Cursor-style plan-and-execute workflows') + print('βœ… Planner loop occurs when plan TEXT β‰  plan OBJECTS in memory') + print('βœ… Solution: Bridge the gap with specialized agents or content parsing') + print('βœ… Working implementation requires PlanAwareMemory + ExecutionPlan objects') + + print('\nπŸš€ Next Steps:') + print('β€’ Try fixed_plan_execute_demo.py for working implementation') + print('β€’ Use PlannerAgent that stores ExecutionPlan objects') + print('β€’ Implement step status tracking for real progress monitoring') + print('β€’ Customize agents for your specific workflow needs') + + +if __name__ == '__main__': + asyncio.run(main()) 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/example.py b/flo_ai/examples/example.py new file mode 100644 index 00000000..60976ff0 --- /dev/null +++ b/flo_ai/examples/example.py @@ -0,0 +1,143 @@ +""" +Example of running a Flo AI workflow with snake_case agent names +Generated from Flo AI Studio +""" + +import asyncio +from flo_ai.arium import AriumBuilder +from flo_ai.tool.base_tool import Tool + +# Set your OpenAI API key +# os.environ["OPENAI_API_KEY"] = "your-api-key-here" + + +# Create a simple web search tool +async def web_search_function(query: str) -> str: + """Simple web search simulation - replace with actual search API""" + return f"Web search results for '{query}': Found relevant articles about the topic including latest developments, applications, and ethical considerations. [This is a mock search - integrate with real search API like Google, Bing, or DuckDuckGo]" + + +# Create the tool instance +web_search_tool = Tool( + name='web_search', + description='Search the web for information on a given topic', + function=web_search_function, + parameters={ + 'query': { + 'type': 'string', + 'description': 'The search query to look up information about', + } + }, +) + +# YAML workflow definition (exported from Flo AI Studio) +workflow_yaml = """ +metadata: + name: New 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 key insights, themes, and important information. + model: + provider: openai + name: gpt-4o-mini + settings: + temperature: 0.7 + max_retries: 3 + reasoning_pattern: DIRECT + - name: researcher + role: Research Specialist + job: Research topics and gather comprehensive information using available tools. + model: + provider: openai + name: gpt-4o + tools: + - web_search + - name: summarizer + role: Summary Generator + job: Create concise, actionable summaries from analysis and content. + model: + provider: openai + name: gpt-4o-mini + routers: + - name: smart_router + type: smart + routing_options: + researcher: If there is not enough information & deep research needs to be done + summarizer: If we have enough information and its time to summarize + model: + provider: openai + name: gpt-4o-mini + settings: + temperature: 0.3 + fallback_strategy: first + workflow: + start: content_analyzer + edges: + - from: content_analyzer + to: [researcher, summarizer] + router: smart_router + - from: researcher + to: [summarizer] + end: + - summarizer +""" + + +async def main(): + """Run the workflow""" + print('πŸš€ Starting Flo AI Workflow...') + print('πŸ“‹ Workflow: Content Analysis with Smart Routing') + print('-' * 50) + + try: + # Create tools dictionary (required format for AriumBuilder) + tools = {'web_search': web_search_tool} + + # Create Arium builder from YAML + builder = AriumBuilder.from_yaml(yaml_str=workflow_yaml, tools=tools) + + # Example input for the workflow + user_input = [ + """I need to understand the current trends in artificial intelligence and machine learning. + Specifically, I'm interested in: + 1. Latest developments in large language models + 2. Applications in healthcare and finance + 3. Ethical considerations and regulations + + Please provide a comprehensive analysis and summary.""" + ] + + print('πŸ“ Input:') + print(user_input[0]) + print('\n' + '=' * 50) + print('πŸ”„ Processing workflow...') + print('=' * 50 + '\n') + + # Build and run the workflow + result = await builder.build_and_run(user_input) + + print('βœ… Workflow Result:') + print('-' * 30) + if isinstance(result, list): + for i, message in enumerate(result): + print(f'{i+1}. {message}') + else: + print(result) + + except Exception as e: + print(f'❌ Error running workflow: {str(e)}') + print('\nπŸ’‘ Make sure you have:') + print('1. Set your OPENAI_API_KEY environment variable') + print('2. Installed flo-ai: pip install flo-ai') + print('3. All required dependencies') + + +if __name__ == '__main__': + asyncio.run(main()) diff --git a/flo_ai/examples/flow_router_example.py b/flo_ai/examples/flow_router_example.py new file mode 100644 index 00000000..2da1e6c6 --- /dev/null +++ b/flo_ai/examples/flow_router_example.py @@ -0,0 +1,454 @@ +""" +Example demonstrating ReflectionRouter for A -> B -> A -> C patterns. + +This example shows how to implement a main -> critic -> main -> final reflection workflow +using the new ReflectionRouter with YAML configuration. +""" + +import asyncio +from flo_ai.arium.builder import AriumBuilder +from flo_ai.llm import OpenAI + +# Example YAML configuration for A -> B -> A -> C flow +MAIN_CRITIC_FLOW_YAML = """ +metadata: + name: main-critic-final-workflow + version: 1.0.0 + description: "A workflow demonstrating A -> B -> A -> C pattern with intelligent flow routing" + +arium: + agents: + - 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. + Be receptive to criticism and use it to refine your output. + model: + provider: openai + name: gpt-4o-mini + settings: + temperature: 0.7 + + - name: critic + role: Critic Agent + job: > + You are a critic agent. Your job is to review the main agent's work and provide + constructive feedback. Analyze the output for: + - Accuracy and correctness + - Completeness and thoroughness + - Clarity and coherence + - Areas for improvement + Provide specific, actionable feedback that the main agent can use to improve. + model: + provider: openai + name: gpt-4o-mini + settings: + temperature: 0.3 + + - name: final_agent + role: Final Agent + job: > + You are the final agent responsible for polishing and finalizing the work. + Take the refined output from the main agent (after critic feedback) and: + - Format it professionally + - Add any final touches or improvements + - Ensure it meets high quality standards + - Provide a polished final result + model: + provider: openai + name: gpt-4o-mini + settings: + temperature: 0.5 + + # Reflection router configuration for A -> B -> A -> C pattern + routers: + - name: main_critic_reflection_router + type: reflection + flow_pattern: [main_agent, critic, main_agent, final_agent] + model: + provider: openai + name: gpt-4o-mini + settings: + temperature: 0.2 + allow_early_exit: false + fallback_strategy: first + + workflow: + start: main_agent + edges: + # Single edge from main_agent using reflection router + # The router will intelligently route to: critic -> main_agent -> final_agent + - from: main_agent + to: [critic, final_agent] # All possible destinations + router: main_critic_reflection_router + - from: critic + to: [main_agent, final_agent] + router: main_critic_reflection_router + - from: final_agent + to: [end] + end: [final_agent] +""" + +# Alternative stricter flow pattern +STRICT_FLOW_YAML = """ +metadata: + name: strict-main-critic-flow + version: 1.0.0 + description: "Strict A -> B -> A -> C flow with no deviations allowed" + +arium: + agents: + - name: writer + role: Content Writer + job: > + You are a content writer. Create initial content based on the user's request. + Focus on getting the core ideas down first, don't worry about perfection. + model: + provider: openai + name: gpt-4o-mini + settings: + temperature: 0.8 + + - name: reviewer + role: Content Reviewer + job: > + You are a content reviewer. Review the writer's work and provide detailed feedback: + - What works well + - What needs improvement + - Specific suggestions for enhancement + - Areas that need clarification + model: + provider: openai + name: gpt-4o-mini + settings: + temperature: 0.2 + + - name: editor + role: Content Editor + job: > + You are the final editor. Take the revised content from the writer and: + - Polish the language and style + - Ensure consistency and flow + - Make final corrections + - Prepare the content for publication + model: + provider: openai + name: gpt-4o-mini + settings: + temperature: 0.3 + + routers: + - name: strict_reflection_router + type: reflection + flow_pattern: [writer, reviewer, writer, editor] + settings: + allow_early_exit: false # Strict adherence to pattern + fallback_strategy: first + + workflow: + start: writer + edges: + - from: writer + to: [reviewer, editor] + router: strict_reflection_router + - from: reviewer + to: [writer, editor] + router: strict_reflection_router + - from: editor + to: [end] + end: [editor] +""" + +# Flexible flow that allows early exit +FLEXIBLE_FLOW_YAML = """ +metadata: + name: flexible-flow-with-early-exit + version: 1.0.0 + description: "Flexible A -> B -> A -> C flow that allows early completion" + +arium: + agents: + - name: analyst + role: Data Analyst + job: > + You are a data analyst. Analyze the given data or question and provide insights. + Create clear, actionable analysis based on the information provided. + model: + provider: openai + name: gpt-4o-mini + settings: + temperature: 0.5 + + - name: validator + role: Analysis Validator + job: > + You are an analysis validator. Review the analyst's work for: + - Logical consistency + - Accuracy of conclusions + - Completeness of analysis + - Potential issues or gaps + If the analysis is solid, you can recommend proceeding directly to completion. + model: + provider: openai + name: gpt-4o-mini + settings: + temperature: 0.1 + + - name: presenter + role: Results Presenter + job: > + You are a results presenter. Take the final analysis and create a professional + presentation of the findings with clear recommendations and next steps. + model: + provider: openai + name: gpt-4o-mini + settings: + temperature: 0.4 + + routers: + - name: flexible_reflection_router + type: reflection + flow_pattern: [analyst, validator, analyst, presenter] + settings: + allow_early_exit: true # Allow skipping steps if appropriate + fallback_strategy: first + + workflow: + start: analyst + edges: + - from: analyst + to: [validator, presenter] + router: flexible_reflection_router + - from: validator + to: [analyst, presenter] + router: flexible_reflection_router + - from: presenter + to: [end] + end: [presenter] +""" + + +async def run_main_critic_flow_example(): + """Example 1: Main -> Critic -> Main -> Final flow""" + print('πŸš€ EXAMPLE 1: Main -> Critic -> Main -> Final Flow') + print('=' * 60) + + # Create workflow from YAML + builder = AriumBuilder.from_yaml( + yaml_str=MAIN_CRITIC_FLOW_YAML, + base_llm=OpenAI(model='gpt-4o-mini', api_key='dummy-key'), # Dummy key for demo + ) + + # Build the workflow + builder.build() + + print('βœ… Workflow built successfully!') + print('πŸ“‹ Reflection Pattern: main_agent β†’ critic β†’ main_agent β†’ final_agent') + print('🎯 The ReflectionRouter will:') + print(' 1. Start with main_agent for initial analysis') + print(' 2. Route to critic for feedback/reflection') + print(' 3. Return to main_agent for improvements') + print(' 4. Finally route to final_agent for polishing') + + # Test input + test_input = 'Write a comprehensive guide on sustainable urban planning' + print(f'\nπŸ“ Test Input: {test_input}') + print('πŸ’‘ This would follow the strict Aβ†’Bβ†’Aβ†’C pattern automatically!') + + +async def run_strict_flow_example(): + """Example 2: Strict flow with no deviations""" + print('\n\n🎯 EXAMPLE 2: Strict Writer -> Reviewer -> Writer -> Editor Flow') + print('=' * 70) + + builder = AriumBuilder.from_yaml( + yaml_str=STRICT_FLOW_YAML, + base_llm=OpenAI(model='gpt-4o-mini', api_key='dummy-key'), + ) + + builder.build() + + print('βœ… Strict workflow built successfully!') + print('πŸ“‹ Flow Pattern: writer β†’ reviewer β†’ writer β†’ editor') + print('πŸ”’ Features:') + print(' β€’ Strict adherence to pattern (allow_early_exit: false)') + print(' β€’ LLM cannot deviate from the Aβ†’Bβ†’Aβ†’C sequence') + print(' β€’ Execution context tracks progress through flow') + + test_input = 'Create a blog post about renewable energy trends in 2024' + print(f'\nπŸ“ Test Input: {test_input}') + + +async def run_flexible_flow_example(): + """Example 3: Flexible flow with early exit option""" + print('\n\n🌟 EXAMPLE 3: Flexible Flow with Early Exit Option') + print('=' * 60) + + builder = AriumBuilder.from_yaml( + yaml_str=FLEXIBLE_FLOW_YAML, + base_llm=OpenAI(model='gpt-4o-mini', api_key='dummy-key'), + ) + + builder.build() + + print('βœ… Flexible workflow built successfully!') + print('πŸ“‹ Flow Pattern: analyst β†’ validator β†’ analyst β†’ presenter') + print('πŸ”“ Features:') + print(' β€’ Flexible routing (allow_early_exit: true)') + print(' β€’ LLM can skip steps if analysis is already sufficient') + print(' β€’ Smart adaptation based on conversation context') + + test_input = 'Analyze the quarterly sales data and identify key trends' + print(f'\nπŸ“ Test Input: {test_input}') + + +def demonstrate_reflection_router_features(): + """Show the key features of ReflectionRouter""" + print('\n\nπŸ“‹ ReflectionRouter Key Features') + print('=' * 50) + + features = """ +🎯 Reflection Pattern Tracking: + β€’ Automatically tracks progress through defined reflection sequence + β€’ Uses execution context (node_visit_count) for intelligent routing + β€’ Prevents infinite loops while allowing intentional revisits + +πŸ“Š Visual Progress Display: + β€’ Shows current position in pattern: β—‹ pending, βœ“ completed + β€’ Displays suggested next step based on reflection pattern + β€’ Provides clear feedback on workflow state + +βš™οΈ Configuration Options: + β€’ allow_early_exit: Enable/disable smart reflection termination + β€’ flow_pattern: Define exact sequence (e.g., [main, critic, main, final]) + β€’ Standard LLM router settings (temperature, fallback_strategy) + +πŸ”„ Execution Context Awareness: + β€’ Tracks how many times each node has been visited + β€’ Calculates expected visits based on reflection pattern position + β€’ Intelligently determines next step in sequence + +πŸ“ YAML Configuration: + ```yaml + routers: + - name: my_reflection_router + type: reflection + flow_pattern: [main_agent, critic, main_agent, final_agent] + settings: + allow_early_exit: false + temperature: 0.2 + ``` + +πŸ›‘οΈ Safety Features: + β€’ Inherits anti-infinite-loop mechanisms from base router + β€’ Provides clear error messages for configuration issues + β€’ Graceful fallback when pattern completion detected +""" + + print(features) + + +def show_yaml_schema(): + """Show complete YAML schema for ReflectionRouter""" + print('\n\nπŸ“„ Complete ReflectionRouter YAML Schema') + print('=' * 50) + + schema = """ +# Complete example with ReflectionRouter +metadata: + name: my-flow-workflow + version: 1.0.0 + description: "A -> B -> A -> C flow pattern example" + +arium: + agents: + - name: main_agent + role: Main Agent + job: "Your main agent job description" + model: + provider: openai + name: gpt-4o-mini + settings: + temperature: 0.7 + + - name: critic + role: Critic + job: "Your critic agent job description" + model: + provider: openai + name: gpt-4o-mini + settings: + temperature: 0.3 + + - name: final_agent + role: Final Agent + job: "Your final agent job description" + model: + provider: openai + name: gpt-4o-mini + + routers: + - name: reflection_router + type: reflection # Router type + flow_pattern: [main_agent, critic, main_agent, final_agent] # A->B->A->C pattern + model: # Optional: LLM for routing decisions + provider: openai + name: gpt-4o-mini + settings: # Optional settings + temperature: 0.2 # Router temperature + allow_early_exit: false # Allow early completion + fallback_strategy: first # first, last, random + + workflow: + start: main_agent + edges: + - from: main_agent + to: [critic, final_agent] # All possible destinations + router: reflection_router # Use reflection router + - from: critic + to: [main_agent, final_agent] + router: reflection_router + - from: final_agent + to: [end] + end: [final_agent] +""" + print(schema) + + +async def main(): + """Run all examples""" + print('🌟 ReflectionRouter Examples - A β†’ B β†’ A β†’ C Pattern Implementation') + print('=' * 80) + print('This example demonstrates the new ReflectionRouter for implementing') + print('structured reflection patterns with intelligent LLM-based routing! πŸŽ‰') + + # Show features and schema first + demonstrate_reflection_router_features() + show_yaml_schema() + + # Run examples + await run_main_critic_flow_example() + await run_strict_flow_example() + await run_flexible_flow_example() + + print('\n\nπŸŽ‰ All ReflectionRouter examples completed!') + print('=' * 80) + print('βœ… ReflectionRouter Benefits:') + print(' β€’ Simple YAML configuration for complex reflection patterns') + print(' β€’ Automatic progress tracking through execution context') + print(' β€’ Intelligent routing decisions based on reflection state') + print(' β€’ Flexible vs strict reflection control options') + print(' β€’ Built-in safety features and loop prevention') + print(' β€’ Easy integration with existing Arium workflows') + + print('\nπŸš€ Try it yourself:') + print(' 1. Define your agents (main, critic, final)') + print(' 2. Create a reflection router with your pattern') + print(' 3. Configure workflow edges') + print(' 4. Run your Aβ†’Bβ†’Aβ†’C reflection workflow!') + + +if __name__ == '__main__': + asyncio.run(main()) diff --git a/flo_ai/examples/plan_execute_router_example.py b/flo_ai/examples/plan_execute_router_example.py new file mode 100644 index 00000000..a03132be --- /dev/null +++ b/flo_ai/examples/plan_execute_router_example.py @@ -0,0 +1,501 @@ +""" +Example demonstrating PlanExecuteRouter for Cursor-style plan-and-execute workflows. + +This example shows how to implement plan-and-execute patterns where tasks are broken down +into sequential steps and executed systematically, similar to how Cursor works. +""" + +import asyncio +from flo_ai.arium.builder import AriumBuilder +from flo_ai.llm import OpenAI + +# Example YAML configuration for Plan-Execute workflow +PLAN_EXECUTE_WORKFLOW_YAML = """ +metadata: + name: plan-execute-development-workflow + version: 1.0.0 + description: "Cursor-style plan-and-execute workflow for software development" + +arium: + agents: + - name: planner + role: Project Planner + job: > + You are a project planner who breaks down complex tasks into detailed, sequential steps. + Create comprehensive execution plans with clear dependencies and assigned agents. + Output your plan as a structured format that can be executed step by step. + model: + provider: openai + name: gpt-4o-mini + settings: + temperature: 0.3 + + - name: developer + role: Software Developer + job: > + You are a software developer who implements code based on specific requirements. + Execute development tasks step by step, focusing on clean, maintainable code. + Report on progress and any issues encountered during implementation. + model: + provider: openai + name: gpt-4o-mini + settings: + temperature: 0.5 + + - name: tester + role: Quality Assurance Tester + job: > + You are a QA tester who validates implementations and ensures quality. + Test code, identify bugs, and verify that requirements are met. + Provide detailed feedback on quality and functionality. + model: + provider: openai + name: gpt-4o-mini + settings: + temperature: 0.2 + + - name: reviewer + role: Code Reviewer + job: > + You are a code reviewer who provides final quality assessment. + Review completed work for best practices, maintainability, and correctness. + Provide final approval or request additional improvements. + model: + provider: openai + name: gpt-4o-mini + settings: + temperature: 0.1 + + # Plan-Execute router configuration + routers: + - name: plan_execute_router + type: plan_execute + agents: + planner: "Creates detailed execution plans by breaking down tasks into sequential steps" + developer: "Implements code and features according to plan specifications" + tester: "Tests implementations and validates functionality" + reviewer: "Reviews and validates completed work for final approval" + model: + provider: openai + name: gpt-4o-mini + settings: + temperature: 0.2 + planner_agent: planner + executor_agent: developer + reviewer_agent: reviewer + + workflow: + start: planner + edges: + - from: planner + to: [developer, tester, reviewer, planner] # All possible destinations + router: plan_execute_router + - from: developer + to: [tester, reviewer, developer, planner] + router: plan_execute_router + - from: tester + to: [developer, reviewer, tester, planner] + router: plan_execute_router + - from: reviewer + to: [end] + end: [reviewer] +""" + +# Simpler Plan-Execute workflow +SIMPLE_PLAN_EXECUTE_YAML = """ +metadata: + name: simple-plan-execute-workflow + version: 1.0.0 + description: "Simple plan-execute pattern for general tasks" + +arium: + agents: + - name: planner + role: Task Planner + job: > + Break down the given task into clear, actionable steps. + Create a detailed plan with dependencies and execution order. + model: + provider: openai + name: gpt-4o-mini + + - name: executor + role: Task Executor + job: > + Execute the steps from the plan one by one. + Focus on completing each step thoroughly before moving to the next. + model: + provider: openai + name: gpt-4o-mini + + - name: validator + role: Quality Validator + job: > + Validate that all steps have been completed correctly. + Ensure the final result meets the original requirements. + model: + provider: openai + name: gpt-4o-mini + + routers: + - name: simple_plan_router + type: plan_execute + agents: + planner: "Creates execution plans" + executor: "Executes plan steps" + validator: "Validates results" + settings: + planner_agent: planner + executor_agent: executor + reviewer_agent: validator + + workflow: + start: planner + edges: + - from: planner + to: [executor, validator, planner] + router: simple_plan_router + - from: executor + to: [validator, executor, planner] + router: simple_plan_router + - from: validator + to: [end] + end: [validator] +""" + +# Research workflow with plan-execute pattern +RESEARCH_PLAN_EXECUTE_YAML = """ +metadata: + name: research-plan-execute-workflow + version: 1.0.0 + description: "Plan-execute workflow for research projects" + +arium: + agents: + - name: research_planner + role: Research Planner + job: > + Create comprehensive research plans by breaking down research questions + into specific investigation steps, data collection tasks, and analysis phases. + model: + provider: openai + name: gpt-4o-mini + + - name: researcher + role: Researcher + job: > + Conduct research according to the plan. Gather information, analyze data, + and document findings for each step of the research plan. + model: + provider: openai + name: gpt-4o-mini + + - name: analyst + role: Data Analyst + job: > + Analyze research data and findings. Identify patterns, draw conclusions, + and prepare analytical insights based on the collected information. + model: + provider: openai + name: gpt-4o-mini + + - name: synthesizer + role: Research Synthesizer + job: > + Synthesize all research findings into a comprehensive final report. + Ensure all research questions are addressed and conclusions are well-supported. + model: + provider: openai + name: gpt-4o-mini + + routers: + - name: research_plan_router + type: plan_execute + agents: + research_planner: "Creates detailed research execution plans" + researcher: "Conducts research and gathers information" + analyst: "Analyzes data and identifies patterns" + synthesizer: "Creates final comprehensive reports" + settings: + planner_agent: research_planner + executor_agent: researcher + reviewer_agent: synthesizer + + workflow: + start: research_planner + edges: + - from: research_planner + to: [researcher, analyst, synthesizer, research_planner] + router: research_plan_router + - from: researcher + to: [analyst, synthesizer, researcher, research_planner] + router: research_plan_router + - from: analyst + to: [synthesizer, analyst, researcher, research_planner] + router: research_plan_router + - from: synthesizer + to: [end] + end: [synthesizer] +""" + + +async def run_development_workflow_example(): + """Example 1: Full development workflow with plan-execute pattern""" + print('πŸš€ EXAMPLE 1: Development Workflow with Plan-Execute Pattern') + print('=' * 65) + + # Create workflow from YAML + builder = AriumBuilder.from_yaml( + yaml_str=PLAN_EXECUTE_WORKFLOW_YAML, + base_llm=OpenAI(model='gpt-4o-mini', api_key='dummy-key'), # Dummy key for demo + ) + + # Build the workflow + builder.build() + + print('βœ… Development workflow built successfully!') + print('πŸ“‹ Plan-Execute Pattern: Planner β†’ Developer β†’ Tester β†’ Reviewer') + print('🎯 The PlanExecuteRouter will:') + print(' 1. Start with planner to create detailed execution plan') + print(' 2. Route to developer for step-by-step implementation') + print(' 3. Route to tester for quality validation') + print(' 4. Route to reviewer for final approval') + print(' 5. Track progress through each step with visual indicators') + + # Test input + test_input = 'Create a REST API for user authentication with JWT tokens' + print(f'\nπŸ“ Test Input: {test_input}') + print('πŸ’‘ Expected plan steps:') + print(' β—‹ Design API endpoints and data models') + print(' β—‹ Implement user registration endpoint') + print(' β—‹ Implement login endpoint with JWT generation') + print(' β—‹ Add authentication middleware') + print(' β—‹ Create comprehensive tests') + print(' β—‹ Review and optimize code') + + +async def run_simple_workflow_example(): + """Example 2: Simple plan-execute workflow""" + print('\n\n🎯 EXAMPLE 2: Simple Plan-Execute Workflow') + print('=' * 50) + + builder = AriumBuilder.from_yaml( + yaml_str=SIMPLE_PLAN_EXECUTE_YAML, + base_llm=OpenAI(model='gpt-4o-mini', api_key='dummy-key'), + ) + + builder.build() + + print('βœ… Simple workflow built successfully!') + print('πŸ“‹ Plan-Execute Pattern: Planner β†’ Executor β†’ Validator') + print('🎯 Features:') + print(' β€’ Automatic task breakdown by planner') + print(' β€’ Sequential step execution') + print(' β€’ Progress tracking with plan state management') + print(' β€’ Quality validation before completion') + + test_input = 'Organize a team building event for 20 people' + print(f'\nπŸ“ Test Input: {test_input}') + print('πŸ’‘ Expected workflow:') + print(' 1. Planner creates detailed event plan') + print(' 2. Executor handles each step (venue, catering, activities)') + print(' 3. Validator ensures everything is properly organized') + + +async def run_research_workflow_example(): + """Example 3: Research workflow with plan-execute pattern""" + print('\n\nπŸ”¬ EXAMPLE 3: Research Plan-Execute Workflow') + print('=' * 50) + + builder = AriumBuilder.from_yaml( + yaml_str=RESEARCH_PLAN_EXECUTE_YAML, + base_llm=OpenAI(model='gpt-4o-mini', api_key='dummy-key'), + ) + + builder.build() + + print('βœ… Research workflow built successfully!') + print( + 'πŸ“‹ Plan-Execute Pattern: Research Planner β†’ Researcher β†’ Analyst β†’ Synthesizer' + ) + print('🎯 Features:') + print(' β€’ Structured research methodology') + print(' β€’ Data collection and analysis phases') + print(' β€’ Comprehensive synthesis and reporting') + print(' β€’ Academic-quality research process') + + test_input = 'Research the impact of remote work on employee productivity' + print(f'\nπŸ“ Test Input: {test_input}') + print('πŸ’‘ Expected research plan:') + print(' β—‹ Literature review on remote work studies') + print(' β—‹ Survey design and data collection') + print(' β—‹ Statistical analysis of productivity metrics') + print(' β—‹ Qualitative analysis of employee feedback') + print(' β—‹ Final report with recommendations') + + +def demonstrate_plan_execute_features(): + """Show the key features of PlanExecuteRouter""" + print('\n\nπŸ“‹ PlanExecuteRouter Key Features') + print('=' * 45) + + features = """ +🎯 Cursor-Style Planning: + β€’ Automatic task breakdown into sequential steps + β€’ Dependency tracking between steps + β€’ Agent assignment for each step + β€’ Progress visualization with status indicators + +πŸ“Š Plan Management: + β€’ ExecutionPlan storage in enhanced memory + β€’ Step status tracking (pending, in_progress, completed, failed) + β€’ Automatic next-step determination + β€’ Failed step recovery and retry logic + +πŸ”„ Execution Flow: + β€’ Phase 1: Planning (create detailed execution plan) + β€’ Phase 2: Execution (complete steps sequentially) + β€’ Phase 3: Review (validate final results) + β€’ Automatic routing between phases + +βš™οΈ Configuration Options: + β€’ agents: Define available agents and their capabilities + β€’ planner_agent: Agent responsible for creating plans + β€’ executor_agent: Default agent for executing steps + β€’ reviewer_agent: Optional agent for final review + +πŸ›‘οΈ Safety Features: + β€’ Prevents infinite loops with step completion tracking + β€’ Handles failed steps with recovery mechanisms + β€’ Memory persistence for plan state + β€’ Execution context awareness + +πŸ“ YAML Configuration: + ```yaml + routers: + - name: plan_execute_router + type: plan_execute + agents: + planner: "Creates execution plans" + developer: "Implements features" + tester: "Validates quality" + settings: + planner_agent: planner + executor_agent: developer + reviewer_agent: tester + ``` + +πŸ”§ Memory Integration: + β€’ Uses PlanAwareMemory for plan storage + β€’ Automatic plan state persistence + β€’ Step result tracking and history + β€’ Context sharing between execution steps +""" + + print(features) + + +def show_yaml_schema(): + """Show complete YAML schema for PlanExecuteRouter""" + print('\n\nπŸ“„ Complete PlanExecuteRouter YAML Schema') + print('=' * 55) + + schema = """ +# Complete example with PlanExecuteRouter +metadata: + name: my-plan-execute-workflow + version: 1.0.0 + description: "Cursor-style plan-and-execute workflow" + +arium: + agents: + - name: planner + role: "Task Planner" + job: "Create detailed execution plans" + model: + provider: openai + name: gpt-4o-mini + + - name: executor + role: "Task Executor" + job: "Execute plan steps systematically" + model: + provider: openai + name: gpt-4o-mini + + - name: reviewer + role: "Quality Reviewer" + job: "Review and validate final results" + model: + provider: openai + name: gpt-4o-mini + + routers: + - name: plan_execute_router + type: plan_execute # Router type + agents: # Required: Available agents + planner: "Creates detailed execution plans" + executor: "Executes individual plan steps" + reviewer: "Reviews final results" + specialist: "Handles specialized tasks" + model: # Optional: LLM for routing + provider: openai + name: gpt-4o-mini + settings: # Optional settings + temperature: 0.2 # Router temperature + planner_agent: planner # Agent for creating plans + executor_agent: executor # Default execution agent + reviewer_agent: reviewer # Optional review agent + max_retries: 3 # Max retries for failed steps + + workflow: + start: planner + edges: + - from: planner + to: [executor, reviewer, specialist, planner] # All possible destinations + router: plan_execute_router + - from: executor + to: [reviewer, specialist, executor, planner] + router: plan_execute_router + - from: reviewer + to: [end] + end: [reviewer] +""" + print(schema) + + +async def main(): + """Run all examples""" + print('🌟 PlanExecuteRouter Examples - Cursor-Style Plan-and-Execute Workflows') + print('=' * 85) + print('This example demonstrates the new PlanExecuteRouter for implementing') + print('sophisticated plan-and-execute patterns with intelligent step tracking! πŸŽ‰') + + # Show features and schema first + demonstrate_plan_execute_features() + show_yaml_schema() + + # Run examples + await run_development_workflow_example() + await run_simple_workflow_example() + await run_research_workflow_example() + + print('\n\nπŸŽ‰ All PlanExecuteRouter examples completed!') + print('=' * 85) + print('βœ… PlanExecuteRouter Benefits:') + print(' β€’ Cursor-style automatic task breakdown and execution') + print(' β€’ Intelligent step-by-step progress tracking') + print(' β€’ Visual progress indicators with plan state management') + print(' β€’ Automatic routing based on execution plan progress') + print(' β€’ Built-in error handling and step retry mechanisms') + print(' β€’ Enhanced memory system with plan persistence') + + print('\nπŸš€ Try it yourself:') + print(' 1. Define your planner, executor, and reviewer agents') + print(' 2. Create a plan-execute router with agent mappings') + print(' 3. Use PlanAwareMemory for plan state persistence') + print(' 4. Run complex tasks with automatic breakdown!') + + +if __name__ == '__main__': + asyncio.run(main()) diff --git a/flo_ai/examples/reflection_router_example.py b/flo_ai/examples/reflection_router_example.py new file mode 100644 index 00000000..2da1e6c6 --- /dev/null +++ b/flo_ai/examples/reflection_router_example.py @@ -0,0 +1,454 @@ +""" +Example demonstrating ReflectionRouter for A -> B -> A -> C patterns. + +This example shows how to implement a main -> critic -> main -> final reflection workflow +using the new ReflectionRouter with YAML configuration. +""" + +import asyncio +from flo_ai.arium.builder import AriumBuilder +from flo_ai.llm import OpenAI + +# Example YAML configuration for A -> B -> A -> C flow +MAIN_CRITIC_FLOW_YAML = """ +metadata: + name: main-critic-final-workflow + version: 1.0.0 + description: "A workflow demonstrating A -> B -> A -> C pattern with intelligent flow routing" + +arium: + agents: + - 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. + Be receptive to criticism and use it to refine your output. + model: + provider: openai + name: gpt-4o-mini + settings: + temperature: 0.7 + + - name: critic + role: Critic Agent + job: > + You are a critic agent. Your job is to review the main agent's work and provide + constructive feedback. Analyze the output for: + - Accuracy and correctness + - Completeness and thoroughness + - Clarity and coherence + - Areas for improvement + Provide specific, actionable feedback that the main agent can use to improve. + model: + provider: openai + name: gpt-4o-mini + settings: + temperature: 0.3 + + - name: final_agent + role: Final Agent + job: > + You are the final agent responsible for polishing and finalizing the work. + Take the refined output from the main agent (after critic feedback) and: + - Format it professionally + - Add any final touches or improvements + - Ensure it meets high quality standards + - Provide a polished final result + model: + provider: openai + name: gpt-4o-mini + settings: + temperature: 0.5 + + # Reflection router configuration for A -> B -> A -> C pattern + routers: + - name: main_critic_reflection_router + type: reflection + flow_pattern: [main_agent, critic, main_agent, final_agent] + model: + provider: openai + name: gpt-4o-mini + settings: + temperature: 0.2 + allow_early_exit: false + fallback_strategy: first + + workflow: + start: main_agent + edges: + # Single edge from main_agent using reflection router + # The router will intelligently route to: critic -> main_agent -> final_agent + - from: main_agent + to: [critic, final_agent] # All possible destinations + router: main_critic_reflection_router + - from: critic + to: [main_agent, final_agent] + router: main_critic_reflection_router + - from: final_agent + to: [end] + end: [final_agent] +""" + +# Alternative stricter flow pattern +STRICT_FLOW_YAML = """ +metadata: + name: strict-main-critic-flow + version: 1.0.0 + description: "Strict A -> B -> A -> C flow with no deviations allowed" + +arium: + agents: + - name: writer + role: Content Writer + job: > + You are a content writer. Create initial content based on the user's request. + Focus on getting the core ideas down first, don't worry about perfection. + model: + provider: openai + name: gpt-4o-mini + settings: + temperature: 0.8 + + - name: reviewer + role: Content Reviewer + job: > + You are a content reviewer. Review the writer's work and provide detailed feedback: + - What works well + - What needs improvement + - Specific suggestions for enhancement + - Areas that need clarification + model: + provider: openai + name: gpt-4o-mini + settings: + temperature: 0.2 + + - name: editor + role: Content Editor + job: > + You are the final editor. Take the revised content from the writer and: + - Polish the language and style + - Ensure consistency and flow + - Make final corrections + - Prepare the content for publication + model: + provider: openai + name: gpt-4o-mini + settings: + temperature: 0.3 + + routers: + - name: strict_reflection_router + type: reflection + flow_pattern: [writer, reviewer, writer, editor] + settings: + allow_early_exit: false # Strict adherence to pattern + fallback_strategy: first + + workflow: + start: writer + edges: + - from: writer + to: [reviewer, editor] + router: strict_reflection_router + - from: reviewer + to: [writer, editor] + router: strict_reflection_router + - from: editor + to: [end] + end: [editor] +""" + +# Flexible flow that allows early exit +FLEXIBLE_FLOW_YAML = """ +metadata: + name: flexible-flow-with-early-exit + version: 1.0.0 + description: "Flexible A -> B -> A -> C flow that allows early completion" + +arium: + agents: + - name: analyst + role: Data Analyst + job: > + You are a data analyst. Analyze the given data or question and provide insights. + Create clear, actionable analysis based on the information provided. + model: + provider: openai + name: gpt-4o-mini + settings: + temperature: 0.5 + + - name: validator + role: Analysis Validator + job: > + You are an analysis validator. Review the analyst's work for: + - Logical consistency + - Accuracy of conclusions + - Completeness of analysis + - Potential issues or gaps + If the analysis is solid, you can recommend proceeding directly to completion. + model: + provider: openai + name: gpt-4o-mini + settings: + temperature: 0.1 + + - name: presenter + role: Results Presenter + job: > + You are a results presenter. Take the final analysis and create a professional + presentation of the findings with clear recommendations and next steps. + model: + provider: openai + name: gpt-4o-mini + settings: + temperature: 0.4 + + routers: + - name: flexible_reflection_router + type: reflection + flow_pattern: [analyst, validator, analyst, presenter] + settings: + allow_early_exit: true # Allow skipping steps if appropriate + fallback_strategy: first + + workflow: + start: analyst + edges: + - from: analyst + to: [validator, presenter] + router: flexible_reflection_router + - from: validator + to: [analyst, presenter] + router: flexible_reflection_router + - from: presenter + to: [end] + end: [presenter] +""" + + +async def run_main_critic_flow_example(): + """Example 1: Main -> Critic -> Main -> Final flow""" + print('πŸš€ EXAMPLE 1: Main -> Critic -> Main -> Final Flow') + print('=' * 60) + + # Create workflow from YAML + builder = AriumBuilder.from_yaml( + yaml_str=MAIN_CRITIC_FLOW_YAML, + base_llm=OpenAI(model='gpt-4o-mini', api_key='dummy-key'), # Dummy key for demo + ) + + # Build the workflow + builder.build() + + print('βœ… Workflow built successfully!') + print('πŸ“‹ Reflection Pattern: main_agent β†’ critic β†’ main_agent β†’ final_agent') + print('🎯 The ReflectionRouter will:') + print(' 1. Start with main_agent for initial analysis') + print(' 2. Route to critic for feedback/reflection') + print(' 3. Return to main_agent for improvements') + print(' 4. Finally route to final_agent for polishing') + + # Test input + test_input = 'Write a comprehensive guide on sustainable urban planning' + print(f'\nπŸ“ Test Input: {test_input}') + print('πŸ’‘ This would follow the strict Aβ†’Bβ†’Aβ†’C pattern automatically!') + + +async def run_strict_flow_example(): + """Example 2: Strict flow with no deviations""" + print('\n\n🎯 EXAMPLE 2: Strict Writer -> Reviewer -> Writer -> Editor Flow') + print('=' * 70) + + builder = AriumBuilder.from_yaml( + yaml_str=STRICT_FLOW_YAML, + base_llm=OpenAI(model='gpt-4o-mini', api_key='dummy-key'), + ) + + builder.build() + + print('βœ… Strict workflow built successfully!') + print('πŸ“‹ Flow Pattern: writer β†’ reviewer β†’ writer β†’ editor') + print('πŸ”’ Features:') + print(' β€’ Strict adherence to pattern (allow_early_exit: false)') + print(' β€’ LLM cannot deviate from the Aβ†’Bβ†’Aβ†’C sequence') + print(' β€’ Execution context tracks progress through flow') + + test_input = 'Create a blog post about renewable energy trends in 2024' + print(f'\nπŸ“ Test Input: {test_input}') + + +async def run_flexible_flow_example(): + """Example 3: Flexible flow with early exit option""" + print('\n\n🌟 EXAMPLE 3: Flexible Flow with Early Exit Option') + print('=' * 60) + + builder = AriumBuilder.from_yaml( + yaml_str=FLEXIBLE_FLOW_YAML, + base_llm=OpenAI(model='gpt-4o-mini', api_key='dummy-key'), + ) + + builder.build() + + print('βœ… Flexible workflow built successfully!') + print('πŸ“‹ Flow Pattern: analyst β†’ validator β†’ analyst β†’ presenter') + print('πŸ”“ Features:') + print(' β€’ Flexible routing (allow_early_exit: true)') + print(' β€’ LLM can skip steps if analysis is already sufficient') + print(' β€’ Smart adaptation based on conversation context') + + test_input = 'Analyze the quarterly sales data and identify key trends' + print(f'\nπŸ“ Test Input: {test_input}') + + +def demonstrate_reflection_router_features(): + """Show the key features of ReflectionRouter""" + print('\n\nπŸ“‹ ReflectionRouter Key Features') + print('=' * 50) + + features = """ +🎯 Reflection Pattern Tracking: + β€’ Automatically tracks progress through defined reflection sequence + β€’ Uses execution context (node_visit_count) for intelligent routing + β€’ Prevents infinite loops while allowing intentional revisits + +πŸ“Š Visual Progress Display: + β€’ Shows current position in pattern: β—‹ pending, βœ“ completed + β€’ Displays suggested next step based on reflection pattern + β€’ Provides clear feedback on workflow state + +βš™οΈ Configuration Options: + β€’ allow_early_exit: Enable/disable smart reflection termination + β€’ flow_pattern: Define exact sequence (e.g., [main, critic, main, final]) + β€’ Standard LLM router settings (temperature, fallback_strategy) + +πŸ”„ Execution Context Awareness: + β€’ Tracks how many times each node has been visited + β€’ Calculates expected visits based on reflection pattern position + β€’ Intelligently determines next step in sequence + +πŸ“ YAML Configuration: + ```yaml + routers: + - name: my_reflection_router + type: reflection + flow_pattern: [main_agent, critic, main_agent, final_agent] + settings: + allow_early_exit: false + temperature: 0.2 + ``` + +πŸ›‘οΈ Safety Features: + β€’ Inherits anti-infinite-loop mechanisms from base router + β€’ Provides clear error messages for configuration issues + β€’ Graceful fallback when pattern completion detected +""" + + print(features) + + +def show_yaml_schema(): + """Show complete YAML schema for ReflectionRouter""" + print('\n\nπŸ“„ Complete ReflectionRouter YAML Schema') + print('=' * 50) + + schema = """ +# Complete example with ReflectionRouter +metadata: + name: my-flow-workflow + version: 1.0.0 + description: "A -> B -> A -> C flow pattern example" + +arium: + agents: + - name: main_agent + role: Main Agent + job: "Your main agent job description" + model: + provider: openai + name: gpt-4o-mini + settings: + temperature: 0.7 + + - name: critic + role: Critic + job: "Your critic agent job description" + model: + provider: openai + name: gpt-4o-mini + settings: + temperature: 0.3 + + - name: final_agent + role: Final Agent + job: "Your final agent job description" + model: + provider: openai + name: gpt-4o-mini + + routers: + - name: reflection_router + type: reflection # Router type + flow_pattern: [main_agent, critic, main_agent, final_agent] # A->B->A->C pattern + model: # Optional: LLM for routing decisions + provider: openai + name: gpt-4o-mini + settings: # Optional settings + temperature: 0.2 # Router temperature + allow_early_exit: false # Allow early completion + fallback_strategy: first # first, last, random + + workflow: + start: main_agent + edges: + - from: main_agent + to: [critic, final_agent] # All possible destinations + router: reflection_router # Use reflection router + - from: critic + to: [main_agent, final_agent] + router: reflection_router + - from: final_agent + to: [end] + end: [final_agent] +""" + print(schema) + + +async def main(): + """Run all examples""" + print('🌟 ReflectionRouter Examples - A β†’ B β†’ A β†’ C Pattern Implementation') + print('=' * 80) + print('This example demonstrates the new ReflectionRouter for implementing') + print('structured reflection patterns with intelligent LLM-based routing! πŸŽ‰') + + # Show features and schema first + demonstrate_reflection_router_features() + show_yaml_schema() + + # Run examples + await run_main_critic_flow_example() + await run_strict_flow_example() + await run_flexible_flow_example() + + print('\n\nπŸŽ‰ All ReflectionRouter examples completed!') + print('=' * 80) + print('βœ… ReflectionRouter Benefits:') + print(' β€’ Simple YAML configuration for complex reflection patterns') + print(' β€’ Automatic progress tracking through execution context') + print(' β€’ Intelligent routing decisions based on reflection state') + print(' β€’ Flexible vs strict reflection control options') + print(' β€’ Built-in safety features and loop prevention') + print(' β€’ Easy integration with existing Arium workflows') + + print('\nπŸš€ Try it yourself:') + print(' 1. Define your agents (main, critic, final)') + print(' 2. Create a reflection router with your pattern') + print(' 3. Configure workflow edges') + print(' 4. Run your Aβ†’Bβ†’Aβ†’C reflection workflow!') + + +if __name__ == '__main__': + asyncio.run(main()) diff --git a/flo_ai/examples/simple_flow_router_demo.py b/flo_ai/examples/simple_flow_router_demo.py new file mode 100644 index 00000000..8b917b60 --- /dev/null +++ b/flo_ai/examples/simple_flow_router_demo.py @@ -0,0 +1,181 @@ +""" +Simple demonstration of ReflectionRouter for A -> B -> A -> C patterns. + +This shows the minimal code needed to implement a main -> critic -> main -> final +reflection pattern using the new ReflectionRouter. +""" + +import asyncio +from flo_ai.arium.builder import AriumBuilder +from flo_ai.arium.memory import MessageMemory +from flo_ai.models.agent import Agent +from flo_ai.llm import OpenAI +from flo_ai.arium.llm_router import create_main_critic_reflection_router + + +async def simple_reflection_demo(): + """Minimal example of A -> B -> A -> C reflection pattern""" + print('πŸš€ Simple ReflectionRouter Demo: A β†’ B β†’ A β†’ C Pattern') + print('=' * 60) + + # Create LLM (use dummy key for demo) + llm = OpenAI(model='gpt-4o-mini', api_key='dummy-key') + + # Create agents + main_agent = Agent( + name='main_agent', + system_prompt='You are the main agent. Analyze tasks and create solutions.', + llm=llm, + ) + + critic = Agent( + name='critic', + system_prompt='You are a critic. Provide constructive feedback to improve work.', + llm=llm, + ) + + final_agent = Agent( + name='final_agent', + system_prompt='You are the final agent. Polish and finalize the work.', + llm=llm, + ) + + # Create reflection router for A -> B -> A -> C pattern + reflection_router = create_main_critic_reflection_router( + main_agent='main_agent', + critic_agent='critic', + final_agent='final_agent', + allow_early_exit=False, # Strict reflection + llm=llm, + ) + + # Build workflow + builder = ( + AriumBuilder() + .with_memory(MessageMemory()) + .add_agents([main_agent, critic, final_agent]) + .start_with(main_agent) + .add_edge(main_agent, [critic, final_agent], reflection_router) + .add_edge(critic, [main_agent, final_agent], reflection_router) + .end_with(final_agent) + ) + + # Build the Arium + arium = builder.build() + + print('βœ… Workflow created successfully!') + print('πŸ“‹ Reflection Pattern: main_agent β†’ critic β†’ main_agent β†’ final_agent') + print('🎯 Router will automatically follow this sequence') + + # Demo input + print('\nπŸ“ Example input: "Create a project plan for a mobile app"') + print('πŸ’‘ The reflection router will:') + print(' Step 1: Route to critic (for feedback/reflection)') + print(' Step 2: Return to main_agent (to incorporate feedback)') + print(' Step 3: Route to final_agent (for final polish)') + + return arium + + +async def programmatic_reflection_example(): + """Show how to create reflection router programmatically""" + print('\n\nπŸ”§ Programmatic ReflectionRouter Creation') + print('=' * 50) + + from flo_ai.arium.llm_router import create_llm_router + + # Method 1: Using the convenience function + create_main_critic_reflection_router( + main_agent='writer', critic_agent='reviewer', final_agent='editor' + ) + print('βœ… Method 1: Convenience function create_main_critic_reflection_router()') + + # Method 2: Using the factory function directly + create_llm_router( + 'reflection', + flow_pattern=['analyst', 'validator', 'analyst', 'presenter'], + allow_early_exit=True, + ) + print('βœ… Method 2: Factory function create_llm_router(type="reflection")') + + # Method 3: Creating ReflectionRouter directly + from flo_ai.arium.llm_router import ReflectionRouter + + ReflectionRouter( + flow_pattern=['main', 'critic', 'main', 'final'], allow_early_exit=False + ) + print('βœ… Method 3: Direct ReflectionRouter instantiation') + + print('\n🎯 All methods create the same Aβ†’Bβ†’Aβ†’C reflection pattern!') + + +def show_minimal_yaml(): + """Show the minimal YAML needed""" + print('\n\nπŸ“„ Minimal YAML Configuration') + print('=' * 35) + + yaml_example = """ +# Minimal ReflectionRouter YAML +arium: + agents: + - name: main_agent + job: "Main agent job" + model: {provider: openai, name: gpt-4o-mini} + - name: critic + job: "Critic job" + model: {provider: openai, name: gpt-4o-mini} + - name: final_agent + job: "Final agent job" + model: {provider: openai, name: gpt-4o-mini} + + routers: + - name: reflection_router + type: reflection + flow_pattern: [main_agent, critic, main_agent, final_agent] + + workflow: + start: main_agent + edges: + - from: main_agent + to: [critic, final_agent] + router: reflection_router + - from: critic + to: [main_agent, final_agent] + router: reflection_router + - from: final_agent + to: [end] + end: [final_agent] +""" + print(yaml_example) + + +async def main(): + """Run the simple demo""" + print('🌟 Simple ReflectionRouter Demo') + print('=' * 35) + print('Demonstrating A β†’ B β†’ A β†’ C reflection pattern with minimal setup\n') + + # Show different creation methods + await programmatic_reflection_example() + + # Show minimal YAML + show_minimal_yaml() + + # Create simple reflection workflow + arium = await simple_reflection_demo() + + print('\n\nπŸŽ‰ Demo completed!') + print('Key takeaways:') + print('βœ… ReflectionRouter makes Aβ†’Bβ†’Aβ†’C patterns trivial to implement') + print('βœ… Works with both YAML and programmatic configuration') + print('βœ… Automatically tracks progress and prevents infinite loops') + print('βœ… Provides intelligent routing based on execution context') + + result = await arium.run( + inputs=['Write a comprehensive guide on sustainable urban planning'] + ) + print(result) + + +if __name__ == '__main__': + asyncio.run(main()) diff --git a/flo_ai/examples/simple_plan_execute_demo.py b/flo_ai/examples/simple_plan_execute_demo.py new file mode 100644 index 00000000..0b4e5607 --- /dev/null +++ b/flo_ai/examples/simple_plan_execute_demo.py @@ -0,0 +1,418 @@ +""" +Simple demonstration of PlanExecuteRouter with actual plan creation and execution. + +This demo shows how the PlanExecuteRouter works with PlanAwareMemory to create, +store, and execute plans step by step, similar to how Cursor works. +""" + +import asyncio +import uuid +from flo_ai.arium.builder import AriumBuilder +from flo_ai.arium.memory import PlanAwareMemory, ExecutionPlan, PlanStep, StepStatus +from flo_ai.llm import OpenAI +from flo_ai.models.agent import Agent +from flo_ai.arium.llm_router import create_plan_execute_router + + +async def demo_plan_aware_memory(): + """Demonstrate PlanAwareMemory functionality""" + print('πŸ“‹ DEMO: PlanAwareMemory - Plan Storage and Management') + print('=' * 55) + + # Create plan-aware memory + memory = PlanAwareMemory() + + # Create a sample execution plan + plan = ExecutionPlan( + id=str(uuid.uuid4()), + title='Build User Authentication System', + description='Create a complete user authentication system with login, registration, and JWT tokens', + steps=[ + PlanStep( + id='step_1', + description='Design database schema for users', + agent='developer', + status=StepStatus.PENDING, + ), + PlanStep( + id='step_2', + description='Implement user registration endpoint', + agent='developer', + dependencies=['step_1'], + status=StepStatus.PENDING, + ), + PlanStep( + id='step_3', + description='Implement login endpoint with JWT generation', + agent='developer', + dependencies=['step_1', 'step_2'], + status=StepStatus.PENDING, + ), + PlanStep( + id='step_4', + description='Add authentication middleware', + agent='developer', + dependencies=['step_3'], + status=StepStatus.PENDING, + ), + PlanStep( + id='step_5', + description='Write comprehensive tests', + agent='tester', + dependencies=['step_4'], + status=StepStatus.PENDING, + ), + PlanStep( + id='step_6', + description='Final code review and optimization', + agent='reviewer', + dependencies=['step_5'], + status=StepStatus.PENDING, + ), + ], + ) + + # Add plan to memory + memory.add_plan(plan) + print(f'βœ… Plan added to memory: {plan.title}') + print(f'πŸ“Š Total steps: {len(plan.steps)}') + + # Show plan progress + def show_progress(): + current_plan = memory.get_current_plan() + if current_plan: + print(f'\nπŸ“‹ Plan Progress: {current_plan.title}') + for step in current_plan.steps: + status_icon = { + StepStatus.PENDING: 'β—‹', + StepStatus.IN_PROGRESS: '⏳', + StepStatus.COMPLETED: 'βœ…', + StepStatus.FAILED: '❌', + }.get(step.status, 'β—‹') + deps = ( + f" (depends on: {', '.join(step.dependencies)})" + if step.dependencies + else '' + ) + print( + f' {status_icon} {step.id}: {step.description} β†’ {step.agent}{deps}' + ) + + show_progress() + + # Simulate step execution + print('\nπŸ”„ Simulating step execution...') + + # Execute step 1 + current_plan = memory.get_current_plan() + next_steps = current_plan.get_next_steps() + if next_steps: + step = next_steps[0] + print(f'\n⏳ Executing: {step.description}') + step.status = StepStatus.IN_PROGRESS + memory.update_plan(current_plan) + + # Simulate completion + step.status = StepStatus.COMPLETED + step.result = ( + 'User table created with id, email, password_hash, created_at fields' + ) + memory.update_plan(current_plan) + print(f'βœ… Completed: {step.description}') + + # Execute step 2 + current_plan = memory.get_current_plan() + next_steps = current_plan.get_next_steps() + if next_steps: + step = next_steps[0] + print(f'\n⏳ Executing: {step.description}') + step.status = StepStatus.COMPLETED + step.result = 'POST /api/register endpoint implemented with validation' + memory.update_plan(current_plan) + print(f'βœ… Completed: {step.description}') + + show_progress() + + # Check what's next + current_plan = memory.get_current_plan() + next_steps = current_plan.get_next_steps() + print(f'\n🎯 Next steps ready for execution: {len(next_steps)}') + for step in next_steps: + print(f' β†’ {step.id}: {step.description} (agent: {step.agent})') + + print(f'\nπŸ“ˆ Plan completion: {current_plan.is_completed()}') + + +async def demo_programmatic_plan_execute(): + """Demonstrate programmatic usage of PlanExecuteRouter""" + print('\n\nπŸ—οΈ DEMO: Programmatic PlanExecuteRouter Usage') + print('=' * 55) + + # Create LLM with dummy key for demo + llm = OpenAI(model='gpt-4o-mini', api_key='dummy-key') + + # Create agents + Agent( + name='planner', + system_prompt="""You are an expert project planner. When given a task, create a detailed execution plan. + +When asked to create a plan, respond with a structured format like this: + +EXECUTION PLAN: [Title] +DESCRIPTION: [Brief description] + +STEPS: +1. step_id: [description] β†’ [agent_name] +2. step_id: [description] β†’ [agent_name] (depends on: step1) +3. step_id: [description] β†’ [agent_name] (depends on: step1, step2) + +Always include clear dependencies and assign appropriate agents.""", + llm=llm, + ) + + Agent( + name='developer', + system_prompt='You are a software developer. Execute development tasks step by step.', + llm=llm, + ) + + Agent( + name='tester', + system_prompt='You are a QA tester. Test implementations and validate functionality.', + llm=llm, + ) + + Agent( + name='reviewer', + system_prompt='You are a code reviewer. Review completed work and provide final validation.', + llm=llm, + ) + + # Create plan-execute router + plan_router = create_plan_execute_router( + planner_agent='planner', + executor_agent='developer', + reviewer_agent='reviewer', + additional_agents={ + 'tester': 'Tests implementations and validates functionality' + }, + llm=llm, + ) + + # Create plan-aware memory + memory = PlanAwareMemory() + + print('βœ… Created agents and PlanExecuteRouter') + print('🎯 Router will coordinate: planner β†’ developer β†’ tester β†’ reviewer') + print('πŸ’Ύ Using PlanAwareMemory for plan state management') + + # Simulate routing decisions + print('\n🧠 Simulating router decision making...') + + # Add a message to trigger planning + memory.add({'role': 'user', 'content': 'Create a TODO app with React and Node.js'}) + + # Test routing with no plan (should route to planner) + try: + next_agent = plan_router(memory) + print(f'πŸ“ Router decision (no plan): {next_agent}') + print(' Expected: planner (to create execution plan)') + + except Exception as e: + print(f'⚠️ Router simulation note: {e}') + print(' (This is expected in demo mode without real LLM calls)') + + print('\nπŸ’‘ In a real scenario:') + print(' 1. Router would route to planner to create execution plan') + print(' 2. Planner creates detailed plan and stores in memory') + print(' 3. Router routes to developer for first development step') + print(' 4. Developer completes step and updates plan status') + print(' 5. Router routes to next step based on plan state') + print(' 6. Process continues until all steps complete') + print(' 7. Router routes to reviewer for final validation') + + +async def demo_yaml_plan_execute(): + """Demonstrate YAML configuration for PlanExecuteRouter""" + print('\n\nπŸ“„ DEMO: YAML Configuration for PlanExecuteRouter') + print('=' * 55) + + yaml_config = """ +metadata: + name: simple-plan-execute-demo + version: 1.0.0 + description: "Demo of plan-execute pattern" + +arium: + agents: + - name: planner + role: Task Planner + job: > + Break down complex tasks into detailed, sequential execution plans. + Create clear steps with dependencies and agent assignments. + model: + provider: openai + name: gpt-4o-mini + + - name: executor + role: Task Executor + job: > + Execute plan steps systematically, one by one. + Report progress and update plan status. + model: + provider: openai + name: gpt-4o-mini + + - name: validator + role: Quality Validator + job: > + Validate completed work and ensure quality standards. + Provide final approval for plan completion. + model: + provider: openai + name: gpt-4o-mini + + routers: + - name: demo_plan_router + type: plan_execute + agents: + planner: "Creates detailed execution plans" + executor: "Executes plan steps systematically" + validator: "Validates final results" + model: + provider: openai + name: gpt-4o-mini + settings: + temperature: 0.2 + planner_agent: planner + executor_agent: executor + reviewer_agent: validator + + workflow: + start: planner + edges: + - from: planner + to: [executor, validator, planner] + router: demo_plan_router + - from: executor + to: [validator, executor, planner] + router: demo_plan_router + - from: validator + to: [end] + end: [validator] +""" + + try: + # Build workflow from YAML + builder = AriumBuilder.from_yaml( + yaml_str=yaml_config, + base_llm=OpenAI(model='gpt-4o-mini', api_key='dummy-key'), + ) + + builder.build() + print('βœ… YAML workflow built successfully!') + print('πŸ“‹ Configured plan-execute pattern with 3 agents') + print('πŸ”„ Router will coordinate planning β†’ execution β†’ validation') + + # Show the workflow structure + print('\nπŸ“Š Workflow Structure:') + print(' Start: planner') + print(' Edges:') + print(' planner β†’ [executor, validator, planner] (plan_execute_router)') + print(' executor β†’ [validator, executor, planner] (plan_execute_router)') + print(' validator β†’ [end]') + print(' End: validator') + + except Exception as e: + print(f'ℹ️ YAML demo note: {e}') + + print('\n🎯 Key YAML Features:') + print(' β€’ type: plan_execute - Enables plan-execute routing') + print(' β€’ agents: Dict mapping agent names to descriptions') + print(' β€’ planner_agent: Agent responsible for creating plans') + print(' β€’ executor_agent: Default agent for executing steps') + print(' β€’ reviewer_agent: Optional agent for final review') + + +def show_memory_integration(): + """Show how PlanExecuteRouter integrates with memory""" + print('\n\nπŸ’Ύ DEMO: Memory Integration with PlanExecuteRouter') + print('=' * 55) + + integration_info = """ +πŸ”„ PlanExecuteRouter Memory Integration: + +1. Plan Creation Phase: + β€’ Router detects no plan in memory + β€’ Routes to planner agent + β€’ Planner creates ExecutionPlan with steps + β€’ Plan stored in PlanAwareMemory + +2. Execution Phase: + β€’ Router checks current plan state + β€’ Identifies next ready steps (dependencies met) + β€’ Routes to appropriate agent for step execution + β€’ Agent updates step status and results + +3. Progress Tracking: + β€’ Plan progress visualized with status indicators: + β—‹ Pending ⏳ In Progress βœ… Completed ❌ Failed + β€’ Dependencies automatically managed + β€’ Failed steps trigger recovery routing + +4. Memory Persistence: + β€’ Plan state persists across agent interactions + β€’ Step results and metadata stored + β€’ Execution context maintained + +5. Completion Handling: + β€’ Router detects when all steps complete + β€’ Routes to reviewer agent (if configured) + β€’ Final validation and workflow completion + +πŸ“Š Memory Structure: +```python +memory = PlanAwareMemory() +memory.add_plan(execution_plan) # Store plan +current_plan = memory.get_current_plan() # Retrieve active plan +next_steps = current_plan.get_next_steps() # Get ready steps +``` + +🎯 Router Intelligence: +β€’ Automatically routes based on plan state +β€’ Handles step dependencies and execution order +β€’ Provides context-aware prompts with progress +β€’ Manages error recovery and retry logic +""" + + print(integration_info) + + +async def main(): + """Run all plan-execute demos""" + print('πŸš€ PlanExecuteRouter Simple Demo') + print('=' * 40) + print('This demo shows the PlanExecuteRouter in action with actual plan') + print('creation, storage, and step-by-step execution tracking! πŸŽ‰\n') + + # Run demos + await demo_plan_aware_memory() + await demo_programmatic_plan_execute() + await demo_yaml_plan_execute() + show_memory_integration() + + print('\n\nπŸŽ‰ PlanExecuteRouter Demo Complete!') + print('=' * 45) + print('βœ… What we demonstrated:') + print(' β€’ PlanAwareMemory for plan storage and state tracking') + print(' β€’ ExecutionPlan with steps, dependencies, and status') + print(' β€’ Programmatic router creation and usage') + print(' β€’ YAML configuration for plan-execute workflows') + print(' β€’ Memory integration and plan state management') + + print('\nπŸš€ Ready to build your own Cursor-style workflows!') + print(' Try the PlanExecuteRouter for complex task automation! 🎯') + + +if __name__ == '__main__': + asyncio.run(main()) diff --git a/flo_ai/examples/simple_reflection_router_demo.py b/flo_ai/examples/simple_reflection_router_demo.py new file mode 100644 index 00000000..7c4c2e40 --- /dev/null +++ b/flo_ai/examples/simple_reflection_router_demo.py @@ -0,0 +1,188 @@ +""" +Simple demonstration of ReflectionRouter for A -> B -> A -> C patterns. + +This shows the minimal code needed to implement a main -> critic -> main -> final +reflection pattern using the new ReflectionRouter. +""" + +import asyncio +from flo_ai.arium.builder import AriumBuilder +from flo_ai.arium.memory import MessageMemory +from flo_ai.models.agent import Agent +from flo_ai.llm import OpenAI +from flo_ai.arium.llm_router import create_main_critic_reflection_router + + +async def simple_reflection_demo(): + """Minimal example of A -> B -> A -> C reflection pattern""" + print('πŸš€ Simple ReflectionRouter Demo: A β†’ B β†’ A β†’ C Pattern') + print('=' * 60) + + # Create LLM (use dummy key for demo) + llm = OpenAI(model='gpt-4o-mini') + + # Create agents + main_agent = Agent( + name='main_agent', + system_prompt='You are the main agent. Analyze tasks and create solutions.', + llm=llm, + ) + + critic = Agent( + name='critic', + system_prompt='You are a critic. Provide constructive feedback to improve work.', + llm=llm, + ) + + final_agent = Agent( + name='final_agent', + system_prompt='You are the final agent. Polish and finalize the work.', + llm=llm, + ) + + # Create reflection router for A -> B -> A -> C pattern + reflection_router = create_main_critic_reflection_router( + main_agent='main_agent', + critic_agent='critic', + final_agent='final_agent', + allow_early_exit=False, # Strict reflection + llm=llm, + ) + + # Build workflow + builder = ( + AriumBuilder() + .with_memory(MessageMemory()) + .add_agents([main_agent, critic, final_agent]) + .start_with(main_agent) + .add_edge(main_agent, [critic, final_agent], reflection_router) + .add_edge(critic, [main_agent, final_agent], reflection_router) + .end_with(final_agent) + ) + + # Build the Arium + arium = builder.build() + + print('βœ… Workflow created successfully!') + print('πŸ“‹ Reflection Pattern: main_agent β†’ critic β†’ main_agent β†’ final_agent') + print('🎯 Router will automatically follow this sequence') + + # Demo input + print('\nπŸ“ Example input: "Create a project plan for a mobile app"') + print('πŸ’‘ The reflection router will:') + print(' Step 1: Route to critic (for feedback/reflection)') + print(' Step 2: Return to main_agent (to incorporate feedback)') + print(' Step 3: Route to final_agent (for final polish)') + + return arium + + +async def programmatic_reflection_example(): + """Show how to create reflection router programmatically""" + print('\n\nπŸ”§ Programmatic ReflectionRouter Creation') + print('=' * 50) + + from flo_ai.arium.llm_router import create_llm_router + + # Method 1: Using the convenience function + create_main_critic_reflection_router( + main_agent='writer', + critic_agent='reviewer', + final_agent='editor', + llm=OpenAI(model='gpt-4o-mini', api_key='dummy-key'), + ) + print('βœ… Method 1: Convenience function create_main_critic_reflection_router()') + + # Method 2: Using the factory function directly + create_llm_router( + 'reflection', + flow_pattern=['analyst', 'validator', 'analyst', 'presenter'], + allow_early_exit=True, + llm=OpenAI(model='gpt-4o-mini', api_key='dummy-key'), + ) + print('βœ… Method 2: Factory function create_llm_router(type="reflection")') + + # Method 3: Creating ReflectionRouter directly + from flo_ai.arium.llm_router import ReflectionRouter + + ReflectionRouter( + flow_pattern=['main', 'critic', 'main', 'final'], + allow_early_exit=False, + llm=OpenAI(model='gpt-4o-mini', api_key='dummy-key'), + ) + print('βœ… Method 3: Direct ReflectionRouter instantiation') + + print('\n🎯 All methods create the same Aβ†’Bβ†’Aβ†’C reflection pattern!') + + +def show_minimal_yaml(): + """Show the minimal YAML needed""" + print('\n\nπŸ“„ Minimal YAML Configuration') + print('=' * 35) + + yaml_example = """ +# Minimal ReflectionRouter YAML +arium: + agents: + - name: main_agent + job: "Main agent job" + model: {provider: openai, name: gpt-4o-mini} + - name: critic + job: "Critic job" + model: {provider: openai, name: gpt-4o-mini} + - name: final_agent + job: "Final agent job" + model: {provider: openai, name: gpt-4o-mini} + + routers: + - name: reflection_router + type: reflection + flow_pattern: [main_agent, critic, main_agent, final_agent] + + workflow: + start: main_agent + edges: + - from: main_agent + to: [critic, final_agent] + router: reflection_router + - from: critic + to: [main_agent, final_agent] + router: reflection_router + - from: final_agent + to: [end] + end: [final_agent] +""" + print(yaml_example) + + +async def main(): + """Run the simple demo""" + print('🌟 Simple ReflectionRouter Demo') + print('=' * 35) + print('Demonstrating A β†’ B β†’ A β†’ C reflection pattern with minimal setup\n') + + # Show different creation methods + await programmatic_reflection_example() + + # Show minimal YAML + show_minimal_yaml() + + # Create simple reflection workflow + arium = await simple_reflection_demo() + + result = await arium.run( + inputs=['Write a comprehensive guide on sustainable urban planning'] + ) + + print('\n\nπŸŽ‰ Demo completed!') + print('Key takeaways:') + print('βœ… ReflectionRouter makes Aβ†’Bβ†’Aβ†’C patterns trivial to implement') + print('βœ… Works with both YAML and programmatic configuration') + print('βœ… Automatically tracks progress and prevents infinite loops') + print('βœ… Provides intelligent routing based on execution context') + + print(result) + + +if __name__ == '__main__': + asyncio.run(main()) diff --git a/flo_ai/examples/simple_working_demo.py b/flo_ai/examples/simple_working_demo.py new file mode 100644 index 00000000..fd939d7f --- /dev/null +++ b/flo_ai/examples/simple_working_demo.py @@ -0,0 +1,335 @@ +""" +Simple Working PlanExecuteRouter Demo + +This version fixes the planner loop by using a different approach: +1. Use a standard workflow without plan storage complexity +2. Demonstrate the routing intelligence with manual plan simulation +3. Show how the router makes decisions based on context +""" + +import asyncio +import os +from flo_ai.models.agent import Agent +from flo_ai.llm import OpenAI +from flo_ai.arium.memory import MessageMemory +from flo_ai.arium.llm_router import create_plan_execute_router +from flo_ai.arium import AriumBuilder + + +async def simple_working_demo(): + """Simple working demo that avoids the planner loop""" + print('βœ… Simple Working PlanExecuteRouter Demo') + print('=' * 45) + + # 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 simple agents focused on their core tasks + planner = Agent( + name='planner', + system_prompt="""You are a project planner. When given a task, create a detailed plan with numbered steps. + +Format your response like this: + +PLAN FOR: [task name] + +EXECUTION STEPS: +1. [First step description] +2. [Second step description] +3. [Third step description] +4. [Final step description] + +NEXT ACTION: The developer should start with step 1. + +Keep it clear and actionable.""", + llm=llm, + ) + + developer = Agent( + name='developer', + system_prompt="""You are a software developer. When given a development task: + +1. Acknowledge what you're implementing +2. Provide implementation details +3. Mention any important considerations +4. State when you've completed the task + +Be specific and thorough in your implementation.""", + llm=llm, + ) + + tester = Agent( + name='tester', + system_prompt="""You are a QA tester. When given something to test: + +1. Acknowledge what you're testing +2. Create test scenarios +3. Identify potential issues +4. Provide test results and recommendations + +Be thorough in your testing approach.""", + llm=llm, + ) + + reviewer = Agent( + name='reviewer', + system_prompt="""You are a senior reviewer. When reviewing work: + +1. Assess overall quality and completeness +2. Check if requirements are met +3. Provide constructive feedback +4. Give final approval or suggest improvements + +Focus on delivering high-quality results.""", + llm=llm, + ) + + print('βœ… Created focused agents: planner, developer, tester, reviewer') + + # Create a simple router using the proper factory function + from typing import Literal + from flo_ai.arium.memory import BaseMemory + + def create_simple_router(): + """Create a simple router with proper type annotations""" + + def router_impl( + memory: BaseMemory, + ) -> Literal['developer', 'tester', 'reviewer']: + """Simple routing logic for demo purposes""" + messages = memory.get() + + # Check the conversation flow + if not messages: + return 'developer' # Start with developer after planner + + last_message = str(messages[-1]) + + # Basic routing logic based on content + if 'PLAN FOR:' in last_message and 'EXECUTION STEPS:' in last_message: + print('πŸ“‹ Plan detected - routing to developer') + return 'developer' + elif ( + 'implemented' in last_message.lower() + or 'development' in last_message.lower() + ): + print('πŸ’» Development complete - routing to tester') + return 'tester' + elif 'test' in last_message.lower() and 'complete' in last_message.lower(): + print('πŸ§ͺ Testing complete - routing to reviewer') + return 'reviewer' + elif len(messages) > 6: # Prevent too many iterations + print('🏁 Workflow complete - ending') + return 'reviewer' + else: + return 'developer' # Default fallback + + # Add required annotations + router_impl.__annotations__ = { + 'memory': BaseMemory, + 'return': Literal['developer', 'tester', 'reviewer'], + } + + return router_impl + + simple_router = create_simple_router() + + # Create memory + memory = MessageMemory() + + # Build workflow with simple routing + # Note: Each edge's 'to' nodes must include all possible router return values + arium = ( + AriumBuilder() + .with_memory(memory) + .add_agents([planner, developer, tester, reviewer]) + .start_with(planner) + .add_edge(planner, [developer, tester, reviewer], simple_router) + .add_edge( + developer, [developer, tester, reviewer], simple_router + ) # Include all possible destinations + .add_edge( + tester, [developer, tester, reviewer], simple_router + ) # Include all possible destinations + .end_with(reviewer) + .build() + ) + + print('βœ… Built simple workflow with basic routing') + + # Task to execute + task = 'Create a simple login endpoint with username and password validation' + + print(f'\nπŸ“‹ Task: {task}') + print('\nπŸ”„ Running simple workflow...') + print(' This demonstrates the routing concept without complex plan storage') + + try: + # Execute workflow + result = await arium.run([task]) + + print('\n' + '=' * 50) + print('πŸŽ‰ SIMPLE WORKFLOW COMPLETED!') + print('=' * 50) + + # Show the conversation flow + if memory.get(): + print('\nπŸ“„ Conversation Flow:') + print('-' * 30) + for i, msg in enumerate(memory.get(), 1): + role = msg.get('role', 'unknown') + content = str(msg.get('content', ''))[:200] + print(f'{i}. {role.upper()}: {content}...') + + # Show final result + if result: + final_result = result[-1] if isinstance(result, list) else result + print('\nπŸ“„ Final Output:') + print('-' * 30) + print(final_result) + + print('\nπŸ’‘ What this demonstrated:') + print(' β€’ Basic plan-execute workflow concept') + print(' β€’ Intelligent routing between phases') + print(' β€’ Planner β†’ Developer β†’ Tester β†’ Reviewer flow') + print(' β€’ How to avoid infinite loops with simple logic') + + except Exception as e: + print(f'\n❌ Error: {e}') + + +async def demonstrate_plan_execute_router(): + """Show the actual PlanExecuteRouter in a controlled way""" + print('\n\nπŸ“Š PlanExecuteRouter Demonstration') + print('=' * 40) + + api_key = os.getenv('OPENAI_API_KEY') + if not api_key: + print('❌ OPENAI_API_KEY not set') + return + + llm = OpenAI(model='gpt-4o', api_key=api_key) + + # Create the actual PlanExecuteRouter + create_plan_execute_router( + planner_agent='planner', + executor_agent='developer', + reviewer_agent='reviewer', + additional_agents={'tester': 'Tests implementations'}, + llm=llm, + ) + + print('βœ… Created actual PlanExecuteRouter') + + # Test router decision making with mock memory + memory = MessageMemory() + + # Simulate different scenarios + scenarios = [ + {'msg': 'Create a user API', 'context': 'Initial request'}, + {'msg': 'Plan created with 4 steps', 'context': 'After planning'}, + {'msg': 'Step 1 implemented successfully', 'context': 'After development'}, + {'msg': 'All tests passed', 'context': 'After testing'}, + ] + + print('\n🧠 Router Decision Making:') + for scenario in scenarios: + memory.add({'role': 'user', 'content': scenario['msg']}) + + try: + # This would make an actual LLM call to decide routing + print(f'\n Context: {scenario["context"]}') + print(f' Message: {scenario["msg"]}') + print(' Router would make intelligent decision based on this context') + # decision = plan_router(memory) # Uncomment to see actual routing + # print(f' Decision: Route to {decision}') + except Exception as e: + print(f' Note: {e}') + + print('\nπŸ’‘ The PlanExecuteRouter makes intelligent routing decisions by:') + print(' β€’ Analyzing conversation context') + print(' β€’ Detecting plan creation vs execution phases') + print(' β€’ Understanding step dependencies and progress') + print(' β€’ Routing to appropriate agents based on workflow state') + + +def show_solution_approaches(): + """Show different approaches to fix the planner loop""" + print('\n\nπŸ”§ Solutions to the Planner Loop Issue') + print('=' * 45) + + solutions = """ +The planner loop happened because the router couldn't detect plan completion. +Here are several solutions: + +1. πŸ“‹ SPECIALIZED PLANNER AGENT (Best Solution) + β€’ PlannerAgent that stores ExecutionPlan objects in PlanAwareMemory + β€’ Router detects when plans exist and switches to execution mode + β€’ See fixed_plan_execute_demo.py for implementation + +2. 🎯 CONTENT-BASED ROUTING (Simple Solution) + β€’ Router analyzes message content to detect phases + β€’ If message contains "PLAN:", route to developer + β€’ If message contains "implemented:", route to tester + β€’ See simple_working_demo.py example above + +3. πŸ”„ LIMITED ITERATIONS (Quick Fix) + β€’ Add max iteration limits to prevent infinite loops + β€’ Router switches phases after X iterations + β€’ Less intelligent but prevents loops + +4. πŸ“Š STATE MANAGEMENT (Advanced Solution) + β€’ Use custom memory with explicit state tracking + β€’ Store workflow phase (planning/executing/reviewing) + β€’ Router uses state to make decisions + +5. 🧠 BETTER PROMPTING (Prompt Engineering) + β€’ Improve router prompts to better detect completion + β€’ Add explicit "PLANNING COMPLETE" markers + β€’ Train router to recognize different phases + +Recommendation: Use approach #1 (Specialized Planner Agent) for production, +approach #2 (Content-Based Routing) for quick demos. +""" + + print(solutions) + + +async def main(): + """Main demo function""" + print('🎯 Working PlanExecuteRouter Demo (Loop Issue Fixed)') + print('This demo shows how to avoid the planner loop issue!\n') + + if not os.getenv('OPENAI_API_KEY'): + print('❌ To run this demo, set your OPENAI_API_KEY environment variable') + print(' export OPENAI_API_KEY=your_key_here') + return + + # Show solution approaches + show_solution_approaches() + + # Run simple working demo + await simple_working_demo() + + # Demonstrate the actual router + await demonstrate_plan_execute_router() + + print('\n\nπŸŽ‰ Demo Complete!') + print('=' * 20) + print('βœ… Demonstrated working plan-execute workflow') + print('βœ… Showed how to avoid planner loops') + print('βœ… Explained multiple solution approaches') + print('\nπŸš€ Try the fixed_plan_execute_demo.py for the complete solution!') + + +if __name__ == '__main__': + asyncio.run(main()) diff --git a/flo_ai/flo_ai/arium/base.py b/flo_ai/flo_ai/arium/base.py index 26943295..c3458354 100644 --- a/flo_ai/flo_ai/arium/base.py +++ b/flo_ai/flo_ai/arium/base.py @@ -103,16 +103,38 @@ def add_edge( if literal_values is None: raise ValueError('Router return type is not a Literal') - invalid_literals = [val for val in literal_values if val not in to_nodes] + # Check if router supports self-reference + supports_self_ref = getattr(router, 'supports_self_reference', False) + + # For self-referencing routers, we need to include the from_node in valid targets + valid_targets = to_nodes.copy() + if supports_self_ref: + valid_targets.append(from_node) + + invalid_literals = [ + val for val in literal_values if val not in valid_targets + ] if invalid_literals: raise ValueError( - f'Router return type includes literal values {invalid_literals} that are not in to_nodes {to_nodes}' + f'Router return type includes literal values {invalid_literals} that are not in valid targets {valid_targets}' ) - if set(literal_values) != set(to_nodes): - raise ValueError( - f'Router return type values {literal_values} do not match to_nodes {to_nodes}' - ) + # For self-referencing routers, allow router options to include from_node + if supports_self_ref: + # Router can return any of: to_nodes + from_node, but must include all to_nodes + missing_targets = [ + node for node in to_nodes if node not in literal_values + ] + if missing_targets: + raise ValueError( + f'Self-referencing router must include all to_nodes {to_nodes}, missing: {missing_targets}' + ) + else: + # Non-self-referencing routers must match exactly + if set(literal_values) != set(to_nodes): + raise ValueError( + f'Router return type values {literal_values} do not match to_nodes {to_nodes}' + ) self.edges[from_node] = Edge( router_fn=router diff --git a/flo_ai/flo_ai/arium/builder.py b/flo_ai/flo_ai/arium/builder.py index fc0821fa..5f0c3b60 100644 --- a/flo_ai/flo_ai/arium/builder.py +++ b/flo_ai/flo_ai/arium/builder.py @@ -230,7 +230,7 @@ def from_yaml( # LLM Router definitions (NEW) routers: - name: content_router - type: smart # smart, task_classifier, conversation_analysis + type: smart # smart, task_classifier, conversation_analysis, reflection, plan_execute routing_options: technical_writer: "Handle technical documentation tasks" creative_writer: "Handle creative writing tasks" @@ -242,6 +242,26 @@ def from_yaml( temperature: 0.3 fallback_strategy: first + # Reflection router for A -> B -> A -> C patterns + - name: main_critic_reflection + type: reflection + flow_pattern: [main_agent, critic, main_agent, final_agent] + settings: + allow_early_exit: false + + # Plan-Execute router for Cursor-style workflows + - name: plan_execute_router + type: plan_execute + agents: + planner: "Creates detailed execution plans" + developer: "Implements code and features" + tester: "Tests implementations" + reviewer: "Reviews final results" + settings: + planner_agent: planner + executor_agent: developer + reviewer_agent: reviewer + workflow: start: content_analyst edges: @@ -321,7 +341,7 @@ def from_yaml( # Method 4: External file reference elif 'yaml_file' in agent_config: - agent_builder = AgentBuilder.from_yaml( + agent_builder: AgentBuilder = AgentBuilder.from_yaml( yaml_file=agent_config['yaml_file'], base_llm=base_llm ) agent = agent_builder.build() @@ -418,9 +438,37 @@ def from_yaml( llm=router_llm, **settings, ) + + elif router_type == 'reflection': + flow_pattern = router_config.get('flow_pattern', []) + if not flow_pattern: + raise ValueError( + f'Reflection router {router_name} must specify flow_pattern' + ) + + router_fn = create_llm_router( + router_type='reflection', + flow_pattern=flow_pattern, + llm=router_llm, + **settings, + ) + + elif router_type == 'plan_execute': + agents = router_config.get('agents', {}) + if not agents: + raise ValueError( + f'Plan-Execute router {router_name} must specify agents' + ) + + router_fn = create_llm_router( + router_type='plan_execute', + agents=agents, + llm=router_llm, + **settings, + ) else: raise ValueError( - f'Unknown router type: {router_type}. Supported types: smart, task_classifier, conversation_analysis' + f'Unknown router type: {router_type}. Supported types: smart, task_classifier, conversation_analysis, reflection, plan_execute' ) yaml_routers[router_name] = router_fn @@ -600,10 +648,8 @@ def _create_agent_from_direct_config( f'Available: {list(available_tools.keys())}' ) - # Extract output schema if present - output_schema = agent_config.get('output_schema') - # Handle parser configuration if present + output_schema = None if 'parser' in agent_config: from flo_ai.formatter.yaml_format_parser import FloYamlParser @@ -612,16 +658,17 @@ def _create_agent_from_direct_config( 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, + agent = ( + AgentBuilder() + .with_name(name) + .with_prompt(job) + .with_llm(llm) + .with_tools(agent_tools) + .with_retries(max_retries) + .with_reasoning(reasoning_pattern) + .with_output_schema(output_schema) + .with_role(role) + .build() ) return agent diff --git a/flo_ai/flo_ai/arium/llm_router.py b/flo_ai/flo_ai/arium/llm_router.py index 900d2cf8..17f98d6c 100644 --- a/flo_ai/flo_ai/arium/llm_router.py +++ b/flo_ai/flo_ai/arium/llm_router.py @@ -6,9 +6,9 @@ """ from abc import ABC, abstractmethod -from typing import Dict, Optional, Callable, Any, Union, get_args +from typing import Dict, Optional, Callable, Any, Union, get_args, List from functools import wraps -from flo_ai.arium.memory import BaseMemory +from flo_ai.arium.memory import BaseMemory, ExecutionPlan, StepStatus from flo_ai.llm.base_llm import BaseLLM from flo_ai.llm import OpenAI from flo_ai.utils.logger import logger @@ -40,6 +40,9 @@ def __init__( self.temperature = temperature self.max_retries = max_retries self.fallback_strategy = fallback_strategy + self.supports_self_reference = ( + False # Most routers don't support self-reference by default + ) @abstractmethod def get_routing_options(self) -> Dict[str, str]: @@ -324,6 +327,381 @@ def get_routing_prompt( return prompt +class ReflectionRouter(BaseLLMRouter): + """ + A router designed for reflection patterns like A -> B -> A -> C. + Commonly used for main -> critic -> main -> final workflows where B is a reflection/critique step. + Uses execution context to determine flow state and make intelligent routing decisions. + """ + + def __init__( + self, + flow_pattern: List[str], + llm: Optional[BaseLLM] = None, + allow_early_exit: bool = False, + **kwargs, + ): + """ + Initialize the reflection router. + + Args: + flow_pattern: List of node names defining the reflection pattern (e.g., ["main", "critic", "main", "final"]) + llm: LLM instance for routing decisions + allow_early_exit: Whether to allow LLM to exit pattern early if appropriate + **kwargs: Additional arguments for BaseLLMRouter + """ + super().__init__(llm=llm, **kwargs) + self.flow_pattern = flow_pattern + self.allow_early_exit = allow_early_exit + self.supports_self_reference = ( + True # ReflectionRouter can return to the same node + ) + + def get_routing_options(self) -> Dict[str, str]: + """Get available routing options based on flow pattern""" + unique_nodes = list( + dict.fromkeys(self.flow_pattern) + ) # Preserve order, remove duplicates + + # Generate descriptions based on reflection pattern + options = {} + for node in unique_nodes: + # Count occurrences and positions in pattern + positions = [i for i, x in enumerate(self.flow_pattern) if x == node] + if len(positions) > 1: + options[node] = ( + f"Step {positions} in the reflection pattern: {' -> '.join(self.flow_pattern)}" + ) + else: + options[node] = ( + f"Step {positions[0] + 1} in the reflection pattern: {' -> '.join(self.flow_pattern)}" + ) + + return options + + def _get_next_step_in_pattern(self, execution_context: dict) -> Optional[str]: + """Determine the next step in the reflection pattern based on execution context""" + if not execution_context: + return self.flow_pattern[0] if self.flow_pattern else None + + visit_counts = execution_context.get('node_visit_count', {}) + current_node = execution_context.get('current_node', '') + + # Find the current position in the pattern + try: + # Find where we are in the reflection pattern + current_step = -1 + for i, node in enumerate(self.flow_pattern): + node_visits = visit_counts.get(node, 0) + + # For nodes that appear multiple times, we need to track which occurrence + if node == current_node: + # Count how many times this node should have been visited at this step + expected_visits = len( + [x for x in self.flow_pattern[: i + 1] if x == node] + ) + if node_visits >= expected_visits: + current_step = i + + # Determine next step + next_step_index = current_step + 1 + if next_step_index < len(self.flow_pattern): + return self.flow_pattern[next_step_index] + else: + # Pattern completed + return None + + except Exception: + # Fallback to first step + return self.flow_pattern[0] if self.flow_pattern else None + + def get_routing_prompt( + self, + memory: BaseMemory, + options: Dict[str, str], + execution_context: dict = None, + ) -> str: + conversation = memory.get() + + # Format conversation history + if isinstance(conversation, list): + conversation_text = '\n'.join( + [str(msg) for msg in conversation[-3:]] + ) # Last 3 messages for flow context + else: + conversation_text = str(conversation) + + # Determine suggested next step based on reflection pattern + suggested_next = self._get_next_step_in_pattern(execution_context) + + # Format options + options_text = '\n'.join( + [f'- {name}: {desc}' for name, desc in options.items()] + ) + + # Add execution context info + context_info = '' + if execution_context: + visit_counts = execution_context.get('node_visit_count', {}) + current_node = execution_context.get('current_node', 'unknown') + iteration = execution_context.get('iteration_count', 0) + + # Show reflection pattern progress + pattern_progress = [] + for i, node in enumerate(self.flow_pattern): + visits = visit_counts.get(node, 0) + expected_visits = len( + [x for x in self.flow_pattern[: i + 1] if x == node] + ) + status = 'βœ“' if visits >= expected_visits else 'β—‹' + pattern_progress.append(f'{status} {node}') + + context_info = f""" +πŸ“‹ REFLECTION PATTERN: {' β†’ '.join(self.flow_pattern)} +πŸ“ CURRENT PROGRESS: {' β†’ '.join(pattern_progress)} +🎯 SUGGESTED NEXT: {suggested_next or 'Pattern Complete'} +πŸ’‘ CURRENT NODE: {current_node} (iteration {iteration}) +""" + + # Create prompt based on whether early exit is allowed + if self.allow_early_exit: + prompt = f"""You are a reflection coordinator managing this workflow pattern: {' β†’ '.join(self.flow_pattern)} + +{context_info} +Available options: +{options_text} + +Recent conversation: +{conversation_text} + +Instructions: +1. Follow the reflection pattern: {' β†’ '.join(self.flow_pattern)} +2. The suggested next step is: {suggested_next or 'Pattern Complete'} +3. You may exit early if the reflection cycle is complete +4. Consider conversation context and reflection progress +5. Respond with ONLY the agent name (no explanations) + +Next agent:""" + else: + prompt = f"""You are a reflection coordinator managing this strict reflection pattern: {' β†’ '.join(self.flow_pattern)} + +{context_info} +Available options: +{options_text} + +Recent conversation: +{conversation_text} + +Instructions: +1. STRICTLY follow the reflection pattern: {' β†’ '.join(self.flow_pattern)} +2. The next step should be: {suggested_next or 'Pattern Complete'} +3. Do not deviate from the pattern unless absolutely necessary +4. Respond with ONLY the agent name (no explanations) + +Next agent:""" + + return prompt + + +class PlanExecuteRouter(BaseLLMRouter): + """ + A router that implements plan-and-execute patterns like Cursor. + Creates execution plans and routes through steps sequentially. + """ + + def __init__( + self, + agents: Dict[str, str], # agent_name -> description mapping + planner_agent: str = 'planner', + executor_agent: str = 'executor', + reviewer_agent: Optional[str] = None, + llm: Optional[BaseLLM] = None, + max_retries: int = 3, + **kwargs, + ): + """ + Initialize the plan-execute router. + + Args: + agents: Dict mapping agent names to their descriptions/capabilities + planner_agent: Name of the agent responsible for creating plans + executor_agent: Name of the agent responsible for executing steps + reviewer_agent: Optional name of the agent responsible for reviewing results + llm: LLM instance for routing decisions + max_retries: Maximum retries for step execution + **kwargs: Additional arguments for BaseLLMRouter + """ + super().__init__(llm=llm, **kwargs) + self.agents = agents + self.planner_agent = planner_agent + self.executor_agent = executor_agent + self.reviewer_agent = reviewer_agent + self.max_retries = max_retries + self.supports_self_reference = ( + True # Can route to same agent for iterative execution + ) + + def get_routing_options(self) -> Dict[str, str]: + """Get available routing options based on configured agents""" + return self.agents + + def get_routing_prompt( + self, + memory: BaseMemory, + options: Dict[str, str], + execution_context: dict = None, + ) -> str: + conversation = memory.get() + + # Format conversation history + if isinstance(conversation, list): + conversation_text = '\n'.join( + [str(msg) for msg in conversation[-3:]] + ) # Last 3 messages for context + else: + conversation_text = str(conversation) + + # Check if we have a plan in memory + current_plan = ( + memory.get_current_plan() if hasattr(memory, 'get_current_plan') else None + ) + + if current_plan is None: + # No plan exists - route to planner + return self._create_planning_prompt(conversation_text, options) + else: + # Plan exists - determine next action based on plan state + return self._create_execution_prompt( + current_plan, conversation_text, options, execution_context + ) + + def _create_planning_prompt( + self, conversation_text: str, options: Dict[str, str] + ) -> str: + """Create prompt for initial planning phase""" + options_text = '\n'.join( + [f'- {name}: {desc}' for name, desc in options.items()] + ) + + prompt = f"""You are coordinating a plan-and-execute workflow. No execution plan exists yet. + +Available agents: +{options_text} + +Recent conversation: +{conversation_text} + +TASK: Create an execution plan by routing to the {self.planner_agent}. + +Instructions: +1. Route to "{self.planner_agent}" to create a detailed execution plan +2. The planner will break down the task into sequential steps +3. Each step will specify which agent should execute it +4. Respond with ONLY the agent name: {self.planner_agent} + +Next agent:""" + + return prompt + + def _create_execution_prompt( + self, + plan: ExecutionPlan, + conversation_text: str, + options: Dict[str, str], + execution_context: dict = None, + ) -> str: + """Create prompt for execution phase based on current plan state""" + + # Get next steps that are ready to execute + next_steps = plan.get_next_steps() + + # Format plan progress + progress_lines = [] + for step in plan.steps: + status_icon = { + StepStatus.PENDING: 'β—‹', + StepStatus.IN_PROGRESS: '⏳', + StepStatus.COMPLETED: 'βœ…', + StepStatus.FAILED: '❌', + StepStatus.SKIPPED: '⏭️', + }.get(step.status, 'β—‹') + progress_lines.append( + f'{status_icon} {step.id}: {step.description} (β†’ {step.agent})' + ) + + progress_text = '\n'.join(progress_lines) + + # Determine what to do next + if plan.is_completed(): + # All steps completed + if self.reviewer_agent and self.reviewer_agent in options: + action = f'Route to {self.reviewer_agent} for final review' + suggested_agent = self.reviewer_agent + else: + action = 'Plan completed - route to any agent for final output' + suggested_agent = next(iter(options.keys())) # First available agent + elif plan.has_failed_steps(): + # Some steps failed - need recovery + failed_steps = [ + step for step in plan.steps if step.status == StepStatus.FAILED + ] + failed_step = failed_steps[0] # Focus on first failed step + action = f"Handle failed step '{failed_step.id}' - route to {failed_step.agent} for retry" + suggested_agent = failed_step.agent + elif next_steps: + # There are steps ready to execute + next_step = next_steps[0] # Execute first ready step + action = f"Execute step '{next_step.id}' - route to {next_step.agent}" + suggested_agent = next_step.agent + else: + # Waiting for dependencies + action = f'Waiting for dependencies - route to {self.executor_agent} for status check' + suggested_agent = self.executor_agent + + options_text = '\n'.join( + [f'- {name}: {desc}' for name, desc in options.items()] + ) + + # Add execution context info + context_info = '' + if execution_context: + current_node = execution_context.get('current_node', 'unknown') + iteration = execution_context.get('iteration_count', 0) + + context_info = f""" +πŸ’‘ EXECUTION CONTEXT: +Current node: {current_node} (iteration {iteration}) +""" + + prompt = f"""You are coordinating plan execution in a plan-and-execute workflow. + +πŸ“‹ EXECUTION PLAN: {plan.title} +{plan.description} + +πŸ“Š CURRENT PROGRESS: +{progress_text} + +🎯 NEXT ACTION: {action} +🎯 SUGGESTED AGENT: {suggested_agent} +{context_info} +Available agents: +{options_text} + +Recent conversation: +{conversation_text} + +Instructions: +1. Follow the execution plan step by step +2. Route to the suggested agent: {suggested_agent} +3. Each agent will execute their assigned step +4. Continue until all steps are completed +5. Respond with ONLY the agent name (no explanations) + +Next agent:""" + + return prompt + + class ConversationAnalysisRouter(BaseLLMRouter): """ A router that analyzes conversation flow and context to make routing decisions. @@ -424,7 +802,7 @@ def create_llm_router(router_type: str, **config) -> Callable[[BaseMemory], str] Factory function to create LLM-powered routers with different configurations. Args: - router_type: Type of router ("smart", "task_classifier", "conversation_analysis") + router_type: Type of router ("smart", "task_classifier", "conversation_analysis", "reflection", "plan_execute") **config: Configuration specific to the router type Returns: @@ -457,6 +835,26 @@ def create_llm_router(router_type: str, **config) -> Callable[[BaseMemory], str] } } ) + + # Reflection router for A -> B -> A -> C patterns + router = create_llm_router( + "reflection", + flow_pattern=["main_agent", "critic", "main_agent", "final_agent"], + allow_early_exit=False + ) + + # Plan-Execute router for Cursor-style workflows + router = create_llm_router( + "plan_execute", + agents={ + "planner": "Creates detailed execution plans", + "developer": "Implements code and features", + "tester": "Tests implementations", + "reviewer": "Reviews final results" + }, + planner_agent="planner", + executor_agent="developer" + ) """ if router_type == 'smart': @@ -481,6 +879,18 @@ def create_llm_router(router_type: str, **config) -> Callable[[BaseMemory], str] router_instance = ConversationAnalysisRouter(**config) + elif router_type == 'reflection': + if 'flow_pattern' not in config: + raise ValueError("ReflectionRouter requires 'flow_pattern' parameter") + + router_instance = ReflectionRouter(**config) + + elif router_type == 'plan_execute': + if 'agents' not in config: + raise ValueError("PlanExecuteRouter requires 'agents' parameter") + + router_instance = PlanExecuteRouter(**config) + else: raise ValueError(f'Unknown router type: {router_type}') @@ -506,6 +916,11 @@ async def router_function(memory: BaseMemory, execution_context: dict = None): # Add proper type annotations for validation router_function.__annotations__ = {'memory': BaseMemory, 'return': literal_type} + # Transfer router instance attributes to the function for validation + router_function.supports_self_reference = getattr( + router_instance, 'supports_self_reference', False + ) + return router_function @@ -658,3 +1073,93 @@ def create_research_analysis_router( }, llm=llm, ) + + +def create_main_critic_reflection_router( + main_agent: str = 'main_agent', + critic_agent: str = 'critic', + final_agent: str = 'final_agent', + allow_early_exit: bool = False, + llm: Optional[BaseLLM] = None, +) -> Callable[[BaseMemory], str]: + """ + Create a router for the A -> B -> A -> C reflection pattern (main -> critic -> main -> final). + + Args: + main_agent: Name of the main agent (appears twice in pattern) + critic_agent: Name of the critic agent for reflection + final_agent: Name of the final agent + allow_early_exit: Whether to allow LLM to exit reflection early if appropriate + llm: LLM instance for routing decisions + + Returns: + Router function for main/critic/final reflection workflows + """ + return create_llm_router( + 'reflection', + flow_pattern=[main_agent, critic_agent, main_agent, final_agent], + allow_early_exit=allow_early_exit, + llm=llm, + ) + + +def create_plan_execute_router( + planner_agent: str = 'planner', + executor_agent: str = 'executor', + reviewer_agent: Optional[str] = None, + additional_agents: Optional[Dict[str, str]] = None, + llm: Optional[BaseLLM] = None, +) -> Callable[[BaseMemory], str]: + """ + Create a router for plan-and-execute workflows like Cursor. + + Args: + planner_agent: Name of the agent responsible for creating plans + executor_agent: Name of the agent responsible for executing steps + reviewer_agent: Optional name of the agent responsible for reviewing results + additional_agents: Additional agents that can be used in execution steps + llm: LLM instance for routing decisions + + Returns: + Router function for plan-execute workflows + """ + agents = { + planner_agent: 'Creates detailed execution plans by breaking down tasks into sequential steps', + executor_agent: 'Executes individual steps from the execution plan', + } + + if reviewer_agent: + agents[reviewer_agent] = 'Reviews and validates completed work' + + if additional_agents: + agents.update(additional_agents) + + return create_llm_router( + 'plan_execute', + agents=agents, + planner_agent=planner_agent, + executor_agent=executor_agent, + reviewer_agent=reviewer_agent, + llm=llm, + ) + + +# Backward compatibility alias +def create_main_critic_flow_router( + main_agent: str = 'main_agent', + critic_agent: str = 'critic', + final_agent: str = 'final_agent', + allow_early_exit: bool = False, + llm: Optional[BaseLLM] = None, +) -> Callable[[BaseMemory], str]: + """ + DEPRECATED: Use create_main_critic_reflection_router instead. + Create a router for the A -> B -> A -> C reflection pattern (main -> critic -> main -> final). + """ + return create_main_critic_reflection_router( + main_agent=main_agent, + critic_agent=critic_agent, + final_agent=final_agent, + allow_early_exit=allow_early_exit, + llm=llm, + ) diff --git a/flo_ai/flo_ai/arium/memory.py b/flo_ai/flo_ai/arium/memory.py index 88e9f113..783ca4ce 100644 --- a/flo_ai/flo_ai/arium/memory.py +++ b/flo_ai/flo_ai/arium/memory.py @@ -1,18 +1,117 @@ from abc import ABC, abstractmethod -from typing import TypeVar, Generic, List, Dict +from typing import TypeVar, Generic, List, Dict, Optional, Any +from dataclasses import dataclass, field +from enum import Enum # Define the generic type variable T = TypeVar('T') +class StepStatus(Enum): + """Status of a plan step""" + + PENDING = 'pending' + IN_PROGRESS = 'in_progress' + COMPLETED = 'completed' + FAILED = 'failed' + SKIPPED = 'skipped' + + +@dataclass +class PlanStep: + """Represents a single step in an execution plan""" + + id: str + description: str + agent: str # Which agent should execute this step + dependencies: List[str] = field(default_factory=list) # Step IDs this depends on + status: StepStatus = StepStatus.PENDING + result: Optional[str] = None + error: Optional[str] = None + metadata: Dict[str, Any] = field(default_factory=dict) + + +@dataclass +class ExecutionPlan: + """Represents a complete execution plan""" + + id: str + title: str + description: str + steps: List[PlanStep] = field(default_factory=list) + created_by: str = 'planner' + status: str = 'active' # active, completed, failed, paused + metadata: Dict[str, Any] = field(default_factory=dict) + + def get_next_steps(self) -> List[PlanStep]: + """Get steps that are ready to execute (pending with no pending dependencies)""" + next_steps = [] + for step in self.steps: + if step.status == StepStatus.PENDING: + # Check if all dependencies are completed + if all( + self.get_step(dep_id).status == StepStatus.COMPLETED + for dep_id in step.dependencies + ): + next_steps.append(step) + return next_steps + + def get_step(self, step_id: str) -> Optional[PlanStep]: + """Get a step by ID""" + for step in self.steps: + if step.id == step_id: + return step + return None + + def mark_step_completed(self, step_id: str, result: str = None): + """Mark a step as completed""" + step = self.get_step(step_id) + if step: + step.status = StepStatus.COMPLETED + step.result = result + + def mark_step_failed(self, step_id: str, error: str = None): + """Mark a step as failed""" + step = self.get_step(step_id) + if step: + step.status = StepStatus.FAILED + step.error = error + + def is_completed(self) -> bool: + """Check if all steps are completed""" + return all(step.status == StepStatus.COMPLETED for step in self.steps) + + def has_failed_steps(self) -> bool: + """Check if any steps have failed""" + return any(step.status == StepStatus.FAILED for step in self.steps) + + class BaseMemory(ABC, Generic[T]): @abstractmethod def add(self, m: T): pass + @abstractmethod def get(self) -> List[T]: pass + # Plan management methods (optional - only implemented by memory classes that support plans) + def add_plan(self, plan: ExecutionPlan): + """Add an execution plan (override in subclasses that support plans)""" + raise NotImplementedError('This memory type does not support plans') + + def get_current_plan(self) -> Optional[ExecutionPlan]: + """Get the current active plan (override in subclasses that support plans)""" + return None + + def update_plan(self, plan: ExecutionPlan): + """Update an existing plan (override in subclasses that support plans)""" + raise NotImplementedError('This memory type does not support plans') + + def get_plan(self, plan_id: str) -> Optional[ExecutionPlan]: + """Get a plan by ID (override in subclasses that support plans)""" + return None + class MessageMemory(BaseMemory[Dict[str, str]]): def __init__(self): @@ -23,3 +122,61 @@ def add(self, message: Dict[str, str]): def get(self) -> List[Dict[str, str]]: return self.messages + + +class PlanAwareMemory(BaseMemory[Dict[str, str]]): + """Enhanced memory that supports both messages and execution plans""" + + def __init__(self): + self.messages = [] + self.plans: Dict[str, ExecutionPlan] = {} + self.current_plan_id: Optional[str] = None + + def add(self, message: Dict[str, str]): + """Add a message to memory""" + self.messages.append(message) + + def get(self) -> List[Dict[str, str]]: + """Get all messages""" + return self.messages + + # Plan management methods + def add_plan(self, plan: ExecutionPlan): + """Add an execution plan and set it as current""" + self.plans[plan.id] = plan + self.current_plan_id = plan.id + + def get_current_plan(self) -> Optional[ExecutionPlan]: + """Get the current active plan""" + if self.current_plan_id and self.current_plan_id in self.plans: + return self.plans[self.current_plan_id] + return None + + def update_plan(self, plan: ExecutionPlan): + """Update an existing plan""" + if plan.id in self.plans: + self.plans[plan.id] = plan + else: + raise ValueError(f'Plan {plan.id} not found in memory') + + def get_plan(self, plan_id: str) -> Optional[ExecutionPlan]: + """Get a plan by ID""" + return self.plans.get(plan_id) + + def set_current_plan(self, plan_id: str): + """Set the current active plan""" + if plan_id in self.plans: + self.current_plan_id = plan_id + else: + raise ValueError(f'Plan {plan_id} not found in memory') + + def get_all_plans(self) -> List[ExecutionPlan]: + """Get all plans""" + return list(self.plans.values()) + + def remove_plan(self, plan_id: str): + """Remove a plan from memory""" + if plan_id in self.plans: + del self.plans[plan_id] + if self.current_plan_id == plan_id: + self.current_plan_id = None diff --git a/flo_ai/flo_ai/llm/anthropic_llm.py b/flo_ai/flo_ai/llm/anthropic_llm.py index 74685397..05c45cf1 100644 --- a/flo_ai/flo_ai/llm/anthropic_llm.py +++ b/flo_ai/flo_ai/llm/anthropic_llm.py @@ -49,7 +49,6 @@ async def generate( 'model': self.model, 'messages': conversation, 'temperature': self.temperature, - 'max_tokens': 8192, **self.kwargs, } diff --git a/flo_ai/flo_ai/llm/gemini_llm.py b/flo_ai/flo_ai/llm/gemini_llm.py index 89151ff0..6b3f2550 100644 --- a/flo_ai/flo_ai/llm/gemini_llm.py +++ b/flo_ai/flo_ai/llm/gemini_llm.py @@ -1,9 +1,8 @@ from typing import Dict, Any, List, Optional from google import genai -import json +from google.genai import types from .base_llm import BaseLLM, ImageMessage from flo_ai.tool.base_tool import Tool -from flo_ai.utils.logger import logger class Gemini(BaseLLM): @@ -27,9 +26,9 @@ async def generate( output_schema: Optional[Dict[str, Any]] = None, ) -> Dict[str, Any]: # Convert messages to Gemini format - # Gemini uses a simple content string format contents = [] system_prompt = '' + for msg in messages: role = msg['role'] message_content = msg['content'] @@ -39,80 +38,57 @@ async def generate( else: contents.append(message_content) - # Add output schema instruction if provided - if output_schema: - contents += f'\nPlease provide your response in JSON format according to this schema:\n{json.dumps(output_schema, indent=2)}\n' - - # Add function information if provided - if functions: - contents += f'\nAvailable functions:\n{json.dumps(functions, indent=2)}\n' - try: # Prepare generation config - generation_config = genai.types.GenerateContentConfig( + generation_config = types.GenerateContentConfig( temperature=self.temperature, system_instruction=system_prompt, **self.kwargs, ) + # Add tools if functions are provided + if functions: + tools = types.Tool(function_declarations=functions) + generation_config.tools = [tools] + + # Add structured output configuration if output_schema is provided + if output_schema: + generation_config.response_mime_type = 'application/json' + generation_config.response_schema = output_schema + # Make the API call response = self.client.models.generate_content( model=self.model, contents=contents, - config=generation_config if generation_config else None, + config=generation_config, ) - # Check if response contains function call information - # For now, we'll assume text response and parse for function calls if needed - response_text = ( - response.text if hasattr(response, 'text') else str(response) - ) - - # Try to detect function calls in the response - # This is a simple implementation - in practice, you might need more sophisticated parsing - if functions and self._is_function_call_response(response_text): - function_call = self._parse_function_call(response_text) - if function_call: + # Check for function call in the response + if ( + functions + and response.candidates + and response.candidates[0].content.parts + ): + part = response.candidates[0].content.parts[0] + if hasattr(part, 'function_call') and part.function_call: + function_call = part.function_call return { - 'content': response_text, - 'function_call': function_call, + 'content': response.text, + 'function_call': { + 'name': function_call.name, + 'arguments': function_call.args, + }, } + # Return regular text response + response_text = ( + response.text if hasattr(response, 'text') else str(response) + ) return {'content': response_text} except Exception as e: raise Exception(f'Error in Gemini API call: {str(e)}') - def _is_function_call_response(self, response_text: str) -> bool: - """Check if the response contains a function call""" - # Simple heuristic - look for function call patterns - return ( - 'function_call' in response_text.lower() - or '(' in response_text - and ')' in response_text - ) - - def _parse_function_call(self, response_text: str) -> Optional[Dict[str, Any]]: - """Parse function call from response text""" - # This is a simplified parser - in practice, you'd want more robust parsing - try: - # Look for JSON-like function call structure - if '{' in response_text and '}' in response_text: - # Extract JSON-like content - start = response_text.find('{') - end = response_text.rfind('}') + 1 - json_str = response_text[start:end] - parsed = json.loads(json_str) - - if 'name' in parsed and 'arguments' in parsed: - return { - 'name': parsed['name'], - 'arguments': json.dumps(parsed['arguments']), - } - except Exception as e: - logger.error(f'Error parsing function call: {str(e)}') - return None - def get_message_content(self, response: Any) -> str: """Extract message content from response""" if isinstance(response, dict): @@ -120,7 +96,7 @@ def get_message_content(self, response: Any) -> str: return str(response) def format_tool_for_llm(self, tool: 'Tool') -> Dict[str, Any]: - """Format a single tool for Gemini's API""" + """Format a single tool for Gemini's function declarations""" return { 'name': tool.name, 'description': tool.description, @@ -142,7 +118,7 @@ def format_tool_for_llm(self, tool: 'Tool') -> Dict[str, Any]: } def format_tools_for_llm(self, tools: List['Tool']) -> List[Dict[str, Any]]: - """Format tools for Gemini's API""" + """Format tools for Gemini's function declarations""" return [self.format_tool_for_llm(tool) for tool in tools] def format_image_in_message(self, image: ImageMessage) -> str: @@ -150,12 +126,12 @@ def format_image_in_message(self, image: ImageMessage) -> str: if image.image_file_path: with open(image.image_file_path, 'rb') as image_file: image_bytes = image_file.read() - return genai.types.Part.from_bytes( + return types.Part.from_bytes( data=image_bytes, mime_type=image.mime_type, ) elif image.image_bytes: - return genai.types.Part.from_bytes( + return types.Part.from_bytes( data=image.image_bytes, mime_type=image.mime_type, ) diff --git a/flo_ai/flo_ai/llm/openai_llm.py b/flo_ai/flo_ai/llm/openai_llm.py index c1dfc24a..6c823bb3 100644 --- a/flo_ai/flo_ai/llm/openai_llm.py +++ b/flo_ai/flo_ai/llm/openai_llm.py @@ -7,7 +7,7 @@ class OpenAI(BaseLLM): def __init__( self, - model='gpt-40-mini', + model='gpt-4o-mini', api_key: str = None, temperature: float = 0.7, base_url: str = None, @@ -53,6 +53,7 @@ async def generate( openai_kwargs = { 'model': self.model, 'messages': messages, + 'temperature': self.temperature, **kwargs, **self.kwargs, } diff --git a/flo_ai/flo_ai/llm/openai_vllm.py b/flo_ai/flo_ai/llm/openai_vllm.py index d0f46b5f..eb4a7c0b 100644 --- a/flo_ai/flo_ai/llm/openai_vllm.py +++ b/flo_ai/flo_ai/llm/openai_vllm.py @@ -18,6 +18,8 @@ def __init__( base_url=base_url, **kwargs, ) + # Store base_url attribute + self.base_url = base_url # overriden async def generate( @@ -37,14 +39,14 @@ async def generate( if messages and messages[0]['role'] == 'system': messages[0]['content'] = ( messages[0]['content'] - + '\n\nPlease provide your response in JSON format according to the specified schema.' + + f'\n\nPlease provide your response in JSON format according to the specified schema. \n\n {output_schema}' ) else: messages.insert( 0, { 'role': 'system', - 'content': 'Please provide your response in JSON format according to the specified schema.', + 'content': f'Please provide your response in JSON format according to the specified schema.\n \n {output_schema}', }, ) @@ -52,6 +54,7 @@ async def generate( vllm_openai_kwargs = { 'model': self.model, 'messages': messages, + 'temperature': self.temperature, **kwargs, **self.kwargs, } diff --git a/flo_ai/flo_ai/llm/vertexai_llm.py b/flo_ai/flo_ai/llm/vertexai_llm.py index 2f1c883c..f4d92527 100644 --- a/flo_ai/flo_ai/llm/vertexai_llm.py +++ b/flo_ai/flo_ai/llm/vertexai_llm.py @@ -17,5 +17,10 @@ def __init__( ): # Initialize only the BaseLLM part to avoid Gemini's client creation BaseLLM.__init__(self, model, api_key, temperature, **kwargs) + + # Store project and location attributes + self.project = project + self.location = location + # Create VertexAI-specific client self.client = genai.Client(project=project, location=location, vertexai=True) 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/flo_ai/pyproject.toml b/flo_ai/pyproject.toml index c7e758a8..1657b35c 100644 --- a/flo_ai/pyproject.toml +++ b/flo_ai/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "flo_ai" -version = "1.0.0-dev5" +version = "1.0.2-dev1" description = "A easy way to create structured AI agents" authors = ["rootflo <*@rootflo.ai>"] license = "MIT" diff --git a/flo_ai/setup.py b/flo_ai/setup.py index c2d924f9..076f8c8e 100644 --- a/flo_ai/setup.py +++ b/flo_ai/setup.py @@ -5,7 +5,7 @@ setuptools.setup( name='flo-ai', - version='1.0.0-dev5', + version='1.0.2-dev1', author='Rootflo', description='Create composable AI agents', long_description=long_description, diff --git a/flo_ai/tests/run_llm_tests.py b/flo_ai/tests/run_llm_tests.py new file mode 100644 index 00000000..5e5a5c5d --- /dev/null +++ b/flo_ai/tests/run_llm_tests.py @@ -0,0 +1,164 @@ +#!/usr/bin/env python3 +""" +Test runner script for all LLM tests in the Flo AI framework. +This script runs comprehensive tests for all LLM implementations. +""" + +import sys +import os +import subprocess +import time + +# Add the flo_ai directory to the path +sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..')) + + +def run_tests_for_llm(test_file, llm_name): + """Run tests for a specific LLM implementation.""" + print(f"\n{'='*60}") + print(f'πŸ§ͺ Running tests for {llm_name}') + print(f"{'='*60}") + + start_time = time.time() + + try: + # Run pytest for the specific test file + result = subprocess.run( + [ + sys.executable, + '-m', + 'pytest', + test_file, + '-v', # Verbose output + '--tb=short', # Short traceback format + '--color=yes', # Colored output + ], + capture_output=True, + text=True, + cwd=os.path.dirname(__file__), + ) + + end_time = time.time() + duration = end_time - start_time + + if result.returncode == 0: + print(f'βœ… {llm_name} tests PASSED in {duration:.2f}s') + print('πŸ“Š Output:') + print(result.stdout) + else: + print(f'❌ {llm_name} tests FAILED in {duration:.2f}s') + print('πŸ“Š Output:') + print(result.stdout) + print('🚨 Errors:') + print(result.stderr) + + return result.returncode == 0, duration + + except Exception as e: + print(f'πŸ’₯ Error running {llm_name} tests: {e}') + return False, 0 + + +def run_all_llm_tests(): + """Run all LLM tests and provide a summary.""" + print('πŸš€ Starting comprehensive LLM test suite for Flo AI') + print(f'πŸ“ Test directory: {os.path.dirname(__file__)}') + print(f'🐍 Python executable: {sys.executable}') + + # Define all LLM test files + test_files = [ + ('test_base_llm.py', 'BaseLLM & ImageMessage'), + ('test_openai_llm.py', 'OpenAI LLM'), + ('test_anthropic_llm.py', 'Anthropic Claude LLM'), + ('test_gemini_llm.py', 'Google Gemini LLM'), + ('test_ollama_llm.py', 'Ollama LLM'), + ('test_vertexai_llm.py', 'Google VertexAI LLM'), + ('test_openai_vllm.py', 'OpenAI VLLM'), + ] + + results = [] + total_tests = 0 + passed_tests = 0 + total_duration = 0 + + for test_file, llm_name in test_files: + test_path = os.path.join(os.path.dirname(__file__), test_file) + + if not os.path.exists(test_path): + print(f'⚠️ Test file not found: {test_file}') + continue + + success, duration = run_tests_for_llm(test_path, llm_name) + results.append((llm_name, success, duration)) + + if success: + passed_tests += 1 + total_tests += 1 + total_duration += duration + + # Print summary + print(f"\n{'='*60}") + print('πŸ“‹ TEST SUMMARY') + print(f"{'='*60}") + + for llm_name, success, duration in results: + status = 'βœ… PASSED' if success else '❌ FAILED' + print(f'{llm_name:<25} {status:<10} {duration:>8.2f}s') + + print('\nπŸ“Š Overall Results:') + print(f' Total LLM implementations: {total_tests}') + print(f' Passed: {passed_tests}') + print(f' Failed: {total_tests - passed_tests}') + print(f' Success rate: {(passed_tests/total_tests)*100:.1f}%') + print(f' Total test time: {total_duration:.2f}s') + + if passed_tests == total_tests: + print('\nπŸŽ‰ All LLM tests passed successfully!') + return 0 + else: + print('\n⚠️ Some LLM tests failed. Please check the output above.') + return 1 + + +def run_specific_llm_test(llm_name): + """Run tests for a specific LLM implementation.""" + test_mapping = { + 'base': 'test_base_llm.py', + 'openai': 'test_openai_llm.py', + 'anthropic': 'test_anthropic_llm.py', + 'gemini': 'test_gemini_llm.py', + 'ollama': 'test_ollama_llm.py', + 'vertexai': 'test_vertexai_llm.py', + 'vllm': 'test_openai_vllm.py', + } + + if llm_name.lower() not in test_mapping: + print(f'❌ Unknown LLM: {llm_name}') + print(f"Available options: {', '.join(test_mapping.keys())}") + return 1 + + test_file = test_mapping[llm_name.lower()] + test_path = os.path.join(os.path.dirname(__file__), test_file) + + if not os.path.exists(test_path): + print(f'❌ Test file not found: {test_file}') + return 1 + + success, duration = run_tests_for_llm(test_path, test_mapping[llm_name.lower()]) + return 0 if success else 1 + + +def main(): + """Main entry point.""" + if len(sys.argv) > 1: + # Run specific LLM test + llm_name = sys.argv[1] + return run_specific_llm_test(llm_name) + else: + # Run all LLM tests + return run_all_llm_tests() + + +if __name__ == '__main__': + exit_code = main() + sys.exit(exit_code) diff --git a/flo_ai/tests/test_anthropic_llm.py b/flo_ai/tests/test_anthropic_llm.py new file mode 100644 index 00000000..975b29f9 --- /dev/null +++ b/flo_ai/tests/test_anthropic_llm.py @@ -0,0 +1,407 @@ +#!/usr/bin/env python3 +""" +Pytest tests for the Anthropic LLM implementation. +""" + +import sys +import os +import pytest +from unittest.mock import Mock, AsyncMock, patch + +# Add the flo_ai directory to the path +sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..')) + +from flo_ai.llm.anthropic_llm import Anthropic +from flo_ai.llm.base_llm import ImageMessage +from flo_ai.tool.base_tool import Tool + + +class TestAnthropic: + """Test class for Anthropic LLM implementation.""" + + def test_anthropic_initialization(self): + """Test Anthropic LLM initialization with different parameters.""" + # Test with minimal parameters + llm = Anthropic() + assert llm.model == 'claude-3-5-sonnet-20240620' + assert llm.api_key is None + assert llm.temperature == 0.7 + assert llm.kwargs == {} + + # Test with custom parameters + llm = Anthropic( + model='claude-3-opus-20240229', + api_key='test-key-123', + temperature=0.5, + max_tokens=1000, + ) + assert llm.model == 'claude-3-opus-20240229' + assert llm.api_key == 'test-key-123' + assert llm.temperature == 0.5 + assert llm.kwargs == {'max_tokens': 1000} + + # Test with base_url + llm = Anthropic(base_url='https://custom.anthropic.com') + assert llm.client.base_url == 'https://custom.anthropic.com' + + def test_anthropic_temperature_handling(self): + """Test temperature parameter handling.""" + # Test default temperature + llm = Anthropic() + assert llm.temperature == 0.7 + + # Test custom temperature + llm = Anthropic(temperature=0.0) + assert llm.temperature == 0.0 + + # Test high temperature + llm = Anthropic(temperature=1.0) + assert llm.temperature == 1.0 + + # Test temperature in kwargs + llm = Anthropic(temperature=0.3, custom_temp=0.8) + assert llm.temperature == 0.3 + assert llm.kwargs['custom_temp'] == 0.8 + + @patch('flo_ai.llm.anthropic_llm.AsyncAnthropic') + def test_anthropic_client_creation(self, mock_async_anthropic): + """Test that AsyncAnthropic client is created correctly.""" + mock_client = Mock() + mock_async_anthropic.return_value = mock_client + + llm = Anthropic(api_key='test-key', base_url='https://custom.com') + + mock_async_anthropic.assert_called_once_with( + api_key='test-key', base_url='https://custom.com' + ) + assert llm.client == mock_client + + @pytest.mark.asyncio + async def test_anthropic_generate_basic(self): + """Test basic generate method without functions or output schema.""" + llm = Anthropic(model='claude-3-5-sonnet-20240620') + + # Mock the client response + mock_content = Mock() + mock_content.text = 'Hello, world!' + mock_content.type = 'text' + + mock_response = Mock() + mock_response.content = [mock_content] + + llm.client = Mock() + llm.client.messages.create = AsyncMock(return_value=mock_response) + + messages = [{'role': 'user', 'content': 'Hello'}] + result = await llm.generate(messages) + + # Verify the API call + llm.client.messages.create.assert_called_once() + call_args = llm.client.messages.create.call_args + + assert call_args[1]['model'] == 'claude-3-5-sonnet-20240620' + assert call_args[1]['messages'] == messages + assert call_args[1]['temperature'] == 0.7 + + # Verify the result + assert result == {'content': 'Hello, world!'} + + @pytest.mark.asyncio + async def test_anthropic_generate_with_system_message(self): + """Test generate method with system message.""" + llm = Anthropic(model='claude-3-5-sonnet-20240620') + + # Mock the client response + mock_content = Mock() + mock_content.text = "I'm a helpful assistant" + mock_content.type = 'text' + + mock_response = Mock() + mock_response.content = [mock_content] + + llm.client = Mock() + llm.client.messages.create = AsyncMock(return_value=mock_response) + + messages = [ + {'role': 'system', 'content': 'You are a helpful assistant'}, + {'role': 'user', 'content': 'Hello'}, + ] + + await llm.generate(messages) + + # Verify system message was passed correctly + call_args = llm.client.messages.create.call_args[1] + assert call_args['system'] == 'You are a helpful assistant' + + # Verify conversation messages don't include system + conversation_messages = call_args['messages'] + assert len(conversation_messages) == 1 + assert conversation_messages[0]['role'] == 'user' + + @pytest.mark.asyncio + async def test_anthropic_generate_with_output_schema(self): + """Test generate method with output schema.""" + llm = Anthropic(model='claude-3-5-sonnet-20240620') + + output_schema = { + 'type': 'object', + 'properties': {'message': {'type': 'string'}, 'count': {'type': 'integer'}}, + } + + # Mock the client response + mock_content = Mock() + mock_content.text = '{"message": "test", "count": 42}' + mock_content.type = 'text' + + mock_response = Mock() + mock_response.content = [mock_content] + + llm.client = Mock() + llm.client.messages.create = AsyncMock(return_value=mock_response) + + messages = [{'role': 'user', 'content': 'Generate JSON'}] + await llm.generate(messages, output_schema=output_schema) + + # Verify system message includes output schema + call_args = llm.client.messages.create.call_args[1] + system_message = call_args['system'] + assert 'Provide output in the following JSON schema' in system_message + assert 'message' in system_message + assert 'count' in system_message + + @pytest.mark.asyncio + async def test_anthropic_generate_with_functions(self): + """Test generate method with functions (tools).""" + llm = Anthropic(model='claude-3-5-sonnet-20240620') + + functions = [ + { + 'type': 'custom', + 'name': 'test_function', + 'description': 'A test function', + 'input_schema': {'type': 'object'}, + } + ] + + # Mock the client response + mock_content = Mock() + mock_content.text = "I'll use the function" + mock_content.type = 'text' + + mock_response = Mock() + mock_response.content = [mock_content] + + llm.client = Mock() + llm.client.messages.create = AsyncMock(return_value=mock_response) + + messages = [{'role': 'user', 'content': 'Use the function'}] + await llm.generate(messages, functions=functions) + + # Verify functions were passed correctly + call_args = llm.client.messages.create.call_args[1] + assert call_args['tools'] == functions + + @pytest.mark.asyncio + async def test_anthropic_generate_with_tool_use(self): + """Test generate method when Claude uses a tool.""" + llm = Anthropic(model='claude-3-5-sonnet-20240620') + + # Mock the client response with tool use + mock_tool_content = Mock() + mock_tool_content.type = 'tool_use' + mock_tool_content.name = 'test_tool' + mock_tool_content.input = {'param': 'value'} + + mock_text_content = Mock() + mock_text_content.text = 'I used the tool' + mock_text_content.type = 'text' + + mock_response = Mock() + mock_response.content = [mock_text_content, mock_tool_content] + + llm.client = Mock() + llm.client.messages.create = AsyncMock(return_value=mock_response) + + messages = [{'role': 'user', 'content': 'Use the tool'}] + result = await llm.generate(messages) + + # Verify tool use was detected and formatted correctly + assert 'function_call' in result + assert result['function_call']['name'] == 'test_tool' + assert result['function_call']['arguments'] == '{"param": "value"}' + + @pytest.mark.asyncio + async def test_anthropic_generate_with_max_tokens(self): + """Test generate method with max_tokens parameter.""" + llm = Anthropic(model='claude-3-5-sonnet-20240620', max_tokens=1000) + + # Mock the client response + mock_content = Mock() + mock_content.text = 'Response with max tokens' + mock_content.type = 'text' + + mock_response = Mock() + mock_response.content = [mock_content] + + llm.client = Mock() + llm.client.messages.create = AsyncMock(return_value=mock_response) + + messages = [{'role': 'user', 'content': 'Hello'}] + await llm.generate(messages) + + # Verify max_tokens was passed + call_args = llm.client.messages.create.call_args[1] + assert call_args['max_tokens'] == 1000 + + @pytest.mark.asyncio + async def test_anthropic_generate_with_kwargs(self): + """Test generate method with additional kwargs.""" + llm = Anthropic(model='claude-3-5-sonnet-20240620', top_p=0.9) + + # Mock the client response + mock_content = Mock() + mock_content.text = 'Response with kwargs' + mock_content.type = 'text' + + mock_response = Mock() + mock_response.content = [mock_content] + + llm.client = Mock() + llm.client.messages.create = AsyncMock(return_value=mock_response) + + messages = [{'role': 'user', 'content': 'Hello'}] + await llm.generate(messages) + + # Verify kwargs were passed through + call_args = llm.client.messages.create.call_args[1] + assert call_args['top_p'] == 0.9 + + def test_anthropic_get_message_content(self): + """Test get_message_content method.""" + llm = Anthropic() + + # Test with dict response + response = {'content': 'Hello, world!'} + result = llm.get_message_content(response) + assert result == 'Hello, world!' + + # Test with string response + result = llm.get_message_content('Direct string') + assert result == 'Direct string' + + # Test with empty content + response = {'content': ''} + result = llm.get_message_content(response) + assert result == '' + + def test_anthropic_format_tool_for_llm(self): + """Test format_tool_for_llm method.""" + llm = Anthropic() + + # Create a mock tool + tool = Tool( + name='test_tool', + description='A test tool', + function=lambda x: x, + parameters={ + 'param1': {'type': 'string', 'description': 'First parameter'}, + 'param2': {'type': 'integer', 'description': 'Second parameter'}, + }, + ) + + formatted = llm.format_tool_for_llm(tool) + + assert formatted['type'] == 'custom' + assert formatted['name'] == 'test_tool' + assert formatted['description'] == 'A test tool' + assert formatted['input_schema']['type'] == 'object' + assert 'param1' in formatted['input_schema']['properties'] + assert 'param2' in formatted['input_schema']['properties'] + assert formatted['input_schema']['required'] == ['param1', 'param2'] + + def test_anthropic_format_tools_for_llm(self): + """Test format_tools_for_llm method.""" + llm = Anthropic() + + # Create mock tools + tool1 = Tool( + name='tool1', + description='First tool', + function=lambda x: x, + parameters={'param': {'type': 'string', 'description': 'Parameter'}}, + ) + + tool2 = Tool( + name='tool2', + description='Second tool', + function=lambda x: x, + parameters={'param': {'type': 'integer', 'description': 'Parameter'}}, + ) + + formatted = llm.format_tools_for_llm([tool1, tool2]) + + assert len(formatted) == 2 + assert formatted[0]['name'] == 'tool1' + assert formatted[1]['name'] == 'tool2' + assert all(tool['type'] == 'custom' for tool in formatted) + + def test_anthropic_format_image_in_message(self): + """Test format_image_in_message method.""" + llm = Anthropic() + + # This method is not implemented yet + image = ImageMessage(image_url='https://example.com/image.jpg') + + with pytest.raises(NotImplementedError): + llm.format_image_in_message(image) + + @pytest.mark.asyncio + async def test_anthropic_generate_error_handling(self): + """Test error handling in generate method.""" + llm = Anthropic(model='claude-3-5-sonnet-20240620') + + # Mock client to raise an exception + llm.client = Mock() + llm.client.messages.create = AsyncMock(side_effect=Exception('API Error')) + + messages = [{'role': 'user', 'content': 'Hello'}] + + with pytest.raises(Exception, match='Error in Claude API call: API Error'): + await llm.generate(messages) + + def test_anthropic_model_parameter_handling(self): + """Test that model parameter is properly handled.""" + test_models = [ + 'claude-3-5-sonnet-20240620', + 'claude-3-opus-20240229', + 'claude-3-sonnet-20240229', + 'claude-3-haiku-20240307', + ] + + for model in test_models: + llm = Anthropic(model=model) + assert llm.model == model + + def test_anthropic_api_key_handling(self): + """Test API key handling.""" + # Test with API key + llm = Anthropic(api_key='secret-key-123') + assert llm.api_key == 'secret-key-123' + + # Test without API key + llm = Anthropic() + assert llm.api_key is None + + # Test with empty string API key + llm = Anthropic(api_key='') + assert llm.api_key == '' + + def test_anthropic_base_url_handling(self): + """Test base URL handling.""" + # Test with base URL + llm = Anthropic(base_url='https://custom.anthropic.com') + assert llm.client.base_url == 'https://custom.anthropic.com' + + # Test without base URL + llm = Anthropic() + assert not hasattr(llm, 'base_url') diff --git a/flo_ai/tests/test_arium_yaml.py b/flo_ai/tests/test_arium_yaml.py index e1935656..3b8c22a7 100644 --- a/flo_ai/tests/test_arium_yaml.py +++ b/flo_ai/tests/test_arium_yaml.py @@ -793,6 +793,11 @@ def test_from_yaml_mixed_configuration_methods(self): mock_llm = Mock() mock_openai.return_value = mock_llm + # Mock for direct agent configuration + mock_direct_agent = Mock(spec=Agent) + mock_direct_agent.name = 'direct_agent' + mock_direct_agent.role = 'Direct Agent' + # Mock for inline YAML config mock_yaml_agent = Mock(spec=Agent) mock_yaml_agent.name = 'yaml_agent' @@ -801,6 +806,24 @@ def test_from_yaml_mixed_configuration_methods(self): mock_file_agent = Mock(spec=Agent) mock_file_agent.name = 'file_agent' + # Mock the AgentBuilder instance for direct configuration + mock_direct_builder = Mock() + mock_direct_builder.with_name.return_value = mock_direct_builder + mock_direct_builder.with_prompt.return_value = mock_direct_builder + mock_direct_builder.with_llm.return_value = mock_direct_builder + mock_direct_builder.with_tools.return_value = mock_direct_builder + mock_direct_builder.with_retries.return_value = mock_direct_builder + mock_direct_builder.with_reasoning.return_value = mock_direct_builder + mock_direct_builder.with_output_schema.return_value = ( + mock_direct_builder + ) + mock_direct_builder.with_role.return_value = mock_direct_builder + mock_direct_builder.build.return_value = mock_direct_agent + + # Mock the AgentBuilder class to return our mocked builder + mock_agent_builder.return_value = mock_direct_builder + + # Mock for inline YAML and file config mock_builder_instance = Mock() mock_builder_instance.build.side_effect = [ mock_yaml_agent, @@ -1015,10 +1038,33 @@ def test_from_yaml_mixed_prebuilt_and_configured_agents(self): mock_llm = Mock() mock_openai.return_value = mock_llm + # Mock for direct agent configuration + mock_direct_agent = Mock(spec=Agent) + mock_direct_agent.name = 'direct_agent' + mock_direct_agent.role = 'Direct Agent' + # Mock for inline YAML config mock_yaml_agent = Mock(spec=Agent) mock_yaml_agent.name = 'yaml_agent' + # Mock the AgentBuilder instance for direct configuration + mock_direct_builder = Mock() + mock_direct_builder.with_name.return_value = mock_direct_builder + mock_direct_builder.with_prompt.return_value = mock_direct_builder + mock_direct_builder.with_llm.return_value = mock_direct_builder + mock_direct_builder.with_tools.return_value = mock_direct_builder + mock_direct_builder.with_retries.return_value = mock_direct_builder + mock_direct_builder.with_reasoning.return_value = mock_direct_builder + mock_direct_builder.with_output_schema.return_value = ( + mock_direct_builder + ) + mock_direct_builder.with_role.return_value = mock_direct_builder + mock_direct_builder.build.return_value = mock_direct_agent + + # Mock the AgentBuilder class to return our mocked builder + mock_agent_builder.return_value = mock_direct_builder + + # Mock for inline YAML config mock_builder_instance = Mock() mock_builder_instance.build.return_value = mock_yaml_agent mock_agent_builder.from_yaml.return_value = mock_builder_instance diff --git a/flo_ai/tests/test_base_llm.py b/flo_ai/tests/test_base_llm.py new file mode 100644 index 00000000..147e3457 --- /dev/null +++ b/flo_ai/tests/test_base_llm.py @@ -0,0 +1,237 @@ +#!/usr/bin/env python3 +""" +Pytest tests for the BaseLLM abstract class and ImageMessage dataclass. +""" + +import sys +import os +import pytest +from unittest.mock import Mock + +# Add the flo_ai directory to the path +sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..')) + +from flo_ai.llm.base_llm import BaseLLM, ImageMessage + + +class MockLLM(BaseLLM): + """Mock LLM implementation for testing BaseLLM functionality""" + + async def generate(self, messages, functions=None): + return {'content': 'Mock response'} + + async def get_function_call(self, response): + if hasattr(response, 'function_call') and response.function_call: + return { + 'name': response.function_call.name, + 'arguments': response.function_call.arguments, + } + elif isinstance(response, dict) and 'function_call' in response: + return { + 'name': response['function_call']['name'], + 'arguments': response['function_call']['arguments'], + } + return None + + def get_message_content(self, response): + if isinstance(response, dict): + return response.get('content', '') + return str(response) + + def format_tool_for_llm(self, tool): + return { + 'name': tool.name, + 'description': tool.description, + 'parameters': {'type': 'object', 'properties': {}}, + } + + def format_tools_for_llm(self, tools): + return [self.format_tool_for_llm(tool) for tool in tools] + + def format_image_in_message(self, image): + return f'image:{image.mime_type}' + + +class TestImageMessage: + """Test class for ImageMessage dataclass.""" + + def test_image_message_creation(self): + """Test creating ImageMessage with different parameters.""" + # Test with image_url + img1 = ImageMessage(image_url='https://example.com/image.jpg') + assert img1.image_url == 'https://example.com/image.jpg' + assert img1.image_bytes is None + assert img1.image_file_path is None + assert img1.image_base64 is None + assert img1.mime_type is None + + # Test with image_bytes + img2 = ImageMessage(image_bytes=b'fake_image_data', mime_type='image/jpeg') + assert img2.image_bytes == b'fake_image_data' + assert img2.mime_type == 'image/jpeg' + + # Test with image_file_path + img3 = ImageMessage(image_file_path='/path/to/image.png', mime_type='image/png') + assert img3.image_file_path == '/path/to/image.png' + assert img3.mime_type == 'image/png' + + # Test with image_base64 + img4 = ImageMessage(image_base64='base64_encoded_string', mime_type='image/gif') + assert img4.image_base64 == 'base64_encoded_string' + assert img4.mime_type == 'image/gif' + + def test_image_message_defaults(self): + """Test ImageMessage with no parameters.""" + img = ImageMessage() + assert img.image_url is None + assert img.image_bytes is None + assert img.image_file_path is None + assert img.image_base64 is None + assert img.mime_type is None + + +class TestBaseLLM: + """Test class for BaseLLM abstract class functionality.""" + + def test_base_llm_initialization(self): + """Test BaseLLM initialization with different parameters.""" + # Test with minimal parameters + llm = MockLLM(model='test-model') + assert llm.model == 'test-model' + assert llm.api_key is None + assert llm.temperature == 0.7 + assert llm.kwargs == {} + + # Test with all parameters + llm = MockLLM( + model='test-model-2', + api_key='test-key', + temperature=0.5, + max_tokens=100, + top_p=0.9, + ) + assert llm.model == 'test-model-2' + assert llm.api_key == 'test-key' + assert llm.temperature == 0.5 + assert llm.kwargs == {'max_tokens': 100, 'top_p': 0.9} + + def test_base_llm_temperature_validation(self): + """Test temperature parameter handling.""" + # Test default temperature + llm = MockLLM(model='test-model') + assert llm.temperature == 0.7 + + # Test custom temperature + llm = MockLLM(model='test-model', temperature=0.0) + assert llm.temperature == 0.0 + + # Test high temperature + llm = MockLLM(model='test-model', temperature=1.0) + assert llm.temperature == 1.0 + + # Test temperature in kwargs + llm = MockLLM(model='test-model', temperature=0.3, custom_temp=0.8) + assert llm.temperature == 0.3 + assert llm.kwargs['custom_temp'] == 0.8 + + def test_base_llm_kwargs_storage(self): + """Test that additional kwargs are properly stored.""" + llm = MockLLM( + model='test-model', + max_tokens=1000, + top_p=0.9, + frequency_penalty=0.1, + presence_penalty=0.1, + ) + + assert 'max_tokens' in llm.kwargs + assert 'top_p' in llm.kwargs + assert 'frequency_penalty' in llm.kwargs + assert 'presence_penalty' in llm.kwargs + assert llm.kwargs['max_tokens'] == 1000 + assert llm.kwargs['top_p'] == 0.9 + + @pytest.mark.asyncio + async def test_base_llm_get_function_call(self): + """Test get_function_call method with different response formats.""" + llm = MockLLM(model='test-model') + + # Test with response object that has function_call attribute + mock_response = Mock() + mock_response.function_call.name = 'test_function' + mock_response.function_call.arguments = '{"param": "value"}' + + result = await llm.get_function_call(mock_response) + assert result == {'name': 'test_function', 'arguments': '{"param": "value"}'} + + # Test with dict response + dict_response = { + 'function_call': { + 'name': 'test_function_2', + 'arguments': '{"param2": "value2"}', + } + } + + result = await llm.get_function_call(dict_response) + assert result == { + 'name': 'test_function_2', + 'arguments': '{"param2": "value2"}', + } + + # Test with response that has no function_call + no_function_response = {'content': 'No function call here'} + result = await llm.get_function_call(no_function_response) + assert result is None + + # Test with response that has function_call but it's None + none_function_response = Mock() + none_function_response.function_call = None + + result = await llm.get_function_call(none_function_response) + assert result is None + + def test_base_llm_abstract_methods(self): + """Test that abstract methods are properly defined.""" + # This should not raise an error since MockLLM implements all abstract methods + llm = MockLLM(model='test-model') + + # Verify all required methods exist + assert hasattr(llm, 'generate') + assert hasattr(llm, 'get_message_content') + assert hasattr(llm, 'format_tool_for_llm') + assert hasattr(llm, 'format_tools_for_llm') + assert hasattr(llm, 'format_image_in_message') + + # Verify they are callable + assert callable(llm.generate) + assert callable(llm.get_message_content) + assert callable(llm.format_tool_for_llm) + assert callable(llm.format_tools_for_llm) + assert callable(llm.format_image_in_message) + + def test_base_llm_cannot_instantiate_abstract(self): + """Test that BaseLLM cannot be instantiated directly.""" + with pytest.raises(TypeError): + BaseLLM(model='test-model') + + def test_base_llm_model_validation(self): + """Test that model parameter is properly set.""" + test_models = ['gpt-4', 'claude-3', 'gemini-pro', 'llama2'] + + for model in test_models: + llm = MockLLM(model=model) + assert llm.model == model + + def test_base_llm_api_key_handling(self): + """Test API key handling.""" + # Test with API key + llm = MockLLM(model='test-model', api_key='secret-key-123') + assert llm.api_key == 'secret-key-123' + + # Test without API key + llm = MockLLM(model='test-model') + assert llm.api_key is None + + # Test with empty string API key + llm = MockLLM(model='test-model', api_key='') + assert llm.api_key == '' diff --git a/flo_ai/tests/test_gemini_llm.py b/flo_ai/tests/test_gemini_llm.py new file mode 100644 index 00000000..50922502 --- /dev/null +++ b/flo_ai/tests/test_gemini_llm.py @@ -0,0 +1,490 @@ +#!/usr/bin/env python3 +""" +Pytest tests for the Gemini LLM implementation. +""" + +import sys +import os +import pytest +from unittest.mock import Mock, patch, mock_open + +# Add the flo_ai directory to the path +sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..')) + +from flo_ai.llm.gemini_llm import Gemini +from flo_ai.llm.base_llm import ImageMessage +from flo_ai.tool.base_tool import Tool + +os.environ['GOOGLE_API_KEY'] = 'test-key-123' + + +class TestGemini: + """Test class for Gemini LLM implementation.""" + + @patch('flo_ai.llm.gemini_llm.genai.Client') + def test_gemini_initialization(self, mock_genai_client): + """Test Gemini LLM initialization with different parameters.""" + mock_client = Mock() + mock_genai_client.return_value = mock_client + + # Test with minimal parameters + llm = Gemini() + assert llm.model == 'gemini-2.5-flash' + assert llm.api_key is None + assert llm.temperature == 0.7 + assert llm.kwargs == {} + + # Test with custom parameters + mock_genai_client.reset_mock() + llm = Gemini( + model='gemini-1.5-pro', + api_key='test-key-123', + temperature=0.5, + max_output_tokens=1000, + ) + assert llm.model == 'gemini-1.5-pro' + assert llm.api_key == 'test-key-123' + assert llm.temperature == 0.5 + assert llm.kwargs == {'max_output_tokens': 1000} + + # Test with base_url (should be ignored as it's not implemented) + mock_genai_client.reset_mock() + llm = Gemini(base_url='https://custom.gemini.com') + # base_url is not stored as an attribute, so we just verify it doesn't crash + + @patch('flo_ai.llm.gemini_llm.genai.Client') + def test_gemini_temperature_handling(self, mock_genai_client): + """Test temperature parameter handling.""" + mock_client = Mock() + mock_genai_client.return_value = mock_client + + # Test default temperature + llm = Gemini() + assert llm.temperature == 0.7 + + # Test custom temperature + mock_genai_client.reset_mock() + llm = Gemini(temperature=0.0) + assert llm.temperature == 0.0 + + # Test high temperature + mock_genai_client.reset_mock() + llm = Gemini(temperature=1.0) + assert llm.temperature == 1.0 + + # Test temperature in kwargs + mock_genai_client.reset_mock() + llm = Gemini(temperature=0.3, custom_temp=0.8) + assert llm.temperature == 0.3 + assert llm.kwargs['custom_temp'] == 0.8 + + @patch('flo_ai.llm.gemini_llm.genai.Client') + def test_gemini_client_creation(self, mock_genai_client): + """Test that genai Client is created correctly.""" + mock_client = Mock() + mock_genai_client.return_value = mock_client + + llm = Gemini(api_key='test-key') + + mock_genai_client.assert_called_once_with(api_key='test-key') + assert llm.client == mock_client + + # Test without API key + mock_genai_client.reset_mock() + llm = Gemini() + + mock_genai_client.assert_called_once_with() + assert llm.client == mock_client + + @pytest.mark.asyncio + @patch('flo_ai.llm.gemini_llm.genai.Client') + async def test_gemini_generate_basic(self, mock_genai_client): + """Test basic generate method without functions or output schema.""" + mock_client = Mock() + mock_genai_client.return_value = mock_client + + llm = Gemini(model='gemini-2.5-flash') + + # Mock the client response + mock_response = Mock() + mock_response.text = 'Hello, world!' + + llm.client = mock_client + llm.client.models.generate_content = Mock(return_value=mock_response) + + messages = [{'role': 'user', 'content': 'Hello'}] + result = await llm.generate(messages) + + # Verify the API call + llm.client.models.generate_content.assert_called_once() + call_args = llm.client.models.generate_content.call_args + + assert call_args[1]['model'] == 'gemini-2.5-flash' + assert call_args[1]['contents'] == ['Hello'] + assert call_args[1]['config'].temperature == 0.7 + + # Verify the result + assert result == {'content': 'Hello, world!'} + + @pytest.mark.asyncio + @patch('flo_ai.llm.gemini_llm.genai.Client') + async def test_gemini_generate_with_system_message(self, mock_genai_client): + """Test generate method with system message.""" + mock_client = Mock() + mock_genai_client.return_value = mock_client + + llm = Gemini(model='gemini-2.5-flash') + + # Mock the client response + mock_response = Mock() + mock_response.text = "I'm a helpful assistant" + + llm.client = mock_client + llm.client.models.generate_content = Mock(return_value=mock_response) + + messages = [ + {'role': 'system', 'content': 'You are a helpful assistant'}, + {'role': 'user', 'content': 'Hello'}, + ] + + await llm.generate(messages) + + # Verify system instruction was passed correctly + call_args = llm.client.models.generate_content.call_args + config = call_args[1]['config'] + assert config.system_instruction == 'You are a helpful assistant\n' + + # Verify contents don't include system message + contents = call_args[1]['contents'] + assert contents == ['Hello'] + + @pytest.mark.asyncio + @patch('flo_ai.llm.gemini_llm.types.GenerateContentConfig') + async def test_gemini_generate_with_output_schema(self, mock_config_class): + """Test generate method with output schema.""" + llm = Gemini(model='gemini-2.5-flash') + + output_schema = { + 'type': 'object', + 'properties': {'message': {'type': 'string'}, 'count': {'type': 'integer'}}, + } + + # Mock the config + mock_config = Mock() + mock_config_class.return_value = mock_config + + # Mock the client response + mock_response = Mock() + mock_response.text = '{"message": "test", "count": 42}' + mock_response.candidates = [] + + llm.client = Mock() + llm.client.models.generate_content = Mock(return_value=mock_response) + + messages = [{'role': 'user', 'content': 'Generate JSON'}] + result = await llm.generate(messages, output_schema=output_schema) + + # Verify structured output was configured + mock_config.response_mime_type = 'application/json' + mock_config.response_schema = output_schema + + # Verify the result + assert result == {'content': '{"message": "test", "count": 42}'} + + @pytest.mark.asyncio + @patch('flo_ai.llm.gemini_llm.types.Tool') + @patch('flo_ai.llm.gemini_llm.types.GenerateContentConfig') + async def test_gemini_generate_with_functions( + self, mock_config_class, mock_tool_class + ): + """Test generate method with functions (tools).""" + llm = Gemini(model='gemini-2.5-flash') + + functions = [ + { + 'name': 'test_function', + 'description': 'A test function', + 'parameters': {'type': 'object'}, + } + ] + + # Mock the tool and config + mock_tool = Mock() + mock_tool_class.return_value = mock_tool + + mock_config = Mock() + mock_config_class.return_value = mock_config + + # Mock the client response + mock_response = Mock() + mock_response.text = "I'll use the function" + mock_response.candidates = [] + + llm.client = Mock() + llm.client.models.generate_content = Mock(return_value=mock_response) + + messages = [{'role': 'user', 'content': 'Use the function'}] + result = await llm.generate(messages, functions=functions) + + # Verify tool was created with function declarations + mock_tool_class.assert_called_once_with(function_declarations=functions) + + # Verify tools were added to config + mock_config.tools = [mock_tool] + + # Verify the result + assert result == {'content': "I'll use the function"} + + @pytest.mark.asyncio + @patch('flo_ai.llm.gemini_llm.types.Tool') + @patch('flo_ai.llm.gemini_llm.types.GenerateContentConfig') + async def test_gemini_generate_with_function_call_detection( + self, mock_config_class, mock_tool_class + ): + """Test generate method with function call detection.""" + llm = Gemini(model='gemini-2.5-flash') + + functions = [ + { + 'name': 'test_function', + 'description': 'A test function', + 'parameters': {'type': 'object'}, + } + ] + + # Mock the tool and config + mock_tool = Mock() + mock_tool_class.return_value = mock_tool + + mock_config = Mock() + mock_config_class.return_value = mock_config + + # Mock the client response with function call + mock_function_call = Mock() + mock_function_call.name = 'test_function' + mock_function_call.args = {'param': 'value'} + + mock_part = Mock() + mock_part.function_call = mock_function_call + + mock_content = Mock() + mock_content.parts = [mock_part] + + mock_candidate = Mock() + mock_candidate.content = mock_content + + mock_response = Mock() + mock_response.text = 'Function called' + mock_response.candidates = [mock_candidate] + + llm.client = Mock() + llm.client.models.generate_content = Mock(return_value=mock_response) + + messages = [{'role': 'user', 'content': 'Use the function'}] + result = await llm.generate(messages, functions=functions) + + # Verify function call was detected and parsed + assert 'function_call' in result + assert result['function_call']['name'] == 'test_function' + assert result['function_call']['arguments'] == {'param': 'value'} + + @pytest.mark.asyncio + async def test_gemini_generate_with_kwargs(self): + """Test generate method with additional kwargs.""" + llm = Gemini(model='gemini-2.5-flash', top_p=0.9, max_output_tokens=1000) + + # Mock the client response + mock_response = Mock() + mock_response.text = 'Response with kwargs' + mock_response.candidates = [] + + llm.client = Mock() + llm.client.models.generate_content = Mock(return_value=mock_response) + + messages = [{'role': 'user', 'content': 'Hello'}] + await llm.generate(messages) + + # Verify kwargs were passed through + call_args = llm.client.models.generate_content.call_args + config = call_args[1]['config'] + assert config.top_p == 0.9 + assert config.max_output_tokens == 1000 + + def test_gemini_get_message_content(self): + """Test get_message_content method.""" + llm = Gemini() + + # Test with dict response + response = {'content': 'Hello, world!'} + result = llm.get_message_content(response) + assert result == 'Hello, world!' + + # Test with string response + result = llm.get_message_content('Direct string') + assert result == 'Direct string' + + # Test with empty content + response = {'content': ''} + result = llm.get_message_content(response) + assert result == '' + + def test_gemini_format_tool_for_llm(self): + """Test format_tool_for_llm method.""" + llm = Gemini() + + # Create a mock tool + tool = Tool( + name='test_tool', + description='A test tool', + function=lambda x: x, + parameters={ + 'param1': {'type': 'string', 'description': 'First parameter'}, + 'param2': {'type': 'integer', 'description': 'Second parameter'}, + }, + ) + + formatted = llm.format_tool_for_llm(tool) + + assert formatted['name'] == 'test_tool' + assert formatted['description'] == 'A test tool' + assert formatted['parameters']['type'] == 'object' + assert 'param1' in formatted['parameters']['properties'] + assert 'param2' in formatted['parameters']['properties'] + assert formatted['parameters']['required'] == ['param1', 'param2'] + + def test_gemini_format_tools_for_llm(self): + """Test format_tools_for_llm method.""" + llm = Gemini() + + # Create mock tools + tool1 = Tool( + name='tool1', + description='First tool', + function=lambda x: x, + parameters={'param': {'type': 'string', 'description': 'Parameter'}}, + ) + + tool2 = Tool( + name='tool2', + description='Second tool', + function=lambda x: x, + parameters={'param': {'type': 'integer', 'description': 'Parameter'}}, + ) + + formatted = llm.format_tools_for_llm([tool1, tool2]) + + assert len(formatted) == 2 + assert formatted[0]['name'] == 'tool1' + assert formatted[1]['name'] == 'tool2' + + def test_gemini_format_image_in_message_file_path(self): + """Test format_image_in_message method with file path.""" + llm = Gemini() + + # Mock file reading + with patch('builtins.open', mock_open(read_data=b'fake_image_data')): + image = ImageMessage( + image_file_path='/path/to/image.jpg', mime_type='image/jpeg' + ) + + result = llm.format_image_in_message(image) + + # Verify genai.types.Part.from_bytes was called + # Note: We can't easily test the genai call without more complex mocking + # but we can verify the method doesn't raise an exception + assert result is not None + + def test_gemini_format_image_in_message_bytes(self): + """Test format_image_in_message method with image bytes.""" + llm = Gemini() + + image = ImageMessage(image_bytes=b'fake_image_data', mime_type='image/png') + + result = llm.format_image_in_message(image) + + # Verify genai.types.Part.from_bytes was called + assert result is not None + + def test_gemini_format_image_in_message_unsupported(self): + """Test format_image_in_message method with unsupported image format.""" + llm = Gemini() + + # Test with image_url (not implemented) + image = ImageMessage(image_url='https://example.com/image.jpg') + + with pytest.raises( + NotImplementedError, + match='Not other way other than file path has been implemented', + ): + llm.format_image_in_message(image) + + # Test with image_base64 (not implemented) + image = ImageMessage(image_base64='base64_string') + + with pytest.raises( + NotImplementedError, + match='Not other way other than file path has been implemented', + ): + llm.format_image_in_message(image) + + @pytest.mark.asyncio + async def test_gemini_generate_error_handling(self): + """Test error handling in generate method.""" + llm = Gemini(model='gemini-2.5-flash') + + # Mock client to raise an exception + llm.client = Mock() + llm.client.models.generate_content = Mock(side_effect=Exception('API Error')) + + messages = [{'role': 'user', 'content': 'Hello'}] + + with pytest.raises(Exception, match='Error in Gemini API call: API Error'): + await llm.generate(messages) + + def test_gemini_model_parameter_handling(self): + """Test that model parameter is properly handled.""" + test_models = [ + 'gemini-2.5-flash', + 'gemini-1.5-pro', + 'gemini-1.5-flash', + 'gemini-pro', + ] + + for model in test_models: + llm = Gemini(model=model) + assert llm.model == model + + def test_gemini_api_key_handling(self): + """Test API key handling.""" + # Test with API key + llm = Gemini(api_key='secret-key-123') + assert llm.api_key == 'secret-key-123' + + # Test without API key + llm = Gemini() + assert llm.api_key is None + + # Test with empty string API key + llm = Gemini(api_key='') + assert llm.api_key == '' + + def test_gemini_generation_config_creation(self): + """Test that generation config is created correctly.""" + llm = Gemini(temperature=0.5, max_output_tokens=1000, top_p=0.9) + + # Mock the client response + mock_response = Mock() + mock_response.text = 'Test response' + mock_response.candidates = [] + + llm.client = Mock() + llm.client.models.generate_content = Mock(return_value=mock_response) + + # We need to patch the types.GenerateContentConfig to test this + with patch('flo_ai.llm.gemini_llm.types.GenerateContentConfig') as mock_config: + mock_config_instance = Mock() + mock_config.return_value = mock_config_instance + + # This would normally be called in generate method + # For testing, we'll just verify the config class exists + assert mock_config is not None diff --git a/flo_ai/tests/test_openai_llm.py b/flo_ai/tests/test_openai_llm.py new file mode 100644 index 00000000..3a069f5a --- /dev/null +++ b/flo_ai/tests/test_openai_llm.py @@ -0,0 +1,325 @@ +#!/usr/bin/env python3 +""" +Pytest tests for the OpenAI LLM implementation. +""" + +import sys +import os +import pytest +from unittest.mock import Mock, AsyncMock, patch + +# Add the flo_ai directory to the path +sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..')) + +from flo_ai.llm.openai_llm import OpenAI +from flo_ai.llm.base_llm import ImageMessage +from flo_ai.tool.base_tool import Tool + +os.environ['OPENAI_API_KEY'] = 'test-key-123' + + +class TestOpenAI: + """Test class for OpenAI LLM implementation.""" + + def test_openai_initialization(self): + """Test OpenAI LLM initialization with different parameters.""" + # Test with minimal parameters + llm = OpenAI() + assert llm.model == 'gpt-4o-mini' + assert llm.api_key is None + assert llm.temperature == 0.7 + assert llm.kwargs == {} + + # Test with custom parameters + llm = OpenAI( + model='gpt-4', api_key='test-key-123', temperature=0.5, max_tokens=1000 + ) + assert llm.model == 'gpt-4' + assert llm.api_key == 'test-key-123' + assert llm.temperature == 0.5 + assert llm.kwargs == {'max_tokens': 1000} + + # Test with base_url + llm = OpenAI(base_url='https://custom.openai.com') + assert llm.client.base_url == 'https://custom.openai.com' + + def test_openai_temperature_handling(self): + """Test temperature parameter handling.""" + # Test default temperature + llm = OpenAI() + assert llm.temperature == 0.7 + + # Test custom temperature + llm = OpenAI(temperature=0.0) + assert llm.temperature == 0.0 + + # Test high temperature + llm = OpenAI(temperature=1.0) + assert llm.temperature == 1.0 + + # Test temperature in kwargs + llm = OpenAI(temperature=0.3, custom_temp=0.8) + assert llm.temperature == 0.3 + assert llm.kwargs['custom_temp'] == 0.8 + + @patch('flo_ai.llm.openai_llm.AsyncOpenAI') + def test_openai_client_creation(self, mock_async_openai): + """Test that AsyncOpenAI client is created correctly.""" + mock_client = Mock() + mock_async_openai.return_value = mock_client + + llm = OpenAI(api_key='test-key', base_url='https://custom.com') + + mock_async_openai.assert_called_once_with( + api_key='test-key', base_url='https://custom.com' + ) + assert llm.client == mock_client + + @pytest.mark.asyncio + async def test_openai_generate_basic(self): + """Test basic generate method without output schema.""" + llm = OpenAI(model='gpt-4o-mini') + + # Mock the client response + mock_response = Mock() + mock_response.choices = [Mock()] + mock_response.choices[0].message = Mock() + mock_response.choices[0].message.content = 'Hello, world!' + + llm.client = Mock() + llm.client.chat.completions.create = AsyncMock(return_value=mock_response) + + messages = [{'role': 'user', 'content': 'Hello'}] + result = await llm.generate(messages) + + # Verify the API call + llm.client.chat.completions.create.assert_called_once() + call_args = llm.client.chat.completions.create.call_args + + assert call_args[1]['model'] == 'gpt-4o-mini' + assert call_args[1]['messages'] == messages + assert call_args[1]['temperature'] == 0.7 + + # Verify the result + assert result == mock_response.choices[0].message + + @pytest.mark.asyncio + async def test_openai_generate_with_output_schema(self): + """Test generate method with output schema.""" + llm = OpenAI(model='gpt-4o-mini') + + output_schema = { + 'title': 'test_response', + 'schema': { + 'type': 'object', + 'properties': { + 'message': {'type': 'string'}, + 'count': {'type': 'integer'}, + }, + }, + } + + # Mock the client response + mock_response = Mock() + mock_response.choices = [Mock()] + mock_response.choices[0].message = Mock() + mock_response.choices[0].message.content = '{"message": "test", "count": 42}' + + llm.client = Mock() + llm.client.chat.completions.create = AsyncMock(return_value=mock_response) + + messages = [{'role': 'user', 'content': 'Generate JSON'}] + await llm.generate(messages, output_schema=output_schema) + + # Verify the API call + call_args = llm.client.chat.completions.create.call_args[1] + + assert call_args['response_format'] == {'type': 'json_object'} + assert call_args['functions'] == [ + {'name': 'test_response', 'parameters': output_schema['schema']} + ] + assert call_args['function_call'] == {'name': 'test_response'} + + # Verify system message was modified + assert len(call_args['messages']) == 2 + assert call_args['messages'][0]['role'] == 'system' + assert 'JSON format' in call_args['messages'][0]['content'] + + @pytest.mark.asyncio + async def test_openai_generate_with_existing_system_message(self): + """Test generate method with existing system message and output schema.""" + llm = OpenAI(model='gpt-4o-mini') + + output_schema = {'title': 'test', 'schema': {'type': 'object'}} + + # Mock the client response + mock_response = Mock() + mock_response.choices = [Mock()] + mock_response.choices[0].message = Mock() + mock_response.choices[0].message.content = '{"test": "value"}' + + llm.client = Mock() + llm.client.chat.completions.create = AsyncMock(return_value=mock_response) + + messages = [ + {'role': 'system', 'content': 'You are a helpful assistant'}, + {'role': 'user', 'content': 'Generate JSON'}, + ] + + await llm.generate(messages, output_schema=output_schema) + + # Verify system message was appended to + call_args = llm.client.chat.completions.create.call_args[1] + system_message = call_args['messages'][0]['content'] + assert 'You are a helpful assistant' in system_message + assert 'JSON format' in system_message + + @pytest.mark.asyncio + async def test_openai_generate_with_kwargs(self): + """Test generate method with additional kwargs.""" + llm = OpenAI(model='gpt-4o-mini', max_tokens=1000, top_p=0.9) + + # Mock the client response + mock_response = Mock() + mock_response.choices = [Mock()] + mock_response.choices[0].message = Mock() + mock_response.choices[0].message.content = 'Response with kwargs' + + llm.client = Mock() + llm.client.chat.completions.create = AsyncMock(return_value=mock_response) + + messages = [{'role': 'user', 'content': 'Hello'}] + await llm.generate(messages, stream=False) + + # Verify kwargs were passed through + call_args = llm.client.chat.completions.create.call_args[1] + assert call_args['max_tokens'] == 1000 + assert call_args['top_p'] == 0.9 + assert not call_args['stream'] + + def test_openai_get_message_content(self): + """Test get_message_content method.""" + llm = OpenAI() + + # Test with string response + result = llm.get_message_content('Hello, world!') + assert result == 'Hello, world!' + + # Test with message object + mock_message = Mock() + mock_message.content = 'Message content' + result = llm.get_message_content(mock_message) + assert result == 'Message content' + + # Test with object without content attribute + mock_obj = Mock() + del mock_obj.content + result = llm.get_message_content(mock_obj) + assert result == str(mock_obj) + + def test_openai_format_tool_for_llm(self): + """Test format_tool_for_llm method.""" + llm = OpenAI() + + # Create a mock tool + tool = Tool( + name='test_tool', + description='A test tool', + function=lambda x: x, + parameters={ + 'param1': {'type': 'string', 'description': 'First parameter'}, + 'param2': {'type': 'integer', 'description': 'Second parameter'}, + }, + ) + + formatted = llm.format_tool_for_llm(tool) + + assert formatted['name'] == 'test_tool' + assert formatted['description'] == 'A test tool' + assert formatted['parameters']['type'] == 'object' + assert 'param1' in formatted['parameters']['properties'] + assert 'param2' in formatted['parameters']['properties'] + assert formatted['parameters']['required'] == ['param1', 'param2'] + + def test_openai_format_tools_for_llm(self): + """Test format_tools_for_llm method.""" + llm = OpenAI() + + # Create mock tools + tool1 = Tool( + name='tool1', + description='First tool', + function=lambda x: x, + parameters={'param': {'type': 'string', 'description': 'Parameter'}}, + ) + + tool2 = Tool( + name='tool2', + description='Second tool', + function=lambda x: x, + parameters={'param': {'type': 'integer', 'description': 'Parameter'}}, + ) + + formatted = llm.format_tools_for_llm([tool1, tool2]) + + assert len(formatted) == 2 + assert formatted[0]['name'] == 'tool1' + assert formatted[1]['name'] == 'tool2' + + def test_openai_format_image_in_message(self): + """Test format_image_in_message method.""" + llm = OpenAI() + + # This method is not implemented yet + image = ImageMessage(image_url='https://example.com/image.jpg') + + with pytest.raises(NotImplementedError): + llm.format_image_in_message(image) + + @pytest.mark.asyncio + async def test_openai_generate_error_handling(self): + """Test error handling in generate method.""" + llm = OpenAI(model='gpt-4o-mini') + + # Mock client to raise an exception + llm.client = Mock() + llm.client.chat.completions.create = AsyncMock( + side_effect=Exception('API Error') + ) + + messages = [{'role': 'user', 'content': 'Hello'}] + + with pytest.raises(Exception, match='API Error'): + await llm.generate(messages) + + def test_openai_model_parameter_handling(self): + """Test that model parameter is properly handled.""" + test_models = ['gpt-4', 'gpt-4o', 'gpt-4o-mini', 'gpt-3.5-turbo'] + + for model in test_models: + llm = OpenAI(model=model) + assert llm.model == model + + def test_openai_api_key_handling(self): + """Test API key handling.""" + # Test with API key + llm = OpenAI(api_key='secret-key-123') + assert llm.api_key == 'secret-key-123' + + # Test without API key + llm = OpenAI() + assert llm.api_key is None + + # Test with empty string API key + llm = OpenAI(api_key='') + assert llm.api_key == '' + + def test_openai_base_url_handling(self): + """Test base URL handling.""" + # Test with base URL + llm = OpenAI(base_url='https://custom.openai.com') + assert llm.client.base_url == 'https://custom.openai.com' + + # Test without base URL + llm = OpenAI() + assert not hasattr(llm, 'base_url') diff --git a/flo_ai/tests/test_openai_vllm.py b/flo_ai/tests/test_openai_vllm.py new file mode 100644 index 00000000..a7938c08 --- /dev/null +++ b/flo_ai/tests/test_openai_vllm.py @@ -0,0 +1,575 @@ +#!/usr/bin/env python3 +""" +Pytest tests for the OpenAI VLLM implementation. +""" + +import sys +import os +import pytest +from unittest.mock import Mock, AsyncMock, patch + +# Add the flo_ai directory to the path +sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..')) + +from flo_ai.llm.openai_vllm import OpenAIVLLM +from flo_ai.llm.base_llm import ImageMessage +from flo_ai.tool.base_tool import Tool + +os.environ['OPENAI_API_KEY'] = 'test-key-123' + + +class TestOpenAIVLLM: + """Test class for OpenAI VLLM implementation.""" + + @patch('flo_ai.llm.openai_llm.AsyncOpenAI') + def test_openai_vllm_initialization(self, mock_async_openai): + """Test OpenAI VLLM initialization with different parameters.""" + mock_client = Mock() + mock_async_openai.return_value = mock_client + + # Test with minimal parameters + llm = OpenAIVLLM(base_url='https://api.vllm.com', model='gpt-4o-mini') + assert llm.model == 'gpt-4o-mini' + assert llm.api_key is None + assert llm.temperature == 0.7 + assert llm.base_url == 'https://api.vllm.com' + assert llm.kwargs == {} + + # Test with custom parameters + mock_async_openai.reset_mock() + llm = OpenAIVLLM( + base_url='https://custom.vllm.com', + model='gpt-4', + api_key='test-key-123', + temperature=0.5, + max_tokens=1000, + ) + assert llm.model == 'gpt-4' + assert llm.api_key == 'test-key-123' + assert llm.temperature == 0.5 + assert llm.base_url == 'https://custom.vllm.com' + assert llm.kwargs == {'max_tokens': 1000} + + # Test with additional kwargs + mock_async_openai.reset_mock() + llm = OpenAIVLLM( + base_url='https://api.vllm.com', + model='gpt-4o-mini', + max_tokens=1000, + top_p=0.9, + ) + assert llm.kwargs == {'max_tokens': 1000, 'top_p': 0.9} + + @patch('flo_ai.llm.openai_llm.AsyncOpenAI') + def test_openai_vllm_temperature_handling(self, mock_async_openai): + """Test temperature parameter handling.""" + mock_client = Mock() + mock_async_openai.return_value = mock_client + + # Test default temperature + llm = OpenAIVLLM(base_url='https://api.vllm.com', model='gpt-4o-mini') + assert llm.temperature == 0.7 + + # Test custom temperature + mock_async_openai.reset_mock() + llm = OpenAIVLLM( + base_url='https://api.vllm.com', model='gpt-4o-mini', temperature=0.0 + ) + assert llm.temperature == 0.0 + + # Test high temperature + mock_async_openai.reset_mock() + llm = OpenAIVLLM( + base_url='https://api.vllm.com', model='gpt-4o-mini', temperature=1.0 + ) + assert llm.temperature == 1.0 + + # Test temperature in kwargs + mock_async_openai.reset_mock() + llm = OpenAIVLLM( + base_url='https://api.vllm.com', + model='gpt-4o-mini', + temperature=0.3, + custom_temp=0.8, + ) + assert llm.temperature == 0.3 + assert llm.kwargs['custom_temp'] == 0.8 + + @patch('flo_ai.llm.openai_llm.AsyncOpenAI') + def test_openai_vllm_client_creation(self, mock_async_openai): + """Test that AsyncOpenAI client is created correctly with VLLM parameters.""" + mock_client = Mock() + mock_async_openai.return_value = mock_client + + llm = OpenAIVLLM( + base_url='https://custom.vllm.com', model='gpt-4o-mini', api_key='test-key' + ) + + mock_async_openai.assert_called_once_with( + api_key='test-key', base_url='https://custom.vllm.com' + ) + assert llm.client == mock_client + + # Test without API key + mock_async_openai.reset_mock() + llm = OpenAIVLLM(base_url='https://api.vllm.com', model='gpt-4o-mini') + + mock_async_openai.assert_called_once_with( + api_key=None, base_url='https://api.vllm.com' + ) + assert llm.client == mock_client + + @pytest.mark.asyncio + @patch('flo_ai.llm.openai_llm.AsyncOpenAI') + async def test_openai_vllm_generate_basic(self, mock_async_openai): + """Test basic generate method without output schema.""" + mock_client = Mock() + mock_async_openai.return_value = mock_client + + llm = OpenAIVLLM(base_url='https://api.vllm.com', model='gpt-4o-mini') + + # Mock the client response + mock_choice = Mock() + mock_choice.message.content = 'Hello, world!' + + mock_response = Mock() + mock_response.choices = [mock_choice] + + llm.client.chat.completions.create = AsyncMock(return_value=mock_response) + + messages = [{'role': 'user', 'content': 'Hello'}] + result = await llm.generate(messages) + + # Verify the API call + llm.client.chat.completions.create.assert_called_once() + call_args = llm.client.chat.completions.create.call_args + + assert call_args[1]['model'] == 'gpt-4o-mini' + assert call_args[1]['messages'] == messages + assert call_args[1]['temperature'] == 0.7 + + # Verify the result + assert result.content == 'Hello, world!' + + @pytest.mark.asyncio + @patch('flo_ai.llm.openai_llm.AsyncOpenAI') + async def test_openai_vllm_generate_with_output_schema(self, mock_async_openai): + """Test generate method with output schema.""" + mock_client = Mock() + mock_async_openai.return_value = mock_client + + llm = OpenAIVLLM(base_url='https://api.vllm.com', model='gpt-4o-mini') + + output_schema = { + 'title': 'test_schema', + 'schema': { + 'type': 'object', + 'properties': { + 'message': {'type': 'string'}, + 'count': {'type': 'integer'}, + }, + }, + } + + # Mock the client response + mock_choice = Mock() + mock_choice.message.content = '{"message": "test", "count": 42}' + + mock_response = Mock() + mock_response.choices = [mock_choice] + + llm.client.chat.completions.create = AsyncMock(return_value=mock_response) + + messages = [{'role': 'user', 'content': 'Generate JSON'}] + result = await llm.generate(messages, output_schema=output_schema) + + # Verify output schema was configured + call_args = llm.client.chat.completions.create.call_args + assert call_args[1]['response_format']['type'] == 'json_schema' + assert call_args[1]['response_format']['json_schema']['name'] == 'test_schema' + + # Verify the result + assert result.content == '{"message": "test", "count": 42}' + + @pytest.mark.asyncio + @patch('flo_ai.llm.openai_llm.AsyncOpenAI') + async def test_openai_vllm_generate_with_existing_system_message( + self, mock_async_openai + ): + """Test generate method with existing system message and output schema.""" + mock_client = Mock() + mock_async_openai.return_value = mock_client + + llm = OpenAIVLLM(base_url='https://api.vllm.com', model='gpt-4o-mini') + + output_schema = {'title': 'test_schema', 'schema': {'type': 'object'}} + + # Mock the client response + mock_choice = Mock() + mock_choice.message.content = '{"result": "success"}' + + mock_response = Mock() + mock_response.choices = [mock_choice] + + llm.client.chat.completions.create = AsyncMock(return_value=mock_response) + + messages = [ + {'role': 'system', 'content': 'You are a helpful assistant'}, + {'role': 'user', 'content': 'Generate JSON'}, + ] + + result = await llm.generate(messages, output_schema=output_schema) + + # Verify system message was updated + call_args = llm.client.chat.completions.create.call_args + updated_messages = call_args[1]['messages'] + assert updated_messages[0]['role'] == 'system' + assert 'JSON format' in updated_messages[0]['content'] + assert 'test_schema' in updated_messages[0]['content'] + + # Verify the result + assert result.content == '{"result": "success"}' + + @pytest.mark.asyncio + @patch('flo_ai.llm.openai_llm.AsyncOpenAI') + async def test_openai_vllm_generate_with_kwargs(self, mock_async_openai): + """Test generate method with additional kwargs.""" + mock_client = Mock() + mock_async_openai.return_value = mock_client + + llm = OpenAIVLLM( + base_url='https://api.vllm.com', + model='gpt-4o-mini', + top_p=0.9, + max_output_tokens=1000, + ) + + # Mock the client response + mock_choice = Mock() + mock_choice.message.content = 'Response with kwargs' + + mock_response = Mock() + mock_response.choices = [mock_choice] + + llm.client.chat.completions.create = AsyncMock(return_value=mock_response) + + messages = [{'role': 'user', 'content': 'Hello'}] + result = await llm.generate(messages) + + # Verify kwargs were passed through + call_args = llm.client.chat.completions.create.call_args + assert call_args[1]['top_p'] == 0.9 + assert call_args[1]['max_output_tokens'] == 1000 + + # Verify the result + assert result.content == 'Response with kwargs' + + def test_openai_vllm_get_message_content(self): + """Test get_message_content method.""" + llm = OpenAIVLLM(base_url='https://api.vllm.com', model='gpt-4o-mini') + + # Test with dict response (should return str representation) + response = {'content': 'Hello, world!'} + result = llm.get_message_content(response) + assert result == "{'content': 'Hello, world!'}" + + # Test with string response + result = llm.get_message_content('Direct string') + assert result == 'Direct string' + + # Test with empty content + response = {'content': ''} + result = llm.get_message_content(response) + assert result == "{'content': ''}" + + # Test with message object that has content attribute + mock_message = Mock() + mock_message.content = 'Message content' + result = llm.get_message_content(mock_message) + assert result == 'Message content' + + def test_openai_vllm_format_tool_for_llm(self): + """Test format_tool_for_llm method.""" + llm = OpenAIVLLM(base_url='https://api.vllm.com', model='gpt-4o-mini') + + # Create a mock tool + tool = Tool( + name='test_tool', + description='A test tool', + function=lambda x: x, + parameters={ + 'param1': {'type': 'string', 'description': 'First parameter'}, + 'param2': {'type': 'integer', 'description': 'Second parameter'}, + }, + ) + + formatted = llm.format_tool_for_llm(tool) + + assert formatted['name'] == 'test_tool' + assert formatted['description'] == 'A test tool' + assert formatted['parameters']['type'] == 'object' + assert 'param1' in formatted['parameters']['properties'] + assert 'param2' in formatted['parameters']['properties'] + assert formatted['parameters']['required'] == ['param1', 'param2'] + + def test_openai_vllm_format_tools_for_llm(self): + """Test format_tools_for_llm method.""" + llm = OpenAIVLLM(base_url='https://api.vllm.com', model='gpt-4o-mini') + + # Create mock tools + tool1 = Tool( + name='tool1', + description='First tool', + function=lambda x: x, + parameters={'param': {'type': 'string', 'description': 'Parameter'}}, + ) + + tool2 = Tool( + name='tool2', + description='Second tool', + function=lambda x: x, + parameters={'param': {'type': 'integer', 'description': 'Parameter'}}, + ) + + formatted = llm.format_tools_for_llm([tool1, tool2]) + + assert len(formatted) == 2 + assert formatted[0]['name'] == 'tool1' + assert formatted[1]['name'] == 'tool2' + + def test_openai_vllm_format_image_in_message(self): + """Test format_image_in_message method.""" + llm = OpenAIVLLM(base_url='https://api.vllm.com', model='gpt-4o-mini') + + # Test with image message + image = ImageMessage(image_url='https://example.com/image.jpg') + + with pytest.raises( + NotImplementedError, match='Not implemented image for LLM OpenAI' + ): + llm.format_image_in_message(image) + + @pytest.mark.asyncio + @patch('flo_ai.llm.openai_llm.AsyncOpenAI') + async def test_openai_vllm_generate_error_handling(self, mock_async_openai): + """Test error handling in generate method.""" + mock_client = Mock() + mock_async_openai.return_value = mock_client + + llm = OpenAIVLLM(base_url='https://api.vllm.com', model='gpt-4o-mini') + + # Mock client to raise an exception + llm.client.chat.completions.create = AsyncMock( + side_effect=Exception('API Error') + ) + + messages = [{'role': 'user', 'content': 'Hello'}] + + with pytest.raises(Exception, match='API Error'): + await llm.generate(messages) + + @patch('flo_ai.llm.openai_llm.AsyncOpenAI') + def test_openai_vllm_model_parameter_handling(self, mock_async_openai): + """Test that model parameter is properly handled.""" + mock_client = Mock() + mock_async_openai.return_value = mock_client + + test_models = ['gpt-4', 'gpt-4o', 'gpt-4o-mini', 'gpt-3.5-turbo'] + + for model in test_models: + mock_async_openai.reset_mock() + llm = OpenAIVLLM(base_url='https://api.vllm.com', model=model) + assert llm.model == model + + @patch('flo_ai.llm.openai_llm.AsyncOpenAI') + def test_openai_vllm_api_key_handling(self, mock_async_openai): + """Test API key handling.""" + mock_client = Mock() + mock_async_openai.return_value = mock_client + + # Test with API key + llm = OpenAIVLLM( + base_url='https://api.vllm.com', + model='gpt-4o-mini', + api_key='secret-key-123', + ) + assert llm.api_key == 'secret-key-123' + + # Test without API key + mock_async_openai.reset_mock() + llm = OpenAIVLLM(base_url='https://api.vllm.com', model='gpt-4o-mini') + assert llm.api_key is None + + # Test with empty string API key + mock_async_openai.reset_mock() + llm = OpenAIVLLM( + base_url='https://api.vllm.com', model='gpt-4o-mini', api_key='' + ) + assert llm.api_key == '' + + @patch('flo_ai.llm.openai_llm.AsyncOpenAI') + def test_openai_vllm_base_url_handling(self, mock_async_openai): + """Test base URL handling.""" + mock_client = Mock() + mock_async_openai.return_value = mock_client + + # Test with base URL + llm = OpenAIVLLM(base_url='https://custom.vllm.com', model='gpt-4o-mini') + assert llm.base_url == 'https://custom.vllm.com' + + # Test with different base URL + mock_async_openai.reset_mock() + llm = OpenAIVLLM(base_url='https://another.vllm.com', model='gpt-4o-mini') + assert llm.base_url == 'https://another.vllm.com' + + @patch('flo_ai.llm.openai_llm.AsyncOpenAI') + def test_openai_vllm_inheritance_from_openai(self, mock_async_openai): + """Test that OpenAIVLLM inherits from OpenAI.""" + mock_client = Mock() + mock_async_openai.return_value = mock_client + + llm = OpenAIVLLM(base_url='https://api.vllm.com', model='gpt-4o-mini') + + # Should inherit from OpenAI + from flo_ai.llm.openai_llm import OpenAI + + assert isinstance(llm, OpenAI) + + # Should have all the methods from OpenAI + assert hasattr(llm, 'generate') + assert hasattr(llm, 'get_message_content') + assert hasattr(llm, 'format_tool_for_llm') + assert hasattr(llm, 'format_tools_for_llm') + assert hasattr(llm, 'format_image_in_message') + + @patch('flo_ai.llm.openai_llm.AsyncOpenAI') + def test_openai_vllm_parameter_combinations(self, mock_async_openai): + """Test various parameter combinations.""" + mock_client = Mock() + mock_async_openai.return_value = mock_client + + # Test with all parameters + llm = OpenAIVLLM( + model='gpt-4', + api_key='key123', + temperature=0.3, + base_url='https://custom.com', + max_tokens=500, + top_p=0.8, + ) + + assert llm.model == 'gpt-4' + assert llm.api_key == 'key123' + assert llm.temperature == 0.3 + assert llm.base_url == 'https://custom.com' + assert llm.kwargs == {'max_tokens': 500, 'top_p': 0.8} + + # Test with minimal parameters + mock_async_openai.reset_mock() + llm = OpenAIVLLM(base_url='https://api.vllm.com', model='gpt-4o-mini') + + assert llm.model == 'gpt-4o-mini' + assert llm.api_key is None + assert llm.temperature == 0.7 + assert llm.base_url == 'https://api.vllm.com' + assert llm.kwargs == {} + + @patch('flo_ai.llm.openai_llm.AsyncOpenAI') + def test_openai_vllm_method_inheritance(self, mock_async_openai): + """Test that OpenAIVLLM inherits all methods from OpenAI.""" + mock_client = Mock() + mock_async_openai.return_value = mock_client + + llm = OpenAIVLLM(base_url='https://api.vllm.com', model='gpt-4o-mini') + + # Test that OpenAIVLLM has all the methods from OpenAI + assert hasattr(llm, 'generate') + assert hasattr(llm, 'get_message_content') + assert hasattr(llm, 'format_tool_for_llm') + assert hasattr(llm, 'format_tools_for_llm') + assert hasattr(llm, 'format_image_in_message') + + # Should be callable + assert callable(llm.generate) + assert callable(llm.get_message_content) + assert callable(llm.format_tool_for_llm) + assert callable(llm.format_tools_for_llm) + assert callable(llm.format_image_in_message) + + @patch('flo_ai.llm.openai_llm.AsyncOpenAI') + def test_openai_vllm_default_values(self, mock_async_openai): + """Test that default values are set correctly.""" + mock_client = Mock() + mock_async_openai.return_value = mock_client + + llm = OpenAIVLLM(base_url='https://api.vllm.com', model='gpt-4o-mini') + + # Default values from OpenAI + assert llm.model == 'gpt-4o-mini' + assert llm.temperature == 0.7 + + # Default values from BaseLLM + assert llm.api_key is None + assert llm.kwargs == {} + + # Default values from OpenAIVLLM + assert llm.base_url == 'https://api.vllm.com' + + @patch('flo_ai.llm.openai_llm.AsyncOpenAI') + def test_openai_vllm_parameter_override(self, mock_async_openai): + """Test that parameters can be overridden after initialization.""" + mock_client = Mock() + mock_async_openai.return_value = mock_client + + llm = OpenAIVLLM(base_url='https://api.vllm.com', model='gpt-4o-mini') + + # Change parameters + llm.model = 'new-model' + llm.temperature = 0.1 + llm.base_url = 'new-url' + + # Verify changes + assert llm.model == 'new-model' + assert llm.temperature == 0.1 + assert llm.base_url == 'new-url' + + @patch('flo_ai.llm.openai_llm.AsyncOpenAI') + def test_openai_vllm_kwargs_storage(self, mock_async_openai): + """Test that additional kwargs are properly stored.""" + mock_client = Mock() + mock_async_openai.return_value = mock_client + + llm = OpenAIVLLM( + base_url='https://api.vllm.com', + model='gpt-4o-mini', + max_tokens=1000, + top_p=0.9, + frequency_penalty=0.1, + presence_penalty=0.1, + ) + + assert 'max_tokens' in llm.kwargs + assert 'top_p' in llm.kwargs + assert 'frequency_penalty' in llm.kwargs + assert 'presence_penalty' in llm.kwargs + assert llm.kwargs['max_tokens'] == 1000 + assert llm.kwargs['top_p'] == 0.9 + + @patch('flo_ai.llm.openai_llm.AsyncOpenAI') + def test_openai_vllm_initialization_order(self, mock_async_openai): + """Test that initialization happens in the correct order.""" + mock_client = Mock() + mock_async_openai.return_value = mock_client + + # This should not raise an error + llm = OpenAIVLLM( + model='test-model', + base_url='https://test.vllm.com', + project='test-project', + location='test-location', + ) + + # Verify all attributes are set correctly + assert llm.model == 'test-model' + assert llm.base_url == 'https://test.vllm.com' + assert llm.client == mock_client diff --git a/flo_ai/tests/test_vertexai_llm.py b/flo_ai/tests/test_vertexai_llm.py new file mode 100644 index 00000000..54f1e79f --- /dev/null +++ b/flo_ai/tests/test_vertexai_llm.py @@ -0,0 +1,584 @@ +#!/usr/bin/env python3 +""" +Pytest tests for the VertexAI LLM implementation. +""" + +import sys +import os +import pytest +from unittest.mock import Mock, patch + +# Add the flo_ai directory to the path +sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..')) + +from flo_ai.llm.vertexai_llm import VertexAI +from flo_ai.tool.base_tool import Tool + +os.environ['GOOGLE_API_KEY'] = 'test-key-123' +os.environ['GOOGLE_PROJECT'] = 'my-project-123' +os.environ['GOOGLE_LOCATION'] = 'us-central1' + + +class TestVertexAI: + """Test class for VertexAI LLM implementation.""" + + @patch('flo_ai.llm.vertexai_llm.genai.Client') + def test_vertexai_initialization(self, mock_genai_client): + """Test VertexAI LLM initialization with different parameters.""" + mock_client = Mock() + mock_genai_client.return_value = mock_client + + # Test with minimal parameters + llm = VertexAI() + assert llm.model == 'gemini-2.5-flash' + assert llm.api_key is None + assert llm.temperature == 0.7 + assert llm.kwargs == {} + assert llm.project is None + assert llm.location is None + + # Test with custom parameters + mock_genai_client.reset_mock() + llm = VertexAI( + model='gemini-1.5-pro', + api_key='test-key-123', + temperature=0.5, + project='my-project-123', + location='us-central1', + ) + assert llm.model == 'gemini-1.5-pro' + assert llm.api_key == 'test-key-123' + assert llm.temperature == 0.5 + assert llm.project == 'my-project-123' + assert llm.location == 'us-central1' + assert llm.kwargs == {} + + # Test with additional kwargs + mock_genai_client.reset_mock() + llm = VertexAI(model='gemini-2.5-flash', max_tokens=1000, top_p=0.9) + assert llm.kwargs == {'max_tokens': 1000, 'top_p': 0.9} + + @patch('flo_ai.llm.vertexai_llm.genai.Client') + def test_vertexai_temperature_handling(self, mock_genai_client): + """Test temperature parameter handling.""" + mock_client = Mock() + mock_genai_client.return_value = mock_client + + # Test default temperature + llm = VertexAI() + assert llm.temperature == 0.7 + + # Test custom temperature + mock_genai_client.reset_mock() + llm = VertexAI(temperature=0.0) + assert llm.temperature == 0.0 + + # Test high temperature + mock_genai_client.reset_mock() + llm = VertexAI(temperature=1.0) + assert llm.temperature == 1.0 + + # Test temperature in kwargs + mock_genai_client.reset_mock() + llm = VertexAI(temperature=0.3, custom_temp=0.8) + assert llm.temperature == 0.3 + assert llm.kwargs['custom_temp'] == 0.8 + + @patch('flo_ai.llm.vertexai_llm.genai.Client') + def test_vertexai_client_creation(self, mock_genai_client): + """Test that genai Client is created correctly with VertexAI parameters.""" + mock_client = Mock() + mock_genai_client.return_value = mock_client + + llm = VertexAI(project='test-project', location='us-central1') + + mock_genai_client.assert_called_once_with( + project='test-project', location='us-central1', vertexai=True + ) + assert llm.client == mock_client + + # Test without project and location + mock_genai_client.reset_mock() + llm = VertexAI() + + mock_genai_client.assert_called_once_with( + project=None, location=None, vertexai=True + ) + assert llm.client == mock_client + + @pytest.mark.asyncio + @patch('flo_ai.llm.vertexai_llm.genai.Client') + async def test_vertexai_generate_basic(self, mock_genai_client): + """Test basic generate method without functions or output schema.""" + mock_client = Mock() + mock_genai_client.return_value = mock_client + + llm = VertexAI(model='gemini-2.5-flash') + + # Mock the client response + mock_response = Mock() + mock_response.text = 'Hello, world!' + mock_response.candidates = [] + + llm.client = mock_client + llm.client.models.generate_content = Mock(return_value=mock_response) + + messages = [{'role': 'user', 'content': 'Hello'}] + result = await llm.generate(messages) + + # Verify the API call + llm.client.models.generate_content.assert_called_once() + call_args = llm.client.models.generate_content.call_args + + assert call_args[1]['model'] == 'gemini-2.5-flash' + assert call_args[1]['contents'] == ['Hello'] + assert call_args[1]['config'].temperature == 0.7 + + # Verify the result + assert result == {'content': 'Hello, world!'} + + @pytest.mark.asyncio + @patch('flo_ai.llm.vertexai_llm.genai.Client') + async def test_vertexai_generate_with_system_message(self, mock_genai_client): + """Test generate method with system message.""" + mock_client = Mock() + mock_genai_client.return_value = mock_client + + llm = VertexAI(model='gemini-2.5-flash') + + # Mock the client response + mock_response = Mock() + mock_response.text = "I'm a helpful assistant" + mock_response.candidates = [] + + llm.client = mock_client + llm.client.models.generate_content = Mock(return_value=mock_response) + + messages = [ + {'role': 'system', 'content': 'You are a helpful assistant'}, + {'role': 'user', 'content': 'Hello'}, + ] + + await llm.generate(messages) + + # Verify system instruction was passed correctly + call_args = llm.client.models.generate_content.call_args + config = call_args[1]['config'] + assert config.system_instruction == 'You are a helpful assistant\n' + + # Verify contents don't include system message + contents = call_args[1]['contents'] + assert contents == ['Hello'] + + @pytest.mark.asyncio + @patch('flo_ai.llm.vertexai_llm.genai.Client') + async def test_vertexai_generate_with_output_schema(self, mock_genai_client): + """Test generate method with output schema.""" + mock_client = Mock() + mock_genai_client.return_value = mock_client + + llm = VertexAI(model='gemini-2.5-flash') + + output_schema = { + 'type': 'object', + 'properties': {'message': {'type': 'string'}, 'count': {'type': 'integer'}}, + } + + # Mock the client response + mock_response = Mock() + mock_response.text = '{"message": "test", "count": 42}' + mock_response.candidates = [] + + llm.client = mock_client + llm.client.models.generate_content = Mock(return_value=mock_response) + + messages = [{'role': 'user', 'content': 'Generate JSON'}] + result = await llm.generate(messages, output_schema=output_schema) + + # Verify the result + assert result == {'content': '{"message": "test", "count": 42}'} + + @pytest.mark.asyncio + @patch('flo_ai.llm.vertexai_llm.genai.Client') + async def test_vertexai_generate_with_functions(self, mock_genai_client): + """Test generate method with functions (tools).""" + mock_client = Mock() + mock_genai_client.return_value = mock_client + + llm = VertexAI(model='gemini-2.5-flash') + + functions = [ + { + 'name': 'test_function', + 'description': 'A test function', + 'parameters': {'type': 'object'}, + } + ] + + # Mock the client response + mock_response = Mock() + mock_response.text = "I'll use the function" + mock_response.candidates = [] + + llm.client = mock_client + llm.client.models.generate_content = Mock(return_value=mock_response) + + messages = [{'role': 'user', 'content': 'Use the function'}] + result = await llm.generate(messages, functions=functions) + + # Verify the result + assert result == {'content': "I'll use the function"} + + @pytest.mark.asyncio + @patch('flo_ai.llm.vertexai_llm.genai.Client') + async def test_vertexai_generate_with_function_call_detection( + self, mock_genai_client + ): + """Test generate method with function call detection.""" + mock_client = Mock() + mock_genai_client.return_value = mock_client + + llm = VertexAI(model='gemini-2.5-flash') + + functions = [ + { + 'name': 'test_function', + 'description': 'A test function', + 'parameters': {'type': 'object'}, + } + ] + + # Mock the client response with function call + mock_function_call = Mock() + mock_function_call.name = 'test_function' + mock_function_call.args = {'param': 'value'} + + mock_part = Mock() + mock_part.function_call = mock_function_call + + mock_content = Mock() + mock_content.parts = [mock_part] + + mock_candidate = Mock() + mock_candidate.content = mock_content + + mock_response = Mock() + mock_response.text = 'Function called' + mock_response.candidates = [mock_candidate] + + llm.client = mock_client + llm.client.models.generate_content = Mock(return_value=mock_response) + + messages = [{'role': 'user', 'content': 'Use the function'}] + result = await llm.generate(messages, functions=functions) + + # Verify function call was detected and parsed + assert 'function_call' in result + assert result['function_call']['name'] == 'test_function' + assert result['function_call']['arguments'] == {'param': 'value'} + + @pytest.mark.asyncio + @patch('flo_ai.llm.vertexai_llm.genai.Client') + async def test_vertexai_generate_with_kwargs(self, mock_genai_client): + """Test generate method with additional kwargs.""" + mock_client = Mock() + mock_genai_client.return_value = mock_client + + llm = VertexAI(model='gemini-2.5-flash', top_p=0.9, max_output_tokens=1000) + + # Mock the client response + mock_response = Mock() + mock_response.text = 'Response with kwargs' + mock_response.candidates = [] + + llm.client = mock_client + llm.client.models.generate_content = Mock(return_value=mock_response) + + messages = [{'role': 'user', 'content': 'Hello'}] + await llm.generate(messages) + + # Verify kwargs were passed through + call_args = llm.client.models.generate_content.call_args + config = call_args[1]['config'] + assert config.top_p == 0.9 + assert config.max_output_tokens == 1000 + + def test_vertexai_get_message_content(self): + """Test get_message_content method.""" + llm = VertexAI() + + # Test with dict response + response = {'content': 'Hello, world!'} + result = llm.get_message_content(response) + assert result == 'Hello, world!' + + # Test with string response + result = llm.get_message_content('Direct string') + assert result == 'Direct string' + + # Test with empty content + response = {'content': ''} + result = llm.get_message_content(response) + assert result == '' + + def test_vertexai_format_tool_for_llm(self): + """Test format_tool_for_llm method.""" + llm = VertexAI() + + # Create a mock tool + tool = Tool( + name='test_tool', + description='A test tool', + function=lambda x: x, + parameters={ + 'param1': {'type': 'string', 'description': 'First parameter'}, + 'param2': {'type': 'integer', 'description': 'Second parameter'}, + }, + ) + + formatted = llm.format_tool_for_llm(tool) + + assert formatted['name'] == 'test_tool' + assert formatted['description'] == 'A test tool' + assert formatted['parameters']['type'] == 'object' + assert 'param1' in formatted['parameters']['properties'] + assert 'param2' in formatted['parameters']['properties'] + assert formatted['parameters']['required'] == ['param1', 'param2'] + + def test_vertexai_format_tools_for_llm(self): + """Test format_tools_for_llm method.""" + llm = VertexAI() + + # Create mock tools + tool1 = Tool( + name='tool1', + description='First tool', + function=lambda x: x, + parameters={'param': {'type': 'string', 'description': 'Parameter'}}, + ) + + tool2 = Tool( + name='tool2', + description='Second tool', + function=lambda x: x, + parameters={'param': {'type': 'integer', 'description': 'Parameter'}}, + ) + + formatted = llm.format_tools_for_llm([tool1, tool2]) + + assert len(formatted) == 2 + assert formatted[0]['name'] == 'tool1' + assert formatted[1]['name'] == 'tool2' + + @patch('flo_ai.llm.vertexai_llm.genai.Client') + def test_vertexai_project_handling(self, mock_genai_client): + """Test project parameter handling.""" + mock_client = Mock() + mock_genai_client.return_value = mock_client + + # Test with project + llm = VertexAI(project='my-project-123') + assert llm.project == 'my-project-123' + + # Test without project + mock_genai_client.reset_mock() + llm = VertexAI() + assert llm.project is None + + # Test with empty string project + mock_genai_client.reset_mock() + llm = VertexAI(project='') + assert llm.project == '' + + @patch('flo_ai.llm.vertexai_llm.genai.Client') + def test_vertexai_location_handling(self, mock_genai_client): + """Test location parameter handling.""" + mock_client = Mock() + mock_genai_client.return_value = mock_client + + # Test with location + llm = VertexAI(location='us-central1') + assert llm.location == 'us-central1' + + # Test without location + mock_genai_client.reset_mock() + llm = VertexAI() + assert llm.location is None + + # Test with empty string location + mock_genai_client.reset_mock() + llm = VertexAI(location='') + assert llm.location == '' + + @patch('flo_ai.llm.vertexai_llm.genai.Client') + def test_vertexai_kwargs_storage(self, mock_genai_client): + """Test that additional kwargs are properly stored.""" + mock_client = Mock() + mock_genai_client.return_value = mock_client + + llm = VertexAI( + max_tokens=1000, top_p=0.9, frequency_penalty=0.1, presence_penalty=0.1 + ) + + assert 'max_tokens' in llm.kwargs + assert 'top_p' in llm.kwargs + assert 'frequency_penalty' in llm.kwargs + assert 'presence_penalty' in llm.kwargs + assert llm.kwargs['max_tokens'] == 1000 + assert llm.kwargs['top_p'] == 0.9 + + @patch('flo_ai.llm.vertexai_llm.genai.Client') + def test_vertexai_base_llm_initialization(self, mock_genai_client): + """Test that BaseLLM is properly initialized.""" + mock_client = Mock() + mock_genai_client.return_value = mock_client + + llm = VertexAI(model='test-model', api_key='test-key', temperature=0.5) + + # These should come from BaseLLM + assert llm.model == 'test-model' + assert llm.api_key == 'test-key' + assert llm.temperature == 0.5 + + @patch('flo_ai.llm.vertexai_llm.genai.Client') + def test_vertexai_vertexai_flag(self, mock_genai_client): + """Test that vertexai=True is always set in client creation.""" + mock_client = Mock() + mock_genai_client.return_value = mock_client + + VertexAI() + + # Verify vertexai=True was passed + call_args = mock_genai_client.call_args[1] + assert call_args['vertexai'] + + @patch('flo_ai.llm.vertexai_llm.genai.Client') + def test_vertexai_parameter_combinations(self, mock_genai_client): + """Test various parameter combinations.""" + mock_client = Mock() + mock_genai_client.return_value = mock_client + + # Test with all parameters + llm = VertexAI( + model='gemini-1.5-pro', + api_key='key123', + temperature=0.3, + project='proj123', + location='us-west1', + max_tokens=500, + top_p=0.8, + ) + + assert llm.model == 'gemini-1.5-pro' + assert llm.api_key == 'key123' + assert llm.temperature == 0.3 + assert llm.project == 'proj123' + assert llm.location == 'us-west1' + assert llm.kwargs == {'max_tokens': 500, 'top_p': 0.8} + + # Test with minimal parameters + mock_genai_client.reset_mock() + llm = VertexAI() + + assert llm.model == 'gemini-2.5-flash' + assert llm.api_key is None + assert llm.temperature == 0.7 + assert llm.project is None + assert llm.location is None + assert llm.kwargs == {} + + @patch('flo_ai.llm.vertexai_llm.genai.Client') + def test_vertexai_method_inheritance(self, mock_genai_client): + """Test that VertexAI inherits all methods from Gemini.""" + mock_client = Mock() + mock_genai_client.return_value = mock_client + + llm = VertexAI() + + # Test that VertexAI has all the methods from Gemini + assert hasattr(llm, 'generate') + assert hasattr(llm, 'get_message_content') + assert hasattr(llm, 'format_tool_for_llm') + assert hasattr(llm, 'format_tools_for_llm') + assert hasattr(llm, 'format_image_in_message') + + @pytest.mark.asyncio + @patch('flo_ai.llm.vertexai_llm.genai.Client') + async def test_vertexai_generate_error_handling(self, mock_genai_client): + """Test error handling in generate method.""" + mock_client = Mock() + mock_genai_client.return_value = mock_client + + llm = VertexAI(model='gemini-2.5-flash') + + # Mock client to raise an exception + llm.client = mock_client + llm.client.models.generate_content = Mock(side_effect=Exception('API Error')) + + messages = [{'role': 'user', 'content': 'Hello'}] + + with pytest.raises(Exception, match='Error in Gemini API call: API Error'): + await llm.generate(messages) + + @patch('flo_ai.llm.vertexai_llm.genai.Client') + def test_vertexai_model_parameter_handling(self, mock_genai_client): + """Test that model parameter is properly handled.""" + mock_client = Mock() + mock_genai_client.return_value = mock_client + + test_models = [ + 'gemini-2.5-flash', + 'gemini-1.5-pro', + 'gemini-1.5-flash', + 'gemini-pro', + ] + + for model in test_models: + mock_genai_client.reset_mock() + llm = VertexAI(model=model) + assert llm.model == model + + @patch('flo_ai.llm.vertexai_llm.genai.Client') + def test_vertexai_api_key_handling(self, mock_genai_client): + """Test API key handling.""" + mock_client = Mock() + mock_genai_client.return_value = mock_client + + # Test with API key + llm = VertexAI(api_key='secret-key-123') + assert llm.api_key == 'secret-key-123' + + # Test without API key + mock_genai_client.reset_mock() + llm = VertexAI() + assert llm.api_key is None + + # Test with empty string API key + mock_genai_client.reset_mock() + llm = VertexAI(api_key='') + assert llm.api_key == '' + + @patch('flo_ai.llm.vertexai_llm.genai.Client') + def test_vertexai_generation_config_creation(self, mock_genai_client): + """Test that generation config is created correctly.""" + mock_client = Mock() + mock_genai_client.return_value = mock_client + + llm = VertexAI(temperature=0.5, max_output_tokens=1000, top_p=0.9) + + # Mock the client response + mock_response = Mock() + mock_response.text = 'Test response' + mock_response.candidates = [] + + llm.client = mock_client + llm.client.models.generate_content = Mock(return_value=mock_response) + + # We need to patch the types.GenerateContentConfig to test this + with patch('flo_ai.llm.gemini_llm.types.GenerateContentConfig') as mock_config: + mock_config_instance = Mock() + mock_config.return_value = mock_config_instance + + # This would normally be called in generate method + # For testing, we'll just verify the config class exists + assert mock_config is not None diff --git a/images/flo-studio-preview.png b/images/flo-studio-preview.png new file mode 100644 index 00000000..0a43992d Binary files /dev/null and b/images/flo-studio-preview.png differ diff --git a/index.html b/index.html deleted file mode 100644 index 473eabe9..00000000 --- a/index.html +++ /dev/null @@ -1,225 +0,0 @@ - - - - - - - - - - -flo-ai | πŸ”₯πŸ”₯πŸ”₯ Simple way to create composable AI agents - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
- -

flo-ai

- - -

- Rootflo -

- -

Composable AI Agentic Workflow

- -

-Rootflo is an alternative to Langgraph, and CrewAI. It lets you easily build composable agentic workflows from using simple components to any size, unlocking the full potential of LLMs. -

- -

- GitHub stars - - GitHub release (latest) - - GitHub commit activity - - License - -
-

- -

-
- Checkout the docs Β» -
-
- Github - β€’ - Website - β€’ - Roadmap -

- -
- -

Flo AI 🌊

- -
-

Build production-ready AI agents and teams with minimal code

-
- -

Flo AI is a Python framework that makes building production-ready AI agents and teams as easy as writing YAML. Think β€œKubernetes for AI Agents” - compose complex AI architectures using pre-built components while maintaining the flexibility to create your own.

- -

✨ Features

- - - -

πŸš€ Quick Start

- -

Installation

- -
pip install flo-ai
-# or using poetry
-poetry add flo-ai
-
- -

Create Your First AI Team in 30 Seconds

- -
from flo_ai import Flo, FloSession
-from langchain_openai import ChatOpenAI
-
-# Define your team in YAML
-yaml_config = """
-apiVersion: flo/alpha-v1
-kind: FloRoutedTeam
-name: research-team
-team:
-    name: ResearchTeam
-    router:
-        name: TeamLead
-        kind: supervisor
-    agents:
-      - name: Researcher
-        role: Research Specialist
-        job: Research latest information on given topics
-        tools:
-          - name: TavilySearchResults
-      - name: Writer
-        role: Content Creator
-        job: Create engaging content from research
-"""
-
-# Set up and run
-llm = ChatOpenAI(temperature=0)
-session = FloSession(llm).register_tool(name="TavilySearchResults", tool=TavilySearchResults())
-flo = Flo.build(session, yaml=yaml_config)
-
-# Start streaming results
-for response in flo.stream("Write about recent AI developments"):
-    print(response)
-
- -

πŸ“– Documentation

- -

Visit our comprehensive documentation for:

- - -

🌟 Why Flo AI?

- -

For AI Engineers

- - -

For Teams

- - -

🎯 Use Cases

- - - -

🀝 Contributing

- -

We love your input! Check out our Contributing Guide to get started. Ways to contribute:

- - - -

πŸ“œ License

- -

Flo AI is MIT Licensed.

- -

πŸ™ Acknowledgments

- -

Built with ❀️ using:

- - -
- -
- Built with ❀️ by the Rootflo team -
Community β€’ - Documentation -
- - - -
- This site is open source. Improve this page. -
- -
- - - - diff --git a/studio/README.md b/studio/README.md new file mode 100644 index 00000000..69ee426c --- /dev/null +++ b/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/studio/package.json b/studio/package.json new file mode 100644 index 00000000..c72e4a1a --- /dev/null +++ b/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/studio/postcss.config.js b/studio/postcss.config.js new file mode 100644 index 00000000..2e7af2b7 --- /dev/null +++ b/studio/postcss.config.js @@ -0,0 +1,6 @@ +export default { + plugins: { + tailwindcss: {}, + autoprefixer: {}, + }, +} diff --git a/studio/src/App.css b/studio/src/App.css new file mode 100644 index 00000000..3baff915 --- /dev/null +++ b/studio/src/App.css @@ -0,0 +1,157 @@ +@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__edge-path { + stroke-width: inherit; +} + +.react-flow__connection-line { + stroke: #3b82f6; + stroke-width: 3px; + stroke-dasharray: 5, 5; + animation: dash 1s linear infinite; +} + +@keyframes dash { + to { + stroke-dashoffset: -10; + } +} + +/* Edge arrows */ +.react-flow__edge .react-flow__edge-path { + stroke-width: inherit; +} + +.react-flow__edge-path { + stroke-linecap: round; +} + +.react-flow__handle { + width: 12px; + height: 12px; + border: 2px solid white; + border-radius: 50%; + transition: all 0.2s ease; +} + +.react-flow__handle:hover { + transform: scale(1.2); + box-shadow: 0 0 0 4px rgba(59, 130, 246, 0.2); +} + +.react-flow__handle-connecting { + background: #3b82f6; + transform: scale(1.3); + box-shadow: 0 0 0 6px rgba(59, 130, 246, 0.3); +} + +.react-flow__handle-valid { + background: #10b981; + transform: scale(1.2); + box-shadow: 0 0 0 4px rgba(16, 185, 129, 0.3); +} + +/* Custom Edge Label Styles */ +.edge-label { + background: white; + border: 1px solid #e5e7eb; + border-radius: 6px; + box-shadow: 0 1px 3px 0 rgba(0, 0, 0, 0.1); + font-size: 12px; + padding: 4px 8px; +} + +.edge-label.selected { + border-color: #3b82f6; + background: #eff6ff; +} + +/* 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); +} + +/* Utility Classes */ +.line-clamp-2 { + display: -webkit-box; + -webkit-line-clamp: 2; + -webkit-box-orient: vertical; + overflow: hidden; +} + +.line-clamp-3 { + display: -webkit-box; + -webkit-line-clamp: 3; + -webkit-box-orient: vertical; + overflow: hidden; +} + +/* YAML Preview Drawer Styles */ +.yaml-drawer { + transition: width 0.3s ease-in-out; +} + +.yaml-drawer-content { + transition: opacity 0.3s ease-in-out; +} + +/* Floating YAML Widget */ +.yaml-widget { + transition: all 0.2s ease-in-out; + animation: float 3s ease-in-out infinite; +} + +.yaml-widget:hover { + transform: translateY(-50%) scale(1.05); +} + +@keyframes float { + 0%, 100% { transform: translateY(-50%) translateX(0); } + 50% { transform: translateY(-50%) translateX(-2px); } +} + +/* Custom Scrollbar for YAML content */ +.yaml-content::-webkit-scrollbar { + width: 8px; +} + +.yaml-content::-webkit-scrollbar-track { + background: #374151; +} + +.yaml-content::-webkit-scrollbar-thumb { + background: #6b7280; + border-radius: 4px; +} + +.yaml-content::-webkit-scrollbar-thumb:hover { + background: #9ca3af; +} diff --git a/studio/src/App.tsx b/studio/src/App.tsx new file mode 100644 index 00000000..a362ad4f --- /dev/null +++ b/studio/src/App.tsx @@ -0,0 +1,110 @@ +import React, { useState } from 'react'; +import { ReactFlowProvider } from 'reactflow'; +import { useDesignerStore } from '@/store/designerStore'; +import { Button } from '@/components/ui/button'; +import { Plus, Settings, Route, Upload, CheckCircle } from 'lucide-react'; +import FlowCanvas from '@/components/flow/FlowCanvas'; +import Sidebar from '@/components/sidebar/Sidebar'; +import AgentEditor from '@/components/editors/AgentEditor'; +import RouterEditor from '@/components/editors/RouterEditor'; +import EdgeEditor from '@/components/editors/EdgeEditor'; +import YamlPreviewDrawer from '@/components/drawer/YamlPreviewDrawer'; +import ImportDialog from '@/components/dialogs/ImportDialog'; +import ValidationPanel from '@/components/panels/ValidationPanel'; +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<{ + showValidation: boolean; + setShowValidation: (show: boolean) => void; +}> = ({ showValidation, setShowValidation }) => { + const { openAgentEditor, openRouterEditor } = useDesignerStore(); + const [isConfigOpen, setIsConfigOpen] = useState(false); + const [isImportOpen, setIsImportOpen] = useState(false); + + return ( + <> +
+
+

Flo AI Studio

+
Visual Workflow Designer
+
+
+ + + + + + +
+
+ setIsConfigOpen(false)} /> + setIsImportOpen(false)} /> + + ); +}; + +function App() { + const [showValidation, setShowValidation] = useState(true); + + return ( +
+ +
+ +
+ + + +
+ {showValidation && } +
+ + {/* Modals */} + + + + + {/* YAML Preview Drawer */} + +
+ ); +} + +export default App; \ No newline at end of file diff --git a/studio/src/components/dialogs/ImportDialog.tsx b/studio/src/components/dialogs/ImportDialog.tsx new file mode 100644 index 00000000..320f4e24 --- /dev/null +++ b/studio/src/components/dialogs/ImportDialog.tsx @@ -0,0 +1,203 @@ +import React, { useState, useRef } from 'react'; +import { + Dialog, + DialogContent, + DialogHeader, + DialogTitle, + DialogFooter, +} from '@/components/ui/dialog'; +import { Button } from '@/components/ui/button'; +import { Textarea } from '@/components/ui/textarea'; +import { Label } from '@/components/ui/label'; +import { Upload, FileText, AlertCircle, CheckCircle } from 'lucide-react'; +import { useDesignerStore } from '@/store/designerStore'; +import { validateAriumYAML, readFileAsText } from '@/utils/yamlImport'; + +interface ImportDialogProps { + isOpen: boolean; + onClose: () => void; +} + +const ImportDialog: React.FC = ({ isOpen, onClose }) => { + const { importFromYAML } = useDesignerStore(); + const [yamlContent, setYamlContent] = useState(''); + const [validationResult, setValidationResult] = useState<{ isValid: boolean; error?: string } | null>(null); + const [isImporting, setIsImporting] = useState(false); + const [importError, setImportError] = useState(null); + const fileInputRef = useRef(null); + + const handleFileUpload = async (event: React.ChangeEvent) => { + const file = event.target.files?.[0]; + if (!file) return; + + try { + const content = await readFileAsText(file); + setYamlContent(content); + setImportError(null); + + // Validate the content + const validation = validateAriumYAML(content); + setValidationResult(validation); + } catch (error) { + setImportError(`Failed to read file: ${error instanceof Error ? error.message : 'Unknown error'}`); + } + }; + + const handleYamlChange = (content: string) => { + setYamlContent(content); + setImportError(null); + + if (content.trim()) { + const validation = validateAriumYAML(content); + setValidationResult(validation); + } else { + setValidationResult(null); + } + }; + + const handleImport = async () => { + if (!yamlContent.trim()) { + setImportError('Please provide YAML content to import'); + return; + } + + setIsImporting(true); + setImportError(null); + + try { + await importFromYAML(yamlContent); + onClose(); + resetDialog(); + } catch (error) { + setImportError(`Import failed: ${error instanceof Error ? error.message : 'Unknown error'}`); + } finally { + setIsImporting(false); + } + }; + + const resetDialog = () => { + setYamlContent(''); + setValidationResult(null); + setImportError(null); + if (fileInputRef.current) { + fileInputRef.current.value = ''; + } + }; + + const handleClose = () => { + resetDialog(); + onClose(); + }; + + return ( + + + + + + Import Workflow from YAML + + + +
+ {/* File Upload */} +
+ +
+ + or paste YAML content below +
+ +
+ + {/* YAML Content */} +
+ +