Added bedrock support and test

Author: wrisa

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

@@ -54,43 +54,66 @@ def on_chat_model_start(
         metadata: dict[str, Any] | None,
         **kwargs: Any,
     ) -> None:
-        invocation_params = kwargs.get("invocation_params")
-        request_model = (
-            invocation_params.get("model_name") if invocation_params else ""
-        )
-        span = self.span_manager.create_llm_span(
+        if "invocation_params" in kwargs:
+            params = (
+                kwargs["invocation_params"].get("params")
+                or kwargs["invocation_params"]
+            )
+        else:
+            params = kwargs
+
+        request_model = "unknown"
+        for model_tag in ("model", "model_id", "model_name", "ls_model_name"):
+            if (model := kwargs.get(model_tag)) is not None:
+                request_model = model
+                break
+            elif (
+                model := (params or {}).get(model_tag)
+            ) is not None:
+                request_model = model
+                break
+            elif (model := (metadata or {}).get(model_tag)) is not None:
+                request_model = model
+                break
+
+        span = self.span_manager.create_chat_span(
             run_id=run_id,
             parent_run_id=parent_run_id,
             request_model=request_model,
         )
 
-        if invocation_params is not None:
-            top_p = invocation_params.get("top_p")
+        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 = invocation_params.get("frequency_penalty")
+            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 = invocation_params.get("presence_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 = invocation_params.get("stop")
+            stop_sequences = params.get("stop")
             if stop_sequences is not None:
                 span.set_attribute(
                     GenAI.GEN_AI_REQUEST_STOP_SEQUENCES, stop_sequences
                 )
-            seed = invocation_params.get("seed")
+            seed = params.get("seed")
             if seed is not None:
                 span.set_attribute(GenAI.GEN_AI_REQUEST_SEED, seed)
-
-        if metadata is not None:
-            max_tokens = metadata.get("ls_max_tokens")
+            temperature = params.get("temperature")
+            if temperature is not None:
+                span.set_attribute(
+                    GenAI.GEN_AI_REQUEST_TEMPERATURE, temperature
+                )
+            max_tokens = params.get("max_completion_tokens")
             if max_tokens is not None:
                 span.set_attribute(GenAI.GEN_AI_REQUEST_MAX_TOKENS, max_tokens)
+
+        if metadata is not None:
             provider = metadata.get("ls_provider")
             if provider is not None:
                 span.set_attribute("gen_ai.provider.name", provider)
@@ -99,6 +122,9 @@ def on_chat_model_start(
                 span.set_attribute(
                     GenAI.GEN_AI_REQUEST_TEMPERATURE, temperature
                 )
+            max_tokens = metadata.get("ls_max_tokens")
+            if max_tokens is not None:
+                span.set_attribute(GenAI.GEN_AI_REQUEST_MAX_TOKENS, max_tokens)
 
     def on_llm_end(
         self,
@@ -124,6 +150,35 @@ def on_llm_end(
                     finish_reason = generation_info.get("finish_reason")
                     if finish_reason is not None:
                         finish_reasons.append(str(finish_reason) or "error")
+                if chat_generation.message:
+                    if (
+                        generation_info is None
+                        and chat_generation.message.response_metadata
+                    ):
+                        finish_reason = (
+                            chat_generation.message.response_metadata.get(
+                                "stopReason"
+                            )
+                        )
+                        if finish_reason is not None and span.is_recording():
+                            finish_reasons.append(finish_reason or "error")
+                    if chat_generation.message.usage_metadata:
+                        input_tokens = (
+                            chat_generation.message.usage_metadata.get(
+                                "input_tokens", 0
+                            )
+                        )
+                        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
+                        )
 
         span.set_attribute(
             GenAI.GEN_AI_RESPONSE_FINISH_REASONS, finish_reasons
@@ -143,22 +198,6 @@ def on_llm_end(
             if response_id is not None:
                 span.set_attribute(GenAI.GEN_AI_RESPONSE_ID, str(response_id))
 
-            # usage
-            usage = llm_output.get("usage") or llm_output.get("token_usage")
-            if usage:
-                prompt_tokens = usage.get("prompt_tokens", 0)
-                completion_tokens = usage.get("completion_tokens", 0)
-                span.set_attribute(
-                    GenAI.GEN_AI_USAGE_INPUT_TOKENS,
-                    int(prompt_tokens) if prompt_tokens is not None else 0,
-                )
-                span.set_attribute(
-                    GenAI.GEN_AI_USAGE_OUTPUT_TOKENS,
-                    int(completion_tokens)
-                    if completion_tokens is not None
-                    else 0,
-                )
-
         # End the LLM span
         self.span_manager.end_span(run_id)
 

instrumentation-genai/opentelemetry-instrumentation-langchain/tests/cassettes/test_langchain_call.yaml renamed to instrumentation-genai/opentelemetry-instrumentation-langchain/tests/cassettes/test_chat_openai_gpt_3_5_turbo_model_llm_call.yaml

@@ -42,7 +42,7 @@ interactions:
       host:
       - api.openai.com
       user-agent:
-      - OpenAI/Python 1.98.0
+      - OpenAI/Python 1.105.0
       x-stainless-arch:
       - arm64
       x-stainless-async:
@@ -52,7 +52,9 @@ interactions:
       x-stainless-os:
       - MacOS
       x-stainless-package-version:
-      - 1.98.0
+      - 1.105.0
+      x-stainless-raw-response:
+      - 'true'
       x-stainless-retry-count:
       - '0'
       x-stainless-runtime:
@@ -65,9 +67,9 @@ interactions:
     body:
       string: |-
         {
-          "id": "chatcmpl-BzXyZbwSwCmuSOKwRx7tnsia3V9lH",
+          "id": "chatcmpl-CBmRYRvfmoAG6EmqDOH4IwkfN02MJ",
           "object": "chat.completion",
-          "created": 1754008311,
+          "created": 1756923860,
           "model": "gpt-3.5-turbo-0125",
           "choices": [
             {
@@ -102,13 +104,13 @@ interactions:
         }
     headers:
       CF-RAY:
-      - 96813c201ae62f46-LAX
+      - 9797488a68f42b63-LAX
       Connection:
       - keep-alive
       Content-Type:
       - application/json
       Date:
-      - Fri, 01 Aug 2025 00:31:51 GMT
+      - Wed, 03 Sep 2025 18:24:21 GMT
       Server:
       - cloudflare
       Set-Cookie: test_set_cookie
@@ -128,29 +130,27 @@ interactions:
       - '822'
       openai-organization: test_openai_org_id
       openai-processing-ms:
-      - '189'
+      - '783'
       openai-project:
-      - proj_GLiYlAc06hF0Fm06IMReZLy4
+      - proj_3o0Aqh32nPiGbrex8BJtPTCm
       openai-version:
       - '2020-10-01'
-      x-envoy-decorator-operation:
-      - router.openai.svc.cluster.local:5004/*
       x-envoy-upstream-service-time:
-      - '211'
+      - '787'
       x-ratelimit-limit-requests:
-      - '10000'
+      - '5000'
       x-ratelimit-limit-tokens:
-      - '200000'
+      - '2000000'
       x-ratelimit-remaining-requests:
-      - '9999'
+      - '4999'
       x-ratelimit-remaining-tokens:
-      - '199982'
+      - '1999982'
       x-ratelimit-reset-requests:
-      - 8.64s
+      - 12ms
       x-ratelimit-reset-tokens:
-      - 5ms
+      - 0s
       x-request-id:
-      - 8a5db9a5-6103-41d4-baf3-813be646bae4
+      - req_c1bd8705b06b4e9180a5d8340e44785c
     status:
       code: 200
       message: OK

instrumentation-genai/opentelemetry-instrumentation-langchain/tests/cassettes/test_us_amazon_nova_lite_v1_0_bedrock_llm_call.yaml

@@ -0,0 +1,89 @@
+interactions:
+- request:
+    body: |-
+      {
+        "messages": [
+          {
+            "role": "user",
+            "content": [
+              {
+                "text": "What is the capital of France?"
+              }
+            ]
+          }
+        ],
+        "system": [
+          {
+            "text": "You are a helpful assistant!"
+          }
+        ],
+        "inferenceConfig": {
+          "maxTokens": 100,
+          "temperature": 0.1
+        }
+      }
+    headers:
+      Content-Length:
+      - '202'
+      Content-Type:
+      - !!binary |
+        YXBwbGljYXRpb24vanNvbg==
+      User-Agent:
+      - !!binary |
+        Qm90bzMvMS40MC4yMiBtZC9Cb3RvY29yZSMxLjQwLjIyIHVhLzIuMSBvcy9tYWNvcyMyNC42LjAg
+        bWQvYXJjaCNhcm02NCBsYW5nL3B5dGhvbiMzLjEzLjUgbWQvcHlpbXBsI0NQeXRob24gbS9aLGIs
+        RCBjZmcvcmV0cnktbW9kZSNsZWdhY3kgQm90b2NvcmUvMS40MC4yMg==
+      X-Amz-Date:
+      - !!binary |
+        MjAyNTA5MDRUMDIyNzM4Wg==
+      amz-sdk-invocation-id:
+      - !!binary |
+        MTMwMjBiMWUtZDhkOC00NTNkLWI1ZjYtY2U5Yjk5ZWQ4Zjg4
+      amz-sdk-request:
+      - !!binary |
+        YXR0ZW1wdD0x
+      authorization:
+      - Bearer test_openai_api_key
+    method: POST
+    uri: https://bedrock-runtime.us-west-2.amazonaws.com/model/arn%3Aaws%3Abedrock%3Aus-west-2%3A906383545488%3Ainference-profile%2Fus.amazon.nova-lite-v1%3A0/converse
+  response:
+    body:
+      string: |-
+        {
+          "metrics": {
+            "latencyMs": 435
+          },
+          "output": {
+            "message": {
+              "content": [
+                {
+                  "text": "The capital of France is Paris. It is not only the capital but also the largest city in the country. Paris is known for its rich history, culture, and landmarks such as the Eiffel Tower, the Louvre Museum, and Notre-Dame Cathedral."
+                }
+              ],
+              "role": "assistant"
+            }
+          },
+          "stopReason": "end_turn",
+          "usage": {
+            "inputTokens": 13,
+            "outputTokens": 50,
+            "totalTokens": 63
+          }
+        }
+    headers:
+      Connection:
+      - keep-alive
+      Content-Length:
+      - '412'
+      Content-Type:
+      - application/json
+      Date:
+      - Thu, 04 Sep 2025 02:27:38 GMT
+      Set-Cookie: test_set_cookie
+      openai-organization: test_openai_org_id
+      x-amzn-RequestId:
+      - 38e8f9b0-89f9-4d78-8da4-bba3948f1889
+    status:
+      code: 200
+      message: OK
+version: 1

instrumentation-genai/opentelemetry-instrumentation-langchain/tests/conftest.py

@@ -3,8 +3,10 @@
 import json
 import os
 
+import boto3
 import pytest
 import yaml
+from langchain_aws import ChatBedrock
 from langchain_openai import ChatOpenAI
 
 from opentelemetry.instrumentation.langchain import LangChainInstrumentor
@@ -15,8 +17,8 @@
 )
 
 
-@pytest.fixture(scope="function", name="llm_model")
-def fixture_llm_model():
+@pytest.fixture(scope="function", name="chat_openai_gpt_3_5_turbo_model")
+def fixture_chat_openai_gpt_3_5_turbo_model():
     llm = ChatOpenAI(
         model="gpt-3.5-turbo",
         temperature=0.1,
@@ -30,6 +32,25 @@ def fixture_llm_model():
     yield llm
 
 
+@pytest.fixture(scope="function", name="us_amazon_nova_lite_v1_0")
+def fixture_us_amazon_nova_lite_v1_0():
+    llm_model_value = "arn:aws:bedrock:us-west-2:906383545488:inference-profile/us.amazon.nova-lite-v1:0"
+    llm = ChatBedrock(
+        model_id=llm_model_value,
+        client=boto3.client(
+            "bedrock-runtime",
+            aws_access_key_id="test_key",
+            aws_secret_access_key="test_secret",
+            region_name="us-west-2",
+            aws_account_id="test_account",
+        ),
+        provider="amazon",
+        temperature=0.1,
+        max_tokens=100,
+    )
+    yield llm
+
+
 @pytest.fixture(scope="function", name="span_exporter")
 def fixture_span_exporter():
     exporter = InMemorySpanExporter()