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

    Fields

    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions