From f948fa886787af05795e90693bca61aaf96947fc Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Fri, 9 Oct 2026 20:52:33 -0700
Subject: [PATCH] test(decisions): move the Databricks OpenAI-format wire row
 to the ai_decide route

---
 .../test_decisions_openai_format_wire.py      | 31 ++++++++++++++++---
 1 file changed, 26 insertions(+), 5 deletions(-)

diff --git a/tests/integration/providers/test_decisions_openai_format_wire.py b/tests/integration/providers/test_decisions_openai_format_wire.py
index f0e7163e264..30d072c2761 100644
--- a/tests/integration/providers/test_decisions_openai_format_wire.py
+++ b/tests/integration/providers/test_decisions_openai_format_wire.py
@@ -88,6 +88,17 @@ _USAGE: Final[dict[str, JsonValue]] = {
     "output_tokens_details": {"reasoning_tokens": 0},
     "total_tokens": _INPUT_TOKENS + _OUTPUT_TOKENS,
 }
+_NO_USAGE: Final[dict[str, JsonValue]] = {
+    "input_tokens": 0,
+    "input_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 0},
+    "output_tokens": 0,
+    "output_tokens_details": {"reasoning_tokens": 0},
+    "total_tokens": 0,
+}
+_AI_DECIDE_ANSWERS: Final[dict[str, JsonValue]] = {
+    **_SYSTEM_ONE_ANSWERS,
+    "defect": {"type": "noul", "probability": 0.93},
+}
 _SDK_QUESTIONS: Final = TypeAdapter(list[dict[str, object]]).validate_python(_QUESTIONS)
 _SPEND_QUERY: Final = (
     "SELECT spend, status, call_type, model_group, custom_llm_provider, api_base, prompt_tokens, completion_tokens, "
@@ -105,15 +116,20 @@ class _Provider:
     wraps_result: bool
     cost_map_key: str | None
     speaks_openai: bool = False
+    ai_decide: bool = False
 
     def upstream_body(self) -> dict[str, JsonValue]:
         if self.speaks_openai:
             return {"model": self.body_model, "input": _INPUT, "questions": _QUESTIONS}
+        if self.ai_decide:
+            return {"state": _INPUT, "questions": _SYSTEM_ONE_QUESTIONS}
         return {"model": self.body_model, "state": _INPUT, "questions": _SYSTEM_ONE_QUESTIONS}
 
     def upstream_reply(self) -> dict[str, JsonValue]:
         if self.speaks_openai:
             return {"model": self.body_model, "answers": _ANSWERS, "usage": _USAGE}
+        if self.ai_decide:
+            return {"response": {"answers": _AI_DECIDE_ANSWERS}, "metadata": {"version": "1.0"}}
         answer: Final[dict[str, JsonValue]] = {
             "model": self.body_model,
             "answers": _SYSTEM_ONE_ANSWERS,
@@ -122,8 +138,13 @@ class _Provider:
         return {"result": answer, "success": True} if self.wraps_result else answer
 
     def litellm_response(self) -> dict[str, JsonValue]:
+        if self.ai_decide:
+            return {"model": self.model, "answers": _ANSWERS, "usage": _NO_USAGE}
         return {"model": self.body_model, "answers": _ANSWERS, "usage": _USAGE}
 
+    def billed_tokens(self) -> tuple[int, int]:
+        return (0, 0) if self.ai_decide else (_INPUT_TOKENS, _OUTPUT_TOKENS)
+
 
 _PROVIDERS: Final = (
     _Provider(
@@ -160,12 +181,13 @@ _PROVIDERS: Final = (
     _Provider("hosted_vllm", "hosted_vllm/Qwen/Qwen3-0.6B", "/v1/systemone", "Qwen/Qwen3-0.6B", None, False, None),
     _Provider(
         "databricks",
-        "databricks/databricks-openjev-qwen35-4b",
-        "/databricks-openjev-qwen35-4b/invocations",
-        "databricks-openjev-qwen35-4b",
+        "databricks/ai_decide",
+        "/api/2.0/ai-functions/ai-decide",
+        "ai_decide",
         _API_KEY,
         False,
         None,
+        ai_decide=True,
     ),
     _Provider(
         "azure_ai", "azure_ai/decision-1", "/providers/microsoft/v1/systemone", "decision-1", _API_KEY, False, None
@@ -351,8 +373,7 @@ def test_each_provider_gets_its_own_path_key_and_body_and_is_billed_from_the_cos
             provider.name,
             model,
             f"{handle.api_base()}{provider.path}",
-            _INPUT_TOKENS,
-            _OUTPUT_TOKENS,
+            *provider.billed_tokens(),
         )
         assert math.isclose(_number(row["spend"]), expected_spend, rel_tol=1e-9), row