app/ai/orchestrator.py

import json
import logging
import re
from typing import Any

from app.ai.providers import (
    ChatProvider,
    InvalidToolCallGenerationError,
    RetryableProviderError,
)
from app.models.domain import AIProviderResponse, ToolCall, ToolResult
from app.tools.registry import ToolRegistry

logger = logging.getLogger(__name__)


class AIOrchestrator:
    def __init__(
        self,
        *,
        primary: ChatProvider,
        fallback: ChatProvider,
        temperature: float,
        max_tool_iterations: int,
    ) -> None:
        self.primary = primary
        self.fallback = fallback
        self.temperature = temperature
        self.max_tool_iterations = max_tool_iterations

    async def chat(
        self,
        *,
        messages: list[dict[str, Any]],
        tools: list[dict[str, Any]] | None = None,
        tool_choice: str | dict[str, Any] | None = None,
        temperature: float = 0.2,
    ) -> AIProviderResponse:
        try:
            return await self.primary.chat(
                messages,
                tools=tools,
                tool_choice=tool_choice,
                temperature=temperature,
            )
        except InvalidToolCallGenerationError:
            logger.warning("Primary provider generated invalid tool call; retrying once")
            try:
                return await self.primary.chat(
                    messages,
                    tools=tools,
                    tool_choice=tool_choice,
                    temperature=max(temperature - 0.2, 0.1),
                )
            except RetryableProviderError:
                logger.warning("Primary retry failed; falling back to OpenRouter")
        except RetryableProviderError:
            logger.warning("Primary provider failed; falling back to OpenRouter")

        return await self.fallback.chat(
            messages,
            tools=tools,
            tool_choice=tool_choice,
            temperature=temperature,
        )

    async def generate_reply(
        self,
        *,
        messages: list[dict[str, Any]],
        tools: list[dict[str, Any]],
        registry: ToolRegistry,
    ) -> str:
        try:
            return await self._run_provider(
                self.primary,
                messages=messages,
                tools=tools,
                registry=registry,
                temperature=self.temperature,
            )
        except InvalidToolCallGenerationError:
            logger.warning("Primary provider generated invalid tool call; retrying once")
            try:
                return await self._run_provider(
                    self.primary,
                    messages=messages,
                    tools=tools,
                    registry=registry,
                    temperature=max(self.temperature - 0.2, 0.1),
                )
            except RetryableProviderError:
                logger.warning("Primary retry failed; falling back to OpenRouter")
        except RetryableProviderError:
            logger.warning("Primary provider failed; falling back to OpenRouter")

        return await self._run_provider(
            self.fallback,
            messages=messages,
            tools=tools,
            registry=registry,
            temperature=self.temperature,
        )

    async def _run_provider(
        self,
        provider: ChatProvider,
        *,
        messages: list[dict[str, Any]],
        tools: list[dict[str, Any]],
        registry: ToolRegistry,
        temperature: float,
    ) -> str:
        working_messages = [dict(message) for message in messages]
        # need to be edited to reponse quackly instead in enter in a for loop
        for _ in range(self.max_tool_iterations + 1):
            response = await provider.chat(
                working_messages,
                tools=tools,
                tool_choice="auto",
                temperature=temperature,
            )
            if not response.tool_calls:
                content = _normalize_brand_name((response.content or "").strip())
                if content:
                    return content
                raise RetryableProviderError(f"{provider.name} returned an empty response")
            working_messages.append(_assistant_tool_message(response))
            for tool_call in response.tool_calls:
                logger.warning("++++++++"+str(tool_call)+ "&&&"+ str(registry))
                result = await _execute_tool_call(registry, tool_call)
                logger.warning("+++++++++++++"+str(result))
                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 _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
        ],
    }


def _normalize_brand_name(text: str) -> str:
    if not text:
        return text

    replacements = {
        "فلسا": "فلزة",
        "فلظ": "فلزة",
        "فلظة": "فلزة",
        "فلز": "فلزة",
        "فلِزة": "فلزة",
        "فلَزة": "فلزة",
        "فلٰزة": "فلزة",
    }

    for wrong, correct in sorted(replacements.items(), key=lambda item: -len(item[0])):
        pattern = rf"(?<![\w\u0600-\u06FF]){re.escape(wrong)}(?![\w\u0600-\u06FF])"
        text = re.sub(pattern, correct, text)
    return text


async def _execute_tool_call(registry: ToolRegistry, tool_call: ToolCall) -> ToolResult:
    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:
        return ToolResult(ok=False, data={}, error=f"Invalid tool arguments: {exc}")

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