fix: include token usage in Databricks streaming responses (#7959)

2026-07-17 12:56:20 +02:00 · 2026-03-18 16:43:35 +01:00
parent 475968db64
commit c2d10cea9b
1 changed files with 12 additions and 0 deletions
@@ -389,6 +389,18 @@ impl Provider for DatabricksProvider {
                .unwrap()
                .insert("stream".to_string(), Value::Bool(true));

+            if let Some(opts) = payload
+                .get_mut("stream_options")
+                .and_then(|v| v.as_object_mut())
+            {
+                opts.entry("include_usage").or_insert(json!(true));
+            } else {
+                payload
+                    .as_object_mut()
+                    .unwrap()
+                    .insert("stream_options".to_string(), json!({"include_usage": true}));
+            }
+
            let mut log = RequestLog::start(model_config, &payload)?;
            let response = self
                .with_retry(|| async {