-
Notifications
You must be signed in to change notification settings - Fork 41
Open
Description
We are using ChatDatabricks to integrate with Azure Open AI models, added as serving endpoints in Databricks. We are using this to stream answers in a chatbot application. We need to extract token usage and model used from the response_metadata (as we are using a model-router) however when streaming, the response_metadata is not populated.
This was partially addressed in PR: #176 which fixed the streaming tokens issue for Databricks Foundation models but it's not working for external model providers. Our app uses Langchain but I've created a very short MRE below to demonstrate:
from databricks_langchain import ChatDatabricks
llm = ChatDatabricks(endpoint=OPENAI_ROUTER, stream_usage=True)
messages = [
("system", "You are a helpful translator. Translate the user sentence to French."),
("human", "I love programming."),
]
print('################## STREAMING ##########################')
for chunk in llm.stream(messages):
print(chunk)
print('################## INVOKING ##########################')
print(llm.invoke(messages))
RESULT:
response_metadata={} for each streamed chunk but it is fully populated with token counts and actual model (not the router name) when invoked.
################## STREAMING ##########################
content='' additional_kwargs={} response_metadata={} id='lc_run--019b02ca-1a7d-7251-a0c4-9a669dd3afa9'
content='J' additional_kwargs={} response_metadata={} id='lc_run--019b02ca-1a7d-7251-a0c4-9a669dd3afa9'
content='’aime' additional_kwargs={} response_metadata={} id='lc_run--019b02ca-1a7d-7251-a0c4-9a669dd3afa9'
content=' la' additional_kwargs={} response_metadata={} id='lc_run--019b02ca-1a7d-7251-a0c4-9a669dd3afa9'
content=' programmation' additional_kwargs={} response_metadata={} id='lc_run--019b02ca-1a7d-7251-a0c4-9a669dd3afa9'
content='.' additional_kwargs={} response_metadata={} id='lc_run--019b02ca-1a7d-7251-a0c4-9a669dd3afa9'
content='' additional_kwargs={} response_metadata={'finish_reason': 'stop'} id='lc_run--019b02ca-1a7d-7251-a0c4-9a669dd3afa9'
content='' additional_kwargs={} response_metadata={} id='lc_run--019b02ca-1a7d-7251-a0c4-9a669dd3afa9' chunk_position='last'
################## INVOKING ##########################
content='J’adore programmer.' additional_kwargs={} response_metadata={'usage': {'prompt_tokens': 27, 'completion_tokens': 207, 'total_tokens': 234}, 'prompt_tokens': 27, 'completion_tokens': 207, 'total_tokens': 234, 'model': 'gpt-5-2025-08-07', 'model_name': 'gpt-5-2025-08-07', 'finish_reason': 'stop'} id='lc_run--019b02ca-2456-7022-b513-0d73cddce31a-0
Metadata
Metadata
Assignees
Labels
No labels