diff --git a/.github/assets/scaledown/provider-after.jpg b/.github/assets/scaledown/provider-after.jpg new file mode 100644 index 00000000000..bcf1c59c96f Binary files /dev/null and b/.github/assets/scaledown/provider-after.jpg differ diff --git a/.github/assets/scaledown/provider-before.jpg b/.github/assets/scaledown/provider-before.jpg new file mode 100644 index 00000000000..a1b89ae13ca Binary files /dev/null and b/.github/assets/scaledown/provider-before.jpg differ diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index 7f01b38ddc8..faa02c8f1a9 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -327,6 +327,7 @@ jobs: tests/unit/llms/sagemaker tests/unit/llms/sail tests/unit/llms/sambanova + tests/unit/llms/scaledown tests/unit/llms/sap tests/unit/llms/scaleway tests/unit/llms/searchapi diff --git a/README.md b/README.md index 141e8232de8..9987f13d127 100644 --- a/README.md +++ b/README.md @@ -407,6 +407,7 @@ Set `LITELLM_PROXY_API_BASE` and `LITELLM_PROXY_API_KEY` and every model call th | [Sagemaker Chat (`sagemaker_chat`)](https://docs.litellm.ai/docs/providers/aws_sagemaker) | ✅ | ✅ | ✅ | | | | | | | | | [Sail (`sail`)](https://docs.litellm.ai/docs/providers/sail) | ✅ | ✅ | ✅ | | | | | | | | | [Sambanova (`sambanova`)](https://docs.litellm.ai/docs/providers/sambanova) | ✅ | ✅ | ✅ | | | | | | | | +| [ScaleDown (`scaledown`)](https://docs.litellm.ai/docs/providers/scaledown) | ✅ | | | | | | | | | | | [Snowflake (`snowflake`)](https://docs.litellm.ai/docs/providers/snowflake) | ✅ | ✅ | ✅ | | | | | | | | | [Strands Decider (`strands_decider`)](https://docs.litellm.ai/docs/providers) | | | | | | | | | | | | [Text Completion Codestral (`text-completion-codestral`)](https://docs.litellm.ai/docs/providers/codestral) | ✅ | ✅ | ✅ | | | | | | | | diff --git a/litellm/__init__.py b/litellm/__init__.py index a6bbd132cbf..554b37f11d4 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -674,6 +674,7 @@ nvidia_nim_models: Set = set() nvidia_riva_models: Set = set() soniox_models: Set = set() sambanova_models: Set = set() +scaledown_models: Final[Set[str]] = set() # Price reloads update this registry through existing references. sambanova_embedding_models: Set = set() novita_models: Set = set() assemblyai_models: Set = set() @@ -756,7 +757,7 @@ def is_openai_finetune_model(key: str) -> bool: return key.startswith("ft:") and not key.count(":") > 1 -def _populate_provider_model_sets(model_cost_map: Dict) -> None: +def _populate_provider_model_sets(model_cost_map: Mapping[str, Mapping[str, object]]) -> None: for key, value in model_cost_map.items(): if value.get("litellm_provider") == "openai" and not is_openai_finetune_model(key): open_ai_chat_completion_models.add(key) @@ -918,6 +919,8 @@ def _populate_provider_model_sets(model_cost_map: Dict) -> None: soniox_models.add(key) elif value.get("litellm_provider") == "sambanova": sambanova_models.add(key) + elif value.get("litellm_provider") == "scaledown": + scaledown_models.add(key) elif value.get("litellm_provider") == "sambanova-embedding-models": sambanova_embedding_models.add(key) elif value.get("litellm_provider") == "novita": @@ -1106,6 +1109,7 @@ model_list = list( | nvidia_riva_models | soniox_models | sambanova_models + | scaledown_models | azure_text_models | novita_models | assemblyai_models @@ -1215,6 +1219,7 @@ def _build_models_by_provider() -> dict: "nvidia_riva": nvidia_riva_models, "soniox": soniox_models, "sambanova": sambanova_models | sambanova_embedding_models, + "scaledown": scaledown_models, "novita": novita_models, "nebius": nebius_models | nebius_embedding_models, "aiml": aiml_models, @@ -2054,6 +2059,9 @@ if TYPE_CHECKING: from .llms.featherless_ai.chat.transformation import ( FeatherlessAIConfig as FeatherlessAIConfig, ) + from .llms.scaledown.chat.transformation import ( + ScaleDownChatConfig as ScaleDownChatConfig, + ) from .llms.cerebras.chat import CerebrasConfig as CerebrasConfig from .llms.nadir.chat.transformation import NadirConfig as NadirConfig from .llms.baseten.chat import BasetenConfig as BasetenConfig diff --git a/litellm/_lazy_imports_registry.py b/litellm/_lazy_imports_registry.py index 9caba6a060a..ce69de3b240 100644 --- a/litellm/_lazy_imports_registry.py +++ b/litellm/_lazy_imports_registry.py @@ -277,6 +277,7 @@ LLM_CONFIG_NAMES: Final = ( "NvidiaNimConfig", "NvidiaNimEmbeddingConfig", "FeatherlessAIConfig", + "ScaleDownChatConfig", "CerebrasConfig", "NadirConfig", "BasetenConfig", @@ -1103,6 +1104,10 @@ _LLM_CONFIGS_IMPORT_MAP: Final = { ".llms.featherless_ai.chat.transformation", "FeatherlessAIConfig", ), + "ScaleDownChatConfig": ( + ".llms.scaledown.chat.transformation", + "ScaleDownChatConfig", + ), "CerebrasConfig": (".llms.cerebras.chat", "CerebrasConfig"), "NadirConfig": (".llms.nadir.chat.transformation", "NadirConfig"), "BasetenConfig": (".llms.baseten.chat", "BasetenConfig"), diff --git a/litellm/constants.py b/litellm/constants.py index a339fb76976..f9c55f59f42 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -825,6 +825,7 @@ LITELLM_CHAT_PROVIDERS: Final = [ "lemonade", "docker_model_runner", "amazon_nova", + "scaledown", ] # Resolving these providers runs an OAuth device flow (their provider info IS the login), so any diff --git a/litellm/llms/scaledown/__init__.py b/litellm/llms/scaledown/__init__.py new file mode 100644 index 00000000000..d571093869e --- /dev/null +++ b/litellm/llms/scaledown/__init__.py @@ -0,0 +1,3 @@ +from .chat.transformation import ScaleDownChatConfig, ScaleDownError + +__all__ = ["ScaleDownChatConfig", "ScaleDownError"] diff --git a/litellm/llms/scaledown/chat/__init__.py b/litellm/llms/scaledown/chat/__init__.py new file mode 100644 index 00000000000..e6ac554eed1 --- /dev/null +++ b/litellm/llms/scaledown/chat/__init__.py @@ -0,0 +1,3 @@ +from .transformation import ScaleDownChatConfig, ScaleDownError + +__all__ = ["ScaleDownChatConfig", "ScaleDownError"] diff --git a/litellm/llms/scaledown/chat/transformation.py b/litellm/llms/scaledown/chat/transformation.py new file mode 100644 index 00000000000..44eb4785615 --- /dev/null +++ b/litellm/llms/scaledown/chat/transformation.py @@ -0,0 +1,648 @@ +import json +import time +import uuid +from collections.abc import Mapping, Sequence +from functools import reduce +from typing import TYPE_CHECKING, Final, NoReturn +from urllib.parse import unquote + +import httpx +from pydantic import JsonValue, TypeAdapter + +from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException +from litellm.secret_managers.main import get_secret_str +from litellm.types.llms.openai import AllMessageValues +from litellm.types.utils import Choices, Message, ModelResponse, Usage + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer + +JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +JSON_MESSAGES: Final = TypeAdapter(list[dict[str, JsonValue]]) + +DEFAULT_API_BASE: Final = "https://api.scaledown.xyz" + +DOMAIN_MODELS: Final = frozenset({"classify", "extract", "summarize", "compress"}) + +DECISIONS_MODELS: Final = frozenset({"decisions"}) + +DECISIONS_UPSTREAM_MODEL: Final = "classify-1" + +NATIVE_PATHS: Final[Mapping[str, str]] = { + "classify": "/classify", + "extract": "/extract", + "summarize": "/summarization/abstractive", + "compress": "/compress/raw/", +} + +DECISION_QUESTION_TYPES: Final = frozenset({"choice", "noul", "score"}) + +EXTRA_BODY_KEYS: Final[Mapping[str, frozenset[str]]] = { + "extract": frozenset({"threshold", "top_n"}), + "summarize": frozenset(), + "compress": frozenset({"compression_rate"}), + "classify": frozenset({"labels"}), + "decisions": frozenset({"questions"}), +} + +# Set by the LiteLLM router and proxy on every call; handled by LiteLLM, never sent upstream. +ROUTER_PARAMS: Final = frozenset({"max_retries", "stream_options"}) + +SCORE_MIN_LEVELS: Final = 2 +SCORE_MAX_LEVELS: Final = 10 + +MAX_ENTITIES: Final = 1000 +MAX_REF_CHAIN: Final = 32 +MAX_SCHEMA_DEPTH: Final = 32 + + +class ScaleDownError(BaseLLMException): + pass + + +def _reject(message: str, status_code: int = 400) -> NoReturn: + """The single place this adapter raises its own errors.""" + raise ScaleDownError(status_code=status_code, message=message) + + +class ScaleDownChatConfig(BaseConfig): + @property + def custom_llm_provider(self) -> str | None: + return "scaledown" + + def validate_environment( + self, + headers: dict[str, str], # mutable-ok: BaseConfig signature + model: str, + messages: list[AllMessageValues], # mutable-ok: BaseConfig signature + optional_params: dict[str, JsonValue], # mutable-ok: BaseConfig signature + litellm_params: dict[str, JsonValue], # mutable-ok: BaseConfig signature + api_key: str | None = None, + api_base: str | None = None, + ) -> dict[str, JsonValue]: # mutable-ok: BaseConfig signature + resolved_key: Final = api_key or get_secret_str("SCALEDOWN_API_KEY") + if resolved_key is None: + _reject("Missing ScaleDown API key. Set SCALEDOWN_API_KEY or pass api_key to the call.", 401) + if not api_key and api_base is not None and _root(api_base) != _trusted_root(): + _reject( + "A custom api_base needs its own api_key: the SCALEDOWN_API_KEY from the environment is only " + "sent to the default host or the host in SCALEDOWN_API_BASE." + ) + return {**headers, "x-api-key": resolved_key, "content-type": "application/json"} + + def get_complete_url( + self, + api_base: str | None, + api_key: str | None, + model: str, + optional_params: dict[str, JsonValue], # mutable-ok: BaseConfig signature + litellm_params: dict[str, JsonValue], # mutable-ok: BaseConfig signature + stream: bool | None = None, + ) -> str: + base: Final = _root(api_base or _trusted_root()) + operation: Final = _operation(model) + if operation in DECISIONS_MODELS: + return f"{base}/v1/scaledown" + return f"{base}{NATIVE_PATHS[operation]}" + + def get_supported_openai_params(self, model: str) -> list[str]: # mutable-ok: BaseConfig signature + operation: Final = _operation(model) + if operation == "extract": + return ["response_format", "stream"] + if operation == "summarize": + return ["max_tokens", "max_completion_tokens", "stream"] + return ["stream"] + + def map_openai_params( + self, + non_default_params: dict[str, JsonValue], # mutable-ok: BaseConfig signature + optional_params: dict[str, JsonValue], # mutable-ok: BaseConfig signature + model: str, + drop_params: bool, + ) -> dict[str, JsonValue]: # mutable-ok: BaseConfig signature + supported: Final = frozenset(self.get_supported_openai_params(model)) + unsupported: Final = sorted(key for key in non_default_params if key not in supported | ROUTER_PARAMS) + if unsupported and not drop_params: + _reject( + f"ScaleDown model '{model}' does not support {unsupported}. " + f"Supported: {sorted(supported)}. Set litellm.drop_params=True to ignore." + ) + mapped: Final = { + ("max_tokens" if key == "max_completion_tokens" else key): value + for key, value in non_default_params.items() + if key in supported + } + return {**optional_params, **mapped} + + def should_fake_stream( + self, model: str | None, stream: bool | None, custom_llm_provider: str | None = None + ) -> bool: + return bool(stream) + + def transform_extra_body( + self, + extra_body: Mapping[str, object], + request: Mapping[str, object], + model: str, + litellm_params: Mapping[str, object], + ) -> Mapping[str, JsonValue]: + # extra_body is merged after guardrails ran on the messages, so it may only add options. + # Anything that carries prompt text or picks the model is refused. + options: Final = JSON_OBJECT.validate_python(extra_body) + operation: Final = _operation(model) + allowed: Final = EXTRA_BODY_KEYS[operation] + unexpected: Final = sorted(set(options) - allowed) + if unexpected: + allowance: Final = f"may only set {sorted(allowed)}" if allowed else "is not accepted" + _reject( + f"extra_body for scaledown/{operation} {allowance}, got {unexpected}. " + "Text and instructions must be passed as messages." + ) + if operation == "compress": + return {"scaledown": {"rate": options["compression_rate"]}} if "compression_rate" in options else {} + if "questions" in options: + _validate_questions(options["questions"]) + if "labels" in options: + _validate_labels(options["labels"]) + return dict(options) + + def sign_request( + self, + headers: dict[str, str], # mutable-ok: BaseConfig signature + optional_params: dict[str, JsonValue], # mutable-ok: BaseConfig signature + request_data: dict[str, JsonValue], # mutable-ok: BaseConfig signature + api_base: str, + api_key: str | None = None, + model: str | None = None, + stream: bool | None = None, + fake_stream: bool | None = None, + ) -> tuple[dict[str, str], bytes | None]: # mutable-ok: BaseConfig signature + # Runs after extra_body is merged, so this is where a decisions body is known to be complete. + if model is not None and _operation(model) in DECISIONS_MODELS: + _require_complete_decisions(request_data) + if model is not None and _operation(model) == "classify": + _validate_labels(request_data.get("labels")) + return headers, None + + def transform_request( + self, + model: str, + messages: list[AllMessageValues], # mutable-ok: BaseConfig signature + optional_params: dict[str, JsonValue], # mutable-ok: BaseConfig signature + litellm_params: dict[str, JsonValue], # mutable-ok: BaseConfig signature + headers: dict[str, str], # mutable-ok: BaseConfig signature + ) -> dict[str, JsonValue]: # mutable-ok: BaseConfig signature + checked_messages: Final = JSON_MESSAGES.validate_python(messages) + operation: Final = _operation(model) + _validate_text_input(operation, checked_messages, optional_params) + if operation == "classify": + return { + "text": _last_user_text(operation, checked_messages), + **({"labels": optional_params["labels"]} if "labels" in optional_params else {}), + } + if operation in DECISIONS_MODELS: + return dict(_decisions_request(checked_messages, optional_params)) + if operation == "extract": + return dict(_extract_request(checked_messages, optional_params)) + if operation == "summarize": + return dict(_summarize_request(checked_messages, optional_params)) + return dict(_compress_request(checked_messages, optional_params)) + + def transform_response( + self, + model: str, + raw_response: httpx.Response, + model_response: ModelResponse, + logging_obj: "LiteLLMLoggingObj", + request_data: dict[str, JsonValue], # mutable-ok: BaseConfig signature + messages: list[AllMessageValues], # mutable-ok: BaseConfig signature + optional_params: dict[str, JsonValue], # mutable-ok: BaseConfig signature + litellm_params: dict[str, JsonValue], # mutable-ok: BaseConfig signature + encoding: "Tokenizer | None", + api_key: str | None = None, + json_mode: bool | None = None, + ) -> ModelResponse: + try: + raw: Final = JSON_OBJECT.validate_json(raw_response.content) + except ValueError: + _reject( + f"ScaleDown returned a non-JSON response: {raw_response.text[:500]}", + raw_response.status_code if raw_response.is_error else 502, + ) + + operation: Final = _operation(model) + if operation in DECISIONS_MODELS: + return _decisions_response(operation, raw, model_response) + return _native_response(operation, raw, model_response, request_data.get("entities")) + + def get_error_class( + self, + error_message: str, + status_code: int, + headers: dict[str, str] | httpx.Headers, # mutable-ok: BaseConfig signature + ) -> BaseLLMException: + return ScaleDownError(status_code=status_code, message=error_message, headers=headers) + + +def _root(api_base: str) -> str: + return api_base.rstrip("/").removesuffix("/v1") + + +def _trusted_root() -> str: + return _root(get_secret_str("SCALEDOWN_API_BASE") or DEFAULT_API_BASE) + + +def _operation(model: str) -> str: + operation: Final = model.split("/", 1)[1] if "/" in model else model + if operation not in DOMAIN_MODELS | DECISIONS_MODELS: + _reject(f"Unknown ScaleDown model '{model}'. Choose from {sorted(DOMAIN_MODELS | DECISIONS_MODELS)}.") + return operation + + +def _text_of(message: Mapping[str, JsonValue]) -> str | None: + content: Final = message.get("content") + if isinstance(content, str): + return content + if isinstance(content, list): + parts: Final = [ + str(part.get("text", "")) for part in content if isinstance(part, dict) and part.get("type") == "text" + ] + return "\n".join(parts) if parts else None + return None + + +def _validate_text_input( + operation: str, messages: Sequence[Mapping[str, JsonValue]], optional_params: Mapping[str, JsonValue] +) -> None: + if any(key in optional_params for key in ("state", "document", "document_mime_type", "text")): + _reject("ScaleDown accepts text through messages only. Native state and document inputs are not supported.") + for message in messages: + if operation in {"classify", "decisions", "extract"} and message.get("role") in {"system", "developer"}: + _reject( + f"ScaleDown {operation} does not support system messages. Put instructions in the field definitions." + ) + if _has_non_text_content(message): + _reject("ScaleDown currently supports text only. Image and document inputs are not supported.") + + +def _has_non_text_content(message: Mapping[str, JsonValue]) -> bool: + content: Final = message.get("content") + return isinstance(content, list) and any( + not isinstance(part, dict) or part.get("type") != "text" for part in content + ) + + +def _last_user_text(operation: str, messages: Sequence[Mapping[str, JsonValue]]) -> str: + message: Final = next((message for message in reversed(messages) if message.get("role") == "user"), None) + text: Final = _text_of(message) if message is not None else None + if not text: + _reject(f"ScaleDown {operation} requires text in the last user message.") + return text + + +def _validate_labels(labels: JsonValue) -> None: + if not isinstance(labels, list) or not labels: + _reject("ScaleDown classify requires a non-empty 'labels' list of objects with 'name' and 'rubric' strings.") + for label in labels: + if not isinstance(label, Mapping) or not all( + isinstance(label.get(key), str) and label[key] for key in ("name", "rubric") + ): + _reject("Each classify label needs non-empty 'name' and 'rubric' strings.") + + +def _child(node: JsonValue, key: str) -> JsonValue: + return node.get(key.replace("~1", "/").replace("~0", "~")) if isinstance(node, Mapping) else None + + +def _extend_ref_path(path: tuple[str, ...], ref: str) -> tuple[str, ...]: + return (*path, ref) + + +def _nullable_branch(schema: Mapping[str, JsonValue]) -> Mapping[str, JsonValue] | None: + if "allOf" in schema: + _reject("ScaleDown extraction does not support allOf. Supply the object's properties directly.") + variants: Final = schema.get("anyOf", schema.get("oneOf")) + if variants is None: + return None + if any(key in schema for key in ("properties", "items", "$ref", "type")) or ( + "anyOf" in schema and "oneOf" in schema + ): + _reject("ScaleDown extraction cannot combine nullable branches with sibling schema definitions.") + if not isinstance(variants, list) or len(variants) != 2: + _reject("ScaleDown extraction supports anyOf/oneOf only for one schema plus null.") + branches: Final = tuple(branch for branch in variants if isinstance(branch, dict)) + non_null: Final = tuple(branch for branch in branches if branch.get("type") != "null") + if len(branches) != 2 or len(non_null) != 1: + _reject("ScaleDown extraction supports anyOf/oneOf only for one schema plus null.") + return { + **non_null[0], + **{key: value for key, value in schema.items() if key not in {"anyOf", "oneOf"}}, + } + + +def _deref( + prop: Mapping[str, JsonValue], root: Mapping[str, JsonValue], path: tuple[str, ...] +) -> tuple[Mapping[str, JsonValue], tuple[str, ...], bool]: + """Follow nullable branches and local `$ref`s (`#/$defs/X`, `#/definitions/X`). + + Returns the schema, the refs followed so far, and whether a ref repeated (a cycle). + The walk is a bounded loop, so a long chain of distinct refs cannot exhaust the stack. + """ + current: Mapping[str, JsonValue] = prop # rebind-ok: bounded walk along a ref chain + followed: tuple[str, ...] = path # rebind-ok: bounded walk along a ref chain + for _ in range(MAX_REF_CHAIN): + nullable: Mapping[str, JsonValue] | None = _nullable_branch(current) + ref: JsonValue = current.get("$ref") + if nullable is not None: + current = nullable + continue + if "$ref" not in current: + return current, followed, False + if not isinstance(ref, str) or not ref.startswith("#"): + _reject("ScaleDown extraction supports only local JSON Pointer $refs.") + pointer: str = unquote(ref[1:]) + if not pointer.startswith("/"): + _reject("ScaleDown extraction supports only local JSON Pointer $refs beginning with #/.") + if any(key in current for key in ("properties", "items", "type")): + _reject("ScaleDown extraction cannot combine $ref with sibling schema definitions.") + if ref in followed: + return current, followed, True + target: JsonValue = reduce(_child, pointer[1:].split("/"), dict(root)) + if not isinstance(target, Mapping): + _reject(f"ScaleDown extraction cannot resolve $ref '{ref}' to a schema object.") + current, followed = target, _extend_ref_path(followed, ref) + _reject(f"The response_format schema follows more than {MAX_REF_CHAIN} chained $refs or nullable branches.") + + +def _build_properties( + properties: Mapping[str, JsonValue], + root: Mapping[str, JsonValue], + path: tuple[str, ...], + budget: int, + depth: int = 0, +) -> tuple[dict[str, JsonValue], int]: # mutable-ok: serialized JSON entity map + """Turn JSON-schema properties into the /extract entity map, spending one budget unit per entity. + + Each property becomes an entity whose value is its description (or its name when there is + none). Nested objects stay nested, arrays of objects become a one-element list holding the + item's entity map, local `$ref`s are followed, and a ref that points back into itself is cut + off at that point instead of expanded. + """ + if depth > MAX_SCHEMA_DEPTH: + _reject(f"The response_format schema nests deeper than {MAX_SCHEMA_DEPTH} levels.") + built: Final[dict[str, JsonValue]] = {} # mutable-ok: filled once while threading the budget through + remaining = budget # rebind-ok: the budget is threaded through the loop + for name, prop in properties.items(): + if remaining <= 0: + _reject( + f"The response_format schema expands to more than {MAX_ENTITIES} entities; " + "remove recursive or heavily repeated $refs." + ) + built[name], remaining = _build_entity(name, prop, root, path, remaining - 1, depth) + return built, remaining + + +def _build_entity( + name: str, prop: JsonValue, root: Mapping[str, JsonValue], path: tuple[str, ...], budget: int, depth: int +) -> tuple[JsonValue, int]: + schema, followed, cyclic = _deref(prop if isinstance(prop, Mapping) else {}, root, path) + leaf: Final = schema.get("description") or name + nested: Final = schema.get("properties") + if cyclic: + return leaf, budget + if isinstance(nested, Mapping): + return _build_properties(nested, root, followed, budget, depth + 1) + raw_items: Final = schema.get("items") + schema_type: Final = schema.get("type") + is_array: Final = schema_type == "array" or (isinstance(schema_type, list) and "array" in schema_type) + if is_array and isinstance(raw_items, Mapping): + items, item_path, item_cyclic = _deref(raw_items, root, followed) + item_properties: Final = items.get("properties") + if not item_cyclic and isinstance(item_properties, Mapping): + built, remaining = _build_properties(item_properties, root, item_path, budget, depth + 1) + return [built], remaining + return leaf, budget + + +def _extract_request( + messages: Sequence[Mapping[str, JsonValue]], optional_params: Mapping[str, JsonValue] +) -> Mapping[str, JsonValue]: + response_format: Final = optional_params.get("response_format") + json_schema: Final = response_format.get("json_schema") if isinstance(response_format, Mapping) else None + if isinstance(json_schema, Mapping) and json_schema.get("strict") is True: + _reject( + "ScaleDown extract does not support strict JSON schemas. Use strict=false and validate the returned fields." + ) + schema: Final = json_schema.get("schema") if isinstance(json_schema, Mapping) else None + properties: Final = _deref(schema, schema, ())[0].get("properties") if isinstance(schema, Mapping) else None + if not isinstance(schema, dict) or not isinstance(properties, dict) or not properties: + _reject( + "ScaleDown extract needs response_format={'type': 'json_schema', ...} with a non-empty " + "'properties' map; each property becomes an entity and its description the extraction hint." + ) + return { + "text": _last_user_text("extract", messages), + "entities": _build_properties(properties, schema, (), MAX_ENTITIES)[0], + **{key: optional_params[key] for key in ("threshold", "top_n") if key in optional_params}, + } + + +def _summarize_request( + messages: Sequence[Mapping[str, JsonValue]], optional_params: Mapping[str, JsonValue] +) -> Mapping[str, JsonValue]: + instructions: Final = [ + text for message in messages if message.get("role") in {"system", "developer"} and (text := _text_of(message)) + ] + return { + "text": _last_user_text("summarize", messages), + **({"instructions": "\n".join(instructions)} if instructions else {}), + **({"max_tokens": optional_params["max_tokens"]} if "max_tokens" in optional_params else {}), + } + + +def _compress_request( + messages: Sequence[Mapping[str, JsonValue]], optional_params: Mapping[str, JsonValue] +) -> Mapping[str, JsonValue]: + prompt: Final = _last_user_text("compress", messages) + last_user_index: Final = max(i for i, m in enumerate(messages) if m.get("role") == "user" and _text_of(m)) + context: Final = "\n\n".join(text for i, m in enumerate(messages) if i != last_user_index and (text := _text_of(m))) + return { + "context": context, + "prompt": prompt, + "scaledown": {"rate": optional_params.get("compression_rate", "auto")}, + } + + +def _decisions_request( + messages: Sequence[Mapping[str, JsonValue]], optional_params: Mapping[str, JsonValue] +) -> Mapping[str, JsonValue]: + questions: Final = optional_params.get("questions") + if questions is not None: + _validate_questions(questions) + return { + "model": DECISIONS_UPSTREAM_MODEL, + "state": {"text": _last_user_text("decisions", messages)}, + **({"questions": questions} if questions is not None else {}), + } + + +def _require_complete_decisions(body: Mapping[str, JsonValue]) -> None: + if not body.get("questions"): + _reject( + "ScaleDown decisions requires a non-empty 'questions' map, passed as an extra parameter, " + 'e.g. {"questions": {"category": {"type": "choice", "criteria": {"billing": "..."}}}}. ' + "Question types are 'choice', 'noul', and 'score'." + ) + + +def _validate_questions(questions: JsonValue) -> None: + if not isinstance(questions, Mapping): + _reject(f"'questions' must be a map of name to question, got {type(questions).__name__}.") + for name, question in questions.items(): + if not isinstance(question, Mapping): + _reject(f"Question '{name}' must be an object.") + question_type = question.get("type") + if not isinstance(question_type, str) or question_type not in DECISION_QUESTION_TYPES: + _reject(f"Question '{name}' has type {question_type!r}; expected one of {sorted(DECISION_QUESTION_TYPES)}.") + criteria = question.get("criteria") + if question_type == "choice" and (not isinstance(criteria, Mapping) or not criteria): + _reject(f"Choice question '{name}' needs a non-empty 'criteria' map of option key to description.") + if question_type == "score" and ( + not isinstance(criteria, list) or not SCORE_MIN_LEVELS <= len(criteria) <= SCORE_MAX_LEVELS + ): + _reject( + f"Score question '{name}' needs 'criteria' as an ordered list of " + f"{SCORE_MIN_LEVELS} to {SCORE_MAX_LEVELS} level descriptions, lowest to highest." + ) + + +def _unwrap(value: JsonValue) -> JsonValue: + wrapper_keys: Final = {"_value", "_span_anchor"} + if isinstance(value, Mapping) and "_value" in value and set(value) <= wrapper_keys: + return value["_value"] + return value + + +def _clean_extraction(value: JsonValue, requested: JsonValue) -> JsonValue: + """Reduce an extraction result to the fields the caller's schema asked for. + + The request's entity map says which keys are real fields, so provider-added keys such as + `_span_anchor` and `{"_value": ..., "_span_anchor": ...}` wrappers are dropped without + touching a requested field that happens to share one of those names. The untouched payload + stays on `_hidden_params["scaledown_response"]`. + """ + if isinstance(requested, Mapping) and isinstance(value, Mapping): + return {key: _clean_extraction(value[key], requested[key]) for key in requested if key in value} + if isinstance(requested, list) and requested and isinstance(value, list): + return [_clean_extraction(item, requested[0]) for item in value] + return _unwrap(value) + + +def _extracted_fields(raw: Mapping[str, JsonValue], requested: JsonValue) -> JsonValue: + entities: Final = raw.get("entities") + matches: Final = ( + tuple(entity for entity in entities if isinstance(entity, Mapping)) if isinstance(entities, list) else () + ) + scalar_names: Final = ( + tuple(name for name, hint in requested.items() if isinstance(hint, str)) + if isinstance(requested, Mapping) + else () + ) + scalars: Final = { + name: next(entity["text"] for entity in matches if entity.get("type") == name and "text" in entity) + for name in scalar_names + if any(entity.get("type") == name and "text" in entity for entity in matches) + } + structured: Final = raw.get("structured_result") + return _clean_extraction({**scalars, **(structured if isinstance(structured, Mapping) else {})}, requested) + + +def _native_response( + operation: str, raw: JsonValue, model_response: ModelResponse, requested: JsonValue = None +) -> ModelResponse: + """Wrap a native /extract, /summarization/abstractive or /compress/raw/ payload. + + Summarize and compress return the upstream payload as JSON on choices[0].message.content. + Extract returns the fields requested by response_format without enforcing schema types. + The native APIs report input tokens only; they return no output token + count, so completion_tokens stays 0 because it is unmeasured, not because nothing was generated. + """ + if not isinstance(raw, Mapping) or not raw: + _reject( + f"ScaleDown '{operation}' returned an empty or malformed response: {json.dumps(raw)[:500]}", + 500, + ) + + results: Final = raw.get("results") + input_tokens: Final = raw.get("input_tokens") or ( + (results.get("original_prompt_tokens") if isinstance(results, Mapping) else None) + or raw.get("original_prompt_tokens") + if operation == "compress" + else None + ) + content: Final = _extracted_fields(raw, requested) if operation == "extract" else raw + + model_response.id = f"chatcmpl-{uuid.uuid4().hex}" # rebind-ok: LiteLLM fills the response object it passes in + model_response.created = int(time.time()) # rebind-ok: LiteLLM fills the response object it passes in + model_response.model = f"scaledown/{operation}" # rebind-ok: LiteLLM fills the response object it passes in + model_response.object = "chat.completion" # rebind-ok: LiteLLM fills the response object it passes in + model_response.choices = [ # rebind-ok: LiteLLM fills the response object it passes in + Choices( + index=0, + message=Message(role="assistant", content=json.dumps(content, separators=(",", ":"))), + finish_reason="stop", + ) + ] + _set_usage(model_response, {"prompt_tokens": input_tokens or 0}) + model_response.hidden_params["scaledown_response"] = raw # pyright: ignore[reportUnknownMemberType] # Inherited hidden_params has untyped values. # rebind-ok: supplied response + return model_response + + +def _decisions_response(operation: str, raw: JsonValue, model_response: ModelResponse) -> ModelResponse: + answers: Final = raw.get("answers") if isinstance(raw, Mapping) else None + if answers is None: + _reject(f"ScaleDown decisions response contained no answers: {json.dumps(raw)[:500]}", 500) + + model_response.id = f"chatcmpl-{uuid.uuid4().hex}" # rebind-ok: LiteLLM fills the response object it passes in + model_response.created = int(time.time()) # rebind-ok: LiteLLM fills the response object it passes in + model_response.model = f"scaledown/{operation}" # rebind-ok: LiteLLM fills the response object it passes in + model_response.object = "chat.completion" # rebind-ok: LiteLLM fills the response object it passes in + model_response.choices = [ # rebind-ok: LiteLLM fills the response object it passes in + Choices( + index=0, + message=Message(role="assistant", content=json.dumps(answers, separators=(",", ":"))), + finish_reason="stop", + ) + ] + + usage: Final = raw.get("usage") if isinstance(raw, dict) else None + counts: Final = usage if isinstance(usage, dict) else {} + _set_usage( + model_response, + { + "prompt_tokens": counts.get("input_tokens", 0), + "completion_tokens": counts.get("output_tokens", 0), + }, + ) + model_response.hidden_params["scaledown_response"] = raw # pyright: ignore[reportUnknownMemberType] # Inherited hidden_params has untyped values. # rebind-ok: supplied response + return model_response + + +def _set_usage(model_response: ModelResponse, usage: Mapping[str, JsonValue]) -> None: + prompt_tokens: Final = _token_count(usage.get("prompt_tokens")) + completion_tokens: Final = _token_count(usage.get("completion_tokens")) + total_tokens: Final = usage.get("total_tokens") + model_response.usage = Usage( # pyright: ignore[reportAttributeAccessIssue] # ModelResponse stores usage as a dynamic Pydantic field. # rebind-ok: fills the supplied response + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + total_tokens=_token_count(total_tokens) if total_tokens is not None else prompt_tokens + completion_tokens, + ) + + +def _token_count(value: JsonValue) -> int: + if value is None: + return 0 + if isinstance(value, int) and value >= 0: + return value + _reject("ScaleDown returned an invalid token count.", 500) diff --git a/litellm/main.py b/litellm/main.py index 88c1257ee2c..c2c6d5affe8 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -1938,6 +1938,39 @@ def _complete_ragflow(ctx: CompletionDispatchContext) -> _CompletionDispatchResu return response +def _complete_scaledown(ctx: CompletionDispatchContext) -> _CompletionDispatchResult: + try: + return base_llm_http_handler.completion( + model=ctx.model, + messages=ctx.messages, + headers=ctx.headers, + model_response=ctx.model_response, + api_key=ctx.api_key, + api_base=ctx.api_base, + acompletion=ctx.acompletion, + logging_obj=ctx.logging, + optional_params=ctx.optional_params, + litellm_params=ctx.litellm_params, + shared_session=ctx.shared_session, + timeout=cast( # cast-ok: CompletionTimeout.resolve normalizes this value + float | httpx.Timeout, ctx.timeout + ), + client=_dispatch_client_http(ctx), + custom_llm_provider=ctx.custom_llm_provider, + encoding=_get_encoding(), + stream=ctx.stream, + provider_config=ctx.provider_config, + ) + except Exception as exc: + ctx.logging.post_call( + input=ctx.messages, + api_key=ctx.api_key, + original_response=str(exc), + additional_args={"headers": ctx.headers}, + ) + raise + + def _complete_xai(ctx: CompletionDispatchContext) -> _CompletionDispatchResult: acompletion: Final = ctx.acompletion api_base: Final = ctx.api_base @@ -5859,6 +5892,8 @@ def completion( elif custom_llm_provider == "ragflow": ## COMPLETION CALL - RAGFlow uses HTTP handler to support custom URL paths response = _complete_ragflow(_dispatch_ctx) + elif custom_llm_provider == "scaledown": + response = _complete_scaledown(_dispatch_ctx) # rebind-ok: shared provider return elif custom_llm_provider == "xai": ## COMPLETION CALL response = _complete_xai(_dispatch_ctx) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index cb910ae12ea..20558be3422 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -71842,5 +71842,51 @@ "image" ], "supports_embedding_image_input": true + }, + "scaledown/extract": { + "input_cost_per_token": 5e-08, + "output_cost_per_token": 0.0, + "litellm_provider": "scaledown", + "mode": "chat", + "supports_response_schema": false, + "supports_system_messages": false, + "supports_vision": false, + "supports_function_calling": false + }, + "scaledown/summarize": { + "input_cost_per_token": 5e-08, + "output_cost_per_token": 0.0, + "litellm_provider": "scaledown", + "mode": "chat", + "supports_system_messages": true, + "supports_vision": false, + "supports_function_calling": false + }, + "scaledown/compress": { + "input_cost_per_token": 5e-08, + "output_cost_per_token": 0.0, + "litellm_provider": "scaledown", + "mode": "chat", + "supports_system_messages": true, + "supports_function_calling": false, + "supports_vision": false + }, + "scaledown/classify": { + "input_cost_per_token": 5e-08, + "output_cost_per_token": 0.0, + "litellm_provider": "scaledown", + "mode": "chat", + "supports_system_messages": false, + "supports_vision": false, + "supports_function_calling": false + }, + "scaledown/decisions": { + "input_cost_per_token": 5e-08, + "output_cost_per_token": 0.0, + "litellm_provider": "scaledown", + "mode": "chat", + "supports_system_messages": false, + "supports_vision": false, + "supports_function_calling": false } } diff --git a/litellm/provider_endpoints_support_backup.json b/litellm/provider_endpoints_support_backup.json index d9a6feb625d..a0e75b1f02d 100644 --- a/litellm/provider_endpoints_support_backup.json +++ b/litellm/provider_endpoints_support_backup.json @@ -2064,6 +2064,24 @@ "search": true } }, + "scaledown": { + "display_name": "ScaleDown (`scaledown`)", + "url": "https://docs.litellm.ai/docs/providers/scaledown", + "endpoints": { + "chat_completions": true, + "messages": false, + "responses": false, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": false, + "interactions": false + } + }, "sambanova": { "display_name": "Sambanova (`sambanova`)", "url": "https://docs.litellm.ai/docs/providers/sambanova", diff --git a/litellm/proxy/public_endpoints/provider_create_fields.json b/litellm/proxy/public_endpoints/provider_create_fields.json index f20369e279a..633e9edf71c 100644 --- a/litellm/proxy/public_endpoints/provider_create_fields.json +++ b/litellm/proxy/public_endpoints/provider_create_fields.json @@ -1,4 +1,22 @@ [ + { + "provider": "SCALEDOWN", + "provider_display_name": "ScaleDown", + "litellm_provider": "scaledown", + "credential_fields": [ + { + "key": "api_key", + "label": "API Key", + "placeholder": null, + "tooltip": null, + "required": true, + "field_type": "password", + "options": null, + "default_value": null + } + ], + "default_model_placeholder": "scaledown/summarize" + }, { "provider": "AIML", "provider_display_name": "AI/ML API", diff --git a/litellm/types/utils.py b/litellm/types/utils.py index d6751610e7a..13fd86b74a7 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -4224,6 +4224,7 @@ class LlmProviders(str, Enum): FRIENDLIAI = "friendliai" FEATHERLESS_AI = "featherless_ai" WATSONX = "watsonx" + SCALEDOWN = "scaledown" WATSONX_TEXT = "watsonx_text" TRITON = "triton" PREDIBASE = "predibase" diff --git a/litellm/utils.py b/litellm/utils.py index e7a171f56c6..6774771dbf9 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -8655,6 +8655,7 @@ class ProviderConfigManager: False, ), LlmProviders.FEATHERLESS_AI: (lambda: litellm.FeatherlessAIConfig(), False), + LlmProviders.SCALEDOWN: (lambda: litellm.ScaleDownChatConfig(), False), LlmProviders.NOVITA: (lambda: litellm.NovitaConfig(), False), LlmProviders.NEBIUS: (lambda: litellm.NebiusConfig(), False), LlmProviders.WANDB: (lambda: litellm.WandbConfig(), False), diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index cb910ae12ea..20558be3422 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -71842,5 +71842,51 @@ "image" ], "supports_embedding_image_input": true + }, + "scaledown/extract": { + "input_cost_per_token": 5e-08, + "output_cost_per_token": 0.0, + "litellm_provider": "scaledown", + "mode": "chat", + "supports_response_schema": false, + "supports_system_messages": false, + "supports_vision": false, + "supports_function_calling": false + }, + "scaledown/summarize": { + "input_cost_per_token": 5e-08, + "output_cost_per_token": 0.0, + "litellm_provider": "scaledown", + "mode": "chat", + "supports_system_messages": true, + "supports_vision": false, + "supports_function_calling": false + }, + "scaledown/compress": { + "input_cost_per_token": 5e-08, + "output_cost_per_token": 0.0, + "litellm_provider": "scaledown", + "mode": "chat", + "supports_system_messages": true, + "supports_function_calling": false, + "supports_vision": false + }, + "scaledown/classify": { + "input_cost_per_token": 5e-08, + "output_cost_per_token": 0.0, + "litellm_provider": "scaledown", + "mode": "chat", + "supports_system_messages": false, + "supports_vision": false, + "supports_function_calling": false + }, + "scaledown/decisions": { + "input_cost_per_token": 5e-08, + "output_cost_per_token": 0.0, + "litellm_provider": "scaledown", + "mode": "chat", + "supports_system_messages": false, + "supports_vision": false, + "supports_function_calling": false } } diff --git a/provider_endpoints_support.json b/provider_endpoints_support.json index 33ac4e4e60b..ea617cf8462 100644 --- a/provider_endpoints_support.json +++ b/provider_endpoints_support.json @@ -2349,6 +2349,24 @@ "search": true } }, + "scaledown": { + "display_name": "ScaleDown (`scaledown`)", + "url": "https://docs.litellm.ai/docs/providers/scaledown", + "endpoints": { + "chat_completions": true, + "messages": false, + "responses": false, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": false, + "moderations": false, + "batches": false, + "rerank": false, + "a2a": false, + "interactions": false + } + }, "sambanova": { "display_name": "Sambanova (`sambanova`)", "url": "https://docs.litellm.ai/docs/providers/sambanova", diff --git a/tests/code_coverage_tests/recursive_detector.py b/tests/code_coverage_tests/recursive_detector.py index 7166891bb1a..cea2b4154a7 100644 --- a/tests/code_coverage_tests/recursive_detector.py +++ b/tests/code_coverage_tests/recursive_detector.py @@ -67,6 +67,7 @@ IGNORE_FUNCTIONS = [ "strict_json_schema", # harness: max depth set (DEFAULT_MAX_RECURSE_DEPTH); fails closed by raising ValueError at the cap. "toml_value", # harness/codex: max depth set (DEFAULT_MAX_RECURSE_DEPTH); fails closed by raising OptionsMismatch at the cap. "with_json_string_leaves", # transitively bounded: only runs on a tree json_string_leaves already walked under the cap. + "_clean_extraction", # ScaleDown: traverses the entity map built under MAX_SCHEMA_DEPTH=32 and MAX_ENTITIES=1000. "json_unrewritable_labels", # max depth set (MAX_STRUCTURED_CONTENT_SCAN_DEPTH); returns the None sentinel at the cap so the caller blocks. "_flatten_form_field", # bounded by the nesting depth of the already-parsed request body (a finite JSON tree, no cycles possible). "_flatten_form_data_field", # bounded by the nesting depth of the already-parsed request body (a finite JSON tree, no cycles possible). diff --git a/tests/unit/llms/scaledown/__init__.py b/tests/unit/llms/scaledown/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/scaledown/chat/__init__.py b/tests/unit/llms/scaledown/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/scaledown/chat/test_scaledown_chat_transformation.py b/tests/unit/llms/scaledown/chat/test_scaledown_chat_transformation.py new file mode 100644 index 00000000000..82f03a2cd80 --- /dev/null +++ b/tests/unit/llms/scaledown/chat/test_scaledown_chat_transformation.py @@ -0,0 +1,1213 @@ +import json +from collections.abc import Mapping +from queue import SimpleQueue +from typing import Final, Literal + +import httpx +import pytest +import respx +from pydantic import JsonValue + +import litellm +from litellm.llms.scaledown.chat.transformation import ( + DECISIONS_UPSTREAM_MODEL, + ScaleDownChatConfig, + ScaleDownError, +) +from litellm.types.utils import ModelResponse + +BASE = "https://api.scaledown.xyz" + + +@pytest.fixture +def config() -> ScaleDownChatConfig: + return ScaleDownChatConfig() + + +@pytest.fixture(autouse=True) +def scaledown_api_key(monkeypatch): + monkeypatch.setenv("SCALEDOWN_API_KEY", "sk-scaledown-test") + monkeypatch.delenv("SCALEDOWN_API_BASE", raising=False) + + +def _transform_response( + config: ScaleDownChatConfig, model: str, payload: dict, request_data: dict | None = None +) -> ModelResponse: + return config.transform_response( + model=model, + raw_response=httpx.Response( + status_code=200, + json=payload, + request=httpx.Request("POST", f"{BASE}/extract"), + ), + model_response=ModelResponse(), + logging_obj=None, + request_data=request_data or {}, + messages=[], + optional_params={}, + litellm_params={}, + encoding=None, + ) + + +def _build_decisions(config: ScaleDownChatConfig, messages: list, optional_params: dict) -> dict: + body = config.transform_request( + model="scaledown/decisions", messages=messages, optional_params=optional_params, litellm_params={}, headers={} + ) + config.sign_request( + headers={}, + optional_params=optional_params, + request_data=body, + api_base=f"{BASE}/v1/scaledown", + model="scaledown/decisions", + ) + return body + + +def test_auth_uses_x_api_key_not_bearer(config): + headers = config.validate_environment( + headers={}, + model="scaledown/extract", + messages=[], + optional_params={}, + litellm_params={}, + api_key="sk-explicit", + ) + + assert headers["x-api-key"] == "sk-explicit" + assert "Authorization" not in headers + + +def test_missing_api_key_is_rejected_before_any_request(config, monkeypatch): + monkeypatch.delenv("SCALEDOWN_API_KEY", raising=False) + + with pytest.raises(ScaleDownError) as exc: + config.validate_environment( + headers={}, + model="scaledown/extract", + messages=[], + optional_params={}, + litellm_params={}, + ) + + assert exc.value.status_code == 401 + + +@pytest.mark.parametrize( + "model, path", + [ + ("scaledown/extract", "/extract"), + ("scaledown/summarize", "/summarization/abstractive"), + ("scaledown/compress", "/compress/raw/"), + ("scaledown/classify", "/classify"), + ("scaledown/decisions", "/v1/scaledown"), + ], +) +def test_each_model_targets_its_native_endpoint(config, model, path): + url = config.get_complete_url(api_base=None, api_key="k", model=model, optional_params={}, litellm_params={}) + + assert url == f"{BASE}{path}" + + +@pytest.mark.parametrize( + "api_base", + ["https://staging.scaledown.xyz", "https://staging.scaledown.xyz/", "https://staging.scaledown.xyz/v1"], +) +def test_api_base_with_or_without_v1_resolves_to_the_same_host(config, api_base): + extract = config.get_complete_url( + api_base=api_base, api_key="k", model="scaledown/extract", optional_params={}, litellm_params={} + ) + decisions = config.get_complete_url( + api_base=api_base, api_key="k", model="scaledown/decisions", optional_params={}, litellm_params={} + ) + + assert extract == "https://staging.scaledown.xyz/extract" + assert decisions == "https://staging.scaledown.xyz/v1/scaledown" + + +def test_decisions_request_carries_state_and_questions_without_chat_keys(config): + questions = { + "category": { + "type": "choice", + "instructions": "Which category?", + "criteria": {"billing": "A charge or refund.", "technical": "A bug."}, + } + } + + body = config.transform_request( + model="scaledown/decisions", + messages=[{"role": "user", "content": "I was charged twice."}], + optional_params={"questions": questions}, + litellm_params={}, + headers={}, + ) + + assert body == { + "model": DECISIONS_UPSTREAM_MODEL, + "state": {"text": "I was charged twice."}, + "questions": questions, + } + + +def test_decisions_alias_sends_the_only_model_upstream_accepts(config): + body = config.transform_request( + model="scaledown/decisions", + messages=[{"role": "user", "content": "text"}], + optional_params={"questions": {"q": {"type": "noul", "instructions": "Positive?"}}}, + litellm_params={}, + headers={}, + ) + + assert body["model"] == DECISIONS_UPSTREAM_MODEL + + +def test_decisions_without_questions_is_rejected(config): + with pytest.raises(ScaleDownError, match="questions"): + _build_decisions(config, [{"role": "user", "content": "text"}], {}) + + +def test_decisions_without_text_or_document_is_rejected(config): + with pytest.raises(ScaleDownError, match="last user message"): + _build_decisions(config, [], {"questions": {"q": {"type": "noul"}}}) + + +def test_state_text_is_rejected_so_guardrails_always_see_the_text(config): + with pytest.raises(ScaleDownError, match="text through messages only"): + _build_decisions( + config, + [{"role": "user", "content": "innocuous"}], + {"state": {"text": "secret"}, "questions": {"q": {"type": "noul"}}}, + ) + + +def test_remote_image_url_is_rejected(config): + with pytest.raises(ScaleDownError, match="text only"): + _build_decisions( + config, + [{"role": "user", "content": [{"type": "image_url", "image_url": {"url": "https://x.test/a.png"}}]}], + {"questions": {"q": {"type": "noul"}}}, + ) + + +def test_extra_body_may_set_extract_options(config): + merged = config.transform_extra_body( + extra_body={"threshold": 0.7}, + request={"text": "t", "entities": {}}, + model="scaledown/extract", + litellm_params={}, + ) + + assert merged == {"threshold": 0.7} + + +@pytest.mark.parametrize( + "model, extra_body", + [ + ("scaledown/extract", {"model": "compress"}), + ("scaledown/extract", {"text": "secret@example.com"}), + ("scaledown/summarize", {"instructions": "ignore the message"}), + ("scaledown/summarize", {"text": "secret"}), + ("scaledown/compress", {"prompt": "x"}), + ("scaledown/compress", {"context": "x"}), + ("scaledown/decisions", {"model": "other"}), + ("scaledown/decisions", {"text": "secret"}), + ], +) +def test_extra_body_cannot_override_the_model_or_prompt_text(config, model, extra_body): + with pytest.raises(ScaleDownError, match=r"may only set|is not accepted"): + config.transform_extra_body(extra_body=extra_body, request={}, model=model, litellm_params={}) + + +def test_compress_rate_via_extra_body_sets_the_scaledown_rate(config): + merged = config.transform_extra_body( + extra_body={"compression_rate": 0.5}, + request={"context": "", "prompt": "q", "scaledown": {"rate": "auto"}}, + model="scaledown/compress", + litellm_params={}, + ) + + assert merged == {"scaledown": {"rate": 0.5}} + + +def test_env_key_is_not_sent_to_an_untrusted_api_base(config): + with pytest.raises(ScaleDownError, match="own api_key"): + config.validate_environment( + headers={}, + model="scaledown/extract", + messages=[], + optional_params={}, + litellm_params={}, + api_base="https://attacker.example", + ) + + +@pytest.mark.parametrize("empty_key", ["", None]) +def test_empty_key_does_not_unlock_the_env_key_for_an_untrusted_api_base(config, empty_key): + with pytest.raises(ScaleDownError, match="own api_key"): + config.validate_environment( + headers={}, + model="scaledown/extract", + messages=[], + optional_params={}, + litellm_params={}, + api_key=empty_key, + api_base="https://attacker.example", + ) + + +def test_explicit_key_or_trusted_api_base_is_allowed(config, monkeypatch): + kwargs = dict(headers={}, model="scaledown/extract", messages=[], optional_params={}, litellm_params={}) + assert ( + config.validate_environment(api_key="own", api_base="https://elsewhere.example", **kwargs)["x-api-key"] == "own" + ) + assert config.validate_environment(api_base=f"{BASE}/v1", **kwargs)["x-api-key"] == "sk-scaledown-test" + monkeypatch.setenv("SCALEDOWN_API_BASE", "https://staging.scaledown.xyz") + assert ( + config.validate_environment(api_base="https://staging.scaledown.xyz/", **kwargs)["x-api-key"] + == "sk-scaledown-test" + ) + + +def test_max_completion_tokens_maps_to_max_tokens(config): + assert config.map_openai_params( + non_default_params={"max_completion_tokens": 50}, + optional_params={}, + model="scaledown/summarize", + drop_params=False, + ) == {"max_tokens": 50} + + +def test_router_params_are_accepted_and_not_forwarded(config): + assert ( + config.map_openai_params( + non_default_params={"max_retries": 2}, optional_params={}, model="scaledown/decisions", drop_params=False + ) + == {} + ) + + +def test_every_model_fakes_a_stream(config): + assert config.should_fake_stream(model="scaledown/compress", stream=True) + assert not config.should_fake_stream(model="scaledown/compress", stream=False) + + +def test_unknown_question_type_is_rejected(config): + with pytest.raises(ScaleDownError, match="expected one of"): + config.transform_request( + model="scaledown/decisions", + messages=[{"role": "user", "content": "t"}], + optional_params={"questions": {"q": {"type": "ranking"}}}, + litellm_params={}, + headers={}, + ) + + +def test_choice_question_without_criteria_is_rejected(config): + with pytest.raises(ScaleDownError, match="non-empty 'criteria' map"): + config.transform_request( + model="scaledown/decisions", + messages=[{"role": "user", "content": "t"}], + optional_params={"questions": {"q": {"type": "choice", "criteria": {}}}}, + litellm_params={}, + headers={}, + ) + + +@pytest.mark.parametrize("levels", [1, 11]) +def test_score_criteria_outside_two_to_ten_is_rejected(config, levels): + with pytest.raises(ScaleDownError, match="ordered list"): + config.transform_request( + model="scaledown/decisions", + messages=[{"role": "user", "content": "t"}], + optional_params={"questions": {"q": {"type": "score", "criteria": ["l"] * levels}}}, + litellm_params={}, + headers={}, + ) + + +@pytest.mark.parametrize("levels", [2, 10]) +def test_score_criteria_within_two_to_ten_is_accepted(config, levels): + criteria = [f"level {index}" for index in range(levels)] + + body = config.transform_request( + model="scaledown/decisions", + messages=[{"role": "user", "content": "t"}], + optional_params={"questions": {"q": {"type": "score", "criteria": criteria}}}, + litellm_params={}, + headers={}, + ) + + assert body["questions"]["q"]["criteria"] == criteria + + +def test_extract_content_is_the_clean_fields_and_raw_payload_is_kept(config): + payload = { + "entities": [], + "structured_result": { + "invoice": { + "vendor": "Northwind", + "vendor_span_anchor": "Invoice from Northwind", + "customer": {"_value": "Ada Lovelace", "_span_anchor": "to Ada Lovelace"}, + "amount": 500, + } + }, + "input_tokens": 171, + } + + requested = {"invoice": {"vendor": "v", "customer": "c", "amount": "a"}} + response = _transform_response(config, "scaledown/extract", payload, {"entities": requested}) + + assert json.loads(response.choices[0].message.content) == { + "invoice": {"vendor": "Northwind", "customer": "Ada Lovelace", "amount": 500} + } + assert response.model == "scaledown/extract" + assert response.usage.prompt_tokens == 171 + assert response.usage.completion_tokens == 0 + assert response._hidden_params["scaledown_response"] == payload + + +def test_extract_cleaning_keeps_requested_fields_named_like_provider_keys(config): + payload = {"structured_result": {"_value": "x", "vendor": "y", "note_span_anchor": "z", "note": "n"}} + requested = {"_value": "v", "vendor": "v", "note_span_anchor": "s", "note": "n"} + + response = _transform_response(config, "scaledown/extract", payload, {"entities": requested}) + + assert json.loads(response.choices[0].message.content) == payload["structured_result"] + + +def test_extract_cleaning_handles_lists_of_objects(config): + payload = {"structured_result": {"items": [{"sku": "a", "sku_span_anchor": "s"}, {"sku": {"_value": "b"}}]}} + + response = _transform_response(config, "scaledown/extract", payload, {"entities": {"items": [{"sku": "code"}]}}) + + assert json.loads(response.choices[0].message.content) == {"items": [{"sku": "a"}, {"sku": "b"}]} + + +def test_long_ref_chain_is_rejected_cleanly_not_with_a_recursion_error(config): + defs = {f"D{i}": {"$ref": f"#/$defs/D{i + 1}"} for i in range(5000)} + defs["D5000"] = {"type": "string"} + schema = {"type": "object", "$defs": defs, "properties": {"f": {"$ref": "#/$defs/D0"}}} + + with pytest.raises(ScaleDownError, match="chained"): + _extract_body(config, schema) + + +def test_deeply_nested_schema_is_rejected_cleanly(config): + nested: dict = {"type": "string"} + for _ in range(200): + nested = {"type": "object", "properties": {"x": nested}} + schema = {"type": "object", "properties": {"root": nested}} + + with pytest.raises(ScaleDownError, match="nests deeper"): + _extract_body(config, schema) + + +def test_more_than_one_image_is_rejected(config): + image = {"type": "image_url", "image_url": {"url": "data:image/png;base64,aGVsbG8="}} + + with pytest.raises(ScaleDownError, match="text only"): + _build_decisions(config, [{"role": "user", "content": [image, image]}], {"questions": {"q": {"type": "noul"}}}) + + +def test_stream_options_is_accepted_and_not_forwarded(config): + assert ( + config.map_openai_params( + non_default_params={"stream_options": {"include_usage": True}}, + optional_params={}, + model="scaledown/summarize", + drop_params=False, + ) + == {} + ) + + +@pytest.mark.parametrize("wrapper", [None, "anyOf", "oneOf"]) +def test_extract_schema_follows_local_refs(config: ScaleDownChatConfig, wrapper: str | None) -> None: + address: Final = {"$ref": "#/$defs/Address"} + schema: Final = { + "type": "object", + "$defs": { + "Address": {"type": "object", "properties": {"city": {"type": "string", "description": "city name"}}} + }, + "properties": { + "address": {wrapper: [address, {"type": "null"}]} if wrapper else address, + "name": {"type": "string"}, + }, + } + + body = config.transform_request( + model="scaledown/extract", + messages=[{"role": "user", "content": "text"}], + optional_params={"response_format": {"type": "json_schema", "json_schema": {"name": "n", "schema": schema}}}, + litellm_params={}, + headers={}, + ) + + assert body["entities"] == {"address": {"city": "city name"}, "name": "name"} + + +@pytest.mark.parametrize( + ("definition", "reference"), + [ + ("Invoice/Details", "#/$defs/Invoice~1Details"), + ("Invoice~Details", "#/$defs/Invoice~0Details"), + ("Invoice~1 Details", "#/%24defs/Invoice~01%20Details"), + ], +) +def test_extract_preserves_fields_from_escaped_refs( + config: ScaleDownChatConfig, definition: str, reference: str +) -> None: + schema: Final = { + "type": "object", + "$defs": {definition: {"type": "object", "properties": {"amount": {"type": "number"}}}}, + "properties": {"invoice": {"$ref": reference}}, + } + + assert _extract_body(config, schema)["entities"] == {"invoice": {"amount": "amount"}} + + +@pytest.mark.parametrize("reference", ["#/$defs/Missing", "https://example.com/schema", "#named-anchor", None, []]) +def test_extract_rejects_unresolvable_refs(config: ScaleDownChatConfig, reference: JsonValue) -> None: + schema: Final = {"type": "object", "properties": {"invoice": {"$ref": reference}}} + + with pytest.raises(ScaleDownError, match=r"\$ref") as exc: + _extract_body(config, schema) + + assert exc.value.status_code == 400 + + +@pytest.mark.parametrize("field_type", ["array", ["array", "null"]]) +def test_extract_nullable_array_keeps_item_fields(config: ScaleDownChatConfig, field_type: str | list[str]) -> None: + schema: Final = { + "type": "object", + "properties": { + "invoices": { + "type": field_type, + "items": {"type": "object", "properties": {"amount": {"type": "number"}}}, + } + }, + } + + assert _extract_body(config, schema)["entities"] == {"invoices": [{"amount": "amount"}]} + + +@pytest.mark.parametrize("wrapper", ["anyOf", "oneOf", "$ref"]) +def test_extract_rejects_sibling_fields_instead_of_dropping_them(config: ScaleDownChatConfig, wrapper: str) -> None: + address: Final = {"type": "object", "properties": {"city": {"type": "string"}}} + reference: Final = {"$ref": "#/$defs/Address"} + field: Final = { + **(reference if wrapper == "$ref" else {wrapper: [address, {"type": "null"}]}), + "properties": {"postal_code": {"type": "string"}}, + } + schema: Final = {"$defs": {"Address": address}, "type": "object", "properties": {"address": field}} + + with pytest.raises(ScaleDownError, match="sibling") as exc: + _extract_body(config, schema) + + assert exc.value.status_code == 400 + + +@pytest.mark.parametrize( + "field", + [ + {"anyOf": [{"type": "object", "properties": {"x": {"type": "string"}}}, {"type": "string"}]}, + {"oneOf": [{"type": "string"}]}, + {"anyOf": [False, {"type": "string"}]}, + {"allOf": [{"type": "object", "properties": {"x": {"type": "string"}}}]}, + ], +) +def test_extract_rejects_unsupported_composition(config: ScaleDownChatConfig, field: dict[str, JsonValue]) -> None: + with pytest.raises(ScaleDownError) as exc: + _extract_body(config, {"type": "object", "properties": {"field": field}}) + + assert exc.value.status_code == 400 + + +@respx.mock +@pytest.mark.parametrize("question_type", [[], {}]) +def test_malformed_question_types_are_request_errors(question_type: list[JsonValue] | dict[str, JsonValue]) -> None: + with pytest.raises(litellm.BadRequestError, match="Question 'q'"): + litellm.completion( + model="scaledown/decisions", + messages=[{"role": "user", "content": "text"}], + questions={"q": {"type": question_type}}, + ) + + assert not respx.calls + + +@respx.mock +def test_unknown_scaledown_model_is_a_request_error() -> None: + with pytest.raises(litellm.BadRequestError, match="Unknown ScaleDown model"): + litellm.completion(model="scaledown/unknown", messages=[{"role": "user", "content": "text"}]) + + assert not respx.calls + + +def _extract_body(config: ScaleDownChatConfig, schema: dict) -> dict: + return config.transform_request( + model="scaledown/extract", + messages=[{"role": "user", "content": "text"}], + optional_params={"response_format": {"type": "json_schema", "json_schema": {"name": "n", "schema": schema}}}, + litellm_params={}, + headers={}, + ) + + +def test_recursive_schema_is_cut_at_the_cycle_instead_of_expanded(config): + children = {name: {"$ref": "#/$defs/Node"} for name in ("a", "b", "c", "d")} + schema = { + "type": "object", + "$defs": {"Node": {"type": "object", "properties": children}}, + "properties": {"root": {"$ref": "#/$defs/Node"}}, + } + + entities = _extract_body(config, schema)["entities"] + + assert entities == {"root": {"a": "a", "b": "b", "c": "c", "d": "d"}} + + +def test_heavily_repeated_acyclic_refs_are_rejected_by_the_entity_budget(config): + defs = { + f"L{level}": {"type": "object", "properties": {name: {"$ref": f"#/$defs/L{level + 1}"} for name in "abcd"}} + for level in range(10) + } + defs["L10"] = {"type": "string", "description": "leaf"} + schema = {"type": "object", "$defs": defs, "properties": {"root": {"$ref": "#/$defs/L0"}}} + + with pytest.raises(ScaleDownError, match="expands to more than"): + _extract_body(config, schema) + + +def test_summarize_response_is_the_native_payload(config): + payload = {"summary": "terse summary", "input_chars": 400, "output_chars": 13, "input_tokens": 78} + + response = _transform_response(config, "scaledown/summarize", payload) + + assert json.loads(response.choices[0].message.content) == payload + assert response.usage.prompt_tokens == 78 + + +def test_compress_response_reports_top_level_input_tokens(config): + payload = { + "results": {"success": True, "compressed_prompt": "short", "original_prompt_tokens": 52}, + "input_tokens": 52, + } + + response = _transform_response(config, "scaledown/compress", payload) + + assert json.loads(response.choices[0].message.content) == payload + assert response.usage.prompt_tokens == 52 + + +def test_compress_response_falls_back_to_original_prompt_tokens(config): + payload = {"results": {"compressed_prompt": "short", "original_prompt_tokens": 900}} + + assert _transform_response(config, "scaledown/compress", payload).usage.prompt_tokens == 900 + + +def test_extract_request_maps_schema_to_entities_including_nesting(config): + schema = { + "type": "object", + "properties": { + "vendor": {"type": "string", "description": "company that issued the invoice"}, + "notes": {"type": "string"}, + "address": {"type": "object", "properties": {"city": {"type": "string", "description": "city name"}}}, + "line_items": { + "type": "array", + "items": {"type": "object", "properties": {"sku": {"type": "string", "description": "item code"}}}, + }, + }, + } + + body = config.transform_request( + model="scaledown/extract", + messages=[{"role": "user", "content": "Invoice from Northwind."}], + optional_params={"response_format": {"type": "json_schema", "json_schema": {"name": "i", "schema": schema}}}, + litellm_params={}, + headers={}, + ) + + assert body == { + "text": "Invoice from Northwind.", + "entities": { + "vendor": "company that issued the invoice", + "notes": "notes", + "address": {"city": "city name"}, + "line_items": [{"sku": "item code"}], + }, + } + + +def test_extract_without_a_schema_is_rejected(config): + with pytest.raises(ScaleDownError, match="response_format"): + config.transform_request( + model="scaledown/extract", + messages=[{"role": "user", "content": "text"}], + optional_params={}, + litellm_params={}, + headers={}, + ) + + +@pytest.mark.parametrize("role", ["system", "developer"]) +def test_summarize_request_splits_instructions_text_and_max_tokens( + config: ScaleDownChatConfig, role: Literal["system", "developer"] +) -> None: + body: Final = config.transform_request( + model="scaledown/summarize", + messages=[ + {"role": role, "content": "Be terse."}, + {"role": "user", "content": "a long document"}, + ], + optional_params={"max_tokens": 40}, + litellm_params={}, + headers={}, + ) + + assert body == {"text": "a long document", "instructions": "Be terse.", "max_tokens": 40} + + +def test_compress_request_puts_earlier_messages_in_context_and_defaults_rate_to_auto(config): + body = config.transform_request( + model="scaledown/compress", + messages=[ + {"role": "system", "content": "background"}, + {"role": "user", "content": "the question"}, + ], + optional_params={}, + litellm_params={}, + headers={}, + ) + + assert body == {"context": "background", "prompt": "the question", "scaledown": {"rate": "auto"}} + + +def test_decisions_response_ignores_the_upstream_cost_field(config): + payload = { + "model": "classify-1", + "answers": { + "category": { + "type": "choice", + "choice": "billing", + "probabilities": {"billing": 0.93, "technical": 0.07}, + "confidence": 0.93, + } + }, + "usage": {"input_tokens": 62, "output_tokens": 1, "cost": 0.00543}, + } + + response = _transform_response(config, "scaledown/decisions", payload) + + assert json.loads(response.choices[0].message.content) == payload["answers"] + assert response.usage.prompt_tokens == 62 + assert response.usage.completion_tokens == 1 + assert "response_cost" not in response._hidden_params + assert response._hidden_params["scaledown_response"] == payload + + +def test_score_answer_survives_the_envelope_intact(config): + answer = { + "type": "score", + "score": 1.43, + "confidence": 0.57, + "legend": {"0": "Cosmetic", "1": "Workaround", "2": "Blocking"}, + "probabilities": {"0": 0.0, "1": 0.57, "2": 0.43}, + } + payload = { + "model": "classify-1", + "answers": {"bug_severity": answer}, + "usage": {"input_tokens": 68, "output_tokens": 1, "cost": 0.0000029}, + } + + response = _transform_response(config, "scaledown/decisions", payload) + + assert json.loads(response.choices[0].message.content)["bug_severity"] == answer + + +def test_decisions_response_without_answers_is_an_error(config): + with pytest.raises(ScaleDownError, match="no answers"): + _transform_response(config, "scaledown/decisions", {"model": "classify-1"}) + + +def test_empty_native_response_is_an_error(config): + with pytest.raises(ScaleDownError, match="empty or malformed"): + _transform_response(config, "scaledown/extract", {}) + + +@pytest.mark.parametrize("upstream_status", [200, 502]) +def test_non_json_response_is_an_upstream_error(config: ScaleDownChatConfig, upstream_status: int) -> None: + raw_response: Final = httpx.Response( + status_code=upstream_status, + text="bad gateway", + request=httpx.Request("POST", f"{BASE}/v1/scaledown"), + ) + + with pytest.raises(ScaleDownError) as exc: + config.transform_response( + model="scaledown/decisions", + raw_response=raw_response, + model_response=ModelResponse(), + logging_obj=None, + request_data={}, + messages=[], + optional_params={}, + litellm_params={}, + encoding=None, + ) + + assert exc.value.status_code == 502 + + +def test_decisions_takes_no_openai_sampling_params(config): + assert config.get_supported_openai_params("scaledown/decisions") == ["stream"] + + +def test_unsupported_param_is_rejected_unless_dropped(config): + with pytest.raises(ScaleDownError, match="does not support"): + config.map_openai_params( + non_default_params={"temperature": 0.7}, + optional_params={}, + model="scaledown/decisions", + drop_params=False, + ) + + assert ( + config.map_openai_params( + non_default_params={"temperature": 0.7}, + optional_params={}, + model="scaledown/decisions", + drop_params=True, + ) + == {} + ) + + +@respx.mock +def test_completion_sends_x_api_key_to_the_native_summarize_endpoint(): + route = respx.post(f"{BASE}/summarization/abstractive").mock( + return_value=httpx.Response(200, json={"summary": "s", "input_tokens": 120}) + ) + + response = litellm.completion( + model="scaledown/summarize", + messages=[ + {"role": "system", "content": "Be terse."}, + {"role": "user", "content": "a long document"}, + ], + ) + + request = route.calls[0].request + assert request.headers["x-api-key"] == "sk-scaledown-test" + assert "authorization" not in request.headers + assert json.loads(request.content) == {"text": "a long document", "instructions": "Be terse."} + assert json.loads(response.choices[0].message.content)["summary"] == "s" + + +@respx.mock +def test_completion_routes_decisions_to_the_decisions_endpoint(): + route = respx.post(f"{BASE}/v1/scaledown").mock( + return_value=httpx.Response( + 200, + json={ + "model": "classify-1", + "answers": { + "category": { + "type": "choice", + "choice": "billing", + "probabilities": {"billing": 0.93, "technical": 0.07}, + "confidence": 0.93, + } + }, + "usage": {"input_tokens": 62, "output_tokens": 1, "cost": 0.00543}, + }, + ) + ) + + response = litellm.completion( + model="scaledown/decisions", + messages=[{"role": "user", "content": "I was charged twice."}], + questions={ + "category": { + "type": "choice", + "criteria": {"billing": "A charge or refund.", "technical": "A bug."}, + } + }, + ) + + sent = json.loads(route.calls[0].request.content) + assert sent["model"] == "classify-1" + assert sent["state"] == {"text": "I was charged twice."} + assert "messages" not in sent + assert json.loads(response.choices[0].message.content)["category"]["choice"] == "billing" + + +@respx.mock +def test_completion_extract_converts_response_format_to_entities(): + schema = { + "type": "json_schema", + "json_schema": { + "name": "invoice", + "schema": { + "type": "object", + "properties": {"vendor": {"type": "string", "description": "company name"}}, + }, + }, + } + extract_result = {"entities": [], "structured_result": {"vendor": "Acme Corp"}, "input_tokens": 20} + route = respx.post(f"{BASE}/extract").mock(return_value=httpx.Response(200, json=extract_result)) + + response = litellm.completion( + model="scaledown/extract", + messages=[{"role": "user", "content": "Acme Corp invoiced $500."}], + response_format=schema, + ) + + assert json.loads(route.calls[0].request.content) == { + "text": "Acme Corp invoiced $500.", + "entities": {"vendor": "company name"}, + } + assert json.loads(response.choices[0].message.content) == {"vendor": "Acme Corp"} + assert response._hidden_params["scaledown_response"] == extract_result + + +@respx.mock +def test_completion_cost_comes_from_input_tokens_not_upstream_usage_cost(): + respx.post(f"{BASE}/v1/scaledown").mock( + return_value=httpx.Response( + 200, + json={ + "model": "classify-1", + "answers": {"q": {"type": "noul", "noul": 0.8}}, + "usage": {"input_tokens": 1000, "output_tokens": 1, "cost": 99.0}, + }, + ) + ) + + response = litellm.completion( + model="scaledown/decisions", + messages=[{"role": "user", "content": "text"}], + questions={"q": {"type": "noul", "instructions": "Positive?"}}, + ) + + cost = litellm.completion_cost(completion_response=response, model="scaledown/decisions") + assert cost == pytest.approx(1000 * litellm.model_cost["scaledown/decisions"]["input_cost_per_token"]) + + +@respx.mock +def test_completion_honors_scaledown_api_base(monkeypatch): + monkeypatch.setenv("SCALEDOWN_API_BASE", "https://staging.scaledown.xyz/v1") + route = respx.post("https://staging.scaledown.xyz/v1/scaledown").mock( + return_value=httpx.Response( + 200, + json={ + "model": "classify-1", + "answers": {"q": {"type": "noul", "noul": 0.8}}, + "usage": {"input_tokens": 10, "output_tokens": 1, "cost": 0.0}, + }, + ) + ) + + litellm.completion( + model="scaledown/decisions", + messages=[{"role": "user", "content": "text"}], + questions={"q": {"type": "noul", "instructions": "Positive?"}}, + ) + + assert route.called + + +@respx.mock +def test_upstream_payment_error_surfaces_the_detail(): + respx.post(f"{BASE}/v1/scaledown").mock(return_value=httpx.Response(402, json={"detail": "Insufficient credits"})) + + with pytest.raises(litellm.exceptions.BadRequestError) as exc: + litellm.completion( + model="scaledown/decisions", + messages=[{"role": "user", "content": "text"}], + questions={"q": {"type": "noul", "instructions": "Positive?"}}, + ) + + assert "Insufficient credits" in str(exc.value) + + +@respx.mock +def test_failed_completion_invokes_post_api_logger() -> None: + events: Final = SimpleQueue[str]() + + def logger_fn(details: Mapping[str, object]) -> None: + if details.get("log_event_type") == "post_api_call": + events.put(str(details.get("original_response"))) + + respx.post(f"{BASE}/v1/scaledown").respond(402, json={"detail": "Insufficient credits"}) + + with pytest.raises(litellm.BadRequestError): + litellm.completion( + model="scaledown/decisions", + messages=[{"role": "user", "content": "text"}], + questions={"q": {"type": "noul", "instructions": "Positive?"}}, + logger_fn=logger_fn, + ) + + assert not events.empty() + assert "Insufficient credits" in events.get_nowait() + + +@respx.mock +def test_malformed_decisions_request_never_reaches_the_network(): + route = respx.post(f"{BASE}/v1/scaledown") + + with pytest.raises(Exception, match="questions"): + litellm.completion( + model="scaledown/decisions", + messages=[{"role": "user", "content": "text"}], + ) + + assert not route.called + + +@respx.mock +def test_questions_passed_through_extra_body_reach_scaledown(): + route = respx.post(f"{BASE}/v1/scaledown").mock( + return_value=httpx.Response( + 200, + json={ + "model": "classify-1", + "answers": {"q": {"type": "noul", "noul": 0.8}}, + "usage": {"input_tokens": 10, "output_tokens": 1}, + }, + ) + ) + questions = {"q": {"type": "noul", "instructions": "Positive?"}} + + litellm.completion( + model="scaledown/decisions", + messages=[{"role": "user", "content": "text"}], + extra_body={"questions": questions}, + ) + + sent = json.loads(route.calls[0].request.content) + assert sent == {"model": DECISIONS_UPSTREAM_MODEL, "state": {"text": "text"}, "questions": questions} + + +@respx.mock +def test_streaming_request_is_served_as_a_single_chunk(): + respx.post(f"{BASE}/compress/raw/").mock( + return_value=httpx.Response(200, json={"compressed_prompt": "c", "original_prompt_tokens": 9}) + ) + + chunks = list( + litellm.completion( + model="scaledown/compress", + messages=[{"role": "user", "content": "the question"}], + stream=True, + ) + ) + + assert "compressed_prompt" in "".join(chunk.choices[0].delta.content or "" for chunk in chunks) + + +@pytest.mark.parametrize( + "structured,expected", + [ + (None, {"vendor": "Northwind"}), + ( + {"invoice": {"amount": 500, "amount_span_anchor": "$500"}}, + {"vendor": "Northwind", "invoice": {"amount": 500}}, + ), + ], +) +def test_extract_merges_first_scalar_match_with_nested_fields(config, structured, expected): + payload = { + "entities": [ + {"type": "vendor", "text": "Northwind", "confidence": 1.0}, + {"type": "vendor", "text": "Contoso", "confidence": 0.8}, + {"type": "unrequested", "text": "ignore"}, + ], + "structured_result": structured, + "input_tokens": 161, + } + response = _transform_response( + config, + "scaledown/extract", + payload, + {"entities": {"vendor": "company name", "missing": "absent field", "invoice": {"amount": "total"}}}, + ) + + assert json.loads(response.choices[0].message.content) == expected + assert response._hidden_params["scaledown_response"] == payload + assert response.usage.prompt_tokens == 161 + + +@respx.mock +def test_strict_extraction_is_rejected_before_network(): + route = respx.post(f"{BASE}/extract") + with pytest.raises(litellm.BadRequestError, match="strict JSON schemas"): + litellm.completion( + model="scaledown/extract", + messages=[{"role": "user", "content": "Invoice from Northwind. Total: $500."}], + response_format={ + "type": "json_schema", + "json_schema": { + "name": "invoice", + "strict": True, + "schema": {"type": "object", "properties": {"amount": {"type": "string"}}, "required": ["amount"]}, + }, + }, + ) + assert not route.called + + +@pytest.mark.parametrize("model", ["classify", "decisions", "extract", "summarize", "compress"]) +@pytest.mark.parametrize( + "content", + [ + [{"type": "image_url", "image_url": {"url": "data:image/png;base64,aGVsbG8="}}], + [ + {"type": "text", "text": "invoice"}, + {"type": "file", "file": {"file_data": "data:application/pdf;base64,aGVsbG8="}}, + ], + ], +) +def test_multimodal_inputs_are_rejected_instead_of_dropped(config, model, content): + with pytest.raises(ScaleDownError, match="text only"): + config.transform_request(f"scaledown/{model}", [{"role": "user", "content": content}], {}, {}, {}) + + +@pytest.mark.parametrize("params", [{"document": "QQ=="}, {"state": {"document": "QQ=="}}, {"text": "override"}]) +def test_native_input_overrides_are_rejected(config, params): + with pytest.raises(ScaleDownError, match="text through messages only"): + config.transform_request("scaledown/decisions", [{"role": "user", "content": "text"}], params, {}, {}) + + +@pytest.mark.parametrize("model", ["classify", "decisions", "extract"]) +def test_unsupported_system_instructions_are_rejected(config, model): + with pytest.raises(ScaleDownError, match="does not support system messages"): + config.transform_request( + f"scaledown/{model}", + [{"role": "system", "content": "instructions"}, {"role": "user", "content": "text"}], + {}, + {}, + {}, + ) + + +@respx.mock +@pytest.mark.parametrize("nested", [False, True]) +def test_classify_uses_native_labels_and_preserves_scores(nested): + payload = {"top_label": "urgent", "scores": {"urgent": 0.9, "routine": 0.1}, "input_tokens": 521} + route = respx.post(f"{BASE}/classify").mock(return_value=httpx.Response(200, json=payload)) + labels = [ + {"name": "urgent", "rubric": "Service is unavailable."}, + {"name": "routine", "rubric": "A cosmetic issue."}, + ] + options = {"extra_body": {"labels": labels}} if nested else {"labels": labels} + + response = litellm.completion( + model="scaledown/classify", + messages=[{"role": "user", "content": [{"type": "text", "text": "Server is down."}]}], + **options, + ) + + assert json.loads(route.calls[0].request.content) == {"text": "Server is down.", "labels": labels} + assert json.loads(response.choices[0].message.content) == payload + assert response.usage.prompt_tokens == 521 + assert litellm.completion_cost(completion_response=response) == pytest.approx( + 521 * litellm.model_cost["scaledown/classify"]["input_cost_per_token"] + ) + + +@respx.mock +@pytest.mark.parametrize("labels", [None, [], ["urgent"], [{"name": "urgent"}]]) +def test_classify_requires_labels_before_network(labels): + route = respx.post(f"{BASE}/classify") + with pytest.raises(litellm.BadRequestError, match=r"labels|rubric"): + litellm.completion( + model="scaledown/classify", messages=[{"role": "user", "content": "text"}], extra_body={"labels": labels} + ) + assert not route.called + + +@pytest.mark.asyncio +@respx.mock +async def test_async_extraction_keeps_flat_fields_and_usage(): + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + + route = respx.post(f"{BASE}/extract").mock( + return_value=httpx.Response( + 200, + json={ + "entities": [{"type": "vendor", "text": "Northwind"}], + "structured_result": None, + "input_tokens": 161, + }, + ) + ) + client = AsyncHTTPHandler(transport=httpx.AsyncHTTPTransport()) + try: + response = await litellm.acompletion( + model="scaledown/extract", + messages=[{"role": "user", "content": "Invoice from Northwind."}], + response_format={ + "type": "json_schema", + "json_schema": { + "name": "invoice", + "strict": False, + "schema": {"properties": {"vendor": {"type": "string"}}}, + }, + }, + client=client, + ) + finally: + await client.close() + assert json.loads(response.choices[0].message.content) == {"vendor": "Northwind"} + assert response.usage.prompt_tokens == 161 + assert json.loads(route.calls[0].request.content) == { + "text": "Invoice from Northwind.", + "entities": {"vendor": "vendor"}, + } + + +@respx.mock +def test_router_stream_includes_usage_without_sending_stream_options_upstream(): + route = respx.post(f"{BASE}/summarization/abstractive").mock( + return_value=httpx.Response(200, json={"summary": "Fixed billing.", "input_tokens": 90}) + ) + router = litellm.Router( + model_list=[ + {"model_name": "summary", "litellm_params": {"model": "scaledown/summarize", "api_key": "test-key"}} + ], + num_retries=0, + ) + chunks = list( + router.completion( + model="summary", + messages=[{"role": "user", "content": "Billing was fixed on Monday."}], + stream=True, + stream_options={"include_usage": True}, + ) + ) + assert ( + json.loads("".join(chunk.choices[0].delta.content or "" for chunk in chunks if chunk.choices))["summary"] + == "Fixed billing." + ) + assert chunks[-1].usage.prompt_tokens == 90 + assert json.loads(route.calls[0].request.content) == {"text": "Billing was fixed on Monday."} + + +@respx.mock +def test_pydantic_response_format_does_not_silently_claim_strict_support(): + from pydantic import BaseModel, ConfigDict + + class Invoice(BaseModel): + model_config = ConfigDict(frozen=True) + vendor: str + + route = respx.post(f"{BASE}/extract") + with pytest.raises(litellm.BadRequestError, match="strict JSON schemas"): + litellm.completion( + model="scaledown/extract", + messages=[{"role": "user", "content": "Invoice from Northwind."}], + response_format=Invoice, + ) + assert not route.called diff --git a/tests/unit/test_scaledown_model_prices.py b/tests/unit/test_scaledown_model_prices.py new file mode 100644 index 00000000000..02c7983f753 --- /dev/null +++ b/tests/unit/test_scaledown_model_prices.py @@ -0,0 +1,46 @@ +from typing import Final + +import pytest + +import litellm +from litellm.types.utils import Usage + +SCALEDOWN_MODELS: Final = ( + "scaledown/classify", + "scaledown/decisions", + "scaledown/extract", + "scaledown/summarize", + "scaledown/compress", +) + +# Source: ScaleDown's standard rate of $0.05 per million input tokens with output unbilled, +# as shown in the ScaleDown usage dashboard (checked 2026-10-09) and its billing service. +INPUT_COST_PER_MILLION_TOKENS: Final = 0.05 + + +@pytest.mark.parametrize("model", SCALEDOWN_MODELS) +def test_scaledown_model_is_priced_on_input_tokens_only(model: str) -> None: + entry: Final = litellm.model_cost[model] + + assert entry["litellm_provider"] == "scaledown" + assert entry["mode"] == "chat" + assert entry["input_cost_per_token"] * 1_000_000 == pytest.approx(INPUT_COST_PER_MILLION_TOKENS) + assert entry["output_cost_per_token"] == 0.0 + + +@pytest.mark.parametrize("model", SCALEDOWN_MODELS) +def test_scaledown_cost_ignores_output_tokens(model: str) -> None: + prompt_cost, completion_cost = litellm.cost_per_token( + model=model, + custom_llm_provider="scaledown", + usage_object=Usage(prompt_tokens=1_000_000, completion_tokens=500_000), + ) + + assert prompt_cost == pytest.approx(INPUT_COST_PER_MILLION_TOKENS) + assert completion_cost == 0.0 + + +def test_every_scaledown_entry_in_the_cost_map_is_covered_here() -> None: + in_map: Final = {key for key in litellm.model_cost if key.startswith("scaledown/")} + + assert in_map == set(SCALEDOWN_MODELS)