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.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

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