Artemis-Eapil/artemis/services/llm.py

# Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Universal LLM Service for Artemis.

Provides connection-pooled model instantiation, robust retries,
thought stream recording, and role-based dynamic dispatching.
"""

import asyncio
from collections.abc import Awaitable, Callable, Coroutine
from dataclasses import replace
import functools
import logging
from pathlib import Path
import random
from typing import Any, Literal, TypeVar, overload
from uuid import uuid4

from langchain_core.language_models.chat_models import BaseChatModel

from artemis.config import (
    AgentNode,
    AgentNodeWithFallback,
    LLMUtilsNode,
    LLMUtilsNodeWithFallback,
    LLMWithFallback,
    get_default_llm_config,
    settings,
)
from artemis.context import ArtemisContext
from artemis.data_engine.engine import _CURRENT_DATA_ENGINE
from artemis.data_engine.trace import CURRENT_TRACE_ID, DataEngineCallbackHandler
from artemis.llm.router import ModelEndpoint, ModelFactory, ModelProvider
from artemis.utils.logger import get_logger

# Logger for internal messages
llm_logger = logging.getLogger(__name__)
# Logger for user-facing messages
user_messages_logger = get_logger(__name__)

T = TypeVar("T")


def _handle_llm_pause_and_resume(last_error: Exception) -> Path:
    """Handles pausing task execution upon persistent failure until resumed."""
    err_msg = str(last_error)
    if len(err_msg) > 1000:
        err_msg = err_msg[:1000] + "... [Truncated by LLM Wrapper]"

    llm_logger.warning(f"LLM Error: {err_msg}. Pausing execution to wait for resume signal...")
    pause_file = Path(settings.TRACES_PATH).parent / ".artemis_paused"
    try:
        pause_file.write_text(f"LLM Error: {err_msg}")
    except Exception:
        pass

    if _CURRENT_DATA_ENGINE:
        _CURRENT_DATA_ENGINE._publish("task_paused", {"error": err_msg})

    return pause_file


def robust_retry_async(func: Callable) -> Callable:
    """Decorator to add robust retry and pause/resume logic to an async LLM call."""

    @functools.wraps(func)
    async def wrapper(*args, **kwargs):
        while True:
            max_retries = 3
            retry_delay = 0.0
            last_error = None
            for attempt in range(max_retries):
                try:
                    return await func(*args, **kwargs)
                except Exception as e:
                    last_error = e
                    err_str = str(e).lower()
                    if attempt == max_retries - 1:
                        break
                    if attempt > 0:
                        retry_delay = 0.5 * (2 ** (attempt - 1)) + random.uniform(0.1, 0.3)
                    if any(
                        err in err_str
                        for err in [
                            "503",
                            "throttled",
                            "overloaded",
                            "429",
                            "quota",
                            "resource_exhausted",
                        ]
                    ):
                        retry_delay = max(retry_delay, 10.0)
                    if retry_delay > 0:
                        llm_logger.warning(f"LLM call failed, retrying in {retry_delay:.2f}s...")

                        if _CURRENT_DATA_ENGINE:
                            _CURRENT_DATA_ENGINE._publish(
                                "llm_retrying",
                                {
                                    "attempt": attempt + 1,
                                    "max_retries": max_retries,
                                    "delay": retry_delay,
                                    "error": str(e),
                                },
                            )
                        await asyncio.sleep(retry_delay)

            pause_file = _handle_llm_pause_and_resume(last_error)
            while pause_file.exists():
                await asyncio.sleep(1)

            llm_logger.info("Resume signal received, retrying LLM call...")

            if _CURRENT_DATA_ENGINE:
                _CURRENT_DATA_ENGINE._publish("task_resumed", {})

    return wrapper


def robust_retry_astream(func: Callable) -> Callable:
    """Decorator to add robust retry and pause/resume logic to an async stream LLM call."""

    @functools.wraps(func)
    async def wrapper(*args, **kwargs):
        while True:
            max_retries = 3
            retry_delay = 0.0
            last_error = None
            for attempt in range(max_retries):
                try:
                    gen = func(*args, **kwargs)
                    first = await gen.__anext__()
                    yield first
                    async for chunk in gen:
                        yield chunk
                    return
                except Exception as e:
                    last_error = e
                    err_str = str(e).lower()
                    if attempt == max_retries - 1:
                        break
                    if attempt > 0:
                        retry_delay = 0.5 * (2 ** (attempt - 1)) + random.uniform(0.1, 0.3)
                    if any(
                        err in err_str
                        for err in [
                            "503",
                            "throttled",
                            "overloaded",
                            "429",
                            "quota",
                            "resource_exhausted",
                        ]
                    ):
                        retry_delay = max(retry_delay, 10.0)
                    if retry_delay > 0:
                        llm_logger.warning(
                            f"LLM stream handshake throttled, retrying in {retry_delay:.2f}s..."
                        )

                        if _CURRENT_DATA_ENGINE:
                            _CURRENT_DATA_ENGINE._publish(
                                "llm_retrying",
                                {
                                    "attempt": attempt + 1,
                                    "max_retries": max_retries,
                                    "delay": retry_delay,
                                    "error": str(e),
                                },
                            )
                        await asyncio.sleep(retry_delay)

            pause_file = _handle_llm_pause_and_resume(last_error)
            while pause_file.exists():
                await asyncio.sleep(1)

            llm_logger.info("Resume signal received, retrying LLM stream...")

            if _CURRENT_DATA_ENGINE:
                _CURRENT_DATA_ENGINE._publish("task_resumed", {})

    return wrapper


def _inject_parent_trace_id(trace_id, *args, **kwargs):
    config = kwargs.get("config") or {}
    if not kwargs.get("config") and len(args) > 1:
        args_list = list(args)
        if isinstance(args_list[1], dict):
            config = args_list[1]
            args_list[1] = config
            args = tuple(args_list)

    if "metadata" not in config:
        config["metadata"] = {}
    config["metadata"]["parent_trace_id"] = str(trace_id)
    kwargs["config"] = config
    return args, kwargs


class RobustChatModelWrapper:
    """Universal LangChain Runnable wrapper ensuring robust execution and telemetry."""

    def __init__(
        self,
        base_model: BaseChatModel,
        ctx: ArtemisContext = None,
        endpoint: ModelEndpoint | None = None,
    ):
        self.base_model = base_model
        self.ctx = ctx
        self.endpoint = endpoint

    def __getattr__(self, item: str):
        return getattr(self.base_model, item)

    def bind_tools(self, *args, **kwargs):
        # Extract tools list from positional args or keyword args
        tools_list = None
        other_args = list(args)
        if other_args:
            tools_list = other_args.pop(0)
        elif "tools" in kwargs:
            tools_list = kwargs.pop("tools")

        processed_tools = list(tools_list) if tools_list is not None else []

        # Determine if the underlying model provider is Google / Gemini
        is_google_provider = False
        if self.endpoint is not None:
            is_google_provider = self.endpoint.provider in (
                ModelProvider.GOOGLE,
                ModelProvider.GEMINI,
            )
        else:
            base_model_cls = getattr(self.base_model, "__class__", None)
            base_model_name = base_model_cls.__name__ if base_model_cls else ""
            is_google_provider = (
                "Google" in base_model_name
                or "Gemini" in base_model_name
                or "VertexAI" in base_model_name
            )

        has_explicit_google_search = any(
            isinstance(t, dict) and "google_search" in t for t in processed_tools
        )
        should_ground = (
            self.endpoint is not None and self.endpoint.enable_grounding
        ) or has_explicit_google_search

        if is_google_provider:
            if should_ground:
                # Add native Google Search Grounding tool if not already present
                if not has_explicit_google_search:
                    processed_tools.append({"google_search": {}})

                # Ensure include_server_side_tool_invocations is True in tool_config
                # to satisfy Gemini API requirements when mixing built-in tools and function calling
                tool_config = kwargs.get("tool_config")
                if tool_config is None:
                    kwargs["tool_config"] = {"include_server_side_tool_invocations": True}
                elif isinstance(tool_config, dict):
                    tool_config = dict(tool_config)
                    tool_config.setdefault("include_server_side_tool_invocations", True)
                    kwargs["tool_config"] = tool_config
        else:
            # If not Gemini/Google (e.g. OpenAI, Anthropic, Ollama, OpenRouter, XAI):
            # Ignore grounding and filter out any Gemini-specific built-in dict tools to avoid provider errors
            processed_tools = [
                t for t in processed_tools if not (isinstance(t, dict) and "google_search" in t)
            ]

        bound = self.base_model.bind_tools(processed_tools, *other_args, **kwargs)
        return RobustChatModelWrapper(bound, self.ctx, endpoint=self.endpoint)

    def with_config(self, *args, **kwargs):
        return RobustChatModelWrapper(
            self.base_model.with_config(*args, **kwargs),
            self.ctx,
            endpoint=self.endpoint,
        )

    def with_structured_output(self, *args, **kwargs):
        if hasattr(self.base_model, "with_structured_output"):
            return RobustChatModelWrapper(
                self.base_model.with_structured_output(*args, **kwargs),
                self.ctx,
                endpoint=self.endpoint,
            )
        raise AttributeError(f"{type(self.base_model)} does not support with_structured_output")

    def __or__(self, other):
        return RobustChatModelWrapper(
            self.base_model.__or__(other), self.ctx, endpoint=self.endpoint
        )

    def __ror__(self, other):
        return RobustChatModelWrapper(
            self.base_model.__ror__(other), self.ctx, endpoint=self.endpoint
        )

    @robust_retry_async
    async def ainvoke(self, *args, **kwargs):
        trace_id = None
        if self.ctx and self.ctx.data_engine:
            trace_id = CURRENT_TRACE_ID.get()
        if trace_id:
            args, kwargs = _inject_parent_trace_id(trace_id, *args, **kwargs)
        return await self.base_model.ainvoke(*args, **kwargs)

    @robust_retry_astream
    async def astream(self, *args, **kwargs):
        trace_id = None
        if self.ctx and self.ctx.data_engine:
            trace_id = CURRENT_TRACE_ID.get()

        if trace_id:
            args, kwargs = _inject_parent_trace_id(trace_id, *args, **kwargs)

        stream_exec_id = uuid4()

        async for chunk in self.base_model.astream(*args, **kwargs):
            if trace_id and self.ctx and self.ctx.data_engine and chunk.content:
                text_to_stream = ""
                thinking_to_stream = ""
                if isinstance(chunk.content, str):
                    text_to_stream = chunk.content
                elif isinstance(chunk.content, list):
                    for item in chunk.content:
                        if isinstance(item, dict):
                            if item.get("type") == "text":
                                text_to_stream += item.get("text", "")
                            elif item.get("type") == "thinking":
                                thinking_to_stream += item.get("thinking", "")

                if text_to_stream:
                    self.ctx.data_engine.stream_output(
                        stream_exec_id, text_to_stream, is_thinking=False
                    )
                if thinking_to_stream:
                    self.ctx.data_engine.stream_output(
                        stream_exec_id, thinking_to_stream, is_thinking=True
                    )
            yield chunk


async def invoke_llm_with_timeout_message[T](
    llm_call: Coroutine[Any, Any, T],
    timeout_seconds: int = 10,
    hard_timeout: int = 180,
) -> T:
    """Send an LLM call and display a countdown / timeout message if delayed."""
    llm_task = asyncio.create_task(llm_call)
    waiter_task = asyncio.create_task(asyncio.sleep(timeout_seconds))

    done, _ = await asyncio.wait({llm_task, waiter_task}, return_when=asyncio.FIRST_COMPLETED)

    if llm_task in done:
        waiter_task.cancel()
        return llm_task.result()
    else:
        user_messages_logger.info("Waiting for LLM call response...")
        start_time = asyncio.get_event_loop().time()

        while True:
            try:
                return await asyncio.wait_for(asyncio.shield(llm_task), timeout=1.0)
            except TimeoutError:
                if llm_task.done():
                    return llm_task.result()

                pause_file = Path(settings.TRACES_PATH).parent / ".artemis_paused"
                if pause_file.exists():
                    start_time = asyncio.get_event_loop().time()
                    continue

                elapsed = asyncio.get_event_loop().time() - start_time
                if elapsed > max(0, hard_timeout - timeout_seconds):
                    user_messages_logger.error(f"LLM call timed out after {hard_timeout} seconds.")
                    raise TimeoutError(f"LLM call timed out after {hard_timeout} seconds.")


# Backward compatible factory functions delegating to ModelFactory
def get_google_llm(
    model_name: str = "gemini-3.6-flash",
    temperature: float | None = None,
    timeout: float | None = None,
    thinking_budget: int | None = None,
    thinking_level: str | None = "medium",
    include_thoughts: bool | None = None,
    enable_grounding: bool = False,
) -> BaseChatModel:
    ep = ModelEndpoint(
        provider=ModelProvider.GOOGLE,
        model_name=model_name,
        temperature=temperature or 0.0,
        timeout_seconds=timeout or 60.0,
        thinking_budget=thinking_budget,
        thinking_level=thinking_level,
        include_thoughts=include_thoughts,
        enable_grounding=enable_grounding,
    )
    return ModelFactory.create_model(ep)


def get_vertex_llm(
    model_name: str = "gemini-3.6-flash",
    temperature: float | None = None,
    timeout: float | None = None,
    thinking_budget: int | None = None,
) -> BaseChatModel:
    ep = ModelEndpoint(
        provider=ModelProvider.VERTEX_AI,
        model_name=model_name,
        temperature=temperature or 0.0,
        timeout_seconds=timeout or 60.0,
        thinking_budget=thinking_budget,
    )
    return ModelFactory.create_model(ep)


def get_openai_llm(
    model_name: str = "o3",
    temperature: float | None = None,
    timeout: float | None = None,
) -> BaseChatModel:
    ep = ModelEndpoint(
        provider=ModelProvider.OPENAI,
        model_name=model_name,
        temperature=temperature or 0.0,
        timeout_seconds=timeout or 60.0,
    )
    return ModelFactory.create_model(ep)


def get_openrouter_llm(
    model_name: str,
    temperature: float | None = None,
    timeout: float | None = None,
) -> BaseChatModel:
    ep = ModelEndpoint(
        provider=ModelProvider.OPENROUTER,
        model_name=model_name,
        temperature=temperature or 0.0,
        timeout_seconds=timeout or 60.0,
    )
    return ModelFactory.create_model(ep)


def get_grok_llm(
    model_name: str,
    temperature: float | None = None,
    timeout: float | None = None,
) -> BaseChatModel:
    ep = ModelEndpoint(
        provider=ModelProvider.XAI,
        model_name=model_name,
        temperature=temperature or 0.0,
        timeout_seconds=timeout or 60.0,
    )
    return ModelFactory.create_model(ep)


def get_anthropic_llm(
    model_name: str,
    temperature: float | None = None,
    timeout: float | None = None,
    thinking_budget: int | None = None,
    reasoning_effort: str | None = None,
) -> BaseChatModel:
    ep = ModelEndpoint(
        provider=ModelProvider.ANTHROPIC,
        model_name=model_name,
        temperature=temperature or 0.0,
        timeout_seconds=timeout or 60.0,
        thinking_budget=thinking_budget,
        reasoning_effort=reasoning_effort,
    )
    return ModelFactory.create_model(ep)


def get_cached_raw_model(
    provider: str,
    model_name: str,
    temperature: float | None = None,
    timeout: float | None = None,
    thinking_budget: int | None = None,
    thinking_level: str | None = None,
    include_thoughts: bool | None = None,
    reasoning_effort: str | None = None,
    enable_grounding: bool = False,
) -> BaseChatModel:
    """Retrieves or instantiates a cached raw LangChain chat model."""
    ep = ModelEndpoint(
        provider=ModelProvider.from_string(provider),
        model_name=model_name,
        temperature=temperature or 0.0,
        timeout_seconds=timeout or 60.0,
        thinking_budget=thinking_budget,
        thinking_level=thinking_level,
        include_thoughts=include_thoughts,
        reasoning_effort=reasoning_effort,
        enable_grounding=enable_grounding,
    )
    return ModelFactory.get_model(ep)


def _resolve_endpoint(
    ctx: ArtemisContext,
    name: str,
    is_utils: bool = False,
    use_fallback: bool = False,
) -> ModelEndpoint:
    """Cleanly resolves a ModelEndpoint from context router or llm_config."""
    # 1. Prefer dynamic ModelRouter if attached to context
    if getattr(ctx, "model_router", None) is not None:
        role_name = str(name).lower()
        if use_fallback:
            fallbacks = ctx.model_router.get_fallbacks(role_name)
            if fallbacks:
                return fallbacks[0]
        return ctx.model_router.get_endpoint(role_name)

    # 2. Fall back to context llm_config
    if getattr(ctx, "llm_config", None) is None:
        try:
            ctx.llm_config = get_default_llm_config()
        except Exception:
            pass

    cfg = ctx.llm_config.get_utils(name) if is_utils else ctx.llm_config.get_agent(name)

    if use_fallback:
        if isinstance(cfg, LLMWithFallback) or (hasattr(cfg, "fallback") and cfg.fallback):
            cfg = cfg.fallback
        else:
            raise ValueError(f"LLM configuration for '{name}' has no fallback!")

    def _get_val(obj, attr, expected_type):
        val = getattr(obj, attr, None)
        return val if isinstance(val, expected_type) else None

    provider_val = getattr(cfg, "provider", "google")
    model_val = getattr(cfg, "model", "gemini-2.5-flash")

    return ModelEndpoint(
        provider=ModelProvider.from_string(provider_val),
        model_name=str(model_val),
        temperature=_get_val(cfg, "temperature", (int, float)) or 0.0,
        timeout_seconds=_get_val(cfg, "timeout", (int, float)) or 60.0,
        thinking_budget=_get_val(cfg, "thinking_budget", int),
        thinking_level=_get_val(cfg, "thinking_level", str),
        reasoning_effort=_get_val(cfg, "reasoning_effort", str),
        include_thoughts=_get_val(cfg, "include_thoughts", bool),
        enable_grounding=_get_val(cfg, "enable_grounding", bool) or False,
    )


@overload
def get_llm(
    ctx: ArtemisContext,
    name: AgentNodeWithFallback,
    *,
    use_fallback: bool = False,
    temperature: float | None = None,
) -> BaseChatModel: ...


@overload
def get_llm(
    ctx: ArtemisContext,
    name: LLMUtilsNode,
    *,
    is_utils: Literal[True],
    temperature: float | None = None,
) -> BaseChatModel: ...


@overload
def get_llm(
    ctx: ArtemisContext,
    name: LLMUtilsNodeWithFallback,
    *,
    is_utils: Literal[True],
    use_fallback: bool = False,
    temperature: float | None = None,
) -> BaseChatModel: ...


def get_llm(
    ctx: ArtemisContext,
    name: AgentNode | LLMUtilsNode | AgentNodeWithFallback,
    is_utils: bool = False,
    use_fallback: bool = False,
    temperature: float | None = None,
) -> BaseChatModel:
    """Resolves and instantiates the appropriate LLM wrapper for the given agent role."""
    endpoint = _resolve_endpoint(ctx, str(name), is_utils=is_utils, use_fallback=use_fallback)
    if temperature is not None:
        endpoint = replace(endpoint, temperature=temperature)
    raw_model = ModelFactory.get_model(endpoint)

    handler = DataEngineCallbackHandler(ctx)
    bound_model = raw_model.with_config(callbacks=[handler])

    return RobustChatModelWrapper(bound_model, ctx, endpoint=endpoint)  # type: ignore


async def with_fallback[T](
    main_call: Callable[[], Awaitable[T]],
    fallback_call: Callable[[], Awaitable[T]],
    none_should_fallback: bool = True,
) -> T:
    try:
        result = await main_call()
        if result is None and none_should_fallback:
            llm_logger.warning("Main LLM inference returned None. Falling back...")
            return await fallback_call()
        return result
    except Exception as e:
        llm_logger.warning(f"❗ Main LLM inference failed: {e}. Falling back...")
        return await fallback_call()