mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-11 08:38:24 +00:00
508a98586e
* fix(databricks): join plain text blocks for non-Claude json_object requests * fix(databricks): join only the last message for json_object and drop cache_control under drop_params Databricks checks for "json" only in the last message, so earlier turns now go out as sent and only the last message, plus any system messages merged into it, is joined into a string. Under drop_params a text block's cache_control is dropped so that message can be joined too, and without it the request still goes out as a list. * fix(databricks): skip dropped assistant turns and blank parts when joining for json_object The json_object join picked the last message before the request transform drops bare assistant turns, so a trailing thinking-only or reasoning-only turn (a Responses input ending in a reasoning item included) left the real last user message as a list and Databricks still returned its 400. The pick now skips the turns that transform drops. The join also ran when only one part had text, which sent whitespace-only parts the existing empty block filter used to drop on requests that already worked. It now joins only two or more non-blank parts, so a request with one non-blank part goes out exactly as before. Without drop_params, a message with a message-level cache_control now goes out as sent instead of being rewritten into one cached block. * fix(databricks): join json_object text parts only on models Databricks hosts Databricks checks for the word "json" only on the models it hosts (databricks-* pay per token names and system.ai.* AI Gateway names). An external model endpoint skips that check and takes a list of text parts today, so its request now goes out as sent instead of joined. The drop_params read now goes through drops_unsupported_params, and a test pins that a text block with a truthy non-string text is kept. * fix(databricks): reject a text part holding a non-string locally with a clear 400 * refactor(databricks): reject a non-string text part at the transform_request boundary * fix(databricks): read pydantic message objects in the non-string text check --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
2194 lines
90 KiB
Python
2194 lines
90 KiB
Python
import itertools
|
|
import json
|
|
import uuid
|
|
from collections.abc import Mapping
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Final, Literal, TypeAlias
|
|
|
|
import anthropic
|
|
import openai
|
|
import pytest
|
|
from integration._support.client import Gateway, eventually
|
|
from integration._support.database import read_rows
|
|
from integration._support.process import owned_proxy_process
|
|
from integration._support.wire import Reply, Request, wire_server
|
|
from integration.translation.chat_completions.cross_model.contracts import assert_completion
|
|
from integration.translation.response_contracts import responses_structure
|
|
from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter
|
|
|
|
_BACKEND: Final = "databricks-glm-5-2"
|
|
_API_KEY: Final = "synthetic-databricks-key"
|
|
_PROMPT: Final = "Summarise the cached briefing in one sentence."
|
|
_JSON_SCHEMA_BACKEND: Final = "databricks-qwen35-122b-a10b"
|
|
_CLAUDE_BACKEND: Final = "databricks-claude-haiku-4-5"
|
|
_JSON_SCHEMA_PROMPT: Final = "Return the requested JSON."
|
|
_JSON_CONTENT: Final = '{"p":{"name":"Ada"}}'
|
|
_PROVIDER_USAGE: Final[Mapping[str, JsonValue]] = {
|
|
"prompt_tokens": 12011,
|
|
"completion_tokens": 8,
|
|
"total_tokens": 12019,
|
|
"cache_read_input_tokens": 12002,
|
|
"cache_creation_input_tokens": 0,
|
|
}
|
|
_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
|
|
_PERSON_SCHEMA: Final[Mapping[str, JsonValue]] = {
|
|
"$defs": {
|
|
"Person": {
|
|
"type": "object",
|
|
"properties": {"name": {"type": "string"}},
|
|
"required": ["name"],
|
|
}
|
|
},
|
|
"type": "object",
|
|
"properties": {"p": {"$ref": "#/$defs/Person"}},
|
|
"required": ["p"],
|
|
}
|
|
_JSON_SCHEMA_RESPONSE_FORMAT: Final[Mapping[str, JsonValue]] = {
|
|
"type": "json_schema",
|
|
"json_schema": {"name": "P", "strict": True, "schema": _PERSON_SCHEMA},
|
|
}
|
|
_MESSAGES_PERSON_SCHEMA: Final[Mapping[str, JsonValue]] = {
|
|
"$defs": {
|
|
"Person": {
|
|
"type": "object",
|
|
"properties": {"name": {"type": "string"}},
|
|
"required": ["name"],
|
|
"additionalProperties": False,
|
|
}
|
|
},
|
|
"type": "object",
|
|
"properties": {"p": {"$ref": "#/$defs/Person"}},
|
|
"required": ["p"],
|
|
"additionalProperties": False,
|
|
}
|
|
_MESSAGES_RESPONSE_FORMAT: Final[Mapping[str, JsonValue]] = {
|
|
"type": "json_schema",
|
|
"json_schema": {
|
|
"name": "structured_output",
|
|
"schema": _MESSAGES_PERSON_SCHEMA,
|
|
"strict": True,
|
|
},
|
|
}
|
|
_JSON_TOOL_CALL_MESSAGE: Final[dict[str, JsonValue]] = {
|
|
"role": "assistant",
|
|
"content": None,
|
|
"tool_calls": [
|
|
{
|
|
"id": "json-tool-call",
|
|
"type": "function",
|
|
"function": {"name": "json_tool_call", "arguments": _JSON_CONTENT},
|
|
}
|
|
],
|
|
}
|
|
|
|
|
|
def test_databricks_responses_reasoning_summary_keeps_reasoning_effort_a_string(gateway: Gateway) -> None:
|
|
def respond(request: Request) -> Reply:
|
|
assert request.method == "POST" and request.target == "/chat/completions"
|
|
assert _JSON_OBJECT.validate_json(request.body) == {
|
|
"model": _BACKEND,
|
|
"messages": [{"role": "user", "content": "Say hello."}],
|
|
"max_tokens": 64,
|
|
"reasoning_effort": "low",
|
|
}
|
|
return Reply(body=_chat_completion_response(_BACKEND, {"role": "assistant", "content": "Hello."}, "stop"))
|
|
|
|
with wire_server(respond) as wire, gateway.scenario() as scenario:
|
|
alias: Final = scenario.model(model=f"databricks/{_BACKEND}", api_base=wire.url, api_key=_API_KEY)
|
|
response: Final = gateway.request(
|
|
"POST",
|
|
"/v1/responses",
|
|
{
|
|
"model": alias,
|
|
"input": "Say hello.",
|
|
"max_output_tokens": 64,
|
|
"reasoning": {"effort": "low", "summary": "auto"},
|
|
},
|
|
)
|
|
assert response.status_code == 200, response.text
|
|
body: Final = _JSON_OBJECT.validate_json(response.content)
|
|
responses_structure(body)
|
|
assert body["usage"] == {
|
|
"input_tokens": 1,
|
|
"output_tokens": 1,
|
|
"total_tokens": 2,
|
|
"input_tokens_details": None,
|
|
"output_tokens_details": None,
|
|
"cost": None,
|
|
}
|
|
assert len(wire.drain()) == 1
|
|
|
|
|
|
def _chat_completion_response(model: str, message: dict[str, JsonValue], finish_reason: str) -> bytes:
|
|
response: Final = {
|
|
"id": "databricks-json-schema-response",
|
|
"object": "chat.completion",
|
|
"created": 1,
|
|
"model": model,
|
|
"choices": [{"index": 0, "message": message, "finish_reason": finish_reason}],
|
|
"usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
|
|
}
|
|
return json.dumps(response).encode()
|
|
|
|
|
|
class _PromptTokensDetails(BaseModel):
|
|
model_config = ConfigDict(extra="ignore")
|
|
cached_tokens: int | None = None
|
|
|
|
|
|
class _Usage(BaseModel):
|
|
model_config = ConfigDict(extra="ignore")
|
|
prompt_tokens: int
|
|
completion_tokens: int
|
|
total_tokens: int
|
|
prompt_tokens_details: _PromptTokensDetails | None = None
|
|
|
|
|
|
class _Delta(BaseModel):
|
|
model_config = ConfigDict(extra="ignore")
|
|
content: str | None = None
|
|
|
|
|
|
class _Choice(BaseModel):
|
|
model_config = ConfigDict(extra="ignore")
|
|
delta: _Delta
|
|
|
|
|
|
class _Chunk(BaseModel):
|
|
model_config = ConfigDict(extra="ignore")
|
|
id: str
|
|
choices: tuple[_Choice, ...]
|
|
usage: _Usage | None = None
|
|
|
|
|
|
def _frame(
|
|
identity: str,
|
|
choices: list[Mapping[str, object]],
|
|
usage: Mapping[str, JsonValue] | None = None,
|
|
model: str = _BACKEND,
|
|
) -> bytes:
|
|
value: Final = {
|
|
"id": identity,
|
|
"object": "chat.completion.chunk",
|
|
"created": 1,
|
|
"model": model,
|
|
"choices": choices,
|
|
**({} if usage is None else {"usage": usage}),
|
|
}
|
|
return b"data: " + json.dumps(value).encode() + b"\n\n"
|
|
|
|
|
|
@pytest.mark.covers("other.provider_wire.databricks.stream_usage_and_cache_reads_reach_client_and_spend_log")
|
|
def test_databricks_stream_final_usage_chunk_reaches_client_and_spend_log(gateway: Gateway) -> None:
|
|
identity: Final = f"databricks-stream-{uuid.uuid4().hex}"
|
|
frames: Final = (
|
|
_frame(
|
|
identity, [{"index": 0, "delta": {"role": "assistant", "content": "The briefing "}, "finish_reason": None}]
|
|
),
|
|
_frame(identity, [{"index": 0, "delta": {"content": "is short."}, "finish_reason": None}]),
|
|
_frame(identity, [{"index": 0, "delta": {}, "finish_reason": "stop"}]),
|
|
_frame(identity, [], usage=_PROVIDER_USAGE),
|
|
b"data: [DONE]\n\n",
|
|
)
|
|
|
|
def respond(request: Request) -> Reply:
|
|
assert request.method == "POST"
|
|
assert request.target == "/chat/completions"
|
|
assert request.headers["authorization"] == f"Bearer {_API_KEY}"
|
|
body: Final = _JSON_OBJECT.validate_json(request.body)
|
|
assert body["model"] == _BACKEND
|
|
assert body["messages"] == [{"role": "user", "content": _PROMPT}]
|
|
assert body["stream"] is True
|
|
return Reply(content_type="text/event-stream", chunks=frames)
|
|
|
|
with wire_server(respond) as wire, gateway.scenario() as scenario:
|
|
model: Final = scenario.model(model=f"databricks/{_BACKEND}", api_base=wire.url, api_key=_API_KEY)
|
|
with gateway.client.stream(
|
|
"POST",
|
|
"/v1/chat/completions",
|
|
json={
|
|
"model": model,
|
|
"messages": [{"role": "user", "content": _PROMPT}],
|
|
"stream": True,
|
|
"stream_options": {"include_usage": True},
|
|
},
|
|
headers={"Authorization": f"Bearer {gateway.key}"},
|
|
) as response:
|
|
assert response.status_code == 200, response.read()
|
|
lines: Final = tuple(line for line in response.iter_lines() if line.startswith("data: "))
|
|
assert lines[-1] == "data: [DONE]", lines
|
|
chunks: Final = tuple(_Chunk.model_validate_json(line.removeprefix("data: ")) for line in lines[:-1])
|
|
assert {chunk.id for chunk in chunks} == {identity}
|
|
assert (
|
|
"".join(choice.delta.content or "" for chunk in chunks for choice in chunk.choices)
|
|
== "The briefing is short."
|
|
)
|
|
usages: Final = tuple(chunk.usage for chunk in chunks if chunk.usage is not None)
|
|
assert len(usages) == 1, lines
|
|
assert (
|
|
usages[0].prompt_tokens,
|
|
usages[0].completion_tokens,
|
|
usages[0].total_tokens,
|
|
usages[0].prompt_tokens_details.cached_tokens if usages[0].prompt_tokens_details is not None else None,
|
|
) == (12011, 8, 12019, 12002), lines
|
|
assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/chat/completions")]
|
|
rows: Final = eventually(
|
|
lambda: read_rows(
|
|
'SELECT prompt_tokens, completion_tokens, total_tokens FROM "LiteLLM_SpendLogs" WHERE request_id=%s',
|
|
(identity,),
|
|
),
|
|
lambda values: len(values) == 1,
|
|
seconds=70,
|
|
)
|
|
assert (rows[0]["prompt_tokens"], rows[0]["completion_tokens"], rows[0]["total_tokens"]) == (12011, 8, 12019)
|
|
|
|
|
|
def _shape_reply(identity: str, stream: bool) -> Reply:
|
|
if stream:
|
|
chunks: Final = (
|
|
_frame(
|
|
identity,
|
|
[{"index": 0, "delta": {"role": "assistant", "content": "ok"}, "finish_reason": None}],
|
|
),
|
|
_frame(identity, [{"index": 0, "delta": {}, "finish_reason": "stop"}]),
|
|
b"data: [DONE]\n\n",
|
|
)
|
|
return Reply(content_type="text/event-stream", chunks=chunks)
|
|
|
|
body: Final = {
|
|
"id": identity,
|
|
"object": "chat.completion",
|
|
"created": 1,
|
|
"model": _BACKEND,
|
|
"choices": [
|
|
{
|
|
"index": 0,
|
|
"message": {"role": "assistant", "content": "ok"},
|
|
"finish_reason": "stop",
|
|
}
|
|
],
|
|
"usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
|
|
}
|
|
return Reply(body=json.dumps(body).encode())
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"case,endpoint,provider_model,request_body,expected_messages,responses_mode,upstream_path",
|
|
[
|
|
pytest.param(
|
|
"chat_nonstream",
|
|
"/v1/chat/completions",
|
|
"databricks-gemma-3-12b",
|
|
{
|
|
"messages": [{"role": "user", "content": [{"type": "text", "text": "Reply in JSON"}]}],
|
|
"response_format": {"type": "json_object"},
|
|
},
|
|
[{"role": "user", "content": "Reply in JSON"}],
|
|
False,
|
|
"/chat/completions",
|
|
id="A1-chat-nonstream",
|
|
),
|
|
pytest.param(
|
|
"chat_stream",
|
|
"/v1/chat/completions",
|
|
"databricks-gemma-3-12b",
|
|
{
|
|
"messages": [{"role": "user", "content": [{"type": "text", "text": "Reply in JSON"}]}],
|
|
"response_format": {"type": "json_object"},
|
|
"stream": True,
|
|
},
|
|
[{"role": "user", "content": "Reply in JSON"}],
|
|
False,
|
|
"/chat/completions",
|
|
id="A2-chat-stream",
|
|
),
|
|
pytest.param(
|
|
"responses-mode-bridge",
|
|
"/v1/chat/completions",
|
|
"databricks-qwen3-next-80b-a3b-instruct",
|
|
{
|
|
"messages": [{"role": "user", "content": "Reply in JSON"}],
|
|
"response_format": {"type": "json_object"},
|
|
},
|
|
[{"role": "user", "content": "Reply in JSON"}],
|
|
True,
|
|
"/chat/completions",
|
|
id="A3-responses-mode",
|
|
),
|
|
pytest.param(
|
|
"responses-mode-bridge-unity-catalog",
|
|
"/v1/chat/completions",
|
|
"system.ai.deepseek-v4-1-flash",
|
|
{
|
|
"messages": [{"role": "user", "content": "Reply in JSON with key a. hi"}],
|
|
"response_format": {"type": "json_object"},
|
|
},
|
|
[{"role": "user", "content": "Reply in JSON with key a. hi"}],
|
|
True,
|
|
"/ai-gateway/mlflow/v1/chat/completions",
|
|
id="A3b-responses-mode-unity-catalog-model",
|
|
),
|
|
pytest.param(
|
|
"messages-endpoint",
|
|
"/v1/messages",
|
|
"databricks-gemma-3-12b",
|
|
{
|
|
"max_tokens": 32,
|
|
"messages": [{"role": "user", "content": [{"type": "text", "text": "Reply in JSON"}]}],
|
|
},
|
|
[{"role": "user", "content": "Reply in JSON"}],
|
|
False,
|
|
"/chat/completions",
|
|
id="A4-messages",
|
|
),
|
|
pytest.param(
|
|
"responses-endpoint",
|
|
"/v1/responses",
|
|
"databricks-gemma-3-12b",
|
|
{"input": [{"role": "user", "content": [{"type": "input_text", "text": "Reply in JSON"}]}]},
|
|
[{"role": "user", "content": "Reply in JSON"}],
|
|
False,
|
|
"/chat/completions",
|
|
id="A5-responses-list-input",
|
|
),
|
|
pytest.param(
|
|
"json-object-two-blocks",
|
|
"/v1/chat/completions",
|
|
"databricks-meta-llama-3-3-70b-instruct",
|
|
{
|
|
"messages": [
|
|
{"role": "user", "content": [{"type": "text", "text": "Reply"}, {"type": "text", "text": " in JSON"}]}
|
|
],
|
|
"response_format": {"type": "json_object"},
|
|
},
|
|
[{"role": "user", "content": "Reply in JSON"}],
|
|
False,
|
|
"/chat/completions",
|
|
id="A6-chat-json-object-two-blocks-joined",
|
|
),
|
|
pytest.param(
|
|
"json-object-two-blocks-stream",
|
|
"/v1/chat/completions",
|
|
"databricks-qwen35-122b-a10b",
|
|
{
|
|
"messages": [
|
|
{"role": "user", "content": [{"type": "text", "text": "Reply"}, {"type": "text", "text": " in JSON"}]}
|
|
],
|
|
"response_format": {"type": "json_object"},
|
|
"stream": True,
|
|
},
|
|
[{"role": "user", "content": "Reply in JSON"}],
|
|
False,
|
|
"/chat/completions",
|
|
id="A7-chat-stream-json-object-two-blocks-joined",
|
|
),
|
|
pytest.param(
|
|
"json-object-history",
|
|
"/v1/chat/completions",
|
|
"databricks-meta-llama-3-3-70b-instruct",
|
|
{
|
|
"messages": [
|
|
{"role": "user", "content": [{"type": "text", "text": "Code: A"}, {"type": "text", "text": "BC"}]},
|
|
{"role": "assistant", "content": "Got it."},
|
|
{"role": "user", "content": [{"type": "text", "text": "Reply"}, {"type": "text", "text": " in JSON"}]},
|
|
],
|
|
"response_format": {"type": "json_object"},
|
|
},
|
|
[
|
|
{"role": "user", "content": [{"type": "text", "text": "Code: A"}, {"type": "text", "text": "BC"}]},
|
|
{"role": "assistant", "content": "Got it."},
|
|
{"role": "user", "content": "Reply in JSON"},
|
|
],
|
|
False,
|
|
"/chat/completions",
|
|
id="A8-chat-json-object-earlier-turns-kept",
|
|
),
|
|
pytest.param(
|
|
"json-object-trailing-system",
|
|
"/v1/chat/completions",
|
|
"databricks-meta-llama-3-3-70b-instruct",
|
|
{
|
|
"messages": [
|
|
{"role": "user", "content": "Say hi"},
|
|
{"role": "system", "content": [{"type": "text", "text": "Alpha"}, {"type": "text", "text": " rule."}]},
|
|
{"role": "system", "content": [{"type": "text", "text": "Reply"}, {"type": "text", "text": " in JSON"}]},
|
|
],
|
|
"response_format": {"type": "json_object"},
|
|
},
|
|
[{"role": "user", "content": "Say hi"}, {"role": "system", "content": "Alpha rule.\n\nReply in JSON"}],
|
|
False,
|
|
"/chat/completions",
|
|
id="A8b-chat-json-object-trailing-system-keeps-break",
|
|
),
|
|
pytest.param(
|
|
"json-object-drop-params-cache-control",
|
|
"/v1/chat/completions",
|
|
"databricks-meta-llama-3-3-70b-instruct",
|
|
{
|
|
"messages": [
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{"type": "text", "text": "Reply"},
|
|
{"type": "text", "text": " in JSON", "cache_control": {"type": "ephemeral"}},
|
|
],
|
|
}
|
|
],
|
|
"response_format": {"type": "json_object"},
|
|
"drop_params": True,
|
|
},
|
|
[{"role": "user", "content": "Reply in JSON"}],
|
|
False,
|
|
"/chat/completions",
|
|
id="A8c-chat-json-object-drop-params-drops-cache-control",
|
|
),
|
|
pytest.param(
|
|
"json-object-trailing-thinking-turn",
|
|
"/v1/chat/completions",
|
|
"databricks-meta-llama-3-3-70b-instruct",
|
|
{
|
|
"messages": [
|
|
{"role": "user", "content": [{"type": "text", "text": "Reply"}, {"type": "text", "text": " in JSON"}]},
|
|
{
|
|
"role": "assistant",
|
|
"content": "",
|
|
"thinking_blocks": [{"type": "thinking", "thinking": "Plan.", "signature": "sig"}],
|
|
},
|
|
],
|
|
"response_format": {"type": "json_object"},
|
|
},
|
|
[{"role": "user", "content": "Reply in JSON"}],
|
|
False,
|
|
"/chat/completions",
|
|
id="A8d-chat-json-object-trailing-thinking-turn-dropped",
|
|
),
|
|
pytest.param(
|
|
"json-object-external-endpoint",
|
|
"/v1/chat/completions",
|
|
"openai-gpt-external",
|
|
{
|
|
"messages": [
|
|
{"role": "user", "content": [{"type": "text", "text": "Reply"}, {"type": "text", "text": " in JSON"}]}
|
|
],
|
|
"response_format": {"type": "json_object"},
|
|
},
|
|
[{"role": "user", "content": [{"type": "text", "text": "Reply"}, {"type": "text", "text": " in JSON"}]}],
|
|
False,
|
|
"/chat/completions",
|
|
id="A8e-chat-json-object-external-endpoint-kept-as-sent",
|
|
),
|
|
pytest.param(
|
|
"json-object-responses-two-parts",
|
|
"/v1/responses",
|
|
"databricks-qwen35-122b-a10b",
|
|
{
|
|
"input": [
|
|
{
|
|
"role": "user",
|
|
"content": [{"type": "input_text", "text": "Reply"}, {"type": "input_text", "text": " in JSON"}],
|
|
}
|
|
],
|
|
"text": {"format": {"type": "json_object"}},
|
|
},
|
|
[{"role": "user", "content": "Reply in JSON"}],
|
|
False,
|
|
"/chat/completions",
|
|
id="A9-responses-json-object-two-parts-joined",
|
|
),
|
|
pytest.param(
|
|
"json-object-responses-trailing-reasoning",
|
|
"/v1/responses",
|
|
"databricks-qwen35-122b-a10b",
|
|
{
|
|
"input": [
|
|
{
|
|
"role": "user",
|
|
"content": [{"type": "input_text", "text": "Reply"}, {"type": "input_text", "text": " in JSON"}],
|
|
},
|
|
{"type": "reasoning", "id": "rs_1", "summary": [{"type": "summary_text", "text": "Plan."}]},
|
|
],
|
|
"text": {"format": {"type": "json_object"}},
|
|
},
|
|
[{"role": "user", "content": "Reply in JSON"}],
|
|
False,
|
|
"/chat/completions",
|
|
id="A9b-responses-json-object-trailing-reasoning-dropped",
|
|
),
|
|
],
|
|
)
|
|
def test_databricks_non_claude_plain_text_content_reaches_upstream_as_string(
|
|
gateway: Gateway,
|
|
case: str,
|
|
endpoint: str,
|
|
provider_model: str,
|
|
request_body: dict[str, JsonValue],
|
|
expected_messages: list[JsonValue],
|
|
responses_mode: bool,
|
|
upstream_path: str,
|
|
) -> None:
|
|
identity: Final = f"databricks-plain-text-{case}-{uuid.uuid4().hex}"
|
|
streaming: Final = request_body.get("stream") is True
|
|
|
|
def respond(request: Request) -> Reply:
|
|
assert request.method == "POST"
|
|
body: Final = _JSON_OBJECT.validate_json(request.body)
|
|
assert body["messages"] == expected_messages, body["messages"]
|
|
return _shape_reply(identity, streaming)
|
|
|
|
with wire_server(respond) as wire, gateway.scenario() as scenario:
|
|
model: Final = scenario.model(
|
|
model=f"databricks/{provider_model}",
|
|
api_base=wire.url,
|
|
api_key=_API_KEY,
|
|
model_info={"mode": "responses"} if responses_mode else None,
|
|
)
|
|
headers: Final = {"anthropic-version": "2023-06-01"} if endpoint == "/v1/messages" else {}
|
|
body: Final = {"model": model, **request_body}
|
|
if streaming:
|
|
with gateway.client.stream(
|
|
"POST", endpoint, json=body, headers={"Authorization": f"Bearer {gateway.key}", **headers}
|
|
) as response:
|
|
assert response.status_code == 200, response.read()
|
|
lines: Final = tuple(line for line in response.iter_lines() if line.startswith("data: "))
|
|
assert lines[-1] == "data: [DONE]", lines
|
|
assert any('"content":"ok"' in line for line in lines), lines
|
|
else:
|
|
response: Final = gateway.request("POST", endpoint, body, headers=headers)
|
|
assert response.status_code == 200, response.text
|
|
assert _JSON_OBJECT.validate_json(response.content)["id"]
|
|
assert [(request.method, request.target) for request in wire.drain()] == [("POST", upstream_path)]
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"case,provider_model,content,expected_content",
|
|
[
|
|
pytest.param(
|
|
"cache-control",
|
|
"databricks-gemma-3-12b",
|
|
[{"type": "text", "text": "Reply in JSON", "cache_control": {"type": "ephemeral"}}],
|
|
[{"type": "text", "text": "Reply in JSON", "cache_control": {"type": "ephemeral"}}],
|
|
id="B1-cache-control-stays-list",
|
|
),
|
|
pytest.param(
|
|
"two-text-blocks",
|
|
"databricks-gemma-3-12b",
|
|
[{"type": "text", "text": "Reply"}, {"type": "text", "text": " in JSON"}],
|
|
[{"type": "text", "text": "Reply"}, {"type": "text", "text": " in JSON"}],
|
|
id="B2-two-text-blocks-stays-list",
|
|
),
|
|
pytest.param(
|
|
"text-and-image",
|
|
"databricks-gemma-3-12b",
|
|
[
|
|
{"type": "text", "text": "Reply in JSON"},
|
|
{"type": "image_url", "image_url": {"url": "https://example.com/image.png"}},
|
|
],
|
|
[
|
|
{"type": "text", "text": "Reply in JSON"},
|
|
{"type": "image_url", "image_url": {"url": "https://example.com/image.png"}},
|
|
],
|
|
id="B3-text-and-image-stays-list",
|
|
),
|
|
pytest.param(
|
|
"claude",
|
|
"databricks-claude-opus-5-5",
|
|
[{"type": "text", "text": "Reply in JSON"}],
|
|
[{"type": "text", "text": "Reply in JSON"}],
|
|
id="B4-claude-stays-list",
|
|
),
|
|
],
|
|
)
|
|
def test_databricks_list_content_stays_list(
|
|
gateway: Gateway,
|
|
case: str,
|
|
provider_model: str,
|
|
content: list[JsonValue],
|
|
expected_content: list[JsonValue],
|
|
) -> None:
|
|
identity: Final = f"databricks-list-content-{case}-{uuid.uuid4().hex}"
|
|
expected_messages: Final = [{"role": "user", "content": expected_content}]
|
|
|
|
def respond(request: Request) -> Reply:
|
|
assert request.method == "POST"
|
|
body: Final = _JSON_OBJECT.validate_json(request.body)
|
|
assert body["messages"] == expected_messages, body["messages"]
|
|
return _shape_reply(identity, False)
|
|
|
|
with wire_server(respond) as wire, gateway.scenario() as scenario:
|
|
model: Final = scenario.model(model=f"databricks/{provider_model}", api_base=wire.url, api_key=_API_KEY)
|
|
response: Final = gateway.request(
|
|
"POST",
|
|
"/v1/chat/completions",
|
|
{"model": model, "messages": [{"role": "user", "content": content}]},
|
|
)
|
|
assert response.status_code == 200, response.text
|
|
assert _JSON_OBJECT.validate_json(response.content)["id"]
|
|
assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/chat/completions")]
|
|
|
|
|
|
def test_databricks_json_schema_refs_are_preserved_in_chat_completions(gateway: Gateway) -> None:
|
|
def respond(request: Request) -> Reply:
|
|
assert request.method == "POST"
|
|
assert request.target == "/chat/completions"
|
|
body: Final = _JSON_OBJECT.validate_json(request.body)
|
|
assert body["model"] == _JSON_SCHEMA_BACKEND
|
|
assert body["messages"] == [{"role": "user", "content": _JSON_SCHEMA_PROMPT}]
|
|
assert body["stream"] is False
|
|
assert body["response_format"] == _JSON_SCHEMA_RESPONSE_FORMAT
|
|
return Reply(
|
|
body=_chat_completion_response(
|
|
_JSON_SCHEMA_BACKEND,
|
|
{"role": "assistant", "content": _JSON_CONTENT},
|
|
"stop",
|
|
)
|
|
)
|
|
|
|
with wire_server(respond) as wire, gateway.scenario() as scenario:
|
|
model: Final = scenario.model(model=f"databricks/{_JSON_SCHEMA_BACKEND}", api_base=wire.url, api_key=_API_KEY)
|
|
response: Final = gateway.request(
|
|
"POST",
|
|
"/v1/chat/completions",
|
|
{
|
|
"model": model,
|
|
"messages": [{"role": "user", "content": _JSON_SCHEMA_PROMPT}],
|
|
"response_format": _JSON_SCHEMA_RESPONSE_FORMAT,
|
|
},
|
|
)
|
|
assert response.status_code == 200, response.text
|
|
body: Final = _JSON_OBJECT.validate_json(response.content)
|
|
assert_completion(body)
|
|
assert body["choices"] == [
|
|
{
|
|
"index": 0,
|
|
"message": {"role": "assistant", "content": _JSON_CONTENT},
|
|
"finish_reason": "stop",
|
|
}
|
|
]
|
|
assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/chat/completions")]
|
|
|
|
|
|
def test_databricks_streaming_json_schema_refs_are_preserved(gateway: Gateway) -> None:
|
|
identity: Final = f"databricks-json-schema-stream-{uuid.uuid4().hex}"
|
|
frames: Final = (
|
|
_frame(
|
|
identity,
|
|
[{"index": 0, "delta": {"role": "assistant", "content": '{"p":'}, "finish_reason": None}],
|
|
model=_JSON_SCHEMA_BACKEND,
|
|
),
|
|
_frame(
|
|
identity,
|
|
[{"index": 0, "delta": {"content": '{"name":"Ada"}}'}, "finish_reason": None}],
|
|
model=_JSON_SCHEMA_BACKEND,
|
|
),
|
|
_frame(identity, [{"index": 0, "delta": {}, "finish_reason": "stop"}], model=_JSON_SCHEMA_BACKEND),
|
|
b"data: [DONE]\n\n",
|
|
)
|
|
|
|
def respond(request: Request) -> Reply:
|
|
assert request.method == "POST"
|
|
assert request.target == "/chat/completions"
|
|
body: Final = _JSON_OBJECT.validate_json(request.body)
|
|
assert body["model"] == _JSON_SCHEMA_BACKEND
|
|
assert body["messages"] == [{"role": "user", "content": _JSON_SCHEMA_PROMPT}]
|
|
assert body["stream"] is True
|
|
assert body["response_format"] == _JSON_SCHEMA_RESPONSE_FORMAT
|
|
return Reply(content_type="text/event-stream", chunks=frames)
|
|
|
|
with wire_server(respond) as wire, gateway.scenario() as scenario:
|
|
model: Final = scenario.model(model=f"databricks/{_JSON_SCHEMA_BACKEND}", api_base=wire.url, api_key=_API_KEY)
|
|
with gateway.client.stream(
|
|
"POST",
|
|
"/v1/chat/completions",
|
|
json={
|
|
"model": model,
|
|
"messages": [{"role": "user", "content": _JSON_SCHEMA_PROMPT}],
|
|
"response_format": _JSON_SCHEMA_RESPONSE_FORMAT,
|
|
"stream": True,
|
|
},
|
|
headers={"Authorization": f"Bearer {gateway.key}"},
|
|
) as response:
|
|
assert response.status_code == 200, response.read()
|
|
lines: Final = tuple(line for line in response.iter_lines() if line.startswith("data: "))
|
|
assert lines[-1] == "data: [DONE]", lines
|
|
chunks: Final = tuple(_Chunk.model_validate_json(line.removeprefix("data: ")) for line in lines[:-1])
|
|
choices: Final = tuple(itertools.chain.from_iterable(chunk.choices for chunk in chunks))
|
|
assert "".join(choice.delta.content or "" for choice in choices) == _JSON_CONTENT
|
|
assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/chat/completions")]
|
|
|
|
|
|
def test_databricks_claude_json_schema_refs_are_preserved_in_tool_parameters(gateway: Gateway) -> None:
|
|
def respond(request: Request) -> Reply:
|
|
assert request.method == "POST"
|
|
assert request.target == "/chat/completions"
|
|
body: Final = _JSON_OBJECT.validate_json(request.body)
|
|
assert body["model"] == _CLAUDE_BACKEND
|
|
assert body["messages"] == [{"role": "user", "content": _JSON_SCHEMA_PROMPT}]
|
|
assert "response_format" not in body
|
|
assert body["tools"] == [
|
|
{
|
|
"type": "function",
|
|
"function": {"name": "json_tool_call", "parameters": _PERSON_SCHEMA},
|
|
}
|
|
]
|
|
assert body["tool_choice"] == {"type": "function", "function": {"name": "json_tool_call"}}
|
|
return Reply(body=_chat_completion_response(_CLAUDE_BACKEND, _JSON_TOOL_CALL_MESSAGE, "tool_calls"))
|
|
|
|
with wire_server(respond) as wire, gateway.scenario() as scenario:
|
|
model: Final = scenario.model(model=f"databricks/{_CLAUDE_BACKEND}", api_base=wire.url, api_key=_API_KEY)
|
|
response: Final = gateway.request(
|
|
"POST",
|
|
"/v1/chat/completions",
|
|
{
|
|
"model": model,
|
|
"messages": [{"role": "user", "content": _JSON_SCHEMA_PROMPT}],
|
|
"response_format": _JSON_SCHEMA_RESPONSE_FORMAT,
|
|
},
|
|
)
|
|
assert response.status_code == 200, response.text
|
|
body: Final = _JSON_OBJECT.validate_json(response.content)
|
|
assert body["choices"] == [
|
|
{
|
|
"index": 0,
|
|
"message": {"role": "assistant", "content": _JSON_CONTENT},
|
|
"finish_reason": "stop",
|
|
}
|
|
]
|
|
assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/chat/completions")]
|
|
|
|
|
|
def test_databricks_responses_json_schema_refs_are_preserved_in_chat_bridge(gateway: Gateway) -> None:
|
|
def respond(request: Request) -> Reply:
|
|
assert request.method == "POST"
|
|
assert request.target == "/chat/completions"
|
|
body: Final = _JSON_OBJECT.validate_json(request.body)
|
|
assert body["model"] == _JSON_SCHEMA_BACKEND
|
|
assert body["messages"] == [{"role": "user", "content": _JSON_SCHEMA_PROMPT}]
|
|
assert body["response_format"] == {
|
|
"type": "json_schema",
|
|
"json_schema": {"name": "P", "schema": _PERSON_SCHEMA, "strict": True},
|
|
}
|
|
return Reply(
|
|
body=_chat_completion_response(
|
|
_JSON_SCHEMA_BACKEND,
|
|
{"role": "assistant", "content": _JSON_CONTENT},
|
|
"stop",
|
|
)
|
|
)
|
|
|
|
with wire_server(respond) as wire, gateway.scenario() as scenario:
|
|
model: Final = scenario.model(model=f"databricks/{_JSON_SCHEMA_BACKEND}", api_base=wire.url, api_key=_API_KEY)
|
|
response: Final = gateway.request(
|
|
"POST",
|
|
"/v1/responses",
|
|
{
|
|
"model": model,
|
|
"input": _JSON_SCHEMA_PROMPT,
|
|
"text": {
|
|
"format": {
|
|
"type": "json_schema",
|
|
"name": "P",
|
|
"strict": True,
|
|
"schema": _PERSON_SCHEMA,
|
|
}
|
|
},
|
|
},
|
|
)
|
|
assert response.status_code == 200, response.text
|
|
assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/chat/completions")]
|
|
|
|
|
|
def test_databricks_messages_json_schema_refs_are_preserved_for_qwen(gateway: Gateway) -> None:
|
|
def respond(request: Request) -> Reply:
|
|
assert request.method == "POST"
|
|
assert request.target == "/chat/completions"
|
|
body: Final = _JSON_OBJECT.validate_json(request.body)
|
|
assert body["model"] == _JSON_SCHEMA_BACKEND
|
|
assert body["response_format"] == _MESSAGES_RESPONSE_FORMAT
|
|
return Reply(
|
|
body=_chat_completion_response(
|
|
_JSON_SCHEMA_BACKEND,
|
|
{"role": "assistant", "content": _JSON_CONTENT},
|
|
"stop",
|
|
)
|
|
)
|
|
|
|
with wire_server(respond) as wire, gateway.scenario() as scenario:
|
|
model: Final = scenario.model(model=f"databricks/{_JSON_SCHEMA_BACKEND}", api_base=wire.url, api_key=_API_KEY)
|
|
response: Final = gateway.request(
|
|
"POST",
|
|
"/v1/messages",
|
|
{
|
|
"model": model,
|
|
"max_tokens": 64,
|
|
"messages": [{"role": "user", "content": _JSON_SCHEMA_PROMPT}],
|
|
"output_format": {"type": "json_schema", "schema": _PERSON_SCHEMA},
|
|
},
|
|
)
|
|
assert response.status_code == 200, response.text
|
|
body: Final = _JSON_OBJECT.validate_json(response.content)
|
|
assert body["content"] == [{"type": "text", "text": _JSON_CONTENT}]
|
|
assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/chat/completions")]
|
|
|
|
|
|
def test_databricks_messages_claude_json_schema_refs_are_preserved_in_tool_parameters(gateway: Gateway) -> None:
|
|
def respond(request: Request) -> Reply:
|
|
assert request.method == "POST"
|
|
assert request.target == "/chat/completions"
|
|
body: Final = _JSON_OBJECT.validate_json(request.body)
|
|
assert body["model"] == _CLAUDE_BACKEND
|
|
assert "response_format" not in body
|
|
assert body["tools"] == [
|
|
{
|
|
"type": "function",
|
|
"function": {"name": "json_tool_call", "parameters": _MESSAGES_PERSON_SCHEMA},
|
|
}
|
|
]
|
|
assert body["tool_choice"] == {"type": "function", "function": {"name": "json_tool_call"}}
|
|
return Reply(body=_chat_completion_response(_CLAUDE_BACKEND, _JSON_TOOL_CALL_MESSAGE, "tool_calls"))
|
|
|
|
with wire_server(respond) as wire, gateway.scenario() as scenario:
|
|
model: Final = scenario.model(model=f"databricks/{_CLAUDE_BACKEND}", api_base=wire.url, api_key=_API_KEY)
|
|
response: Final = gateway.request(
|
|
"POST",
|
|
"/v1/messages",
|
|
{
|
|
"model": model,
|
|
"max_tokens": 64,
|
|
"messages": [{"role": "user", "content": _JSON_SCHEMA_PROMPT}],
|
|
"output_format": {"type": "json_schema", "schema": _PERSON_SCHEMA},
|
|
},
|
|
)
|
|
assert response.status_code == 200, response.text
|
|
body: Final = _JSON_OBJECT.validate_json(response.content)
|
|
assert body["content"] == [{"type": "text", "text": _JSON_CONTENT}]
|
|
assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/chat/completions")]
|
|
|
|
|
|
_Endpoint: TypeAlias = Literal["chat", "responses", "messages"]
|
|
_PATHS: Final[Mapping[_Endpoint, str]] = {
|
|
"chat": "/v1/chat/completions",
|
|
"responses": "/v1/responses",
|
|
"messages": "/v1/messages",
|
|
}
|
|
_OUTAGE: Final = "upstream-outage"
|
|
_OUTAGE_MESSAGE: Final = "The service is temporarily unavailable"
|
|
_EXPECTED_JSON: Final[JsonValue] = _JSON_OBJECT.validate_json(_JSON_CONTENT)
|
|
_CONTENT_DELTAS: Final[tuple[Mapping[str, JsonValue], ...]] = (
|
|
{"role": "assistant", "content": '{"p":'},
|
|
{"content": '{"name":"Ada"}}'},
|
|
)
|
|
_TOOL_CALL_DELTAS: Final[tuple[Mapping[str, JsonValue], ...]] = (
|
|
{
|
|
"role": "assistant",
|
|
"content": None,
|
|
"tool_calls": [
|
|
{
|
|
"index": 0,
|
|
"id": "json-tool-call",
|
|
"type": "function",
|
|
"function": {"name": "json_tool_call", "arguments": ""},
|
|
}
|
|
],
|
|
},
|
|
{"tool_calls": [{"index": 0, "function": {"arguments": _JSON_CONTENT}}]},
|
|
)
|
|
|
|
|
|
class _ChatMessage(BaseModel):
|
|
model_config = ConfigDict(extra="ignore", frozen=True)
|
|
content: str | None = None
|
|
|
|
|
|
class _ChatChoice(BaseModel):
|
|
model_config = ConfigDict(extra="ignore", frozen=True)
|
|
message: _ChatMessage
|
|
|
|
|
|
class _ChatCompletion(BaseModel):
|
|
model_config = ConfigDict(extra="ignore", frozen=True)
|
|
choices: tuple[_ChatChoice, ...]
|
|
|
|
|
|
class _OutputPart(BaseModel):
|
|
model_config = ConfigDict(extra="ignore", frozen=True)
|
|
type: str
|
|
text: str | None = None
|
|
|
|
|
|
class _OutputItem(BaseModel):
|
|
model_config = ConfigDict(extra="ignore", frozen=True)
|
|
type: str
|
|
content: tuple[_OutputPart, ...] = ()
|
|
|
|
|
|
class _ResponsesBody(BaseModel):
|
|
model_config = ConfigDict(extra="ignore", frozen=True)
|
|
output: tuple[_OutputItem, ...]
|
|
|
|
|
|
class _ResponsesEvent(BaseModel):
|
|
model_config = ConfigDict(extra="ignore", frozen=True)
|
|
type: str
|
|
response: _ResponsesBody | None = None
|
|
|
|
|
|
class _MessagesBlock(BaseModel):
|
|
model_config = ConfigDict(extra="ignore", frozen=True)
|
|
type: str
|
|
text: str | None = None
|
|
|
|
|
|
class _MessagesBody(BaseModel):
|
|
model_config = ConfigDict(extra="ignore", frozen=True)
|
|
content: tuple[_MessagesBlock, ...]
|
|
|
|
|
|
class _MessagesDelta(BaseModel):
|
|
model_config = ConfigDict(extra="ignore", frozen=True)
|
|
type: str | None = None
|
|
text: str | None = None
|
|
|
|
|
|
class _MessagesEvent(BaseModel):
|
|
model_config = ConfigDict(extra="ignore", frozen=True)
|
|
type: str
|
|
delta: _MessagesDelta | None = None
|
|
|
|
|
|
class _StreamDelta(BaseModel):
|
|
model_config = ConfigDict(extra="ignore", frozen=True)
|
|
content: str | None = None
|
|
|
|
|
|
class _StreamChoice(BaseModel):
|
|
model_config = ConfigDict(extra="ignore", frozen=True)
|
|
delta: _StreamDelta
|
|
|
|
|
|
class _StreamChunk(BaseModel):
|
|
model_config = ConfigDict(extra="ignore", frozen=True)
|
|
choices: tuple[_StreamChoice, ...]
|
|
|
|
|
|
def _json_schema_format(schema: Mapping[str, JsonValue]) -> dict[str, JsonValue]:
|
|
return {"type": "json_schema", "json_schema": {"name": "P", "strict": True, "schema": dict(schema)}}
|
|
|
|
|
|
def _client_body(
|
|
endpoint: _Endpoint,
|
|
model: str,
|
|
stream: bool,
|
|
prompt: str = _JSON_SCHEMA_PROMPT,
|
|
) -> dict[str, JsonValue]:
|
|
match endpoint:
|
|
case "chat":
|
|
return {
|
|
"model": model,
|
|
"messages": [{"role": "user", "content": prompt}],
|
|
"response_format": _json_schema_format(_PERSON_SCHEMA),
|
|
"stream": stream,
|
|
}
|
|
case "responses":
|
|
return {
|
|
"model": model,
|
|
"input": prompt,
|
|
"text": {
|
|
"format": {"type": "json_schema", "name": "P", "strict": True, "schema": dict(_PERSON_SCHEMA)}
|
|
},
|
|
"stream": stream,
|
|
}
|
|
case "messages":
|
|
return {
|
|
"model": model,
|
|
"max_tokens": 64,
|
|
"messages": [{"role": "user", "content": prompt}],
|
|
"output_format": {"type": "json_schema", "schema": dict(_PERSON_SCHEMA)},
|
|
"stream": stream,
|
|
}
|
|
|
|
|
|
def _expected_outbound(endpoint: _Endpoint, backend: str) -> dict[str, JsonValue]:
|
|
schema: Final = dict(_MESSAGES_PERSON_SCHEMA if endpoint == "messages" else _PERSON_SCHEMA)
|
|
if "claude" in backend:
|
|
return _JSON_OBJECT.validate_python(
|
|
{
|
|
"tools": [{"type": "function", "function": {"name": "json_tool_call", "parameters": schema}}],
|
|
"tool_choice": {"type": "function", "function": {"name": "json_tool_call"}},
|
|
}
|
|
)
|
|
name: Final = "structured_output" if endpoint == "messages" else "P"
|
|
return _JSON_OBJECT.validate_python(
|
|
{"response_format": {"type": "json_schema", "json_schema": {"name": name, "schema": schema, "strict": True}}}
|
|
)
|
|
|
|
|
|
def _assert_outbound(request: Request, endpoint: _Endpoint, backend: str, stream: bool) -> None:
|
|
assert (request.method, request.target) == ("POST", "/chat/completions")
|
|
body: Final = _JSON_OBJECT.validate_json(request.body)
|
|
expected: Final = _expected_outbound(endpoint, backend)
|
|
assert {key: body.get(key) for key in expected} == expected
|
|
assert (body["model"], body.get("stream") is True) == (backend, stream)
|
|
assert "claude" not in backend or "response_format" not in body
|
|
|
|
|
|
def _upstream_reply(backend: str, stream: bool) -> Reply:
|
|
claude: Final = "claude" in backend
|
|
finish_reason: Final = "tool_calls" if claude else "stop"
|
|
if not stream:
|
|
message: Final[dict[str, JsonValue]] = (
|
|
_JSON_TOOL_CALL_MESSAGE if claude else {"role": "assistant", "content": _JSON_CONTENT}
|
|
)
|
|
return Reply(body=_chat_completion_response(backend, message, finish_reason))
|
|
identity: Final = f"databricks-json-schema-{uuid.uuid4().hex}"
|
|
deltas: Final = _TOOL_CALL_DELTAS if claude else _CONTENT_DELTAS
|
|
frames: Final = (
|
|
*(_frame(identity, [{"index": 0, "delta": delta, "finish_reason": None}], model=backend) for delta in deltas),
|
|
_frame(identity, [{"index": 0, "delta": {}, "finish_reason": finish_reason}], model=backend),
|
|
b"data: [DONE]\n\n",
|
|
)
|
|
return Reply(content_type="text/event-stream", chunks=frames)
|
|
|
|
|
|
def _outage_or_reply(request: Request) -> Reply:
|
|
if _OUTAGE.encode() in request.body:
|
|
error: Final = {"error_code": "TEMPORARILY_UNAVAILABLE", "message": _OUTAGE_MESSAGE}
|
|
return Reply(status=503, body=json.dumps(error).encode())
|
|
body: Final = _JSON_OBJECT.validate_json(request.body)
|
|
return _upstream_reply(str(body["model"]), body.get("stream") is True)
|
|
|
|
|
|
def _item_text(item: _OutputItem) -> str:
|
|
return "".join(part.text or "" for part in item.content) if item.type == "message" else ""
|
|
|
|
|
|
def _responses_text(body: _ResponsesBody) -> str:
|
|
return "".join(_item_text(item) for item in body.output)
|
|
|
|
|
|
def _messages_delta_text(delta: _MessagesDelta) -> str:
|
|
return delta.text or "" if delta.type == "text_delta" else ""
|
|
|
|
|
|
def _final_text(endpoint: _Endpoint, content: bytes) -> str:
|
|
match endpoint:
|
|
case "chat":
|
|
return _ChatCompletion.model_validate_json(content).choices[0].message.content or ""
|
|
case "responses":
|
|
return _responses_text(_ResponsesBody.model_validate_json(content))
|
|
case "messages":
|
|
return "".join(block.text or "" for block in _MessagesBody.model_validate_json(content).content)
|
|
|
|
|
|
def _streamed_text(endpoint: _Endpoint, data: tuple[str, ...]) -> str:
|
|
match endpoint:
|
|
case "chat":
|
|
assert data[-1] == "[DONE]", data
|
|
chunks: Final = tuple(_StreamChunk.model_validate_json(line) for line in data[:-1])
|
|
choices: Final = tuple(itertools.chain.from_iterable(chunk.choices for chunk in chunks))
|
|
return "".join(choice.delta.content or "" for choice in choices)
|
|
case "responses":
|
|
events: Final = tuple(_ResponsesEvent.model_validate_json(line) for line in data if line != "[DONE]")
|
|
completed: Final = tuple(event.response for event in events if event.type == "response.completed")
|
|
assert len(completed) == 1 and completed[0] is not None, data
|
|
return _responses_text(completed[0])
|
|
case "messages":
|
|
message_events: Final = tuple(_MessagesEvent.model_validate_json(line) for line in data)
|
|
assert message_events[-1].type == "message_stop", data
|
|
deltas: Final = tuple(event.delta for event in message_events if event.delta is not None)
|
|
return "".join(_messages_delta_text(delta) for delta in deltas)
|
|
|
|
|
|
def _outcome(status: int, text: str) -> tuple[int, JsonValue]:
|
|
if status == 200:
|
|
return status, _JSON_OBJECT.validate_json(text)
|
|
return status, _OUTAGE_MESSAGE if _OUTAGE_MESSAGE in text else text
|
|
|
|
|
|
def _call(gateway: Gateway, endpoint: _Endpoint, body: Mapping[str, JsonValue]) -> tuple[int, str]:
|
|
if body["stream"] is not True:
|
|
response: Final = gateway.request("POST", _PATHS[endpoint], body)
|
|
return response.status_code, _final_text(endpoint, response.content) if response.is_success else response.text
|
|
with gateway.client.stream(
|
|
"POST", _PATHS[endpoint], json=dict(body), headers={"Authorization": f"Bearer {gateway.key}"}
|
|
) as streamed:
|
|
if not streamed.is_success:
|
|
return streamed.status_code, streamed.read().decode()
|
|
data: Final = tuple(line.removeprefix("data: ") for line in streamed.iter_lines() if line.startswith("data: "))
|
|
return 200, _streamed_text(endpoint, data)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("endpoint", "backend", "stream"),
|
|
[
|
|
pytest.param("chat", _CLAUDE_BACKEND, True, id="chat-stream-claude"),
|
|
pytest.param("responses", _CLAUDE_BACKEND, False, id="responses-claude"),
|
|
pytest.param("responses", _JSON_SCHEMA_BACKEND, True, id="responses-stream-qwen"),
|
|
pytest.param("responses", _CLAUDE_BACKEND, True, id="responses-stream-claude"),
|
|
pytest.param("messages", _JSON_SCHEMA_BACKEND, True, id="messages-stream-qwen"),
|
|
pytest.param("messages", _CLAUDE_BACKEND, True, id="messages-stream-claude"),
|
|
],
|
|
)
|
|
def test_databricks_json_schema_refs_are_preserved_on_every_endpoint(
|
|
gateway: Gateway, endpoint: _Endpoint, backend: str, stream: bool
|
|
) -> None:
|
|
with wire_server(lambda _request: _upstream_reply(backend, stream)) as wire, gateway.scenario() as scenario:
|
|
model: Final = scenario.model(model=f"databricks/{backend}", api_base=wire.url, api_key=_API_KEY)
|
|
status, text = _call(gateway, endpoint, _client_body(endpoint, model, stream))
|
|
received: Final = wire.drain()
|
|
assert len(received) == 1, received
|
|
_assert_outbound(received[0], endpoint, backend, stream)
|
|
assert _outcome(status, text) == (200, _EXPECTED_JSON), text
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"response_format",
|
|
[
|
|
pytest.param(
|
|
_json_schema_format(
|
|
{
|
|
"definitions": {
|
|
"Person": {"type": "object", "properties": {"name": {"type": "string"}}, "required": ["name"]}
|
|
},
|
|
"type": "object",
|
|
"properties": {"p": {"$ref": "#/definitions/Person"}},
|
|
"required": ["p"],
|
|
}
|
|
),
|
|
id="definitions-ref",
|
|
),
|
|
pytest.param(
|
|
_json_schema_format(
|
|
{
|
|
"type": "object",
|
|
"properties": {
|
|
"name": {"type": "string"},
|
|
"children": {"type": "array", "items": {"$ref": "#"}},
|
|
},
|
|
"required": ["name", "children"],
|
|
}
|
|
),
|
|
id="recursive-root-ref",
|
|
),
|
|
pytest.param(
|
|
_json_schema_format({"type": "object", "properties": {"$ref": {"type": "string"}}, "required": ["$ref"]}),
|
|
id="property-named-ref",
|
|
),
|
|
pytest.param(
|
|
_json_schema_format({"type": "object", "properties": {"name": {"type": "string"}}, "required": ["name"]}),
|
|
id="flat-schema",
|
|
),
|
|
pytest.param({"type": "json_object"}, id="json-object"),
|
|
],
|
|
)
|
|
def test_databricks_response_format_reaches_upstream_as_sent(
|
|
gateway: Gateway, response_format: dict[str, JsonValue]
|
|
) -> None:
|
|
with (
|
|
wire_server(lambda _request: _upstream_reply(_JSON_SCHEMA_BACKEND, False)) as wire,
|
|
gateway.scenario() as scenario,
|
|
):
|
|
model: Final = scenario.model(model=f"databricks/{_JSON_SCHEMA_BACKEND}", api_base=wire.url, api_key=_API_KEY)
|
|
response: Final = gateway.request(
|
|
"POST",
|
|
"/v1/chat/completions",
|
|
{
|
|
"model": model,
|
|
"messages": [{"role": "user", "content": _JSON_SCHEMA_PROMPT}],
|
|
"response_format": response_format,
|
|
},
|
|
)
|
|
received: Final = wire.drain()
|
|
assert response.status_code == 200, response.text
|
|
assert _final_text("chat", response.content) == _JSON_CONTENT
|
|
assert [(request.method, request.target) for request in received] == [("POST", "/chat/completions")]
|
|
assert _JSON_OBJECT.validate_json(received[0].body)["response_format"] == response_format
|
|
|
|
|
|
def test_databricks_json_schema_without_schema_returns_the_upstream_error(gateway: Gateway) -> None:
|
|
response_format: Final[dict[str, JsonValue]] = {"type": "json_schema", "json_schema": {"name": "P", "strict": True}}
|
|
upstream_error: Final = {"error_code": "INVALID_PARAMETER_VALUE", "message": "json_schema.schema is required"}
|
|
|
|
def respond(_request: Request) -> Reply:
|
|
return Reply(status=400, body=json.dumps(upstream_error).encode())
|
|
|
|
with wire_server(respond) as wire, gateway.scenario() as scenario:
|
|
model: Final = scenario.model(model=f"databricks/{_JSON_SCHEMA_BACKEND}", api_base=wire.url, api_key=_API_KEY)
|
|
response: Final = gateway.request(
|
|
"POST",
|
|
"/v1/chat/completions",
|
|
{
|
|
"model": model,
|
|
"messages": [{"role": "user", "content": _JSON_SCHEMA_PROMPT}],
|
|
"response_format": response_format,
|
|
},
|
|
)
|
|
received: Final = wire.drain()
|
|
assert (response.status_code, "json_schema.schema is required" in response.text) == (400, True), response.text
|
|
assert [(request.method, request.target) for request in received] == [("POST", "/chat/completions")]
|
|
assert _JSON_OBJECT.validate_json(received[0].body)["response_format"] == response_format
|
|
|
|
|
|
def _assert_burst_outbound(request: Request) -> None:
|
|
body: Final = _JSON_OBJECT.validate_json(request.body)
|
|
endpoint: Final[_Endpoint] = "messages" if body["response_format"] == _MESSAGES_RESPONSE_FORMAT else "chat"
|
|
_assert_outbound(request, endpoint, _JSON_SCHEMA_BACKEND, body.get("stream") is True)
|
|
|
|
|
|
def _burst_request(index: int) -> tuple[_Endpoint, bool, bool]:
|
|
endpoints: Final[tuple[_Endpoint, ...]] = ("chat", "responses", "messages")
|
|
return endpoints[index % 3], index % 2 == 1, index % 5 == 2
|
|
|
|
|
|
def test_databricks_json_schema_burst_through_an_upstream_outage_keeps_every_ref(gateway: Gateway) -> None:
|
|
plan: Final = tuple(_burst_request(index) for index in range(24))
|
|
recovery: Final[tuple[_Endpoint, ...]] = ("chat", "responses", "messages")
|
|
with wire_server(_outage_or_reply) as wire, gateway.scenario() as scenario:
|
|
model: Final = scenario.model(model=f"databricks/{_JSON_SCHEMA_BACKEND}", api_base=wire.url, api_key=_API_KEY)
|
|
endpoints: Final = tuple(endpoint for endpoint, _stream, _down in plan)
|
|
bodies: Final = tuple(
|
|
_client_body(endpoint, model, stream, f"{_JSON_SCHEMA_PROMPT} {index} {_OUTAGE if down else ''}")
|
|
for index, (endpoint, stream, down) in enumerate(plan)
|
|
)
|
|
with ThreadPoolExecutor(max_workers=12) as pool:
|
|
results: Final = tuple(pool.map(_call, itertools.repeat(gateway), endpoints, bodies))
|
|
recovered: Final = tuple(
|
|
_call(gateway, endpoint, _client_body(endpoint, model, False)) for endpoint in recovery
|
|
)
|
|
received: Final = wire.drain()
|
|
assert len(received) == len(plan) + len(recovery), received
|
|
for request in received:
|
|
_assert_burst_outbound(request)
|
|
expected: Final = tuple(
|
|
(503, _OUTAGE_MESSAGE) if down else (200, _EXPECTED_JSON) for _endpoint, _stream, down in plan
|
|
)
|
|
assert tuple(_outcome(*result) for result in results) == expected
|
|
assert tuple(_outcome(*result) for result in recovered) == ((200, _EXPECTED_JSON),) * len(recovery)
|
|
|
|
|
|
_GEMMA_BACKEND: Final = "databricks-gemma-3-12b"
|
|
_NO_CACHE: Final[dict[str, JsonValue]] = {"no-cache": True}
|
|
|
|
|
|
def _one_block(text: str) -> list[JsonValue]:
|
|
return [{"type": "text", "text": text}]
|
|
|
|
|
|
def _plain_reply(_request: Request) -> Reply:
|
|
return _shape_reply(f"databricks-plain-{uuid.uuid4().hex}", False)
|
|
|
|
|
|
def _sent_messages(request: Request) -> JsonValue:
|
|
return _JSON_OBJECT.validate_json(request.body)["messages"]
|
|
|
|
|
|
def _tool_call(identity: str, city: str) -> dict[str, JsonValue]:
|
|
return {
|
|
"id": identity,
|
|
"type": "function",
|
|
"function": {"name": "get_weather", "arguments": json.dumps({"city": city})},
|
|
}
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"messages,expected_messages",
|
|
[
|
|
pytest.param(
|
|
[
|
|
{"role": "system", "content": _one_block("Alpha rule.")},
|
|
{"role": "system", "content": _one_block("Bravo rule.")},
|
|
{"role": "user", "content": _one_block("Reply in JSON")},
|
|
],
|
|
[
|
|
{"role": "system", "content": "Alpha rule.\n\nBravo rule."},
|
|
{"role": "user", "content": "Reply in JSON"},
|
|
],
|
|
id="C1-two-one-block-system-messages",
|
|
),
|
|
pytest.param(
|
|
[
|
|
{"role": "system", "content": _one_block("Alpha rule.")},
|
|
{"role": "developer", "content": _one_block("Bravo rule.")},
|
|
{"role": "user", "content": "Reply in JSON"},
|
|
],
|
|
[
|
|
{"role": "system", "content": "Alpha rule.\n\nBravo rule."},
|
|
{"role": "user", "content": "Reply in JSON"},
|
|
],
|
|
id="C2-system-then-developer-one-block",
|
|
),
|
|
pytest.param(
|
|
[
|
|
{"role": "system", "content": _one_block("Alpha rule.")},
|
|
{"role": "system", "content": "Bravo rule."},
|
|
{"role": "user", "content": "Reply in JSON"},
|
|
],
|
|
[
|
|
{"role": "system", "content": "Alpha rule.\n\nBravo rule."},
|
|
{"role": "user", "content": "Reply in JSON"},
|
|
],
|
|
id="C3-one-block-then-string-system",
|
|
),
|
|
pytest.param(
|
|
[
|
|
{"role": "system", "content": _one_block("Alpha rule.")},
|
|
{
|
|
"role": "system",
|
|
"content": [{"type": "text", "text": "Bravo", "cache_control": {"type": "ephemeral"}}],
|
|
},
|
|
{"role": "user", "content": "Reply in JSON"},
|
|
],
|
|
[
|
|
{
|
|
"role": "system",
|
|
"content": [
|
|
{"type": "text", "text": "Alpha rule."},
|
|
{"type": "text", "text": "Bravo", "cache_control": {"type": "ephemeral"}},
|
|
],
|
|
},
|
|
{"role": "user", "content": "Reply in JSON"},
|
|
],
|
|
id="C4-one-block-then-cache-marked-system-stays-parts",
|
|
),
|
|
pytest.param(
|
|
[
|
|
{"role": "user", "content": _one_block("Weather in Paris and Rome?")},
|
|
{
|
|
"role": "assistant",
|
|
"content": _one_block("Checking both."),
|
|
"tool_calls": [_tool_call("call-paris", "Paris"), _tool_call("call-rome", "Rome")],
|
|
},
|
|
{"role": "tool", "tool_call_id": "call-paris", "content": _one_block("18C")},
|
|
{"role": "tool", "tool_call_id": "call-rome", "content": _one_block("22C")},
|
|
],
|
|
[
|
|
{"role": "user", "content": "Weather in Paris and Rome?"},
|
|
{"role": "assistant", "content": "Checking both.", "tool_calls": [_tool_call("call-paris", "Paris")]},
|
|
{"role": "tool", "tool_call_id": "call-paris", "content": "18C"},
|
|
{"role": "assistant", "tool_calls": [_tool_call("call-rome", "Rome")]},
|
|
{"role": "tool", "tool_call_id": "call-rome", "content": "22C"},
|
|
],
|
|
id="C5-parallel-tool-results-split-and-collapsed",
|
|
),
|
|
pytest.param(
|
|
[
|
|
{"role": "user", "content": [{"type": "text", "text": "Reply in JSON", "cache_control": None}]},
|
|
{"role": "assistant", "content": [{"type": "text", "text": "ok", "annotations": []}]},
|
|
{"role": "user", "content": _one_block(" ")},
|
|
{"role": "user", "content": _one_block("x" * 5120)},
|
|
],
|
|
[
|
|
{"role": "user", "content": [{"type": "text", "text": "Reply in JSON", "cache_control": None}]},
|
|
{"role": "assistant", "content": [{"type": "text", "text": "ok", "annotations": []}]},
|
|
{"role": "user"},
|
|
{"role": "user", "content": "x" * 5120},
|
|
],
|
|
id="F1-extra-keys-stay-list-blank-dropped-5kb-collapsed",
|
|
),
|
|
],
|
|
)
|
|
def test_databricks_non_claude_message_shapes_reach_upstream(
|
|
gateway: Gateway, messages: list[JsonValue], expected_messages: list[JsonValue]
|
|
) -> None:
|
|
with wire_server(_plain_reply) as wire, gateway.scenario() as scenario:
|
|
model: Final = scenario.model(model=f"databricks/{_GEMMA_BACKEND}", api_base=wire.url, api_key=_API_KEY)
|
|
response: Final = gateway.request(
|
|
"POST",
|
|
"/v1/chat/completions",
|
|
{"model": model, "messages": messages, "cache": _NO_CACHE},
|
|
)
|
|
received: Final = wire.drain()
|
|
assert response.status_code == 200, response.text
|
|
assert [(request.method, request.target) for request in received] == [("POST", "/chat/completions")]
|
|
assert _sent_messages(received[0]) == expected_messages
|
|
|
|
|
|
def test_databricks_responses_instructions_and_one_block_developer_reach_upstream_joined(gateway: Gateway) -> None:
|
|
with wire_server(_plain_reply) as wire, gateway.scenario() as scenario:
|
|
model: Final = scenario.model(model=f"databricks/{_GEMMA_BACKEND}", api_base=wire.url, api_key=_API_KEY)
|
|
response: Final = gateway.request(
|
|
"POST",
|
|
"/v1/responses",
|
|
{
|
|
"model": model,
|
|
"instructions": "Alpha rule.",
|
|
"input": [
|
|
{"role": "developer", "content": [{"type": "input_text", "text": "Bravo rule."}]},
|
|
{"role": "user", "content": [{"type": "input_text", "text": "Reply in JSON"}]},
|
|
],
|
|
},
|
|
)
|
|
received: Final = wire.drain()
|
|
assert response.status_code == 200, response.text
|
|
assert [(request.method, request.target) for request in received] == [("POST", "/chat/completions")]
|
|
assert _sent_messages(received[0]) == [
|
|
{"role": "system", "content": "Alpha rule.\n\nBravo rule."},
|
|
{"role": "user", "content": "Reply in JSON"},
|
|
]
|
|
|
|
|
|
def test_databricks_fallback_sends_the_collapsed_string_to_both_deployments(gateway: Gateway) -> None:
|
|
marker: Final = uuid.uuid4().hex
|
|
fallback_backend: Final = "databricks-qwen3-next-80b-a3b-instruct"
|
|
|
|
def respond(request: Request) -> Reply:
|
|
if _JSON_OBJECT.validate_json(request.body)["model"] == _GEMMA_BACKEND:
|
|
return Reply(status=500, body=b'{"error_code": "INTERNAL_ERROR", "message": "primary deployment is down"}')
|
|
return _shape_reply(f"databricks-fallback-{marker}", False)
|
|
|
|
with wire_server(respond) as wire, gateway.scenario() as scenario:
|
|
primary: Final = scenario.model(model=f"databricks/{_GEMMA_BACKEND}", api_base=wire.url, api_key=_API_KEY)
|
|
fallback: Final = scenario.model(model=f"databricks/{fallback_backend}", api_base=wire.url, api_key=_API_KEY)
|
|
response: Final = gateway.request(
|
|
"POST",
|
|
"/v1/chat/completions",
|
|
{
|
|
"model": primary,
|
|
"messages": [{"role": "user", "content": _one_block(f"Reply in JSON {marker}")}],
|
|
"response_format": {"type": "json_object"},
|
|
"fallbacks": [fallback],
|
|
"num_retries": 0,
|
|
"cache": _NO_CACHE,
|
|
},
|
|
)
|
|
received: Final = wire.drain()
|
|
assert response.status_code == 200, response.text
|
|
assert _JSON_OBJECT.validate_json(response.content)["id"] == f"databricks-fallback-{marker}"
|
|
assert [_JSON_OBJECT.validate_json(request.body)["model"] for request in received] == [
|
|
_GEMMA_BACKEND,
|
|
fallback_backend,
|
|
]
|
|
assert [_sent_messages(request) for request in received] == [
|
|
[{"role": "user", "content": f"Reply in JSON {marker}"}],
|
|
] * 2
|
|
|
|
|
|
def _proxy_root(gateway: Gateway) -> str:
|
|
return str(gateway.client.base_url).rstrip("/")
|
|
|
|
|
|
def test_databricks_openai_sdk_sync_one_block_reaches_upstream_as_string_and_logs_spend(gateway: Gateway) -> None:
|
|
identity: Final = f"databricks-openai-sdk-{uuid.uuid4().hex}"
|
|
with (
|
|
wire_server(lambda _request: _shape_reply(identity, False)) as wire,
|
|
gateway.scenario() as scenario,
|
|
openai.OpenAI(base_url=f"{_proxy_root(gateway)}/v1", api_key=gateway.key, max_retries=0) as client,
|
|
):
|
|
model: Final = scenario.model(model=f"databricks/{_GEMMA_BACKEND}", api_base=wire.url, api_key=_API_KEY)
|
|
completion: Final = client.chat.completions.create(
|
|
model=model,
|
|
messages=[{"role": "user", "content": [{"type": "text", "text": "Reply in JSON"}]}],
|
|
response_format={"type": "json_object"},
|
|
extra_body={"cache": _NO_CACHE},
|
|
)
|
|
received: Final = wire.drain()
|
|
assert (completion.id, completion.choices[0].message.content) == (identity, "ok")
|
|
assert [_sent_messages(request) for request in received] == [[{"role": "user", "content": "Reply in JSON"}]]
|
|
rows: Final = eventually(
|
|
lambda: read_rows('SELECT status, call_type FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (identity,)),
|
|
lambda values: len(values) == 1,
|
|
seconds=70,
|
|
)
|
|
assert (rows[0]["status"], rows[0]["call_type"]) == ("success", "acompletion")
|
|
|
|
|
|
async def test_databricks_openai_sdk_async_stream_one_block_reaches_upstream_as_string(gateway: Gateway) -> None:
|
|
identity: Final = f"databricks-openai-async-{uuid.uuid4().hex}"
|
|
with wire_server(lambda _request: _shape_reply(identity, True)) as wire, gateway.scenario() as scenario:
|
|
model: Final = scenario.model(model=f"databricks/{_GEMMA_BACKEND}", api_base=wire.url, api_key=_API_KEY)
|
|
async with openai.AsyncOpenAI(
|
|
base_url=f"{_proxy_root(gateway)}/v1", api_key=gateway.key, max_retries=0
|
|
) as client:
|
|
stream: Final = await client.chat.completions.create(
|
|
model=model,
|
|
messages=[{"role": "user", "content": [{"type": "text", "text": "Reply in JSON"}]}],
|
|
response_format={"type": "json_object"},
|
|
stream=True,
|
|
extra_body={"cache": _NO_CACHE},
|
|
)
|
|
chunks: Final = [chunk async for chunk in stream]
|
|
received: Final = wire.drain()
|
|
assert {chunk.id for chunk in chunks} == {identity}
|
|
assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks if chunk.choices) == "ok"
|
|
assert [_sent_messages(request) for request in received] == [[{"role": "user", "content": "Reply in JSON"}]]
|
|
|
|
|
|
async def test_databricks_anthropic_sdk_async_stream_one_block_system_and_user_reach_upstream_as_strings(
|
|
gateway: Gateway,
|
|
) -> None:
|
|
identity: Final = f"databricks-anthropic-async-{uuid.uuid4().hex}"
|
|
with wire_server(lambda _request: _shape_reply(identity, True)) as wire, gateway.scenario() as scenario:
|
|
model: Final = scenario.model(model=f"databricks/{_GEMMA_BACKEND}", api_base=wire.url, api_key=_API_KEY)
|
|
async with anthropic.AsyncAnthropic(
|
|
base_url=_proxy_root(gateway), api_key=gateway.key, max_retries=0
|
|
) as client:
|
|
async with client.messages.stream(
|
|
model=model,
|
|
max_tokens=32,
|
|
system=[{"type": "text", "text": "Reply in JSON."}],
|
|
messages=[{"role": "user", "content": [{"type": "text", "text": "What is 2+2?"}]}],
|
|
) as stream:
|
|
final: Final = await stream.get_final_message()
|
|
received: Final = wire.drain()
|
|
assert [block.text for block in final.content if block.type == "text"] == ["ok"]
|
|
assert [_sent_messages(request) for request in received] == [
|
|
[{"role": "system", "content": "Reply in JSON."}, {"role": "user", "content": "What is 2+2?"}]
|
|
]
|
|
|
|
|
|
def test_databricks_one_block_burst_through_an_upstream_outage_reaches_upstream_as_strings(gateway: Gateway) -> None:
|
|
plan: Final = tuple(_burst_request(index) for index in range(24))
|
|
|
|
def respond(request: Request) -> Reply:
|
|
if _OUTAGE.encode() in request.body:
|
|
error: Final = {"error_code": "TEMPORARILY_UNAVAILABLE", "message": _OUTAGE_MESSAGE}
|
|
return Reply(status=503, body=json.dumps(error).encode())
|
|
body: Final = _JSON_OBJECT.validate_json(request.body)
|
|
return _upstream_reply(_JSON_SCHEMA_BACKEND, body.get("stream") is True)
|
|
|
|
def one_block_body(endpoint: _Endpoint, model: str, stream: bool, prompt: str) -> dict[str, JsonValue]:
|
|
body: Final = _client_body(endpoint, model, stream, prompt)
|
|
match endpoint:
|
|
case "responses":
|
|
return {**body, "input": [{"role": "user", "content": [{"type": "input_text", "text": prompt}]}]}
|
|
case "chat" | "messages":
|
|
return {**body, "messages": [{"role": "user", "content": _one_block(prompt)}]}
|
|
|
|
with wire_server(respond) as wire, gateway.scenario() as scenario:
|
|
model: Final = scenario.model(model=f"databricks/{_JSON_SCHEMA_BACKEND}", api_base=wire.url, api_key=_API_KEY)
|
|
prompts: Final = tuple(
|
|
f"{_JSON_SCHEMA_PROMPT} {index} {_OUTAGE if down else ''}".strip()
|
|
for index, (_e, _s, down) in enumerate(plan)
|
|
)
|
|
bodies: Final = tuple(
|
|
one_block_body(endpoint, model, stream, prompt) for (endpoint, stream, _down), prompt in zip(plan, prompts)
|
|
)
|
|
with ThreadPoolExecutor(max_workers=12) as pool:
|
|
results: Final = tuple(pool.map(_call, itertools.repeat(gateway), (e for e, _s, _d in plan), bodies))
|
|
received: Final = wire.drain()
|
|
assert sorted(json.dumps(_sent_messages(request)) for request in received) == sorted(
|
|
json.dumps([{"role": "user", "content": prompt}]) for prompt in prompts
|
|
)
|
|
expected: Final = tuple(
|
|
(503, _OUTAGE_MESSAGE) if down else (200, _EXPECTED_JSON) for _endpoint, _stream, down in plan
|
|
)
|
|
assert tuple(_outcome(*result) for result in results) == expected
|
|
|
|
|
|
def test_databricks_transform_request_shows_the_collapsed_string(gateway: Gateway, tmp_path: Path) -> None:
|
|
environment: Final = {
|
|
"DATABRICKS_API_BASE": "https://databricks.invalid/serving-endpoints",
|
|
"DATABRICKS_API_KEY": _API_KEY,
|
|
}
|
|
with owned_proxy_process(gateway, tmp_path, environment) as owned:
|
|
response: Final = owned.gateway.request(
|
|
"POST",
|
|
"/utils/transform_request",
|
|
{
|
|
"call_type": "completion",
|
|
"request_body": {
|
|
"model": f"databricks/{_GEMMA_BACKEND}",
|
|
"messages": [{"role": "user", "content": _one_block("Reply in JSON")}],
|
|
},
|
|
},
|
|
)
|
|
assert response.status_code == 200, response.text
|
|
body: Final = _JSON_OBJECT.validate_json(response.content)
|
|
assert body["error"] is None, body
|
|
assert body["raw_request_api_base"] == "https://databricks.invalid/serving-endpoints/chat/completions"
|
|
sent: Final = body["raw_request_body"]
|
|
assert isinstance(sent, dict), body
|
|
assert sent["messages"] == [{"role": "user", "content": "Reply in JSON"}]
|
|
|
|
|
|
def test_databricks_non_string_text_block_is_rejected_before_reaching_upstream(gateway: Gateway) -> None:
|
|
with wire_server(_plain_reply) as wire, gateway.scenario() as scenario:
|
|
model: Final = scenario.model(model=f"databricks/{_GEMMA_BACKEND}", api_base=wire.url, api_key=_API_KEY)
|
|
response: Final = gateway.request(
|
|
"POST",
|
|
"/v1/chat/completions",
|
|
{
|
|
"model": model,
|
|
"messages": [{"role": "user", "content": [{"type": "text", "text": 42}]}],
|
|
"cache": _NO_CACHE,
|
|
},
|
|
)
|
|
received: Final = wire.drain()
|
|
assert response.status_code == 400, response.text
|
|
assert "A text content block needs a string 'text', got int" in response.text
|
|
assert received == ()
|
|
|
|
|
|
_JSON_ANSWER_FRAGMENTS: Final = ('{"p":', '{"name":"Ada"}', "}")
|
|
_PROSE_FRAGMENTS: Final = ("I will check ", "the weather first.")
|
|
_WEATHER_FRAGMENTS: Final = ('{"city": ', '"Paris"}')
|
|
_WEATHER_PARAMETERS: Final[Mapping[str, JsonValue]] = {
|
|
"type": "object",
|
|
"properties": {"city": {"type": "string"}},
|
|
"required": ["city"],
|
|
}
|
|
_FORCED_JSON_TOOL_CHOICE: Final[Mapping[str, JsonValue]] = {"type": "function", "function": {"name": "json_tool_call"}}
|
|
_JSON_VALUE: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue)
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class _Reasoning:
|
|
text: str
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class _Text:
|
|
fragments: tuple[str, ...]
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class _ToolUse:
|
|
index: int
|
|
name: str
|
|
fragments: tuple[str, ...]
|
|
|
|
|
|
_Part: TypeAlias = _Reasoning | _Text | _ToolUse
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class _Answer:
|
|
text: str
|
|
tools: tuple[tuple[str, JsonValue], ...]
|
|
finish: str
|
|
thought: bool
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class _Shape:
|
|
parts: tuple[_Part, ...]
|
|
upstream_finish: str
|
|
thinking: bool
|
|
user_tools: bool
|
|
json_schema: bool
|
|
expected: _Answer
|
|
|
|
|
|
_THINKING: Final = _Reasoning("The caller wants Ada.")
|
|
_PROSE: Final = _Text(_PROSE_FRAGMENTS)
|
|
_WEATHER: Final = ("get_weather", _JSON_VALUE.validate_json("".join(_WEATHER_FRAGMENTS)))
|
|
_JSON_ONLY: Final = _Answer(_JSON_CONTENT, (), "stop", True)
|
|
_JSON_AND_WEATHER: Final = _Answer(_JSON_CONTENT, (_WEATHER,), "tool_calls", True)
|
|
_PROSE_AND_WEATHER: Final = _Answer("".join(_PROSE_FRAGMENTS), (_WEATHER,), "tool_calls", True)
|
|
|
|
|
|
def _json_tool(index: int) -> _ToolUse:
|
|
return _ToolUse(index, "json_tool_call", _JSON_ANSWER_FRAGMENTS)
|
|
|
|
|
|
def _weather_tool(index: int) -> _ToolUse:
|
|
return _ToolUse(index, "get_weather", _WEATHER_FRAGMENTS)
|
|
|
|
|
|
_SHAPES: Final[Mapping[str, _Shape]] = {
|
|
"json-only": _Shape((_json_tool(0),), "tool_calls", False, False, True, _Answer(_JSON_CONTENT, (), "stop", False)),
|
|
"prose-then-json": _Shape((_THINKING, _PROSE, _json_tool(0)), "tool_calls", True, False, True, _JSON_ONLY),
|
|
"json-then-user-tool": _Shape(
|
|
(_THINKING, _json_tool(0), _weather_tool(1)), "tool_calls", True, True, True, _JSON_AND_WEATHER
|
|
),
|
|
"user-tool-then-json": _Shape(
|
|
(_THINKING, _weather_tool(0), _json_tool(1)), "tool_calls", True, True, True, _JSON_AND_WEATHER
|
|
),
|
|
"prose-user-tool-json": _Shape(
|
|
(_THINKING, _PROSE, _weather_tool(0), _json_tool(1)), "tool_calls", True, True, True, _JSON_AND_WEATHER
|
|
),
|
|
"text-answer": _Shape((_THINKING, _Text(_JSON_ANSWER_FRAGMENTS)), "stop", True, False, True, _JSON_ONLY),
|
|
"prose-then-user-tool": _Shape(
|
|
(_THINKING, _PROSE, _weather_tool(0)), "tool_calls", True, True, True, _PROSE_AND_WEATHER
|
|
),
|
|
"no-json-schema-prose-then-user-tool": _Shape(
|
|
(_THINKING, _PROSE, _weather_tool(0)), "tool_calls", True, True, False, _PROSE_AND_WEATHER
|
|
),
|
|
}
|
|
|
|
|
|
def _reasoning_block(text: str) -> JsonValue:
|
|
return {"type": "reasoning", "summary": [{"type": "summary_text", "text": text, "signature": "signature"}]}
|
|
|
|
|
|
def _text_delta(fragment: str) -> Mapping[str, JsonValue]:
|
|
return {"content": fragment}
|
|
|
|
|
|
def _part_deltas(part: _Part) -> tuple[Mapping[str, JsonValue], ...]:
|
|
match part:
|
|
case _Reasoning(text=text):
|
|
return ({"content": [_reasoning_block(text)]},)
|
|
case _Text(fragments=fragments):
|
|
return tuple(_text_delta(fragment) for fragment in fragments)
|
|
case _ToolUse(index=index, name=name, fragments=fragments):
|
|
opening: Final[Mapping[str, JsonValue]] = {
|
|
"tool_calls": [
|
|
{
|
|
"index": index,
|
|
"id": f"toolu_{name}",
|
|
"type": "function",
|
|
"function": {"name": name, "arguments": ""},
|
|
}
|
|
]
|
|
}
|
|
arguments: Final[tuple[Mapping[str, JsonValue], ...]] = tuple(
|
|
{"tool_calls": [{"index": index, "function": {"arguments": fragment}}]} for fragment in fragments
|
|
)
|
|
return (opening, *arguments)
|
|
|
|
|
|
def _answer_message(parts: tuple[_Part, ...]) -> JsonValue:
|
|
reasoning: Final = tuple(_reasoning_block(part.text) for part in parts if isinstance(part, _Reasoning))
|
|
texts: Final[tuple[JsonValue, ...]] = tuple(
|
|
{"type": "text", "text": "".join(part.fragments)} for part in parts if isinstance(part, _Text)
|
|
)
|
|
tool_calls: Final[tuple[JsonValue, ...]] = tuple(
|
|
{
|
|
"id": f"toolu_{part.name}",
|
|
"type": "function",
|
|
"function": {"name": part.name, "arguments": "".join(part.fragments)},
|
|
}
|
|
for part in parts
|
|
if isinstance(part, _ToolUse)
|
|
)
|
|
return {"role": "assistant", "content": [*reasoning, *texts], "tool_calls": list(tool_calls)}
|
|
|
|
|
|
def _answer_reply(shape: _Shape, stream: bool) -> Reply:
|
|
identity: Final = f"databricks-json-answer-{uuid.uuid4().hex}"
|
|
if not stream:
|
|
completion: Final = {
|
|
"id": identity,
|
|
"object": "chat.completion",
|
|
"created": 1,
|
|
"model": _CLAUDE_BACKEND,
|
|
"choices": [{"index": 0, "message": _answer_message(shape.parts), "finish_reason": shape.upstream_finish}],
|
|
"usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
|
|
}
|
|
return Reply(body=json.dumps(completion).encode())
|
|
deltas: Final = ({"role": "assistant"}, *itertools.chain.from_iterable(_part_deltas(part) for part in shape.parts))
|
|
frames: Final = (
|
|
*(
|
|
_frame(identity, [{"index": 0, "delta": delta, "finish_reason": None}], model=_CLAUDE_BACKEND)
|
|
for delta in deltas
|
|
),
|
|
_frame(identity, [{"index": 0, "delta": {}, "finish_reason": shape.upstream_finish}], model=_CLAUDE_BACKEND),
|
|
b"data: [DONE]\n\n",
|
|
)
|
|
return Reply(content_type="text/event-stream", chunks=frames)
|
|
|
|
|
|
def _answer_request(
|
|
endpoint: _Endpoint, model: str, stream: bool, shape: _Shape, marker: str
|
|
) -> Mapping[str, JsonValue]:
|
|
prompt: Final = f"{_JSON_SCHEMA_PROMPT} {marker}"
|
|
match endpoint:
|
|
case "chat":
|
|
return {
|
|
"model": model,
|
|
"messages": [{"role": "user", "content": prompt}],
|
|
"max_tokens": 2048,
|
|
"stream": stream,
|
|
**({"response_format": _json_schema_format(_PERSON_SCHEMA)} if shape.json_schema else {}),
|
|
**(
|
|
{
|
|
"tools": [
|
|
{
|
|
"type": "function",
|
|
"function": {"name": "get_weather", "parameters": dict(_WEATHER_PARAMETERS)},
|
|
}
|
|
]
|
|
}
|
|
if shape.user_tools
|
|
else {}
|
|
),
|
|
**({"reasoning_effort": "low"} if shape.thinking else {}),
|
|
}
|
|
case "responses":
|
|
return {
|
|
"model": model,
|
|
"input": prompt,
|
|
"max_output_tokens": 2048,
|
|
"stream": stream,
|
|
**(
|
|
{
|
|
"text": {
|
|
"format": {
|
|
"type": "json_schema",
|
|
"name": "P",
|
|
"strict": True,
|
|
"schema": dict(_PERSON_SCHEMA),
|
|
}
|
|
}
|
|
}
|
|
if shape.json_schema
|
|
else {}
|
|
),
|
|
**(
|
|
{"tools": [{"type": "function", "name": "get_weather", "parameters": dict(_WEATHER_PARAMETERS)}]}
|
|
if shape.user_tools
|
|
else {}
|
|
),
|
|
**({"reasoning": {"effort": "low"}} if shape.thinking else {}),
|
|
}
|
|
case "messages":
|
|
return {
|
|
"model": model,
|
|
"max_tokens": 2048,
|
|
"messages": [{"role": "user", "content": prompt}],
|
|
"stream": stream,
|
|
**(
|
|
{"output_format": {"type": "json_schema", "schema": dict(_PERSON_SCHEMA)}}
|
|
if shape.json_schema
|
|
else {}
|
|
),
|
|
**(
|
|
{"tools": [{"name": "get_weather", "input_schema": dict(_WEATHER_PARAMETERS)}]}
|
|
if shape.user_tools
|
|
else {}
|
|
),
|
|
**({"thinking": {"type": "enabled", "budget_tokens": 1024}} if shape.thinking else {}),
|
|
}
|
|
|
|
|
|
class _OutboundFunction(BaseModel):
|
|
model_config = ConfigDict(extra="ignore", frozen=True)
|
|
name: str
|
|
|
|
|
|
class _OutboundTool(BaseModel):
|
|
model_config = ConfigDict(extra="ignore", frozen=True)
|
|
function: _OutboundFunction
|
|
|
|
|
|
class _OutboundBody(BaseModel):
|
|
model_config = ConfigDict(extra="ignore", frozen=True)
|
|
model: str
|
|
stream: bool | None = None
|
|
tools: tuple[_OutboundTool, ...] = ()
|
|
tool_choice: JsonValue = None
|
|
|
|
|
|
def _assert_answer_outbound(request: Request, shape: _Shape, stream: bool, marker: str) -> None:
|
|
assert (request.method, request.target) == ("POST", "/chat/completions")
|
|
assert marker.encode() in request.body
|
|
body: Final = _OutboundBody.model_validate_json(request.body)
|
|
tool_names: Final = (
|
|
*(("get_weather",) if shape.user_tools else ()),
|
|
*(("json_tool_call",) if shape.json_schema else ()),
|
|
)
|
|
assert (body.model, body.stream is True, sorted(tool.function.name for tool in body.tools)) == (
|
|
_CLAUDE_BACKEND,
|
|
stream,
|
|
sorted(tool_names),
|
|
)
|
|
assert (body.tool_choice == _FORCED_JSON_TOOL_CHOICE) == (shape.json_schema and not shape.thinking), body
|
|
|
|
|
|
class _ChatToolFunction(BaseModel):
|
|
model_config = ConfigDict(extra="ignore", frozen=True)
|
|
name: str | None = None
|
|
arguments: str | None = None
|
|
|
|
|
|
class _ChatToolCallDelta(BaseModel):
|
|
model_config = ConfigDict(extra="ignore", frozen=True)
|
|
index: int
|
|
function: _ChatToolFunction | None = None
|
|
|
|
|
|
class _ChatAnswerDelta(BaseModel):
|
|
model_config = ConfigDict(extra="ignore", frozen=True)
|
|
content: str | None = None
|
|
reasoning_content: str | None = None
|
|
tool_calls: tuple[_ChatToolCallDelta, ...] | None = None
|
|
|
|
|
|
class _ChatAnswerStreamChoice(BaseModel):
|
|
model_config = ConfigDict(extra="ignore", frozen=True)
|
|
delta: _ChatAnswerDelta
|
|
finish_reason: str | None = None
|
|
|
|
|
|
class _ChatAnswerChunk(BaseModel):
|
|
model_config = ConfigDict(extra="ignore", frozen=True)
|
|
choices: tuple[_ChatAnswerStreamChoice, ...]
|
|
|
|
|
|
class _ChatToolCall(BaseModel):
|
|
model_config = ConfigDict(extra="ignore", frozen=True)
|
|
function: _ChatToolFunction
|
|
|
|
|
|
class _ChatAnswerMessage(BaseModel):
|
|
model_config = ConfigDict(extra="ignore", frozen=True)
|
|
content: str | None = None
|
|
reasoning_content: str | None = None
|
|
tool_calls: tuple[_ChatToolCall, ...] | None = None
|
|
|
|
|
|
class _ChatAnswerChoice(BaseModel):
|
|
model_config = ConfigDict(extra="ignore", frozen=True)
|
|
message: _ChatAnswerMessage
|
|
finish_reason: str | None = None
|
|
|
|
|
|
class _ChatAnswer(BaseModel):
|
|
model_config = ConfigDict(extra="ignore", frozen=True)
|
|
choices: tuple[_ChatAnswerChoice, ...]
|
|
|
|
|
|
class _ResponsesAnswerPart(BaseModel):
|
|
model_config = ConfigDict(extra="ignore", frozen=True)
|
|
type: str | None = None
|
|
text: str | None = None
|
|
|
|
|
|
class _ResponsesAnswerItem(BaseModel):
|
|
model_config = ConfigDict(extra="ignore", frozen=True)
|
|
type: str
|
|
name: str | None = None
|
|
arguments: str | None = None
|
|
content: tuple[_ResponsesAnswerPart, ...] | None = None
|
|
|
|
|
|
class _ResponsesAnswer(BaseModel):
|
|
model_config = ConfigDict(extra="ignore", frozen=True)
|
|
status: str | None = None
|
|
output: tuple[_ResponsesAnswerItem, ...]
|
|
|
|
|
|
class _ResponsesAnswerEvent(BaseModel):
|
|
model_config = ConfigDict(extra="ignore", frozen=True)
|
|
type: str
|
|
response: _ResponsesAnswer | None = None
|
|
|
|
|
|
class _MessagesAnswerBlock(BaseModel):
|
|
model_config = ConfigDict(extra="ignore", frozen=True)
|
|
type: str
|
|
text: str | None = None
|
|
name: str | None = None
|
|
input: JsonValue = None
|
|
|
|
|
|
class _MessagesAnswer(BaseModel):
|
|
model_config = ConfigDict(extra="ignore", frozen=True)
|
|
content: tuple[_MessagesAnswerBlock, ...]
|
|
stop_reason: str | None = None
|
|
|
|
|
|
class _MessagesAnswerDelta(BaseModel):
|
|
model_config = ConfigDict(extra="ignore", frozen=True)
|
|
type: str | None = None
|
|
text: str | None = None
|
|
partial_json: str | None = None
|
|
stop_reason: str | None = None
|
|
|
|
|
|
class _MessagesAnswerEvent(BaseModel):
|
|
model_config = ConfigDict(extra="ignore", frozen=True)
|
|
type: str
|
|
index: int | None = None
|
|
content_block: _MessagesAnswerBlock | None = None
|
|
delta: _MessagesAnswerDelta | None = None
|
|
|
|
|
|
def _only(values: tuple[str, ...]) -> str:
|
|
assert len(values) == 1, values
|
|
return values[0]
|
|
|
|
|
|
def _arguments(arguments: str | None) -> JsonValue:
|
|
return _JSON_VALUE.validate_json(arguments or "{}")
|
|
|
|
|
|
def _chat_answer(body: _ChatAnswer) -> _Answer:
|
|
assert len(body.choices) == 1, body
|
|
message: Final = body.choices[0].message
|
|
return _Answer(
|
|
text=message.content or "",
|
|
tools=tuple(
|
|
(call.function.name or "", _arguments(call.function.arguments)) for call in message.tool_calls or ()
|
|
),
|
|
finish=body.choices[0].finish_reason or "",
|
|
thought=bool(message.reasoning_content),
|
|
)
|
|
|
|
|
|
def _chat_stream_tool(fragments: tuple[_ChatToolCallDelta, ...]) -> tuple[str, JsonValue]:
|
|
functions: Final = tuple(fragment.function for fragment in fragments if fragment.function is not None)
|
|
names: Final = tuple(function.name for function in functions if function.name)
|
|
return _only(names), _arguments("".join(function.arguments or "" for function in functions))
|
|
|
|
|
|
def _chat_stream_answer(data: tuple[str, ...]) -> _Answer:
|
|
assert data[-1] == "[DONE]", data
|
|
choices: Final = tuple(
|
|
itertools.chain.from_iterable(_ChatAnswerChunk.model_validate_json(line).choices for line in data[:-1])
|
|
)
|
|
deltas: Final = tuple(choice.delta for choice in choices)
|
|
tool_deltas: Final = tuple(itertools.chain.from_iterable(delta.tool_calls or () for delta in deltas))
|
|
indexes: Final = tuple(dict.fromkeys(tool_delta.index for tool_delta in tool_deltas))
|
|
assert indexes == tuple(range(len(indexes))), indexes
|
|
return _Answer(
|
|
text="".join(delta.content or "" for delta in deltas),
|
|
tools=tuple(
|
|
_chat_stream_tool(tuple(tool_delta for tool_delta in tool_deltas if tool_delta.index == index))
|
|
for index in indexes
|
|
),
|
|
finish=_only(tuple(choice.finish_reason for choice in choices if choice.finish_reason)),
|
|
thought=any(delta.reasoning_content for delta in deltas),
|
|
)
|
|
|
|
|
|
def _responses_item_text(item: _ResponsesAnswerItem) -> str:
|
|
return "".join(part.text or "" for part in item.content or ()) if item.type == "message" else ""
|
|
|
|
|
|
def _responses_answer(body: _ResponsesAnswer) -> _Answer:
|
|
return _Answer(
|
|
text="".join(_responses_item_text(item) for item in body.output),
|
|
tools=tuple(
|
|
(item.name or "", _arguments(item.arguments)) for item in body.output if item.type == "function_call"
|
|
),
|
|
finish=body.status or "",
|
|
thought=any(item.type == "reasoning" for item in body.output),
|
|
)
|
|
|
|
|
|
def _responses_stream_answer(data: tuple[str, ...]) -> _Answer:
|
|
events: Final = tuple(_ResponsesAnswerEvent.model_validate_json(line) for line in data if line != "[DONE]")
|
|
completed: Final = tuple(event.response for event in events if event.type == "response.completed")
|
|
assert len(completed) == 1 and completed[0] is not None, data
|
|
return _responses_answer(completed[0])
|
|
|
|
|
|
def _messages_answer(body: _MessagesAnswer) -> _Answer:
|
|
return _Answer(
|
|
text="".join(block.text or "" for block in body.content if block.type == "text"),
|
|
tools=tuple((block.name or "", block.input) for block in body.content if block.type == "tool_use"),
|
|
finish=body.stop_reason or "",
|
|
thought=any(block.type == "thinking" for block in body.content),
|
|
)
|
|
|
|
|
|
def _messages_streamed_block(
|
|
start: _MessagesAnswerEvent, events: tuple[_MessagesAnswerEvent, ...]
|
|
) -> _MessagesAnswerBlock:
|
|
assert start.content_block is not None, start
|
|
block_events: Final = tuple(
|
|
event for event in events if event.type == "content_block_delta" and event.index == start.index
|
|
)
|
|
deltas: Final = tuple(event.delta for event in block_events if event.delta is not None)
|
|
partial_json: Final = "".join(delta.partial_json or "" for delta in deltas)
|
|
return start.content_block.model_copy(
|
|
update={
|
|
"text": "".join(delta.text or "" for delta in deltas),
|
|
"input": _JSON_VALUE.validate_json(partial_json) if partial_json else start.content_block.input,
|
|
}
|
|
)
|
|
|
|
|
|
def _messages_stream_answer(data: tuple[str, ...]) -> _Answer:
|
|
events: Final = tuple(_MessagesAnswerEvent.model_validate_json(line) for line in data)
|
|
assert events[-1].type == "message_stop", data
|
|
starts: Final = tuple(event for event in events if event.type == "content_block_start")
|
|
stop_reasons: Final = tuple(
|
|
event.delta.stop_reason for event in events if event.type == "message_delta" and event.delta is not None
|
|
)
|
|
return _messages_answer(
|
|
_MessagesAnswer(
|
|
content=tuple(_messages_streamed_block(start, events) for start in starts),
|
|
stop_reason=_only(tuple(reason for reason in stop_reasons if reason)),
|
|
)
|
|
)
|
|
|
|
|
|
def _final_answer(endpoint: _Endpoint, payload: bytes) -> _Answer:
|
|
match endpoint:
|
|
case "chat":
|
|
return _chat_answer(_ChatAnswer.model_validate_json(payload))
|
|
case "responses":
|
|
return _responses_answer(_ResponsesAnswer.model_validate_json(payload))
|
|
case "messages":
|
|
return _messages_answer(_MessagesAnswer.model_validate_json(payload))
|
|
|
|
|
|
def _streamed_answer(endpoint: _Endpoint, data: tuple[str, ...]) -> _Answer:
|
|
match endpoint:
|
|
case "chat":
|
|
return _chat_stream_answer(data)
|
|
case "responses":
|
|
return _responses_stream_answer(data)
|
|
case "messages":
|
|
return _messages_stream_answer(data)
|
|
|
|
|
|
def _call_for_answer(gateway: Gateway, endpoint: _Endpoint, body: Mapping[str, JsonValue]) -> _Answer:
|
|
stream: Final = body["stream"] is True
|
|
if not stream:
|
|
response: Final = gateway.request("POST", _PATHS[endpoint], body)
|
|
assert response.status_code == 200, response.text
|
|
return _final_answer(endpoint, response.content)
|
|
with gateway.client.stream(
|
|
"POST", _PATHS[endpoint], json=dict(body), headers={"Authorization": f"Bearer {gateway.key}"}
|
|
) as streamed:
|
|
assert streamed.status_code == 200, streamed.read()
|
|
data: Final = tuple(line.removeprefix("data: ") for line in streamed.iter_lines() if line.startswith("data: "))
|
|
return _streamed_answer(endpoint, data)
|
|
|
|
|
|
_ANTHROPIC_STOP_REASONS: Final[Mapping[str, str]] = {"stop": "end_turn", "tool_calls": "tool_use"}
|
|
|
|
|
|
def _endpoint_answer(endpoint: _Endpoint, answer: _Answer) -> _Answer:
|
|
match endpoint:
|
|
case "chat":
|
|
return answer
|
|
case "responses":
|
|
return _Answer(answer.text, answer.tools, "completed", answer.thought)
|
|
case "messages":
|
|
return _Answer(answer.text, answer.tools, _ANTHROPIC_STOP_REASONS[answer.finish], answer.thought)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("endpoint", "stream", "shape"),
|
|
[
|
|
pytest.param("chat", True, "json-only", id="chat-stream-json-only"),
|
|
pytest.param("chat", True, "prose-then-json", id="chat-stream-prose-then-json"),
|
|
pytest.param("chat", True, "json-then-user-tool", id="chat-stream-json-then-user-tool"),
|
|
pytest.param("chat", True, "user-tool-then-json", id="chat-stream-user-tool-then-json"),
|
|
pytest.param("chat", True, "prose-user-tool-json", id="chat-stream-prose-user-tool-json"),
|
|
pytest.param("chat", True, "text-answer", id="chat-stream-text-answer"),
|
|
pytest.param("chat", True, "prose-then-user-tool", id="chat-stream-prose-then-user-tool"),
|
|
pytest.param("chat", True, "no-json-schema-prose-then-user-tool", id="chat-stream-no-json-schema"),
|
|
pytest.param("responses", True, "prose-then-json", id="responses-stream-prose-then-json"),
|
|
pytest.param("responses", True, "json-then-user-tool", id="responses-stream-json-then-user-tool"),
|
|
pytest.param("responses", True, "prose-user-tool-json", id="responses-stream-prose-user-tool-json"),
|
|
pytest.param("responses", True, "prose-then-user-tool", id="responses-stream-prose-then-user-tool"),
|
|
pytest.param("messages", True, "json-only", id="messages-stream-json-only"),
|
|
pytest.param("messages", True, "prose-then-json", id="messages-stream-prose-then-json"),
|
|
pytest.param("messages", True, "json-then-user-tool", id="messages-stream-json-then-user-tool"),
|
|
pytest.param("messages", True, "prose-user-tool-json", id="messages-stream-prose-user-tool-json"),
|
|
pytest.param("messages", True, "prose-then-user-tool", id="messages-stream-prose-then-user-tool"),
|
|
pytest.param("messages", True, "no-json-schema-prose-then-user-tool", id="messages-stream-no-json-schema"),
|
|
pytest.param("chat", False, "prose-user-tool-json", id="chat-prose-user-tool-json"),
|
|
pytest.param("responses", False, "prose-user-tool-json", id="responses-prose-user-tool-json"),
|
|
pytest.param("messages", False, "prose-user-tool-json", id="messages-prose-user-tool-json"),
|
|
],
|
|
)
|
|
def test_databricks_claude_json_schema_answer_reaches_every_endpoint_as_text(
|
|
gateway: Gateway, endpoint: _Endpoint, stream: bool, shape: str
|
|
) -> None:
|
|
scripted: Final = _SHAPES[shape]
|
|
marker: Final = uuid.uuid4().hex
|
|
with wire_server(lambda _request: _answer_reply(scripted, stream)) as wire, gateway.scenario() as scenario:
|
|
model: Final = scenario.model(model=f"databricks/{_CLAUDE_BACKEND}", api_base=wire.url, api_key=_API_KEY)
|
|
answer: Final = _call_for_answer(gateway, endpoint, _answer_request(endpoint, model, stream, scripted, marker))
|
|
received: Final = wire.drain()
|
|
assert len(received) == 1, received
|
|
_assert_answer_outbound(received[0], scripted, stream, marker)
|
|
assert answer == _endpoint_answer(endpoint, scripted.expected)
|