[ChatDatabricks] Token usage counts are inflated when using streaming

[lucaboulard]

Description

When using ChatDatabricks inside a LangGraph graph with stream_mode=["messages", "values"],
the usage_metadata on the final AIMessage in the graph state yields significantly
inflated token counts compared to the correct values returned by a plain .invoke() call
on the same prompt.

The root cause is in _stream(): usage is emitted twice with cumulative values:

1. Embedded in each intermediate chunk via _convert_dict_to_message_chunk(..., usage=usage) — each chunk carries the cumulative usage up to that point
2. Again in a dedicated final usage-only chunk via _build_usage_chunk_from_completions(final_usage)

Since AIMessageChunk.__add__ sums usage_metadata across all chunks, the final
AIMessage reconstructed by LangGraph ends up with heavily inflated counts.
The same bug is present in _astream().

Steps to Reproduce

from typing import Annotated
from typing_extensions import TypedDict
from functools import reduce

from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages
from langchain_core.messages import HumanMessage
from databricks_langchain import ChatDatabricks

class State(TypedDict):
    messages: Annotated[list, add_messages]

llm = ChatDatabricks(
    endpoint="<your-endpoint>",
    temperature=0,
    max_tokens=512,
)

def call_llm(state: State) -> State:
    return {"messages": [llm.invoke(state["messages"])]}

builder = StateGraph(State)
builder.add_node("llm", call_llm)
builder.add_edge(START, "llm")
builder.add_edge("llm", END)
graph = builder.compile()

prompt = "Explain what Unity Catalog does in 3 bullet points."

# --- invoke: correct usage metadata---
final_state = graph.invoke({"messages": [HumanMessage(content=prompt)]})
print("invoke usage:", final_state["messages"][-1].usage_metadata)

# --- stream: inflated usage metadata---
graph_state = None
for mode, event in graph.stream(
    {"messages": [HumanMessage(content=prompt)]},
    stream_mode=["messages", "values"],
):
    if mode == "values":
        graph_state = event

print("stream usage:", graph_state["messages"][-1].usage_metadata)

Expected and actual behaviors

Both invoke and streaming calls should return consistent token counts for the same prompt and response. eg:

invoke usage: {'input_tokens': 30,   'output_tokens': 220, 'total_tokens': 250}
stream usage: {'input_tokens': 30,   'output_tokens': 220, 'total_tokens': 250}

Contrary you get something like:

invoke usage: {'input_tokens': 30,   'output_tokens': 220, 'total_tokens': 250}
stream usage: {'input_tokens': 2160, 'output_tokens': 440, 'total_tokens': 490}

Note: the figures above are illustrative. Actual values were obtained using
claude-sonnet-4-5 as the served model and a prompt written in Italian.

input_tokens is multiplied by approximately the number of streamed chunks, confirming that intermediate chunks carry cumulative values that get summed on top of the already-correct final usage chunk.

The bug is triggered specifically by stream_mode=["messages", "values"]. Using stream_mode="values" alone does not trigger
it, because in that case LangGraph does not enable token-level streaming and the
node falls back to llm.invoke() internally.

Suggested Fix

Strip usage_metadata from all intermediate chunks in _stream() and _astream(),
preserving it only on the final dedicated usage chunk.
This can be implemented cleanly with a one-chunk lookahead buffer, as done in the following example where a subclass of ChatDatabricks is created

class PatchedChatDatabricks(ChatDatabricks):
  
  def _stream(self, messages, stop=None, run_manager=None, **kwargs):
      prev = None
      for chunk in super()._stream(messages, stop=stop, run_manager=run_manager, **kwargs):
          if prev is not None:
              prev.message.usage_metadata = None  # strip intermediate usage
              yield prev
          prev = chunk
      if prev is not None:
          yield prev  # last chunk: usage_metadata preserved intact

The same pattern applies to _astream() using async for.

Environment

• databricks-langchain version: 0.19.0
• langchain-core version: 1.4.0
• langgraph version: 1.2.0
• Databricks Serverless Environment version: 5

[s-a-nersisyan]

Thanks for raising this, I almost got a heart attack today seeing trace with $300 cost (token count was inflated to like 100 mln tokens haha)