mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-17 04:34:25 +00:00
fix(prompt-templates): treat initial_prompt_value and final_prompt_value as optional (#39762)
* feat(dd_span_tagger): emit litellm_user_email span tag for JWT-authenticated requests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(dd_span_tagger): use dotted litellm.user_email tag for consistency Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * feat(model_hub): surface model_info.description in model group info and Model Hub UI Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(router): aggregate model group description without in-place mutation Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(prompt-templates): treat initial_prompt_value and final_prompt_value as optional completion() only stores a custom prompt template key when its value is truthy, and custom_prompt() already defaults both prompt values to "". Seven call sites read them back with direct indexing, so a template that sets only roles raises KeyError instead of rendering, on ollama, petals, vllm, predibase, codestral, anthropic text completion, and response_schema_prompt. Five other providers already read defensively, so the same config works on sagemaker and 500s on ollama. Read both keys with .get(key, ""), matching those five providers and the defaults custom_prompt() already declares. Templates that set all three keys render exactly as before. * test(prompt-templates): drop the explanatory docstrings from the optional-keys regression test * fix(health): skip background health check DB writes when the latest-row read fails Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(prompt-templates): read every custom prompt template field through one typed helper Templates that omit roles, initial_prompt_value, or final_prompt_value now render with empty defaults across every provider that reads them, instead of failing with KeyError or TypeError. A field with the wrong type fails with a validation error naming it. * fix(prompt-templates): keep templates that rendered on main working and fail bad fields with a 400 Only the roles a request uses are checked, null fields mean unset, and extra role keys are ignored, so every template main rendered renders the same. A wrong-type field now raises a BadRequestError naming it. initial_prompt_value is declared litellm-owned so it stays out of provider request bodies, like final_prompt_value. The template tests move into the mapped factory test file, and the predibase test asserts the real rendered prompt * fix(prompt-templates): keep bad-template failover and render non-mapping role entries like main A wrong-type template field now raises a 500 InternalServerError naming the field, so the router still retries and cools down the bad deployment. A roles value or role entry that is not a mapping renders unwrapped as it did on main, and role keys are only checked for roles a request uses. * fix(prompt-templates): warn when a custom prompt template roles value is not a mapping * fix(prompt-templates): keep deployments in rotation on a bad template and reject a template with nothing to render * test(integration): audit custom prompt templates across endpoints, providers, and chaos --------- Co-authored-by: milan <milan@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: yassin <yassin@berri.ai> Co-authored-by: ryan-crabbe-berri <ryan@berri.ai> Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
Marcus
GitHub
parent
43be1c2ae7
commit
642a2ad934
17 changed files with 1720 additions and 80 deletions
|
|
@ -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 = "",
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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", ""),
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
542
tests/integration/providers/test_custom_prompt_template_chaos.py
Normal file
542
tests/integration/providers/test_custom_prompt_template_chaos.py
Normal file
|
|
@ -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": "<sys>", "post_message": "</sys>"},
|
||||
"user": {"pre_message": "<user>", "post_message": "</user>"},
|
||||
}
|
||||
_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"<sys>{_SYSTEM}</sys><user>{_question(marker)}</user>"
|
||||
|
||||
|
||||
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, "<first>")
|
||||
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, "<first>"),
|
||||
lambda found: "x-litellm-cache-key" in found.headers,
|
||||
seconds=20,
|
||||
)
|
||||
changed: Final = _cached_chat(candidate, marker, field, "<second>")
|
||||
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": "<i>",
|
||||
"final_prompt_value": "<f>",
|
||||
}
|
||||
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": "<schema>", "post_message": "</schema>"}}},
|
||||
"claude-2.1": {"roles": _ROLES},
|
||||
"claude-instant-1.2": {"initial_prompt_value": "<i>", "final_prompt_value": "<f>"},
|
||||
}
|
||||
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("<schema>"))
|
||||
assert len(wrapped) == 1, texts
|
||||
assert wrapped[0].endswith("</schema>") 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"<sys>{_SYSTEM}</sys><user>{_question('anthropic-text-roles')}</user>",
|
||||
"claude-instant-1.2": f"<i>{_SYSTEM}{_question('anthropic-text-values')}<f>",
|
||||
}
|
||||
701
tests/integration/providers/test_custom_prompt_template_wire.py
Normal file
701
tests/integration/providers/test_custom_prompt_template_wire.py
Normal file
|
|
@ -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 = "<bos>"
|
||||
_FINAL: Final = "<assistant>"
|
||||
_ROLES: Final[dict[str, JsonValue]] = {
|
||||
"system": {"pre_message": "<sys>", "post_message": "</sys>"},
|
||||
"user": {"pre_message": "<user>", "post_message": "</user>"},
|
||||
"assistant": {"pre_message": "<asst>", "post_message": "</asst>"},
|
||||
}
|
||||
_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"<sys>{_SYSTEM}</sys><user>{_QUESTION}</user>"
|
||||
_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"<user>{_QUESTION}</user>")
|
||||
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="<override>")
|
||||
assert code == 200, text
|
||||
_assert_only_rendered(wire, f"<override>{_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": "<s>"})
|
||||
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": [],
|
||||
}
|
||||
|
|
@ -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 = "<start>"
|
||||
TEMPLATE_FINAL = "<end>"
|
||||
|
||||
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 == "<start>Be terse.[INST] hi [/INST]hello[INST] sky colour? [/INST]<end>"
|
||||
|
||||
|
||||
@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": "<s>"}, id="bos-token-only"),
|
||||
pytest.param({"eos_token": "</s>"}, 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"
|
||||
|
|
|
|||
|
|
@ -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": "<user>", "post_message": "</user>"}},
|
||||
"initial_prompt_value": "<s>",
|
||||
"final_prompt_value": "</s>",
|
||||
}
|
||||
}
|
||||
},
|
||||
headers={},
|
||||
)
|
||||
|
||||
assert request_data["inputs"] == "custom-prompt"
|
||||
assert request_data["inputs"] == "<s><user>hello</user></s>"
|
||||
|
||||
|
||||
def test_predibase_get_complete_url_stream_and_non_stream():
|
||||
|
|
|
|||
|
|
@ -199,6 +199,7 @@ OPTION_NAMES: Final = (
|
|||
"prompt_label",
|
||||
"litellm_system_prompt",
|
||||
"custom_prompt_dict",
|
||||
"initial_prompt_value",
|
||||
"roles",
|
||||
"final_prompt_value",
|
||||
"bos_token",
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue