fix(langchain): report token usage histogram (#3059)

Co-authored-by: Nir Gazit <[email protected]>

Author: LakshmiPriyaSujith

packages/opentelemetry-instrumentation-langchain/opentelemetry/instrumentation/langchain/callback_handler.py

@@ -34,6 +34,7 @@
 from opentelemetry.instrumentation.langchain.span_utils import (
     SpanHolder,
     _set_span_attribute,
+    extract_model_name_from_response_metadata,
     set_chat_request,
     set_chat_response,
     set_chat_response_usage,
@@ -446,7 +447,8 @@ def on_llm_end(
             id = response.llm_output.get("id")
             if id is not None and id != "":
                 _set_span_attribute(span, GEN_AI_RESPONSE_ID, id)
-
+        if model_name is None:
+            model_name = extract_model_name_from_response_metadata(response)
         token_usage = (response.llm_output or {}).get("token_usage") or (
             response.llm_output or {}
         ).get("usage")
@@ -495,7 +497,7 @@ def on_llm_end(
                         SpanAttributes.LLM_RESPONSE_MODEL: model_name or "unknown",
                     },
                 )
-        set_chat_response_usage(span, response)
+        set_chat_response_usage(span, response, self.token_histogram, token_usage is None, model_name)
         if should_emit_events():
             self._emit_llm_end_events(response)
         else:

packages/opentelemetry-instrumentation-langchain/opentelemetry/instrumentation/langchain/span_utils.py

@@ -15,6 +15,7 @@
     CallbackFilteredJSONEncoder,
     should_send_prompts,
 )
+from opentelemetry.metrics import Histogram
 from opentelemetry.semconv_ai import (
     SpanAttributes,
 )
@@ -271,13 +272,18 @@ def set_chat_response(span: Span, response: LLMResult) -> None:
             i += 1
 
 
-def set_chat_response_usage(span: Span, response: LLMResult):
+def set_chat_response_usage(
+    span: Span,
+    response: LLMResult,
+    token_histogram: Histogram,
+    record_token_usage: bool,
+    model_name: str
+) -> None:
     input_tokens = 0
     output_tokens = 0
     total_tokens = 0
     cache_read_tokens = 0
 
-    i = 0
     for generations in response.generations:
         for generation in generations:
             if (
@@ -302,7 +308,6 @@ def set_chat_response_usage(span: Span, response: LLMResult):
                         "input_token_details", {}
                     )
                     cache_read_tokens += input_token_details.get("cache_read", 0)
-            i += 1
 
     if (
         input_tokens > 0
@@ -330,6 +335,38 @@ def set_chat_response_usage(span: Span, response: LLMResult):
             SpanAttributes.LLM_USAGE_CACHE_READ_INPUT_TOKENS,
             cache_read_tokens,
         )
+        if record_token_usage:
+            if input_tokens > 0:
+                token_histogram.record(
+                    input_tokens,
+                    attributes={
+                        SpanAttributes.LLM_SYSTEM: "Langchain",
+                        SpanAttributes.LLM_TOKEN_TYPE: "input",
+                        SpanAttributes.LLM_RESPONSE_MODEL: model_name,
+                    },
+                )
+
+            if output_tokens > 0:
+                token_histogram.record(
+                    output_tokens,
+                    attributes={
+                        SpanAttributes.LLM_SYSTEM: "Langchain",
+                        SpanAttributes.LLM_TOKEN_TYPE: "output",
+                        SpanAttributes.LLM_RESPONSE_MODEL: model_name,
+                    },
+                )
+
+
+def extract_model_name_from_response_metadata(response: LLMResult) -> str:
+    for generations in response.generations:
+        for generation in generations:
+            if (
+                getattr(generation, "message", None)
+                and getattr(generation.message, "response_metadata", None)
+                and (model_name := generation.message.response_metadata.get("model_name"))
+            ):
+                return model_name
+    return "unknown"
 
 
 def _set_chat_tool_calls(

packages/opentelemetry-instrumentation-langchain/tests/metrics/cassettes/test_langchain_metrics/test_llm_chain_metrics_with_none_llm_output.yaml

@@ -0,0 +1,103 @@
+interactions:
+- request:
+    body: '{"messages": [{"content": "What is a good name for a company that makes
+      colorful socks?", "role": "user"}], "model": "gpt-3.5-turbo", "n": 1, "stream":
+      false, "temperature": 0.0}'
+    headers:
+      accept:
+      - application/json
+      accept-encoding:
+      - gzip, deflate
+      connection:
+      - keep-alive
+      content-length:
+      - '178'
+      content-type:
+      - application/json
+      host:
+      - api.openai.com
+      traceparent:
+      - 00-d77972e6e3b0d96f4a65edea48fc5e34-92f98f134382afae-01
+      user-agent:
+      - OpenAI/Python 1.45.1
+      x-stainless-arch:
+      - arm64
+      x-stainless-async:
+      - 'false'
+      x-stainless-lang:
+      - python
+      x-stainless-os:
+      - MacOS
+      x-stainless-package-version:
+      - 1.45.1
+      x-stainless-runtime:
+      - CPython
+      x-stainless-runtime-version:
+      - 3.12.1
+    method: POST
+    uri: https://api.openai.com/v1/chat/completions
+  response:
+    body:
+      string: !!binary |
+        H4sIAAAAAAAAA2xQPU/DMBTc8yssz02VDwI0WydExVAB6gBCkeO8JgbHtuyX8lH1vyOHtElVFg93
+        vnt3tw8IoaKiOaG8YchbI8Pl/c/u6mm92m1eeJpmD5s1cyq7W8VLI1s68wpdvgPHo2rOdWskoNDq
+        j+YWGIJ3jW+SRbSIk9vrnmh1BdLLaoNhOs9C7GypwyhOskHZaMHB0Zy8BoQQsu9fn1FV8EVzEs2O
+        SAvOsRpofvpECLVaeoQy54RDppDORpJrhaD62I9MqFJ/kufGAqvc9JeFbeeYT6k6KQf8cDordW2s
+        Lt3An/CtUMI1hQXmtPInHGpDe/YQEPLW1+vOElNjdWuwQP0ByhsmQzs6DjqS6cChRiYnmiN+ZlZU
+        gExIN1mHcsYbqEZlFEyaXZ78z+KvnVD1hUswOFH37RDaYitUDdZY0e/dL3kIfgEAAP//AwDTDjnS
+        bgIAAA==
+    headers:
+      CF-Cache-Status:
+      - DYNAMIC
+      CF-RAY:
+      - 8d38f65e8c5807d6-ATL
+      Connection:
+      - keep-alive
+      Content-Encoding:
+      - gzip
+      Content-Type:
+      - application/json
+      Date:
+      - Wed, 16 Oct 2024 15:08:07 GMT
+      Server:
+      - cloudflare
+      Set-Cookie:
+      - __cf_bm=OSPGzzsIkijMjLSfTJjIkSiCxrDnWuQJoJfatCLq_os-1729091287-1.0.1.1-VsHqFTiK6ZfwZayLMdlW8YQ1RAnQm5rP76aV3S2QPBtjrbPWclwVCrqNC0SMuxcZxXGHj2IP1r8OafsbY_cwXA;
+        path=/; expires=Wed, 16-Oct-24 15:38:07 GMT; domain=.api.openai.com; HttpOnly;
+        Secure; SameSite=None
+      - _cfuvid=kMO6YhIaggKVe4UZFi5l9RM0fjUpZYSwQRUPEBZNJr8-1729091287006-0.0.1.1-604800000;
+        path=/; domain=.api.openai.com; HttpOnly; Secure; SameSite=None
+      Transfer-Encoding:
+      - chunked
+      X-Content-Type-Options:
+      - nosniff
+      access-control-expose-headers:
+      - X-Request-ID
+      alt-svc:
+      - h3=":443"; ma=86400
+      openai-organization:
+      - traceloop
+      openai-processing-ms:
+      - '160'
+      openai-version:
+      - '2020-10-01'
+      strict-transport-security:
+      - max-age=31536000; includeSubDomains; preload
+      x-ratelimit-limit-requests:
+      - '5000'
+      x-ratelimit-limit-tokens:
+      - '4000000'
+      x-ratelimit-remaining-requests:
+      - '4999'
+      x-ratelimit-remaining-tokens:
+      - '3999968'
+      x-ratelimit-reset-requests:
+      - 12ms
+      x-ratelimit-reset-tokens:
+      - 0s
+      x-request-id:
+      - req_1f22ead323fed2094d4f3f5c0187205d
+    status:
+      code: 200
+      message: OK
+version: 1

packages/opentelemetry-instrumentation-langchain/tests/metrics/test_langchain_metrics.py

@@ -1,3 +1,4 @@
+from unittest.mock import patch
 import pytest
 from langchain.chains import LLMChain
 from langchain.prompts import PromptTemplate
@@ -114,3 +115,66 @@ def test_llm_chain_streaming_metrics(instrument_legacy, reader, llm):
 
     assert found_token_metric is True
     assert found_duration_metric is True
+
+
+def verify_token_metrics(data_points):
+    for data_point in data_points:
+        assert data_point.attributes[SpanAttributes.LLM_TOKEN_TYPE] in [
+            "output",
+            "input",
+        ]
+        assert data_point.sum > 0
+        assert data_point.attributes[SpanAttributes.LLM_SYSTEM] == "Langchain"
+
+
+def verify_duration_metrics(data_points):
+    assert any(data_point.count > 0 for data_point in data_points)
+    assert any(data_point.sum > 0 for data_point in data_points)
+    for data_point in data_points:
+        assert data_point.attributes[SpanAttributes.LLM_SYSTEM] == "Langchain"
+
+
+def verify_langchain_metrics(reader):
+    metrics_data = reader.get_metrics_data()
+    resource_metrics = metrics_data.resource_metrics
+    assert len(resource_metrics) > 0
+
+    found_token_metric = False
+    found_duration_metric = False
+
+    for rm in resource_metrics:
+        for sm in rm.scope_metrics:
+            for metric in sm.metrics:
+                if metric.name == Meters.LLM_TOKEN_USAGE:
+                    found_token_metric = True
+                    verify_token_metrics(metric.data.data_points)
+
+                if metric.name == Meters.LLM_OPERATION_DURATION:
+                    found_duration_metric = True
+                    verify_duration_metrics(metric.data.data_points)
+
+    return found_token_metric, found_duration_metric
+
+
+@pytest.mark.vcr
+def test_llm_chain_metrics_with_none_llm_output(instrument_legacy, reader, chain, llm):
+    """
+    This test verifies that the metrics system correctly handles edge cases where the
+    LLM response contains a None value in the llm_output field, ensuring that token
+    usage and operation duration metrics are still properly recorded.
+    """
+    original_generate = llm._generate
+
+    # Create a patched version that returns results with None llm_output
+    def patched_generate(*args, **kwargs):
+        result = original_generate(*args, **kwargs)
+        result.llm_output = None
+        return result
+
+    with patch.object(llm, '_generate', side_effect=patched_generate):
+        chain.run(product="colorful socks")
+
+    found_token_metric, found_duration_metric = verify_langchain_metrics(reader)
+
+    assert found_token_metric is True, "Token usage metrics not found"
+    assert found_duration_metric is True, "Operation duration metrics not found"