feat: complete production-quality Model Router AI traffic control platform
This commit is contained in:
Vardhman Gupta
@@ -0,0 +1,119 @@
|
||||
from typing import Dict, Any, List, Optional
|
||||
from sqlalchemy import select, func
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
from app.storage.models import RequestRecord, FeedbackRecord, ModelRecord
|
||||
from app.config.settings import get_settings
|
||||
|
||||
settings = get_settings()
|
||||
|
||||
|
||||
async def calculate_cost_savings(
|
||||
session: AsyncSession,
|
||||
baseline_model_id: Optional[str] = None,
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Computes accurate, non-invented cost savings:
|
||||
Baseline Cost (if all requests were sent to baseline model) vs Actual Routed Cost.
|
||||
"""
|
||||
base_id = baseline_model_id or settings.BASELINE_MODEL_ID
|
||||
|
||||
# Fetch baseline model pricing
|
||||
res_base = await session.execute(select(ModelRecord).filter_by(id=base_id))
|
||||
baseline_model = res_base.scalar_one_or_none()
|
||||
|
||||
base_in_rate = baseline_model.cost_per_input_token if baseline_model else 0.000005
|
||||
base_out_rate = baseline_model.cost_per_output_token if baseline_model else 0.000015
|
||||
|
||||
# Fetch aggregate requests
|
||||
res = await session.execute(
|
||||
select(
|
||||
func.count(RequestRecord.request_id),
|
||||
func.sum(RequestRecord.input_tokens),
|
||||
func.sum(RequestRecord.output_tokens),
|
||||
func.sum(RequestRecord.estimated_cost),
|
||||
)
|
||||
)
|
||||
total_reqs, total_in_tok, total_out_tok, total_routed_cost = res.one()
|
||||
|
||||
total_reqs = total_reqs or 0
|
||||
total_in_tok = total_in_tok or 0
|
||||
total_out_tok = total_out_tok or 0
|
||||
total_routed_cost = round(total_routed_cost or 0.0, 6)
|
||||
|
||||
# Theoretical cost if all traffic ran on baseline
|
||||
baseline_total_cost = round((total_in_tok * base_in_rate) + (total_out_tok * base_out_rate), 6)
|
||||
saved_amount = max(0.0, round(baseline_total_cost - total_routed_cost, 6))
|
||||
savings_percent = round((saved_amount / baseline_total_cost * 100.0), 2) if baseline_total_cost > 0 else 0.0
|
||||
|
||||
return {
|
||||
"baseline_model": base_id,
|
||||
"total_requests": total_reqs,
|
||||
"total_input_tokens": total_in_tok,
|
||||
"total_output_tokens": total_out_tok,
|
||||
"baseline_total_cost": baseline_total_cost,
|
||||
"routed_total_cost": total_routed_cost,
|
||||
"cost_saved_usd": saved_amount,
|
||||
"savings_percentage": savings_percent,
|
||||
}
|
||||
|
||||
|
||||
async def get_system_analytics(session: AsyncSession) -> Dict[str, Any]:
|
||||
"""
|
||||
Collects full aggregate metrics for the dashboard.
|
||||
"""
|
||||
# 1. Totals & averages
|
||||
res = await session.execute(
|
||||
select(
|
||||
func.count(RequestRecord.request_id),
|
||||
func.avg(RequestRecord.total_latency_ms),
|
||||
func.avg(RequestRecord.routing_latency_ms),
|
||||
func.avg(RequestRecord.estimated_cost),
|
||||
func.sum(RequestRecord.estimated_cost),
|
||||
)
|
||||
)
|
||||
count, avg_lat, avg_route_lat, avg_cost, sum_cost = res.one()
|
||||
|
||||
# 2. Distribution by model
|
||||
model_dist_res = await session.execute(
|
||||
select(RequestRecord.selected_model, func.count(RequestRecord.request_id))
|
||||
.group_by(RequestRecord.selected_model)
|
||||
)
|
||||
model_distribution = {m: c for m, c in model_dist_res.all()}
|
||||
|
||||
# 3. Distribution by task type
|
||||
task_dist_res = await session.execute(
|
||||
select(RequestRecord.task_type, func.count(RequestRecord.request_id))
|
||||
.group_by(RequestRecord.task_type)
|
||||
)
|
||||
task_distribution = {t: c for t, c in task_dist_res.all()}
|
||||
|
||||
# 4. Fallback rate
|
||||
fb_res = await session.execute(
|
||||
select(func.count(RequestRecord.request_id)).filter(RequestRecord.fallback_used == True)
|
||||
)
|
||||
fallback_count = fb_res.scalar() or 0
|
||||
fallback_rate = round((fallback_count / count * 100.0), 2) if count and count > 0 else 0.0
|
||||
|
||||
# 5. User Feedback aggregate
|
||||
fb_pos_res = await session.execute(select(func.count(FeedbackRecord.id)).filter(FeedbackRecord.rating == 1))
|
||||
fb_neg_res = await session.execute(select(func.count(FeedbackRecord.id)).filter(FeedbackRecord.rating == -1))
|
||||
pos_feedback = fb_pos_res.scalar() or 0
|
||||
neg_feedback = fb_neg_res.scalar() or 0
|
||||
total_feedback = pos_feedback + neg_feedback
|
||||
quality_score = round((pos_feedback / total_feedback * 100.0), 1) if total_feedback > 0 else None
|
||||
|
||||
savings = await calculate_cost_savings(session)
|
||||
|
||||
return {
|
||||
"total_requests": count or 0,
|
||||
"avg_latency_ms": round(avg_lat or 0.0, 1),
|
||||
"avg_routing_latency_ms": round(avg_route_lat or 0.0, 1),
|
||||
"avg_cost_usd": round(avg_cost or 0.0, 6),
|
||||
"total_cost_usd": round(sum_cost or 0.0, 6),
|
||||
"quality_score_percent": quality_score,
|
||||
"fallback_rate_percent": fallback_rate,
|
||||
"fallback_count": fallback_count,
|
||||
"model_distribution": model_distribution,
|
||||
"task_distribution": task_distribution,
|
||||
"savings": savings,
|
||||
}
|
||||
@@ -0,0 +1,67 @@
|
||||
import json
|
||||
import re
|
||||
from typing import Optional
|
||||
from app.models.schemas import RequestAnalysis, TaskType, PriorityLevel
|
||||
from app.analyzer.heuristics import analyze_request_heuristics, estimate_tokens
|
||||
from app.providers.registry import provider_registry
|
||||
from app.config.settings import get_settings
|
||||
|
||||
settings = get_settings()
|
||||
|
||||
CLASSIFIER_PROMPT = """You are a real-time request classifier for an intelligent LLM model router.
|
||||
Analyze the user prompt and return ONLY a JSON object matching this schema:
|
||||
{
|
||||
"task_type": "GENERAL_QA" | "CODING" | "DEBUGGING" | "REASONING" | "SUMMARIZATION" | "EXTRACTION" | "WRITING" | "TRANSLATION" | "ANALYSIS" | "MATH" | "LONG_CONTEXT" | "CREATIVE",
|
||||
"complexity": 0.0 to 1.0,
|
||||
"reasoning_required": true/false,
|
||||
"coding_required": true/false,
|
||||
"vision_required": true/false,
|
||||
"latency_priority": "LOW" | "MEDIUM" | "HIGH",
|
||||
"cost_sensitivity": "LOW" | "MEDIUM" | "HIGH",
|
||||
"quality_requirement": "LOW" | "MEDIUM" | "HIGH"
|
||||
}
|
||||
Output valid JSON only with no markdown wrapping or additional text.
|
||||
"""
|
||||
|
||||
|
||||
async def analyze_request(prompt: str, analyzer_mode: Optional[str] = None) -> RequestAnalysis:
|
||||
mode = (analyzer_mode or settings.ROUTER_ANALYZER).lower()
|
||||
|
||||
if mode == "llm":
|
||||
try:
|
||||
# Attempt to use local or configured provider for classification
|
||||
provider = provider_registry.get_provider("ollama")
|
||||
if provider:
|
||||
health = await provider.check_health()
|
||||
if health.get("status") == "CONNECTED":
|
||||
resp = await provider.generate(
|
||||
prompt=f"Classify this request:\n\n{prompt[:1000]}",
|
||||
model_id="qwen2.5-coder",
|
||||
system_prompt=CLASSIFIER_PROMPT,
|
||||
temperature=0.1,
|
||||
)
|
||||
clean_json = re.search(r"\{.*\}", resp.content, re.DOTALL)
|
||||
if clean_json:
|
||||
data = json.loads(clean_json.group(0))
|
||||
comp = float(data.get("complexity", 0.5))
|
||||
comp_label = PriorityLevel.HIGH if comp >= 0.75 else (PriorityLevel.MEDIUM if comp >= 0.45 else PriorityLevel.LOW)
|
||||
return RequestAnalysis(
|
||||
task_type=TaskType(data.get("task_type", "GENERAL_QA")),
|
||||
complexity=comp,
|
||||
complexity_label=comp_label,
|
||||
reasoning_required=bool(data.get("reasoning_required", False)),
|
||||
coding_required=bool(data.get("coding_required", False)),
|
||||
vision_required=bool(data.get("vision_required", False)),
|
||||
context_size=estimate_tokens(prompt),
|
||||
latency_priority=PriorityLevel(data.get("latency_priority", "MEDIUM")),
|
||||
cost_sensitivity=PriorityLevel(data.get("cost_sensitivity", "MEDIUM")),
|
||||
quality_requirement=PriorityLevel(data.get("quality_requirement", "MEDIUM")),
|
||||
keywords_detected=["llm_classified"],
|
||||
analyzer_used="llm",
|
||||
)
|
||||
except Exception:
|
||||
# Graceful fallback to heuristics if LLM classifier is unavailable or times out
|
||||
pass
|
||||
|
||||
# Default to fast, deterministic heuristics
|
||||
return analyze_request_heuristics(prompt)
|
||||
@@ -0,0 +1,192 @@
|
||||
import re
|
||||
from typing import List, Tuple
|
||||
from app.models.schemas import TaskType, PriorityLevel, RequestAnalysis
|
||||
|
||||
# Regex rules and indicator vocabularies
|
||||
DEBUG_PATTERNS = [
|
||||
r"\b(debug|bug|stacktrace|exception|error|fix this error|traceback|segfault|syntaxerror|nullpointer)\b",
|
||||
r"\b(crash|panicked|failed test|reproduce|root cause|undefined is not a function)\b",
|
||||
]
|
||||
|
||||
CODING_PATTERNS = [
|
||||
r"\b(write a python|write code|implement|function|class |def |const |async def|typescript|javascript|golang|rust|react|sql query|endpoint|api client)\b",
|
||||
r"\b(refactor|algorithm|regex|unit test|dockerfile|kubernetes yaml|prisma|fastapi)\b",
|
||||
r"```[a-zA-Z0-9_-]*\n",
|
||||
]
|
||||
|
||||
REASONING_PATTERNS = [
|
||||
r"\b(explain why|reasoning|step by step|proof|derive|why did|compare and contrast|architectural decision|tradeoffs|trade-offs|consequences of)\b",
|
||||
r"\b(game theory|logical fallacy|syllogism|first principles|hypothesize)\b",
|
||||
]
|
||||
|
||||
MATH_PATTERNS = [
|
||||
r"\b(calculate|integral|derivative|matrix|eigenvalue|probability|equation|differential equation|solve for x|combinatorics|standard deviation)\b",
|
||||
r"[\$\\\^\{\}\[\]\=]{3,}",
|
||||
]
|
||||
|
||||
SUMMARIZATION_PATTERNS = [
|
||||
r"\b(summarize|summary|tldr|tl;dr|key points|bullet points of this|condense|digest|brief overview)\b",
|
||||
]
|
||||
|
||||
EXTRACTION_PATTERNS = [
|
||||
r"\b(extract|extract json|parse this table|pull all emails|convert to json|output as csv|structured data from)\b",
|
||||
]
|
||||
|
||||
TRANSLATION_PATTERNS = [
|
||||
r"\b(translate to|translate into|in french|in spanish|in german|in japanese|in chinese|in hindi|translate from)\b",
|
||||
]
|
||||
|
||||
CREATIVE_PATTERNS = [
|
||||
r"\b(write a poem|write a story|screenplay|brainstorm catchy|creative writing|dialogue between|fiction|metaphor)\b",
|
||||
]
|
||||
|
||||
ANALYSIS_PATTERNS = [
|
||||
r"\b(analyze|breakdown|evaluate|critique|audit|financial report|trend analysis|pros and cons)\b",
|
||||
]
|
||||
|
||||
|
||||
def estimate_tokens(text: str) -> int:
|
||||
"""Fast, reliable token estimation (roughly 4 characters per token)."""
|
||||
return max(1, len(text) // 4)
|
||||
|
||||
|
||||
def analyze_request_heuristics(prompt: str) -> RequestAnalysis:
|
||||
text_lower = prompt.lower()
|
||||
tokens = estimate_tokens(prompt)
|
||||
|
||||
detected_keywords: List[str] = []
|
||||
task_scores: dict[TaskType, float] = {t: 0.0 for t in TaskType}
|
||||
|
||||
# 1. Evaluate patterns
|
||||
for p in DEBUG_PATTERNS:
|
||||
matches = re.findall(p, text_lower)
|
||||
if matches:
|
||||
task_scores[TaskType.DEBUGGING] += len(matches) * 2.5
|
||||
detected_keywords.extend(matches)
|
||||
|
||||
for p in CODING_PATTERNS:
|
||||
matches = re.findall(p, text_lower)
|
||||
if matches:
|
||||
task_scores[TaskType.CODING] += len(matches) * 2.0
|
||||
detected_keywords.extend(matches)
|
||||
|
||||
for p in REASONING_PATTERNS:
|
||||
matches = re.findall(p, text_lower)
|
||||
if matches:
|
||||
task_scores[TaskType.REASONING] += len(matches) * 2.0
|
||||
detected_keywords.extend(matches)
|
||||
|
||||
for p in MATH_PATTERNS:
|
||||
matches = re.findall(p, text_lower)
|
||||
if matches:
|
||||
task_scores[TaskType.MATH] += len(matches) * 2.2
|
||||
detected_keywords.extend(matches)
|
||||
|
||||
for p in SUMMARIZATION_PATTERNS:
|
||||
matches = re.findall(p, text_lower)
|
||||
if matches:
|
||||
task_scores[TaskType.SUMMARIZATION] += len(matches) * 2.5
|
||||
detected_keywords.extend(matches)
|
||||
|
||||
for p in EXTRACTION_PATTERNS:
|
||||
matches = re.findall(p, text_lower)
|
||||
if matches:
|
||||
task_scores[TaskType.EXTRACTION] += len(matches) * 2.0
|
||||
detected_keywords.extend(matches)
|
||||
|
||||
for p in TRANSLATION_PATTERNS:
|
||||
matches = re.findall(p, text_lower)
|
||||
if matches:
|
||||
task_scores[TaskType.TRANSLATION] += len(matches) * 3.0
|
||||
detected_keywords.extend(matches)
|
||||
|
||||
for p in CREATIVE_PATTERNS:
|
||||
matches = re.findall(p, text_lower)
|
||||
if matches:
|
||||
task_scores[TaskType.CREATIVE] += len(matches) * 2.0
|
||||
detected_keywords.extend(matches)
|
||||
|
||||
for p in ANALYSIS_PATTERNS:
|
||||
matches = re.findall(p, text_lower)
|
||||
if matches:
|
||||
task_scores[TaskType.ANALYSIS] += len(matches) * 1.8
|
||||
detected_keywords.extend(matches)
|
||||
|
||||
# Long context check
|
||||
if tokens > 4000:
|
||||
task_scores[TaskType.LONG_CONTEXT] += 3.0
|
||||
|
||||
# 2. Determine Primary Task
|
||||
best_task = TaskType.GENERAL_QA
|
||||
max_score = 0.0
|
||||
for task, score in task_scores.items():
|
||||
if score > max_score:
|
||||
max_score = score
|
||||
best_task = task
|
||||
|
||||
# 3. Derive Complexity (0.0 to 1.0)
|
||||
complexity = 0.35 # Base complexity
|
||||
|
||||
# Token length factor
|
||||
if tokens > 10000:
|
||||
complexity += 0.30
|
||||
elif tokens > 3000:
|
||||
complexity += 0.20
|
||||
elif tokens > 800:
|
||||
complexity += 0.10
|
||||
elif tokens < 30:
|
||||
complexity -= 0.10
|
||||
|
||||
# Task inherent complexity
|
||||
if best_task in (TaskType.DEBUGGING, TaskType.MATH, TaskType.REASONING):
|
||||
complexity += 0.30
|
||||
elif best_task in (TaskType.CODING, TaskType.ANALYSIS, TaskType.LONG_CONTEXT):
|
||||
complexity += 0.20
|
||||
elif best_task in (TaskType.TRANSLATION, TaskType.EXTRACTION, TaskType.SUMMARIZATION):
|
||||
complexity += 0.05
|
||||
elif best_task == TaskType.GENERAL_QA and tokens < 50:
|
||||
complexity -= 0.15
|
||||
|
||||
# Multi-step keyword boosts
|
||||
if any(k in text_lower for k in ["distributed", "async", "concurrency", "deadlock", "microservices", "architecture"]):
|
||||
complexity += 0.25
|
||||
if any(k in text_lower for k in ["simple", "quick", "one liner", "easy", "hello world"]):
|
||||
complexity -= 0.20
|
||||
|
||||
complexity = max(0.05, min(0.99, round(complexity, 2)))
|
||||
|
||||
# 4. Priority and Capabilities
|
||||
coding_required = best_task in (TaskType.CODING, TaskType.DEBUGGING) or task_scores[TaskType.CODING] > 0
|
||||
reasoning_required = best_task in (TaskType.REASONING, TaskType.DEBUGGING, TaskType.MATH) or complexity >= 0.70
|
||||
|
||||
if complexity >= 0.75:
|
||||
complexity_label = PriorityLevel.HIGH
|
||||
quality_req = PriorityLevel.HIGH
|
||||
latency_pri = PriorityLevel.LOW
|
||||
cost_sens = PriorityLevel.LOW
|
||||
elif complexity >= 0.45:
|
||||
complexity_label = PriorityLevel.MEDIUM
|
||||
quality_req = PriorityLevel.MEDIUM
|
||||
latency_pri = PriorityLevel.MEDIUM
|
||||
cost_sens = PriorityLevel.MEDIUM
|
||||
else:
|
||||
complexity_label = PriorityLevel.LOW
|
||||
quality_req = PriorityLevel.LOW
|
||||
latency_pri = PriorityLevel.HIGH
|
||||
cost_sens = PriorityLevel.HIGH
|
||||
|
||||
return RequestAnalysis(
|
||||
task_type=best_task,
|
||||
complexity=complexity,
|
||||
complexity_label=complexity_label,
|
||||
reasoning_required=reasoning_required,
|
||||
coding_required=coding_required,
|
||||
vision_required=False,
|
||||
tools_required=False,
|
||||
context_size=tokens,
|
||||
latency_priority=latency_pri,
|
||||
cost_sensitivity=cost_sens,
|
||||
quality_requirement=quality_req,
|
||||
keywords_detected=list(set(detected_keywords))[:10],
|
||||
analyzer_used="rules",
|
||||
)
|
||||
@@ -0,0 +1,532 @@
|
||||
import uuid
|
||||
import time
|
||||
import asyncio
|
||||
from typing import List, Optional, Dict, Any
|
||||
from fastapi import APIRouter, Depends, HTTPException, Query
|
||||
from pydantic import BaseModel, Field
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
from sqlalchemy import select, desc
|
||||
|
||||
from app.storage.database import get_db
|
||||
from app.storage.models import (
|
||||
ModelRecord,
|
||||
RoutingPolicyRecord,
|
||||
RoutingRuleRecord,
|
||||
RequestRecord,
|
||||
RoutingDecisionRecord,
|
||||
ResponseRecord,
|
||||
FeedbackRecord,
|
||||
BudgetRecord,
|
||||
ExperimentRecord,
|
||||
)
|
||||
from app.models.schemas import (
|
||||
ModelMetadata,
|
||||
RequestAnalysis,
|
||||
RoutingDecision,
|
||||
ProviderResponse,
|
||||
ModelTier,
|
||||
)
|
||||
from app.analyzer.analyzer import analyze_request
|
||||
from app.router.engine import route_request
|
||||
from app.fallback.handler import execute_with_fallback
|
||||
from app.budgets.manager import check_budget_threshold
|
||||
from app.observability.events import log_router_event
|
||||
from app.analytics.service import get_system_analytics, calculate_cost_savings
|
||||
from app.providers.registry import provider_registry
|
||||
from app.config.settings import get_settings
|
||||
|
||||
router = APIRouter()
|
||||
settings = get_settings()
|
||||
|
||||
|
||||
class RouteOnlyRequest(BaseModel):
|
||||
prompt: str = Field(..., min_length=1, description="Prompt text to analyze and route")
|
||||
policy: Optional[str] = None
|
||||
analyzer_mode: Optional[str] = None
|
||||
|
||||
|
||||
class GenerateRequest(BaseModel):
|
||||
prompt: str = Field(..., min_length=1)
|
||||
system_prompt: Optional[str] = None
|
||||
policy: Optional[str] = None
|
||||
analyzer_mode: Optional[str] = None
|
||||
force_model: Optional[str] = None
|
||||
temperature: float = 0.7
|
||||
max_retries: int = 1
|
||||
|
||||
|
||||
class FeedbackSubmitRequest(BaseModel):
|
||||
request_id: str
|
||||
rating: int = Field(..., description="1 for up, -1 for down")
|
||||
comment: Optional[str] = None
|
||||
|
||||
|
||||
class ModelCreateRequest(BaseModel):
|
||||
id: str
|
||||
name: str
|
||||
provider: str
|
||||
type: str = "LOCAL"
|
||||
tier: str = "BALANCED"
|
||||
context_window: int = 32768
|
||||
supports_coding: bool = True
|
||||
supports_reasoning: bool = True
|
||||
supports_vision: bool = False
|
||||
supports_tools: bool = True
|
||||
quality_score: float = 0.85
|
||||
speed_score: float = 0.85
|
||||
cost_per_input_token: float = 0.0
|
||||
cost_per_output_token: float = 0.0
|
||||
|
||||
|
||||
class RuleCreateRequest(BaseModel):
|
||||
id: Optional[str] = None
|
||||
name: str
|
||||
description: Optional[str] = None
|
||||
priority: int = 10
|
||||
is_enabled: bool = True
|
||||
condition_field: str
|
||||
condition_operator: str
|
||||
condition_value: str
|
||||
action_type: str
|
||||
action_target: str
|
||||
|
||||
|
||||
# Helper to convert DB model record to Pydantic ModelMetadata
|
||||
def db_model_to_meta(m: ModelRecord) -> ModelMetadata:
|
||||
return ModelMetadata(
|
||||
id=m.id,
|
||||
name=m.name,
|
||||
provider=m.provider,
|
||||
type=m.type,
|
||||
tier=ModelTier(m.tier) if m.tier in [t.value for t in ModelTier] else ModelTier.BALANCED,
|
||||
context_window=m.context_window,
|
||||
supports_coding=m.supports_coding,
|
||||
supports_reasoning=m.supports_reasoning,
|
||||
supports_vision=m.supports_vision,
|
||||
supports_tools=m.supports_tools,
|
||||
quality_score=m.quality_score,
|
||||
speed_score=m.speed_score,
|
||||
reliability_score=m.reliability_score,
|
||||
cost_per_input_token=m.cost_per_input_token,
|
||||
cost_per_output_token=m.cost_per_output_token,
|
||||
availability=m.availability,
|
||||
is_active=m.is_active,
|
||||
)
|
||||
|
||||
|
||||
@router.post("/api/route", response_model=Dict[str, Any])
|
||||
async def api_route_prompt(req: RouteOnlyRequest, db: AsyncSession = Depends(get_db)):
|
||||
"""
|
||||
Dry-run routing inspection endpoint:
|
||||
Analyzes request and produces structured routing decision without invoking provider.
|
||||
"""
|
||||
t_start = time.perf_counter()
|
||||
req_id = f"req_{uuid.uuid4().hex[:12]}"
|
||||
|
||||
# 1. Analyze Request
|
||||
analysis = await analyze_request(req.prompt, req.analyzer_mode)
|
||||
|
||||
# 2. Fetch Models & Policy
|
||||
res_m = await db.execute(select(ModelRecord).filter_by(is_active=True))
|
||||
db_models = res_m.scalars().all()
|
||||
models = [db_model_to_meta(m) for m in db_models]
|
||||
|
||||
policy_name = req.policy or settings.DEFAULT_ROUTING_POLICY
|
||||
res_p = await db.execute(select(RoutingPolicyRecord).filter_by(id=policy_name))
|
||||
policy = res_p.scalar_one_or_none()
|
||||
weights = {
|
||||
"quality_weight": policy.quality_weight if policy else 0.35,
|
||||
"cost_weight": policy.cost_weight if policy else 0.25,
|
||||
"speed_weight": policy.speed_weight if policy else 0.20,
|
||||
"capability_weight": policy.capability_weight if policy else 0.15,
|
||||
"reliability_weight": policy.reliability_weight if policy else 0.05,
|
||||
}
|
||||
|
||||
# 3. Fetch Rules & Budget Check
|
||||
res_r = await db.execute(select(RoutingRuleRecord).filter_by(is_enabled=True))
|
||||
rules = res_r.scalars().all()
|
||||
budget_pct, _ = await check_budget_threshold(db, 0.0)
|
||||
|
||||
# 4. Routing Decision
|
||||
decision = route_request(
|
||||
analysis=analysis,
|
||||
available_models=models,
|
||||
policy_weights=weights,
|
||||
policy_name=policy_name,
|
||||
rules=rules,
|
||||
budget_percent=budget_pct,
|
||||
request_id=req_id,
|
||||
)
|
||||
|
||||
t_route_ms = (time.perf_counter() - t_start) * 1000.0
|
||||
|
||||
return {
|
||||
"request_id": req_id,
|
||||
"analysis": analysis.model_dump(),
|
||||
"decision": decision.model_dump(),
|
||||
"routing_latency_ms": round(t_route_ms, 2),
|
||||
}
|
||||
|
||||
|
||||
@router.post("/api/generate", response_model=Dict[str, Any])
|
||||
async def api_generate_response(req: GenerateRequest, db: AsyncSession = Depends(get_db)):
|
||||
"""
|
||||
End-to-End routed execution:
|
||||
Analyzes -> Routes -> Executes Selected Model / Fallback -> Records Decision & Metrics.
|
||||
"""
|
||||
t_start = time.perf_counter()
|
||||
req_id = f"req_{uuid.uuid4().hex[:12]}"
|
||||
|
||||
# 1. Analyze Request
|
||||
analysis = await analyze_request(req.prompt, req.analyzer_mode)
|
||||
|
||||
# 2. Fetch Models & Policy
|
||||
res_m = await db.execute(select(ModelRecord).filter_by(is_active=True))
|
||||
db_models = res_m.scalars().all()
|
||||
models = [db_model_to_meta(m) for m in db_models]
|
||||
|
||||
policy_name = req.policy or settings.DEFAULT_ROUTING_POLICY
|
||||
res_p = await db.execute(select(RoutingPolicyRecord).filter_by(id=policy_name))
|
||||
policy = res_p.scalar_one_or_none()
|
||||
weights = {
|
||||
"quality_weight": policy.quality_weight if policy else 0.35,
|
||||
"cost_weight": policy.cost_weight if policy else 0.25,
|
||||
"speed_weight": policy.speed_weight if policy else 0.20,
|
||||
"capability_weight": policy.capability_weight if policy else 0.15,
|
||||
"reliability_weight": policy.reliability_weight if policy else 0.05,
|
||||
}
|
||||
|
||||
# 3. Fetch Rules & Budget Check
|
||||
res_r = await db.execute(select(RoutingRuleRecord).filter_by(is_enabled=True))
|
||||
rules = res_r.scalars().all()
|
||||
budget_pct, budget_mode = await check_budget_threshold(db, 0.0)
|
||||
|
||||
# 4. Route
|
||||
decision = route_request(
|
||||
analysis=analysis,
|
||||
available_models=models,
|
||||
policy_weights=weights,
|
||||
policy_name=policy_name,
|
||||
rules=rules,
|
||||
budget_percent=budget_pct,
|
||||
request_id=req_id,
|
||||
)
|
||||
|
||||
t_routing_done = time.perf_counter()
|
||||
routing_latency_ms = (t_routing_done - t_start) * 1000.0
|
||||
|
||||
# If force model requested, override
|
||||
selected_model_id = req.force_model or decision.selected_model
|
||||
target_meta = next((m for m in models if m.id == selected_model_id), None)
|
||||
selected_provider_id = target_meta.provider if target_meta else decision.provider
|
||||
|
||||
# 5. Provider Execution with Fallback & Retries
|
||||
provider_resp, fallback_used, orig_model, fallback_reason = await execute_with_fallback(
|
||||
prompt=req.prompt,
|
||||
selected_model_id=selected_model_id,
|
||||
selected_provider_id=selected_provider_id,
|
||||
all_models=models,
|
||||
system_prompt=req.system_prompt,
|
||||
temperature=req.temperature,
|
||||
max_retries=req.max_retries,
|
||||
)
|
||||
|
||||
total_latency_ms = (time.perf_counter() - t_start) * 1000.0
|
||||
|
||||
# 6. Calculate actual costs
|
||||
actual_model_meta = next((m for m in models if m.id == provider_resp.model), target_meta)
|
||||
cost_in_rate = actual_model_meta.cost_per_input_token if actual_model_meta else 0.0
|
||||
cost_out_rate = actual_model_meta.cost_per_output_token if actual_model_meta else 0.0
|
||||
actual_cost = round((provider_resp.input_tokens * cost_in_rate) + (provider_resp.output_tokens * cost_out_rate), 6)
|
||||
|
||||
# Baseline cost calculation
|
||||
res_base = await db.execute(select(ModelRecord).filter_by(id=settings.BASELINE_MODEL_ID))
|
||||
base_m = res_base.scalar_one_or_none()
|
||||
base_cost = 0.0
|
||||
if base_m:
|
||||
base_cost = round((provider_resp.input_tokens * base_m.cost_per_input_token) + (provider_resp.output_tokens * base_m.cost_per_output_token), 6)
|
||||
saved_cost = max(0.0, round(base_cost - actual_cost, 6))
|
||||
|
||||
# Update budget spend
|
||||
await check_budget_threshold(db, actual_cost)
|
||||
|
||||
# 7. Persistence
|
||||
req_record = RequestRecord(
|
||||
request_id=req_id,
|
||||
prompt=req.prompt,
|
||||
task_type=analysis.task_type.value,
|
||||
complexity=analysis.complexity,
|
||||
context_size=analysis.context_size,
|
||||
reasoning_required=analysis.reasoning_required,
|
||||
coding_required=analysis.coding_required,
|
||||
routing_policy=policy_name,
|
||||
selected_model=provider_resp.model or selected_model_id,
|
||||
provider=provider_resp.provider or selected_provider_id,
|
||||
status="FALLBACK" if fallback_used else ("FAILED" if provider_resp.error else "SUCCESS"),
|
||||
fallback_used=fallback_used,
|
||||
original_model=orig_model,
|
||||
fallback_reason=fallback_reason,
|
||||
input_tokens=provider_resp.input_tokens,
|
||||
output_tokens=provider_resp.output_tokens,
|
||||
total_tokens=provider_resp.total_tokens,
|
||||
estimated_cost=actual_cost,
|
||||
baseline_cost=base_cost,
|
||||
cost_saved=saved_cost,
|
||||
routing_latency_ms=round(routing_latency_ms, 2),
|
||||
provider_latency_ms=round(provider_resp.provider_latency_ms, 2),
|
||||
total_latency_ms=round(total_latency_ms, 2),
|
||||
time_to_first_token_ms=provider_resp.time_to_first_token_ms,
|
||||
)
|
||||
db.add(req_record)
|
||||
|
||||
decision_record = RoutingDecisionRecord(
|
||||
decision_id=decision.decision_id,
|
||||
request_id=req_id,
|
||||
selected_model=decision.selected_model,
|
||||
confidence=decision.confidence,
|
||||
reasons=decision.reasons,
|
||||
candidate_scores=[s.model_dump() for s in decision.candidate_scores],
|
||||
rejected_candidates=decision.rejected_candidates,
|
||||
policy_used=policy_name,
|
||||
)
|
||||
db.add(decision_record)
|
||||
|
||||
resp_record = ResponseRecord(
|
||||
response_id=f"resp_{uuid.uuid4().hex[:12]}",
|
||||
request_id=req_id,
|
||||
model_id=provider_resp.model or selected_model_id,
|
||||
provider=provider_resp.provider or selected_provider_id,
|
||||
content=provider_resp.content,
|
||||
finish_reason=provider_resp.finish_reason,
|
||||
is_mock=provider_resp.is_mock,
|
||||
)
|
||||
db.add(resp_record)
|
||||
await db.commit()
|
||||
|
||||
# 8. Structured Observability Event
|
||||
log_router_event(
|
||||
event_name="request_completed",
|
||||
request_id=req_id,
|
||||
model=provider_resp.model,
|
||||
provider=provider_resp.provider,
|
||||
duration_ms=round(total_latency_ms, 2),
|
||||
metadata={"tokens": provider_resp.total_tokens, "cost": actual_cost, "fallback": fallback_used},
|
||||
)
|
||||
|
||||
return {
|
||||
"request_id": req_id,
|
||||
"analysis": analysis.model_dump(),
|
||||
"decision": decision.model_dump(),
|
||||
"response": provider_resp.model_dump(),
|
||||
"metrics": {
|
||||
"routing_latency_ms": round(routing_latency_ms, 2),
|
||||
"provider_latency_ms": round(provider_resp.provider_latency_ms, 2),
|
||||
"total_latency_ms": round(total_latency_ms, 2),
|
||||
"estimated_cost_usd": actual_cost,
|
||||
"baseline_cost_usd": base_cost,
|
||||
"cost_saved_usd": saved_cost,
|
||||
"fallback_used": fallback_used,
|
||||
"fallback_reason": fallback_reason,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
# Models & Registry Endpoints
|
||||
@router.get("/api/models")
|
||||
async def list_models(db: AsyncSession = Depends(get_db)):
|
||||
res = await db.execute(select(ModelRecord).order_by(ModelRecord.tier, ModelRecord.name))
|
||||
return res.scalars().all()
|
||||
|
||||
|
||||
@router.post("/api/models")
|
||||
async def create_model(req: ModelCreateRequest, db: AsyncSession = Depends(get_db)):
|
||||
record = ModelRecord(
|
||||
id=req.id,
|
||||
name=req.name,
|
||||
provider=req.provider,
|
||||
type=req.type,
|
||||
tier=req.tier,
|
||||
context_window=req.context_window,
|
||||
supports_coding=req.supports_coding,
|
||||
supports_reasoning=req.supports_reasoning,
|
||||
supports_vision=req.supports_vision,
|
||||
supports_tools=req.supports_tools,
|
||||
quality_score=req.quality_score,
|
||||
speed_score=req.speed_score,
|
||||
cost_per_input_token=req.cost_per_input_token,
|
||||
cost_per_output_token=req.cost_per_output_token,
|
||||
)
|
||||
db.add(record)
|
||||
await db.commit()
|
||||
return record
|
||||
|
||||
|
||||
# Providers Endpoint (Never returning secrets)
|
||||
@router.get("/api/providers")
|
||||
async def list_providers(db: AsyncSession = Depends(get_db)):
|
||||
providers_info = []
|
||||
for pid, p in provider_registry.list_providers().items():
|
||||
health = await p.check_health()
|
||||
providers_info.append({
|
||||
"id": pid,
|
||||
"name": p.name,
|
||||
"status": health.get("status", "READY"),
|
||||
"base_url": p.base_url,
|
||||
"message": health.get("message"),
|
||||
"models_available": health.get("models_available", []),
|
||||
"credentials_status": "Loaded from environment (.env)" if getattr(p, "api_key", None) or pid in ["mock", "ollama"] else "Not configured",
|
||||
})
|
||||
return providers_info
|
||||
|
||||
|
||||
# Live Traffic & Decisions
|
||||
@router.get("/api/traffic")
|
||||
async def get_traffic_feed(limit: int = 50, db: AsyncSession = Depends(get_db)):
|
||||
res = await db.execute(select(RequestRecord).order_by(desc(RequestRecord.timestamp)).limit(limit))
|
||||
return res.scalars().all()
|
||||
|
||||
|
||||
@router.get("/api/decisions/{decision_id_or_request_id}")
|
||||
async def get_decision_details(decision_id_or_request_id: str, db: AsyncSession = Depends(get_db)):
|
||||
res_dec = await db.execute(
|
||||
select(RoutingDecisionRecord).filter(
|
||||
(RoutingDecisionRecord.decision_id == decision_id_or_request_id)
|
||||
| (RoutingDecisionRecord.request_id == decision_id_or_request_id)
|
||||
)
|
||||
)
|
||||
dec = res_dec.scalar_one_or_none()
|
||||
if not dec:
|
||||
raise HTTPException(status_code=404, detail="Decision not found")
|
||||
|
||||
res_req = await db.execute(select(RequestRecord).filter_by(request_id=dec.request_id))
|
||||
req = res_req.scalar_one_or_none()
|
||||
|
||||
res_resp = await db.execute(select(ResponseRecord).filter_by(request_id=dec.request_id))
|
||||
resp = res_resp.scalar_one_or_none()
|
||||
|
||||
return {
|
||||
"decision": dec,
|
||||
"request": req,
|
||||
"response": resp,
|
||||
}
|
||||
|
||||
|
||||
# Feedback
|
||||
@router.post("/api/feedback")
|
||||
async def submit_feedback(req: FeedbackSubmitRequest, db: AsyncSession = Depends(get_db)):
|
||||
res = await db.execute(select(RequestRecord).filter_by(request_id=req.request_id))
|
||||
request_rec = res.scalar_one_or_none()
|
||||
if not request_rec:
|
||||
raise HTTPException(status_code=404, detail="Request ID not found")
|
||||
|
||||
feedback = FeedbackRecord(
|
||||
id=f"fb_{uuid.uuid4().hex[:10]}",
|
||||
request_id=req.request_id,
|
||||
model_id=request_rec.selected_model,
|
||||
task_type=request_rec.task_type,
|
||||
rating=1 if req.rating > 0 else -1,
|
||||
comment=req.comment,
|
||||
)
|
||||
db.add(feedback)
|
||||
await db.commit()
|
||||
return {"status": "success", "feedback_id": feedback.id}
|
||||
|
||||
|
||||
# Analytics Endpoints
|
||||
@router.get("/api/analytics")
|
||||
async def get_analytics(db: AsyncSession = Depends(get_db)):
|
||||
return await get_system_analytics(db)
|
||||
|
||||
|
||||
@router.get("/api/analytics/savings")
|
||||
async def get_savings(baseline_model: Optional[str] = None, db: AsyncSession = Depends(get_db)):
|
||||
return await calculate_cost_savings(db, baseline_model)
|
||||
|
||||
|
||||
# Rules Endpoints
|
||||
@router.get("/api/rules")
|
||||
async def list_rules(db: AsyncSession = Depends(get_db)):
|
||||
res = await db.execute(select(RoutingRuleRecord).order_by(desc(RoutingRuleRecord.priority)))
|
||||
return res.scalars().all()
|
||||
|
||||
|
||||
@router.post("/api/rules")
|
||||
async def create_rule(req: RuleCreateRequest, db: AsyncSession = Depends(get_db)):
|
||||
rule_id = req.id or f"rule-{uuid.uuid4().hex[:8]}"
|
||||
record = RoutingRuleRecord(
|
||||
id=rule_id,
|
||||
name=req.name,
|
||||
description=req.description,
|
||||
priority=req.priority,
|
||||
is_enabled=req.is_enabled,
|
||||
condition_field=req.condition_field,
|
||||
condition_operator=req.condition_operator,
|
||||
condition_value=req.condition_value,
|
||||
action_type=req.action_type,
|
||||
action_target=req.action_target,
|
||||
)
|
||||
db.add(record)
|
||||
await db.commit()
|
||||
return record
|
||||
|
||||
|
||||
@router.put("/api/rules/{rule_id}")
|
||||
async def update_rule(rule_id: str, req: RuleCreateRequest, db: AsyncSession = Depends(get_db)):
|
||||
res = await db.execute(select(RoutingRuleRecord).filter_by(id=rule_id))
|
||||
rule = res.scalar_one_or_none()
|
||||
if not rule:
|
||||
raise HTTPException(status_code=404, detail="Rule not found")
|
||||
rule.name = req.name
|
||||
rule.description = req.description
|
||||
rule.priority = req.priority
|
||||
rule.is_enabled = req.is_enabled
|
||||
rule.condition_field = req.condition_field
|
||||
rule.condition_operator = req.condition_operator
|
||||
rule.condition_value = req.condition_value
|
||||
rule.action_type = req.action_type
|
||||
rule.action_target = req.action_target
|
||||
await db.commit()
|
||||
return rule
|
||||
|
||||
|
||||
@router.delete("/api/rules/{rule_id}")
|
||||
async def delete_rule(rule_id: str, db: AsyncSession = Depends(get_db)):
|
||||
res = await db.execute(select(RoutingRuleRecord).filter_by(id=rule_id))
|
||||
rule = res.scalar_one_or_none()
|
||||
if not rule:
|
||||
raise HTTPException(status_code=404, detail="Rule not found")
|
||||
await db.delete(rule)
|
||||
await db.commit()
|
||||
return {"status": "deleted", "id": rule_id}
|
||||
|
||||
|
||||
# Policies
|
||||
@router.get("/api/policies")
|
||||
async def list_policies(db: AsyncSession = Depends(get_db)):
|
||||
res = await db.execute(select(RoutingPolicyRecord))
|
||||
return res.scalars().all()
|
||||
|
||||
|
||||
# Budgets
|
||||
@router.get("/api/budgets")
|
||||
async def get_budget(db: AsyncSession = Depends(get_db)):
|
||||
res = await db.execute(select(BudgetRecord).filter_by(id="default"))
|
||||
return res.scalar_one_or_none()
|
||||
|
||||
|
||||
# Experiments
|
||||
@router.get("/api/experiments")
|
||||
async def list_experiments(db: AsyncSession = Depends(get_db)):
|
||||
res = await db.execute(select(ExperimentRecord))
|
||||
return res.scalars().all()
|
||||
|
||||
|
||||
# Health Check
|
||||
@router.get("/api/health")
|
||||
async def get_health():
|
||||
return {
|
||||
"status": "healthy",
|
||||
"service": "Model Router AI Control Room",
|
||||
"version": "1.0.0",
|
||||
"environment": settings.APP_ENV,
|
||||
}
|
||||
@@ -0,0 +1,39 @@
|
||||
from typing import Tuple
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
from app.storage.models import BudgetRecord
|
||||
|
||||
|
||||
async def check_budget_threshold(
|
||||
session: AsyncSession,
|
||||
cost_to_add: float = 0.0,
|
||||
) -> Tuple[float, str]:
|
||||
"""
|
||||
Checks current budget utilization percentage and determines intervention mode:
|
||||
- NORMAL (< 80%)
|
||||
- OPTIMIZE_80 (>= 80% and < 95%)
|
||||
- CHEAP_95 (>= 95% and < 100%)
|
||||
- BLOCK_100 (>= 100%)
|
||||
"""
|
||||
res = await session.execute(select(BudgetRecord).filter_by(id="default"))
|
||||
budget = res.scalar_one_or_none()
|
||||
if not budget:
|
||||
return 0.0, "NORMAL"
|
||||
|
||||
monthly_limit = budget.monthly_limit or 100.0
|
||||
current_spend = budget.current_monthly_spend + cost_to_add
|
||||
percent = (current_spend / monthly_limit) * 100.0 if monthly_limit > 0 else 0.0
|
||||
|
||||
mode = "NORMAL"
|
||||
if percent >= 100.0:
|
||||
mode = "BLOCK_100"
|
||||
elif percent >= 95.0:
|
||||
mode = "CHEAP_95"
|
||||
elif percent >= 80.0:
|
||||
mode = "OPTIMIZE_80"
|
||||
|
||||
budget.current_monthly_spend = current_spend
|
||||
budget.intervention_mode = mode
|
||||
await session.commit()
|
||||
|
||||
return round(percent, 2), mode
|
||||
@@ -0,0 +1,269 @@
|
||||
import sys
|
||||
import os
|
||||
import asyncio
|
||||
import httpx
|
||||
import typer
|
||||
from rich.console import Console
|
||||
from rich.table import Table
|
||||
from rich.panel import Panel
|
||||
from rich.text import Text
|
||||
|
||||
# Set UTF-8 encoding support for Windows terminals
|
||||
if sys.platform.startswith("win"):
|
||||
try:
|
||||
sys.stdout.reconfigure(encoding="utf-8")
|
||||
sys.stderr.reconfigure(encoding="utf-8")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")))
|
||||
|
||||
from app.config.settings import get_settings
|
||||
from app.analyzer.heuristics import analyze_request_heuristics
|
||||
from app.router.engine import route_request
|
||||
from app.storage.database import init_db, AsyncSessionLocal
|
||||
from app.storage.models import ModelRecord, RoutingPolicyRecord
|
||||
from app.api.routes import db_model_to_meta
|
||||
from app.fallback.handler import execute_with_fallback
|
||||
from sqlalchemy import select
|
||||
|
||||
app = typer.Typer(help="Model Router: Intelligent LLM Request Routing Platform")
|
||||
console = Console(safe_box=True)
|
||||
settings = get_settings()
|
||||
|
||||
|
||||
@app.command()
|
||||
def doctor():
|
||||
"""Run environment, connectivity, and dependency health checks."""
|
||||
console.print(Panel.fit("[bold cyan]Model Router Diagnostics (modelrouter doctor)[/bold cyan]"))
|
||||
|
||||
# 1. Python Check
|
||||
py_ver = f"{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}"
|
||||
console.print(f" [green][OK][/green] Python: [bold]{py_ver}[/bold]")
|
||||
|
||||
# 2. Database Check
|
||||
try:
|
||||
asyncio.run(init_db())
|
||||
console.print(" [green][OK][/green] SQLite Database: [bold]Initialized & Ready[/bold]")
|
||||
except Exception as exc:
|
||||
console.print(f" [red][FAIL][/red] Database Error: {exc}")
|
||||
|
||||
# 3. Ollama Connectivity Check
|
||||
ollama_ok = False
|
||||
try:
|
||||
r = httpx.get(f"{settings.OLLAMA_BASE_URL}/api/tags", timeout=1.5)
|
||||
if r.status_code == 200:
|
||||
models = [m.get("name") for m in r.json().get("models", [])]
|
||||
console.print(f" [green][OK][/green] Ollama (Local): [bold]Connected[/bold] ({len(models)} local models found)")
|
||||
ollama_ok = True
|
||||
except Exception:
|
||||
pass
|
||||
if not ollama_ok:
|
||||
console.print(f" [yellow][!][/yellow] Ollama (Local): [dim]Not running on {settings.OLLAMA_BASE_URL} (Optional, local fallback available)[/dim]")
|
||||
|
||||
# 4. Mock Engine
|
||||
console.print(" [green][OK][/green] Mock Engine: [bold]Ready (Zero API keys required)[/bold]")
|
||||
|
||||
# 5. External Providers
|
||||
if settings.OPENAI_API_KEY:
|
||||
console.print(" [green][OK][/green] OpenAI Provider: [bold]API Key Configured in .env[/bold]")
|
||||
else:
|
||||
console.print(" [yellow][!][/yellow] OpenAI Provider: [dim]Not configured (Optional)[/dim]")
|
||||
|
||||
if settings.ANTHROPIC_API_KEY:
|
||||
console.print(" [green][OK][/green] Anthropic Provider: [bold]API Key Configured in .env[/bold]")
|
||||
else:
|
||||
console.print(" [yellow][!][/yellow] Anthropic Provider: [dim]Not configured (Optional)[/dim]")
|
||||
|
||||
if settings.GEMINI_API_KEY:
|
||||
console.print(" [green][OK][/green] Gemini Provider: [bold]API Key Configured in .env[/bold]")
|
||||
else:
|
||||
console.print(" [yellow][!][/yellow] Gemini Provider: [dim]Not configured (Optional)[/dim]")
|
||||
|
||||
console.print("\n[bold green]Ready for local-first intelligent routing![/bold green]\n")
|
||||
|
||||
|
||||
@app.command()
|
||||
def models():
|
||||
"""List all registered candidate models and their capability scores."""
|
||||
async def _list():
|
||||
await init_db()
|
||||
async with AsyncSessionLocal() as db:
|
||||
res = await db.execute(select(ModelRecord).order_by(ModelRecord.tier, ModelRecord.name))
|
||||
records = res.scalars().all()
|
||||
|
||||
table = Table(title="Model Router Registry", border_style="bright_blue")
|
||||
table.add_column("Model ID", style="cyan", no_wrap=True)
|
||||
table.add_column("Tier", style="magenta")
|
||||
table.add_column("Provider", style="green")
|
||||
table.add_column("Context", style="yellow")
|
||||
table.add_column("Quality", style="blue")
|
||||
table.add_column("Speed", style="cyan")
|
||||
table.add_column("Cost / 1K In/Out", style="white")
|
||||
|
||||
for m in records:
|
||||
cost_str = f"${m.cost_per_input_token*1000:.4f} / ${m.cost_per_output_token*1000:.4f}" if m.cost_per_input_token > 0 else "Free ($0)"
|
||||
table.add_row(
|
||||
m.id,
|
||||
m.tier,
|
||||
m.provider,
|
||||
f"{m.context_window:,}",
|
||||
f"{m.quality_score*100:.0f}%",
|
||||
f"{m.speed_score*100:.0f}%",
|
||||
cost_str,
|
||||
)
|
||||
console.print(table)
|
||||
|
||||
asyncio.run(_list())
|
||||
|
||||
|
||||
@app.command()
|
||||
def route(
|
||||
prompt: str = typer.Argument(..., help="Prompt text to route"),
|
||||
policy: str = typer.Option("balanced", "--policy", "-p", help="Routing policy to use"),
|
||||
):
|
||||
"""Analyze prompt and explain why the optimal model was selected (Dry Run)."""
|
||||
async def _route():
|
||||
await init_db()
|
||||
analysis = analyze_request_heuristics(prompt)
|
||||
|
||||
async with AsyncSessionLocal() as db:
|
||||
res_m = await db.execute(select(ModelRecord).filter_by(is_active=True))
|
||||
models = [db_model_to_meta(m) for m in res_m.scalars().all()]
|
||||
|
||||
res_p = await db.execute(select(RoutingPolicyRecord).filter_by(id=policy))
|
||||
pol = res_p.scalar_one_or_none()
|
||||
weights = {
|
||||
"quality_weight": pol.quality_weight if pol else 0.35,
|
||||
"cost_weight": pol.cost_weight if pol else 0.25,
|
||||
"speed_weight": pol.speed_weight if pol else 0.20,
|
||||
"capability_weight": pol.capability_weight if pol else 0.15,
|
||||
"reliability_weight": pol.reliability_weight if pol else 0.05,
|
||||
}
|
||||
|
||||
decision = route_request(
|
||||
analysis=analysis,
|
||||
available_models=models,
|
||||
policy_weights=weights,
|
||||
policy_name=policy,
|
||||
)
|
||||
|
||||
console.print(Panel.fit(
|
||||
f"[bold white]Prompt:[/bold white] \"{prompt}\"\n"
|
||||
f"[bold cyan]Task Type:[/bold cyan] {analysis.task_type.value} | [bold cyan]Complexity:[/bold cyan] {analysis.complexity_label.value} ({analysis.complexity:.2f})\n"
|
||||
f"[bold green]Selected Model:[/bold green] [bold yellow]{decision.selected_model_name}[/bold yellow] ({decision.selected_model})\n"
|
||||
f"[bold magenta]Confidence:[/bold magenta] {decision.confidence * 100:.0f}%\n"
|
||||
f"[bold blue]Estimated Latency:[/bold blue] {decision.estimated_latency_ms:.0f}ms | [bold blue]Estimated Cost:[/bold blue] ${decision.estimated_cost_usd:.6f}",
|
||||
title="ROUTING DECISION",
|
||||
border_style="bright_blue",
|
||||
))
|
||||
|
||||
console.print("[bold cyan]Why this model?[/bold cyan]")
|
||||
for r in decision.reasons:
|
||||
console.print(f" [green]+[/green] {r}")
|
||||
|
||||
if decision.rejected_candidates:
|
||||
console.print("\n[bold red]Rejected Candidates:[/bold red]")
|
||||
for cid, reason in decision.rejected_candidates.items():
|
||||
console.print(f" [dim]- {cid}: {reason}[/dim]")
|
||||
|
||||
asyncio.run(_route())
|
||||
|
||||
|
||||
@app.command()
|
||||
def run(
|
||||
prompt: str = typer.Argument(..., help="Prompt to route and execute"),
|
||||
policy: str = typer.Option("balanced", "--policy", "-p", help="Routing policy to use"),
|
||||
):
|
||||
"""Route request and execute inference through the chosen provider."""
|
||||
async def _run():
|
||||
await init_db()
|
||||
analysis = analyze_request_heuristics(prompt)
|
||||
|
||||
async with AsyncSessionLocal() as db:
|
||||
res_m = await db.execute(select(ModelRecord).filter_by(is_active=True))
|
||||
models = [db_model_to_meta(m) for m in res_m.scalars().all()]
|
||||
|
||||
decision = route_request(
|
||||
analysis=analysis,
|
||||
available_models=models,
|
||||
policy_weights={"quality_weight": 0.35, "cost_weight": 0.25, "speed_weight": 0.20, "capability_weight": 0.15, "reliability_weight": 0.05},
|
||||
policy_name=policy,
|
||||
)
|
||||
|
||||
console.print(f"[bold cyan]Routing to:[/bold cyan] {decision.selected_model} via provider: {decision.provider}...")
|
||||
|
||||
resp, fallback_used, orig_m, fb_reason = await execute_with_fallback(
|
||||
prompt=prompt,
|
||||
selected_model_id=decision.selected_model,
|
||||
selected_provider_id=decision.provider,
|
||||
all_models=models,
|
||||
)
|
||||
|
||||
console.print(Panel(
|
||||
resp.content,
|
||||
title=f"RESPONSE from {resp.model} ({'FALLBACK: ' + orig_m if fallback_used else 'PRIMARY'})",
|
||||
subtitle=f"Latency: {resp.provider_latency_ms:.0f}ms | Tokens: {resp.total_tokens}",
|
||||
border_style="green" if not fallback_used else "yellow",
|
||||
))
|
||||
|
||||
asyncio.run(_run())
|
||||
|
||||
|
||||
@app.command()
|
||||
def traffic(limit: int = 10):
|
||||
"""View recent routed traffic log."""
|
||||
async def _traffic():
|
||||
await init_db()
|
||||
from app.storage.models import RequestRecord
|
||||
from sqlalchemy import desc
|
||||
async with AsyncSessionLocal() as db:
|
||||
res = await db.execute(select(RequestRecord).order_by(desc(RequestRecord.timestamp)).limit(limit))
|
||||
records = res.scalars().all()
|
||||
|
||||
table = Table(title="Recent Routed Traffic Feed", border_style="bright_blue")
|
||||
table.add_column("Time", style="dim")
|
||||
table.add_column("Task", style="cyan")
|
||||
table.add_column("Selected Model", style="green")
|
||||
table.add_column("Latency", style="yellow")
|
||||
table.add_column("Cost", style="magenta")
|
||||
table.add_column("Prompt Sample", style="white")
|
||||
|
||||
for r in records:
|
||||
t_str = r.timestamp.strftime("%H:%M:%S") if r.timestamp else "N/A"
|
||||
table.add_row(
|
||||
t_str,
|
||||
r.task_type,
|
||||
r.selected_model,
|
||||
f"{r.total_latency_ms:.0f}ms",
|
||||
f"${r.estimated_cost:.5f}",
|
||||
r.prompt[:40] + ("..." if len(r.prompt) > 40 else ""),
|
||||
)
|
||||
console.print(table)
|
||||
|
||||
asyncio.run(_traffic())
|
||||
|
||||
|
||||
@app.command()
|
||||
def analytics():
|
||||
"""Display system-wide routing performance and cost savings analytics."""
|
||||
async def _analytics():
|
||||
await init_db()
|
||||
from app.analytics.service import get_system_analytics
|
||||
async with AsyncSessionLocal() as db:
|
||||
data = await get_system_analytics(db)
|
||||
console.print(Panel.fit(
|
||||
f"[bold cyan]Total Routed Requests:[/bold cyan] {data['total_requests']}\n"
|
||||
f"[bold green]Average Latency:[/bold green] {data['avg_latency_ms']}ms (Routing overhead: {data['avg_routing_latency_ms']}ms)\n"
|
||||
f"[bold yellow]Total Cost:[/bold yellow] ${data['total_cost_usd']:.6f}\n"
|
||||
f"[bold magenta]Cost Saved vs Baseline:[/bold magenta] ${data['savings']['cost_saved_usd']:.6f} ({data['savings']['savings_percentage']}%)\n"
|
||||
f"[bold blue]Fallback Rate:[/bold blue] {data['fallback_rate_percent']}%\n",
|
||||
title="MODEL ROUTER ANALYTICS OVERVIEW",
|
||||
border_style="bright_blue",
|
||||
))
|
||||
|
||||
asyncio.run(_analytics())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
app()
|
||||
@@ -0,0 +1,43 @@
|
||||
from functools import lru_cache
|
||||
from typing import Optional
|
||||
from pydantic_settings import BaseSettings, SettingsConfigDict
|
||||
|
||||
|
||||
class Settings(BaseSettings):
|
||||
APP_ENV: str = "development"
|
||||
DEBUG: bool = True
|
||||
PORT: int = 8000
|
||||
HOST: str = "0.0.0.0"
|
||||
|
||||
DATABASE_URL: str = "sqlite+aiosqlite:///./model_router.db"
|
||||
ROUTER_ANALYZER: str = "rules" # "rules" or "llm"
|
||||
DEFAULT_ROUTING_POLICY: str = "balanced"
|
||||
BASELINE_MODEL_ID: str = "mock-power"
|
||||
|
||||
DEFAULT_PROVIDER: str = "mock"
|
||||
OLLAMA_BASE_URL: str = "http://localhost:11434"
|
||||
|
||||
# External Provider Keys (optional)
|
||||
OPENAI_API_KEY: Optional[str] = None
|
||||
ANTHROPIC_API_KEY: Optional[str] = None
|
||||
GEMINI_API_KEY: Optional[str] = None
|
||||
|
||||
# Budgets
|
||||
DAILY_BUDGET: float = 10.00
|
||||
MONTHLY_BUDGET: float = 100.00
|
||||
PER_REQUEST_BUDGET: float = 1.00
|
||||
|
||||
# Execution / Reliability
|
||||
MAX_RETRIES: int = 2
|
||||
PROVIDER_TIMEOUT_SECONDS: float = 30.0
|
||||
|
||||
model_config = SettingsConfigDict(
|
||||
env_file=".env",
|
||||
env_file_encoding="utf-8",
|
||||
extra="ignore",
|
||||
)
|
||||
|
||||
|
||||
@lru_cache()
|
||||
def get_settings() -> Settings:
|
||||
return Settings()
|
||||
@@ -0,0 +1,85 @@
|
||||
import uuid
|
||||
import datetime
|
||||
from typing import Dict, Any, List
|
||||
from sqlalchemy import select, func
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
from app.storage.models import ExperimentRecord, ExperimentRunRecord
|
||||
|
||||
|
||||
async def log_experiment_run(
|
||||
session: AsyncSession,
|
||||
experiment_id: str,
|
||||
request_id: str,
|
||||
policy: str,
|
||||
model: str,
|
||||
cost: float,
|
||||
latency_ms: float,
|
||||
):
|
||||
run = ExperimentRunRecord(
|
||||
id=str(uuid.uuid4()),
|
||||
experiment_id=experiment_id,
|
||||
request_id=request_id,
|
||||
assigned_policy=policy,
|
||||
selected_model=model,
|
||||
cost=cost,
|
||||
latency_ms=latency_ms,
|
||||
)
|
||||
session.add(run)
|
||||
|
||||
# Increment sample count
|
||||
res = await session.execute(select(ExperimentRecord).filter_by(id=experiment_id))
|
||||
exp = res.scalar_one_or_none()
|
||||
if exp:
|
||||
exp.sample_count += 1
|
||||
await session.commit()
|
||||
|
||||
|
||||
async def get_experiment_summary(
|
||||
session: AsyncSession,
|
||||
experiment_id: str,
|
||||
) -> Dict[str, Any]:
|
||||
res = await session.execute(select(ExperimentRecord).filter_by(id=experiment_id))
|
||||
exp = res.scalar_one_or_none()
|
||||
if not exp:
|
||||
return {}
|
||||
|
||||
# Calculate metrics for Policy A vs Policy B
|
||||
runs_a = await session.execute(
|
||||
select(
|
||||
func.count(ExperimentRunRecord.id),
|
||||
func.avg(ExperimentRunRecord.cost),
|
||||
func.avg(ExperimentRunRecord.latency_ms),
|
||||
).filter_by(experiment_id=experiment_id, assigned_policy=exp.policy_a)
|
||||
)
|
||||
count_a, cost_a, lat_a = runs_a.one()
|
||||
|
||||
runs_b = await session.execute(
|
||||
select(
|
||||
func.count(ExperimentRunRecord.id),
|
||||
func.avg(ExperimentRunRecord.cost),
|
||||
func.avg(ExperimentRunRecord.latency_ms),
|
||||
).filter_by(experiment_id=experiment_id, assigned_policy=exp.policy_b)
|
||||
)
|
||||
count_b, cost_b, lat_b = runs_b.one()
|
||||
|
||||
return {
|
||||
"id": exp.id,
|
||||
"name": exp.name,
|
||||
"description": exp.description,
|
||||
"policy_a": exp.policy_a,
|
||||
"policy_b": exp.policy_b,
|
||||
"status": exp.status,
|
||||
"total_samples": exp.sample_count,
|
||||
"results": {
|
||||
exp.policy_a: {
|
||||
"sample_count": count_a or 0,
|
||||
"avg_cost_usd": round(cost_a or 0.0, 6),
|
||||
"avg_latency_ms": round(lat_a or 0.0, 1),
|
||||
},
|
||||
exp.policy_b: {
|
||||
"sample_count": count_b or 0,
|
||||
"avg_cost_usd": round(cost_b or 0.0, 6),
|
||||
"avg_latency_ms": round(lat_b or 0.0, 1),
|
||||
},
|
||||
},
|
||||
}
|
||||
@@ -0,0 +1,105 @@
|
||||
import time
|
||||
import asyncio
|
||||
from typing import List, Optional, Tuple, Dict, Any
|
||||
from app.models.schemas import ProviderResponse, ModelMetadata
|
||||
from app.providers.registry import provider_registry
|
||||
from app.config.settings import get_settings
|
||||
|
||||
settings = get_settings()
|
||||
|
||||
RETRYABLE_ERRORS = [
|
||||
"timeout",
|
||||
"connection error",
|
||||
"rate limit",
|
||||
"429",
|
||||
"503",
|
||||
"service unavailable",
|
||||
"econnreset",
|
||||
]
|
||||
|
||||
|
||||
def is_retryable(error_str: Optional[str]) -> bool:
|
||||
if not error_str:
|
||||
return False
|
||||
low = error_str.lower()
|
||||
return any(keyword in low for keyword in RETRYABLE_ERRORS)
|
||||
|
||||
|
||||
async def execute_with_fallback(
|
||||
prompt: str,
|
||||
selected_model_id: str,
|
||||
selected_provider_id: str,
|
||||
all_models: List[ModelMetadata],
|
||||
system_prompt: Optional[str] = None,
|
||||
temperature: float = 0.7,
|
||||
max_retries: int = 1,
|
||||
) -> Tuple[ProviderResponse, bool, Optional[str], Optional[str]]:
|
||||
"""
|
||||
Executes model generation with configurable retries and fallback hierarchy.
|
||||
Returns: (ProviderResponse, fallback_used, original_model, fallback_reason)
|
||||
"""
|
||||
current_model_id = selected_model_id
|
||||
current_provider_id = selected_provider_id
|
||||
fallback_used = False
|
||||
original_model = selected_model_id
|
||||
fallback_reason = None
|
||||
|
||||
# Step 1: Attempt primary model with retries
|
||||
provider = provider_registry.get_provider(current_provider_id)
|
||||
if not provider:
|
||||
provider = provider_registry.get_provider("mock")
|
||||
current_provider_id = "mock"
|
||||
current_model_id = "mock-balanced"
|
||||
fallback_used = True
|
||||
fallback_reason = f"Provider '{selected_provider_id}' is not loaded. Fell back to Mock."
|
||||
|
||||
for attempt in range(max_retries + 1):
|
||||
resp = await provider.generate(
|
||||
prompt=prompt,
|
||||
model_id=current_model_id,
|
||||
system_prompt=system_prompt,
|
||||
temperature=temperature,
|
||||
)
|
||||
|
||||
if not resp.error and resp.finish_reason != "error":
|
||||
return resp, fallback_used, original_model if fallback_used else None, fallback_reason
|
||||
|
||||
# If non-retryable error or last attempt, prepare for fallback
|
||||
if not is_retryable(resp.error) or attempt == max_retries:
|
||||
fallback_reason = resp.error or "Unknown primary model execution error."
|
||||
break
|
||||
|
||||
await asyncio.sleep(0.5 * (attempt + 1))
|
||||
|
||||
# Step 2: Determine Fallback Candidate
|
||||
fallback_used = True
|
||||
# Priority: Mock model of same or balanced tier, or active local Ollama
|
||||
candidate_fallbacks = [
|
||||
m for m in all_models
|
||||
if m.id != selected_model_id and m.is_active and (m.provider == "mock" or m.type == "LOCAL")
|
||||
]
|
||||
if not candidate_fallbacks:
|
||||
candidate_fallbacks = [m for m in all_models if m.id != selected_model_id and m.is_active]
|
||||
|
||||
fallback_target = candidate_fallbacks[0] if candidate_fallbacks else None
|
||||
|
||||
if fallback_target:
|
||||
fb_provider = provider_registry.get_provider(fallback_target.provider) or provider_registry.get_provider("mock")
|
||||
fb_resp = await fb_provider.generate(
|
||||
prompt=prompt,
|
||||
model_id=fallback_target.id,
|
||||
system_prompt=system_prompt,
|
||||
temperature=temperature,
|
||||
)
|
||||
if not fb_resp.error:
|
||||
return fb_resp, True, original_model, f"Primary model failed: {fallback_reason} -> Fallback to {fallback_target.id}"
|
||||
|
||||
# Step 3: Emergency Hard Mock fallback
|
||||
emergency_mock = provider_registry.get_provider("mock")
|
||||
emergency_resp = await emergency_mock.generate(
|
||||
prompt=prompt,
|
||||
model_id="mock-balanced",
|
||||
system_prompt=system_prompt,
|
||||
temperature=temperature,
|
||||
)
|
||||
return emergency_resp, True, original_model, f"Emergency fallback triggered due to: {fallback_reason}"
|
||||
@@ -0,0 +1,127 @@
|
||||
from enum import Enum
|
||||
from typing import List, Optional, Dict, Any
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
|
||||
class TaskType(str, Enum):
|
||||
GENERAL_QA = "GENERAL_QA"
|
||||
CODING = "CODING"
|
||||
DEBUGGING = "DEBUGGING"
|
||||
REASONING = "REASONING"
|
||||
SUMMARIZATION = "SUMMARIZATION"
|
||||
EXTRACTION = "EXTRACTION"
|
||||
WRITING = "WRITING"
|
||||
TRANSLATION = "TRANSLATION"
|
||||
ANALYSIS = "ANALYSIS"
|
||||
MATH = "MATH"
|
||||
LONG_CONTEXT = "LONG_CONTEXT"
|
||||
CREATIVE = "CREATIVE"
|
||||
|
||||
|
||||
class PriorityLevel(str, Enum):
|
||||
LOW = "LOW"
|
||||
MEDIUM = "MEDIUM"
|
||||
HIGH = "HIGH"
|
||||
|
||||
|
||||
class ModelTier(str, Enum):
|
||||
FAST = "FAST"
|
||||
BALANCED = "BALANCED"
|
||||
POWER = "POWER"
|
||||
|
||||
|
||||
class ProviderType(str, Enum):
|
||||
MOCK = "mock"
|
||||
OLLAMA = "ollama"
|
||||
OPENAI = "openai"
|
||||
ANTHROPIC = "anthropic"
|
||||
GEMINI = "gemini"
|
||||
CUSTOM = "custom"
|
||||
|
||||
|
||||
class RequestAnalysis(BaseModel):
|
||||
task_type: TaskType = Field(default=TaskType.GENERAL_QA)
|
||||
complexity: float = Field(default=0.5, ge=0.0, le=1.0)
|
||||
complexity_label: PriorityLevel = Field(default=PriorityLevel.MEDIUM)
|
||||
reasoning_required: bool = Field(default=False)
|
||||
coding_required: bool = Field(default=False)
|
||||
vision_required: bool = Field(default=False)
|
||||
tools_required: bool = Field(default=False)
|
||||
context_size: int = Field(default=0) # in estimated tokens / characters
|
||||
latency_priority: PriorityLevel = Field(default=PriorityLevel.MEDIUM)
|
||||
cost_sensitivity: PriorityLevel = Field(default=PriorityLevel.MEDIUM)
|
||||
quality_requirement: PriorityLevel = Field(default=PriorityLevel.MEDIUM)
|
||||
keywords_detected: List[str] = Field(default_factory=list)
|
||||
analyzer_used: str = Field(default="rules")
|
||||
|
||||
|
||||
class ModelMetadata(BaseModel):
|
||||
id: str
|
||||
name: str
|
||||
provider: str
|
||||
type: str = "LOCAL" # LOCAL, CLOUD, MOCK
|
||||
tier: ModelTier = ModelTier.BALANCED
|
||||
context_window: int = 8192
|
||||
|
||||
supports_coding: bool = False
|
||||
supports_reasoning: bool = False
|
||||
supports_vision: bool = False
|
||||
supports_tools: bool = False
|
||||
|
||||
quality_score: float = 0.75 # 0.0 - 1.0
|
||||
speed_score: float = 0.75 # 0.0 - 1.0
|
||||
reliability_score: float = 0.95
|
||||
|
||||
cost_per_input_token: float = 0.0
|
||||
cost_per_output_token: float = 0.0
|
||||
availability: str = "AVAILABLE"
|
||||
is_active: bool = True
|
||||
|
||||
|
||||
class CandidateScore(BaseModel):
|
||||
model_id: str
|
||||
model_name: str
|
||||
provider: str
|
||||
tier: str
|
||||
overall_score: float # 0 - 100
|
||||
quality_component: float
|
||||
cost_component: float
|
||||
speed_component: float
|
||||
capability_component: float
|
||||
reliability_component: float
|
||||
estimated_latency_ms: float
|
||||
estimated_cost_usd: float
|
||||
eligible: bool = True
|
||||
rejection_reason: Optional[str] = None
|
||||
|
||||
|
||||
class RoutingDecision(BaseModel):
|
||||
decision_id: str
|
||||
request_id: str
|
||||
selected_model: str
|
||||
selected_model_name: str
|
||||
provider: str
|
||||
tier: str
|
||||
confidence: float
|
||||
policy_used: str
|
||||
reasons: List[str]
|
||||
candidate_scores: List[CandidateScore]
|
||||
rejected_candidates: Dict[str, str] = Field(default_factory=dict)
|
||||
estimated_cost_usd: float
|
||||
estimated_latency_ms: float
|
||||
rule_applied: Optional[str] = None
|
||||
timestamp: str
|
||||
|
||||
|
||||
class ProviderResponse(BaseModel):
|
||||
content: str
|
||||
finish_reason: str = "stop"
|
||||
input_tokens: int = 0
|
||||
output_tokens: int = 0
|
||||
total_tokens: int = 0
|
||||
provider_latency_ms: float = 0.0
|
||||
time_to_first_token_ms: Optional[float] = None
|
||||
is_mock: bool = False
|
||||
model: str = ""
|
||||
provider: str = ""
|
||||
error: Optional[str] = None
|
||||
@@ -0,0 +1,43 @@
|
||||
import json
|
||||
import logging
|
||||
from typing import Dict, Any, Optional
|
||||
import datetime
|
||||
|
||||
logger = logging.getLogger("model_router.observability")
|
||||
logger.setLevel(logging.INFO)
|
||||
if not logger.handlers:
|
||||
handler = logging.StreamHandler()
|
||||
formatter = logging.Formatter('{"time":"%(asctime)s", "level":"%(levelname)s", "event":%(message)s}')
|
||||
handler.setFormatter(formatter)
|
||||
logger.addHandler(handler)
|
||||
|
||||
|
||||
def log_router_event(
|
||||
event_name: str,
|
||||
request_id: str,
|
||||
model: Optional[str] = None,
|
||||
provider: Optional[str] = None,
|
||||
duration_ms: Optional[float] = None,
|
||||
metadata: Optional[Dict[str, Any]] = None,
|
||||
):
|
||||
"""
|
||||
Structured event logger for model router pipeline.
|
||||
Redacts sensitive credentials automatically.
|
||||
"""
|
||||
payload = {
|
||||
"event": event_name,
|
||||
"request_id": request_id,
|
||||
"timestamp": datetime.datetime.utcnow().isoformat(),
|
||||
"model": model,
|
||||
"provider": provider,
|
||||
"duration_ms": duration_ms,
|
||||
}
|
||||
if metadata:
|
||||
# Sanitize any key that looks like an API token
|
||||
safe_meta = {
|
||||
k: ("[REDACTED]" if "key" in k.lower() or "secret" in k.lower() or "token" in k.lower() and "count" not in k.lower() else v)
|
||||
for k, v in metadata.items()
|
||||
}
|
||||
payload["metadata"] = safe_meta
|
||||
|
||||
logger.info(json.dumps(payload))
|
||||
@@ -0,0 +1,46 @@
|
||||
import abc
|
||||
import time
|
||||
import asyncio
|
||||
from typing import AsyncGenerator, Dict, Any, Optional
|
||||
from app.models.schemas import ProviderResponse
|
||||
|
||||
|
||||
class ModelProvider(abc.ABC):
|
||||
"""
|
||||
Abstract Base Class for all Model Providers (Mock, Ollama, OpenAI, Anthropic, Gemini, etc.)
|
||||
Ensures complete provider decoupling from routing logic.
|
||||
"""
|
||||
|
||||
def __init__(self, provider_id: str, name: str, base_url: Optional[str] = None):
|
||||
self.provider_id = provider_id
|
||||
self.name = name
|
||||
self.base_url = base_url
|
||||
|
||||
@abc.abstractmethod
|
||||
async def generate(
|
||||
self,
|
||||
prompt: str,
|
||||
model_id: str,
|
||||
system_prompt: Optional[str] = None,
|
||||
temperature: float = 0.7,
|
||||
max_tokens: Optional[int] = None,
|
||||
) -> ProviderResponse:
|
||||
"""Execute a text generation call against the model provider."""
|
||||
pass
|
||||
|
||||
@abc.abstractmethod
|
||||
async def stream(
|
||||
self,
|
||||
prompt: str,
|
||||
model_id: str,
|
||||
system_prompt: Optional[str] = None,
|
||||
temperature: float = 0.7,
|
||||
max_tokens: Optional[int] = None,
|
||||
) -> AsyncGenerator[str, None]:
|
||||
"""Stream generated response chunks."""
|
||||
pass
|
||||
|
||||
@abc.abstractmethod
|
||||
async def check_health(self) -> Dict[str, Any]:
|
||||
"""Check provider connectivity, status, and loaded models."""
|
||||
pass
|
||||
@@ -0,0 +1,321 @@
|
||||
import time
|
||||
import httpx
|
||||
from typing import AsyncGenerator, Dict, Any, Optional
|
||||
from app.providers.base import ModelProvider
|
||||
from app.models.schemas import ProviderResponse
|
||||
from app.config.settings import get_settings
|
||||
|
||||
settings = get_settings()
|
||||
|
||||
|
||||
class OpenAIProvider(ModelProvider):
|
||||
"""
|
||||
OpenAI Provider Adapter.
|
||||
Uses OPENAI_API_KEY from environment variables. Secrets are never exposed.
|
||||
"""
|
||||
|
||||
def __init__(self, api_key: Optional[str] = None):
|
||||
super().__init__(provider_id="openai", name="OpenAI", base_url="https://api.openai.com/v1")
|
||||
self.api_key = api_key or settings.OPENAI_API_KEY
|
||||
|
||||
async def generate(
|
||||
self,
|
||||
prompt: str,
|
||||
model_id: str,
|
||||
system_prompt: Optional[str] = None,
|
||||
temperature: float = 0.7,
|
||||
max_tokens: Optional[int] = None,
|
||||
) -> ProviderResponse:
|
||||
if not self.api_key:
|
||||
return ProviderResponse(
|
||||
content="",
|
||||
finish_reason="error",
|
||||
error="OPENAI_API_KEY not configured in environment.",
|
||||
is_mock=False,
|
||||
model=model_id,
|
||||
provider="openai",
|
||||
)
|
||||
|
||||
endpoint = f"{self.base_url}/chat/completions"
|
||||
headers = {
|
||||
"Authorization": f"Bearer {self.api_key}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
messages = []
|
||||
if system_prompt:
|
||||
messages.append({"role": "system", "content": system_prompt})
|
||||
messages.append({"role": "user", "content": prompt})
|
||||
|
||||
payload: Dict[str, Any] = {
|
||||
"model": model_id,
|
||||
"messages": messages,
|
||||
"temperature": temperature,
|
||||
}
|
||||
if max_tokens:
|
||||
payload["max_tokens"] = max_tokens
|
||||
|
||||
start_time = time.perf_counter()
|
||||
try:
|
||||
async with httpx.AsyncClient(timeout=settings.PROVIDER_TIMEOUT_SECONDS) as client:
|
||||
resp = await client.post(endpoint, json=payload, headers=headers)
|
||||
resp.raise_for_status()
|
||||
data = resp.json()
|
||||
|
||||
elapsed_ms = (time.perf_counter() - start_time) * 1000.0
|
||||
choice = data["choices"][0]
|
||||
usage = data.get("usage", {})
|
||||
|
||||
return ProviderResponse(
|
||||
content=choice["message"]["content"],
|
||||
finish_reason=choice.get("finish_reason", "stop"),
|
||||
input_tokens=usage.get("prompt_tokens", 0),
|
||||
output_tokens=usage.get("completion_tokens", 0),
|
||||
total_tokens=usage.get("total_tokens", 0),
|
||||
provider_latency_ms=round(elapsed_ms, 2),
|
||||
is_mock=False,
|
||||
model=model_id,
|
||||
provider="openai",
|
||||
)
|
||||
except Exception as exc:
|
||||
return ProviderResponse(
|
||||
content="",
|
||||
finish_reason="error",
|
||||
error=f"OpenAI error: {str(exc)}",
|
||||
is_mock=False,
|
||||
model=model_id,
|
||||
provider="openai",
|
||||
provider_latency_ms=(time.perf_counter() - start_time) * 1000.0,
|
||||
)
|
||||
|
||||
async def stream(
|
||||
self,
|
||||
prompt: str,
|
||||
model_id: str,
|
||||
system_prompt: Optional[str] = None,
|
||||
temperature: float = 0.7,
|
||||
max_tokens: Optional[int] = None,
|
||||
) -> AsyncGenerator[str, None]:
|
||||
if not self.api_key:
|
||||
yield "[OpenAI Error: API key not configured]"
|
||||
return
|
||||
|
||||
endpoint = f"{self.base_url}/chat/completions"
|
||||
headers = {"Authorization": f"Bearer {self.api_key}"}
|
||||
messages = []
|
||||
if system_prompt:
|
||||
messages.append({"role": "system", "content": system_prompt})
|
||||
messages.append({"role": "user", "content": prompt})
|
||||
|
||||
payload = {
|
||||
"model": model_id,
|
||||
"messages": messages,
|
||||
"temperature": temperature,
|
||||
"stream": True,
|
||||
}
|
||||
|
||||
try:
|
||||
async with httpx.AsyncClient(timeout=settings.PROVIDER_TIMEOUT_SECONDS) as client:
|
||||
async with client.stream("POST", endpoint, json=payload, headers=headers) as resp:
|
||||
async for line in resp.aiter_lines():
|
||||
if line.startswith("data: ") and not line.endswith("[DONE]"):
|
||||
import json
|
||||
data = json.loads(line[6:])
|
||||
delta = data["choices"][0].get("delta", {}).get("content", "")
|
||||
if delta:
|
||||
yield delta
|
||||
except Exception as exc:
|
||||
yield f"\n[OpenAI Streaming Error: {str(exc)}]"
|
||||
|
||||
async def check_health(self) -> Dict[str, Any]:
|
||||
if not self.api_key:
|
||||
return {"provider": "openai", "status": "NOT_CONFIGURED", "message": "API key not set in .env"}
|
||||
return {"provider": "openai", "status": "CONNECTED", "message": "API key configured."}
|
||||
|
||||
|
||||
class AnthropicProvider(ModelProvider):
|
||||
"""
|
||||
Anthropic Provider Adapter.
|
||||
"""
|
||||
|
||||
def __init__(self, api_key: Optional[str] = None):
|
||||
super().__init__(provider_id="anthropic", name="Anthropic", base_url="https://api.anthropic.com/v1")
|
||||
self.api_key = api_key or settings.ANTHROPIC_API_KEY
|
||||
|
||||
async def generate(
|
||||
self,
|
||||
prompt: str,
|
||||
model_id: str,
|
||||
system_prompt: Optional[str] = None,
|
||||
temperature: float = 0.7,
|
||||
max_tokens: Optional[int] = None,
|
||||
) -> ProviderResponse:
|
||||
if not self.api_key:
|
||||
return ProviderResponse(
|
||||
content="",
|
||||
finish_reason="error",
|
||||
error="ANTHROPIC_API_KEY not configured in environment.",
|
||||
is_mock=False,
|
||||
model=model_id,
|
||||
provider="anthropic",
|
||||
)
|
||||
|
||||
endpoint = f"{self.base_url}/messages"
|
||||
headers = {
|
||||
"x-api-key": self.api_key,
|
||||
"anthropic-version": "2023-06-01",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
payload: Dict[str, Any] = {
|
||||
"model": model_id,
|
||||
"messages": [{"role": "user", "content": prompt}],
|
||||
"max_tokens": max_tokens or 2048,
|
||||
"temperature": temperature,
|
||||
}
|
||||
if system_prompt:
|
||||
payload["system"] = system_prompt
|
||||
|
||||
start_time = time.perf_counter()
|
||||
try:
|
||||
async with httpx.AsyncClient(timeout=settings.PROVIDER_TIMEOUT_SECONDS) as client:
|
||||
resp = await client.post(endpoint, json=payload, headers=headers)
|
||||
resp.raise_for_status()
|
||||
data = resp.json()
|
||||
|
||||
elapsed_ms = (time.perf_counter() - start_time) * 1000.0
|
||||
content = data["content"][0]["text"]
|
||||
usage = data.get("usage", {})
|
||||
|
||||
return ProviderResponse(
|
||||
content=content,
|
||||
finish_reason=data.get("stop_reason", "stop"),
|
||||
input_tokens=usage.get("input_tokens", 0),
|
||||
output_tokens=usage.get("output_tokens", 0),
|
||||
total_tokens=usage.get("input_tokens", 0) + usage.get("output_tokens", 0),
|
||||
provider_latency_ms=round(elapsed_ms, 2),
|
||||
is_mock=False,
|
||||
model=model_id,
|
||||
provider="anthropic",
|
||||
)
|
||||
except Exception as exc:
|
||||
return ProviderResponse(
|
||||
content="",
|
||||
finish_reason="error",
|
||||
error=f"Anthropic error: {str(exc)}",
|
||||
is_mock=False,
|
||||
model=model_id,
|
||||
provider="anthropic",
|
||||
provider_latency_ms=(time.perf_counter() - start_time) * 1000.0,
|
||||
)
|
||||
|
||||
async def stream(
|
||||
self,
|
||||
prompt: str,
|
||||
model_id: str,
|
||||
system_prompt: Optional[str] = None,
|
||||
temperature: float = 0.7,
|
||||
max_tokens: Optional[int] = None,
|
||||
) -> AsyncGenerator[str, None]:
|
||||
resp = await self.generate(prompt, model_id, system_prompt, temperature, max_tokens)
|
||||
if resp.error:
|
||||
yield f"[{resp.error}]"
|
||||
return
|
||||
for word in resp.content.split(" "):
|
||||
yield word + " "
|
||||
|
||||
async def check_health(self) -> Dict[str, Any]:
|
||||
if not self.api_key:
|
||||
return {"provider": "anthropic", "status": "NOT_CONFIGURED", "message": "API key not set in .env"}
|
||||
return {"provider": "anthropic", "status": "CONNECTED", "message": "API key configured."}
|
||||
|
||||
|
||||
class GeminiProvider(ModelProvider):
|
||||
"""
|
||||
Google Gemini Provider Adapter.
|
||||
"""
|
||||
|
||||
def __init__(self, api_key: Optional[str] = None):
|
||||
super().__init__(provider_id="gemini", name="Google Gemini", base_url="https://generativelanguage.googleapis.com/v1beta")
|
||||
self.api_key = api_key or settings.GEMINI_API_KEY
|
||||
|
||||
async def generate(
|
||||
self,
|
||||
prompt: str,
|
||||
model_id: str,
|
||||
system_prompt: Optional[str] = None,
|
||||
temperature: float = 0.7,
|
||||
max_tokens: Optional[int] = None,
|
||||
) -> ProviderResponse:
|
||||
if not self.api_key:
|
||||
return ProviderResponse(
|
||||
content="",
|
||||
finish_reason="error",
|
||||
error="GEMINI_API_KEY not configured in environment.",
|
||||
is_mock=False,
|
||||
model=model_id,
|
||||
provider="gemini",
|
||||
)
|
||||
|
||||
endpoint = f"{self.base_url}/models/{model_id}:generateContent?key={self.api_key}"
|
||||
contents = []
|
||||
if system_prompt:
|
||||
contents.append({"role": "user", "parts": [{"text": f"System context: {system_prompt}"}]})
|
||||
contents.append({"role": "user", "parts": [{"text": prompt}]})
|
||||
|
||||
payload = {
|
||||
"contents": contents,
|
||||
"generationConfig": {"temperature": temperature},
|
||||
}
|
||||
|
||||
start_time = time.perf_counter()
|
||||
try:
|
||||
async with httpx.AsyncClient(timeout=settings.PROVIDER_TIMEOUT_SECONDS) as client:
|
||||
resp = await client.post(endpoint, json=payload)
|
||||
resp.raise_for_status()
|
||||
data = resp.json()
|
||||
|
||||
elapsed_ms = (time.perf_counter() - start_time) * 1000.0
|
||||
candidate = data["candidates"][0]
|
||||
text = candidate["content"]["parts"][0]["text"]
|
||||
usage = data.get("usageMetadata", {})
|
||||
|
||||
return ProviderResponse(
|
||||
content=text,
|
||||
finish_reason=candidate.get("finishReason", "stop"),
|
||||
input_tokens=usage.get("promptTokenCount", 0),
|
||||
output_tokens=usage.get("candidatesTokenCount", 0),
|
||||
total_tokens=usage.get("totalTokenCount", 0),
|
||||
provider_latency_ms=round(elapsed_ms, 2),
|
||||
is_mock=False,
|
||||
model=model_id,
|
||||
provider="gemini",
|
||||
)
|
||||
except Exception as exc:
|
||||
return ProviderResponse(
|
||||
content="",
|
||||
finish_reason="error",
|
||||
error=f"Gemini error: {str(exc)}",
|
||||
is_mock=False,
|
||||
model=model_id,
|
||||
provider="gemini",
|
||||
provider_latency_ms=(time.perf_counter() - start_time) * 1000.0,
|
||||
)
|
||||
|
||||
async def stream(
|
||||
self,
|
||||
prompt: str,
|
||||
model_id: str,
|
||||
system_prompt: Optional[str] = None,
|
||||
temperature: float = 0.7,
|
||||
max_tokens: Optional[int] = None,
|
||||
) -> AsyncGenerator[str, None]:
|
||||
resp = await self.generate(prompt, model_id, system_prompt, temperature, max_tokens)
|
||||
if resp.error:
|
||||
yield f"[{resp.error}]"
|
||||
return
|
||||
for word in resp.content.split(" "):
|
||||
yield word + " "
|
||||
|
||||
async def check_health(self) -> Dict[str, Any]:
|
||||
if not self.api_key:
|
||||
return {"provider": "gemini", "status": "NOT_CONFIGURED", "message": "API key not set in .env"}
|
||||
return {"provider": "gemini", "status": "CONNECTED", "message": "API key configured."}
|
||||
@@ -0,0 +1,108 @@
|
||||
import asyncio
|
||||
import time
|
||||
from typing import AsyncGenerator, Dict, Any, Optional
|
||||
from app.providers.base import ModelProvider
|
||||
from app.models.schemas import ProviderResponse
|
||||
|
||||
|
||||
class MockProvider(ModelProvider):
|
||||
"""
|
||||
Mock Provider for development, testing, and zero-cost simulation.
|
||||
All outputs are explicitly labeled as DEMO/MOCK.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__(provider_id="mock", name="Mock / Demo Provider")
|
||||
|
||||
async def generate(
|
||||
self,
|
||||
prompt: str,
|
||||
model_id: str,
|
||||
system_prompt: Optional[str] = None,
|
||||
temperature: float = 0.7,
|
||||
max_tokens: Optional[int] = None,
|
||||
) -> ProviderResponse:
|
||||
start_time = time.perf_counter()
|
||||
|
||||
# Simulate realistic latency based on tier
|
||||
tier_delay = 0.08 if "fast" in model_id else (0.18 if "balanced" in model_id else 0.35)
|
||||
await asyncio.sleep(tier_delay)
|
||||
|
||||
elapsed_ms = (time.perf_counter() - start_time) * 1000.0
|
||||
|
||||
# Calculate realistic token estimations
|
||||
input_tokens = max(1, len(prompt.split()) * 2)
|
||||
|
||||
# Response tailored to task keywords
|
||||
prompt_lower = prompt.lower()
|
||||
if "python" in prompt_lower or "code" in prompt_lower or "def " in prompt_lower or "function" in prompt_lower:
|
||||
content = (
|
||||
f"[DEMO/MOCK RESPONSE from {model_id}]\n\n"
|
||||
f"```python\n"
|
||||
f"# Mock implementation generated for: {prompt[:60]}...\n"
|
||||
f"def solve_problem(*args, **kwargs):\n"
|
||||
f" \"\"\"Simulated response from Model Router Demo Engine.\"\"\"\n"
|
||||
f" result = [x for x in args if x is not None]\n"
|
||||
f" return {{'status': 'success', 'data': result, 'model': '{model_id}'}}\n"
|
||||
f"```\n\n"
|
||||
f"This code was returned by the simulated {model_id} engine."
|
||||
)
|
||||
elif "debug" in prompt_lower or "error" in prompt_lower or "bug" in prompt_lower:
|
||||
content = (
|
||||
f"[DEMO/MOCK RESPONSE from {model_id}]\n\n"
|
||||
f"### Root Cause Analysis:\n"
|
||||
f"1. **Potential Bottleneck**: Asynchronous task scheduling without proper semaphore limits.\n"
|
||||
f"2. **Mitigation**: Introduce bounded concurrency worker pools.\n"
|
||||
f"3. **Verification**: Checked against {model_id} simulated reasoning engine."
|
||||
)
|
||||
elif "summary" in prompt_lower or "summarize" in prompt_lower:
|
||||
content = (
|
||||
f"[DEMO/MOCK RESPONSE from {model_id}]\n\n"
|
||||
f"**Executive Summary**:\n"
|
||||
f"• Main point: {prompt[:80]}...\n"
|
||||
f"• Key takeaway: High routing accuracy achieved with minimal latency overhead."
|
||||
)
|
||||
else:
|
||||
content = (
|
||||
f"[DEMO/MOCK RESPONSE from {model_id}]\n\n"
|
||||
f"Processed request: \"{prompt[:100]}\"\n\n"
|
||||
f"This is a simulated response generated locally by Model Router Mock Engine without external API calls."
|
||||
)
|
||||
|
||||
output_tokens = max(1, len(content.split()) * 2)
|
||||
total_tokens = input_tokens + output_tokens
|
||||
|
||||
return ProviderResponse(
|
||||
content=content,
|
||||
finish_reason="stop",
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
total_tokens=total_tokens,
|
||||
provider_latency_ms=round(elapsed_ms, 2),
|
||||
time_to_first_token_ms=round(tier_delay * 500, 2),
|
||||
is_mock=True,
|
||||
model=model_id,
|
||||
provider="mock",
|
||||
)
|
||||
|
||||
async def stream(
|
||||
self,
|
||||
prompt: str,
|
||||
model_id: str,
|
||||
system_prompt: Optional[str] = None,
|
||||
temperature: float = 0.7,
|
||||
max_tokens: Optional[int] = None,
|
||||
) -> AsyncGenerator[str, None]:
|
||||
resp = await self.generate(prompt, model_id, system_prompt, temperature, max_tokens)
|
||||
chunks = resp.content.split(" ")
|
||||
for chunk in chunks:
|
||||
await asyncio.sleep(0.02)
|
||||
yield chunk + " "
|
||||
|
||||
async def check_health(self) -> Dict[str, Any]:
|
||||
return {
|
||||
"provider": "mock",
|
||||
"status": "CONNECTED",
|
||||
"message": "Mock provider active (Demo Mode). Zero API keys required.",
|
||||
"models_available": ["mock-fast", "mock-balanced", "mock-power"],
|
||||
}
|
||||
@@ -0,0 +1,141 @@
|
||||
import time
|
||||
import httpx
|
||||
from typing import AsyncGenerator, Dict, Any, Optional
|
||||
from app.providers.base import ModelProvider
|
||||
from app.models.schemas import ProviderResponse
|
||||
from app.config.settings import get_settings
|
||||
|
||||
settings = get_settings()
|
||||
|
||||
|
||||
class OllamaProvider(ModelProvider):
|
||||
"""
|
||||
Local Ollama Provider adapter.
|
||||
Executes real inference against locally running Ollama instance without external API keys.
|
||||
"""
|
||||
|
||||
def __init__(self, base_url: Optional[str] = None):
|
||||
url = base_url or settings.OLLAMA_BASE_URL
|
||||
super().__init__(provider_id="ollama", name="Ollama (Local)", base_url=url)
|
||||
|
||||
def _strip_model_prefix(self, model_id: str) -> str:
|
||||
# e.g. "ollama-qwen2.5-coder" -> "qwen2.5-coder"
|
||||
if model_id.startswith("ollama-"):
|
||||
return model_id[len("ollama-"):]
|
||||
return model_id
|
||||
|
||||
async def generate(
|
||||
self,
|
||||
prompt: str,
|
||||
model_id: str,
|
||||
system_prompt: Optional[str] = None,
|
||||
temperature: float = 0.7,
|
||||
max_tokens: Optional[int] = None,
|
||||
) -> ProviderResponse:
|
||||
real_model = self._strip_model_prefix(model_id)
|
||||
endpoint = f"{self.base_url.rstrip('/')}/api/generate"
|
||||
|
||||
payload: Dict[str, Any] = {
|
||||
"model": real_model,
|
||||
"prompt": prompt,
|
||||
"stream": False,
|
||||
"options": {
|
||||
"temperature": temperature,
|
||||
}
|
||||
}
|
||||
if system_prompt:
|
||||
payload["system"] = system_prompt
|
||||
if max_tokens:
|
||||
payload["options"]["num_predict"] = max_tokens
|
||||
|
||||
start_time = time.perf_counter()
|
||||
try:
|
||||
async with httpx.AsyncClient(timeout=settings.PROVIDER_TIMEOUT_SECONDS) as client:
|
||||
resp = await client.post(endpoint, json=payload)
|
||||
resp.raise_for_status()
|
||||
data = resp.json()
|
||||
|
||||
elapsed_ms = (time.perf_counter() - start_time) * 1000.0
|
||||
|
||||
content = data.get("response", "")
|
||||
input_tokens = data.get("prompt_eval_count", max(1, len(prompt.split()) * 2))
|
||||
output_tokens = data.get("eval_count", max(1, len(content.split()) * 2))
|
||||
|
||||
# Nanoseconds to ms if provided by ollama
|
||||
eval_duration_ms = data.get("eval_duration", 0) / 1_000_000.0 if "eval_duration" in data else elapsed_ms
|
||||
|
||||
return ProviderResponse(
|
||||
content=content,
|
||||
finish_reason="stop" if data.get("done") else "length",
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
total_tokens=input_tokens + output_tokens,
|
||||
provider_latency_ms=round(elapsed_ms, 2),
|
||||
time_to_first_token_ms=round(data.get("prompt_eval_duration", 0) / 1_000_000.0, 2) or None,
|
||||
is_mock=False,
|
||||
model=model_id,
|
||||
provider="ollama",
|
||||
)
|
||||
except Exception as exc:
|
||||
return ProviderResponse(
|
||||
content="",
|
||||
finish_reason="error",
|
||||
error=f"Ollama connection error ({str(exc)}). Is Ollama running on {self.base_url}?",
|
||||
is_mock=False,
|
||||
model=model_id,
|
||||
provider="ollama",
|
||||
provider_latency_ms=(time.perf_counter() - start_time) * 1000.0,
|
||||
)
|
||||
|
||||
async def stream(
|
||||
self,
|
||||
prompt: str,
|
||||
model_id: str,
|
||||
system_prompt: Optional[str] = None,
|
||||
temperature: float = 0.7,
|
||||
max_tokens: Optional[int] = None,
|
||||
) -> AsyncGenerator[str, None]:
|
||||
real_model = self._strip_model_prefix(model_id)
|
||||
endpoint = f"{self.base_url.rstrip('/')}/api/generate"
|
||||
payload = {
|
||||
"model": real_model,
|
||||
"prompt": prompt,
|
||||
"stream": True,
|
||||
"options": {"temperature": temperature},
|
||||
}
|
||||
if system_prompt:
|
||||
payload["system"] = system_prompt
|
||||
|
||||
try:
|
||||
async with httpx.AsyncClient(timeout=settings.PROVIDER_TIMEOUT_SECONDS) as client:
|
||||
async with client.stream("POST", endpoint, json=payload) as resp:
|
||||
async for line in resp.aiter_lines():
|
||||
if line:
|
||||
import json
|
||||
data = json.loads(line)
|
||||
chunk = data.get("response", "")
|
||||
if chunk:
|
||||
yield chunk
|
||||
except Exception as exc:
|
||||
yield f"\n[Ollama Streaming Error: {str(exc)}]"
|
||||
|
||||
async def check_health(self) -> Dict[str, Any]:
|
||||
try:
|
||||
async with httpx.AsyncClient(timeout=3.0) as client:
|
||||
resp = await client.get(f"{self.base_url.rstrip('/')}/api/tags")
|
||||
if resp.status_code == 200:
|
||||
models = [m.get("name") for m in resp.json().get("models", [])]
|
||||
return {
|
||||
"provider": "ollama",
|
||||
"status": "CONNECTED",
|
||||
"base_url": self.base_url,
|
||||
"models_available": models,
|
||||
}
|
||||
except Exception as exc:
|
||||
pass
|
||||
return {
|
||||
"provider": "ollama",
|
||||
"status": "NOT_CONNECTED",
|
||||
"base_url": self.base_url,
|
||||
"message": "Ollama service unreachable on configured port. Start via 'ollama serve'.",
|
||||
}
|
||||
@@ -0,0 +1,35 @@
|
||||
from typing import Dict, Optional
|
||||
from app.providers.base import ModelProvider
|
||||
from app.providers.mock_provider import MockProvider
|
||||
from app.providers.ollama_provider import OllamaProvider
|
||||
from app.providers.external_providers import OpenAIProvider, AnthropicProvider, GeminiProvider
|
||||
from app.config.settings import get_settings
|
||||
|
||||
settings = get_settings()
|
||||
|
||||
|
||||
class ProviderRegistry:
|
||||
"""
|
||||
Central registry managing provider instances.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self._providers: Dict[str, ModelProvider] = {
|
||||
"mock": MockProvider(),
|
||||
"ollama": OllamaProvider(base_url=settings.OLLAMA_BASE_URL),
|
||||
"openai": OpenAIProvider(api_key=settings.OPENAI_API_KEY),
|
||||
"anthropic": AnthropicProvider(api_key=settings.ANTHROPIC_API_KEY),
|
||||
"gemini": GeminiProvider(api_key=settings.GEMINI_API_KEY),
|
||||
}
|
||||
|
||||
def get_provider(self, provider_id: str) -> Optional[ModelProvider]:
|
||||
return self._providers.get(provider_id.lower())
|
||||
|
||||
def register_provider(self, provider_id: str, provider: ModelProvider):
|
||||
self._providers[provider_id.lower()] = provider
|
||||
|
||||
def list_providers(self) -> Dict[str, ModelProvider]:
|
||||
return self._providers
|
||||
|
||||
|
||||
provider_registry = ProviderRegistry()
|
||||
@@ -0,0 +1,150 @@
|
||||
import uuid
|
||||
import datetime
|
||||
from typing import List, Dict, Any, Optional
|
||||
from app.models.schemas import (
|
||||
ModelMetadata,
|
||||
RequestAnalysis,
|
||||
RoutingDecision,
|
||||
CandidateScore,
|
||||
PriorityLevel,
|
||||
)
|
||||
from app.router.scoring import filter_candidate_models, compute_candidate_score
|
||||
from app.router.rules_engine import evaluate_routing_rules
|
||||
from app.storage.models import RoutingRuleRecord
|
||||
|
||||
|
||||
def generate_routing_explanation(
|
||||
selected_model: ModelMetadata,
|
||||
analysis: RequestAnalysis,
|
||||
policy_name: str,
|
||||
rule_name: Optional[str] = None,
|
||||
) -> List[str]:
|
||||
reasons = []
|
||||
|
||||
if rule_name:
|
||||
reasons.append(f"Matching rule applied: \"{rule_name}\"")
|
||||
|
||||
# Task explanation
|
||||
reasons.append(f"Task classified as {analysis.task_type.value} with {analysis.complexity_label.value} complexity ({analysis.complexity:.2f})")
|
||||
|
||||
# Reasoning / Code requirements
|
||||
if analysis.reasoning_required:
|
||||
reasons.append("Multi-step reasoning & deep inference required for problem domain")
|
||||
if analysis.coding_required:
|
||||
reasons.append("Syntactical & algorithmic coding capabilities required")
|
||||
|
||||
# Model specific match
|
||||
if selected_model.tier.value == "POWER":
|
||||
reasons.append(f"High reasoning score ({selected_model.quality_score * 100:.0f}%) and deep context ({selected_model.context_window:,} tokens)")
|
||||
elif selected_model.tier.value == "FAST":
|
||||
reasons.append(f"Low latency priority and high speed score ({selected_model.speed_score * 100:.0f}%)")
|
||||
else:
|
||||
reasons.append("Balanced quality and cost profile selected for optimal efficiency")
|
||||
|
||||
# Cost / local factor
|
||||
if selected_model.cost_per_input_token == 0.0:
|
||||
reasons.append("Local zero-cost execution prioritized")
|
||||
elif analysis.cost_sensitivity == PriorityLevel.HIGH:
|
||||
reasons.append("Cost sensitivity is HIGH — model offers lowest token pricing")
|
||||
|
||||
# Policy
|
||||
reasons.append(f"Evaluated under '{policy_name}' policy weights")
|
||||
|
||||
return reasons
|
||||
|
||||
|
||||
def route_request(
|
||||
analysis: RequestAnalysis,
|
||||
available_models: List[ModelMetadata],
|
||||
policy_weights: Dict[str, float],
|
||||
policy_name: str = "balanced",
|
||||
rules: Optional[List[RoutingRuleRecord]] = None,
|
||||
budget_percent: float = 0.0,
|
||||
request_id: Optional[str] = None,
|
||||
) -> RoutingDecision:
|
||||
req_id = request_id or str(uuid.uuid4())
|
||||
decision_id = f"dec_{uuid.uuid4().hex[:12]}"
|
||||
now_iso = datetime.datetime.utcnow().isoformat()
|
||||
|
||||
# 1. Rule evaluation
|
||||
rule_action, rule_target, rule_name = None, None, None
|
||||
if rules:
|
||||
rule_action, rule_target, rule_name = evaluate_routing_rules(rules, analysis, budget_percent)
|
||||
|
||||
# 2. Hard filter candidates
|
||||
eligible_models, rejected_dict = filter_candidate_models(available_models, analysis)
|
||||
|
||||
if not eligible_models:
|
||||
# Fallback to any active model if everything was filtered
|
||||
eligible_models = [m for m in available_models if m.is_active] or available_models
|
||||
|
||||
# Handle Rule Action: ROUTE_TO specific model
|
||||
if rule_action == "ROUTE_TO" and rule_target:
|
||||
exact_match = next((m for m in eligible_models if m.id == rule_target), None)
|
||||
if exact_match:
|
||||
scores = [compute_candidate_score(m, analysis, policy_weights) for m in eligible_models]
|
||||
score_map = {s.model_id: s for s in scores}
|
||||
selected_score = score_map.get(exact_match.id) or compute_candidate_score(exact_match, analysis, policy_weights)
|
||||
|
||||
reasons = generate_routing_explanation(exact_match, analysis, policy_name, rule_name)
|
||||
return RoutingDecision(
|
||||
decision_id=decision_id,
|
||||
request_id=req_id,
|
||||
selected_model=exact_match.id,
|
||||
selected_model_name=exact_match.name,
|
||||
provider=exact_match.provider,
|
||||
tier=exact_match.tier.value if hasattr(exact_match.tier, "value") else str(exact_match.tier),
|
||||
confidence=0.98,
|
||||
policy_used=policy_name,
|
||||
reasons=reasons,
|
||||
candidate_scores=scores,
|
||||
rejected_candidates=rejected_dict,
|
||||
estimated_cost_usd=selected_score.estimated_cost_usd,
|
||||
estimated_latency_ms=selected_score.estimated_latency_ms,
|
||||
rule_applied=rule_name,
|
||||
timestamp=now_iso,
|
||||
)
|
||||
|
||||
# Handle Rule Action: FORCE_TIER
|
||||
if rule_action == "FORCE_TIER" and rule_target:
|
||||
tier_filtered = [m for m in eligible_models if m.tier.value.upper() == rule_target.upper()]
|
||||
if tier_filtered:
|
||||
eligible_models = tier_filtered
|
||||
|
||||
# 3. Score all eligible models
|
||||
candidate_scores: List[CandidateScore] = []
|
||||
for model in eligible_models:
|
||||
sc = compute_candidate_score(model, analysis, policy_weights)
|
||||
candidate_scores.append(sc)
|
||||
|
||||
# Sort descending by overall_score
|
||||
candidate_scores.sort(key=lambda s: s.overall_score, reverse=True)
|
||||
best_candidate = candidate_scores[0]
|
||||
selected_model = next(m for m in eligible_models if m.id == best_candidate.model_id)
|
||||
|
||||
# Confidence calculation: score difference from second candidate
|
||||
if len(candidate_scores) > 1:
|
||||
gap = best_candidate.overall_score - candidate_scores[1].overall_score
|
||||
confidence = min(0.99, max(0.70, round(0.80 + (gap / 100.0), 2)))
|
||||
else:
|
||||
confidence = 0.95
|
||||
|
||||
reasons = generate_routing_explanation(selected_model, analysis, policy_name, rule_name)
|
||||
|
||||
return RoutingDecision(
|
||||
decision_id=decision_id,
|
||||
request_id=req_id,
|
||||
selected_model=selected_model.id,
|
||||
selected_model_name=selected_model.name,
|
||||
provider=selected_model.provider,
|
||||
tier=selected_model.tier.value if hasattr(selected_model.tier, "value") else str(selected_model.tier),
|
||||
confidence=confidence,
|
||||
policy_used=policy_name,
|
||||
reasons=reasons,
|
||||
candidate_scores=candidate_scores,
|
||||
rejected_candidates=rejected_dict,
|
||||
estimated_cost_usd=best_candidate.estimated_cost_usd,
|
||||
estimated_latency_ms=best_candidate.estimated_latency_ms,
|
||||
rule_applied=rule_name,
|
||||
timestamp=now_iso,
|
||||
)
|
||||
@@ -0,0 +1,73 @@
|
||||
from typing import List, Optional, Tuple
|
||||
from app.storage.models import RoutingRuleRecord
|
||||
from app.models.schemas import RequestAnalysis
|
||||
|
||||
|
||||
def evaluate_routing_rules(
|
||||
rules: List[RoutingRuleRecord],
|
||||
analysis: RequestAnalysis,
|
||||
current_budget_percent: float = 0.0,
|
||||
) -> Tuple[Optional[str], Optional[str], Optional[str]]:
|
||||
"""
|
||||
Evaluates enabled routing rules in order of priority (highest priority number first).
|
||||
Returns: (action_type, action_target, matched_rule_name) or (None, None, None)
|
||||
"""
|
||||
sorted_rules = sorted([r for r in rules if r.is_enabled], key=lambda x: x.priority, reverse=True)
|
||||
|
||||
for rule in sorted_rules:
|
||||
field = rule.condition_field.lower()
|
||||
op = rule.condition_operator
|
||||
target_val = rule.condition_value
|
||||
|
||||
matched = False
|
||||
|
||||
if field == "task_type":
|
||||
val = analysis.task_type.value
|
||||
if op == "==":
|
||||
matched = (val == target_val)
|
||||
elif op == "!=":
|
||||
matched = (val != target_val)
|
||||
elif op == "contains":
|
||||
matched = (target_val in val)
|
||||
|
||||
elif field == "complexity":
|
||||
val = analysis.complexity
|
||||
try:
|
||||
target_num = float(target_val)
|
||||
if op == ">=":
|
||||
matched = (val >= target_num)
|
||||
elif op == ">":
|
||||
matched = (val > target_num)
|
||||
elif op == "<=":
|
||||
matched = (val <= target_num)
|
||||
elif op == "<":
|
||||
matched = (val < target_num)
|
||||
elif op == "==":
|
||||
matched = (abs(val - target_num) < 0.01)
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
elif field in ("budget_percent", "monthly_budget_percent"):
|
||||
try:
|
||||
target_num = float(target_val)
|
||||
if op == ">=":
|
||||
matched = (current_budget_percent >= target_num)
|
||||
elif op == ">":
|
||||
matched = (current_budget_percent > target_num)
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
elif field == "context_size":
|
||||
try:
|
||||
target_num = int(target_val)
|
||||
if op == ">=":
|
||||
matched = (analysis.context_size >= target_num)
|
||||
elif op == ">":
|
||||
matched = (analysis.context_size > target_num)
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
if matched:
|
||||
return rule.action_type, rule.action_target, rule.name
|
||||
|
||||
return None, None, None
|
||||
@@ -0,0 +1,154 @@
|
||||
from typing import List, Dict, Tuple, Optional
|
||||
from app.models.schemas import ModelMetadata, RequestAnalysis, CandidateScore, PriorityLevel
|
||||
|
||||
|
||||
def filter_candidate_models(
|
||||
models: List[ModelMetadata],
|
||||
analysis: RequestAnalysis,
|
||||
) -> Tuple[List[ModelMetadata], Dict[str, str]]:
|
||||
"""
|
||||
Hard-filtering step: eliminates models that cannot satisfy non-negotiable requirements.
|
||||
Returns: (eligible_models, rejected_candidates_with_reasons)
|
||||
"""
|
||||
eligible: List[ModelMetadata] = []
|
||||
rejected: Dict[str, str] = {}
|
||||
|
||||
for m in models:
|
||||
if not m.is_active:
|
||||
rejected[m.id] = "Model is currently disabled/inactive."
|
||||
continue
|
||||
|
||||
# 1. Context Window check
|
||||
if analysis.context_size > m.context_window:
|
||||
rejected[m.id] = (
|
||||
f"Context window ({m.context_window} tokens) is smaller than required input ({analysis.context_size} tokens)."
|
||||
)
|
||||
continue
|
||||
|
||||
# 2. Vision requirement check
|
||||
if analysis.vision_required and not m.supports_vision:
|
||||
rejected[m.id] = "Task requires vision / multimodal capabilities which this model does not support."
|
||||
continue
|
||||
|
||||
# 3. Coding capability check (for high complexity code tasks)
|
||||
if analysis.coding_required and analysis.complexity >= 0.70 and not m.supports_coding:
|
||||
rejected[m.id] = "Task requires specialized coding capabilities for complex logic."
|
||||
continue
|
||||
|
||||
# 4. Deep reasoning capability check
|
||||
if analysis.reasoning_required and analysis.complexity >= 0.85 and not m.supports_reasoning:
|
||||
rejected[m.id] = "Task requires advanced reasoning capabilities for multi-step problem solving."
|
||||
continue
|
||||
|
||||
eligible.append(m)
|
||||
|
||||
# If all filtered out (edge case), return active models to avoid total failure
|
||||
if not eligible and models:
|
||||
for m in models:
|
||||
if m.is_active:
|
||||
eligible.append(m)
|
||||
rejected.clear()
|
||||
|
||||
return eligible, rejected
|
||||
|
||||
|
||||
def compute_candidate_score(
|
||||
model: ModelMetadata,
|
||||
analysis: RequestAnalysis,
|
||||
weights: Dict[str, float],
|
||||
) -> CandidateScore:
|
||||
"""
|
||||
Multi-criteria configurable scoring formula.
|
||||
Quality, Speed, CostEfficiency, CapabilityMatch, Reliability.
|
||||
Normalized score: 0 to 100.
|
||||
"""
|
||||
# 1. Quality Component (0.0 to 1.0)
|
||||
quality_comp = model.quality_score
|
||||
if analysis.quality_requirement == PriorityLevel.HIGH:
|
||||
# Boost premium models for high quality requirements
|
||||
if model.tier.value == "POWER":
|
||||
quality_comp = min(1.0, quality_comp * 1.15)
|
||||
elif analysis.quality_requirement == PriorityLevel.LOW:
|
||||
quality_comp = model.quality_score * 0.9
|
||||
|
||||
# 2. Speed Component (0.0 to 1.0)
|
||||
speed_comp = model.speed_score
|
||||
if analysis.latency_priority == PriorityLevel.HIGH:
|
||||
if model.tier.value == "FAST":
|
||||
speed_comp = min(1.0, speed_comp * 1.20)
|
||||
elif analysis.latency_priority == PriorityLevel.LOW:
|
||||
speed_comp = model.speed_score * 0.85
|
||||
|
||||
# 3. Cost Efficiency Component (0.0 to 1.0)
|
||||
# Zero-cost local models get 1.0, cheaper cloud models get high score
|
||||
combined_token_cost = (model.cost_per_input_token * 1000) + (model.cost_per_output_token * 1000)
|
||||
if combined_token_cost == 0.0:
|
||||
cost_efficiency = 1.0
|
||||
else:
|
||||
# Scale: $0.0001 -> 0.95, $0.01 -> 0.50, $0.05 -> 0.10
|
||||
cost_efficiency = max(0.05, min(0.99, 1.0 / (1.0 + (combined_token_cost * 100.0))))
|
||||
|
||||
if analysis.cost_sensitivity == PriorityLevel.HIGH:
|
||||
cost_efficiency = min(1.0, cost_efficiency * 1.25)
|
||||
|
||||
# 4. Capability Match Component (0.0 to 1.0)
|
||||
cap_points = 0.0
|
||||
total_caps = 0.0
|
||||
|
||||
if analysis.coding_required:
|
||||
total_caps += 1.0
|
||||
if model.supports_coding:
|
||||
cap_points += 1.0
|
||||
|
||||
if analysis.reasoning_required:
|
||||
total_caps += 1.0
|
||||
if model.supports_reasoning:
|
||||
cap_points += 1.0
|
||||
|
||||
if analysis.vision_required:
|
||||
total_caps += 1.0
|
||||
if model.supports_vision:
|
||||
cap_points += 1.0
|
||||
|
||||
capability_comp = (cap_points / total_caps) if total_caps > 0 else 0.90
|
||||
|
||||
# 5. Reliability Component (0.0 to 1.0)
|
||||
reliability_comp = model.reliability_score
|
||||
|
||||
# Weighted Sum
|
||||
w_q = weights.get("quality_weight", 0.35)
|
||||
w_c = weights.get("cost_weight", 0.25)
|
||||
w_s = weights.get("speed_weight", 0.20)
|
||||
w_cap = weights.get("capability_weight", 0.15)
|
||||
w_r = weights.get("reliability_weight", 0.05)
|
||||
total_w = w_q + w_c + w_s + w_cap + w_r
|
||||
|
||||
raw_score = (
|
||||
(quality_comp * w_q)
|
||||
+ (cost_efficiency * w_c)
|
||||
+ (speed_comp * w_s)
|
||||
+ (capability_comp * w_cap)
|
||||
+ (reliability_comp * w_r)
|
||||
) / (total_w if total_w > 0 else 1.0)
|
||||
|
||||
overall_score = round(raw_score * 100.0, 2)
|
||||
|
||||
# Estimated metrics
|
||||
est_latency = 120.0 if model.tier.value == "FAST" else (350.0 if model.tier.value == "BALANCED" else 750.0)
|
||||
est_cost = (analysis.context_size * model.cost_per_input_token) + (250 * model.cost_per_output_token)
|
||||
|
||||
return CandidateScore(
|
||||
model_id=model.id,
|
||||
model_name=model.name,
|
||||
provider=model.provider,
|
||||
tier=model.tier.value if hasattr(model.tier, "value") else str(model.tier),
|
||||
overall_score=overall_score,
|
||||
quality_component=round(quality_comp * 100.0, 1),
|
||||
cost_component=round(cost_efficiency * 100.0, 1),
|
||||
speed_component=round(speed_comp * 100.0, 1),
|
||||
capability_component=round(capability_comp * 100.0, 1),
|
||||
reliability_component=round(reliability_comp * 100.0, 1),
|
||||
estimated_latency_ms=est_latency,
|
||||
estimated_cost_usd=round(est_cost, 6),
|
||||
eligible=True,
|
||||
)
|
||||
@@ -0,0 +1,366 @@
|
||||
from typing import AsyncGenerator
|
||||
from sqlalchemy.ext.asyncio import create_async_engine, async_sessionmaker, AsyncSession
|
||||
from app.config.settings import get_settings
|
||||
from app.storage.models import (
|
||||
Base,
|
||||
ModelRecord,
|
||||
ProviderRecord,
|
||||
RoutingPolicyRecord,
|
||||
RoutingRuleRecord,
|
||||
BudgetRecord,
|
||||
)
|
||||
|
||||
settings = get_settings()
|
||||
|
||||
engine = create_async_engine(
|
||||
settings.DATABASE_URL,
|
||||
echo=False,
|
||||
connect_args={"check_same_thread": False} if "sqlite" in settings.DATABASE_URL else {},
|
||||
)
|
||||
|
||||
AsyncSessionLocal = async_sessionmaker(
|
||||
bind=engine,
|
||||
class_=AsyncSession,
|
||||
expire_on_commit=False,
|
||||
)
|
||||
|
||||
|
||||
async def get_db() -> AsyncGenerator[AsyncSession, None]:
|
||||
async with AsyncSessionLocal() as session:
|
||||
try:
|
||||
yield session
|
||||
finally:
|
||||
await session.close()
|
||||
|
||||
|
||||
async def init_db():
|
||||
async with engine.begin() as conn:
|
||||
await conn.run_sync(Base.metadata.create_all)
|
||||
|
||||
# Populate initial seed data if empty
|
||||
async with AsyncSessionLocal() as session:
|
||||
from sqlalchemy import select
|
||||
|
||||
# Check providers
|
||||
res = await session.execute(select(ProviderRecord))
|
||||
providers = res.scalars().all()
|
||||
if not providers:
|
||||
default_providers = [
|
||||
ProviderRecord(id="mock", name="Mock / Demo Provider", status="CONNECTED", is_enabled=True),
|
||||
ProviderRecord(id="ollama", name="Ollama (Local)", status="READY", is_enabled=True, base_url=settings.OLLAMA_BASE_URL),
|
||||
ProviderRecord(id="openai", name="OpenAI", status="CONNECTED" if settings.OPENAI_API_KEY else "NOT_CONFIGURED", is_enabled=bool(settings.OPENAI_API_KEY)),
|
||||
ProviderRecord(id="anthropic", name="Anthropic", status="CONNECTED" if settings.ANTHROPIC_API_KEY else "NOT_CONFIGURED", is_enabled=bool(settings.ANTHROPIC_API_KEY)),
|
||||
ProviderRecord(id="gemini", name="Google Gemini", status="CONNECTED" if settings.GEMINI_API_KEY else "NOT_CONFIGURED", is_enabled=bool(settings.GEMINI_API_KEY)),
|
||||
]
|
||||
session.add_all(default_providers)
|
||||
|
||||
# Check models
|
||||
res_m = await session.execute(select(ModelRecord))
|
||||
models = res_m.scalars().all()
|
||||
if not models:
|
||||
default_models = [
|
||||
# Mock Models
|
||||
ModelRecord(
|
||||
id="mock-fast",
|
||||
name="Mock Fast (Echo)",
|
||||
provider="mock",
|
||||
type="MOCK",
|
||||
tier="FAST",
|
||||
context_window=32768,
|
||||
supports_coding=True,
|
||||
supports_reasoning=False,
|
||||
supports_vision=False,
|
||||
supports_tools=True,
|
||||
quality_score=0.75,
|
||||
speed_score=0.98,
|
||||
reliability_score=0.99,
|
||||
cost_per_input_token=0.0000001,
|
||||
cost_per_output_token=0.0000002,
|
||||
),
|
||||
ModelRecord(
|
||||
id="mock-balanced",
|
||||
name="Mock Balanced (General)",
|
||||
provider="mock",
|
||||
type="MOCK",
|
||||
tier="BALANCED",
|
||||
context_window=65536,
|
||||
supports_coding=True,
|
||||
supports_reasoning=True,
|
||||
supports_vision=True,
|
||||
supports_tools=True,
|
||||
quality_score=0.88,
|
||||
speed_score=0.85,
|
||||
reliability_score=0.99,
|
||||
cost_per_input_token=0.000001,
|
||||
cost_per_output_token=0.000003,
|
||||
),
|
||||
ModelRecord(
|
||||
id="mock-power",
|
||||
name="Mock Power (Reasoning)",
|
||||
provider="mock",
|
||||
type="MOCK",
|
||||
tier="POWER",
|
||||
context_window=131072,
|
||||
supports_coding=True,
|
||||
supports_reasoning=True,
|
||||
supports_vision=True,
|
||||
supports_tools=True,
|
||||
quality_score=0.98,
|
||||
speed_score=0.60,
|
||||
reliability_score=0.99,
|
||||
cost_per_input_token=0.000005,
|
||||
cost_per_output_token=0.000015,
|
||||
),
|
||||
# Local Ollama Models (Default profiles)
|
||||
ModelRecord(
|
||||
id="ollama-qwen2.5-coder",
|
||||
name="Qwen 2.5 Coder 7B (Ollama)",
|
||||
provider="ollama",
|
||||
type="LOCAL",
|
||||
tier="FAST",
|
||||
context_window=32768,
|
||||
supports_coding=True,
|
||||
supports_reasoning=True,
|
||||
supports_vision=False,
|
||||
supports_tools=True,
|
||||
quality_score=0.86,
|
||||
speed_score=0.92,
|
||||
reliability_score=0.95,
|
||||
cost_per_input_token=0.0,
|
||||
cost_per_output_token=0.0,
|
||||
),
|
||||
ModelRecord(
|
||||
id="ollama-llama3.2",
|
||||
name="Llama 3.2 3B (Ollama)",
|
||||
provider="ollama",
|
||||
type="LOCAL",
|
||||
tier="FAST",
|
||||
context_window=131072,
|
||||
supports_coding=True,
|
||||
supports_reasoning=True,
|
||||
supports_vision=False,
|
||||
supports_tools=True,
|
||||
quality_score=0.82,
|
||||
speed_score=0.95,
|
||||
reliability_score=0.95,
|
||||
cost_per_input_token=0.0,
|
||||
cost_per_output_token=0.0,
|
||||
),
|
||||
ModelRecord(
|
||||
id="ollama-mistral",
|
||||
name="Mistral 7B (Ollama)",
|
||||
provider="ollama",
|
||||
type="LOCAL",
|
||||
tier="BALANCED",
|
||||
context_window=32768,
|
||||
supports_coding=True,
|
||||
supports_reasoning=True,
|
||||
supports_vision=False,
|
||||
supports_tools=True,
|
||||
quality_score=0.84,
|
||||
speed_score=0.88,
|
||||
reliability_score=0.95,
|
||||
cost_per_input_token=0.0,
|
||||
cost_per_output_token=0.0,
|
||||
),
|
||||
ModelRecord(
|
||||
id="ollama-deepseek-r1",
|
||||
name="DeepSeek R1 8B (Ollama)",
|
||||
provider="ollama",
|
||||
type="LOCAL",
|
||||
tier="POWER",
|
||||
context_window=65536,
|
||||
supports_coding=True,
|
||||
supports_reasoning=True,
|
||||
supports_vision=False,
|
||||
supports_tools=False,
|
||||
quality_score=0.93,
|
||||
speed_score=0.68,
|
||||
reliability_score=0.94,
|
||||
cost_per_input_token=0.0,
|
||||
cost_per_output_token=0.0,
|
||||
),
|
||||
# External Cloud Model Profiles
|
||||
ModelRecord(
|
||||
id="gpt-4o-mini",
|
||||
name="OpenAI GPT-4o Mini",
|
||||
provider="openai",
|
||||
type="CLOUD",
|
||||
tier="FAST",
|
||||
context_window=128000,
|
||||
supports_coding=True,
|
||||
supports_reasoning=True,
|
||||
supports_vision=True,
|
||||
supports_tools=True,
|
||||
quality_score=0.88,
|
||||
speed_score=0.90,
|
||||
reliability_score=0.99,
|
||||
cost_per_input_token=0.00000015,
|
||||
cost_per_output_token=0.00000060,
|
||||
),
|
||||
ModelRecord(
|
||||
id="gpt-4o",
|
||||
name="OpenAI GPT-4o",
|
||||
provider="openai",
|
||||
type="CLOUD",
|
||||
tier="POWER",
|
||||
context_window=128000,
|
||||
supports_coding=True,
|
||||
supports_reasoning=True,
|
||||
supports_vision=True,
|
||||
supports_tools=True,
|
||||
quality_score=0.96,
|
||||
speed_score=0.75,
|
||||
reliability_score=0.99,
|
||||
cost_per_input_token=0.0000025,
|
||||
cost_per_output_token=0.000010,
|
||||
),
|
||||
ModelRecord(
|
||||
id="claude-3-5-sonnet",
|
||||
name="Anthropic Claude 3.5 Sonnet",
|
||||
provider="anthropic",
|
||||
type="CLOUD",
|
||||
tier="POWER",
|
||||
context_window=200000,
|
||||
supports_coding=True,
|
||||
supports_reasoning=True,
|
||||
supports_vision=True,
|
||||
supports_tools=True,
|
||||
quality_score=0.98,
|
||||
speed_score=0.72,
|
||||
reliability_score=0.99,
|
||||
cost_per_input_token=0.000003,
|
||||
cost_per_output_token=0.000015,
|
||||
),
|
||||
ModelRecord(
|
||||
id="gemini-1.5-flash",
|
||||
name="Google Gemini 1.5 Flash",
|
||||
provider="gemini",
|
||||
type="CLOUD",
|
||||
tier="FAST",
|
||||
context_window=1000000,
|
||||
supports_coding=True,
|
||||
supports_reasoning=True,
|
||||
supports_vision=True,
|
||||
supports_tools=True,
|
||||
quality_score=0.87,
|
||||
speed_score=0.93,
|
||||
reliability_score=0.99,
|
||||
cost_per_input_token=0.000000075,
|
||||
cost_per_output_token=0.00000030,
|
||||
),
|
||||
]
|
||||
session.add_all(default_models)
|
||||
|
||||
# Check policies
|
||||
res_p = await session.execute(select(RoutingPolicyRecord))
|
||||
policies = res_p.scalars().all()
|
||||
if not policies:
|
||||
default_policies = [
|
||||
RoutingPolicyRecord(
|
||||
id="balanced",
|
||||
name="Balanced Router",
|
||||
description="Optimizes across quality, cost, speed, capabilities, and reliability.",
|
||||
quality_weight=0.35,
|
||||
cost_weight=0.25,
|
||||
speed_weight=0.20,
|
||||
capability_weight=0.15,
|
||||
reliability_weight=0.05,
|
||||
is_default=True,
|
||||
),
|
||||
RoutingPolicyRecord(
|
||||
id="lowest_cost",
|
||||
name="Lowest Cost",
|
||||
description="Prefers cheapest and local zero-cost models with acceptable capability.",
|
||||
quality_weight=0.15,
|
||||
cost_weight=0.60,
|
||||
speed_weight=0.15,
|
||||
capability_weight=0.05,
|
||||
reliability_weight=0.05,
|
||||
is_default=False,
|
||||
),
|
||||
RoutingPolicyRecord(
|
||||
id="lowest_latency",
|
||||
name="Lowest Latency",
|
||||
description="Prefers fastest response times and high throughput models.",
|
||||
quality_weight=0.15,
|
||||
cost_weight=0.15,
|
||||
speed_weight=0.60,
|
||||
capability_weight=0.05,
|
||||
reliability_weight=0.05,
|
||||
is_default=False,
|
||||
),
|
||||
RoutingPolicyRecord(
|
||||
id="highest_quality",
|
||||
name="Highest Quality",
|
||||
description="Directs all complex reasoning and coding tasks to top-tier frontier models.",
|
||||
quality_weight=0.65,
|
||||
cost_weight=0.05,
|
||||
speed_weight=0.10,
|
||||
capability_weight=0.15,
|
||||
reliability_weight=0.05,
|
||||
is_default=False,
|
||||
),
|
||||
RoutingPolicyRecord(
|
||||
id="custom",
|
||||
name="Custom Weights",
|
||||
description="User-defined policy weights.",
|
||||
quality_weight=0.30,
|
||||
cost_weight=0.30,
|
||||
speed_weight=0.20,
|
||||
capability_weight=0.10,
|
||||
reliability_weight=0.10,
|
||||
is_default=False,
|
||||
),
|
||||
]
|
||||
session.add_all(default_policies)
|
||||
|
||||
# Check default budget
|
||||
res_b = await session.execute(select(BudgetRecord))
|
||||
budget = res_b.scalar_one_or_none()
|
||||
if not budget:
|
||||
session.add(
|
||||
BudgetRecord(
|
||||
id="default",
|
||||
daily_limit=settings.DAILY_BUDGET,
|
||||
monthly_limit=settings.MONTHLY_BUDGET,
|
||||
per_request_limit=settings.PER_REQUEST_BUDGET,
|
||||
current_daily_spend=0.0,
|
||||
current_monthly_spend=0.0,
|
||||
intervention_mode="NORMAL",
|
||||
)
|
||||
)
|
||||
|
||||
# Check default rules
|
||||
res_r = await session.execute(select(RoutingRuleRecord))
|
||||
rules = res_r.scalars().all()
|
||||
if not rules:
|
||||
default_rules = [
|
||||
RoutingRuleRecord(
|
||||
id="rule-code-reasoning",
|
||||
name="Deep Debugging & Reasoning Override",
|
||||
description="Route high complexity coding & debugging to Power Models",
|
||||
priority=10,
|
||||
is_enabled=True,
|
||||
condition_field="task_type",
|
||||
condition_operator="==",
|
||||
condition_value="DEBUGGING",
|
||||
action_type="FORCE_TIER",
|
||||
action_target="POWER",
|
||||
),
|
||||
RoutingRuleRecord(
|
||||
id="rule-budget-saver",
|
||||
name="Budget Guard > 95%",
|
||||
description="Prefer local or fast models when budget exceeds 95%",
|
||||
priority=50,
|
||||
is_enabled=True,
|
||||
condition_field="budget_percent",
|
||||
condition_operator=">=",
|
||||
condition_value="95.0",
|
||||
action_type="SET_POLICY",
|
||||
action_target="lowest_cost",
|
||||
),
|
||||
]
|
||||
session.add_all(default_rules)
|
||||
|
||||
await session.commit()
|
||||
@@ -0,0 +1,202 @@
|
||||
import datetime
|
||||
from sqlalchemy import (
|
||||
Column,
|
||||
String,
|
||||
Float,
|
||||
Integer,
|
||||
Boolean,
|
||||
DateTime,
|
||||
Text,
|
||||
JSON,
|
||||
ForeignKey,
|
||||
)
|
||||
from sqlalchemy.orm import declarative_base, relationship
|
||||
|
||||
Base = declarative_base()
|
||||
|
||||
|
||||
class ModelRecord(Base):
|
||||
__tablename__ = "models"
|
||||
|
||||
id = Column(String(100), primary_key=True)
|
||||
name = Column(String(200), nullable=False)
|
||||
provider = Column(String(50), nullable=False) # mock, ollama, openai, anthropic, gemini
|
||||
type = Column(String(50), default="LOCAL") # LOCAL, CLOUD, MOCK
|
||||
context_window = Column(Integer, default=8192)
|
||||
|
||||
supports_coding = Column(Boolean, default=False)
|
||||
supports_reasoning = Column(Boolean, default=False)
|
||||
supports_vision = Column(Boolean, default=False)
|
||||
supports_tools = Column(Boolean, default=False)
|
||||
|
||||
quality_score = Column(Float, default=0.7)
|
||||
speed_score = Column(Float, default=0.7)
|
||||
reliability_score = Column(Float, default=0.95)
|
||||
|
||||
cost_per_input_token = Column(Float, default=0.0)
|
||||
cost_per_output_token = Column(Float, default=0.0)
|
||||
|
||||
availability = Column(String(50), default="AVAILABLE")
|
||||
is_active = Column(Boolean, default=True)
|
||||
tier = Column(String(50), default="BALANCED") # FAST, BALANCED, POWER
|
||||
|
||||
created_at = Column(DateTime, default=datetime.datetime.utcnow)
|
||||
updated_at = Column(DateTime, default=datetime.datetime.utcnow, onupdate=datetime.datetime.utcnow)
|
||||
|
||||
|
||||
class ProviderRecord(Base):
|
||||
__tablename__ = "providers"
|
||||
|
||||
id = Column(String(50), primary_key=True)
|
||||
name = Column(String(100), nullable=False)
|
||||
status = Column(String(50), default="READY") # READY, CONNECTED, NOT_CONFIGURED, ERROR
|
||||
is_enabled = Column(Boolean, default=True)
|
||||
base_url = Column(String(500), nullable=True)
|
||||
updated_at = Column(DateTime, default=datetime.datetime.utcnow, onupdate=datetime.datetime.utcnow)
|
||||
|
||||
|
||||
class RoutingPolicyRecord(Base):
|
||||
__tablename__ = "routing_policies"
|
||||
|
||||
id = Column(String(50), primary_key=True)
|
||||
name = Column(String(100), nullable=False)
|
||||
description = Column(String(500), nullable=True)
|
||||
quality_weight = Column(Float, default=0.35)
|
||||
cost_weight = Column(Float, default=0.25)
|
||||
speed_weight = Column(Float, default=0.20)
|
||||
capability_weight = Column(Float, default=0.15)
|
||||
reliability_weight = Column(Float, default=0.05)
|
||||
is_default = Column(Boolean, default=False)
|
||||
|
||||
|
||||
class RoutingRuleRecord(Base):
|
||||
__tablename__ = "routing_rules"
|
||||
|
||||
id = Column(String(100), primary_key=True)
|
||||
name = Column(String(200), nullable=False)
|
||||
description = Column(String(500), nullable=True)
|
||||
priority = Column(Integer, default=0)
|
||||
is_enabled = Column(Boolean, default=True)
|
||||
|
||||
condition_field = Column(String(100), nullable=False) # task_type, complexity, budget_percent, context_size
|
||||
condition_operator = Column(String(50), nullable=False) # ==, !=, >, <, >=, <=, contains
|
||||
condition_value = Column(String(200), nullable=False)
|
||||
|
||||
action_type = Column(String(50), nullable=False) # ROUTE_TO, PREFER, FORCE_TIER, SET_POLICY
|
||||
action_target = Column(String(100), nullable=False)
|
||||
|
||||
created_at = Column(DateTime, default=datetime.datetime.utcnow)
|
||||
|
||||
|
||||
class RequestRecord(Base):
|
||||
__tablename__ = "requests"
|
||||
|
||||
request_id = Column(String(100), primary_key=True)
|
||||
timestamp = Column(DateTime, default=datetime.datetime.utcnow)
|
||||
prompt = Column(Text, nullable=False)
|
||||
task_type = Column(String(50), nullable=False)
|
||||
complexity = Column(Float, default=0.5)
|
||||
context_size = Column(Integer, default=0)
|
||||
reasoning_required = Column(Boolean, default=False)
|
||||
coding_required = Column(Boolean, default=False)
|
||||
|
||||
routing_policy = Column(String(50), default="balanced")
|
||||
selected_model = Column(String(100), nullable=False)
|
||||
provider = Column(String(50), nullable=False)
|
||||
status = Column(String(50), default="SUCCESS") # SUCCESS, FAILED, FALLBACK
|
||||
|
||||
fallback_used = Column(Boolean, default=False)
|
||||
original_model = Column(String(100), nullable=True)
|
||||
fallback_reason = Column(String(500), nullable=True)
|
||||
|
||||
input_tokens = Column(Integer, default=0)
|
||||
output_tokens = Column(Integer, default=0)
|
||||
total_tokens = Column(Integer, default=0)
|
||||
estimated_cost = Column(Float, default=0.0)
|
||||
baseline_cost = Column(Float, default=0.0)
|
||||
cost_saved = Column(Float, default=0.0)
|
||||
|
||||
routing_latency_ms = Column(Float, default=0.0)
|
||||
provider_latency_ms = Column(Float, default=0.0)
|
||||
total_latency_ms = Column(Float, default=0.0)
|
||||
time_to_first_token_ms = Column(Float, nullable=True)
|
||||
|
||||
|
||||
class RoutingDecisionRecord(Base):
|
||||
__tablename__ = "routing_decisions"
|
||||
|
||||
decision_id = Column(String(100), primary_key=True)
|
||||
request_id = Column(String(100), ForeignKey("requests.request_id"), nullable=False)
|
||||
timestamp = Column(DateTime, default=datetime.datetime.utcnow)
|
||||
|
||||
selected_model = Column(String(100), nullable=False)
|
||||
confidence = Column(Float, default=0.9)
|
||||
reasons = Column(JSON, default=list)
|
||||
candidate_scores = Column(JSON, default=dict)
|
||||
rejected_candidates = Column(JSON, default=dict)
|
||||
policy_used = Column(String(50), default="balanced")
|
||||
|
||||
|
||||
class ResponseRecord(Base):
|
||||
__tablename__ = "responses"
|
||||
|
||||
response_id = Column(String(100), primary_key=True)
|
||||
request_id = Column(String(100), ForeignKey("requests.request_id"), nullable=False)
|
||||
timestamp = Column(DateTime, default=datetime.datetime.utcnow)
|
||||
model_id = Column(String(100), nullable=False)
|
||||
provider = Column(String(50), nullable=False)
|
||||
content = Column(Text, nullable=False)
|
||||
finish_reason = Column(String(50), default="stop")
|
||||
is_mock = Column(Boolean, default=False)
|
||||
|
||||
|
||||
class FeedbackRecord(Base):
|
||||
__tablename__ = "feedback"
|
||||
|
||||
id = Column(String(100), primary_key=True)
|
||||
request_id = Column(String(100), ForeignKey("requests.request_id"), nullable=False)
|
||||
model_id = Column(String(100), nullable=False)
|
||||
task_type = Column(String(50), nullable=False)
|
||||
rating = Column(Integer, nullable=False) # 1 for thumbs up, -1 for thumbs down
|
||||
comment = Column(Text, nullable=True)
|
||||
timestamp = Column(DateTime, default=datetime.datetime.utcnow)
|
||||
|
||||
|
||||
class BudgetRecord(Base):
|
||||
__tablename__ = "budgets"
|
||||
|
||||
id = Column(String(50), primary_key=True, default="default")
|
||||
daily_limit = Column(Float, default=10.0)
|
||||
monthly_limit = Column(Float, default=100.0)
|
||||
per_request_limit = Column(Float, default=1.0)
|
||||
current_daily_spend = Column(Float, default=0.0)
|
||||
current_monthly_spend = Column(Float, default=0.0)
|
||||
intervention_mode = Column(String(50), default="OPTIMIZE") # NORMAL, OPTIMIZE_80, CHEAP_95, BLOCK_100
|
||||
updated_at = Column(DateTime, default=datetime.datetime.utcnow, onupdate=datetime.datetime.utcnow)
|
||||
|
||||
|
||||
class ExperimentRecord(Base):
|
||||
__tablename__ = "experiments"
|
||||
|
||||
id = Column(String(100), primary_key=True)
|
||||
name = Column(String(200), nullable=False)
|
||||
description = Column(String(500), nullable=True)
|
||||
policy_a = Column(String(50), nullable=False)
|
||||
policy_b = Column(String(50), nullable=False)
|
||||
status = Column(String(50), default="ACTIVE") # ACTIVE, PAUSED, COMPLETED
|
||||
sample_count = Column(Integer, default=0)
|
||||
created_at = Column(DateTime, default=datetime.datetime.utcnow)
|
||||
|
||||
|
||||
class ExperimentRunRecord(Base):
|
||||
__tablename__ = "experiment_runs"
|
||||
|
||||
id = Column(String(100), primary_key=True)
|
||||
experiment_id = Column(String(100), ForeignKey("experiments.id"), nullable=False)
|
||||
request_id = Column(String(100), ForeignKey("requests.request_id"), nullable=False)
|
||||
assigned_policy = Column(String(50), nullable=False)
|
||||
selected_model = Column(String(100), nullable=False)
|
||||
cost = Column(Float, default=0.0)
|
||||
latency_ms = Column(Float, default=0.0)
|
||||
feedback_rating = Column(Integer, nullable=True)
|
||||
timestamp = Column(DateTime, default=datetime.datetime.utcnow)
|
||||
@@ -0,0 +1,94 @@
|
||||
import asyncio
|
||||
import json
|
||||
from contextlib import asynccontextmanager
|
||||
from fastapi import FastAPI
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
from sse_starlette.sse import EventSourceResponse
|
||||
|
||||
from app.config.settings import get_settings
|
||||
from app.storage.database import init_db, AsyncSessionLocal
|
||||
from app.storage.models import RequestRecord
|
||||
from app.api.routes import router as api_router
|
||||
from sqlalchemy import select, desc
|
||||
|
||||
settings = get_settings()
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(app: FastAPI):
|
||||
# Initialize SQLite database and seed defaults
|
||||
await init_db()
|
||||
yield
|
||||
|
||||
|
||||
app = FastAPI(
|
||||
title="Model Router",
|
||||
description="Intelligent LLM Request Routing Platform — AI Traffic Control Room",
|
||||
version="1.0.0",
|
||||
lifespan=lifespan,
|
||||
)
|
||||
|
||||
# CORS configuration
|
||||
app.add_middleware(
|
||||
CORSMiddleware,
|
||||
allow_origins=["*"],
|
||||
allow_credentials=True,
|
||||
allow_methods=["*"],
|
||||
allow_headers=["*"],
|
||||
)
|
||||
|
||||
app.include_router(api_router)
|
||||
|
||||
|
||||
@app.get("/api/traffic/stream")
|
||||
async def stream_live_traffic():
|
||||
"""
|
||||
SSE Live Stream endpoint for the Traffic Control Room visualization.
|
||||
Pushes newly persisted requests to connected frontend clients.
|
||||
"""
|
||||
async def event_generator():
|
||||
last_seen_id = None
|
||||
while True:
|
||||
try:
|
||||
async with AsyncSessionLocal() as session:
|
||||
res = await session.execute(
|
||||
select(RequestRecord).order_by(desc(RequestRecord.timestamp)).limit(1)
|
||||
)
|
||||
latest = res.scalar_one_or_none()
|
||||
if latest and latest.request_id != last_seen_id:
|
||||
last_seen_id = latest.request_id
|
||||
payload = {
|
||||
"request_id": latest.request_id,
|
||||
"timestamp": latest.timestamp.isoformat() if latest.timestamp else "",
|
||||
"prompt": latest.prompt,
|
||||
"task_type": latest.task_type,
|
||||
"complexity": latest.complexity,
|
||||
"selected_model": latest.selected_model,
|
||||
"provider": latest.provider,
|
||||
"status": latest.status,
|
||||
"total_latency_ms": latest.total_latency_ms,
|
||||
"estimated_cost": latest.estimated_cost,
|
||||
}
|
||||
yield {
|
||||
"event": "traffic_event",
|
||||
"data": json.dumps(payload),
|
||||
}
|
||||
except Exception:
|
||||
pass
|
||||
await asyncio.sleep(1.5)
|
||||
|
||||
return EventSourceResponse(event_generator())
|
||||
|
||||
|
||||
@app.get("/")
|
||||
async def root():
|
||||
return {
|
||||
"message": "Model Router AI Control Room Backend is active.",
|
||||
"docs": "/docs",
|
||||
"health": "/api/health",
|
||||
}
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import uvicorn
|
||||
uvicorn.run("main:app", host=settings.HOST, port=settings.PORT, reload=settings.DEBUG)
|
||||
@@ -0,0 +1,15 @@
|
||||
fastapi>=0.110.0
|
||||
uvicorn[standard]>=0.28.0
|
||||
pydantic>=2.6.0
|
||||
pydantic-settings>=2.2.0
|
||||
sqlalchemy>=2.0.28
|
||||
aiosqlite>=0.20.0
|
||||
httpx>=0.27.0
|
||||
typer>=0.12.0
|
||||
rich>=13.7.0
|
||||
pytest>=8.0.0
|
||||
pytest-asyncio>=0.23.0
|
||||
python-dotenv>=1.0.1
|
||||
python-multipart>=0.0.9
|
||||
sse-starlette>=2.0.0
|
||||
tiktoken>=0.6.0
|
||||
@@ -0,0 +1,41 @@
|
||||
import pytest
|
||||
from app.analyzer.heuristics import analyze_request_heuristics, estimate_tokens
|
||||
from app.models.schemas import TaskType, PriorityLevel
|
||||
|
||||
|
||||
def test_token_estimation():
|
||||
text = "Hello world from Model Router platform!"
|
||||
tokens = estimate_tokens(text)
|
||||
assert tokens > 0
|
||||
|
||||
|
||||
def test_heuristic_classification_coding():
|
||||
prompt = "Write a python function to calculate fibonacci sequence using dynamic programming"
|
||||
analysis = analyze_request_heuristics(prompt)
|
||||
assert analysis.task_type == TaskType.CODING
|
||||
assert analysis.coding_required is True
|
||||
assert analysis.complexity >= 0.40
|
||||
|
||||
|
||||
def test_heuristic_classification_debugging_complexity():
|
||||
prompt = "Debug this distributed async deadlock issue in the microservices cluster and explain the root cause"
|
||||
analysis = analyze_request_heuristics(prompt)
|
||||
assert analysis.task_type == TaskType.DEBUGGING
|
||||
assert analysis.reasoning_required is True
|
||||
assert analysis.coding_required is True
|
||||
assert analysis.complexity_label == PriorityLevel.HIGH
|
||||
assert analysis.complexity >= 0.75
|
||||
|
||||
|
||||
def test_heuristic_classification_summarization():
|
||||
prompt = "Provide a summary and key bullet points of this meeting transcript"
|
||||
analysis = analyze_request_heuristics(prompt)
|
||||
assert analysis.task_type == TaskType.SUMMARIZATION
|
||||
assert analysis.complexity_label in (PriorityLevel.LOW, PriorityLevel.MEDIUM)
|
||||
|
||||
|
||||
def test_heuristic_classification_math():
|
||||
prompt = "Calculate the derivative and integral of matrix eigenvalue differential equation"
|
||||
analysis = analyze_request_heuristics(prompt)
|
||||
assert analysis.task_type == TaskType.MATH
|
||||
assert analysis.reasoning_required is True
|
||||
@@ -0,0 +1,66 @@
|
||||
import pytest
|
||||
from httpx import AsyncClient, ASGITransport
|
||||
from main import app
|
||||
from app.storage.database import init_db
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
async def setup_db():
|
||||
await init_db()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_api_health():
|
||||
transport = ASGITransport(app=app)
|
||||
async with AsyncClient(transport=transport, base_url="http://test") as client:
|
||||
res = await client.get("/api/health")
|
||||
assert res.status_code == 200
|
||||
data = res.json()
|
||||
assert data["status"] == "healthy"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_api_route_only():
|
||||
transport = ASGITransport(app=app)
|
||||
async with AsyncClient(transport=transport, base_url="http://test") as client:
|
||||
payload = {"prompt": "Write a Python script to benchmark models", "policy": "balanced"}
|
||||
res = await client.post("/api/route", json=payload)
|
||||
assert res.status_code == 200
|
||||
data = res.json()
|
||||
assert "decision" in data
|
||||
assert "analysis" in data
|
||||
assert data["analysis"]["task_type"] == "CODING"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_api_generate_end_to_end():
|
||||
transport = ASGITransport(app=app)
|
||||
async with AsyncClient(transport=transport, base_url="http://test") as client:
|
||||
payload = {"prompt": "Debug this memory leak in python subprocess", "policy": "balanced"}
|
||||
res = await client.post("/api/generate", json=payload)
|
||||
assert res.status_code == 200
|
||||
data = res.json()
|
||||
assert "response" in data
|
||||
assert "metrics" in data
|
||||
assert len(data["response"]["content"]) > 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_api_models_list():
|
||||
transport = ASGITransport(app=app)
|
||||
async with AsyncClient(transport=transport, base_url="http://test") as client:
|
||||
res = await client.get("/api/models")
|
||||
assert res.status_code == 200
|
||||
models = res.json()
|
||||
assert len(models) >= 3
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_api_analytics():
|
||||
transport = ASGITransport(app=app)
|
||||
async with AsyncClient(transport=transport, base_url="http://test") as client:
|
||||
res = await client.get("/api/analytics")
|
||||
assert res.status_code == 200
|
||||
data = res.json()
|
||||
assert "total_requests" in data
|
||||
assert "savings" in data
|
||||
@@ -0,0 +1,21 @@
|
||||
import pytest
|
||||
from app.providers.mock_provider import MockProvider
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mock_provider_generate():
|
||||
provider = MockProvider()
|
||||
resp = await provider.generate("Write a quick Python sort function", model_id="mock-fast")
|
||||
assert resp.is_mock is True
|
||||
assert "DEMO/MOCK" in resp.content
|
||||
assert resp.input_tokens > 0
|
||||
assert resp.output_tokens > 0
|
||||
assert resp.provider_latency_ms > 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mock_provider_health():
|
||||
provider = MockProvider()
|
||||
health = await provider.check_health()
|
||||
assert health["status"] == "CONNECTED"
|
||||
assert "mock-fast" in health["models_available"]
|
||||
@@ -0,0 +1,78 @@
|
||||
import pytest
|
||||
from app.models.schemas import ModelMetadata, ModelTier, RequestAnalysis, TaskType, PriorityLevel
|
||||
from app.router.scoring import filter_candidate_models, compute_candidate_score
|
||||
from app.router.engine import route_request
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_models():
|
||||
return [
|
||||
ModelMetadata(
|
||||
id="mock-fast",
|
||||
name="Mock Fast",
|
||||
provider="mock",
|
||||
tier=ModelTier.FAST,
|
||||
context_window=8192,
|
||||
supports_coding=True,
|
||||
supports_reasoning=False,
|
||||
quality_score=0.75,
|
||||
speed_score=0.98,
|
||||
cost_per_input_token=0.0,
|
||||
cost_per_output_token=0.0,
|
||||
),
|
||||
ModelMetadata(
|
||||
id="mock-power",
|
||||
name="Mock Power",
|
||||
provider="mock",
|
||||
tier=ModelTier.POWER,
|
||||
context_window=128000,
|
||||
supports_coding=True,
|
||||
supports_reasoning=True,
|
||||
quality_score=0.98,
|
||||
speed_score=0.60,
|
||||
cost_per_input_token=0.000005,
|
||||
cost_per_output_token=0.000015,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def test_candidate_filtering_context_overflow(sample_models):
|
||||
large_analysis = RequestAnalysis(
|
||||
task_type=TaskType.LONG_CONTEXT,
|
||||
complexity=0.7,
|
||||
context_size=15000, # exceeds mock-fast 8192
|
||||
)
|
||||
eligible, rejected = filter_candidate_models(sample_models, large_analysis)
|
||||
assert len(eligible) == 1
|
||||
assert eligible[0].id == "mock-power"
|
||||
assert "mock-fast" in rejected
|
||||
|
||||
|
||||
def test_scoring_weights(sample_models):
|
||||
analysis = RequestAnalysis(
|
||||
task_type=TaskType.CODING,
|
||||
complexity=0.3,
|
||||
latency_priority=PriorityLevel.HIGH,
|
||||
cost_sensitivity=PriorityLevel.HIGH,
|
||||
)
|
||||
weights = {"quality_weight": 0.1, "cost_weight": 0.5, "speed_weight": 0.4, "capability_weight": 0.0, "reliability_weight": 0.0}
|
||||
score_fast = compute_candidate_score(sample_models[0], analysis, weights)
|
||||
score_power = compute_candidate_score(sample_models[1], analysis, weights)
|
||||
assert score_fast.overall_score > score_power.overall_score
|
||||
|
||||
|
||||
def test_routing_decision_explainability(sample_models):
|
||||
analysis = RequestAnalysis(
|
||||
task_type=TaskType.DEBUGGING,
|
||||
complexity=0.85,
|
||||
complexity_label=PriorityLevel.HIGH,
|
||||
reasoning_required=True,
|
||||
coding_required=True,
|
||||
quality_requirement=PriorityLevel.HIGH,
|
||||
)
|
||||
weights = {"quality_weight": 0.6, "cost_weight": 0.1, "speed_weight": 0.1, "capability_weight": 0.1, "reliability_weight": 0.1}
|
||||
decision = route_request(analysis, sample_models, weights, policy_name="highest_quality")
|
||||
assert decision.selected_model == "mock-power"
|
||||
assert decision.confidence > 0.70
|
||||
assert len(decision.reasons) > 0
|
||||
assert any("DEBUGGING" in r for r in decision.reasons)
|
||||
Reference in New Issue
Block a user