fix: improve cross-device evidence, request diagnostics and prompt efficiency

artemis/services/llm.py

@@ -23,13 +23,14 @@ from collections.abc import Awaitable, Callable, Coroutine
 from contextvars import ContextVar, Token
 from contextlib import aclosing
 import functools
+import hashlib
 import logging
 from pathlib import Path
 import re
 import sys
 import time
 from typing import Any, Literal, TypeVar, overload
-from uuid import uuid4
+from uuid import UUID, uuid4
 
 from langchain_core.language_models.chat_models import BaseChatModel
 from langchain_core.messages import AIMessage, BaseMessage, HumanMessage
@@ -46,6 +47,7 @@ from artemis.config import (
 )
 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.reliability import (
     CircuitBreaker,
     FailureCategory,
@@ -58,6 +60,7 @@ from artemis.llm.reliability import (
 from artemis.llm.messages import adapt_responses_messages
 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.structured import (
     ParseFailure,
     StructuredOutputError,
@@ -335,6 +338,8 @@ async def _run_with_recovery[T](
     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.
 
@@ -370,16 +375,33 @@ async def _run_with_recovery[T](
                     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):
+                except (KeyboardInterrupt, asyncio.CancelledError) as e:
+                    observation.finish("cancelled", error=e, category="cancelled")
                     raise
-                except LLMCallError:
+                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:
@@ -393,7 +415,7 @@ async def _run_with_recovery[T](
                         _record_llm_event(
                             "llm_gave_up",
                             {
-                                "error": str(e)[:1000],
+                                "error": failure_text,
                                 "category": failure.category.value,
                                 "retryable": False,
                             },
@@ -412,7 +434,7 @@ async def _run_with_recovery[T](
                         break
                     delay = policy.delay_for(attempt)
                     _record_llm_retry(
-                        str(e),
+                        failure_text,
                         delay,
                         attempt=total_attempts,
                         max_retries=max_attempts - 1,
@@ -421,11 +443,21 @@ async def _run_with_recovery[T](
                     llm_logger.warning(
                         f"LLM call failed ({failure.category.value}), retrying in {delay:.2f}s..."
                     )
-                    await asyncio.sleep(delay)
+                    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():
@@ -493,10 +525,16 @@ class RobustChatModelWrapper:
         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)
@@ -520,13 +558,18 @@ class RobustChatModelWrapper:
                 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)
+        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):
@@ -534,21 +577,42 @@ class RobustChatModelWrapper:
             # 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
+            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
+            self.base_model.__ror__(other),
+            self.ctx,
+            endpoint=self.endpoint,
+            role=self.role,
+            is_fallback=self.is_fallback,
         )
 
     def _provider_value(self) -> str | None:
@@ -572,6 +636,16 @@ class RobustChatModelWrapper:
 
     def _prepare_model_call(self, args: tuple, kwargs: dict) -> tuple[tuple, dict]:
         model_name = self.endpoint.model_name if self.endpoint else None
+        engine = getattr(self.ctx, "data_engine", None)
+        session = getattr(self.ctx, "exploration_id", None) or getattr(
+            engine, "current_session_id", None
+        )
+        if isinstance(session, (str, UUID)) and str(session):
+            identity = f"{session}:{self.role or 'model'}:{model_name}"
+            kwargs = {
+                "prompt_cache_key": "artemis:" + hashlib.sha256(identity.encode()).hexdigest()[:40],
+                **kwargs,
+            }
         if (
             args
             and isinstance(args[0], (list, tuple))
@@ -596,17 +670,42 @@ class RobustChatModelWrapper:
             engine = _get_current_data_engine()
         record_llm_usage(engine, response, source=self._endpoint_key())
 
+    def _diagnostics(self, args, kwargs):
+        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"),
+        }
+
+    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.base_model.ainvoke, *args, **kwargs),
+            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,
         )
-        self._meter_usage(response)
         return response
 
     async def complete(self, *args, **kwargs):
@@ -619,6 +718,8 @@ class RobustChatModelWrapper:
         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)
@@ -628,6 +729,8 @@ class RobustChatModelWrapper:
             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
@@ -662,6 +765,9 @@ class RobustChatModelWrapper:
         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)
                     if full_response is None:
                         full_response = chunk
                     else:
@@ -1004,7 +1110,9 @@ def get_llm(
     handler = DataEngineCallbackHandler(ctx)
     bound_model = raw_model.with_config(callbacks=[handler])
 
-    return RobustChatModelWrapper(bound_model, ctx, endpoint=endpoint)  # type: ignore
+    return RobustChatModelWrapper(
+        bound_model, ctx, endpoint=endpoint, role=str(name), is_fallback=use_fallback
+    )  # type: ignore
 
 
 async def with_fallback[T](