OpenTelemetry span for each prompt and chain, with purpose

README.md

@@ -451,6 +451,18 @@ async def llm_default_model(datasette, purpose, actor):
     return None  # Use config defaults
 ```
 
+## OpenTelemetry
+
+Each `prompt()` and `chain()` call emits one `datasette_llm.prompt` span (scope `datasette_llm`). For a chain it covers the whole iteration and is current while you iterate, so `llm_prompt_context` hook work nests under it. If you stop iterating early, close the generator (e.g. `async with contextlib.aclosing(chain.responses()) as responses:`), otherwise the span stays current in your task and only ends, marked `datasette_llm.outcome="abandoned"`, when the generator is garbage-collected. Exceptions raised in your own loop body are not recorded on the span. Attributes:
+
+- `datasette_llm.kind`: `"prompt"` or `"chain"`
+- `datasette_llm.purpose`: the purpose passed to `llm.model()`, or `"none"`
+- `gen_ai.request.model`: the model ID
+- `datasette_llm.actor.present`: whether an actor was attached (never the actor itself)
+- `error.type`: the exception class name, if the call failed (`llm.PauseChain` is not an error)
+
+No prompt or response text, actor details or error messages are recorded. This plugin depends only on `opentelemetry-api`, so the span is a no-op until an OpenTelemetry SDK and tracer provider are configured. Versions of `llm` that include OpenTelemetry support add their own `chat` and `execute_tool` spans as children of this one.
+
 ## Related Plugins
 
 - **[datasette-secrets](https://github.com/datasette/datasette-secrets)**: Secure API key storage (required dependency)

datasette_llm/__init__.py

@@ -9,6 +9,7 @@
 from datasette_secrets import Secret, get_secret
 
 from . import hookspecs
+from .telemetry import prompt_span
 
 # Register our plugin hooks
 pm.add_hookspecs(hookspecs)
@@ -247,8 +248,12 @@ def __init__(
         context_factories,
         result: PromptResult,
         group: Optional[Group] = None,
+        model_id: Optional[str] = None,
+        actor: Optional[dict] = None,
     ):
         self._chain_response = chain_response
+        self._model_id = model_id
+        self._actor = actor
         self._context_factories = context_factories
         self._result = result
         self._result.is_chain = True
@@ -263,6 +268,11 @@ async def responses(self):
             self._prepared = True
 
         async with AsyncExitStack() as stack:
+            # Entered first so it is current for the hooks and the whole
+            # iteration, and sees any exception they raise
+            stack.enter_context(
+                prompt_span("chain", self._result.purpose, self._model_id, self._actor)
+            )
             for factory in self._context_factories:
                 if factory is not None:
                     ctx = factory(self._result)
@@ -346,6 +356,9 @@ async def prompt(self, prompt_text: str, **kwargs) -> llm_library.AsyncResponse:
 
         # Enter all context managers, then execute prompt, then exit
         async with AsyncExitStack() as stack:
+            stack.enter_context(
+                prompt_span("prompt", self._purpose, self._model_id, self._actor)
+            )
             for factory in context_factories:
                 if factory is not None:
                     ctx = factory(result)
@@ -389,6 +402,8 @@ def chain(self, prompt_text: str, **kwargs):
             context_factories=context_factories,
             result=result,
             group=self._group,
+            model_id=self._model_id,
+            actor=self._actor,
         )
 
 
@@ -472,6 +487,11 @@ async def prompt(self, prompt_text: str, **kwargs) -> llm_library.AsyncResponse:
 
         # Enter all context managers, then execute prompt, then exit
         async with AsyncExitStack() as stack:
+            # prompt() awaits the response, forcing the whole model call,
+            # so this span covers the hooks and the call
+            stack.enter_context(
+                prompt_span("prompt", self._purpose, self._model.model_id, self._actor)
+            )
             for factory in context_factories:
                 if factory is not None:
                     # Each factory takes the result object and returns a context manager

datasette_llm/telemetry.py

@@ -0,0 +1,75 @@
+"""OpenTelemetry for datasette-llm: a span per prompt()/chain(), API-only."""
+
+from contextlib import contextmanager
+import contextvars
+from importlib.metadata import PackageNotFoundError, version
+
+import llm as llm_library
+from opentelemetry import context, trace
+from opentelemetry.trace import SpanKind, Status, StatusCode
+
+try:
+    _version = version("datasette-llm")
+except PackageNotFoundError:  # pragma: no cover
+    _version = None
+
+tracer = trace.get_tracer("datasette_llm", _version)
+
+SPAN_PROMPT = "datasette_llm.prompt"
+KIND = "datasette_llm.kind"
+PURPOSE = "datasette_llm.purpose"
+ACTOR_PRESENT = "datasette_llm.actor.present"
+GEN_AI_REQUEST_MODEL = "gen_ai.request.model"
+ERROR_TYPE = "error.type"
+OUTCOME = "datasette_llm.outcome"
+NO_PURPOSE = "none"
+# llm.PauseChain (llm>=0.32a3) suspends a turn: an outcome, not an error
+_NOT_ERRORS = getattr(llm_library, "PauseChain", ())
+# Tells exit in our own Context apart from an abandoned chain's finalizer
+_owner = contextvars.ContextVar("datasette_llm_prompt_span")
+
+
+@contextmanager
+def prompt_span(kind, purpose, model_id, actor):
+    "Current datasette_llm.prompt span; no text, actor ids or error messages."
+    span = tracer.start_span(
+        SPAN_PROMPT,
+        kind=SpanKind.INTERNAL,
+        attributes={
+            KIND: kind,
+            PURPOSE: NO_PURPOSE if purpose is None else purpose,
+            GEN_AI_REQUEST_MODEL: model_id,
+            ACTOR_PRESENT: actor is not None,
+        },
+    )
+    if not span.is_recording():
+        # No SDK: don't attach a context an abandoned chain could leak
+        yield
+        return
+    previous = context.get_current()
+    marker = _owner.set(None)
+    context.attach(trace.set_span_in_context(span, previous))
+    try:
+        yield
+    except GeneratorExit:
+        # The consumer stopped iterating a chain early: not an error
+        span.set_attribute(OUTCOME, "abandoned")
+        raise
+    except _NOT_ERRORS:
+        raise
+    except Exception as ex:
+        span.set_attribute(ERROR_TYPE, type(ex).__qualname__)
+        span.set_status(Status(StatusCode.ERROR))
+        raise
+    finally:
+        # Not context.detach(), which logs an error when an abandoned chain
+        # is finalized in another contextvars.Context. In our own Context,
+        # restore the previous context (even over one a hook leaked); in a
+        # finalizer's Context there is nothing of ours to restore.
+        try:
+            _owner.reset(marker)
+        except ValueError:
+            pass
+        else:
+            context.attach(previous)
+        span.end()

pyproject.toml

@@ -12,14 +12,16 @@ requires-python = ">=3.10"
 dependencies = [
     "datasette>=1.0a23",
     "llm>=0.32a1",
-    "datasette-secrets"
+    "datasette-secrets",
+    "opentelemetry-api>=1.37",
 ]
 
 [dependency-groups]
 dev = [
     "pytest",
     "pytest-asyncio",
     "llm-echo>=0.5a0",
+    "opentelemetry-sdk>=1.37",
 ]
 
 [build-system]

tests/test_telemetry.py

@@ -0,0 +1,176 @@
+import asyncio
+from contextlib import asynccontextmanager, contextmanager
+import gc
+import json
+
+from datasette import hookimpl
+from datasette.app import Datasette
+from datasette.plugins import pm
+import llm
+from opentelemetry import context, trace
+from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.sdk.trace.export import SimpleSpanProcessor
+from opentelemetry.sdk.trace.export.in_memory_span_exporter import (
+    InMemorySpanExporter,
+)
+from opentelemetry.trace import StatusCode
+import pytest
+
+from datasette_llm import LLM, telemetry
+
+TOOL_PROMPT = json.dumps(
+    {"tool_calls": [{"name": "example", "arguments": {"input": "x"}}], "prompt": "hi"}
+)
+
+
+def example(input: str) -> str:
+    return f"Example output for {input}"
+
+
+@pytest.fixture(scope="session")
+def _exporter():
+    # The global tracer provider can only be set once per process
+    exporter = InMemorySpanExporter()
+    provider = TracerProvider()
+    provider.add_span_processor(SimpleSpanProcessor(exporter))
+    trace.set_tracer_provider(provider)
+    return exporter
+
+
+@pytest.fixture
+def exporter(_exporter):
+    _exporter.clear()
+    return _exporter
+
+
+def spans(exporter, name="datasette_llm.prompt"):
+    return [s for s in exporter.get_finished_spans() if s.name == name]
+
+
+@contextmanager
+def prompt_hook(on_enter):
+    class Plugin:
+        __name__ = "test_telemetry_hook"
+
+        @hookimpl
+        def llm_prompt_context(self, datasette, model_id, prompt, purpose):
+            @asynccontextmanager
+            async def ctx(result):
+                on_enter()
+                yield
+
+            return ctx
+
+    plugin = Plugin()
+    pm.register(plugin)
+    try:
+        yield
+    finally:
+        pm.unregister(plugin)
+
+
+def open_hook_span():
+    trace.get_tracer("test").start_span("hook").end()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "conversation,purpose,actor", [(False, "agent", None), (True, None, {"id": "a1"})]
+)
+@pytest.mark.parametrize("kind", ["prompt", "chain"])
+async def test_prompt_and_chain_spans(exporter, conversation, purpose, actor, kind):
+    llm_ = LLM(Datasette(memory=True))
+    model = await llm_.model("echo", purpose=purpose, actor=actor)
+    target = model.conversation() if conversation else model
+    current = []
+    with prompt_hook(open_hook_span):
+        if kind == "prompt":
+            await (await target.prompt("Hello")).text()
+        else:
+            chain = target.chain(TOOL_PROMPT, tools=[example])
+            async for response in chain.responses():
+                await response.text()
+                current.append(trace.get_current_span().get_span_context())
+
+    (span,) = spans(exporter)
+    # Exactly these attributes: no actor id, no text
+    assert dict(span.attributes) == {
+        "datasette_llm.kind": kind,
+        "datasette_llm.purpose": purpose or "none",
+        "datasette_llm.actor.present": actor is not None,
+        "gen_ai.request.model": "echo",
+    }
+    assert span.status.status_code == StatusCode.UNSET
+    # Hook work nests under the span; a chain's span is current while iterating
+    (hook,) = spans(exporter, "hook")
+    assert hook.parent.span_id == span.context.span_id
+    if kind == "chain":
+        assert current == [span.context] * 2
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("close", [True, False])
+async def test_chain_abandoned(exporter, caplog, close):
+    model = await LLM(Datasette(memory=True)).model("echo")
+    gen = model.chain(TOOL_PROMPT, tools=[example]).responses()
+    async for response in gen:
+        break
+    if close:
+        await gen.aclose()
+        assert not trace.get_current_span().get_span_context().is_valid
+    else:
+        # asyncio finalizes the generator in another task and Context; the
+        # span stays current in this task's context, by design (see README)
+        del gen
+        gc.collect()
+        await asyncio.sleep(0.01)
+    (span,) = spans(exporter)
+    assert span.status.status_code == StatusCode.UNSET
+    assert span.attributes["datasette_llm.outcome"] == "abandoned"
+    assert "Failed to detach context" not in caplog.text
+
+
+@pytest.mark.asyncio
+async def test_hook_leaked_context_does_not_reach_caller(exporter):
+    leaked = trace.get_tracer("test").start_span("leaked")
+    attach = lambda: context.attach(trace.set_span_in_context(leaked))
+    model = await LLM(Datasette(memory=True)).model("echo")
+    with prompt_hook(attach):
+        await model.prompt("Hello")
+    assert trace.get_current_span() is not leaked
+
+
+@pytest.mark.asyncio
+async def test_no_sdk_attaches_nothing(exporter, monkeypatch):
+    monkeypatch.setattr(telemetry, "tracer", trace.NoOpTracer())
+    model = await LLM(Datasette(memory=True)).model("echo")
+    before = context.get_current()
+    async for response in model.chain(TOOL_PROMPT, tools=[example]).responses():
+        assert context.get_current() == before
+    assert spans(exporter) == []
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("kind", ["prompt", "chain"])
+@pytest.mark.parametrize("exception", [ValueError, llm.PauseChain])
+async def test_hook_exception(exporter, kind, exception):
+    def boom():
+        raise exception("secret message")
+
+    model = await LLM(Datasette(memory=True)).model("echo", purpose="agent")
+    with prompt_hook(boom), pytest.raises(exception):
+        if kind == "prompt":
+            await model.prompt("Hello")
+        else:
+            async for response in model.chain("Hello").responses():
+                pass
+
+    (span,) = spans(exporter)
+    assert not span.events and not span.status.description
+    if exception is ValueError:
+        assert span.attributes["error.type"] == "ValueError"
+        assert span.status.status_code == StatusCode.ERROR
+    else:
+        # PauseChain suspends a turn; it is not an error
+        assert "error.type" not in span.attributes
+        assert span.status.status_code == StatusCode.UNSET