langchain: link chat models to the prompt that built their input

instrumentation/opentelemetry-instrumentation-genai-langchain/.changelog/939.added

@@ -0,0 +1 @@
+Link chat spans to the prompt template that produced their input when the prompt runs outside a chain or several prompts share a parent, set `gen_ai.prompt.version` from `prompt_version` or LangChain Hub metadata, and detect every `BasePromptTemplate` subclass

instrumentation/opentelemetry-instrumentation-genai-langchain/src/opentelemetry/instrumentation/genai/langchain/_execution_context.py

@@ -18,6 +18,9 @@
 from wrapt import wrap_function_wrapper
 
 from opentelemetry.context import attach, detach, get_current
+from opentelemetry.instrumentation.genai.langchain._prompt_input import (
+    _wrap_prompt_input,
+)
 from opentelemetry.instrumentation.genai.langchain._run_context import (
     _astart_run,
     _RunScope,
@@ -230,6 +233,26 @@ def _patch(
             )
 
     def instrument(self) -> None:
+        # Innermost on the streams, so the stream boundary below chooses the
+        # run id first and the prompt input reuses it.
+        for method, stream, asynchronous in (
+            ("invoke", False, False),
+            ("ainvoke", False, True),
+            ("stream", True, False),
+            ("astream", True, True),
+        ):
+            self._patch(
+                "langchain_core.language_models.chat_models",
+                "BaseChatModel",
+                method,
+                partial(
+                    _wrap_prompt_input,
+                    invocations=self._invocations,
+                    stream=stream,
+                    asynchronous=asynchronous,
+                ),
+            )
+
         for module_name, class_name, method in _METHODS:
             self._patch(module_name, class_name, method, self._wrap)
 

instrumentation/opentelemetry-instrumentation-genai-langchain/src/opentelemetry/instrumentation/genai/langchain/_prompt_input.py

@@ -0,0 +1,132 @@
+# Copyright The OpenTelemetry Authors
+# SPDX-License-Identifier: Apache-2.0
+
+"""Link a chat model call to the prompt template that produced its input.
+
+The run tree alone cannot always tell: a prompt formatted outside any chain
+has no parent run to publish to, and two prompts formatted under one parent
+reach their models in any order. The callback handler records the prompt value
+each prompt run returns; these wrappers find that value in a chat model's
+input and record its template under the run id chosen here for the call.
+"""
+
+from __future__ import annotations
+
+from collections.abc import AsyncIterator, Callable, Iterator
+from functools import wraps
+from typing import Any
+from uuid import UUID
+
+from langchain_core.prompt_values import PromptValue
+
+from opentelemetry.instrumentation.genai.langchain._run_context import (
+    _new_run_id,
+    _signature,
+)
+from opentelemetry.instrumentation.genai.langchain.invocation_manager import (
+    _InvocationManager,
+)
+
+__all__ = ["_wrap_prompt_input"]
+
+_Call = tuple[UUID | None, tuple[Any, ...], dict[str, Any]]
+
+
+def _wrap_prompt_input(
+    original: Callable[..., Any],
+    invocations: _InvocationManager,
+    stream: bool,
+    asynchronous: bool,
+) -> Callable[..., Any]:
+    parameters = _signature(original)
+
+    def prepare(
+        instance: Any, args: tuple[Any, ...], kwargs: dict[str, Any]
+    ) -> _Call:
+        bound = parameters.bind(instance, *args, **kwargs)
+        value = bound.arguments.get("input")
+        prompt_context = (
+            invocations.get_prompt_value_context(value)
+            if isinstance(value, PromptValue)
+            else None
+        )
+        if prompt_context is None:
+            return None, args, kwargs
+        # The model starts its run with this id, so the handler finds the
+        # template when the run starts.
+        config = dict(bound.arguments.get("config") or {})
+        run_id: UUID = config.get("run_id") or _new_run_id()
+        config["run_id"] = run_id
+        bound.arguments["config"] = config
+        invocations.set_model_prompt_context(run_id, prompt_context)
+        return run_id, bound.args[1:], bound.kwargs
+
+    def forget(run_id: UUID | None) -> None:
+        # The handler takes the template when the run starts; this drops it
+        # when the run never started.
+        if run_id is not None:
+            invocations.pop_model_prompt_context(run_id)
+
+    @wraps(original)
+    def invoke(
+        wrapped: Callable[..., Any],
+        instance: Any,
+        args: tuple[Any, ...],
+        kwargs: dict[str, Any],
+    ) -> Any:
+        run_id, args, kwargs = prepare(instance, args, kwargs)
+        try:
+            return wrapped(*args, **kwargs)
+        finally:
+            forget(run_id)
+
+    @wraps(original)
+    async def ainvoke(
+        wrapped: Callable[..., Any],
+        instance: Any,
+        args: tuple[Any, ...],
+        kwargs: dict[str, Any],
+    ) -> Any:
+        run_id, args, kwargs = prepare(instance, args, kwargs)
+        try:
+            return await wrapped(*args, **kwargs)
+        finally:
+            forget(run_id)
+
+    # The stream wrappers are generators, so the input is read on the first
+    # advancement, like the model's own stream, and an unread stream records
+    # nothing.
+    @wraps(original)
+    def sync_stream(
+        wrapped: Callable[..., Any],
+        instance: Any,
+        args: tuple[Any, ...],
+        kwargs: dict[str, Any],
+    ) -> Iterator[Any]:
+        run_id, args, kwargs = prepare(instance, args, kwargs)
+        try:
+            yield from wrapped(*args, **kwargs)
+        finally:
+            forget(run_id)
+
+    @wraps(original)
+    async def async_stream(
+        wrapped: Callable[..., Any],
+        instance: Any,
+        args: tuple[Any, ...],
+        kwargs: dict[str, Any],
+    ) -> AsyncIterator[Any]:
+        run_id, args, kwargs = prepare(instance, args, kwargs)
+        chunks = wrapped(*args, **kwargs)
+        try:
+            async for chunk in chunks:
+                yield chunk
+        finally:
+            forget(run_id)
+            close = getattr(chunks, "aclose", None)
+            if close is not None:
+                await close()
+
+    if stream:
+        return async_stream if asynchronous else sync_stream
+    return ainvoke if asynchronous else invoke

instrumentation/opentelemetry-instrumentation-genai-langchain/src/opentelemetry/instrumentation/genai/langchain/callback_handler.py

@@ -16,6 +16,7 @@
     GenerationChunk,
     LLMResult,
 )
+from langchain_core.prompt_values import PromptValue
 
 from opentelemetry.instrumentation.genai.langchain.agent_context import (
     claim_agent,
@@ -75,14 +76,17 @@
 _PROMPT_TEMPLATE_TYPES = {"ChatPromptTemplate", "PromptTemplate"}
 
 
-def _prompt_template_type(runnable: Mapping[str, Any]) -> str | None:
+def _class_name(runnable: Mapping[str, Any]) -> str | None:
     runnable_id = runnable.get("id")
     if not isinstance(runnable_id, Sequence) or isinstance(runnable_id, str):
         return None
-    runnable_id = cast(Sequence[str], runnable_id)
-    template_type = runnable_id[-1] if runnable_id else None
-    if not isinstance(template_type, str):
-        return None
+    runnable_id = cast(Sequence[object], runnable_id)
+    class_name = runnable_id[-1] if runnable_id else None
+    return class_name if isinstance(class_name, str) else None
+
+
+def _prompt_template_type(runnable: Mapping[str, Any]) -> str | None:
+    template_type = _class_name(runnable)
     return template_type if template_type in _PROMPT_TEMPLATE_TYPES else None
 
 
@@ -94,15 +98,26 @@ def _string_sequence(value: object) -> list[str]:
     ]
 
 
+def _string_metadata(
+    metadata: Mapping[str, Any] | None, *keys: str
+) -> str | None:
+    for key in keys:
+        value = (metadata or {}).get(key)
+        if isinstance(value, (str, int)) and not isinstance(value, bool):
+            return str(value)
+    return None
+
+
 def _extract_prompt_context(
     runnable: Mapping[str, Any],
     inputs: object,
     metadata: Mapping[str, Any] | None,
     *,
     capture_variables: bool = True,
+    run_type: str | None = None,
 ) -> _PromptContext | None:
-    template_type = _prompt_template_type(runnable)
-    if template_type is None:
+    # LangChain reports every BasePromptTemplate subclass with run_type "prompt".
+    if _prompt_template_type(runnable) is None and run_type != "prompt":
         return None
 
     raw_template_kwargs = runnable.get("kwargs")
@@ -114,17 +129,27 @@ def _extract_prompt_context(
     name = (metadata or {}).get("prompt_name")
     if not isinstance(name, str):
         name = template_kwargs.get("name")
+    if not isinstance(name, str):
+        # Prompts pulled from LangChain Hub carry their repo and commit.
+        name = _string_metadata(metadata, "lc_hub_repo")
     if not isinstance(name, str):
         runnable_name = runnable.get("name")
+        # An unnamed template reports its class name, which names nothing.
         name = (
             runnable_name
             if isinstance(runnable_name, str)
             and runnable_name not in _PROMPT_TEMPLATE_TYPES
+            and runnable_name != _class_name(runnable)
             else None
         )
+    version = (
+        _string_metadata(metadata, "prompt_version", "lc_hub_commit_hash")
+        if name is not None
+        else None
+    )
 
     if not capture_variables:
-        return _PromptContext(name=name, variables={})
+        return _PromptContext(name=name, variables={}, version=version)
 
     required_variables = _string_sequence(
         template_kwargs.get("input_variables")
@@ -172,7 +197,9 @@ def _extract_prompt_context(
                 continue
             effective_variables[key] = value
 
-    return _PromptContext(name=name, variables=effective_variables)
+    return _PromptContext(
+        name=name, variables=effective_variables, version=version
+    )
 
 
 def _conversation_id(metadata: dict[str, Any] | None) -> str | None:
@@ -410,6 +437,7 @@ def on_chain_start(
             inputs,
             metadata,
             capture_variables=capture_content,
+            run_type=kwargs.get("run_type"),
         ):
             self._invocation_manager.set_prompt_context(run_id, prompt_context)
 
@@ -421,7 +449,14 @@ def on_chain_end(
         parent_run_id: UUID | None = None,
         **kwargs: Any,
     ) -> Any:
-        self._invocation_manager.publish_prompt_context(run_id)
+        prompt_context = self._invocation_manager.publish_prompt_context(
+            run_id
+        )
+        if prompt_context is not None and isinstance(outputs, PromptValue):
+            # The value carries its template to whichever model it reaches.
+            self._invocation_manager.set_prompt_value_context(
+                outputs, prompt_context
+            )
         invocation = self._invocation_manager.get_invocation(run_id=run_id)
         if invocation is None or not isinstance(
             invocation, (WorkflowInvocation, LocalAgentInvocation)
@@ -553,8 +588,14 @@ def on_chat_model_start(
         parent_context = self._invocation_manager.get_parent_context(
             parent_run_id
         )
-        prompt_context = self._invocation_manager.consume_prompt_context(
-            parent_run_id
+        # Consume the sibling prompt either way, so a later model call under the
+        # same parent cannot take it; the model's own input wins when known.
+        sibling_prompt_context = (
+            self._invocation_manager.consume_prompt_context(parent_run_id)
+        )
+        prompt_context = (
+            self._invocation_manager.pop_model_prompt_context(run_id)
+            or sibling_prompt_context
         )
         llm_invocation = self._telemetry_handler.inference(
             provider,
@@ -565,6 +606,7 @@ def on_chat_model_start(
         )
         if prompt_context is not None:
             llm_invocation.prompt_name = prompt_context.name
+            llm_invocation.prompt_version = prompt_context.version
             llm_invocation.prompt_variables = prompt_context.variables
         llm_invocation.input_messages = input_messages
         llm_invocation.top_p = top_p