python/src/server/services/llm/base.py

import asyncio
import time
from typing import cast

from ...config.logfire_config import get_logger

logger = get_logger(__name__)


# --- Mock Classes ---
class MockMessage:
    def __init__(self, content):
        self.content = content
        self.reasoning_content = None


class MockChoice:
    def __init__(self, content):
        self.message = MockMessage(content)


class MockResponse:
    def __init__(self, content):
        self.choices = [MockChoice(content)]
        self.usage = None


class MockCompletions:
    def __init__(self, provider_name):
        self.provider_name = provider_name

    async def create(self, *args, **kwargs):
        logger.info(f"MockLLMClient ({self.provider_name}) received request: {kwargs}")
        messages = kwargs.get("messages", [])
        last_user_content = ""
        for m in reversed(messages):
            if m.get("role") == "user":
                last_user_content = m.get("content", "")
                break

        # Original Mock Context-Aware Responses Restored
        response_content = f"✨ [Mock] Magic Content for: {last_user_content[:30]}..."
        if "pitch" in last_user_content.lower() or "job" in last_user_content.lower():
            response_content = "[ENGLISH PITCH]\nHi there, I noticed you're hiring...\n\n[CHINESE PITCH]\n您好，這是一份模擬的銷售信件..."
        elif "image" in last_user_content.lower() or "nana" in last_user_content.lower():
            response_content = "A beautiful futuristic city with glowing lights"

        return MockResponse(response_content)


class MockChat:
    def __init__(self, provider_name):
        self.completions = MockCompletions(provider_name)


class MockLLMClient:
    def __init__(self, provider_name="mock"):
        self.chat = MockChat(provider_name)
        self.models = None

    async def close(self):
        try:
            from ..token_usage_service import TokenUsageService

            # Simulation of usage logging
            asyncio.create_task(
                TokenUsageService.log_usage(
                    request_id=f"mock-{int(time.time())}",
                    user_id="mock-user-001",
                    model="mock-gpt-4",
                    provider="mock",
                    input_tokens=50,
                    output_tokens=100,
                    context_type="mock_generation",
                )
            )
        except Exception:
            pass

    async def aclose(self):
        await self.close()


# --- Tracking Classes ---
class UsageTrackingCompletions:
    def __init__(self, original_completions, context):
        self._original = original_completions
        self._context = context

    async def create(self, *args, **kwargs):
        import os

        import openai

        from ...utils.retry_utils import retry_with_backoff
        from ..credential_service import credential_service

        forced_tier_str = await credential_service.get_credential("forced_fallback_tier")
        try:
            forced_tier = int(forced_tier_str) if forced_tier_str else 0
        except Exception:
            forced_tier = 0

        async def _execute_on_hf(model_name: str):
            hf_token = await credential_service.get_credential("HF_TOKEN")
            if not hf_token:
                raise ValueError("HF_TOKEN not configured for Tier 2 fallback")
            hf_model = "google/gemma-1.1-2b-it"
            client = openai.AsyncOpenAI(api_key=hf_token, base_url="https://api-inference.huggingface.co/v1/")
            try:
                kwargs_copy = kwargs.copy()
                kwargs_copy["model"] = hf_model
                return await client.chat.completions.create(*args, **kwargs_copy)
            finally:
                await client.close()

        async def _execute_on_ollama():
            from .clients import _get_optimal_ollama_instance
            url = await _get_optimal_ollama_instance("chat", False, None)
            client = openai.AsyncOpenAI(api_key="ollama", base_url=url)
            try:
                kwargs_copy = kwargs.copy()
                kwargs_copy["model"] = "gemma3"
                return await client.chat.completions.create(*args, **kwargs_copy)
            finally:
                await client.close()

        @retry_with_backoff(max_retries=5, initial_delay=2.0)
        async def _execute(override_key: str | None = None):
            original_client = self._original._client
            original_api_key = original_client.api_key
            original_headers = getattr(original_client, "default_headers", {})

            try:
                if override_key:
                    original_client.api_key = override_key
                    if "x-goog-api-key" in original_headers:
                        new_headers = dict(original_headers)
                        new_headers["x-goog-api-key"] = override_key
                        original_client.default_headers = new_headers

                return await self._original.create(*args, **kwargs)
            finally:
                if override_key:
                    original_client.api_key = original_api_key
                    original_client.default_headers = original_headers

        # Scan for Lean 4 context
        is_lean = False
        proof_context = ""
        for m in kwargs.get("messages", []):
            if isinstance(m, dict):
                content = m.get("content", "") or ""
            else:
                content = getattr(m, "content", "") or ""
            if "lean 4" in content.lower() or "lake build" in content.lower() or "theorem" in content.lower():
                is_lean = True
                proof_context += content + "\n"

        retry_count = 0
        if "extra_body" in kwargs and isinstance(kwargs["extra_body"], dict):
            retry_count = kwargs["extra_body"].get("retry_count", 0)

        if forced_tier == 2:
            logger.info("Forced Tier 2 Fallback (Hugging Face) by Human Operator")
            credential_service.set_active_tier(2)
            response = await _execute_on_hf(kwargs.get("model", ""))
        elif forced_tier == 3:
            logger.info("Forced Tier 3 Fallback (Ollama) by Human Operator")
            credential_service.set_active_tier(3)
            response = await _execute_on_ollama()
        elif is_lean:
            from .hybrid_router import hybrid_router
            if hybrid_router.should_escalate_to_cloud(proof_context, retry_count):
                logger.info("Hybrid Router: Escalating Lean proof task to Tier 1 Cloud")
                try:
                    response = await _execute()
                    credential_service.set_active_tier(1)
                except Exception:
                    logger.warning("Tier 1 Cloud failed for escalated Lean task, trying Tier 3")
                    credential_service.set_active_tier(3)
                    response = await _execute_on_ollama()
            else:
                logger.info("Hybrid Router: Routing Lean proof task to Tier 3 Ollama (Local)")
                try:
                    credential_service.set_active_tier(3)
                    response = await _execute_on_ollama()
                except Exception:
                    logger.warning("Local Tier 3 failed for Lean task, falling back to Tier 1")
                    response = await _execute()
                    credential_service.set_active_tier(1)
        else:
            try:
                from .hybrid_router import hybrid_router
                if hybrid_router.is_query_simple_and_offline(kwargs.get("messages", [])):
                    logger.info("Hybrid Router: Routing simple query to Tier 3 Ollama (Local)")
                    credential_service.set_active_tier(3)
                    response = await _execute_on_ollama()
                else:
                    response = await _execute()
                    credential_service.set_active_tier(1)
            except Exception as e:
                err_msg = str(e)
                provider = self._context.get("provider", "unknown")
                logger.warning(f"Tier 1 (or simple query local) failed: {err_msg}")

                if forced_tier == 1:
                    raise e

                # Connection Error -> Go straight to Tier 3
                if isinstance(e, openai.APIConnectionError) or "connect" in err_msg.lower():
                    logger.error("Connection error. Bypassing Tier 2, falling back directly to Tier 3 (Ollama)...")
                    try:
                        credential_service.set_active_tier(3)
                        response = await _execute_on_ollama()
                    except Exception as ollama_e:
                        logger.error(f"Tier 3 (Ollama) fallback failed: {ollama_e}")
                        raise ollama_e from e
                # Authentication or Rate Limit Error -> Try Tier 2 (HF)
                elif isinstance(e, (openai.AuthenticationError, openai.RateLimitError)) or "429" in err_msg or "401" in err_msg:
                    try:
                        if provider == "google":
                            primary_key = os.getenv("GEMINI_API_KEY")
                            google_key_backup = os.getenv("GOOGLE_API_KEY")
                            if google_key_backup and google_key_backup != primary_key:
                                logger.warning("⚠️ Primary GEMINI_API_KEY exhausted. Rotating to backup...")
                                response = await _execute(override_key=google_key_backup)
                                credential_service.set_active_tier(1)
                                return response
                    except Exception as backup_e:
                        logger.error(f"Backup key failed: {backup_e}")
                        e = backup_e

                    logger.warning("Attempting Tier 2 (Hugging Face) fallback...")
                    try:
                        credential_service.set_active_tier(2)
                        response = await _execute_on_hf(kwargs.get("model", ""))
                    except Exception as hf_e:
                        logger.error(f"Tier 2 (HF) failed: {hf_e}. Falling back to Tier 3 (Ollama)...")
                        try:
                            credential_service.set_active_tier(3)
                            response = await _execute_on_ollama()
                        except Exception as ollama_e:
                            logger.error(f"Tier 3 (Ollama) failed: {ollama_e}")
                            raise ollama_e from hf_e
                else:
                    logger.warning("Unhandled Tier 1 error. Trying Tier 3 (Ollama) fallback...")
                    try:
                        credential_service.set_active_tier(3)
                        response = await _execute_on_ollama()
                    except Exception as last_e:
                        logger.error(f"Tier 3 fallback failed: {last_e}")
                        raise e from None

        try:
            if hasattr(response, "usage") and response.usage:
                model = kwargs.get("model", "unknown")
                from ..token_usage_service import TokenUsageService

                # Use ensure_future to not block response (Restored from Original)
                asyncio.ensure_future(
                    TokenUsageService.log_usage(
                        request_id=str(self._context.get("request_id", "")),
                        user_id=cast(str | None, self._context.get("user_id")),
                        model=str(model),
                        provider=str(self._context.get("provider", "unknown")),
                        input_tokens=int(response.usage.prompt_tokens),
                        output_tokens=int(response.usage.completion_tokens),
                        context_type="llm_client_call",
                    )
                )
        except Exception as e:
            logger.warning(f"Failed to log token usage: {e}")
        return response


class UsageTrackingChat:
    def __init__(self, original_chat, context):
        self._original = original_chat
        self.completions = UsageTrackingCompletions(original_chat.completions, context)

    def __getattr__(self, name):
        return getattr(self._original, name)


class UsageTrackingClient:
    def __init__(self, original_client, user_id, request_id, provider):
        self._original = original_client
        self._context = {"user_id": user_id, "request_id": request_id, "provider": provider}
        self.chat = UsageTrackingChat(original_client.chat, self._context)

    def __getattr__(self, name):
        return getattr(self._original, name)