[opentelemetry-instrumentation-genai-dspy] Add instrumentation for inference invocations

instrumentation/opentelemetry-instrumentation-genai-dspy/.changelog/683.added

@@ -0,0 +1 @@
+Add instrumentation for ``dspy.LM``.

instrumentation/opentelemetry-instrumentation-genai-dspy/README.rst

@@ -8,9 +8,9 @@ OpenTelemetry DSPy Instrumentation
 
 This package provides OpenTelemetry instrumentation for the
 `DSPy framework <https://pypi.org/project/dspy/>`_, emitting Generative AI
-semantic conventions for DSPy Tool executions, ReAct agent loops
-(supporting both ``dspy.ReAct`` and ``dspy.ReActV2``), and ``dspy.Retrieve``
-retrieval operations.
+semantic conventions for DSPy language models (``dspy.LM``), Tool executions,
+ReAct agent loops (supporting both ``dspy.ReAct`` and ``dspy.ReActV2``), and
+``dspy.Retrieve`` retrieval operations.
 
 Installation
 ------------

instrumentation/opentelemetry-instrumentation-genai-dspy/src/opentelemetry/instrumentation/genai/dspy/patch.py

@@ -7,7 +7,8 @@
 
 import inspect
 import sys
-from collections.abc import Awaitable, Callable, Sequence
+from collections.abc import Awaitable, Callable, Mapping, Sequence
+from contextvars import ContextVar
 from copy import copy, deepcopy
 from importlib import import_module
 from typing import TYPE_CHECKING, Any, cast
@@ -21,13 +22,22 @@
 
 from opentelemetry.instrumentation.genai.dspy.utils import (
     SENTINEL_TOOL_NAMES,
+    _safe_float,
+    _safe_int,
+    _safe_stop_sequences,
+    apply_usage_to_invocation,
     extract_input_content,
+    extract_lm_input_messages,
+    extract_lm_output_messages,
     extract_output_content,
     prepare_tool_definitions,
+    resolve_provider,
+    resolve_request_model,
 )
 from opentelemetry.instrumentation.utils import unwrap
 from opentelemetry.util.genai.handler import TelemetryHandler
 from opentelemetry.util.genai.invocation import (
+    InferenceInvocation,
     LocalAgentInvocation,
     RetrievalInvocation,
     ToolInvocation,
@@ -40,6 +50,8 @@
 
 if TYPE_CHECKING:
     from dspy.adapters.types.tool import Tool
+    from dspy.clients.lm import LM
+    from dspy.core.types import LMResponse
     from dspy.primitives.module import Module
     from dspy.primitives.prediction import Prediction
     from dspy.retrievers.retrieve import Retrieve
@@ -50,6 +62,10 @@
 _REACT_V2_MODULE = "dspy.predict.react_v2"
 _REACT_V2_CLASS = "ReActV2"
 
+_current_lm_history_entry: ContextVar[Any] = ContextVar(
+    "_current_lm_history_entry", default=None
+)
+
 
 if TYPE_CHECKING:
     _BoundFunctionWrapper = BoundFunctionWrapper[Any, Any]
@@ -217,6 +233,26 @@ def patch_dspy(handler: TelemetryHandler) -> None:
                 _react_aforward(handler, "dspy.ReActV2"),
             )
 
+    if hasattr(dspy, "LM"):
+        lm_module = dspy.LM.__module__
+        lm_name = dspy.LM.__name__
+        _wrap_function(
+            lm_module,
+            f"{lm_name}.__call__",
+            _lm_call(handler),
+        )
+        _wrap_function(
+            lm_module,
+            f"{lm_name}.acall",
+            _lm_acall(handler),
+        )
+        if hasattr(dspy.LM, "update_history"):
+            _wrap_function(
+                lm_module,
+                f"{lm_name}.update_history",
+                _lm_update_history(),
+            )
+
 
 def unpatch_dspy() -> None:
     """Remove patches from DSPy classes."""
@@ -238,6 +274,226 @@ def unpatch_dspy() -> None:
         unwrap(react_v2_cls, "forward")
         unwrap(react_v2_cls, "aforward")
 
+    if hasattr(dspy, "LM"):
+        unwrap(dspy.LM, "__call__")
+        unwrap(dspy.LM, "acall")
+        if hasattr(dspy.LM, "update_history"):
+            unwrap(dspy.LM, "update_history")
+
+
+def _start_lm_invocation(
+    handler: TelemetryHandler,
+    instance: LM,
+    args: tuple[Any, ...],
+    kwargs: dict[str, Any],
+) -> InferenceInvocation:
+    provider = resolve_provider(instance)
+    request_model = resolve_request_model(instance)
+
+    invocation = handler.inference(
+        provider=provider,
+        request_model=request_model,
+    )
+
+    merged_kwargs = {**getattr(instance, "kwargs", {}), **kwargs}
+
+    invocation.temperature = _safe_float(merged_kwargs.get("temperature"))
+    invocation.max_tokens = _safe_int(merged_kwargs.get("max_tokens"))
+    invocation.top_p = _safe_float(
+        merged_kwargs.get("top_p")
+        if merged_kwargs.get("top_p") is not None
+        else merged_kwargs.get("p")
+    )
+    invocation.frequency_penalty = _safe_float(
+        merged_kwargs.get("frequency_penalty")
+    )
+    invocation.presence_penalty = _safe_float(
+        merged_kwargs.get("presence_penalty")
+    )
+    invocation.seed = _safe_int(merged_kwargs.get("seed"))
+
+    invocation.stop_sequences = _safe_stop_sequences(merged_kwargs.get("stop"))
+
+    choice_count = _safe_int(merged_kwargs.get("n"))
+    if choice_count is not None and choice_count != 1:
+        invocation.request_choice_count = choice_count
+
+    if handler.should_capture_content():
+        invocation.input_messages = extract_lm_input_messages(args, kwargs)
+
+    return invocation
+
+
+def _get_field(obj: Any, key: str) -> Any:
+    if isinstance(obj, Mapping):
+        return cast(Mapping[str, Any], obj).get(key)
+    return getattr(obj, key, None)
+
+
+def _lm_update_history() -> Callable[..., Any]:
+    """Capture the history entry in context to isolate concurrent calls from shared instance.history."""
+
+    def traced_method(
+        wrapped: Callable[..., Any],
+        instance: LM,
+        args: tuple[Any, ...],
+        kwargs: dict[str, Any],
+    ) -> Any:
+        entry = args[0] if args else kwargs.get("entry")
+        if entry is not None:
+            _current_lm_history_entry.set(entry)
+        return wrapped(*args, **kwargs)
+
+    return traced_method
+
+
+def _set_lm_invocation_response(
+    handler: TelemetryHandler,
+    invocation: InferenceInvocation,
+    instance: LM,
+    result: LMResponse | list[dict[str, Any] | str],
+    history_entry: Any = None,
+) -> None:
+    if not isinstance(result, list):
+        if result.model:
+            invocation.response_model_name = str(result.model)
+        if result.response_id:
+            invocation.response_id = str(result.response_id)
+
+        usage_dict = result.usage_as_dict()
+        if usage_dict:
+            apply_usage_to_invocation(invocation, usage_dict)
+
+        finish_reasons = [
+            out.finish_reason for out in result.outputs if out.finish_reason
+        ]
+        if finish_reasons:
+            invocation.finish_reasons = finish_reasons
+
+        if handler.should_capture_content():
+            invocation.output_messages = extract_lm_output_messages(result)
+        return
+
+    # DSPy 3.x LM calls return a legacy list by default unless experimental=True
+    # or an LMRequest is used. When available, use history_entry captured from
+    # LM.update_history for per-call isolation; fall back to instance.history[-1].
+    entry = history_entry
+    if entry is None:
+        history: Sequence[Mapping[str, Any]] | None = getattr(
+            instance, "history", None
+        )
+        if isinstance(history, Sequence) and history:
+            entry = history[-1]
+
+    choice_finish_reasons: list[str | None] | None = None
+    if entry is not None:
+        resp_model = _get_field(entry, "response_model") or _get_field(
+            entry, "model"
+        )
+        if resp_model:
+            invocation.response_model_name = str(resp_model)
+
+        usage = _get_field(entry, "usage")
+        if isinstance(usage, Mapping):
+            apply_usage_to_invocation(
+                invocation, cast(Mapping[str, Any], usage)
+            )
+
+        resp_obj: Any = _get_field(entry, "response")
+        if resp_obj is not None:
+            resp_id = _get_field(resp_obj, "id")
+            if resp_id:
+                invocation.response_id = str(resp_id)
+
+            choices = _get_field(resp_obj, "choices")
+            if isinstance(choices, Sequence) and choices:
+                choice_finish_reasons = []
+                for choice in cast(Sequence[object], choices):
+                    fr = _get_field(choice, "finish_reason")
+                    choice_finish_reasons.append(
+                        str(fr) if fr is not None else None
+                    )
+                invocation_finish_reasons = [
+                    fr for fr in choice_finish_reasons if fr is not None
+                ]
+                if invocation_finish_reasons:
+                    invocation.finish_reasons = invocation_finish_reasons
+            else:
+                outputs = _get_field(resp_obj, "outputs")
+                if isinstance(outputs, Sequence) and outputs:
+                    choice_finish_reasons = []
+                    for out in cast(Sequence[object], outputs):
+                        fr = _get_field(out, "finish_reason")
+                        choice_finish_reasons.append(
+                            str(fr) if fr is not None else None
+                        )
+                    invocation_finish_reasons = [
+                        fr for fr in choice_finish_reasons if fr is not None
+                    ]
+                    if invocation_finish_reasons:
+                        invocation.finish_reasons = invocation_finish_reasons
+
+    if handler.should_capture_content():
+        invocation.output_messages = extract_lm_output_messages(
+            result, finish_reasons=choice_finish_reasons
+        )
+
+
+def _lm_call(handler: TelemetryHandler) -> Callable[..., Any]:
+    def traced_method(
+        wrapped: Callable[..., Any],
+        instance: LM,
+        args: tuple[Any, ...],
+        kwargs: dict[str, Any],
+    ) -> Any:
+        invocation = _start_lm_invocation(handler, instance, args, kwargs)
+        # Isolate per-call history metadata from concurrent calls on the shared LM instance.
+        token = _current_lm_history_entry.set(None)
+        try:
+            with invocation:
+                result = wrapped(*args, **kwargs)
+                history_entry = _current_lm_history_entry.get()
+                _set_lm_invocation_response(
+                    handler,
+                    invocation,
+                    instance,
+                    result,
+                    history_entry=history_entry,
+                )
+                return result
+        finally:
+            _current_lm_history_entry.reset(token)
+
+    return traced_method
+
+
+def _lm_acall(handler: TelemetryHandler) -> Callable[..., Any]:
+    async def traced_method(
+        wrapped: Callable[..., Awaitable[Any]],
+        instance: LM,
+        args: tuple[Any, ...],
+        kwargs: dict[str, Any],
+    ) -> Any:
+        invocation = _start_lm_invocation(handler, instance, args, kwargs)
+        # Isolate per-call history metadata from concurrent calls on the shared LM instance.
+        token = _current_lm_history_entry.set(None)
+        try:
+            with invocation:
+                result = await wrapped(*args, **kwargs)
+                history_entry = _current_lm_history_entry.get()
+                _set_lm_invocation_response(
+                    handler,
+                    invocation,
+                    instance,
+                    result,
+                    history_entry=history_entry,
+                )
+                return result
+        finally:
+            _current_lm_history_entry.reset(token)
+
+    return traced_method
+
 
 def _extract_tool_arguments(
     instance: Tool,
@@ -430,10 +686,7 @@ def _extract_retrieval_k(
     if k is None and hasattr(instance, "k"):
         k = getattr(instance, "k", None)
     if k is not None:
-        try:
-            return int(k)
-        except (ValueError, TypeError):
-            return None
+        return _safe_int(k)
     return None