mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-11 08:38:24 +00:00
Merge remote-tracking branch 'origin/main' into litellm_vertex_identity_labels
This commit is contained in:
mateo-berri
commit
186f151711
34 changed files with 1071 additions and 649 deletions
1
.github/CODEOWNERS
vendored
1
.github/CODEOWNERS
vendored
|
|
@ -2,3 +2,4 @@
|
|||
/litellm/model_prices_and_context_window_backup.json @mateo-berri @ryan-crabbe-berri @kerry-berri @joshua-berri
|
||||
/helm/ @yassin-berriai
|
||||
/terraform/ @yassin-berriai
|
||||
/.github/workflows/zizmor.yml @yuneng-berri
|
||||
|
|
|
|||
5
.github/workflows/cost-map-guard.yml
vendored
5
.github/workflows/cost-map-guard.yml
vendored
|
|
@ -1,7 +1,7 @@
|
|||
name: Cost map guard
|
||||
|
||||
on: # zizmor: ignore[dangerous-triggers] runs the base branch's code only; the PR's cost map files are read as data and never executed
|
||||
pull_request_target:
|
||||
on:
|
||||
pull_request:
|
||||
branches:
|
||||
- main
|
||||
- "litellm_**"
|
||||
|
|
@ -19,6 +19,7 @@ jobs:
|
|||
steps:
|
||||
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
|
||||
with:
|
||||
ref: ${{ github.event.pull_request.base.sha }}
|
||||
persist-credentials: false
|
||||
- name: Fetch the pull request head and its merge base
|
||||
id: revisions
|
||||
|
|
|
|||
6
.github/workflows/spend-logs-guard.yml
vendored
6
.github/workflows/spend-logs-guard.yml
vendored
|
|
@ -1,7 +1,7 @@
|
|||
name: Spend logs guard
|
||||
|
||||
on: # zizmor: ignore[dangerous-triggers] runs the base branch's code only; the PR's diff is read as data and never executed
|
||||
pull_request_target:
|
||||
on:
|
||||
pull_request:
|
||||
branches:
|
||||
- main
|
||||
|
||||
|
|
@ -19,7 +19,7 @@ jobs:
|
|||
steps:
|
||||
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
|
||||
with:
|
||||
ref: ${{ github.event.pull_request.base.ref }}
|
||||
ref: ${{ github.event.pull_request.base.sha }}
|
||||
persist-credentials: false
|
||||
- name: Fetch the pull request head, its merge base and its reviews
|
||||
id: revisions
|
||||
|
|
|
|||
17
.github/workflows/zizmor.yml
vendored
17
.github/workflows/zizmor.yml
vendored
|
|
@ -27,10 +27,15 @@ jobs:
|
|||
with:
|
||||
persist-credentials: false
|
||||
|
||||
- name: Install zizmor
|
||||
run: |
|
||||
curl -fsSL --retry 3 -o "$RUNNER_TEMP/zizmor.tar.gz" \
|
||||
https://github.com/zizmorcore/zizmor/releases/download/v1.31.1/zizmor-x86_64-unknown-linux-gnu.tar.gz
|
||||
echo "d185f726bdaf92df78273217dfdd24ebcd6b0ef48cf3ce776ff5762b6e138c99 $RUNNER_TEMP/zizmor.tar.gz" | sha256sum -c -
|
||||
tar xzf "$RUNNER_TEMP/zizmor.tar.gz" -C "$RUNNER_TEMP" zizmor
|
||||
|
||||
- name: Run zizmor
|
||||
uses: zizmorcore/zizmor-action@5f14fd08f7cf1cb1609c1e344975f152c7ee938d # v0.5.6
|
||||
with:
|
||||
version: "1.24.1"
|
||||
min-severity: medium
|
||||
advanced-security: false
|
||||
annotations: true
|
||||
env:
|
||||
GH_TOKEN: ${{ github.token }}
|
||||
run: |
|
||||
"$RUNNER_TEMP/zizmor" --no-ignores --no-config --min-severity=medium --format=github .
|
||||
|
|
|
|||
|
|
@ -1337,7 +1337,6 @@ all_embedding_models = (
|
|||
)
|
||||
|
||||
####### IMAGE GENERATION MODELS ###################
|
||||
openai_image_generation_models = ["dall-e-2", "dall-e-3"]
|
||||
|
||||
####### VIDEO GENERATION MODELS ###################
|
||||
openai_video_generation_models = ["sora-2"]
|
||||
|
|
|
|||
|
|
@ -764,7 +764,6 @@ ANTHROPIC_WEB_SEARCH_TOOL_MAX_USES: Final = {
|
|||
# Used for web search interception across providers
|
||||
LITELLM_WEB_SEARCH_TOOL_NAME: Final = "litellm_web_search"
|
||||
|
||||
DEFAULT_IMAGE_ENDPOINT_MODEL: Final = "dall-e-2"
|
||||
DEFAULT_VIDEO_ENDPOINT_MODEL: Final = "sora-2"
|
||||
|
||||
DEFAULT_GOOGLE_VIDEO_DURATION_SECONDS: Final = get_env_int("DEFAULT_GOOGLE_VIDEO_DURATION_SECONDS", 8)
|
||||
|
|
|
|||
|
|
@ -1592,14 +1592,6 @@ def completion_cost(
|
|||
for part in ws_tier_parts
|
||||
)
|
||||
|
||||
if (
|
||||
(call_type == "aimage_generation" or call_type == "image_generation")
|
||||
and model is not None
|
||||
and isinstance(model, str)
|
||||
and len(model) == 0
|
||||
and custom_llm_provider == "azure"
|
||||
):
|
||||
model = "dall-e-2" # for dall-e-2, azure expects an empty model name
|
||||
# Handle Inputs to completion_cost
|
||||
prompt_tokens = 0
|
||||
prompt_characters: int | None = None
|
||||
|
|
@ -1674,7 +1666,8 @@ def completion_cost(
|
|||
raw_candidates if custom_pricing is True else _priced_candidates_before_unpriced_rules(raw_candidates)
|
||||
)
|
||||
|
||||
for idx, (model, candidate_pricing_provider) in enumerate(potential_model_names):
|
||||
for idx, (candidate_model, candidate_pricing_provider) in enumerate(potential_model_names):
|
||||
model = candidate_model
|
||||
try:
|
||||
if verbose_logger.isEnabledFor(logging.DEBUG):
|
||||
verbose_logger.debug("selected model name for cost calculation: %s", model)
|
||||
|
|
|
|||
|
|
@ -16,7 +16,6 @@ import litellm
|
|||
|
||||
# client is imported from litellm as it's a decorator
|
||||
from litellm import client
|
||||
from litellm.constants import DEFAULT_IMAGE_ENDPOINT_MODEL
|
||||
from litellm.constants import request_timeout as DEFAULT_REQUEST_TIMEOUT
|
||||
from litellm.exceptions import LiteLLMUnknownProvider
|
||||
from litellm.litellm_core_utils.get_llm_provider_logic import (
|
||||
|
|
@ -299,17 +298,14 @@ def image_generation(
|
|||
if custom_llm_provider == "chatgpt" or (model is not None and model.startswith("chatgpt/"))
|
||||
else None
|
||||
)
|
||||
dynamic_api_key: str | None = None
|
||||
if model is not None or custom_llm_provider is not None:
|
||||
model, custom_llm_provider, dynamic_api_key, api_base = get_llm_provider(
|
||||
model=model,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
api_base=api_base,
|
||||
litellm_params=chatgpt_params,
|
||||
)
|
||||
else:
|
||||
model = "dall-e-2"
|
||||
custom_llm_provider = "openai" # default to dall-e-2 on openai
|
||||
if model is None and custom_llm_provider is None:
|
||||
raise litellm.BadRequestError(message="model is required", model="", llm_provider="")
|
||||
model, custom_llm_provider, dynamic_api_key, api_base = get_llm_provider(
|
||||
model=model,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
api_base=api_base,
|
||||
litellm_params=chatgpt_params,
|
||||
)
|
||||
model_response.hidden_params["model"] = model
|
||||
openai_params: Final = [
|
||||
"user",
|
||||
|
|
@ -712,7 +708,7 @@ async def aimage_variation(*args, **kwargs) -> ImageResponse:
|
|||
@client
|
||||
def image_variation(
|
||||
image: FileTypes,
|
||||
model: str = "dall-e-2", # set to dall-e-2 by default - like OpenAI.
|
||||
model: str | None = None,
|
||||
n: int = 1,
|
||||
response_format: Literal["url", "b64_json"] = "url",
|
||||
size: str | None = None,
|
||||
|
|
@ -728,7 +724,7 @@ def image_variation(
|
|||
litellm_params: Final = get_litellm_params(**kwargs)
|
||||
# get the custom llm provider
|
||||
model, custom_llm_provider, dynamic_api_key, api_base = get_llm_provider(
|
||||
model=model,
|
||||
model=model or "",
|
||||
custom_llm_provider=litellm_params.get("custom_llm_provider", None),
|
||||
api_base=litellm_params.get("api_base", None),
|
||||
api_key=litellm_params.get("api_key", None),
|
||||
|
|
@ -747,7 +743,7 @@ def image_variation(
|
|||
response: ImageResponse | None = None
|
||||
|
||||
provider_config: Final = ProviderConfigManager.get_provider_model_info(
|
||||
model=model or "", # openai defaults to dall-e-2
|
||||
model=model,
|
||||
provider=llm_provider,
|
||||
)
|
||||
|
||||
|
|
@ -772,7 +768,7 @@ def image_variation(
|
|||
model_response=model_response,
|
||||
api_key=api_key,
|
||||
api_base=api_base,
|
||||
model=model,
|
||||
model=model or None,
|
||||
image=image,
|
||||
timeout=litellm_params.get("timeout", None),
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
|
|
@ -790,7 +786,7 @@ def image_variation(
|
|||
model_response=model_response,
|
||||
api_key=api_key,
|
||||
api_base=api_base,
|
||||
model=model,
|
||||
model=model or None,
|
||||
image=image,
|
||||
timeout=litellm_params.get("timeout", None) or DEFAULT_REQUEST_TIMEOUT,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
|
|
@ -901,7 +897,7 @@ def image_edit(
|
|||
# get llm provider logic
|
||||
litellm_params: Final = GenericLiteLLMParams(**kwargs)
|
||||
model, custom_llm_provider, _, _ = get_llm_provider(
|
||||
model=model or DEFAULT_IMAGE_ENDPOINT_MODEL,
|
||||
model=model or "",
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -437,7 +437,6 @@ def get_llm_provider(
|
|||
model in litellm.open_ai_chat_completion_models
|
||||
or "ft:gpt-3.5-turbo" in model
|
||||
or "ft:gpt-4" in model # catches ft:gpt-4-0613, ft:gpt-4o
|
||||
or model in litellm.openai_image_generation_models
|
||||
or model.startswith("gpt-image")
|
||||
or model in litellm.openai_video_generation_models
|
||||
):
|
||||
|
|
|
|||
|
|
@ -13,13 +13,12 @@ from openai import (
|
|||
BadRequestError,
|
||||
OpenAI,
|
||||
)
|
||||
from pydantic import BaseModel, TypeAdapter
|
||||
from pydantic import BaseModel
|
||||
|
||||
import litellm
|
||||
from litellm.constants import AZURE_OPERATION_POLLING_TIMEOUT, DEFAULT_MAX_RETRIES
|
||||
from litellm.constants import DEFAULT_MAX_RETRIES
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.litellm_core_utils.logging_utils import speech_request_body, track_llm_api_timing
|
||||
from litellm.litellm_core_utils.url_utils import SSRFError, assert_same_origin
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
AsyncHTTPHandler,
|
||||
HTTPHandler,
|
||||
|
|
@ -34,7 +33,6 @@ from litellm.types.utils import (
|
|||
from litellm.utils import (
|
||||
CustomStreamWrapper,
|
||||
convert_to_model_response_object,
|
||||
modify_url,
|
||||
)
|
||||
|
||||
from ...types.llms.openai import HttpxBinaryResponseContent
|
||||
|
|
@ -59,8 +57,6 @@ from .image_generation import (
|
|||
)
|
||||
from .image_generation.http_utils import azure_deployment_image_generation_json_body
|
||||
|
||||
_DECODED_JSON: Final = TypeAdapter(object)
|
||||
|
||||
|
||||
class AzureOpenAIAssistantsAPIConfig:
|
||||
"""
|
||||
|
|
@ -860,11 +856,6 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
|
|||
headers: dict,
|
||||
deployment_name: str | None = None,
|
||||
) -> httpx.Response:
|
||||
"""
|
||||
Implemented for azure dall-e-2 image gen calls
|
||||
|
||||
Alternative to needing a custom transport implementation
|
||||
"""
|
||||
if client is None:
|
||||
_params: Final = {}
|
||||
if timeout is not None:
|
||||
|
|
@ -881,95 +872,6 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
|
|||
else:
|
||||
async_handler = client
|
||||
|
||||
if (
|
||||
"images/generations" in api_base
|
||||
and api_version
|
||||
in [ # dall-e-3 starts from `2023-12-01-preview` so we should be able to avoid conflict
|
||||
"2023-06-01-preview",
|
||||
"2023-07-01-preview",
|
||||
"2023-08-01-preview",
|
||||
"2023-09-01-preview",
|
||||
"2023-10-01-preview",
|
||||
]
|
||||
): # CREATE + POLL for azure dall-e-2 calls
|
||||
api_base = modify_url(original_url=api_base, new_path="/openai/images/generations:submit")
|
||||
|
||||
data.pop(
|
||||
"model", None
|
||||
) # REMOVE 'model' from dall-e-2 arg https://learn.microsoft.com/en-us/azure/ai-services/openai/reference#request-a-generated-image-dall-e-2-preview
|
||||
response = await async_handler.post(
|
||||
url=api_base,
|
||||
data=json.dumps(data),
|
||||
headers=headers,
|
||||
)
|
||||
if "operation-location" in response.headers:
|
||||
operation_location_url: Final = response.headers["operation-location"]
|
||||
else:
|
||||
raise AzureOpenAIError(status_code=500, message=response.text)
|
||||
# Reject polling URLs that don't share an origin with ``api_base``.
|
||||
# Without this an upstream-controlled or attacker-controlled
|
||||
# value would receive the operator's Azure API key in the
|
||||
# request headers below. VERIA-51.
|
||||
try:
|
||||
assert_same_origin(operation_location_url, api_base)
|
||||
except SSRFError as ssrf_err:
|
||||
raise AzureOpenAIError(
|
||||
status_code=502,
|
||||
message=f"Rejected polling URL: {ssrf_err}",
|
||||
)
|
||||
response = await async_handler.get(
|
||||
url=operation_location_url,
|
||||
headers=headers,
|
||||
)
|
||||
|
||||
await response.aread()
|
||||
|
||||
timeout_secs: Final[int] = AZURE_OPERATION_POLLING_TIMEOUT
|
||||
start_time: Final = time.time()
|
||||
if "status" not in response.json():
|
||||
# Don't reflect the raw response body — when the polling
|
||||
# URL points at an internal JSON API (cloud metadata
|
||||
# service etc.) reflecting it here turns Blind SSRF into
|
||||
# Full-Read SSRF. VERIA-51.
|
||||
raise AzureOpenAIError(
|
||||
status_code=502,
|
||||
message="Polling response missing 'status' field",
|
||||
)
|
||||
while response.json()["status"] not in ["succeeded", "failed"]:
|
||||
if time.time() - start_time > timeout_secs:
|
||||
raise AzureOpenAIError(status_code=408, message="Operation polling timed out.")
|
||||
|
||||
await asyncio.sleep(int(response.headers.get("retry-after") or 10))
|
||||
response = await async_handler.get(
|
||||
url=operation_location_url,
|
||||
headers=headers,
|
||||
)
|
||||
await response.aread()
|
||||
|
||||
if response.json()["status"] == "failed":
|
||||
error_data: Final = response.json()
|
||||
# Preserve Azure error details (e.g. content_policy_violation,
|
||||
# inner_error, content_filter_results) as structured body so
|
||||
# exception_type() can route them correctly.
|
||||
_error_body: Final = error_data.get("error", error_data)
|
||||
_error_msg: Final = (
|
||||
_error_body.get("message", "Image generation failed")
|
||||
if isinstance(_error_body, dict)
|
||||
else json.dumps(error_data)
|
||||
)
|
||||
raise AzureOpenAIError(
|
||||
status_code=400,
|
||||
message=_error_msg,
|
||||
body=error_data,
|
||||
)
|
||||
|
||||
result: Final = _DECODED_JSON.validate_python(response.json()["result"])
|
||||
return httpx.Response(
|
||||
status_code=200,
|
||||
headers=response.headers,
|
||||
content=json.dumps(result).encode("utf-8"),
|
||||
request=httpx.Request(method="POST", url="https://api.openai.com/v1"),
|
||||
)
|
||||
request_json: Final = azure_deployment_image_generation_json_body(api_base, data, deployment_name)
|
||||
return await async_handler.post(
|
||||
url=api_base,
|
||||
|
|
@ -988,11 +890,6 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
|
|||
headers: dict,
|
||||
deployment_name: str | None = None,
|
||||
) -> httpx.Response:
|
||||
"""
|
||||
Implemented for azure dall-e-2 image gen calls
|
||||
|
||||
Alternative to needing a custom transport implementation
|
||||
"""
|
||||
if client is None:
|
||||
_params: Final = {}
|
||||
if timeout is not None:
|
||||
|
|
@ -1006,87 +903,6 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
|
|||
else:
|
||||
sync_handler = client
|
||||
|
||||
if (
|
||||
"images/generations" in api_base
|
||||
and api_version
|
||||
in [ # dall-e-3 starts from `2023-12-01-preview` so we should be able to avoid conflict
|
||||
"2023-06-01-preview",
|
||||
"2023-07-01-preview",
|
||||
"2023-08-01-preview",
|
||||
"2023-09-01-preview",
|
||||
"2023-10-01-preview",
|
||||
]
|
||||
): # CREATE + POLL for azure dall-e-2 calls
|
||||
api_base = modify_url(original_url=api_base, new_path="/openai/images/generations:submit")
|
||||
|
||||
data.pop(
|
||||
"model", None
|
||||
) # REMOVE 'model' from dall-e-2 arg https://learn.microsoft.com/en-us/azure/ai-services/openai/reference#request-a-generated-image-dall-e-2-preview
|
||||
response = sync_handler.post(
|
||||
url=api_base,
|
||||
data=json.dumps(data),
|
||||
headers=headers,
|
||||
)
|
||||
if "operation-location" in response.headers:
|
||||
operation_location_url: Final = response.headers["operation-location"]
|
||||
else:
|
||||
raise AzureOpenAIError(status_code=500, message=response.text)
|
||||
try:
|
||||
assert_same_origin(operation_location_url, api_base)
|
||||
except SSRFError as ssrf_err:
|
||||
raise AzureOpenAIError(
|
||||
status_code=502,
|
||||
message=f"Rejected polling URL: {ssrf_err}",
|
||||
)
|
||||
response = sync_handler.get(
|
||||
url=operation_location_url,
|
||||
headers=headers,
|
||||
)
|
||||
|
||||
response.read()
|
||||
|
||||
timeout_secs: Final[int] = AZURE_OPERATION_POLLING_TIMEOUT
|
||||
start_time: Final = time.time()
|
||||
if "status" not in response.json():
|
||||
raise AzureOpenAIError(
|
||||
status_code=502,
|
||||
message="Polling response missing 'status' field",
|
||||
)
|
||||
while response.json()["status"] not in ["succeeded", "failed"]:
|
||||
if time.time() - start_time > timeout_secs:
|
||||
raise AzureOpenAIError(status_code=408, message="Operation polling timed out.")
|
||||
|
||||
time.sleep(int(response.headers.get("retry-after") or 10))
|
||||
response = sync_handler.get(
|
||||
url=operation_location_url,
|
||||
headers=headers,
|
||||
)
|
||||
response.read()
|
||||
|
||||
if response.json()["status"] == "failed":
|
||||
error_data: Final = response.json()
|
||||
# Preserve Azure error details (e.g. content_policy_violation,
|
||||
# inner_error, content_filter_results) as structured body so
|
||||
# exception_type() can route them correctly.
|
||||
_error_body: Final = error_data.get("error", error_data)
|
||||
_error_msg: Final = (
|
||||
_error_body.get("message", "Image generation failed")
|
||||
if isinstance(_error_body, dict)
|
||||
else json.dumps(error_data)
|
||||
)
|
||||
raise AzureOpenAIError(
|
||||
status_code=400,
|
||||
message=_error_msg,
|
||||
body=error_data,
|
||||
)
|
||||
|
||||
result: Final = _DECODED_JSON.validate_python(response.json()["result"])
|
||||
return httpx.Response(
|
||||
status_code=200,
|
||||
headers=response.headers,
|
||||
content=json.dumps(result).encode("utf-8"),
|
||||
request=httpx.Request(method="POST", url="https://api.openai.com/v1"),
|
||||
)
|
||||
request_json: Final = azure_deployment_image_generation_json_body(api_base, data, deployment_name)
|
||||
return sync_handler.post(
|
||||
url=api_base,
|
||||
|
|
@ -1193,10 +1009,10 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
|
|||
deployment_name=model,
|
||||
)
|
||||
|
||||
provider_config: Final = get_azure_image_generation_config(data.get("model", "dall-e-2"))
|
||||
provider_config: Final = get_azure_image_generation_config(data.get("model", ""))
|
||||
if provider_config is not None:
|
||||
return provider_config.transform_image_generation_response(
|
||||
model=data.get("model", "dall-e-2"),
|
||||
model=data.get("model", ""),
|
||||
raw_response=httpx_response,
|
||||
model_response=model_response or ImageResponse(),
|
||||
logging_obj=logging_obj,
|
||||
|
|
@ -1329,10 +1145,10 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
|
|||
headers=request_headers,
|
||||
deployment_name=model,
|
||||
)
|
||||
provider_config: Final = get_azure_image_generation_config(data.get("model", "dall-e-2"))
|
||||
provider_config: Final = get_azure_image_generation_config(data.get("model", ""))
|
||||
if isinstance(provider_config, AzureFoundryMAIImageGenerationConfig):
|
||||
return provider_config.transform_image_generation_response(
|
||||
model=data.get("model", "dall-e-2"),
|
||||
model=data.get("model", ""),
|
||||
raw_response=httpx_response,
|
||||
model_response=model_response or ImageResponse(),
|
||||
logging_obj=logging_obj,
|
||||
|
|
|
|||
|
|
@ -4,14 +4,10 @@ from litellm.llms.base_llm.image_generation.transformation import (
|
|||
BaseImageGenerationConfig,
|
||||
)
|
||||
|
||||
from .dall_e_2_transformation import AzureDallE2ImageGenerationConfig
|
||||
from .dall_e_3_transformation import AzureDallE3ImageGenerationConfig
|
||||
from .gpt_transformation import AzureGPTImageGenerationConfig
|
||||
from .http_utils import azure_deployment_image_generation_json_body
|
||||
|
||||
__all__ = [
|
||||
"AzureDallE2ImageGenerationConfig",
|
||||
"AzureDallE3ImageGenerationConfig",
|
||||
"AzureGPTImageGenerationConfig",
|
||||
"azure_deployment_image_generation_json_body",
|
||||
]
|
||||
|
|
@ -21,11 +17,7 @@ def get_azure_image_generation_config(model: str) -> BaseImageGenerationConfig:
|
|||
model = model.lower()
|
||||
model = model.replace("-", "")
|
||||
model = model.replace("_", "")
|
||||
if model == "" or "dalle2" in model: # empty model is dall-e-2
|
||||
return AzureDallE2ImageGenerationConfig()
|
||||
elif "dalle3" in model:
|
||||
return AzureDallE3ImageGenerationConfig()
|
||||
elif AzureFoundryMAIImageGenerationConfig.is_mai_model(model):
|
||||
if AzureFoundryMAIImageGenerationConfig.is_mai_model(model):
|
||||
return AzureFoundryMAIImageGenerationConfig()
|
||||
else:
|
||||
verbose_logger.debug(
|
||||
|
|
|
|||
|
|
@ -1,7 +0,0 @@
|
|||
from litellm.llms.openai.image_generation import DallE2ImageGenerationConfig
|
||||
|
||||
|
||||
class AzureDallE2ImageGenerationConfig(DallE2ImageGenerationConfig):
|
||||
"""
|
||||
Azure dall-e-2 image generation config
|
||||
"""
|
||||
|
|
@ -1,7 +0,0 @@
|
|||
from litellm.llms.openai.image_generation import DallE3ImageGenerationConfig
|
||||
|
||||
|
||||
class AzureDallE3ImageGenerationConfig(DallE3ImageGenerationConfig):
|
||||
"""
|
||||
Azure dall-e-3 image generation config
|
||||
"""
|
||||
|
|
@ -3,15 +3,11 @@ from litellm.llms.base_llm.image_generation.transformation import (
|
|||
BaseImageGenerationConfig,
|
||||
)
|
||||
|
||||
from .dall_e_2_transformation import AzureFoundryDallE2ImageGenerationConfig
|
||||
from .dall_e_3_transformation import AzureFoundryDallE3ImageGenerationConfig
|
||||
from .flux_transformation import AzureFoundryFluxImageGenerationConfig
|
||||
from .gpt_transformation import AzureFoundryGPTImageGenerationConfig
|
||||
from .mai_transformation import AzureFoundryMAIImageGenerationConfig
|
||||
|
||||
__all__ = [
|
||||
"AzureFoundryDallE2ImageGenerationConfig",
|
||||
"AzureFoundryDallE3ImageGenerationConfig",
|
||||
"AzureFoundryFluxImageGenerationConfig",
|
||||
"AzureFoundryGPTImageGenerationConfig",
|
||||
"AzureFoundryMAIImageGenerationConfig",
|
||||
|
|
@ -22,11 +18,7 @@ def get_azure_ai_image_generation_config(model: str) -> BaseImageGenerationConfi
|
|||
model = model.lower()
|
||||
model = model.replace("-", "")
|
||||
model = model.replace("_", "")
|
||||
if model == "" or "dalle2" in model: # empty model is dall-e-2
|
||||
return AzureFoundryDallE2ImageGenerationConfig()
|
||||
elif "dalle3" in model:
|
||||
return AzureFoundryDallE3ImageGenerationConfig()
|
||||
elif AzureFoundryMAIImageGenerationConfig.is_mai_model(model):
|
||||
if AzureFoundryMAIImageGenerationConfig.is_mai_model(model):
|
||||
return AzureFoundryMAIImageGenerationConfig()
|
||||
elif "flux" in model:
|
||||
return AzureFoundryFluxImageGenerationConfig()
|
||||
|
|
|
|||
|
|
@ -1,7 +0,0 @@
|
|||
from litellm.llms.openai.image_generation import DallE2ImageGenerationConfig
|
||||
|
||||
|
||||
class AzureFoundryDallE2ImageGenerationConfig(DallE2ImageGenerationConfig):
|
||||
"""
|
||||
Azure dall-e-2 image generation config
|
||||
"""
|
||||
|
|
@ -1,7 +0,0 @@
|
|||
from litellm.llms.openai.image_generation import DallE3ImageGenerationConfig
|
||||
|
||||
|
||||
class AzureFoundryDallE3ImageGenerationConfig(DallE3ImageGenerationConfig):
|
||||
"""
|
||||
Azure dall-e-3 image generation config
|
||||
"""
|
||||
|
|
@ -561,11 +561,8 @@ class BaseLLMAIOHTTPHandler:
|
|||
organization: str | None = None,
|
||||
headers: dict | None = None,
|
||||
) -> ImageResponse:
|
||||
if model is None:
|
||||
raise ValueError("model is required for non-openai image variations")
|
||||
|
||||
provider_config: Final = ProviderConfigManager.get_provider_image_variation_config(
|
||||
model=model, # openai defaults to dall-e-2
|
||||
model=model or "",
|
||||
provider=LlmProviders(custom_llm_provider),
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -1,31 +1,5 @@
|
|||
from typing import Final
|
||||
|
||||
from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig
|
||||
|
||||
from .dalle2_transformation import DallE2ImageEditConfig
|
||||
from .transformation import OpenAIImageEditConfig
|
||||
|
||||
__all__ = [
|
||||
"DallE2ImageEditConfig",
|
||||
"OpenAIImageEditConfig",
|
||||
"get_openai_image_edit_config",
|
||||
]
|
||||
|
||||
|
||||
def get_openai_image_edit_config(model: str) -> BaseImageEditConfig:
|
||||
"""
|
||||
Get the appropriate OpenAI image edit config based on the model.
|
||||
|
||||
Args:
|
||||
model: The model name (e.g., "dall-e-2", "gpt-image-1")
|
||||
|
||||
Returns:
|
||||
The appropriate config instance for the model
|
||||
"""
|
||||
model_normalized: Final = model.lower().replace("-", "").replace("_", "")
|
||||
|
||||
if model_normalized == "dalle2":
|
||||
return DallE2ImageEditConfig()
|
||||
else:
|
||||
# Default to standard OpenAI config for gpt-image-1 and other models
|
||||
return OpenAIImageEditConfig()
|
||||
|
|
|
|||
|
|
@ -1,98 +0,0 @@
|
|||
from io import BufferedReader
|
||||
from typing import TYPE_CHECKING, Any, Final, cast
|
||||
|
||||
from httpx._types import RequestFiles
|
||||
|
||||
import litellm
|
||||
from litellm.images.utils import ImageEditRequestUtils
|
||||
from litellm.types.images.main import ImageEditRequestParams
|
||||
from litellm.types.llms.openai import FileTypes
|
||||
from litellm.types.router import GenericLiteLLMParams
|
||||
|
||||
from .transformation import OpenAIImageEditConfig
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
|
||||
|
||||
LiteLLMLoggingObj = _LiteLLMLoggingObj
|
||||
else:
|
||||
LiteLLMLoggingObj = Any
|
||||
|
||||
|
||||
class DallE2ImageEditConfig(OpenAIImageEditConfig):
|
||||
"""
|
||||
DALL-E-2 specific configuration for image edit API.
|
||||
|
||||
DALL-E-2 only supports editing a single image (not an array).
|
||||
Uses "image" field name instead of "image[]".
|
||||
"""
|
||||
|
||||
def transform_image_edit_request(
|
||||
self,
|
||||
model: str,
|
||||
prompt: str | None,
|
||||
image: FileTypes | None,
|
||||
image_edit_optional_request_params: dict,
|
||||
litellm_params: GenericLiteLLMParams,
|
||||
headers: dict,
|
||||
) -> tuple[dict, RequestFiles]:
|
||||
"""
|
||||
Transform image edit request for DALL-E-2.
|
||||
|
||||
DALL-E-2 only accepts a single image with field name "image" (not "image[]").
|
||||
"""
|
||||
request_params: Final = {
|
||||
"model": model,
|
||||
**image_edit_optional_request_params,
|
||||
}
|
||||
if image is not None:
|
||||
request_params["image"] = image
|
||||
if prompt is not None:
|
||||
request_params["prompt"] = prompt
|
||||
|
||||
request: Final = ImageEditRequestParams(**request_params)
|
||||
request_dict: Final = cast(dict, request)
|
||||
|
||||
#########################################################
|
||||
# Separate images and masks as `files` and send other parameters as `data`
|
||||
#########################################################
|
||||
_image_list: Final = request_dict.get("image")
|
||||
_mask = request_dict.get("mask")
|
||||
data_without_files: Final = {k: v for k, v in request_dict.items() if k not in ["image", "mask"]}
|
||||
files_list: Final[list[tuple[str, Any]]] = []
|
||||
|
||||
# Handle image parameter - DALL-E-2 only supports single image
|
||||
if _image_list is not None:
|
||||
image_list: Final = [_image_list] if not isinstance(_image_list, list) else _image_list
|
||||
|
||||
# Validate only one image is provided
|
||||
if len(image_list) > 1:
|
||||
raise litellm.BadRequestError(
|
||||
message="DALL-E-2 only supports editing a single image. Please provide one image.",
|
||||
model=model,
|
||||
llm_provider="openai",
|
||||
)
|
||||
|
||||
# Use "image" field name (singular) for DALL-E-2
|
||||
for _image in image_list:
|
||||
if _image is not None:
|
||||
self._add_image_to_files(
|
||||
files_list=files_list,
|
||||
image=_image,
|
||||
field_name="image",
|
||||
)
|
||||
|
||||
# Handle mask parameter if provided
|
||||
if _mask is not None:
|
||||
# Handle case where mask can be a list (extract first mask)
|
||||
if isinstance(_mask, list):
|
||||
_mask = _mask[0] if _mask else None
|
||||
|
||||
if _mask is not None:
|
||||
mask_content_type: Final[str] = ImageEditRequestUtils.get_image_content_type(_mask)
|
||||
if isinstance(_mask, BufferedReader):
|
||||
files_list.append(("mask", (_mask.name, _mask, mask_content_type)))
|
||||
else:
|
||||
files_list.append(("mask", ("mask.png", _mask, mask_content_type)))
|
||||
|
||||
return data_without_files, files_list
|
||||
|
|
@ -96,7 +96,7 @@ class OpenAIImageEditConfig(BaseImageEditConfig):
|
|||
"""
|
||||
# Build request params, only including non-None values
|
||||
request_params: Final = {
|
||||
"model": model,
|
||||
**({"model": model} if model else {}),
|
||||
**image_edit_optional_request_params,
|
||||
}
|
||||
if image is not None:
|
||||
|
|
|
|||
|
|
@ -1,9 +1,3 @@
|
|||
from litellm.llms.base_llm.image_generation.transformation import (
|
||||
BaseImageGenerationConfig,
|
||||
)
|
||||
|
||||
from .dall_e_2_transformation import DallE2ImageGenerationConfig
|
||||
from .dall_e_3_transformation import DallE3ImageGenerationConfig
|
||||
from .gpt_transformation import GPTImageGenerationConfig
|
||||
from .guardrail_translation import (
|
||||
OpenAIImageGenerationHandler,
|
||||
|
|
@ -11,18 +5,7 @@ from .guardrail_translation import (
|
|||
)
|
||||
|
||||
__all__ = [
|
||||
"DallE2ImageGenerationConfig",
|
||||
"DallE3ImageGenerationConfig",
|
||||
"GPTImageGenerationConfig",
|
||||
"OpenAIImageGenerationHandler",
|
||||
"guardrail_translation_mappings",
|
||||
]
|
||||
|
||||
|
||||
def get_openai_image_generation_config(model: str) -> BaseImageGenerationConfig:
|
||||
if model.startswith("dall-e-2") or model == "": # empty model is dall-e-2
|
||||
return DallE2ImageGenerationConfig()
|
||||
elif model.startswith("dall-e-3"):
|
||||
return DallE3ImageGenerationConfig()
|
||||
else:
|
||||
return GPTImageGenerationConfig()
|
||||
|
|
|
|||
|
|
@ -1,81 +0,0 @@
|
|||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
import httpx
|
||||
|
||||
from litellm.llms.base_llm.image_generation.transformation import (
|
||||
BaseImageGenerationConfig,
|
||||
)
|
||||
from litellm.types.llms.openai import OpenAIImageGenerationOptionalParams
|
||||
from litellm.types.utils import ImageResponse
|
||||
from litellm.utils import convert_to_model_response_object
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.logging import Logging as LiteLLMLoggingObj
|
||||
|
||||
from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer
|
||||
|
||||
|
||||
class DallE2ImageGenerationConfig(BaseImageGenerationConfig):
|
||||
"""
|
||||
OpenAI dall-e-2 image generation config
|
||||
"""
|
||||
|
||||
def get_supported_openai_params(self, model: str) -> list[OpenAIImageGenerationOptionalParams]:
|
||||
return ["n", "response_format", "quality", "size", "user"]
|
||||
|
||||
def map_openai_params(
|
||||
self,
|
||||
non_default_params: dict,
|
||||
optional_params: dict,
|
||||
model: str,
|
||||
drop_params: bool,
|
||||
) -> dict:
|
||||
supported_params: Final = self.get_supported_openai_params(model)
|
||||
for k in non_default_params:
|
||||
if k not in optional_params:
|
||||
if k in supported_params:
|
||||
optional_params[k] = non_default_params[k]
|
||||
elif drop_params:
|
||||
pass
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Parameter {k} is not supported for model {model}. Supported parameters are {supported_params}. Set drop_params=True to drop unsupported parameters."
|
||||
)
|
||||
|
||||
return optional_params
|
||||
|
||||
def transform_image_generation_response(
|
||||
self,
|
||||
model: str,
|
||||
raw_response: httpx.Response,
|
||||
model_response: ImageResponse,
|
||||
logging_obj: "LiteLLMLoggingObj",
|
||||
request_data: dict,
|
||||
optional_params: dict,
|
||||
litellm_params: dict,
|
||||
encoding: "Tokenizer | None",
|
||||
api_key: str | None = None,
|
||||
json_mode: bool | None = None,
|
||||
) -> ImageResponse:
|
||||
response: Final = raw_response.json()
|
||||
|
||||
stringified_response: Final = response
|
||||
## LOGGING
|
||||
logging_obj.post_call(
|
||||
input=request_data.get("prompt", ""),
|
||||
api_key=api_key,
|
||||
additional_args={"complete_input_dict": request_data},
|
||||
original_response=stringified_response,
|
||||
)
|
||||
image_response: Final[ImageResponse] = convert_to_model_response_object(
|
||||
response_object=stringified_response,
|
||||
model_response_object=model_response,
|
||||
response_type="image_generation",
|
||||
)
|
||||
|
||||
# set optional params
|
||||
image_response.size = optional_params.get("size", "1024x1024") # default is always 1024x1024
|
||||
image_response.quality = optional_params.get("quality", "standard") # always standard for dall-e-2
|
||||
image_response.output_format = optional_params.get("output_format", "png") # always png for dall-e-2
|
||||
|
||||
return image_response
|
||||
|
|
@ -1,81 +0,0 @@
|
|||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
import httpx
|
||||
|
||||
from litellm.llms.base_llm.image_generation.transformation import (
|
||||
BaseImageGenerationConfig,
|
||||
)
|
||||
from litellm.types.llms.openai import OpenAIImageGenerationOptionalParams
|
||||
from litellm.types.utils import ImageResponse
|
||||
from litellm.utils import convert_to_model_response_object
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.logging import Logging as LiteLLMLoggingObj
|
||||
|
||||
from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer
|
||||
|
||||
|
||||
class DallE3ImageGenerationConfig(BaseImageGenerationConfig):
|
||||
"""
|
||||
OpenAI dall-e-3 image generation config
|
||||
"""
|
||||
|
||||
def get_supported_openai_params(self, model: str) -> list[OpenAIImageGenerationOptionalParams]:
|
||||
return ["n", "response_format", "quality", "size", "user", "style"]
|
||||
|
||||
def map_openai_params(
|
||||
self,
|
||||
non_default_params: dict,
|
||||
optional_params: dict,
|
||||
model: str,
|
||||
drop_params: bool,
|
||||
) -> dict:
|
||||
supported_params: Final = self.get_supported_openai_params(model)
|
||||
for k in non_default_params:
|
||||
if k not in optional_params:
|
||||
if k in supported_params:
|
||||
optional_params[k] = non_default_params[k]
|
||||
elif drop_params:
|
||||
pass
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Parameter {k} is not supported for model {model}. Supported parameters are {supported_params}. Set drop_params=True to drop unsupported parameters."
|
||||
)
|
||||
|
||||
return optional_params
|
||||
|
||||
def transform_image_generation_response(
|
||||
self,
|
||||
model: str,
|
||||
raw_response: httpx.Response,
|
||||
model_response: ImageResponse,
|
||||
logging_obj: "LiteLLMLoggingObj",
|
||||
request_data: dict,
|
||||
optional_params: dict,
|
||||
litellm_params: dict,
|
||||
encoding: "Tokenizer | None",
|
||||
api_key: str | None = None,
|
||||
json_mode: bool | None = None,
|
||||
) -> ImageResponse:
|
||||
response: Final = raw_response.json()
|
||||
|
||||
stringified_response: Final = response
|
||||
## LOGGING
|
||||
logging_obj.post_call(
|
||||
input=request_data.get("prompt", ""),
|
||||
api_key=api_key,
|
||||
additional_args={"complete_input_dict": request_data},
|
||||
original_response=stringified_response,
|
||||
)
|
||||
image_response: Final[ImageResponse] = convert_to_model_response_object(
|
||||
response_object=stringified_response,
|
||||
model_response_object=model_response,
|
||||
response_type="image_generation",
|
||||
)
|
||||
|
||||
# set optional params
|
||||
image_response.size = optional_params.get("size", "1024x1024") # default is always 1024x1024
|
||||
image_response.quality = optional_params.get("quality", "hd") # always hd for dall-e-3
|
||||
image_response.output_format = optional_params.get("output_format", "png") # always png for dall-e-3
|
||||
|
||||
return image_response
|
||||
|
|
@ -151,7 +151,7 @@ class OpenAIImageVariationsHandler:
|
|||
) -> ImageResponse:
|
||||
try:
|
||||
provider_config: Final = ProviderConfigManager.get_provider_image_variation_config(
|
||||
model=model or "", # openai defaults to dall-e-2
|
||||
model=model or "",
|
||||
provider=LlmProviders.OPENAI,
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -7664,7 +7664,6 @@ def validate_environment(
|
|||
model in litellm.open_ai_chat_completion_models
|
||||
or model in litellm.open_ai_text_completion_models
|
||||
or model in litellm.open_ai_embedding_models
|
||||
or model in litellm.openai_image_generation_models
|
||||
or model.startswith("gpt-image")
|
||||
):
|
||||
if "OPENAI_API_KEY" in os.environ:
|
||||
|
|
@ -10311,11 +10310,9 @@ class ProviderConfigManager:
|
|||
provider: LlmProviders,
|
||||
) -> BaseImageGenerationConfig | None:
|
||||
if LlmProviders.OPENAI == provider:
|
||||
from litellm.llms.openai.image_generation import (
|
||||
get_openai_image_generation_config,
|
||||
)
|
||||
from litellm.llms.openai.image_generation import GPTImageGenerationConfig
|
||||
|
||||
return get_openai_image_generation_config(model)
|
||||
return GPTImageGenerationConfig()
|
||||
elif LlmProviders.AZURE == provider:
|
||||
from litellm.llms.azure.image_generation import (
|
||||
get_azure_image_generation_config,
|
||||
|
|
@ -10546,9 +10543,9 @@ class ProviderConfigManager:
|
|||
provider: LlmProviders,
|
||||
) -> BaseImageEditConfig | None:
|
||||
if LlmProviders.OPENAI == provider:
|
||||
from litellm.llms.openai.image_edit import get_openai_image_edit_config
|
||||
from litellm.llms.openai.image_edit import OpenAIImageEditConfig
|
||||
|
||||
return get_openai_image_edit_config(model=model)
|
||||
return OpenAIImageEditConfig()
|
||||
elif LlmProviders.HOSTED_VLLM == provider:
|
||||
from litellm.llms.hosted_vllm.image_edit import get_hosted_vllm_image_edit_config
|
||||
|
||||
|
|
|
|||
|
|
@ -26,6 +26,20 @@ _DECISIONS_PROBE_REPLY: Final = JsonResponse(
|
|||
"usage": {"input_tokens": 10, "output_tokens": 1},
|
||||
},
|
||||
)
|
||||
_EXTRACTION_PROBE_REPLY: Final = JsonResponse(
|
||||
content_type="application/json",
|
||||
body={
|
||||
"result": {
|
||||
"model": "clairic",
|
||||
"answers": {"reachable": {"type": "text", "value": "health", "confidence": 0.9}},
|
||||
"truncated": False,
|
||||
"usage": {"input_tokens": 10, "output_tokens": 0},
|
||||
},
|
||||
"success": True,
|
||||
"errors": [],
|
||||
"messages": [],
|
||||
},
|
||||
)
|
||||
_CONFIGURED_PROBE_REPLY: Final = JsonResponse(
|
||||
content_type="application/json",
|
||||
body={
|
||||
|
|
@ -287,6 +301,29 @@ def test_decisions_mode_health_check_resolves_the_mode_from_the_cost_map_and_sen
|
|||
]
|
||||
|
||||
|
||||
def test_extraction_mode_health_check_resolves_the_mode_from_the_cost_map_and_sends_the_extraction_probe(
|
||||
gateway: Gateway,
|
||||
) -> None:
|
||||
with gateway.scenario() as scenario:
|
||||
handle: Final = register_scenario(f"health-extraction-{uuid.uuid4().hex[:12]}", _EXTRACTION_PROBE_REPLY)
|
||||
scenario.cleanups.callback(delete_scenario, handle)
|
||||
model: Final = scenario.model(
|
||||
model="cloudflare/@cf/cloudflare/clairic", api_base=handle.api_base(), api_key="synthetic-cloudflare-key"
|
||||
)
|
||||
report: Final = _health_report(gateway, model)
|
||||
assert (report["healthy_count"], report["unhealthy_count"]) == (1, 0), report
|
||||
assert _probes_sent_to(gateway, handle) == [
|
||||
(
|
||||
f"/{handle.scenario_id}/ai/run/@cf/cloudflare/clairic",
|
||||
{
|
||||
"model": "clairic",
|
||||
"state": os.environ.get("DEFAULT_HEALTH_CHECK_PROMPT", "test from litellm"),
|
||||
"fields": {"reachable": {"type": "text", "instructions": "Any word from the text"}},
|
||||
},
|
||||
)
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mode", ("decisions", "evaluation"), ids=("decisions", "legacy-evaluation"))
|
||||
def test_decisions_mode_health_check_sends_the_configured_state_and_questions(gateway: Gateway, mode: str) -> None:
|
||||
with gateway.scenario() as scenario:
|
||||
|
|
|
|||
|
|
@ -6,7 +6,7 @@ from concurrent.futures import ThreadPoolExecutor
|
|||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass
|
||||
from hashlib import sha256
|
||||
from itertools import chain
|
||||
from itertools import chain, count
|
||||
from pathlib import Path
|
||||
from typing import Final, Literal, TypeVar
|
||||
|
||||
|
|
@ -50,6 +50,7 @@ def _llm_gateway(tmp_path_factory: pytest.TempPathFactory) -> Iterator[tuple[Gat
|
|||
directory: Final = tmp_path_factory.mktemp("mcp-llm-gateway")
|
||||
stock: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()))
|
||||
config: Final = directory / "stored-responses.yaml"
|
||||
(directory / f"{STREAM_BLOCK_MODULE}.py").write_text(STREAM_BLOCK_SOURCE)
|
||||
config.write_text(
|
||||
yaml.safe_dump(
|
||||
{
|
||||
|
|
@ -57,8 +58,17 @@ def _llm_gateway(tmp_path_factory: pytest.TempPathFactory) -> Iterator[tuple[Gat
|
|||
"general_settings": {**object_value(stock["general_settings"]), "store_prompts_in_spend_logs": True},
|
||||
"litellm_settings": {
|
||||
**object_value(stock["litellm_settings"]),
|
||||
"cost_margin_config": {"hosted_vllm": {"fixed_amount": FIXED_MARGIN}},
|
||||
"cost_margin_config": {
|
||||
"hosted_vllm": {"fixed_amount": FIXED_MARGIN},
|
||||
"deepseek": {"fixed_amount": FIXED_MARGIN},
|
||||
},
|
||||
},
|
||||
"guardrails": [
|
||||
{
|
||||
"guardrail_name": STREAM_BLOCK,
|
||||
"litellm_params": {"guardrail": f"{STREAM_BLOCK_MODULE}.BlockGuard", "mode": "post_call"},
|
||||
}
|
||||
],
|
||||
}
|
||||
)
|
||||
)
|
||||
|
|
@ -79,7 +89,7 @@ def gateway(_llm_gateway: tuple[Gateway, Mapping[str, str]]) -> Iterator[Gateway
|
|||
yield private
|
||||
|
||||
|
||||
Surface = Literal["chat", "responses", "messages", "messages_bridge"]
|
||||
Surface = Literal["chat", "responses", "messages", "messages_bridge", "responses_bridge"]
|
||||
SURFACES: Final[tuple[Surface, ...]] = ("chat", "responses", "messages", "messages_bridge")
|
||||
ADD: Final = {"a": 2, "b": 3}
|
||||
ANSWER: Final = "the sum is 5"
|
||||
|
|
@ -88,6 +98,19 @@ AUTO: Final = {**GATEWAY_REF, "require_approval": "never"}
|
|||
OUTAGE: Final = "bridge-outage"
|
||||
DIRECT: Final = "answer-without-tools"
|
||||
FIXED_MARGIN: Final = 0.25
|
||||
STREAM_BLOCK: Final = "mcp-stream-block"
|
||||
STREAM_BLOCK_MODULE: Final = "mcp_stream_block"
|
||||
STREAM_BLOCK_DETAIL: Final = "blocked by output policy"
|
||||
STREAM_BLOCK_SOURCE: Final = (
|
||||
"from fastapi import HTTPException\n\n"
|
||||
"from litellm.integrations.custom_guardrail import CustomGuardrail\n\n\n"
|
||||
"class BlockGuard(CustomGuardrail):\n"
|
||||
" async def async_post_call_streaming_iterator_hook(self, user_api_key_dict, response, request_data):\n"
|
||||
" async for _ in response:\n"
|
||||
" pass\n"
|
||||
f' raise HTTPException(status_code=400, detail="{STREAM_BLOCK_DETAIL}")\n'
|
||||
" yield\n"
|
||||
)
|
||||
JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
|
||||
JSON_VALUE: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue)
|
||||
|
||||
|
|
@ -128,6 +151,8 @@ class Billing:
|
|||
first: tuple[int, int] = (10, 5)
|
||||
follow_up: tuple[int, int] = (10, 5)
|
||||
follow_up_fails: bool = False
|
||||
follow_up_breaks: bool = False
|
||||
tool_rounds: int = 1
|
||||
|
||||
|
||||
EVEN_ROUNDS: Final = Billing()
|
||||
|
|
@ -138,6 +163,10 @@ def _fixed_turn(tool: str) -> Callable[[Mapping[str, JsonValue]], Turn]:
|
|||
return lambda _: Turn(tool, json.dumps(ADD), ANSWER)
|
||||
|
||||
|
||||
def _slow_turn(tool: str, seconds: float) -> Callable[[Mapping[str, JsonValue]], Turn]:
|
||||
return lambda _: Turn(tool, json.dumps({"seconds": seconds}), ANSWER)
|
||||
|
||||
|
||||
def _echoing_turn(body: Mapping[str, JsonValue]) -> Turn:
|
||||
names: Final = _tool_names(body)
|
||||
return Turn(names[0] if names else "", json.dumps({"query": _prompt(body)}), _tool_result_text(body) or "")
|
||||
|
|
@ -183,6 +212,50 @@ def _item_events(index: int, item: Mapping[str, JsonValue]) -> tuple[dict[str, J
|
|||
)
|
||||
|
||||
|
||||
def _event_stream(events: Sequence[Mapping[str, JsonValue]]) -> Reply:
|
||||
return Reply(
|
||||
content_type="text/event-stream",
|
||||
chunks=tuple(
|
||||
f"event: {event['type']}\ndata: {json.dumps({**event, 'sequence_number': number})}\n\n".encode()
|
||||
for number, event in enumerate(events)
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _broken_follow_up(target: str, identity: str, tokens: tuple[int, int]) -> Reply:
|
||||
outage: Final = {"message": FOLLOW_UP_OUTAGE, "type": "server_error", "code": None}
|
||||
if target.endswith("/chat/completions"):
|
||||
delta: Final = {"role": "assistant", "content": "the sum"}
|
||||
started: Final = {
|
||||
"id": f"chatcmpl-{identity}",
|
||||
"object": "chat.completion.chunk",
|
||||
"created": 1,
|
||||
"model": "gpt-4o-mini",
|
||||
"choices": [{"index": 0, "delta": delta, "finish_reason": None}],
|
||||
}
|
||||
return Reply(
|
||||
content_type="text/event-stream",
|
||||
chunks=tuple(f"data: {json.dumps(frame)}\n\n".encode() for frame in (started, {"error": outage})),
|
||||
)
|
||||
input_tokens, output_tokens = tokens
|
||||
failed: Final[dict[str, JsonValue]] = {
|
||||
"id": f"resp_{identity}",
|
||||
"object": "response",
|
||||
"created_at": 1,
|
||||
"status": "failed",
|
||||
"model": "gpt-4o-mini",
|
||||
"output": [],
|
||||
"error": {"code": "server_error", "message": FOLLOW_UP_OUTAGE},
|
||||
"usage": {"input_tokens": input_tokens, "output_tokens": output_tokens, "total_tokens": sum(tokens)},
|
||||
}
|
||||
return _event_stream(
|
||||
(
|
||||
{"type": "response.created", "response": {**failed, "status": "in_progress", "error": None, "usage": None}},
|
||||
{"type": "response.failed", "response": failed},
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _responses_stream(response: Mapping[str, JsonValue], items: Sequence[Mapping[str, JsonValue]]) -> Reply:
|
||||
item_events: Final = tuple(chain.from_iterable(_item_events(index, item) for index, item in enumerate(items)))
|
||||
events: Final = (
|
||||
|
|
@ -191,13 +264,7 @@ def _responses_stream(response: Mapping[str, JsonValue], items: Sequence[Mapping
|
|||
*item_events,
|
||||
{"type": "response.completed", "response": response},
|
||||
)
|
||||
return Reply(
|
||||
content_type="text/event-stream",
|
||||
chunks=tuple(
|
||||
f"event: {event['type']}\ndata: {json.dumps({**event, 'sequence_number': number})}\n\n".encode()
|
||||
for number, event in enumerate(events)
|
||||
),
|
||||
)
|
||||
return _event_stream(events)
|
||||
|
||||
|
||||
def _block_ids(message: JsonValue, kind: str, field: str) -> frozenset[str]:
|
||||
|
|
@ -252,6 +319,8 @@ def _anthropic_tool_result_rule(body: Mapping[str, JsonValue]) -> Reply | None:
|
|||
def _model_double(
|
||||
plan: Callable[[Mapping[str, JsonValue]], Turn], billing: Billing = EVEN_ROUNDS, parallel: int = 1
|
||||
) -> Callable[[Request], Reply]:
|
||||
served: Final = count()
|
||||
|
||||
def respond(request: Request) -> Reply:
|
||||
if request.method == "GET" and request.target.endswith("/models"):
|
||||
return _json({"object": "list", "data": []})
|
||||
|
|
@ -260,11 +329,12 @@ def _model_double(
|
|||
outage: Final = {"error": {"message": "scripted provider outage", "type": "server_error", "code": None}}
|
||||
return Reply(status=500, body=json.dumps(outage).encode())
|
||||
turn: Final = plan(body)
|
||||
call_number: Final = next(served)
|
||||
answered: Final = _has_tool_result(body)
|
||||
if answered and billing.follow_up_fails:
|
||||
failure: Final = {"error": {"message": FOLLOW_UP_OUTAGE, "type": "server_error", "code": None}}
|
||||
return Reply(status=500, body=json.dumps(failure).encode())
|
||||
done: Final = answered or DIRECT in _prompt(body)
|
||||
done: Final = DIRECT in _prompt(body) or (answered and call_number >= billing.tool_rounds)
|
||||
calls: Final = range(1, parallel + 1)
|
||||
prompt_tokens, completion_tokens = billing.follow_up if answered else billing.first
|
||||
identity: Final = uuid.uuid4().hex[:12]
|
||||
|
|
@ -273,9 +343,15 @@ def _model_double(
|
|||
"completion_tokens": completion_tokens,
|
||||
"total_tokens": prompt_tokens + completion_tokens,
|
||||
}
|
||||
if answered and billing.follow_up_breaks:
|
||||
return _broken_follow_up(request.target, identity, (prompt_tokens, completion_tokens))
|
||||
if request.target.endswith("/chat/completions"):
|
||||
requested: Final = tuple(
|
||||
{"id": f"call_{call}", "type": "function", "function": {"name": turn.tool, "arguments": turn.arguments}}
|
||||
{
|
||||
"id": f"call_{identity}_{call}",
|
||||
"type": "function",
|
||||
"function": {"name": turn.tool, "arguments": turn.arguments},
|
||||
}
|
||||
for call in calls
|
||||
)
|
||||
message: Final = (
|
||||
|
|
@ -460,6 +536,7 @@ def _upstream_model(surface: Surface) -> str:
|
|||
"responses": "openai/responses/gpt-4o-mini",
|
||||
"messages": "anthropic/claude-sonnet-4-5",
|
||||
"messages_bridge": "hosted_vllm/gpt-4o-mini",
|
||||
"responses_bridge": "deepseek/gpt-4o-mini",
|
||||
}[surface]
|
||||
|
||||
|
||||
|
|
@ -475,11 +552,15 @@ def _rig(
|
|||
billing: Billing = EVEN_ROUNDS,
|
||||
prices: Mapping[str, JsonValue] | None = None,
|
||||
parallel: int = 1,
|
||||
slow_tool_seconds: float | None = None,
|
||||
) -> Iterator[Rig]:
|
||||
alias: Final = "llm" + uuid.uuid4().hex[:8]
|
||||
plan: Final = (
|
||||
_fixed_turn(f"{alias}-add") if slow_tool_seconds is None else _slow_turn(f"{alias}-slow", slow_tool_seconds)
|
||||
)
|
||||
with (
|
||||
mcp_peer() as peer,
|
||||
wire_server(_model_double(_fixed_turn(f"{alias}-add"), billing, parallel)) as wire,
|
||||
mcp_peer(rich=slow_tool_seconds is not None) as peer,
|
||||
wire_server(_model_double(plan, billing, parallel)) as wire,
|
||||
gateway.scenario() as scenario,
|
||||
):
|
||||
server_id: Final = register_mcp(scenario, peer, alias)
|
||||
|
|
@ -1048,6 +1129,262 @@ def test_bridged_next_turn_replays_the_mcp_turn_answer(gateway: Gateway) -> None
|
|||
assert ANSWER in _assistant_texts(opening), opening
|
||||
|
||||
|
||||
StreamedUpstream = Literal["messages_bridge", "responses_bridge"]
|
||||
STREAMED_UPSTREAMS: Final[tuple[StreamedUpstream, ...]] = ("messages_bridge", "responses_bridge")
|
||||
STREAMED_IDS: Final = ("native", "bridged")
|
||||
SLOW_TOOL_SECONDS: Final = 5.0
|
||||
|
||||
|
||||
def _events(response: httpx.Response) -> tuple[dict[str, JsonValue], ...]:
|
||||
return tuple(
|
||||
JSON_OBJECT.validate_json(line.removeprefix("data: "))
|
||||
for line in response.text.splitlines()
|
||||
if line.startswith("data: ") and line != "data: [DONE]"
|
||||
)
|
||||
|
||||
|
||||
def _completed(response: httpx.Response) -> dict[str, JsonValue]:
|
||||
finished: Final = tuple(
|
||||
object_value(event["response"]) for event in _events(response) if event["type"] == "response.completed"
|
||||
)
|
||||
assert len(finished) == 1, response.text
|
||||
return finished[0]
|
||||
|
||||
|
||||
def _model_rows(call_id: str, settled: Callable[[list[dict[str, JsonValue]]], bool]) -> list[dict[str, JsonValue]]:
|
||||
return eventually(
|
||||
lambda: read_rows(
|
||||
"SELECT request_id, status, spend, prompt_tokens, completion_tokens "
|
||||
'FROM "LiteLLM_SpendLogs" WHERE litellm_call_id=%s AND call_type NOT IN (%s, %s)',
|
||||
(call_id, "call_mcp_tool", "list_mcp_tools"),
|
||||
),
|
||||
settled,
|
||||
seconds=70,
|
||||
return_last_on_timeout=True,
|
||||
)
|
||||
|
||||
|
||||
def _billed(rows: Sequence[Mapping[str, JsonValue]]) -> float:
|
||||
return sum(float(str(row["spend"])) for row in rows)
|
||||
|
||||
|
||||
def _bills_at_least(expected: float) -> Callable[[list[dict[str, JsonValue]]], bool]:
|
||||
return lambda rows: bool(rows) and _billed(rows) >= expected - 1e-9
|
||||
|
||||
|
||||
def _statuses(rows: Sequence[Mapping[str, JsonValue]]) -> list[str]:
|
||||
return sorted(str(row["status"]) for row in rows)
|
||||
|
||||
|
||||
def _logs_a_failure_billing(expected: float) -> Callable[[list[dict[str, JsonValue]]], bool]:
|
||||
return lambda rows: "failure" in _statuses(rows) and _billed(rows) >= expected - 1e-9
|
||||
|
||||
|
||||
def _limited_key(rig: Rig) -> str:
|
||||
return rig.scenario.key(object_permission={"mcp_servers": [rig.server_id]}, max_parallel_requests=1)
|
||||
|
||||
|
||||
def _answers_directly(rig: Rig, key: str) -> int:
|
||||
return rig.send(key, [AUTO], prompt=f"{DIRECT} {rig.alias}").status_code
|
||||
|
||||
|
||||
STREAMED_ROUNDS: Final = (
|
||||
pytest.param("messages_bridge", 1, id="native"),
|
||||
pytest.param("responses_bridge", 1, id="bridged"),
|
||||
pytest.param("messages_bridge", 2, id="native-two-tool-rounds"),
|
||||
pytest.param("responses_bridge", 2, id="bridged-two-tool-rounds"),
|
||||
pytest.param("messages_bridge", 0, id="native-no-tool-call"),
|
||||
pytest.param("responses_bridge", 0, id="bridged-no-tool-call"),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("upstream", "tool_rounds"), STREAMED_ROUNDS)
|
||||
def test_streaming_responses_auto_execute_bills_every_round_once_on_one_success_row(
|
||||
gateway: Gateway, upstream: StreamedUpstream, tool_rounds: int
|
||||
) -> None:
|
||||
rounds: Final = (FIRST_ROUND, *(FOLLOW_UP_ROUND,) * tool_rounds)
|
||||
billing: Final = Billing(FIRST_ROUND, FOLLOW_UP_ROUND, tool_rounds=max(tool_rounds, 1))
|
||||
with _rig(gateway, "responses", upstream=upstream, billing=billing, prices=PRICES) as rig:
|
||||
prompt: Final = None if tool_rounds else f"{DIRECT} {rig.alias}"
|
||||
response: Final = rig.send(_granted_key(rig), [AUTO], prompt=prompt, stream=True)
|
||||
assert response.status_code == 200, response.text
|
||||
completed: Final = _completed(response)
|
||||
assert rig.final_text(completed) == ANSWER, response.text
|
||||
assert len(rig.upstream_tools()) == len(rounds)
|
||||
assert len(_peer_add_calls(rig.peer)) == tool_rounds
|
||||
expected: Final = _price(*rounds) + FIXED_MARGIN
|
||||
rows: Final = _model_rows(response.headers["x-litellm-call-id"], _bills_at_least(expected))
|
||||
assert [row["status"] for row in rows] == ["success"], rows
|
||||
row: Final = rows[0]
|
||||
assert (row["prompt_tokens"], row["completion_tokens"]) == (
|
||||
sum(prompt_tokens for prompt_tokens, _ in rounds),
|
||||
sum(completion_tokens for _, completion_tokens in rounds),
|
||||
), row
|
||||
usage: Final = object_value(completed["usage"])
|
||||
assert (usage["input_tokens"], usage["output_tokens"]) == (row["prompt_tokens"], row["completion_tokens"])
|
||||
assert float(str(row["spend"])) == pytest.approx(expected), row
|
||||
|
||||
|
||||
@pytest.mark.parametrize("upstream", STREAMED_UPSTREAMS, ids=STREAMED_IDS)
|
||||
def test_streaming_responses_rejected_follow_up_bills_round_one_once(
|
||||
gateway: Gateway, upstream: StreamedUpstream
|
||||
) -> None:
|
||||
billing: Final = Billing(FIRST_ROUND, FOLLOW_UP_ROUND, follow_up_fails=True)
|
||||
with _rig(gateway, "responses", upstream=upstream, billing=billing, prices=PRICES) as rig:
|
||||
response: Final = rig.send(_granted_key(rig), [AUTO], stream=True)
|
||||
assert FOLLOW_UP_OUTAGE in response.text, response.text
|
||||
assert ANSWER not in response.text, response.text
|
||||
assert len(rig.upstream_tools()) == 2
|
||||
expected: Final = _price(FIRST_ROUND) + FIXED_MARGIN
|
||||
rows: Final = _model_rows(response.headers["x-litellm-call-id"], _bills_at_least(expected))
|
||||
assert _billed(rows) == pytest.approx(expected), rows
|
||||
|
||||
|
||||
BROKEN_FOLLOW_UPS: Final = (
|
||||
pytest.param("messages_bridge", FOLLOW_UP_ROUND, id="native"),
|
||||
pytest.param("responses_bridge", None, id="bridged"),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("upstream", "reported_usage"), BROKEN_FOLLOW_UPS)
|
||||
def test_streaming_responses_follow_up_broken_mid_stream_bills_round_one_once_and_itself_on_a_failure_row(
|
||||
gateway: Gateway, upstream: StreamedUpstream, reported_usage: tuple[int, int] | None
|
||||
) -> None:
|
||||
billing: Final = Billing(FIRST_ROUND, FOLLOW_UP_ROUND, follow_up_breaks=True)
|
||||
with _rig(gateway, "responses", upstream=upstream, billing=billing, prices=PRICES) as rig:
|
||||
response: Final = rig.send(_granted_key(rig), [AUTO], stream=True)
|
||||
assert ANSWER not in response.text, response.text
|
||||
assert len(rig.upstream_tools()) == 2
|
||||
round_one: Final = _price(FIRST_ROUND) + FIXED_MARGIN
|
||||
rows: Final = _model_rows(response.headers["x-litellm-call-id"], _logs_a_failure_billing(round_one))
|
||||
assert _statuses(rows) == ["failure", "success"], rows
|
||||
(success,) = (row for row in rows if row["status"] == "success")
|
||||
(failure,) = (row for row in rows if row["status"] == "failure")
|
||||
assert (success["prompt_tokens"], success["completion_tokens"]) == FIRST_ROUND, rows
|
||||
assert float(str(success["spend"])) == pytest.approx(round_one), rows
|
||||
failed_round: Final = (int(str(failure["prompt_tokens"])), int(str(failure["completion_tokens"])))
|
||||
assert min(failed_round) > 0, rows
|
||||
assert reported_usage in (None, failed_round), rows
|
||||
assert float(str(failure["spend"])) == pytest.approx(_price(failed_round)), rows
|
||||
|
||||
|
||||
@pytest.mark.parametrize("upstream", STREAMED_UPSTREAMS, ids=STREAMED_IDS)
|
||||
def test_streaming_responses_blocked_by_a_post_call_guardrail_bills_every_round_once(
|
||||
gateway: Gateway, upstream: StreamedUpstream
|
||||
) -> None:
|
||||
billing: Final = Billing(FIRST_ROUND, FOLLOW_UP_ROUND)
|
||||
with _rig(gateway, "responses", upstream=upstream, billing=billing, prices=PRICES) as rig:
|
||||
key: Final = _granted_key(rig)
|
||||
one_round: Final = rig.send(key, [], prompt=f"{DIRECT} {rig.alias}", stream=True, guardrails=[STREAM_BLOCK])
|
||||
assert STREAM_BLOCK_DETAIL in one_round.text, one_round.text
|
||||
one_round_rows: Final = _model_rows(
|
||||
one_round.headers["x-litellm-call-id"], _logs_a_failure_billing(_price(FIRST_ROUND) + FIXED_MARGIN)
|
||||
)
|
||||
rig.wire.drain()
|
||||
response: Final = rig.send(key, [AUTO], stream=True, guardrails=[STREAM_BLOCK])
|
||||
assert STREAM_BLOCK_DETAIL in response.text, response.text
|
||||
assert len(rig.upstream_tools()) == 2 and len(_peer_add_calls(rig.peer)) == 1
|
||||
expected: Final = _price(FIRST_ROUND, FOLLOW_UP_ROUND) + FIXED_MARGIN
|
||||
rows: Final = _model_rows(response.headers["x-litellm-call-id"], _logs_a_failure_billing(expected))
|
||||
assert _statuses(rows) == _statuses(one_round_rows), (rows, one_round_rows)
|
||||
assert _billed(rows) == pytest.approx(expected), rows
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("surface", "stream"),
|
||||
(("responses", True), ("responses", False), ("chat", False)),
|
||||
ids=("responses-stream", "responses", "chat"),
|
||||
)
|
||||
def test_auto_executed_tool_call_runs_inside_the_requests_one_parallel_slot(
|
||||
gateway: Gateway, surface: Surface, stream: bool
|
||||
) -> None:
|
||||
with _rig(gateway, surface, slow_tool_seconds=SLOW_TOOL_SECONDS) as rig:
|
||||
key: Final = _limited_key(rig)
|
||||
with ThreadPoolExecutor(max_workers=1) as pool:
|
||||
running: Final = pool.submit(rig.send, key, [AUTO], None, stream=stream)
|
||||
eventually(lambda: _peer_add_calls(rig.peer), lambda found: len(found) == 1, seconds=30)
|
||||
during: Final = _answers_directly(rig, key)
|
||||
served: Final = running.result(timeout=90)
|
||||
assert during == 429
|
||||
assert served.status_code == 200 and ANSWER in served.text, served.text
|
||||
eventually(lambda: _answers_directly(rig, key), lambda status: status == 200, seconds=30)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("upstream", STREAMED_UPSTREAMS, ids=STREAMED_IDS)
|
||||
def test_streaming_responses_hung_up_during_the_tool_call_bills_each_served_round_once_and_frees_the_slot(
|
||||
gateway: Gateway, upstream: StreamedUpstream
|
||||
) -> None:
|
||||
billing: Final = Billing(FIRST_ROUND, FOLLOW_UP_ROUND)
|
||||
with _rig(
|
||||
gateway,
|
||||
"responses",
|
||||
upstream=upstream,
|
||||
billing=billing,
|
||||
prices=PRICES,
|
||||
slow_tool_seconds=SLOW_TOOL_SECONDS,
|
||||
) as rig:
|
||||
key: Final = _limited_key(rig)
|
||||
with rig.gateway.client.stream(
|
||||
"POST",
|
||||
"/v1/responses",
|
||||
headers={"Authorization": f"Bearer {key}"},
|
||||
json={"model": rig.model, "input": f"add {rig.alias}", "tools": [AUTO], "stream": True},
|
||||
timeout=90,
|
||||
) as hung_up:
|
||||
assert hung_up.status_code == 200
|
||||
call_id: Final = hung_up.headers["x-litellm-call-id"]
|
||||
eventually(lambda: _peer_add_calls(rig.peer), lambda found: len(found) == 1, seconds=30)
|
||||
eventually(lambda: _answers_directly(rig, key), lambda status: status == 200, seconds=60)
|
||||
posted: Final = tuple(
|
||||
JSON_OBJECT.validate_json(request.body) for request in rig.wire.drain() if request.method == "POST"
|
||||
)
|
||||
served: Final = tuple(
|
||||
FOLLOW_UP_ROUND if _has_tool_result(body) else FIRST_ROUND for body in posted if DIRECT not in _prompt(body)
|
||||
)
|
||||
expected: Final = _price(*served) + FIXED_MARGIN
|
||||
rows: Final = _model_rows(call_id, _bills_at_least(expected))
|
||||
assert served and rows, (served, rows)
|
||||
assert _billed(rows) == pytest.approx(expected), (served, rows)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("upstream", STREAMED_UPSTREAMS, ids=STREAMED_IDS)
|
||||
def test_streaming_responses_repeated_from_the_response_cache_bills_nothing(
|
||||
gateway: Gateway, upstream: StreamedUpstream
|
||||
) -> None:
|
||||
billing: Final = Billing(FIRST_ROUND, FOLLOW_UP_ROUND)
|
||||
with _rig(gateway, "responses", upstream=upstream, billing=billing, prices=PRICES) as rig:
|
||||
key: Final = _granted_key(rig)
|
||||
paid: Final = rig.send(key, [AUTO], stream=True)
|
||||
assert paid.status_code == 200 and rig.final_text(_completed(paid)) == ANSWER, paid.text
|
||||
paid_price: Final = _price(FIRST_ROUND, FOLLOW_UP_ROUND) + FIXED_MARGIN
|
||||
_model_rows(paid.headers["x-litellm-call-id"], _bills_at_least(paid_price))
|
||||
assert len(rig.upstream_tools()) == 2
|
||||
cached: Final = rig.send(key, [AUTO], stream=True)
|
||||
assert cached.status_code == 200 and rig.final_text(_completed(cached)) == ANSWER, cached.text
|
||||
assert rig.upstream_tools() == ()
|
||||
rows: Final = _model_rows(
|
||||
cached.headers["x-litellm-call-id"], lambda found: "success" in {row["status"] for row in found}
|
||||
)
|
||||
assert "success" in {row["status"] for row in rows} and _billed(rows) == 0.0, rows
|
||||
|
||||
|
||||
def test_bridged_next_turn_after_a_streamed_mcp_turn_replays_its_answer(gateway: Gateway) -> None:
|
||||
billing: Final = Billing(FIRST_ROUND, FOLLOW_UP_ROUND)
|
||||
with _rig(gateway, "responses", upstream="responses_bridge", billing=billing, prices=PRICES) as rig:
|
||||
key: Final = _granted_key(rig)
|
||||
first: Final = rig.send(key, [AUTO], stream=True)
|
||||
assert first.status_code == 200, first.text
|
||||
first_id: Final = string_value(_completed(first)["id"])
|
||||
_model_rows(first.headers["x-litellm-call-id"], _bills_at_least(_price(FIRST_ROUND, FOLLOW_UP_ROUND)))
|
||||
rig.wire.drain()
|
||||
second: Final = rig.send(key, [AUTO], previous_response_id=first_id, cache={"no-cache": True})
|
||||
assert second.status_code == 200, second.text
|
||||
opening: Final = next(
|
||||
JSON_OBJECT.validate_json(request.body) for request in rig.wire.drain() if request.method == "POST"
|
||||
)
|
||||
assert ANSWER in _assistant_texts(opening), opening
|
||||
|
||||
|
||||
Bridge = Literal["chat", "responses", "messages"]
|
||||
BRIDGES: Final[tuple[Bridge, ...]] = ("chat", "responses")
|
||||
Client = Literal["sync", "async"]
|
||||
|
|
|
|||
|
|
@ -0,0 +1,418 @@
|
|||
import json
|
||||
import math
|
||||
import re
|
||||
import uuid
|
||||
from collections.abc import Iterator, Mapping
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Final, Literal
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import yaml
|
||||
from integration._support.client import JSON_OBJECT, Gateway, eventually, object_value
|
||||
from integration._support.database import read_rows
|
||||
from integration._support.process import OwnedProxy, WorkerHealthcheck, graceful_stop_seconds, owned_proxy_process
|
||||
from integration._support.wire import Reply, Request, Wire, wire_server
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
|
||||
Route = Literal["/v1/extraction", "/extraction"]
|
||||
|
||||
_API_KEY: Final = "synthetic-cloudflare-key"
|
||||
_CLAIRIC: Final = "cloudflare/@cf/cloudflare/clairic"
|
||||
_CLAIRIC_PATH: Final = "/ai/run/@cf/cloudflare/clairic"
|
||||
_INPUT_TOKENS: Final = 412
|
||||
_EMAIL: Final = "billing@acmetools.example"
|
||||
_SCENARIO: Final = re.compile(r"scenario-([a-z-]+)")
|
||||
_JSON_ROWS: Final = TypeAdapter(list[dict[str, JsonValue]])
|
||||
_FIELDS: Final[dict[str, JsonValue]] = {
|
||||
"vendor": {"type": "text", "instructions": "Who sent the invoice?"},
|
||||
"total": {"type": "number", "unit": "USD"},
|
||||
}
|
||||
_SPEND_QUERY: Final = (
|
||||
'SELECT spend, status, call_type, model_group, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" '
|
||||
"WHERE request_id = %s"
|
||||
)
|
||||
_OWNED_TIMEOUT: Final = int(2 * graceful_stop_seconds() + 120)
|
||||
_REQUEST_GUARDRAIL: Final = "extraction-mask-email-in"
|
||||
_RESPONSE_GUARDRAIL: Final = "extraction-mask-email-out"
|
||||
_OWNED_MODEL: Final = "invoice-reader"
|
||||
|
||||
|
||||
def _state(marker: str, scenario: str = "plain") -> str:
|
||||
return f"Invoice from Acme Tools, total $174. scenario-{scenario} marker-{marker}"
|
||||
|
||||
|
||||
def _vendor_span(state: str) -> list[JsonValue]:
|
||||
start: Final = state.index("Acme Tools")
|
||||
return [start, start + len("Acme Tools")]
|
||||
|
||||
|
||||
def _email_span(state: str) -> list[JsonValue]:
|
||||
start: Final = state.index(_EMAIL)
|
||||
return [start, start + len(_EMAIL)]
|
||||
|
||||
|
||||
def _answers(state: str) -> dict[str, JsonValue]:
|
||||
return {
|
||||
"vendor": {"type": "text", "value": "Acme Tools", "confidence": 0.97, "span": _vendor_span(state)},
|
||||
"total": {"type": "number", "value": 174, "confidence": 0.95},
|
||||
}
|
||||
|
||||
|
||||
def _envelope(result: Mapping[str, JsonValue]) -> Reply:
|
||||
return Reply(body=json.dumps({"result": result, "success": True, "errors": [], "messages": []}).encode())
|
||||
|
||||
|
||||
def _clairic_reply(request: Request) -> Reply:
|
||||
state: Final = JSON_OBJECT.validate_json(request.body)["state"]
|
||||
assert isinstance(state, str), request.body
|
||||
return _envelope(
|
||||
{
|
||||
"model": "clairic",
|
||||
"answers": _answers(state),
|
||||
"truncated": False,
|
||||
"usage": {"input_tokens": _INPUT_TOKENS, "output_tokens": 0},
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _expected_response(model: str, state: str) -> dict[str, JsonValue]:
|
||||
return {
|
||||
"model": model,
|
||||
"answers": _answers(state),
|
||||
"truncated": False,
|
||||
"usage": {"input_tokens": _INPUT_TOKENS, "output_tokens": 0},
|
||||
}
|
||||
|
||||
|
||||
def _upstream_body(request: Request) -> dict[str, JsonValue]:
|
||||
assert (request.method, request.target) == ("POST", _CLAIRIC_PATH), (request.method, request.target)
|
||||
assert request.headers["authorization"] == f"Bearer {_API_KEY}", request.headers
|
||||
return JSON_OBJECT.validate_json(request.body)
|
||||
|
||||
|
||||
def _body(model: str, marker: str, **extra: JsonValue) -> dict[str, JsonValue]:
|
||||
return {"model": model, "state": _state(marker), "fields": _FIELDS, **extra}
|
||||
|
||||
|
||||
def _extract(
|
||||
gateway: Gateway,
|
||||
model: str,
|
||||
marker: str,
|
||||
*,
|
||||
route: Route = "/v1/extraction",
|
||||
key: str | None = None,
|
||||
**extra: JsonValue,
|
||||
) -> httpx.Response:
|
||||
return gateway.request("POST", route, _body(model, marker, **extra), key=key)
|
||||
|
||||
|
||||
def _number(value: JsonValue) -> float:
|
||||
assert isinstance(value, (int, float)) and not isinstance(value, bool), value
|
||||
return float(value)
|
||||
|
||||
|
||||
def _expected_spend() -> float:
|
||||
cost_map: Final = JSON_OBJECT.validate_json(Path("model_prices_and_context_window.json").read_bytes())
|
||||
priced: Final = cost_map.get(_CLAIRIC)
|
||||
assert priced is not None, f"{_CLAIRIC} has no cost map entry"
|
||||
return _INPUT_TOKENS * _number(object_value(priced)["input_cost_per_token"])
|
||||
|
||||
|
||||
def _spend_row(call_id: str) -> dict[str, JsonValue]:
|
||||
rows: Final = eventually(lambda: read_rows(_SPEND_QUERY, (call_id,)), lambda found: len(found) == 1, seconds=70)
|
||||
return rows[0]
|
||||
|
||||
|
||||
def _assert_billed(gateway: Gateway, response: httpx.Response, model: str) -> None:
|
||||
expected: Final = _expected_spend()
|
||||
assert math.isclose(float(response.headers["x-litellm-response-cost"]), expected, rel_tol=1e-9), response.headers
|
||||
call_id: Final = response.headers["x-litellm-call-id"]
|
||||
row: Final = _spend_row(call_id)
|
||||
assert (row["status"], row["call_type"], row["model_group"], row["prompt_tokens"], row["completion_tokens"]) == (
|
||||
"success",
|
||||
"aextraction",
|
||||
model,
|
||||
_INPUT_TOKENS,
|
||||
0,
|
||||
), row
|
||||
assert math.isclose(_number(row["spend"]), expected, rel_tol=1e-9), row
|
||||
listed: Final = gateway.request("GET", "/spend/logs", params={"request_id": call_id})
|
||||
assert listed.status_code == 200, listed.text
|
||||
(reported,) = _JSON_ROWS.validate_json(listed.content)
|
||||
assert (reported["request_id"], reported["model_group"]) == (call_id, model), reported
|
||||
assert math.isclose(_number(reported["spend"]), expected, rel_tol=1e-9), reported
|
||||
|
||||
|
||||
@pytest.mark.parametrize("route", ("/v1/extraction", "/extraction"))
|
||||
def test_extraction_reaches_clairic_on_both_routes_and_is_billed(gateway: Gateway, route: Route) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
with wire_server(_clairic_reply) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=_CLAIRIC, api_base=wire.url, api_key=_API_KEY)
|
||||
response: Final = _extract(gateway, model, marker, route=route, metadata={"tags": [f"audit-{marker}"]})
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.json() == _expected_response(model, _state(marker))
|
||||
(request,) = wire.drain()
|
||||
assert _upstream_body(request) == {"model": "clairic", "state": _state(marker), "fields": _FIELDS}
|
||||
_assert_billed(gateway, response, model)
|
||||
|
||||
|
||||
_INVALID_BODIES: Final[tuple[tuple[str, dict[str, JsonValue]], ...]] = (
|
||||
("unknown_type", {"fields": {"vendor": {"type": "paragraph"}}}),
|
||||
("list_without_items", {"fields": {"items": {"type": "list"}}}),
|
||||
("remote_image", {"images": ["https://images.example/receipt.png"]}),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("label", "change"), _INVALID_BODIES, ids=[label for label, _ in _INVALID_BODIES])
|
||||
def test_invalid_extraction_bodies_are_refused_before_any_upstream_call(
|
||||
gateway: Gateway, label: str, change: dict[str, JsonValue]
|
||||
) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
with wire_server(_clairic_reply) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=_CLAIRIC, api_base=wire.url, api_key=_API_KEY)
|
||||
response: Final = gateway.request("POST", "/v1/extraction", {**_body(model, marker), **change})
|
||||
assert response.status_code == 400, (label, response.text)
|
||||
assert "Invalid Extraction request" in response.text, (label, response.text)
|
||||
assert wire.drain() == (), label
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("param", "value"), (("temperature", 0.2), ("stream", True)), ids=("temperature", "stream"))
|
||||
def test_unsupported_top_level_params_are_refused_unless_the_deployment_drops_params(
|
||||
gateway: Gateway, param: str, value: JsonValue
|
||||
) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
with wire_server(_clairic_reply) as wire, gateway.scenario() as scenario:
|
||||
strict: Final = scenario.model(model=_CLAIRIC, api_base=wire.url, api_key=_API_KEY)
|
||||
dropping: Final = scenario.model(model=_CLAIRIC, api_base=wire.url, api_key=_API_KEY, drop_params=True)
|
||||
refused: Final = _extract(gateway, strict, marker, **{param: value})
|
||||
assert refused.status_code == 400, refused.text
|
||||
assert param in refused.text, refused.text
|
||||
assert wire.drain() == ()
|
||||
dropped: Final = _extract(gateway, dropping, marker, **{param: value})
|
||||
assert dropped.status_code == 200, dropped.text
|
||||
assert dropped.headers["content-type"].startswith("application/json"), dropped.headers
|
||||
assert dropped.json() == _expected_response(dropping, _state(marker))
|
||||
(request,) = wire.drain()
|
||||
assert _upstream_body(request) == {"model": "clairic", "state": _state(marker), "fields": _FIELDS}
|
||||
_assert_billed(gateway, dropped, dropping)
|
||||
|
||||
|
||||
def test_key_checks_on_extraction_match_chat(gateway: Gateway) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
with wire_server(_clairic_reply) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=_CLAIRIC, api_base=wire.url, api_key=_API_KEY)
|
||||
anonymous: Final = gateway.client.post("/v1/extraction", json=_body(model, marker))
|
||||
assert anonymous.status_code == 401, anonymous.text
|
||||
restricted: Final = scenario.key(models=[f"other-{uuid.uuid4().hex}"])
|
||||
refused: Final = _extract(gateway, model, marker, key=restricted)
|
||||
chat_refused: Final = gateway.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{"model": model, "messages": [{"role": "user", "content": marker}]},
|
||||
key=restricted,
|
||||
)
|
||||
assert 400 <= refused.status_code < 500, refused.text
|
||||
assert refused.status_code == chat_refused.status_code, (refused.text, chat_refused.text)
|
||||
assert wire.drain() == ()
|
||||
allowed: Final = scenario.key(models=[model], max_budget=1e-06)
|
||||
first: Final = _extract(gateway, model, marker, key=allowed)
|
||||
assert first.status_code == 200, first.text
|
||||
assert len(wire.drain()) == 1
|
||||
blocked, upstream_calls = eventually(
|
||||
lambda: (_extract(gateway, model, uuid.uuid4().hex, key=allowed), len(wire.drain())),
|
||||
lambda answer: answer[0].status_code != 200,
|
||||
seconds=70,
|
||||
)
|
||||
chat_blocked: Final = gateway.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{"model": model, "messages": [{"role": "user", "content": marker}]},
|
||||
key=allowed,
|
||||
)
|
||||
assert 400 <= blocked.status_code < 500, blocked.text
|
||||
assert blocked.status_code == chat_blocked.status_code, (blocked.text, chat_blocked.text)
|
||||
assert upstream_calls == 0
|
||||
assert wire.drain() == ()
|
||||
|
||||
|
||||
def test_a_small_key_tpm_limit_serves_extraction_and_counts_its_input_tokens(gateway: Gateway) -> None:
|
||||
with wire_server(_clairic_reply) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=_CLAIRIC, api_base=wire.url, api_key=_API_KEY)
|
||||
key: Final = scenario.key(models=[model], tpm_limit=100)
|
||||
first: Final = _extract(gateway, model, uuid.uuid4().hex, key=key)
|
||||
assert first.status_code == 200, first.text
|
||||
assert len(wire.drain()) == 1
|
||||
limited, upstream_calls = eventually(
|
||||
lambda: (_extract(gateway, model, uuid.uuid4().hex, key=key), len(wire.drain())),
|
||||
lambda answer: answer[0].status_code != 200,
|
||||
seconds=70,
|
||||
)
|
||||
assert limited.status_code == 429, limited.text
|
||||
assert upstream_calls == 0
|
||||
|
||||
|
||||
def test_cloudflare_errors_keep_their_status_and_log_an_unbilled_failure(gateway: Gateway) -> None:
|
||||
marker: Final = uuid.uuid4().hex
|
||||
refusal: Final = json.dumps(
|
||||
{
|
||||
"result": None,
|
||||
"success": False,
|
||||
"errors": [{"code": 5006, "message": f"min above max {marker}"}],
|
||||
"messages": [],
|
||||
}
|
||||
).encode()
|
||||
with wire_server(lambda _: Reply(status=422, body=refusal)) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(model=_CLAIRIC, api_base=wire.url, api_key=_API_KEY)
|
||||
response: Final = _extract(gateway, model, marker)
|
||||
assert response.status_code == 422, response.text
|
||||
assert f"min above max {marker}" in response.text, response.text
|
||||
assert len(wire.drain()) >= 1
|
||||
row: Final = _spend_row(response.headers["x-litellm-call-id"])
|
||||
assert (row["status"], row["call_type"], row["model_group"], _number(row["spend"])) == (
|
||||
"failure",
|
||||
"aextraction",
|
||||
model,
|
||||
0.0,
|
||||
), row
|
||||
|
||||
|
||||
def _owned_reply(request: Request) -> Reply:
|
||||
body: Final = JSON_OBJECT.validate_json(request.body)
|
||||
state: Final = body["state"]
|
||||
assert isinstance(state, str), request.body
|
||||
scenario: Final = _SCENARIO.search(state)
|
||||
assert scenario is not None, state
|
||||
if scenario.group(1) != "mask-out":
|
||||
return _clairic_reply(request)
|
||||
return _envelope(
|
||||
{
|
||||
"model": "clairic",
|
||||
"answers": {
|
||||
"vendor": {"type": "text", "value": "Acme Tools", "confidence": 0.97, "span": _vendor_span(state)},
|
||||
"contact_email": {"type": "text", "value": _EMAIL, "confidence": 0.9, "span": _email_span(state)},
|
||||
},
|
||||
"truncated": False,
|
||||
"state_text": state,
|
||||
"usage": {"input_tokens": _INPUT_TOKENS, "output_tokens": 0},
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _owned_config(upstream_url: str) -> dict[str, JsonValue]:
|
||||
base: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()))
|
||||
email_filter: Final[list[JsonValue]] = [{"pattern_type": "prebuilt", "pattern_name": "email", "action": "MASK"}]
|
||||
return {
|
||||
**base,
|
||||
"guardrails": [
|
||||
{
|
||||
"guardrail_name": _REQUEST_GUARDRAIL,
|
||||
"litellm_params": {"guardrail": "litellm_content_filter", "mode": "pre_call", "patterns": email_filter},
|
||||
},
|
||||
{
|
||||
"guardrail_name": _RESPONSE_GUARDRAIL,
|
||||
"litellm_params": {
|
||||
"guardrail": "litellm_content_filter",
|
||||
"mode": "post_call",
|
||||
"patterns": email_filter,
|
||||
},
|
||||
},
|
||||
],
|
||||
"model_list": [
|
||||
{
|
||||
"model_name": _OWNED_MODEL,
|
||||
"litellm_params": {"model": _CLAIRIC, "api_base": upstream_url, "api_key": _API_KEY},
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Rig:
|
||||
proxy: OwnedProxy
|
||||
upstream: Wire
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]:
|
||||
with wire_server(_owned_reply) as upstream:
|
||||
directory: Final = tmp_path_factory.mktemp("clairic-extraction-owned")
|
||||
config_path: Final = directory / "clairic-extraction.yaml"
|
||||
config_path.write_text(yaml.safe_dump(_owned_config(upstream.url)))
|
||||
with owned_proxy_process(
|
||||
session_gateway,
|
||||
directory,
|
||||
{},
|
||||
config=config_path,
|
||||
workers=2,
|
||||
extra_arguments=("--max_tokens", "256", "--temperature", "0.2"),
|
||||
healthcheck=WorkerHealthcheck.LENIENT,
|
||||
) as proxy:
|
||||
yield Rig(proxy, upstream)
|
||||
|
||||
|
||||
def _owned_extract(rig: Rig, state: str, **extra: JsonValue) -> httpx.Response:
|
||||
return rig.proxy.gateway.request(
|
||||
"POST", "/v1/extraction", {"model": _OWNED_MODEL, "state": state, "fields": _FIELDS, **extra}
|
||||
)
|
||||
|
||||
|
||||
def _applied_guardrails(response: httpx.Response) -> frozenset[str]:
|
||||
return frozenset(name.strip() for name in response.headers.get("x-litellm-applied-guardrails", "").split(","))
|
||||
|
||||
|
||||
@pytest.mark.timeout(_OWNED_TIMEOUT)
|
||||
def test_cli_generation_defaults_never_reach_an_extraction_call(rig: Rig) -> None:
|
||||
state: Final = _state(uuid.uuid4().hex)
|
||||
response: Final = _owned_extract(rig, state)
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.json() == _expected_response(_OWNED_MODEL, state)
|
||||
(request,) = rig.upstream.drain()
|
||||
assert _upstream_body(request) == {"model": "clairic", "state": state, "fields": _FIELDS}
|
||||
|
||||
|
||||
@pytest.mark.timeout(_OWNED_TIMEOUT)
|
||||
def test_a_request_guardrail_masks_the_state_before_clairic_and_drops_spans_into_it(rig: Rig) -> None:
|
||||
state: Final = f"{_state(uuid.uuid4().hex, 'mask-in')} Questions go to {_EMAIL}."
|
||||
response: Final = _owned_extract(rig, state, guardrails=[_REQUEST_GUARDRAIL])
|
||||
assert response.status_code == 200, response.text
|
||||
assert _REQUEST_GUARDRAIL in _applied_guardrails(response), response.headers
|
||||
(request,) = rig.upstream.drain()
|
||||
sent: Final = _upstream_body(request)
|
||||
sent_state: Final = sent["state"]
|
||||
assert isinstance(sent_state, str) and _EMAIL not in sent_state, sent
|
||||
assert sent_state.startswith(state.removesuffix(f"{_EMAIL}.")) and "REDACTED" in sent_state, sent
|
||||
assert sent == {"model": "clairic", "state": sent_state, "fields": _FIELDS}
|
||||
answers: Final = object_value(response.json()["answers"])
|
||||
assert object_value(answers["vendor"]) == {"type": "text", "value": "Acme Tools", "confidence": 0.97}, answers
|
||||
|
||||
|
||||
@pytest.mark.timeout(_OWNED_TIMEOUT)
|
||||
def test_a_response_guardrail_masks_text_answers_and_state_text(rig: Rig) -> None:
|
||||
state: Final = f"{_state(uuid.uuid4().hex, 'mask-out')} Questions go to {_EMAIL}."
|
||||
fields: Final[dict[str, JsonValue]] = {**_FIELDS, "contact_email": {"type": "text"}}
|
||||
response: Final = rig.proxy.gateway.request(
|
||||
"POST",
|
||||
"/v1/extraction",
|
||||
{
|
||||
"model": _OWNED_MODEL,
|
||||
"state": state,
|
||||
"fields": fields,
|
||||
"guardrails": [_RESPONSE_GUARDRAIL],
|
||||
"options": {"return_state_text": True},
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
assert _RESPONSE_GUARDRAIL in _applied_guardrails(response), response.headers
|
||||
(request,) = rig.upstream.drain()
|
||||
assert _upstream_body(request)["state"] == state
|
||||
body: Final = response.json()
|
||||
state_text: Final = body["state_text"]
|
||||
assert isinstance(state_text, str) and _EMAIL not in state_text and "REDACTED" in state_text, body
|
||||
answers: Final = object_value(body["answers"])
|
||||
contact: Final = object_value(answers["contact_email"])
|
||||
assert "span" not in contact, answers
|
||||
masked: Final = contact["value"]
|
||||
assert isinstance(masked, str) and _EMAIL not in masked and "REDACTED" in masked, answers
|
||||
assert object_value(answers["vendor"]) == {"type": "text", "value": "Acme Tools", "confidence": 0.97}, answers
|
||||
|
|
@ -1,16 +1,113 @@
|
|||
import json
|
||||
import os
|
||||
import uuid
|
||||
from itertools import product
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import yaml
|
||||
from integration._support.client import Gateway, eventually, object_value
|
||||
from integration._support.process import graceful_stop_seconds, owned_proxy
|
||||
from integration._support.redis_process import owned_redis
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
from redis import Redis
|
||||
|
||||
import litellm
|
||||
from litellm.caching.in_memory_cache import InMemoryCache
|
||||
from litellm.caching.redis_cache import RedisCache
|
||||
from litellm.constants import GLOBAL_PROXY_SPEND_CACHE_KEY
|
||||
from litellm.constants import GLOBAL_PROXY_SPEND_CACHE_KEY, LITELLM_PROXY_BUDGET_NAME
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
from litellm.proxy.proxy_server import update_cache
|
||||
|
||||
BUDGET_MODEL: Final = "proxy-budget-model"
|
||||
CONFIG: Final = TypeAdapter(dict[str, JsonValue])
|
||||
SECTION: Final = TypeAdapter(dict[str, JsonValue])
|
||||
OWNED_PROXY_SECONDS: Final = 2 * graceful_stop_seconds() + 120
|
||||
|
||||
|
||||
def _proxy_budget_config(gateway: Gateway, directory: Path, max_budget: float) -> Path:
|
||||
loaded: Final = CONFIG.validate_python(
|
||||
yaml.safe_load((Path(__file__).resolve().parents[1] / "proxy_config.yaml").read_text())
|
||||
)
|
||||
model: Final[JsonValue] = {
|
||||
"model_name": BUDGET_MODEL,
|
||||
"litellm_params": {
|
||||
"model": "openai/gpt-4o-mini",
|
||||
"api_key": "integration-provider-key",
|
||||
"api_base": f"{gateway.upstream_url}/v1",
|
||||
"input_cost_per_token": 0.001,
|
||||
"output_cost_per_token": 0.002,
|
||||
},
|
||||
}
|
||||
path: Final = directory / "proxy_budget_config.yaml"
|
||||
path.write_text(
|
||||
yaml.safe_dump(
|
||||
{
|
||||
**loaded,
|
||||
"model_list": [model],
|
||||
"general_settings": {
|
||||
**SECTION.validate_python(loaded["general_settings"]),
|
||||
"user_api_key_cache_ttl": 600,
|
||||
},
|
||||
"litellm_settings": {
|
||||
**SECTION.validate_python(loaded["litellm_settings"]),
|
||||
"max_budget": max_budget,
|
||||
"budget_duration": "30d",
|
||||
},
|
||||
}
|
||||
)
|
||||
)
|
||||
return path
|
||||
|
||||
|
||||
def _chat(proxy: Gateway, key: str, marker: str) -> httpx.Response:
|
||||
return proxy.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{"model": BUDGET_MODEL, "max_tokens": 20, "messages": [{"role": "user", "content": marker}]},
|
||||
key=key,
|
||||
)
|
||||
|
||||
|
||||
def _served_id(proxy: Gateway, key: str) -> str:
|
||||
response: Final = _chat(proxy, key, f"proxy budget {uuid.uuid4().hex}")
|
||||
assert response.status_code == 200, response.text
|
||||
identity: Final = response.json()["id"]
|
||||
assert isinstance(identity, str), response.text
|
||||
return identity
|
||||
|
||||
|
||||
def _logged_spend(proxy: Gateway, request_ids: tuple[str, ...]) -> tuple[float, ...]:
|
||||
rows: Final = eventually(
|
||||
lambda: proxy.rows(
|
||||
'SELECT spend FROM "LiteLLM_SpendLogs" WHERE request_id = ANY(%s::text[])',
|
||||
("{" + ",".join(request_ids) + "}",),
|
||||
),
|
||||
lambda found: len(found) == len(request_ids),
|
||||
seconds=70,
|
||||
)
|
||||
return tuple(float(str(row["spend"])) for row in rows)
|
||||
|
||||
|
||||
def _proxy_budget_spend(proxy: Gateway) -> float | None:
|
||||
rows: Final = proxy.rows('SELECT spend FROM "LiteLLM_UserTable" WHERE user_id = %s', (LITELLM_PROXY_BUDGET_NAME,))
|
||||
return float(str(rows[0]["spend"])) if rows else None
|
||||
|
||||
|
||||
def _counter_key(redis: Redis) -> str:
|
||||
found: Final = tuple(str(key, "utf-8") for key in redis.scan_iter(match=f"*{GLOBAL_PROXY_SPEND_CACHE_KEY}"))
|
||||
assert len(found) == 1, found
|
||||
return found[0]
|
||||
|
||||
|
||||
def _upstream_saw(proxy: Gateway, marker: str) -> bool:
|
||||
with httpx.Client(base_url=proxy.upstream_url, timeout=5, trust_env=False) as upstream:
|
||||
response: Final = upstream.get("/__observations")
|
||||
response.raise_for_status()
|
||||
return marker in json.dumps(object_value(response.json())["requests"])
|
||||
|
||||
|
||||
async def _add_request_cost(cache: UserApiKeyCache, response_cost: float) -> None:
|
||||
await update_cache(
|
||||
|
|
@ -56,3 +153,63 @@ async def test_global_proxy_spend_adds_every_instance_cost_and_stays_cleared_aft
|
|||
assert await redis_cache.async_get_cache(key=GLOBAL_PROXY_SPEND_CACHE_KEY) is None
|
||||
finally:
|
||||
await redis_cache.async_delete_cache(key=GLOBAL_PROXY_SPEND_CACHE_KEY)
|
||||
|
||||
|
||||
@pytest.mark.timeout(OWNED_PROXY_SECONDS)
|
||||
def test_a_key_with_no_user_spends_toward_the_proxy_max_budget(gateway: Gateway, tmp_path: Path) -> None:
|
||||
with (
|
||||
owned_proxy(gateway, tmp_path, {}, config=_proxy_budget_config(gateway, tmp_path, max_budget=0.01)) as owned,
|
||||
owned.scenario() as scenario,
|
||||
):
|
||||
key: Final = scenario.key(models=[BUDGET_MODEL])
|
||||
(spent,) = _logged_spend(owned, (_served_id(owned, key),))
|
||||
assert spent > 0.01, spent
|
||||
eventually(
|
||||
lambda: _proxy_budget_spend(owned),
|
||||
lambda recorded: recorded is not None and recorded > 0.01,
|
||||
seconds=70,
|
||||
)
|
||||
|
||||
marker: Final = f"over the proxy budget {uuid.uuid4().hex}"
|
||||
refused: Final = _chat(owned, key, marker)
|
||||
|
||||
assert refused.status_code != 200, refused.text
|
||||
assert object_value(refused.json()["error"])["type"] == "budget_exceeded", refused.text
|
||||
assert not _upstream_saw(owned, marker), refused.text
|
||||
|
||||
|
||||
@pytest.mark.timeout(2 * OWNED_PROXY_SECONDS)
|
||||
def test_two_proxies_add_every_request_to_one_proxy_spend_counter(gateway: Gateway, tmp_path: Path) -> None:
|
||||
config: Final = _proxy_budget_config(gateway, tmp_path, max_budget=1000.0)
|
||||
with owned_redis(tmp_path / "redis") as cache, Redis(host=cache.host, port=cache.port) as redis:
|
||||
shared_redis: Final = {"REDIS_HOST": cache.host, "REDIS_PORT": str(cache.port)}
|
||||
with owned_proxy(gateway, tmp_path / "first", shared_redis, config=config) as first:
|
||||
assert first.database_url is not None
|
||||
with (
|
||||
owned_proxy(
|
||||
gateway, tmp_path / "second", {**shared_redis, "DATABASE_URL": first.database_url}, config=config
|
||||
) as second,
|
||||
first.scenario() as scenario,
|
||||
):
|
||||
user_key: Final = scenario.key(user_id=scenario.user(), models=[BUDGET_MODEL])
|
||||
userless_key: Final = scenario.key(models=[BUDGET_MODEL])
|
||||
counter: Final = _counter_key(redis)
|
||||
ttl_before: Final = redis.ttl(counter)
|
||||
|
||||
request_ids: Final = tuple(
|
||||
_served_id(proxy, key) for key, proxy in product((user_key, userless_key), (first, second))
|
||||
)
|
||||
total: Final = sum(_logged_spend(first, request_ids))
|
||||
persisted: Final = eventually(
|
||||
lambda: _proxy_budget_spend(first),
|
||||
lambda spend: spend is not None and spend == pytest.approx(total),
|
||||
seconds=70,
|
||||
)
|
||||
|
||||
assert persisted == pytest.approx(total)
|
||||
eventually(
|
||||
lambda: redis.get(counter),
|
||||
lambda shared: shared is not None and float(str(shared, "utf-8")) == pytest.approx(total),
|
||||
seconds=70,
|
||||
)
|
||||
assert 0 < redis.ttl(counter) <= ttl_before
|
||||
|
|
|
|||
|
|
@ -105,6 +105,24 @@ def test_openai_image_edit_sync_accepts_bytesio_images(respx_mock: respx.MockRou
|
|||
assert _multipart_image_parts(route.calls[0].request) == (_FIRST_IMAGE,)
|
||||
|
||||
|
||||
def test_openai_image_edit_without_a_model_leaves_model_out_of_the_request(
|
||||
respx_mock: respx.MockRouter, httpx_transport: None
|
||||
) -> None:
|
||||
route: Final = respx_mock.post("https://api.openai.com/v1/images/edits").mock(
|
||||
return_value=httpx.Response(200, json=_EDIT_RESPONSE)
|
||||
)
|
||||
|
||||
litellm.image_edit(
|
||||
prompt="combine the reference images",
|
||||
image=io.BytesIO(_FIRST_IMAGE),
|
||||
custom_llm_provider="openai",
|
||||
api_key="fake-key",
|
||||
)
|
||||
|
||||
assert route.call_count == 1
|
||||
assert b'name="model"' not in route.calls[0].request.read()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_openai_image_edit_accepts_bytesio_images(respx_mock: respx.MockRouter, httpx_transport: None) -> None:
|
||||
route: Final = respx_mock.post("https://api.openai.com/v1/images/edits").mock(
|
||||
|
|
|
|||
|
|
@ -300,7 +300,7 @@ async def test_aimage_generation_ignores_a_non_string_base_url_like_the_sync_pat
|
|||
)
|
||||
|
||||
response: Final = await litellm.aimage_generation(
|
||||
model="dall-e-3",
|
||||
model="gpt-image-2.5-sunburst",
|
||||
prompt="a red circle",
|
||||
api_base=api_base,
|
||||
base_url=httpx.URL(api_base),
|
||||
|
|
@ -688,3 +688,24 @@ async def test_chatgpt_images_forward_client_login_and_workspace(
|
|||
assert "cookie" not in request.headers
|
||||
assert "x-litellm-api-key" not in request.headers
|
||||
assert json.loads(request.content) == {"model": "gpt-image-2", "prompt": "A blue circle"}
|
||||
|
||||
|
||||
def test_image_generation_without_a_model_or_provider_fails_before_calling_any_provider(
|
||||
respx_mock: respx.MockRouter,
|
||||
) -> None:
|
||||
with pytest.raises(litellm.BadRequestError, match="model is required"):
|
||||
litellm.image_generation(prompt="a red circle", api_key="fake_openai_api_key")
|
||||
|
||||
assert not respx_mock.calls
|
||||
|
||||
|
||||
def test_topaz_image_variation_works_without_a_model(respx_mock: respx.MockRouter) -> None:
|
||||
route: Final = respx_mock.post("https://api.topazlabs.com/image/v1/enhance").mock(
|
||||
return_value=httpx.Response(200, content=b"enhanced-image", headers={"content-type": "image/png"})
|
||||
)
|
||||
|
||||
response: Final = litellm.image_variation(image=b"source-image", custom_llm_provider="topaz", api_key="fake-key")
|
||||
|
||||
assert route.call_count == 1
|
||||
assert response.data is not None
|
||||
assert response.data[0].b64_json == base64.b64encode(b"enhanced-image").decode()
|
||||
|
|
|
|||
|
|
@ -20,27 +20,9 @@ from litellm.llms.azure.image_generation.http_utils import (
|
|||
azure_deployment_image_generation_json_body,
|
||||
)
|
||||
from litellm.llms.custom_httpx.http_handler import HTTPHandler
|
||||
from litellm.llms.azure.image_generation import (
|
||||
AzureDallE3ImageGenerationConfig,
|
||||
get_azure_image_generation_config,
|
||||
)
|
||||
from litellm.utils import get_optional_params_image_gen
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"received_model, expected_config",
|
||||
[
|
||||
("dall-e-3", AzureDallE3ImageGenerationConfig),
|
||||
("dalle-3", AzureDallE3ImageGenerationConfig),
|
||||
("openai_dall_e_3", AzureDallE3ImageGenerationConfig),
|
||||
],
|
||||
)
|
||||
def test_azure_image_generation_config(received_model, expected_config):
|
||||
assert isinstance(
|
||||
get_azure_image_generation_config(received_model), expected_config
|
||||
)
|
||||
|
||||
|
||||
def test_azure_deployment_image_generation_json_body():
|
||||
"""Deployment-scoped Azure image URL must not send ``model`` in JSON."""
|
||||
api = (
|
||||
|
|
|
|||
|
|
@ -265,7 +265,7 @@ class TestImageGenerationEntryPointHeaders:
|
|||
test_headers = {"cf-aig-authorization": "Bearer my-secret"}
|
||||
|
||||
await litellm.aimage_generation(
|
||||
model="dall-e-3",
|
||||
model="gpt-image-2.5-sunburst",
|
||||
prompt="A white cat",
|
||||
extra_headers=test_headers,
|
||||
client=mock_openai_client,
|
||||
|
|
|
|||
|
|
@ -46,12 +46,12 @@ API_FUNCTION_PARAMS = [
|
|||
(
|
||||
"image_generation",
|
||||
False,
|
||||
{"model": "dall-e-3", "prompt": "A beautiful sunset over mountains"},
|
||||
{"model": "gpt-image-2.5-sunburst", "prompt": "A beautiful sunset over mountains"},
|
||||
),
|
||||
(
|
||||
"image_generation",
|
||||
True,
|
||||
{"model": "dall-e-3", "prompt": "A beautiful sunset over mountains"},
|
||||
{"model": "gpt-image-2.5-sunburst", "prompt": "A beautiful sunset over mountains"},
|
||||
),
|
||||
(
|
||||
"speech",
|
||||
|
|
@ -84,11 +84,7 @@ async def test_openai_client_reuse(function_name, is_async, args):
|
|||
"""
|
||||
|
||||
# Determine which client class to mock based on whether the test is async
|
||||
client_path = (
|
||||
"litellm.llms.openai.openai.AsyncOpenAI"
|
||||
if is_async
|
||||
else "litellm.llms.openai.openai.OpenAI"
|
||||
)
|
||||
client_path = "litellm.llms.openai.openai.AsyncOpenAI" if is_async else "litellm.llms.openai.openai.OpenAI"
|
||||
|
||||
# Create the appropriate patches
|
||||
with (
|
||||
|
|
@ -98,9 +94,7 @@ async def test_openai_client_reuse(function_name, is_async, args):
|
|||
):
|
||||
# Setup the mock to return None first time (cache miss) then a client for subsequent calls
|
||||
mock_client = MagicMock()
|
||||
mock_get_cache.side_effect = [None] + [
|
||||
mock_client
|
||||
] * 9 # First call returns None, rest return the mock client
|
||||
mock_get_cache.side_effect = [None] + [mock_client] * 9 # First call returns None, rest return the mock client
|
||||
|
||||
# Make 10 API calls
|
||||
for _ in range(10):
|
||||
|
|
@ -118,9 +112,9 @@ async def test_openai_client_reuse(function_name, is_async, args):
|
|||
pass
|
||||
|
||||
# Verify client was created only once
|
||||
assert (
|
||||
mock_client_class.call_count == 1
|
||||
), f"{'Async' if is_async else ''}OpenAI client should be created only once"
|
||||
assert mock_client_class.call_count == 1, (
|
||||
f"{'Async' if is_async else ''}OpenAI client should be created only once"
|
||||
)
|
||||
|
||||
# Verify the client was cached
|
||||
assert mock_set_cache.call_count == 1, "Client should be cached once"
|
||||
|
|
@ -144,12 +138,8 @@ def test_precomputed_init_params_match_inspect_signature():
|
|||
_OPENAI_INIT_PARAMS,
|
||||
)
|
||||
|
||||
expected_openai = tuple(
|
||||
p for p in inspect.signature(OpenAI.__init__).parameters if p != "self"
|
||||
)
|
||||
expected_azure = tuple(
|
||||
p for p in inspect.signature(AzureOpenAI.__init__).parameters if p != "self"
|
||||
)
|
||||
expected_openai = tuple(p for p in inspect.signature(OpenAI.__init__).parameters if p != "self")
|
||||
expected_azure = tuple(p for p in inspect.signature(AzureOpenAI.__init__).parameters if p != "self")
|
||||
|
||||
assert _OPENAI_INIT_PARAMS == expected_openai
|
||||
assert _AZURE_OPENAI_INIT_PARAMS == expected_azure
|
||||
|
|
@ -500,7 +490,10 @@ async def test_sdk_api_factories_keep_httpx_transport_and_request_timeouts(api:
|
|||
moved: Final = mock.get(host="sdk-default.example", path="/v1/moved-models").mock(
|
||||
return_value=httpx.Response(
|
||||
200,
|
||||
json={"object": "list", "data": [{"id": "sdk-compat", "object": "model", "created": 0, "owned_by": "t"}]},
|
||||
json={
|
||||
"object": "list",
|
||||
"data": [{"id": "sdk-compat", "object": "model", "created": 0, "owned_by": "t"}],
|
||||
},
|
||||
)
|
||||
)
|
||||
sdk_client: Final = _sdk_api_client(api, is_async, timeout)
|
||||
|
|
@ -542,7 +535,10 @@ async def test_azure_gateway_and_image_variations_use_the_callers_async_session(
|
|||
assert azure_client is not None
|
||||
assert azure_client._client is session
|
||||
response: Final = await aimage_variation(
|
||||
image=BytesIO(b"image-bytes"), api_key="transport-only", api_base="https://sdk-default.example/v1"
|
||||
image=BytesIO(b"image-bytes"),
|
||||
custom_llm_provider="openai",
|
||||
api_key="transport-only",
|
||||
api_base="https://sdk-default.example/v1",
|
||||
)
|
||||
assert response.data[0].url == "https://example.com/image.png"
|
||||
assert not session.is_closed
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue