mantisdk 0.1.0__py3-none-any.whl

mantisdk/llm_proxy.py

@@ -0,0 +1,1890 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import ast
+import asyncio
+import json
+import logging
+import os
+import re
+import tempfile
+import threading
+import time
+from contextlib import asynccontextmanager
+from contextvars import ContextVar
+from datetime import datetime
+from typing import (
+    Any,
+    AsyncGenerator,
+    Awaitable,
+    Callable,
+    Dict,
+    Iterable,
+    List,
+    Literal,
+    Optional,
+    Sequence,
+    Tuple,
+    Type,
+    TypedDict,
+    Union,
+    cast,
+)
+
+import litellm
+import opentelemetry.trace as trace_api
+import yaml
+from fastapi import Request, Response
+from fastapi.responses import StreamingResponse
+from litellm.integrations.custom_logger import CustomLogger
+from litellm.integrations.opentelemetry import OpenTelemetry, OpenTelemetryConfig
+from litellm.proxy.proxy_server import app, save_worker_config  # pyright: ignore[reportUnknownVariableType]
+from litellm.types.utils import CallTypes
+from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
+from opentelemetry.sdk.resources import Resource
+from opentelemetry.sdk.trace import ReadableSpan, SpanContext
+from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult
+from opentelemetry.trace import Link, Status
+from opentelemetry.util.types import Attributes
+from starlette.middleware.base import BaseHTTPMiddleware
+from starlette.types import Scope
+
+from mantisdk.semconv import LightningResourceAttributes
+from mantisdk.types import LLM, ProxyLLM
+from mantisdk.utils.server_launcher import (
+    LaunchMode,
+    PythonServerLauncher,
+    PythonServerLauncherArgs,
+    noop_context,
+)
+
+from .store.base import LightningStore
+
+logger = logging.getLogger(__name__)
+
+# Context variable to store HTTP request headers for LiteLLM callback access
+_request_headers_context: ContextVar[Optional[Dict[str, str]]] = ContextVar("request_headers", default=None)
+
+__all__ = [
+    "LLMProxy",
+]
+
+
+class ModelConfig(TypedDict):
+    """LiteLLM model registration entry.
+
+    This mirrors the items in LiteLLM's `model_list` section.
+
+    Attributes:
+        model_name: Logical model name exposed by the proxy.
+        litellm_params: Parameters passed to LiteLLM for this model
+            (e.g., backend model id, api_base, additional options).
+    """  # Google style kept concise.
+
+    model_name: str
+    litellm_params: Dict[str, Any]
+
+
+def _get_pre_call_data(args: Any, kwargs: Any) -> Dict[str, Any]:
+    """Extract LiteLLM request payload from hook args.
+
+    The LiteLLM logger hooks receive `(*args, **kwargs)` whose third positional
+    argument or `data=` kwarg contains the request payload.
+
+    Args:
+        args: Positional arguments from the hook.
+        kwargs: Keyword arguments from the hook.
+
+    Returns:
+        The request payload dict.
+
+    Raises:
+        ValueError: If the payload cannot be located or is not a dict.
+    """
+    if kwargs.get("data"):
+        data = kwargs["data"]
+    elif len(args) >= 3:
+        data = args[2]
+    else:
+        raise ValueError(f"Unable to get request data from args or kwargs: {args}, {kwargs}")
+    if not isinstance(data, dict):
+        raise ValueError(f"Request data is not a dictionary: {data}")
+    return cast(Dict[str, Any], data)
+
+
+def _reset_litellm_logging_worker() -> None:
+    """Reset LiteLLM's global logging worker to the current event loop.
+
+    LiteLLM keeps a module-level ``GLOBAL_LOGGING_WORKER`` singleton that owns an
+    ``asyncio.Queue``. The queue is bound to the event loop where it was created.
+    When the proxy is restarted, Uvicorn spins up a brand new event loop in a new
+    thread. If the existing logging worker (and its queue) are reused, LiteLLM
+    raises ``RuntimeError: <Queue ...> is bound to a different event loop`` the
+    next time it tries to log. Recreating the worker ensures that LiteLLM will
+    lazily initialise a fresh queue on the new loop.
+    """
+
+    # ``GLOBAL_LOGGING_WORKER`` is imported in a few LiteLLM modules at runtime.
+    # Update any already-imported references so future calls use the fresh worker.
+    try:
+        import litellm.utils as litellm_utils
+        from litellm.litellm_core_utils import logging_worker as litellm_logging_worker
+
+        litellm_logging_worker.GLOBAL_LOGGING_WORKER = litellm_logging_worker.LoggingWorker()
+        litellm_utils.GLOBAL_LOGGING_WORKER = litellm_logging_worker.GLOBAL_LOGGING_WORKER  # type: ignore[reportAttributeAccessIssue]
+    except Exception:  # pragma: no cover - best-effort hygiene
+        logger.warning("Unable to propagate LiteLLM logging worker reset.", exc_info=True)
+
+
+def _reset_litellm_logging_callback_manager() -> None:
+    """Reset LiteLLM's global callback manager.
+
+    To get rid of the warning message: "Cannot add callback - would exceed MAX_CALLBACKS limit of 30."
+    when litellm is restarted multiple times in the same process.
+
+    It does not respect existing input/output callbacks.
+    """
+
+    try:
+        litellm.logging_callback_manager._reset_all_callbacks()  # pyright: ignore[reportPrivateUsage]
+    except Exception:  # pragma: no cover - best-effort hygiene
+        logger.warning("Unable to reset LiteLLM logging callback manager.", exc_info=True)
+
+
+class AddReturnTokenIds(CustomLogger):
+    """LiteLLM logger hook to request token ids from vLLM.
+
+    This mutates the outgoing request payload to include `return_token_ids=True`
+    for backends that support token id return (e.g., vLLM).
+
+    See also:
+        [vLLM PR #22587](https://github.com/vllm-project/vllm/pull/22587)
+    """
+
+    async def async_pre_call_hook(self, *args: Any, **kwargs: Any) -> Optional[Union[Exception, str, Dict[str, Any]]]:
+        """Async pre-call hook to adjust request payload.
+
+        Args:
+            args: Positional args from LiteLLM.
+            kwargs: Keyword args from LiteLLM.
+
+        Returns:
+            Either an updated payload dict or an Exception to short-circuit.
+        """
+        try:
+            data = _get_pre_call_data(args, kwargs)
+        except Exception as e:
+            return e
+
+        # Ensure token ids are requested from the backend when supported.
+        return {**data, "return_token_ids": True}
+
+
+class AddLogprobs(CustomLogger):
+    """LiteLLM logger hook to request logprobs from vLLM.
+
+    This mutates the outgoing request payload to include `logprobs=1`
+    for backends that support logprobs return (e.g., vLLM).
+    """
+
+    async def async_pre_call_hook(self, *args: Any, **kwargs: Any) -> Optional[Union[Exception, str, Dict[str, Any]]]:
+        """Async pre-call hook to adjust request payload."""
+        try:
+            data = _get_pre_call_data(args, kwargs)
+        except Exception as e:
+            return e
+
+        # Ensure logprobs are requested from the backend when supported.
+        return {**data, "logprobs": 1}
+
+
+class SpanWithExtraAttributes(ReadableSpan):
+    """Wrapper around ReadableSpan that adds extra span attributes.
+    
+    Since ReadableSpan is immutable, this wrapper intercepts the attributes
+    property to include additional attributes for Langfuse integration
+    (environment, tags, etc.).
+    """
+    
+    def __init__(self, wrapped_span: ReadableSpan, extra_attributes: Dict[str, Any]):
+        """Initialize wrapper with a span and extra attributes to inject.
+        
+        Args:
+            wrapped_span: The original ReadableSpan to wrap.
+            extra_attributes: Dictionary of extra attributes to add.
+                These should use Langfuse conventions (e.g., "langfuse.environment").
+        """
+        self._wrapped = wrapped_span
+        self._extra_attributes = extra_attributes
+    
+    @property
+    def name(self) -> str:
+        return self._wrapped.name
+    
+    @property
+    def context(self) -> Optional[SpanContext]:
+        return self._wrapped.context
+    
+    def get_span_context(self) -> SpanContext:
+        return self._wrapped.get_span_context()
+    
+    @property
+    def parent(self) -> Optional[SpanContext]:
+        return self._wrapped.parent
+    
+    @property
+    def start_time(self) -> int:
+        return self._wrapped.start_time
+    
+    @property
+    def end_time(self) -> int:
+        return self._wrapped.end_time
+    
+    @property
+    def status(self) -> Status:
+        return self._wrapped.status
+    
+    @property
+    def attributes(self) -> Attributes:
+        """Return original attributes merged with extra attributes."""
+        original_attrs = self._wrapped.attributes or {}
+        # Create a merged dict with original attrs and extra attrs
+        merged = dict(original_attrs)
+        merged.update(self._extra_attributes)
+        return merged
+    
+    @property
+    def events(self) -> tuple:
+        return self._wrapped.events
+    
+    @property
+    def links(self) -> tuple:
+        return self._wrapped.links
+    
+    @property
+    def resource(self) -> Resource:
+        return self._wrapped.resource
+    
+    @property
+    def instrumentation_scope(self):
+        return self._wrapped.instrumentation_scope
+    
+    # Also expose _resource for direct modification if needed
+    @property
+    def _resource(self) -> Resource:
+        return self._wrapped._resource  # pyright: ignore[reportPrivateUsage]
+    
+    @_resource.setter
+    def _resource(self, value: Resource):
+        self._wrapped._resource = value  # pyright: ignore[reportPrivateUsage]
+    
+    def __getattr__(self, name: str):
+        """Delegate all other attribute access to the wrapped span.
+        
+        This ensures OTLPSpanExporter can access all private attributes (_events, _links, etc.)
+        that it needs for serialization without us having to enumerate them all.
+        """
+        return getattr(self._wrapped, name)
+
+
+class LightningSpanExporter(SpanExporter):
+    """Buffered OTEL span exporter with subtree flushing and training-store sink.
+
+    Design:
+
+    * Spans are buffered until a root span's entire subtree is available.
+    * A private event loop on a daemon thread runs async flush logic.
+    * Rollout/attempt/sequence metadata is reconstructed by merging headers
+      from any span within a subtree.
+
+    Thread-safety:
+
+    * Buffer access is protected by a re-entrant lock.
+    * Export is synchronous to the caller yet schedules an async flush on the
+      internal loop, then waits for completion.
+    """
+
+    def __init__(
+        self,
+        _store: Optional[LightningStore] = None,
+        otlp_endpoint: Optional[str] = None,
+        otlp_headers: Optional[Dict[str, str]] = None,
+    ):
+        self._store: Optional[LightningStore] = _store  # this is only for testing purposes
+        self._otlp_endpoint: Optional[str] = otlp_endpoint  # Direct OTLP export endpoint
+        self._buffer: List[ReadableSpan] = []
+        self._lock: Optional[threading.Lock] = None
+        self._loop_lock_pid: Optional[int] = None
+
+        # Single dedicated event loop running in a daemon thread.
+        # This decouples OTEL SDK threads from our async store I/O.
+        # Deferred creation until first use.
+        self._loop: Optional[asyncio.AbstractEventLoop] = None
+        self._loop_thread: Optional[threading.Thread] = None
+
+        # Initialize OTLP exporter with custom endpoint and headers if provided
+        if otlp_endpoint:
+            self._otlp_exporter = OTLPSpanExporter(endpoint=otlp_endpoint, headers=otlp_headers or {})
+        else:
+            self._otlp_exporter = OTLPSpanExporter()
+
+    def _ensure_loop(self) -> asyncio.AbstractEventLoop:
+        """Lazily initialize the event loop and thread on first use.
+
+        Returns:
+            asyncio.AbstractEventLoop: The initialized event loop.
+        """
+        self._clear_loop_and_lock()
+        if self._loop is None:
+            self._loop = asyncio.new_event_loop()
+            self._loop_thread = threading.Thread(target=self._run_loop, name="LightningSpanExporterLoop", daemon=True)
+            self._loop_thread.start()
+        return self._loop
+
+    def _ensure_lock(self) -> threading.Lock:
+        """Lazily initialize the lock on first use.
+
+        Returns:
+            threading.Lock: The initialized lock.
+        """
+        self._clear_loop_and_lock()
+        if self._lock is None:
+            self._lock = threading.Lock()
+        return self._lock
+
+    def _clear_loop_and_lock(self) -> None:
+        """Clear the loop and lock.
+        This happens if the exporter was used in a process then used in another process.
+
+        This should only happen in CI.
+        """
+        if os.getpid() != self._loop_lock_pid:
+            logger.warning("Loop and lock are not owned by the current process. Clearing them.")
+            self._loop = None
+            self._loop_thread = None
+            self._lock = None
+            self._loop_lock_pid = os.getpid()
+        elif self._loop_lock_pid is None:
+            self._loop_lock_pid = os.getpid()
+
+    def _run_loop(self) -> None:
+        """Run the private asyncio loop forever on the exporter thread."""
+        assert self._loop is not None, "Loop should be initialized before thread starts"
+        asyncio.set_event_loop(self._loop)
+        self._loop.run_forever()
+
+    def shutdown(self) -> None:
+        """Shut down the exporter event loop.
+
+        Safe to call at process exit.
+
+        """
+        if self._loop is None:
+            return
+
+        try:
+
+            def _stop():
+                assert self._loop is not None
+                self._loop.stop()
+
+            self._loop.call_soon_threadsafe(_stop)
+            if self._loop_thread is not None:
+                self._loop_thread.join(timeout=2.0)
+            self._loop.close()
+        except Exception:
+            logger.exception("Error during exporter shutdown")
+
+    def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
+        """Export spans via buffered subtree flush.
+
+        Appends spans to the internal buffer, then triggers an async flush on the
+        private event loop. Blocks until that flush completes.
+
+        Args:
+            spans: Sequence of spans to export.
+
+        Returns:
+            SpanExportResult: SUCCESS on flush success, else FAILURE.
+        """
+        # Buffer append under lock to protect against concurrent exporters.
+        with self._ensure_lock():
+            for span in spans:
+                self._buffer.append(span)
+            default_endpoint = self._otlp_exporter._endpoint  # pyright: ignore[reportPrivateUsage]
+            try:
+                self._maybe_flush()
+            except Exception as e:
+                logger.exception("Export flush failed: %s", e)
+                return SpanExportResult.FAILURE
+            finally:
+                self._otlp_exporter._endpoint = default_endpoint  # pyright: ignore[reportPrivateUsage]
+
+        return SpanExportResult.SUCCESS
+
+    def _get_job_id_from_store(self, store: Any) -> Optional[str]:
+        """Get the job_id from the store's listeners (if any InsightTracker is attached)."""
+        if hasattr(store, "listeners"):
+            for listener in store.listeners:
+                if hasattr(listener, "job_id"):
+                    return listener.job_id
+        return None
+
+    def _get_tracing_metadata_from_rollout(
+        self, store: Any, rollout_id: str
+    ) -> Tuple[Optional[str], Optional[List[str]], Optional[str]]:
+        """Fetch tracing metadata (environment, tags, session_id) from a rollout.
+        
+        Args:
+            store: The LightningStore instance.
+            rollout_id: The rollout ID to fetch.
+            
+        Returns:
+            Tuple of (environment, tags, session_id). All may be None if not set.
+        """
+        print(f"[TracingMetadata] Fetching metadata for rollout {rollout_id}")
+        logger.info(f"[TracingMetadata] Fetching metadata for rollout {rollout_id}")
+        try:
+            loop = self._ensure_loop()
+            get_rollout_task = store.get_rollout_by_id(rollout_id)
+            fut = asyncio.run_coroutine_threadsafe(get_rollout_task, loop)
+            rollout = fut.result(timeout=5.0)  # Short timeout for metadata fetch
+            
+            if rollout is None:
+                logger.warning(f"[TracingMetadata] Rollout {rollout_id} not found in store")
+                return None, None, None
+            
+            logger.info(f"[TracingMetadata] Rollout {rollout_id} found, metadata={rollout.metadata}")
+            
+            if rollout.metadata is None:
+                logger.warning(f"[TracingMetadata] Rollout {rollout_id} has no metadata (None)")
+                return None, None, None
+            
+            if not rollout.metadata:
+                logger.warning(f"[TracingMetadata] Rollout {rollout_id} has empty metadata dict")
+                return None, None, None
+            
+            environment = rollout.metadata.get("environment")
+            tags = rollout.metadata.get("tags")
+            session_id = rollout.metadata.get("session_id")
+            
+            logger.info(f"[TracingMetadata] Rollout {rollout_id}: environment={environment}, tags={tags}, session_id={session_id}")
+            
+            return environment, tags, session_id
+        except Exception as e:
+            logger.warning(f"[TracingMetadata] Failed to fetch rollout metadata for {rollout_id}: {e}")
+            import traceback
+            traceback.print_exc()
+            return None, None, None
+
+    def _maybe_flush(self):
+        """Flush ready subtrees from the buffer.
+
+        Strategy:
+            We consider a subtree "ready" if we can identify a root span. We
+            then take that root and all its descendants out of the buffer and
+            try to reconstruct rollout/attempt/sequence headers by merging any
+            span's `metadata.requester_custom_headers` within the subtree.
+
+        Span types:
+            - Rollout spans: Have rollout_id/attempt_id/sequence_id headers
+            - Job spans: No rollout context, tagged with job_id for experiment tracking
+
+        Direct OTLP mode:
+            When `otlp_endpoint` is configured, spans are exported directly to the
+            endpoint without requiring a store or header validation.
+
+        Raises:
+            None directly. Logs and skips malformed spans.
+
+        """
+        # Iterate over current roots. Each iteration pops a whole subtree.
+        for root_span_id in self._get_root_span_ids():
+            subtree_spans = self._pop_subtrees(root_span_id)
+            if not subtree_spans:
+                continue
+
+            # Merge all custom headers found in the subtree.
+            # This must happen BEFORE the direct OTLP check so both paths can apply tags.
+            headers_merged: Dict[str, Any] = {}
+
+            for span in subtree_spans:
+                if span.attributes is None:
+                    continue
+                headers_str = span.attributes.get("metadata.requester_custom_headers")
+                if headers_str is None:
+                    continue
+                if not isinstance(headers_str, str):
+                    logger.debug(f"metadata.requester_custom_headers is not a string: {headers_str}")
+                    continue
+                if not headers_str.strip():
+                    continue
+                try:
+                    # Use literal_eval to parse the stringified dict safely.
+                    headers = ast.literal_eval(headers_str)
+                except Exception as e:
+                    logger.debug(f"Failed to parse metadata.requester_custom_headers: {e}")
+                    continue
+                if isinstance(headers, dict):
+                    headers_merged.update(cast(Dict[str, Any], headers))
+
+            # Extract rollout context if available
+            rollout_id = headers_merged.get("x-rollout-id")
+            attempt_id = headers_merged.get("x-attempt-id")
+            sequence_id = headers_merged.get("x-sequence-id")
+            
+            # Determine if we're using OTLP export (either direct or store-based)
+            store = self._store or get_active_llm_proxy().get_store()
+            otlp_enabled = bool(self._otlp_endpoint) or (store and store.capabilities.get("otlp_traces", False))
+
+            has_rollout_context = (
+                rollout_id
+                and attempt_id
+                and sequence_id
+                and isinstance(rollout_id, str)
+                and isinstance(attempt_id, str)
+                and isinstance(sequence_id, str)
+                and sequence_id.isdigit()
+            )
+
+            if has_rollout_context:
+                # Rollout-scoped spans: tag with rollout/attempt/sequence
+                sequence_id_decimal = int(sequence_id)
+                print(f"[TracingMetadata] Processing rollout {rollout_id} with {len(subtree_spans)} spans, otlp_enabled={otlp_enabled}")
+                logger.info(f"[TracingMetadata] Processing rollout {rollout_id} with {len(subtree_spans)} spans, otlp_enabled={otlp_enabled}")
+                
+                # Fetch tracing metadata (environment, tags, session_id) from the rollout
+                environment, tags, session_id = self._get_tracing_metadata_from_rollout(store, rollout_id)
+                logger.info(f"[TracingMetadata] Fetched: environment={environment}, tags={tags}, session_id={session_id}")
+
+                if otlp_enabled:
+                    # Build resource attributes for Mantisdk metadata
+                    resource_attrs: Dict[str, Any] = {
+                        LightningResourceAttributes.ROLLOUT_ID.value: rollout_id,
+                        LightningResourceAttributes.ATTEMPT_ID.value: attempt_id,
+                        LightningResourceAttributes.SPAN_SEQUENCE_ID.value: sequence_id_decimal,
+                        LightningResourceAttributes.SPAN_TYPE.value: "rollout",
+                    }
+                    
+                    # Build span attributes for Langfuse-expected metadata
+                    # Per Langfuse docs, use langfuse.* namespace for environment and tags
+                    span_extra_attrs: Dict[str, Any] = {}
+                    if session_id:
+                        span_extra_attrs["session.id"] = session_id
+                        logger.info(f"[TracingMetadata] Setting session.id={session_id}")
+                    if environment:
+                        span_extra_attrs["langfuse.environment"] = environment
+                        logger.info(f"[TracingMetadata] Setting langfuse.environment={environment}")
+                    
+                    # Extract call type from headers (set by @gepa.judge, @gepa.agent decorators)
+                    call_type = headers_merged.get("x-mantis-call-type")
+                    
+                    # Build final tags list, including call_type if present
+                    final_tags = list(tags) if tags else []
+                    if call_type and call_type not in final_tags:
+                        final_tags.append(call_type)
+                    
+                    if final_tags:
+                        # Insight's OTEL ingestion expects tags on the *resource* under `langfuse.trace.tags`.
+                        # Span-level `langfuse.tags` is not reliably ingested into `traces.tags`.
+                        resource_attrs["langfuse.trace.tags"] = final_tags
+                        # Keep span-level tags too for backwards-compat/debuggability
+                        span_extra_attrs["langfuse.tags"] = final_tags
+                        logger.info(f"[TracingMetadata] Setting langfuse.trace.tags={final_tags}")
+                    
+                    # Prepare spans for export
+                    spans_to_export: List[ReadableSpan] = []
+                    for span in subtree_spans:
+                        # Add resource attributes
+                        span._resource = span._resource.merge(  # pyright: ignore[reportPrivateUsage]
+                            Resource.create(resource_attrs)
+                        )
+                        # Wrap with extra span attributes if we have any
+                        if span_extra_attrs:
+                            wrapped_span = SpanWithExtraAttributes(span, span_extra_attrs)
+                            spans_to_export.append(wrapped_span)
+                        else:
+                            spans_to_export.append(span)
+                    
+                    export_result = self._otlp_exporter.export(spans_to_export)
+                    if export_result != SpanExportResult.SUCCESS:
+                        logger.error(f"Failed to export rollout spans via OTLP. Result: {export_result}")
+                else:
+                    # The old way: store does not support OTLP endpoint
+                    for span in subtree_spans:
+                        loop = self._ensure_loop()
+                        add_otel_span_task = store.add_otel_span(
+                            rollout_id=rollout_id,
+                            attempt_id=attempt_id,
+                            sequence_id=sequence_id_decimal,
+                            readable_span=span,
+                        )
+                        fut = asyncio.run_coroutine_threadsafe(add_otel_span_task, loop)
+                        fut.result()
+
+            elif otlp_enabled:
+                # Job-scoped spans (no rollout context): tag with job_id for experiment tracking
+                job_id = self._get_job_id_from_store(store)
+                
+                # Extract Mantis tracing metadata from x-mantis-* headers
+                mantis_session_id = headers_merged.get("x-mantis-session-id")
+                mantis_environment = headers_merged.get("x-mantis-environment")
+                mantis_tags_str = headers_merged.get("x-mantis-tags")
+                mantis_call_type = headers_merged.get("x-mantis-call-type")
+                mantis_tags = None
+                if mantis_tags_str:
+                    try:
+                        mantis_tags = ast.literal_eval(mantis_tags_str) if mantis_tags_str.startswith("[") else None
+                    except Exception:
+                        pass
+                
+                # Add call_type to tags if present
+                if mantis_call_type:
+                    if mantis_tags is None:
+                        mantis_tags = []
+                    if mantis_call_type not in mantis_tags:
+                        mantis_tags.append(mantis_call_type)
+                
+                # Build span attributes for Mantis/Langfuse metadata
+                span_extra_attrs: Dict[str, Any] = {}
+                job_resource_attrs: Dict[str, Any] = {}
+                
+                if mantis_session_id:
+                    span_extra_attrs["session.id"] = mantis_session_id
+                    logger.info(f"[TracingMetadata] Job span: session.id={mantis_session_id}")
+                if mantis_environment:
+                    span_extra_attrs["langfuse.environment"] = mantis_environment
+                    logger.info(f"[TracingMetadata] Job span: langfuse.environment={mantis_environment}")
+                if mantis_tags:
+                    # Set tags as RESOURCE attributes (required for Insight ingestion)
+                    job_resource_attrs["langfuse.trace.tags"] = mantis_tags
+                    # Also set as span attributes for backwards-compat
+                    span_extra_attrs["langfuse.tags"] = mantis_tags
+                    logger.info(f"[TracingMetadata] Job span: langfuse.trace.tags={mantis_tags}")
+                
+                # Prepare spans for export
+                spans_to_export: List[ReadableSpan] = []
+                for span in subtree_spans:
+                    if job_id:
+                        span._resource = span._resource.merge(  # pyright: ignore[reportPrivateUsage]
+                            Resource.create(
+                                {
+                                    LightningResourceAttributes.JOB_ID.value: job_id,
+                                    LightningResourceAttributes.SPAN_TYPE.value: "job",
+                                }
+                            )
+                        )
+                    # Merge job-level langfuse resource attrs (tags) if present
+                    try:
+                        if "job_resource_attrs" in locals() and job_resource_attrs:
+                            span._resource = span._resource.merge(  # pyright: ignore[reportPrivateUsage]
+                                Resource.create(job_resource_attrs)
+                            )
+                    except Exception:
+                        pass
+                    # Wrap with extra span attributes if we have any
+                    if span_extra_attrs:
+                        wrapped_span = SpanWithExtraAttributes(span, span_extra_attrs)
+                        spans_to_export.append(wrapped_span)
+                    else:
+                        spans_to_export.append(span)
+                
+                export_result = self._otlp_exporter.export(spans_to_export)
+                if export_result != SpanExportResult.SUCCESS:
+                    logger.error(f"Failed to export job spans via OTLP. Result: {export_result}")
+                else:
+                    logger.debug(f"Exported {len(spans_to_export)} job-scoped spans (job_id={job_id}, has_mantis_metadata={bool(span_extra_attrs)})")
+            else:
+                # No OTLP and no rollout context - skip with warning
+                logger.debug(
+                    f"Skipping {len(subtree_spans)} spans: no rollout context and OTLP not enabled"
+                )
+
+    def _get_root_span_ids(self) -> Iterable[int]:
+        """Yield span_ids for root spans currently in the buffer.
+
+        A root span is defined as one with `parent is None`.
+
+        Yields:
+            int: Span id for each root span found.
+        """
+        for span in self._buffer:
+            if span.parent is None:
+                span_context = span.get_span_context()
+                if span_context is not None:
+                    yield span_context.span_id
+
+    def _get_subtrees(self, root_span_id: int) -> Iterable[int]:
+        """Yield span_ids in the subtree rooted at `root_span_id`.
+
+        Depth-first traversal over the current buffer.
+
+        Args:
+            root_span_id: The span id of the root.
+
+        Yields:
+            int: Span ids including the root and all descendants found.
+        """
+        # Yield the root span id first.
+        yield root_span_id
+        for span in self._buffer:
+            # Check whether the span's parent is the root_span_id.
+            if span.parent is not None and span.parent.span_id == root_span_id:
+                span_context = span.get_span_context()
+                if span_context is not None:
+                    # Recursively get child spans.
+                    yield from self._get_subtrees(span_context.span_id)
+
+    def _pop_subtrees(self, root_span_id: int) -> List[ReadableSpan]:
+        """Remove and return the subtree for a particular root from the buffer.
+
+        Args:
+            root_span_id: Root span id identifying the subtree.
+
+        Returns:
+            list[ReadableSpan]: Spans that were part of the subtree. Order follows buffer order.
+        """
+        subtree_span_ids = set(self._get_subtrees(root_span_id))
+        subtree_spans: List[ReadableSpan] = []
+        new_buffer: List[ReadableSpan] = []
+        for span in self._buffer:
+            span_context = span.get_span_context()
+            if span_context is not None and span_context.span_id in subtree_span_ids:
+                subtree_spans.append(span)
+            else:
+                new_buffer.append(span)
+        # Replace buffer with remaining spans to avoid re-processing.
+        self._buffer = new_buffer
+        return subtree_spans
+
+
+class LightningOpenTelemetry(OpenTelemetry):
+    """OpenTelemetry integration that exports spans to the Lightning store.
+
+    Responsibilities:
+
+    * Ensures each request is annotated with a per-attempt sequence id so spans
+      are ordered deterministically even with clock skew across nodes.
+    * Uses [`LightningSpanExporter`][mantisdk.llm_proxy.LightningSpanExporter] to persist spans for analytics and training.
+    * Adds Mantisdk-specific attributes (session_id, tags, environment) to spans.
+
+    Args:
+        otlp_endpoint: Optional OTLP endpoint URL for direct trace export to external
+            collectors (e.g., Langfuse/Insight). When set, spans are exported directly without
+            requiring a store or rollout/attempt headers.
+        otlp_headers: Optional dict of HTTP headers for OTLP authentication (e.g., Basic Auth).
+    """
+
+    def __init__(self, otlp_endpoint: Optional[str] = None, otlp_headers: Optional[Dict[str, str]] = None):
+        exporter = LightningSpanExporter(otlp_endpoint=otlp_endpoint, otlp_headers=otlp_headers)
+        config = OpenTelemetryConfig(exporter=exporter)
+
+        # Check for tracer initialization
+        if _check_tracer_provider():
+            logger.error("Tracer is already initialized. OpenTelemetry may not work as expected.")
+
+        super().__init__(config=config)  # pyright: ignore[reportUnknownMemberType]
+        
+        # Store exporter reference for debugging
+        self._custom_exporter = exporter
+    
+    def _init_tracing(self, tracer_provider):
+        """Override to ensure our span processor is added even when reusing existing TracerProvider.
+        
+        LiteLLM's parent _init_tracing reuses existing TracerProviders but doesn't add
+        our custom span processor to them. We override to force adding our processor.
+        """
+        from opentelemetry import trace as otel_trace_api
+        from opentelemetry.sdk.trace import TracerProvider as TracerProviderSDK
+        from opentelemetry.trace import SpanKind
+        
+        # Call parent to set up tracer
+        super()._init_tracing(tracer_provider)  # pyright: ignore[reportUnknownMemberType]
+        
+        # If an existing provider was reused, add our span processor to it
+        current_provider = otel_trace_api.get_tracer_provider()
+        if isinstance(current_provider, TracerProviderSDK):
+            # Check if our processor is already added
+            has_our_processor = False
+            if hasattr(current_provider, "_active_span_processor"):
+                active_processor = current_provider._active_span_processor  # pyright: ignore[reportPrivateUsage]
+                if hasattr(active_processor, "_span_processors"):
+                    for proc in active_processor._span_processors:  # pyright: ignore[reportPrivateUsage]
+                        if hasattr(proc, "_exporter") and isinstance(proc._exporter, LightningSpanExporter):  # pyright: ignore[reportPrivateUsage]
+                            has_our_processor = True
+                            break
+            
+            if not has_our_processor:
+                # Add our span processor
+                span_processor = self._get_span_processor()
+                current_provider.add_span_processor(span_processor)
+        
+        self.span_kind = SpanKind
+    
+    def _get_span_processor(self, dynamic_headers: Optional[dict] = None):
+        """Override to ensure our custom exporter is used.
+        
+        LiteLLM's parent class checks if OTEL_EXPORTER has export() method and wraps it
+        in SimpleSpanProcessor. We override to add logging and ensure this happens.
+        """
+        from opentelemetry.sdk.trace.export import SimpleSpanProcessor
+        
+        # Use our custom exporter directly
+        if hasattr(self.OTEL_EXPORTER, "export"):
+            processor = SimpleSpanProcessor(self.OTEL_EXPORTER)
+            return processor
+        
+        # Fallback to parent implementation
+        return super()._get_span_processor(dynamic_headers)  # pyright: ignore[reportUnknownMemberType]
+
+    def set_attributes(self, span: Any, kwargs: Dict[str, Any], response_obj: Optional[Any]) -> None:
+        """Override to add Mantisdk-specific attributes from metadata.
+        
+        Extracts session_id, tags, and environment from kwargs["metadata"] and sets
+        them as OTEL span attributes for visibility in Insight/Langfuse.
+        
+        Also extracts extra_headers from kwargs and sets them as metadata.requester_custom_headers
+        so the proxy can read them for tagging.
+        """
+        # Call parent implementation first
+        super().set_attributes(span, kwargs, response_obj)  # pyright: ignore[reportUnknownMemberType]
+        
+        # Extract extra_headers from kwargs and set as metadata.requester_custom_headers
+        # This allows the proxy to read x-mantis-* headers for tagging
+        # Check multiple sources: extra_headers (from OpenAI SDK), kwargs["headers"], and context variable
+        extra_headers = kwargs.get("extra_headers") or kwargs.get("extraHeaders")
+        request_headers = kwargs.get("headers") or {}
+        context_headers = _request_headers_context.get() or {}
+        
+        # Merge all sources: extra_headers, request_headers, and context headers
+        merged_headers = {}
+        if extra_headers and isinstance(extra_headers, dict):
+            merged_headers.update(extra_headers)
+        if request_headers and isinstance(request_headers, dict):
+            # Extract x-mantis-* headers from HTTP request headers passed by LiteLLM
+            for key, value in request_headers.items():
+                if isinstance(key, str) and key.lower().startswith("x-mantis-"):
+                    merged_headers[key] = value
+        # Also check context variable (set by MantisHeadersMiddleware)
+        if context_headers:
+            merged_headers.update(context_headers)
+        
+        if merged_headers:
+            # LiteLLM stores this as a stringified dict in metadata.requester_custom_headers
+            span.set_attribute("metadata.requester_custom_headers", str(merged_headers))
+        
+        # Extract Mantisdk tracing metadata from kwargs
+        metadata = kwargs.get("metadata", {}) if kwargs else {}
+        if not metadata:
+            return
+        
+        # Set session_id as span attribute (standard Langfuse/Insight attribute)
+        session_id = metadata.get("session_id")
+        if session_id:
+            span.set_attribute("session.id", session_id)
+        
+        # Set tags as span attribute
+        tags = metadata.get("tags")
+        if tags and isinstance(tags, list):
+            # Set as JSON array for Langfuse compatibility
+            import json
+            span.set_attribute("tags", json.dumps(tags))
+            # Also set individual tag attributes for filtering
+            for i, tag in enumerate(tags):
+                span.set_attribute(f"tag.{i}", str(tag))
+        
+        # Set environment as span attribute
+        environment = metadata.get("environment")
+        if environment:
+            span.set_attribute("environment", environment)
+
+    async def async_pre_call_deployment_hook(
+        self, kwargs: Dict[str, Any], call_type: Optional[CallTypes] = None
+    ) -> Optional[Dict[str, Any]]:
+        """The root span is sometimes missing (e.g., when Anthropic endpoint is used).
+        It is created in an auth module in LiteLLM. If it's missing, we create it here.
+        """
+        if "metadata" not in kwargs or "litellm_parent_otel_span" not in kwargs["metadata"]:
+            parent_otel_span = self.create_litellm_proxy_request_started_span(  # type: ignore
+                start_time=datetime.now(),
+                headers=kwargs.get("headers", {}),
+            )
+            updated_metadata = {**kwargs.get("metadata", {}), "litellm_parent_otel_span": parent_otel_span}
+
+            return {**kwargs, "metadata": updated_metadata}
+        else:
+            return kwargs
+
+
+class RolloutAttemptMiddleware(BaseHTTPMiddleware):
+    """
+    Rewrites /rollout/{rid}/attempt/{aid}/... -> /...
+    and injects x-rollout-id, x-attempt-id, x-sequence-id headers.
+
+    LLMProxy can update store later without rebuilding middleware.
+    """
+
+    async def dispatch(self, request: Request, call_next: Callable[[Request], Awaitable[Response]]) -> Response:
+        # Decode rollout and attempt from the URL prefix. Example:
+        #   /rollout/r123/attempt/a456/v1/chat/completions
+        # becomes
+        #   /v1/chat/completions
+        # while adding request-scoped headers for trace attribution.
+        path = request.url.path
+
+        match = re.match(r"^/rollout/([^/]+)/attempt/([^/]+)(/.*)?$", path)
+        if match:
+            rollout_id = match.group(1)
+            attempt_id = match.group(2)
+            new_path = match.group(3) if match.group(3) is not None else "/"
+
+            # Rewrite the ASGI scope path so downstream sees a clean OpenAI path.
+            request.scope["path"] = new_path
+            request.scope["raw_path"] = new_path.encode()
+
+            store = get_active_llm_proxy().get_store()
+            if store is not None:
+                # Allocate a monotonic sequence id per (rollout, attempt).
+                sequence_id = await store.get_next_span_sequence_id(rollout_id, attempt_id)
+
+                # Inject headers so downstream components and exporters can retrieve them.
+                request.scope["headers"] = list(request.scope["headers"]) + [
+                    (b"x-rollout-id", rollout_id.encode()),
+                    (b"x-attempt-id", attempt_id.encode()),
+                    (b"x-sequence-id", str(sequence_id).encode()),
+                ]
+            else:
+                logger.warning("Store is not set. Skipping sequence id allocation and header injection.")
+
+        response = await call_next(request)
+        return response
+
+
+class MantisHeadersMiddleware(BaseHTTPMiddleware):
+    """Middleware to intercept x-mantis-* HTTP headers and store them in context.
+    
+    This allows LiteLLM callbacks to access custom headers (like x-mantis-call-type)
+    even if LiteLLM doesn't pass them through kwargs.
+    """
+    
+    async def dispatch(self, request: Request, call_next: Callable[[Request], Awaitable[Response]]) -> Response:
+        # Extract x-mantis-* headers from HTTP request
+        mantis_headers: Dict[str, str] = {}
+        for header_name, header_value in request.headers.items():
+            if isinstance(header_name, str) and header_name.lower().startswith("x-mantis-"):
+                mantis_headers[header_name] = header_value
+        
+        # Store in context variable for callback access
+        if mantis_headers:
+            _request_headers_context.set(mantis_headers)
+        
+        response = await call_next(request)
+        return response
+
+
+class MessageInspectionMiddleware(BaseHTTPMiddleware):
+    """Middleware to inspect the request and response bodies.
+
+    It's for debugging purposes. Add it via "message_inspection" middleware alias.
+    """
+
+    async def dispatch(self, request: Request, call_next: Callable[[Request], Awaitable[Response]]) -> Response:
+        ti = time.time()
+        logger.info(f"Received request with scope: {request.scope}")
+        logger.info(f"Received request with body: {await request.body()}")
+        response = await call_next(request)
+        elapsed = time.time() - ti
+        logger.info(f"Response to request took {elapsed} seconds")
+        logger.info(f"Received response with status code: {response.status_code}")
+        logger.info(f"Received response with body: {response.body}")
+        return response
+
+
+class StreamConversionMiddleware(BaseHTTPMiddleware):
+    """Middleware to convert streaming responses to non-streaming responses.
+
+    Useful for backend that only supports non-streaming responses.
+
+    LiteLLM's OpenTelemetry is also buggy with streaming responses.
+    The conversion will hopefully bypass the bug.
+    """
+
+    async def dispatch(self, request: Request, call_next: Callable[[Request], Awaitable[Response]]) -> Response:
+        # Only process POST requests to completion endpoints
+        if request.method != "POST":
+            return await call_next(request)
+
+        # Check if it's a chat completions or messages endpoint
+        endpoint_format: Literal["openai", "anthropic", "unknown"] = "unknown"
+        if request.url.path.endswith("/chat/completions") or "/chat/completions?" in request.url.path:
+            endpoint_format = "openai"
+        elif request.url.path.endswith("/messages") or "/messages?" in request.url.path:
+            endpoint_format = "anthropic"
+        else:
+            endpoint_format = "unknown"
+
+        if endpoint_format == "unknown":
+            # Directly bypass the middleware
+            return await call_next(request)
+
+        # Read the request body
+        try:
+            json_body = await request.json()
+        except json.JSONDecodeError:
+            logger.warning(f"Request body is not valid JSON: {request.body}")
+            return await call_next(request)
+
+        # Check if streaming is requested
+        is_streaming = json_body.get("stream", False)
+
+        # Simple case: no streaming requested, just return the response
+        if not is_streaming:
+            return await call_next(request)
+
+        # Now the stream case
+        return await self._handle_stream_case(request, json_body, endpoint_format, call_next)
+
+    async def _handle_stream_case(
+        self,
+        request: Request,
+        json_body: Dict[str, Any],
+        endpoint_format: Literal["openai", "anthropic"],
+        call_next: Callable[[Request], Awaitable[Response]],
+    ) -> Response:
+        # 1) Modify the request body to force stream=False
+        modified_json = dict(json_body)
+        modified_json["stream"] = False
+        modified_body = json.dumps(modified_json).encode("utf-8")
+
+        # 2) Build a new scope + receive that yields our modified body
+        scope: Scope = dict(request.scope)
+        # rewrite headers for accept/content-length
+        new_headers: List[Tuple[bytes, bytes]] = []
+        saw_accept = False
+        for k, v in scope["headers"]:
+            kl = k.lower()
+            if kl == b"accept":
+                saw_accept = True
+                new_headers.append((k, b"application/json"))
+            elif kl == b"content-length":
+                # replace with new length
+                continue
+            else:
+                new_headers.append((k, v))
+        if not saw_accept:
+            new_headers.append((b"accept", b"application/json"))
+        new_headers.append((b"content-length", str(len(modified_body)).encode("ascii")))
+        scope["headers"] = new_headers
+
+        # Directly modify the request body
+        # Creating a new request won't work because request is cached in the base class
+        request._body = modified_body  # type: ignore
+
+        response = await call_next(request)
+
+        buffered: Optional[bytes] = None
+        # 4) If OK, buffer the response body (it should be JSON because we forced stream=False)
+        if 200 <= response.status_code < 300:
+            try:
+                if hasattr(response, "body_iterator"):
+                    # Buffer body safely
+                    body_chunks: List[bytes] = []
+                    async for chunk in response.body_iterator:  # type: ignore
+                        body_chunks.append(chunk)  # type: ignore
+                    buffered = b"".join(body_chunks)
+                else:
+                    buffered = response.body  # type: ignore
+
+                data = json.loads(buffered or b"{}")
+
+                if endpoint_format == "anthropic":
+                    return StreamingResponse(
+                        self.anthropic_stream_generator(data),
+                        media_type="text/event-stream",
+                        headers={"Cache-Control": "no-cache", "Connection": "keep-alive", "X-Accel-Buffering": "no"},
+                    )
+                else:
+                    # openai format
+                    return StreamingResponse(
+                        self.openai_stream_generator(data),
+                        media_type="text/event-stream",
+                        headers={"Cache-Control": "no-cache", "Connection": "keep-alive", "X-Accel-Buffering": "no"},
+                    )
+            except Exception as e:
+                # If anything goes wrong, fall back to non-streaming JSON
+                logger.exception(f"Error converting to stream; returning non-stream response: {e}")
+                # Rebuild the consumed response
+                return Response(
+                    content=buffered if buffered is not None else b"",
+                    status_code=response.status_code,
+                    headers=dict(response.headers),
+                    media_type=response.media_type,
+                    background=response.background,
+                )
+        else:
+            return response
+
+    async def anthropic_stream_generator(self, original_response: Dict[str, Any]):
+        """Generate Anthropic SSE-formatted chunks from complete content blocks
+
+        This is a dirty hack for Anthropic-style streaming from non-streaming response.
+        The sse format is subject to change based on Anthropic's implementation.
+        If so, try to use `MessageInspectionMiddleware` to inspect the update and fix accordingly.
+        """
+        # Anthropic format - handle multiple content blocks (text + tool_use)
+        content_blocks: List[Dict[str, Any]] = original_response.get("content", [])
+        message_id = original_response.get("id", f"msg_{int(time.time() * 1000)}")
+        model = original_response.get("model", "claude")
+
+        # Send message_start event
+        message_start: Dict[str, Any] = {
+            "type": "message_start",
+            "message": {
+                "id": message_id,
+                "type": "message",
+                "role": "assistant",
+                "content": [],
+                "model": model,
+                "stop_reason": None,
+                "stop_sequence": None,
+                "usage": original_response.get("usage", {"input_tokens": 0, "output_tokens": 0}),
+            },
+        }
+        yield f"event: message_start\ndata: {json.dumps(message_start)}\n\n"
+
+        # Send ping to keep connection alive
+        ping = {"type": "ping"}
+        yield f"event: ping\ndata: {json.dumps(ping)}\n\n"
+
+        # Process each content block
+        for block_index, block in enumerate(content_blocks):
+            block_type = block.get("type", "text")
+
+            if block_type == "text":
+                # Handle text block
+                content = block.get("text", "")
+
+                # Send content_block_start event
+                content_block_start = {
+                    "type": "content_block_start",
+                    "index": block_index,
+                    "content_block": {"type": "text", "text": ""},
+                }
+                yield f"event: content_block_start\ndata: {json.dumps(content_block_start)}\n\n"
+
+                # Stream text content in chunks
+                if content:
+                    words = content.split()
+                    chunk_size = 5
+
+                    for i in range(0, len(words), chunk_size):
+                        chunk_words = words[i : i + chunk_size]
+                        text_chunk = " ".join(chunk_words)
+
+                        # Add space after chunk unless it's the last one
+                        if i + chunk_size < len(words):
+                            text_chunk += " "
+
+                        content_block_delta = {
+                            "type": "content_block_delta",
+                            "index": block_index,
+                            "delta": {"type": "text_delta", "text": text_chunk},
+                        }
+                        yield f"event: content_block_delta\ndata: {json.dumps(content_block_delta)}\n\n"
+                        await asyncio.sleep(0.02)
+
+                # Send content_block_stop event
+                content_block_stop = {"type": "content_block_stop", "index": block_index}
+                yield f"event: content_block_stop\ndata: {json.dumps(content_block_stop)}\n\n"
+
+            elif block_type == "tool_use":
+                # Handle tool_use block
+                tool_name = block.get("name", "")
+                tool_input = block.get("input", {})
+                tool_id = block.get("id", f"toolu_{int(time.time() * 1000)}")
+
+                # Send content_block_start event for tool use
+                content_block_start: Dict[str, Any] = {
+                    "type": "content_block_start",
+                    "index": block_index,
+                    "content_block": {"type": "tool_use", "id": tool_id, "name": tool_name, "input": {}},
+                }
+                yield f"event: content_block_start\ndata: {json.dumps(content_block_start)}\n\n"
+
+                # Stream tool input as JSON string chunks
+                input_json = json.dumps(tool_input)
+                chunk_size = 20  # characters per chunk for JSON
+
+                for i in range(0, len(input_json), chunk_size):
+                    json_chunk = input_json[i : i + chunk_size]
+
+                    content_block_delta = {
+                        "type": "content_block_delta",
+                        "index": block_index,
+                        "delta": {"type": "input_json_delta", "partial_json": json_chunk},
+                    }
+                    yield f"event: content_block_delta\ndata: {json.dumps(content_block_delta)}\n\n"
+                    await asyncio.sleep(0.01)
+
+                # Send content_block_stop event
+                content_block_stop = {"type": "content_block_stop", "index": block_index}
+                yield f"event: content_block_stop\ndata: {json.dumps(content_block_stop)}\n\n"
+
+        # Send message_delta event with stop reason
+        message_delta = {
+            "type": "message_delta",
+            "delta": {"stop_reason": original_response.get("stop_reason", "end_turn"), "stop_sequence": None},
+            "usage": {"output_tokens": original_response.get("usage", {}).get("output_tokens", 0)},
+        }
+        yield f"event: message_delta\ndata: {json.dumps(message_delta)}\n\n"
+
+        # Send message_stop event
+        message_stop = {"type": "message_stop"}
+        yield f"event: message_stop\ndata: {json.dumps(message_stop)}\n\n"
+
+    async def openai_stream_generator(self, response_json: Dict[str, Any]) -> AsyncGenerator[str, Any]:
+        """
+        Convert a *complete* OpenAI chat.completions choice into a stream of
+        OpenAI-compatible SSE chunks.
+
+        This emits:
+
+          - an initial delta with the role ("assistant"),
+          - a sequence of deltas for message.content (split into small chunks),
+          - deltas for any tool_calls (including id/name and chunked arguments),
+          - a terminal chunk with finish_reason,
+          - and finally the literal '[DONE]'.
+
+        Notes:
+
+        - We only handle a *single* choice (index 0 typically).
+        - We purposefully don't attempt to stream logprobs.
+        - Chunking strategy is simple and conservative to avoid splitting
+          multi-byte characters: we slice on spaces where possible, then fall
+          back to fixed-size substrings.
+        """
+        choice = cast(Dict[str, Any], (response_json.get("choices") or [{}])[0])
+        model = response_json.get("model", "unknown")
+        created: int = int(time.time())
+        index: int = choice.get("index", 0)
+
+        message: Dict[str, Any] = choice.get("message", {}) or {}
+        role: str = message.get("role", "assistant")
+        content: str = message.get("content") or ""
+        tool_calls: List[Any] = message.get("tool_calls") or []
+        finish_reason: Optional[str] = choice.get(
+            "finish_reason"
+        )  # e.g., "stop", "length", "tool_calls", "content_filter"
+
+        def sse_chunk(obj: Dict[str, Any]) -> str:
+            return f"data: {json.dumps(obj, ensure_ascii=False)}\n\n"
+
+        # 1) initial chunk with the role
+        yield sse_chunk(
+            {
+                "id": f"chatcmpl-{created}",
+                "object": "chat.completion.chunk",
+                "created": created,
+                "model": model,
+                "choices": [{"index": index, "delta": {"role": role}, "finish_reason": None}],
+            }
+        )
+
+        # 2) stream textual content as small deltas
+        async def stream_content(text: str):
+            if not text:
+                return
+            # prefer splitting on spaces in ~20–40 char pieces
+            approx = 28
+            start = 0
+            n = len(text)
+            while start < n:
+                end = min(start + approx, n)
+                if end < n:
+                    # try to break on a space going forward
+                    space = text.rfind(" ", start, end)
+                    if space > start:
+                        end = space + 1
+                delta_text = text[start:end]
+                start = end
+                if not delta_text:
+                    break
+                yield sse_chunk(
+                    {
+                        "id": f"chatcmpl-{created}",
+                        "object": "chat.completion.chunk",
+                        "created": created,
+                        "model": model,
+                        "choices": [{"index": index, "delta": {"content": delta_text}, "finish_reason": None}],
+                    }
+                )
+                # tiny pause helps some UIs animate smoothly; keep very small
+                await asyncio.sleep(0.0)
+
+        async for piece in stream_content(content):  # type: ignore[misc]
+            yield piece  # pass through the produced chunks
+
+        # 3) stream tool_calls if present (id/name first, then arguments piecemeal)
+        for tc_index, tc in enumerate(tool_calls):
+            tc_type = tc.get("type", "function")
+            tc_id = tc.get("id") or f"call_{created}_{tc_index}"
+            fn: Dict[str, Any] = (tc.get("function") or {}) if tc_type == "function" else {}
+            fn_name: str = fn.get("name", "")
+            fn_args: str = fn.get("arguments", "") or ""
+
+            # (a) delta that announces the tool call id/type/name
+            yield sse_chunk(
+                {
+                    "id": f"chatcmpl-{created}",
+                    "object": "chat.completion.chunk",
+                    "created": created,
+                    "model": model,
+                    "choices": [
+                        {
+                            "index": index,
+                            "delta": {
+                                "tool_calls": [
+                                    {"index": tc_index, "id": tc_id, "type": tc_type, "function": {"name": fn_name}}
+                                ]
+                            },
+                            "finish_reason": None,
+                        }
+                    ],
+                }
+            )
+
+            # (b) stream arguments in small substrings
+            arg_chunk_size = 40
+            for pos in range(0, len(fn_args), arg_chunk_size):
+                partial = fn_args[pos : pos + arg_chunk_size]
+                yield sse_chunk(
+                    {
+                        "id": f"chatcmpl-{created}",
+                        "object": "chat.completion.chunk",
+                        "created": created,
+                        "model": model,
+                        "choices": [
+                            {
+                                "index": index,
+                                "delta": {"tool_calls": [{"index": tc_index, "function": {"arguments": partial}}]},
+                                "finish_reason": None,
+                            }
+                        ],
+                    }
+                )
+                await asyncio.sleep(0.0)
+
+        # 4) terminal chunk with finish_reason (default to "stop" if missing)
+        yield sse_chunk(
+            {
+                "id": f"chatcmpl-{created}",
+                "object": "chat.completion.chunk",
+                "created": created,
+                "model": model,
+                "choices": [
+                    {
+                        "index": index,
+                        "delta": {},
+                        "finish_reason": finish_reason or ("tool_calls" if tool_calls else "stop"),
+                    }
+                ],
+            }
+        )
+
+        # 5) literal DONE sentinel
+        yield "data: [DONE]\n\n"
+
+
+_MIDDLEWARE_REGISTRY: Dict[str, Type[BaseHTTPMiddleware]] = {
+    "rollout_attempt": RolloutAttemptMiddleware,
+    "stream_conversion": StreamConversionMiddleware,
+    "message_inspection": MessageInspectionMiddleware,
+    "mantis_headers": MantisHeadersMiddleware,
+}
+
+
+_CALLBACK_REGISTRY = {
+    "return_token_ids": AddReturnTokenIds,
+    "logprobs": AddLogprobs,
+    "opentelemetry": LightningOpenTelemetry,
+}
+
+
+class LLMProxy:
+    """Host a LiteLLM OpenAI-compatible proxy bound to a LightningStore.
+
+    The proxy:
+
+    * Serves an OpenAI-compatible API via uvicorn.
+    * Adds rollout/attempt routing and headers via middleware.
+    * Registers OTEL export and token-id callbacks.
+    * Writes a LiteLLM worker config file with `model_list` and settings.
+
+    Lifecycle:
+
+    * [`start()`][mantisdk.LLMProxy.start] writes config, starts uvicorn server in a thread, and waits until ready.
+    * [`stop()`][mantisdk.LLMProxy.stop] tears down the server and removes the temp config file.
+    * [`restart()`][mantisdk.LLMProxy.restart] convenience wrapper to stop then start.
+
+    !!! note
+
+        As the LLM Proxy sets up an OpenTelemetry tracer, it's recommended to run it in a different
+        process from the main runner (i.e., tracer from agents). See `launch_mode` for how to change that.
+
+    !!! warning
+
+        By default (or when "stream_conversion" middleware is enabled), the LLM Proxy will convert OpenAI and Anthropic requests with `stream=True`
+        to a non-streaming request before going through the LiteLLM proxy. This is because the OpenTelemetry tracer provided by
+        LiteLLM is buggy with streaming responses. You can disable this by removing the "stream_conversion" middleware.
+        In that case, you might lose some tracing information like token IDs.
+
+    !!! danger
+
+        Do not run LLM proxy in the same process as the main runner. It's easy to cause conflicts in the tracer provider
+        with tracers like [`AgentOpsTracer`][mantisdk.AgentOpsTracer].
+
+    Args:
+        port: TCP port to bind. Will bind to a random port if not provided.
+        model_list: LiteLLM `model_list` entries.
+        store: LightningStore used for span sequence and persistence.
+        host: Publicly reachable host used in resource endpoints. See `host` of `launcher_args` for more details.
+        litellm_config: Extra LiteLLM proxy config merged with `model_list`.
+        num_retries: Default LiteLLM retry count injected into `litellm_settings`.
+        num_workers: Number of workers to run in the server. Only applicable for "mp" launch mode. Ignored if launcher_args is provided.
+            When `num_workers > 1`, the server will be run using [gunicorn](https://gunicorn.org/).
+        launch_mode: Launch mode for the server. Defaults to "mp". Cannot be used together with launcher_args. Ignored if launcher_args is provided.
+            It's recommended to use `launch_mode="mp"` to launch the proxy, which will launch the server in a separate process.
+            `launch_mode="thread"` can also be used if used in caution. It will launch the server in a separate thread.
+            `launch_mode="asyncio"` launches the server in the current thread as an asyncio task.
+            It is NOT recommended because it often causes hanging requests. Only use it if you know what you are doing.
+        launcher_args: Arguments for the server launcher. If this is provided, host, port, and launch_mode will be ignored. Cannot be used together with port, host, and launch_mode.
+        middlewares: List of FastAPI middleware classes or strings to register. You can specify the class aliases or classes that have been imported.
+            If not provided, the default middlewares (RolloutAttemptMiddleware and StreamConversionMiddleware) will be used.
+            Available middleware aliases are: "rollout_attempt", "stream_conversion", "message_inspection".
+            Middlewares are the **first layer** of request processing. They are applied to all requests before the LiteLLM proxy.
+        callbacks: List of LiteLLM callback classes or strings to register. You can specify the class aliases or classes that have been imported.
+            If not provided, the default callbacks (AddReturnTokenIds and LightningOpenTelemetry) will be used.
+            Available callback aliases are: "return_token_ids", "opentelemetry", "logprobs".
+        otlp_endpoint: Optional OTLP endpoint URL for direct trace export to external
+            collectors (e.g., Langfuse/Insight). When set, spans are exported directly without
+            requiring rollout/attempt headers. Format: "http://host:port/api/public/otel/v1/traces"
+        otlp_headers: Optional dict of HTTP headers for OTLP authentication.
+            For Langfuse/Insight, use Basic Auth: {"Authorization": "Basic base64(publicKey:secretKey)"}
+    """
+
+    def __init__(
+        self,
+        port: int | None = None,
+        model_list: List[ModelConfig] | None = None,
+        store: Optional[LightningStore] = None,
+        host: str | None = None,
+        litellm_config: Dict[str, Any] | None = None,
+        num_retries: int = 0,
+        num_workers: int = 1,
+        launch_mode: LaunchMode = "mp",
+        launcher_args: PythonServerLauncherArgs | None = None,
+        middlewares: Sequence[Union[Type[BaseHTTPMiddleware], str]] | None = None,
+        callbacks: Sequence[Union[Type[CustomLogger], str]] | None = None,
+        otlp_endpoint: Optional[str] = None,
+        otlp_headers: Optional[Dict[str, str]] = None,
+    ):
+        self.store = store
+        self._otlp_endpoint = otlp_endpoint
+        self._otlp_headers = otlp_headers
+        
+        # Log OTLP configuration for diagnostics
+        if otlp_endpoint:
+            logger.info(f"LLMProxy initialized with OTLP endpoint: {otlp_endpoint}")
+            if otlp_headers:
+                logger.info(f"LLMProxy OTLP headers configured: {list(otlp_headers.keys())}")
+        else:
+            logger.debug("LLMProxy initialized without OTLP endpoint (will use store-based export)")
+
+        if launcher_args is not None and (
+            port is not None or host is not None or launch_mode != "mp" or num_workers != 1
+        ):
+            raise ValueError("port, host, launch_mode, and num_workers cannot be set when launcher_args is provided.")
+
+        self.server_launcher_args = launcher_args or PythonServerLauncherArgs(
+            port=port,
+            host=host,
+            launch_mode=launch_mode,
+            n_workers=num_workers,
+            # NOTE: This /health endpoint can be slow sometimes because it actually probes the backend LLM service.
+            healthcheck_url="/health",
+            startup_timeout=60.0,
+        )
+
+        if self.server_launcher_args.healthcheck_url is None:
+            logger.warning("healthcheck_url is not set. LLM Proxy will not be checked for healthiness after starting.")
+
+        self.model_list = model_list or []
+        self.litellm_config = litellm_config or {}
+
+        # Ensure num_retries is present inside the litellm_settings block.
+        self.litellm_config.setdefault("litellm_settings", {})
+        self.litellm_config["litellm_settings"].setdefault("num_retries", num_retries)
+        self.server_launcher = PythonServerLauncher(app, self.server_launcher_args, noop_context())
+
+        self._config_file = None
+
+        self.middlewares: List[Type[BaseHTTPMiddleware]] = []
+        if middlewares is None:
+            middlewares = ["mantis_headers", "rollout_attempt", "stream_conversion"]
+        for middleware in middlewares:
+            if isinstance(middleware, str):
+                if middleware not in _MIDDLEWARE_REGISTRY:
+                    raise ValueError(
+                        f"Invalid middleware alias: {middleware}. Available aliases are: {list(_MIDDLEWARE_REGISTRY.keys())}"
+                    )
+                middleware = _MIDDLEWARE_REGISTRY[middleware]
+                self.middlewares.append(middleware)
+            else:
+                self.middlewares.append(middleware)
+
+        self.callbacks: List[Type[CustomLogger]] = []
+        if callbacks is None:
+            callbacks = ["return_token_ids", "opentelemetry"]
+        for callback in callbacks:
+            if isinstance(callback, str):
+                if callback not in _CALLBACK_REGISTRY:
+                    raise ValueError(
+                        f"Invalid callback alias: {callback}. Available aliases are: {list(_CALLBACK_REGISTRY.keys())}"
+                    )
+                callback = _CALLBACK_REGISTRY[callback]
+                self.callbacks.append(callback)
+            else:
+                self.callbacks.append(callback)
+
+    def get_store(self) -> Optional[LightningStore]:
+        """Get the store used by the proxy.
+
+        Returns:
+            The store used by the proxy.
+        """
+        return self.store
+
+    def set_store(self, store: LightningStore) -> None:
+        """Set the store for the proxy.
+
+        Args:
+            store: The store to use for the proxy.
+        """
+        self.store = store
+
+    def update_model_list(self, model_list: List[ModelConfig]) -> None:
+        """Replace the in-memory model list.
+
+        Args:
+            model_list: New list of model entries.
+        """
+        self.model_list = model_list
+        logger.info(f"Updating LLMProxy model list to: {model_list}")
+        # Do nothing if the server is not running.
+
+    def initialize(self):
+        """Initialize global middleware and LiteLLM callbacks.
+
+        Installs:
+
+        * A FastAPI middleware that rewrites /rollout/{rid}/attempt/{aid}/... paths,
+        injects rollout/attempt/sequence headers, and forwards downstream.
+        * LiteLLM callbacks for token ids and OpenTelemetry export.
+
+        The middleware can only be installed once because once the FastAPI app has started,
+        the middleware cannot be changed any more.
+
+        This function does not start any server. It only wires global hooks.
+        """
+        if self.store is None:
+            raise ValueError("Store is not set. Please set the store before initializing the LLMProxy.")
+
+        if _global_llm_proxy is not None:
+            logger.warning("A global LLMProxy is already set. Overwriting it with the new instance.")
+
+        # Patch for LiteLLM v1.80.6+: https://github.com/BerriAI/litellm/issues/17243
+        os.environ["USE_OTEL_LITELLM_REQUEST_SPAN"] = "true"
+
+        # Set the global LLMProxy reference for middleware/exporter access.
+        set_active_llm_proxy(self)
+
+        # Install middleware if it's not already installed.
+        installation_status: Dict[Any, bool] = {}
+        for mw in app.user_middleware:
+            installation_status[mw.cls] = True
+
+        for mw in self.middlewares:
+            if mw not in installation_status:
+                logger.info(f"Adding middleware {mw} to the FastAPI app.")
+                app.add_middleware(mw)
+            else:
+                logger.info(f"Middleware {mw} is already installed. Will not install a new one.")
+
+        if not initialize_llm_callbacks(self.callbacks, otlp_endpoint=self._otlp_endpoint, otlp_headers=self._otlp_headers):
+            # If it's not the first time to initialize the callbacks, also
+            # reset LiteLLM's logging worker so its asyncio.Queue binds to the new loop.
+            _reset_litellm_logging_worker()
+
+    @asynccontextmanager
+    async def _serve_context(self) -> AsyncGenerator[None, None]:
+        """Context manager to serve the proxy server.
+
+        See [`start`][mantisdk.LLMProxy.start] and [`stop`][mantisdk.LLMProxy.stop] for more details.
+        """
+
+        if not self.store:
+            raise ValueError("Store is not set. Please set the store before starting the LLMProxy.")
+
+        # Initialize global middleware and callbacks.
+        self.initialize()
+
+        # Persist a temp worker config for LiteLLM and point the proxy at it.
+        self._config_file = tempfile.NamedTemporaryFile(suffix=".yaml", delete=False).name
+        with open(self._config_file, "w") as fp:
+            yaml.safe_dump(
+                {
+                    "model_list": self.model_list,
+                    **self.litellm_config,
+                },
+                fp,
+            )
+
+        save_worker_config(config=self._config_file)
+
+        # NOTE: When running the _serve_context in current process, you might encounter the following problems:
+        # Problem 1: in litellm worker, <Queue at 0x70f1d028cd90 maxsize=50000> is bound to a different event loop
+        # Problem 2: Proxy has conflicted opentelemetry setup with the main process.
+
+        # Ready
+        logger.info("LLMProxy preparation is done. Will start the server.")
+        yield
+
+        # Clean up
+
+        logger.info("LLMProxy server is cleaning up.")
+
+        # Remove worker config to avoid stale references.
+        if self._config_file and os.path.exists(self._config_file):
+            os.unlink(self._config_file)
+
+        logger.info("LLMProxy server finishes.")
+
+    async def start(self):
+        """Start the proxy server thread and initialize global wiring.
+
+        Side effects:
+
+        * Sets the module-level global store for middleware/exporter access.
+        * Calls `initialize()` once to register middleware and callbacks.
+        * Writes a temporary YAML config consumed by LiteLLM worker.
+        * Launches uvicorn in a daemon thread and waits for readiness.
+        """
+        # Refresh the serve context
+        self.server_launcher.serve_context = self._serve_context()
+
+        if self.store is None:
+            raise ValueError("Store is not set. Please set the store before starting the LLMProxy.")
+
+        store_capabilities = self.store.capabilities
+        if self.server_launcher.args.launch_mode == "mp" and not store_capabilities.get("zero_copy", False):
+            raise RuntimeError(
+                "The store does not support zero-copy. Please use another store, or use asyncio or thread mode to launch the server."
+            )
+        elif self.server_launcher.args.launch_mode == "thread" and not store_capabilities.get("thread_safe", False):
+            raise RuntimeError(
+                "The store is not thread-safe. Please use another store, or use asyncio mode to launch the server."
+            )
+        elif self.server_launcher.args.launch_mode == "asyncio" and not store_capabilities.get("async_safe", False):
+            raise RuntimeError("The store is not async-safe. Please use another store.")
+
+        logger.info(
+            f"Starting LLMProxy server in {self.server_launcher.args.launch_mode} mode with store capabilities: {store_capabilities}"
+        )
+
+        await self.server_launcher.start()
+
+    async def stop(self):
+        """Stop the proxy server and clean up temporary artifacts.
+
+        This is a best-effort graceful shutdown with a bounded join timeout.
+        """
+        if not self.is_running():
+            logger.warning("LLMProxy is not running. Nothing to stop.")
+            return
+
+        await self.server_launcher.stop()
+
+    async def restart(self, *, _port: int | None = None) -> None:
+        """Restart the proxy if running, else start it.
+
+        Convenience wrapper calling `stop()` followed by `start()`.
+        """
+        logger.info("Restarting LLMProxy server...")
+        if self.is_running():
+            await self.stop()
+        if _port is not None:
+            self.server_launcher_args.port = _port
+        await self.start()
+
+    def is_running(self) -> bool:
+        """Return whether the uvicorn server is active.
+
+        Returns:
+            bool: True if server was started and did not signal exit.
+        """
+        return self.server_launcher.is_running()
+
+    def as_resource(
+        self,
+        rollout_id: str | None = None,
+        attempt_id: str | None = None,
+        model: str | None = None,
+        sampling_parameters: Dict[str, Any] | None = None,
+    ) -> LLM:
+        """Create an `LLM` resource pointing at this proxy with rollout context.
+
+        The returned endpoint is:
+            `http://{host}:{port}/rollout/{rollout_id}/attempt/{attempt_id}`
+
+        Args:
+            rollout_id: Rollout identifier used for span attribution. If None, will instantiate a ProxyLLM resource.
+            attempt_id: Attempt identifier used for span attribution. If None, will instantiate a ProxyLLM resource.
+            model: Logical model name to use. If omitted and exactly one model
+                is configured or all models have the same name, that model is used.
+            sampling_parameters: Optional default sampling parameters.
+
+        Returns:
+            LLM: Configured resource ready for OpenAI-compatible calls.
+
+        Raises:
+            ValueError: If `model` is omitted and zero or multiple models are configured.
+        """
+        if model is None:
+            if len(self.model_list) == 1:
+                model = self.model_list[0]["model_name"]
+            elif len(self.model_list) == 0:
+                raise ValueError("No models found in model_list. Please specify the model.")
+            else:
+                first_model_name = self.model_list[0]["model_name"]
+                if all(model_config["model_name"] == first_model_name for model_config in self.model_list):
+                    model = first_model_name
+                else:
+                    raise ValueError(
+                        f"Multiple models found in model_list: {self.model_list}. Please specify the model."
+                    )
+
+        if rollout_id is None and attempt_id is None:
+            return ProxyLLM(
+                endpoint=self.server_launcher.access_endpoint,
+                model=model,
+                sampling_parameters=dict(sampling_parameters or {}),
+            )
+        elif rollout_id is not None and attempt_id is not None:
+            return LLM(
+                endpoint=f"{self.server_launcher.access_endpoint}/rollout/{rollout_id}/attempt/{attempt_id}",
+                model=model,
+                sampling_parameters=dict(sampling_parameters or {}),
+            )
+        else:
+            raise ValueError("Either rollout_id and attempt_id must be provided, or neither.")
+
+
+_global_llm_proxy: Optional[LLMProxy] = None
+_callbacks_before_litellm_start: Optional[List[Any]] = None
+
+
+def get_active_llm_proxy() -> LLMProxy:
+    """Get the current global LLMProxy instance.
+
+    Returns:
+        Optional[LLMProxy]: The current LLMProxy if set, else None.
+    """
+    if _global_llm_proxy is None:
+        raise ValueError("Global LLMProxy is not set. Please call llm_proxy.start() first.")
+    return _global_llm_proxy
+
+
+def set_active_llm_proxy(proxy: LLMProxy) -> None:
+    """Set the current global LLMProxy instance.
+
+    Args:
+        proxy: The LLMProxy instance to set as global.
+    """
+    global _global_llm_proxy
+    _global_llm_proxy = proxy
+
+
+def initialize_llm_callbacks(
+    callback_classes: List[Type[CustomLogger]],
+    otlp_endpoint: Optional[str] = None,
+    otlp_headers: Optional[Dict[str, str]] = None,
+) -> bool:
+    """Restore `litellm.callbacks` to a state that is just initialized by mantisdk.
+
+    When litellm is restarted multiple times in the same process, more and more callbacks
+    will be appended to `litellm.callbacks`, which may exceed the MAX_CALLBACKS limit.
+    This function remembers the initial state of `litellm.callbacks` and always restore to that state.
+
+    Args:
+        callback_classes: List of callback classes to register.
+        otlp_endpoint: Optional OTLP endpoint URL for direct trace export.
+        otlp_headers: Optional dict of HTTP headers for OTLP authentication.
+
+    Returns:
+        Whether the callbacks are initialized for the first time.
+    """
+    global _callbacks_before_litellm_start
+
+    def _instantiate_callback(cls: Type[CustomLogger]) -> CustomLogger:
+        """Instantiate callback with appropriate arguments."""
+        if cls is LightningOpenTelemetry:
+            return LightningOpenTelemetry(otlp_endpoint=otlp_endpoint, otlp_headers=otlp_headers)
+        return cls()
+
+    if _callbacks_before_litellm_start is None:
+        litellm.callbacks.extend([_instantiate_callback(cls) for cls in callback_classes])  # type: ignore
+        _callbacks_before_litellm_start = [*litellm.callbacks]  # type: ignore
+        return True
+
+    else:
+        # Put whatever is missing in the new callback classes to the existing callbacks.
+        for cls in callback_classes:
+            if not any(isinstance(cb, cls) for cb in _callbacks_before_litellm_start):
+                logger.info(f"Adding missing callback {cls} to the existing callbacks.")
+                _callbacks_before_litellm_start.append(_instantiate_callback(cls))
+
+    _reset_litellm_logging_callback_manager()
+
+    if LightningOpenTelemetry in callback_classes:
+        # Check if tracer provider is malformed due to global tracer clear in tests.
+        if not _check_tracer_provider():
+            logger.warning(
+                "Global tracer provider might have been cleared outside. Re-initializing OpenTelemetry callback."
+            )
+            _callbacks_before_litellm_start = [
+                cb for cb in _callbacks_before_litellm_start if not isinstance(cb, LightningOpenTelemetry)
+            ] + [LightningOpenTelemetry(otlp_endpoint=otlp_endpoint, otlp_headers=otlp_headers)]
+        else:
+            logger.debug("Global tracer provider is valid. Reusing existing OpenTelemetry callback.")
+    # Otherwise, we just skip the check for opentelemetry and use the existing callback.
+
+    litellm.callbacks.clear()  # type: ignore
+    litellm.callbacks.extend(_callbacks_before_litellm_start)  # type: ignore
+    return False
+
+
+def _check_tracer_provider() -> bool:
+    """Check if the global tracer provider is properly initialized.
+
+    We don't guarantee the tracer provider is our tracer provider.
+
+    Returns:
+        bool: True if the tracer provider is valid, else False.
+    """
+    if (
+        hasattr(trace_api, "_TRACER_PROVIDER")
+        and trace_api._TRACER_PROVIDER is not None  # pyright: ignore[reportPrivateUsage]
+    ):
+        return True
+    return False