Skip to content

ChatDatabricks: UsageMetadata inflates input_tokens when prompt caching is used #446

Description

@anguyen

UsageMetadata emitted by ChatDatabricks double-counts cached input tokens in UsageMetadata.input_tokens. An example of this (adapted from the LangChain docs) is:

import os

import requests
from databricks.sdk import WorkspaceClient
from databricks_langchain import ChatDatabricks
from langchain_core.messages import AnyMessage, HumanMessage, SystemMessage

def create_model() -> ChatDatabricks:
    workspace_client = WorkspaceClient(
        host=os.getenv("DATABRICKS_HOST"),
        client_id=os.getenv("DATABRICKS_CLIENT_ID"),
        client_secret=os.getenv("DATABRICKS_CLIENT_SECRET"),
        auth_type="oauth-m2m",
    )
    chat_databricks = ChatDatabricks(
        workspace_client=workspace_client,
        endpoint=os.getenv("DATABRICKS_MODEL_ENDPOINT"),
        temperature=0,
        max_tokens=1024,
    )
    return chat_databricks

def create_messages() -> list[AnyMessage]:
    get_response = requests.get("https://raw.githubusercontent.com/langchain-ai/langchain/b476fdb54aa6e6f5f0b24a68c2f4a94e43b369f9/README.md")
    get_response.raise_for_status()
    readme = get_response.text

    return [
        SystemMessage(content=[
            {
                "type": "text",
                "text": "You are a technology expert.",
            },
            {
                "type": "text",
                "text": f"{readme}",
                "cache_control": {"type": "ephemeral"},
            },
        ]),
        HumanMessage(content="What is Langchain?"),
    ]

def main():
    chat_databricks = create_model()
    messages = create_messages()
    
    databricks_response = chat_databricks.invoke(messages)
    print(databricks_response.usage_metadata)

if __name__ == "__main__":
    main()

When run using Claude Sonnet 4.5 as the model, the result is:

{'input_tokens': 5654, 'output_tokens': 399, 'total_tokens': 3232, 'input_token_details': {'cache_creation': 2821, 'cache_read': 0}}

As you can see, the input_tokens plus output_tokens does not equal total_tokens. Subtracting the cache_creation count from input_tokens fixes this discrepancy.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions