scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-24 05:45:11 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 72

test(decisions): move the Databricks OpenAI-format wire row to the ai_decide route

Browse source
This commit is contained in:
mateo-berri 2026-10-09 20:52:33 -07:00
parent d4ded5398e
commit f948fa8867
1 changed files with 26 additions and 5 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

31
tests/integration/providers/test_decisions_openai_format_wire.py
View file

@ -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