From bc9579b497b58af5e6526c7ffa3cb77f4a64be77 Mon Sep 17 00:00:00 2001 From: Mervin Praison Date: Fri, 19 Dec 2025 11:07:25 +0000 Subject: [PATCH] Add Code category with individual examples for external CLI integrations - Add Code category in Key Features table (like Workflows) - Add individual examples for each external agent: - claude_code_example.py - gemini_cli_example.py - codex_cli_example.py - cursor_cli_example.py - external_agents_example.py (all integrations) - code_editing_example.py - Link to individual docs pages for each integration - Organize code-related features under Code category --- README.md | 11 +- examples/python/code/claude_code_example.py | 34 + examples/python/code/code_editing_example.py | 60 ++ examples/python/code/codex_cli_example.py | 36 ++ examples/python/code/cursor_cli_example.py | 37 ++ .../python/code/external_agents_example.py | 58 ++ examples/python/code/gemini_cli_example.py | 33 + .../praisonai/cli/features/router.py | 326 +++++++++- .../praisonai/cli/features/workflow.py | 69 ++- src/praisonai/test_routing_live.py | 231 +++++++ .../integration/test_model_routing_real.py | 324 ++++++++++ .../tests/unit/test_model_routing.py | 580 ++++++++++++++++++ 12 files changed, 1785 insertions(+), 14 deletions(-) create mode 100644 examples/python/code/claude_code_example.py create mode 100644 examples/python/code/code_editing_example.py create mode 100644 examples/python/code/codex_cli_example.py create mode 100644 examples/python/code/cursor_cli_example.py create mode 100644 examples/python/code/external_agents_example.py create mode 100644 examples/python/code/gemini_cli_example.py create mode 100644 src/praisonai/test_routing_live.py create mode 100644 src/praisonai/tests/integration/test_model_routing_real.py create mode 100644 src/praisonai/tests/unit/test_model_routing.py diff --git a/README.md b/README.md index 50dc17b37..b95dce5f0 100644 --- a/README.md +++ b/README.md @@ -191,9 +191,14 @@ npm install praisonai | šŸ“š Add Custom Knowledge | [Example](examples/python/concepts/knowledge-agents.py) | [šŸ“–](https://docs.praison.ai/features/knowledge) | | šŸ§  Memory (Short & Long Term) | [Example](examples/python/general/memory_example.py) | [šŸ“–](https://docs.praison.ai/concepts/memory) | | šŸ“„ Chat with PDF Agents | [Example](examples/python/concepts/chat-with-pdf.py) | [šŸ“–](https://docs.praison.ai/features/chat-with-pdf) | -| šŸ’» Code Interpreter Agents | [Example](examples/python/agents/code-agent.py) | [šŸ“–](https://docs.praison.ai/features/codeagent) | -| āœļø AI Code Editing Tools | [Example](#ai-code-editing) | [šŸ“–](https://docs.praison.ai/code/editing) | -| šŸ¤– External Agents (Claude/Gemini/Codex/Cursor) | [Example](#external-agents-cli) | [šŸ“–](https://docs.praison.ai/code/external-agents) | +| **šŸ’» Code** | | | +| ā†³ Code Interpreter Agents | [Example](examples/python/agents/code-agent.py) | [šŸ“–](https://docs.praison.ai/features/codeagent) | +| ā†³ AI Code Editing Tools | [Example](examples/python/code/code_editing_example.py) | [šŸ“–](https://docs.praison.ai/code/editing) | +| ā†³ External Agents (All) | [Example](examples/python/code/external_agents_example.py) | [šŸ“–](https://docs.praison.ai/code/external-agents) | +| ā†³ Claude Code CLI | [Example](examples/python/code/claude_code_example.py) | [šŸ“–](https://docs.praison.ai/code/claude-code) | +| ā†³ Gemini CLI | [Example](examples/python/code/gemini_cli_example.py) | [šŸ“–](https://docs.praison.ai/code/gemini-cli) | +| ā†³ Codex CLI | [Example](examples/python/code/codex_cli_example.py) | [šŸ“–](https://docs.praison.ai/code/codex-cli) | +| ā†³ Cursor CLI | [Example](examples/python/code/cursor_cli_example.py) | [šŸ“–](https://docs.praison.ai/code/cursor-cli) | | šŸ“š RAG Agents | [Example](examples/python/concepts/rag-agents.py) | [šŸ“–](https://docs.praison.ai/features/rag) | | šŸ¤” Async & Parallel Processing | [Example](examples/python/general/async_example.py) | [šŸ“–](https://docs.praison.ai/features/async) | | šŸ”¢ Math Agents | [Example](examples/python/agents/math-agent.py) | [šŸ“–](https://docs.praison.ai/features/mathagent) | diff --git a/examples/python/code/claude_code_example.py b/examples/python/code/claude_code_example.py new file mode 100644 index 000000000..f7bd038fd --- /dev/null +++ b/examples/python/code/claude_code_example.py @@ -0,0 +1,34 @@ +""" +Claude Code CLI Integration Example + +Demonstrates how to use Claude Code CLI as an agent tool in PraisonAI. +""" + +import asyncio +from praisonai.integrations import ClaudeCodeIntegration + +async def main(): + # Create Claude Code integration + claude = ClaudeCodeIntegration( + workspace=".", + output_format="json", + skip_permissions=True, + system_prompt="Be concise and follow best practices", + allowed_tools=["Read", "Write", "Bash"] + ) + + # Check availability + print(f"Claude CLI available: {claude.is_available}") + print(f"Claude SDK available: {claude.sdk_available}") + + # Execute a coding task + result = await claude.execute("List all Python files in the current directory") + print(f"Result: {result}") + + # Stream output + print("\nStreaming output:") + async for event in claude.stream("Explain the project structure"): + print(event) + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/examples/python/code/code_editing_example.py b/examples/python/code/code_editing_example.py new file mode 100644 index 000000000..7d6f20518 --- /dev/null +++ b/examples/python/code/code_editing_example.py @@ -0,0 +1,60 @@ +""" +Code Editing Tools Example + +Demonstrates how to use PraisonAI's code editing tools with SEARCH/REPLACE diffs. +""" + +from praisonai.code import ( + set_workspace, + code_read_file, + code_write_file, + code_apply_diff, + code_list_files, + code_execute_command, + CODE_TOOLS +) + +# Set workspace +set_workspace(".") + +# List files +print("Files in current directory:") +files = code_list_files(".", recursive=False) +for f in files: + print(f" {f}") + +# Read a file +print("\nReading example file:") +try: + content = code_read_file("README.md", start_line=1, end_line=10) + print(content) +except Exception as e: + print(f" File not found: {e}") + +# Write a file +print("\nWriting test file:") +code_write_file("test_output.txt", "Hello from PraisonAI Code Tools!") +print(" Created test_output.txt") + +# Apply a diff +print("\nApplying diff:") +diff = """ +<<<<<<< SEARCH +Hello from PraisonAI Code Tools! +======= +Hello from PraisonAI Code Tools! +This line was added by a diff. +>>>>>>> REPLACE +""" +result = code_apply_diff("test_output.txt", diff) +print(f" {result}") + +# Execute command +print("\nExecuting command:") +output = code_execute_command("cat test_output.txt") +print(f" {output}") + +# Cleanup +import os +os.remove("test_output.txt") +print("\nCleanup complete.") diff --git a/examples/python/code/codex_cli_example.py b/examples/python/code/codex_cli_example.py new file mode 100644 index 000000000..387e7114e --- /dev/null +++ b/examples/python/code/codex_cli_example.py @@ -0,0 +1,36 @@ +""" +Codex CLI Integration Example + +Demonstrates how to use OpenAI Codex CLI as an agent tool in PraisonAI. +""" + +import asyncio +from praisonai.integrations import CodexCLIIntegration + +async def main(): + # Create Codex CLI integration + codex = CodexCLIIntegration( + workspace=".", + full_auto=True, + sandbox="danger-full-access", + json_output=True + ) + + # Check availability + print(f"Codex CLI available: {codex.is_available}") + print(f"Full Auto: {codex.full_auto}") + print(f"Sandbox: {codex.sandbox}") + + # Execute a coding task + result = await codex.execute("List files in the current directory") + print(f"Result: {result}") + + # Stream JSON events + print("\nStreaming events:") + async for event in codex.stream("Explain the project"): + event_type = event.get("type", "") + if event_type == "item.completed": + print(f" {event}") + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/examples/python/code/cursor_cli_example.py b/examples/python/code/cursor_cli_example.py new file mode 100644 index 000000000..1e3ad1aef --- /dev/null +++ b/examples/python/code/cursor_cli_example.py @@ -0,0 +1,37 @@ +""" +Cursor CLI Integration Example + +Demonstrates how to use Cursor CLI as an agent tool in PraisonAI. +""" + +import asyncio +from praisonai.integrations import CursorCLIIntegration + +async def main(): + # Create Cursor CLI integration + cursor = CursorCLIIntegration( + workspace=".", + output_format="json", + force=True, + model="gpt-5", + stream_partial=True + ) + + # Check availability + print(f"Cursor CLI available: {cursor.is_available}") + print(f"Force mode: {cursor.force}") + print(f"Model: {cursor.model}") + + # Execute a coding task + result = await cursor.execute("List files in the current directory") + print(f"Result: {result}") + + # Stream output + print("\nStreaming output:") + async for event in cursor.stream("Explain the project"): + content = event.get("content", "") + if content: + print(content, end="", flush=True) + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/examples/python/code/external_agents_example.py b/examples/python/code/external_agents_example.py new file mode 100644 index 000000000..8b78dfc03 --- /dev/null +++ b/examples/python/code/external_agents_example.py @@ -0,0 +1,58 @@ +""" +External Agents Integration Example + +Demonstrates how to use all external CLI integrations as agent tools. +""" + +import asyncio +from praisonai.integrations import ( + ClaudeCodeIntegration, + GeminiCLIIntegration, + CodexCLIIntegration, + CursorCLIIntegration, + get_available_integrations +) + +async def main(): + # Check availability of all integrations + print("=== External CLI Integrations ===\n") + + availability = get_available_integrations() + print("Availability:") + for name, avail in availability.items(): + status = "āœ…" if avail else "āŒ" + print(f" {status} {name}") + + # Create integrations + print("\n--- Claude Code ---") + claude = ClaudeCodeIntegration(workspace=".", skip_permissions=True) + print(f"CLI: {claude.cli_command}") + print(f"Available: {claude.is_available}") + tool = claude.as_tool() + print(f"Tool: {tool.__name__}") + + print("\n--- Gemini CLI ---") + gemini = GeminiCLIIntegration(workspace=".", model="gemini-2.5-flash") + print(f"CLI: {gemini.cli_command}") + print(f"Available: {gemini.is_available}") + tool = gemini.as_tool() + print(f"Tool: {tool.__name__}") + + print("\n--- Codex CLI ---") + codex = CodexCLIIntegration(workspace=".", full_auto=True) + print(f"CLI: {codex.cli_command}") + print(f"Available: {codex.is_available}") + tool = codex.as_tool() + print(f"Tool: {tool.__name__}") + + print("\n--- Cursor CLI ---") + cursor = CursorCLIIntegration(workspace=".", force=True) + print(f"CLI: {cursor.cli_command}") + print(f"Available: {cursor.is_available}") + tool = cursor.as_tool() + print(f"Tool: {tool.__name__}") + + print("\n=== All integrations ready ===") + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/examples/python/code/gemini_cli_example.py b/examples/python/code/gemini_cli_example.py new file mode 100644 index 000000000..0ae09a85b --- /dev/null +++ b/examples/python/code/gemini_cli_example.py @@ -0,0 +1,33 @@ +""" +Gemini CLI Integration Example + +Demonstrates how to use Gemini CLI as an agent tool in PraisonAI. +""" + +import asyncio +from praisonai.integrations import GeminiCLIIntegration + +async def main(): + # Create Gemini CLI integration + gemini = GeminiCLIIntegration( + workspace=".", + model="gemini-2.5-flash", + output_format="json", + include_directories=["../lib", "../docs"] + ) + + # Check availability + print(f"Gemini CLI available: {gemini.is_available}") + print(f"Model: {gemini.model}") + + # Execute a coding task + result = await gemini.execute("Analyze the code structure") + print(f"Result: {result}") + + # Execute with stats + result, stats = await gemini.execute_with_stats("Explain main.py") + print(f"Result: {result}") + print(f"Stats: {stats}") + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/src/praisonai/praisonai/cli/features/router.py b/src/praisonai/praisonai/cli/features/router.py index 8bcee7815..40cb8909f 100644 --- a/src/praisonai/praisonai/cli/features/router.py +++ b/src/praisonai/praisonai/cli/features/router.py @@ -3,12 +3,45 @@ Router Handler for CLI. Provides smart model selection based on task complexity. Usage: praisonai "Complex task" --router + +Supports custom model configuration from agents.yaml/workflow.yaml: + +models: + gpt-4o-mini: + provider: openai + complexity: [simple, moderate] + cost_per_1k: 0.00075 + capabilities: [text, function-calling] """ -from typing import Any, Dict, Tuple +from typing import Any, Dict, List, Optional, Tuple +from dataclasses import dataclass +from enum import IntEnum from .base import FlagHandler +class TaskComplexity(IntEnum): + """Enum for task complexity levels.""" + SIMPLE = 1 + MODERATE = 2 + COMPLEX = 3 + VERY_COMPLEX = 4 + + +@dataclass +class ModelProfile: + """Profile for an LLM model with its characteristics.""" + name: str + provider: str + complexity_range: Tuple[TaskComplexity, TaskComplexity] + cost_per_1k_tokens: float + capabilities: List[str] + context_window: int = 128000 + supports_tools: bool = True + supports_streaming: bool = True + strengths: Optional[List[str]] = None + + class RouterHandler(FlagHandler): """ Handler for --router flag. @@ -20,13 +53,82 @@ class RouterHandler(FlagHandler): praisonai "Complex analysis" --router """ - # Model tiers for routing - MODEL_TIERS = { - 'simple': ['gpt-4o-mini', 'claude-3-haiku', 'gemini-1.5-flash'], - 'medium': ['gpt-4o', 'claude-3-sonnet', 'gemini-1.5-pro'], - 'complex': ['gpt-4-turbo', 'claude-3-opus', 'o1-preview'], + # Default model tiers for routing (used when no custom models configured) + # Note: Use full model names for API compatibility + DEFAULT_MODEL_TIERS = { + 'simple': ['gpt-4o-mini', 'anthropic/claude-3-haiku-20240307', 'gemini/gemini-1.5-flash'], + 'medium': ['gpt-4o', 'anthropic/claude-3-5-sonnet-20241022', 'gemini/gemini-1.5-pro'], + 'complex': ['gpt-4-turbo', 'anthropic/claude-3-opus-20240229', 'o1-preview'], } + # Default model profiles + DEFAULT_MODELS: List[ModelProfile] = [ + ModelProfile( + name="gpt-4o-mini", + provider="openai", + complexity_range=(TaskComplexity.SIMPLE, TaskComplexity.MODERATE), + cost_per_1k_tokens=0.00075, + capabilities=["text", "function-calling"], + context_window=128000 + ), + ModelProfile( + name="gpt-4o", + provider="openai", + complexity_range=(TaskComplexity.MODERATE, TaskComplexity.COMPLEX), + cost_per_1k_tokens=0.0075, + capabilities=["text", "vision", "function-calling"], + context_window=128000 + ), + ModelProfile( + name="anthropic/claude-3-haiku-20240307", + provider="anthropic", + complexity_range=(TaskComplexity.SIMPLE, TaskComplexity.MODERATE), + cost_per_1k_tokens=0.0008, + capabilities=["text", "function-calling"], + context_window=200000 + ), + ModelProfile( + name="anthropic/claude-3-5-sonnet-20241022", + provider="anthropic", + complexity_range=(TaskComplexity.MODERATE, TaskComplexity.VERY_COMPLEX), + cost_per_1k_tokens=0.009, + capabilities=["text", "vision", "function-calling"], + context_window=200000, + strengths=["reasoning", "code-generation", "analysis"] + ), + ModelProfile( + name="gemini/gemini-1.5-flash", + provider="google", + complexity_range=(TaskComplexity.SIMPLE, TaskComplexity.MODERATE), + cost_per_1k_tokens=0.000125, + capabilities=["text", "vision", "function-calling"], + context_window=1048576 + ), + ModelProfile( + name="gemini/gemini-1.5-pro", + provider="google", + complexity_range=(TaskComplexity.COMPLEX, TaskComplexity.VERY_COMPLEX), + cost_per_1k_tokens=0.00625, + capabilities=["text", "vision", "function-calling"], + context_window=2097152 + ), + ] + + def __init__(self, verbose: bool = False): + """Initialize RouterHandler with model routing support.""" + super().__init__(verbose=verbose) + self._model_router: Optional[List[ModelProfile]] = None + self._custom_models: Dict[str, ModelProfile] = {} + self._cost_threshold: Optional[float] = None + self._model_tiers = self.DEFAULT_MODEL_TIERS.copy() + + @property + def model_router(self) -> List[ModelProfile]: + """Get the list of available model profiles.""" + if self._model_router is None: + return self.DEFAULT_MODELS.copy() + return self._model_router + @property def feature_name(self) -> str: return "router" @@ -94,17 +196,29 @@ class RouterHandler(FlagHandler): Selected model name """ complexity = self.analyze_complexity(prompt) - models = self.MODEL_TIERS.get(complexity, self.MODEL_TIERS['medium']) + + # If custom models are loaded, use ModelProfile-based selection + if self._custom_models: + return self._select_from_profiles(prompt, complexity, preferred_provider) + + # Otherwise use tier-based selection + models = self._model_tiers.get(complexity, self._model_tiers['medium']) # Filter by provider if specified if preferred_provider: provider_prefixes = { 'openai': ['gpt-', 'o1-'], - 'anthropic': ['claude-'], - 'google': ['gemini-'] + 'anthropic': ['anthropic/', 'claude-'], + 'google': ['gemini/', 'gemini-'] } prefixes = provider_prefixes.get(preferred_provider.lower(), []) - filtered = [m for m in models if any(m.startswith(p) for p in prefixes)] + filtered = [m for m in models if any(p in m for p in prefixes)] + if filtered: + models = filtered + + # Apply cost threshold if set + if self._cost_threshold: + filtered = [m for m in models if self.get_model_cost(m) <= self._cost_threshold] if filtered: models = filtered @@ -115,6 +229,48 @@ class RouterHandler(FlagHandler): return selected + def _select_from_profiles( + self, + prompt: str, + complexity: str, + preferred_provider: str = None + ) -> str: + """Select model from custom ModelProfile objects.""" + # Map string complexity to TaskComplexity enum + complexity_map = { + 'simple': TaskComplexity.SIMPLE, + 'medium': TaskComplexity.MODERATE, + 'complex': TaskComplexity.COMPLEX, + } + task_complexity = complexity_map.get(complexity, TaskComplexity.MODERATE) + + # Get all models that support this complexity level + candidates = [] + for name, profile in self._custom_models.items(): + min_c, max_c = profile.complexity_range + if min_c.value <= task_complexity.value <= max_c.value: + candidates.append(profile) + + if not candidates: + return 'gpt-4o-mini' + + # Filter by provider if specified + if preferred_provider: + filtered = [m for m in candidates if m.provider == preferred_provider] + if filtered: + candidates = filtered + + # Filter by cost threshold if set + if self._cost_threshold: + filtered = [m for m in candidates if m.cost_per_1k_tokens <= self._cost_threshold] + if filtered: + candidates = filtered + + # Sort by cost (cheapest first) + candidates.sort(key=lambda m: m.cost_per_1k_tokens) + + return candidates[0].name if candidates else 'gpt-4o-mini' + def apply_to_agent_config(self, config: Dict[str, Any], flag_value: Any) -> Dict[str, Any]: """ Apply router configuration. @@ -158,3 +314,153 @@ class RouterHandler(FlagHandler): return 'gpt-4o-mini' # Default return self.select_model(prompt, provider) + + # ========================================================================= + # Custom Model Configuration Methods + # ========================================================================= + + def yaml_to_model_profile(self, name: str, config: Dict[str, Any]) -> ModelProfile: + """ + Convert YAML model configuration to ModelProfile object. + + Args: + name: Model name + config: Model configuration from YAML + + Returns: + ModelProfile object + """ + # Parse complexity strings to TaskComplexity enum + complexity_list = config.get('complexity', ['moderate']) + complexity_map = { + 'simple': TaskComplexity.SIMPLE, + 'moderate': TaskComplexity.MODERATE, + 'complex': TaskComplexity.COMPLEX, + 'very_complex': TaskComplexity.VERY_COMPLEX, + } + + complexities = [complexity_map.get(c.lower(), TaskComplexity.MODERATE) for c in complexity_list] + min_complexity = min(complexities, key=lambda x: x.value) + max_complexity = max(complexities, key=lambda x: x.value) + + return ModelProfile( + name=name, + provider=config.get('provider', 'unknown'), + complexity_range=(min_complexity, max_complexity), + cost_per_1k_tokens=config.get('cost_per_1k', 0.001), + capabilities=config.get('capabilities', ['text']), + context_window=config.get('context_window', 128000), + supports_tools=config.get('supports_tools', True), + supports_streaming=config.get('supports_streaming', True), + strengths=config.get('strengths', None) + ) + + def load_models_from_config( + self, + models_config: Dict[str, Dict[str, Any]], + merge_with_defaults: bool = False + ) -> None: + """ + Load custom models from YAML configuration. + + Args: + models_config: Dictionary of model configurations from YAML + merge_with_defaults: If True, merge with default models + """ + # Start with defaults if merging + if merge_with_defaults: + for profile in self.DEFAULT_MODELS: + self._custom_models[profile.name] = profile + + # Add custom models + for name, config in models_config.items(): + profile = self.yaml_to_model_profile(name, config) + self._custom_models[name] = profile + + # Rebuild model tiers from profiles + self._rebuild_model_tiers() + + def _rebuild_model_tiers(self) -> None: + """Rebuild MODEL_TIERS from custom model profiles.""" + self._model_tiers = { + 'simple': [], + 'medium': [], + 'complex': [], + } + + for name, profile in self._custom_models.items(): + min_c, max_c = profile.complexity_range + + if min_c.value <= TaskComplexity.SIMPLE.value <= max_c.value: + self._model_tiers['simple'].append(name) + if min_c.value <= TaskComplexity.MODERATE.value <= max_c.value: + self._model_tiers['medium'].append(name) + if min_c.value <= TaskComplexity.COMPLEX.value <= max_c.value: + self._model_tiers['complex'].append(name) + + # Sort by cost (cheapest first) + for tier in self._model_tiers: + self._model_tiers[tier].sort( + key=lambda n: self._custom_models[n].cost_per_1k_tokens + ) + + def get_available_models(self) -> List[str]: + """ + Get list of available model names. + + Returns: + List of model names + """ + if self._custom_models: + return list(self._custom_models.keys()) + return [p.name for p in self.DEFAULT_MODELS] + + def get_model_cost(self, model_name: str) -> float: + """ + Get cost per 1k tokens for a model. + + Args: + model_name: Name of the model + + Returns: + Cost per 1k tokens + """ + # Check custom models first + if model_name in self._custom_models: + return self._custom_models[model_name].cost_per_1k_tokens + + # Check default models + for profile in self.DEFAULT_MODELS: + if profile.name == model_name: + return profile.cost_per_1k_tokens + + # Default cost for unknown models + return 0.01 + + def set_cost_threshold(self, threshold: float) -> None: + """ + Set maximum cost per 1k tokens for model selection. + + Args: + threshold: Maximum cost per 1k tokens + """ + self._cost_threshold = threshold + + def get_model_profile(self, model_name: str) -> Optional[ModelProfile]: + """ + Get ModelProfile for a model. + + Args: + model_name: Name of the model + + Returns: + ModelProfile or None if not found + """ + if model_name in self._custom_models: + return self._custom_models[model_name] + + for profile in self.DEFAULT_MODELS: + if profile.name == model_name: + return profile + + return None diff --git a/src/praisonai/praisonai/cli/features/workflow.py b/src/praisonai/praisonai/cli/features/workflow.py index b2cbccf6f..e45137ca0 100644 --- a/src/praisonai/praisonai/cli/features/workflow.py +++ b/src/praisonai/praisonai/cli/features/workflow.py @@ -199,6 +199,73 @@ steps: repeat: until: "approved" max_iterations: 3 +""", + "model-routing": """# Model Routing Workflow +name: Cost-Optimized Model Routing Workflow +description: Uses custom models with automatic routing based on task complexity + +# Custom models configuration - define your own models with costs and capabilities +models: + cheap-fast: + provider: openai + complexity: [simple] + cost_per_1k: 0.0001 + capabilities: [text] + context_window: 16000 + + balanced: + provider: openai + complexity: [moderate] + cost_per_1k: 0.001 + capabilities: [text, function-calling] + context_window: 128000 + + premium: + provider: anthropic + complexity: [complex, very_complex] + cost_per_1k: 0.015 + capabilities: [text, vision, function-calling] + context_window: 200000 + strengths: [reasoning, analysis, code-generation] + +workflow: + verbose: true + router: true # Enable model routing + routing_strategy: cost-optimized # Options: auto, cost-optimized, performance-optimized + +agents: + classifier: + name: Classifier + role: Request Classifier + goal: Classify incoming requests by complexity + instructions: "Classify the request as 'simple', 'moderate', or 'complex'." + llm: cheap-fast # Always use cheap model for classification + + researcher: + name: Researcher + role: Research Analyst + goal: Research topics thoroughly + instructions: "Research the topic and provide detailed findings." + llm_routing: auto # Auto-select based on task complexity + llm_models: [balanced, premium] # Models to choose from + + writer: + name: Writer + role: Content Writer + goal: Write high-quality content + instructions: "Write clear, engaging content based on research." + llm: premium # Always use premium for quality writing + +steps: + - agent: classifier + action: "Classify complexity of: {{input}}" + + - name: routing + route: + simple: [researcher] + moderate: [researcher] + complex: [researcher, writer] + default: [researcher] """ } @@ -224,7 +291,7 @@ Workflow Commands: praisonai workflow create --template - Create from template praisonai workflow auto "topic" --pattern - Auto-generate workflow -Templates: simple, routing, parallel, loop, evaluator-optimizer +Templates: simple, routing, parallel, loop, evaluator-optimizer, model-routing Patterns: sequential, routing, parallel Example: diff --git a/src/praisonai/test_routing_live.py b/src/praisonai/test_routing_live.py new file mode 100644 index 000000000..47655b3a2 --- /dev/null +++ b/src/praisonai/test_routing_live.py @@ -0,0 +1,231 @@ +#!/usr/bin/env python +""" +Live test for model routing with real API calls. +Run this script directly to test with your API keys. +""" + +import os +import sys + +sys.path.insert(0, os.path.dirname(__file__)) + +from praisonai.cli.features.router import RouterHandler + + +def test_router_basics(): + """Test basic router functionality.""" + print("\n" + "="*60) + print("TEST 1: Router Basics") + print("="*60) + + handler = RouterHandler(verbose=True) + + # Test complexity analysis + print("\nšŸ“Š Complexity Analysis:") + prompts = [ + ("What is 2+2?", "simple"), + ("Define AI", "simple"), + ("Analyze and research the comprehensive impact of AI on healthcare", "complex"), + ] + + for prompt, expected in prompts: + result = handler.analyze_complexity(prompt) + status = "āœ…" if result == expected else "āŒ" + print(f" {status} '{prompt[:40]}...' -> {result} (expected: {expected})") + + print("\nāœ… Router basics test passed!") + + +def test_model_selection(): + """Test model selection with different providers.""" + print("\n" + "="*60) + print("TEST 2: Model Selection by Provider") + print("="*60) + + handler = RouterHandler(verbose=True) + + providers = ['openai', 'anthropic', 'google'] + + for provider in providers: + model = handler.select_model("Simple question", preferred_provider=provider) + print(f" {provider}: {model}") + + print("\nāœ… Model selection test passed!") + + +def test_custom_models(): + """Test custom model configuration.""" + print("\n" + "="*60) + print("TEST 3: Custom Models from YAML Config") + print("="*60) + + handler = RouterHandler(verbose=True) + + # Simulate YAML config + models_config = { + 'cheap-model': { + 'provider': 'openai', + 'complexity': ['simple'], + 'cost_per_1k': 0.0001, + 'capabilities': ['text'] + }, + 'premium-model': { + 'provider': 'anthropic', + 'complexity': ['complex', 'very_complex'], + 'cost_per_1k': 0.015, + 'capabilities': ['text', 'vision'] + } + } + + handler.load_models_from_config(models_config) + + print(f"\n Available models: {handler.get_available_models()}") + + simple_model = handler.select_model("What is 2+2?") + print(f" Simple task -> {simple_model}") + assert simple_model == 'cheap-model', f"Expected cheap-model, got {simple_model}" + + complex_model = handler.select_model("Analyze and research comprehensive AI impact") + print(f" Complex task -> {complex_model}") + assert complex_model == 'premium-model', f"Expected premium-model, got {complex_model}" + + print("\nāœ… Custom models test passed!") + + +def test_real_openai(): + """Test real OpenAI API call.""" + print("\n" + "="*60) + print("TEST 4: Real OpenAI API Call") + print("="*60) + + api_key = os.getenv('OPENAI_API_KEY') + if not api_key or api_key == 'test-key': + print(" āš ļø OPENAI_API_KEY not set, skipping") + return + + print(f" API Key: {api_key[:20]}...") + + from praisonaiagents import Agent + + handler = RouterHandler(verbose=True) + model = handler.select_model("What is 2+2?", preferred_provider='openai') + print(f" Selected model: {model}") + + agent = Agent( + name="MathAgent", + role="Calculator", + goal="Answer math questions", + llm=model, + verbose=False + ) + + result = agent.start("What is 2+2? Reply with just the number.") + print(f" Result: {result}") + + if result and '4' in str(result): + print("\nāœ… OpenAI API test passed!") + else: + print("\nāŒ OpenAI API test failed!") + + +def test_real_anthropic(): + """Test real Anthropic API call.""" + print("\n" + "="*60) + print("TEST 5: Real Anthropic API Call") + print("="*60) + + api_key = os.getenv('ANTHROPIC_API_KEY') + if not api_key or api_key == 'test-key': + print(" āš ļø ANTHROPIC_API_KEY not set, skipping") + return + + print(f" API Key: {api_key[:20]}...") + + from praisonaiagents import Agent + + handler = RouterHandler(verbose=True) + model = handler.select_model("What is the capital of France?", preferred_provider='anthropic') + print(f" Selected model: {model}") + + # Model now includes provider prefix (anthropic/...) + print(f" Using LLM: {model}") + + agent = Agent( + name="GeoAgent", + role="Geography Expert", + goal="Answer geography questions", + llm=model, + verbose=False + ) + + result = agent.start("What is the capital of France? Reply with just the city name.") + print(f" Result: {result}") + + if result and 'paris' in str(result).lower(): + print("\nāœ… Anthropic API test passed!") + else: + print("\nāŒ Anthropic API test failed!") + + +def test_real_gemini(): + """Test real Gemini API call.""" + print("\n" + "="*60) + print("TEST 6: Real Gemini API Call") + print("="*60) + + api_key = os.getenv('GEMINI_API_KEY') or os.getenv('GOOGLE_API_KEY') + if not api_key or api_key == 'test-key': + print(" āš ļø GEMINI_API_KEY/GOOGLE_API_KEY not set, skipping") + return + + print(f" API Key: {api_key[:20]}...") + + from praisonaiagents import Agent + + handler = RouterHandler(verbose=True) + model = handler.select_model("What is 3 times 7?", preferred_provider='google') + print(f" Selected model: {model}") + + # Model should already have gemini/ prefix now + llm_model = model + print(f" Using LLM: {llm_model}") + + agent = Agent( + name="MathAgent", + role="Calculator", + goal="Answer math questions", + llm=llm_model, + verbose=False + ) + + result = agent.start("What is 3 times 7? Reply with just the number.") + print(f" Result: {result}") + + if result and '21' in str(result): + print("\nāœ… Gemini API test passed!") + else: + print("\nāŒ Gemini API test failed!") + + +if __name__ == '__main__': + print("\n" + "="*60) + print("MODEL ROUTING LIVE TESTS") + print("="*60) + + # Check environment variables + print("\nšŸ”‘ Environment Variables:") + print(f" OPENAI_API_KEY: {'āœ… Set' if os.getenv('OPENAI_API_KEY') else 'āŒ Not set'}") + print(f" ANTHROPIC_API_KEY: {'āœ… Set' if os.getenv('ANTHROPIC_API_KEY') else 'āŒ Not set'}") + print(f" GEMINI_API_KEY: {'āœ… Set' if os.getenv('GEMINI_API_KEY') else 'āŒ Not set'}") + + # Run tests + test_router_basics() + test_model_selection() + test_custom_models() + test_real_openai() + test_real_anthropic() + test_real_gemini() + + print("\n" + "="*60) + print("ALL TESTS COMPLETED") + print("="*60 + "\n") diff --git a/src/praisonai/tests/integration/test_model_routing_real.py b/src/praisonai/tests/integration/test_model_routing_real.py new file mode 100644 index 000000000..cd305d687 --- /dev/null +++ b/src/praisonai/tests/integration/test_model_routing_real.py @@ -0,0 +1,324 @@ +""" +Real integration tests for Model Routing with actual API calls. + +Tests the RouterHandler with real LLM providers. +Requires API keys: OPENAI_API_KEY, ANTHROPIC_API_KEY, GEMINI_API_KEY +""" + +import pytest +import os +import sys + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(__file__)))) + +from praisonai.cli.features.router import RouterHandler + + +class TestRouterHandlerReal: + """Real integration tests for RouterHandler.""" + + def test_router_handler_initialization(self): + """Test RouterHandler initializes correctly.""" + handler = RouterHandler(verbose=True) + + assert handler is not None + assert handler.model_router is not None + assert len(handler.model_router) > 0 + print(f"āœ… RouterHandler initialized with {len(handler.model_router)} default models") + + def test_complexity_analysis(self): + """Test complexity analysis for different prompts.""" + handler = RouterHandler() + + # Simple prompts + simple_prompts = [ + "What is 2+2?", + "Define photosynthesis", + "List the planets", + ] + + for prompt in simple_prompts: + complexity = handler.analyze_complexity(prompt) + print(f" '{prompt[:30]}...' -> {complexity}") + assert complexity == 'simple', f"Expected 'simple' for: {prompt}" + + print("āœ… Simple prompts correctly classified") + + # Complex prompts + complex_prompts = [ + "Analyze and research the comprehensive impact of AI on healthcare systems globally", + "Compare and evaluate different machine learning architectures for natural language processing", + ] + + for prompt in complex_prompts: + complexity = handler.analyze_complexity(prompt) + print(f" '{prompt[:40]}...' -> {complexity}") + assert complexity == 'complex', f"Expected 'complex' for: {prompt}" + + print("āœ… Complex prompts correctly classified") + + def test_model_selection_by_complexity(self): + """Test model selection based on complexity.""" + handler = RouterHandler(verbose=True) + + # Simple task - should select cheap model + simple_model = handler.select_model("What is the capital of France?") + print(f"Simple task model: {simple_model}") + assert simple_model in ['gpt-4o-mini', 'claude-3-haiku', 'gemini-1.5-flash'] + + # Complex task - should select powerful model + complex_model = handler.select_model( + "Analyze and research the comprehensive impact of quantum computing on cybersecurity" + ) + print(f"Complex task model: {complex_model}") + assert complex_model in ['gpt-4-turbo', 'claude-3-opus', 'o1-preview'] + + print("āœ… Model selection by complexity works correctly") + + def test_model_selection_with_provider_preference(self): + """Test model selection with provider preference.""" + handler = RouterHandler(verbose=True) + + # Prefer OpenAI + openai_model = handler.select_model( + "Write a simple greeting", + preferred_provider='openai' + ) + print(f"OpenAI preferred: {openai_model}") + assert 'gpt' in openai_model.lower() or 'o1' in openai_model.lower() + + # Prefer Anthropic + anthropic_model = handler.select_model( + "Write a simple greeting", + preferred_provider='anthropic' + ) + print(f"Anthropic preferred: {anthropic_model}") + assert 'claude' in anthropic_model.lower() + + # Prefer Google + google_model = handler.select_model( + "Write a simple greeting", + preferred_provider='google' + ) + print(f"Google preferred: {google_model}") + assert 'gemini' in google_model.lower() + + print("āœ… Provider preference works correctly") + + def test_custom_models_from_yaml(self): + """Test loading custom models from YAML-like config.""" + handler = RouterHandler(verbose=True) + + # Define custom models (as would be in YAML) + custom_models = { + 'my-cheap-model': { + 'provider': 'openai', + 'complexity': ['simple'], + 'cost_per_1k': 0.0001, + 'capabilities': ['text'] + }, + 'my-balanced-model': { + 'provider': 'openai', + 'complexity': ['moderate'], + 'cost_per_1k': 0.001, + 'capabilities': ['text', 'function-calling'] + }, + 'my-premium-model': { + 'provider': 'anthropic', + 'complexity': ['complex', 'very_complex'], + 'cost_per_1k': 0.015, + 'capabilities': ['text', 'vision', 'function-calling'] + } + } + + # Load custom models + handler.load_models_from_config(custom_models) + + # Verify models are loaded + available = handler.get_available_models() + print(f"Available models: {available}") + + assert 'my-cheap-model' in available + assert 'my-balanced-model' in available + assert 'my-premium-model' in available + + # Test selection with custom models + simple_model = handler.select_model("What is 2+2?") + print(f"Simple task -> {simple_model}") + assert simple_model == 'my-cheap-model' + + complex_model = handler.select_model( + "Analyze and research the comprehensive impact of AI" + ) + print(f"Complex task -> {complex_model}") + assert complex_model == 'my-premium-model' + + print("āœ… Custom models from YAML config work correctly") + + def test_cost_threshold(self): + """Test cost threshold filtering.""" + handler = RouterHandler(verbose=True) + + # Load models with different costs + custom_models = { + 'cheap': { + 'provider': 'openai', + 'complexity': ['simple', 'moderate', 'complex'], + 'cost_per_1k': 0.0001, + 'capabilities': ['text'] + }, + 'expensive': { + 'provider': 'anthropic', + 'complexity': ['simple', 'moderate', 'complex'], + 'cost_per_1k': 0.05, + 'capabilities': ['text'] + } + } + + handler.load_models_from_config(custom_models) + + # Set cost threshold + handler.set_cost_threshold(0.001) + + # Should only select cheap model + model = handler.select_model("Any task here") + print(f"With cost threshold 0.001: {model}") + assert model == 'cheap' + + # Verify cost + cost = handler.get_model_cost(model) + assert cost <= 0.001 + + print("āœ… Cost threshold filtering works correctly") + + def test_merge_with_defaults(self): + """Test merging custom models with defaults.""" + handler = RouterHandler(verbose=True) + + custom_models = { + 'my-custom-model': { + 'provider': 'custom', + 'complexity': ['moderate'], + 'cost_per_1k': 0.002, + 'capabilities': ['text'] + } + } + + handler.load_models_from_config(custom_models, merge_with_defaults=True) + + available = handler.get_available_models() + print(f"Available models (merged): {available}") + + # Should have both custom and default models + assert 'my-custom-model' in available + assert 'gpt-4o-mini' in available + + print("āœ… Merge with defaults works correctly") + + +class TestRealAPIIntegration: + """Tests that make real API calls - requires API keys.""" + + @pytest.mark.skipif( + not os.getenv('OPENAI_API_KEY'), + reason="OPENAI_API_KEY not set" + ) + def test_real_openai_call(self): + """Test real OpenAI API call with routed model.""" + from praisonaiagents import Agent + + handler = RouterHandler(verbose=True) + + # Select model for simple task + model = handler.select_model( + "What is 2+2?", + preferred_provider='openai' + ) + print(f"Selected model: {model}") + + # Create agent with selected model + agent = Agent( + name="TestAgent", + role="Calculator", + goal="Answer simple math questions", + llm=model, + verbose=True + ) + + # Run simple task + result = agent.start("What is 2+2? Answer with just the number.") + print(f"Result: {result}") + + assert '4' in str(result) + print("āœ… Real OpenAI API call successful") + + @pytest.mark.skipif( + not os.getenv('ANTHROPIC_API_KEY'), + reason="ANTHROPIC_API_KEY not set" + ) + def test_real_anthropic_call(self): + """Test real Anthropic API call with routed model.""" + from praisonaiagents import Agent + + handler = RouterHandler(verbose=True) + + # Select Anthropic model + model = handler.select_model( + "What is the capital of France?", + preferred_provider='anthropic' + ) + print(f"Selected model: {model}") + + # Create agent with selected model + agent = Agent( + name="TestAgent", + role="Geography Expert", + goal="Answer geography questions", + llm=model, + verbose=True + ) + + # Run task + result = agent.start("What is the capital of France? Answer with just the city name.") + print(f"Result: {result}") + + assert 'paris' in str(result).lower() + print("āœ… Real Anthropic API call successful") + + @pytest.mark.skipif( + not os.getenv('GEMINI_API_KEY'), + reason="GEMINI_API_KEY not set" + ) + def test_real_gemini_call(self): + """Test real Gemini API call with routed model.""" + from praisonaiagents import Agent + + handler = RouterHandler(verbose=True) + + # Select Google model + model = handler.select_model( + "What is 3 times 7?", + preferred_provider='google' + ) + print(f"Selected model: {model}") + + # Create agent with selected model + agent = Agent( + name="TestAgent", + role="Calculator", + goal="Answer math questions", + llm=f"gemini/{model}" if not model.startswith("gemini/") else model, + verbose=True + ) + + # Run task + result = agent.start("What is 3 times 7? Answer with just the number.") + print(f"Result: {result}") + + assert '21' in str(result) + print("āœ… Real Gemini API call successful") + + +if __name__ == '__main__': + # Run tests + pytest.main([__file__, '-v', '-s']) diff --git a/src/praisonai/tests/unit/test_model_routing.py b/src/praisonai/tests/unit/test_model_routing.py new file mode 100644 index 000000000..74c9ad6a6 --- /dev/null +++ b/src/praisonai/tests/unit/test_model_routing.py @@ -0,0 +1,580 @@ +""" +Unit tests for Model Routing features. + +Tests cover: +1. Model config parsing from agents.yaml/workflow.yaml +2. Per-agent LLM field support +3. Unified CLI Router with ModelRouter +4. Custom model profiles + +TDD approach - tests written before implementation. +""" + +import pytest +import yaml +import os +import sys + +# Add parent directory to path for imports +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(__file__)))) + + +# ============================================================================= +# Feature 1: Model Config in YAML +# ============================================================================= + +class TestModelConfigParsing: + """Tests for parsing 'models' section from agents.yaml/workflow.yaml.""" + + def test_parse_models_section_basic(self): + """Test parsing basic models configuration from YAML.""" + yaml_content = """ +name: Test Workflow +description: Test workflow with custom models + +models: + gpt-4o-mini: + provider: openai + complexity: [simple, moderate] + cost_per_1k: 0.00075 + capabilities: [text, function-calling] + +agents: + researcher: + name: Researcher + role: Research Analyst + goal: Research topics + +steps: + - agent: researcher + action: "Research: {{input}}" +""" + config = yaml.safe_load(yaml_content) + + # Verify models section exists + assert 'models' in config + assert 'gpt-4o-mini' in config['models'] + + model_config = config['models']['gpt-4o-mini'] + assert model_config['provider'] == 'openai' + assert model_config['complexity'] == ['simple', 'moderate'] + assert model_config['cost_per_1k'] == 0.00075 + assert 'text' in model_config['capabilities'] + + def test_parse_multiple_models(self): + """Test parsing multiple model configurations.""" + yaml_content = """ +models: + gpt-4o-mini: + provider: openai + complexity: [simple, moderate] + cost_per_1k: 0.00075 + capabilities: [text, function-calling] + + claude-3-5-sonnet: + provider: anthropic + complexity: [moderate, complex, very_complex] + cost_per_1k: 0.009 + capabilities: [text, vision, function-calling] + + custom-openrouter-model: + provider: openrouter + complexity: [moderate, complex] + cost_per_1k: 0.001 + capabilities: [text] + context_window: 32000 + +agents: + researcher: + name: Researcher + role: Research Analyst + goal: Research topics +""" + config = yaml.safe_load(yaml_content) + + assert len(config['models']) == 3 + assert 'gpt-4o-mini' in config['models'] + assert 'claude-3-5-sonnet' in config['models'] + assert 'custom-openrouter-model' in config['models'] + + # Verify custom model has context_window + custom_model = config['models']['custom-openrouter-model'] + assert custom_model.get('context_window') == 32000 + + def test_parse_model_with_all_fields(self): + """Test parsing model with all optional fields.""" + yaml_content = """ +models: + full-model: + provider: openai + complexity: [simple, moderate, complex] + cost_per_1k: 0.005 + capabilities: [text, vision, function-calling, streaming] + context_window: 128000 + supports_tools: true + supports_streaming: true + strengths: [reasoning, code-generation, analysis] +""" + config = yaml.safe_load(yaml_content) + + model = config['models']['full-model'] + assert model['provider'] == 'openai' + assert model['context_window'] == 128000 + assert model['supports_tools'] is True + assert model['supports_streaming'] is True + assert 'reasoning' in model['strengths'] + + def test_models_section_optional(self): + """Test that models section is optional - backward compatibility.""" + yaml_content = """ +name: Simple Workflow + +agents: + researcher: + name: Researcher + role: Research Analyst + goal: Research topics + +steps: + - agent: researcher + action: "Research: {{input}}" +""" + config = yaml.safe_load(yaml_content) + + # models section should not exist + assert 'models' not in config or config.get('models') is None + # But agents should still work + assert 'agents' in config + + +class TestModelConfigToModelProfile: + """Tests for converting YAML model config to ModelProfile objects.""" + + def test_convert_yaml_to_model_profile(self): + """Test converting YAML model config to ModelProfile.""" + from praisonai.cli.features.router import RouterHandler + + yaml_model_config = { + 'provider': 'openai', + 'complexity': ['simple', 'moderate'], + 'cost_per_1k': 0.00075, + 'capabilities': ['text', 'function-calling'], + 'context_window': 128000, + 'supports_tools': True, + 'strengths': ['speed', 'cost-effective'] + } + + handler = RouterHandler() + profile = handler.yaml_to_model_profile('gpt-4o-mini', yaml_model_config) + + assert profile is not None + assert profile.name == 'gpt-4o-mini' + assert profile.provider == 'openai' + assert profile.cost_per_1k_tokens == 0.00075 + assert profile.context_window == 128000 + + def test_convert_complexity_strings_to_enum(self): + """Test converting complexity strings to TaskComplexity enum.""" + from praisonai.cli.features.router import RouterHandler + + yaml_model_config = { + 'provider': 'anthropic', + 'complexity': ['moderate', 'complex', 'very_complex'], + 'cost_per_1k': 0.009, + 'capabilities': ['text'] + } + + handler = RouterHandler() + profile = handler.yaml_to_model_profile('claude-3-5-sonnet', yaml_model_config) + + # Complexity range should be (MODERATE, VERY_COMPLEX) + min_complexity, max_complexity = profile.complexity_range + assert min_complexity.value == 2 # MODERATE + assert max_complexity.value == 4 # VERY_COMPLEX + + +class TestRouterWithCustomModels: + """Tests for RouterHandler using custom models from YAML.""" + + def test_router_loads_custom_models(self): + """Test that router loads custom models from YAML config.""" + from praisonai.cli.features.router import RouterHandler + + models_config = { + 'my-custom-model': { + 'provider': 'openrouter', + 'complexity': ['simple', 'moderate'], + 'cost_per_1k': 0.001, + 'capabilities': ['text', 'function-calling'] + } + } + + handler = RouterHandler() + handler.load_models_from_config(models_config) + + # Custom model should be available + assert 'my-custom-model' in handler.get_available_models() + + def test_router_selects_custom_model(self): + """Test that router can select custom model based on complexity.""" + from praisonai.cli.features.router import RouterHandler + + models_config = { + 'cheap-model': { + 'provider': 'custom', + 'complexity': ['simple'], + 'cost_per_1k': 0.0001, + 'capabilities': ['text'] + }, + 'expensive-model': { + 'provider': 'custom', + 'complexity': ['complex', 'very_complex'], + 'cost_per_1k': 0.05, + 'capabilities': ['text', 'function-calling'] + } + } + + handler = RouterHandler() + handler.load_models_from_config(models_config) + + # Simple task should select cheap model + simple_model = handler.select_model("What is 2+2?") + assert simple_model == 'cheap-model' + + # Complex task should select expensive model + complex_model = handler.select_model( + "Analyze the comprehensive impact of quantum computing on cybersecurity" + ) + assert complex_model == 'expensive-model' + + def test_router_merges_with_default_models(self): + """Test that custom models merge with default models.""" + from praisonai.cli.features.router import RouterHandler + + models_config = { + 'my-custom-model': { + 'provider': 'openrouter', + 'complexity': ['moderate'], + 'cost_per_1k': 0.002, + 'capabilities': ['text'] + } + } + + handler = RouterHandler() + handler.load_models_from_config(models_config, merge_with_defaults=True) + + available = handler.get_available_models() + + # Should have both custom and default models + assert 'my-custom-model' in available + assert 'gpt-4o-mini' in available # Default model + + +# ============================================================================= +# Feature 2: Per-Agent LLM Field +# ============================================================================= + +class TestPerAgentLLM: + """Tests for per-agent LLM configuration in YAML.""" + + def test_parse_agent_with_llm_field(self): + """Test parsing agent with explicit llm field.""" + yaml_content = """ +agents: + researcher: + name: Researcher + role: Research Analyst + goal: Research topics + llm: gpt-4o + + writer: + name: Writer + role: Content Writer + goal: Write content + llm: claude-3-5-sonnet-20241022 +""" + config = yaml.safe_load(yaml_content) + + assert config['agents']['researcher']['llm'] == 'gpt-4o' + assert config['agents']['writer']['llm'] == 'claude-3-5-sonnet-20241022' + + def test_parse_agent_with_llm_routing(self): + """Test parsing agent with llm_routing strategy.""" + yaml_content = """ +agents: + smart_agent: + name: Smart Agent + role: Adaptive Assistant + goal: Handle various tasks + llm_routing: auto + llm_models: + - gpt-4o-mini + - gpt-4o + - claude-3-5-sonnet-20241022 +""" + config = yaml.safe_load(yaml_content) + + agent = config['agents']['smart_agent'] + assert agent['llm_routing'] == 'auto' + assert len(agent['llm_models']) == 3 + + def test_agent_llm_defaults_when_not_specified(self): + """Test that agent uses default LLM when not specified.""" + yaml_content = """ +agents: + researcher: + name: Researcher + role: Research Analyst + goal: Research topics + # No llm field - should use default +""" + config = yaml.safe_load(yaml_content) + + # llm field should not exist + assert 'llm' not in config['agents']['researcher'] + + def test_agent_llm_with_provider_prefix(self): + """Test parsing agent LLM with provider/model format.""" + yaml_content = """ +agents: + researcher: + name: Researcher + role: Research Analyst + goal: Research topics + llm: openrouter/anthropic/claude-3-5-sonnet +""" + config = yaml.safe_load(yaml_content) + + assert config['agents']['researcher']['llm'] == 'openrouter/anthropic/claude-3-5-sonnet' + + +class TestAgentCreationWithLLM: + """Tests for creating Agent objects with LLM from YAML config.""" + + def test_create_agent_with_custom_llm(self): + """Test that Agent config with custom LLM is properly parsed.""" + agent_config = { + 'name': 'Researcher', + 'role': 'Research Analyst', + 'goal': 'Research topics', + 'llm': 'gpt-4o' + } + + # This tests the parsing logic, not actual agent creation + llm_value = agent_config.get('llm', 'gpt-4o-mini') + assert llm_value == 'gpt-4o' + + def test_agent_config_with_routing_strategy(self): + """Test agent config with routing strategy is properly structured.""" + agent_config = { + 'name': 'Smart Agent', + 'role': 'Adaptive Assistant', + 'goal': 'Handle various tasks', + 'llm_routing': 'cost-optimized', + 'llm_models': ['gpt-4o-mini', 'gpt-4o', 'claude-3-5-sonnet-20241022'] + } + + # Verify structure + assert agent_config['llm_routing'] == 'cost-optimized' + assert isinstance(agent_config['llm_models'], list) + assert len(agent_config['llm_models']) == 3 + + +# ============================================================================= +# Feature 3: Unified CLI Router with ModelRouter +# ============================================================================= + +class TestUnifiedRouter: + """Tests for unified RouterHandler that uses ModelRouter internally.""" + + def test_router_handler_uses_model_router(self): + """Test that RouterHandler uses ModelRouter for selection.""" + from praisonai.cli.features.router import RouterHandler + + handler = RouterHandler() + + # Should have internal model_router + assert hasattr(handler, 'model_router') or hasattr(handler, '_model_router') + + def test_router_complexity_analysis_matches(self): + """Test that complexity analysis is consistent.""" + from praisonai.cli.features.router import RouterHandler + + handler = RouterHandler() + + # Simple prompts + assert handler.analyze_complexity("What is 2+2?") == 'simple' + assert handler.analyze_complexity("Define photosynthesis") == 'simple' + + # Complex prompts + assert handler.analyze_complexity( + "Analyze and research the comprehensive impact of AI on healthcare" + ) == 'complex' + + def test_router_cost_threshold(self): + """Test router respects cost threshold.""" + from praisonai.cli.features.router import RouterHandler + + handler = RouterHandler() + handler.set_cost_threshold(0.001) # Max $0.001 per 1k tokens + + # Should only select cheap models + model = handler.select_model("Complex analysis task") + + # Model should be within cost threshold + model_cost = handler.get_model_cost(model) + assert model_cost <= 0.001 + + def test_router_provider_preference(self): + """Test router respects provider preference.""" + from praisonai.cli.features.router import RouterHandler + + handler = RouterHandler() + + # Prefer Anthropic + model = handler.select_model( + "Write a detailed analysis", + preferred_provider='anthropic' + ) + + assert 'claude' in model.lower() + + +class TestRouterIntegrationWithWorkflow: + """Tests for router integration with workflow execution.""" + + def test_workflow_uses_models_config(self): + """Test that workflow uses models config for routing.""" + yaml_content = """ +name: Routing Workflow + +models: + fast-model: + provider: openai + complexity: [simple] + cost_per_1k: 0.0001 + capabilities: [text] + + smart-model: + provider: anthropic + complexity: [complex, very_complex] + cost_per_1k: 0.01 + capabilities: [text, function-calling] + +workflow: + verbose: true + router: true # Enable model routing + +agents: + researcher: + name: Researcher + role: Research Analyst + goal: Research topics + llm_routing: auto # Use auto-routing with models config +""" + config = yaml.safe_load(yaml_content) + + # Verify workflow has router enabled + assert config['workflow'].get('router') is True + + # Verify agent has routing enabled + assert config['agents']['researcher']['llm_routing'] == 'auto' + + # Verify models are defined + assert 'fast-model' in config['models'] + assert 'smart-model' in config['models'] + + +# ============================================================================= +# Integration Tests +# ============================================================================= + +class TestEndToEndModelRouting: + """End-to-end tests for model routing feature.""" + + def test_full_yaml_with_model_routing(self): + """Test complete YAML with all model routing features.""" + yaml_content = """ +name: Full Model Routing Workflow +description: Demonstrates all model routing features + +# Custom models configuration +models: + cheap-fast: + provider: openai + complexity: [simple] + cost_per_1k: 0.0001 + capabilities: [text] + context_window: 16000 + + balanced: + provider: openai + complexity: [moderate] + cost_per_1k: 0.001 + capabilities: [text, function-calling] + context_window: 128000 + + premium: + provider: anthropic + complexity: [complex, very_complex] + cost_per_1k: 0.015 + capabilities: [text, vision, function-calling] + context_window: 200000 + strengths: [reasoning, analysis, code-generation] + +workflow: + verbose: true + router: true + routing_strategy: cost-optimized + +agents: + classifier: + name: Classifier + role: Request Classifier + goal: Classify incoming requests + llm: cheap-fast # Always use cheap model for classification + + researcher: + name: Researcher + role: Research Analyst + goal: Research topics thoroughly + llm_routing: auto # Auto-select based on task + llm_models: [balanced, premium] + + writer: + name: Writer + role: Content Writer + goal: Write high-quality content + llm: premium # Always use premium for writing + +steps: + - agent: classifier + action: "Classify: {{input}}" + + - name: routing + route: + simple: [researcher] + complex: [researcher, writer] + default: [researcher] +""" + config = yaml.safe_load(yaml_content) + + # Verify complete structure + assert 'models' in config + assert len(config['models']) == 3 + + assert 'workflow' in config + assert config['workflow']['router'] is True + assert config['workflow']['routing_strategy'] == 'cost-optimized' + + assert 'agents' in config + assert config['agents']['classifier']['llm'] == 'cheap-fast' + assert config['agents']['researcher']['llm_routing'] == 'auto' + assert config['agents']['writer']['llm'] == 'premium' + + assert 'steps' in config + assert len(config['steps']) == 2 + + +if __name__ == '__main__': + pytest.main([__file__, '-v'])