langchain: fix token usage from EM endpoints (#265)

Signed-off-by: Bryan Qiu <bryan.qiu@databricks.com>
Co-authored-by: Jake Steelman <jake.steelman@insurica.com>

Author: bbqiu

integrations/langchain/src/databricks_langchain/chat_models.py

@@ -690,6 +690,18 @@ def _stream(
                             logprobs=generation_info.get("logprobs"),
                         )
                     yield generation_chunk
+                elif chunk.usage and stream_usage:
+                    # Some models send a final chunk that does not have
+                    # a delta or choices, but does have usage info
+                    if not usage_chunk_emitted:
+                        input_tokens = getattr(chunk.usage, "prompt_tokens", None)
+                        output_tokens = getattr(chunk.usage, "completion_tokens", None)
+                        if input_tokens is not None and output_tokens is not None:
+                            final_usage = {
+                                "input_tokens": input_tokens,
+                                "output_tokens": output_tokens,
+                                "total_tokens": input_tokens + output_tokens,
+                            }
 
             # Emit special usage chunk at end of stream
             if stream_usage and final_usage and not usage_chunk_emitted:

integrations/langchain/tests/integration_tests/test_chat_models.py

@@ -13,7 +13,6 @@
 from typing import Annotated
 
 import pytest
-from langchain.agents import AgentExecutor, create_tool_calling_agent
 from langchain_core.callbacks.base import BaseCallbackHandler
 from langchain_core.messages import (
     AIMessage,
@@ -378,30 +377,6 @@ def multiply(a: int, b: int) -> int:
     return a * b
 
 
-@pytest.mark.foundation_models
-@pytest.mark.parametrize("model", _FOUNDATION_MODELS)
-def test_chat_databricks_agent_executor(model):
-    model = ChatDatabricks(
-        model=model,
-        temperature=0,
-        max_tokens=100,
-    )
-    tools = [add, multiply]
-    prompt = ChatPromptTemplate.from_messages(
-        [
-            ("system", "You are a helpful assistant"),
-            ("human", "{input}"),
-            ("placeholder", "{agent_scratchpad}"),
-        ]
-    )
-
-    agent = create_tool_calling_agent(model, tools, prompt)
-    agent_executor = AgentExecutor(agent=agent, tools=tools)
-
-    response = agent_executor.invoke({"input": "What is (10 + 5) * 3?"})
-    assert "45" in response["output"]
-
-
 @pytest.mark.foundation_models
 @pytest.mark.parametrize("model", _FOUNDATION_MODELS)
 def test_chat_databricks_langgraph(model):
@@ -784,7 +759,7 @@ def test_chat_databricks_custom_outputs():
 
 
 def test_chat_databricks_custom_outputs_stream():
-    llm = ChatDatabricks(model="agents_ml-bbqiu-codegen", use_responses_api=True)
+    llm = ChatDatabricks(model="agents_ml-bbqiu-mcp-openai", use_responses_api=True)
     response = llm.stream(
         "What is the 10th fibonacci number?",
         custom_inputs={"key": "value"},
@@ -820,3 +795,83 @@ def test_chat_databricks_token_count():
         last_chunk.usage_metadata["total_tokens"]
         == last_chunk.usage_metadata["input_tokens"] + last_chunk.usage_metadata["output_tokens"]
     )
+
+
+def test_chat_databricks_gpt5_stream_with_usage():
+    """
+    Test GPT-5 streaming with usage metadata.
+
+    GPT-5 sends a final chunk with only usage data (no choices/delta).
+    This test verifies that the usage metadata is correctly extracted from that final chunk.
+
+    Example final chunk from GPT-5:
+    ChatCompletionChunk(
+        id='chatcmpl-...',
+        choices=[],  # Empty choices array
+        created=...,
+        model='gpt-5-2025-08-07',
+        object='chat.completion.chunk',
+        usage=CompletionUsage(
+            completion_tokens=267,
+            prompt_tokens=4861,
+            total_tokens=5128,
+            ...
+        )
+    )
+    """
+    from databricks.sdk import WorkspaceClient
+
+    # Use dogfood profile to access GPT-5
+    workspace_client = WorkspaceClient(profile=DATABRICKS_CLI_PROFILE)
+
+    llm = ChatDatabricks(
+        endpoint="gpt-5",
+        workspace_client=workspace_client,
+        max_tokens=100,
+        stream_usage=True,
+    )
+
+    # Stream a simple query
+    chunks = list(llm.stream("hello"))
+
+    # Verify we get chunks
+    assert len(chunks) > 0, "Expected at least one chunk from GPT-5 streaming"
+
+    # Find content chunks (non-empty content)
+    content_chunks = [chunk for chunk in chunks if chunk.content != ""]
+    assert len(content_chunks) > 0, "Expected at least one content chunk"
+
+    # Find usage chunks (empty content with usage_metadata)
+    usage_chunks = [
+        chunk for chunk in chunks if chunk.content == "" and chunk.usage_metadata is not None
+    ]
+
+    # Should have exactly ONE usage chunk from the final usage-only chunk
+    assert len(usage_chunks) == 1, (
+        f"Expected exactly 1 usage chunk from GPT-5 final chunk, got {len(usage_chunks)}"
+    )
+
+    # Verify usage chunk has correct metadata structure
+    usage_chunk = usage_chunks[0]
+    assert isinstance(usage_chunk, AIMessageChunk)
+    assert usage_chunk.content == ""
+    assert "input_tokens" in usage_chunk.usage_metadata
+    assert "output_tokens" in usage_chunk.usage_metadata
+    assert "total_tokens" in usage_chunk.usage_metadata
+
+    # Verify token counts are positive
+    assert usage_chunk.usage_metadata["input_tokens"] > 0, (
+        f"Expected positive input_tokens, got {usage_chunk.usage_metadata['input_tokens']}"
+    )
+    assert usage_chunk.usage_metadata["output_tokens"] > 0, (
+        f"Expected positive output_tokens, got {usage_chunk.usage_metadata['output_tokens']}"
+    )
+
+    # Verify total_tokens equals sum of input and output
+    expected_total = (
+        usage_chunk.usage_metadata["input_tokens"] + usage_chunk.usage_metadata["output_tokens"]
+    )
+    assert usage_chunk.usage_metadata["total_tokens"] == expected_total, (
+        f"Expected total_tokens ({usage_chunk.usage_metadata['total_tokens']}) "
+        f"to equal input_tokens + output_tokens ({expected_total})"
+    )

integrations/langchain/tests/unit_tests/test_chat_models.py

@@ -350,6 +350,193 @@ def test_chat_model_stream_no_duplicate_usage_chunks():
         assert len(usage_chunks) == 1, f"Expected exactly 1 usage chunk, got {len(usage_chunks)}"
 
 
+def test_chat_model_stream_usage_only_final_chunk():
+    """Test that a final chunk with only usage data (no choices) correctly emits usage metadata."""
+    from unittest.mock import Mock, patch
+
+    mock_usage = Mock()
+    mock_usage.prompt_tokens = 15
+    mock_usage.completion_tokens = 10
+
+    # Simulate GPT-5 streaming behavior: content chunks followed by usage-only chunk
+    mock_chunks = [
+        Mock(
+            choices=[
+                Mock(
+                    delta=Mock(
+                        role="assistant",
+                        content="Hello",
+                        model_dump=Mock(return_value={"role": "assistant", "content": "Hello"}),
+                    ),
+                    finish_reason=None,
+                    logprobs=None,
+                )
+            ],
+            usage=None,
+        ),
+        Mock(
+            choices=[
+                Mock(
+                    delta=Mock(
+                        role="assistant",
+                        content=" world",
+                        model_dump=Mock(return_value={"role": "assistant", "content": " world"}),
+                    ),
+                    finish_reason="stop",
+                    logprobs=None,
+                )
+            ],
+            usage=None,
+        ),
+        # Final chunk with ONLY usage data, no choices/delta
+        Mock(
+            choices=[],
+            usage=mock_usage,
+        ),
+    ]
+
+    # Verify mock structure matches GPT-5 behavior
+    # Final chunk has empty choices list and usage data (no delta)
+    assert len(mock_chunks[2].choices) == 0
+    assert mock_chunks[2].usage is not None
+
+    with patch("databricks_langchain.chat_models.get_openai_client") as mock_get_client:
+        mock_client = Mock()
+        mock_get_client.return_value = mock_client
+        mock_client.chat.completions.create.return_value = iter(mock_chunks)
+
+        llm = ChatDatabricks(model="test-model")
+        messages = [HumanMessage(content="Hello")]
+
+        chunks = list(llm.stream(messages, stream_usage=True))
+
+        # Should get content chunks plus one usage chunk
+        content_chunks = [chunk for chunk in chunks if chunk.content != ""]
+        assert len(content_chunks) == 2
+        assert content_chunks[0].content == "Hello"
+        assert content_chunks[1].content == " world"
+
+        # Should emit exactly ONE usage chunk
+        usage_chunks = [
+            chunk for chunk in chunks if chunk.content == "" and chunk.usage_metadata is not None
+        ]
+        assert len(usage_chunks) == 1, f"Expected exactly 1 usage chunk, got {len(usage_chunks)}"
+
+        # Verify usage chunk has correct metadata
+        usage_chunk = usage_chunks[0]
+        assert isinstance(usage_chunk, AIMessageChunk)
+        assert usage_chunk.content == ""
+        assert usage_chunk.usage_metadata["input_tokens"] == 15
+        assert usage_chunk.usage_metadata["output_tokens"] == 10
+        assert usage_chunk.usage_metadata["total_tokens"] == 25
+
+
+def test_chat_model_stream_usage_only_chunk_missing_tokens():
+    """Test that a usage-only chunk with missing token data doesn't emit usage metadata."""
+    from unittest.mock import Mock, patch
+
+    mock_usage = Mock()
+    mock_usage.prompt_tokens = None  # Missing prompt_tokens
+    mock_usage.completion_tokens = 10
+
+    mock_chunks = [
+        Mock(
+            choices=[
+                Mock(
+                    delta=Mock(
+                        role="assistant",
+                        content="Hello",
+                        model_dump=Mock(return_value={"role": "assistant", "content": "Hello"}),
+                    ),
+                    finish_reason="stop",
+                    logprobs=None,
+                )
+            ],
+            usage=None,
+        ),
+        # Final chunk with usage data but missing prompt_tokens
+        Mock(
+            choices=[],
+            usage=mock_usage,
+        ),
+    ]
+
+    with patch("databricks_langchain.chat_models.get_openai_client") as mock_get_client:
+        mock_client = Mock()
+        mock_get_client.return_value = mock_client
+        mock_client.chat.completions.create.return_value = iter(mock_chunks)
+
+        llm = ChatDatabricks(model="test-model")
+        messages = [HumanMessage(content="Hello")]
+
+        chunks = list(llm.stream(messages, stream_usage=True))
+
+        # Should get content chunks but NO usage chunk (due to missing tokens)
+        content_chunks = [chunk for chunk in chunks if chunk.content != ""]
+        assert len(content_chunks) == 1
+
+        # Should NOT emit a usage chunk when tokens are missing
+        usage_chunks = [
+            chunk for chunk in chunks if chunk.content == "" and chunk.usage_metadata is not None
+        ]
+        assert len(usage_chunks) == 0, (
+            f"Expected 0 usage chunks when tokens are missing, got {len(usage_chunks)}"
+        )
+
+
+def test_chat_model_stream_usage_only_chunk_stream_usage_false():
+    """Test that a usage-only chunk is ignored when stream_usage=False."""
+    from unittest.mock import Mock, patch
+
+    mock_usage = Mock()
+    mock_usage.prompt_tokens = 15
+    mock_usage.completion_tokens = 10
+
+    mock_chunks = [
+        Mock(
+            choices=[
+                Mock(
+                    delta=Mock(
+                        role="assistant",
+                        content="Hello",
+                        model_dump=Mock(return_value={"role": "assistant", "content": "Hello"}),
+                    ),
+                    finish_reason="stop",
+                    logprobs=None,
+                )
+            ],
+            usage=None,
+        ),
+        # Final chunk with usage data
+        Mock(
+            choices=[],
+            usage=mock_usage,
+        ),
+    ]
+
+    with patch("databricks_langchain.chat_models.get_openai_client") as mock_get_client:
+        mock_client = Mock()
+        mock_get_client.return_value = mock_client
+        mock_client.chat.completions.create.return_value = iter(mock_chunks)
+
+        llm = ChatDatabricks(model="test-model")
+        messages = [HumanMessage(content="Hello")]
+
+        chunks = list(llm.stream(messages, stream_usage=False))
+
+        # Should get content chunks only
+        content_chunks = [chunk for chunk in chunks if chunk.content != ""]
+        assert len(content_chunks) == 1
+
+        # Should NOT emit a usage chunk when stream_usage=False
+        usage_chunks = [
+            chunk for chunk in chunks if chunk.content == "" and chunk.usage_metadata is not None
+        ]
+        assert len(usage_chunks) == 0, (
+            f"Expected 0 usage chunks when stream_usage=False, got {len(usage_chunks)}"
+        )
+
+
 class GetWeather(BaseModel):
     """Get the current weather in a given location"""