fix(langchain): token usage reporting (#2074)

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

@@ -153,9 +153,30 @@ def _set_chat_response(span: Span, response: LLMResult) -> None:
     if not should_send_prompts():
         return
 
+    input_tokens = 0
+    output_tokens = 0
+    total_tokens = 0
+
     i = 0
     for generations in response.generations:
         for generation in generations:
+            if (
+                hasattr(generation, "message")
+                and hasattr(generation.message, "usage_metadata")
+                and generation.message.usage_metadata is not None
+            ):
+                input_tokens += (
+                    generation.message.usage_metadata.get("input_tokens")
+                    or generation.message.usage_metadata.get("prompt_tokens")
+                    or 0
+                )
+                output_tokens += (
+                    generation.message.usage_metadata.get("output_tokens")
+                    or generation.message.usage_metadata.get("completion_tokens")
+                    or 0
+                )
+                total_tokens = input_tokens + output_tokens
+
             prefix = f"{SpanAttributes.LLM_COMPLETIONS}.{i}"
             if hasattr(generation, "text") and generation.text != "":
                 span.set_attribute(
@@ -201,6 +222,20 @@ def _set_chat_response(span: Span, response: LLMResult) -> None:
                     )
             i += 1
 
+    if input_tokens > 0 or output_tokens > 0 or total_tokens > 0:
+        span.set_attribute(
+            SpanAttributes.LLM_USAGE_PROMPT_TOKENS,
+            input_tokens,
+        )
+        span.set_attribute(
+            SpanAttributes.LLM_USAGE_COMPLETION_TOKENS,
+            output_tokens,
+        )
+        span.set_attribute(
+            SpanAttributes.LLM_USAGE_TOTAL_TOKENS,
+            total_tokens,
+        )
+
 
 class TraceloopCallbackHandler(BaseCallbackHandler):
     def __init__(self, tracer: Tracer) -> None:
@@ -481,13 +516,19 @@ def on_llm_end(
 
         span = self._get_span(run_id)
 
-        token_usage = (response.llm_output or {}).get("token_usage")
+        token_usage = (response.llm_output or {}).get("token_usage") or (
+            response.llm_output or {}
+        ).get("usage")
         if token_usage is not None:
-            prompt_tokens = token_usage.get("prompt_tokens") or token_usage.get(
-                "input_token_count"
+            prompt_tokens = (
+                token_usage.get("prompt_tokens")
+                or token_usage.get("input_token_count")
+                or token_usage.get("input_tokens")
             )
-            completion_tokens = token_usage.get("completion_tokens") or token_usage.get(
-                "generated_token_count"
+            completion_tokens = (
+                token_usage.get("completion_tokens")
+                or token_usage.get("generated_token_count")
+                or token_usage.get("output_tokens")
             )
             total_tokens = token_usage.get("total_tokens") or (
                 prompt_tokens + completion_tokens

packages/opentelemetry-instrumentation-langchain/tests/test_chains.py

@@ -53,15 +53,32 @@ def test_sequential_chain(exporter):
         "SequentialChain.workflow",
     ] == [span.name for span in spans]
 
-    workflow_span = next(span for span in spans if span.name == "SequentialChain.workflow")
-    task_spans = [span for span in spans if span.name in ["synopsis.task", "LLMChain.task"]]
+    workflow_span = next(
+        span for span in spans if span.name == "SequentialChain.workflow"
+    )
+    task_spans = [
+        span for span in spans if span.name in ["synopsis.task", "LLMChain.task"]
+    ]
     llm_spans = [span for span in spans if span.name == "OpenAI.completion"]
 
     assert workflow_span.attributes[SpanAttributes.TRACELOOP_SPAN_KIND] == "workflow"
-    assert workflow_span.attributes[SpanAttributes.TRACELOOP_ENTITY_NAME] == "SequentialChain"
-    assert all(span.attributes[SpanAttributes.TRACELOOP_SPAN_KIND] == "task" for span in task_spans)
-    assert all(span.attributes[SpanAttributes.TRACELOOP_WORKFLOW_NAME] == "SequentialChain" for span in spans)
-    assert all(span.attributes[SpanAttributes.TRACELOOP_ENTITY_PATH] in ["synopsis", "LLMChain"] for span in llm_spans)
+    assert (
+        workflow_span.attributes[SpanAttributes.TRACELOOP_ENTITY_NAME]
+        == "SequentialChain"
+    )
+    assert all(
+        span.attributes[SpanAttributes.TRACELOOP_SPAN_KIND] == "task"
+        for span in task_spans
+    )
+    assert all(
+        span.attributes[SpanAttributes.TRACELOOP_WORKFLOW_NAME] == "SequentialChain"
+        for span in spans
+    )
+    assert all(
+        span.attributes[SpanAttributes.TRACELOOP_ENTITY_PATH]
+        in ["synopsis", "LLMChain"]
+        for span in llm_spans
+    )
 
     synopsis_span = next(span for span in spans if span.name == "synopsis.task")
     review_span = next(span for span in spans if span.name == "LLMChain.task")
@@ -200,12 +217,14 @@ def test_stream(exporter):
     chunks = list(runnable.stream({"product": "colorful socks"}))
     spans = exporter.get_finished_spans()
 
-    assert [
-        "PromptTemplate.task",
-        "ChatCohere.chat",
-        "StrOutputParser.task",
-        "RunnableSequence.workflow",
-    ] == [span.name for span in spans]
+    assert set(
+        [
+            "PromptTemplate.task",
+            "StrOutputParser.task",
+            "ChatCohere.chat",
+            "RunnableSequence.workflow",
+        ]
+    ) == set([span.name for span in spans])
     assert len(chunks) == 62
 
 
@@ -223,10 +242,12 @@ async def test_astream(exporter):
         chunks.append(chunk)
     spans = exporter.get_finished_spans()
 
-    assert [
-        "PromptTemplate.task",
-        "ChatCohere.chat",
-        "StrOutputParser.task",
-        "RunnableSequence.workflow",
-    ] == [span.name for span in spans]
+    assert set(
+        [
+            "PromptTemplate.task",
+            "ChatCohere.chat",
+            "StrOutputParser.task",
+            "RunnableSequence.workflow",
+        ]
+    ) == set([span.name for span in spans])
     assert len(chunks) == 144

packages/opentelemetry-instrumentation-langchain/tests/test_lcel.py

@@ -36,12 +36,14 @@ class Joke(BaseModel):
 
     spans = exporter.get_finished_spans()
 
-    assert [
-        "ChatPromptTemplate.task",
-        "ChatOpenAI.chat",
-        "JsonOutputFunctionsParser.task",
-        "ThisIsATestChain.workflow",
-    ] == [span.name for span in spans]
+    assert set(
+        [
+            "ChatPromptTemplate.task",
+            "JsonOutputFunctionsParser.task",
+            "ChatOpenAI.chat",
+            "ThisIsATestChain.workflow",
+        ]
+    ) == set([span.name for span in spans])
 
     workflow_span = next(
         span for span in spans if span.name == "ThisIsATestChain.workflow"

packages/opentelemetry-instrumentation-langchain/tests/test_llms.py

@@ -90,6 +90,9 @@ def test_openai(exporter):
     assert (
         openai_span.attributes[f"{SpanAttributes.LLM_COMPLETIONS}.0.role"]
     ) == "assistant"
+    assert openai_span.attributes[SpanAttributes.LLM_USAGE_PROMPT_TOKENS] == 24
+    assert openai_span.attributes[SpanAttributes.LLM_USAGE_COMPLETION_TOKENS] == 26
+    assert openai_span.attributes[SpanAttributes.LLM_USAGE_TOTAL_TOKENS] == 50
 
 
 @pytest.mark.vcr
@@ -113,12 +116,14 @@ class Joke(BaseModel):
 
     spans = exporter.get_finished_spans()
 
-    assert [
-        "ChatPromptTemplate.task",
-        "ChatOpenAI.chat",
-        "JsonOutputFunctionsParser.task",
-        "RunnableSequence.workflow",
-    ] == [span.name for span in spans]
+    assert set(
+        [
+            "ChatPromptTemplate.task",
+            "JsonOutputFunctionsParser.task",
+            "ChatOpenAI.chat",
+            "RunnableSequence.workflow",
+        ]
+    ) == set([span.name for span in spans])
 
     openai_span = next(span for span in spans if span.name == "ChatOpenAI.chat")
 
@@ -169,6 +174,9 @@ class Joke(BaseModel):
         )
         == response
     )
+    assert openai_span.attributes[SpanAttributes.LLM_USAGE_PROMPT_TOKENS] == 76
+    assert openai_span.attributes[SpanAttributes.LLM_USAGE_COMPLETION_TOKENS] == 35
+    assert openai_span.attributes[SpanAttributes.LLM_USAGE_TOTAL_TOKENS] == 111
 
 
 @pytest.mark.vcr
@@ -214,6 +222,9 @@ def test_anthropic(exporter):
     assert (
         anthropic_span.attributes[f"{SpanAttributes.LLM_COMPLETIONS}.0.role"]
     ) == "assistant"
+    assert anthropic_span.attributes[SpanAttributes.LLM_USAGE_PROMPT_TOKENS] == 19
+    assert anthropic_span.attributes[SpanAttributes.LLM_USAGE_COMPLETION_TOKENS] == 22
+    assert anthropic_span.attributes[SpanAttributes.LLM_USAGE_TOTAL_TOKENS] == 41
     output = json.loads(
         workflow_span.attributes[SpanAttributes.TRACELOOP_ENTITY_OUTPUT]
     )
@@ -286,6 +297,9 @@ def test_bedrock(exporter):
     assert (
         bedrock_span.attributes[f"{SpanAttributes.LLM_COMPLETIONS}.0.role"]
     ) == "assistant"
+    assert bedrock_span.attributes[SpanAttributes.LLM_USAGE_PROMPT_TOKENS] == 16
+    assert bedrock_span.attributes[SpanAttributes.LLM_USAGE_COMPLETION_TOKENS] == 27
+    assert bedrock_span.attributes[SpanAttributes.LLM_USAGE_TOTAL_TOKENS] == 43
     output = json.loads(
         workflow_span.attributes[SpanAttributes.TRACELOOP_ENTITY_OUTPUT]
     )