""" Multi-model orchestration tools for Clover - A terminal assistant for AI-powered project management """ import asyncio import json import os import sys import threading import time from concurrent.futures import ThreadPoolExecutor, as_completed from dataclasses import dataclass from enum import Enum from typing import Any, Dict, List, Optional, Tuple # Add the current directory to Python path for imports sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) from config.settings import load_config from models.api_client import APIClient class TaskType(Enum): """Enum for different task types""" CODE_GENERATION = "code_generation" CODE_REVIEW = "code_review" DOCUMENTATION = "documentation" TESTING = "testing" ANALYSIS = "analysis" SUMMARIZATION = "summarization" DEBUGGING = "debugging" REFACTORING = "refactoring" SECURITY_ANALYSIS = "security_analysis" PERFORMANCE_OPTIMIZATION = "performance_optimization" class ModelCapability(Enum): """Enum for model capabilities""" FAST_RESPONSE = "fast_response" HIGH_QUALITY = "high_quality" CODE_SPECIALIZED = "code_specialized" COST_EFFECTIVE = "cost_effective" LARGE_CONTEXT = "large_context" MULTILINGUAL = "multilingual" @dataclass class ModelProfile: """Profile for a language model with its capabilities and costs""" name: str capabilities: List[ModelCapability] cost_per_1k_tokens: float max_context_length: int avg_response_time: float quality_score: float specializations: List[str] available: bool = True @dataclass class Task: """Represents a task to be executed by a model""" id: str task_type: TaskType prompt: str context: str = "" priority: int = 1 # 1 = high, 2 = medium, 3 = low max_tokens: int = 1000 timeout: int = 300 requires_capabilities: List[ModelCapability] = None callback: callable = None metadata: Dict[str, Any] = None class ModelOrchestrator: """Orchestrate tasks across multiple language models for optimal performance and cost""" def __init__(self): self.config = load_config() self.api_client = APIClient() self.models = {} self.task_queue = [] self.results_cache = {} self.executor = ThreadPoolExecutor(max_workers=self.config.get("threads", 5)) self.lock = threading.Lock() # Initialize model profiles self._initialize_model_profiles() # Task routing rules self.task_routing = { TaskType.CODE_GENERATION: [ ModelCapability.CODE_SPECIALIZED, ModelCapability.HIGH_QUALITY, ], TaskType.CODE_REVIEW: [ ModelCapability.CODE_SPECIALIZED, ModelCapability.HIGH_QUALITY, ], TaskType.DOCUMENTATION: [ ModelCapability.HIGH_QUALITY, ModelCapability.MULTILINGUAL, ], TaskType.TESTING: [ ModelCapability.CODE_SPECIALIZED, ModelCapability.FAST_RESPONSE, ], TaskType.ANALYSIS: [ ModelCapability.HIGH_QUALITY, ModelCapability.LARGE_CONTEXT, ], TaskType.SUMMARIZATION: [ ModelCapability.FAST_RESPONSE, ModelCapability.COST_EFFECTIVE, ], TaskType.DEBUGGING: [ ModelCapability.CODE_SPECIALIZED, ModelCapability.HIGH_QUALITY, ], TaskType.REFACTORING: [ ModelCapability.CODE_SPECIALIZED, ModelCapability.HIGH_QUALITY, ], TaskType.SECURITY_ANALYSIS: [ ModelCapability.CODE_SPECIALIZED, ModelCapability.HIGH_QUALITY, ], TaskType.PERFORMANCE_OPTIMIZATION: [ ModelCapability.CODE_SPECIALIZED, ModelCapability.HIGH_QUALITY, ], } def _initialize_model_profiles(self): """Initialize model profiles with capabilities and characteristics""" # Define common model profiles model_profiles = [ ModelProfile( name="gpt-4", capabilities=[ ModelCapability.HIGH_QUALITY, ModelCapability.CODE_SPECIALIZED, ModelCapability.LARGE_CONTEXT, ], cost_per_1k_tokens=0.03, max_context_length=8192, avg_response_time=3.0, quality_score=0.95, specializations=["general", "code", "analysis"], ), ModelProfile( name="gpt-3.5-turbo", capabilities=[ ModelCapability.FAST_RESPONSE, ModelCapability.COST_EFFECTIVE, ModelCapability.CODE_SPECIALIZED, ], cost_per_1k_tokens=0.002, max_context_length=4096, avg_response_time=1.5, quality_score=0.85, specializations=["general", "code", "summarization"], ), ModelProfile( name="claude-3-opus", capabilities=[ ModelCapability.HIGH_QUALITY, ModelCapability.LARGE_CONTEXT, ModelCapability.MULTILINGUAL, ], cost_per_1k_tokens=0.015, max_context_length=100000, avg_response_time=2.5, quality_score=0.93, specializations=["analysis", "writing", "reasoning"], ), ModelProfile( name="claude-3-sonnet", capabilities=[ ModelCapability.HIGH_QUALITY, ModelCapability.COST_EFFECTIVE, ModelCapability.CODE_SPECIALIZED, ], cost_per_1k_tokens=0.003, max_context_length=100000, avg_response_time=2.0, quality_score=0.90, specializations=["code", "analysis", "general"], ), ModelProfile( name="qwen2.5-coder:7b", capabilities=[ ModelCapability.CODE_SPECIALIZED, ModelCapability.FAST_RESPONSE, ModelCapability.COST_EFFECTIVE, ], cost_per_1k_tokens=0.0, # Assuming local model max_context_length=32768, avg_response_time=1.0, quality_score=0.80, specializations=["code", "debugging", "refactoring"], ), ModelProfile( name="qwen3-coder:30b", capabilities=[ ModelCapability.CODE_SPECIALIZED, ModelCapability.HIGH_QUALITY, ModelCapability.LARGE_CONTEXT, ], cost_per_1k_tokens=0.0, # Assuming local model max_context_length=32768, avg_response_time=2.5, quality_score=0.88, specializations=["code", "architecture", "analysis"], ), ModelProfile( name="llama2-70b", capabilities=[ ModelCapability.HIGH_QUALITY, ModelCapability.LARGE_CONTEXT, ModelCapability.MULTILINGUAL, ], cost_per_1k_tokens=0.0, max_context_length=4096, avg_response_time=3.0, quality_score=0.82, specializations=["general", "reasoning", "analysis"], ), ] # Store models by name for profile in model_profiles: self.models[profile.name] = profile def select_optimal_model( self, task: Task, available_models: List[str] = None ) -> str: """ Select the optimal model for a given task based on requirements and optimization criteria Args: task (Task): Task to be executed available_models (List[str]): List of available model names (None for all) Returns: str: Name of the selected model """ try: # Filter available models candidate_models = {} for name, profile in self.models.items(): if available_models is None or name in available_models: if profile.available: candidate_models[name] = profile if not candidate_models: # Fallback to configured default model return self.config.get("model", "qwen2.5-coder:7b") # Get required capabilities for task type required_caps = task.requires_capabilities or self.task_routing.get( task.task_type, [] ) # Score models based on multiple criteria model_scores = {} for name, profile in candidate_models.items(): score = 0.0 # Capability matching (40% weight) capability_score = 0 if required_caps: matching_caps = len(set(profile.capabilities) & set(required_caps)) capability_score = matching_caps / len(required_caps) else: capability_score = 1.0 # No specific requirements score += capability_score * 0.4 # Quality score (30% weight) score += profile.quality_score * 0.3 # Cost efficiency (15% weight) - lower cost is better max_cost = max(p.cost_per_1k_tokens for p in candidate_models.values()) cost_score = ( 1.0 - (profile.cost_per_1k_tokens / max_cost) if max_cost > 0 else 1.0 ) score += cost_score * 0.15 # Response time (10% weight) - faster is better max_time = max(p.avg_response_time for p in candidate_models.values()) time_score = 1.0 - (profile.avg_response_time / max_time) score += time_score * 0.1 # Context length bonus (5% weight) if len(task.prompt + task.context) > 4000: if profile.max_context_length >= 8000: score += 0.05 model_scores[name] = score # Select model with highest score best_model = max(model_scores, key=model_scores.get) return best_model except Exception as e: print(f"Error in model selection: {e}") return self.config.get("model", "qwen2.5-coder:7b") def estimate_cost(self, task: Task, model_name: str) -> float: """ Estimate the cost of executing a task with a specific model Args: task (Task): Task to estimate cost for model_name (str): Name of the model to use Returns: float: Estimated cost in USD """ try: if model_name not in self.models: return 0.0 profile = self.models[model_name] # Estimate token count (rough approximation: 4 characters per token) input_tokens = len(task.prompt + task.context) / 4 output_tokens = task.max_tokens total_tokens = input_tokens + output_tokens estimated_cost = (total_tokens / 1000) * profile.cost_per_1k_tokens return estimated_cost except Exception as e: return 0.0 def execute_task(self, task: Task, model_name: str = None) -> Dict[str, Any]: """ Execute a single task with the specified or optimal model Args: task (Task): Task to execute model_name (str): Specific model to use (None for auto-selection) Returns: Dict containing execution results """ try: start_time = time.time() # Select model if not specified if model_name is None: model_name = self.select_optimal_model(task) # Estimate cost estimated_cost = self.estimate_cost(task, model_name) # Check cache first cache_key = f"{task.task_type.value}_{hash(task.prompt + task.context)}" if cache_key in self.results_cache: cached_result = self.results_cache[cache_key] cached_result["from_cache"] = True return cached_result # Prepare messages for API call messages = [] if task.context: messages.append({"role": "system", "content": task.context}) messages.append({"role": "user", "content": task.prompt}) # Execute the task response = self.api_client.chat_completion( messages=messages, model=model_name, max_tokens=task.max_tokens ) execution_time = time.time() - start_time # Process response if "error" in response: result = { "task_id": task.id, "success": False, "error": response["error"], "model_used": model_name, "execution_time": execution_time, "estimated_cost": estimated_cost, } else: # Extract response content content = "" if "choices" in response and len(response["choices"]) > 0: content = response["choices"][0]["message"]["content"] result = { "task_id": task.id, "success": True, "response": content, "model_used": model_name, "execution_time": execution_time, "estimated_cost": estimated_cost, "from_cache": False, } # Cache successful results self.results_cache[cache_key] = result.copy() # Call callback if provided if task.callback: try: task.callback(result) except Exception as e: print(f"Error in task callback: {e}") return result except Exception as e: return { "task_id": task.id, "success": False, "error": f"Error executing task: {str(e)}", "model_used": model_name, "execution_time": 0, "estimated_cost": 0, } def execute_batch( self, tasks: List[Task], parallel: bool = True ) -> List[Dict[str, Any]]: """ Execute multiple tasks, optionally in parallel Args: tasks (List[Task]): List of tasks to execute parallel (bool): Whether to execute tasks in parallel Returns: List of execution results """ try: if not parallel: # Sequential execution results = [] for task in tasks: result = self.execute_task(task) results.append(result) return results # Parallel execution results = [None] * len(tasks) # Submit all tasks future_to_index = {} for i, task in enumerate(tasks): model_name = self.select_optimal_model(task) future = self.executor.submit(self.execute_task, task, model_name) future_to_index[future] = i # Collect results as they complete for future in as_completed( future_to_index.keys(), timeout=max(t.timeout for t in tasks) ): index = future_to_index[future] try: result = future.result() results[index] = result except Exception as e: results[index] = { "task_id": tasks[index].id, "success": False, "error": f"Task execution failed: {str(e)}", "model_used": "unknown", "execution_time": 0, "estimated_cost": 0, } return results except Exception as e: # Return error results for all tasks return [ { "task_id": task.id, "success": False, "error": f"Batch execution failed: {str(e)}", "model_used": "unknown", "execution_time": 0, "estimated_cost": 0, } for task in tasks ] def optimize_task_distribution(self, tasks: List[Task]) -> Dict[str, List[Task]]: """ Optimize distribution of tasks across available models Args: tasks (List[Task]): List of tasks to distribute Returns: Dict mapping model names to lists of tasks """ try: distribution = {} # Sort tasks by priority sorted_tasks = sorted(tasks, key=lambda t: t.priority) for task in sorted_tasks: # Select optimal model for this task model_name = self.select_optimal_model(task) if model_name not in distribution: distribution[model_name] = [] distribution[model_name].append(task) return distribution except Exception as e: # Fallback: assign all tasks to default model default_model = self.config.get("model", "qwen2.5-coder:7b") return {default_model: tasks} def get_model_stats(self) -> Dict[str, Any]: """ Get statistics about model usage and performance Returns: Dict containing model statistics """ try: stats = { "available_models": len( [m for m in self.models.values() if m.available] ), "total_models": len(self.models), "cache_size": len(self.results_cache), "model_profiles": {}, } for name, profile in self.models.items(): stats["model_profiles"][name] = { "available": profile.available, "capabilities": [cap.value for cap in profile.capabilities], "cost_per_1k_tokens": profile.cost_per_1k_tokens, "max_context_length": profile.max_context_length, "quality_score": profile.quality_score, "specializations": profile.specializations, } return stats except Exception as e: return {"error": f"Error getting model stats: {str(e)}"} def clear_cache(self): """Clear the results cache""" with self.lock: self.results_cache.clear() def update_model_availability(self, model_name: str, available: bool): """ Update model availability status Args: model_name (str): Name of the model available (bool): Whether the model is available """ if model_name in self.models: self.models[model_name].available = available # Convenience functions for common orchestration tasks def model_selector( task_type: TaskType, prompt: str, context: str = "", **kwargs ) -> str: """ Choose best LLM for specific sub-task based on cost/speed Args: task_type (TaskType): Type of task prompt (str): Task prompt context (str): Additional context **kwargs: Additional task parameters Returns: str: Selected model name """ try: orchestrator = ModelOrchestrator() task = Task( id="selector_task", task_type=task_type, prompt=prompt, context=context, **kwargs, ) return orchestrator.select_optimal_model(task) except Exception as e: print(f"Error in model selection: {e}") config = load_config() return config.get("model", "qwen2.5-coder:7b") def task_orchestrator(tasks: List[Dict[str, Any]]) -> List[Dict[str, Any]]: """ Schedule tools to appropriate model providers Args: tasks (List[Dict]): List of task dictionaries Returns: List of execution results """ try: orchestrator = ModelOrchestrator() # Convert dict tasks to Task objects task_objects = [] for i, task_dict in enumerate(tasks): task = Task( id=task_dict.get("id", f"task_{i}"), task_type=TaskType(task_dict.get("task_type", "analysis")), prompt=task_dict.get("prompt", ""), context=task_dict.get("context", ""), priority=task_dict.get("priority", 2), max_tokens=task_dict.get("max_tokens", 1000), timeout=task_dict.get("timeout", 300), ) task_objects.append(task) return orchestrator.execute_batch(task_objects) except Exception as e: return [{"error": f"Error in task orchestration: {str(e)}"}] def cost_optimizer(tasks: List[Task], budget: float = None) -> Dict[str, Any]: """ Track and optimize API costs across operations Args: tasks (List[Task]): List of tasks to optimize budget (float): Optional budget constraint Returns: Dict containing cost optimization results """ try: orchestrator = ModelOrchestrator() # Calculate costs for different model assignments optimization_results = { "total_tasks": len(tasks), "model_assignments": {}, "total_estimated_cost": 0.0, "budget": budget, "within_budget": True, } total_cost = 0.0 for task in tasks: # Get optimal model for this task optimal_model = orchestrator.select_optimal_model(task) estimated_cost = orchestrator.estimate_cost(task, optimal_model) optimization_results["model_assignments"][task.id] = { "model": optimal_model, "estimated_cost": estimated_cost, } total_cost += estimated_cost optimization_results["total_estimated_cost"] = total_cost if budget is not None: optimization_results["within_budget"] = total_cost <= budget if total_cost > budget: # Try to optimize by using cheaper models print( f"Cost {total_cost:.4f} exceeds budget {budget:.4f}, optimizing..." ) # Re-assign tasks to more cost-effective models adjusted_cost = 0.0 for task in tasks: # Find the most cost-effective model that can handle the task cheapest_model = min( orchestrator.models.keys(), key=lambda m: orchestrator.models[m].cost_per_1k_tokens, ) cost = orchestrator.estimate_cost(task, cheapest_model) optimization_results["model_assignments"][task.id] = { "model": cheapest_model, "estimated_cost": cost, "optimized": True, } adjusted_cost += cost optimization_results["adjusted_cost"] = adjusted_cost optimization_results["cost_savings"] = total_cost - adjusted_cost return optimization_results except Exception as e: return {"error": f"Error in cost optimization: {str(e)}"} # Example usage if __name__ == "__main__": # Example: Create and execute tasks orchestrator = ModelOrchestrator() # Create sample tasks tasks = [ Task( id="code_gen_1", task_type=TaskType.CODE_GENERATION, prompt="Write a Python function to calculate factorial", priority=1, ), Task( id="doc_gen_1", task_type=TaskType.DOCUMENTATION, prompt="Generate documentation for a REST API", priority=2, ), Task( id="analysis_1", task_type=TaskType.ANALYSIS, prompt="Analyze the complexity of this algorithm", context="def bubble_sort(arr): ...", priority=3, ), ] # Execute tasks results = orchestrator.execute_batch(tasks) for result in results: print( f"Task {result['task_id']}: {'Success' if result['success'] else 'Failed'}" ) print(f"Model: {result['model_used']}, Time: {result['execution_time']:.2f}s")