Spaces:
Sleeping
Sleeping
File size: 2,650 Bytes
0f336cf | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 | import asyncio
import os
import random
from typing import Any, Dict, Optional
from schemas.enums import LogicalModel
from models.key_rotator import APIKeyRotationEngine
from telemetry.resource_manager import ColonyResourceManager
class ModelManager:
"""Universal Model Gateway mapping logical models (mdl_fst, mdl_adv) to dynamic physical models."""
def __init__(self, key_rotator: APIKeyRotationEngine, resource_manager: ColonyResourceManager):
self.key_rotator = key_rotator
self.resource_manager = resource_manager
# Model mappings loaded from environment variable with fallbacks
self.logical_models = {
LogicalModel.MDL_FST.value: os.getenv("MDL_FST", "gemini-2.5-flash"),
LogicalModel.MDL_ADV.value: os.getenv("MDL_ADV", "llama-3.3-70b-versatile"),
}
self.total_calls = 0
async def generate_response(
self,
logical_model: LogicalModel,
prompt: str,
system_prompt: Optional[str] = None,
vision_input: Optional[str] = None,
) -> Dict[str, Any]:
self.total_calls += 1
physical_model = self.logical_models.get(logical_model.value, "gemini-2.5-flash")
provider = "groq" if "llama" in physical_model.lower() or "groq" in physical_model.lower() else "gemini"
key = self.key_rotator.acquire_key(provider)
if not key:
# Fallback to any provider if primary unavailable
provider = "gemini" if provider == "groq" else "groq"
key = self.key_rotator.acquire_key(provider)
try:
await asyncio.sleep(0.05) # Simulated LLM latency execution
tokens_used = len(prompt.split()) + random.randint(50, 150)
if provider == "gemini":
self.resource_manager.record_gemini_tokens(tokens_used)
else:
self.resource_manager.record_groq_tokens(tokens_used)
response_text = f"Synthesized Response via [{logical_model.value} -> {physical_model}]: Processed input prompt ({len(prompt)} chars)."
if key:
self.key_rotator.release_key(key, success=True)
return {
"logical_model": logical_model.value,
"physical_model": physical_model,
"provider": provider,
"key_used": key.key_id if key else "unauthenticated",
"tokens_used": tokens_used,
"content": response_text,
}
except Exception as e:
if key:
self.key_rotator.release_key(key, success=False, error_msg=str(e))
raise e
|