Databricks AI Gateway

from openai import OpenAI
from lago_agent_sdk import LagoSDK
import json

sdk = LagoSDK(api_key="<YOUR_LAGO_API_KEY>", default_subscription_id="sub_acme")

tags = json.dumps({"lago_subscription": "sub_acme", "team": "search"})
client = sdk.wrap(OpenAI(
    api_key="<DATABRICKS_TOKEN>",
    base_url="https://<workspace>.cloud.databricks.com/ai-gateway/mlflow/v1",
    default_headers={"Databricks-Ai-Gateway-Request-Tags": tags},
))

client.chat.completions.create(
    model="system.ai.llama-4-maverick",
    messages=[{"role": "user", "content": "Hello"}],
    max_tokens=400,
)
sdk.flush()

from anthropic import Anthropic

client = sdk.wrap(Anthropic(
    api_key="unused",                      # Databricks supplies the real credential
    base_url="https://<workspace>.cloud.databricks.com/ai-gateway/anthropic",
    default_headers={
        "Authorization": f"Bearer {databricks_token}",
        "Databricks-Model-Provider-Service": "<unity-catalog-credential>",
        "Databricks-Ai-Gateway-Request-Tags": tags,
    },
))

client.messages.create(
    model="claude-sonnet-4-5",
    max_tokens=400,
    messages=[{"role": "user", "content": "Hello"}],
)

from openai import OpenAI

# Same class as hosted — only the base_url differs.
client = sdk.wrap(OpenAI(
    api_key="<DATABRICKS_TOKEN>",
    base_url="https://<workspace>.cloud.databricks.com/ai-gateway/openai/v1",
    default_headers={
        "Databricks-Model-Provider-Service": "<unity-catalog-credential>",
        "Databricks-Ai-Gateway-Request-Tags": tags,
    },
))

client.chat.completions.create(model="gpt-4o", messages=[{"role": "user", "content": "Hello"}])

from lago_agent_sdk.gateway.databricks import DatabricksSource

source = DatabricksSource(
    host="https://<workspace>.cloud.databricks.com",
    token="<DATABRICKS_TOKEN>",
    warehouse_id="<WAREHOUSE_ID>",
)
# ...or DatabricksSource.from_env(), which reads
# DATABRICKS_HOST, DATABRICKS_TOKEN and DATABRICKS_WAREHOUSE_ID.

rows = list(source.read_usage("7 days"))

counts = sdk.backfill_databricks(rows, default_subscription="sub_acme")
assert sdk.flush(timeout=30.0)
print(counts)   # {'cost': 60, 'tokens': 88, 'skipped': 0}

import { DatabricksSource } from "lago-agent-sdk/gateway";

const source = new DatabricksSource(
  "https://<workspace>.cloud.databricks.com",
  process.env.DATABRICKS_TOKEN!,
  process.env.DATABRICKS_WAREHOUSE_ID!,
);
// ...or DatabricksSource.fromEnv(), which reads the same three variables.

const rows = await source.readUsage("7 days");
const counts = await sdk.backfillDatabricks(rows, "7 days", {
  defaultSubscription: "sub_acme",
});
await sdk.flush(30_000);