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 contextvars import ContextVar, Token
from contextlib import aclosing
import functools
import logging
from pathlib import Path
import re
import sys
import time
from typing import Any, Literal, TypeVar, overload
from uuid import uuid4

from langchain_core.language_models.chat_models import BaseChatModel
from langchain_core.messages import AIMessage, BaseMessage, HumanMessage

from artemis.config import (
    PAUSE_FILE,
    AgentNode,
    AgentNodeWithFallback,
    LLMUtilsNode,
    LLMUtilsNodeWithFallback,
    LLMWithFallback,
    get_default_llm_config,
    settings,
)
from artemis.context import ArtemisContext
from artemis.data_engine.trace import CURRENT_TRACE_ID, DataEngineCallbackHandler
from artemis.llm.diagnostics import ACTIVE_ATTEMPT, RequestAttempt, error_text, input_images
from artemis.llm.cache_affinity import cache_affinity_options
from artemis.llm.reliability import (
    CircuitBreaker,
    FailureCategory,
    LLMCallError,
    LLMExhaustedError,
    LLMPermanentError,
    classify_failure,
    retry_policy_for,
)
from artemis.llm.messages import adapt_responses_messages
from artemis.llm.prefix_cache import PromptPrefixTracker
from artemis.llm.router import ModelEndpoint, ModelFactory, ModelProvider
from artemis.services.token_meter import record_llm_usage
from artemis.llm.usage import usage_from_message
from artemis.llm.stream_guard import ToolArgumentGuard, ToolArgumentOverflow
from artemis.llm.structured import (
    ParseFailure,
    StructuredOutputError,
    content_to_text,
    parse_structured,
)
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")

# Retained as an override seam for unit tests. At runtime the active engine must
# be resolved dynamically because DataEngine.start_session() updates the module
# global after this module has already been imported.
_CURRENT_DATA_ENGINE = None

# Provider retry logs do not carry an Artemis trace id.  Keep the active model
# request in async-local state so observable SDK retries and the terminal error
# are emitted as one lifecycle, even when several model calls run concurrently.
_ACTIVE_LLM_REQUEST: ContextVar[dict[str, Any] | None] = ContextVar(
    "active_llm_request", default=None
)


def _begin_llm_request(provider: str | None) -> Token:
    return _ACTIVE_LLM_REQUEST.set(
        {
            "request_id": str(uuid4()),
            "provider": provider,
            "started_at": time.time(),
            "retries": [],
        }
    )


# Whether the current async context has a configured fallback model waiting
# behind this call (set by with_fallback).  When a fallback exists, exhausting
# retries hands over to it immediately instead of pausing the whole task.
_FALLBACK_AVAILABLE: ContextVar[bool] = ContextVar("llm_fallback_available", default=False)
_FAILED_ENDPOINTS: ContextVar[dict[str, Exception] | None] = ContextVar(
    "llm_failed_endpoints", default=None
)

# Shared circuit breaker keyed by "provider:model".  Transient failures open
# it; while open, concurrent calls wait out the cooldown instead of hammering
# a provider that is already melting down.
_ENDPOINT_BREAKER = CircuitBreaker(threshold=3, cooldown_seconds=30.0)

# Endpoints observed to not support streaming.  Detected once, loudly, then
# remembered so subsequent calls go straight to non-streaming invocation.
_NON_STREAMING_ENDPOINTS: set[str] = set()


def _get_current_data_engine():
    if _CURRENT_DATA_ENGINE is not None:
        return _CURRENT_DATA_ENGINE
    engine_module = sys.modules.get("artemis.data_engine.engine")
    return getattr(engine_module, "_CURRENT_DATA_ENGINE", None) if engine_module else None


def _record_llm_retry(
    error: str,
    delay: float,
    *,
    attempt: int | None = None,
    max_retries: int | None = None,
    provider: str | None = None,
    source: str = "artemis",
) -> None:
    """Persist and publish one recoverable LLM retry for UI transparency."""
    request = _ACTIVE_LLM_REQUEST.get()
    timestamp = time.time()
    request_id = request.get("request_id") if request else None
    provider = provider or (request.get("provider") if request else None)

    engine = _get_current_data_engine()
    if not engine or not getattr(engine, "current_session_id", None):
        return

    error_text = str(error)
    if len(error_text) > 1000:
        error_text = error_text[:1000] + "... [Truncated by LLM Retry Telemetry]"

    payload = {
        "error": error_text,
        "delay": float(delay),
        "attempt": attempt,
        "max_retries": max_retries,
        "provider": provider,
        "source": source,
        "recoverable": True,
        "request_id": request_id,
        "scheduled_at": timestamp,
    }
    payload = {key: value for key, value in payload.items() if value is not None}
    step_id = getattr(engine, "current_step_id", None)
    parent_trace_id = CURRENT_TRACE_ID.get()
    try:
        trace_id = engine.record_trace(
            type="llm_call",
            name="llm_retry",
            payload=payload,
            step_id=step_id,
            parent_trace_id=parent_trace_id,
            status="retrying",
        )
    except Exception as trace_error:
        llm_logger.warning("Failed to persist LLM retry trace: %s", trace_error)
        trace_id = None

    if request is not None:
        request["retries"].append(
            {
                **payload,
                "trace_id": str(trace_id) if trace_id else None,
                "timestamp": timestamp,
            }
        )

    try:
        engine._publish(
            "llm_retrying",
            {
                **payload,
                "step_id": str(step_id) if step_id else None,
                "trace_id": str(trace_id) if trace_id else None,
                "timestamp": timestamp,
            },
        )
    except Exception as publish_error:
        # Retry visibility is best-effort and must never break the LLM retry itself.
        llm_logger.warning("Failed to publish LLM retry event: %s", publish_error)


def _record_llm_event(name: str, payload: dict[str, Any], *, status: str = "retrying") -> None:
    """Persist and publish one degradation/lifecycle event on the LLM IO path.

    Every degradation (fallback, stream downgrade, stream reset, giving up)
    must be observable: WARNING-level logging is handled by callers, this
    records the structured trace + live event, best-effort.
    """
    request = _ACTIVE_LLM_REQUEST.get()
    engine = _get_current_data_engine()
    if not engine or not getattr(engine, "current_session_id", None):
        return

    full_payload = dict(payload)
    if request is not None:
        full_payload.setdefault("request_id", request.get("request_id"))
        if request.get("provider"):
            full_payload.setdefault("provider", request.get("provider"))
    full_payload = {key: value for key, value in full_payload.items() if value is not None}

    step_id = getattr(engine, "current_step_id", None)
    trace_id = None
    try:
        trace_id = engine.record_trace(
            type="llm_call",
            name=name,
            payload=full_payload,
            step_id=step_id,
            parent_trace_id=CURRENT_TRACE_ID.get(),
            status=status,
        )
    except Exception as trace_error:
        llm_logger.warning("Failed to persist LLM event trace %s: %s", name, trace_error)
    try:
        engine._publish(
            name,
            {
                **full_payload,
                "step_id": str(step_id) if step_id else None,
                "trace_id": str(trace_id) if trace_id else None,
                "timestamp": time.time(),
            },
        )
    except Exception as publish_error:
        llm_logger.warning("Failed to publish LLM event %s: %s", name, publish_error)


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 = PAUSE_FILE
    try:
        pause_file.write_text(f"LLM Error: {err_msg}", encoding="utf-8")
    except (OSError, UnicodeError):
        pass

    current_engine = _get_current_data_engine()
    if current_engine:
        step_id = getattr(current_engine, "current_step_id", None)
        timestamp = time.time()
        request = _ACTIVE_LLM_REQUEST.get()
        failure_payload: dict[str, Any] = {"error": err_msg, "pause": True}
        if request:
            failure_payload.update(
                {
                    "request_id": request["request_id"],
                    "provider": request.get("provider"),
                    "waited_seconds": max(0.0, timestamp - request["started_at"]),
                    # Only SDK-observed retries are included. Other providers
                    # may retry internally, but their attempt data is opaque.
                    "retries": list(request["retries"]),
                }
            )
            failure_payload = {
                key: value for key, value in failure_payload.items() if value is not None
            }
        try:
            # Persist exhausted LLM retries as a failed call so the same error
            # card is available both in the live stream and after a refresh.
            current_engine.record_trace(
                type="llm_call",
                name="llm_pause",
                payload=failure_payload,
                step_id=step_id,
                status="failed",
            )
        except Exception as trace_error:
            llm_logger.warning("Failed to persist paused LLM error trace: %s", trace_error)

        current_engine._publish(
            "task_paused",
            {
                **failure_payload,
                "step_id": str(step_id) if step_id else None,
                "timestamp": timestamp,
            },
        )

    return pause_file


async def _wait_for_resume(pause_file: Path) -> bool:
    """Wait for the pause file to be cleared. Returns False on deadline."""
    deadline = float(getattr(settings, "LLM_PAUSE_TIMEOUT_SECONDS", 0.0) or 0.0)
    waited = 0.0
    while pause_file.exists():
        if deadline > 0 and waited >= deadline:
            return False
        await asyncio.sleep(1)
        waited += 1
    return True


async def _wait_for_breaker(key: str) -> None:
    """Wait out an open circuit instead of hammering a melting provider."""
    waited = 0.0
    warned = False
    while not _ENDPOINT_BREAKER.allow(key):
        if not warned:
            warned = True
            llm_logger.warning(f"LLM circuit open for {key}; waiting for cooldown...")
        step = min(max(_ENDPOINT_BREAKER.open_remaining(key), 0.5), 2.0)
        await asyncio.sleep(step)
        waited += step
        if waited >= 120.0:
            # Never deadlock behind a stuck half-open trial; proceed anyway.
            break


async def _run_with_recovery[T](
    call_once: Callable[[], Awaitable[T]],
    *,
    provider: str | None = None,
    endpoint_key: str | None = None,
    request_timeout: float | None = None,
    max_retries: int | None = None,
    diagnostics: dict | None = None,
    engine=None,
) -> T:
    """Execute one LLM call under the classified retry / pause / breaker policy.

    Failure handling is decided per category (see artemis.llm.reliability):

    - Retryable categories back off per policy; each retry is recorded for UI
      transparency via _record_llm_retry.
    - Non-retryable categories (auth, bad request) raise LLMPermanentError
      immediately: retrying cannot help and pausing would hang the task.
    - When retryable attempts are exhausted: if a fallback model is waiting
      (with_fallback), raise LLMExhaustedError so it takes over immediately;
      otherwise pause the task (bounded by settings.LLM_PAUSE_TIMEOUT_SECONDS)
      and retry from scratch on resume.
    """
    while True:
        request_token = _begin_llm_request(provider)
        last_error: Exception | None = None
        last_failure = None
        attempts: dict[FailureCategory, int] = {}
        pause_file: Path | None = None
        try:
            while True:
                failed_endpoints = _FAILED_ENDPOINTS.get()
                previous = (
                    failed_endpoints.get(endpoint_key)
                    if failed_endpoints is not None and endpoint_key
                    else None
                )
                if previous is not None:
                    last_error, last_failure = previous, classify_failure(previous)
                    if not last_failure.retryable:
                        raise LLMPermanentError(str(previous), failure=last_failure, cause=previous)
                    break
                if endpoint_key and not attempts:
                    await _wait_for_breaker(endpoint_key)
                timeout = request_timeout or settings.OPENAI_TIMEOUT_SECONDS
                request = _ACTIVE_LLM_REQUEST.get() or {}
                observation = RequestAttempt(
                    engine or _get_current_data_engine(),
                    request_id=request.get("request_id"),
                    attempt=sum(attempts.values()) + 1,
                    timeout=timeout,
                    max_retries=max_retries
                    if max_retries is not None
                    else settings.OPENAI_MAX_RETRIES,
                    metadata=diagnostics or {"provider": provider},
                )
                observation_token = ACTIVE_ATTEMPT.set(observation)
                try:
                    async with asyncio.timeout(request_timeout or settings.OPENAI_TIMEOUT_SECONDS):
                        result = await call_once()
                except (KeyboardInterrupt, asyncio.CancelledError) as e:
                    observation.finish("cancelled", error=e, category="cancelled")
                    raise
                except LLMCallError as e:
                    observation.finish("failed", error=e, category=e.failure.category.value)
                    # Already classified terminal by a nested recovery layer.
                    raise
                except Exception as e:
                    failure = classify_failure(e)
                    observation.finish("failed", error=e, category=failure.category.value)
                    failure_text = error_text(e, timeout)
                    if failure.category is FailureCategory.CANCELLED:
                        raise
                    if endpoint_key:
                        _ENDPOINT_BREAKER.record_failure(endpoint_key, failure)
                    if not failure.retryable:
                        if failed_endpoints is not None and endpoint_key:
                            failed_endpoints[endpoint_key] = e
                        llm_logger.warning(
                            f"LLM call failed permanently ({failure.category.value}): {e}"
                        )
                        _record_llm_event(
                            "llm_gave_up",
                            {
                                "error": failure_text,
                                "category": failure.category.value,
                                "retryable": False,
                            },
                            status="failed",
                        )
                        raise LLMPermanentError(str(e), failure=failure, cause=e) from e
                    last_error, last_failure = e, failure
                    attempt = attempts.get(failure.category, 0) + 1
                    attempts[failure.category] = attempt
                    policy = retry_policy_for(failure.category)
                    total_attempts = sum(attempts.values())
                    max_attempts = (
                        max_retries + 1 if max_retries is not None else policy.max_attempts
                    )
                    if failure.category is FailureCategory.INVALID_OUTPUT:
                        # One fresh response can recover a generation loop. Repeated
                        # runaway output must reach the existing fallback/pause path,
                        # even when ordinary network retries are configured higher.
                        max_attempts = min(
                            max_attempts, total_attempts + max(0, policy.max_attempts - attempt)
                        )
                    if total_attempts >= max_attempts:
                        break
                    delay = policy.delay_for(attempt)
                    _record_llm_retry(
                        failure_text,
                        delay,
                        attempt=total_attempts,
                        max_retries=max_attempts - 1,
                        provider=provider,
                    )
                    llm_logger.warning(
                        f"LLM call failed ({failure.category.value}), retrying in {delay:.2f}s..."
                    )
                    observation.retry(delay)
                    try:
                        await asyncio.sleep(delay)
                    except asyncio.CancelledError:
                        # Guidance may cancel a retry wait while the task continues.
                        # Close this attempt so it cannot look active next to the new request.
                        observation.finish("cancelled")
                        raise
                else:
                    observation.finish("completed", response=result)
                    if endpoint_key:
                        _ENDPOINT_BREAKER.record_success(endpoint_key)
                    return result
                finally:
                    ACTIVE_ATTEMPT.reset(observation_token)

            # Retryable attempts exhausted.
            if _FALLBACK_AVAILABLE.get():
                if failed_endpoints is not None and endpoint_key:
                    failed_endpoints[endpoint_key] = last_error
                _record_llm_event(
                    "llm_gave_up",
                    {
                        "error": str(last_error)[:1000],
                        "category": last_failure.category.value,
                        "retryable": True,
                        "handover": "fallback",
                    },
                    status="failed",
                )
                raise LLMExhaustedError(
                    str(last_error), failure=last_failure, cause=last_error
                ) from last_error
            pause_file = _handle_llm_pause_and_resume(last_error)
        finally:
            _ACTIVE_LLM_REQUEST.reset(request_token)

        if not await _wait_for_resume(pause_file):
            llm_logger.error(
                "LLM pause deadline"
                f" ({settings.LLM_PAUSE_TIMEOUT_SECONDS:.0f}s) exceeded; giving up."
            )
            try:
                pause_file.unlink(missing_ok=True)
            except OSError:
                pass
            raise LLMExhaustedError(
                str(last_error), failure=last_failure, cause=last_error
            ) from last_error

        llm_logger.info("Resume signal received, retrying LLM call...")
        if failed_endpoints is not None and endpoint_key:
            failed_endpoints.pop(endpoint_key, None)
        current_engine = _get_current_data_engine()
        if current_engine:
            current_engine._publish("task_resumed", {})


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,
        role: str | None = None,
        is_fallback: bool = False,
        structured_mode: str | None = None,
    ):
        self.base_model = base_model
        self.ctx = ctx
        self.endpoint = endpoint
        self.role = role
        self.is_fallback = is_fallback
        self.structured_mode = structured_mode

    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 []

        if self.endpoint is not None and self.endpoint.enable_grounding:
            if not any(
                isinstance(t, dict) and t.get("type") in ("web_search", "web_search_preview")
                for t in processed_tools
            ):
                processed_tools.append({"type": "web_search"})

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

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

    def with_structured_output(self, *args, **kwargs):
        if hasattr(self.base_model, "with_structured_output"):
            # Compatible Responses endpoints may support functions but not text.format=json_schema.
            # Keep schema validation in LangChain/Pydantic without changing the API transport.
            kwargs.setdefault("method", "function_calling")
            include_raw = kwargs.get("include_raw", False)
            kwargs["include_raw"] = True
            return RobustChatModelWrapper(
                self.base_model.with_structured_output(*args, **kwargs),
                self.ctx,
                endpoint=self.endpoint,
                role=self.role,
                is_fallback=self.is_fallback,
                structured_mode="raw" if include_raw else "parsed",
            )
        raise AttributeError(f"{type(self.base_model)} does not support with_structured_output")

    def __or__(self, other):
        if self.structured_mode:
            from langchain_core.runnables import RunnableLambda

            return RunnableLambda(self.ainvoke) | other
        return RobustChatModelWrapper(
            self.base_model.__or__(other),
            self.ctx,
            endpoint=self.endpoint,
            role=self.role,
            is_fallback=self.is_fallback,
        )

    def __ror__(self, other):
        if self.structured_mode:
            from langchain_core.runnables import RunnableLambda

            return other | RunnableLambda(self.ainvoke)
        return RobustChatModelWrapper(
            self.base_model.__ror__(other),
            self.ctx,
            endpoint=self.endpoint,
            role=self.role,
            is_fallback=self.is_fallback,
        )

    def _provider_value(self) -> str | None:
        provider = getattr(self.endpoint, "provider", None)
        if provider is None:
            return None
        return str(getattr(provider, "value", provider))

    def _endpoint_key(self) -> str:
        if self.endpoint is not None:
            return f"{self._provider_value()}:{self.endpoint.model_name}:{self.endpoint.api_base or settings.OPENAI_BASE_URL}:{self.endpoint.reasoning_effort}"
        return f"{type(self.base_model).__name__}:{id(self.base_model)}"

    def _traced_call(self, args: tuple, kwargs: dict) -> tuple[tuple, dict, Any]:
        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 args, kwargs, trace_id

    def _prepare_model_call(self, args: tuple, kwargs: dict) -> tuple[tuple, dict]:
        model_name = self.endpoint.model_name if self.endpoint else None
        kwargs = cache_affinity_options(self.ctx, self.role, model_name, kwargs)
        if (
            args
            and isinstance(args[0], (list, tuple))
            and all(isinstance(message, BaseMessage) for message in args[0])
        ):
            return (adapt_responses_messages(args[0], model_name), *args[1:]), kwargs
        if (
            "input" in kwargs
            and isinstance(kwargs["input"], (list, tuple))
            and all(isinstance(message, BaseMessage) for message in kwargs["input"])
        ):
            return args, {
                **kwargs,
                "input": adapt_responses_messages(kwargs["input"], model_name),
            }
        return args, kwargs

    def _meter_usage(self, response) -> None:
        """Record measured token usage for this call (M0 metering, best-effort)."""
        engine = self.ctx.data_engine if self.ctx else None
        if engine is None:
            engine = _get_current_data_engine()
        record_llm_usage(engine, response, source=self._endpoint_key())

    def _diagnostics(self, args, kwargs):
        prefix = None
        if self.ctx is not None:
            tracker = getattr(self.ctx, "prompt_prefix_tracker", None)
            if not isinstance(tracker, PromptPrefixTracker):
                tracker = PromptPrefixTracker()
                self.ctx.prompt_prefix_tracker = tracker
            bound_options = getattr(self.base_model, "kwargs", {})
            options = {
                **(bound_options if isinstance(bound_options, dict) else {}),
                **{
                    key: kwargs[key]
                    for key in ("tools", "tool_choice", "response_format")
                    if key in kwargs
                },
            }
            identity = f"{kwargs.get('prompt_cache_key')}:{self.role}:{self._endpoint_key()}"
            prefix = tracker.observe(identity, args[0] if args else kwargs.get("input"), options)
            ledger = getattr(self.ctx, "transcript_ledger", None)
            if ledger is not None:
                prefix["context_phase"] = getattr(ledger, "context_phase", 0)
                prefix["last_compaction"] = getattr(ledger, "last_compaction", None)
        return {
            "provider": self._provider_value(),
            "model": self.endpoint.model_name if self.endpoint else None,
            "reasoning_effort": self.endpoint.reasoning_effort if self.endpoint else None,
            "max_output_tokens": self.endpoint.max_tokens if self.endpoint else None,
            "role": self.role,
            "is_fallback": self.is_fallback,
            "images": input_images(args, kwargs),
            "prompt_cache_key": kwargs.get("prompt_cache_key"),
            "prefix_cache": prefix,
        }

    async def _invoke_attempt(self, *args, **kwargs):
        response = await self.base_model.ainvoke(*args, **kwargs)
        self._meter_usage(response)
        observation = ACTIVE_ATTEMPT.get()
        if observation is not None and (usage := usage_from_message(response)) is not None:
            observation.data["usage"] = usage
        if self.structured_mode == "parsed" and isinstance(response, dict) and "raw" in response:
            if response.get("parsing_error") is not None:
                raise response["parsing_error"]
            return response["parsed"]
        return response

    async def ainvoke(self, *args, **kwargs):
        args, kwargs = self._prepare_model_call(args, kwargs)
        args, kwargs, _ = self._traced_call(args, kwargs)
        response = await _run_with_recovery(
            functools.partial(self._invoke_attempt, *args, **kwargs),
            provider=self._provider_value(),
            endpoint_key=self._endpoint_key(),
            request_timeout=self.endpoint.timeout_seconds if self.endpoint else None,
            max_retries=self.endpoint.max_retries if self.endpoint else None,
            diagnostics=self._diagnostics(args, kwargs),
            engine=self.ctx.data_engine if self.ctx else None,
        )
        return response

    async def complete(self, *args, **kwargs):
        """Single completion entry point: returns the full final message.

        Streaming is a transport detail handled internally — provider chunks
        are forwarded to the live UI as deltas and accumulated here, so
        downstream consumers only ever see one complete message (or a typed
        LLMCallError).  A mid-stream failure discards the partial output,
        signals the UI to drop it, and retries the whole call per policy;
        partial or duplicated chunks can never reach message history.
        """
        if self.structured_mode:
            return await self.ainvoke(*args, **kwargs)
        args, kwargs = self._prepare_model_call(args, kwargs)
        args, kwargs, _ = self._traced_call(args, kwargs)
        emit_deltas = bool(self.ctx and self.ctx.data_engine)
        response = await _run_with_recovery(
            functools.partial(self._complete_attempt, emit_deltas, *args, **kwargs),
            provider=self._provider_value(),
            endpoint_key=self._endpoint_key(),
            request_timeout=self.endpoint.timeout_seconds if self.endpoint else None,
            max_retries=self.endpoint.max_retries if self.endpoint else None,
            diagnostics=self._diagnostics(args, kwargs),
            engine=self.ctx.data_engine if self.ctx else None,
        )
        self._meter_usage(response)
        return response

    def _emit_stream_delta(self, stream_exec_id, chunk) -> None:
        content = getattr(chunk, "content", None)
        if not content:
            return
        text_to_stream = ""
        thinking_to_stream = ""
        if isinstance(content, str):
            text_to_stream = content
        elif isinstance(content, list):
            for item in 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)

    async def _complete_attempt(self, emit_deltas: bool, *args, **kwargs):
        endpoint_key = self._endpoint_key()
        if endpoint_key in _NON_STREAMING_ENDPOINTS:
            return await self.base_model.ainvoke(*args, **kwargs)

        stream_exec_id = uuid4()
        full_response = None
        argument_guard = ToolArgumentGuard()
        try:
            async with aclosing(self.base_model.astream(*args, **kwargs)) as stream:
                async for chunk in stream:
                    observation = ACTIVE_ATTEMPT.get()
                    if observation is not None:
                        observation.chunk(chunk)
                    argument_guard.observe(chunk)
                    if full_response is None:
                        full_response = chunk
                    else:
                        full_response = full_response + chunk
                    if emit_deltas:
                        self._emit_stream_delta(stream_exec_id, chunk)
        except (KeyboardInterrupt, asyncio.CancelledError):
            raise
        except Exception as stream_error:
            if full_response is not None or isinstance(stream_error, ToolArgumentOverflow):
                # Mid-stream failure: the partial output must never reach
                # message history. Tell the UI to discard it, then let the
                # recovery policy retry the whole call with a fresh stream id.
                llm_logger.warning(
                    f"LLM stream broke mid-response ({stream_error}); discarding partial output."
                )
                _record_llm_event(
                    "llm_stream_reset",
                    {
                        "stream_exec_id": str(stream_exec_id),
                        "action": "discard",
                        "reason": "mid_stream_failure",
                        "error": str(stream_error)[:500],
                        "category": classify_failure(stream_error).category.value,
                        "message": "本次模型输出异常，未完成的输出已丢弃；正在按恢复策略处理。",
                    },
                )
                raise
            if _is_stream_unsupported_error(stream_error):
                _NON_STREAMING_ENDPOINTS.add(endpoint_key)
                llm_logger.warning(
                    f"Endpoint {endpoint_key} does not support streaming"
                    f" ({stream_error}); switching to non-streaming calls."
                )
                _record_llm_event(
                    "llm_stream_downgrade",
                    {"endpoint": endpoint_key, "error": str(stream_error)[:500]},
                )
                return await self.base_model.ainvoke(*args, **kwargs)
            raise
        if full_response is None:
            llm_logger.warning("LLM stream yielded no chunks; using non-streaming call.")
            return await self.base_model.ainvoke(*args, **kwargs)
        return full_response

    async def astream(self, *args, **kwargs):
        """Deprecated compatibility shim: yields exactly one final message.

        The old chunk-level astream retry could re-deliver already-yielded
        chunks after a mid-stream failure, corrupting accumulated message
        history. Streaming now happens inside complete(); live token deltas
        still reach the UI through the data engine.
        """
        yield await self.complete(*args, **kwargs)


def _is_stream_unsupported_error(error: BaseException) -> bool:
    """Detect 'this endpoint/model cannot stream' as opposed to a transient failure.

    A positive here permanently disables streaming for the endpoint in this
    process, so a bare AttributeError/TypeError from unrelated code must not
    qualify — the message has to actually implicate streaming.
    """
    if isinstance(error, NotImplementedError):
        return True
    message = str(error).lower()
    if isinstance(error, (AttributeError, TypeError)):
        return "stream" in message
    return "stream" in message and any(
        marker in message
        for marker in ("not support", "unsupported", "not implemented", "not available")
    )


async def acomplete(llm, *args, **kwargs):
    """Get one complete response from any chat model object.

    This is the single call-shape agents should use. For gateway-managed
    models (RobustChatModelWrapper) it delegates to complete(), which owns
    streaming, classified retries, and telemetry. For raw models and test
    doubles it preserves the legacy accumulate-stream-else-ainvoke shape.
    """
    if isinstance(llm, RobustChatModelWrapper):
        return await llm.complete(*args, **kwargs)

    try:
        full_response = None
        async with aclosing(llm.astream(*args, **kwargs)) as stream:
            async for chunk in stream:
                if full_response is None:
                    full_response = chunk
                else:
                    full_response = full_response + chunk
        if full_response is not None:
            return full_response
    except (KeyboardInterrupt, asyncio.CancelledError):
        raise
    except Exception as stream_error:
        llm_logger.warning(
            f"astream unavailable on {type(llm).__name__} ({stream_error}); using ainvoke."
        )
    return await llm.ainvoke(*args, **kwargs)


async def acomplete_structured(
    llm,
    messages: list,
    *,
    schema=None,
    correction_attempts: int = 1,
):
    """Complete and parse a JSON response, with one corrective re-ask.

    On a parse failure the model is shown its own output's parse error and
    asked to re-emit valid JSON (up to ``correction_attempts`` times) — giving
    the model a chance to see its mistake instead of feeding garbage
    downstream. If it still fails, raises StructuredOutputError (never
    returns raw text masquerading as parsed data). Each repair round emits an
    llm_parse_repair telemetry event.
    """
    response = await acomplete(llm, messages)
    text = content_to_text(getattr(response, "content", ""))
    parsed = parse_structured(text, schema=schema)
    attempt = 0
    while isinstance(parsed, ParseFailure) and attempt < correction_attempts:
        attempt += 1
        llm_logger.warning(
            f"Structured output parse failed ({parsed.error});"
            f" asking the model to correct itself (attempt {attempt})..."
        )
        _record_llm_event(
            "llm_parse_repair",
            {"error": parsed.error[:500], "attempt": attempt},
        )
        correction_messages = [
            *messages,
            response if isinstance(response, BaseMessage) else AIMessage(content=text),
            HumanMessage(
                content=(
                    "Your previous reply could not be parsed as the required"
                    f" JSON ({parsed.error}). Re-emit ONLY the corrected JSON"
                    " payload, with no surrounding prose or code fences."
                )
            ),
        ]
        response = await acomplete(llm, correction_messages)
        text = content_to_text(getattr(response, "content", ""))
        parsed = parse_structured(text, schema=schema)
    if isinstance(parsed, ParseFailure):
        _record_llm_event(
            "llm_gave_up",
            {
                "error": f"structured output unparseable: {parsed.error[:400]}",
                "category": "bad_request",
                "retryable": False,
            },
            status="failed",
        )
        raise StructuredOutputError(parsed)
    return parsed


async def invoke_llm_with_timeout_message[T](
    llm_call: Coroutine[Any, Any, T],
    timeout_seconds: int = 10,
    hard_timeout: int | None = None,
) -> 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))
    try:
        done, _ = await asyncio.wait({llm_task, waiter_task}, return_when=asyncio.FIRST_COMPLETED)

        if llm_task in done:
            return llm_task.result()

        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 = PAUSE_FILE
                if pause_file.exists():
                    start_time = asyncio.get_event_loop().time()
                    continue

                elapsed = asyncio.get_event_loop().time() - start_time
                if hard_timeout is not None and 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.")
    except BaseException:
        if not llm_task.done():
            llm_task.cancel()
        await asyncio.gather(llm_task, return_exceptions=True)
        raise
    finally:
        if not waiter_task.done():
            waiter_task.cancel()
        await asyncio.gather(waiter_task, return_exceptions=True)


# Backward compatible factory functions delegating to ModelFactory


def get_openai_llm(
    model_name: str | None = None,
    temperature: float | None = None,
    timeout: float | None = None,
) -> BaseChatModel:
    ep = ModelEndpoint(
        provider=ModelProvider.OPENAI,
        model_name=model_name or settings.OPENAI_MODEL,
        temperature=temperature,
        timeout_seconds=timeout if timeout is not None else settings.OPENAI_TIMEOUT_SECONDS,
    )
    return ModelFactory.create_model(ep)


def get_cached_raw_model(
    provider: str,
    model_name: str,
    temperature: float | None = None,
    timeout: float | 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 or settings.OPENAI_MODEL,
        temperature=temperature,
        timeout_seconds=timeout if timeout is not None else settings.OPENAI_TIMEOUT_SECONDS,
        reasoning_effort=reasoning_effort
        if reasoning_effort is not None
        else settings.OPENAI_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 llm_config."""
    if getattr(ctx, "llm_config", None) is None:
        try:
            ctx.llm_config = get_default_llm_config()
        except Exception as exc:  # pylint: disable=broad-exception-caught
            # Config loading has open-ended failure modes; surface the real cause
            # instead of letting the attribute access below fail with a misleading
            # AttributeError on the still-missing config.
            raise RuntimeError(
                f"Default LLM config could not be loaded while resolving '{name}': {exc}"
            ) from exc

    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", "openai")
    model_val = getattr(cfg, "model", settings.OPENAI_MODEL)

    return ModelEndpoint(
        provider=ModelProvider.from_string(provider_val),
        model_name=str(model_val),
        temperature=_get_val(cfg, "temperature", (int, float)),
        timeout_seconds=_get_val(cfg, "timeout", (int, float)) or settings.OPENAI_TIMEOUT_SECONDS,
        max_retries=_get_val(cfg, "max_retries", int)
        if _get_val(cfg, "max_retries", int) is not None
        else settings.OPENAI_MAX_RETRIES,
        max_tokens=_get_val(cfg, "max_tokens", int),
        api_base=_get_val(cfg, "api_base", str),
        reasoning_effort=_get_val(cfg, "reasoning_effort", str),
        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 = endpoint.model_copy(update={"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, role=str(name), is_fallback=use_fallback
    )  # type: ignore


async def with_fallback[T](
    main_call: Callable[[], Awaitable[T]],
    fallback_call: Callable[[], Awaitable[T]],
    none_should_fallback: bool = True,
) -> T:
    """Share failed endpoint identities so an identical fallback cannot repeat retries."""
    token = _FAILED_ENDPOINTS.set({})
    try:
        return await _with_fallback(main_call, fallback_call, none_should_fallback)
    finally:
        _FAILED_ENDPOINTS.reset(token)


async def _with_fallback[T](
    main_call: Callable[[], Awaitable[T]],
    fallback_call: Callable[[], Awaitable[T]],
    none_should_fallback: bool = True,
) -> T:
    """Run main_call, switching to fallback_call only when it can actually help.

    Falling back is an explicit, observed decision: the failure is classified
    and only categories where a different endpoint might succeed trigger the
    fallback (a bad request would just hide the bug inside a weaker model's
    output). Every switch is logged at WARNING and recorded as an
    llm_fallback telemetry event. While main_call runs, the recovery layer
    knows a fallback exists and hands over immediately on retry exhaustion
    instead of pausing the task.
    """

    def _switch(reason: str, category: str | None, error: str | None) -> None:
        llm_logger.warning(
            f"❗ Main LLM inference failed ({reason}"
            f"{f': {error}' if error else ''}). Falling back..."
        )
        _record_llm_event(
            "llm_fallback",
            {
                "reason": reason,
                "category": category,
                "error": error[:500] if error else None,
            },
        )

    # The contextvar must be reset BEFORE fallback_call runs: the fallback has
    # no further fallback behind it, so its own exhaustion should pause.
    fallback_token = _FALLBACK_AVAILABLE.set(True)
    try:
        result = await main_call()
    except (KeyboardInterrupt, asyncio.CancelledError):
        _FALLBACK_AVAILABLE.reset(fallback_token)
        raise
    except LLMCallError as e:
        _FALLBACK_AVAILABLE.reset(fallback_token)
        if not e.failure.should_fallback:
            raise
        _switch("terminal_error", e.failure.category.value, str(e))
        return await fallback_call()
    except Exception as e:
        _FALLBACK_AVAILABLE.reset(fallback_token)
        failure = classify_failure(e)
        if failure.category is FailureCategory.CANCELLED or not failure.should_fallback:
            raise
        _switch("error", failure.category.value, str(e))
        return await fallback_call()
    else:
        _FALLBACK_AVAILABLE.reset(fallback_token)

    if result is None and none_should_fallback:
        _switch("empty_result", None, None)
        return await fallback_call()
    return result