models/model_manager.py

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