392bcfa2ef
Major enhancements to Clover CLI: ✨ New Features: - AI agent with multi-turn conversation capabilities - Tool calling system with 11+ tools for file operations, Git, linting, etc. - Step-by-step AI assistance with play-by-play commentary - Enhanced interactive mode with better UX 🔧 Core Components Added: - ai_agent.py: Main AI agent with conversation management - models/: API client and model management system - Comprehensive tool system for development tasks 🛠️ Tools Available: - File operations (create, read, update, delete) - Command execution with safety checks - Git operations (status, diff, commit, push) - Code linting and formatting - Project structure analysis - Security scanning and dependency management 💡 User Experience: - Real-time tool execution summaries - File creation with full path visibility - Error handling and retry mechanisms - Clean conversation flow until task completion 🧹 Repository Cleanup: - Added comprehensive .gitignore - Removed __pycache__ directories and build artifacts - Organized project structure The AI can now actually create files, run commands, and work through complex development tasks step-by-step with full transparency.
764 lines
25 KiB
Python
764 lines
25 KiB
Python
"""
|
|
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")
|