diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py
index 76adadb2e1f..07534e70a27 100644
--- a/litellm/litellm_core_utils/prompt_templates/factory.py
+++ b/litellm/litellm_core_utils/prompt_templates/factory.py
@@ -9,10 +9,11 @@ import xml.etree.ElementTree as ET
from collections.abc import Container, Iterable, Iterator, Mapping, Sequence
from enum import Enum
from types import MappingProxyType
-from typing import Any, Final, TypeAlias, TypedDict, cast, overload
+from typing import Annotated, Any, Final, TypeAlias, TypedDict, cast, overload
from jinja2.sandbox import ImmutableSandboxedEnvironment
-from pydantic import BaseModel, JsonValue, TypeAdapter
+from pydantic import BaseModel, BeforeValidator, ConfigDict, JsonValue, TypeAdapter, ValidationError, model_validator
+from pydantic_core import ErrorDetails, PydanticCustomError
import litellm
import litellm.types
@@ -5565,7 +5566,7 @@ def function_call_prompt(
]
-def response_schema_prompt(model: str, response_schema: dict) -> str:
+def response_schema_prompt(model: str, response_schema: dict, llm_provider: str) -> str:
"""
Decides if a user-defined custom prompt or default needs to be used
@@ -5581,11 +5582,8 @@ def response_schema_prompt(model: str, response_schema: dict) -> str:
custom_prompt_details = litellm.custom_prompt_dict["response_schema_prompt"]
if custom_prompt_details is not None:
- return custom_prompt(
- role_dict=custom_prompt_details["roles"],
- initial_prompt_value=custom_prompt_details["initial_prompt_value"],
- final_prompt_value=custom_prompt_details["final_prompt_value"],
- messages=response_schema_as_message,
+ return render_custom_prompt_template(
+ template=custom_prompt_details, messages=response_schema_as_message, model=model, llm_provider=llm_provider
)
else:
return default_response_schema_prompt(response_schema=response_schema)
@@ -5605,8 +5603,95 @@ def default_response_schema_prompt(response_schema: dict) -> str:
# Custom prompt template
+def _mapping_or_none(value: object) -> object:
+ if value is None or issubclass(type(value), Mapping):
+ return value
+ verbose_logger.warning(
+ "Ignoring a custom prompt template roles value of type %s because it is not a mapping, "
+ "so those messages render without role markers",
+ type(value).__name__,
+ )
+ return None
+
+
+class CustomPromptRole(BaseModel):
+ model_config = ConfigDict(frozen=True)
+
+ pre_message: str | None = None
+ post_message: str | None = None
+
+
+class CustomPromptRoles(BaseModel):
+ model_config = ConfigDict(frozen=True, title="custom prompt template")
+
+ roles: Mapping[str, Annotated[CustomPromptRole | None, BeforeValidator(_mapping_or_none)]]
+
+
+class CustomPromptTemplate(BaseModel):
+ model_config = ConfigDict(frozen=True, title="custom prompt template")
+
+ roles: Annotated[Mapping[object, object] | None, BeforeValidator(_mapping_or_none)] = None
+ initial_prompt_value: str | None = None
+ final_prompt_value: str | None = None
+
+ @model_validator(mode="after")
+ def _sets_a_rendered_field(self) -> "CustomPromptTemplate":
+ if self.roles is None and self.initial_prompt_value is None and self.final_prompt_value is None:
+ raise PydanticCustomError(
+ "empty_custom_prompt_template",
+ "it sets none of roles, initial_prompt_value, or final_prompt_value",
+ )
+ return self
+
+
+def render_custom_prompt_template(
+ template: Mapping[str, object],
+ messages: Sequence[Mapping[str, object]],
+ model: str,
+ llm_provider: str,
+) -> str:
+ try:
+ return _render_custom_prompt_template(template=template, messages=messages)
+ except ValidationError as error:
+ raise litellm.APIConnectionError(
+ message=f"The custom prompt template for {model} is invalid: {_describe_template_errors(error)}",
+ model=model,
+ llm_provider=llm_provider,
+ ) from error
+
+
+def _describe_template_errors(error: ValidationError) -> str:
+ return "; ".join(
+ _describe_template_error(detail) for detail in error.errors(include_url=False, include_input=False)
+ )
+
+
+def _describe_template_error(detail: ErrorDetails) -> str:
+ location: Final = ".".join(str(part) for part in detail["loc"])
+ return f"{location}: {detail['msg']}" if location else detail["msg"]
+
+
+def _render_custom_prompt_template(template: Mapping[str, object], messages: Sequence[Mapping[str, object]]) -> str:
+ validated: Final = CustomPromptTemplate.model_validate(template)
+ message_roles: Final = frozenset(message.get("role") for message in messages)
+ template_roles: Final = validated.roles or {}
+ used_roles: Final = CustomPromptRoles.model_validate(
+ {"roles": {role: entry for role, entry in template_roles.items() if role in message_roles}}
+ )
+ return custom_prompt(
+ role_dict={
+ role: {"pre_message": entry.pre_message or "", "post_message": entry.post_message or ""}
+ for role, entry in used_roles.roles.items()
+ if entry is not None
+ },
+ initial_prompt_value=validated.initial_prompt_value or "",
+ final_prompt_value=validated.final_prompt_value or "",
+ messages=list(messages),
+ )
+
+
def custom_prompt(
- role_dict: dict,
+ role_dict: Mapping[str, Mapping[str, str]],
messages: list,
initial_prompt_value: str = "",
final_prompt_value: str = "",
diff --git a/litellm/llms/anthropic/completion/transformation.py b/litellm/llms/anthropic/completion/transformation.py
index e3ad961ed76..73e94e85f04 100644
--- a/litellm/llms/anthropic/completion/transformation.py
+++ b/litellm/llms/anthropic/completion/transformation.py
@@ -14,8 +14,8 @@ import httpx
import litellm
from litellm.constants import DEFAULT_MAX_TOKENS
from litellm.litellm_core_utils.prompt_templates.factory import (
- custom_prompt,
prompt_factory,
+ render_custom_prompt_template,
)
from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator
from litellm.llms.base_llm.chat.transformation import (
@@ -244,11 +244,8 @@ class AnthropicTextConfig(BaseConfig):
if model in custom_prompt_dict:
# check if the model has a registered custom prompt
model_prompt_details: Final = custom_prompt_dict[model]
- prompt = custom_prompt(
- role_dict=model_prompt_details["roles"],
- initial_prompt_value=model_prompt_details["initial_prompt_value"],
- final_prompt_value=model_prompt_details["final_prompt_value"],
- messages=messages,
+ prompt = render_custom_prompt_template(
+ template=model_prompt_details, messages=messages, model=model, llm_provider="anthropic"
)
else:
prompt = prompt_factory(model=model, messages=messages, custom_llm_provider="anthropic")
diff --git a/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py
index 1b3b0c344df..f026e4f9a99 100644
--- a/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py
+++ b/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py
@@ -14,9 +14,9 @@ from litellm.litellm_core_utils.get_litellm_params import stored_control_options
from litellm.litellm_core_utils.logging_utils import track_llm_api_timing
from litellm.litellm_core_utils.prompt_templates.factory import (
cohere_message_pt,
- custom_prompt,
deepseek_r1_pt,
prompt_factory,
+ render_custom_prompt_template,
)
from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException
from litellm.llms.bedrock.chat.invoke_handler import make_call, make_sync_call
@@ -601,7 +601,7 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM):
return cast(litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL, provider)
return None
- def convert_messages_to_prompt(self, model, messages, provider, custom_prompt_dict) -> tuple[str, list | None]:
+ def convert_messages_to_prompt(self, model: str, messages, provider, custom_prompt_dict) -> tuple[str, list | None]:
# handle anthropic prompts and amazon titan prompts
prompt = ""
chat_history: list | None = None
@@ -609,11 +609,8 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM):
if model in custom_prompt_dict:
# check if the model has a registered custom prompt
model_prompt_details: Final = custom_prompt_dict[model]
- prompt = custom_prompt(
- role_dict=model_prompt_details["roles"],
- initial_prompt_value=model_prompt_details.get("initial_prompt_value", ""),
- final_prompt_value=model_prompt_details.get("final_prompt_value", ""),
- messages=messages,
+ prompt = render_custom_prompt_template(
+ template=model_prompt_details, messages=messages, model=model, llm_provider="bedrock"
)
return prompt, None
## ELSE
diff --git a/litellm/llms/codestral/completion/handler.py b/litellm/llms/codestral/completion/handler.py
index f8486d3b274..08799c3a0a5 100644
--- a/litellm/llms/codestral/completion/handler.py
+++ b/litellm/llms/codestral/completion/handler.py
@@ -13,8 +13,8 @@ import litellm
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
from litellm.litellm_core_utils.logging_utils import track_llm_api_timing
from litellm.litellm_core_utils.prompt_templates.factory import (
- custom_prompt,
prompt_factory,
+ render_custom_prompt_template,
)
from litellm.llms.custom_httpx.http_handler import (
AsyncHTTPHandler,
@@ -261,11 +261,8 @@ class CodestralTextCompletion:
if model in custom_prompt_dict:
# check if the model has a registered custom prompt
model_prompt_details: Final = custom_prompt_dict[model]
- prompt = custom_prompt(
- role_dict=model_prompt_details["roles"],
- initial_prompt_value=model_prompt_details["initial_prompt_value"],
- final_prompt_value=model_prompt_details["final_prompt_value"],
- messages=messages,
+ prompt = render_custom_prompt_template(
+ template=model_prompt_details, messages=messages, model=model, llm_provider="codestral"
)
else:
prompt = prompt_factory(model=model, messages=messages)
diff --git a/litellm/llms/ollama/completion/transformation.py b/litellm/llms/ollama/completion/transformation.py
index 703d45acdf2..ac41310fb61 100644
--- a/litellm/llms/ollama/completion/transformation.py
+++ b/litellm/llms/ollama/completion/transformation.py
@@ -16,8 +16,8 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
)
from litellm.litellm_core_utils.prompt_templates.factory import (
convert_to_ollama_image,
- custom_prompt,
ollama_pt,
+ render_custom_prompt_template,
)
from litellm.litellm_core_utils.prompt_templates.image_handling import (
async_inline_remote_media,
@@ -437,11 +437,8 @@ class OllamaConfig(BaseConfig):
if model in custom_prompt_dict:
# check if the model has a registered custom prompt
model_prompt_details: Final = custom_prompt_dict[model]
- ollama_prompt = custom_prompt(
- role_dict=model_prompt_details["roles"],
- initial_prompt_value=model_prompt_details["initial_prompt_value"],
- final_prompt_value=model_prompt_details["final_prompt_value"],
- messages=messages,
+ ollama_prompt = render_custom_prompt_template(
+ template=model_prompt_details, messages=messages, model=model, llm_provider="ollama"
)
elif text_completion_request: # handle `/completions` requests
ollama_prompt = get_str_from_messages(messages=messages)
diff --git a/litellm/llms/petals/completion/handler.py b/litellm/llms/petals/completion/handler.py
index 425d3729d7c..a2ccd7eaf08 100644
--- a/litellm/llms/petals/completion/handler.py
+++ b/litellm/llms/petals/completion/handler.py
@@ -5,8 +5,8 @@ from typing import Final
import litellm
from litellm.litellm_core_utils.prompt_templates.factory import (
- custom_prompt,
prompt_factory,
+ render_custom_prompt_template,
)
from litellm.llms.custom_httpx.http_handler import (
AsyncHTTPHandler,
@@ -43,11 +43,8 @@ def completion(
if model in litellm.custom_prompt_dict:
# check if the model has a registered custom prompt
model_prompt_details: Final = litellm.custom_prompt_dict[model]
- prompt = custom_prompt(
- role_dict=model_prompt_details["roles"],
- initial_prompt_value=model_prompt_details["initial_prompt_value"],
- final_prompt_value=model_prompt_details["final_prompt_value"],
- messages=messages,
+ prompt = render_custom_prompt_template(
+ template=model_prompt_details, messages=messages, model=model, llm_provider="petals"
)
else:
prompt = prompt_factory(model=model, messages=messages)
diff --git a/litellm/llms/predibase/chat/transformation.py b/litellm/llms/predibase/chat/transformation.py
index 7bd05c6e30b..0dd7e27722e 100644
--- a/litellm/llms/predibase/chat/transformation.py
+++ b/litellm/llms/predibase/chat/transformation.py
@@ -10,8 +10,8 @@ import litellm
from litellm.constants import DEFAULT_MAX_TOKENS
from litellm.litellm_core_utils.core_helpers import map_finish_reason
from litellm.litellm_core_utils.prompt_templates.factory import (
- custom_prompt,
prompt_factory,
+ render_custom_prompt_template,
)
from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException
from litellm.types.llms.openai import AllMessageValues
@@ -277,11 +277,8 @@ class PredibaseConfig(BaseConfig):
custom_prompt_dict: Final = litellm_params.get("custom_prompt_dict", {})
if model in custom_prompt_dict:
model_prompt_details: Final = custom_prompt_dict[model]
- prompt = custom_prompt(
- role_dict=model_prompt_details["roles"],
- initial_prompt_value=model_prompt_details["initial_prompt_value"],
- final_prompt_value=model_prompt_details["final_prompt_value"],
- messages=messages,
+ prompt = render_custom_prompt_template(
+ template=model_prompt_details, messages=messages, model=model, llm_provider="predibase"
)
else:
prompt = prompt_factory(model=model, messages=messages)
diff --git a/litellm/llms/sagemaker/completion/transformation.py b/litellm/llms/sagemaker/completion/transformation.py
index 274f884de0b..a69846c816f 100644
--- a/litellm/llms/sagemaker/completion/transformation.py
+++ b/litellm/llms/sagemaker/completion/transformation.py
@@ -15,8 +15,8 @@ from httpx._models import Headers, Response
import litellm
from litellm.litellm_core_utils.asyncify import asyncify
from litellm.litellm_core_utils.prompt_templates.factory import (
- custom_prompt,
prompt_factory,
+ render_custom_prompt_template,
)
from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException
from litellm.types.llms.openai import AllMessageValues
@@ -122,20 +122,14 @@ class SagemakerConfig(BaseConfig):
if model in custom_prompt_dict:
# check if the model has a registered custom prompt
model_prompt_details = custom_prompt_dict[model]
- prompt = custom_prompt(
- role_dict=model_prompt_details.get("roles", None),
- initial_prompt_value=model_prompt_details.get("initial_prompt_value", ""),
- final_prompt_value=model_prompt_details.get("final_prompt_value", ""),
- messages=messages,
+ prompt = render_custom_prompt_template(
+ template=model_prompt_details, messages=messages, model=model, llm_provider="sagemaker"
)
elif hf_model_name in custom_prompt_dict:
# check if the base huggingface model has a registered custom prompt
model_prompt_details = custom_prompt_dict[hf_model_name]
- prompt = custom_prompt(
- role_dict=model_prompt_details.get("roles", None),
- initial_prompt_value=model_prompt_details.get("initial_prompt_value", ""),
- final_prompt_value=model_prompt_details.get("final_prompt_value", ""),
- messages=messages,
+ prompt = render_custom_prompt_template(
+ template=model_prompt_details, messages=messages, model=model, llm_provider="sagemaker"
)
else:
if hf_model_name is None:
diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py
index 895d59965b7..7b832b321da 100644
--- a/litellm/llms/vertex_ai/gemini/transformation.py
+++ b/litellm/llms/vertex_ai/gemini/transformation.py
@@ -1154,6 +1154,7 @@ def transform_request_body(
user_response_schema_message: Final = response_schema_prompt(
model=model,
response_schema=optional_params.get("response_schema"),
+ llm_provider=custom_llm_provider,
)
messages.append({"role": "user", "content": user_response_schema_message})
optional_params.pop("response_schema")
diff --git a/litellm/llms/vllm/completion/handler.py b/litellm/llms/vllm/completion/handler.py
index 78e6c74c2f7..f330f0928d0 100644
--- a/litellm/llms/vllm/completion/handler.py
+++ b/litellm/llms/vllm/completion/handler.py
@@ -5,8 +5,8 @@ from typing import Final
import httpx
from litellm.litellm_core_utils.prompt_templates.factory import (
- custom_prompt,
prompt_factory,
+ render_custom_prompt_template,
)
from litellm.utils import ModelResponse, Usage
@@ -56,11 +56,8 @@ def completion(
if model in custom_prompt_dict:
# check if the model has a registered custom prompt
model_prompt_details: Final = custom_prompt_dict[model]
- prompt = custom_prompt(
- role_dict=model_prompt_details["roles"],
- initial_prompt_value=model_prompt_details["initial_prompt_value"],
- final_prompt_value=model_prompt_details["final_prompt_value"],
- messages=messages,
+ prompt = render_custom_prompt_template(
+ template=model_prompt_details, messages=messages, model=model, llm_provider="vllm"
)
else:
prompt = prompt_factory(model=model, messages=messages)
@@ -144,11 +141,8 @@ def batch_completions(model: str, messages: list, optional_params=None, custom_p
# check if the model has a registered custom prompt
model_prompt_details: Final = custom_prompt_dict[model]
for message in messages:
- prompt = custom_prompt(
- role_dict=model_prompt_details["roles"],
- initial_prompt_value=model_prompt_details["initial_prompt_value"],
- final_prompt_value=model_prompt_details["final_prompt_value"],
- messages=message,
+ prompt = render_custom_prompt_template(
+ template=model_prompt_details, messages=message, model=model, llm_provider="vllm"
)
prompts.append(prompt)
else:
diff --git a/litellm/llms/watsonx/common_utils.py b/litellm/llms/watsonx/common_utils.py
index fb1fa8b9ffb..6d3686e9f85 100644
--- a/litellm/llms/watsonx/common_utils.py
+++ b/litellm/llms/watsonx/common_utils.py
@@ -141,7 +141,7 @@ async def _aconvert_watsonx_messages_core(
model_prompt_dict: Final = custom_prompt_dict[model]
return ptf.custom_prompt(
messages=messages,
- role_dict=model_prompt_dict.get("role_dict", model_prompt_dict.get("roles")),
+ role_dict=model_prompt_dict.get("role_dict", model_prompt_dict.get("roles", {})),
initial_prompt_value=model_prompt_dict.get("initial_prompt_value", ""),
final_prompt_value=model_prompt_dict.get("final_prompt_value", ""),
bos_token=model_prompt_dict.get("bos_token", ""),
@@ -173,7 +173,7 @@ def _convert_watsonx_messages_core(
model_prompt_dict: Final = custom_prompt_dict[model]
return ptf.custom_prompt(
messages=messages,
- role_dict=model_prompt_dict.get("role_dict", model_prompt_dict.get("roles")),
+ role_dict=model_prompt_dict.get("role_dict", model_prompt_dict.get("roles", {})),
initial_prompt_value=model_prompt_dict.get("initial_prompt_value", ""),
final_prompt_value=model_prompt_dict.get("final_prompt_value", ""),
bos_token=model_prompt_dict.get("bos_token", ""),
diff --git a/litellm/types/litellm_params.py b/litellm/types/litellm_params.py
index 096ad1c30dd..8387afa22bb 100644
--- a/litellm/types/litellm_params.py
+++ b/litellm/types/litellm_params.py
@@ -260,6 +260,7 @@ class PromptOptions:
prompt_label: str | None = None
litellm_system_prompt: str | None = None
custom_prompt_dict: Mapping[str, object] | None = None
+ initial_prompt_value: str | None = None
roles: Mapping[str, object] | None = None
final_prompt_value: str | None = None
bos_token: str | None = None
diff --git a/tests/integration/providers/test_custom_prompt_template_chaos.py b/tests/integration/providers/test_custom_prompt_template_chaos.py
new file mode 100644
index 00000000000..c0da612af23
--- /dev/null
+++ b/tests/integration/providers/test_custom_prompt_template_chaos.py
@@ -0,0 +1,542 @@
+import asyncio
+import json
+import re
+import threading
+import uuid
+from collections.abc import Callable, Mapping
+from dataclasses import dataclass
+from pathlib import Path
+from queue import SimpleQueue
+from types import MappingProxyType
+from typing import Final, Literal
+from urllib.parse import urlsplit
+
+import httpx
+import psutil
+import pytest
+import yaml
+from integration._support.client import Gateway, eventually
+from integration._support.database import read_rows
+from integration._support.process import graceful_stop_seconds, owned_proxy_process
+from integration._support.wire import Reply, Request, Wire, wire_server
+from pydantic import JsonValue, TypeAdapter
+
+_BACKEND: Final = "llama3-custom-template-chaos"
+_API_KEY: Final = "synthetic-ollama-key"
+_CONFIG_MODEL: Final = "ollama-custom-template-chaos"
+_SYSTEM: Final = "Answer in one word."
+_ROLES: Final[dict[str, JsonValue]] = {
+ "system": {"pre_message": "", "post_message": ""},
+ "user": {"pre_message": "", "post_message": ""},
+}
+_REJECTION: Final = "initial_prompt_value: Input should be a valid string"
+_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
+_MARKER: Final = re.compile(r"marker-([0-9a-f]{32})")
+_STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]")
+_OWNED_PROXY_CELL_SECONDS: Final = 2 * graceful_stop_seconds() + 120
+
+Endpoint = Literal["chat", "messages", "responses"]
+
+
+@dataclass(frozen=True, slots=True)
+class _Call:
+ endpoint: Endpoint
+ stream: bool
+ marker: str
+ rejected: bool = False
+
+
+@dataclass(frozen=True, slots=True)
+class _Served:
+ call: _Call
+ status: int
+ text: str
+
+
+def _question(marker: str) -> str:
+ return f"What colour is the sky? marker-{marker}"
+
+
+def _prompt(marker: str) -> str:
+ return f"{_SYSTEM}{_question(marker)}"
+
+
+def _path(endpoint: Endpoint) -> str:
+ match endpoint:
+ case "chat":
+ return "/v1/chat/completions"
+ case "messages":
+ return "/v1/messages"
+ case "responses":
+ return "/v1/responses"
+
+
+def _body(model: str, call: _Call) -> dict[str, JsonValue]:
+ common: Final[dict[str, JsonValue]] = {
+ "model": model,
+ "stream": call.stream,
+ "num_retries": 0,
+ "cache": {"no-cache": True},
+ **({"initial_prompt_value": 5} if call.rejected else {}),
+ }
+ match call.endpoint:
+ case "chat":
+ return {
+ **common,
+ "messages": [
+ {"role": "system", "content": _SYSTEM},
+ {"role": "user", "content": _question(call.marker)},
+ ],
+ }
+ case "messages":
+ return {
+ **common,
+ "max_tokens": 64,
+ "system": _SYSTEM,
+ "messages": [{"role": "user", "content": _question(call.marker)}],
+ }
+ case "responses":
+ return {**common, "store": False, "instructions": _SYSTEM, "input": _question(call.marker)}
+
+
+def _generate_reply(marker: str, stream: bool, drop_connection: bool = False) -> Reply:
+ done: Final = {
+ "model": _BACKEND,
+ "created_at": "2026-10-15T00:00:00Z",
+ "response": "",
+ "done": True,
+ "done_reason": "stop",
+ "prompt_eval_count": 30,
+ "eval_count": 12,
+ }
+ if not stream:
+ return Reply(
+ body=json.dumps({**done, "response": f"blue marker-{marker}"}).encode(), drop_connection=drop_connection
+ )
+ frames: Final = (
+ *(
+ {"model": _BACKEND, "created_at": "2026-10-15T00:00:00Z", "response": piece, "done": False}
+ for piece in ("blue ", f"marker-{marker}")
+ ),
+ done,
+ )
+ return Reply(
+ content_type="application/x-ndjson",
+ chunks=tuple(json.dumps(frame).encode() + b"\n" for frame in frames),
+ drop_connection=drop_connection,
+ )
+
+
+def _is_generate(request: Request) -> bool:
+ return (request.method, request.target) == ("POST", "/api/generate")
+
+
+def _marker_of(request: Request) -> str:
+ found: Final = _MARKER.search(request.body.decode())
+ assert found is not None, request.body
+ return found.group(1)
+
+
+def _echo(request: Request) -> Reply:
+ if not _is_generate(request):
+ return Reply(body=b"{}")
+ stream: Final = _JSON_OBJECT.validate_json(request.body).get("stream") is True
+ return _generate_reply(_marker_of(request), stream)
+
+
+def _assert_each_prompt_rendered_once(received: tuple[Request, ...], markers: frozenset[str]) -> None:
+ generates: Final = tuple(request for request in received if _is_generate(request))
+ assert sorted(_marker_of(request) for request in generates) == sorted(markers)
+ for request in generates:
+ body: Final = _JSON_OBJECT.validate_json(request.body)
+ assert body["prompt"] == _prompt(_marker_of(request)), body["prompt"]
+ assert "roles" not in body and "initial_prompt_value" not in body, sorted(body)
+
+
+def _response_id(served: _Served) -> str | None:
+ if served.call.endpoint == "responses":
+ return None
+ if not served.call.stream:
+ identity: Final = _JSON_OBJECT.validate_json(served.text)["id"]
+ assert isinstance(identity, str)
+ return identity
+ for line in served.text.splitlines():
+ if not line.startswith("data: ") or line == "data: [DONE]":
+ continue
+ payload: Final = _JSON_OBJECT.validate_json(line.removeprefix("data: "))
+ if served.call.endpoint == "chat":
+ first: Final = payload["id"]
+ assert isinstance(first, str)
+ return first
+ if payload.get("type") == "message_start":
+ message: Final = payload["message"]
+ assert isinstance(message, dict) and isinstance(message["id"], str)
+ return message["id"]
+ raise AssertionError(served.text)
+
+
+def _spend_statuses(model: str, expected: int) -> MappingProxyType[str, str]:
+ rows: Final = eventually(
+ lambda: read_rows('SELECT request_id, status FROM "LiteLLM_SpendLogs" WHERE model_group=%s', (model,)),
+ lambda found: len(found) >= expected,
+ seconds=70,
+ )
+ statuses: Final = MappingProxyType({str(row["request_id"]): str(row["status"]) for row in rows})
+ assert len(statuses) == len(rows) == expected, rows
+ return statuses
+
+
+def _successes(statuses: Mapping[str, str]) -> frozenset[str]:
+ return frozenset(identity for identity, status in statuses.items() if status == "success")
+
+
+async def _send(client: httpx.AsyncClient, key: str, model: str, call: _Call) -> _Served:
+ async with client.stream(
+ "POST",
+ _path(call.endpoint),
+ json=_body(model, call),
+ headers={"Authorization": f"Bearer {key}", "anthropic-version": "2023-06-01"},
+ ) as response:
+ raw: Final = await response.aread()
+ return _Served(call=call, status=response.status_code, text=raw.decode())
+
+
+async def _burst(
+ base_url: str, key: str, model: str, calls: tuple[_Call, ...], *, tolerate_transport_errors: bool = False
+) -> tuple[_Served, ...]:
+ async with httpx.AsyncClient(base_url=base_url, timeout=60, trust_env=False) as client:
+ results: Final = await asyncio.gather(
+ *(_send(client, key, model, call) for call in calls), return_exceptions=tolerate_transport_errors
+ )
+ for result in results:
+ assert not isinstance(result, BaseException) or isinstance(result, httpx.TransportError), repr(result)
+ return tuple(result for result in results if isinstance(result, _Served))
+
+
+def _calls(count: int, endpoints: tuple[Endpoint, ...], stream: Callable[[int], bool]) -> tuple[_Call, ...]:
+ return tuple(
+ _Call(endpoint=endpoints[index % len(endpoints)], stream=stream(index), marker=uuid.uuid4().hex)
+ for index in range(count)
+ )
+
+
+def _assert_answered_with_its_own_marker(served: _Served) -> None:
+ assert served.status == 200, served.text
+ assert set(_MARKER.findall(served.text)) == {served.call.marker}, served.text
+
+
+def _assert_rejected_by_name(served: _Served) -> None:
+ assert served.status == 500, served.text
+ assert _REJECTION in served.text and "marker-" not in served.text, served.text
+
+
+async def test_concurrent_burst_renders_each_template_and_rejects_only_the_bad_requests(gateway: Gateway) -> None:
+ calls: Final = tuple(
+ _Call(endpoint=call.endpoint, stream=call.stream, marker=call.marker, rejected=index % 5 == 4)
+ for index, call in enumerate(_calls(30, ("chat", "messages", "responses"), lambda index: index % 2 == 0))
+ )
+ accepted: Final = tuple(call for call in calls if not call.rejected)
+ with wire_server(_echo) as wire, gateway.scenario() as scenario:
+ model: Final = scenario.model(model=f"ollama/{_BACKEND}", api_base=wire.url, api_key=_API_KEY, roles=_ROLES)
+ served: Final = await _burst(str(gateway.client.base_url), gateway.key, model, calls)
+ assert len(served) == 30
+ for item in served:
+ if item.call.rejected:
+ _assert_rejected_by_name(item)
+ else:
+ _assert_answered_with_its_own_marker(item)
+ _assert_each_prompt_rendered_once(wire.drain(), frozenset(call.marker for call in accepted))
+ answered: Final = tuple(item for item in served if not item.call.rejected)
+ known: Final = frozenset(identity for identity in map(_response_id, answered) if identity is not None)
+ assert len(known) == 16, known
+ statuses: Final = _spend_statuses(model, 30)
+ assert len(_successes(statuses)) == len(accepted), statuses
+ assert known <= _successes(statuses)
+
+
+async def test_dropped_ollama_connections_fail_their_callers_and_the_rest_keep_their_template(
+ gateway: Gateway,
+) -> None:
+ calls: Final = _calls(12, ("chat",), lambda index: index % 2 == 1)
+ dropped: Final = frozenset(call.marker for index, call in enumerate(calls) if index % 3 == 0)
+
+ def respond(request: Request) -> Reply:
+ if not _is_generate(request):
+ return Reply(body=b"{}")
+ marker: Final = _marker_of(request)
+ stream: Final = _JSON_OBJECT.validate_json(request.body).get("stream") is True
+ return _generate_reply(marker, stream, drop_connection=marker in dropped)
+
+ with wire_server(respond) as wire, gateway.scenario() as scenario:
+ model: Final = scenario.model(model=f"ollama/{_BACKEND}", api_base=wire.url, api_key=_API_KEY, roles=_ROLES)
+ served: Final = await _burst(str(gateway.client.base_url), gateway.key, model, calls)
+ assert len(served) == 12
+ for item in served:
+ if item.call.marker in dropped:
+ assert item.status == 500, item.text
+ assert "APIConnectionError" in item.text and "marker-" not in item.text, item.text
+ else:
+ _assert_answered_with_its_own_marker(item)
+ recovery: Final = _Call(endpoint="chat", stream=False, marker=uuid.uuid4().hex)
+ (recovered,) = await _burst(str(gateway.client.base_url), gateway.key, model, (recovery,))
+ _assert_answered_with_its_own_marker(recovered)
+ _assert_each_prompt_rendered_once(wire.drain(), frozenset(call.marker for call in (*calls, recovery)))
+ answered: Final = tuple(item for item in served if item.call.marker not in dropped)
+ survivors: Final = frozenset(
+ identity for identity in map(_response_id, (*answered, recovered)) if identity is not None
+ )
+ assert len(survivors) == 9, survivors
+ statuses: Final = _spend_statuses(model, 13)
+ assert _successes(statuses) == survivors, statuses
+ assert sum(status == "failure" for status in statuses.values()) == len(dropped), statuses
+
+
+def _owned_config(
+ wires: tuple[Wire, ...],
+ tmp_path: Path,
+ template: Mapping[str, JsonValue],
+ litellm_settings: Mapping[str, JsonValue],
+ router_settings: Mapping[str, JsonValue] = MappingProxyType({}),
+) -> Path:
+ base: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
+ deployments: Final = [
+ {
+ "model_name": _CONFIG_MODEL,
+ "litellm_params": {
+ "model": f"ollama/{_BACKEND}",
+ "api_base": wire.url,
+ "api_key": _API_KEY,
+ "order": order,
+ **template,
+ },
+ }
+ for order, wire in enumerate(wires, start=1)
+ ]
+ config: Final = {
+ **base,
+ "model_list": deployments,
+ "litellm_settings": {**base.get("litellm_settings", {}), **litellm_settings},
+ "router_settings": {**base.get("router_settings", {}), **router_settings},
+ }
+ path: Final = tmp_path / "ollama-custom-template-chaos.yaml"
+ path.write_text(yaml.safe_dump(config))
+ return path
+
+
+def _open_upstream_connections(pid: int, upstream: str) -> int:
+ port: Final = urlsplit(upstream).port
+ return sum(
+ 1
+ for connection in psutil.Process(pid).net_connections(kind="tcp")
+ if connection.status == psutil.CONN_ESTABLISHED and connection.raddr and connection.raddr.port == port
+ )
+
+
+@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS)
+async def test_worker_sigkill_mid_burst_leaves_the_sibling_rendering_templates(
+ gateway: Gateway, tmp_path: Path
+) -> None:
+ calls: Final = _calls(20, ("chat",), lambda _: False)
+ release: Final = threading.Event()
+ held_markers: Final[SimpleQueue[str]] = SimpleQueue()
+
+ def held(request: Request) -> Reply:
+ if not _is_generate(request):
+ return Reply(body=b"{}")
+ held_markers.put(_marker_of(request))
+ assert release.wait(timeout=60), "The burst was never released"
+ return _echo(request)
+
+ with wire_server(held) as wire:
+ path: Final = _owned_config((wire,), tmp_path, {"roles": _ROLES}, {})
+ with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned:
+ candidate: Final = owned.gateway
+ workers: Final = eventually(
+ lambda: tuple(int(pid) for pid in _STARTED_WORKER.findall(owned.log.read_text())),
+ lambda pids: len(pids) == 2,
+ seconds=30,
+ )
+ burst: Final = asyncio.create_task(
+ _burst(
+ str(candidate.client.base_url), candidate.key, _CONFIG_MODEL, calls, tolerate_transport_errors=True
+ )
+ )
+ await asyncio.to_thread(eventually, held_markers.qsize, lambda size: size == 20, 60)
+ held_by: Final = MappingProxyType({pid: _open_upstream_connections(pid, wire.url) for pid in workers})
+ assert sum(held_by.values()) == 20, held_by
+ victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__)
+ with owned.replacing_worker(victim_pid) as kill:
+ kill.freeze()
+ await asyncio.to_thread(kill)
+ release.set()
+ served: Final = await burst
+ assert held_by[survivor_pid] >= 10, held_by
+ assert len(served) == held_by[survivor_pid], (held_by, len(served))
+ for item in served:
+ _assert_answered_with_its_own_marker(item)
+ follow_up: Final = _Call(endpoint="chat", stream=False, marker=uuid.uuid4().hex)
+ (answered,) = await _burst(str(candidate.client.base_url), candidate.key, _CONFIG_MODEL, (follow_up,))
+ _assert_answered_with_its_own_marker(answered)
+ _assert_each_prompt_rendered_once(wire.drain(), frozenset(call.marker for call in (*calls, follow_up)))
+
+
+@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS)
+async def test_rejected_templates_never_cool_down_the_preferred_deployment(gateway: Gateway, tmp_path: Path) -> None:
+ rejected: Final = tuple(
+ _Call(endpoint="chat", stream=False, marker=uuid.uuid4().hex, rejected=True) for _ in range(4)
+ )
+ valid: Final = _calls(8, ("chat",), lambda _: False)
+ with wire_server(_echo) as preferred, wire_server(_echo) as standby:
+ path: Final = _owned_config(
+ (preferred, standby), tmp_path, {"roles": _ROLES}, {}, {"disable_cooldowns": False, "allowed_fails": 0}
+ )
+ with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned:
+ base_url: Final = str(owned.gateway.client.base_url)
+ for served in await _burst(base_url, owned.gateway.key, _CONFIG_MODEL, rejected):
+ _assert_rejected_by_name(served)
+ for served in await _burst(base_url, owned.gateway.key, _CONFIG_MODEL, valid):
+ _assert_answered_with_its_own_marker(served)
+ _assert_each_prompt_rendered_once(preferred.drain(), frozenset(call.marker for call in valid))
+ assert not any(_is_generate(request) for request in standby.drain())
+
+
+def _cached_chat(candidate: Gateway, marker: str, field: str, value: str) -> httpx.Response:
+ return candidate.request(
+ "POST",
+ "/v1/chat/completions",
+ {
+ "model": _CONFIG_MODEL,
+ "messages": [{"role": "system", "content": _SYSTEM}, {"role": "user", "content": _question(marker)}],
+ field: value,
+ },
+ )
+
+
+def _served_from_cache(candidate: Gateway, field: str) -> bool:
+ marker: Final = uuid.uuid4().hex
+ first: Final = _cached_chat(candidate, marker, field, "")
+ assert first.status_code == 200, first.text
+ assert "x-litellm-cache-key" not in first.headers, sorted(first.headers)
+ eventually(
+ lambda: _cached_chat(candidate, marker, field, ""),
+ lambda found: "x-litellm-cache-key" in found.headers,
+ seconds=20,
+ )
+ changed: Final = _cached_chat(candidate, marker, field, "")
+ assert changed.status_code == 200, changed.text
+ return "x-litellm-cache-key" in changed.headers
+
+
+@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS)
+def test_initial_and_final_prompt_values_key_the_provider_specific_cache_alike(
+ gateway: Gateway, tmp_path: Path
+) -> None:
+ template: Final[dict[str, JsonValue]] = {
+ "roles": _ROLES,
+ "initial_prompt_value": "",
+ "final_prompt_value": "",
+ }
+ with wire_server(_echo) as wire:
+ path: Final = _owned_config(
+ (wire,), tmp_path, template, {"enable_caching_on_provider_specific_optional_params": True}
+ )
+ with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned:
+ initial: Final = _served_from_cache(owned.gateway, "initial_prompt_value")
+ final: Final = _served_from_cache(owned.gateway, "final_prompt_value")
+ assert initial == final, {"initial_prompt_value": initial, "final_prompt_value": final}
+
+
+def _gemini_reply(_: Request) -> Reply:
+ candidate: Final[dict[str, JsonValue]] = {
+ "content": {"role": "model", "parts": [{"text": '{"colour": "blue"}'}]},
+ "finishReason": "STOP",
+ }
+ usage: Final[dict[str, JsonValue]] = {"promptTokenCount": 30, "candidatesTokenCount": 12, "totalTokenCount": 42}
+ return Reply(body=json.dumps({"candidates": [candidate], "usageMetadata": usage}).encode())
+
+
+def _texts(contents: JsonValue) -> tuple[str, ...]:
+ assert isinstance(contents, list), contents
+ parts: Final = tuple(part for content in contents if isinstance(content, dict) for part in _parts(content))
+ return tuple(text for part in parts if isinstance(text := part.get("text"), str))
+
+
+def _parts(content: dict[str, JsonValue]) -> tuple[dict[str, JsonValue], ...]:
+ parts: Final = content.get("parts")
+ assert isinstance(parts, list), content
+ return tuple(part for part in parts if isinstance(part, dict))
+
+
+def _anthropic_text_reply(_: Request) -> Reply:
+ completion: Final[dict[str, JsonValue]] = {
+ "type": "completion",
+ "id": "compl_scripted",
+ "completion": "blue",
+ "stop_reason": "stop_sequence",
+ }
+ return Reply(body=json.dumps(completion).encode())
+
+
+def _boot_deployment(name: str, model: str, wire: Wire) -> dict[str, JsonValue]:
+ return {"model_name": name, "litellm_params": {"model": model, "api_base": wire.url, "api_key": _API_KEY}}
+
+
+def _boot_chat(candidate: Gateway, model: str, extra: Mapping[str, JsonValue]) -> None:
+ response: Final = candidate.request(
+ "POST",
+ "/v1/chat/completions",
+ {
+ "model": model,
+ "messages": [{"role": "system", "content": _SYSTEM}, {"role": "user", "content": _question(model)}],
+ "cache": {"no-cache": True},
+ **extra,
+ },
+ )
+ assert response.status_code == 200, response.text
+
+
+def _anthropic_text_prompts(wire: Wire) -> MappingProxyType[str, JsonValue]:
+ bodies: Final = tuple(_JSON_OBJECT.validate_json(request.body) for request in wire.drain())
+ return MappingProxyType({str(body["model"]): body["prompt"] for body in bodies})
+
+
+@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS)
+def test_boot_time_custom_prompt_dict_renders_partial_templates(gateway: Gateway, tmp_path: Path) -> None:
+ templates: Final[dict[str, JsonValue]] = {
+ "response_schema_prompt": {"roles": {"user": {"pre_message": "", "post_message": ""}}},
+ "claude-2.1": {"roles": _ROLES},
+ "claude-instant-1.2": {"initial_prompt_value": "", "final_prompt_value": ""},
+ }
+ schema: Final[dict[str, JsonValue]] = {
+ "type": "object",
+ "properties": {"colour": {"type": "string"}},
+ "required": ["colour"],
+ }
+ with wire_server(_gemini_reply) as gemini, wire_server(_anthropic_text_reply) as anthropic_text:
+ base: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
+ deployments: Final = [
+ _boot_deployment("gemini-schema", "gemini/scripted-template-model", gemini),
+ _boot_deployment("anthropic-text-roles", "anthropic_text/claude-2.1", anthropic_text),
+ _boot_deployment("anthropic-text-values", "anthropic_text/claude-instant-1.2", anthropic_text),
+ ]
+ settings: Final = {**base.get("litellm_settings", {}), "custom_prompt_dict": templates}
+ path: Final = tmp_path / "boot-time-custom-prompt-dict.yaml"
+ path.write_text(yaml.safe_dump({**base, "model_list": deployments, "litellm_settings": settings}))
+ with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned:
+ _boot_chat(
+ owned.gateway,
+ "gemini-schema",
+ {"response_format": {"type": "json_schema", "json_schema": {"name": "sky", "schema": schema}}},
+ )
+ posts: Final = tuple(request for request in gemini.drain() if request.method == "POST")
+ assert len(posts) == 1, [request.target for request in posts]
+ texts: Final = _texts(_JSON_OBJECT.validate_json(posts[0].body)["contents"])
+ wrapped: Final = tuple(text for text in texts if text.startswith(""))
+ assert len(wrapped) == 1, texts
+ assert wrapped[0].endswith("") and "colour" in wrapped[0], wrapped[0]
+ _boot_chat(owned.gateway, "anthropic-text-roles", {})
+ _boot_chat(owned.gateway, "anthropic-text-values", {})
+ assert _anthropic_text_prompts(anthropic_text) == {
+ "claude-2.1": f"{_SYSTEM}{_question('anthropic-text-roles')}",
+ "claude-instant-1.2": f"{_SYSTEM}{_question('anthropic-text-values')}",
+ }
diff --git a/tests/integration/providers/test_custom_prompt_template_wire.py b/tests/integration/providers/test_custom_prompt_template_wire.py
new file mode 100644
index 00000000000..49a47a8c5e5
--- /dev/null
+++ b/tests/integration/providers/test_custom_prompt_template_wire.py
@@ -0,0 +1,701 @@
+import json
+import uuid
+from collections.abc import Callable, Iterator, Mapping
+from contextlib import contextmanager
+from dataclasses import dataclass
+from typing import Final
+
+import anthropic
+import openai
+import pytest
+from integration._support.client import Gateway, Scenario, eventually, object_value, string_value
+from integration._support.database import read_rows
+from integration._support.wire import Reply, Request, Wire, wire_server
+from pydantic import JsonValue, TypeAdapter
+
+_BACKEND: Final = "llama3-custom-template"
+_API_KEY: Final = "synthetic-ollama-key"
+_SYSTEM: Final = "Answer in one word."
+_QUESTION: Final = "What colour is the sky?"
+_ANSWER: Final = "blue"
+_INITIAL: Final = ""
+_FINAL: Final = ""
+_ROLES: Final[dict[str, JsonValue]] = {
+ "system": {"pre_message": "", "post_message": ""},
+ "user": {"pre_message": "", "post_message": ""},
+ "assistant": {"pre_message": "", "post_message": ""},
+}
+_ROLES_ONLY: Final[dict[str, JsonValue]] = {"roles": _ROLES}
+_VALUES_ONLY: Final[dict[str, JsonValue]] = {"initial_prompt_value": _INITIAL, "final_prompt_value": _FINAL}
+_FULL: Final[dict[str, JsonValue]] = {**_ROLES_ONLY, **_VALUES_ONLY}
+_ROLES_PROMPT: Final = f"{_SYSTEM}{_QUESTION}"
+_VALUES_PROMPT: Final = f"{_INITIAL}{_SYSTEM}{_QUESTION}{_FINAL}"
+_FULL_PROMPT: Final = f"{_INITIAL}{_ROLES_PROMPT}{_FINAL}"
+_TEMPLATES: Final[Mapping[str, tuple[dict[str, JsonValue], str]]] = {
+ "roles_only": (_ROLES_ONLY, _ROLES_PROMPT),
+ "values_only": (_VALUES_ONLY, _VALUES_PROMPT),
+}
+_TEMPLATE_FIELDS: Final = frozenset({"roles", "initial_prompt_value", "final_prompt_value", "bos_token", "eos_token"})
+_AWS: Final[dict[str, JsonValue]] = {
+ "aws_access_key_id": "AKIASYNTHETICKEY0001",
+ "aws_secret_access_key": "synthetic-secret-key-for-testing",
+ "aws_region_name": "us-east-1",
+}
+_NO_CACHE: Final[dict[str, JsonValue]] = {"no-cache": True}
+_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
+_JSON_LIST: Final = TypeAdapter(list[JsonValue])
+
+
+def _generate_reply(text: str) -> Reply:
+ return Reply(
+ body=json.dumps(
+ {
+ "model": _BACKEND,
+ "created_at": "2026-10-15T00:00:00Z",
+ "response": text,
+ "done": True,
+ "done_reason": "stop",
+ "prompt_eval_count": 30,
+ "eval_count": 12,
+ }
+ ).encode()
+ )
+
+
+def _streamed_reply(text: str) -> Reply:
+ frames: Final = (
+ {"model": _BACKEND, "created_at": "2026-10-15T00:00:00Z", "response": text, "done": False},
+ {
+ "model": _BACKEND,
+ "created_at": "2026-10-15T00:00:00Z",
+ "response": "",
+ "done": True,
+ "done_reason": "stop",
+ "prompt_eval_count": 30,
+ "eval_count": 12,
+ },
+ )
+ return Reply(
+ content_type="application/x-ndjson", chunks=tuple(json.dumps(frame).encode() + b"\n" for frame in frames)
+ )
+
+
+def _is_generate(request: Request) -> bool:
+ return (request.method, request.target) == ("POST", "/api/generate")
+
+
+def _echo_answer(request: Request) -> Reply:
+ stream: Final = _JSON_OBJECT.validate_json(request.body).get("stream") is True
+ return _streamed_reply(_ANSWER) if stream else _generate_reply(_ANSWER)
+
+
+@contextmanager
+def _ollama_server(respond: Callable[[Request], Reply] = _echo_answer) -> Iterator[Wire]:
+ with wire_server(lambda request: respond(request) if _is_generate(request) else Reply(body=b"{}")) as wire:
+ yield wire
+
+
+def _generate_bodies(wire: Wire) -> tuple[dict[str, JsonValue], ...]:
+ return tuple(_JSON_OBJECT.validate_json(request.body) for request in wire.drain() if _is_generate(request))
+
+
+def _assert_rendered(body: dict[str, JsonValue], prompt: str, *, stream: bool) -> None:
+ assert body["model"] == _BACKEND
+ assert body["prompt"] == prompt, body["prompt"]
+ assert body["stream"] is stream
+ options: Final = body.get("options", {})
+ assert isinstance(options, dict)
+ assert not _TEMPLATE_FIELDS & (set(body) | set(options)), sorted(body)
+
+
+def _assert_only_rendered(wire: Wire, prompt: str, *, stream: bool = False) -> None:
+ bodies: Final = _generate_bodies(wire)
+ assert len(bodies) == 1, bodies
+ _assert_rendered(bodies[0], prompt, stream=stream)
+
+
+def _ollama_model(scenario: Scenario, wire: Wire, template: Mapping[str, JsonValue]) -> str:
+ return scenario.model(model=f"ollama/{_BACKEND}", api_base=wire.url, api_key=_API_KEY, **template)
+
+
+def _chat_messages() -> list[dict[str, JsonValue]]:
+ return [{"role": "system", "content": _SYSTEM}, {"role": "user", "content": _QUESTION}]
+
+
+def _spend_row(identity: str) -> dict[str, JsonValue]:
+ rows: Final = eventually(
+ lambda: read_rows(
+ 'SELECT model_group, status, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" WHERE request_id=%s',
+ (identity,),
+ ),
+ lambda found: len(found) == 1,
+ seconds=70,
+ )
+ return rows[0]
+
+
+def _model_spend_rows(model: str, expected: int) -> list[dict[str, JsonValue]]:
+ rows: Final = eventually(
+ lambda: read_rows(
+ 'SELECT request_id, status, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" WHERE model_group=%s',
+ (model,),
+ ),
+ lambda found: len(found) >= expected,
+ seconds=70,
+ )
+ assert len({row["request_id"] for row in rows}) == len(rows) == expected, rows
+ return rows
+
+
+def _billed(model: str) -> dict[str, JsonValue]:
+ return {"model_group": model, "status": "success", "prompt_tokens": 30, "completion_tokens": 12}
+
+
+def _openai_client(gateway: Gateway) -> openai.OpenAI:
+ return openai.OpenAI(base_url=str(gateway.client.base_url) + "/v1", api_key=gateway.key, max_retries=0)
+
+
+def _async_openai_client(gateway: Gateway) -> openai.AsyncOpenAI:
+ return openai.AsyncOpenAI(base_url=str(gateway.client.base_url) + "/v1", api_key=gateway.key, max_retries=0)
+
+
+def _anthropic_client(gateway: Gateway) -> anthropic.Anthropic:
+ return anthropic.Anthropic(base_url=str(gateway.client.base_url), api_key=gateway.key, max_retries=0)
+
+
+def _async_anthropic_client(gateway: Gateway) -> anthropic.AsyncAnthropic:
+ return anthropic.AsyncAnthropic(base_url=str(gateway.client.base_url), api_key=gateway.key, max_retries=0)
+
+
+def _post(gateway: Gateway, path: str, body: Mapping[str, JsonValue]) -> tuple[int, str]:
+ response: Final = gateway.request("POST", path, {**body, "cache": _NO_CACHE})
+ return response.status_code, response.text
+
+
+def _chat(gateway: Gateway, model: str, **extra: JsonValue) -> tuple[int, str]:
+ return _post(gateway, "/v1/chat/completions", {"model": model, "messages": _chat_messages(), **extra})
+
+
+def _chat_content(text: str) -> str:
+ choices: Final = _JSON_LIST.validate_python(_JSON_OBJECT.validate_json(text)["choices"])
+ message: Final = object_value(object_value(choices[0])["message"])
+ return string_value(message["content"])
+
+
+def test_openai_sdk_chat_renders_a_roles_only_template(gateway: Gateway) -> None:
+ with _ollama_server() as wire, gateway.scenario() as scenario:
+ model: Final = _ollama_model(scenario, wire, _ROLES_ONLY)
+ completion: Final = _openai_client(gateway).chat.completions.create(
+ model=model,
+ messages=_chat_messages(), # pyright: ignore[reportArgumentType] # plain JSON messages
+ extra_body={"cache": _NO_CACHE},
+ )
+ assert completion.choices[0].message.content == _ANSWER
+ assert completion.choices[0].finish_reason == "stop"
+ assert completion.usage is not None
+ assert (completion.usage.prompt_tokens, completion.usage.completion_tokens) == (30, 12)
+ _assert_only_rendered(wire, _ROLES_PROMPT)
+ assert _spend_row(completion.id) == _billed(model)
+
+
+async def test_async_openai_sdk_chat_stream_renders_a_roles_only_template(gateway: Gateway) -> None:
+ with _ollama_server() as wire, gateway.scenario() as scenario:
+ model: Final = _ollama_model(scenario, wire, _ROLES_ONLY)
+ stream: Final = await _async_openai_client(gateway).chat.completions.create(
+ model=model,
+ messages=_chat_messages(), # pyright: ignore[reportArgumentType] # plain JSON messages
+ stream=True,
+ stream_options={"include_usage": True},
+ extra_body={"cache": _NO_CACHE},
+ )
+ chunks: Final = [chunk async for chunk in stream]
+ assert {chunk.id for chunk in chunks} == {chunks[0].id}
+ assert "".join(choice.delta.content or "" for chunk in chunks for choice in chunk.choices) == _ANSWER
+ usages: Final = [chunk.usage for chunk in chunks if chunk.usage is not None]
+ assert [(usage.prompt_tokens, usage.completion_tokens) for usage in usages] == [(30, 12)]
+ _assert_only_rendered(wire, _ROLES_PROMPT, stream=True)
+ assert _spend_row(chunks[0].id) == _billed(model)
+
+
+def test_openai_sdk_completions_render_a_roles_only_template(gateway: Gateway) -> None:
+ with _ollama_server() as wire, gateway.scenario() as scenario:
+ model: Final = _ollama_model(scenario, wire, _ROLES_ONLY)
+ completion: Final = _openai_client(gateway).completions.create(
+ model=model, prompt=_QUESTION, extra_body={"cache": _NO_CACHE}
+ )
+ assert completion.choices[0].text == _ANSWER
+ _assert_only_rendered(wire, f"{_QUESTION}")
+ assert _spend_row(completion.id) == _billed(model)
+
+
+def test_anthropic_sdk_messages_render_a_roles_only_template(gateway: Gateway) -> None:
+ with _ollama_server() as wire, gateway.scenario() as scenario:
+ model: Final = _ollama_model(scenario, wire, _ROLES_ONLY)
+ message: Final = _anthropic_client(gateway).messages.create(
+ model=model,
+ max_tokens=64,
+ system=_SYSTEM,
+ messages=[{"role": "user", "content": _QUESTION}],
+ extra_body={"cache": _NO_CACHE},
+ )
+ assert message.stop_reason == "end_turn"
+ assert [(block.type, getattr(block, "text", None)) for block in message.content] == [("text", _ANSWER)]
+ assert (message.usage.input_tokens, message.usage.output_tokens) == (30, 12)
+ _assert_only_rendered(wire, _ROLES_PROMPT)
+ assert _spend_row(message.id) == _billed(model)
+
+
+async def test_async_anthropic_sdk_messages_stream_render_a_roles_only_template(gateway: Gateway) -> None:
+ with _ollama_server() as wire, gateway.scenario() as scenario:
+ model: Final = _ollama_model(scenario, wire, _ROLES_ONLY)
+ async with _async_anthropic_client(gateway).messages.stream(
+ model=model,
+ max_tokens=64,
+ system=_SYSTEM,
+ messages=[{"role": "user", "content": _QUESTION}],
+ extra_body={"cache": _NO_CACHE},
+ ) as stream:
+ text: Final = "".join([piece async for piece in stream.text_stream])
+ final: Final = await stream.get_final_message()
+ assert text == _ANSWER
+ assert final.stop_reason == "end_turn"
+ _assert_only_rendered(wire, _ROLES_PROMPT, stream=True)
+ assert _spend_row(final.id) == _billed(model)
+
+
+def test_openai_sdk_responses_render_a_roles_only_template(gateway: Gateway) -> None:
+ with _ollama_server() as wire, gateway.scenario() as scenario:
+ model: Final = _ollama_model(scenario, wire, _ROLES_ONLY)
+ response: Final = _openai_client(gateway).responses.create(
+ model=model, instructions=_SYSTEM, input=_QUESTION, store=False, extra_body={"cache": _NO_CACHE}
+ )
+ assert response.output_text == _ANSWER
+ assert response.usage is not None and (response.usage.input_tokens, response.usage.output_tokens) == (30, 12)
+ _assert_only_rendered(wire, _ROLES_PROMPT)
+ rows: Final = _model_spend_rows(model, 1)
+ assert (rows[0]["status"], rows[0]["prompt_tokens"], rows[0]["completion_tokens"]) == ("success", 30, 12)
+
+
+async def test_async_openai_sdk_responses_stream_render_a_roles_only_template(gateway: Gateway) -> None:
+ with _ollama_server() as wire, gateway.scenario() as scenario:
+ model: Final = _ollama_model(scenario, wire, _ROLES_ONLY)
+ stream: Final = await _async_openai_client(gateway).responses.create(
+ model=model,
+ instructions=_SYSTEM,
+ input=_QUESTION,
+ store=False,
+ stream=True,
+ extra_body={"cache": _NO_CACHE},
+ )
+ events: Final = [event async for event in stream]
+ assert "".join(event.delta for event in events if event.type == "response.output_text.delta") == _ANSWER
+ completed: Final = [event for event in events if event.type == "response.completed"]
+ assert len(completed) == 1 and completed[0].response.output_text == _ANSWER
+ _assert_only_rendered(wire, _ROLES_PROMPT, stream=True)
+ assert _model_spend_rows(model, 1)[0]["status"] == "success"
+
+
+def test_prompt_values_only_template_wraps_the_bare_messages(gateway: Gateway) -> None:
+ with _ollama_server() as wire, gateway.scenario() as scenario:
+ model: Final = _ollama_model(scenario, wire, _VALUES_ONLY)
+ code, text = _chat(gateway, model)
+ assert code == 200, text
+ assert _chat_content(text) == _ANSWER
+ _assert_only_rendered(wire, _VALUES_PROMPT)
+
+
+def test_full_template_renders_every_part(gateway: Gateway) -> None:
+ with _ollama_server() as wire, gateway.scenario() as scenario:
+ model: Final = _ollama_model(scenario, wire, _FULL)
+ code, text = _chat(gateway, model)
+ assert code == 200, text
+ assert _chat_content(text) == _ANSWER
+ _assert_only_rendered(wire, _FULL_PROMPT)
+
+
+@dataclass(frozen=True, slots=True)
+class _Member:
+ deployment: Callable[[str], dict[str, JsonValue]]
+ reply: Callable[[str], Reply]
+ prompt_field: str
+ target: Callable[[str], bool]
+
+
+def _codestral_reply(text: str) -> Reply:
+ return Reply(
+ body=json.dumps(
+ {
+ "id": "cmpl-scripted",
+ "object": "chat.completion",
+ "model": "codestral-latest",
+ "created": 1,
+ "choices": [{"index": 0, "message": {"role": "assistant", "content": text}, "finish_reason": "stop"}],
+ "usage": {"prompt_tokens": 30, "completion_tokens": 12, "total_tokens": 42},
+ }
+ ).encode()
+ )
+
+
+def _predibase_reply(text: str) -> Reply:
+ return Reply(body=json.dumps({"generated_text": text, "details": {"finish_reason": "eos_token"}}).encode())
+
+
+def _sagemaker_reply(text: str) -> Reply:
+ return Reply(body=json.dumps([{"generated_text": text}]).encode())
+
+
+def _bedrock_meta_reply(text: str) -> Reply:
+ return Reply(
+ body=json.dumps(
+ {"generation": text, "prompt_token_count": 30, "generation_token_count": 12, "stop_reason": "stop"}
+ ).encode()
+ )
+
+
+_FAMILY: Final[Mapping[str, _Member]] = {
+ "codestral": _Member(
+ deployment=lambda url: {
+ "model": "text-completion-codestral/codestral-latest",
+ "api_base": f"{url}/v1/fim/completions",
+ "api_key": "scripted-key",
+ },
+ reply=_codestral_reply,
+ prompt_field="prompt",
+ target=lambda target: target == "/v1/fim/completions",
+ ),
+ "predibase": _Member(
+ deployment=lambda url: {
+ "model": "predibase/llama-3-8b-instruct",
+ "api_base": url,
+ "api_key": "scripted-key",
+ "predibase_tenant_id": "scripted-tenant",
+ },
+ reply=_predibase_reply,
+ prompt_field="inputs",
+ target=lambda target: target.startswith("/scripted-tenant/deployments/v2/llms/llama-3-8b-instruct/generate"),
+ ),
+ "sagemaker": _Member(
+ deployment=lambda url: {
+ "model": "sagemaker/scripted-endpoint",
+ "sagemaker_base_url": f"{url}/endpoints/scripted-endpoint/invocations",
+ **_AWS,
+ },
+ reply=_sagemaker_reply,
+ prompt_field="inputs",
+ target=lambda target: target == "/endpoints/scripted-endpoint/invocations",
+ ),
+ "bedrock_invoke": _Member(
+ deployment=lambda url: {"model": "bedrock/invoke/meta.llama3-8b-instruct-v1:0", "api_base": url, **_AWS},
+ reply=_bedrock_meta_reply,
+ prompt_field="prompt",
+ target=lambda target: target.endswith("/invoke"),
+ ),
+}
+
+
+@pytest.mark.parametrize("template", sorted(_TEMPLATES))
+@pytest.mark.parametrize("member", sorted(_FAMILY))
+def test_family_member_renders_a_partial_template(gateway: Gateway, member: str, template: str) -> None:
+ spec: Final = _FAMILY[member]
+ fields, prompt = _TEMPLATES[template]
+ with wire_server(lambda _: spec.reply(_ANSWER)) as wire, gateway.scenario() as scenario:
+ model: Final = scenario.model(**spec.deployment(wire.url), **fields)
+ code, text = _chat(gateway, model)
+ received: Final = tuple(request for request in wire.drain() if request.method == "POST")
+ assert code == 200, text
+ assert _chat_content(text) == _ANSWER
+ assert len(received) == 1, [request.target for request in received]
+ assert spec.target(received[0].target), received[0].target
+ body: Final = _JSON_OBJECT.validate_json(received[0].body)
+ assert body[spec.prompt_field] == prompt, body
+ assert not _TEMPLATE_FIELDS & set(body), sorted(body)
+
+
+@pytest.mark.parametrize("provider", ("vllm", "petals"))
+def test_in_process_family_member_fails_alone_and_the_rig_keeps_serving(gateway: Gateway, provider: str) -> None:
+ with _ollama_server() as wire, gateway.scenario() as scenario:
+ broken: Final = scenario.model(model=f"{provider}/scripted-model", **_ROLES_ONLY)
+ healthy: Final = _ollama_model(scenario, wire, _ROLES_ONLY)
+ failed, failure = _chat(gateway, broken)
+ assert failed >= 400, failure
+ code, text = _chat(gateway, healthy)
+ assert code == 200, text
+ _assert_only_rendered(wire, _ROLES_PROMPT)
+
+
+def _watsonx_reply(text: str) -> Reply:
+ result: Final[dict[str, JsonValue]] = {
+ "generated_text": text,
+ "generated_token_count": 12,
+ "input_token_count": 30,
+ "stop_reason": "eos_token",
+ }
+ return Reply(body=json.dumps({"model_id": "ibm/granite", "results": [result]}).encode())
+
+
+def test_watsonx_text_deployment_with_a_template_still_answers(gateway: Gateway) -> None:
+ with wire_server(lambda _: _watsonx_reply(_ANSWER)) as wire, gateway.scenario() as scenario:
+ model: Final = scenario.model(
+ model="watsonx_text/ibm/granite-chat-scripted",
+ api_base=wire.url,
+ token="scripted-token",
+ project_id="p",
+ **_ROLES_ONLY,
+ )
+ code, text = _chat(gateway, model)
+ assert code == 200, text
+ assert _chat_content(text) == _ANSWER
+ posts: Final = tuple(request for request in wire.drain() if request.method == "POST")
+ assert len(posts) == 1, [request.target for request in posts]
+ body: Final = _JSON_OBJECT.validate_json(posts[0].body)
+ assert _QUESTION in string_value(body["input"]), body
+ assert not _TEMPLATE_FIELDS & set(body), sorted(body)
+
+
+def _openai_chat_reply(text: str) -> Reply:
+ return Reply(
+ body=json.dumps(
+ {
+ "id": "chatcmpl-scripted",
+ "object": "chat.completion",
+ "created": 1,
+ "model": "scripted-compat",
+ "choices": [{"index": 0, "message": {"role": "assistant", "content": text}, "finish_reason": "stop"}],
+ "usage": {"prompt_tokens": 30, "completion_tokens": 12, "total_tokens": 42},
+ }
+ ).encode()
+ )
+
+
+def _converse_reply(text: str) -> Reply:
+ return Reply(
+ body=json.dumps(
+ {
+ "output": {"message": {"role": "assistant", "content": [{"text": text}]}},
+ "stopReason": "end_turn",
+ "usage": {"inputTokens": 30, "outputTokens": 12, "totalTokens": 42},
+ "metrics": {"latencyMs": 1},
+ }
+ ).encode()
+ )
+
+
+_FORWARDING: Final[Mapping[str, tuple[Callable[[str], dict[str, JsonValue]], Callable[[str], Reply]]]] = {
+ "openai_compatible": (
+ lambda url: {"model": "openai/scripted-compat", "api_base": f"{url}/v1", "api_key": "scripted-key"},
+ _openai_chat_reply,
+ ),
+ "bedrock_converse": (
+ lambda url: {"model": "bedrock/converse/meta.llama3-8b-instruct-v1:0", "api_base": url, **_AWS},
+ _converse_reply,
+ ),
+}
+
+
+def _keys_anywhere(value: JsonValue) -> frozenset[str]:
+ if isinstance(value, dict):
+ nested: Final = frozenset(key for child in value.values() for key in _keys_anywhere(child))
+ return frozenset(value) | nested
+ if isinstance(value, list):
+ return frozenset(key for child in value for key in _keys_anywhere(child))
+ return frozenset()
+
+
+@pytest.mark.parametrize("source", ("deployment", "request"))
+@pytest.mark.parametrize("member", sorted(_FORWARDING))
+def test_template_fields_never_reach_a_provider_that_does_not_render_them(
+ gateway: Gateway, member: str, source: str
+) -> None:
+ deployment, reply = _FORWARDING[member]
+ with wire_server(lambda _: reply(_ANSWER)) as wire, gateway.scenario() as scenario:
+ model: Final = scenario.model(**deployment(wire.url), **(_FULL if source == "deployment" else {}))
+ code, text = _chat(gateway, model, **(_FULL if source == "request" else {}))
+ assert code == 200, text
+ assert _chat_content(text) == _ANSWER
+ posts: Final = tuple(request for request in wire.drain() if request.method == "POST")
+ assert len(posts) == 1, [request.target for request in posts]
+ body: Final = _JSON_OBJECT.validate_json(posts[0].body)
+ assert not _TEMPLATE_FIELDS & _keys_anywhere(body), body
+ assert _INITIAL not in posts[0].body.decode() and _FINAL not in posts[0].body.decode(), body
+
+
+@pytest.mark.parametrize("value", (5, [_INITIAL]), ids=("int", "list"))
+@pytest.mark.parametrize("source", ("deployment", "request"))
+def test_non_string_initial_prompt_value_is_rejected_by_name(gateway: Gateway, source: str, value: JsonValue) -> None:
+ with _ollama_server() as wire, gateway.scenario() as scenario:
+ template: Final[dict[str, JsonValue]] = {**_FULL, "initial_prompt_value": value}
+ broken: Final = _ollama_model(scenario, wire, template if source == "deployment" else _FULL)
+ healthy: Final = _ollama_model(scenario, wire, _FULL)
+ failed, failure = _chat(gateway, broken, **({"initial_prompt_value": value} if source == "request" else {}))
+ assert failed == 500, failure
+ assert "initial_prompt_value: Input should be a valid string" in failure, failure
+ assert _generate_bodies(wire) == ()
+ code, text = _chat(gateway, healthy)
+ assert code == 200, text
+ _assert_only_rendered(wire, _FULL_PROMPT)
+
+
+@pytest.mark.parametrize("value", ("", None), ids=("empty", "null"))
+def test_blank_initial_prompt_value_in_the_request_drops_the_deployment_value(
+ gateway: Gateway, value: JsonValue
+) -> None:
+ with _ollama_server() as wire, gateway.scenario() as scenario:
+ model: Final = _ollama_model(scenario, wire, _FULL)
+ code, text = _chat(gateway, model, initial_prompt_value=value)
+ assert code == 200, text
+ _assert_only_rendered(wire, f"{_ROLES_PROMPT}{_FINAL}")
+
+
+def test_request_initial_prompt_value_overrides_the_deployment_value(gateway: Gateway) -> None:
+ with _ollama_server() as wire, gateway.scenario() as scenario:
+ model: Final = _ollama_model(scenario, wire, _FULL)
+ code, text = _chat(gateway, model, initial_prompt_value="")
+ assert code == 200, text
+ _assert_only_rendered(wire, f"{_ROLES_PROMPT}{_FINAL}")
+
+
+def test_five_kilobyte_initial_prompt_value_leads_the_prompt(gateway: Gateway) -> None:
+ large: Final = "x" * 5120
+ with _ollama_server() as wire, gateway.scenario() as scenario:
+ model: Final = _ollama_model(scenario, wire, _ROLES_ONLY)
+ code, text = _chat(gateway, model, initial_prompt_value=large)
+ assert code == 200, text
+ _assert_only_rendered(wire, f"{large}{_ROLES_PROMPT}")
+
+
+def test_identical_requests_each_render_and_bill_once(gateway: Gateway) -> None:
+ with _ollama_server() as wire, gateway.scenario() as scenario:
+ model: Final = _ollama_model(scenario, wire, _ROLES_ONLY)
+ first, first_text = _chat(gateway, model)
+ second, second_text = _chat(gateway, model)
+ assert (first, second) == (200, 200), (first_text, second_text)
+ bodies: Final = _generate_bodies(wire)
+ assert len(bodies) == 2, bodies
+ for body in bodies:
+ _assert_rendered(body, _ROLES_PROMPT, stream=False)
+ identities: Final = {string_value(_JSON_OBJECT.validate_json(text)["id"]) for text in (first_text, second_text)}
+ assert len(identities) == 2
+ for identity in identities:
+ assert _spend_row(identity) == _billed(model)
+
+
+def test_template_with_only_a_bos_token_names_the_fields_it_lacks(gateway: Gateway) -> None:
+ with _ollama_server() as wire, gateway.scenario() as scenario:
+ model: Final = _ollama_model(scenario, wire, {"bos_token": ""})
+ code, text = _chat(gateway, model)
+ assert code == 500, text
+ assert "it sets none of roles, initial_prompt_value, or final_prompt_value" in text, text
+ assert _generate_bodies(wire) == ()
+
+
+def test_roles_that_are_not_a_mapping_render_without_role_markers(gateway: Gateway) -> None:
+ with _ollama_server() as wire, gateway.scenario() as scenario:
+ model: Final = _ollama_model(scenario, wire, {**_VALUES_ONLY, "roles": json.dumps(_ROLES)})
+ code, text = _chat(gateway, model)
+ assert code == 200, text
+ _assert_only_rendered(wire, _VALUES_PROMPT)
+
+
+def test_provider_error_reaches_the_caller_after_one_attempt(gateway: Gateway) -> None:
+ missing: Final = Reply(status=404, body=json.dumps({"error": f"model '{_BACKEND}' not found"}).encode())
+ with _ollama_server(lambda _: missing) as wire, gateway.scenario() as scenario:
+ model: Final = _ollama_model(scenario, wire, _ROLES_ONLY)
+ code, text = _chat(gateway, model)
+ assert code == 404, text
+ assert f"model '{_BACKEND}' not found" in text, text
+ bodies: Final = _generate_bodies(wire)
+ assert len(bodies) == 1, bodies
+ _assert_rendered(bodies[0], _ROLES_PROMPT, stream=False)
+
+
+def _template_keys(fields: Mapping[str, JsonValue]) -> tuple[bool, bool]:
+ return "initial_prompt_value" in fields, "final_prompt_value" in fields
+
+
+def test_pass_through_body_treats_initial_and_final_prompt_values_alike(gateway: Gateway) -> None:
+ with wire_server(lambda _: _generate_reply(_ANSWER)) as wire, gateway.scenario() as scenario:
+ path: Final = f"/integration-template-passthrough-{uuid.uuid4().hex}"
+ created: Final = gateway.post(
+ "/config/pass_through_endpoint", {"path": path, "target": f"{wire.url}/api/generate"}
+ )
+ endpoint: Final = object_value(_JSON_LIST.validate_python(created["endpoints"])[0])
+ scenario.cleanups.callback(
+ lambda: gateway.request(
+ "DELETE", "/config/pass_through_endpoint", params={"endpoint_id": str(endpoint["id"])}
+ )
+ )
+ response: Final = eventually(
+ lambda: gateway.request("POST", path, {"model": _BACKEND, "prompt": _QUESTION, "stream": False, **_FULL}),
+ lambda answered: answered.status_code != 404,
+ seconds=30,
+ )
+ assert response.status_code == 200, response.text
+ posts: Final = tuple(request for request in wire.drain() if request.method == "POST")
+ forwarded: Final = _JSON_OBJECT.validate_json(posts[-1].body)
+ assert (forwarded["model"], forwarded["prompt"]) == (_BACKEND, _QUESTION), forwarded
+ initial, final = _template_keys(forwarded)
+ assert initial == final, sorted(forwarded)
+
+
+def _a2a_reply(request: Request) -> Reply:
+ result: Final[dict[str, JsonValue]] = {
+ "kind": "message",
+ "role": "agent",
+ "messageId": f"{uuid.uuid4().hex}-out",
+ "parts": [{"kind": "text", "text": _ANSWER}],
+ }
+ outbound: Final = _JSON_OBJECT.validate_json(request.body)
+ return Reply(body=json.dumps({"jsonrpc": "2.0", "id": outbound["id"], "result": result}).encode())
+
+
+def test_a2a_params_treat_initial_and_final_prompt_values_alike(gateway: Gateway) -> None:
+ marker: Final = uuid.uuid4().hex
+ agent_name: Final = f"template-agent-{marker}"
+
+ def peer(request: Request) -> Reply:
+ if request.method == "GET":
+ return Reply(body=json.dumps(_agent_card(wire.url, agent_name)).encode())
+ return _a2a_reply(request)
+
+ with wire_server(peer) as wire, gateway.scenario() as scenario:
+ created: Final = gateway.post(
+ "/v1/agents", {"agent_name": agent_name, "agent_card_params": _agent_card(wire.url, agent_name)}
+ )
+ agent_id: Final = string_value(created["agent_id"])
+ scenario.cleanups.callback(lambda: gateway.request("DELETE", f"/v1/agents/{agent_id}"))
+ message: Final[dict[str, JsonValue]] = {
+ "kind": "message",
+ "role": "user",
+ "messageId": f"{marker}-in",
+ "parts": [{"kind": "text", "text": _QUESTION}],
+ }
+ invoked: Final = gateway.request(
+ "POST",
+ f"/a2a/{agent_id}",
+ {
+ "jsonrpc": "2.0",
+ "id": marker,
+ "method": "message/send",
+ "params": {"message": message, "initial_prompt_value": _INITIAL, "final_prompt_value": _FINAL},
+ },
+ )
+ assert invoked.status_code == 200, invoked.text
+ assert object_value(invoked.json()["result"])["parts"] == [{"kind": "text", "text": _ANSWER}], invoked.text
+ posts: Final = tuple(request for request in wire.drain() if request.method == "POST")
+ assert len(posts) == 1, [request.target for request in posts]
+ params: Final = object_value(_JSON_OBJECT.validate_json(posts[0].body)["params"])
+ assert _template_keys(params) == (False, False), sorted(params)
+
+
+def _agent_card(url: str, name: str) -> dict[str, JsonValue]:
+ return {
+ "protocolVersion": "0.3",
+ "name": name,
+ "description": "Scripted A2A peer",
+ "version": "1.0.0",
+ "url": url + "/",
+ "capabilities": {"streaming": True},
+ "defaultInputModes": ["text"],
+ "defaultOutputModes": ["text"],
+ "skills": [],
+ }
diff --git a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py
index ad216548d08..c878f0e149d 100644
--- a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py
+++ b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py
@@ -3,7 +3,9 @@ import json
import logging
import os
import re
+import sys
import time
+import types
from collections.abc import Callable, Mapping, Sequence
from copy import deepcopy
from typing import Final, List, TypeAlias
@@ -69,6 +71,17 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
from litellm.llms.vertex_ai.gemini.transformation import gemini_convert_messages_with_history
from litellm.types.llms.bedrock import BedrockConverseReasoningContentBlock
from litellm.types.llms.openai import AllMessageValues
+from litellm.litellm_core_utils.prompt_templates.factory import response_schema_prompt
+from litellm.llms.anthropic.completion.transformation import AnthropicTextConfig
+from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation import AmazonInvokeConfig
+from litellm.llms.codestral.completion.handler import CodestralTextCompletion
+from litellm.llms.ollama.completion.transformation import OllamaConfig
+from litellm.llms.petals.completion import handler as petals_handler
+from litellm.llms.predibase.chat.transformation import PredibaseConfig
+from litellm.llms.sagemaker.completion.transformation import SagemakerConfig
+from litellm.llms.vllm.completion import handler as vllm_handler
+from litellm.llms.watsonx.common_utils import _convert_watsonx_messages_core
+from litellm.types.utils import ModelResponse, TextCompletionResponse
READ_FILE: Final[list[object]] = [{"name": "read_file", "description": "Read the contents of a file."}]
@@ -7871,3 +7884,334 @@ def test_anthropic_messages_pt_merges_a_long_assistant_run_in_linear_time():
growth: Final = _merged_assistant_run_cost(4000) / _merged_assistant_run_cost(1000)
assert growth < 8.0, f"cost grew {growth:.1f}x for 4x the assistant run"
+
+
+TEMPLATE_MODEL = "partial-template-model"
+TEMPLATE_ROLES = {"user": {"pre_message": "[INST] ", "post_message": " [/INST]"}}
+TEMPLATE_MESSAGES = [{"role": "user", "content": "hi"}]
+TEMPLATE_RENDERED_MESSAGES = "[INST] hi [/INST]"
+TEMPLATE_BARE_MESSAGES = "hi"
+
+TEMPLATE_RESPONSE_SCHEMA = {"type": "object"}
+TEMPLATE_RENDERED_SCHEMA = f"[INST] {TEMPLATE_RESPONSE_SCHEMA} [/INST]"
+TEMPLATE_BARE_SCHEMA = str(TEMPLATE_RESPONSE_SCHEMA)
+
+TEMPLATE_INITIAL = ""
+TEMPLATE_FINAL = ""
+
+PARTIAL_TEMPLATE = {"roles": TEMPLATE_ROLES}
+FULL_TEMPLATE = {"roles": TEMPLATE_ROLES, "initial_prompt_value": TEMPLATE_INITIAL, "final_prompt_value": TEMPLATE_FINAL}
+ROLELESS_TEMPLATE = {"initial_prompt_value": TEMPLATE_INITIAL, "final_prompt_value": TEMPLATE_FINAL}
+
+
+class _PromptCaptured(Exception):
+ def __init__(self, prompt):
+ super().__init__(prompt)
+ self.prompt = prompt
+
+
+class _CapturingLogging:
+ def pre_call(self, input, api_key, additional_args=None, **kwargs):
+ raise _PromptCaptured(input)
+
+ def post_call(self, *args, **kwargs):
+ pass
+
+
+def _captured_prompt(call):
+ with pytest.raises(_PromptCaptured) as excinfo:
+ call()
+ return excinfo.value.prompt
+
+
+def _stub_vllm_import(monkeypatch):
+ class _SamplingParams:
+ def __init__(self, **kwargs):
+ self.kwargs = kwargs
+
+ class _LLM:
+ def __init__(self, model):
+ self.model = model
+
+ def generate(self, prompts, sampling_params):
+ raise _PromptCaptured(prompts)
+
+ stub = types.ModuleType("vllm")
+ stub.LLM = _LLM
+ stub.SamplingParams = _SamplingParams
+
+ monkeypatch.setitem(sys.modules, "vllm", stub)
+ monkeypatch.setattr(vllm_handler, "llm", None)
+
+
+def _render_ollama(template, monkeypatch, messages=TEMPLATE_MESSAGES):
+ request = OllamaConfig().transform_request(
+ model=TEMPLATE_MODEL,
+ messages=messages,
+ optional_params={},
+ litellm_params={"custom_prompt_dict": {TEMPLATE_MODEL: template}},
+ headers={},
+ )
+ return request["prompt"]
+
+
+def _render_predibase(template, monkeypatch):
+ request = PredibaseConfig().transform_request(
+ model=TEMPLATE_MODEL,
+ messages=TEMPLATE_MESSAGES,
+ optional_params={},
+ litellm_params={"custom_prompt_dict": {TEMPLATE_MODEL: template}},
+ headers={},
+ )
+ return request["inputs"]
+
+
+def _render_anthropic_text(template, monkeypatch):
+ monkeypatch.setattr(litellm, "custom_prompt_dict", {TEMPLATE_MODEL: template})
+ return AnthropicTextConfig()._get_anthropic_text_prompt_from_messages(
+ messages=TEMPLATE_MESSAGES, model=TEMPLATE_MODEL
+ )
+
+
+def _render_response_schema_prompt(template, monkeypatch):
+ monkeypatch.setattr(litellm, "custom_prompt_dict", {f"{TEMPLATE_MODEL}/response_schema_prompt": template})
+ return response_schema_prompt(
+ model=TEMPLATE_MODEL, response_schema=TEMPLATE_RESPONSE_SCHEMA, llm_provider="vertex_ai"
+ )
+
+
+def _render_petals(template, monkeypatch):
+ monkeypatch.setattr(litellm, "custom_prompt_dict", {TEMPLATE_MODEL: template})
+ return _captured_prompt(
+ lambda: petals_handler.completion(
+ model=TEMPLATE_MODEL,
+ messages=TEMPLATE_MESSAGES,
+ api_base="http://petals.invalid",
+ model_response=ModelResponse(),
+ print_verbose=lambda *args, **kwargs: None,
+ encoding=None,
+ logging_obj=_CapturingLogging(),
+ optional_params={},
+ )
+ )
+
+
+def _render_codestral(template, monkeypatch):
+ return _captured_prompt(
+ lambda: CodestralTextCompletion().completion(
+ model=TEMPLATE_MODEL,
+ messages=TEMPLATE_MESSAGES,
+ api_base="http://codestral.invalid",
+ custom_prompt_dict={TEMPLATE_MODEL: template},
+ model_response=TextCompletionResponse(),
+ print_verbose=lambda *args, **kwargs: None,
+ encoding=None,
+ api_key="sk-not-used",
+ logging_obj=_CapturingLogging(),
+ optional_params={},
+ timeout=1.0,
+ )
+ )
+
+
+def _render_vllm_completion(template, monkeypatch):
+ _stub_vllm_import(monkeypatch)
+ return _captured_prompt(
+ lambda: vllm_handler.completion(
+ model=TEMPLATE_MODEL,
+ messages=TEMPLATE_MESSAGES,
+ model_response=ModelResponse(),
+ print_verbose=lambda *args, **kwargs: None,
+ encoding=None,
+ logging_obj=_CapturingLogging(),
+ optional_params={},
+ custom_prompt_dict={TEMPLATE_MODEL: template},
+ )
+ )
+
+
+def _render_vllm_batch(template, monkeypatch):
+ _stub_vllm_import(monkeypatch)
+ prompts = _captured_prompt(
+ lambda: vllm_handler.batch_completions(
+ model=TEMPLATE_MODEL,
+ messages=[TEMPLATE_MESSAGES],
+ optional_params={},
+ custom_prompt_dict={TEMPLATE_MODEL: template},
+ )
+ )
+ assert len(prompts) == 1
+ return prompts[0]
+
+
+def _render_bedrock_invoke(template, monkeypatch):
+ prompt, _ = AmazonInvokeConfig().convert_messages_to_prompt(
+ model=TEMPLATE_MODEL, messages=TEMPLATE_MESSAGES, provider="meta", custom_prompt_dict={TEMPLATE_MODEL: template}
+ )
+ return prompt
+
+
+def _render_sagemaker(template, monkeypatch):
+ return SagemakerConfig()._transform_prompt(
+ model=TEMPLATE_MODEL,
+ messages=TEMPLATE_MESSAGES,
+ custom_prompt_dict={TEMPLATE_MODEL: template},
+ hf_model_name=None,
+ )
+
+
+def _render_watsonx(template, monkeypatch):
+ return _convert_watsonx_messages_core(
+ model=TEMPLATE_MODEL,
+ messages=TEMPLATE_MESSAGES,
+ provider="ibm",
+ custom_prompt_dict={TEMPLATE_MODEL: template},
+ apply_template_fn=None,
+ )
+
+
+TEMPLATE_CALL_SITES = [
+ pytest.param(_render_ollama, TEMPLATE_RENDERED_MESSAGES, TEMPLATE_BARE_MESSAGES, id="ollama"),
+ pytest.param(_render_petals, TEMPLATE_RENDERED_MESSAGES, TEMPLATE_BARE_MESSAGES, id="petals"),
+ pytest.param(_render_vllm_completion, TEMPLATE_RENDERED_MESSAGES, TEMPLATE_BARE_MESSAGES, id="vllm-completion"),
+ pytest.param(_render_vllm_batch, TEMPLATE_RENDERED_MESSAGES, TEMPLATE_BARE_MESSAGES, id="vllm-batch-completions"),
+ pytest.param(_render_predibase, TEMPLATE_RENDERED_MESSAGES, TEMPLATE_BARE_MESSAGES, id="predibase"),
+ pytest.param(_render_codestral, TEMPLATE_RENDERED_MESSAGES, TEMPLATE_BARE_MESSAGES, id="codestral"),
+ pytest.param(_render_anthropic_text, TEMPLATE_RENDERED_MESSAGES, TEMPLATE_BARE_MESSAGES, id="anthropic-text"),
+ pytest.param(
+ _render_response_schema_prompt, TEMPLATE_RENDERED_SCHEMA, TEMPLATE_BARE_SCHEMA, id="response-schema-prompt"
+ ),
+ pytest.param(_render_bedrock_invoke, TEMPLATE_RENDERED_MESSAGES, TEMPLATE_BARE_MESSAGES, id="bedrock-invoke"),
+ pytest.param(_render_sagemaker, TEMPLATE_RENDERED_MESSAGES, TEMPLATE_BARE_MESSAGES, id="sagemaker"),
+ pytest.param(_render_watsonx, TEMPLATE_RENDERED_MESSAGES, TEMPLATE_BARE_MESSAGES, id="watsonx"),
+]
+
+
+@pytest.mark.parametrize(("render", "rendered_body", "bare_body"), TEMPLATE_CALL_SITES)
+def test_template_without_the_optional_prompt_values_renders(render, rendered_body, bare_body, monkeypatch):
+ assert render(PARTIAL_TEMPLATE, monkeypatch) == rendered_body
+
+
+@pytest.mark.parametrize(("render", "rendered_body", "bare_body"), TEMPLATE_CALL_SITES)
+def test_template_with_the_optional_prompt_values_still_applies_them(render, rendered_body, bare_body, monkeypatch):
+ assert render(FULL_TEMPLATE, monkeypatch) == TEMPLATE_INITIAL + rendered_body + TEMPLATE_FINAL
+
+
+@pytest.mark.parametrize(("render", "rendered_body", "bare_body"), TEMPLATE_CALL_SITES)
+def test_template_without_roles_renders_messages_unwrapped(render, rendered_body, bare_body, monkeypatch):
+ assert render(ROLELESS_TEMPLATE, monkeypatch) == TEMPLATE_INITIAL + bare_body + TEMPLATE_FINAL
+
+
+def test_template_wraps_only_the_roles_it_defines_across_a_conversation(monkeypatch):
+ messages = [
+ {"role": "system", "content": "Be terse."},
+ {"role": "user", "content": "hi"},
+ {"role": "assistant", "content": "hello"},
+ {"role": "user", "content": "sky colour?"},
+ ]
+
+ rendered = _render_ollama(FULL_TEMPLATE, monkeypatch, messages=messages)
+
+ assert rendered == "Be terse.[INST] hi [/INST]hello[INST] sky colour? [/INST]"
+
+
+@pytest.mark.parametrize(
+ "template",
+ [
+ pytest.param({"roles": {**TEMPLATE_ROLES, "tool": None}}, id="null-role-entry"),
+ pytest.param({"roles": {"user": {**TEMPLATE_ROLES["user"], "priority": 1}}}, id="extra-role-key"),
+ pytest.param({"roles": {**TEMPLATE_ROLES, "tool": {"pre_message": 5}}}, id="unused-role-wrong-type"),
+ pytest.param({"roles": {**TEMPLATE_ROLES, 1: {"pre_message": "x"}}}, id="unused-non-string-role"),
+ pytest.param({"roles": TEMPLATE_ROLES, "initial_prompt_value": None, "bos_token": 7}, id="null-and-unread-keys"),
+ ],
+)
+def test_template_shapes_main_rendered_keep_rendering(template, monkeypatch):
+ assert _render_ollama(template, monkeypatch) == TEMPLATE_RENDERED_MESSAGES
+
+
+@pytest.mark.parametrize(
+ "template",
+ [
+ pytest.param({"roles": {"user": "[INST]"}}, id="string-role-entry"),
+ pytest.param({"roles": {"user": ["[INST]"]}}, id="list-role-entry"),
+ pytest.param({"roles": ["user"]}, id="list-roles"),
+ pytest.param({"roles": '{"user": {"pre_message": "[INST]"}}'}, id="json-string-roles"),
+ ],
+)
+def test_roles_that_are_not_a_mapping_render_unwrapped_with_a_warning(template, monkeypatch, caplog):
+ with caplog.at_level(logging.WARNING, logger="LiteLLM"):
+ rendered = _render_ollama({**template, "final_prompt_value": TEMPLATE_FINAL}, monkeypatch)
+
+ assert rendered == TEMPLATE_BARE_MESSAGES + TEMPLATE_FINAL
+ assert "not a mapping" in caplog.text
+
+
+def test_null_role_entry_renders_that_role_unwrapped_without_a_warning(monkeypatch, caplog):
+ with caplog.at_level(logging.WARNING, logger="LiteLLM"):
+ rendered = _render_ollama({"roles": {"user": None}}, monkeypatch)
+
+ assert rendered == TEMPLATE_BARE_MESSAGES
+ assert "not a mapping" not in caplog.text
+
+
+@pytest.mark.parametrize(
+ ("template", "field"),
+ [
+ pytest.param({"initial_prompt_value": 5}, "initial_prompt_value", id="initial-prompt-value"),
+ pytest.param({"final_prompt_value": ["x"]}, "final_prompt_value", id="final-prompt-value"),
+ pytest.param({"roles": {"user": {"pre_message": 5}}}, "roles.user.pre_message", id="pre-message"),
+ ],
+)
+def test_wrong_type_template_field_fails_as_a_server_error_naming_it(template, field):
+ with pytest.raises(litellm.APIConnectionError) as excinfo:
+ litellm.completion(
+ model=f"ollama/{TEMPLATE_MODEL}",
+ messages=TEMPLATE_MESSAGES,
+ api_base="http://ollama.invalid",
+ **template,
+ )
+
+ assert excinfo.value.status_code == 500
+ assert f"{field}: Input should be a valid string" in excinfo.value.message
+
+
+@pytest.mark.parametrize(
+ "template",
+ [
+ pytest.param({"bos_token": ""}, id="bos-token-only"),
+ pytest.param({"eos_token": ""}, id="eos-token-only"),
+ pytest.param({"roles": '{"user": {"pre_message": "[INST]"}}'}, id="json-string-roles-only"),
+ ],
+)
+def test_template_that_sets_nothing_the_prompt_uses_fails_as_a_server_error(template):
+ with pytest.raises(litellm.APIConnectionError) as excinfo:
+ litellm.completion(
+ model=f"ollama/{TEMPLATE_MODEL}",
+ messages=TEMPLATE_MESSAGES,
+ api_base="http://ollama.invalid",
+ **template,
+ )
+
+ assert excinfo.value.status_code == 500
+ assert "is invalid: it sets none of roles, initial_prompt_value, or final_prompt_value" in excinfo.value.message
+
+
+def test_wrong_type_template_field_in_requests_leaves_the_deployments_in_rotation():
+ router = litellm.Router(
+ model_list=[
+ {
+ "model_name": "template-group",
+ "litellm_params": {"model": f"ollama/{TEMPLATE_MODEL}", "api_base": "http://ollama.invalid"},
+ }
+ for _ in range(2)
+ ],
+ num_retries=0,
+ allowed_fails=0,
+ )
+
+ for _ in range(2):
+ with pytest.raises(litellm.APIConnectionError):
+ router.completion(model="template-group", messages=TEMPLATE_MESSAGES, initial_prompt_value=5)
+ response = router.completion(model="template-group", messages=TEMPLATE_MESSAGES, mock_response="Blue")
+
+ assert response.choices[0].message.content == "Blue"
diff --git a/tests/unit/llms/test_predibase_transformation.py b/tests/unit/llms/test_predibase_transformation.py
index d6425488be3..20d547c3b63 100644
--- a/tests/unit/llms/test_predibase_transformation.py
+++ b/tests/unit/llms/test_predibase_transformation.py
@@ -39,14 +39,9 @@ def test_predibase_transform_request_non_stream():
assert "stream" not in request_data["parameters"]
-def test_predibase_transform_request_custom_prompt(monkeypatch):
+def test_predibase_transform_request_custom_prompt():
config = PredibaseConfig()
- monkeypatch.setattr(
- "litellm.llms.predibase.chat.transformation.custom_prompt",
- lambda **kwargs: "custom-prompt",
- )
-
request_data = config.transform_request(
model="predibase-model",
messages=[{"role": "user", "content": "hello"}],
@@ -54,16 +49,16 @@ def test_predibase_transform_request_custom_prompt(monkeypatch):
litellm_params={
"custom_prompt_dict": {
"predibase-model": {
- "roles": {},
- "initial_prompt_value": "",
- "final_prompt_value": "",
+ "roles": {"user": {"pre_message": "", "post_message": ""}},
+ "initial_prompt_value": "",
+ "final_prompt_value": "",
}
}
},
headers={},
)
- assert request_data["inputs"] == "custom-prompt"
+ assert request_data["inputs"] == "hello"
def test_predibase_get_complete_url_stream_and_non_stream():
diff --git a/tests/unit/types/test_litellm_params.py b/tests/unit/types/test_litellm_params.py
index d27dc733c78..122d811e491 100644
--- a/tests/unit/types/test_litellm_params.py
+++ b/tests/unit/types/test_litellm_params.py
@@ -199,6 +199,7 @@ OPTION_NAMES: Final = (
"prompt_label",
"litellm_system_prompt",
"custom_prompt_dict",
+ "initial_prompt_value",
"roles",
"final_prompt_value",
"bos_token",