379 lines
14 KiB
Python
379 lines
14 KiB
Python
|
#!/usr/bin/env python3
|
||
|
|
"""
|
||
|
|
Task routing system for local-first AI fallbacks.
|
||
|
|
Issue #483: [AUDIT][RISK] Maintain local-first fallbacks for all cloud AI
|
||
|
|
"""
|
||
|
|
import json
|
||
|
|
import time
|
||
|
|
import argparse
|
||
|
|
from typing import Dict, List, Any, Optional
|
||
|
|
from dataclasses import dataclass
|
||
|
|
from enum import Enum
|
||
|
|
import requests
|
||
|
|
|
||
|
|
class TaskType(Enum):
|
||
|
|
"""Types of tasks that can be routed."""
|
||
|
|
CODE_GENERATION = "code_generation"
|
||
|
|
WEB_SEARCH = "web_search"
|
||
|
|
DOCUMENT_ANALYSIS = "document_analysis"
|
||
|
|
CREATIVE_WRITING = "creative_writing"
|
||
|
|
DATA_ANALYSIS = "data_analysis"
|
||
|
|
QUESTION_ANSWERING = "question_answering"
|
||
|
|
SUMMARIZATION = "summarization"
|
||
|
|
|
||
|
|
class Priority(Enum):
|
||
|
|
"""Routing priority."""
|
||
|
|
LOCAL_FIRST = "local-first"
|
||
|
|
QUALITY_FIRST = "quality-first"
|
||
|
|
BALANCED = "balanced"
|
||
|
|
COST_FIRST = "cost-first"
|
||
|
|
|
||
|
|
@dataclass
|
||
|
|
class TaskResult:
|
||
|
|
"""Result from a task execution."""
|
||
|
|
provider: str
|
||
|
|
model: str
|
||
|
|
result: str
|
||
|
|
quality_score: float
|
||
|
|
execution_time: float
|
||
|
|
cost: float
|
||
|
|
timestamp: str
|
||
|
|
|
||
|
|
class LocalModel:
|
||
|
|
"""Interface to local Ollama models."""
|
||
|
|
|
||
|
|
def __init__(self, name: str, endpoint: str = "http://localhost:11434"):
|
||
|
|
self.name = name
|
||
|
|
self.endpoint = endpoint
|
||
|
|
self.available = self._check_availability()
|
||
|
|
|
||
|
|
def _check_availability(self) -> bool:
|
||
|
|
"""Check if model is available."""
|
||
|
|
try:
|
||
|
|
response = requests.get(f"{self.endpoint}/api/tags", timeout=5)
|
||
|
|
if response.status_code == 200:
|
||
|
|
models = [m["name"] for m in response.json().get("models", [])]
|
||
|
|
return self.name in models
|
||
|
|
except:
|
||
|
|
pass
|
||
|
|
return False
|
||
|
|
|
||
|
|
def execute(self, prompt: str, max_tokens: int = 100) -> Optional[str]:
|
||
|
|
"""Execute a prompt on the local model."""
|
||
|
|
if not self.available:
|
||
|
|
return None
|
||
|
|
|
||
|
|
try:
|
||
|
|
payload = {
|
||
|
|
"model": self.name,
|
||
|
|
"prompt": prompt,
|
||
|
|
"stream": False,
|
||
|
|
"options": {"num_predict": max_tokens}
|
||
|
|
}
|
||
|
|
|
||
|
|
response = requests.post(
|
||
|
|
f"{self.endpoint}/api/generate",
|
||
|
|
json=payload,
|
||
|
|
timeout=30
|
||
|
|
)
|
||
|
|
|
||
|
|
if response.status_code == 200:
|
||
|
|
return response.json().get("response", "")
|
||
|
|
except Exception as e:
|
||
|
|
print(f"Error executing local model {self.name}: {e}")
|
||
|
|
|
||
|
|
return None
|
||
|
|
|
||
|
|
class ExternalService:
|
||
|
|
"""Interface to external AI services."""
|
||
|
|
|
||
|
|
def __init__(self, name: str, api_key: Optional[str] = None):
|
||
|
|
self.name = name
|
||
|
|
self.api_key = api_key
|
||
|
|
self.available = self._check_availability()
|
||
|
|
|
||
|
|
def _check_availability(self) -> bool:
|
||
|
|
"""Check if service is available."""
|
||
|
|
# Simplified check - in reality would test actual API
|
||
|
|
if self.name == "perplexity":
|
||
|
|
return bool(self.api_key)
|
||
|
|
return False
|
||
|
|
|
||
|
|
def execute(self, prompt: str) -> Optional[str]:
|
||
|
|
"""Execute a prompt on the external service."""
|
||
|
|
if not self.available:
|
||
|
|
return None
|
||
|
|
|
||
|
|
# Simplified implementation
|
||
|
|
# In reality, would call actual API
|
||
|
|
return f"External {self.name} response to: {prompt[:50]}..."
|
||
|
|
|
||
|
|
class TaskRouter:
|
||
|
|
"""Routes tasks to appropriate providers based on priority."""
|
||
|
|
|
||
|
|
def __init__(self):
|
||
|
|
self.local_models = {
|
||
|
|
"hermes4": LocalModel("hermes4"),
|
||
|
|
"llama3-8b": LocalModel("llama3-8b"),
|
||
|
|
"mistral-7b": LocalModel("mistral-7b")
|
||
|
|
}
|
||
|
|
|
||
|
|
self.external_services = {
|
||
|
|
"perplexity": ExternalService("perplexity")
|
||
|
|
}
|
||
|
|
|
||
|
|
# Task capability matrix
|
||
|
|
self.capabilities = {
|
||
|
|
TaskType.CODE_GENERATION: {
|
||
|
|
"local": ["hermes4", "llama3-8b"],
|
||
|
|
"external": [],
|
||
|
|
"quality_local": 0.8,
|
||
|
|
"quality_external": 0.9
|
||
|
|
},
|
||
|
|
TaskType.WEB_SEARCH: {
|
||
|
|
"local": [],
|
||
|
|
"external": ["perplexity"],
|
||
|
|
"quality_local": 0.3,
|
||
|
|
"quality_external": 0.95
|
||
|
|
},
|
||
|
|
TaskType.DOCUMENT_ANALYSIS: {
|
||
|
|
"local": ["hermes4", "llama3-8b", "mistral-7b"],
|
||
|
|
"external": [],
|
||
|
|
"quality_local": 0.85,
|
||
|
|
"quality_external": 0.9
|
||
|
|
},
|
||
|
|
TaskType.CREATIVE_WRITING: {
|
||
|
|
"local": ["hermes4", "mistral-7b"],
|
||
|
|
"external": [],
|
||
|
|
"quality_local": 0.75,
|
||
|
|
"quality_external": 0.85
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
def route_task(self, task_type: TaskType, prompt: str, priority: Priority = Priority.BALANCED) -> Optional[TaskResult]:
|
||
|
|
"""Route a task based on priority."""
|
||
|
|
print(f"Routing {task_type.value} with {priority.value} priority...")
|
||
|
|
|
||
|
|
if priority == Priority.LOCAL_FIRST:
|
||
|
|
return self._route_local_first(task_type, prompt)
|
||
|
|
elif priority == Priority.QUALITY_FIRST:
|
||
|
|
return self._route_quality_first(task_type, prompt)
|
||
|
|
elif priority == Priority.COST_FIRST:
|
||
|
|
return self._route_cost_first(task_type, prompt)
|
||
|
|
else: # BALANCED
|
||
|
|
return self._route_balanced(task_type, prompt)
|
||
|
|
|
||
|
|
def _route_local_first(self, task_type: TaskType, prompt: str) -> Optional[TaskResult]:
|
||
|
|
"""Try local models first, fallback to external."""
|
||
|
|
print(" Trying local models first...")
|
||
|
|
|
||
|
|
# Try local models
|
||
|
|
for model_name in self.capabilities.get(task_type, {}).get("local", []):
|
||
|
|
model = self.local_models.get(model_name)
|
||
|
|
if model and model.available:
|
||
|
|
print(f" Trying {model_name}...")
|
||
|
|
start_time = time.time()
|
||
|
|
result = model.execute(prompt)
|
||
|
|
exec_time = time.time() - start_time
|
||
|
|
|
||
|
|
if result:
|
||
|
|
return TaskResult(
|
||
|
|
provider="local",
|
||
|
|
model=model_name,
|
||
|
|
result=result,
|
||
|
|
quality_score=self.capabilities[task_type]["quality_local"],
|
||
|
|
execution_time=exec_time,
|
||
|
|
cost=0.0,
|
||
|
|
timestamp=time.strftime("%Y-%m-%dT%H:%M:%SZ")
|
||
|
|
)
|
||
|
|
|
||
|
|
# Fallback to external
|
||
|
|
print(" Local failed, trying external...")
|
||
|
|
for service_name in self.capabilities.get(task_type, {}).get("external", []):
|
||
|
|
service = self.external_services.get(service_name)
|
||
|
|
if service and service.available:
|
||
|
|
start_time = time.time()
|
||
|
|
result = service.execute(prompt)
|
||
|
|
exec_time = time.time() - start_time
|
||
|
|
|
||
|
|
if result:
|
||
|
|
return TaskResult(
|
||
|
|
provider="external",
|
||
|
|
model=service_name,
|
||
|
|
result=result,
|
||
|
|
quality_score=self.capabilities[task_type]["quality_external"],
|
||
|
|
execution_time=exec_time,
|
||
|
|
cost=0.01, # Estimated cost
|
||
|
|
timestamp=time.strftime("%Y-%m-%dT%H:%M:%SZ")
|
||
|
|
)
|
||
|
|
|
||
|
|
return None
|
||
|
|
|
||
|
|
def _route_quality_first(self, task_type: TaskType, prompt: str) -> Optional[TaskResult]:
|
||
|
|
"""Choose provider with highest quality."""
|
||
|
|
print(" Choosing highest quality provider...")
|
||
|
|
|
||
|
|
best_result = None
|
||
|
|
best_quality = 0
|
||
|
|
|
||
|
|
# Check external first (usually higher quality)
|
||
|
|
for service_name in self.capabilities.get(task_type, {}).get("external", []):
|
||
|
|
service = self.external_services.get(service_name)
|
||
|
|
if service and service.available:
|
||
|
|
quality = self.capabilities[task_type]["quality_external"]
|
||
|
|
if quality > best_quality:
|
||
|
|
start_time = time.time()
|
||
|
|
result = service.execute(prompt)
|
||
|
|
exec_time = time.time() - start_time
|
||
|
|
|
||
|
|
if result:
|
||
|
|
best_quality = quality
|
||
|
|
best_result = TaskResult(
|
||
|
|
provider="external",
|
||
|
|
model=service_name,
|
||
|
|
result=result,
|
||
|
|
quality_score=quality,
|
||
|
|
execution_time=exec_time,
|
||
|
|
cost=0.01,
|
||
|
|
timestamp=time.strftime("%Y-%m-%dT%H:%M:%SZ")
|
||
|
|
)
|
||
|
|
|
||
|
|
# Check local models
|
||
|
|
for model_name in self.capabilities.get(task_type, {}).get("local", []):
|
||
|
|
model = self.local_models.get(model_name)
|
||
|
|
if model and model.available:
|
||
|
|
quality = self.capabilities[task_type]["quality_local"]
|
||
|
|
if quality > best_quality:
|
||
|
|
start_time = time.time()
|
||
|
|
result = model.execute(prompt)
|
||
|
|
exec_time = time.time() - start_time
|
||
|
|
|
||
|
|
if result:
|
||
|
|
best_quality = quality
|
||
|
|
best_result = TaskResult(
|
||
|
|
provider="local",
|
||
|
|
model=model_name,
|
||
|
|
result=result,
|
||
|
|
quality_score=quality,
|
||
|
|
execution_time=exec_time,
|
||
|
|
cost=0.0,
|
||
|
|
timestamp=time.strftime("%Y-%m-%dT%H:%M:%SZ")
|
||
|
|
)
|
||
|
|
|
||
|
|
return best_result
|
||
|
|
|
||
|
|
def _route_cost_first(self, task_type: TaskType, prompt: str) -> Optional[TaskResult]:
|
||
|
|
"""Choose cheapest provider."""
|
||
|
|
print(" Choosing cheapest provider...")
|
||
|
|
|
||
|
|
# Try local first (free)
|
||
|
|
for model_name in self.capabilities.get(task_type, {}).get("local", []):
|
||
|
|
model = self.local_models.get(model_name)
|
||
|
|
if model and model.available:
|
||
|
|
start_time = time.time()
|
||
|
|
result = model.execute(prompt)
|
||
|
|
exec_time = time.time() - start_time
|
||
|
|
|
||
|
|
if result:
|
||
|
|
return TaskResult(
|
||
|
|
provider="local",
|
||
|
|
model=model_name,
|
||
|
|
result=result,
|
||
|
|
quality_score=self.capabilities[task_type]["quality_local"],
|
||
|
|
execution_time=exec_time,
|
||
|
|
cost=0.0,
|
||
|
|
timestamp=time.strftime("%Y-%m-%dT%H:%M:%SZ")
|
||
|
|
)
|
||
|
|
|
||
|
|
# Fallback to external (paid)
|
||
|
|
for service_name in self.capabilities.get(task_type, {}).get("external", []):
|
||
|
|
service = self.external_services.get(service_name)
|
||
|
|
if service and service.available:
|
||
|
|
start_time = time.time()
|
||
|
|
result = service.execute(prompt)
|
||
|
|
exec_time = time.time() - start_time
|
||
|
|
|
||
|
|
if result:
|
||
|
|
return TaskResult(
|
||
|
|
provider="external",
|
||
|
|
model=service_name,
|
||
|
|
result=result,
|
||
|
|
quality_score=self.capabilities[task_type]["quality_external"],
|
||
|
|
execution_time=exec_time,
|
||
|
|
cost=0.01,
|
||
|
|
timestamp=time.strftime("%Y-%m-%dT%H:%M:%SZ")
|
||
|
|
)
|
||
|
|
|
||
|
|
return None
|
||
|
|
|
||
|
|
def _route_balanced(self, task_type: TaskType, prompt: str) -> Optional[TaskResult]:
|
||
|
|
"""Balance quality and cost."""
|
||
|
|
print(" Balancing quality and cost...")
|
||
|
|
|
||
|
|
# Simple heuristic: if local quality is within 10% of external, use local
|
||
|
|
local_quality = self.capabilities.get(task_type, {}).get("quality_local", 0)
|
||
|
|
external_quality = self.capabilities.get(task_type, {}).get("quality_external", 0)
|
||
|
|
|
||
|
|
if local_quality >= external_quality * 0.9:
|
||
|
|
# Local quality is good enough
|
||
|
|
return self._route_local_first(task_type, prompt)
|
||
|
|
else:
|
||
|
|
# External quality is significantly better
|
||
|
|
return self._route_quality_first(task_type, prompt)
|
||
|
|
|
||
|
|
def main():
|
||
|
|
parser = argparse.ArgumentParser(description="Route tasks to local or external AI providers")
|
||
|
|
parser.add_argument("--task", required=True, choices=[t.value for t in TaskType],
|
||
|
|
help="Task type")
|
||
|
|
parser.add_argument("--prompt", required=True, help="Task prompt")
|
||
|
|
parser.add_argument("--priority", default="balanced",
|
||
|
|
choices=[p.value for p in Priority],
|
||
|
|
help="Routing priority")
|
||
|
|
parser.add_argument("--output", help="Output file for results")
|
||
|
|
|
||
|
|
args = parser.parse_args()
|
||
|
|
|
||
|
|
# Create router
|
||
|
|
router = TaskRouter()
|
||
|
|
|
||
|
|
# Route task
|
||
|
|
task_type = TaskType(args.task)
|
||
|
|
priority = Priority(args.priority)
|
||
|
|
|
||
|
|
result = router.route_task(task_type, args.prompt, priority)
|
||
|
|
|
||
|
|
if result:
|
||
|
|
print(f"\nTask routed to: {result.provider} ({result.model})")
|
||
|
|
print(f"Quality score: {result.quality_score:.2f}")
|
||
|
|
print(f"Execution time: {result.execution_time:.2f}s")
|
||
|
|
print(f"Cost: ${result.cost:.4f}")
|
||
|
|
print(f"Result preview: {result.result[:100]}...")
|
||
|
|
|
||
|
|
# Save results if requested
|
||
|
|
if args.output:
|
||
|
|
with open(args.output, 'w') as f:
|
||
|
|
json.dump({
|
||
|
|
"task_type": task_type.value,
|
||
|
|
"prompt": args.prompt,
|
||
|
|
"priority": priority.value,
|
||
|
|
"result": {
|
||
|
|
"provider": result.provider,
|
||
|
|
"model": result.model,
|
||
|
|
"quality_score": result.quality_score,
|
||
|
|
"execution_time": result.execution_time,
|
||
|
|
"cost": result.cost,
|
||
|
|
"timestamp": result.timestamp,
|
||
|
|
"result": result.result
|
||
|
|
}
|
||
|
|
}, f, indent=2)
|
||
|
|
print(f"Results saved to {args.output}")
|
||
|
|
else:
|
||
|
|
print("Failed to route task - no available providers")
|
||
|
|
return 1
|
||
|
|
|
||
|
|
return 0
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
import sys
|
||
|
|
sys.exit(main())
|