app/ai/router.py

import json
import logging
from typing import Any

import litellm

from app.models.domain import AIProviderResponse, ToolCall
from app.tools.registry import ToolRegistry

logger = logging.getLogger(__name__)


class LiteLLMOrchestration:
    """Provides AI orchestration using LiteLLM Router for automatic API management"""

    def __init__(
        self,
        groq_api_key: str,
        groq_model: str,
        openrouter_api_key: str | None = None,
        openrouter_model: str | None = None,
        groq_api_key_2: str | None = None,
        groq_model_2: str | None = None,
        google_ai_api_key: str | None = None,
        google_ai_model: str | None = None,
        max_tool_iterations: int = 3,
        temperature: float = 0.2,
    ) -> None:
        self.name = "litellm-orchestration"
        self.groq_api_key = groq_api_key
        self.groq_model = groq_model
        self.groq_api_key_2 = groq_api_key_2
        self.groq_model_2 = groq_model_2
        self.google_ai_api_key = google_ai_api_key
        self.google_ai_model = google_ai_model
        self.openrouter_api_key = openrouter_api_key
        self.openrouter_model = openrouter_model
        self.max_tool_iterations = max_tool_iterations
        self.temperature = temperature
        self._initialized = False

    @staticmethod
    def _litellm_model(model: str, provider: str) -> str:
        """Ensure model name has the litellm provider prefix.

        litellm strips the provider prefix before sending to the API,
        so 'groq/openai/gpt-oss-120b' sends 'openai/gpt-oss-120b' to Groq.

        Also strips non-litellm prefixes (e.g. 'google/' from OpenRouter
        model IDs) before applying the correct litellm provider prefix.
        """
        for prefix in ("google",):
            if model.startswith(f"{prefix}/"):
                model = model[len(prefix) + 1 :]
                break

        if model.startswith(f"{provider}/"):
            return model
        return f"{provider}/{model}"

    async def _ensure_initialized(self):
        if not self._initialized:
            self.router = litellm.Router(
                model_list=[
                    {
                        "model_name": "groq",
                        "litellm_params": {
                            "model": self._litellm_model(self.groq_model, "groq"),
                            "api_key": self.groq_api_key,
                            "api_base": "https://api.groq.com/openai/v1",
                        },
                    },
                    {
                        "model_name": "groq_2",
                        "litellm_params": {
                            "model": self._litellm_model(self.groq_model_2 or self.groq_model, "groq"),
                            "api_key": self.groq_api_key_2 or self.groq_api_key,
                            "api_base": "https://api.groq.com/openai/v1",
                        },
                    },
                    {
                        "model_name": "google_ai_studio",
                        "litellm_params": {
                            "model": self._litellm_model(self.google_ai_model or self.openrouter_model, "gemini"),
                            "api_key": self.google_ai_api_key or self.openrouter_api_key,
                            "api_base": "https://generativelanguage.googleapis.com/v1beta",
                        },
                    },
                    {
                        "model_name": "openrouter",
                        "litellm_params": {
                            "model": self._litellm_model(self.openrouter_model, "openrouter"),
                            "api_key": self.openrouter_api_key,
                            "api_base": "https://openrouter.ai/api/v1",
                        },
                    },
                ],
                fallbacks=[{"groq": ["groq_2", "google_ai_studio", "openrouter"]},],
                num_retries=0,
                retry_policy={
                    "TimeoutErrorRetries": 0,
                    "RateLimitErrorRetries": 0,
                    "InternalServerErrorRetries": 0,
                },
                routing_strategy="latency-based-routing",
                set_verbose=False,
            )
            self._initialized = True

    async def chat(
        self,
        *,
        messages: list[dict[str, Any]],
        tools: list[dict[str, Any]] | None = None,
        temperature: float | None = None,
        tool_choice: str | None = "auto",
    ) -> AIProviderResponse:
        if temperature is None:
            temperature = self.temperature

        await self._ensure_initialized()

        kwargs = {
            "messages": messages,
            "model": "groq",
            "temperature": temperature,
        }

        if tools:
            kwargs["tools"] = tools
            kwargs["tool_choice"] = tool_choice

        try:
            response = await self.router.acompletion(**kwargs)
            return self._adapt_response(response)
        except Exception as exc:
            logger.error("LiteLLM Router error: %s", exc)
            raise RuntimeError(f"AI request failed: {exc}") from exc

    async def generate_reply(
        self,
        *,
        messages: list[dict[str, Any]],
        tools: list[dict[str, Any]],
        registry: ToolRegistry,
    ) -> str:
        working_messages = [dict(msg) for msg in messages]

        for _ in range(self.max_tool_iterations + 1):
            response = await self.chat(
                messages=working_messages,
                tools=tools,
                temperature=self.temperature,
            )

            if not response.tool_calls:
                content = (response.content or "").strip()
                if content:
                    return content
                continue

            working_messages.append(self._assistant_tool_message(response))
            for tool_call in response.tool_calls:
                result = await self._execute_tool_call(registry, tool_call)
                if result.suppress_llm_reply:
                    return result.error or ""
                working_messages.append(
                    {
                        "role": "tool",
                        "tool_call_id": tool_call.id,
                        "name": tool_call.name,
                        "content": json.dumps(result.to_payload(), ensure_ascii=False),
                    }
                )

        return (
            "I found that this request needs extra checking. "
            "A support team member will follow up with you shortly."
        )

    def _adapt_response(self, response: Any) -> AIProviderResponse:
        raw_message: dict[str, Any] = {}

        if hasattr(response, "choices") and response.choices:
            choice = response.choices[0]
            if hasattr(choice, "message") and choice.message:
                if hasattr(choice.message, "model_dump"):
                    raw_message = choice.message.model_dump(exclude_none=True)
                elif isinstance(choice.message, dict):
                    raw_message = choice.message
                else:
                    raw_message = {
                        "content": getattr(choice.message, "content", None),
                        "tool_calls": getattr(choice.message, "tool_calls", None),
                    }

        tool_calls = []
        for tool_call in raw_message.get("tool_calls") or []:
            if hasattr(tool_call, "function"):
                function = tool_call.function
                name = function.name
                arguments = function.arguments or "{}"
                id_val = getattr(tool_call, "id", None) or f"tool-call-{name}"
            else:
                function = tool_call.get("function") or {}
                name = function.get("name", "")
                arguments = function.get("arguments") or "{}"
                id_val = tool_call.get("id") or f"tool-call-{name}"

            tool_calls.append(
                ToolCall(
                    id=id_val,
                    name=name,
                    arguments=arguments,
                )
            )

        return AIProviderResponse(
            content=raw_message.get("content"),
            tool_calls=tool_calls,
            raw_message=raw_message,
        )

    @staticmethod
    def _assistant_tool_message(response: AIProviderResponse) -> dict[str, Any]:
        if response.raw_message:
            return response.raw_message
        return {
            "role": "assistant",
            "content": response.content,
            "tool_calls": [
                {
                    "id": tool_call.id,
                    "type": "function",
                    "function": {"name": tool_call.name, "arguments": tool_call.arguments},
                }
                for tool_call in response.tool_calls
            ],
        }

    @staticmethod
    async def _execute_tool_call(
        registry: ToolRegistry, tool_call: ToolCall
    ) -> Any:
        try:
            arguments = json.loads(tool_call.arguments or "{}")
            if not isinstance(arguments, dict):
                raise ValueError("Tool arguments must be a JSON object")
        except (json.JSONDecodeError, ValueError) as exc:
            from app.models.domain import ToolResult

            return ToolResult(
                ok=False, data={}, error=f"Invalid tool arguments: {exc}"
            )

        return await registry.execute(tool_call.name, arguments)