Removed telemetry from inference callbacks and added calls to genai utils apis instead.

Author: wrisa

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

@@ -48,8 +48,7 @@
 from opentelemetry.instrumentation.langchain.package import _instruments
 from opentelemetry.instrumentation.langchain.version import __version__
 from opentelemetry.instrumentation.utils import unwrap
-from opentelemetry.semconv.schemas import Schemas
-from opentelemetry.trace import get_tracer
+from opentelemetry.util.genai.handler import get_telemetry_handler
 
 
 class LangChainInstrumentor(BaseInstrumentor):
@@ -72,15 +71,10 @@ def _instrument(self, **kwargs: Any):
         Enable Langchain instrumentation.
         """
         tracer_provider = kwargs.get("tracer_provider")
-        tracer = get_tracer(
-            __name__,
-            __version__,
-            tracer_provider,
-            schema_url=Schemas.V1_37_0.value,
-        )
 
+        telemetry_handler = get_telemetry_handler(tracer_provider=tracer_provider)
         otel_callback_handler = OpenTelemetryLangChainCallbackHandler(
-            tracer=tracer,
+            telemetry_handler=telemetry_handler
         )
 
         wrap_function_wrapper(

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

@@ -21,12 +21,10 @@
 from langchain_core.messages import BaseMessage  # type: ignore
 from langchain_core.outputs import LLMResult  # type: ignore
 
-from opentelemetry.instrumentation.langchain.span_manager import _SpanManager
-from opentelemetry.semconv._incubating.attributes import (
-    gen_ai_attributes as GenAI,
-)
-from opentelemetry.trace import Tracer
+from opentelemetry.instrumentation.langchain.invocation_manager import _InvocationManager
 
+from opentelemetry.util.genai.handler import TelemetryHandler
+from opentelemetry.util.genai.types import LLMInvocation, Text, InputMessage, OutputMessage, Error
 
 class OpenTelemetryLangChainCallbackHandler(BaseCallbackHandler):  # type: ignore[misc]
     """
@@ -35,13 +33,11 @@ class OpenTelemetryLangChainCallbackHandler(BaseCallbackHandler):  # type: ignor
 
     def __init__(
         self,
-        tracer: Tracer,
+        telemetry_handler: TelemetryHandler
     ) -> None:
         super().__init__()  # type: ignore
-
-        self.span_manager = _SpanManager(
-            tracer=tracer,
-        )
+        self._telemetry_handler = telemetry_handler
+        self._invocation_manager = _InvocationManager()
 
     def on_chat_model_start(
         self,
@@ -82,59 +78,68 @@ def on_chat_model_start(
         if request_model == "unknown":
             return
 
-        span = self.span_manager.create_chat_span(
-            run_id=run_id,
-            parent_run_id=parent_run_id,
-            request_model=request_model,
-        )
-
-        if params is not None:
-            top_p = params.get("top_p")
-            if top_p is not None:
-                span.set_attribute(GenAI.GEN_AI_REQUEST_TOP_P, top_p)
-            frequency_penalty = params.get("frequency_penalty")
-            if frequency_penalty is not None:
-                span.set_attribute(
-                    GenAI.GEN_AI_REQUEST_FREQUENCY_PENALTY, frequency_penalty
-                )
-            presence_penalty = params.get("presence_penalty")
-            if presence_penalty is not None:
-                span.set_attribute(
-                    GenAI.GEN_AI_REQUEST_PRESENCE_PENALTY, presence_penalty
-                )
-            stop_sequences = params.get("stop")
-            if stop_sequences is not None:
-                span.set_attribute(
-                    GenAI.GEN_AI_REQUEST_STOP_SEQUENCES, stop_sequences
-                )
-            seed = params.get("seed")
-            if seed is not None:
-                span.set_attribute(GenAI.GEN_AI_REQUEST_SEED, seed)
-            # ChatOpenAI
-            temperature = params.get("temperature")
-            if temperature is not None:
-                span.set_attribute(
-                    GenAI.GEN_AI_REQUEST_TEMPERATURE, temperature
-                )
-            # ChatOpenAI
-            max_tokens = params.get("max_completion_tokens")
-            if max_tokens is not None:
-                span.set_attribute(GenAI.GEN_AI_REQUEST_MAX_TOKENS, max_tokens)
-
+        # TODO: uncomment and modify following after PR #3862 is merged
+        # if params is not None:
+        #     top_p = params.get("top_p")
+        #     if top_p is not None:
+        #         span.set_attribute(GenAI.GEN_AI_REQUEST_TOP_P, top_p)
+        #     frequency_penalty = params.get("frequency_penalty")
+        #     if frequency_penalty is not None:
+        #         span.set_attribute(
+        #             GenAI.GEN_AI_REQUEST_FREQUENCY_PENALTY, frequency_penalty
+        #         )
+        #     presence_penalty = params.get("presence_penalty")
+        #     if presence_penalty is not None:
+        #         span.set_attribute(
+        #             GenAI.GEN_AI_REQUEST_PRESENCE_PENALTY, presence_penalty
+        #         )
+        #     stop_sequences = params.get("stop")
+        #     if stop_sequences is not None:
+        #         span.set_attribute(
+        #             GenAI.GEN_AI_REQUEST_STOP_SEQUENCES, stop_sequences
+        #         )
+        #     seed = params.get("seed")
+        #     if seed is not None:
+        #         span.set_attribute(GenAI.GEN_AI_REQUEST_SEED, seed)
+        #     # ChatOpenAI
+        #     temperature = params.get("temperature")
+        #     if temperature is not None:
+        #         span.set_attribute(
+        #             GenAI.GEN_AI_REQUEST_TEMPERATURE, temperature
+        #         )
+        #     # ChatOpenAI
+        #     max_tokens = params.get("max_completion_tokens")
+        #     if max_tokens is not None:
+        #         span.set_attribute(GenAI.GEN_AI_REQUEST_MAX_TOKENS, max_tokens)
+        #
+        provider = "unknown"
         if metadata is not None:
             provider = metadata.get("ls_provider")
-            if provider is not None:
-                span.set_attribute("gen_ai.provider.name", provider)
-            # ChatBedrock
-            temperature = metadata.get("ls_temperature")
-            if temperature is not None:
-                span.set_attribute(
-                    GenAI.GEN_AI_REQUEST_TEMPERATURE, temperature
-                )
-            # ChatBedrock
-            max_tokens = metadata.get("ls_max_tokens")
-            if max_tokens is not None:
-                span.set_attribute(GenAI.GEN_AI_REQUEST_MAX_TOKENS, max_tokens)
+
+        # TODO: uncomment and modify following after PR #3862 is merged
+
+        #     # ChatBedrock
+        #     temperature = metadata.get("ls_temperature")
+        #     if temperature is not None:
+        #         span.set_attribute(
+        #             GenAI.GEN_AI_REQUEST_TEMPERATURE, temperature
+        #         )
+        #     # ChatBedrock
+        #     max_tokens = metadata.get("ls_max_tokens")
+        #     if max_tokens is not None:
+        #         span.set_attribute(GenAI.GEN_AI_REQUEST_MAX_TOKENS, max_tokens)
+
+        input_messages: list[InputMessage] = []
+        for sub_messages in messages:
+            for message in sub_messages:
+                content = get_property_value(message, "content")
+                role = get_property_value(message, "type")
+                parts = [Text(content=content, type="text")]
+                input_messages.append(InputMessage(parts=parts, role=role))
+
+        llm_invocation = LLMInvocation(request_model=request_model, input_messages=input_messages, provider=provider)
+        llm_invocation = self._telemetry_handler.start_llm(invocation=llm_invocation)
+        self._invocation_manager.add_invocation_state(run_id=run_id, parent_run_id=parent_run_id, invocation=llm_invocation)
 
     def on_llm_end(
         self,
@@ -144,25 +149,28 @@ def on_llm_end(
         parent_run_id: UUID | None,
         **kwargs: Any,
     ) -> None:
-        span = self.span_manager.get_span(run_id)
+        invocation = self._invocation_manager.get_invocation(run_id=run_id)
 
-        if span is None:
-            # If the span does not exist, we cannot set attributes or end it
+        if invocation is None or not isinstance(invocation, LLMInvocation):
+            # If the invocation does not exist, we cannot set attributes or end it
             return
 
-        finish_reasons: list[str] = []
+        llm_invocation: LLMInvocation = invocation
+
+        output_messages: list[OutputMessage] = []
         for generation in getattr(response, "generations", []):  # type: ignore
             for chat_generation in generation:
+                # Get finish reason
                 generation_info = getattr(
                     chat_generation, "generation_info", None
                 )
                 if generation_info is not None:
                     finish_reason = generation_info.get(
                         "finish_reason", "unknown"
                     )
-                    if finish_reason is not None:
-                        finish_reasons.append(str(finish_reason))
+
                 if chat_generation.message:
+                    # Get finish reason if generation_info is None above
                     if (
                         generation_info is None
                         and chat_generation.message.response_metadata
@@ -172,46 +180,46 @@ def on_llm_end(
                                 "stopReason", "unknown"
                             )
                         )
-                        if finish_reason is not None:
-                            finish_reasons.append(str(finish_reason))
+
+                    # Get message content
+                    parts = [Text(content=get_property_value(chat_generation.message, "content"), type="text")]
+                    role = get_property_value(chat_generation.message, "type")
+                    output_message = OutputMessage(role=role, parts=parts, finish_reason=finish_reason)
+                    output_messages.append(output_message)
+
+                    # Get token usage if available
                     if chat_generation.message.usage_metadata:
                         input_tokens = (
                             chat_generation.message.usage_metadata.get(
                                 "input_tokens", 0
                             )
                         )
+                        llm_invocation.input_tokens = input_tokens
+
                         output_tokens = (
                             chat_generation.message.usage_metadata.get(
                                 "output_tokens", 0
                             )
                         )
-                        span.set_attribute(
-                            GenAI.GEN_AI_USAGE_INPUT_TOKENS, input_tokens
-                        )
-                        span.set_attribute(
-                            GenAI.GEN_AI_USAGE_OUTPUT_TOKENS, output_tokens
-                        )
+                        llm_invocation.output_tokens = output_tokens
 
-        span.set_attribute(
-            GenAI.GEN_AI_RESPONSE_FINISH_REASONS, finish_reasons
-        )
+        llm_invocation.output_messages = output_messages
 
         llm_output = getattr(response, "llm_output", None)  # type: ignore
         if llm_output is not None:
             response_model = llm_output.get("model_name") or llm_output.get(
                 "model"
             )
             if response_model is not None:
-                span.set_attribute(
-                    GenAI.GEN_AI_RESPONSE_MODEL, str(response_model)
-                )
+                llm_invocation.response_model_name = str(response_model)
 
             response_id = llm_output.get("id")
             if response_id is not None:
-                span.set_attribute(GenAI.GEN_AI_RESPONSE_ID, str(response_id))
+                llm_invocation.response_id = str(response_id)
 
-        # End the LLM span
-        self.span_manager.end_span(run_id)
+        invocation = self._telemetry_handler.stop_llm(invocation=invocation)
+        if not invocation.span.is_recording():
+            self._invocation_manager.delete_invocation_state(run_id=run_id)
 
     def on_llm_error(
         self,
@@ -221,4 +229,22 @@ def on_llm_error(
         parent_run_id: UUID | None,
         **kwargs: Any,
     ) -> None:
-        self.span_manager.handle_error(error, run_id)
+        invocation = self._invocation_manager.get_invocation(run_id=run_id)
+
+        if invocation is None or not isinstance(invocation, LLMInvocation):
+            # If the invocation does not exist, we cannot set attributes or end it
+            return
+
+        invocation: LLMInvocation = invocation
+
+        error = Error(message=str(error), type=type(error))
+        invocation = self._telemetry_handler.fail_llm(invocation=invocation, error=error)
+        if not invocation.span.is_recording():
+            self._invocation_manager.delete_invocation_state(run_id=run_id)
+
+
+def get_property_value(obj, property_name):
+    if isinstance(obj, dict):
+        return obj.get(property_name, None)
+
+    return getattr(obj, property_name, None)

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

@@ -0,0 +1,62 @@
+# Copyright The OpenTelemetry Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from dataclasses import dataclass, field
+from typing import Dict, List, Optional
+from uuid import UUID
+
+from opentelemetry.util.genai.types import GenAIInvocation
+
+__all__ = ["_InvocationManager"]
+
+
+@dataclass
+class _InvocationState:
+    invocation: GenAIInvocation
+    children: List[UUID] = field(default_factory=lambda: list())
+
+
+class _InvocationManager:
+    def __init__(
+        self,
+    ) -> None:
+        # Map from run_id -> _InvocationState, to keep track of invocations and parent/child relationships
+        # TODO: Use weak references or a TTL cache to avoid memory leaks in long-running processes. See #3735
+        self._invocations: Dict[UUID, _InvocationState] = {}
+
+    def add_invocation_state(
+        self,
+        run_id: UUID,
+        parent_run_id: Optional[UUID],
+        invocation: GenAIInvocation,
+    ):
+        if parent_run_id is not None and parent_run_id in self._invocations:
+            parent_invocation_state = self._invocations[parent_run_id]
+            parent_invocation_state.children.append(run_id)
+
+        invocation_state = _InvocationState(invocation=invocation)
+        self._invocations[run_id] = invocation_state
+
+
+    def get_invocation(self, run_id: UUID) -> Optional[GenAIInvocation]:
+        invocation_state = self._invocations.get(run_id)
+        return invocation_state.invocation if invocation_state else None
+
+    def delete_invocation_state(self, run_id: UUID) -> None:
+        invocation_state = self._invocations[run_id]
+        for child_id in invocation_state.children:
+            child_invocation_state = self._invocations.get(child_id)
+            if child_invocation_state:
+                del self._invocations[child_id]
+        del self._invocations[run_id]