otelzero 0.2.0__py3-none-any.whl

otelzero/hooks/llama_index.py

@@ -0,0 +1,724 @@
+"""LlamaIndex instrumentation: wrap the workflow runtime every agent starts through.
+
+`LlamaIndex <https://github.com/run-llama/llama_index>`_ builds every modern
+agent -- ``FunctionAgent``, ``ReActAgent``, ``CodeActAgent`` and the multi-agent
+``AgentWorkflow`` -- as a *workflow*. Each of those classes overrides ``run`` only
+to package its arguments into a start event and then calls ``super().run(...)``,
+so one wrap on ``workflows.workflow.Workflow.run`` captures all of them, plus any
+workflow a user writes by hand. ``instance`` is the agent itself.
+
+Two things make this hook's shape different from the others:
+
+* **``run()`` returns before the run happens.** It is synchronous: it schedules
+  an asyncio task and hands back a ``WorkflowHandler``, so a ``with`` block
+  around the call would close the span before the agent had done anything. The
+  run is opened with :func:`~otelzero.hooks.base.start_agent_run` instead and
+  closed from a done-callback on the handler's future -- which is safe from the
+  event loop, because ``send_log`` hands the POST to a daemon thread.
+* **The result carries no tokens.** ``AgentOutput`` has no usage and its ``raw``
+  holds only the last LLM response, so token counts come from LlamaIndex's own
+  instrumentation dispatcher: one root event handler sums the ``LLMChatEndEvent``
+  / ``LLMCompletionEndEvent`` payloads a run produces. Predict and
+  structured-predict events are deliberately *not* counted -- each wraps an inner
+  chat call that is already counted.
+
+With content capture opted in, a run records the user's turn, the agent's system
+prompt and its output. The output comes from the last chat the run made, whose
+history holds each tool call and result, and falls back to the final answer.
+
+A workflow started from inside a tool or step folds into the run it is part of
+(as a LangGraph subgraph folds into its parent), because the run state lives in a
+:class:`~contextvars.ContextVar` that the run task inherits.
+
+Known limitations: a human-in-the-loop run that is never resumed never closes; an
+OpenAI model streaming without ``stream_options={"include_usage": True}`` reports
+no tokens (``FunctionAgent`` streams by default); LLM events that arrive after a
+run has closed are dropped.
+"""
+
+import asyncio
+import contextvars
+import importlib
+import inspect
+import sys
+import threading
+from typing import Any, Optional, Tuple
+
+import wrapt
+
+from otelzero import semconv
+from otelzero.hooks import _content
+from otelzero.hooks._usage import UsageAccumulator, extract_usage
+from otelzero.hooks.base import _ImportInstrumenter, install_hook, start_agent_run
+
+# otelzero's own marker so a chokepoint is never wrapped twice -- ``Workflow`` is
+# re-exported from two watched modules, and two installs both see the import.
+_MARKER = "_otelzero_wrapped"
+
+# The run in progress on this context. A wrapper that finds a live one passes
+# straight through, which is how a workflow started from inside a tool folds into
+# the run it belongs to. The run task copies this context when it is created, so
+# the LLM events it dispatches find the same state.
+_active_run: contextvars.ContextVar = contextvars.ContextVar(
+    "otelzero_llama_index_run", default=None
+)
+
+# Strong references to futures otelzero had to create itself, so the loop does
+# not drop them mid-flight.
+_pending_futures: set = set()
+
+# Only these two events are counted. ``LLMPredictEndEvent`` and
+# ``LLMStructuredPredictEndEvent`` wrap an inner chat call that already emits one
+# of these, so counting them too would double every predict-based run.
+_COUNTED_EVENTS = ("LLMChatEndEvent", "LLMCompletionEndEvent")
+
+# Name the usage handler registers itself under, so a second install finds it.
+_HANDLER_NAME = "OtelzeroLlamaIndexUsageHandler"
+
+# ``run``'s own parameters, which are never the run's intent text.
+_NOT_INTENT = ("run_id", "ctx", "start_event", "stepwise", "checkpoint_callback")
+
+# LlamaIndex's placeholder when a model object cannot name itself.
+_UNKNOWN_MODEL = "unknown"
+
+
+def _attr_or_key(obj: Any, key: str) -> Any:
+    """Read ``key`` off a dict or an object attribute; ``None`` when absent."""
+    if obj is None:
+        return None
+    if isinstance(obj, dict):
+        return obj.get(key)
+    return getattr(obj, key, None)
+
+
+def _message_text(value: Any) -> Optional[str]:
+    """The text of a prompt or ``ChatMessage``, whichever shape it arrives in."""
+    if isinstance(value, str):
+        return value or None
+    content = _attr_or_key(value, "content")
+    return content if isinstance(content, str) and content else None
+
+
+# --- Reading the agent -------------------------------------------------------
+
+
+def _is_agent_workflow(instance: Any) -> bool:
+    """Whether this is the multi-agent ``AgentWorkflow``, duck-typed.
+
+    Importing ``AgentWorkflow`` to check would drag ``llama_index.core`` into the
+    import of a module that may only be ``workflows``, so the pair of attributes
+    that only it has is the test.
+    """
+    return getattr(instance, "agents", None) is not None and (
+        getattr(instance, "root_agent", None) is not None
+    )
+
+
+def _root_agent(instance: Any) -> Any:
+    """The agent an ``AgentWorkflow`` starts with; ``root_agent`` names it."""
+    root = getattr(instance, "root_agent", None)
+    agents = getattr(instance, "agents", None)
+    if isinstance(root, str) and isinstance(agents, dict):
+        return agents.get(root)
+    return root
+
+
+def _agent_name(instance: Any) -> str:
+    """Human label for the run: the agent's name, its root agent's, or its class."""
+    if _is_agent_workflow(instance):
+        root = getattr(instance, "root_agent", None)
+        if isinstance(root, str) and root:
+            return root
+        name = getattr(root, "name", None)
+        if isinstance(name, str) and name:
+            return name
+    name = getattr(instance, "name", None)
+    if isinstance(name, str) and name:
+        return name
+    return type(instance).__name__
+
+
+def _agent_names(instance: Any) -> list:
+    """The members of an ``AgentWorkflow``, for the run's ``details``."""
+    agents = getattr(instance, "agents", None)
+    if isinstance(agents, dict):
+        return sorted(str(name) for name in agents)
+    return []
+
+
+def _llama_index_model_id(llm: Any) -> Optional[str]:
+    """Best-effort model identifier for a LlamaIndex ``LLM``.
+
+    Most integrations carry the id on ``model`` (OpenAI, Anthropic) or
+    ``model_name`` (Hugging Face, Ollama), and LlamaIndex fills either with the
+    literal ``"unknown"`` when the class does not know -- which is not a model
+    name, so it is skipped. ``metadata`` is the general fallback but is a
+    computed property: OpenAI's raises for a model id it does not recognise, so
+    reading it cannot be allowed to break the run.
+    """
+    if llm is None:
+        return None
+    if isinstance(llm, str):
+        return llm or None
+    for attr in ("model", "model_name"):
+        value = getattr(llm, attr, None)
+        if isinstance(value, str) and value and value != _UNKNOWN_MODEL:
+            return value
+    try:
+        name = getattr(getattr(llm, "metadata", None), "model_name", None)
+    except Exception:  # pragma: no cover - defensive: metadata is computed
+        return None
+    if isinstance(name, str) and name and name != _UNKNOWN_MODEL:
+        return name
+    return None
+
+
+def _agent_llm(instance: Any) -> Any:
+    """The LLM a run will use: the agent's, or an ``AgentWorkflow``'s root one."""
+    llm = getattr(instance, "llm", None)
+    if llm is None and _is_agent_workflow(instance):
+        return getattr(_root_agent(instance), "llm", None)
+    return llm
+
+
+def _tool_names(instance: Any) -> list:
+    """Names of the tools the agent was given.
+
+    A LlamaIndex tool names itself on ``metadata.name``; a bare callable passed
+    to ``FunctionAgent(tools=[...])`` is wrapped into one by the time the agent
+    holds it, but an unfamiliar shape falls back to the function's own name
+    rather than erroring inside the wrapper.
+    """
+    try:
+        names = set()
+        for tool in getattr(instance, "tools", None) or []:
+            name = (
+                getattr(getattr(tool, "metadata", None), "name", None)
+                or getattr(tool, "name", None)
+                or getattr(tool, "__name__", None)
+            )
+            if name:
+                names.add(str(name))
+        return sorted(names)
+    except Exception:  # pragma: no cover - defensive
+        return []
+
+
+def _run_hint(args, kwargs) -> Optional[str]:
+    """Best-effort intent text for the session title.
+
+    An agent packs the user's turn into the start event before delegating to
+    ``Workflow.run``, so that is the first place to look; a hand-written workflow
+    is run with its own keyword arguments instead, and the first string among
+    them is the closest thing to an intent.
+    """
+    # ``run``'s positional parameters have moved between releases -- ``ctx`` and
+    # ``start_event`` have swapped places -- so every argument is offered to the
+    # same read rather than one position being assumed to be the start event.
+    for source in (kwargs.get("start_event"), *args, kwargs):
+        text = _message_text(_attr_or_key(source, "user_msg"))
+        if text:
+            return text
+        history = _attr_or_key(source, "chat_history")
+        if isinstance(history, (list, tuple)) and history:
+            text = _message_text(history[-1])
+            if text:
+                return text
+    for key, value in kwargs.items():
+        if key in _NOT_INTENT:
+            continue
+        if isinstance(value, str) and value.strip():
+            return value
+    return None
+
+
+def _system_prompt(instance: Any) -> Optional[str]:
+    """The agent's system prompt, or an ``AgentWorkflow`` root agent's."""
+    prompt = getattr(instance, "system_prompt", None)
+    if prompt is None and _is_agent_workflow(instance):
+        prompt = getattr(_root_agent(instance), "system_prompt", None)
+    return prompt if isinstance(prompt, str) else None
+
+
+# --- Message content ---------------------------------------------------------
+
+
+def _role_of(message: Any) -> Optional[str]:
+    role = _attr_or_key(message, "role")
+    return getattr(role, "value", role)
+
+
+def _chat_message(message: Any) -> Optional[dict]:
+    """One LlamaIndex ``ChatMessage`` as a semconv message.
+
+    Tool traffic is not on the fields ``_content`` reads: results name their
+    call in ``additional_kwargs``, and calls are in ``additional_kwargs`` or, on
+    newer provider integrations, in ``ToolCallBlock`` blocks.
+    """
+    extra = _attr_or_key(message, "additional_kwargs") or {}
+    text = _attr_or_key(message, "content")
+    role = _role_of(message)
+    if role == "tool":
+        call_id = extra.get("tool_call_id") or extra.get("tool_use_id")
+        return _content.message("tool", [_content.tool_call_response_part(call_id, text)])
+    parts = [_content.text_part(text)]
+    seen = set()
+    blocks = [
+        b
+        for b in _attr_or_key(message, "blocks") or []
+        if _attr_or_key(b, "block_type") == "tool_call"
+    ]
+    for call in [*blocks, *(extra.get("tool_calls") or [])]:
+        function = _attr_or_key(call, "function")
+        source = function if function is not None else call
+        call_id = (
+            _attr_or_key(call, "tool_call_id")
+            or _attr_or_key(call, "tool_id")
+            or _attr_or_key(call, "id")
+        )
+        if call_id is not None and call_id in seen:
+            continue
+        seen.add(call_id)
+        parts.append(
+            _content.tool_call_part(
+                call_id,
+                _attr_or_key(source, "tool_name") or _attr_or_key(source, "name"),
+                _attr_or_key(source, "tool_kwargs") or _attr_or_key(source, "arguments"),
+            )
+        )
+    return _content.message(role or "assistant", parts)
+
+
+def _output_messages(state: "_RunState", result: Any) -> Any:
+    """A run's output: everything after the user's last turn, else its answer.
+
+    ``None`` unless content capture is opted in. Reads the run's last chat, the
+    only place holding the tool hops; a workflow that never chatted, or an
+    unfamiliar message shape, reports the final answer instead.
+    """
+    if not _content.enabled():
+        return None
+    try:
+        history, response = state.last_chat
+        turns = list(history)
+        start = max((i for i, m in enumerate(turns) if _role_of(m) == "user"), default=-1)
+        shaped = [_chat_message(m) for m in [*turns[start + 1 :], response.message]]
+        messages = _content.cap_messages([m for m in shaped if m])
+    except Exception:
+        messages = None
+    return messages or _content.text_of(_attr_or_key(result, "response") or result)
+
+
+# --- Token usage, off LlamaIndex's instrumentation dispatcher ----------------
+
+
+def _response_usage(response: Any) -> Tuple[Any, Optional[dict]]:
+    """``(source, usage)`` for an LLM response -- the first spelling that hits.
+
+    Providers put usage in four different places and LlamaIndex passes the
+    payload through untouched, so all four are tried in order of specificity;
+    the last is OpenAI's flat ``prompt_tokens``/``completion_tokens`` sitting
+    directly in ``additional_kwargs``. The source object is returned so the
+    caller can dedupe on the payload the numbers came from.
+    """
+    raw = getattr(response, "raw", None)
+    extra = getattr(response, "additional_kwargs", None)
+    for source in (
+        _attr_or_key(raw, "usage"),
+        _attr_or_key(raw, "usage_metadata"),
+        _attr_or_key(extra, "usage"),
+        extra,
+    ):
+        usage = extract_usage(source)
+        if usage:
+            return source, usage
+    return None, None
+
+
+def _response_model(response: Any) -> Optional[str]:
+    """Which model served a response, as the provider reported it."""
+    raw = getattr(response, "raw", None)
+    for key in ("model", "model_name"):
+        value = _attr_or_key(raw, key)
+        if isinstance(value, str) and value:
+            return value
+    return None
+
+
+class _RunState:
+    """Everything a run accumulates between its start and its done-callback.
+
+    Lives in a :class:`~contextvars.ContextVar` rather than on the handler, so
+    the LLM events a run's task dispatches find it without the dispatcher having
+    to know which run it is serving. The lock is not optional: LlamaIndex runs a
+    synchronous LLM call in a worker thread, so its end event can be dispatched
+    off the loop thread while the loop is starting the next one.
+    """
+
+    __slots__ = (
+        "model",
+        "fallback_model",
+        "closed",
+        "last_chat",
+        "_acc",
+        "_seen",
+        "_refs",
+        "_lock",
+    )
+
+    def __init__(self, fallback_model: Optional[str] = None):
+        self.model: Optional[str] = None
+        self.fallback_model = fallback_model
+        self.closed = False
+        # (messages, response) of the run's latest chat, when content is opted in.
+        self.last_chat = None
+        self._acc = UsageAccumulator()
+        self._seen = set()
+        # Strong references to the payloads already counted: ids are only unique
+        # among live objects, so releasing one could let a later payload reuse
+        # its id and be skipped.
+        self._refs = []
+        self._lock = threading.Lock()
+
+    def record(self, response: Any) -> None:
+        """Fold one LLM response into the run, counting each payload once.
+
+        A ``CustomLLM`` builds ``chat()`` on top of ``complete()`` and emits an
+        end event for each, both carrying the *same* provider payload -- so the
+        payload's identity, not the event, is what decides whether the tokens
+        have already been counted.
+        """
+        source, usage = _response_usage(response)
+        raw = getattr(response, "raw", None)
+        fallback = source if source is not None else response
+        anchor = raw if raw is not None else fallback
+        model = _response_model(response)
+        with self._lock:
+            if self.closed or id(anchor) in self._seen:
+                return
+            self._seen.add(id(anchor))
+            self._refs.append(anchor)
+            self._acc.add(usage)
+            if self.model is None and model:
+                self.model = model
+
+    def close(self) -> Optional[dict]:
+        """Stop accepting events and hand back the run's total usage."""
+        with self._lock:
+            self.closed = True
+            return self._acc.result()
+
+
+def _record_llm_event(event: Any) -> None:
+    """Dispatcher callback: count an LLM end event against the run in progress.
+
+    Matched by class name so this module imports nothing from LlamaIndex, and
+    silent on anything unexpected -- the dispatcher calls its handlers without a
+    guard, so raising here would surface inside the user's agent.
+    """
+    try:
+        state = _active_run.get()
+        if state is None or state.closed:
+            return
+        if type(event).__name__ not in _COUNTED_EVENTS:
+            return
+        response = getattr(event, "response", None)
+        if response is not None:
+            state.record(response)
+            messages = getattr(event, "messages", None)
+            if messages is not None and _content.enabled():
+                state.last_chat = (messages, response)
+    except Exception:  # pragma: no cover - defensive
+        pass
+
+
+def _instrumentation_module():
+    """LlamaIndex's instrumentation package, whichever of the two is in play.
+
+    It was split out of ``llama_index.core`` into its own distribution, and the
+    old path is now a re-export of the new one. Whichever the host application
+    already imported is the one holding the live root dispatcher, so it is
+    preferred over importing the other.
+    """
+    for name in ("llama_index.core.instrumentation", "llama_index_instrumentation"):
+        module = sys.modules.get(name)
+        if module is not None:
+            return module
+    return importlib.import_module("llama_index_instrumentation")
+
+
+def _build_usage_handler(module):
+    """Subclass LlamaIndex's ``BaseEventHandler`` with otelzero's counter.
+
+    Built here rather than at module scope because the base class is a pydantic
+    model inside LlamaIndex, which otelzero must not import to be installed. The
+    subclass holds no state -- a pydantic model would need it declared as a field,
+    and the run state belongs to the run, not to the handler.
+    """
+    base = importlib.import_module(f"{module.__name__}.event_handlers").BaseEventHandler
+
+    class _UsageEventHandler(base):  # type: ignore[misc, valid-type]
+        @classmethod
+        def class_name(cls) -> str:
+            return _HANDLER_NAME
+
+        def handle(self, event, **kwargs) -> None:
+            _record_llm_event(event)
+
+    return _UsageEventHandler
+
+
+def _ensure_usage_handler() -> None:
+    """Register the usage handler on the root dispatcher, exactly once.
+
+    Deliberately *not* done at install time: the hook is instrumented from inside
+    another module's import, where importing LlamaIndex's instrumentation package
+    would re-enter the import machinery. Waiting for the first run also means an
+    application that never runs a workflow is never touched.
+    """
+    try:
+        module = _instrumentation_module()
+        dispatcher = module.get_dispatcher()
+        for existing in list(getattr(dispatcher, "event_handlers", None) or []):
+            namer = getattr(existing, "class_name", None)
+            if callable(namer) and namer() == _HANDLER_NAME:
+                return
+        dispatcher.add_event_handler(_build_usage_handler(module)())
+    except Exception:  # pragma: no cover - usage is best-effort
+        pass
+
+
+# --- Closing a deferred run --------------------------------------------------
+
+
+def _run_future(handler: Any):
+    """The future that completes when the run does, or ``None``.
+
+    ``WorkflowHandler`` *is* an ``asyncio.Future`` on older cores and wraps one
+    on current releases. Nothing else is expected, but an awaitable is still
+    driven rather than dropped -- with a strong reference held, since the loop
+    keeps only a weak one.
+    """
+    if isinstance(handler, asyncio.Future):
+        return handler
+    task = getattr(handler, "_result_task", None)
+    if isinstance(task, asyncio.Future):
+        return task
+    if inspect.isawaitable(handler):
+        task = asyncio.ensure_future(handler)
+        _pending_futures.add(task)
+        task.add_done_callback(_pending_futures.discard)
+        return task
+    return None
+
+
+def _is_user_cancellation(exc: BaseException) -> bool:
+    """Whether a run ended because the caller cancelled it, not because it broke.
+
+    Matched by class name up the MRO so the exception class does not have to be
+    imported (and so a version that moves it still matches).
+    """
+    return any(cls.__name__ == "WorkflowCancelledByUser" for cls in type(exc).__mro__)
+
+
+def _stopped_cleanly(handler: Any) -> bool:
+    """Whether a cancelled handler had already produced its stop event.
+
+    ``cancel_run()`` cancels the task, but so does the run finishing while
+    something else is still awaiting it; a stop event means the workflow reached
+    its end, so the run is reported as finished rather than interrupted.
+    """
+    getter = getattr(handler, "get_stop_event", None)
+    if not callable(getter):
+        return False
+    try:
+        return getter() is not None
+    except Exception:  # pragma: no cover - defensive
+        return False
+
+
+def _finish(run, state: _RunState, *, exc=None, interrupted: bool = False) -> None:
+    """Attach what the run accumulated, then close it (idempotently)."""
+    run.set_usage(state.close())
+    # The LLM's class name is a last resort: a model named before the run is what
+    # was asked for, and RunContext keeps that one.
+    run.set_model(state.model or state.fallback_model)
+    run.finish(exc=exc, interrupted=interrupted)
+
+
+def _attach_completion(handler: Any, run, state: _RunState) -> None:
+    """Close the run when its future completes, whichever way it completes.
+
+    A done-callback, not a watcher task: a task awaiting the handler would be
+    cancelled at ``asyncio.run`` teardown, which would both cancel the user's run
+    and mislabel a finished-but-undelivered one as interrupted. Registered here,
+    inside the wrapper, so it inherits the caller's ``run_context`` -- the
+    terminal event reads the user and session at send time.
+    """
+
+    def _done(future) -> None:
+        if future.cancelled():
+            _finish(run, state, interrupted=not _stopped_cleanly(handler))
+            return
+        exc = future.exception()
+        if exc is None:
+            run.set_output(_output_messages(state, future.result()))
+            _finish(run, state)
+        elif _is_user_cancellation(exc):
+            _finish(run, state, interrupted=True)
+        else:
+            _finish(run, state, exc=exc)
+
+    future = _run_future(handler)
+    if future is None:
+        # Nothing to wait on -- close now rather than leave the run dangling.
+        _finish(run, state)
+        return
+    try:
+        future.add_done_callback(_done)
+    except Exception:  # pragma: no cover - loop already gone
+        _finish(run, state)
+
+
+# --- The hook ----------------------------------------------------------------
+
+
+def wrap_workflow_module(module, telemetry) -> None:
+    """Instrument ``Workflow.run`` on an imported module, if it is LlamaIndex's.
+
+    ``workflows`` is a generic distribution name, so the class is identified by
+    the signature of the method being wrapped rather than by module name alone.
+    """
+    workflow_cls = getattr(module, "Workflow", None)
+    original = getattr(workflow_cls, "run", None)
+    if original is None or getattr(original, _MARKER, False):
+        return
+    try:
+        if "start_event" not in inspect.signature(original).parameters:
+            return
+    except (TypeError, ValueError):  # pragma: no cover - unsignable callable
+        return
+
+    @wrapt.decorator
+    def run_wrapper(wrapped, instance, args, kwargs):
+        active = _active_run.get()
+        if active is not None and not active.closed:
+            # A workflow started from inside a tool or step is part of the run
+            # already in progress, not a second agent run. A *closed* state means
+            # the parent is over, so a late background run opens its own.
+            return wrapped(*args, **kwargs)
+
+        _ensure_usage_handler()
+        agent_name = _agent_name(instance)
+        llm = _agent_llm(instance)
+        model = _llama_index_model_id(llm)
+        tools = _tool_names(instance)
+        details = {"model": model, "tools": tools}
+        if _is_agent_workflow(instance):
+            details["agents"] = _agent_names(instance)
+        run = start_agent_run(
+            telemetry,
+            semconv.span_name(semconv.OP_INVOKE_AGENT, agent_name),
+            {
+                semconv.GEN_AI_OPERATION_NAME: semconv.OP_INVOKE_AGENT,
+                semconv.GEN_AI_AGENT_NAME: agent_name,
+                semconv.GEN_AI_REQUEST_MODEL: model,
+                semconv.GEN_AI_TOOL_DEFINITIONS: semconv.tool_definitions(tools),
+            },
+            start_event="agent_start",
+            end_event="agent_end",
+            error_event="agent_error",
+            log_fields={
+                "agent_name": agent_name,
+                "agent_role": agent_name,
+                "model": model,
+            },
+            start_details=details,
+            session_name_hint=_run_hint(args, kwargs),
+        )
+        run.set_input(_content.messages_from(_run_hint(args, kwargs)))
+        run.set_system(_system_prompt(instance))
+        # The LLM's class name identifies a model object that cannot name itself
+        # (LlamaIndex's mocks, a hand-written LLM); kept aside as a last resort.
+        fallback = None if model or llm is None else type(llm).__name__
+        state = _RunState(fallback_model=fallback)
+        token = _active_run.set(state)
+        try:
+            # The span is current only around the call that schedules the run:
+            # that is where the task copies the context, which is what makes the
+            # framework's own spans -- and any nested workflow -- children of it.
+            with run.activate():
+                handler = wrapped(*args, **kwargs)
+        except BaseException as exc:
+            # run() itself failed (no running loop, bad arguments): there will be
+            # no run task, so this wrapper is the only place that can report it.
+            _finish(run, state, exc=exc)
+            raise
+        else:
+            _attach_completion(handler, run, state)
+            return handler
+        finally:
+            _active_run.reset(token)
+
+    wrapped = run_wrapper(original)
+    setattr(wrapped, _MARKER, True)
+    workflow_cls.run = wrapped
+
+
+class LlamaIndexImportHook(_ImportInstrumenter):
+    """Intercept LlamaIndex imports and instrument the workflow runtime.
+
+    Watches both homes of the ``Workflow`` class: ``workflows.workflow`` is where
+    it lives now that the engine ships as its own distribution, and
+    ``llama_index.core.workflow.workflow`` is where a pre-split core kept it. On
+    a current release the second is a re-export of the first, so the wrapped-once
+    marker makes watching it free.
+    """
+
+    WATCHED = {
+        "workflows.workflow": "_wrap_workflow",
+        "llama_index.core.workflow.workflow": "_wrap_workflow",
+    }
+
+    def _wrap_workflow(self, module):
+        try:
+            wrap_workflow_module(module, self.telemetry)
+        except Exception as exc:  # pragma: no cover - never break a user import
+            print(f"[INSTRUMENT] LlamaIndex instrumentation skipped: {exc}")
+
+
+def install_llama_index_monitoring(
+    server_url: str = "http://localhost:8000",
+    user: Optional[str] = None,
+    api_key: Optional[str] = None,
+    session: Optional[str] = None,
+    name: Optional[str] = None,
+):
+    """
+    Install LlamaIndex monitoring instrumentation.
+
+    Instruments ``workflows.workflow.Workflow.run`` -- the method every LlamaIndex
+    agent (``FunctionAgent``, ``ReActAgent``, ``CodeActAgent``, ``AgentWorkflow``)
+    and every hand-written ``Workflow`` starts through -- so each run streams
+    start/end/error telemetry with zero changes to the agent code. Token usage is
+    summed from LlamaIndex's own instrumentation events. A workflow started from
+    inside a tool is folded into the run it belongs to rather than reported as a
+    second agent. Call this *before* ``llama_index`` is imported.
+
+    Args:
+        server_url: URL of the monitoring server
+        user: Name of the executing user
+        api_key: Ingest key tying events to an owner (for the iOS app)
+        session: Per-run id used to group events in the app
+        name: Explicit session title (else derived from the run's first message)
+    """
+    return install_hook(
+        LlamaIndexImportHook,
+        framework="LlamaIndex",
+        label="LlamaIndex",
+        server_url=server_url,
+        user=user,
+        api_key=api_key,
+        session=session,
+        name=name,
+    )