mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-29 06:33:28 +00:00
Merge pull request #47517 from BerriAI/patch_4fe6df3e
chore(release): v1.106.1 patch
This commit is contained in:
commit
978b4d1e4e
107 changed files with 10309 additions and 861 deletions
1
.github/workflows/test-unit.yml
vendored
1
.github/workflows/test-unit.yml
vendored
|
|
@ -670,6 +670,7 @@ jobs:
|
|||
- shard: proxy-extras
|
||||
test-path: >-
|
||||
tests/unit/litellm_proxy_extras
|
||||
tests/unit/db_scripts
|
||||
python-version: "3.12"
|
||||
workers: 2
|
||||
reruns: 0
|
||||
|
|
|
|||
|
|
@ -19,6 +19,7 @@ NONNEG_NUMBER: JsonSchema = {"type": "number", "minimum": 0}
|
|||
NONNEG_INTEGER: JsonSchema = {"type": "integer", "minimum": 0}
|
||||
BOOLEAN: JsonSchema = {"type": "boolean"}
|
||||
STRING: JsonSchema = {"type": "string"}
|
||||
NON_BLANK_STRING: Final[JsonSchema] = {"type": "string", "pattern": r"\S"}
|
||||
TIME_WINDOW: Final[JsonSchema] = {"type": "string", "pattern": r"^([01]\d|2[0-3]):[0-5]\d-([01]\d|2[0-3]):[0-5]\d$"}
|
||||
WEEKDAY_PATTERN: Final = (
|
||||
r"(?i)^(mon|monday|tue|tues|tuesday|wed|wednesday|thu|thur|thurs|thursday|fri|friday|sat|saturday|sun|sunday)$"
|
||||
|
|
@ -103,6 +104,16 @@ OBJECT_KEYS: dict[str, JsonSchema] = {
|
|||
"type": "object",
|
||||
"description": "Provider-internal routing hints (e.g. bedrock_invocation_schema).",
|
||||
},
|
||||
"token_count_model": {
|
||||
"type": "object",
|
||||
"required": ["model"],
|
||||
"additionalProperties": False,
|
||||
"properties": {
|
||||
"model": NON_BLANK_STRING,
|
||||
"region": NON_BLANK_STRING,
|
||||
},
|
||||
"description": "Model that counts tokens for this one, and the AWS region its count API runs in.",
|
||||
},
|
||||
}
|
||||
|
||||
ARRAY_KEYS: dict[str, JsonSchema] = {
|
||||
|
|
@ -189,6 +200,10 @@ INTEGER_KEYS: dict[str, JsonSchema] = {
|
|||
**NONNEG_INTEGER,
|
||||
"description": "Smallest prefix the provider will actually cache; absent means the provider default applies.",
|
||||
},
|
||||
"token_count_framing": {
|
||||
**NONNEG_INTEGER,
|
||||
"description": "Tokens added once per request to a local count estimate, for the provider's chat framing.",
|
||||
},
|
||||
"tpm": {**NONNEG_INTEGER, "description": "Provider default tokens-per-minute limit."},
|
||||
"rpm": {**NONNEG_INTEGER, "description": "Provider default requests-per-minute limit."},
|
||||
}
|
||||
|
|
@ -244,6 +259,11 @@ def cost_schema(key: str) -> JsonSchema:
|
|||
|
||||
def string_key_schemas(modes: tuple) -> dict[str, JsonSchema]:
|
||||
return {
|
||||
"token_count_tokenizer": {
|
||||
"type": "string",
|
||||
"description": "Tokenizer litellm counts this model's requests with when the provider cannot count them.",
|
||||
"enum": ["kimi_k3"],
|
||||
},
|
||||
"litellm_provider": {
|
||||
"type": "string",
|
||||
"description": "LiteLLM provider slug; one of https://docs.litellm.ai/docs/providers.",
|
||||
|
|
|
|||
201
db_scripts/strip_vector_store_internal_metadata.py
Normal file
201
db_scripts/strip_vector_store_internal_metadata.py
Normal file
|
|
@ -0,0 +1,201 @@
|
|||
import argparse
|
||||
import os
|
||||
import sys
|
||||
from collections.abc import Iterator, Sequence
|
||||
from dataclasses import dataclass
|
||||
from itertools import accumulate, repeat, takewhile
|
||||
from typing import Final, Literal, TextIO
|
||||
|
||||
import httpx
|
||||
from pydantic import BaseModel, ConfigDict, JsonValue, ValidationError
|
||||
|
||||
from litellm.vector_stores.utils import is_litellm_written_metadata_key
|
||||
|
||||
DEFAULT_API_BASE: Final = "https://api.openai.com/v1"
|
||||
LIST_PAGE_SIZE: Final = "100"
|
||||
|
||||
|
||||
class VectorStore(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="ignore")
|
||||
|
||||
id: str
|
||||
metadata: dict[str, JsonValue] | None = None
|
||||
|
||||
|
||||
class VectorStorePage(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="ignore")
|
||||
|
||||
data: tuple[VectorStore, ...]
|
||||
has_more: bool = False
|
||||
|
||||
|
||||
class OpenAIErrorBody(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="ignore")
|
||||
|
||||
message: str = ""
|
||||
|
||||
|
||||
class OpenAIErrorResponse(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="ignore")
|
||||
|
||||
error: OpenAIErrorBody
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Options:
|
||||
api_key: str
|
||||
api_base: str
|
||||
vector_store_ids: tuple[str, ...]
|
||||
extra_keys: frozenset[str]
|
||||
apply: bool
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class StoreResult:
|
||||
store_id: str
|
||||
matched_keys: tuple[str, ...]
|
||||
status: Literal["dry_run", "updated", "failed"]
|
||||
error: str = ""
|
||||
|
||||
|
||||
def is_proxy_written_key(key: str, extra_keys: frozenset[str]) -> bool:
|
||||
return is_litellm_written_metadata_key(key) or key in extra_keys
|
||||
|
||||
|
||||
def matching_keys(store: VectorStore, extra_keys: frozenset[str]) -> tuple[str, ...]:
|
||||
return tuple(sorted(key for key in (store.metadata or {}) if is_proxy_written_key(key, extra_keys)))
|
||||
|
||||
|
||||
def fetch_page(client: httpx.Client, after: str | None) -> VectorStorePage:
|
||||
params: Final = {"limit": LIST_PAGE_SIZE} if after is None else {"limit": LIST_PAGE_SIZE, "after": after}
|
||||
return VectorStorePage.model_validate_json(client.get("/vector_stores", params=params).raise_for_status().content)
|
||||
|
||||
|
||||
def fetch_next_page(client: httpx.Client, page: VectorStorePage | None) -> VectorStorePage | None:
|
||||
if page is None or not page.has_more or not page.data:
|
||||
return None
|
||||
return fetch_page(client, after=page.data[-1].id)
|
||||
|
||||
|
||||
def list_all_stores(client: httpx.Client) -> Iterator[VectorStore]:
|
||||
pages: Final = accumulate(
|
||||
repeat(None), lambda page, _: fetch_next_page(client, page), initial=fetch_page(client, after=None)
|
||||
)
|
||||
for page in takewhile(lambda page: page is not None, pages):
|
||||
yield from page.data
|
||||
|
||||
|
||||
def retrieve_stores(client: httpx.Client, store_ids: tuple[str, ...]) -> Iterator[VectorStore]:
|
||||
for store_id in store_ids:
|
||||
yield VectorStore.model_validate_json(client.get(f"/vector_stores/{store_id}").raise_for_status().content)
|
||||
|
||||
|
||||
def describe_http_error(response: httpx.Response) -> str:
|
||||
try:
|
||||
message: Final = OpenAIErrorResponse.model_validate_json(response.content).error.message
|
||||
except ValidationError:
|
||||
return f"HTTP {response.status_code}"
|
||||
return f"HTTP {response.status_code}: {message}"
|
||||
|
||||
|
||||
def delete_keys(client: httpx.Client, store_id: str, keys: tuple[str, ...]) -> StoreResult:
|
||||
try:
|
||||
response: Final = client.post(f"/vector_stores/{store_id}", json={"metadata": dict.fromkeys(keys)})
|
||||
except httpx.TransportError as error:
|
||||
return StoreResult(store_id=store_id, matched_keys=keys, status="failed", error=type(error).__name__)
|
||||
if response.is_success:
|
||||
return StoreResult(store_id=store_id, matched_keys=keys, status="updated")
|
||||
return StoreResult(store_id=store_id, matched_keys=keys, status="failed", error=describe_http_error(response))
|
||||
|
||||
|
||||
def clean_store(client: httpx.Client, store: VectorStore, options: Options) -> StoreResult | None:
|
||||
keys: Final = matching_keys(store, options.extra_keys)
|
||||
if not keys:
|
||||
return None
|
||||
if not options.apply:
|
||||
return StoreResult(store_id=store.id, matched_keys=keys, status="dry_run")
|
||||
return delete_keys(client, store.id, keys)
|
||||
|
||||
|
||||
def format_result(result: StoreResult) -> str:
|
||||
key_list: Final = ", ".join(result.matched_keys)
|
||||
count: Final = len(result.matched_keys)
|
||||
match result.status:
|
||||
case "dry_run":
|
||||
return f"{result.store_id}: would delete {count} keys: {key_list}"
|
||||
case "updated":
|
||||
return f"{result.store_id}: deleted {count} keys: {key_list}"
|
||||
case "failed":
|
||||
return f"{result.store_id}: failed to delete {count} keys ({result.error}): {key_list}"
|
||||
|
||||
|
||||
def clean_and_report(
|
||||
client: httpx.Client, stores: Iterator[VectorStore], options: Options, out: TextIO
|
||||
) -> Iterator[StoreResult | None]:
|
||||
for store in stores:
|
||||
result: Final = clean_store(client, store, options)
|
||||
if result is not None:
|
||||
print(format_result(result), file=out)
|
||||
yield result
|
||||
|
||||
|
||||
def run(options: Options, client: httpx.Client, out: TextIO) -> int:
|
||||
stores: Final = (
|
||||
retrieve_stores(client, options.vector_store_ids) if options.vector_store_ids else list_all_stores(client)
|
||||
)
|
||||
scanned: Final = tuple(clean_and_report(client, stores, options, out))
|
||||
results: Final = tuple(result for result in scanned if result is not None)
|
||||
updated: Final = sum(1 for result in results if result.status == "updated")
|
||||
failed: Final = sum(1 for result in results if result.status == "failed")
|
||||
mode: Final = "applied" if options.apply else "dry run, pass --apply to delete"
|
||||
print(
|
||||
f"scanned={len(scanned)} matched={len(results)} updated={updated} failed={failed} ({mode})",
|
||||
file=out,
|
||||
)
|
||||
return 1 if failed else 0
|
||||
|
||||
|
||||
def parse_options(argv: Sequence[str] | None) -> Options:
|
||||
parser: Final = argparse.ArgumentParser(
|
||||
description="Remove LiteLLM proxy-written keys from OpenAI vector store metadata. Dry run by default."
|
||||
)
|
||||
parser.add_argument("--api-key", default=os.environ.get("OPENAI_API_KEY"), help="defaults to $OPENAI_API_KEY")
|
||||
parser.add_argument("--api-base", default=DEFAULT_API_BASE)
|
||||
parser.add_argument("--vector-store-id", action="append", default=[], help="repeatable, defaults to every store")
|
||||
parser.add_argument("--extra-key", action="append", default=[], help="repeatable, another key name to delete")
|
||||
parser.add_argument("--apply", action="store_true", help="delete the matching keys instead of listing them")
|
||||
args: Final = parser.parse_args(argv)
|
||||
if not args.api_key:
|
||||
parser.error("--api-key or OPENAI_API_KEY is required")
|
||||
return Options(
|
||||
api_key=args.api_key,
|
||||
api_base=args.api_base.rstrip("/"),
|
||||
vector_store_ids=tuple(args.vector_store_id),
|
||||
extra_keys=frozenset(args.extra_key),
|
||||
apply=args.apply,
|
||||
)
|
||||
|
||||
|
||||
def main(argv: Sequence[str] | None = None, transport: httpx.BaseTransport | None = None) -> int:
|
||||
options: Final = parse_options(argv)
|
||||
with httpx.Client(
|
||||
base_url=options.api_base,
|
||||
headers={"Authorization": f"Bearer {options.api_key}"},
|
||||
timeout=60.0,
|
||||
transport=transport,
|
||||
) as client:
|
||||
try:
|
||||
return run(options, client, sys.stdout)
|
||||
except httpx.HTTPStatusError as error:
|
||||
sys.stderr.write(f"could not read vector stores: {describe_http_error(error.response)}\n")
|
||||
return 2
|
||||
except httpx.TransportError as error:
|
||||
sys.stderr.write(f"could not reach {options.api_base}: {type(error).__name__}\n")
|
||||
return 2
|
||||
except ValidationError:
|
||||
sys.stderr.write(f"unexpected vector store response shape from {options.api_base}\n")
|
||||
return 2
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
|
|
@ -74,7 +74,7 @@ if TYPE_CHECKING:
|
|||
ChatCompletionToolReferenceObject,
|
||||
OpenAIMessageContentListBlock,
|
||||
)
|
||||
from litellm.types.utils import Choices
|
||||
from litellm.types.utils import Choices, Message
|
||||
|
||||
|
||||
_CHAT_COMPLETION_FIELDS: Final = frozenset((*ModelResponse.model_fields, "usage"))
|
||||
|
|
@ -308,6 +308,67 @@ def _as_chat_reasoning_items(
|
|||
return cast(list[ChatCompletionReasoningItem], list(reasoning_items))
|
||||
|
||||
|
||||
def _shift_annotation(annotation: ChatCompletionAnnotation, offset: int) -> ChatCompletionAnnotation:
|
||||
shifted: Final = {
|
||||
key: value + offset if key in ("start_index", "end_index") and isinstance(value, int) else value
|
||||
for key, value in annotation.items()
|
||||
}
|
||||
return cast(ChatCompletionAnnotation, shifted)
|
||||
|
||||
|
||||
def _message_annotations(message: "Message") -> Sequence[ChatCompletionAnnotation]:
|
||||
annotations: Sequence[ChatCompletionAnnotation] | None = getattr(message, "annotations", None)
|
||||
return annotations or ()
|
||||
|
||||
|
||||
def _message_reasoning_content(message: "Message") -> str | None:
|
||||
reasoning_content: str | None = getattr(message, "reasoning_content", None)
|
||||
return reasoning_content
|
||||
|
||||
|
||||
def _message_reasoning_items(message: "Message") -> Sequence[ChatCompletionReasoningItem]:
|
||||
reasoning_items: Sequence[ChatCompletionReasoningItem] | None = getattr(message, "reasoning_items", None)
|
||||
return reasoning_items or ()
|
||||
|
||||
|
||||
def _shifted_annotations(messages: Sequence["Message"], offsets: Sequence[int]) -> Iterator[ChatCompletionAnnotation]:
|
||||
for message, offset in zip(messages, offsets):
|
||||
for annotation in _message_annotations(message):
|
||||
yield _shift_annotation(annotation, offset)
|
||||
|
||||
|
||||
def _fold_choices_into_one(
|
||||
choices: Sequence["Choices"],
|
||||
tool_calls: Sequence[Mapping[str, object]],
|
||||
trailing_reasoning_content: str | None,
|
||||
trailing_reasoning_item: _BuiltReasoningItem | None,
|
||||
) -> "Choices":
|
||||
"""Clients that only read choices[0] must still see every text and tool call."""
|
||||
from litellm.types.utils import Choices, Message
|
||||
|
||||
messages: Final = tuple(choice.message for choice in choices)
|
||||
texts: Final = tuple(str(message.content or "") for message in messages)
|
||||
offsets: Final = (0, *accumulate(len(text) for text in texts))
|
||||
annotations: Final = tuple(_shifted_annotations(messages, offsets))
|
||||
reasoning_contents: Final = (*(_message_reasoning_content(m) for m in messages), trailing_reasoning_content)
|
||||
reasoning_items: Final = (
|
||||
*chain.from_iterable(_message_reasoning_items(m) for m in messages),
|
||||
*(() if trailing_reasoning_item is None else (trailing_reasoning_item,)),
|
||||
)
|
||||
return Choices(
|
||||
message=Message(
|
||||
role=messages[0].role if messages else "assistant",
|
||||
content="".join(texts) if messages else None,
|
||||
annotations=list(annotations) or None,
|
||||
tool_calls=list(tool_calls) or None,
|
||||
reasoning_content=" ".join(value for value in reasoning_contents if value) or None,
|
||||
reasoning_items=_as_chat_reasoning_items(reasoning_items),
|
||||
),
|
||||
finish_reason="tool_calls" if tool_calls else "stop",
|
||||
index=0,
|
||||
)
|
||||
|
||||
|
||||
_ToolChoiceT = TypeVar("_ToolChoiceT")
|
||||
|
||||
|
||||
|
|
@ -951,47 +1012,16 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
|
|||
else:
|
||||
pass # don't fail request if item in list is not supported
|
||||
|
||||
if accumulated_tool_calls and choices:
|
||||
last_choice: Final = choices[-1]
|
||||
last_reasoning_content: Final = getattr(last_choice.message, "reasoning_content", None)
|
||||
last_reasoning_items: Final = getattr(last_choice.message, "reasoning_items", None)
|
||||
merged_reasoning_content: Final = (
|
||||
" ".join(value for value in (last_reasoning_content, reasoning_content) if value) or None
|
||||
)
|
||||
merged_reasoning_items: Final = _as_chat_reasoning_items(
|
||||
(
|
||||
*(last_reasoning_items or ()),
|
||||
*(() if pending_reasoning_item is None else (pending_reasoning_item,)),
|
||||
)
|
||||
)
|
||||
merged_message: Final = Message(
|
||||
role=last_choice.message.role,
|
||||
content=last_choice.message.content,
|
||||
annotations=getattr(last_choice.message, "annotations", None),
|
||||
if not choices and not accumulated_tool_calls:
|
||||
return choices
|
||||
return [
|
||||
_fold_choices_into_one(
|
||||
choices=choices,
|
||||
tool_calls=accumulated_tool_calls,
|
||||
reasoning_content=merged_reasoning_content,
|
||||
reasoning_items=merged_reasoning_items,
|
||||
trailing_reasoning_content=reasoning_content,
|
||||
trailing_reasoning_item=pending_reasoning_item,
|
||||
)
|
||||
return [
|
||||
*choices[:-1],
|
||||
Choices(message=merged_message, finish_reason="tool_calls", index=last_choice.index),
|
||||
]
|
||||
|
||||
if accumulated_tool_calls:
|
||||
msg = Message(
|
||||
content=None,
|
||||
tool_calls=accumulated_tool_calls,
|
||||
reasoning_content=reasoning_content,
|
||||
reasoning_items=cast(
|
||||
list[ChatCompletionReasoningItem] | None,
|
||||
([pending_reasoning_item] if pending_reasoning_item is not None else None),
|
||||
),
|
||||
)
|
||||
choices.append(Choices(message=msg, finish_reason="tool_calls", index=index))
|
||||
reasoning_content = None
|
||||
pending_reasoning_item = None
|
||||
|
||||
return choices
|
||||
]
|
||||
|
||||
@staticmethod
|
||||
def _build_empty_incomplete_choice(
|
||||
|
|
|
|||
|
|
@ -1588,6 +1588,7 @@ OUTPUT_TOKEN_CEILING_PARAMS: Final = frozenset({"max_tokens", "max_completion_to
|
|||
CLIENT_OUTPUT_CEILING_METADATA_KEY: Final = "_client_output_ceiling"
|
||||
CONSUMED_REQUEST_TAGS_METADATA_KEY: Final = "_consumed_request_tags"
|
||||
ROUTING_REQUEST_TAGS_METADATA_KEY: Final = "_routing_request_tags"
|
||||
VECTOR_STORE_CALLER_METADATA_KWARG: Final = "litellm_vector_store_caller_metadata"
|
||||
ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY: Final = "_litellm_router_usage_counted_tokens"
|
||||
INTERNAL_CALL_ORIGIN_METADATA_KEY: Final = "internal_call_origin"
|
||||
SESSION_ID_GENERATED_METADATA_KEY: Final = "litellm_session_id_generated"
|
||||
|
|
|
|||
|
|
@ -219,7 +219,7 @@ class GitLabTemplateManager:
|
|||
elif value.replace(".", "").isdigit():
|
||||
try:
|
||||
result[key] = float(value)
|
||||
except Exception:
|
||||
except ValueError:
|
||||
result[key] = value
|
||||
else:
|
||||
result[key] = value.strip("\"'")
|
||||
|
|
|
|||
|
|
@ -3,9 +3,9 @@
|
|||
import copy
|
||||
import logging
|
||||
import re
|
||||
from collections.abc import Collection, Iterable, Mapping
|
||||
from collections.abc import Collection, Iterable, Iterator, Mapping
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, cast
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, cast, runtime_checkable
|
||||
|
||||
import httpx
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
|
|
@ -52,6 +52,28 @@ def safe_divide_seconds(seconds: float, denominator: float, default: float | Non
|
|||
_DROP_PARAMS_BOOL: Final = TypeAdapter(bool)
|
||||
|
||||
|
||||
def as_str_mapping(value: object) -> Mapping[str, object] | None:
|
||||
return value if isinstance(value, Mapping) else None
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class _SupportsIter(Protocol):
|
||||
def __iter__(self) -> Iterator[object]: ...
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class _SupportsGetItem(Protocol):
|
||||
def __getitem__(self, index: int, /) -> object: ...
|
||||
|
||||
|
||||
def iter_items(value: object) -> Iterable[object]:
|
||||
if isinstance(value, _SupportsIter):
|
||||
return value
|
||||
if isinstance(value, _SupportsGetItem):
|
||||
return iter(value)
|
||||
raise TypeError(f"'{type(value).__name__}' object is not iterable")
|
||||
|
||||
|
||||
def normalize_drop_params(value: object) -> bool | None:
|
||||
if value is None or isinstance(value, bool):
|
||||
return value
|
||||
|
|
|
|||
|
|
@ -2100,10 +2100,8 @@ def is_unsignable_thinking_block(block: object) -> bool:
|
|||
Anthropic verifies the signature cryptographically, so a block with a null,
|
||||
empty, or missing signature (e.g. from an open-source reasoning model) is
|
||||
rejected with a 400, and so is a block whose signature or data carries
|
||||
another provider's encrypted reasoning. It also rejects a `thinking` block
|
||||
whose text is empty or whitespace-only ("each thinking block must contain
|
||||
thinking"), regardless of signature, e.g. when a `thinking_blocks` history
|
||||
item from a non-Anthropic reasoning provider is replayed through this path.
|
||||
another provider's encrypted reasoning. A signed block is kept whatever its
|
||||
text, since ``display: omitted`` returns signed blocks with empty ``thinking``.
|
||||
`redacted_thinking` blocks carry no signature and are always kept.
|
||||
"""
|
||||
if is_encrypted_reasoning_block(block):
|
||||
|
|
@ -2114,10 +2112,7 @@ def is_unsignable_thinking_block(block: object) -> bool:
|
|||
if mapping.get("type") != "thinking":
|
||||
return False
|
||||
signature: Final = mapping.get("signature")
|
||||
if not (isinstance(signature, str) and len(signature) > 0):
|
||||
return True
|
||||
thinking_text: Final = mapping.get("thinking")
|
||||
return not (isinstance(thinking_text, str) and len(thinking_text.strip()) > 0)
|
||||
return not (isinstance(signature, str) and len(signature) > 0)
|
||||
|
||||
|
||||
def strip_encrypted_reasoning_from_messages(
|
||||
|
|
|
|||
|
|
@ -2774,13 +2774,8 @@ def anthropic_messages_pt(
|
|||
if not isinstance(m, dict):
|
||||
continue
|
||||
# handle thinking blocks
|
||||
thinking_block = cast(str, m.get("thinking", ""))
|
||||
text_block = cast(str, m.get("text", ""))
|
||||
if (
|
||||
m.get("type", "") == "thinking"
|
||||
and len(thinking_block) > 0
|
||||
and not is_unsignable_thinking_block(m)
|
||||
): # don't pass empty text blocks. anthropic api raises errors.
|
||||
if m.get("type", "") == "thinking" and not is_unsignable_thinking_block(m):
|
||||
anthropic_message: ChatCompletionThinkingBlock | AnthropicMessagesTextParam = cast(
|
||||
ChatCompletionThinkingBlock, m
|
||||
)
|
||||
|
|
@ -5781,3 +5776,25 @@ def resolve_structured_messages(
|
|||
if structured:
|
||||
return [msg if isinstance(msg, dict) else msg.model_dump() for msg in structured]
|
||||
return None
|
||||
|
||||
|
||||
def is_flat_responses_tool(tool: Mapping[str, object]) -> bool:
|
||||
tool_type: Final = tool.get("type")
|
||||
return tool_type in ("function", "custom") and tool_type not in tool and "input_schema" not in tool
|
||||
|
||||
|
||||
def chat_completion_tools_for(tool: Mapping[str, object]) -> tuple[Mapping[str, object], ...]:
|
||||
from litellm.llms.anthropic.pass_through.adapters.transformation import LiteLLMAnthropicMessagesAdapter
|
||||
from litellm.responses.litellm_completion_transformation.transformation import (
|
||||
LiteLLMCompletionResponsesConfig,
|
||||
)
|
||||
|
||||
if is_flat_responses_tool(tool):
|
||||
chat_tools, _ = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools(
|
||||
[dict(tool)] # pyright: ignore[reportArgumentType] # a JSON dict in a flat Responses tool shape
|
||||
)
|
||||
return tuple(chat_tools)
|
||||
anthropic_tools, _ = LiteLLMAnthropicMessagesAdapter().translate_anthropic_tools_to_openai(
|
||||
[dict(tool)] # pyright: ignore[reportArgumentType] # a JSON dict in an Anthropic or chat tool shape
|
||||
)
|
||||
return tuple(anthropic_tools)
|
||||
|
|
|
|||
|
|
@ -265,8 +265,7 @@ def _assistant_part_renders(part: object) -> bool:
|
|||
"""A text part always renders: the converter pads empty text with a placeholder."""
|
||||
part_type: Final = message_field(part, "type")
|
||||
if part_type == "thinking":
|
||||
thinking: Final = message_field(part, "thinking")
|
||||
return isinstance(thinking, str) and bool(thinking) and _thinking_block_renders(part)
|
||||
return _thinking_block_renders(part)
|
||||
return part_type in _RENDERED_ASSISTANT_PART_TYPES or (
|
||||
isinstance(part_type, str) and part_type.endswith("_tool_result")
|
||||
)
|
||||
|
|
|
|||
|
|
@ -329,6 +329,10 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
|
|||
def _resolved_provider(self) -> str:
|
||||
return self.custom_llm_provider or "anthropic"
|
||||
|
||||
@property
|
||||
def _carries_thinking_block_binding(self) -> bool:
|
||||
return type(self) is AnthropicConfig
|
||||
|
||||
@classmethod
|
||||
def get_config(cls, *, model: str | None = None):
|
||||
config: Final = super().get_config()
|
||||
|
|
@ -1656,7 +1660,10 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
|
|||
else:
|
||||
optional_params["thinking"] = value
|
||||
AnthropicModelInfo.translate_legacy_thinking_for_adaptive_model(
|
||||
model=model, optional_params=optional_params, custom_llm_provider=self._resolved_provider
|
||||
model=model,
|
||||
optional_params=optional_params,
|
||||
custom_llm_provider=self._resolved_provider,
|
||||
carry_block_binding=self._carries_thinking_block_binding,
|
||||
)
|
||||
elif param == "reasoning_effort":
|
||||
# Accept both string ("low") and dict ({"effort": "low",
|
||||
|
|
|
|||
|
|
@ -106,7 +106,7 @@ _INFERENCE_PROFILE_MINOR_RE: Final = re.compile(r":\d+$")
|
|||
_DATED_RELEASE_SUFFIX_RE: Final = re.compile(r"-\d{8}$")
|
||||
_DOTTED_VERSION_RE: Final = re.compile(r"(\d)\.(\d)")
|
||||
_CLAUDE_CODE_BILLING_HEADER_PREFIX: Final = "x-anthropic-billing-header:"
|
||||
_CLAUDE_CODE_OBJECT_MAPPING_ADAPTER: Final = TypeAdapter(dict[object, object])
|
||||
_OBJECT_MAPPING_ADAPTER: Final = TypeAdapter(dict[object, object])
|
||||
_CLAUDE_CODE_OBJECT_LIST_ADAPTER: Final = TypeAdapter(list[object])
|
||||
|
||||
|
||||
|
|
@ -155,13 +155,18 @@ def is_claude_code_user_agent(user_agent: str) -> bool:
|
|||
return user_agent.startswith(_CLAUDE_CODE_USER_AGENT_PREFIXES)
|
||||
|
||||
|
||||
def _validated_claude_code_mapping(value: object) -> dict[object, object] | None:
|
||||
def validated_object_mapping(value: object) -> dict[object, object] | None:
|
||||
try:
|
||||
return _CLAUDE_CODE_OBJECT_MAPPING_ADAPTER.validate_python(value)
|
||||
return _OBJECT_MAPPING_ADAPTER.validate_python(value)
|
||||
except ValidationError:
|
||||
return None
|
||||
|
||||
|
||||
def thinking_block_binding(thinking: Mapping[object, object] | None) -> dict[str, object]:
|
||||
block_binding: Final = thinking.get("block_binding") if thinking is not None else None
|
||||
return {} if block_binding is None else {"block_binding": block_binding}
|
||||
|
||||
|
||||
def _validated_claude_code_list(value: object) -> list[object] | None:
|
||||
try:
|
||||
return _CLAUDE_CODE_OBJECT_LIST_ADAPTER.validate_python(value)
|
||||
|
|
@ -194,7 +199,7 @@ def _claude_code_billing_texts(system: object) -> tuple[str, ...] | None:
|
|||
blocks: Final = _validated_claude_code_list(system)
|
||||
if blocks is None:
|
||||
return None
|
||||
block_mappings: Final = tuple(_validated_claude_code_mapping(block) for block in blocks)
|
||||
block_mappings: Final = tuple(validated_object_mapping(block) for block in blocks)
|
||||
if any(block is None for block in block_mappings):
|
||||
return None
|
||||
text_values: Final = tuple(
|
||||
|
|
@ -227,7 +232,7 @@ def is_claude_code_one_shot_subagent_request(
|
|||
tools: object,
|
||||
user_agent: str | None,
|
||||
) -> bool:
|
||||
only_message: Final = _validated_claude_code_mapping(messages[0]) if len(messages) == 1 else None
|
||||
only_message: Final = validated_object_mapping(messages[0]) if len(messages) == 1 else None
|
||||
return (
|
||||
user_agent is not None
|
||||
and is_claude_code_user_agent(user_agent)
|
||||
|
|
@ -914,6 +919,7 @@ class AnthropicModelInfo(BaseLLMModelInfo):
|
|||
model: str,
|
||||
optional_params: MutableMapping[str, object], # mutable-ok: in-place out-param like the sibling helpers
|
||||
custom_llm_provider: str,
|
||||
carry_block_binding: bool = False,
|
||||
) -> None:
|
||||
"""Translate legacy ``thinking.type=enabled`` to adaptive for the
|
||||
adaptive-thinking models that reject it (4.7+ and the 5 families).
|
||||
|
|
@ -926,21 +932,21 @@ class AnthropicModelInfo(BaseLLMModelInfo):
|
|||
return
|
||||
if AnthropicModelInfo._supports_legacy_thinking(model, custom_llm_provider):
|
||||
return
|
||||
thinking: Final = optional_params.get("thinking")
|
||||
if not isinstance(thinking, dict) or thinking.get("type") != "enabled":
|
||||
thinking: Final = validated_object_mapping(optional_params.get("thinking"))
|
||||
if thinking is None or thinking.get("type") != "enabled":
|
||||
return
|
||||
|
||||
budget_tokens: Final = thinking.get("budget_tokens") or 0
|
||||
effort: Final = AnthropicModelInfo._legacy_budget_to_effort(
|
||||
model=model,
|
||||
budget_tokens=int(thinking.get("budget_tokens") or 0),
|
||||
budget_tokens=int(budget_tokens) if isinstance(budget_tokens, (int, float, str)) else 0,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
)
|
||||
existing_output_config: Final = optional_params.get("output_config")
|
||||
display: Final = thinking.get("display")
|
||||
if display in ("summarized", "omitted"):
|
||||
optional_params["thinking"] = {"type": "adaptive", "display": display}
|
||||
else:
|
||||
optional_params["thinking"] = {"type": "adaptive"}
|
||||
supported_display: Final = {"display": display} if display in ("summarized", "omitted") else {}
|
||||
carried_block_binding: Final = thinking_block_binding(thinking) if carry_block_binding else {}
|
||||
optional_params["thinking"] = {**carried_block_binding, "type": "adaptive", **supported_display}
|
||||
optional_params["output_config"] = {
|
||||
"effort": effort,
|
||||
**(existing_output_config if isinstance(existing_output_config, dict) else MappingProxyType({})),
|
||||
|
|
@ -1608,6 +1614,20 @@ def is_anthropic_invalid_thinking_block_error(error_text: str) -> bool:
|
|||
return "must contain thinking" in lower
|
||||
|
||||
|
||||
def is_anthropic_thinking_binding_mismatch_error(error_text: str) -> bool:
|
||||
"""messages.N.content.M: Invalid `signature` in `thinking` block. The block is bound to a different conversation."""
|
||||
return "bound to a different conversation" in error_text.lower()
|
||||
|
||||
|
||||
def anthropic_request_errors_on_thinking_binding_mismatch(request_body: object) -> bool:
|
||||
"""``thinking.block_binding.prefix_mismatch_behavior: error`` is the caller asking
|
||||
for the 400, so a retry that strips thinking must not override it."""
|
||||
body: Final = validated_object_mapping(request_body)
|
||||
thinking: Final = validated_object_mapping(body.get("thinking")) if body is not None else None
|
||||
binding: Final = validated_object_mapping(thinking.get("block_binding")) if thinking is not None else None
|
||||
return binding is not None and binding.get("prefix_mismatch_behavior") == "error"
|
||||
|
||||
|
||||
def strip_thinking_blocks_from_anthropic_messages(messages: Sequence[object]) -> list[object]:
|
||||
"""
|
||||
Return a new message list with thinking / redacted_thinking content blocks removed
|
||||
|
|
@ -1678,7 +1698,7 @@ def strip_empty_content_blocks_from_anthropic_messages(
|
|||
) -> list[Any]:
|
||||
"""
|
||||
Return a new message list with empty or whitespace-only ``{"type": "text"}``
|
||||
and ``{"type": "thinking"}`` content blocks removed.
|
||||
content blocks and empty *unsigned* ``{"type": "thinking"}`` blocks removed.
|
||||
|
||||
Anthropic's API rejects requests containing such blocks with
|
||||
``"messages: text content blocks must be non-empty"`` and
|
||||
|
|
@ -1694,8 +1714,13 @@ def strip_empty_content_blocks_from_anthropic_messages(
|
|||
on the unified ``/v1/messages`` path. ``/v1/chat/completions`` already
|
||||
handles this in ``anthropic_messages_pt``; this helper provides the
|
||||
equivalent guarantee for the native Anthropic Messages path.
|
||||
``redacted_thinking`` blocks are never touched: they carry opaque
|
||||
``data`` instead of thinking text.
|
||||
|
||||
A thinking block with empty text but a non-empty ``signature`` is kept:
|
||||
it is the default response shape of Claude Opus 4.8 / Fable 5 / Fable 5.1
|
||||
(``thinking.display`` omitted), the signature is the model's preserved
|
||||
reasoning, and Anthropic accepts it on input. ``redacted_thinking``
|
||||
blocks are never touched: they carry opaque ``data`` instead of thinking
|
||||
text.
|
||||
|
||||
Messages whose content is a list and becomes empty after stripping are
|
||||
omitted, matching :func:`strip_thinking_blocks_from_anthropic_messages`.
|
||||
|
|
@ -1708,7 +1733,7 @@ def strip_empty_content_blocks_from_anthropic_messages(
|
|||
out.append(m)
|
||||
continue
|
||||
content = m["content"]
|
||||
filtered = [b for b in content if not _is_empty_text_block(b) and not is_empty_thinking_block(b)]
|
||||
filtered = [b for b in content if not _is_empty_text_block(b) and not is_empty_unsigned_thinking_block(b)]
|
||||
if len(filtered) == len(content):
|
||||
out.append(m)
|
||||
elif filtered:
|
||||
|
|
@ -1726,11 +1751,11 @@ def _is_empty_text_block(block: object) -> bool:
|
|||
def is_empty_thinking_block(block: object) -> bool:
|
||||
"""
|
||||
True for a ``{"type": "thinking"}`` content block whose thinking text is
|
||||
missing, not a string, or empty/whitespace-only after ``.strip()``.
|
||||
Anthropic rejects such blocks with ``"each thinking block must contain
|
||||
thinking"`` (whitespace-only included, verified live), regardless of any
|
||||
signature they carry. ``redacted_thinking`` blocks are a different type
|
||||
and always return False.
|
||||
missing, not a string, or empty/whitespace-only after ``.strip()``,
|
||||
whatever its signature. This is a shape check only; whether such a block
|
||||
may be sent to Anthropic depends on the signature, see
|
||||
:func:`is_empty_unsigned_thinking_block`. ``redacted_thinking`` blocks are
|
||||
a different type and always return False.
|
||||
"""
|
||||
if not isinstance(block, dict) or block.get("type") != "thinking":
|
||||
return False
|
||||
|
|
@ -1742,15 +1767,16 @@ def is_empty_unsigned_thinking_block(block: object) -> bool:
|
|||
"""
|
||||
True for an empty ``{"type": "thinking"}`` block carrying no signature.
|
||||
|
||||
The emit-side predicate: response paths drop a thinking block only when it
|
||||
holds nothing the client could need. A signature-only block is a real
|
||||
provider response (Bedrock Converse under adaptive thinking emits a
|
||||
reasoning block with empty text and only a signature) and the client needs
|
||||
the signature to replay reasoning across tool-use turns, so it must be
|
||||
emitted. Request paths keep using :func:`is_empty_thinking_block`:
|
||||
Anthropic rejects empty thinking blocks in request history regardless of
|
||||
signature, and the inbound strip self-heals a replayed signature-only
|
||||
block.
|
||||
Used on both the emit and the request side: a thinking block is dropped
|
||||
only when it holds nothing the provider or the client could use. A
|
||||
signature-only block is a real provider response (Claude Opus 4.8 /
|
||||
Fable 5 / Fable 5.1 return it by default, and Bedrock Converse under
|
||||
adaptive thinking emits a reasoning block with empty text and only a
|
||||
signature); the client needs the signature to replay reasoning across
|
||||
tool-use turns, and Anthropic accepts the replayed block and binds it to
|
||||
its prefix. Only an empty block with no signature (a turn a
|
||||
non-Anthropic reasoning model served through the bridge) is rejected
|
||||
with ``"each thinking block must contain thinking"`` and must go.
|
||||
"""
|
||||
if not isinstance(block, dict) or not is_empty_thinking_block(block):
|
||||
return False
|
||||
|
|
@ -1884,7 +1910,7 @@ def _replayed_server_tool_use(block: object) -> _ReplayedServerToolUse | None:
|
|||
|
||||
|
||||
def tool_call_is_rebuilt_as_server_tool_use(tool_call_id: object, provider_specific_fields: object) -> bool:
|
||||
fields: Final = _validated_claude_code_mapping(provider_specific_fields)
|
||||
fields: Final = validated_object_mapping(provider_specific_fields)
|
||||
if not isinstance(tool_call_id, str) or fields is None:
|
||||
return False
|
||||
return (
|
||||
|
|
|
|||
|
|
@ -4,11 +4,14 @@ Anthropic Token Counter implementation using the CountTokens API.
|
|||
|
||||
from typing import Any, Final
|
||||
|
||||
from pydantic import JsonValue
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.exceptions import AuthenticationError
|
||||
from litellm.llms.anthropic.count_tokens.handler import AnthropicCountTokensHandler
|
||||
from litellm.llms.base_llm.base_utils import BaseTokenCounter
|
||||
from litellm.types.utils import LlmProviders, TokenCountResponse
|
||||
from litellm.types.utils import LlmProviders, TokenCountMessagesFormat, TokenCountResponse
|
||||
|
||||
# Global handler instance - reuse across all token counting requests
|
||||
anthropic_count_tokens_handler: Final = AnthropicCountTokensHandler()
|
||||
|
|
@ -30,8 +33,10 @@ class AnthropicTokenCounter(BaseTokenCounter):
|
|||
contents: list[dict[str, object]] | None,
|
||||
deployment: dict[str, Any] | None = None,
|
||||
request_model: str = "",
|
||||
tools: list[dict[str, Any]] | None = None,
|
||||
system: Any | None = None,
|
||||
tools: list[dict[str, JsonValue]] | None = None,
|
||||
system: JsonValue = None,
|
||||
messages_format: TokenCountMessagesFormat = "anthropic",
|
||||
output_format: JsonValue = None,
|
||||
) -> TokenCountResponse | None:
|
||||
"""
|
||||
Count tokens using Anthropic's CountTokens API.
|
||||
|
|
@ -57,14 +62,24 @@ class AnthropicTokenCounter(BaseTokenCounter):
|
|||
|
||||
try:
|
||||
auth_header: Final = await AnthropicModelInfo.aget_auth_header(
|
||||
api_key=litellm_params.get("api_key"),
|
||||
api_key=litellm_params.get("api_key") or litellm.anthropic_key or litellm.api_key,
|
||||
api_base=api_base,
|
||||
litellm_params=litellm_params,
|
||||
allow_workload_identity=True,
|
||||
)
|
||||
if auth_header is None:
|
||||
verbose_logger.warning("No Anthropic credential found for token counting")
|
||||
return None
|
||||
return TokenCountResponse(
|
||||
total_tokens=0,
|
||||
request_model=request_model,
|
||||
model_used=model_to_use,
|
||||
tokenizer_type="anthropic_api",
|
||||
error=True,
|
||||
error_message=(
|
||||
f"No Anthropic API key is set for model {request_model or model_to_use}, so tokens can't be "
|
||||
"counted. Set api_key on the deployment or ANTHROPIC_API_KEY in the proxy environment."
|
||||
),
|
||||
status_code=401,
|
||||
)
|
||||
|
||||
result: Final = await anthropic_count_tokens_handler.handle_count_tokens_request(
|
||||
model=model_to_use,
|
||||
|
|
|
|||
|
|
@ -461,7 +461,7 @@ class LiteLLMAnthropicMessagesAdapter:
|
|||
*,
|
||||
custom_llm_provider: str | None = None,
|
||||
preserve_midturn_system: bool = False,
|
||||
) -> list:
|
||||
) -> list[AllMessageValues]:
|
||||
new_messages: Final[list[AllMessageValues]] = []
|
||||
replayable_messages: Final = strip_encrypted_reasoning_blocks_from_anthropic_messages(messages)
|
||||
leading_count: Final = next(
|
||||
|
|
|
|||
106
litellm/llms/anthropic/pass_through/convertible_blocks.py
Normal file
106
litellm/llms/anthropic/pass_through/convertible_blocks.py
Normal file
|
|
@ -0,0 +1,106 @@
|
|||
from collections.abc import Iterator, Mapping, Sequence
|
||||
from itertools import chain
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
from pydantic import JsonValue, TypeAdapter, ValidationError
|
||||
from typing_extensions import assert_never
|
||||
|
||||
from litellm.llms.anthropic.pass_through.messages.utils import normalize_native_messages
|
||||
from litellm.types.utils import TokenCountMessagesFormat
|
||||
|
||||
_CONVERTIBLE_BLOCK_TYPES: Final = MappingProxyType(
|
||||
{
|
||||
"user": frozenset({"text", "image", "document", "tool_result"}),
|
||||
"assistant": frozenset({"text", "tool_use", "thinking", "redacted_thinking"}),
|
||||
}
|
||||
)
|
||||
_SOURCED_BLOCK_TYPES: Final = frozenset({"image", "document"})
|
||||
_TOOL_RESULT_BLOCK_TYPES: Final = frozenset({"text", "image", "document"})
|
||||
|
||||
|
||||
_JSON_MESSAGES: Final = TypeAdapter(tuple[Mapping[str, JsonValue], ...])
|
||||
|
||||
|
||||
def _json_messages(messages: Sequence[Mapping[str, object]]) -> tuple[Mapping[str, JsonValue], ...]:
|
||||
try:
|
||||
return _JSON_MESSAGES.validate_python(messages)
|
||||
except ValidationError:
|
||||
return ()
|
||||
|
||||
|
||||
def _source_label(block_type: str, source: JsonValue) -> str | None:
|
||||
if not isinstance(source, dict):
|
||||
return f"{block_type} without a source"
|
||||
match source.get("type"):
|
||||
case "base64" if source.get("data"):
|
||||
return None
|
||||
case "url" if source.get("url"):
|
||||
return None
|
||||
case "base64" | "url" as source_type:
|
||||
return f"{block_type} with an empty {source_type} source"
|
||||
case str() as source_type:
|
||||
return f"{block_type} with a {source_type} source"
|
||||
case _:
|
||||
return f"{block_type} with an untyped source"
|
||||
|
||||
|
||||
def _block_label(block: JsonValue, convertible: frozenset[str]) -> str | None:
|
||||
if not isinstance(block, dict):
|
||||
return None
|
||||
block_type: Final = block.get("type")
|
||||
if not isinstance(block_type, str):
|
||||
return "a block without a type"
|
||||
if block_type not in convertible:
|
||||
return block_type
|
||||
return _source_label(block_type, block.get("source")) if block_type in _SOURCED_BLOCK_TYPES else None
|
||||
|
||||
|
||||
def _tool_result_labels(block: JsonValue) -> Iterator[str]:
|
||||
if not isinstance(block, dict) or block.get("type") != "tool_result":
|
||||
return iter(())
|
||||
inner: Final = block.get("content")
|
||||
if not isinstance(inner, list):
|
||||
return iter(())
|
||||
return (
|
||||
f"{label} inside a tool_result"
|
||||
for item in inner
|
||||
if (label := _block_label(item, _TOOL_RESULT_BLOCK_TYPES)) is not None
|
||||
)
|
||||
|
||||
|
||||
def _content_labels(block: JsonValue, convertible: frozenset[str]) -> Iterator[str]:
|
||||
label: Final = _block_label(block, convertible)
|
||||
return chain(() if label is None else (label,), _tool_result_labels(block))
|
||||
|
||||
|
||||
def _message_labels(message: Mapping[str, JsonValue]) -> Iterator[str]:
|
||||
role: Final = message.get("role")
|
||||
convertible: Final = _CONVERTIBLE_BLOCK_TYPES.get(role) if isinstance(role, str) else None
|
||||
content: Final = message.get("content")
|
||||
if convertible is None or not isinstance(content, list):
|
||||
return iter(())
|
||||
return chain.from_iterable(_content_labels(block, convertible) for block in content)
|
||||
|
||||
|
||||
def unconvertible_anthropic_blocks(messages: Sequence[Mapping[str, object]]) -> tuple[str, ...]:
|
||||
normalized: Final = _json_messages(normalize_native_messages(messages))
|
||||
return tuple(sorted(set(chain.from_iterable(_message_labels(message) for message in normalized))))
|
||||
|
||||
|
||||
def unconvertible_blocks_error(
|
||||
messages: Sequence[Mapping[str, object]] | None, messages_format: TokenCountMessagesFormat
|
||||
) -> str | None:
|
||||
match messages_format:
|
||||
case "openai":
|
||||
return None
|
||||
case "anthropic":
|
||||
labels: Final = unconvertible_anthropic_blocks(messages or ())
|
||||
return (
|
||||
"count_tokens can't convert these content blocks for this model, so the count would leave them out: "
|
||||
+ ", ".join(labels)
|
||||
if labels
|
||||
else None
|
||||
)
|
||||
case _:
|
||||
assert_never(messages_format)
|
||||
|
|
@ -80,6 +80,10 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig):
|
|||
def _resolved_provider(self) -> str:
|
||||
return self.custom_llm_provider or "anthropic"
|
||||
|
||||
@property
|
||||
def _carries_thinking_block_binding(self) -> bool:
|
||||
return type(self) is AnthropicMessagesConfig
|
||||
|
||||
def get_supported_anthropic_messages_params(self, model: str) -> list:
|
||||
return [
|
||||
"messages",
|
||||
|
|
@ -607,6 +611,7 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig):
|
|||
model=model,
|
||||
optional_params=anthropic_messages_optional_request_params,
|
||||
custom_llm_provider=self._resolved_provider,
|
||||
carry_block_binding=self._carries_thinking_block_binding,
|
||||
)
|
||||
|
||||
self._translate_adaptive_effort_for_non_adaptive_model(
|
||||
|
|
|
|||
|
|
@ -160,6 +160,13 @@ def prepare_native_messages(
|
|||
)
|
||||
|
||||
|
||||
def normalize_native_messages(messages: Sequence[object]) -> list[dict[str, object]]:
|
||||
"""Apply the block rewrites every /v1/messages request gets before it reaches a provider."""
|
||||
return flatten_unencrypted_web_search_results_in_anthropic_messages(
|
||||
sanitize_tool_use_ids_in_anthropic_messages(strip_empty_content_blocks_from_anthropic_messages(messages))
|
||||
)
|
||||
|
||||
|
||||
@lru_cache(maxsize=1)
|
||||
def anthropic_messages_optional_param_keys() -> frozenset[str]:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -495,6 +495,19 @@ class LiteLLMAnthropicToResponsesAPIAdapter:
|
|||
result["summary"] = "detailed"
|
||||
return result
|
||||
|
||||
@staticmethod
|
||||
def system_blocks_to_instructions(system: object) -> str:
|
||||
if not isinstance(system, list):
|
||||
return ""
|
||||
return "\n".join(
|
||||
text
|
||||
for block in system
|
||||
if isinstance(block, dict)
|
||||
and block.get("type") == "text"
|
||||
and isinstance(text := block.get("text"), str)
|
||||
and text
|
||||
)
|
||||
|
||||
def translate_request(
|
||||
self,
|
||||
anthropic_request: AnthropicMessagesRequest,
|
||||
|
|
@ -544,9 +557,7 @@ class LiteLLMAnthropicToResponsesAPIAdapter:
|
|||
if isinstance(system, str):
|
||||
responses_kwargs["instructions"] = system
|
||||
elif isinstance(system, list):
|
||||
responses_kwargs["instructions"] = "\n".join(
|
||||
filter(None, (b.get("text", "") for b in system if isinstance(b, dict) and b.get("type") == "text"))
|
||||
)
|
||||
responses_kwargs["instructions"] = self.system_blocks_to_instructions(system)
|
||||
|
||||
# max_tokens -> max_output_tokens
|
||||
max_tokens: Final = anthropic_request.get("max_tokens")
|
||||
|
|
|
|||
|
|
@ -29,7 +29,7 @@ class AzureAIAnthropicCountTokensHandler(AzureAIAnthropicCountTokensConfig):
|
|||
self,
|
||||
model: str,
|
||||
messages: list[dict[str, Any]],
|
||||
api_key: str,
|
||||
api_key: str | None,
|
||||
api_base: str,
|
||||
litellm_params: dict[str, object] | None = None,
|
||||
timeout: float | httpx.Timeout | None = None,
|
||||
|
|
@ -79,6 +79,15 @@ class AzureAIAnthropicCountTokensHandler(AzureAIAnthropicCountTokensConfig):
|
|||
api_key=api_key,
|
||||
litellm_params=litellm_params,
|
||||
)
|
||||
if not {"x-api-key", "Authorization"} & headers.keys():
|
||||
raise AnthropicError(
|
||||
status_code=401,
|
||||
message=(
|
||||
f"No Azure AI API key or Entra ID credential is set for model {model}, so tokens can't be "
|
||||
"counted. Set api_key or an Entra ID credential on the deployment, or AZURE_AI_API_KEY or "
|
||||
"AZURE_API_KEY in the proxy environment."
|
||||
),
|
||||
)
|
||||
|
||||
# Use LiteLLM's async httpx client
|
||||
async_client: Final = get_async_httpx_client(llm_provider=litellm.LlmProviders.AZURE_AI)
|
||||
|
|
|
|||
|
|
@ -5,17 +5,35 @@ Azure AI Anthropic Token Counter implementation using the CountTokens API.
|
|||
import os
|
||||
from typing import Any, Final
|
||||
|
||||
from pydantic import JsonValue
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.llms.azure_ai.anthropic.count_tokens.handler import (
|
||||
AzureAIAnthropicCountTokensHandler,
|
||||
)
|
||||
from litellm.llms.base_llm.base_utils import BaseTokenCounter
|
||||
from litellm.types.utils import LlmProviders, TokenCountResponse
|
||||
from litellm.types.utils import LlmProviders, TokenCountMessagesFormat, TokenCountResponse
|
||||
|
||||
# Global handler instance - reuse across all token counting requests
|
||||
azure_ai_anthropic_count_tokens_handler: Final = AzureAIAnthropicCountTokensHandler()
|
||||
|
||||
|
||||
def _missing_api_base_response(request_model: str, model_to_use: str) -> TokenCountResponse:
|
||||
return TokenCountResponse(
|
||||
total_tokens=0,
|
||||
request_model=request_model,
|
||||
model_used=model_to_use,
|
||||
tokenizer_type="azure_ai_anthropic_api",
|
||||
error=True,
|
||||
error_message=(
|
||||
f"No Azure AI API base is set for model {request_model or model_to_use}, so tokens can't be counted. "
|
||||
"Set api_base on the deployment or AZURE_AI_API_BASE or AZURE_API_BASE in the proxy environment."
|
||||
),
|
||||
status_code=500,
|
||||
)
|
||||
|
||||
|
||||
class AzureAIAnthropicTokenCounter(BaseTokenCounter):
|
||||
"""Token counter implementation for Azure AI Anthropic provider using the CountTokens API."""
|
||||
|
||||
|
|
@ -34,6 +52,8 @@ class AzureAIAnthropicTokenCounter(BaseTokenCounter):
|
|||
request_model: str = "",
|
||||
tools: list[dict[str, object]] | None = None,
|
||||
system: object | None = None,
|
||||
messages_format: TokenCountMessagesFormat = "anthropic",
|
||||
output_format: JsonValue = None,
|
||||
) -> TokenCountResponse | None:
|
||||
"""
|
||||
Count tokens using Azure AI Anthropic's CountTokens API.
|
||||
|
|
@ -61,18 +81,15 @@ class AzureAIAnthropicTokenCounter(BaseTokenCounter):
|
|||
if not api_key:
|
||||
api_key = os.getenv("AZURE_AI_API_KEY")
|
||||
|
||||
# Get API base from deployment config or environment
|
||||
api_base = litellm_params.get("api_base")
|
||||
if not api_base:
|
||||
api_base = os.getenv("AZURE_AI_API_BASE")
|
||||
|
||||
if not api_key:
|
||||
verbose_logger.warning("No Azure AI API key found for token counting")
|
||||
return None
|
||||
api_base: Final = (
|
||||
litellm_params.get("api_base")
|
||||
or os.getenv("AZURE_AI_API_BASE")
|
||||
or os.getenv("AZURE_API_BASE")
|
||||
or litellm.api_base
|
||||
)
|
||||
|
||||
if not api_base:
|
||||
verbose_logger.warning("No Azure AI API base found for token counting")
|
||||
return None
|
||||
return _missing_api_base_response(request_model, model_to_use)
|
||||
|
||||
try:
|
||||
result: Final = await azure_ai_anthropic_count_tokens_handler.handle_count_tokens_request(
|
||||
|
|
|
|||
|
|
@ -24,7 +24,7 @@ class AzureAIAnthropicCountTokensConfig(AnthropicCountTokensConfig):
|
|||
|
||||
def get_required_headers(
|
||||
self,
|
||||
api_key: str,
|
||||
api_key: str | None,
|
||||
litellm_params: dict[str, object] | None = None,
|
||||
) -> dict[str, str]:
|
||||
"""
|
||||
|
|
@ -45,12 +45,12 @@ class AzureAIAnthropicCountTokensConfig(AnthropicCountTokensConfig):
|
|||
"Content-Type": "application/json",
|
||||
"anthropic-version": "2023-06-01",
|
||||
"anthropic-beta": ANTHROPIC_TOKEN_COUNTING_BETA_VERSION,
|
||||
"x-api-key": api_key, # Azure AI Anthropic requires this header
|
||||
**({"x-api-key": api_key} if api_key else {}),
|
||||
}
|
||||
|
||||
# Also set up Azure auth headers for flexibility
|
||||
litellm_params = litellm_params or {}
|
||||
if "api_key" not in litellm_params:
|
||||
if api_key and "api_key" not in litellm_params:
|
||||
litellm_params["api_key"] = api_key
|
||||
|
||||
litellm_params_obj: Final = GenericLiteLLMParams.model_validate(litellm_params)
|
||||
|
|
@ -58,8 +58,9 @@ class AzureAIAnthropicCountTokensConfig(AnthropicCountTokensConfig):
|
|||
# Get Azure auth headers (api-key or Authorization)
|
||||
azure_headers = BaseAzureLLM.base_validate_azure_environment(headers={}, litellm_params=litellm_params_obj)
|
||||
|
||||
# Merge Azure auth headers
|
||||
headers.update(azure_headers)
|
||||
if "api-key" in azure_headers:
|
||||
headers.setdefault("x-api-key", azure_headers["api-key"])
|
||||
|
||||
return headers
|
||||
|
||||
|
|
|
|||
|
|
@ -22,6 +22,10 @@ class AzureAnthropicMessagesConfig(AnthropicMessagesConfig):
|
|||
def custom_llm_provider(self) -> str | None:
|
||||
return "azure_ai"
|
||||
|
||||
@property
|
||||
def _carries_thinking_block_binding(self) -> bool:
|
||||
return type(self) is AzureAnthropicMessagesConfig
|
||||
|
||||
def should_strip_billing_metadata(self) -> bool:
|
||||
return True
|
||||
|
||||
|
|
|
|||
|
|
@ -39,6 +39,10 @@ class AzureAnthropicConfig(AnthropicConfig):
|
|||
def custom_llm_provider(self) -> str | None:
|
||||
return "azure_ai"
|
||||
|
||||
@property
|
||||
def _carries_thinking_block_binding(self) -> bool:
|
||||
return type(self) is AzureAnthropicConfig
|
||||
|
||||
def should_strip_billing_metadata(self) -> bool:
|
||||
return True
|
||||
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
from abc import ABC, abstractmethod
|
||||
from collections.abc import AsyncIterator
|
||||
from typing import TYPE_CHECKING, Any
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
|
||||
import httpx
|
||||
|
||||
|
|
@ -179,16 +179,28 @@ class BaseAnthropicMessagesConfig(ABC):
|
|||
"""
|
||||
return 2
|
||||
|
||||
def should_retry_anthropic_messages_on_http_error(self, e: httpx.HTTPStatusError, litellm_params: dict) -> bool:
|
||||
def should_retry_anthropic_messages_on_http_error(
|
||||
self, e: httpx.HTTPStatusError, litellm_params: dict, request_body: object = None
|
||||
) -> bool:
|
||||
"""
|
||||
When True, async_anthropic_messages_handler will transform the request body
|
||||
and issue one more attempt (bounded by max_retry_on_anthropic_messages_http_error).
|
||||
A binding-mismatch 400 is not retried when the body's ``thinking.block_binding``
|
||||
asked for it, since the retry would strip the very blocks the caller pinned.
|
||||
"""
|
||||
from litellm.llms.anthropic.common_utils import (
|
||||
anthropic_request_errors_on_thinking_binding_mismatch,
|
||||
is_anthropic_invalid_thinking_block_error,
|
||||
is_anthropic_thinking_binding_mismatch_error,
|
||||
)
|
||||
|
||||
return e.response.status_code == 400 and is_anthropic_invalid_thinking_block_error(e.response.text)
|
||||
if e.response.status_code != 400:
|
||||
return False
|
||||
error_text: Final[str] = e.response.text
|
||||
caller_asked_for_this_error: Final = anthropic_request_errors_on_thinking_binding_mismatch(
|
||||
request_body
|
||||
) and is_anthropic_thinking_binding_mismatch_error(error_text)
|
||||
return not caller_asked_for_this_error and is_anthropic_invalid_thinking_block_error(error_text)
|
||||
|
||||
def transform_anthropic_messages_request_on_http_error(self, e: httpx.HTTPStatusError, request_data: dict) -> dict:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -9,11 +9,11 @@ from collections.abc import Mapping
|
|||
from typing import Any, Final
|
||||
|
||||
from openai.lib import _parsing, _pydantic # pyright: ignore[reportPrivateUsage] # SDK parser internals
|
||||
from pydantic import BaseModel
|
||||
from pydantic import BaseModel, JsonValue
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolCallChunk
|
||||
from litellm.types.utils import Message, ProviderSpecificModelInfo, TokenCountResponse
|
||||
from litellm.types.utils import Message, ProviderSpecificModelInfo, TokenCountMessagesFormat, TokenCountResponse
|
||||
|
||||
|
||||
class BaseTokenCounter(ABC):
|
||||
|
|
@ -27,6 +27,8 @@ class BaseTokenCounter(ABC):
|
|||
request_model: str = "",
|
||||
tools: list[dict[str, Any]] | None = None,
|
||||
system: Any | None = None,
|
||||
messages_format: TokenCountMessagesFormat = "openai",
|
||||
output_format: JsonValue = None,
|
||||
) -> TokenCountResponse | None:
|
||||
pass
|
||||
|
||||
|
|
|
|||
|
|
@ -12,7 +12,7 @@ from litellm.litellm_core_utils.prompt_templates.image_handling import (
|
|||
convert_url_to_base64,
|
||||
)
|
||||
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
|
||||
from litellm.llms.anthropic.common_utils import AnthropicModelInfo
|
||||
from litellm.llms.anthropic.common_utils import AnthropicModelInfo, thinking_block_binding
|
||||
from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation import (
|
||||
AmazonInvokeConfig,
|
||||
)
|
||||
|
|
@ -67,6 +67,10 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig):
|
|||
def custom_llm_provider(self) -> str | None:
|
||||
return "bedrock"
|
||||
|
||||
@property
|
||||
def _carries_thinking_block_binding(self) -> bool:
|
||||
return type(self) is AmazonAnthropicClaudeConfig
|
||||
|
||||
def should_strip_billing_metadata(self) -> bool:
|
||||
return True
|
||||
|
||||
|
|
@ -115,7 +119,10 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig):
|
|||
model = original_model
|
||||
|
||||
AnthropicModelInfo.translate_legacy_thinking_for_adaptive_model(
|
||||
model=original_model, optional_params=optional_params, custom_llm_provider="bedrock"
|
||||
model=original_model,
|
||||
optional_params=optional_params,
|
||||
custom_llm_provider="bedrock",
|
||||
carry_block_binding=self._carries_thinking_block_binding,
|
||||
)
|
||||
translated_thinking: Final = optional_params.get("thinking")
|
||||
if (
|
||||
|
|
@ -123,8 +130,11 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig):
|
|||
and isinstance(translated_thinking, dict)
|
||||
and translated_thinking.get("type") == "adaptive"
|
||||
):
|
||||
thinking_with_display: Final[AnthropicThinkingParam] = {"type": "adaptive", "display": "updates"}
|
||||
optional_params["thinking"] = thinking_with_display
|
||||
optional_params["thinking"] = {
|
||||
**(thinking_block_binding(translated_thinking) if self._carries_thinking_block_binding else {}),
|
||||
"type": "adaptive",
|
||||
"display": "updates",
|
||||
}
|
||||
|
||||
# The stub model hides the original model from the parent's forced-tool-use backstop
|
||||
response_format_tool_choice: Final = optional_params.get("tool_choice")
|
||||
|
|
|
|||
|
|
@ -16,7 +16,7 @@ from litellm.llms.bedrock.count_tokens.handler import BedrockCountTokensHandler
|
|||
from litellm.llms.bedrock.count_tokens.mantle_handler import BedrockMantleCountTokensHandler
|
||||
from litellm.llms.bedrock_mantle.common_utils import is_mantle_claude_model
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
|
||||
from litellm.types.utils import LiteLLMPydanticObjectBase, LlmProviders, TokenCountResponse
|
||||
from litellm.types.utils import LiteLLMPydanticObjectBase, LlmProviders, TokenCountMessagesFormat, TokenCountResponse
|
||||
|
||||
RUNTIME_TOKENIZER_TYPE: Final = "bedrock_api"
|
||||
MANTLE_TOKENIZER_TYPE: Final = "bedrock_mantle_api"
|
||||
|
|
@ -106,6 +106,8 @@ class BedrockTokenCounter(BaseTokenCounter):
|
|||
request_model: str = "",
|
||||
tools: Sequence[Mapping[str, object]] | None = None,
|
||||
system: object | None = None,
|
||||
messages_format: TokenCountMessagesFormat = "anthropic",
|
||||
output_format: JsonValue = None,
|
||||
) -> TokenCountResponse | None:
|
||||
if not messages:
|
||||
return None
|
||||
|
|
|
|||
502
litellm/llms/bedrock/count_tokens/kimi_k3.py
Normal file
502
litellm/llms/bedrock/count_tokens/kimi_k3.py
Normal file
|
|
@ -0,0 +1,502 @@
|
|||
"""
|
||||
Exact Bedrock token counts for Kimi K3, which has no CountTokens support.
|
||||
|
||||
Renders the Converse body litellm sends through Moonshot's Kimi K3 chat template and counts it with the
|
||||
official tokenizer, downloaded once from Hugging Face at a pinned revision and checked against its SHA-256.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import copy
|
||||
import hashlib
|
||||
import json
|
||||
import threading
|
||||
import time
|
||||
from collections.abc import Callable, Iterable, Mapping, Sequence
|
||||
from concurrent.futures import Future
|
||||
from dataclasses import dataclass
|
||||
from enum import Enum
|
||||
from itertools import chain, groupby
|
||||
from pathlib import Path
|
||||
from typing import Annotated, Final, Literal, TypeAlias
|
||||
|
||||
import regex
|
||||
import tiktoken
|
||||
from pydantic import BaseModel, ConfigDict, Field, JsonValue, ValidationError, model_validator
|
||||
from typing_extensions import Self
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.litellm_core_utils.prompt_templates.factory import (
|
||||
chat_completion_tools_for,
|
||||
is_flat_responses_tool,
|
||||
)
|
||||
from litellm.llms.anthropic.pass_through.adapters.transformation import LiteLLMAnthropicMessagesAdapter
|
||||
from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam
|
||||
|
||||
KIMI_K3_REPO: Final = "moonshotai/Kimi-K3"
|
||||
KIMI_K3_REVISION: Final = "f831ab66814297da540d832a5235f8e904f29d06"
|
||||
KIMI_K3_RANKS_FILE: Final = "tiktoken.model"
|
||||
KIMI_K3_RANKS_SHA256: Final = "b6c497a7469b33ced9c38afb1ad6e47f03f5e5dc05f15930799210ec050c5103"
|
||||
_PATTERN: Final = "|".join(
|
||||
(
|
||||
r"[\p{Han}]+",
|
||||
r"[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}&&[^\p{Han}]]*[\p{Ll}\p{Lm}\p{Lo}\p{M}&&[^\p{Han}]]+(?i:'s|'t|'re|'ve|'m|'ll|'d)?",
|
||||
r"[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}&&[^\p{Han}]]+[\p{Ll}\p{Lm}\p{Lo}\p{M}&&[^\p{Han}]]*(?i:'s|'t|'re|'ve|'m|'ll|'d)?",
|
||||
r"\p{N}{1,3}",
|
||||
r" ?[^\s\p{L}\p{N}]+[\r\n]*",
|
||||
r"\s*[\r\n]+",
|
||||
r"\s+(?!\S)",
|
||||
r"\s+",
|
||||
)
|
||||
)
|
||||
# Ends chunks only where the pattern above always splits, so their counts sum to the whole text's count while BPE
|
||||
# stays linear. Only a run of over 1024 characters with no such split, like 2000 spaces, is cut mid-piece.
|
||||
_CHUNK: Final = regex.compile(
|
||||
r"(?s).{1,1024}(?:(?<=\S)(?= \S)|(?<=\p{N})(?=\P{N})|(?<=[^\s\p{N}])(?=\p{N})|(?<=\p{L})(?=[^\p{L}\p{M}'])"
|
||||
r"|(?<=[\r\n])(?=\S)|\Z)|.{1024}\p{N}*"
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class KimiK3TokenizerSource:
|
||||
ranks_path: Callable[[], str]
|
||||
sha256: str
|
||||
|
||||
|
||||
def _hugging_face_ranks_path() -> str:
|
||||
import huggingface_hub
|
||||
|
||||
return huggingface_hub.hf_hub_download(
|
||||
repo_id=KIMI_K3_REPO,
|
||||
filename=KIMI_K3_RANKS_FILE,
|
||||
revision=KIMI_K3_REVISION,
|
||||
local_files_only=litellm.disable_hf_tokenizer_download is True,
|
||||
)
|
||||
|
||||
|
||||
HUGGING_FACE_SOURCE: Final = KimiK3TokenizerSource(ranks_path=_hugging_face_ranks_path, sha256=KIMI_K3_RANKS_SHA256)
|
||||
|
||||
|
||||
def _loaded(source: KimiK3TokenizerSource) -> tiktoken.Encoding | str:
|
||||
try:
|
||||
data: Final = Path(source.ranks_path()).read_bytes()
|
||||
except Exception as e: # noqa: BLE001 # a download or read failure of any kind falls back to the estimate
|
||||
return type(e).__name__
|
||||
if hashlib.sha256(data).hexdigest() != source.sha256:
|
||||
return "SHA-256 mismatch"
|
||||
rows: Final = (line.split() for line in data.splitlines() if line)
|
||||
return tiktoken.Encoding(
|
||||
name="kimi_k3",
|
||||
pat_str=_PATTERN,
|
||||
mergeable_ranks={base64.b64decode(token): int(rank) for token, rank in rows},
|
||||
special_tokens={},
|
||||
)
|
||||
|
||||
|
||||
def load_kimi_k3_encoding(source: KimiK3TokenizerSource) -> tiktoken.Encoding | None:
|
||||
loaded: Final = _loaded(source)
|
||||
if isinstance(loaded, tiktoken.Encoding):
|
||||
return loaded
|
||||
verbose_logger.warning(
|
||||
"Kimi K3 tokenizer unavailable (%s), so Kimi K3 token counts are estimated until a later count loads it. "
|
||||
"Allow downloads of %s from Hugging Face repo %s at revision %s, or put it in the Hugging Face cache",
|
||||
loaded,
|
||||
KIMI_K3_RANKS_FILE,
|
||||
KIMI_K3_REPO,
|
||||
KIMI_K3_REVISION,
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
def _started_load(source: KimiK3TokenizerSource) -> "Future[tiktoken.Encoding | None]":
|
||||
encoding: Final = Future[tiktoken.Encoding | None]()
|
||||
threading.Thread(
|
||||
target=lambda: encoding.set_result(load_kimi_k3_encoding(source)), name="litellm-kimi-k3-tokenizer", daemon=True
|
||||
).start()
|
||||
return encoding
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _EncodingLoad:
|
||||
encoding: "Future[tiktoken.Encoding | None]"
|
||||
started_at: float
|
||||
is_retry: bool
|
||||
|
||||
def retry_due(self, now: float, retry_after_seconds: float) -> bool:
|
||||
return self.encoding.done() and self.encoding.result() is None and now - self.started_at >= retry_after_seconds
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class KimiK3Encodings:
|
||||
source: KimiK3TokenizerSource
|
||||
retry_after_seconds: float
|
||||
clock: Callable[[], float] = time.monotonic
|
||||
current: _EncodingLoad | None = None
|
||||
|
||||
def _load(self) -> _EncodingLoad:
|
||||
now: Final = self.clock()
|
||||
current: Final = self.current
|
||||
if current is not None and not current.retry_due(now, self.retry_after_seconds):
|
||||
return current
|
||||
load: Final = _EncodingLoad(encoding=_started_load(self.source), started_at=now, is_retry=current is not None)
|
||||
self.current = load
|
||||
return load
|
||||
|
||||
async def encoding(self, wait_seconds: float) -> tiktoken.Encoding | None:
|
||||
load: Final = self._load()
|
||||
remaining: Final = 0 if load.is_retry else load.started_at + wait_seconds - self.clock()
|
||||
if load.encoding.done() or remaining <= 0:
|
||||
return load.encoding.result() if load.encoding.done() else None
|
||||
await asyncio.wait((asyncio.wrap_future(load.encoding),), timeout=remaining)
|
||||
if load.encoding.done():
|
||||
return load.encoding.result()
|
||||
verbose_logger.warning(
|
||||
"Kimi K3 tokenizer still loading after %s s, so Kimi K3 token counts are estimated until it loads. "
|
||||
"Check that %s can be downloaded from Hugging Face repo %s at revision %s, or put it in the Hugging Face cache",
|
||||
wait_seconds,
|
||||
KIMI_K3_RANKS_FILE,
|
||||
KIMI_K3_REPO,
|
||||
KIMI_K3_REVISION,
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
class _Converse(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="forbid")
|
||||
|
||||
|
||||
class _Text(_Converse):
|
||||
text: str
|
||||
|
||||
|
||||
class _CachePoint(_Converse):
|
||||
cache_point: JsonValue = Field(alias="cachePoint")
|
||||
|
||||
|
||||
class _ToolUse(_Converse):
|
||||
tool_use_id: str = Field(alias="toolUseId")
|
||||
name: str
|
||||
arguments: Mapping[str, JsonValue] = Field(alias="input")
|
||||
|
||||
|
||||
class _ToolUseBlock(_Converse):
|
||||
tool_use: _ToolUse = Field(alias="toolUse")
|
||||
|
||||
|
||||
class _ToolResult(_Converse):
|
||||
tool_use_id: str = Field(alias="toolUseId")
|
||||
content: tuple[_Text, ...]
|
||||
status: str | None = None
|
||||
|
||||
|
||||
class _ToolResultBlock(_Converse):
|
||||
tool_result: _ToolResult = Field(alias="toolResult")
|
||||
|
||||
|
||||
class _ReasoningText(_Converse):
|
||||
text: str
|
||||
signature: str | None = None
|
||||
|
||||
|
||||
class _Reasoning(_Converse):
|
||||
reasoning_text: _ReasoningText = Field(alias="reasoningText")
|
||||
|
||||
|
||||
class _ReasoningBlock(_Converse):
|
||||
reasoning_content: _Reasoning = Field(alias="reasoningContent")
|
||||
|
||||
|
||||
class _UserMessage(_Converse):
|
||||
role: Literal["user"]
|
||||
content: tuple[_Text | _ToolResultBlock | _CachePoint, ...]
|
||||
|
||||
|
||||
class _AssistantMessage(_Converse):
|
||||
role: Literal["assistant"]
|
||||
content: tuple[_Text | _ToolUseBlock | _ReasoningBlock | _CachePoint, ...]
|
||||
|
||||
def calls(self) -> tuple[_ToolUse, ...]:
|
||||
return tuple(block.tool_use for block in self.content if isinstance(block, _ToolUseBlock))
|
||||
|
||||
|
||||
class _InputSchema(_Converse):
|
||||
definition: JsonValue = Field(alias="json")
|
||||
|
||||
|
||||
class _ToolSpec(_Converse):
|
||||
name: str
|
||||
description: str | None = None
|
||||
input_schema: _InputSchema = Field(alias="inputSchema")
|
||||
|
||||
def declared(self) -> Mapping[str, JsonValue]:
|
||||
described: Final[Mapping[str, JsonValue]] = (
|
||||
{} if self.description is None else {"description": self.description}
|
||||
)
|
||||
return {"name": self.name, **described, "parameters": self.input_schema.definition}
|
||||
|
||||
|
||||
class _ToolSpecBlock(_Converse):
|
||||
tool_spec: _ToolSpec = Field(alias="toolSpec")
|
||||
|
||||
|
||||
class _ToolConfig(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="ignore")
|
||||
tools: tuple[_ToolSpecBlock | _CachePoint, ...]
|
||||
|
||||
|
||||
_Message: TypeAlias = Annotated[_UserMessage | _AssistantMessage, Field(discriminator="role")]
|
||||
|
||||
|
||||
def _result_ids(message: _Message) -> frozenset[str]:
|
||||
if isinstance(message, _AssistantMessage):
|
||||
return frozenset()
|
||||
return frozenset(block.tool_result.tool_use_id for block in message.content if isinstance(block, _ToolResultBlock))
|
||||
|
||||
|
||||
def _call_ids(message: _Message | None) -> frozenset[str]:
|
||||
if not isinstance(message, _AssistantMessage):
|
||||
return frozenset()
|
||||
return frozenset(call.tool_use_id for call in message.calls())
|
||||
|
||||
|
||||
class _ConverseRequest(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="ignore")
|
||||
messages: tuple[_Message, ...]
|
||||
system: tuple[_Text | _CachePoint, ...] = ()
|
||||
tool_config: _ToolConfig | None = Field(default=None, alias="toolConfig")
|
||||
|
||||
@model_validator(mode="after")
|
||||
def _results_answer_the_previous_calls(self) -> Self:
|
||||
pairs: Final = zip((None, *self.messages), self.messages, strict=False)
|
||||
if not all(_result_ids(message) <= _call_ids(previous) for previous, message in pairs):
|
||||
raise ValueError("a toolResult does not answer a toolUse of the previous assistant message")
|
||||
return self
|
||||
|
||||
|
||||
class _Control(Enum):
|
||||
OPEN = "<|open|>"
|
||||
CLOSE = "<|close|>"
|
||||
SEP = "<|sep|>"
|
||||
END_OF_MSG = "<|end_of_msg|>"
|
||||
|
||||
|
||||
_Segment: TypeAlias = str | _Control
|
||||
_Segments: TypeAlias = tuple[_Segment, ...]
|
||||
|
||||
|
||||
def _open(tag: str, attrs: Iterable[tuple[str, str]] = ()) -> _Segments:
|
||||
escaped: Final = (
|
||||
(f" {key}", '="', value.replace("&", "&").replace('"', """), '"') for key, value in attrs
|
||||
)
|
||||
return (_Control.OPEN, tag, *chain.from_iterable(escaped), _Control.SEP)
|
||||
|
||||
|
||||
def _close(tag: str) -> _Segments:
|
||||
return (_Control.CLOSE, tag, _Control.SEP)
|
||||
|
||||
|
||||
def _message(attrs: Iterable[tuple[str, str]], body: Iterable[_Segment]) -> _Segments:
|
||||
return (*_open("message", attrs), *body, *_close("message"), _Control.END_OF_MSG)
|
||||
|
||||
|
||||
def _tool_declare(specs: Sequence[_ToolSpec]) -> _Segments:
|
||||
tools: Final = json.dumps(
|
||||
[{"type": "function", "function": spec.declared()} for spec in specs],
|
||||
ensure_ascii=False,
|
||||
separators=(",", ":"),
|
||||
sort_keys=True,
|
||||
)
|
||||
return _message(
|
||||
(("role", "system"), ("type", "tool-declare")),
|
||||
(f"# Tools\nHere are the available tools, described in JSONSchema.\n\n```json\n{tools}\n```",),
|
||||
)
|
||||
|
||||
|
||||
def _text(blocks: Iterable[_Converse]) -> str:
|
||||
return "".join(block.text for block in blocks if isinstance(block, _Text))
|
||||
|
||||
|
||||
def _json_type(value: JsonValue) -> str:
|
||||
match value:
|
||||
case bool():
|
||||
return "boolean"
|
||||
case None:
|
||||
return "null"
|
||||
case int() | float():
|
||||
return "number"
|
||||
case str():
|
||||
return "string"
|
||||
case dict():
|
||||
return "object"
|
||||
case _:
|
||||
return "array"
|
||||
|
||||
|
||||
def _argument(key: str, value: JsonValue) -> _Segments:
|
||||
text: Final = value if isinstance(value, str) else json.dumps(value, ensure_ascii=False)
|
||||
return (*_open("argument", (("key", key), ("type", _json_type(value)))), text, *_close("argument"))
|
||||
|
||||
|
||||
def _call(index: int, call: _ToolUse) -> _Segments:
|
||||
arguments: Final = chain.from_iterable(_argument(key, value) for key, value in call.arguments.items())
|
||||
return (*_open("call", (("tool", call.name), ("index", str(index)))), *arguments, *_close("call"))
|
||||
|
||||
|
||||
def _tool_calls(calls: Sequence[_ToolUse]) -> _Segments:
|
||||
if not calls:
|
||||
return ()
|
||||
rendered: Final = chain.from_iterable(_call(index, call) for index, call in enumerate(calls, start=1))
|
||||
return (*_open("tools"), *rendered, *_close("tools"))
|
||||
|
||||
|
||||
def _assistant_message(message: _AssistantMessage) -> _Segments:
|
||||
reasoning: Final = "".join(
|
||||
block.reasoning_content.reasoning_text.text for block in message.content if isinstance(block, _ReasoningBlock)
|
||||
)
|
||||
return _message(
|
||||
(("role", "assistant"),),
|
||||
(
|
||||
*_open("think"),
|
||||
reasoning,
|
||||
*_close("think"),
|
||||
*_open("response"),
|
||||
_text(message.content),
|
||||
*_close("response"),
|
||||
*_tool_calls(message.calls()),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _tool_messages(results: Sequence[_ToolResult], previous: _AssistantMessage) -> _Segments:
|
||||
names: Final = {call.tool_use_id: call.name for call in previous.calls()}
|
||||
return tuple(
|
||||
chain.from_iterable(
|
||||
_message(
|
||||
(("role", "tool"), ("tool", names[result.tool_use_id]), ("index", str(index))),
|
||||
(_text(result.content),),
|
||||
)
|
||||
for index, result in enumerate(results, start=1)
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _user_run(blocks: Sequence[_Text | _ToolResultBlock], previous: _Message | None) -> _Segments:
|
||||
results: Final = tuple(block.tool_result for block in blocks if isinstance(block, _ToolResultBlock))
|
||||
if results and isinstance(previous, _AssistantMessage):
|
||||
return _tool_messages(results, previous)
|
||||
return _message((("role", "user"),), (_text(blocks),))
|
||||
|
||||
|
||||
def _user_messages(message: _UserMessage, previous: _Message | None) -> _Segments:
|
||||
blocks: Final = tuple(block for block in message.content if not isinstance(block, _CachePoint))
|
||||
runs: Final = groupby(blocks, key=lambda block: isinstance(block, _ToolResultBlock))
|
||||
return tuple(chain.from_iterable(_user_run(tuple(run), previous) for _, run in runs))
|
||||
|
||||
|
||||
def _turn(previous: _Message | None, message: _Message) -> _Segments:
|
||||
if isinstance(message, _AssistantMessage):
|
||||
return _assistant_message(message)
|
||||
return _user_messages(message, previous)
|
||||
|
||||
|
||||
def _rendered(request: _ConverseRequest) -> _Segments:
|
||||
specs: Final = (
|
||||
()
|
||||
if request.tool_config is None
|
||||
else tuple(block.tool_spec for block in request.tool_config.tools if isinstance(block, _ToolSpecBlock))
|
||||
)
|
||||
system: Final = _text(request.system)
|
||||
turns: Final = chain.from_iterable(
|
||||
_turn(previous, message) for previous, message in zip((None, *request.messages), request.messages, strict=False)
|
||||
)
|
||||
return (
|
||||
*(_tool_declare(specs) if specs else ()),
|
||||
*(_message((("role", "system"),), (system,)) if system else ()),
|
||||
*turns,
|
||||
*_open("message", (("role", "assistant"),)),
|
||||
*_open("think"),
|
||||
)
|
||||
|
||||
|
||||
def _text_tokens(encoding: tiktoken.Encoding, text: str) -> int:
|
||||
spans: Final = (match.span() for match in _CHUNK.finditer(text))
|
||||
return sum(len(encoding.encode_ordinary(text[start:end])) for start, end in spans)
|
||||
|
||||
|
||||
class _TextPart(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="ignore")
|
||||
type: Literal["text"]
|
||||
|
||||
|
||||
class _TextOnlyMessage(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="ignore")
|
||||
content: str | tuple[_TextPart, ...] | None = None
|
||||
|
||||
|
||||
def _text_only(message: AllMessageValues) -> bool:
|
||||
try:
|
||||
_TextOnlyMessage.model_validate(message)
|
||||
except ValidationError:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def _as_converse_input(tool: ChatCompletionToolParam) -> tuple[Mapping[str, object], ...]:
|
||||
return chat_completion_tools_for(tool) if is_flat_responses_tool(tool) else (tool,)
|
||||
|
||||
|
||||
def _converse_request(
|
||||
model: str, messages: Sequence[AllMessageValues], tools: Sequence[ChatCompletionToolParam] | None
|
||||
) -> _ConverseRequest | None:
|
||||
if not all(_text_only(message) for message in messages):
|
||||
return None
|
||||
try:
|
||||
chat_tools: Final = tuple(
|
||||
chat_tool for tool in tools or () for chat_tool in _as_converse_input(tool)
|
||||
) # comprehension-ok: flattens each tool's translation
|
||||
return _ConverseRequest.model_validate(
|
||||
litellm.AmazonConverseConfig().transform_request(
|
||||
model=model,
|
||||
messages=copy.deepcopy(list(messages)),
|
||||
optional_params={"tools": copy.deepcopy(list(chat_tools))} if chat_tools else {},
|
||||
litellm_params={},
|
||||
headers={},
|
||||
)
|
||||
)
|
||||
except Exception: # noqa: BLE001 # a body litellm cannot send to Bedrock, or cannot render, falls back to the estimate
|
||||
return None
|
||||
|
||||
|
||||
def kimi_k3_tokens(
|
||||
encoding: tiktoken.Encoding,
|
||||
model: str,
|
||||
messages: Sequence[AllMessageValues],
|
||||
tools: Sequence[ChatCompletionToolParam] | None,
|
||||
) -> int | None:
|
||||
request: Final = _converse_request(model, messages, tools)
|
||||
if request is None:
|
||||
return None
|
||||
return sum(
|
||||
1 if isinstance(segment, _Control) else _text_tokens(encoding, segment) for segment in _rendered(request)
|
||||
)
|
||||
|
||||
|
||||
def kimi_k3_anthropic_tokens(encoding: tiktoken.Encoding, model: str, fields: Mapping[str, object]) -> int | None:
|
||||
body: Final = {**fields, "model": model, "max_tokens": 1}
|
||||
try:
|
||||
translated, _ = LiteLLMAnthropicMessagesAdapter().translate_anthropic_to_openai(
|
||||
copy.deepcopy(body), # pyright: ignore[reportArgumentType] # raw /v1/messages fields
|
||||
custom_llm_provider="bedrock",
|
||||
)
|
||||
except Exception: # noqa: BLE001 # a request the adapter cannot translate falls back to the estimate
|
||||
return None
|
||||
match translated:
|
||||
case {"tools": tools}:
|
||||
return kimi_k3_tokens(encoding, model, translated["messages"], tools)
|
||||
case _:
|
||||
return kimi_k3_tokens(encoding, model, translated["messages"], None)
|
||||
|
||||
|
||||
def kimi_k3_template_tokens(message_count: int, has_tools: bool) -> int:
|
||||
return 14 + 13 * max(message_count - 1, 0) + (40 if has_tools else 0)
|
||||
|
|
@ -17,7 +17,11 @@ from litellm.llms.anthropic.chat.transformation import (
|
|||
DROP_UNSUPPORTED_OUTPUT_CONFIG_WARNING,
|
||||
AnthropicConfig,
|
||||
)
|
||||
from litellm.llms.anthropic.common_utils import AnthropicModelInfo
|
||||
from litellm.llms.anthropic.common_utils import (
|
||||
AnthropicModelInfo,
|
||||
thinking_block_binding,
|
||||
validated_object_mapping,
|
||||
)
|
||||
from litellm.llms.anthropic.pass_through.messages.transformation import (
|
||||
AnthropicMessagesConfig,
|
||||
)
|
||||
|
|
@ -49,7 +53,6 @@ from litellm.types.llms.anthropic import (
|
|||
ANTHROPIC_BETA_HEADER_VALUES,
|
||||
ANTHROPIC_FINE_GRAINED_TOOL_STREAMING_BETA_HEADER,
|
||||
ANTHROPIC_TOOL_SEARCH_BETA_HEADER,
|
||||
AnthropicThinkingParam,
|
||||
)
|
||||
from litellm.types.llms.bedrock import BedrockInvokeAnthropicMessagesRequest
|
||||
from litellm.types.llms.openai import AllMessageValues
|
||||
|
|
@ -82,6 +85,10 @@ class AmazonAnthropicClaudeMessagesConfig(
|
|||
def custom_llm_provider(self) -> str | None:
|
||||
return "bedrock"
|
||||
|
||||
@property
|
||||
def _carries_thinking_block_binding(self) -> bool:
|
||||
return type(self) is AmazonAnthropicClaudeMessagesConfig
|
||||
|
||||
@property
|
||||
def beta_headers_provider(self) -> str:
|
||||
return self.custom_llm_provider or "bedrock"
|
||||
|
|
@ -349,18 +356,14 @@ class AmazonAnthropicClaudeMessagesConfig(
|
|||
if not isinstance(output_config, dict):
|
||||
output_config = {}
|
||||
output_config.setdefault("effort", self._effort_from_thinking_budget(budget_tokens))
|
||||
thinking: Final = anthropic_messages_request.get("thinking")
|
||||
display: Final = thinking.get("display") if isinstance(thinking, dict) else None
|
||||
thinking: Final = validated_object_mapping(anthropic_messages_request.get("thinking"))
|
||||
display: Final = thinking.get("display") if thinking is not None else None
|
||||
anthropic_messages_request["output_config"] = output_config
|
||||
if display is None:
|
||||
adaptive_thinking: Final[AnthropicThinkingParam] = {"type": "adaptive"}
|
||||
anthropic_messages_request["thinking"] = adaptive_thinking
|
||||
else:
|
||||
adaptive_thinking_with_display: Final[AnthropicThinkingParam] = {
|
||||
"type": "adaptive",
|
||||
"display": display,
|
||||
}
|
||||
anthropic_messages_request["thinking"] = adaptive_thinking_with_display
|
||||
anthropic_messages_request["thinking"] = {
|
||||
**(thinking_block_binding(thinking) if self._carries_thinking_block_binding else {}),
|
||||
"type": "adaptive",
|
||||
**({} if display is None else {"display": display}),
|
||||
}
|
||||
verbose_logger.debug(
|
||||
"Bedrock clear_thinking_20251015: injected adaptive thinking with effort=%s for model=%s",
|
||||
output_config.get("effort"),
|
||||
|
|
@ -701,8 +704,11 @@ class AmazonAnthropicClaudeMessagesConfig(
|
|||
and isinstance(translated_thinking, dict)
|
||||
and translated_thinking.get("type") == "adaptive"
|
||||
):
|
||||
thinking_with_display: Final[AnthropicThinkingParam] = {"type": "adaptive", "display": "updates"}
|
||||
anthropic_messages_request["thinking"] = thinking_with_display
|
||||
anthropic_messages_request["thinking"] = {
|
||||
**(thinking_block_binding(translated_thinking) if self._carries_thinking_block_binding else {}),
|
||||
"type": "adaptive",
|
||||
"display": "updates",
|
||||
}
|
||||
self._normalize_system_role_messages(anthropic_messages_request, model=model)
|
||||
#########################################################
|
||||
############## BEDROCK Invoke SPECIFIC TRANSFORMATION ###
|
||||
|
|
|
|||
|
|
@ -2091,7 +2091,7 @@ class BaseLLMHTTPHandler:
|
|||
# str when the caller passes a pre-serialized (unsigned) body to avoid
|
||||
# re-dumping; bytes when a provider signed the request (e.g. Bedrock).
|
||||
signed_json_body: str | bytes | None,
|
||||
request_body: dict,
|
||||
request_body: dict[str, object],
|
||||
stream: bool,
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
provider_config: BaseAnthropicMessagesConfig,
|
||||
|
|
@ -2118,7 +2118,7 @@ class BaseLLMHTTPHandler:
|
|||
except httpx.HTTPStatusError as e:
|
||||
hit_max_attempt = attempt_idx + 1 == max_attempts
|
||||
should_retry = provider_config.should_retry_anthropic_messages_on_http_error(
|
||||
e=e, litellm_params=litellm_params_dict
|
||||
e=e, litellm_params=litellm_params_dict, request_body=request_body
|
||||
)
|
||||
if should_retry and not hit_max_attempt:
|
||||
if logging_obj.baseline_cache_context is not None:
|
||||
|
|
|
|||
|
|
@ -6,7 +6,7 @@ from collections.abc import Mapping, Sequence
|
|||
from typing import Any, Final
|
||||
|
||||
import httpx
|
||||
from pydantic import TypeAdapter
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
|
||||
import litellm
|
||||
from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH
|
||||
|
|
@ -14,7 +14,7 @@ from litellm.llms.base_llm.base_utils import BaseLLMModelInfo, BaseTokenCounter
|
|||
from litellm.llms.base_llm.chat.transformation import BaseLLMException
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
from litellm.types.llms.openai import AllMessageValues
|
||||
from litellm.types.utils import TokenCountResponse
|
||||
from litellm.types.utils import TokenCountMessagesFormat, TokenCountResponse
|
||||
|
||||
GEMINI_IMAGE_ASPECT_RATIOS: Final[dict[str, float]] = {
|
||||
"1:1": 1 / 1,
|
||||
|
|
@ -495,21 +495,54 @@ class GoogleAIStudioTokenCounter(BaseTokenCounter):
|
|||
request_model: str = "",
|
||||
tools: list[dict[str, object]] | None = None,
|
||||
system: object | None = None,
|
||||
messages_format: TokenCountMessagesFormat = "openai",
|
||||
output_format: JsonValue = None,
|
||||
) -> TokenCountResponse | None:
|
||||
import copy
|
||||
|
||||
from litellm.llms.gemini.count_tokens.handler import GoogleAIStudioTokenCounter
|
||||
from litellm.llms.vertex_ai.count_tokens.context import (
|
||||
VertexGeminiCountContextError,
|
||||
count_context_error_response,
|
||||
gemini_count_contents,
|
||||
inference_request_view,
|
||||
provider_error_response,
|
||||
unconvertible_messages_error,
|
||||
vertex_gemini_count_context,
|
||||
)
|
||||
|
||||
request_view: Final = inference_request_view(messages or [], system, messages_format)
|
||||
count_context: Final = vertex_gemini_count_context(system=request_view.system, tools=tools)
|
||||
if isinstance(count_context, VertexGeminiCountContextError):
|
||||
return count_context_error_response(count_context, request_model, model_to_use)
|
||||
blocks_error: Final = unconvertible_messages_error(request_view.messages, messages_format)
|
||||
if blocks_error is not None:
|
||||
return count_context_error_response(blocks_error, request_model, model_to_use)
|
||||
|
||||
deployment = deployment or {}
|
||||
count_tokens_params_request: Final = copy.deepcopy(deployment.get("litellm_params", {}))
|
||||
resolved_contents: Final = (
|
||||
contents
|
||||
if contents is not None
|
||||
else await gemini_count_contents(
|
||||
request_view.messages, messages_format, model_to_use, "gemini", count_tokens_params_request
|
||||
)
|
||||
)
|
||||
if isinstance(resolved_contents, VertexGeminiCountContextError):
|
||||
return count_context_error_response(resolved_contents, request_model, model_to_use)
|
||||
count_tokens_params: Final = {
|
||||
"model": model_to_use,
|
||||
"contents": contents,
|
||||
"contents": resolved_contents,
|
||||
"system_instruction": count_context.system_instruction,
|
||||
"tools": count_context.tools,
|
||||
}
|
||||
count_tokens_params_request.update(count_tokens_params)
|
||||
result: Final = await GoogleAIStudioTokenCounter().acount_tokens(
|
||||
**count_tokens_params_request,
|
||||
)
|
||||
try:
|
||||
result: Final = await GoogleAIStudioTokenCounter().acount_tokens(
|
||||
**count_tokens_params_request,
|
||||
)
|
||||
except litellm.APIError as error:
|
||||
return provider_error_response(error, request_model, model_to_use)
|
||||
|
||||
if result is not None:
|
||||
return TokenCountResponse(
|
||||
|
|
|
|||
|
|
@ -1,9 +1,10 @@
|
|||
from collections.abc import Mapping, Sequence
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
|
||||
import httpx
|
||||
|
||||
import litellm
|
||||
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, get_async_httpx_client
|
||||
from litellm.types.utils import LlmProviders
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
|
@ -28,12 +29,12 @@ class GoogleAIStudioTokenCounter:
|
|||
"""
|
||||
import copy
|
||||
|
||||
from google.genai.types import FunctionResponse
|
||||
|
||||
# Handle None or empty contents
|
||||
if not contents:
|
||||
return contents
|
||||
|
||||
from google.genai.types import FunctionResponse
|
||||
|
||||
cleaned_contents: Final = copy.deepcopy(contents)
|
||||
|
||||
for content in cleaned_contents:
|
||||
|
|
@ -54,6 +55,28 @@ class GoogleAIStudioTokenCounter:
|
|||
base_url: Final = api_base or "https://generativelanguage.googleapis.com"
|
||||
return f"{base_url}/v1beta/models/{model}:countTokens"
|
||||
|
||||
def __init__(self, client: AsyncHTTPHandler | None = None) -> None:
|
||||
super().__init__()
|
||||
self._client: Final = client
|
||||
|
||||
def _build_count_tokens_request_body(
|
||||
self,
|
||||
model: str,
|
||||
contents: object,
|
||||
system_instruction: Mapping[str, object] | None,
|
||||
tools: Sequence[Mapping[str, object]] | None,
|
||||
) -> Mapping[str, object]:
|
||||
if system_instruction is None and not tools:
|
||||
return {"contents": contents}
|
||||
|
||||
generate_content_request: Final = {
|
||||
"model": model if model.startswith("models/") else f"models/{model}",
|
||||
"contents": contents,
|
||||
**({"systemInstruction": dict(system_instruction)} if system_instruction is not None else {}),
|
||||
**({"tools": [dict(tool) for tool in tools]} if tools else {}),
|
||||
}
|
||||
return {"generateContentRequest": generate_content_request}
|
||||
|
||||
async def validate_environment(
|
||||
self,
|
||||
api_base: str | None = None,
|
||||
|
|
@ -84,6 +107,8 @@ class GoogleAIStudioTokenCounter:
|
|||
api_key: str | None = None,
|
||||
api_base: str | None = None,
|
||||
timeout: float | httpx.Timeout | None = None,
|
||||
system_instruction: Mapping[str, object] | None = None,
|
||||
tools: Sequence[Mapping[str, object]] | None = None,
|
||||
**kwargs: object,
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
|
|
@ -130,14 +155,19 @@ class GoogleAIStudioTokenCounter:
|
|||
|
||||
# Prepare request body - clean up contents to remove unsupported fields
|
||||
cleaned_contents: Final = self._clean_contents_for_gemini_api(contents)
|
||||
request_body: Final = {"contents": cleaned_contents}
|
||||
request_body: Final = self._build_count_tokens_request_body(
|
||||
model=model,
|
||||
contents=cleaned_contents,
|
||||
system_instruction=system_instruction,
|
||||
tools=tools,
|
||||
)
|
||||
|
||||
async_httpx_client: Final = get_async_httpx_client(
|
||||
llm_provider=LlmProviders.GEMINI,
|
||||
async_httpx_client: Final = (
|
||||
self._client if self._client is not None else get_async_httpx_client(llm_provider=LlmProviders.GEMINI)
|
||||
)
|
||||
|
||||
try:
|
||||
response: Final = await async_httpx_client.post(url=url, headers=headers, json=request_body)
|
||||
response: Final = await async_httpx_client.post(url=url, headers=headers, json=dict(request_body))
|
||||
|
||||
# Check for HTTP errors
|
||||
response.raise_for_status()
|
||||
|
|
|
|||
|
|
@ -5,8 +5,8 @@ Uses httpx for HTTP requests to OpenAI's /v1/responses/input_tokens endpoint.
|
|||
"""
|
||||
|
||||
import json
|
||||
from collections.abc import Sequence
|
||||
from typing import Any, Final
|
||||
from collections.abc import Mapping, Sequence
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
|
||||
|
|
@ -31,9 +31,10 @@ class OpenAICountTokensHandler(OpenAICountTokensConfig):
|
|||
api_key: str,
|
||||
api_base: str | None = None,
|
||||
timeout: float | httpx.Timeout | None = None,
|
||||
tools: list[dict[str, object]] | None = None,
|
||||
tools: Sequence[Mapping[str, object]] | None = None,
|
||||
instructions: str | None = None,
|
||||
) -> dict[str, Any]:
|
||||
text: Mapping[str, object] | None = None,
|
||||
) -> dict[str, object]:
|
||||
"""
|
||||
Handle a token counting request to OpenAI's Responses API.
|
||||
|
||||
|
|
@ -53,6 +54,7 @@ class OpenAICountTokensHandler(OpenAICountTokensConfig):
|
|||
input=input,
|
||||
tools=tools,
|
||||
instructions=instructions,
|
||||
text=text,
|
||||
)
|
||||
|
||||
endpoint_url: Final = self.get_openai_count_tokens_endpoint(api_base)
|
||||
|
|
|
|||
|
|
@ -2,24 +2,47 @@
|
|||
OpenAI Token Counter implementation using the Responses API /input_tokens endpoint.
|
||||
"""
|
||||
|
||||
import os
|
||||
from collections.abc import Mapping, Sequence
|
||||
from typing import Any, Final
|
||||
|
||||
import openai
|
||||
from pydantic import JsonValue
|
||||
from typing_extensions import assert_never
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.litellm_core_utils.asyncify import asyncify
|
||||
from litellm.llms.anthropic.pass_through.convertible_blocks import unconvertible_blocks_error
|
||||
from litellm.llms.anthropic.pass_through.responses_adapters.transformation import (
|
||||
LiteLLMAnthropicToResponsesAPIAdapter,
|
||||
)
|
||||
from litellm.llms.base_llm.base_utils import BaseTokenCounter
|
||||
from litellm.llms.openai.common_utils import OpenAIError
|
||||
from litellm.llms.openai.common_utils import OpenAIError, get_openai_credentials
|
||||
from litellm.llms.openai.responses.count_tokens.handler import (
|
||||
OpenAICountTokensHandler,
|
||||
)
|
||||
from litellm.llms.openai.responses.count_tokens.transformation import (
|
||||
OpenAICountTokensConfig,
|
||||
ResponsesCountRequestParts,
|
||||
)
|
||||
from litellm.types.utils import LlmProviders, TokenCountResponse
|
||||
from litellm.llms.openai.workload_identity import (
|
||||
OpenAIWorkloadIdentityConfig,
|
||||
get_workload_identity_bearer_token,
|
||||
resolve_openai_workload_identity_config,
|
||||
)
|
||||
from litellm.types.utils import LlmProviders, TokenCountMessagesFormat, TokenCountResponse
|
||||
|
||||
# Global handler instance - reuse across all token counting requests
|
||||
openai_count_tokens_handler: Final = OpenAICountTokensHandler()
|
||||
|
||||
|
||||
def _workload_identity_config(litellm_params: Mapping[str, object]) -> OpenAIWorkloadIdentityConfig | None:
|
||||
api_base: Final = litellm_params.get("api_base")
|
||||
return resolve_openai_workload_identity_config(
|
||||
api_key=None,
|
||||
api_base=api_base if isinstance(api_base, str) else None,
|
||||
)
|
||||
|
||||
|
||||
class OpenAITokenCounter(BaseTokenCounter):
|
||||
"""Token counter implementation for OpenAI provider using the Responses API."""
|
||||
|
||||
|
|
@ -29,6 +52,37 @@ class OpenAITokenCounter(BaseTokenCounter):
|
|||
) -> bool:
|
||||
return custom_llm_provider == LlmProviders.OPENAI.value
|
||||
|
||||
@staticmethod
|
||||
def _system_instructions(system: object) -> str | None:
|
||||
if system is None or isinstance(system, str):
|
||||
return system
|
||||
return LiteLLMAnthropicToResponsesAPIAdapter.system_blocks_to_instructions(system) or None
|
||||
|
||||
@staticmethod
|
||||
def _responses_request_parts(
|
||||
model_to_use: str,
|
||||
messages: Sequence[Mapping[str, object]],
|
||||
tools: Sequence[Mapping[str, object]] | None,
|
||||
system: object,
|
||||
messages_format: TokenCountMessagesFormat,
|
||||
output_format: JsonValue,
|
||||
) -> ResponsesCountRequestParts:
|
||||
match messages_format:
|
||||
case "anthropic":
|
||||
return OpenAICountTokensConfig.anthropic_messages_to_responses_input(
|
||||
model=model_to_use, messages=messages, system=system, tools=tools, output_format=output_format
|
||||
)
|
||||
case "openai":
|
||||
input_items, instructions = OpenAICountTokensConfig.messages_to_responses_input(
|
||||
[dict(message) for message in messages]
|
||||
)
|
||||
fallback_instructions: Final = OpenAITokenCounter._system_instructions(system)
|
||||
return OpenAICountTokensConfig.count_request_parts(
|
||||
input_items, instructions if instructions is not None else fallback_instructions, tools
|
||||
)
|
||||
case _:
|
||||
assert_never(messages_format)
|
||||
|
||||
async def count_tokens(
|
||||
self,
|
||||
model_to_use: str,
|
||||
|
|
@ -38,6 +92,8 @@ class OpenAITokenCounter(BaseTokenCounter):
|
|||
request_model: str = "",
|
||||
tools: list[dict[str, object]] | None = None,
|
||||
system: object = None,
|
||||
messages_format: TokenCountMessagesFormat = "openai",
|
||||
output_format: JsonValue = None,
|
||||
) -> TokenCountResponse | None:
|
||||
"""
|
||||
Count tokens using OpenAI's Responses API /input_tokens endpoint.
|
||||
|
|
@ -45,39 +101,76 @@ class OpenAITokenCounter(BaseTokenCounter):
|
|||
if not messages:
|
||||
return None
|
||||
|
||||
malformed_error: Final = (
|
||||
OpenAICountTokensConfig.malformed_anthropic_messages_error(messages)
|
||||
if messages_format == "anthropic"
|
||||
else None
|
||||
)
|
||||
request_error: Final = malformed_error or unconvertible_blocks_error(messages, messages_format)
|
||||
if request_error is not None:
|
||||
return TokenCountResponse(
|
||||
total_tokens=0,
|
||||
request_model=request_model,
|
||||
model_used=model_to_use,
|
||||
tokenizer_type="openai_api",
|
||||
error=True,
|
||||
error_message=request_error,
|
||||
status_code=400,
|
||||
rejected_before_provider_call=True,
|
||||
)
|
||||
|
||||
deployment = deployment or {}
|
||||
litellm_params: Final = deployment.get("litellm_params", {})
|
||||
|
||||
# Get OpenAI API key from deployment config or environment
|
||||
api_key = litellm_params.get("api_key")
|
||||
if not api_key:
|
||||
api_key = os.getenv("OPENAI_API_KEY")
|
||||
credentials: Final = get_openai_credentials(
|
||||
api_base=litellm_params.get("api_base"), api_key=litellm_params.get("api_key")
|
||||
)
|
||||
credential: Final = credentials.api_key or _workload_identity_config(litellm_params)
|
||||
|
||||
if not api_key:
|
||||
verbose_logger.warning("No OpenAI API key found for token counting")
|
||||
return None
|
||||
if credential is None:
|
||||
return TokenCountResponse(
|
||||
total_tokens=0,
|
||||
request_model=request_model,
|
||||
model_used=model_to_use,
|
||||
tokenizer_type="openai_api",
|
||||
error=True,
|
||||
error_message=(
|
||||
f"No OpenAI API key is set for model {request_model or model_to_use}, so tokens can't be "
|
||||
"counted. Set api_key or workload identity on the deployment, or OPENAI_API_KEY in the proxy "
|
||||
"environment."
|
||||
),
|
||||
status_code=401,
|
||||
)
|
||||
|
||||
api_base: Final = litellm_params.get("api_base")
|
||||
api_base: Final = credentials.api_base
|
||||
|
||||
# Convert chat messages to Responses API input format
|
||||
input_items, instructions = OpenAICountTokensConfig.messages_to_responses_input(messages)
|
||||
|
||||
# Use system param if instructions not extracted from messages
|
||||
if instructions is None and system is not None:
|
||||
instructions = system if isinstance(system, str) else str(system)
|
||||
input_items, instructions, counted_tools, text = self._responses_request_parts(
|
||||
model_to_use=model_to_use,
|
||||
messages=messages,
|
||||
tools=tools,
|
||||
system=system,
|
||||
messages_format=messages_format,
|
||||
output_format=output_format,
|
||||
)
|
||||
|
||||
# If no input items were produced (e.g., system-only messages), fall back to local counting
|
||||
if not input_items:
|
||||
return None
|
||||
|
||||
try:
|
||||
api_key: Final = (
|
||||
credential
|
||||
if isinstance(credential, str)
|
||||
else await asyncify(get_workload_identity_bearer_token)(credential)
|
||||
)
|
||||
result: Final = await openai_count_tokens_handler.handle_count_tokens_request(
|
||||
model=model_to_use,
|
||||
input=input_items if input_items is not None else [],
|
||||
input=input_items,
|
||||
api_key=api_key,
|
||||
api_base=api_base,
|
||||
tools=tools,
|
||||
tools=counted_tools,
|
||||
instructions=instructions,
|
||||
text=text,
|
||||
)
|
||||
|
||||
if result is not None:
|
||||
|
|
@ -99,6 +192,17 @@ class OpenAITokenCounter(BaseTokenCounter):
|
|||
error_message=e.message,
|
||||
status_code=e.status_code,
|
||||
)
|
||||
except openai.AuthenticationError as e:
|
||||
verbose_logger.warning("OpenAI workload identity token exchange failed: %s", e.message)
|
||||
return TokenCountResponse(
|
||||
total_tokens=0,
|
||||
request_model=request_model,
|
||||
model_used=model_to_use,
|
||||
tokenizer_type="openai_api",
|
||||
error=True,
|
||||
error_message=e.message,
|
||||
status_code=401,
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_logger.warning("Error calling OpenAI CountTokens API: %s", e)
|
||||
return TokenCountResponse(
|
||||
|
|
|
|||
|
|
@ -5,10 +5,18 @@ This module handles the transformation of requests to OpenAI's /v1/responses/inp
|
|||
"""
|
||||
|
||||
from collections.abc import Mapping, Sequence
|
||||
from typing import Any, Final, Literal
|
||||
from typing import Any, Final, Literal, TypeAlias
|
||||
|
||||
from pydantic import JsonValue, TypeAdapter, ValidationError
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
from litellm.llms.anthropic.pass_through.messages.utils import normalize_native_messages
|
||||
from litellm.llms.anthropic.pass_through.responses_adapters.transformation import (
|
||||
LiteLLMAnthropicToResponsesAPIAdapter,
|
||||
)
|
||||
from litellm.responses.utils import ResponsesAPIRequestUtils
|
||||
from litellm.types.llms.anthropic import AnthropicMessagesRequest, AnthropicOutputSchema
|
||||
|
||||
|
||||
class ResponsesInputTextPart(TypedDict):
|
||||
type: ReadOnly[Literal["input_text"]]
|
||||
|
|
@ -31,6 +39,30 @@ ResponsesInputPart = ResponsesInputTextPart | ResponsesInputImagePart | Response
|
|||
|
||||
ResponsesContentRole = Literal["user", "assistant"]
|
||||
|
||||
ResponsesCountRequestParts: TypeAlias = tuple[
|
||||
tuple[Mapping[str, object], ...],
|
||||
str | None,
|
||||
tuple[Mapping[str, object], ...] | None,
|
||||
Mapping[str, object] | None,
|
||||
]
|
||||
_RESPONSES_COUNT_REQUEST_PARTS: Final = TypeAdapter(ResponsesCountRequestParts)
|
||||
_TRANSLATED_REQUEST: Final = TypeAdapter(Mapping[str, object])
|
||||
|
||||
|
||||
class _AnthropicCountMessage(TypedDict):
|
||||
role: ReadOnly[str]
|
||||
|
||||
|
||||
_ANTHROPIC_COUNT_MESSAGES: Final = TypeAdapter(tuple[_AnthropicCountMessage, ...])
|
||||
_ANTHROPIC_OUTPUT_SCHEMA: Final = TypeAdapter(AnthropicOutputSchema)
|
||||
|
||||
|
||||
def _anthropic_output_schema(output_format: JsonValue) -> AnthropicOutputSchema | None:
|
||||
try:
|
||||
return _ANTHROPIC_OUTPUT_SCHEMA.validate_python(output_format)
|
||||
except ValidationError:
|
||||
return None
|
||||
|
||||
|
||||
def _chat_image_block_to_responses_part(image_url: object) -> ResponsesInputImagePart | None:
|
||||
url: Final = image_url.get("url") if isinstance(image_url, Mapping) else image_url
|
||||
|
|
@ -118,8 +150,9 @@ class OpenAICountTokensConfig:
|
|||
self,
|
||||
model: str,
|
||||
input: str | Sequence[object],
|
||||
tools: list[dict[str, Any]] | None = None,
|
||||
tools: Sequence[Mapping[str, object]] | None = None,
|
||||
instructions: str | None = None,
|
||||
text: Mapping[str, object] | None = None,
|
||||
) -> dict[str, object]:
|
||||
"""
|
||||
Transform request to OpenAI Responses API token counting format.
|
||||
|
|
@ -135,7 +168,10 @@ class OpenAICountTokensConfig:
|
|||
request["instructions"] = instructions
|
||||
|
||||
if tools is not None:
|
||||
request["tools"] = self._transform_tools_for_responses_api(tools)
|
||||
request["tools"] = self._transform_tools_for_responses_api([dict(tool) for tool in tools])
|
||||
|
||||
if text is not None:
|
||||
request["text"] = text
|
||||
|
||||
return request
|
||||
|
||||
|
|
@ -183,7 +219,7 @@ class OpenAICountTokensConfig:
|
|||
@staticmethod
|
||||
def messages_to_responses_input(
|
||||
messages: list[dict[str, Any]],
|
||||
) -> tuple:
|
||||
) -> tuple[list[dict[str, object]], str | None]:
|
||||
"""
|
||||
Convert standard chat messages format to OpenAI Responses API input format.
|
||||
|
||||
|
|
@ -246,3 +282,47 @@ class OpenAICountTokensConfig:
|
|||
|
||||
instructions: Final = "\n".join(instructions_parts) if instructions_parts else None
|
||||
return input_items, instructions
|
||||
|
||||
@staticmethod
|
||||
def malformed_anthropic_messages_error(messages: Sequence[Mapping[str, object]]) -> str | None:
|
||||
try:
|
||||
_ANTHROPIC_COUNT_MESSAGES.validate_python(messages)
|
||||
except ValidationError as e:
|
||||
return "; ".join(
|
||||
f"messages.{'.'.join(str(part) for part in error['loc'])}: {error['msg']}" for error in e.errors()
|
||||
)
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def anthropic_messages_to_responses_input(
|
||||
model: str,
|
||||
messages: Sequence[Mapping[str, object]],
|
||||
system: object,
|
||||
tools: Sequence[Mapping[str, object]] | None,
|
||||
output_format: JsonValue = None,
|
||||
) -> ResponsesCountRequestParts:
|
||||
"""Convert with the adapter /v1/messages inference uses, so the count sees the same input it bills."""
|
||||
anthropic_request: Final = AnthropicMessagesRequest(
|
||||
model=model,
|
||||
messages=normalize_native_messages(messages),
|
||||
system=system if isinstance(system, (str, list)) else None,
|
||||
tools=[dict(tool) for tool in tools] if tools else None,
|
||||
output_format=_anthropic_output_schema(output_format),
|
||||
)
|
||||
translated: Final = _TRANSLATED_REQUEST.validate_python(
|
||||
LiteLLMAnthropicToResponsesAPIAdapter().translate_request(
|
||||
anthropic_request, include_encrypted_reasoning=False
|
||||
)
|
||||
)
|
||||
return OpenAICountTokensConfig.count_request_parts(
|
||||
ResponsesAPIRequestUtils._restore_encrypted_content_item_ids_in_input(translated.get("input")), # pyright: ignore[reportPrivateUsage] # the input restore inference runs
|
||||
translated.get("instructions"),
|
||||
translated.get("tools"),
|
||||
translated.get("text"),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def count_request_parts(
|
||||
input_items: object, instructions: object, tools: object, text: object = None
|
||||
) -> ResponsesCountRequestParts:
|
||||
return _RESPONSES_COUNT_REQUEST_PARTS.validate_python((input_items, instructions, tools, text))
|
||||
|
|
|
|||
|
|
@ -7,7 +7,7 @@ from types import MappingProxyType
|
|||
from typing import Any, Final, Literal, cast, get_type_hints
|
||||
|
||||
import httpx
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
from pydantic import JsonValue, TypeAdapter, ValidationError
|
||||
from typing_extensions import NotRequired, ReadOnly, TypedDict
|
||||
|
||||
import litellm
|
||||
|
|
@ -22,7 +22,7 @@ from litellm.types.llms.vertex_ai import (
|
|||
PartType,
|
||||
Schema,
|
||||
)
|
||||
from litellm.types.utils import TokenCountResponse
|
||||
from litellm.types.utils import TokenCountMessagesFormat, TokenCountResponse
|
||||
from litellm.utils import supports_response_schema, supports_system_messages
|
||||
|
||||
GEMINI_VIDEO_METADATA_KEYS: Final = MappingProxyType(
|
||||
|
|
@ -1271,6 +1271,8 @@ class VertexAITokenCounter(BaseTokenCounter):
|
|||
request_model: str = "",
|
||||
tools: list[dict[str, object]] | None = None,
|
||||
system: object | None = None,
|
||||
messages_format: TokenCountMessagesFormat = "openai",
|
||||
output_format: JsonValue = None,
|
||||
) -> TokenCountResponse | None:
|
||||
import copy
|
||||
|
||||
|
|
@ -1342,30 +1344,55 @@ class VertexAITokenCounter(BaseTokenCounter):
|
|||
original_response=result,
|
||||
)
|
||||
else:
|
||||
from litellm.llms.vertex_ai.count_tokens.handler import VertexAITokenCounter
|
||||
from litellm.llms.vertex_ai.gemini.transformation import (
|
||||
gemini_convert_messages_with_history, # pyright: ignore[reportPrivateUsage] # shared helper already used by gemini/chat, context_caching, and vertex_and_google_ai_studio_gemini
|
||||
from litellm.llms.vertex_ai.count_tokens.context import (
|
||||
VertexGeminiCountContextError,
|
||||
count_context_error_response,
|
||||
gemini_count_contents,
|
||||
inference_request_view,
|
||||
provider_error_response,
|
||||
unconvertible_messages_error,
|
||||
vertex_gemini_count_context,
|
||||
)
|
||||
from litellm.llms.vertex_ai.count_tokens.handler import VertexAITokenCounter
|
||||
|
||||
request_view: Final = inference_request_view(messages or [], system, messages_format)
|
||||
count_context: Final = vertex_gemini_count_context(system=request_view.system, tools=tools)
|
||||
if isinstance(count_context, VertexGeminiCountContextError):
|
||||
return count_context_error_response(count_context, request_model, model_to_use)
|
||||
blocks_error: Final = unconvertible_messages_error(request_view.messages, messages_format)
|
||||
if blocks_error is not None:
|
||||
return count_context_error_response(blocks_error, request_model, model_to_use)
|
||||
|
||||
resolved_contents: Final = (
|
||||
contents if contents is not None else gemini_convert_messages_with_history(messages=messages or [])
|
||||
contents
|
||||
if contents is not None
|
||||
else await gemini_count_contents(
|
||||
request_view.messages, messages_format, model_to_use, "vertex_ai", count_tokens_params_request
|
||||
)
|
||||
)
|
||||
if isinstance(resolved_contents, VertexGeminiCountContextError):
|
||||
return count_context_error_response(resolved_contents, request_model, model_to_use)
|
||||
|
||||
count_tokens_params: Final = {
|
||||
"model": model_to_use,
|
||||
"contents": resolved_contents,
|
||||
"system_instruction": count_context.system_instruction,
|
||||
"tools": count_context.tools,
|
||||
}
|
||||
count_tokens_params_request.update(count_tokens_params)
|
||||
result = await VertexAITokenCounter().acount_tokens(
|
||||
**count_tokens_params_request,
|
||||
)
|
||||
try:
|
||||
result = await VertexAITokenCounter().acount_tokens(
|
||||
**count_tokens_params_request,
|
||||
)
|
||||
except litellm.APIError as error:
|
||||
return provider_error_response(error, request_model, model_to_use)
|
||||
|
||||
if result is not None and "totalTokens" in result:
|
||||
if "totalTokens" in result:
|
||||
return TokenCountResponse(
|
||||
total_tokens=result["totalTokens"],
|
||||
request_model=request_model,
|
||||
model_used=model_to_use,
|
||||
tokenizer_type=result.get("tokenizer_used", ""),
|
||||
tokenizer_type=result.get("tokenizer_used", "gemini"),
|
||||
original_response=result,
|
||||
)
|
||||
|
||||
|
|
|
|||
242
litellm/llms/vertex_ai/count_tokens/context.py
Normal file
242
litellm/llms/vertex_ai/count_tokens/context.py
Normal file
|
|
@ -0,0 +1,242 @@
|
|||
import copy
|
||||
from collections.abc import Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Final, Literal
|
||||
|
||||
from pydantic import JsonValue, TypeAdapter, ValidationError
|
||||
from typing_extensions import assert_never
|
||||
|
||||
from litellm.litellm_core_utils.prompt_templates.factory import chat_completion_tools_for
|
||||
from litellm.types.llms.anthropic import AnthropicMessagesTextParam
|
||||
from litellm.types.llms.openai import AllMessageValues
|
||||
from litellm.types.llms.vertex_ai import ContentType
|
||||
from litellm.types.utils import TokenCountMessagesFormat, TokenCountResponse
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.exceptions import APIError
|
||||
from litellm.llms.anthropic.pass_through.context_management.result import PolyfillResult
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class VertexGeminiCountContext:
|
||||
system_instruction: Mapping[str, object] | None
|
||||
tools: tuple[Mapping[str, object], ...] | None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class VertexGeminiCountContextError:
|
||||
error_message: str
|
||||
|
||||
|
||||
_SYSTEM_BLOCKS: Final = TypeAdapter(tuple[AnthropicMessagesTextParam, ...])
|
||||
_TOOLS: Final = TypeAdapter(tuple[Mapping[str, JsonValue], ...])
|
||||
_VERTEX_TOOL_KEYS: Final = MappingProxyType(
|
||||
{
|
||||
"function_declarations": "functionDeclarations",
|
||||
"url_context": "urlContext",
|
||||
"code_execution": "codeExecution",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _system_parts(
|
||||
system: object | None,
|
||||
) -> tuple[Mapping[str, JsonValue], ...] | VertexGeminiCountContextError:
|
||||
if system is None:
|
||||
return ()
|
||||
if isinstance(system, str):
|
||||
return ({"text": system},) if system else ()
|
||||
try:
|
||||
blocks: Final = _SYSTEM_BLOCKS.validate_python(system)
|
||||
except ValidationError:
|
||||
return VertexGeminiCountContextError(error_message="Gemini system instructions only support text blocks.")
|
||||
return tuple({"text": block["text"]} for block in blocks)
|
||||
|
||||
|
||||
_INVALID_TOOLS: Final = VertexGeminiCountContextError(
|
||||
error_message="tools must be a list of objects, and each function tool needs a JSON object schema."
|
||||
)
|
||||
|
||||
|
||||
def _has_object_schema(tool: Mapping[str, JsonValue]) -> bool:
|
||||
function: Final = tool.get("function")
|
||||
schema: Final = (
|
||||
function.get("parameters") if isinstance(function, dict) else tool.get("input_schema", tool.get("parameters"))
|
||||
)
|
||||
return schema is None or isinstance(schema, dict)
|
||||
|
||||
|
||||
def _is_function_tool(tool: Mapping[str, JsonValue]) -> bool:
|
||||
return isinstance(tool.get("name"), str) or "input_schema" in tool or "function" in tool
|
||||
|
||||
|
||||
def _map_function_tools(
|
||||
tools: object | None,
|
||||
) -> tuple[Mapping[str, object], ...] | None | VertexGeminiCountContextError:
|
||||
try:
|
||||
tool_values: Final = _TOOLS.validate_python(tools or ())
|
||||
except ValidationError:
|
||||
return _INVALID_TOOLS
|
||||
function_tools: Final = tuple(tool for tool in tool_values if _is_function_tool(tool))
|
||||
if not function_tools:
|
||||
return None
|
||||
if not all(_has_object_schema(tool) for tool in function_tools):
|
||||
return _INVALID_TOOLS
|
||||
|
||||
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexGeminiConfig
|
||||
|
||||
chat_tools: Final = tuple(
|
||||
chat_tool for tool in function_tools for chat_tool in chat_completion_tools_for(tool)
|
||||
) # comprehension-ok: flattens each tool's translation
|
||||
try:
|
||||
mapped_tools: Final = VertexGeminiConfig()._map_function( # pyright: ignore[reportPrivateUsage] # the tool mapping inference uses
|
||||
value=[dict(tool) for tool in chat_tools],
|
||||
optional_params={},
|
||||
)
|
||||
except (AttributeError, TypeError, ValueError):
|
||||
return _INVALID_TOOLS
|
||||
return tuple(
|
||||
MappingProxyType({_VERTEX_TOOL_KEYS.get(key, key): value for key, value in tool.items()})
|
||||
for tool in mapped_tools
|
||||
)
|
||||
|
||||
|
||||
def vertex_gemini_count_context(
|
||||
system: object | None,
|
||||
tools: object | None,
|
||||
) -> VertexGeminiCountContext | VertexGeminiCountContextError:
|
||||
system_parts: Final = _system_parts(system)
|
||||
if isinstance(system_parts, VertexGeminiCountContextError):
|
||||
return system_parts
|
||||
function_tools: Final = _map_function_tools(tools)
|
||||
if isinstance(function_tools, VertexGeminiCountContextError):
|
||||
return function_tools
|
||||
return VertexGeminiCountContext(
|
||||
system_instruction=MappingProxyType({"parts": system_parts}) if system_parts else None,
|
||||
tools=function_tools,
|
||||
)
|
||||
|
||||
|
||||
def unconvertible_messages_error(
|
||||
messages: Sequence[Mapping[str, object]] | None, messages_format: TokenCountMessagesFormat
|
||||
) -> VertexGeminiCountContextError | None:
|
||||
from litellm.llms.anthropic.pass_through.convertible_blocks import unconvertible_blocks_error
|
||||
|
||||
error_message: Final = unconvertible_blocks_error(messages, messages_format)
|
||||
return None if error_message is None else VertexGeminiCountContextError(error_message=error_message)
|
||||
|
||||
|
||||
def count_context_error_response(
|
||||
error: VertexGeminiCountContextError, request_model: str, model_used: str
|
||||
) -> TokenCountResponse:
|
||||
return TokenCountResponse(
|
||||
total_tokens=0,
|
||||
request_model=request_model,
|
||||
model_used=model_used,
|
||||
tokenizer_type="gemini",
|
||||
error=True,
|
||||
error_message=error.error_message,
|
||||
status_code=400,
|
||||
rejected_before_provider_call=True,
|
||||
)
|
||||
|
||||
|
||||
def provider_error_response(error: "APIError", request_model: str, model_used: str) -> TokenCountResponse:
|
||||
return TokenCountResponse(
|
||||
total_tokens=0,
|
||||
request_model=request_model,
|
||||
model_used=model_used,
|
||||
tokenizer_type="gemini",
|
||||
error=True,
|
||||
error_message=error.message,
|
||||
status_code=error.status_code,
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class InferenceRequestView:
|
||||
messages: Sequence[Mapping[str, object]]
|
||||
system: object | None
|
||||
|
||||
|
||||
def _client_compaction(messages: Sequence[Mapping[str, object]], system: object | None) -> "PolyfillResult | None":
|
||||
from litellm.llms.anthropic.pass_through.context_management.editors.compact import (
|
||||
apply_client_compaction_block_history,
|
||||
)
|
||||
|
||||
chat_messages: Final = [dict(message) for message in messages]
|
||||
if system is None or isinstance(system, str):
|
||||
return apply_client_compaction_block_history(messages=chat_messages, system=system)
|
||||
try:
|
||||
blocks: Final = _SYSTEM_BLOCKS.validate_python(system)
|
||||
except ValidationError:
|
||||
return None
|
||||
return apply_client_compaction_block_history(messages=chat_messages, system=[dict(block) for block in blocks])
|
||||
|
||||
|
||||
def inference_request_view(
|
||||
messages: Sequence[Mapping[str, object]], system: object | None, messages_format: TokenCountMessagesFormat
|
||||
) -> InferenceRequestView:
|
||||
"""Apply the history rewrites /v1/messages runs before it converts a body for Gemini."""
|
||||
from litellm.llms.anthropic.pass_through.messages.utils import normalize_native_messages
|
||||
|
||||
match messages_format:
|
||||
case "openai":
|
||||
return InferenceRequestView(messages=messages, system=system)
|
||||
case "anthropic":
|
||||
normalized: Final = normalize_native_messages(messages)
|
||||
compacted: Final = _client_compaction(normalized, system)
|
||||
if compacted is None:
|
||||
return InferenceRequestView(messages=normalized, system=system)
|
||||
return InferenceRequestView(messages=compacted.messages, system=compacted.system)
|
||||
case _:
|
||||
assert_never(messages_format)
|
||||
|
||||
|
||||
def _chat_messages(
|
||||
messages: Sequence[Mapping[str, object]],
|
||||
messages_format: TokenCountMessagesFormat,
|
||||
model: str,
|
||||
custom_llm_provider: Literal["vertex_ai", "gemini"],
|
||||
) -> Sequence[AllMessageValues]:
|
||||
from litellm.llms.anthropic.pass_through.adapters.transformation import LiteLLMAnthropicMessagesAdapter
|
||||
|
||||
match messages_format:
|
||||
case "anthropic":
|
||||
return LiteLLMAnthropicMessagesAdapter().translate_anthropic_messages_to_openai(
|
||||
messages=list(messages), # pyright: ignore[reportArgumentType] # raw Anthropic body, read the way /v1/messages inference reads it
|
||||
model=model,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
)
|
||||
case "openai":
|
||||
return copy.deepcopy(list(messages)) # pyright: ignore[reportReturnType] # chat messages from the request, copied because the AI Studio conversion rewrites them in place
|
||||
case _:
|
||||
assert_never(messages_format)
|
||||
|
||||
|
||||
async def gemini_count_contents(
|
||||
messages: Sequence[Mapping[str, object]],
|
||||
messages_format: TokenCountMessagesFormat,
|
||||
model: str,
|
||||
custom_llm_provider: Literal["vertex_ai", "gemini"],
|
||||
litellm_params: Mapping[str, object],
|
||||
) -> Sequence[ContentType] | VertexGeminiCountContextError:
|
||||
from litellm.exceptions import BadRequestError
|
||||
from litellm.litellm_core_utils.asyncify import asyncify
|
||||
from litellm.litellm_core_utils.prompt_templates.image_handling import async_inline_remote_media
|
||||
from litellm.llms.gemini.chat.transformation import GoogleAIStudioGeminiConfig
|
||||
from litellm.llms.vertex_ai.gemini.transformation import should_inline_remote_media
|
||||
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexGeminiConfig
|
||||
|
||||
config: Final = GoogleAIStudioGeminiConfig() if custom_llm_provider == "gemini" else VertexGeminiConfig()
|
||||
try:
|
||||
chat_messages: Final = await async_inline_remote_media(
|
||||
list(_chat_messages(messages, messages_format, model, custom_llm_provider)),
|
||||
should_inline=should_inline_remote_media(custom_llm_provider),
|
||||
)
|
||||
return await asyncify(config._transform_messages)( # pyright: ignore[reportPrivateUsage] # the message transform inference runs
|
||||
messages=chat_messages, model=model, litellm_params=dict(litellm_params)
|
||||
)
|
||||
except BadRequestError as error:
|
||||
return VertexGeminiCountContextError(error_message=error.message)
|
||||
|
|
@ -1,3 +1,4 @@
|
|||
from collections.abc import Mapping, Sequence
|
||||
from typing import Any, Final
|
||||
|
||||
from litellm.llms.gemini.count_tokens.handler import GoogleAIStudioTokenCounter
|
||||
|
|
@ -5,6 +6,19 @@ from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
|
|||
|
||||
|
||||
class VertexAITokenCounter(GoogleAIStudioTokenCounter, VertexBase):
|
||||
def _build_count_tokens_request_body(
|
||||
self,
|
||||
model: str,
|
||||
contents: object,
|
||||
system_instruction: Mapping[str, object] | None,
|
||||
tools: Sequence[Mapping[str, object]] | None,
|
||||
) -> Mapping[str, object]:
|
||||
return {
|
||||
"contents": contents,
|
||||
**({"systemInstruction": dict(system_instruction)} if system_instruction else {}),
|
||||
**({"tools": [dict(tool) for tool in tools]} if tools else {}),
|
||||
}
|
||||
|
||||
async def validate_environment(
|
||||
self,
|
||||
api_base: str | None = None,
|
||||
|
|
@ -16,29 +30,29 @@ class VertexAITokenCounter(GoogleAIStudioTokenCounter, VertexBase):
|
|||
"""
|
||||
Returns a Tuple of headers and url for the Vertex AI countTokens endpoint.
|
||||
"""
|
||||
litellm_params = litellm_params or {}
|
||||
vertex_credentials: Final = self.get_vertex_ai_credentials(litellm_params=litellm_params)
|
||||
vertex_project = self.get_vertex_ai_project(litellm_params=litellm_params)
|
||||
vertex_location: Final = self.get_vertex_ai_location(litellm_params=litellm_params)
|
||||
_auth_header, vertex_project = await self._ensure_access_token_async(
|
||||
resolved_litellm_params: Final = litellm_params or {}
|
||||
vertex_credentials: Final = self.get_vertex_ai_credentials(litellm_params=resolved_litellm_params)
|
||||
vertex_project: Final = self.get_vertex_ai_project(litellm_params=resolved_litellm_params)
|
||||
vertex_location: Final = self.get_vertex_ai_location(litellm_params=resolved_litellm_params)
|
||||
_auth_header, resolved_vertex_project = await self._ensure_access_token_async(
|
||||
credentials=vertex_credentials,
|
||||
project_id=vertex_project,
|
||||
custom_llm_provider="vertex_ai",
|
||||
)
|
||||
|
||||
auth_header, api_base = self._get_token_and_url(
|
||||
auth_header, url = self._get_token_and_url(
|
||||
model=model,
|
||||
gemini_api_key=None,
|
||||
auth_header=_auth_header,
|
||||
vertex_project=vertex_project,
|
||||
vertex_project=resolved_vertex_project,
|
||||
vertex_location=vertex_location,
|
||||
vertex_credentials=vertex_credentials,
|
||||
stream=False,
|
||||
custom_llm_provider="vertex_ai",
|
||||
api_base=None,
|
||||
api_base=api_base,
|
||||
mode="count_tokens",
|
||||
)
|
||||
headers = {
|
||||
"Authorization": f"Bearer {auth_header}",
|
||||
}
|
||||
return headers, api_base
|
||||
return headers, url
|
||||
|
|
|
|||
|
|
@ -6,7 +6,7 @@ Why separate file? Make it easy to see how transformation works
|
|||
|
||||
import os
|
||||
import re
|
||||
from collections.abc import Mapping
|
||||
from collections.abc import Callable, Mapping
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, cast
|
||||
from urllib.parse import quote
|
||||
|
||||
|
|
@ -1311,6 +1311,10 @@ def _vertex_inlines(media: RemoteMedia) -> bool:
|
|||
)
|
||||
|
||||
|
||||
def should_inline_remote_media(custom_llm_provider: str) -> Callable[[RemoteMedia], bool]:
|
||||
return _ai_studio_inlines if custom_llm_provider == "gemini" else _vertex_inlines
|
||||
|
||||
|
||||
async def async_transform_request_body(
|
||||
gemini_api_key: str | None,
|
||||
messages: list[AllMessageValues],
|
||||
|
|
@ -1353,7 +1357,7 @@ async def async_transform_request_body(
|
|||
)
|
||||
|
||||
inlined_messages: Final = await async_inline_remote_media(
|
||||
messages, should_inline=_ai_studio_inlines if custom_llm_provider == "gemini" else _vertex_inlines
|
||||
messages, should_inline=should_inline_remote_media(custom_llm_provider)
|
||||
)
|
||||
|
||||
if _openai_messages_may_need_sync_gcs_metadata_fetch(inlined_messages):
|
||||
|
|
|
|||
|
|
@ -20077,7 +20077,13 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": false
|
||||
"supports_vision": false,
|
||||
"token_count_framing": 97
|
||||
},
|
||||
"databricks/system.ai.deepseek-v4-flash-0731": {
|
||||
"litellm_provider": "databricks",
|
||||
"mode": "chat",
|
||||
"token_count_framing": 97
|
||||
},
|
||||
"databricks/databricks-deepseek-v4-pro-0813": {
|
||||
"cache_creation_input_token_cost": 1.31999e-06,
|
||||
|
|
@ -21068,7 +21074,13 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"token_count_framing": 100
|
||||
},
|
||||
"databricks/system.ai.kimi-k3": {
|
||||
"litellm_provider": "databricks",
|
||||
"mode": "chat",
|
||||
"token_count_framing": 100
|
||||
},
|
||||
"databricks/databricks-llama-4-maverick": {
|
||||
"cache_creation_input_token_cost": 5.0001e-07,
|
||||
|
|
@ -77877,7 +77889,9 @@
|
|||
"supports_response_schema": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"token_count_framing": 64,
|
||||
"token_count_tokenizer": "kimi_k3"
|
||||
},
|
||||
"global.moonshotai.kimi-k3": {
|
||||
"supports_regex_lookaround": false,
|
||||
|
|
@ -77907,7 +77921,9 @@
|
|||
"supports_response_schema": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"token_count_framing": 64,
|
||||
"token_count_tokenizer": "kimi_k3"
|
||||
},
|
||||
"us.moonshotai.kimi-k3": {
|
||||
"supports_regex_lookaround": false,
|
||||
|
|
@ -77937,7 +77953,9 @@
|
|||
"supports_response_schema": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"token_count_framing": 64,
|
||||
"token_count_tokenizer": "kimi_k3"
|
||||
},
|
||||
"openrouter/xiaomi/mimo-v2.6-flash": {
|
||||
"cache_read_input_token_cost": 2.8e-09,
|
||||
|
|
|
|||
|
|
@ -2450,6 +2450,20 @@ async def fire_mcp_tool_call_failure_logging(
|
|||
)
|
||||
|
||||
|
||||
def _single_scoped_server(
|
||||
mcp_servers: Sequence[str] | None, allowed_mcp_servers: Sequence[MCPServer]
|
||||
) -> MCPServer | None:
|
||||
"""The one server a ``/{server_name}/mcp`` path or ``x-mcp-servers`` header scoped this call to.
|
||||
|
||||
A bare tool name on such a request resolves against that server, the way a REST ``server_id``
|
||||
does, instead of against the per-process tool-name map, which stays empty for a server whose
|
||||
upstream only lists tools with the caller's own credentials until this process has listed it.
|
||||
"""
|
||||
if mcp_servers is None or len(mcp_servers) != 1 or len(allowed_mcp_servers) != 1:
|
||||
return None
|
||||
return allowed_mcp_servers[0]
|
||||
|
||||
|
||||
@client
|
||||
@catalog_operation(lambda: global_mcp_server_manager)
|
||||
async def call_mcp_tool(
|
||||
|
|
@ -2511,6 +2525,7 @@ async def call_mcp_tool(
|
|||
detail="User not allowed to call this tool.",
|
||||
)
|
||||
|
||||
scoped_server: Final = _single_scoped_server(mcp_servers, allowed_mcp_servers)
|
||||
# Delegate to execute_mcp_tool for execution
|
||||
response = await execute_mcp_tool(
|
||||
name=name,
|
||||
|
|
@ -2527,6 +2542,7 @@ async def call_mcp_tool(
|
|||
input_responses=input_responses,
|
||||
request_state=request_state,
|
||||
incoming_bearer_token=incoming_bearer_token,
|
||||
requested_server_id=scoped_server.server_id if scoped_server is not None else None,
|
||||
**kwargs,
|
||||
)
|
||||
except Exception as e:
|
||||
|
|
|
|||
|
|
@ -2,10 +2,12 @@
|
|||
Unified /v1/messages endpoint - (Anthropic Spec)
|
||||
"""
|
||||
|
||||
from collections.abc import Mapping
|
||||
from typing import Final
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException, Request, Response
|
||||
from fastapi.responses import JSONResponse
|
||||
from pydantic import JsonValue
|
||||
|
||||
import litellm
|
||||
from litellm.anthropic_interface.exceptions import (
|
||||
|
|
@ -27,6 +29,7 @@ from litellm.proxy.common_request_processing import (
|
|||
create_response,
|
||||
log_llm_api_exception,
|
||||
proxy_exception_from_http_exception,
|
||||
proxy_exception_from_route_error,
|
||||
resolve_litellm_call_id,
|
||||
)
|
||||
from litellm.proxy.common_utils.error_body_call_id import error_body_call_id
|
||||
|
|
@ -41,7 +44,7 @@ from litellm.proxy.common_utils.openai_error_payload import (
|
|||
openai_error_type,
|
||||
with_litellm_call_id,
|
||||
)
|
||||
from litellm.types.utils import TokenCountResponse
|
||||
from litellm.proxy.common_utils.token_count_settings import count_source_fields
|
||||
|
||||
router: Final = APIRouter()
|
||||
|
||||
|
|
@ -289,6 +292,14 @@ async def anthropic_response(
|
|||
)
|
||||
|
||||
|
||||
def _requested_output_format(data: Mapping[str, JsonValue]) -> JsonValue:
|
||||
output_format: Final = data.get("output_format")
|
||||
output_config: Final = data.get("output_config")
|
||||
if isinstance(output_format, dict) or not isinstance(output_config, dict):
|
||||
return output_format
|
||||
return output_config.get("format")
|
||||
|
||||
|
||||
@router.post(
|
||||
"/v1/messages/count_tokens",
|
||||
tags=["[beta] Anthropic Messages Token Counting"],
|
||||
|
|
@ -318,7 +329,7 @@ async def count_tokens(
|
|||
|
||||
Returns: {"input_tokens": <number>}
|
||||
"""
|
||||
from litellm.proxy.proxy_server import token_counter as internal_token_counter
|
||||
from litellm.proxy.proxy_server import count_request_tokens
|
||||
|
||||
litellm_call_id: Final = resolve_litellm_call_id(request.headers.get("x-litellm-call-id"))
|
||||
try:
|
||||
|
|
@ -330,10 +341,10 @@ async def count_tokens(
|
|||
messages: Final = data.get("messages", [])
|
||||
|
||||
if not model_name:
|
||||
raise HTTPException(status_code=400, detail={"error": "model parameter is required"})
|
||||
raise HTTPException(status_code=400, detail="model parameter is required")
|
||||
|
||||
if not messages:
|
||||
raise HTTPException(status_code=400, detail={"error": "messages parameter is required"})
|
||||
raise HTTPException(status_code=400, detail="messages parameter is required")
|
||||
|
||||
# Create TokenCountRequest for the internal endpoint
|
||||
from litellm.proxy._types import TokenCountRequest
|
||||
|
|
@ -345,35 +356,21 @@ async def count_tokens(
|
|||
system=data.get("system"),
|
||||
)
|
||||
|
||||
# Call the internal token counter function with direct request flag set to False
|
||||
token_response: Final = await internal_token_counter(
|
||||
token_response: Final = await count_request_tokens(
|
||||
request=token_request,
|
||||
call_endpoint=True,
|
||||
messages_format="anthropic",
|
||||
output_format=_requested_output_format(data),
|
||||
)
|
||||
_token_response_dict: dict = {}
|
||||
if isinstance(token_response, TokenCountResponse):
|
||||
_token_response_dict = token_response.model_dump()
|
||||
elif isinstance(token_response, dict):
|
||||
_token_response_dict = token_response
|
||||
|
||||
# Convert the internal response to Anthropic API format
|
||||
return {"input_tokens": _token_response_dict.get("total_tokens", 0)}
|
||||
|
||||
except HTTPException:
|
||||
raise
|
||||
except HTTPException as e:
|
||||
return _anthropic_error_json_response(proxy_exception_from_http_exception(e, {}), request)
|
||||
except ProxyException as e:
|
||||
status_code: Final = int(e.code) if e.code and e.code.isdigit() else 500
|
||||
detail: Final = AnthropicExceptionMapping.transform_to_anthropic_error(
|
||||
status_code=status_code,
|
||||
raw_message=e.message,
|
||||
)
|
||||
raise HTTPException(
|
||||
status_code=status_code,
|
||||
detail=detail,
|
||||
)
|
||||
return _anthropic_error_json_response(e, request)
|
||||
except Exception as e:
|
||||
log_llm_api_exception(e, litellm_call_id)
|
||||
raise HTTPException(status_code=500, detail={"error": f"Internal server error: {e}"})
|
||||
return _anthropic_error_json_response(proxy_exception_from_route_error(e), request)
|
||||
else:
|
||||
return {"input_tokens": token_response.total_tokens, **count_source_fields(token_response)}
|
||||
|
||||
|
||||
@router.post(
|
||||
|
|
|
|||
|
|
@ -4597,6 +4597,24 @@ def check_model_access_helper(
|
|||
_check_model_access_helper: Final = check_model_access_helper
|
||||
|
||||
|
||||
def _owning_team_names_for_team_deployment(
|
||||
model: str,
|
||||
llm_router: Router | None,
|
||||
team_id: str | None,
|
||||
team_model_aliases: Mapping[str, str] | None,
|
||||
) -> tuple[str, ...]:
|
||||
if llm_router is None or team_id is None or not model.startswith(f"model_name_{team_id}_"):
|
||||
return ()
|
||||
deployment: Final = llm_router.get_deployment_by_model_group_name(model_group_name=model)
|
||||
public_name: Final = (
|
||||
deployment.model_info.team_public_model_name
|
||||
if deployment is not None and deployment.model_info.team_id == team_id
|
||||
else None
|
||||
)
|
||||
legacy_aliases: Final = tuple(alias for alias, target in (team_model_aliases or {}).items() if target == model)
|
||||
return (*((public_name,) if public_name else ()), *legacy_aliases)
|
||||
|
||||
|
||||
def _can_object_call_model(
|
||||
model: str | list[str],
|
||||
llm_router: Router | None,
|
||||
|
|
@ -4668,6 +4686,13 @@ def _can_object_call_model(
|
|||
*((global_or_router_alias_target,) if global_or_router_alias_target else ()),
|
||||
)
|
||||
)
|
||||
owning_team_names: Final = (
|
||||
()
|
||||
if key_alias_applied
|
||||
else _owning_team_names_for_team_deployment(
|
||||
model=model, llm_router=llm_router, team_id=team_id, team_model_aliases=team_model_aliases
|
||||
)
|
||||
)
|
||||
|
||||
## check model access for alias + underlying model - allow if either is in allowed models
|
||||
for m in potential_models:
|
||||
|
|
@ -4679,6 +4704,11 @@ def _can_object_call_model(
|
|||
team_id=team_id,
|
||||
):
|
||||
return True
|
||||
if any(
|
||||
check_model_access_helper(model=name, llm_router=llm_router, models=models, team_id=team_id)
|
||||
for name in owning_team_names
|
||||
):
|
||||
return True
|
||||
|
||||
internal_message: Final = (
|
||||
f"{object_type} not allowed to access model. This {object_type} can only access models={models}. "
|
||||
|
|
|
|||
|
|
@ -637,6 +637,27 @@ def proxy_exception_from_http_exception(exc: HTTPException, headers: dict[str, s
|
|||
)
|
||||
|
||||
|
||||
def proxy_exception_from_route_error(exc: Exception) -> ProxyException:
|
||||
if isinstance(exc, ProxyException):
|
||||
return exc
|
||||
if isinstance(exc, HTTPException):
|
||||
return proxy_exception_from_http_exception(exc, {})
|
||||
error_status: Final = error_status_code(exc, status.HTTP_500_INTERNAL_SERVER_ERROR)
|
||||
message: Final = attribute_of(exc, "message", str(exc))
|
||||
carried_code: Final = attribute_of(exc, "code")
|
||||
provider_fields: Final = attribute_of(exc, "provider_specific_fields")
|
||||
return ProxyException(
|
||||
message=redact_internal_details_from_client_message(
|
||||
strip_bug_report_notice(message) if isinstance(message, str) else str(exc)
|
||||
),
|
||||
type=openai_error_type(exc, error_status),
|
||||
param=openai_error_param(exc),
|
||||
code=error_status,
|
||||
openai_code=carried_code if isinstance(carried_code, str) else None,
|
||||
provider_specific_fields=provider_fields if isinstance(provider_fields, dict) else None,
|
||||
)
|
||||
|
||||
|
||||
def _collect_response_file_search_vector_store_ids(data: Mapping[str, object]) -> set[str]:
|
||||
vector_store_ids: Final[set[str]] = set()
|
||||
tools: Final = data.get("tools")
|
||||
|
|
|
|||
164
litellm/proxy/common_utils/token_count_settings.py
Normal file
164
litellm/proxy/common_utils/token_count_settings.py
Normal file
|
|
@ -0,0 +1,164 @@
|
|||
from collections.abc import Callable, Mapping
|
||||
from dataclasses import dataclass, replace
|
||||
from types import MappingProxyType
|
||||
from typing import Final, Literal
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, ValidationError
|
||||
|
||||
import litellm
|
||||
from litellm.anthropic_interface.exceptions.exception_mapping_utils import AnthropicExceptionMapping
|
||||
from litellm.litellm_core_utils.core_helpers import as_str_mapping
|
||||
from litellm.proxy._types import ProxyException
|
||||
from litellm.types.token_count import TokenCountSourceFields, TokenCountTokenizer
|
||||
from litellm.types.utils import TokenCountResponse
|
||||
|
||||
CANNOT_COUNT_MODEL_STATUSES: Final = frozenset({400, 404})
|
||||
TOKENIZER_LOAD_WAIT_SECONDS: Final = 5.0
|
||||
TOKENIZER_LOAD_REUSE_SECONDS: Final = 60.0
|
||||
TOKENIZER_LOAD_MAX_RUNNING: Final = 4
|
||||
|
||||
|
||||
class _UpstreamErrorDetail(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="allow")
|
||||
message: str
|
||||
|
||||
|
||||
class _UpstreamErrorEnvelope(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="allow")
|
||||
type: Literal["error"]
|
||||
error: _UpstreamErrorDetail
|
||||
|
||||
|
||||
def _joined(message: str, hint: str) -> str:
|
||||
return f"{message.rstrip().rstrip('.')}. {hint}"
|
||||
|
||||
|
||||
def _with_hint(message: str, hint: str) -> str:
|
||||
try:
|
||||
envelope: Final = _UpstreamErrorEnvelope.model_validate_json(message)
|
||||
except ValidationError:
|
||||
return _joined(AnthropicExceptionMapping.extract_error_message(message), hint)
|
||||
detail: Final = envelope.error.model_copy(update={"message": _joined(envelope.error.message, hint)})
|
||||
return envelope.model_copy(update={"error": detail}).model_dump_json()
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class TokenCountModel:
|
||||
model: str
|
||||
region: str | None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class TokenCountSettings:
|
||||
framing: int | None
|
||||
count_model: TokenCountModel | None
|
||||
tokenizer: TokenCountTokenizer | None = None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ProviderCountFailure:
|
||||
status_code: int
|
||||
message: str
|
||||
rejected_before_provider_call: bool = False
|
||||
|
||||
@property
|
||||
def cannot_count_model(self) -> bool:
|
||||
return not self.rejected_before_provider_call and self.status_code in CANNOT_COUNT_MODEL_STATUSES
|
||||
|
||||
def with_hint(self, hint: str) -> "ProviderCountFailure":
|
||||
return replace(self, message=_with_hint(self.message, hint))
|
||||
|
||||
def to_proxy_exception(self) -> ProxyException:
|
||||
return ProxyException(
|
||||
message=self.message,
|
||||
type="token_counting_error",
|
||||
param="model",
|
||||
code=self.status_code,
|
||||
)
|
||||
|
||||
|
||||
def cost_map_entry(model: str, custom_llm_provider: str | None = None) -> Mapping[str, object]:
|
||||
try:
|
||||
return litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider)
|
||||
except Exception: # noqa: BLE001 # get_model_info raises bare Exception for unmapped models
|
||||
return MappingProxyType({})
|
||||
|
||||
|
||||
def _framing(value: object) -> int | None:
|
||||
match value:
|
||||
case bool():
|
||||
return None
|
||||
case int() if value >= 0:
|
||||
return value
|
||||
case float() if value.is_integer() and value >= 0:
|
||||
return int(value)
|
||||
case _:
|
||||
return None
|
||||
|
||||
|
||||
def _non_blank(value: object) -> str | None:
|
||||
return value if isinstance(value, str) and value.strip() else None
|
||||
|
||||
|
||||
def _count_model(value: Mapping[str, object]) -> TokenCountModel | None:
|
||||
model: Final = _non_blank(value.get("model"))
|
||||
return None if model is None else TokenCountModel(model=model, region=_non_blank(value.get("region")))
|
||||
|
||||
|
||||
def _geography(region: str) -> str:
|
||||
return region.rsplit("-", 2)[0]
|
||||
|
||||
|
||||
def _within_deployment_geography(
|
||||
count_model: TokenCountModel | None, deployment_region: Callable[[], str]
|
||||
) -> TokenCountModel | None:
|
||||
if count_model is None or count_model.region is None or count_model.model.startswith("global."):
|
||||
return count_model
|
||||
if _geography(deployment_region()) == _geography(count_model.region):
|
||||
return count_model
|
||||
return replace(count_model, region=None)
|
||||
|
||||
|
||||
def _resolved_count_model(
|
||||
configured: object, default: object, deployment_region: Callable[[], str]
|
||||
) -> TokenCountModel | None:
|
||||
configured_mapping: Final = as_str_mapping(configured)
|
||||
if configured_mapping is not None:
|
||||
return _count_model(configured_mapping)
|
||||
default_mapping: Final = as_str_mapping(default)
|
||||
return (
|
||||
None
|
||||
if default_mapping is None
|
||||
else _within_deployment_geography(_count_model(default_mapping), deployment_region)
|
||||
)
|
||||
|
||||
|
||||
def _tokenizer(value: object) -> TokenCountTokenizer | None:
|
||||
match value:
|
||||
case "kimi_k3":
|
||||
return "kimi_k3"
|
||||
case _:
|
||||
return None
|
||||
|
||||
|
||||
def resolve_token_count_settings(
|
||||
configured: Mapping[str, object], default: Mapping[str, object], deployment_region: Callable[[], str]
|
||||
) -> TokenCountSettings:
|
||||
framing: Final = _framing(configured.get("token_count_framing"))
|
||||
return TokenCountSettings(
|
||||
framing=framing if framing is not None else _framing(default.get("token_count_framing")),
|
||||
count_model=_resolved_count_model(
|
||||
configured.get("token_count_model"), default.get("token_count_model"), deployment_region
|
||||
),
|
||||
tokenizer=_tokenizer(default.get("token_count_tokenizer")),
|
||||
)
|
||||
|
||||
|
||||
def count_source_fields(response: TokenCountResponse) -> TokenCountSourceFields:
|
||||
match response.count_source:
|
||||
case None:
|
||||
return {}
|
||||
case "token_count_model":
|
||||
return {"litellm_count_source": "token_count_model", "litellm_count_model": response.model_used}
|
||||
case source:
|
||||
return {"litellm_count_source": source}
|
||||
|
|
@ -1,22 +1,142 @@
|
|||
from collections.abc import Mapping
|
||||
from itertools import chain
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
from fastapi import APIRouter, Depends, Request, Response
|
||||
from fastapi.responses import ORJSONResponse
|
||||
from fastapi import APIRouter, Depends, HTTPException, Request, Response
|
||||
from fastapi.responses import JSONResponse, ORJSONResponse
|
||||
from pydantic import BaseModel, ConfigDict, Field, JsonValue, TypeAdapter, ValidationError, field_validator
|
||||
|
||||
from litellm.litellm_core_utils.json_validation_rule import normalize_json_schema_types
|
||||
from litellm.proxy._types import *
|
||||
from litellm.proxy._types import ProxyException, TokenCountRequest
|
||||
from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth
|
||||
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
|
||||
from litellm.proxy.common_request_processing import (
|
||||
ProxyBaseLLMRequestProcessing,
|
||||
log_llm_api_exception,
|
||||
proxy_exception_from_http_exception,
|
||||
resolve_litellm_call_id,
|
||||
)
|
||||
from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # legacy module exports
|
||||
_read_request_body, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
|
||||
read_request_body,
|
||||
)
|
||||
from litellm.proxy.common_utils.token_count_settings import count_source_fields
|
||||
from litellm.types.llms.vertex_ai import TokenCountDetailsResponse
|
||||
|
||||
|
||||
class _GoogleCountTokensTextPart(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
text: str
|
||||
|
||||
|
||||
class _GoogleCountTokensSystemInstruction(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
parts: tuple[_GoogleCountTokensTextPart, ...]
|
||||
|
||||
@field_validator("parts", mode="before")
|
||||
@classmethod
|
||||
def normalize_parts(cls, parts: object) -> object:
|
||||
return (parts,) if isinstance(parts, Mapping) else parts
|
||||
|
||||
|
||||
class _GoogleCountTokensFunctionDeclaration(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, populate_by_name=True)
|
||||
name: str
|
||||
description: str | None = None
|
||||
parameters: JsonValue | None = None
|
||||
parameters_json_schema: JsonValue | None = Field(default=None, alias="parametersJsonSchema")
|
||||
|
||||
|
||||
class _GoogleCountTokensTool(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, populate_by_name=True)
|
||||
function_declarations: tuple[_GoogleCountTokensFunctionDeclaration, ...] = Field(
|
||||
default=(), alias="functionDeclarations"
|
||||
)
|
||||
|
||||
|
||||
class _GoogleCountTokensRequest(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="ignore", populate_by_name=True)
|
||||
contents: tuple[Mapping[str, JsonValue], ...] = ()
|
||||
system_instruction: _GoogleCountTokensSystemInstruction | None = Field(default=None, alias="systemInstruction")
|
||||
tools: tuple[_GoogleCountTokensTool, ...] | None = ()
|
||||
|
||||
|
||||
router: Final = APIRouter(
|
||||
tags=["google genai endpoints"],
|
||||
)
|
||||
|
||||
|
||||
def _input_schema(declaration: _GoogleCountTokensFunctionDeclaration) -> object:
|
||||
if declaration.parameters_json_schema is not None:
|
||||
return declaration.parameters_json_schema
|
||||
if declaration.parameters is not None:
|
||||
return normalize_json_schema_types(declaration.parameters)
|
||||
return {"type": "object", "properties": {}}
|
||||
|
||||
|
||||
_GOOGLE_RPC_STATUS: Final = MappingProxyType(
|
||||
{
|
||||
400: "INVALID_ARGUMENT",
|
||||
401: "UNAUTHENTICATED",
|
||||
403: "PERMISSION_DENIED",
|
||||
404: "NOT_FOUND",
|
||||
409: "ABORTED",
|
||||
429: "RESOURCE_EXHAUSTED",
|
||||
499: "CANCELLED",
|
||||
501: "UNIMPLEMENTED",
|
||||
503: "UNAVAILABLE",
|
||||
504: "DEADLINE_EXCEEDED",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _google_error_response(request: Request, error: ProxyException) -> JSONResponse:
|
||||
from litellm.proxy.proxy_server import (
|
||||
_close_dangling_otel_server_span, # pyright: ignore[reportPrivateUsage] # proxy_server keeps the span-close helper private; error JSONResponses returned by the route must stamp the OTel server span like the global ProxyException handler does
|
||||
)
|
||||
|
||||
status_code: Final = int(error.code) if error.code is not None and error.code.isdigit() else 500
|
||||
_close_dangling_otel_server_span(request, status_code, exc=error)
|
||||
status: Final = _GOOGLE_RPC_STATUS.get(status_code, "INTERNAL" if status_code >= 500 else "FAILED_PRECONDITION")
|
||||
return JSONResponse(
|
||||
status_code=status_code,
|
||||
content={"error": {"code": status_code, "message": error.message, "status": status}},
|
||||
headers=error.headers,
|
||||
)
|
||||
|
||||
|
||||
def _google_count_tokens_request(model_name: str, data: object) -> TokenCountRequest:
|
||||
from litellm.google_genai.adapters.transformation import GoogleGenAIAdapter
|
||||
|
||||
validated_data: Final = TypeAdapter(_GoogleCountTokensRequest).validate_python(data)
|
||||
contents: Final = [dict(content) for content in validated_data.contents]
|
||||
messages: Final = (
|
||||
GoogleGenAIAdapter().translate_generate_content_to_completion(model_name, contents).get("messages", [])
|
||||
)
|
||||
system: Final = (
|
||||
[{"type": "text", "text": part.text} for part in validated_data.system_instruction.parts]
|
||||
if validated_data.system_instruction is not None
|
||||
else None
|
||||
)
|
||||
declarations: Final = chain.from_iterable(tool.function_declarations for tool in validated_data.tools or ())
|
||||
tools: Final = [
|
||||
{
|
||||
"name": declaration.name,
|
||||
**({"description": declaration.description} if declaration.description is not None else {}),
|
||||
"input_schema": _input_schema(declaration),
|
||||
}
|
||||
for declaration in declarations
|
||||
]
|
||||
return TokenCountRequest(
|
||||
model=model_name,
|
||||
contents=contents,
|
||||
messages=messages,
|
||||
system=system,
|
||||
tools=tools or None,
|
||||
)
|
||||
|
||||
|
||||
@router.post(
|
||||
"/v1beta/models/{model_name:path}:generateContent",
|
||||
dependencies=[Depends(user_api_key_auth)],
|
||||
|
|
@ -168,36 +288,39 @@ async def google_count_tokens(request: Request, model_name: str):
|
|||
}
|
||||
```
|
||||
"""
|
||||
from litellm.google_genai.adapters.transformation import GoogleGenAIAdapter
|
||||
from litellm.proxy.common_utils.http_parsing_utils import read_request_body
|
||||
from litellm.proxy.proxy_server import token_counter as internal_token_counter
|
||||
|
||||
data: Final = await read_request_body(request=request)
|
||||
contents: Final = data.get("contents", [])
|
||||
# Create TokenCountRequest for the internal endpoint
|
||||
from litellm.proxy._types import TokenCountRequest
|
||||
try:
|
||||
token_request: Final = _google_count_tokens_request(model_name, data)
|
||||
except ValidationError:
|
||||
return _google_error_response(
|
||||
request,
|
||||
ProxyException(
|
||||
message="Invalid Google countTokens request body", type="invalid_request_error", param=None, code=400
|
||||
),
|
||||
)
|
||||
|
||||
# Translate contents to openai format messages using the adapter
|
||||
messages = GoogleGenAIAdapter().translate_generate_content_to_completion(model_name, contents).get("messages", [])
|
||||
|
||||
token_request: Final = TokenCountRequest(
|
||||
model=model_name,
|
||||
contents=contents,
|
||||
messages=messages, # compatibility when use openai-like endpoint
|
||||
)
|
||||
|
||||
# Call the internal token counter function with direct request flag set to False
|
||||
token_response: Final = await internal_token_counter(
|
||||
request=token_request,
|
||||
call_endpoint=True,
|
||||
)
|
||||
if token_response is None:
|
||||
return TokenCountDetailsResponse(totalTokens=0, promptTokensDetails=[])
|
||||
try:
|
||||
token_response: Final = await internal_token_counter(request=token_request, call_endpoint=True)
|
||||
except ProxyException as error:
|
||||
return _google_error_response(request, error)
|
||||
except HTTPException as error:
|
||||
return _google_error_response(request, proxy_exception_from_http_exception(error, {}))
|
||||
except Exception as error:
|
||||
log_llm_api_exception(error, resolve_litellm_call_id(request.headers.get("x-litellm-call-id")))
|
||||
return _google_error_response(
|
||||
request,
|
||||
ProxyException(message="Internal server error", type="internal_server_error", param=None, code=500),
|
||||
)
|
||||
original_response: Final[dict] = token_response.original_response or {}
|
||||
return TokenCountDetailsResponse(
|
||||
totalTokens=original_response.get("totalTokens") or token_response.total_tokens or 0,
|
||||
promptTokensDetails=original_response.get("promptTokensDetails", []),
|
||||
)
|
||||
counted: Final[TokenCountDetailsResponse] = {
|
||||
"totalTokens": original_response.get("totalTokens") or token_response.total_tokens or 0,
|
||||
"promptTokensDetails": original_response.get("promptTokensDetails", []),
|
||||
**count_source_fields(token_response),
|
||||
}
|
||||
return counted
|
||||
|
||||
|
||||
# ============================================================
|
||||
|
|
|
|||
|
|
@ -31,6 +31,7 @@ from litellm.constants import (
|
|||
SESSION_ID_GENERATED_METADATA_KEY,
|
||||
SESSION_ID_OMITTED_METADATA_KEY,
|
||||
TYPESAFE_DANGER_ZONE_HEADER_PREFIX,
|
||||
VECTOR_STORE_CALLER_METADATA_KWARG,
|
||||
X_LITELLM_DISABLE_CALLBACKS,
|
||||
)
|
||||
from litellm.integrations.custom_guardrail import without_server_streaming_classification
|
||||
|
|
@ -2099,7 +2100,7 @@ def refresh_proxy_server_request_body_snapshot(
|
|||
if not isinstance(proxy_server_request, dict):
|
||||
return
|
||||
_body_snapshot_exclude: Final = (
|
||||
frozenset({"secret_fields", "proxy_server_request", "litellm_logging_obj"})
|
||||
frozenset({"secret_fields", "proxy_server_request", "litellm_logging_obj", VECTOR_STORE_CALLER_METADATA_KWARG})
|
||||
| _TRANSPORT_ONLY_CREDENTIAL_KEYS
|
||||
| _CALLBACK_CREDENTIAL_KEYS
|
||||
)
|
||||
|
|
|
|||
|
|
@ -47,6 +47,7 @@ from typing import (
|
|||
)
|
||||
|
||||
import anyio
|
||||
import anyio.lowlevel
|
||||
import websockets
|
||||
import websockets.exceptions
|
||||
from pydantic import BaseModel, ConfigDict, Json, JsonValue, TypeAdapter, ValidationError, with_config
|
||||
|
|
@ -87,6 +88,13 @@ from litellm.litellm_core_utils.litellm_logging import (
|
|||
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
|
||||
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
|
||||
from litellm.litellm_core_utils.token_counter import offload_token_count
|
||||
from litellm.llms.bedrock.count_tokens.kimi_k3 import (
|
||||
HUGGING_FACE_SOURCE,
|
||||
KimiK3Encodings,
|
||||
kimi_k3_anthropic_tokens,
|
||||
kimi_k3_template_tokens,
|
||||
kimi_k3_tokens,
|
||||
)
|
||||
from litellm.proxy._types import (
|
||||
UI_TEAM_ID,
|
||||
CallbackDelete,
|
||||
|
|
@ -136,6 +144,15 @@ from litellm.proxy.common_utils.realtime_utils import ( # noqa: F401, RUF100 #
|
|||
_realtime_request_body, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
|
||||
realtime_request_body,
|
||||
)
|
||||
from litellm.proxy.common_utils.token_count_settings import (
|
||||
TOKENIZER_LOAD_MAX_RUNNING,
|
||||
TOKENIZER_LOAD_REUSE_SECONDS,
|
||||
TOKENIZER_LOAD_WAIT_SECONDS,
|
||||
ProviderCountFailure,
|
||||
TokenCountSettings,
|
||||
cost_map_entry,
|
||||
resolve_token_count_settings,
|
||||
)
|
||||
from litellm.proxy.management_helpers.auto_router_availability import AutoRouterCatalogEntry, build_auto_router_catalog
|
||||
from litellm.router_utils.access_windows import access_windows_config_error
|
||||
from litellm.router_utils.add_retry_fallback_headers import (
|
||||
|
|
@ -165,6 +182,7 @@ from litellm.types.utils import (
|
|||
ModelResponseStream,
|
||||
StreamingChoices,
|
||||
TextCompletionResponse,
|
||||
TokenCountMessagesFormat,
|
||||
TokenCountResponse,
|
||||
echoed_cost_map_pricing_fields,
|
||||
is_server_derived_pricing_key,
|
||||
|
|
@ -289,6 +307,7 @@ from litellm.constants import (
|
|||
CLI_SSO_SESSION_TTL_SECONDS,
|
||||
DAYS_IN_A_MONTH,
|
||||
DEFAULT_HEALTH_CHECK_INTERVAL,
|
||||
DEFAULT_MAX_LRU_CACHE_SIZE,
|
||||
DEFAULT_MODEL_CREATED_AT_TIME,
|
||||
GLOBAL_PROXY_SPEND_CACHE_KEY,
|
||||
LITELLM_PROXY_ADMIN_NAME,
|
||||
|
|
@ -318,6 +337,7 @@ from litellm.litellm_core_utils.agentic_loop_settings import (
|
|||
)
|
||||
from litellm.litellm_core_utils.audio_utils.utils import resolve_speech_media_type
|
||||
from litellm.litellm_core_utils.core_helpers import (
|
||||
as_str_mapping,
|
||||
drop_params_flag,
|
||||
get_litellm_metadata_from_kwargs,
|
||||
get_parent_otel_span_from_kwargs,
|
||||
|
|
@ -960,7 +980,13 @@ from litellm.types.secret_managers.main import (
|
|||
KeyManagementSettings,
|
||||
KeyManagementSystem,
|
||||
)
|
||||
from litellm.types.utils import CredentialItem, CustomHuggingfaceTokenizer, RawRequestTypedDict, StandardLoggingPayload
|
||||
from litellm.types.utils import (
|
||||
CredentialItem,
|
||||
CustomHuggingfaceTokenizer,
|
||||
RawRequestTypedDict,
|
||||
SelectTokenizerResponse,
|
||||
StandardLoggingPayload,
|
||||
)
|
||||
from litellm.types.utils import ModelInfo as ModelMapInfo
|
||||
from litellm.utils import add_custom_logger_callback_to_specific_event
|
||||
|
||||
|
|
@ -14206,52 +14232,362 @@ async def _try_provider_token_count(
|
|||
provider_counter: "BaseTokenCounter",
|
||||
custom_llm_provider: str | None,
|
||||
model_to_use: str,
|
||||
region: str | None,
|
||||
messages: list | None,
|
||||
contents: list | None,
|
||||
deployment: dict[str, object] | None,
|
||||
deployment: Mapping[str, object] | None,
|
||||
request_model: str,
|
||||
tools: list | None = None,
|
||||
system: str | None = None,
|
||||
) -> Optional["TokenCountResponse"]:
|
||||
"""Attempt provider-specific token counting. Returns result on success, None to fall through to local counting."""
|
||||
messages_format: TokenCountMessagesFormat = "openai",
|
||||
output_format: JsonValue = None,
|
||||
) -> Union["TokenCountResponse", ProviderCountFailure, None]:
|
||||
"""Attempt provider-specific token counting. Returns None to fall through to local counting."""
|
||||
if not provider_counter.should_use_token_counting_api(custom_llm_provider=custom_llm_provider):
|
||||
return None
|
||||
deployment_in_region: Final = (
|
||||
None
|
||||
if deployment is None
|
||||
else dict(deployment)
|
||||
if region is None
|
||||
else {
|
||||
**deployment,
|
||||
"litellm_params": {**(as_str_mapping(deployment.get("litellm_params")) or {}), "aws_region_name": region},
|
||||
}
|
||||
)
|
||||
try:
|
||||
result: Final = await provider_counter.count_tokens(
|
||||
model_to_use=model_to_use or "",
|
||||
messages=messages,
|
||||
contents=contents,
|
||||
deployment=deployment,
|
||||
deployment=deployment_in_region,
|
||||
request_model=request_model,
|
||||
tools=tools,
|
||||
system=system,
|
||||
messages_format=messages_format,
|
||||
output_format=output_format,
|
||||
)
|
||||
except httpx.HTTPStatusError as e:
|
||||
error_message: Final = getattr(e, "message", None) or str(e)
|
||||
status_code: Final = getattr(e, "status_code", None) or e.response.status_code
|
||||
raise ProxyException(
|
||||
message=error_message,
|
||||
type="token_counting_error",
|
||||
param="model",
|
||||
code=status_code,
|
||||
return ProviderCountFailure(
|
||||
status_code=getattr(e, "status_code", None) or e.response.status_code,
|
||||
message=getattr(e, "message", None) or str(e),
|
||||
)
|
||||
if result is not None and result.error is True:
|
||||
if litellm.disable_token_counter is True:
|
||||
raise ProxyException(
|
||||
message=result.error_message or "Token counting failed",
|
||||
type="token_counting_error",
|
||||
param="model",
|
||||
code=result.status_code or 500,
|
||||
)
|
||||
verbose_proxy_logger.warning(
|
||||
"Provider token counting failed (%s): %s. Falling back to local tokenizer.",
|
||||
result.status_code,
|
||||
result.error_message,
|
||||
return ProviderCountFailure(
|
||||
status_code=result.status_code or 500,
|
||||
message=result.error_message or "Token counting failed",
|
||||
rejected_before_provider_call=result.rejected_before_provider_call,
|
||||
)
|
||||
return None
|
||||
return result
|
||||
|
||||
|
||||
async def _provider_token_count(
|
||||
count_with: Callable[[str, str | None], Awaitable[Union["TokenCountResponse", ProviderCountFailure, None]]],
|
||||
model_to_use: str,
|
||||
settings: TokenCountSettings,
|
||||
local_estimate_allowed: bool,
|
||||
) -> Union["TokenCountResponse", ProviderCountFailure, None]:
|
||||
count_model: Final = settings.count_model
|
||||
counted: Final = await (
|
||||
count_with(model_to_use, None) if count_model is None else count_with(count_model.model, count_model.region)
|
||||
)
|
||||
if counted is None:
|
||||
return None
|
||||
if isinstance(counted, TokenCountResponse):
|
||||
if count_model is None:
|
||||
return counted.model_copy(update={"count_source": "provider"})
|
||||
return counted.model_copy(update={"count_source": "token_count_model", "model_used": count_model.model})
|
||||
failure: Final = (
|
||||
counted
|
||||
if count_model is None or count_model.region is None
|
||||
else counted.with_hint(
|
||||
f"This model counts tokens with {count_model.model} in {count_model.region}. "
|
||||
"If that region is not enabled for the AWS account, enable it or set token_count_model in model_info"
|
||||
)
|
||||
)
|
||||
if local_estimate_allowed:
|
||||
return failure
|
||||
raise failure.to_proxy_exception()
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _LocalTokenEstimate:
|
||||
total_tokens: int
|
||||
tokenizer_type: str
|
||||
|
||||
|
||||
def _is_fork_guard_error(error: BaseException) -> bool:
|
||||
from litellm.rust_bridge.fork_guard import ForkedAfterNativeRuntimeStarted, ProcessReservedForForking
|
||||
|
||||
return isinstance(error, (ForkedAfterNativeRuntimeStarted, ProcessReservedForForking))
|
||||
|
||||
|
||||
def _is_tokenizer_panic(error: BaseException) -> bool:
|
||||
return type(error).__module__ == "pyo3_runtime" and type(error).__name__ == "PanicException"
|
||||
|
||||
|
||||
def _answers_provider_error(error: BaseException) -> bool:
|
||||
return (isinstance(error, Exception) or _is_tokenizer_panic(error)) and not _is_fork_guard_error(error)
|
||||
|
||||
|
||||
def _provider_error_after_failed_estimate(
|
||||
provider_failure: ProviderCountFailure, error: BaseException
|
||||
) -> ProxyException:
|
||||
verbose_proxy_logger.warning(
|
||||
"Local token estimate failed after the provider's count error %s, answering that error: %s",
|
||||
provider_failure.status_code,
|
||||
type(error).__name__,
|
||||
)
|
||||
return provider_failure.to_proxy_exception()
|
||||
|
||||
|
||||
async def _estimate_or_provider_error(estimate: Awaitable[int], provider_failure: ProviderCountFailure) -> int:
|
||||
try:
|
||||
return await estimate
|
||||
except BaseException as e:
|
||||
if not _answers_provider_error(e):
|
||||
raise
|
||||
raise _provider_error_after_failed_estimate(provider_failure, e) from e
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _CustomTokenizerKey:
|
||||
identifier: str
|
||||
revision: str
|
||||
auth_token: str | None
|
||||
|
||||
|
||||
class _TokenizerLoadsBusy(Exception):
|
||||
pass
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _TokenizerLoadOutcome:
|
||||
tokenizer: SelectTokenizerResponse | BaseException
|
||||
finished_at: float
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _SharedTokenizerLoad:
|
||||
outcome: asyncio.Task[_TokenizerLoadOutcome]
|
||||
started_at: float
|
||||
|
||||
def reusable(self, now: float) -> bool:
|
||||
return not self.outcome.done() or now - self.outcome.result().finished_at < TOKENIZER_LOAD_REUSE_SECONDS
|
||||
|
||||
|
||||
async def _loaded_tokenizer(
|
||||
model: str, custom_tokenizer: CustomHuggingfaceTokenizer | None, threads: anyio.CapacityLimiter
|
||||
) -> SelectTokenizerResponse | BaseException:
|
||||
try:
|
||||
return await asyncify(litellm.utils._select_tokenizer, limiter=threads)(
|
||||
model=model, custom_tokenizer=custom_tokenizer
|
||||
)
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except BaseException as e: # noqa: BLE001 # every count sharing this load decides what it answers, pyo3 panics included
|
||||
return e
|
||||
|
||||
|
||||
async def _tokenizer_load_outcome(
|
||||
model: str, custom_tokenizer: CustomHuggingfaceTokenizer | None, threads: anyio.CapacityLimiter
|
||||
) -> _TokenizerLoadOutcome:
|
||||
tokenizer: Final = await _loaded_tokenizer(model, custom_tokenizer, threads)
|
||||
return _TokenizerLoadOutcome(tokenizer=tokenizer, finished_at=asyncio.get_running_loop().time())
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class _TokenizerLoads:
|
||||
threads: anyio.CapacityLimiter
|
||||
running: int = 0
|
||||
|
||||
def start(
|
||||
self, model: str, custom_tokenizer: CustomHuggingfaceTokenizer | None, loop: asyncio.AbstractEventLoop
|
||||
) -> _SharedTokenizerLoad | None:
|
||||
if self.running >= self.threads.total_tokens:
|
||||
return None
|
||||
self.running += 1
|
||||
outcome: Final = loop.create_task(_tokenizer_load_outcome(model, custom_tokenizer, self.threads))
|
||||
outcome.add_done_callback(self._finished)
|
||||
return _SharedTokenizerLoad(outcome=outcome, started_at=loop.time())
|
||||
|
||||
def _finished(self, _outcome: asyncio.Task[_TokenizerLoadOutcome]) -> None:
|
||||
self.running -= 1
|
||||
|
||||
|
||||
_TOKENIZER_LOADS: Final = anyio.lowlevel.RunVar[_TokenizerLoads]("litellm_proxy_tokenizer_loads")
|
||||
|
||||
|
||||
def _tokenizer_loads_for_this_loop() -> _TokenizerLoads:
|
||||
existing: Final = _TOKENIZER_LOADS.get(None)
|
||||
if existing is not None:
|
||||
return existing
|
||||
created: Final = _TokenizerLoads(threads=anyio.CapacityLimiter(TOKENIZER_LOAD_MAX_RUNNING))
|
||||
_TOKENIZER_LOADS.set(created)
|
||||
return created
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class _TokenizerLoadSlot:
|
||||
model: str
|
||||
custom_tokenizer: CustomHuggingfaceTokenizer | None
|
||||
current: _SharedTokenizerLoad | None = None
|
||||
|
||||
def load(self, loads: _TokenizerLoads, loop: asyncio.AbstractEventLoop) -> _SharedTokenizerLoad | None:
|
||||
current: Final = self.current
|
||||
if current is not None and current.reusable(loop.time()):
|
||||
return current
|
||||
self.current = loads.start(self.model, self.custom_tokenizer, loop)
|
||||
return self.current
|
||||
|
||||
|
||||
@lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE)
|
||||
def _tokenizer_load_slot(
|
||||
model: str, custom_tokenizer: _CustomTokenizerKey | None, loop: asyncio.AbstractEventLoop
|
||||
) -> _TokenizerLoadSlot:
|
||||
return _TokenizerLoadSlot(
|
||||
model=model,
|
||||
custom_tokenizer=None
|
||||
if custom_tokenizer is None
|
||||
else CustomHuggingfaceTokenizer(
|
||||
identifier=custom_tokenizer.identifier,
|
||||
revision=custom_tokenizer.revision,
|
||||
auth_token=custom_tokenizer.auth_token,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
async def _tokenizer_after_provider_error(
|
||||
model: str, custom_tokenizer: CustomHuggingfaceTokenizer | None
|
||||
) -> SelectTokenizerResponse | BaseException:
|
||||
loop: Final = asyncio.get_running_loop()
|
||||
try:
|
||||
load: Final = _tokenizer_load_slot(
|
||||
model,
|
||||
None
|
||||
if custom_tokenizer is None
|
||||
else _CustomTokenizerKey(
|
||||
identifier=custom_tokenizer["identifier"],
|
||||
revision=custom_tokenizer["revision"],
|
||||
auth_token=custom_tokenizer["auth_token"],
|
||||
),
|
||||
loop,
|
||||
).load(_tokenizer_loads_for_this_loop(), loop)
|
||||
if load is None:
|
||||
return _TokenizerLoadsBusy()
|
||||
outcome: Final = await asyncio.wait_for(
|
||||
asyncio.shield(load.outcome),
|
||||
timeout=max(load.started_at + TOKENIZER_LOAD_WAIT_SECONDS - loop.time(), 0),
|
||||
)
|
||||
except Exception as e: # noqa: BLE001 # a tokenizer that cannot be read, loaded in time, or loaded at all answers the provider's error
|
||||
return e
|
||||
return outcome.tokenizer
|
||||
|
||||
|
||||
async def _tokenizer_or_provider_error(
|
||||
model: str, custom_tokenizer: CustomHuggingfaceTokenizer | None, provider_failure: ProviderCountFailure
|
||||
) -> SelectTokenizerResponse:
|
||||
tokenizer: Final = await _tokenizer_after_provider_error(model, custom_tokenizer)
|
||||
if not isinstance(tokenizer, BaseException):
|
||||
return tokenizer
|
||||
if not _answers_provider_error(tokenizer):
|
||||
raise tokenizer.with_traceback(None)
|
||||
raise _provider_error_after_failed_estimate(provider_failure, tokenizer) from tokenizer
|
||||
|
||||
|
||||
async def _local_token_estimate(
|
||||
model: str,
|
||||
prompt: str | None,
|
||||
messages: Sequence[AllMessageValues] | None,
|
||||
tools: list[ChatCompletionToolParam] | None,
|
||||
custom_tokenizer: CustomHuggingfaceTokenizer | None,
|
||||
provider_failure: ProviderCountFailure | None,
|
||||
) -> _LocalTokenEstimate:
|
||||
tokenizer: Final = (
|
||||
await asyncify(litellm.utils._select_tokenizer)(model=model, custom_tokenizer=custom_tokenizer)
|
||||
if provider_failure is None
|
||||
else await _tokenizer_or_provider_error(model, custom_tokenizer, provider_failure)
|
||||
)
|
||||
count: Final = offload_token_count(litellm.token_counter)(
|
||||
model=model,
|
||||
text=prompt,
|
||||
messages=messages,
|
||||
tools=tools,
|
||||
custom_tokenizer=tokenizer,
|
||||
)
|
||||
total_tokens: Final = await (
|
||||
count if provider_failure is None else _estimate_or_provider_error(count, provider_failure)
|
||||
)
|
||||
return _LocalTokenEstimate(total_tokens=total_tokens, tokenizer_type=str(tokenizer["type"]))
|
||||
|
||||
|
||||
_KIMI_K3_ENCODINGS: Final = KimiK3Encodings(
|
||||
source=HUGGING_FACE_SOURCE, retry_after_seconds=TOKENIZER_LOAD_REUSE_SECONDS
|
||||
)
|
||||
|
||||
|
||||
async def _kimi_k3_count(
|
||||
request: TokenCountRequest,
|
||||
messages_format: TokenCountMessagesFormat,
|
||||
model: str,
|
||||
messages: Sequence[AllMessageValues],
|
||||
tools: Sequence[ChatCompletionToolParam] | None,
|
||||
) -> int | None:
|
||||
encoding: Final = await _KIMI_K3_ENCODINGS.encoding(TOKENIZER_LOAD_WAIT_SECONDS)
|
||||
if encoding is None:
|
||||
return None
|
||||
match messages_format:
|
||||
case "anthropic":
|
||||
return await offload_token_count(kimi_k3_anthropic_tokens)(
|
||||
encoding, model, request.model_dump(include={"messages", "system", "tools"}, exclude_none=True)
|
||||
)
|
||||
case "openai":
|
||||
return await offload_token_count(kimi_k3_tokens)(encoding, model, messages, tools)
|
||||
|
||||
|
||||
class _TokenCountDeployment(TypedDict, total=False):
|
||||
litellm_params: ReadOnly[Mapping[str, object]]
|
||||
model_info: ReadOnly[Mapping[str, object]]
|
||||
|
||||
|
||||
_TOKEN_COUNT_DEPLOYMENT: Final = TypeAdapter(_TokenCountDeployment)
|
||||
|
||||
|
||||
def _model_group_deployment(router: Router | None, model_group: str) -> Mapping[str, object] | None:
|
||||
group: Final = () if router is None else router._get_all_deployments(model_name=model_group)
|
||||
return group[0] if group else None
|
||||
|
||||
|
||||
def _non_empty_str(value: object) -> str | None:
|
||||
return value if isinstance(value, str) and value else None
|
||||
|
||||
|
||||
def _deployment_aws_region(deployment: _TokenCountDeployment) -> str:
|
||||
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
|
||||
|
||||
params: Final = deployment.get("litellm_params", {})
|
||||
return BaseAWSLLM()._get_aws_region_name(
|
||||
{"aws_region_name": params.get("aws_region_name")}, model=_non_empty_str(params.get("model"))
|
||||
)
|
||||
|
||||
|
||||
def _token_count_settings_model(deployment: _TokenCountDeployment, request_model: str) -> tuple[str, str | None]:
|
||||
model_info: Final = deployment.get("model_info", {})
|
||||
params: Final = deployment.get("litellm_params", {})
|
||||
base_model: Final = _non_empty_str(model_info.get("base_model")) or _non_empty_str(params.get("base_model"))
|
||||
if base_model is not None:
|
||||
return base_model, None
|
||||
model: Final = _non_empty_str(params.get("model"))
|
||||
if model is None:
|
||||
return request_model, None
|
||||
return model, _non_empty_str(params.get("custom_llm_provider"))
|
||||
|
||||
|
||||
_PROVIDERS_COUNTING_WITHOUT_MESSAGES: Final = frozenset(
|
||||
{litellm.LlmProviders.GEMINI.value, litellm.LlmProviders.VERTEX_AI.value}
|
||||
)
|
||||
|
||||
|
||||
def _system_message(system: object) -> ChatCompletionSystemMessage | None:
|
||||
if not isinstance(system, (str, list)) or not system:
|
||||
return None
|
||||
|
|
@ -14274,8 +14610,15 @@ async def token_counter(request: TokenCountRequest, call_endpoint: bool = False)
|
|||
Returns:
|
||||
TokenCountResponse
|
||||
"""
|
||||
global llm_router
|
||||
return await count_request_tokens(request=request, call_endpoint=call_endpoint)
|
||||
|
||||
|
||||
async def count_request_tokens(
|
||||
request: TokenCountRequest,
|
||||
call_endpoint: bool = False,
|
||||
messages_format: TokenCountMessagesFormat = "openai",
|
||||
output_format: JsonValue = None,
|
||||
) -> TokenCountResponse:
|
||||
prompt: Final = request.prompt
|
||||
messages: Final = request.messages
|
||||
contents: Final = request.contents
|
||||
|
|
@ -14311,8 +14654,16 @@ async def token_counter(request: TokenCountRequest, call_endpoint: bool = False)
|
|||
litellm_model_name = litellm_model_name.split("/", 1)[1]
|
||||
|
||||
model_to_use: str = (
|
||||
litellm_model_name or request.model
|
||||
) # use litellm model name, if it's not avalable then fallback to request.model
|
||||
litellm_model_name if isinstance(litellm_model_name, str) and litellm_model_name else request.model
|
||||
)
|
||||
deployment_view: Final = _TOKEN_COUNT_DEPLOYMENT.validate_python(
|
||||
deployment or _model_group_deployment(llm_router, request.model) or {}
|
||||
)
|
||||
settings: Final = await asyncify(resolve_token_count_settings)(
|
||||
configured=deployment_view.get("model_info", {}),
|
||||
default=await asyncify(cost_map_entry)(*_token_count_settings_model(deployment_view, request.model)),
|
||||
deployment_region=partial(_deployment_aws_region, deployment_view),
|
||||
)
|
||||
|
||||
# Try provider-specific token counting first - only for non-direct requests (from provider endpoints)
|
||||
provider_counter: BaseTokenCounter | None = None
|
||||
|
|
@ -14323,20 +14674,33 @@ async def token_counter(request: TokenCountRequest, call_endpoint: bool = False)
|
|||
if _model is not None:
|
||||
model_to_use = _model
|
||||
|
||||
if provider_counter is not None:
|
||||
result: Final = await _try_provider_token_count(
|
||||
provider_counter=provider_counter,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
conversation_empty: Final = not messages and not contents
|
||||
prompt_only: Final = bool(prompt) and conversation_empty
|
||||
prompt_as_message: Final = prompt_only and custom_llm_provider in _PROVIDERS_COUNTING_WITHOUT_MESSAGES
|
||||
ladder: Final = (
|
||||
None
|
||||
if provider_counter is None
|
||||
else await _provider_token_count(
|
||||
count_with=partial(
|
||||
_try_provider_token_count,
|
||||
provider_counter,
|
||||
custom_llm_provider,
|
||||
messages=[{"role": "user", "content": prompt}] if prompt_as_message else messages,
|
||||
contents=None if prompt_as_message else contents,
|
||||
deployment=deployment,
|
||||
request_model=request.model,
|
||||
tools=tools,
|
||||
system=system,
|
||||
messages_format=messages_format,
|
||||
output_format=output_format,
|
||||
),
|
||||
model_to_use=model_to_use,
|
||||
messages=messages,
|
||||
contents=contents,
|
||||
deployment=deployment,
|
||||
request_model=request.model,
|
||||
tools=tools,
|
||||
system=system,
|
||||
settings=settings,
|
||||
local_estimate_allowed=litellm.disable_token_counter is not True,
|
||||
)
|
||||
if result is not None:
|
||||
return result
|
||||
)
|
||||
if isinstance(ladder, TokenCountResponse):
|
||||
return ladder
|
||||
|
||||
# Check if token counter is disabled before fallback
|
||||
if litellm.disable_token_counter is True:
|
||||
|
|
@ -14354,14 +14718,11 @@ async def token_counter(request: TokenCountRequest, call_endpoint: bool = False)
|
|||
CustomHuggingfaceTokenizer | None,
|
||||
model_info.get("custom_tokenizer", None),
|
||||
)
|
||||
_tokenizer_used: Final = await asyncify(litellm.utils.select_tokenizer)(
|
||||
model=model_to_use, custom_tokenizer=custom_tokenizer
|
||||
)
|
||||
|
||||
tokenizer_used: Final = str(_tokenizer_used["type"])
|
||||
system_message: Final = _system_message(system)
|
||||
prompt_turn: Final = prompt_only and not (system_message is None and not tools)
|
||||
typed_messages: Final = cast( # cast-ok: request messages are raw chat-shaped dicts that token_counter normalizes
|
||||
Sequence[AllMessageValues] | None, messages
|
||||
Sequence[AllMessageValues] | None,
|
||||
[{"role": "user", "content": prompt}] if prompt_turn else (None if prompt_only else messages),
|
||||
)
|
||||
counted_messages: Final = (
|
||||
typed_messages if typed_messages is None or system_message is None else (system_message, *typed_messages)
|
||||
|
|
@ -14369,18 +14730,38 @@ async def token_counter(request: TokenCountRequest, call_endpoint: bool = False)
|
|||
counted_tools: Final = cast( # cast-ok: raw OpenAI or Anthropic tool dicts, both of which token_counter formats
|
||||
list[ChatCompletionToolParam] | None, tools if counted_messages is not None else None
|
||||
)
|
||||
total_tokens: Final = await offload_token_count(litellm.token_counter)(
|
||||
framing: Final = settings.framing if counted_messages is not None and settings.framing is not None else 0
|
||||
kimi_k3_messages: Final = counted_messages if settings.tokenizer == "kimi_k3" else None
|
||||
exact: Final = (
|
||||
await _kimi_k3_count(request, messages_format, model_to_use, kimi_k3_messages, counted_tools)
|
||||
if kimi_k3_messages is not None and custom_tokenizer is None
|
||||
else None
|
||||
)
|
||||
if exact is not None:
|
||||
return TokenCountResponse(
|
||||
total_tokens=exact + framing,
|
||||
request_model=request.model,
|
||||
model_used=model_to_use,
|
||||
tokenizer_type="huggingface_tokenizer",
|
||||
count_source="local_estimate",
|
||||
)
|
||||
estimate: Final = await _local_token_estimate(
|
||||
model=model_to_use,
|
||||
text=prompt,
|
||||
prompt=prompt if counted_messages is None else None,
|
||||
messages=counted_messages,
|
||||
tools=counted_tools,
|
||||
custom_tokenizer=_tokenizer_used,
|
||||
custom_tokenizer=custom_tokenizer,
|
||||
provider_failure=ladder,
|
||||
)
|
||||
template: Final = (
|
||||
0 if kimi_k3_messages is None else kimi_k3_template_tokens(len(kimi_k3_messages), bool(counted_tools))
|
||||
)
|
||||
return TokenCountResponse(
|
||||
total_tokens=total_tokens,
|
||||
total_tokens=estimate.total_tokens + template + framing,
|
||||
request_model=request.model,
|
||||
model_used=model_to_use,
|
||||
tokenizer_type=tokenizer_used,
|
||||
tokenizer_type=estimate.tokenizer_type,
|
||||
count_source="local_estimate",
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -36,6 +36,7 @@ from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # leg
|
|||
read_request_body,
|
||||
safe_set_request_parsed_body,
|
||||
)
|
||||
from litellm.proxy.common_utils.token_count_settings import count_source_fields
|
||||
from litellm.proxy.route_llm_request import raise_if_required_body_param_missing
|
||||
from litellm.types.llms.base import LiteLLMBaseModel
|
||||
from litellm.types.llms.openai import (
|
||||
|
|
@ -44,6 +45,7 @@ from litellm.types.llms.openai import (
|
|||
ResponsesAPIResponse,
|
||||
)
|
||||
from litellm.types.responses.main import DeleteResponseResult
|
||||
from litellm.types.token_count import TokenCountSourceFields
|
||||
from litellm.types.utils import TokenCountResponse
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
|
@ -1067,7 +1069,7 @@ class _ResponsesApiErrorBody(TypedDict):
|
|||
error: ReadOnly[_ResponsesApiErrorDetail]
|
||||
|
||||
|
||||
class _ResponsesInputTokensResult(TypedDict):
|
||||
class _ResponsesInputTokensResult(TokenCountSourceFields):
|
||||
object: ReadOnly[str]
|
||||
input_tokens: ReadOnly[int]
|
||||
|
||||
|
|
@ -1198,6 +1200,7 @@ async def responses_input_tokens(
|
|||
result: Final[_ResponsesInputTokensResult] = {
|
||||
"object": "response.input_tokens",
|
||||
"input_tokens": token_response.total_tokens,
|
||||
**count_source_fields(token_response),
|
||||
}
|
||||
return result
|
||||
|
||||
|
|
|
|||
|
|
@ -9,12 +9,14 @@ from typing import (
|
|||
|
||||
from fastapi import APIRouter, Depends, HTTPException, Request, Response
|
||||
|
||||
from litellm.constants import VECTOR_STORE_CALLER_METADATA_KWARG
|
||||
from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook import (
|
||||
LiteLLM_ManagedVectorStore,
|
||||
)
|
||||
from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
|
||||
from litellm.proxy.common_utils.http_parsing_utils import read_raw_json_body
|
||||
from litellm.proxy.utils import jsonify_object
|
||||
from litellm.proxy.vector_store_endpoints.utils import (
|
||||
assert_proxy_admin_for_vector_store_index_management,
|
||||
|
|
@ -23,6 +25,7 @@ from litellm.proxy.vector_store_endpoints.utils import (
|
|||
)
|
||||
from litellm.repositories.table_repositories import ManagedVectorStoreIndexRepository
|
||||
from litellm.types.vector_stores import IndexCreateRequest, IndexListResponse
|
||||
from litellm.vector_stores.utils import VectorStoreWriteBody, callers_metadata_snapshot
|
||||
from litellm.vector_stores.vector_store_registry import VectorStoreIndexRegistry
|
||||
|
||||
router: Final = APIRouter()
|
||||
|
|
@ -213,6 +216,9 @@ async def vector_store_create(
|
|||
)
|
||||
|
||||
data: Final = await read_request_body(request=request)
|
||||
data[VECTOR_STORE_CALLER_METADATA_KWARG] = callers_metadata_snapshot(
|
||||
await read_raw_json_body(request), VectorStoreWriteBody.model_validate(data)
|
||||
)
|
||||
|
||||
# Check for target_model_names parameter
|
||||
target_model_names: Final = data.pop("target_model_names", None)
|
||||
|
|
@ -448,6 +454,9 @@ async def vector_store_update(
|
|||
data = await read_request_body(request=request) # rebind-ok: pre-existing rebinding on a rename-only line
|
||||
if "vector_store_id" not in data:
|
||||
data["vector_store_id"] = vector_store_id
|
||||
data[VECTOR_STORE_CALLER_METADATA_KWARG] = callers_metadata_snapshot(
|
||||
await read_raw_json_body(request), VectorStoreWriteBody.model_validate(data)
|
||||
)
|
||||
|
||||
data = await _update_request_data_with_litellm_managed_vector_store_registry(
|
||||
data=data, vector_store_id=vector_store_id, user_api_key_dict=user_api_key_dict
|
||||
|
|
|
|||
|
|
@ -732,10 +732,15 @@ ANTHROPIC_API_ONLY_HEADERS: Final = { # fails if calling anthropic on vertex ai
|
|||
}
|
||||
|
||||
|
||||
class AnthropicThinkingBlockBinding(TypedDict, total=False):
|
||||
prefix_mismatch_behavior: ReadOnly[Literal["error", "drop_block"]]
|
||||
|
||||
|
||||
class AnthropicThinkingParam(TypedDict, total=False):
|
||||
type: ReadOnly[Literal["enabled", "adaptive", "disabled"]]
|
||||
budget_tokens: int
|
||||
display: ReadOnly[Literal["summarized", "omitted", "updates"]]
|
||||
block_binding: ReadOnly[AnthropicThinkingBlockBinding]
|
||||
|
||||
|
||||
class ANTHROPIC_HOSTED_TOOLS(str, Enum):
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ from typing_extensions import (
|
|||
)
|
||||
|
||||
from litellm.types.llms.openai import ChatCompletionFileObject, EmbeddingInput
|
||||
from litellm.types.token_count import TokenCountSourceFields
|
||||
|
||||
# Gemini supports nested-list inputs (e.g. [["text", "image"]]) as an explicit
|
||||
# opt-in for combined embeddings — a provider-specific extension of the
|
||||
|
|
@ -319,12 +320,13 @@ class UsageMetadata(TypedDict, total=False):
|
|||
candidatesTokensDetails: list[PromptTokensDetails] # Alternative key name used in some responses
|
||||
|
||||
|
||||
class TokenCountDetailsResponse(TypedDict):
|
||||
class TokenCountDetailsResponse(TokenCountSourceFields):
|
||||
"""
|
||||
Response structure for token count details with modality breakdown.
|
||||
|
||||
Example:
|
||||
{'totalTokens': 12, 'promptTokensDetails': [{'modality': 'TEXT', 'tokenCount': 12}]}
|
||||
{'totalTokens': 12, 'promptTokensDetails': [{'modality': 'TEXT', 'tokenCount': 12}],
|
||||
'litellm_count_source': 'provider'}
|
||||
"""
|
||||
|
||||
totalTokens: int
|
||||
|
|
|
|||
16
litellm/types/token_count.py
Normal file
16
litellm/types/token_count.py
Normal file
|
|
@ -0,0 +1,16 @@
|
|||
from typing import Literal, TypeAlias
|
||||
|
||||
from typing_extensions import NotRequired, ReadOnly, TypedDict
|
||||
|
||||
TokenCountSource: TypeAlias = Literal["provider", "token_count_model", "local_estimate"]
|
||||
TokenCountTokenizer: TypeAlias = Literal["kimi_k3"]
|
||||
|
||||
|
||||
class TokenCountModelInfo(TypedDict):
|
||||
model: ReadOnly[str]
|
||||
region: ReadOnly[NotRequired[str]]
|
||||
|
||||
|
||||
class TokenCountSourceFields(TypedDict, total=False):
|
||||
litellm_count_source: ReadOnly[TokenCountSource]
|
||||
litellm_count_model: ReadOnly[str]
|
||||
|
|
@ -9,6 +9,7 @@ from typing import (
|
|||
Any,
|
||||
Final,
|
||||
Literal,
|
||||
TypeAlias,
|
||||
get_args,
|
||||
)
|
||||
|
||||
|
|
@ -56,6 +57,7 @@ from litellm.types.llms.base import (
|
|||
LiteLLMPydanticObjectBase,
|
||||
)
|
||||
from litellm.types.mcp import MCPServerCostInfo
|
||||
from litellm.types.token_count import TokenCountModelInfo, TokenCountSource, TokenCountTokenizer
|
||||
from litellm.types.workload_identity import (
|
||||
ANTHROPIC_WIF_KWARGS_KEYS,
|
||||
OPENAI_WIF_KWARGS_KEYS,
|
||||
|
|
@ -227,6 +229,9 @@ class ProviderSpecificModelInfo(TypedDict, total=False):
|
|||
bedrock_output_config_effort_ceiling: Literal["low", "medium", "high", "max", "xhigh"] | None
|
||||
bedrock_converse_supports_strict_tools: bool | None
|
||||
supports_regex_lookaround: ReadOnly[bool | None]
|
||||
token_count_framing: ReadOnly[int | None]
|
||||
token_count_model: ReadOnly[TokenCountModelInfo | None]
|
||||
token_count_tokenizer: ReadOnly[TokenCountTokenizer | None]
|
||||
|
||||
|
||||
class SearchContextCostPerQuery(TypedDict, total=False):
|
||||
|
|
@ -3958,11 +3963,14 @@ DEPLOYMENT_SCOPED_PRICING_FIELDS: Final[frozenset[str]] = frozenset({"off_peak_p
|
|||
|
||||
DEPLOYMENT_SCOPED_CAPABILITY_FIELDS: Final[frozenset[str]] = frozenset({"supports_regex_lookaround"})
|
||||
|
||||
DEPLOYMENT_SCOPED_TOKEN_COUNT_FIELDS: Final[frozenset[str]] = frozenset({"token_count_framing", "token_count_model"})
|
||||
|
||||
SHARED_BACKEND_MODEL_INFO_FIELDS: Final[frozenset[str]] = (
|
||||
frozenset(ModelInfoBase.__required_keys__ | ModelInfoBase.__optional_keys__)
|
||||
- frozenset(CustomPricingLiteLLMParams.model_fields)
|
||||
- DEPLOYMENT_SCOPED_PRICING_FIELDS
|
||||
- DEPLOYMENT_SCOPED_CAPABILITY_FIELDS
|
||||
- DEPLOYMENT_SCOPED_TOKEN_COUNT_FIELDS
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -4364,11 +4372,15 @@ class LiteLLMLoggingBaseClass:
|
|||
pass
|
||||
|
||||
|
||||
TokenCountMessagesFormat: TypeAlias = Literal["openai", "anthropic"]
|
||||
|
||||
|
||||
class TokenCountResponse(LiteLLMPydanticObjectBase):
|
||||
total_tokens: int
|
||||
request_model: str
|
||||
model_used: str
|
||||
tokenizer_type: str
|
||||
count_source: TokenCountSource | None = None
|
||||
original_response: dict | None = None
|
||||
"""
|
||||
Original Response from upstream API call - if an API call was made for token counting
|
||||
|
|
@ -4379,6 +4391,7 @@ class TokenCountResponse(LiteLLMPydanticObjectBase):
|
|||
HTTP status code from the token counting API (e.g., 200 for success, 429 for rate limit, 400 for bad request)
|
||||
"""
|
||||
status_code: int | None = None
|
||||
rejected_before_provider_call: bool = Field(default=False, exclude=True)
|
||||
|
||||
|
||||
class CustomHuggingfaceTokenizer(TypedDict):
|
||||
|
|
|
|||
|
|
@ -6429,6 +6429,9 @@ def get_model_info_helper(
|
|||
uses_embed_content=_model_info.get("uses_embed_content", None),
|
||||
supports_image_size=_model_info.get("supports_image_size", None),
|
||||
supported_audio_formats=_model_info.get("supported_audio_formats", None),
|
||||
token_count_framing=_model_info.get("token_count_framing", None),
|
||||
token_count_model=_model_info.get("token_count_model", None),
|
||||
token_count_tokenizer=_model_info.get("token_count_tokenizer", None),
|
||||
vertex_ai_audio_api=_model_info.get("vertex_ai_audio_api", None),
|
||||
)
|
||||
for cost_key, cost_value in _model_info.items():
|
||||
|
|
|
|||
|
|
@ -1,11 +1,92 @@
|
|||
import copy
|
||||
from collections.abc import Container, Mapping
|
||||
from typing import Any, Final, cast, get_type_hints
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
|
||||
|
||||
from litellm.constants import ROUTING_REQUEST_TAGS_METADATA_KEY, VECTOR_STORE_CALLER_METADATA_KWARG
|
||||
from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig
|
||||
from litellm.types.vector_stores import (
|
||||
VectorStoreCreateOptionalRequestParams,
|
||||
VectorStoreSearchOptionalRequestParams,
|
||||
)
|
||||
|
||||
LITELLM_WRITTEN_METADATA_KEYS: Final = frozenset(
|
||||
{
|
||||
"requester_ip_address",
|
||||
"user_agent",
|
||||
"endpoint",
|
||||
"litellm_api_version",
|
||||
"deployment",
|
||||
"deployment_model_name",
|
||||
"api_base",
|
||||
"agent_id",
|
||||
"actor_agent_id",
|
||||
"target_agent_id",
|
||||
"billing_agent_id",
|
||||
"agent_execution_mode",
|
||||
"verified_human_user_id",
|
||||
"litellm_caller_session_id",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def is_litellm_written_metadata_key(key: str) -> bool:
|
||||
return key == "user_api_key" or key.startswith("user_api_key_") or key in LITELLM_WRITTEN_METADATA_KEYS
|
||||
|
||||
|
||||
def _is_callers_key(key: object) -> bool:
|
||||
return not isinstance(key, str) or not is_litellm_written_metadata_key(key)
|
||||
|
||||
|
||||
def _sent_by_proxy_or_router(params: Mapping[str, object], metadata: Container[object]) -> bool:
|
||||
return params.get("proxy_server_request") is not None or ROUTING_REQUEST_TAGS_METADATA_KEY in metadata
|
||||
|
||||
|
||||
def _callers_metadata(metadata: Mapping[object, object]) -> Mapping[object, object]:
|
||||
return {key: value for key, value in metadata.items() if _is_callers_key(key)}
|
||||
|
||||
|
||||
_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object])
|
||||
|
||||
|
||||
class VectorStoreWriteBody(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="ignore")
|
||||
|
||||
metadata: object = None
|
||||
|
||||
|
||||
def _json_object(raw: str | bytes) -> Mapping[str, object] | None:
|
||||
try:
|
||||
return _JSON_OBJECT.validate_json(raw)
|
||||
except ValidationError:
|
||||
return None
|
||||
|
||||
|
||||
def _decoded_metadata(metadata: object) -> object:
|
||||
if not isinstance(metadata, str):
|
||||
return metadata
|
||||
decoded: Final = _json_object(metadata)
|
||||
return metadata if decoded is None else decoded
|
||||
|
||||
|
||||
def callers_metadata_snapshot(raw_json_body: bytes | None, parsed_body: VectorStoreWriteBody) -> object:
|
||||
raw_body: Final = None if raw_json_body is None else _json_object(raw_json_body)
|
||||
metadata: Final = copy.deepcopy(parsed_body.metadata) if raw_body is None else raw_body.get("metadata")
|
||||
return _decoded_metadata(metadata)
|
||||
|
||||
|
||||
def _metadata_for_provider(params: Mapping[str, object]) -> object:
|
||||
if VECTOR_STORE_CALLER_METADATA_KWARG in params:
|
||||
return params[VECTOR_STORE_CALLER_METADATA_KWARG]
|
||||
metadata: Final = params.get("metadata")
|
||||
if not isinstance(metadata, Mapping):
|
||||
return metadata
|
||||
request_metadata: Final = cast(Mapping[object, object], metadata) # cast-ok: request metadata is a JSON object
|
||||
if not _sent_by_proxy_or_router(params, request_metadata):
|
||||
return request_metadata
|
||||
return _callers_metadata(request_metadata)
|
||||
|
||||
|
||||
class VectorStoreRequestUtils:
|
||||
"""Helper utils for constructing Vector Store search requests"""
|
||||
|
|
@ -48,7 +129,11 @@ class VectorStoreRequestUtils:
|
|||
Returns:
|
||||
VectorStoreCreateOptionalRequestParams instance with only the valid parameters
|
||||
"""
|
||||
valid_keys: Final = get_type_hints(VectorStoreCreateOptionalRequestParams).keys()
|
||||
valid_keys: Final = get_type_hints(VectorStoreCreateOptionalRequestParams).keys() - {"metadata"}
|
||||
filtered_params: Final = {k: v for k, v in params.items() if k in valid_keys and v is not None}
|
||||
|
||||
return cast(VectorStoreCreateOptionalRequestParams, filtered_params)
|
||||
metadata: Final = _metadata_for_provider(params)
|
||||
if metadata is None:
|
||||
return cast(VectorStoreCreateOptionalRequestParams, filtered_params)
|
||||
return cast( # cast-ok: same params as above plus the metadata the provider should store
|
||||
VectorStoreCreateOptionalRequestParams, {**filtered_params, "metadata": metadata}
|
||||
)
|
||||
|
|
|
|||
|
|
@ -20077,7 +20077,13 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": false
|
||||
"supports_vision": false,
|
||||
"token_count_framing": 97
|
||||
},
|
||||
"databricks/system.ai.deepseek-v4-flash-0731": {
|
||||
"litellm_provider": "databricks",
|
||||
"mode": "chat",
|
||||
"token_count_framing": 97
|
||||
},
|
||||
"databricks/databricks-deepseek-v4-pro-0813": {
|
||||
"cache_creation_input_token_cost": 1.31999e-06,
|
||||
|
|
@ -21068,7 +21074,13 @@
|
|||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"token_count_framing": 100
|
||||
},
|
||||
"databricks/system.ai.kimi-k3": {
|
||||
"litellm_provider": "databricks",
|
||||
"mode": "chat",
|
||||
"token_count_framing": 100
|
||||
},
|
||||
"databricks/databricks-llama-4-maverick": {
|
||||
"cache_creation_input_token_cost": 5.0001e-07,
|
||||
|
|
@ -77877,7 +77889,9 @@
|
|||
"supports_response_schema": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"token_count_framing": 64,
|
||||
"token_count_tokenizer": "kimi_k3"
|
||||
},
|
||||
"global.moonshotai.kimi-k3": {
|
||||
"supports_regex_lookaround": false,
|
||||
|
|
@ -77907,7 +77921,9 @@
|
|||
"supports_response_schema": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"token_count_framing": 64,
|
||||
"token_count_tokenizer": "kimi_k3"
|
||||
},
|
||||
"us.moonshotai.kimi-k3": {
|
||||
"supports_regex_lookaround": false,
|
||||
|
|
@ -77937,7 +77953,9 @@
|
|||
"supports_response_schema": true,
|
||||
"supports_system_messages": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
"supports_vision": true,
|
||||
"token_count_framing": 64,
|
||||
"token_count_tokenizer": "kimi_k3"
|
||||
},
|
||||
"openrouter/xiaomi/mimo-v2.6-flash": {
|
||||
"cache_read_input_token_cost": 2.8e-09,
|
||||
|
|
|
|||
|
|
@ -1219,6 +1219,18 @@
|
|||
"additionalProperties": false
|
||||
}
|
||||
},
|
||||
"token_count_framing": {
|
||||
"type": "integer",
|
||||
"minimum": 0,
|
||||
"description": "Tokens added once per request to a local count estimate, for the provider's chat framing."
|
||||
},
|
||||
"token_count_tokenizer": {
|
||||
"type": "string",
|
||||
"description": "Tokenizer litellm counts this model's requests with when the provider cannot count them.",
|
||||
"enum": [
|
||||
"kimi_k3"
|
||||
]
|
||||
},
|
||||
"tpm": {
|
||||
"type": "integer",
|
||||
"minimum": 0,
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
[project]
|
||||
name = "litellm"
|
||||
version = "1.106.0"
|
||||
version = "1.106.1"
|
||||
description = "Library to easily interface with LLM API providers"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.10, <3.15"
|
||||
|
|
@ -373,7 +373,7 @@ litellm-enterprise = { workspace = true }
|
|||
members = ["enterprise", "litellm-proxy-extras"]
|
||||
|
||||
[tool.commitizen]
|
||||
version = "1.106.0"
|
||||
version = "1.106.1"
|
||||
version_files = [
|
||||
"pyproject.toml:^version",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -0,0 +1,215 @@
|
|||
import json
|
||||
from collections.abc import Iterator
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import anthropic
|
||||
import openai
|
||||
import pytest
|
||||
from integration._support.client import Gateway, gateway_from_environment
|
||||
from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process
|
||||
from integration._support.wire import Reply, Request, wire_server
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
|
||||
pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
|
||||
|
||||
_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
|
||||
_MESSAGES: Final[list[JsonValue]] = [{"role": "user", "content": "Count this synthetic request."}]
|
||||
_PROVIDER_KEY_VARIABLES: Final = (
|
||||
"ANTHROPIC_API_KEY",
|
||||
"ANTHROPIC_AUTH_TOKEN",
|
||||
"AZURE_AI_API_KEY",
|
||||
"AZURE_API_KEY",
|
||||
"AZURE_OPENAI_API_KEY",
|
||||
"AZURE_AD_TOKEN",
|
||||
"AZURE_CLIENT_ID",
|
||||
"AZURE_CLIENT_SECRET",
|
||||
"AZURE_TENANT_ID",
|
||||
"OPENAI_API_KEY",
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def provider_error_config(tmp_path_factory: pytest.TempPathFactory) -> Path:
|
||||
config: Final = tmp_path_factory.mktemp("count-tokens-provider-errors") / "config.yaml"
|
||||
base: Final = Path(__file__).resolve().parents[1] / "proxy_config.yaml"
|
||||
config.write_text(base.read_text() + "\nlitellm_settings:\n disable_token_counter: true\n")
|
||||
return config
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def keyless_proxy(tmp_path_factory: pytest.TempPathFactory, provider_error_config: Path) -> Iterator[OwnedProxy]:
|
||||
with (
|
||||
gateway_from_environment() as gateway,
|
||||
owned_proxy_process(
|
||||
gateway,
|
||||
tmp_path_factory.mktemp("count-tokens-error-shapes"),
|
||||
{},
|
||||
config=provider_error_config,
|
||||
remove_environment=_PROVIDER_KEY_VARIABLES,
|
||||
workers=1,
|
||||
) as owned,
|
||||
):
|
||||
yield owned
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def azure_environment_key_proxy(tmp_path_factory: pytest.TempPathFactory) -> Iterator[OwnedProxy]:
|
||||
with (
|
||||
gateway_from_environment() as gateway,
|
||||
owned_proxy_process(
|
||||
gateway,
|
||||
tmp_path_factory.mktemp("count-tokens-azure-environment-key"),
|
||||
{"AZURE_API_KEY": "synthetic-azure-env-key"},
|
||||
remove_environment=tuple(name for name in _PROVIDER_KEY_VARIABLES if name != "AZURE_API_KEY"),
|
||||
workers=1,
|
||||
) as owned,
|
||||
):
|
||||
yield owned
|
||||
|
||||
|
||||
def _anthropic(gateway: Gateway) -> anthropic.Anthropic:
|
||||
return anthropic.Anthropic(
|
||||
api_key=gateway.key, base_url=str(gateway.client.base_url), max_retries=0, http_client=gateway.client
|
||||
)
|
||||
|
||||
|
||||
def _never_called(request: Request) -> Reply:
|
||||
raise AssertionError(f"count reached the provider without a key: {request.target}")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("deployment", "missing"),
|
||||
(
|
||||
({"model": "anthropic/claude-sonnet-4-6"}, "No Anthropic API key is set for model "),
|
||||
({"model": "azure_ai/claude-sonnet-4-6"}, "No Azure AI API key or Entra ID credential is set for model "),
|
||||
),
|
||||
ids=("anthropic", "azure_ai_anthropic"),
|
||||
)
|
||||
def test_anthropic_count_without_a_provider_credential_is_an_anthropic_401(
|
||||
keyless_proxy: OwnedProxy, deployment: dict[str, JsonValue], missing: str
|
||||
) -> None:
|
||||
gateway: Final = keyless_proxy.gateway
|
||||
with wire_server(_never_called) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(**deployment, api_base=wire.url, api_key=None, num_retries=0)
|
||||
with pytest.raises(anthropic.AuthenticationError) as raised:
|
||||
_anthropic(gateway).messages.count_tokens(model=model, messages=_MESSAGES)
|
||||
assert raised.value.body == {
|
||||
"type": "error",
|
||||
"error": {"type": "authentication_error", "message": raised.value.body["error"]["message"]},
|
||||
}
|
||||
assert missing in raised.value.message, raised.value.message
|
||||
assert wire.drain() == ()
|
||||
|
||||
|
||||
def test_azure_ai_anthropic_count_sends_the_azure_environment_key_as_x_api_key(
|
||||
azure_environment_key_proxy: OwnedProxy,
|
||||
) -> None:
|
||||
gateway: Final = azure_environment_key_proxy.gateway
|
||||
|
||||
def count(request: Request) -> Reply:
|
||||
assert request.headers.get("x-api-key") == "synthetic-azure-env-key", request.headers
|
||||
return Reply(body=json.dumps({"input_tokens": 58}).encode())
|
||||
|
||||
with wire_server(count) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(
|
||||
model="azure_ai/claude-sonnet-4-6",
|
||||
api_base=wire.url,
|
||||
api_key=None,
|
||||
num_retries=0,
|
||||
)
|
||||
result: Final = _anthropic(gateway).messages.count_tokens(model=model, messages=_MESSAGES)
|
||||
assert result.input_tokens == 58
|
||||
|
||||
|
||||
def test_azure_ai_anthropic_count_authenticates_with_an_entra_token(keyless_proxy: OwnedProxy) -> None:
|
||||
gateway: Final = keyless_proxy.gateway
|
||||
|
||||
def count(request: Request) -> Reply:
|
||||
assert request.headers.get("authorization") == "Bearer synthetic-entra-token", request.headers
|
||||
assert "x-api-key" not in request.headers, request.headers
|
||||
return Reply(body=json.dumps({"input_tokens": 61}).encode())
|
||||
|
||||
with wire_server(count) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(
|
||||
model="azure_ai/claude-sonnet-4-6",
|
||||
api_base=wire.url,
|
||||
api_key=None,
|
||||
azure_ad_token="synthetic-entra-token",
|
||||
num_retries=0,
|
||||
)
|
||||
result: Final = _anthropic(gateway).messages.count_tokens(model=model, messages=_MESSAGES)
|
||||
assert result.input_tokens == 61
|
||||
assert len(wire.drain()) == 1
|
||||
|
||||
|
||||
def test_anthropic_count_request_errors_use_the_anthropic_envelope(keyless_proxy: OwnedProxy) -> None:
|
||||
gateway: Final = keyless_proxy.gateway
|
||||
upstream_error: Final = {"type": "error", "error": {"type": "invalid_request_error", "message": "scripted reject"}}
|
||||
with (
|
||||
wire_server(lambda _request: Reply(status=400, body=json.dumps(upstream_error).encode())) as wire,
|
||||
gateway.scenario() as scenario,
|
||||
):
|
||||
model: Final = scenario.model(
|
||||
model="anthropic/claude-sonnet-4-6", api_base=wire.url, api_key="synthetic-key", num_retries=0
|
||||
)
|
||||
missing_messages: Final = gateway.request("POST", "/v1/messages/count_tokens", {"model": model})
|
||||
assert missing_messages.status_code == 400, missing_messages.text
|
||||
assert _JSON_OBJECT.validate_json(missing_messages.content) == {
|
||||
"type": "error",
|
||||
"error": {"type": "invalid_request_error", "message": "messages parameter is required"},
|
||||
}
|
||||
|
||||
rejected: Final = gateway.request(
|
||||
"POST", "/v1/messages/count_tokens", {"model": model, "messages": _MESSAGES}
|
||||
)
|
||||
assert rejected.status_code == 400, rejected.text
|
||||
body: Final = _JSON_OBJECT.validate_json(rejected.content)
|
||||
assert set(body) <= {"type", "error", "request_id"}, body
|
||||
assert body["error"] == upstream_error["error"], body
|
||||
assert len(wire.drain()) == 1
|
||||
|
||||
|
||||
def test_openai_input_tokens_without_a_provider_key_is_an_openai_401(keyless_proxy: OwnedProxy) -> None:
|
||||
gateway: Final = keyless_proxy.gateway
|
||||
with wire_server(_never_called) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(
|
||||
model="openai/gpt-4o-mini", api_base=f"{wire.url}/v1", api_key=None, num_retries=0
|
||||
)
|
||||
client: Final = openai.OpenAI(
|
||||
api_key=gateway.key, base_url=str(gateway.client.base_url), max_retries=0, http_client=gateway.client
|
||||
)
|
||||
with pytest.raises(openai.AuthenticationError) as raised:
|
||||
client.responses.input_tokens.count(model=model, input="Count this synthetic request.")
|
||||
assert f"No OpenAI API key is set for model {model}" in raised.value.message
|
||||
assert wire.drain() == ()
|
||||
|
||||
|
||||
def test_google_count_tokens_errors_use_the_google_envelope(keyless_proxy: OwnedProxy) -> None:
|
||||
gateway: Final = keyless_proxy.gateway
|
||||
upstream_error: Final = {
|
||||
"error": {"code": 400, "message": "API key not valid. Please pass a valid API key.", "status": "INVALID_ARGUMENT"}
|
||||
}
|
||||
with (
|
||||
wire_server(lambda _request: Reply(status=400, body=json.dumps(upstream_error).encode())) as wire,
|
||||
gateway.scenario() as scenario,
|
||||
):
|
||||
model: Final = scenario.model(
|
||||
model="gemini/gemini-2.5-flash", api_base=wire.url, api_key="synthetic-key", num_retries=0
|
||||
)
|
||||
rejected: Final = gateway.request(
|
||||
"POST",
|
||||
f"/v1beta/models/{model}:countTokens",
|
||||
{"contents": [{"role": "user", "parts": [{"text": "Count this synthetic request."}]}]},
|
||||
)
|
||||
assert rejected.status_code == 400, rejected.text
|
||||
error: Final = _JSON_OBJECT.validate_json(rejected.content)["error"]
|
||||
assert isinstance(error, dict) and error["code"] == 400 and error["status"] == "INVALID_ARGUMENT", error
|
||||
assert "API key not valid" in str(error["message"]), error
|
||||
assert len(wire.drain()) == 1
|
||||
|
||||
malformed: Final = gateway.request("POST", f"/v1beta/models/{model}:countTokens", {"contents": "not-a-list"})
|
||||
assert malformed.status_code == 400, malformed.text
|
||||
assert _JSON_OBJECT.validate_json(malformed.content) == {
|
||||
"error": {"code": 400, "message": "Invalid Google countTokens request body", "status": "INVALID_ARGUMENT"}
|
||||
}
|
||||
|
|
@ -0,0 +1,393 @@
|
|||
import json
|
||||
from collections.abc import Iterator, Mapping
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from integration._support.client import Scenario, gateway_from_environment
|
||||
from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process
|
||||
from integration._support.wire import Reply, Request, Wire, wire_server
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
|
||||
pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
|
||||
|
||||
_COUNT: Final = 4242
|
||||
_TOOL_SCHEMA: Final[dict[str, JsonValue]] = {"type": "object", "properties": {}}
|
||||
_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def counting_proxy(tmp_path_factory: pytest.TempPathFactory) -> Iterator[OwnedProxy]:
|
||||
with (
|
||||
gateway_from_environment() as gateway,
|
||||
owned_proxy_process(gateway, tmp_path_factory.mktemp("openai-count-tools"), {}, workers=1) as owned,
|
||||
):
|
||||
yield owned
|
||||
|
||||
|
||||
def _deployment(scenario: Scenario, api_base: str) -> str:
|
||||
return scenario.model(
|
||||
model="openai/gpt-4o-mini",
|
||||
api_base=f"{api_base}/v1",
|
||||
api_key="sk-integration-openai",
|
||||
)
|
||||
|
||||
|
||||
def _count_response(_request: Request) -> Reply:
|
||||
return Reply(body=json.dumps({"input_tokens": _COUNT}).encode())
|
||||
|
||||
|
||||
_PDF_BASE64: Final = "JVBERi0xLjQK"
|
||||
_PNG_BASE64: Final = "aW1hZ2U="
|
||||
_LOOKUP_TOOL: Final[dict[str, JsonValue]] = {
|
||||
"name": "lookup",
|
||||
"description": "Find a record",
|
||||
"input_schema": {"type": "object", "properties": {"query": {"type": "string"}}},
|
||||
}
|
||||
_RESPONSES_REPLY: Final[Mapping[str, JsonValue]] = {
|
||||
"id": "resp_count_parity",
|
||||
"object": "response",
|
||||
"created_at": 1,
|
||||
"status": "completed",
|
||||
"model": "gpt-4o-mini",
|
||||
"output": [
|
||||
{
|
||||
"type": "message",
|
||||
"id": "msg_count_parity",
|
||||
"status": "completed",
|
||||
"role": "assistant",
|
||||
"content": [{"type": "output_text", "text": "OK", "annotations": []}],
|
||||
}
|
||||
],
|
||||
"usage": {"input_tokens": 10, "output_tokens": 1, "total_tokens": 11},
|
||||
}
|
||||
_WRAPPED_REASONING: Final = "litellm_enc:bW9kZWxfaWQ6ZGVwbG95bWVudC0x;gAAAAoriginal"
|
||||
_ANSWER_SCHEMA: Final[dict[str, JsonValue]] = {"type": "object", "properties": {"answer": {"type": "string"}}}
|
||||
_ANTHROPIC_BODIES: Final[Mapping[str, Mapping[str, JsonValue]]] = {
|
||||
"structured-output": {
|
||||
"output_config": {"format": {"type": "json_schema", "schema": _ANSWER_SCHEMA}},
|
||||
"messages": [{"role": "user", "content": "Answer in JSON."}],
|
||||
},
|
||||
"affinity-wrapped-reasoning": {
|
||||
"messages": [
|
||||
{"role": "user", "content": "hi"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{"type": "redacted_thinking", "data": f"litellm_encrypted_reasoning:{_WRAPPED_REASONING}"},
|
||||
{"type": "text", "text": "OK"},
|
||||
],
|
||||
},
|
||||
{"role": "user", "content": "again"},
|
||||
],
|
||||
},
|
||||
"base64-pdf": {
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "Summarize the attached contract."},
|
||||
{
|
||||
"type": "document",
|
||||
"source": {"type": "base64", "media_type": "application/pdf", "data": _PDF_BASE64},
|
||||
},
|
||||
],
|
||||
}
|
||||
],
|
||||
},
|
||||
"base64-image": {
|
||||
"system": [{"type": "text", "text": "Describe images briefly."}],
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": _PNG_BASE64}},
|
||||
{"type": "text", "text": "What is this?"},
|
||||
],
|
||||
}
|
||||
],
|
||||
},
|
||||
"replayed-web-search": {
|
||||
"messages": [
|
||||
{"role": "user", "content": "What is the newest LiteLLM release?"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{
|
||||
"type": "server_tool_use",
|
||||
"id": "srvtoolu_01",
|
||||
"name": "web_search",
|
||||
"input": {"query": "newest LiteLLM release"},
|
||||
},
|
||||
{
|
||||
"type": "web_search_tool_result",
|
||||
"tool_use_id": "srvtoolu_01",
|
||||
"content": [
|
||||
{
|
||||
"type": "web_search_result",
|
||||
"url": "https://docs.litellm.ai/release_notes",
|
||||
"title": "Release notes",
|
||||
"encrypted_content": "",
|
||||
"snippet": "LiteLLM ships a weekly release candidate and a stable line.",
|
||||
}
|
||||
],
|
||||
},
|
||||
{"type": "text", "text": "The newest release is in the release notes."},
|
||||
],
|
||||
},
|
||||
{"role": "user", "content": "Thanks."},
|
||||
],
|
||||
},
|
||||
"tool-turns": {
|
||||
"system": "Use the lookup tool.",
|
||||
"tools": [_LOOKUP_TOOL],
|
||||
"messages": [
|
||||
{"role": "user", "content": "Find record 14."},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{"type": "text", "text": "Looking it up."},
|
||||
{"type": "tool_use", "id": "toolu_01", "name": "lookup", "input": {"query": "14"}},
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [{"type": "tool_result", "tool_use_id": "toolu_01", "content": "14. Indemnification."}],
|
||||
},
|
||||
],
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _responses_peer(request: Request) -> Reply:
|
||||
if request.target == "/v1/responses/input_tokens":
|
||||
return _count_response(request)
|
||||
if request.target == "/v1/responses":
|
||||
return Reply(body=json.dumps(_RESPONSES_REPLY).encode())
|
||||
return Reply(status=404, body=json.dumps({"error": request.target}).encode())
|
||||
|
||||
|
||||
def _body_for(wire_requests: tuple[Request, ...], target: str) -> dict[str, JsonValue]:
|
||||
matching: Final = tuple(request for request in wire_requests if request.target == target)
|
||||
assert len(matching) == 1, [request.target for request in wire_requests]
|
||||
return _JSON_OBJECT.validate_json(matching[0].body)
|
||||
|
||||
|
||||
def _counted_fields(body: Mapping[str, JsonValue]) -> Mapping[str, JsonValue]:
|
||||
return {key: body[key] for key in ("input", "instructions", "tools", "text") if body.get(key) is not None}
|
||||
|
||||
|
||||
def _count_body(wire: Wire) -> dict[str, JsonValue]:
|
||||
requests: Final = wire.drain()
|
||||
assert len(requests) == 1, requests
|
||||
request: Final = requests[0]
|
||||
assert request.method == "POST"
|
||||
assert request.target == "/v1/responses/input_tokens", request.target
|
||||
return _JSON_OBJECT.validate_json(request.body)
|
||||
|
||||
|
||||
def test_openai_count_tokens_uses_deployment_api_base(counting_proxy: OwnedProxy) -> None:
|
||||
gateway: Final = counting_proxy.gateway
|
||||
with wire_server(_count_response) as wire, gateway.scenario() as scenario:
|
||||
model: Final = _deployment(scenario, wire.url)
|
||||
response: Final = gateway.request(
|
||||
"POST",
|
||||
"/v1/messages/count_tokens",
|
||||
{"model": model, "messages": [{"role": "user", "content": "hello"}]},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
assert _JSON_OBJECT.validate_json(response.content)["input_tokens"] == _COUNT
|
||||
assert "tools" not in _count_body(wire)
|
||||
|
||||
|
||||
def test_google_count_tokens_maps_anthropic_tools_to_openai(counting_proxy: OwnedProxy) -> None:
|
||||
gateway: Final = counting_proxy.gateway
|
||||
with wire_server(_count_response) as wire, gateway.scenario() as scenario:
|
||||
model: Final = _deployment(scenario, wire.url)
|
||||
response: Final = gateway.request(
|
||||
"POST",
|
||||
f"/v1beta/models/{model}:countTokens",
|
||||
{
|
||||
"contents": [{"role": "user", "parts": [{"text": "hello"}]}],
|
||||
"tools": [
|
||||
{
|
||||
"functionDeclarations": [
|
||||
{
|
||||
"name": "lookup",
|
||||
"description": "Find a record",
|
||||
"parameters": _TOOL_SCHEMA,
|
||||
}
|
||||
]
|
||||
}
|
||||
],
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
assert _JSON_OBJECT.validate_json(response.content)["totalTokens"] == _COUNT
|
||||
assert _count_body(wire)["tools"] == [
|
||||
{
|
||||
"type": "function",
|
||||
"name": "lookup",
|
||||
"description": "Find a record",
|
||||
"parameters": _TOOL_SCHEMA,
|
||||
"strict": False,
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
def test_messages_count_tokens_maps_anthropic_tools_to_openai(counting_proxy: OwnedProxy) -> None:
|
||||
gateway: Final = counting_proxy.gateway
|
||||
with wire_server(_count_response) as wire, gateway.scenario() as scenario:
|
||||
model: Final = _deployment(scenario, wire.url)
|
||||
response: Final = gateway.request(
|
||||
"POST",
|
||||
"/v1/messages/count_tokens",
|
||||
{
|
||||
"model": model,
|
||||
"messages": [{"role": "user", "content": "hello"}],
|
||||
"tools": [
|
||||
{
|
||||
"name": "lookup",
|
||||
"description": "Find a record",
|
||||
"input_schema": _TOOL_SCHEMA,
|
||||
}
|
||||
],
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
assert _JSON_OBJECT.validate_json(response.content)["input_tokens"] == _COUNT
|
||||
assert _count_body(wire)["tools"] == [
|
||||
{
|
||||
"type": "function",
|
||||
"name": "lookup",
|
||||
"description": "Find a record",
|
||||
"parameters": _TOOL_SCHEMA,
|
||||
"strict": False,
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
def test_utils_token_counter_keeps_chat_tool_conversion(counting_proxy: OwnedProxy) -> None:
|
||||
gateway: Final = counting_proxy.gateway
|
||||
with wire_server(_count_response) as wire, gateway.scenario() as scenario:
|
||||
model: Final = _deployment(scenario, wire.url)
|
||||
response: Final = gateway.request(
|
||||
"POST",
|
||||
"/utils/token_counter",
|
||||
{
|
||||
"model": model,
|
||||
"messages": [{"role": "user", "content": "hello"}],
|
||||
"tools": [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "lookup",
|
||||
"description": "Find a record",
|
||||
"parameters": _TOOL_SCHEMA,
|
||||
},
|
||||
}
|
||||
],
|
||||
},
|
||||
params={"call_endpoint": "true"},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
assert _JSON_OBJECT.validate_json(response.content)["total_tokens"] == _COUNT
|
||||
assert _count_body(wire)["tools"] == [
|
||||
{
|
||||
"type": "function",
|
||||
"name": "lookup",
|
||||
"description": "Find a record",
|
||||
"parameters": _TOOL_SCHEMA,
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("case", sorted(_ANTHROPIC_BODIES))
|
||||
def test_messages_count_tokens_sends_the_input_v1_messages_inference_sends(
|
||||
counting_proxy: OwnedProxy, case: str
|
||||
) -> None:
|
||||
gateway: Final = counting_proxy.gateway
|
||||
with wire_server(_responses_peer) as wire, gateway.scenario() as scenario:
|
||||
model: Final = _deployment(scenario, wire.url)
|
||||
body: Final = {"model": model, **_ANTHROPIC_BODIES[case]}
|
||||
generated: Final = gateway.request("POST", "/v1/messages", {**body, "max_tokens": 16})
|
||||
counted: Final = gateway.request("POST", "/v1/messages/count_tokens", body)
|
||||
assert generated.status_code == 200, generated.text
|
||||
assert counted.status_code == 200, counted.text
|
||||
assert _JSON_OBJECT.validate_json(counted.content)["input_tokens"] == _COUNT
|
||||
requests: Final = wire.drain()
|
||||
inference_body: Final = _body_for(requests, "/v1/responses")
|
||||
count_body: Final = _body_for(requests, "/v1/responses/input_tokens")
|
||||
assert _counted_fields(count_body) == _counted_fields(inference_body)
|
||||
assert count_body["model"] == inference_body["model"]
|
||||
|
||||
|
||||
_UNCONVERTIBLE_BODY: Final[Mapping[str, JsonValue]] = {
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "Summarize the notes."},
|
||||
{"type": "document", "source": {"type": "text", "media_type": "text/plain", "data": "notes"}},
|
||||
{"type": "search_result", "source": "https://example.com/notes", "title": "Notes", "content": []},
|
||||
],
|
||||
},
|
||||
{"role": "assistant", "content": [{"type": "tool_use", "id": "toolu_01", "name": "search", "input": {}}]},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "tool_result",
|
||||
"tool_use_id": "toolu_01",
|
||||
"content": [
|
||||
{
|
||||
"type": "search_result",
|
||||
"source": "https://example.com/hits",
|
||||
"title": "Hits",
|
||||
"content": [{"type": "text", "text": "hit"}],
|
||||
}
|
||||
],
|
||||
}
|
||||
],
|
||||
},
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
def test_messages_count_tokens_rejects_blocks_it_cannot_convert(counting_proxy: OwnedProxy) -> None:
|
||||
gateway: Final = counting_proxy.gateway
|
||||
with wire_server(_responses_peer) as wire, gateway.scenario() as scenario:
|
||||
model: Final = _deployment(scenario, wire.url)
|
||||
response: Final = gateway.request("POST", "/v1/messages/count_tokens", {"model": model, **_UNCONVERTIBLE_BODY})
|
||||
assert response.status_code == 400, response.text
|
||||
assert "document with a text source, search_result, search_result inside a tool_result" in response.text
|
||||
assert wire.drain() == ()
|
||||
|
||||
|
||||
def test_messages_count_tokens_without_a_role_answers_400_before_calling_openai(counting_proxy: OwnedProxy) -> None:
|
||||
gateway: Final = counting_proxy.gateway
|
||||
with wire_server(_count_response) as wire, gateway.scenario() as scenario:
|
||||
model: Final = _deployment(scenario, wire.url)
|
||||
response: Final = gateway.request(
|
||||
"POST",
|
||||
"/v1/messages/count_tokens",
|
||||
{"model": model, "messages": [{"role": "user", "content": "hello"}, {"content": "no role"}]},
|
||||
)
|
||||
assert response.status_code == 400, response.text
|
||||
assert "messages.1.role: Field required" in response.text
|
||||
assert wire.drain() == ()
|
||||
|
||||
|
||||
def test_google_count_tokens_sends_the_system_instruction_text(counting_proxy: OwnedProxy) -> None:
|
||||
gateway: Final = counting_proxy.gateway
|
||||
with wire_server(_count_response) as wire, gateway.scenario() as scenario:
|
||||
model: Final = _deployment(scenario, wire.url)
|
||||
response: Final = gateway.request(
|
||||
"POST",
|
||||
f"/v1beta/models/{model}:countTokens",
|
||||
{
|
||||
"systemInstruction": {"parts": [{"text": "Be brief."}, {"text": "Answer in English."}]},
|
||||
"contents": [{"role": "user", "parts": [{"text": "hello"}]}],
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
assert _count_body(wire)["instructions"] == "Be brief.\nAnswer in English."
|
||||
|
|
@ -3,8 +3,9 @@ import uuid
|
|||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from integration._support.client import Gateway
|
||||
from integration._support.client import Gateway, list_value, object_value
|
||||
from integration._support.wire import Reply, Request, wire_server
|
||||
from pydantic import JsonValue
|
||||
|
||||
|
||||
@pytest.mark.covers("other.provider_wire.responses_bridge.max_output_tokens_incomplete_maps_to_length")
|
||||
|
|
@ -572,3 +573,256 @@ def test_chat_over_responses_deployment_merges_function_call_followed_by_message
|
|||
}
|
||||
], response.text
|
||||
assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/responses")]
|
||||
|
||||
|
||||
_WEATHER_TOOL: Final[dict[str, JsonValue]] = {
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_weather",
|
||||
"description": "Get the weather for a city.",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {"city": {"type": "string"}},
|
||||
"required": ["city"],
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _message_item(item_id: str, text: str) -> dict[str, JsonValue]:
|
||||
return {
|
||||
"type": "message",
|
||||
"id": item_id,
|
||||
"status": "completed",
|
||||
"role": "assistant",
|
||||
"content": [{"type": "output_text", "text": text, "annotations": []}],
|
||||
}
|
||||
|
||||
|
||||
_WEATHER_CALL: Final[dict[str, JsonValue]] = {
|
||||
"type": "function_call",
|
||||
"id": "fc_1",
|
||||
"call_id": "call_1",
|
||||
"name": "get_weather",
|
||||
"arguments": '{"city":"Paris"}',
|
||||
"status": "completed",
|
||||
}
|
||||
|
||||
_TWO_MESSAGES_THEN_CALL: Final = (
|
||||
_message_item("msg_first", "First message."),
|
||||
_message_item("msg_second", "Second message."),
|
||||
_WEATHER_CALL,
|
||||
)
|
||||
|
||||
|
||||
def test_chat_over_responses_deployment_folds_two_messages_and_function_call_into_one_choice(
|
||||
gateway: Gateway,
|
||||
) -> None:
|
||||
identity: Final = "responses-bridge-two-messages-" + uuid.uuid4().hex
|
||||
|
||||
def respond(request: Request) -> Reply:
|
||||
if request.method == "GET" and request.target == "/v1/models":
|
||||
return Reply(body=b'{"object":"list","data":[]}')
|
||||
assert request.method == "POST" and request.target == "/responses", request.target
|
||||
return Reply(
|
||||
body=json.dumps(
|
||||
{
|
||||
"id": "resp_two_messages",
|
||||
"object": "response",
|
||||
"created_at": 1789788253,
|
||||
"status": "completed",
|
||||
"model": "gpt-6-sol",
|
||||
"output": list(_TWO_MESSAGES_THEN_CALL),
|
||||
"usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15},
|
||||
}
|
||||
).encode()
|
||||
)
|
||||
|
||||
with wire_server(respond) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(
|
||||
model="openai/responses/gpt-6-sol", api_base=wire.url, api_key="synthetic-openai-key"
|
||||
)
|
||||
response: Final = gateway.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{
|
||||
"model": model,
|
||||
"messages": [{"role": "user", "content": f"What is the weather in Paris? {identity}"}],
|
||||
"tools": [_WEATHER_TOOL],
|
||||
"cache": {"no-cache": True},
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
body: Final = response.json()
|
||||
assert body["choices"] == [
|
||||
{
|
||||
"finish_reason": "tool_calls",
|
||||
"index": 0,
|
||||
"message": {
|
||||
"role": "assistant",
|
||||
"content": "First message.Second message.",
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "fc_1",
|
||||
"type": "function",
|
||||
"function": {"name": "get_weather", "arguments": '{"city":"Paris"}'},
|
||||
"index": 0,
|
||||
}
|
||||
],
|
||||
},
|
||||
}
|
||||
], response.text
|
||||
assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/responses")]
|
||||
|
||||
|
||||
def _responses_frame(event: str, payload: dict[str, JsonValue]) -> bytes:
|
||||
return f"event: {event}\ndata: {json.dumps(payload, separators=(',', ':'))}\n\n".encode()
|
||||
|
||||
|
||||
def _message_item_frames(output_index: int, item: dict[str, JsonValue]) -> tuple[bytes, ...]:
|
||||
item_id: Final = item["id"]
|
||||
text: Final = object_value(list_value(item["content"])[0])["text"]
|
||||
return (
|
||||
_responses_frame(
|
||||
"response.output_item.added",
|
||||
{
|
||||
"type": "response.output_item.added",
|
||||
"output_index": output_index,
|
||||
"item": {**item, "status": "in_progress", "content": []},
|
||||
},
|
||||
),
|
||||
_responses_frame(
|
||||
"response.output_text.delta",
|
||||
{
|
||||
"type": "response.output_text.delta",
|
||||
"item_id": item_id,
|
||||
"output_index": output_index,
|
||||
"content_index": 0,
|
||||
"delta": text,
|
||||
},
|
||||
),
|
||||
_responses_frame(
|
||||
"response.output_item.done",
|
||||
{"type": "response.output_item.done", "output_index": output_index, "item": item},
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _function_call_frames(output_index: int, item: dict[str, JsonValue]) -> tuple[bytes, ...]:
|
||||
return (
|
||||
_responses_frame(
|
||||
"response.output_item.added",
|
||||
{
|
||||
"type": "response.output_item.added",
|
||||
"output_index": output_index,
|
||||
"item": {**item, "status": "in_progress", "arguments": ""},
|
||||
},
|
||||
),
|
||||
_responses_frame(
|
||||
"response.function_call_arguments.delta",
|
||||
{
|
||||
"type": "response.function_call_arguments.delta",
|
||||
"item_id": item["id"],
|
||||
"output_index": output_index,
|
||||
"delta": item["arguments"],
|
||||
},
|
||||
),
|
||||
_responses_frame(
|
||||
"response.function_call_arguments.done",
|
||||
{
|
||||
"type": "response.function_call_arguments.done",
|
||||
"item_id": item["id"],
|
||||
"output_index": output_index,
|
||||
"arguments": item["arguments"],
|
||||
},
|
||||
),
|
||||
_responses_frame(
|
||||
"response.output_item.done",
|
||||
{"type": "response.output_item.done", "output_index": output_index, "item": item},
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _sse_payloads(lines: list[str]) -> list[dict[str, JsonValue]]:
|
||||
return [json.loads(line[6:]) for line in lines if line.startswith("data: ") and line != "data: [DONE]"]
|
||||
|
||||
|
||||
def test_chat_stream_over_responses_deployment_folds_two_messages_and_function_call_into_one_choice(
|
||||
gateway: Gateway,
|
||||
) -> None:
|
||||
identity: Final = "responses-bridge-two-messages-stream-" + uuid.uuid4().hex
|
||||
response_payload: Final[dict[str, JsonValue]] = {
|
||||
"id": "resp_two_messages_stream",
|
||||
"object": "response",
|
||||
"created_at": 1789788253,
|
||||
"model": "gpt-6-sol",
|
||||
"status": "in_progress",
|
||||
"output": [],
|
||||
}
|
||||
|
||||
def respond(request: Request) -> Reply:
|
||||
if request.method == "GET" and request.target == "/v1/models":
|
||||
return Reply(body=b'{"object":"list","data":[]}')
|
||||
assert request.method == "POST" and request.target == "/responses", request.target
|
||||
assert json.loads(request.body)["stream"] is True, request.body
|
||||
return Reply(
|
||||
content_type="text/event-stream",
|
||||
chunks=(
|
||||
_responses_frame("response.created", {"type": "response.created", "response": response_payload}),
|
||||
*_message_item_frames(0, _TWO_MESSAGES_THEN_CALL[0]),
|
||||
*_message_item_frames(1, _TWO_MESSAGES_THEN_CALL[1]),
|
||||
*_function_call_frames(2, _TWO_MESSAGES_THEN_CALL[2]),
|
||||
_responses_frame(
|
||||
"response.completed",
|
||||
{
|
||||
"type": "response.completed",
|
||||
"response": {
|
||||
**response_payload,
|
||||
"status": "completed",
|
||||
"output": list(_TWO_MESSAGES_THEN_CALL),
|
||||
"usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15},
|
||||
},
|
||||
},
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
with wire_server(respond) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(
|
||||
model="openai/responses/gpt-6-sol", api_base=wire.url, api_key="synthetic-openai-key"
|
||||
)
|
||||
response: Final = gateway.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{
|
||||
"model": model,
|
||||
"messages": [{"role": "user", "content": f"What is the weather in Paris? {identity}"}],
|
||||
"tools": [_WEATHER_TOOL],
|
||||
"stream": True,
|
||||
"cache": {"no-cache": True},
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
chunks: Final = _sse_payloads(list(response.iter_lines()))
|
||||
choices: Final = [object_value(choice) for chunk in chunks for choice in list_value(chunk["choices"])]
|
||||
assert [choice["index"] for choice in choices] == [0] * len(choices), chunks
|
||||
deltas: Final = [object_value(choice["delta"]) for choice in choices]
|
||||
assert (
|
||||
"".join(str(delta["content"]) for delta in deltas if delta.get("content"))
|
||||
== "First message.Second message."
|
||||
)
|
||||
tool_calls: Final = [
|
||||
object_value(call)
|
||||
for delta in deltas
|
||||
if delta.get("tool_calls")
|
||||
for call in list_value(delta["tool_calls"])
|
||||
]
|
||||
assert [
|
||||
object_value(call["function"])["name"] for call in tool_calls if object_value(call["function"]).get("name")
|
||||
] == ["get_weather"], chunks
|
||||
assert (
|
||||
"".join(str(object_value(call["function"]).get("arguments") or "") for call in tool_calls)
|
||||
== '{"city":"Paris"}'
|
||||
), chunks
|
||||
assert [choice["finish_reason"] for choice in choices if choice.get("finish_reason")] == ["tool_calls"], chunks
|
||||
assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/responses")]
|
||||
|
|
|
|||
|
|
@ -0,0 +1,630 @@
|
|||
import asyncio
|
||||
import json
|
||||
import threading
|
||||
from collections.abc import Mapping, Sequence
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from contextlib import ExitStack
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
from uuid import uuid4
|
||||
|
||||
import anthropic
|
||||
import httpx
|
||||
import pytest
|
||||
from integration._support.client import Gateway, Scenario
|
||||
from integration._support.process import graceful_stop_seconds
|
||||
from integration._support.vertex import service_account_json
|
||||
from integration._support.wire import Reply, Request, Wire, wire_server
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
|
||||
pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
|
||||
|
||||
_MODEL: Final = "gemini-4-argon"
|
||||
_PROJECT: Final = "scripted-project"
|
||||
_LOCATION: Final = "us-east5"
|
||||
_COUNT: Final = 777
|
||||
_UNIT: Final = " red blue green yellow orange purple black white silver gold"
|
||||
_USER: Final = f"<filler>{_UNIT * 100}</filler>\nReturn exactly OK."
|
||||
_MESSAGES: Final[JsonValue] = [{"role": "user", "content": _USER}]
|
||||
_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
|
||||
_JSON_LIST: Final = TypeAdapter(list[JsonValue])
|
||||
_REPLY: Final[Mapping[str, JsonValue]] = {
|
||||
"candidates": [
|
||||
{
|
||||
"content": {"role": "model", "parts": [{"text": "OK"}]},
|
||||
"finishReason": "STOP",
|
||||
}
|
||||
],
|
||||
"usageMetadata": {"promptTokenCount": 10, "candidatesTokenCount": 1, "totalTokenCount": 11},
|
||||
}
|
||||
|
||||
|
||||
class _GeminiPeer:
|
||||
def __init__(self) -> None:
|
||||
self.status: int = 200
|
||||
self.requests: tuple[Request, ...] = ()
|
||||
self._lock: Final = threading.Lock()
|
||||
|
||||
def set_status(self, status: int) -> None:
|
||||
self.status = status
|
||||
|
||||
def __call__(self, request: Request) -> Reply:
|
||||
with self._lock:
|
||||
self.requests = (*self.requests, request)
|
||||
status: Final = self.status
|
||||
if request.target.endswith(":countTokens"):
|
||||
return Reply(status=status, body=json.dumps({"totalTokens": _COUNT}).encode())
|
||||
if request.target.endswith(":generateContent"):
|
||||
return Reply(body=json.dumps(_REPLY).encode())
|
||||
return Reply(status=404, body=json.dumps({"error": request.target}).encode())
|
||||
|
||||
def drain(self) -> tuple[Request, ...]:
|
||||
with self._lock:
|
||||
requests: Final = self.requests
|
||||
self.requests = ()
|
||||
return requests
|
||||
|
||||
|
||||
def _deployment(gateway: Gateway, scenario: Scenario, api_base: str) -> str:
|
||||
return scenario.model(
|
||||
model=f"vertex_ai/{_MODEL}",
|
||||
api_base=api_base,
|
||||
api_key=None,
|
||||
vertex_project=_PROJECT,
|
||||
vertex_location=_LOCATION,
|
||||
vertex_credentials=service_account_json(_PROJECT, gateway.upstream_url.rstrip("/")),
|
||||
)
|
||||
|
||||
|
||||
def _messages(model: str, system: JsonValue | None, tools: JsonValue | None) -> Mapping[str, JsonValue]:
|
||||
return {
|
||||
"model": model,
|
||||
"messages": _MESSAGES,
|
||||
**({"system": system} if system is not None else {}),
|
||||
**({"tools": tools} if tools is not None else {}),
|
||||
}
|
||||
|
||||
|
||||
def _json_body(request: Request) -> Mapping[str, JsonValue]:
|
||||
return _JSON_OBJECT.validate_json(request.body)
|
||||
|
||||
|
||||
def _count_requests(wire: Wire) -> tuple[Request, ...]:
|
||||
return tuple(request for request in wire.drain() if request.target.endswith(":countTokens"))
|
||||
|
||||
|
||||
def _count_tokens_tools(generated_tools: JsonValue) -> JsonValue:
|
||||
return _JSON_LIST.validate_python(
|
||||
[
|
||||
{
|
||||
{
|
||||
"function_declarations": "functionDeclarations",
|
||||
"url_context": "urlContext",
|
||||
"code_execution": "codeExecution",
|
||||
}.get(key, key): value
|
||||
for key, value in tool.items()
|
||||
}
|
||||
for tool in TypeAdapter(list[Mapping[str, JsonValue]]).validate_python(generated_tools)
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
def _count_tokens_request_from_generate_content(
|
||||
generated_body: Mapping[str, JsonValue],
|
||||
) -> Mapping[str, JsonValue]:
|
||||
system_instruction: Final = generated_body.get("systemInstruction", generated_body.get("system_instruction"))
|
||||
generated_tools: Final = generated_body.get("tools")
|
||||
return MappingProxyType(
|
||||
{
|
||||
"contents": generated_body["contents"],
|
||||
**({"systemInstruction": system_instruction} if system_instruction is not None else {}),
|
||||
**({"tools": _count_tokens_tools(generated_tools)} if generated_tools is not None else {}),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _tools(description: str) -> JsonValue:
|
||||
return _JSON_LIST.validate_python(
|
||||
[
|
||||
{
|
||||
"name": "noop",
|
||||
"description": description,
|
||||
"input_schema": {"type": "object", "properties": {}},
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("large_system", "large_tool"),
|
||||
((False, False), (True, False), (False, True), (True, True)),
|
||||
ids=("baseline", "large-system", "large-tool", "large-system-and-tool"),
|
||||
)
|
||||
def test_messages_and_count_tokens_send_matching_gemini_context(
|
||||
gateway: Gateway, large_system: bool, large_tool: bool
|
||||
) -> None:
|
||||
peer: Final = _GeminiPeer()
|
||||
with wire_server(peer) as wire, gateway.scenario() as scenario:
|
||||
model: Final = _deployment(gateway, scenario, wire.url)
|
||||
system: Final = (
|
||||
f"<cache-buster>{uuid4().hex}</cache-buster>\n"
|
||||
"This is a synthetic context-budget diagnostic. Return exactly OK." + (_UNIT * 100 if large_system else "")
|
||||
if large_system
|
||||
else None
|
||||
)
|
||||
tools: Final = _tools(_UNIT * 100) if large_tool else None
|
||||
body: Final = _messages(model, system, tools)
|
||||
generated: Final = gateway.request("POST", "/v1/messages", {**body, "max_tokens": 16})
|
||||
counted: Final = gateway.request("POST", "/v1/messages/count_tokens", body)
|
||||
assert generated.status_code == 200, generated.text
|
||||
assert counted.status_code == 200, counted.text
|
||||
assert _JSON_OBJECT.validate_json(counted.content)["input_tokens"] == _COUNT, counted.text
|
||||
requests: Final = peer.drain()
|
||||
generated_body: Final = _json_body(
|
||||
next(request for request in requests if request.target.endswith(":generateContent"))
|
||||
)
|
||||
count_body: Final = _json_body(next(request for request in requests if request.target.endswith(":countTokens")))
|
||||
assert count_body == _count_tokens_request_from_generate_content(generated_body)
|
||||
|
||||
|
||||
_LOOKUP_TOOL: Final[JsonValue] = {
|
||||
"name": "lookup",
|
||||
"description": "Look up a value",
|
||||
"input_schema": {"type": "object", "properties": {"query": {"type": "string"}}},
|
||||
}
|
||||
_CONTENT_BLOCK_BODIES: Final[Mapping[str, Mapping[str, JsonValue]]] = MappingProxyType(
|
||||
{
|
||||
"base64-pdf": {
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "Summarize the document."},
|
||||
{
|
||||
"type": "document",
|
||||
"source": {"type": "base64", "media_type": "application/pdf", "data": "JVBERi0xLjQK"},
|
||||
},
|
||||
],
|
||||
}
|
||||
],
|
||||
},
|
||||
"base64-image": {
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "aW1hZ2U="}},
|
||||
{"type": "text", "text": "What is this?"},
|
||||
],
|
||||
}
|
||||
],
|
||||
},
|
||||
"tool-turns": {
|
||||
"tools": [_LOOKUP_TOOL],
|
||||
"messages": [
|
||||
{"role": "user", "content": "Look up the value."},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{"type": "text", "text": "Looking it up."},
|
||||
{"type": "tool_use", "id": "toolu_01", "name": "lookup", "input": {"query": "value"}},
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [{"type": "tool_result", "tool_use_id": "toolu_01", "content": "42"}],
|
||||
},
|
||||
],
|
||||
},
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _ai_studio_deployment(scenario: Scenario, api_base: str, **parameters: JsonValue) -> str:
|
||||
return scenario.model(model=f"gemini/{_MODEL}", api_base=api_base, api_key="scripted-gemini-key", **parameters)
|
||||
|
||||
|
||||
def _counted_contents(count_body: Mapping[str, JsonValue]) -> JsonValue:
|
||||
if "contents" in count_body:
|
||||
return count_body["contents"]
|
||||
return _JSON_OBJECT.validate_python(count_body["generateContentRequest"])["contents"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("provider", ("vertex_ai", "gemini"))
|
||||
@pytest.mark.parametrize("case", sorted(_CONTENT_BLOCK_BODIES))
|
||||
def test_count_tokens_sends_the_contents_messages_inference_sends(gateway: Gateway, provider: str, case: str) -> None:
|
||||
peer: Final = _GeminiPeer()
|
||||
with wire_server(peer) as wire, gateway.scenario() as scenario:
|
||||
model: Final = (
|
||||
_deployment(gateway, scenario, wire.url)
|
||||
if provider == "vertex_ai"
|
||||
else _ai_studio_deployment(scenario, wire.url)
|
||||
)
|
||||
body: Final = {"model": model, **_CONTENT_BLOCK_BODIES[case]}
|
||||
generated: Final = gateway.request("POST", "/v1/messages", {**body, "max_tokens": 16})
|
||||
counted: Final = gateway.request("POST", "/v1/messages/count_tokens", body)
|
||||
assert generated.status_code == 200, generated.text
|
||||
assert counted.status_code == 200, counted.text
|
||||
assert _JSON_OBJECT.validate_json(counted.content)["input_tokens"] == _COUNT, counted.text
|
||||
requests: Final = peer.drain()
|
||||
generated_body: Final = _json_body(
|
||||
next(request for request in requests if request.target.endswith(":generateContent"))
|
||||
)
|
||||
count_body: Final = _json_body(next(request for request in requests if request.target.endswith(":countTokens")))
|
||||
assert _counted_contents(count_body) == generated_body["contents"]
|
||||
|
||||
|
||||
_UNCONVERTIBLE_MESSAGES: Final[JsonValue] = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "Summarize the notes."},
|
||||
{"type": "document", "source": {"type": "text", "media_type": "text/plain", "data": "notes"}},
|
||||
{"type": "search_result", "source": "https://example.com/notes", "title": "Notes", "content": []},
|
||||
],
|
||||
},
|
||||
{"role": "assistant", "content": [{"type": "tool_use", "id": "toolu_01", "name": "search", "input": {}}]},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "tool_result",
|
||||
"tool_use_id": "toolu_01",
|
||||
"content": [
|
||||
{
|
||||
"type": "search_result",
|
||||
"source": "https://example.com/hits",
|
||||
"title": "Hits",
|
||||
"content": [{"type": "text", "text": "hit"}],
|
||||
}
|
||||
],
|
||||
}
|
||||
],
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("provider", ("vertex_ai", "gemini"))
|
||||
def test_count_tokens_rejects_blocks_it_cannot_convert(gateway: Gateway, provider: str) -> None:
|
||||
peer: Final = _GeminiPeer()
|
||||
with wire_server(peer) as wire, gateway.scenario() as scenario:
|
||||
model: Final = (
|
||||
_deployment(gateway, scenario, wire.url)
|
||||
if provider == "vertex_ai"
|
||||
else _ai_studio_deployment(scenario, wire.url)
|
||||
)
|
||||
response: Final = gateway.request(
|
||||
"POST", "/v1/messages/count_tokens", {"model": model, "messages": _UNCONVERTIBLE_MESSAGES}
|
||||
)
|
||||
assert response.status_code == 400, response.text
|
||||
assert "document with a text source, search_result, search_result inside a tool_result" in response.text
|
||||
assert peer.drain() == ()
|
||||
|
||||
|
||||
def test_count_tokens_keeps_each_text_block_in_system_instruction(gateway: Gateway) -> None:
|
||||
peer: Final = _GeminiPeer()
|
||||
with wire_server(peer) as wire, gateway.scenario() as scenario:
|
||||
model: Final = _deployment(gateway, scenario, wire.url)
|
||||
system: Final[JsonValue] = [
|
||||
{"type": "text", "text": "First instruction"},
|
||||
{"type": "text", "text": "Second instruction"},
|
||||
]
|
||||
response: Final = gateway.request("POST", "/v1/messages/count_tokens", _messages(model, system, None))
|
||||
assert response.status_code == 200, response.text
|
||||
requests: Final = _count_requests(wire)
|
||||
assert len(requests) == 1
|
||||
assert _json_body(requests[0]) == {
|
||||
"contents": [{"role": "user", "parts": [{"text": _USER}]}],
|
||||
"systemInstruction": {"parts": [{"text": "First instruction"}, {"text": "Second instruction"}]},
|
||||
}
|
||||
|
||||
|
||||
def test_count_tokens_without_system_or_tools_sends_contents_only(gateway: Gateway) -> None:
|
||||
peer: Final = _GeminiPeer()
|
||||
with wire_server(peer) as wire, gateway.scenario() as scenario:
|
||||
model: Final = _deployment(gateway, scenario, wire.url)
|
||||
response: Final = gateway.request("POST", "/v1/messages/count_tokens", _messages(model, None, None))
|
||||
assert response.status_code == 200, response.text
|
||||
requests: Final = _count_requests(wire)
|
||||
assert len(requests) == 1
|
||||
assert _json_body(requests[0]) == {"contents": [{"role": "user", "parts": [{"text": _USER}]}]}
|
||||
|
||||
|
||||
_SDK_COUNT_ARGS: Final[Mapping[str, JsonValue]] = MappingProxyType(
|
||||
{
|
||||
"messages": [{"role": "user", "content": _USER}],
|
||||
"system": "Follow the system instruction.",
|
||||
"tools": [
|
||||
{
|
||||
"name": "noop",
|
||||
"description": "Return no output",
|
||||
"input_schema": {"type": "object", "properties": {}},
|
||||
}
|
||||
],
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _sync_sdk_count_tokens(gateway: Gateway, model: str) -> int:
|
||||
client: Final = anthropic.Anthropic(base_url=str(gateway.client.base_url), api_key=gateway.key, max_retries=0)
|
||||
return client.messages.count_tokens(model=model, **_SDK_COUNT_ARGS).input_tokens # pyright: ignore[reportArgumentType] # SDK TypedDicts reject a JsonValue mapping
|
||||
|
||||
|
||||
async def _async_sdk_count_tokens(gateway: Gateway, model: str) -> int:
|
||||
client: Final = anthropic.AsyncAnthropic(base_url=str(gateway.client.base_url), api_key=gateway.key, max_retries=0)
|
||||
return (await client.messages.count_tokens(model=model, **_SDK_COUNT_ARGS)).input_tokens # pyright: ignore[reportArgumentType] # SDK TypedDicts reject a JsonValue mapping
|
||||
|
||||
|
||||
@pytest.mark.parametrize("asynchronous", (False, True), ids=("sync-sdk", "async-sdk"))
|
||||
def test_anthropic_sdk_count_tokens_forwards_system_and_tools(gateway: Gateway, asynchronous: bool) -> None:
|
||||
peer: Final = _GeminiPeer()
|
||||
with wire_server(peer) as wire, gateway.scenario() as scenario:
|
||||
model: Final = _deployment(gateway, scenario, wire.url)
|
||||
input_tokens: Final = (
|
||||
asyncio.run(_async_sdk_count_tokens(gateway, model))
|
||||
if asynchronous
|
||||
else _sync_sdk_count_tokens(gateway, model)
|
||||
)
|
||||
assert input_tokens == _COUNT
|
||||
requests: Final = _count_requests(wire)
|
||||
assert len(requests) == 1
|
||||
request_body: Final = _json_body(requests[0])
|
||||
assert request_body["systemInstruction"] == {"parts": [{"text": "Follow the system instruction."}]}
|
||||
assert request_body["tools"] == [
|
||||
{
|
||||
"functionDeclarations": [
|
||||
{"name": "noop", "description": "Return no output", "parameters": {"type": "object"}}
|
||||
]
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
def test_google_count_tokens_forwards_system_and_tools_to_vertex(gateway: Gateway) -> None:
|
||||
peer: Final = _GeminiPeer()
|
||||
with wire_server(peer) as wire, gateway.scenario() as scenario:
|
||||
model: Final = _deployment(gateway, scenario, wire.url)
|
||||
body: Final[Mapping[str, JsonValue]] = {
|
||||
"contents": [
|
||||
{"role": "user", "parts": [{"text": _USER}]},
|
||||
{"role": "model", "parts": [{"functionCall": {"name": "noop", "args": {"query": "value"}}}]},
|
||||
{
|
||||
"role": "user",
|
||||
"parts": [{"functionResponse": {"name": "noop", "response": {"result": "ok"}}}],
|
||||
},
|
||||
],
|
||||
"systemInstruction": {"parts": [{"text": "Follow this instruction."}]},
|
||||
"tools": [
|
||||
{
|
||||
"functionDeclarations": [
|
||||
{
|
||||
"name": "noop",
|
||||
"description": "Return no output",
|
||||
"parameters": {"type": "OBJECT", "properties": {}},
|
||||
}
|
||||
]
|
||||
}
|
||||
],
|
||||
}
|
||||
response: Final = gateway.request("POST", f"/v1beta/models/{model}:countTokens", body)
|
||||
assert response.status_code == 200, response.text
|
||||
assert _JSON_OBJECT.validate_json(response.content)["totalTokens"] == _COUNT, response.text
|
||||
requests: Final = _count_requests(wire)
|
||||
assert len(requests) == 1
|
||||
request_body: Final = _json_body(requests[0])
|
||||
assert request_body["contents"] == body["contents"]
|
||||
assert request_body["systemInstruction"] == body["systemInstruction"]
|
||||
assert request_body["tools"] == [
|
||||
{
|
||||
"functionDeclarations": [
|
||||
{
|
||||
"name": "noop",
|
||||
"description": "Return no output",
|
||||
"parameters": {"type": "object"},
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
def test_responses_input_tokens_maps_flat_function_tools_to_vertex(gateway: Gateway) -> None:
|
||||
peer: Final = _GeminiPeer()
|
||||
with wire_server(peer) as wire, gateway.scenario() as scenario:
|
||||
model: Final = _deployment(gateway, scenario, wire.url)
|
||||
body: Final[Mapping[str, JsonValue]] = {
|
||||
"model": model,
|
||||
"instructions": "Follow the system instruction.",
|
||||
"input": [
|
||||
{
|
||||
"type": "function_call",
|
||||
"call_id": "call_lookup",
|
||||
"name": "lookup",
|
||||
"arguments": '{"query":"value"}',
|
||||
},
|
||||
{
|
||||
"type": "function_call_output",
|
||||
"call_id": "call_lookup",
|
||||
"output": "result",
|
||||
},
|
||||
],
|
||||
"tools": [
|
||||
{
|
||||
"type": "function",
|
||||
"name": "lookup",
|
||||
"description": "Look up a value",
|
||||
"parameters": {"type": "object", "properties": {}},
|
||||
}
|
||||
],
|
||||
}
|
||||
response: Final = gateway.request("POST", "/v1/responses/input_tokens", body)
|
||||
assert response.status_code == 200, response.text
|
||||
assert _JSON_OBJECT.validate_json(response.content)["input_tokens"] == _COUNT, response.text
|
||||
requests: Final = _count_requests(wire)
|
||||
assert len(requests) == 1
|
||||
request_body: Final = _json_body(requests[0])
|
||||
assert request_body["tools"] == [
|
||||
{
|
||||
"functionDeclarations": [
|
||||
{
|
||||
"name": "lookup",
|
||||
"description": "Look up a value",
|
||||
"parameters": {"type": "object"},
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
def test_utils_token_counter_maps_openai_function_tools_to_vertex(gateway: Gateway) -> None:
|
||||
peer: Final = _GeminiPeer()
|
||||
with wire_server(peer) as wire, gateway.scenario() as scenario:
|
||||
model: Final = _deployment(gateway, scenario, wire.url)
|
||||
tool_schema: Final[JsonValue] = {"type": "object", "properties": {}}
|
||||
body: Final[Mapping[str, JsonValue]] = {
|
||||
"model": model,
|
||||
"messages": _MESSAGES,
|
||||
"system": "Follow the system instruction.",
|
||||
"tools": [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "noop",
|
||||
"description": "Return no output",
|
||||
"parameters": tool_schema,
|
||||
},
|
||||
}
|
||||
],
|
||||
}
|
||||
response: Final = gateway.request("POST", "/utils/token_counter", body, params={"call_endpoint": "true"})
|
||||
assert response.status_code == 200, response.text
|
||||
assert _JSON_OBJECT.validate_json(response.content)["total_tokens"] == _COUNT, response.text
|
||||
requests: Final = _count_requests(wire)
|
||||
assert len(requests) == 1
|
||||
request_body: Final = _json_body(requests[0])
|
||||
assert request_body["systemInstruction"] == {"parts": [{"text": "Follow the system instruction."}]}
|
||||
assert request_body["tools"] == [
|
||||
{
|
||||
"functionDeclarations": [
|
||||
{
|
||||
"name": "noop",
|
||||
"description": "Return no output",
|
||||
"parameters": {"type": "object"},
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
def test_count_tokens_surfaces_peer_errors_and_rejects_non_text_system(gateway: Gateway) -> None:
|
||||
peer: Final = _GeminiPeer()
|
||||
peer.set_status(400)
|
||||
with wire_server(peer) as wire, gateway.scenario() as scenario:
|
||||
model: Final = _deployment(gateway, scenario, wire.url)
|
||||
response: Final = gateway.request("POST", "/v1/messages/count_tokens", _messages(model, None, None))
|
||||
assert response.status_code == 400, response.text
|
||||
assert len(_count_requests(wire)) == 1
|
||||
|
||||
invalid_peer: Final = _GeminiPeer()
|
||||
with wire_server(invalid_peer) as invalid_wire, gateway.scenario() as scenario:
|
||||
invalid_model: Final = _deployment(gateway, scenario, invalid_wire.url)
|
||||
invalid_response: Final = gateway.request(
|
||||
"POST",
|
||||
"/v1/messages/count_tokens",
|
||||
_messages(
|
||||
invalid_model,
|
||||
[{"type": "image", "source": {"type": "base64", "data": "AAAA"}}],
|
||||
None,
|
||||
),
|
||||
)
|
||||
assert invalid_response.status_code == 400, invalid_response.text
|
||||
assert _count_requests(invalid_wire) == ()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"tools",
|
||||
(
|
||||
[{"name": "noop", "input_schema": {"type": "object", "properties": {"query": "not-a-schema"}}}],
|
||||
[{"name": "noop", "input_schema": "not-a-schema"}],
|
||||
),
|
||||
ids=("malformed-property", "string-schema"),
|
||||
)
|
||||
def test_count_tokens_rejects_unconvertible_tools_without_request(gateway: Gateway, tools: JsonValue) -> None:
|
||||
peer: Final = _GeminiPeer()
|
||||
with wire_server(peer) as wire, gateway.scenario() as scenario:
|
||||
model: Final = _deployment(gateway, scenario, wire.url)
|
||||
response: Final = gateway.request("POST", "/v1/messages/count_tokens", _messages(model, None, tools))
|
||||
assert response.status_code == 400, response.text
|
||||
assert "tools" in response.text
|
||||
assert _count_requests(wire) == ()
|
||||
|
||||
|
||||
def _count_wave(clients: Sequence[httpx.Client], body: Mapping[str, JsonValue]) -> tuple[tuple[int, JsonValue], ...]:
|
||||
def count(client: httpx.Client) -> tuple[int, JsonValue]:
|
||||
response: Final = client.post("/v1/messages/count_tokens", json=dict(body))
|
||||
payload: Final = _JSON_OBJECT.validate_json(response.content)
|
||||
return response.status_code, payload.get("input_tokens")
|
||||
|
||||
with ThreadPoolExecutor(max_workers=len(clients)) as executor:
|
||||
return tuple(executor.map(count, clients))
|
||||
|
||||
|
||||
def test_concurrent_count_waves_recover_after_vertex_count_outage(gateway: Gateway) -> None:
|
||||
peer: Final = _GeminiPeer()
|
||||
with wire_server(peer) as wire, gateway.scenario() as scenario, ExitStack() as stack:
|
||||
model: Final = _deployment(gateway, scenario, wire.url)
|
||||
clients: Final = tuple(
|
||||
stack.enter_context(
|
||||
httpx.Client(
|
||||
base_url=str(gateway.client.base_url),
|
||||
headers={"Authorization": f"Bearer {gateway.key}"},
|
||||
timeout=30,
|
||||
trust_env=False,
|
||||
)
|
||||
)
|
||||
for _ in range(3)
|
||||
)
|
||||
body: Final = _messages(model, None, None)
|
||||
assert _count_wave(clients, body) == ((200, _COUNT), (200, _COUNT), (200, _COUNT))
|
||||
peer.set_status(503)
|
||||
assert _count_wave(clients, body) == ((503, None), (503, None), (503, None))
|
||||
peer.set_status(200)
|
||||
assert _count_wave(clients, body) == ((200, _COUNT), (200, _COUNT), (200, _COUNT))
|
||||
requests: Final = _count_requests(wire)
|
||||
assert len(requests) == 9
|
||||
assert (
|
||||
tuple(_json_body(request) for request in requests)
|
||||
== ({"contents": [{"role": "user", "parts": [{"text": _USER}]}]},) * 9
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("conversation", "call_endpoint"),
|
||||
(
|
||||
({"messages": _MESSAGES}, True),
|
||||
({"contents": [{"role": "user", "parts": [{"text": _USER}]}]}, True),
|
||||
({"messages": _MESSAGES}, False),
|
||||
),
|
||||
ids=("provider-count-messages", "provider-count-contents", "local-count-messages"),
|
||||
)
|
||||
@pytest.mark.parametrize("drop_params", (False, True), ids=("no-drop", "drop"))
|
||||
def test_utils_token_counter_counts_the_conversation_regardless_of_prompt_and_drop_params(
|
||||
gateway: Gateway, conversation: Mapping[str, JsonValue], call_endpoint: bool, drop_params: bool
|
||||
) -> None:
|
||||
peer: Final = _GeminiPeer()
|
||||
params: Final = {"call_endpoint": "true" if call_endpoint else "false"}
|
||||
with wire_server(peer) as wire, gateway.scenario() as scenario:
|
||||
model: Final = _ai_studio_deployment(scenario, wire.url, drop_params=drop_params)
|
||||
mixed: Final = gateway.request(
|
||||
"POST", "/utils/token_counter", {"model": model, "prompt": "hi", **conversation}, params=params
|
||||
)
|
||||
mixed_requests: Final = _count_requests(wire)
|
||||
alone: Final = gateway.request("POST", "/utils/token_counter", {"model": model, **conversation}, params=params)
|
||||
alone_requests: Final = _count_requests(wire)
|
||||
assert mixed.status_code == 200, mixed.text
|
||||
assert alone.status_code == 200, alone.text
|
||||
mixed_tokens: Final = TypeAdapter(int).validate_python(_JSON_OBJECT.validate_json(mixed.content)["total_tokens"])
|
||||
alone_tokens: Final = TypeAdapter(int).validate_python(_JSON_OBJECT.validate_json(alone.content)["total_tokens"])
|
||||
assert mixed_tokens == alone_tokens > 0
|
||||
assert len(mixed_requests) == len(alone_requests) == int(call_endpoint)
|
||||
assert [_counted_contents(_json_body(request)) for request in mixed_requests] == [
|
||||
_counted_contents(_json_body(request)) for request in alone_requests
|
||||
]
|
||||
|
|
@ -24,6 +24,8 @@ model_list:
|
|||
model: anthropic/claude-opus-5-5
|
||||
api_base: http://127.0.0.1:8191
|
||||
api_key: synthetic-anthropic-key
|
||||
model_info:
|
||||
token_count_framing: 0
|
||||
- model_name: anthropic/claude-sonnet-4-6
|
||||
litellm_params:
|
||||
model: anthropic/claude-sonnet-4-6
|
||||
|
|
|
|||
|
|
@ -399,9 +399,11 @@ def test_in_flight_count_tokens_does_not_reserve_key_budget_away_from_a_completi
|
|||
[{"role": "user", "content": prompt}]
|
||||
]
|
||||
assert count.status_code == 200, count.text
|
||||
assert count.json() == {"totalTokens": 12, "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 12}]}, (
|
||||
count.text
|
||||
)
|
||||
assert count.json() == {
|
||||
"totalTokens": 12,
|
||||
"promptTokensDetails": [{"modality": "TEXT", "tokenCount": 12}],
|
||||
"litellm_count_source": "provider",
|
||||
}, count.text
|
||||
provider_calls: Final = wire.drain()
|
||||
assert [(call.method, call.target) for call in provider_calls] == [
|
||||
("POST", "/v1beta/models/gemini-3.8-flash:countTokens")
|
||||
|
|
|
|||
|
|
@ -5019,4 +5019,62 @@ def test_transform_request_drop_params_in_litellm_params_gates_the_prompt_cache_
|
|||
litellm_logging_obj=Mock(),
|
||||
)
|
||||
|
||||
assert "prompt_cache_breakpoint" not in result["input"][0]["content"][0]
|
||||
assert "prompt_cache_breakpoint" not in result["input"][0]["content"][0]
|
||||
|
||||
|
||||
def test_convert_response_output_folds_two_messages_and_function_call_into_one_choice() -> None:
|
||||
first: Final = ResponseOutputMessage.model_validate(
|
||||
{
|
||||
"id": "msg_first",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"status": "completed",
|
||||
"content": [{"type": "output_text", "text": "First message.", "annotations": []}],
|
||||
}
|
||||
)
|
||||
second_annotation: Final = {
|
||||
"type": "url_citation",
|
||||
"start_index": 0,
|
||||
"end_index": 6,
|
||||
"title": "Second",
|
||||
"url": "https://example.com/second",
|
||||
}
|
||||
second: Final = ResponseOutputMessage.model_validate(
|
||||
{
|
||||
"id": "msg_second",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"status": "completed",
|
||||
"content": [{"type": "output_text", "text": "Second message.", "annotations": [second_annotation]}],
|
||||
}
|
||||
)
|
||||
function_call: Final = ResponseFunctionToolCall(
|
||||
id="fc_1",
|
||||
type="function_call",
|
||||
status="completed",
|
||||
arguments='{"city":"Paris"}',
|
||||
call_id="call_1",
|
||||
name="get_weather",
|
||||
)
|
||||
|
||||
choices: Final = LiteLLMResponsesTransformationHandler._convert_response_output_to_choices( # pyright: ignore[reportPrivateUsage] # exercises the converter directly
|
||||
(first, second, function_call)
|
||||
)
|
||||
|
||||
assert len(choices) == 1
|
||||
assert choices[0].index == 0
|
||||
assert choices[0].finish_reason == "tool_calls"
|
||||
assert choices[0].message.content == "First message.Second message."
|
||||
assert choices[0].message.annotations == [
|
||||
{**second_annotation, "start_index": len("First message."), "end_index": len("First message.") + 6}
|
||||
]
|
||||
tool_calls: Final = choices[0].message.tool_calls
|
||||
assert tool_calls is not None
|
||||
assert [call.model_dump(exclude_none=True) for call in tool_calls] == [
|
||||
{
|
||||
"id": "fc_1",
|
||||
"type": "function",
|
||||
"function": {"name": "get_weather", "arguments": '{"city":"Paris"}'},
|
||||
"index": 0,
|
||||
}
|
||||
]
|
||||
|
|
|
|||
0
tests/unit/db_scripts/__init__.py
Normal file
0
tests/unit/db_scripts/__init__.py
Normal file
|
|
@ -0,0 +1,198 @@
|
|||
import importlib.util
|
||||
import json
|
||||
import sys
|
||||
from collections.abc import Mapping
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
_SCRIPT: Final = Path(__file__).resolve().parents[3] / "db_scripts" / "strip_vector_store_internal_metadata.py"
|
||||
_spec: Final = importlib.util.spec_from_file_location("strip_vector_store_internal_metadata", _SCRIPT)
|
||||
strip_script: Final = importlib.util.module_from_spec(_spec)
|
||||
sys.modules[_spec.name] = strip_script
|
||||
_spec.loader.exec_module(strip_script)
|
||||
|
||||
POLLUTED: Final = {
|
||||
"id": "vs_1",
|
||||
"metadata": {
|
||||
"user_api_key": "hash-value-1",
|
||||
"user_api_key_hash": "hash-value-1",
|
||||
"user_api_key_user_email": "member@example.com",
|
||||
"endpoint": "http://proxy/v1/vector_stores",
|
||||
"deployment": "openai/gpt-5",
|
||||
"team": "llmproxy",
|
||||
"trace_id": "trace-123",
|
||||
},
|
||||
}
|
||||
CLEAN: Final = {"id": "vs_2", "metadata": {"team": "llmproxy"}}
|
||||
NO_METADATA: Final = {"id": "vs_3", "metadata": None}
|
||||
|
||||
|
||||
def _fake_openai(
|
||||
stores: tuple[Mapping[str, object], ...], page_size: int = 100, post_status: int = 200
|
||||
) -> tuple[httpx.MockTransport, list[httpx.Request]]:
|
||||
seen: Final[list[httpx.Request]] = []
|
||||
by_id: Final = {store["id"]: store for store in stores}
|
||||
|
||||
def handle(request: httpx.Request) -> httpx.Response:
|
||||
seen.append(request)
|
||||
path: Final = request.url.path.removeprefix("/v1")
|
||||
if request.method == "POST":
|
||||
if post_status != 200:
|
||||
return httpx.Response(post_status, json={"error": {"message": "upstream rejected the update"}})
|
||||
return httpx.Response(200, json=by_id[path.removeprefix("/vector_stores/")])
|
||||
if path == "/vector_stores":
|
||||
after: Final = request.url.params.get("after")
|
||||
start: Final = 0 if after is None else [store["id"] for store in stores].index(after) + 1
|
||||
page: Final = stores[start : start + page_size]
|
||||
return httpx.Response(200, json={"data": list(page), "has_more": start + page_size < len(stores)})
|
||||
return httpx.Response(200, json=by_id[path.removeprefix("/vector_stores/")])
|
||||
|
||||
return httpx.MockTransport(handle), seen
|
||||
|
||||
|
||||
def _run(args: list[str], transport: httpx.MockTransport) -> int:
|
||||
return strip_script.main(["--api-key", "sk-test", *args], transport=transport)
|
||||
|
||||
|
||||
def _posts(seen: list[httpx.Request]) -> list[tuple[str, dict[str, object]]]:
|
||||
return [(request.url.path, json.loads(request.content)) for request in seen if request.method == "POST"]
|
||||
|
||||
|
||||
def test_dry_run_lists_matching_key_names_without_values_and_sends_no_update(capsys):
|
||||
transport, seen = _fake_openai((POLLUTED, CLEAN))
|
||||
|
||||
exit_code: Final = _run([], transport)
|
||||
|
||||
out: Final = capsys.readouterr().out
|
||||
assert exit_code == 0
|
||||
assert (
|
||||
"vs_1: would delete 5 keys: deployment, endpoint, user_api_key, user_api_key_hash, user_api_key_user_email"
|
||||
in out
|
||||
)
|
||||
assert "scanned=2 matched=1 updated=0 failed=0" in out
|
||||
assert "hash-value-1" not in out
|
||||
assert "member@example.com" not in out
|
||||
assert _posts(seen) == []
|
||||
assert {request.headers["authorization"] for request in seen} == {"Bearer sk-test"}
|
||||
|
||||
|
||||
def test_apply_nulls_exactly_the_proxy_written_keys_and_leaves_caller_keys_alone(capsys):
|
||||
transport, seen = _fake_openai((POLLUTED, CLEAN, NO_METADATA))
|
||||
|
||||
exit_code: Final = _run(["--apply"], transport)
|
||||
|
||||
assert exit_code == 0
|
||||
assert _posts(seen) == [
|
||||
(
|
||||
"/v1/vector_stores/vs_1",
|
||||
{
|
||||
"metadata": {
|
||||
"deployment": None,
|
||||
"endpoint": None,
|
||||
"user_api_key": None,
|
||||
"user_api_key_hash": None,
|
||||
"user_api_key_user_email": None,
|
||||
}
|
||||
},
|
||||
)
|
||||
]
|
||||
assert "scanned=3 matched=1 updated=1 failed=0" in capsys.readouterr().out
|
||||
|
||||
|
||||
def test_extra_key_is_deleted_alongside_the_built_in_keys():
|
||||
transport, seen = _fake_openai((POLLUTED,))
|
||||
|
||||
_run(["--apply", "--extra-key", "trace_id"], transport)
|
||||
|
||||
assert _posts(seen)[0][1]["metadata"]["trace_id"] is None
|
||||
assert "team" not in _posts(seen)[0][1]["metadata"]
|
||||
|
||||
|
||||
def test_listing_follows_has_more_with_the_last_store_id(capsys):
|
||||
transport, seen = _fake_openai((CLEAN, POLLUTED, NO_METADATA), page_size=2)
|
||||
|
||||
_run([], transport)
|
||||
|
||||
assert [request.url.params.get("after") for request in seen] == [None, "vs_1"]
|
||||
assert "scanned=3 matched=1" in capsys.readouterr().out
|
||||
|
||||
|
||||
def test_listing_walks_more_pages_than_the_interpreter_recursion_limit(capsys):
|
||||
stores: Final = tuple({"id": f"vs_{index}", "metadata": None} for index in range(sys.getrecursionlimit() + 50))
|
||||
transport, _ = _fake_openai(stores, page_size=1)
|
||||
|
||||
assert _run([], transport) == 0
|
||||
assert f"scanned={len(stores)} matched=0" in capsys.readouterr().out
|
||||
|
||||
|
||||
def test_non_string_metadata_values_are_scanned_without_printing_them(capsys):
|
||||
transport, _ = _fake_openai(({"id": "vs_9", "metadata": {"user_api_key": "secret-value", "count": 7}},))
|
||||
|
||||
exit_code: Final = _run([], transport)
|
||||
|
||||
out: Final = capsys.readouterr().out
|
||||
assert exit_code == 0
|
||||
assert "vs_9: would delete 1 keys: user_api_key" in out
|
||||
assert "secret-value" not in out
|
||||
|
||||
|
||||
def test_an_unexpected_response_shape_exits_non_zero_without_echoing_metadata(capsys):
|
||||
def malformed(request: httpx.Request) -> httpx.Response:
|
||||
return httpx.Response(200, json={"data": [{"metadata": {"user_api_key": "secret-value"}}], "has_more": False})
|
||||
|
||||
exit_code: Final = _run([], httpx.MockTransport(malformed))
|
||||
|
||||
captured: Final = capsys.readouterr()
|
||||
assert exit_code == 2
|
||||
assert "unexpected vector store response shape" in captured.err
|
||||
assert "secret-value" not in captured.out + captured.err
|
||||
|
||||
|
||||
def test_vector_store_id_retrieves_only_the_named_stores_without_listing(capsys):
|
||||
transport, seen = _fake_openai((POLLUTED, CLEAN, NO_METADATA))
|
||||
|
||||
_run(["--vector-store-id", "vs_2", "--vector-store-id", "vs_1"], transport)
|
||||
|
||||
assert [request.url.path for request in seen] == ["/v1/vector_stores/vs_2", "/v1/vector_stores/vs_1"]
|
||||
assert "scanned=2 matched=1" in capsys.readouterr().out
|
||||
|
||||
|
||||
def test_a_failed_update_is_reported_and_exits_non_zero(capsys):
|
||||
transport, _ = _fake_openai((POLLUTED,), post_status=400)
|
||||
|
||||
exit_code: Final = _run(["--apply"], transport)
|
||||
|
||||
out: Final = capsys.readouterr().out
|
||||
assert exit_code == 1
|
||||
assert "vs_1: failed to delete 5 keys (HTTP 400: upstream rejected the update)" in out
|
||||
assert "updated=0 failed=1" in out
|
||||
|
||||
|
||||
def test_a_listing_error_exits_non_zero_without_updating_anything():
|
||||
seen: Final[list[httpx.Request]] = []
|
||||
|
||||
def reject(request: httpx.Request) -> httpx.Response:
|
||||
seen.append(request)
|
||||
return httpx.Response(401, json={"error": {"message": "bad key"}})
|
||||
|
||||
assert _run(["--apply"], httpx.MockTransport(reject)) == 2
|
||||
assert [request.method for request in seen] == ["GET"]
|
||||
|
||||
|
||||
def test_an_unreachable_api_exits_non_zero():
|
||||
def unreachable(request: httpx.Request) -> httpx.Response:
|
||||
raise httpx.ConnectError("connection refused", request=request)
|
||||
|
||||
assert _run([], httpx.MockTransport(unreachable)) == 2
|
||||
|
||||
|
||||
def test_an_api_key_is_required(monkeypatch):
|
||||
monkeypatch.delenv("OPENAI_API_KEY", raising=False)
|
||||
|
||||
with pytest.raises(SystemExit) as exit_info:
|
||||
strip_script.main([])
|
||||
|
||||
assert exit_info.value.code == 2
|
||||
|
|
@ -32,7 +32,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import (
|
|||
sanitize_messages_for_tool_calling,
|
||||
THOUGHT_SIGNATURE_SEPARATOR,
|
||||
)
|
||||
from litellm.types.llms.openai import ChatCompletionToolMessage
|
||||
from litellm.types.llms.openai import ChatCompletionAssistantMessage, ChatCompletionToolMessage
|
||||
from litellm.utils import (
|
||||
_invalidate_model_cost_lowercase_map,
|
||||
function_setup,
|
||||
|
|
@ -4130,44 +4130,42 @@ def test_anthropic_messages_pt_drops_a_system_message_with_no_text():
|
|||
assert [m["role"] for m in result] == ["user", "assistant"]
|
||||
|
||||
|
||||
def test_anthropic_messages_pt_drops_empty_but_signed_thinking_block():
|
||||
"""
|
||||
Anthropic rejects a `thinking` block whose `thinking` text is empty, even
|
||||
when it carries a valid-looking signature, with:
|
||||
400 messages.N.content.M.thinking: each thinking block must contain thinking
|
||||
This shape is reachable via cross-provider replay of a `thinking_blocks`
|
||||
history item (see PR #36033), so `is_unsignable_thinking_block()` must
|
||||
also check the thinking text, not just the signature.
|
||||
"""
|
||||
@pytest.mark.parametrize(
|
||||
"assistant_message",
|
||||
(
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "4",
|
||||
"thinking_blocks": [{"type": "thinking", "thinking": "", "signature": "sig_abc123_looks_valid"}],
|
||||
},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{"type": "thinking", "thinking": "", "signature": "sig_abc123_looks_valid"},
|
||||
{"type": "text", "text": "4"},
|
||||
],
|
||||
},
|
||||
),
|
||||
ids=("thinking_blocks", "inline_content"),
|
||||
)
|
||||
def test_anthropic_messages_pt_keeps_empty_but_signed_thinking_block(
|
||||
assistant_message: ChatCompletionAssistantMessage,
|
||||
) -> None:
|
||||
"""``display: omitted`` returns signed thinking blocks with empty text, and they
|
||||
must reach the next turn unchanged."""
|
||||
from litellm.litellm_core_utils.prompt_templates.factory import (
|
||||
anthropic_messages_pt,
|
||||
)
|
||||
|
||||
messages = [
|
||||
{"role": "user", "content": "What's 2+2?"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "4",
|
||||
"thinking_blocks": [
|
||||
{
|
||||
"type": "thinking",
|
||||
"thinking": "",
|
||||
"signature": "sig_abc123_looks_valid",
|
||||
}
|
||||
],
|
||||
},
|
||||
]
|
||||
|
||||
result = anthropic_messages_pt(
|
||||
messages=messages,
|
||||
model="claude-sonnet-4-5-20250929",
|
||||
messages=[{"role": "user", "content": "What's 2+2?"}, assistant_message],
|
||||
model="claude-opus-5-5",
|
||||
llm_provider="anthropic",
|
||||
)
|
||||
|
||||
assistant_msg = result[1]
|
||||
assert isinstance(assistant_msg["content"], list)
|
||||
content_types = [block.get("type") for block in assistant_msg["content"]]
|
||||
assert "thinking" not in content_types, "empty-text thinking block must be dropped even though it has a signature"
|
||||
assert assistant_msg["content"][0] == {"type": "thinking", "thinking": "", "signature": "sig_abc123_looks_valid"}
|
||||
|
||||
|
||||
def test_anthropic_messages_pt_keeps_non_empty_signed_thinking_block():
|
||||
|
|
@ -4281,29 +4279,31 @@ def test_anthropic_messages_pt_drops_unsigned_thinking_block():
|
|||
assert "thinking" not in content_types, "unsigned thinking block must still be dropped"
|
||||
|
||||
|
||||
def test_is_unsignable_thinking_block_treats_whitespace_only_as_empty():
|
||||
def test_anthropic_messages_pt_drops_unsigned_empty_thinking_block():
|
||||
"""
|
||||
Edge case: a `thinking` field that is present but whitespace-only (e.g.
|
||||
a single trailing newline forwarded from another provider's empty
|
||||
reasoning summary) is functionally empty and Anthropic's API will still
|
||||
reject it with "each thinking block must contain thinking". We treat it
|
||||
the same as a fully empty string and drop the block.
|
||||
|
||||
The check lives in the shared `is_unsignable_thinking_block` helper, which
|
||||
`_drop_unsignable_thinking_blocks` calls standalone, so the whitespace-aware
|
||||
test has to hold there rather than only at the factory call site.
|
||||
An empty block without a signature is the shape Claude Code replays after a
|
||||
non-Anthropic turn, and Anthropic rejects it, so only the missing signature
|
||||
decides whether a block is dropped.
|
||||
"""
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
is_unsignable_thinking_block,
|
||||
from litellm.litellm_core_utils.prompt_templates.factory import (
|
||||
anthropic_messages_pt,
|
||||
)
|
||||
|
||||
whitespace_only_block = {
|
||||
"type": "thinking",
|
||||
"thinking": " \n\t ",
|
||||
"signature": "sig_abc123_looks_valid",
|
||||
}
|
||||
result = anthropic_messages_pt(
|
||||
messages=[
|
||||
{"role": "user", "content": "What's 2+2?"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [{"type": "thinking", "thinking": "", "signature": ""}, {"type": "text", "text": "4"}],
|
||||
},
|
||||
],
|
||||
model="claude-opus-5-5",
|
||||
llm_provider="anthropic",
|
||||
)
|
||||
|
||||
assert is_unsignable_thinking_block(whitespace_only_block) is True
|
||||
assistant_msg = result[1]
|
||||
assert isinstance(assistant_msg["content"], list)
|
||||
assert [block.get("type") for block in assistant_msg["content"]] == ["text"]
|
||||
|
||||
|
||||
_CONTENT_LESS_USER_MESSAGES: Final = ({"role": "user"}, {"role": "user", "content": None})
|
||||
|
|
|
|||
|
|
@ -337,6 +337,7 @@ def test_flagged_placement_keeps_a_system_whose_empty_assistant_follower_ends_th
|
|||
},
|
||||
{"role": "assistant", "content": None, "thinking_blocks": [{"type": "redacted_thinking", "data": "x"}]},
|
||||
{"role": "assistant", "content": [{"type": "thinking", "thinking": "hm", "signature": "s"}]},
|
||||
{"role": "assistant", "content": [{"type": "thinking", "thinking": "", "signature": "s"}]},
|
||||
{"role": "assistant", "content": None, "function_call": {"name": "f", "arguments": "{}"}},
|
||||
litellm.Message(
|
||||
role="assistant",
|
||||
|
|
@ -349,6 +350,7 @@ def test_flagged_placement_keeps_a_system_whose_empty_assistant_follower_ends_th
|
|||
"signed-thinking-block",
|
||||
"redacted-thinking-block",
|
||||
"signed-thinking-part",
|
||||
"signed-empty-thinking-part",
|
||||
"function-call",
|
||||
"pydantic-tool-calls",
|
||||
],
|
||||
|
|
|
|||
|
|
@ -0,0 +1,192 @@
|
|||
import asyncio
|
||||
import json
|
||||
from collections.abc import Mapping
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
|
||||
from litellm.llms.anthropic.pass_through.adapters.transformation import LiteLLMAnthropicMessagesAdapter
|
||||
from litellm.llms.anthropic.pass_through.convertible_blocks import (
|
||||
unconvertible_anthropic_blocks,
|
||||
unconvertible_blocks_error,
|
||||
)
|
||||
from litellm.llms.anthropic.pass_through.responses_adapters.transformation import (
|
||||
LiteLLMAnthropicToResponsesAPIAdapter,
|
||||
)
|
||||
from litellm.llms.vertex_ai.count_tokens.context import gemini_count_contents
|
||||
from litellm.types.llms.anthropic import AllAnthropicPassThroughMessageValues
|
||||
|
||||
_JSON_OBJECTS: Final = TypeAdapter(list[dict[str, JsonValue]])
|
||||
|
||||
_PDF: Final = {"type": "base64", "media_type": "application/pdf", "data": "JVBERi0xLjQK"}
|
||||
_PNG: Final = {"type": "base64", "media_type": "image/png", "data": "aW1hZ2U="}
|
||||
|
||||
_CONVERTED_BLOCKS: Final[tuple[tuple[str, Mapping[str, object]], ...]] = (
|
||||
("user", {"type": "text", "text": "hi"}),
|
||||
("user", {"type": "image", "source": _PNG}),
|
||||
("user", {"type": "image", "source": {"type": "url", "url": "https://example.com/a.png"}}),
|
||||
("user", {"type": "document", "source": _PDF}),
|
||||
("user", {"type": "document", "source": {"type": "url", "url": "https://example.com/a.pdf"}}),
|
||||
("user", {"type": "tool_result", "tool_use_id": "toolu_01", "content": "42"}),
|
||||
("assistant", {"type": "text", "text": "hi"}),
|
||||
("assistant", {"type": "tool_use", "id": "toolu_01", "name": "lookup", "input": {}}),
|
||||
("assistant", {"type": "thinking", "thinking": "Let me think.", "signature": "sig"}),
|
||||
)
|
||||
|
||||
_DROPPED_BLOCKS: Final[tuple[tuple[str, Mapping[str, object], str], ...]] = (
|
||||
("user", {"type": "search_result", "source": "https://example.com", "title": "t", "content": []}, "search_result"),
|
||||
("user", {"type": "document", "source": {"type": "text", "data": "notes"}}, "document with a text source"),
|
||||
("user", {"type": "image", "source": {"type": "file", "file_id": "file_01"}}, "image with a file source"),
|
||||
("user", {"type": "image", "source": {"type": "base64", "data": ""}}, "image with an empty base64 source"),
|
||||
("user", {"type": "document"}, "document without a source"),
|
||||
("user", {"text": "untyped"}, "a block without a type"),
|
||||
("user", {"type": "tool_use", "id": "toolu_01", "name": "lookup", "input": {}}, "tool_use"),
|
||||
(
|
||||
"assistant",
|
||||
{"type": "server_tool_use", "id": "srvtoolu_01", "name": "web_search", "input": {}},
|
||||
"server_tool_use",
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _adapter_messages(role: str, block: Mapping[str, object]) -> list[AllAnthropicPassThroughMessageValues]:
|
||||
return [{"role": role, "content": [dict(block)]}] # pyright: ignore[reportReturnType] # sends block shapes the TypedDicts leave out on purpose
|
||||
|
||||
|
||||
def _converted_by_both_adapters(role: str, block: Mapping[str, object]) -> tuple[bool, bool]:
|
||||
messages: Final = _adapter_messages(role, block)
|
||||
chat: Final = _JSON_OBJECTS.validate_python(
|
||||
LiteLLMAnthropicMessagesAdapter().translate_anthropic_messages_to_openai(messages=messages)
|
||||
)
|
||||
responses: Final = LiteLLMAnthropicToResponsesAPIAdapter().translate_messages_to_responses_input(messages)
|
||||
chat_kept: Final = any(
|
||||
message.get("content") or message.get("tool_calls") or message.get("thinking_blocks") for message in chat
|
||||
)
|
||||
return chat_kept, bool(responses)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("role", "block"), _CONVERTED_BLOCKS)
|
||||
def test_blocks_both_conversions_keep_are_accepted(role: str, block: Mapping[str, object]) -> None:
|
||||
assert _converted_by_both_adapters(role, block) == (True, True)
|
||||
assert unconvertible_anthropic_blocks([{"role": role, "content": [block]}]) == ()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("role", "block", "label"), _DROPPED_BLOCKS)
|
||||
def test_blocks_both_conversions_drop_are_named(role: str, block: Mapping[str, object], label: str) -> None:
|
||||
assert _converted_by_both_adapters(role, block) == (False, False)
|
||||
assert unconvertible_anthropic_blocks([{"role": role, "content": [block]}]) == (label,)
|
||||
|
||||
|
||||
def test_labels_are_listed_once_in_order_and_string_content_is_accepted() -> None:
|
||||
messages: Final = [
|
||||
{"role": "user", "content": "plain text"},
|
||||
{"role": "user", "content": [{"type": "search_result"}, {"type": "container_upload"}]},
|
||||
{"role": "assistant", "content": [{"type": "search_result"}]},
|
||||
]
|
||||
assert unconvertible_anthropic_blocks(messages) == ("container_upload", "search_result")
|
||||
|
||||
|
||||
def test_redacted_thinking_is_accepted_like_other_replayed_thinking() -> None:
|
||||
messages: Final = [{"role": "assistant", "content": [{"type": "redacted_thinking", "data": "opaque"}]}]
|
||||
assert unconvertible_anthropic_blocks(messages) == ()
|
||||
|
||||
|
||||
def test_messages_that_are_not_json_are_left_to_the_count() -> None:
|
||||
messages: Final = [{"role": "user", "content": [{"type": "search_result", "source": object()}]}]
|
||||
assert unconvertible_anthropic_blocks(messages) == ()
|
||||
|
||||
|
||||
def _tool_turn(inner: Mapping[str, object]) -> tuple[Mapping[str, object], ...]:
|
||||
return (
|
||||
{"role": "assistant", "content": [{"type": "tool_use", "id": "toolu_01", "name": "lookup", "input": {}}]},
|
||||
{"role": "user", "content": [{"type": "tool_result", "tool_use_id": "toolu_01", "content": [dict(inner)]}]},
|
||||
)
|
||||
|
||||
|
||||
def _sent_by_each_counter(inner: Mapping[str, object]) -> tuple[str, str]:
|
||||
messages: Final = _tool_turn(inner)
|
||||
gemini: Final = asyncio.run(gemini_count_contents(messages, "anthropic", "gemini-2.5-flash", "gemini", {}))
|
||||
openai: Final = LiteLLMAnthropicToResponsesAPIAdapter().translate_messages_to_responses_input(
|
||||
[dict(message) for message in messages] # pyright: ignore[reportArgumentType] # sends block shapes the TypedDicts leave out on purpose
|
||||
)
|
||||
return json.dumps(gemini, default=str), json.dumps(openai, default=str)
|
||||
|
||||
|
||||
_KEPT_INSIDE_TOOL_RESULT: Final[tuple[tuple[Mapping[str, object], str], ...]] = (
|
||||
({"type": "text", "text": "forty-two"}, "forty-two"),
|
||||
({"type": "image", "source": _PNG}, _PNG["data"]),
|
||||
({"type": "document", "source": _PDF}, _PDF["data"]),
|
||||
)
|
||||
|
||||
_DROPPED_INSIDE_TOOL_RESULT: Final[tuple[tuple[Mapping[str, object], str, str], ...]] = (
|
||||
(
|
||||
{
|
||||
"type": "search_result",
|
||||
"source": "https://example.com",
|
||||
"title": "t",
|
||||
"content": [{"type": "text", "text": "search-snippet"}],
|
||||
},
|
||||
"search-snippet",
|
||||
"search_result inside a tool_result",
|
||||
),
|
||||
(
|
||||
{"type": "document", "source": {"type": "text", "data": "text-notes"}},
|
||||
"text-notes",
|
||||
"document with a text source inside a tool_result",
|
||||
),
|
||||
({"type": "tool_reference", "tool_name": "find_docs"}, "find_docs", "tool_reference inside a tool_result"),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("inner", "marker"), _KEPT_INSIDE_TOOL_RESULT)
|
||||
def test_tool_result_content_both_counters_send_is_accepted(inner: Mapping[str, object], marker: str) -> None:
|
||||
gemini, openai = _sent_by_each_counter(inner)
|
||||
|
||||
assert marker in gemini
|
||||
assert marker in openai
|
||||
assert unconvertible_anthropic_blocks(_tool_turn(inner)) == ()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("inner", "marker", "label"), _DROPPED_INSIDE_TOOL_RESULT)
|
||||
def test_tool_result_content_both_counters_drop_is_named(inner: Mapping[str, object], marker: str, label: str) -> None:
|
||||
gemini, openai = _sent_by_each_counter(inner)
|
||||
|
||||
assert marker not in gemini
|
||||
assert marker not in openai
|
||||
assert unconvertible_anthropic_blocks(_tool_turn(inner)) == (label,)
|
||||
|
||||
|
||||
def _replayed_web_search(encrypted_content: str) -> tuple[Mapping[str, object], ...]:
|
||||
return (
|
||||
{"role": "user", "content": "What is new?"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{"type": "server_tool_use", "id": "srvtoolu_01", "name": "web_search", "input": {"query": "news"}},
|
||||
{
|
||||
"type": "web_search_tool_result",
|
||||
"tool_use_id": "srvtoolu_01",
|
||||
"content": [
|
||||
{
|
||||
"type": "web_search_result",
|
||||
"url": "https://docs.litellm.ai/release_notes",
|
||||
"title": "Release notes",
|
||||
"encrypted_content": encrypted_content,
|
||||
"snippet": "LiteLLM ships a weekly release candidate.",
|
||||
}
|
||||
],
|
||||
},
|
||||
],
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def test_a_replayed_web_search_inference_sends_as_text_is_accepted() -> None:
|
||||
assert unconvertible_blocks_error(_replayed_web_search(""), "anthropic") is None
|
||||
|
||||
|
||||
def test_an_encrypted_web_search_inference_cannot_send_is_named() -> None:
|
||||
error: Final = unconvertible_blocks_error(_replayed_web_search("opaque"), "anthropic")
|
||||
assert error is not None
|
||||
assert error.endswith("server_tool_use, web_search_tool_result")
|
||||
|
|
@ -14,8 +14,9 @@ import json
|
|||
import os
|
||||
import sys
|
||||
import threading
|
||||
from collections.abc import Mapping
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
from types import MappingProxyType, SimpleNamespace
|
||||
from typing import Final
|
||||
from unittest.mock import patch
|
||||
|
||||
|
|
@ -24,6 +25,9 @@ import pytest
|
|||
|
||||
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../..")))
|
||||
|
||||
from litellm.llms.base_llm.anthropic_messages.transformation import ( # noqa: E402 # sys.path must be patched before importing litellm
|
||||
BaseAnthropicMessagesConfig,
|
||||
)
|
||||
from litellm.proxy._types import SpecialHeaders # noqa: E402 # sys.path must be patched before importing litellm
|
||||
|
||||
# Fake tokens for testing (not real secrets)
|
||||
|
|
@ -1332,6 +1336,12 @@ class TestPassthroughAuthToken:
|
|||
assert url == "https://custom.example.com/v1/messages"
|
||||
|
||||
|
||||
_BOUND_THINKING_ERROR: Final = (
|
||||
"messages.1.content.0: Invalid `signature` in `thinking` block. The block is bound to a different conversation."
|
||||
)
|
||||
_NULL_SIGNATURE_ERROR: Final = "messages.1.content.0.thinking.signature.str: Input should be a valid string"
|
||||
|
||||
|
||||
class TestAnthropicThinkingSignatureSelfHeal:
|
||||
"""Helpers for retrying after invalid thinking blocks in replayed history:
|
||||
invalid encrypted signatures, and blocks with empty thinking text."""
|
||||
|
|
@ -1422,11 +1432,11 @@ class TestAnthropicThinkingSignatureSelfHeal:
|
|||
assert is_empty_unsigned_thinking_block({"type": "redacted_thinking", "data": "opaque"}) is False
|
||||
assert is_empty_unsigned_thinking_block("not a dict") is False
|
||||
|
||||
def test_strip_empty_content_blocks_drops_empty_thinking_blocks(self):
|
||||
def test_strip_empty_content_blocks_drops_empty_unsigned_thinking_blocks(self):
|
||||
"""LIT-6357 ingestion half: an assistant tool-loop turn carrying an
|
||||
empty (even signed) thinking block keeps its tool_use blocks and loses
|
||||
the poison; whitespace-only counts as empty; a non-empty thinking block
|
||||
and redacted_thinking are untouched."""
|
||||
empty, unsigned thinking block keeps its tool_use blocks and loses the
|
||||
poison; whitespace-only counts as empty; a non-empty thinking block and
|
||||
redacted_thinking are untouched."""
|
||||
from litellm.llms.anthropic.common_utils import (
|
||||
strip_empty_content_blocks_from_anthropic_messages,
|
||||
)
|
||||
|
|
@ -1436,7 +1446,7 @@ class TestAnthropicThinkingSignatureSelfHeal:
|
|||
{"role": "user", "content": "weather?"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [{"type": "thinking", "thinking": "", "signature": "sig_abc"}, tu],
|
||||
"content": [{"type": "thinking", "thinking": "", "signature": ""}, tu],
|
||||
},
|
||||
{
|
||||
"role": "assistant",
|
||||
|
|
@ -1455,6 +1465,36 @@ class TestAnthropicThinkingSignatureSelfHeal:
|
|||
assert out[2]["content"][0]["thinking"] == "real plan"
|
||||
assert len(msgs[1]["content"]) == 2
|
||||
|
||||
def test_strip_empty_content_blocks_keeps_signed_empty_thinking_blocks(self):
|
||||
"""A thinking block with empty text and a signature is the default
|
||||
response shape of Claude Opus 4.8 / Fable 5 / Fable 5.1 (thinking.display
|
||||
omitted): the signature IS the preserved reasoning, Anthropic accepts it
|
||||
on input and binds it to its prefix. Dropping it silently discards the
|
||||
model's reasoning on every replayed turn and leaves preserved-thinking
|
||||
enforcement (thinking.block_binding) nothing to check. The message and
|
||||
block order must survive untouched."""
|
||||
from litellm.llms.anthropic.common_utils import (
|
||||
strip_empty_content_blocks_from_anthropic_messages,
|
||||
)
|
||||
|
||||
tu = {"type": "tool_use", "id": "toolu_01A", "name": "get_weather", "input": {"city": "Paris"}}
|
||||
signed = {"type": "thinking", "thinking": "", "signature": "EqQBCkYIBxgCKkD"}
|
||||
msgs = [
|
||||
{"role": "user", "content": "weather?"},
|
||||
{"role": "assistant", "content": [signed, tu]},
|
||||
{"role": "user", "content": [{"type": "tool_result", "tool_use_id": "toolu_01A", "content": "sunny"}]},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [signed, {"type": "text", "text": ""}, {"type": "text", "text": "Sunny."}],
|
||||
},
|
||||
]
|
||||
out = strip_empty_content_blocks_from_anthropic_messages(msgs)
|
||||
assert len(out) == 4
|
||||
assert out[1] is msgs[1]
|
||||
assert [b["type"] for b in out[3]["content"]] == ["thinking", "text"]
|
||||
assert out[3]["content"][0] is signed
|
||||
assert out[3]["content"][1]["text"] == "Sunny."
|
||||
|
||||
def test_strip_thinking_blocks_from_anthropic_messages(self):
|
||||
from litellm.llms.anthropic.common_utils import (
|
||||
strip_thinking_blocks_from_anthropic_messages,
|
||||
|
|
@ -1967,6 +2007,49 @@ class TestAnthropicThinkingSignatureSelfHeal:
|
|||
assert "thinking" not in data
|
||||
assert data["messages"] == []
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("thinking", "error_message", "expected"),
|
||||
(
|
||||
(
|
||||
{"type": "adaptive", "block_binding": {"prefix_mismatch_behavior": "error"}},
|
||||
_BOUND_THINKING_ERROR,
|
||||
False,
|
||||
),
|
||||
(
|
||||
{"type": "adaptive", "block_binding": {"prefix_mismatch_behavior": "error"}},
|
||||
_NULL_SIGNATURE_ERROR,
|
||||
True,
|
||||
),
|
||||
(
|
||||
{"type": "adaptive", "block_binding": {"prefix_mismatch_behavior": "drop_block"}},
|
||||
_BOUND_THINKING_ERROR,
|
||||
True,
|
||||
),
|
||||
({"type": "adaptive"}, _BOUND_THINKING_ERROR, True),
|
||||
),
|
||||
ids=("error", "error_on_unrelated_bad_block", "drop_block", "unbound"),
|
||||
)
|
||||
def test_anthropic_messages_retry_honors_thinking_block_binding(
|
||||
self, thinking: Mapping[str, object], error_message: str, expected: bool
|
||||
) -> None:
|
||||
from litellm.llms.anthropic.pass_through.messages.transformation import (
|
||||
AnthropicMessagesConfig,
|
||||
)
|
||||
|
||||
req: Final = httpx.Request("POST", "https://api.anthropic.com/v1/messages")
|
||||
resp: Final = httpx.Response(
|
||||
400,
|
||||
request=req,
|
||||
text=json.dumps({"type": "error", "error": {"type": "invalid_request_error", "message": error_message}}),
|
||||
)
|
||||
err: Final = httpx.HTTPStatusError("bad", request=req, response=resp)
|
||||
|
||||
should_retry = AnthropicMessagesConfig().should_retry_anthropic_messages_on_http_error(
|
||||
err, {}, request_body={"model": "claude-opus-5-5", "thinking": thinking}
|
||||
)
|
||||
|
||||
assert should_retry is expected
|
||||
|
||||
|
||||
class TestClaudeOpus48AdaptiveThinking:
|
||||
"""Opus 4.8 requires adaptive thinking (``thinking.type='adaptive'`` +
|
||||
|
|
@ -4092,6 +4175,117 @@ def test_shared_legacy_thinking_translation_preserves_supported_display(
|
|||
assert optional_params["thinking"] == expected_thinking
|
||||
|
||||
|
||||
LEGACY_THINKING_WITH_BLOCK_BINDING: Final = MappingProxyType(
|
||||
{
|
||||
"type": "enabled",
|
||||
"budget_tokens": 2048,
|
||||
"block_binding": {"prefix_mismatch_behavior": "error"},
|
||||
"future_key": 1,
|
||||
}
|
||||
)
|
||||
ADAPTIVE_THINKING_WITH_BLOCK_BINDING: Final = MappingProxyType(
|
||||
{"type": "adaptive", "block_binding": {"prefix_mismatch_behavior": "error"}}
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("local_model_cost_map")
|
||||
@pytest.mark.parametrize(
|
||||
("model", "custom_llm_provider", "expected_thinking"),
|
||||
(
|
||||
("claude-opus-4-8", "anthropic", ADAPTIVE_THINKING_WITH_BLOCK_BINDING),
|
||||
("claude-opus-4-8", "azure_ai", ADAPTIVE_THINKING_WITH_BLOCK_BINDING),
|
||||
("invoke/us.anthropic.claude-opus-4-8", "bedrock", ADAPTIVE_THINKING_WITH_BLOCK_BINDING),
|
||||
("converse/us.anthropic.claude-opus-4-8", "bedrock", {"type": "adaptive"}),
|
||||
("mantle/anthropic.claude-opus-4-8", "bedrock", {"type": "adaptive"}),
|
||||
("claude_platform/claude-opus-4-8", "bedrock", {"type": "adaptive"}),
|
||||
("claude-opus-4-8", "vertex_ai", {"type": "adaptive"}),
|
||||
("databricks-claude-opus-4-8", "databricks", {"type": "adaptive"}),
|
||||
),
|
||||
ids=(
|
||||
"anthropic",
|
||||
"azure_ai",
|
||||
"bedrock_invoke",
|
||||
"bedrock_converse",
|
||||
"bedrock_mantle",
|
||||
"bedrock_claude_platform",
|
||||
"vertex_ai",
|
||||
"databricks",
|
||||
),
|
||||
)
|
||||
def test_chat_legacy_thinking_translation_carries_only_block_binding_where_supported(
|
||||
model: str, custom_llm_provider: str, expected_thinking: Mapping[str, object]
|
||||
) -> None:
|
||||
from litellm.utils import get_optional_params
|
||||
|
||||
optional_params: Final = get_optional_params(
|
||||
model=model,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
thinking=dict(LEGACY_THINKING_WITH_BLOCK_BINDING),
|
||||
)
|
||||
|
||||
assert optional_params["thinking"] == expected_thinking
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("local_model_cost_map", "local_beta_headers_config")
|
||||
@pytest.mark.parametrize(
|
||||
("model", "custom_llm_provider", "expected_thinking"),
|
||||
(
|
||||
("claude-opus-4-8", "anthropic", ADAPTIVE_THINKING_WITH_BLOCK_BINDING),
|
||||
("claude-opus-4-8", "azure_ai", ADAPTIVE_THINKING_WITH_BLOCK_BINDING),
|
||||
("us.anthropic.claude-opus-4-8", "bedrock", ADAPTIVE_THINKING_WITH_BLOCK_BINDING),
|
||||
("mantle/anthropic.claude-opus-4-8", "bedrock", {"type": "adaptive"}),
|
||||
("anthropic.claude-opus-4-8", "bedrock_mantle", {"type": "adaptive"}),
|
||||
("claude_platform/claude-opus-4-8", "bedrock", {"type": "adaptive"}),
|
||||
("claude-opus-4-8", "vertex_ai", {"type": "adaptive"}),
|
||||
),
|
||||
ids=(
|
||||
"anthropic",
|
||||
"azure_ai",
|
||||
"bedrock_invoke",
|
||||
"bedrock_mantle",
|
||||
"bedrock_mantle_provider",
|
||||
"bedrock_claude_platform",
|
||||
"vertex_ai",
|
||||
),
|
||||
)
|
||||
def test_messages_legacy_thinking_translation_carries_only_block_binding_where_supported(
|
||||
model: str, custom_llm_provider: str, expected_thinking: Mapping[str, object]
|
||||
) -> None:
|
||||
from litellm.types.utils import LlmProviders
|
||||
from litellm.utils import ProviderConfigManager
|
||||
|
||||
config: Final = ProviderConfigManager.get_provider_anthropic_messages_config(
|
||||
model=model, provider=LlmProviders(custom_llm_provider)
|
||||
)
|
||||
assert config is not None
|
||||
|
||||
assert _thinking_sent_for_legacy_messages_request(config, model) == expected_thinking
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("local_model_cost_map", "local_beta_headers_config")
|
||||
def test_openai_compatible_messages_passthrough_keeps_legacy_thinking_without_block_binding() -> None:
|
||||
from litellm.llms.openai_like.messages.transformation import OpenAILikeAnthropicMessagesConfig
|
||||
|
||||
assert _thinking_sent_for_legacy_messages_request(OpenAILikeAnthropicMessagesConfig(), "claude-opus-4-8") == {
|
||||
"type": "adaptive"
|
||||
}
|
||||
|
||||
|
||||
def _thinking_sent_for_legacy_messages_request(config: BaseAnthropicMessagesConfig, model: str) -> object:
|
||||
from litellm.types.router import GenericLiteLLMParams
|
||||
|
||||
return config.transform_anthropic_messages_request(
|
||||
model=model,
|
||||
messages=[{"role": "user", "content": "Reply with OK"}],
|
||||
anthropic_messages_optional_request_params={
|
||||
"max_tokens": 512,
|
||||
"thinking": dict(LEGACY_THINKING_WITH_BLOCK_BINDING),
|
||||
},
|
||||
litellm_params=GenericLiteLLMParams(),
|
||||
headers={},
|
||||
).get("thinking")
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("local_model_cost_map", "local_beta_headers_config")
|
||||
@pytest.mark.parametrize("action", (None, "tool_addition", "tool_removal"))
|
||||
@pytest.mark.parametrize("explicit_beta", (False, True))
|
||||
|
|
|
|||
|
|
@ -352,3 +352,48 @@ def test_count_token_counters_select_their_own_provider(
|
|||
|
||||
assert counter.should_use_token_counting_api(custom_llm_provider=provider) is True
|
||||
assert counter.should_use_token_counting_api(custom_llm_provider="unknown") is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_global_anthropic_key_counts_like_chat_would_authenticate(monkeypatch):
|
||||
for name in ("ANTHROPIC_API_KEY", "ANTHROPIC_AUTH_TOKEN", "ANTHROPIC_API_BASE", "ANTHROPIC_BASE_URL"):
|
||||
monkeypatch.delenv(name, raising=False)
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
monkeypatch.setattr(litellm, "anthropic_key", FAKE_REGULAR_KEY)
|
||||
|
||||
with respx.mock(assert_all_called=True) as router:
|
||||
route = router.post("https://api.anthropic.com/v1/messages/count_tokens").mock(
|
||||
return_value=httpx.Response(200, json={"input_tokens": 13})
|
||||
)
|
||||
result = await AnthropicTokenCounter().count_tokens(
|
||||
model_to_use="claude-sonnet-4-5",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
contents=None,
|
||||
deployment={"litellm_params": {"model": "anthropic/claude-sonnet-4-5"}},
|
||||
request_model="claude-sonnet-4-5",
|
||||
)
|
||||
|
||||
assert result is not None
|
||||
assert result.total_tokens == 13
|
||||
assert route.calls.last.request.headers["x-api-key"] == FAKE_REGULAR_KEY
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_no_anthropic_key_returns_a_401_count_error(monkeypatch):
|
||||
for name in ("ANTHROPIC_API_KEY", "ANTHROPIC_AUTH_TOKEN"):
|
||||
monkeypatch.delenv(name, raising=False)
|
||||
monkeypatch.setattr(litellm, "anthropic_key", None)
|
||||
monkeypatch.setattr(litellm, "api_key", None)
|
||||
|
||||
result = await AnthropicTokenCounter().count_tokens(
|
||||
model_to_use="claude-sonnet-4-5",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
contents=None,
|
||||
deployment={"litellm_params": {"model": "anthropic/claude-sonnet-4-5"}},
|
||||
request_model="claude-sonnet-4-5",
|
||||
)
|
||||
|
||||
assert result is not None
|
||||
assert result.error is True
|
||||
assert result.status_code == 401
|
||||
assert result.total_tokens == 0
|
||||
|
|
|
|||
0
tests/unit/llms/azure_ai/anthropic/__init__.py
Normal file
0
tests/unit/llms/azure_ai/anthropic/__init__.py
Normal file
|
|
@ -0,0 +1,217 @@
|
|||
import json
|
||||
import operator
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import respx
|
||||
from pydantic import ValidationError
|
||||
|
||||
import litellm
|
||||
from litellm.llms.azure_ai.anthropic.count_tokens.token_counter import AzureAIAnthropicTokenCounter
|
||||
from litellm.types.utils import TokenCountResponse
|
||||
|
||||
API_BASE: Final = "https://my-resource.services.ai.azure.com"
|
||||
COUNT_URL: Final = f"{API_BASE}/anthropic/v1/messages/count_tokens"
|
||||
|
||||
|
||||
async def _count(monkeypatch: pytest.MonkeyPatch, response: httpx.Response) -> TokenCountResponse | None:
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
with respx.mock:
|
||||
respx.post(COUNT_URL).mock(return_value=response)
|
||||
return await AzureAIAnthropicTokenCounter().count_tokens(
|
||||
model_to_use="claude-sonnet-4-5",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
contents=None,
|
||||
deployment={"litellm_params": {"api_key": "azure-key", "api_base": API_BASE}},
|
||||
request_model="azure_ai/claude-sonnet-4-5",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("payload", "expected_total"),
|
||||
[
|
||||
({"input_tokens": 15}, 15),
|
||||
({"input_tokens": "15"}, 15),
|
||||
({"input_tokens": 15.0}, 15),
|
||||
({"input_tokens": True}, 1),
|
||||
({"input_tokens": -3}, -3),
|
||||
({}, 0),
|
||||
({"type": "error", "error": {"message": "429 rate limit"}}, 0),
|
||||
({"input_tokens": 7, "context_management": {"original_input_tokens": 9}, "nested": {"a": [1, None]}}, 7),
|
||||
],
|
||||
)
|
||||
async def test_count_reports_upstream_input_tokens_and_keeps_the_raw_payload(
|
||||
monkeypatch: pytest.MonkeyPatch, payload: dict[str, object], expected_total: int
|
||||
) -> None:
|
||||
result: Final = await _count(monkeypatch, httpx.Response(200, json=payload))
|
||||
|
||||
assert result is not None
|
||||
assert repr(result.model_dump()) == repr(
|
||||
{
|
||||
"total_tokens": expected_total,
|
||||
"request_model": "azure_ai/claude-sonnet-4-5",
|
||||
"model_used": "claude-sonnet-4-5",
|
||||
"tokenizer_type": "azure_ai_anthropic_api",
|
||||
"count_source": None,
|
||||
"original_response": payload,
|
||||
"error": False,
|
||||
"error_message": None,
|
||||
"status_code": None,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"rejected_count",
|
||||
[
|
||||
None,
|
||||
15.5,
|
||||
"Request timed out",
|
||||
"429 Too Many Requests",
|
||||
[15],
|
||||
[*range(403), "Request timed out"],
|
||||
{"value": 429},
|
||||
],
|
||||
)
|
||||
async def test_non_integer_input_tokens_becomes_an_error_response_carrying_the_validation_message(
|
||||
monkeypatch: pytest.MonkeyPatch, rejected_count: object
|
||||
) -> None:
|
||||
with pytest.raises(ValidationError) as rejection:
|
||||
TokenCountResponse.model_validate(
|
||||
{
|
||||
"total_tokens": rejected_count,
|
||||
"request_model": "azure_ai/claude-sonnet-4-5",
|
||||
"model_used": "claude-sonnet-4-5",
|
||||
"tokenizer_type": "azure_ai_anthropic_api",
|
||||
}
|
||||
)
|
||||
|
||||
result: Final = await _count(monkeypatch, httpx.Response(200, json={"input_tokens": rejected_count}))
|
||||
|
||||
assert result is not None
|
||||
assert [error["loc"] for error in rejection.value.errors()] == [("total_tokens",)]
|
||||
assert repr(result.model_dump()) == repr(
|
||||
{
|
||||
"total_tokens": 0,
|
||||
"request_model": "azure_ai/claude-sonnet-4-5",
|
||||
"model_used": "claude-sonnet-4-5",
|
||||
"tokenizer_type": "azure_ai_anthropic_api",
|
||||
"count_source": None,
|
||||
"original_response": None,
|
||||
"error": True,
|
||||
"error_message": str(rejection.value),
|
||||
"status_code": 500,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"payload",
|
||||
[[*({"input_tokens": position} for position in range(403)), "429"], "input_tokens", 15, 0, False],
|
||||
)
|
||||
async def test_count_payload_that_is_not_an_object_becomes_an_error_response_naming_the_failed_lookup(
|
||||
monkeypatch: pytest.MonkeyPatch, payload: object
|
||||
) -> None:
|
||||
with pytest.raises(AttributeError) as failed_lookup:
|
||||
operator.attrgetter("get")(payload)
|
||||
|
||||
result: Final = await _count(monkeypatch, httpx.Response(200, json=payload))
|
||||
|
||||
assert result is not None
|
||||
assert repr(result.model_dump()) == repr(
|
||||
{
|
||||
"total_tokens": 0,
|
||||
"request_model": "azure_ai/claude-sonnet-4-5",
|
||||
"model_used": "claude-sonnet-4-5",
|
||||
"tokenizer_type": "azure_ai_anthropic_api",
|
||||
"count_source": None,
|
||||
"original_response": None,
|
||||
"error": True,
|
||||
"error_message": str(failed_lookup.value),
|
||||
"status_code": 500,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_null_count_payload_falls_back_to_local_counting(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
assert await _count(monkeypatch, httpx.Response(200, content=b"null")) is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("upstream_status", "upstream_text"),
|
||||
[(429, "rate limit exceeded"), (403, "forbidden"), (408, "Request timed out")],
|
||||
)
|
||||
async def test_upstream_error_status_becomes_an_error_response_with_that_status_and_text(
|
||||
monkeypatch: pytest.MonkeyPatch, upstream_status: int, upstream_text: str
|
||||
) -> None:
|
||||
result: Final = await _count(monkeypatch, httpx.Response(upstream_status, text=upstream_text))
|
||||
|
||||
assert result is not None
|
||||
assert repr(result.model_dump()) == repr(
|
||||
{
|
||||
"total_tokens": 0,
|
||||
"request_model": "azure_ai/claude-sonnet-4-5",
|
||||
"model_used": "claude-sonnet-4-5",
|
||||
"tokenizer_type": "azure_ai_anthropic_api",
|
||||
"count_source": None,
|
||||
"original_response": None,
|
||||
"error": True,
|
||||
"error_message": upstream_text,
|
||||
"status_code": upstream_status,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_undecodable_count_body_becomes_a_500_error_response_carrying_the_decoder_message(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
with pytest.raises(json.JSONDecodeError) as undecodable:
|
||||
json.loads("not json")
|
||||
|
||||
result: Final = await _count(monkeypatch, httpx.Response(200, text="not json"))
|
||||
|
||||
assert result is not None
|
||||
assert repr(result.model_dump()) == repr(
|
||||
{
|
||||
"total_tokens": 0,
|
||||
"request_model": "azure_ai/claude-sonnet-4-5",
|
||||
"model_used": "claude-sonnet-4-5",
|
||||
"tokenizer_type": "azure_ai_anthropic_api",
|
||||
"count_source": None,
|
||||
"original_response": None,
|
||||
"error": True,
|
||||
"error_message": f"CountTokens processing error: {undecodable.value}",
|
||||
"status_code": 500,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("set_base", ("env", "global"), ids=("AZURE_API_BASE", "litellm.api_base"))
|
||||
async def test_count_uses_the_shared_api_base_chat_would_use(monkeypatch: pytest.MonkeyPatch, set_base: str) -> None:
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
monkeypatch.delenv("AZURE_AI_API_BASE", raising=False)
|
||||
monkeypatch.delenv("AZURE_API_BASE", raising=False)
|
||||
if set_base == "env":
|
||||
monkeypatch.setenv("AZURE_API_BASE", API_BASE)
|
||||
else:
|
||||
monkeypatch.setattr(litellm, "api_base", API_BASE)
|
||||
with respx.mock:
|
||||
route: Final = respx.post(COUNT_URL).mock(return_value=httpx.Response(200, json={"input_tokens": 6}))
|
||||
result: Final = await AzureAIAnthropicTokenCounter().count_tokens(
|
||||
model_to_use="claude-sonnet-4-5",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
contents=None,
|
||||
deployment={"litellm_params": {"api_key": "azure-key"}},
|
||||
request_model="azure_ai/claude-sonnet-4-5",
|
||||
)
|
||||
assert result is not None
|
||||
assert result.total_tokens == 6
|
||||
assert route.called
|
||||
|
|
@ -747,7 +747,7 @@ def test_bedrock_chat_invoke_tool_based_response_format_still_upgrades_legacy_th
|
|||
"schema": {"type": "object", "properties": {"result": {"type": "string"}}},
|
||||
},
|
||||
},
|
||||
"thinking": {"type": "enabled", "budget_tokens": 4096},
|
||||
"thinking": {"type": "enabled", "budget_tokens": 4096, "block_binding": BLOCK_BINDING},
|
||||
"max_tokens": 8192,
|
||||
},
|
||||
optional_params={},
|
||||
|
|
@ -756,10 +756,34 @@ def test_bedrock_chat_invoke_tool_based_response_format_still_upgrades_legacy_th
|
|||
)
|
||||
|
||||
assert "tools" in result
|
||||
assert result["thinking"] == {"type": "adaptive"}
|
||||
assert result["thinking"] == {"type": "adaptive", "block_binding": BLOCK_BINDING}
|
||||
assert result["output_config"] == {"effort": "high"}
|
||||
|
||||
|
||||
def test_bedrock_chat_invoke_adaptive_thinking_keeps_the_forced_response_format_tool(local_model_cost_map):
|
||||
"""Anthropic accepts a forced tool_choice next to adaptive thinking, and dropping
|
||||
it lets the model answer in prose instead of the requested JSON."""
|
||||
result = AmazonAnthropicClaudeConfig().map_openai_params(
|
||||
non_default_params={
|
||||
"response_format": {
|
||||
"type": "json_schema",
|
||||
"json_schema": {
|
||||
"name": "test_schema",
|
||||
"schema": {"type": "object", "properties": {"result": {"type": "string"}}},
|
||||
},
|
||||
},
|
||||
"thinking": {"type": "adaptive"},
|
||||
"max_tokens": 8192,
|
||||
},
|
||||
optional_params={},
|
||||
model="us.anthropic.claude-opus-5",
|
||||
drop_params=False,
|
||||
)
|
||||
|
||||
assert result["tool_choice"] == {"type": "tool", "name": result["tools"][0]["name"]}
|
||||
assert result["thinking"] == {"type": "adaptive"}
|
||||
|
||||
|
||||
def test_bedrock_chat_invoke_response_format_stub_still_upgrades_legacy_thinking(local_model_cost_map):
|
||||
"""Regression: the tool-based ``response_format`` path swaps in a Claude 3 stub
|
||||
model before the shared Anthropic mapping, which hid the adaptive-only model
|
||||
|
|
@ -974,6 +998,71 @@ def test_bedrock_chat_invoke_preserves_display_when_translating_legacy_thinking(
|
|||
assert ANTHROPIC_THINKING_DISPLAY_UPDATES_BETA_HEADER in result.get("anthropic_beta", [])
|
||||
|
||||
|
||||
BLOCK_BINDING: Final = {"prefix_mismatch_behavior": "error"}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"requested_thinking",
|
||||
(
|
||||
{"type": "adaptive", "display": "updates", "block_binding": BLOCK_BINDING, "future_key": 1},
|
||||
{
|
||||
"type": "enabled",
|
||||
"budget_tokens": 2048,
|
||||
"display": "updates",
|
||||
"block_binding": BLOCK_BINDING,
|
||||
"future_key": 1,
|
||||
},
|
||||
),
|
||||
ids=("adaptive", "legacy"),
|
||||
)
|
||||
def test_bedrock_chat_invoke_keeps_block_binding_with_display_updates(
|
||||
local_model_cost_map, local_beta_headers_config, requested_thinking
|
||||
) -> None:
|
||||
config: Final = AmazonAnthropicClaudeConfig()
|
||||
model: Final = "us.anthropic.claude-opus-5"
|
||||
optional_params: Final = config.map_openai_params(
|
||||
non_default_params={"max_tokens": 512, "thinking": requested_thinking},
|
||||
optional_params={},
|
||||
model=model,
|
||||
drop_params=False,
|
||||
)
|
||||
result: Final = config.transform_request(
|
||||
model=model,
|
||||
messages=[{"role": "user", "content": "Reply with OK"}],
|
||||
optional_params=optional_params,
|
||||
litellm_params={},
|
||||
headers={},
|
||||
)
|
||||
|
||||
assert result.get("thinking") == {"type": "adaptive", "display": "updates", "block_binding": BLOCK_BINDING}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"config_path",
|
||||
(
|
||||
"litellm.llms.bedrock.chat.mantle.transformation.AmazonMantleConfig",
|
||||
"litellm.llms.bedrock_mantle.chat.claude_transformation.BedrockMantleClaudeChatConfig",
|
||||
),
|
||||
ids=("bedrock_mantle_route", "bedrock_mantle_provider"),
|
||||
)
|
||||
def test_bedrock_chat_mantle_rebuilds_display_updates_without_block_binding(
|
||||
local_model_cost_map, local_beta_headers_config, config_path: str
|
||||
) -> None:
|
||||
module_name, _, class_name = config_path.rpartition(".")
|
||||
config: Final = getattr(importlib.import_module(module_name), class_name)()
|
||||
optional_params: Final = config.map_openai_params(
|
||||
non_default_params={
|
||||
"max_tokens": 512,
|
||||
"thinking": {"type": "adaptive", "display": "updates", "block_binding": BLOCK_BINDING, "future_key": 1},
|
||||
},
|
||||
optional_params={},
|
||||
model="anthropic.claude-opus-4-8",
|
||||
drop_params=False,
|
||||
)
|
||||
|
||||
assert optional_params.get("thinking") == {"type": "adaptive", "display": "updates"}
|
||||
|
||||
|
||||
FINE_GRAINED_TOOL_STREAMING_BETA: Final = "fine-grained-tool-streaming-2025-05-14"
|
||||
EAGER_TOOL_SCHEMA: Final = {"type": "object", "properties": {"path": {"type": "string"}}, "required": ["path"]}
|
||||
|
||||
|
|
|
|||
524
tests/unit/llms/bedrock/count_tokens/test_bedrock_kimi_k3.py
Normal file
524
tests/unit/llms/bedrock/count_tokens/test_bedrock_kimi_k3.py
Normal file
|
|
@ -0,0 +1,524 @@
|
|||
import asyncio
|
||||
import base64
|
||||
import hashlib
|
||||
import json
|
||||
import logging
|
||||
import re
|
||||
import threading
|
||||
from collections.abc import Iterable
|
||||
from itertools import product
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import respx
|
||||
import tiktoken
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.llms.bedrock.count_tokens.kimi_k3 import (
|
||||
KimiK3Encodings,
|
||||
KimiK3TokenizerSource,
|
||||
kimi_k3_anthropic_tokens,
|
||||
kimi_k3_tokens,
|
||||
load_kimi_k3_encoding,
|
||||
)
|
||||
from litellm.responses.litellm_completion_transformation.transformation import LiteLLMCompletionResponsesConfig
|
||||
|
||||
MODEL: Final = "us.moonshotai.kimi-k3"
|
||||
_CONTROL: Final = re.compile(r"<\|(?:open|close|sep|end_of_msg)\|>")
|
||||
_WEATHER: Final = {
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_weather",
|
||||
"description": "Get the weather in °C",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {"city": {"type": "string"}, "days": {"type": "integer"}},
|
||||
"required": ["city"],
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _call(call_id: str, arguments: dict[str, object]) -> dict[str, object]:
|
||||
return {"id": call_id, "type": "function", "function": {"name": "get_weather", "arguments": json.dumps(arguments)}}
|
||||
|
||||
|
||||
def _written_ranks(path: Path, tokens: Iterable[bytes]) -> KimiK3TokenizerSource:
|
||||
path.write_bytes(b"\n".join(base64.b64encode(token) + b" %d" % rank for rank, token in enumerate(tokens)))
|
||||
return KimiK3TokenizerSource(ranks_path=lambda: str(path), sha256=hashlib.sha256(path.read_bytes()).hexdigest())
|
||||
|
||||
|
||||
def _single_bytes() -> tuple[bytes, ...]:
|
||||
return tuple(bytes((value,)) for value in range(256))
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def byte_source(tmp_path_factory: pytest.TempPathFactory) -> KimiK3TokenizerSource:
|
||||
return _written_ranks(tmp_path_factory.mktemp("bytes") / "tiktoken.model", _single_bytes())
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def byte_encoding(byte_source: KimiK3TokenizerSource) -> tiktoken.Encoding:
|
||||
loaded: Final = load_kimi_k3_encoding(byte_source)
|
||||
assert loaded is not None
|
||||
return loaded
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def pair_encoding(tmp_path_factory: pytest.TempPathFactory) -> tiktoken.Encoding:
|
||||
pairs: Final = (bytes(pair) for pair in product(range(256), repeat=2))
|
||||
source: Final = _written_ranks(tmp_path_factory.mktemp("pairs") / "tiktoken.model", (*_single_bytes(), *pairs))
|
||||
loaded: Final = load_kimi_k3_encoding(source)
|
||||
assert loaded is not None
|
||||
return loaded
|
||||
|
||||
|
||||
def _template_count(encoding: tiktoken.Encoding, rendered: str) -> int:
|
||||
texts: Final = _CONTROL.split(rendered)
|
||||
return len(_CONTROL.findall(rendered)) + sum(len(encoding.encode_ordinary(text)) for text in texts)
|
||||
|
||||
|
||||
_GENERATION_PROMPT: Final = '<|open|>message role="assistant"<|sep|><|open|>think<|sep|>'
|
||||
|
||||
|
||||
# Expected renders follow the Kimi K3 chat template in moonshotai/Kimi-K3 at revision f831ab6, which matched
|
||||
# Bedrock's prompt_tokens for us.moonshotai.kimi-k3 in 19 live calls on 2026-10-16
|
||||
@pytest.mark.parametrize(
|
||||
("messages", "tools", "rendered"),
|
||||
[
|
||||
pytest.param(
|
||||
[{"role": "user", "content": "Hello there"}],
|
||||
None,
|
||||
'<|open|>message role="user"<|sep|>Hello there<|close|>message<|sep|><|end_of_msg|>' + _GENERATION_PROMPT,
|
||||
id="user",
|
||||
),
|
||||
pytest.param(
|
||||
[
|
||||
{
|
||||
"role": "system",
|
||||
"content": [{"type": "text", "text": "Be terse."}, {"type": "text", "text": "Answer in French."}],
|
||||
},
|
||||
{"role": "user", "content": [{"type": "text", "text": "Hi"}, {"type": "text", "text": " there"}]},
|
||||
],
|
||||
None,
|
||||
'<|open|>message role="system"<|sep|>Be terse.Answer in French.<|close|>message<|sep|><|end_of_msg|>'
|
||||
'<|open|>message role="user"<|sep|>Hi there<|close|>message<|sep|><|end_of_msg|>' + _GENERATION_PROMPT,
|
||||
id="text_parts_join",
|
||||
),
|
||||
pytest.param(
|
||||
[
|
||||
{"role": "user", "content": "What is 2+2?"},
|
||||
{"role": "assistant", "content": "4"},
|
||||
{"role": "user", "content": "And 3+3?"},
|
||||
],
|
||||
None,
|
||||
'<|open|>message role="user"<|sep|>What is 2+2?<|close|>message<|sep|><|end_of_msg|>'
|
||||
'<|open|>message role="assistant"<|sep|><|open|>think<|sep|><|close|>think<|sep|>'
|
||||
"<|open|>response<|sep|>4<|close|>response<|sep|><|close|>message<|sep|><|end_of_msg|>"
|
||||
'<|open|>message role="user"<|sep|>And 3+3?<|close|>message<|sep|><|end_of_msg|>' + _GENERATION_PROMPT,
|
||||
id="assistant_turn",
|
||||
),
|
||||
pytest.param(
|
||||
[
|
||||
{"role": "user", "content": "Weather in Paris and Oslo?"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"tool_calls": [_call("c1", {"city": "Paris", "days": 3}), _call("c2", {"city": "Oslo"})],
|
||||
},
|
||||
{"role": "tool", "tool_call_id": "c2", "content": "Snow"},
|
||||
{"role": "tool", "tool_call_id": "c1", "content": "Sun"},
|
||||
],
|
||||
[_WEATHER],
|
||||
'<|open|>message role="system" type="tool-declare"<|sep|># Tools\n'
|
||||
"Here are the available tools, described in JSONSchema.\n\n```json\n"
|
||||
'[{"function":{"description":"Get the weather in °C","name":"get_weather","parameters":{"properties":'
|
||||
'{"city":{"type":"string"},"days":{"type":"integer"}},"required":["city"],"type":"object"}},'
|
||||
'"type":"function"}]\n```<|close|>message<|sep|><|end_of_msg|>'
|
||||
'<|open|>message role="user"<|sep|>Weather in Paris and Oslo?<|close|>message<|sep|><|end_of_msg|>'
|
||||
'<|open|>message role="assistant"<|sep|><|open|>think<|sep|><|close|>think<|sep|>'
|
||||
"<|open|>response<|sep|><|close|>response<|sep|><|open|>tools<|sep|>"
|
||||
'<|open|>call tool="get_weather" index="1"<|sep|>'
|
||||
'<|open|>argument key="city" type="string"<|sep|>Paris<|close|>argument<|sep|>'
|
||||
'<|open|>argument key="days" type="number"<|sep|>3<|close|>argument<|sep|><|close|>call<|sep|>'
|
||||
'<|open|>call tool="get_weather" index="2"<|sep|>'
|
||||
'<|open|>argument key="city" type="string"<|sep|>Oslo<|close|>argument<|sep|><|close|>call<|sep|>'
|
||||
"<|close|>tools<|sep|><|close|>message<|sep|><|end_of_msg|>"
|
||||
'<|open|>message role="tool" tool="get_weather" index="1"<|sep|>Sun<|close|>message<|sep|><|end_of_msg|>'
|
||||
'<|open|>message role="tool" tool="get_weather" index="2"<|sep|>Snow<|close|>message<|sep|><|end_of_msg|>'
|
||||
+ _GENERATION_PROMPT,
|
||||
id="tool_results_out_of_call_order",
|
||||
),
|
||||
pytest.param(
|
||||
[
|
||||
{"role": "user", "content": "Weather?"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"tool_calls": [
|
||||
_call(
|
||||
"c1", {"a&b": 'x"y', "flag": True, "none": None, "list": [1, "zé"], "object": {"k": True}}
|
||||
)
|
||||
],
|
||||
},
|
||||
{"role": "tool", "tool_call_id": "c1", "content": "Rain"},
|
||||
],
|
||||
[_WEATHER],
|
||||
'<|open|>message role="system" type="tool-declare"<|sep|># Tools\n'
|
||||
"Here are the available tools, described in JSONSchema.\n\n```json\n"
|
||||
'[{"function":{"description":"Get the weather in °C","name":"get_weather","parameters":{"properties":'
|
||||
'{"city":{"type":"string"},"days":{"type":"integer"}},"required":["city"],"type":"object"}},'
|
||||
'"type":"function"}]\n```<|close|>message<|sep|><|end_of_msg|>'
|
||||
'<|open|>message role="user"<|sep|>Weather?<|close|>message<|sep|><|end_of_msg|>'
|
||||
'<|open|>message role="assistant"<|sep|><|open|>think<|sep|><|close|>think<|sep|>'
|
||||
"<|open|>response<|sep|><|close|>response<|sep|><|open|>tools<|sep|>"
|
||||
'<|open|>call tool="get_weather" index="1"<|sep|>'
|
||||
'<|open|>argument key="a&b" type="string"<|sep|>x"y<|close|>argument<|sep|>'
|
||||
'<|open|>argument key="flag" type="boolean"<|sep|>true<|close|>argument<|sep|>'
|
||||
'<|open|>argument key="none" type="null"<|sep|>null<|close|>argument<|sep|>'
|
||||
'<|open|>argument key="list" type="array"<|sep|>[1, "zé"]<|close|>argument<|sep|>'
|
||||
'<|open|>argument key="object" type="object"<|sep|>{"k": true}<|close|>argument<|sep|>'
|
||||
"<|close|>call<|sep|><|close|>tools<|sep|><|close|>message<|sep|><|end_of_msg|>"
|
||||
'<|open|>message role="tool" tool="get_weather" index="1"<|sep|>Rain<|close|>message<|sep|><|end_of_msg|>'
|
||||
+ _GENERATION_PROMPT,
|
||||
id="argument_types_and_escaping",
|
||||
),
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("encoding_fixture", ["byte_encoding", "pair_encoding"])
|
||||
def test_count_is_the_kimi_k3_chat_template_render_of_the_bedrock_request(
|
||||
request: pytest.FixtureRequest,
|
||||
encoding_fixture: str,
|
||||
messages: list[dict[str, object]],
|
||||
tools: list[dict[str, object]] | None,
|
||||
rendered: str,
|
||||
):
|
||||
encoding: Final[tiktoken.Encoding] = request.getfixturevalue(encoding_fixture)
|
||||
|
||||
assert kimi_k3_tokens(encoding, MODEL, messages, tools) == _template_count(encoding, rendered)
|
||||
|
||||
|
||||
def _parallel_calls(count: int) -> list[dict[str, object]]:
|
||||
ids: Final = tuple(f"c{number}" for number in range(1, count + 1))
|
||||
return [
|
||||
{"role": "user", "content": "Weather?"},
|
||||
{"role": "assistant", "content": None, "tool_calls": [_call(call_id, {"city": "Paris"}) for call_id in ids]},
|
||||
*({"role": "tool", "tool_call_id": call_id, "content": "Sun"} for call_id in ids),
|
||||
]
|
||||
|
||||
|
||||
def test_tenth_parallel_call_and_its_result_are_numbered_ten(byte_encoding: tiktoken.Encoding):
|
||||
nine: Final = kimi_k3_tokens(byte_encoding, MODEL, _parallel_calls(9), [_WEATHER])
|
||||
ten: Final = kimi_k3_tokens(byte_encoding, MODEL, _parallel_calls(10), [_WEATHER])
|
||||
assert nine is not None and ten is not None
|
||||
|
||||
assert ten - nine == _template_count(
|
||||
byte_encoding,
|
||||
'<|open|>call tool="get_weather" index="10"<|sep|>'
|
||||
'<|open|>argument key="city" type="string"<|sep|>Paris<|close|>argument<|sep|><|close|>call<|sep|>'
|
||||
'<|open|>message role="tool" tool="get_weather" index="10"<|sep|>Sun<|close|>message<|sep|><|end_of_msg|>',
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"text",
|
||||
[
|
||||
pytest.param("The quick brown fox, jumps over 12345 lazy dogs. " * 60, id="prose"),
|
||||
pytest.param("a" + "1234567890" * 250 + "1", id="long_digit_run"),
|
||||
pytest.param("--- " * 600, id="symbols"),
|
||||
pytest.param("def f(x):\n return x # it's 42\n\n" * 80, id="code"),
|
||||
pytest.param("word next\t\ttab\n \n" * 120, id="whitespace"),
|
||||
pytest.param("naïve café, “quotes” & <tags> 中文汉字,日本語のテキスト 😀🎉\r\n" * 40, id="unicode"),
|
||||
pytest.param("We're sure they'd say it's DON'T or I'M " * 60, id="contractions"),
|
||||
pytest.param("abc123def4567ghi89 " * 120, id="letters_and_digits"),
|
||||
],
|
||||
)
|
||||
def test_long_text_counts_what_the_tokenizer_counts_for_the_whole_text(pair_encoding: tiktoken.Encoding, text: str):
|
||||
def counted(content: str) -> int | None:
|
||||
return kimi_k3_tokens(pair_encoding, MODEL, [{"role": "user", "content": content}], None)
|
||||
|
||||
short: Final = counted("x")
|
||||
assert short is not None
|
||||
|
||||
assert counted(text) == short - len(pair_encoding.encode_ordinary("x")) + len(pair_encoding.encode_ordinary(text))
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("messages", "tools"),
|
||||
[
|
||||
pytest.param(
|
||||
[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "What is this?"},
|
||||
{"type": "image_url", "image_url": {"url": "data:image/png;base64,iVBORw0KGgo="}},
|
||||
],
|
||||
}
|
||||
],
|
||||
None,
|
||||
id="image",
|
||||
),
|
||||
pytest.param(
|
||||
[{"role": "user", "content": "Weather?"}, {"role": "tool", "tool_call_id": "c1", "content": "Sun"}],
|
||||
[_WEATHER],
|
||||
id="tool_result_without_its_call",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_request_the_template_cannot_render_is_not_counted(
|
||||
pair_encoding: tiktoken.Encoding, messages: list[dict[str, object]], tools: list[dict[str, object]] | None
|
||||
):
|
||||
assert kimi_k3_tokens(pair_encoding, MODEL, messages, tools) is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"tool",
|
||||
[
|
||||
pytest.param(
|
||||
{
|
||||
"type": "function",
|
||||
"name": "get_weather",
|
||||
"description": "Get the weather in °C",
|
||||
"parameters": _WEATHER["function"]["parameters"],
|
||||
},
|
||||
id="responses_api",
|
||||
),
|
||||
pytest.param(
|
||||
{
|
||||
"name": "get_weather",
|
||||
"description": "Get the weather in °C",
|
||||
"input_schema": _WEATHER["function"]["parameters"],
|
||||
},
|
||||
id="anthropic_messages",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_tool_in_another_api_shape_counts_like_the_same_chat_completions_tool(
|
||||
pair_encoding: tiktoken.Encoding, tool: dict[str, object]
|
||||
):
|
||||
messages: Final = [{"role": "user", "content": "Weather in Paris?"}]
|
||||
|
||||
counted: Final = kimi_k3_tokens(pair_encoding, MODEL, messages, [tool])
|
||||
|
||||
assert counted is not None
|
||||
assert counted == kimi_k3_tokens(pair_encoding, MODEL, messages, [_WEATHER])
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"tool",
|
||||
[
|
||||
pytest.param({"type": "web_search_20250305", "name": "web_search"}, id="anthropic_web_search"),
|
||||
pytest.param({"type": "web_search"}, id="responses_web_search"),
|
||||
pytest.param(
|
||||
{
|
||||
"type": "mcp",
|
||||
"server_label": "docs",
|
||||
"server_url": "https://mcp.example/mcp",
|
||||
"require_approval": "never",
|
||||
},
|
||||
id="mcp",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_tool_bedrock_inference_drops_counts_like_the_function_tools_alone(
|
||||
pair_encoding: tiktoken.Encoding, tool: dict[str, object]
|
||||
):
|
||||
messages: Final = [{"role": "user", "content": "Weather in Paris?"}]
|
||||
|
||||
counted: Final = kimi_k3_tokens(pair_encoding, MODEL, messages, [_WEATHER, tool])
|
||||
|
||||
assert counted is not None
|
||||
assert counted == kimi_k3_tokens(pair_encoding, MODEL, messages, [_WEATHER])
|
||||
|
||||
|
||||
def test_responses_custom_tool_counts_like_the_function_tool_inference_sends(pair_encoding: tiktoken.Encoding):
|
||||
messages: Final = [{"role": "user", "content": "Weather in Paris?"}]
|
||||
custom_tool: Final = {"type": "custom", "name": "grammar", "description": "Free-form SQL"}
|
||||
sent_to_inference, _ = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools(
|
||||
[dict(custom_tool)]
|
||||
)
|
||||
|
||||
without_it: Final = kimi_k3_tokens(pair_encoding, MODEL, messages, [_WEATHER])
|
||||
|
||||
counted: Final = kimi_k3_tokens(pair_encoding, MODEL, messages, [_WEATHER, custom_tool])
|
||||
|
||||
assert counted is not None and without_it is not None
|
||||
assert counted > without_it
|
||||
assert counted == kimi_k3_tokens(pair_encoding, MODEL, messages, [_WEATHER, *sent_to_inference])
|
||||
|
||||
|
||||
def test_counting_a_request_with_an_image_url_never_downloads_the_image(
|
||||
pair_encoding: tiktoken.Encoding, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
|
||||
):
|
||||
monkeypatch.setattr(litellm, "user_url_validation", False)
|
||||
image: Final = respx_mock.get("https://images.example/cat.png").mock(return_value=httpx.Response(200))
|
||||
messages: Final = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "What is this?"},
|
||||
{"type": "image_url", "image_url": {"url": "https://images.example/cat.png"}},
|
||||
],
|
||||
}
|
||||
]
|
||||
|
||||
assert kimi_k3_tokens(pair_encoding, MODEL, messages, None) is None
|
||||
assert not image.called
|
||||
|
||||
|
||||
def test_anthropic_request_counts_like_the_same_openai_request(pair_encoding: tiktoken.Encoding):
|
||||
anthropic: Final = kimi_k3_anthropic_tokens(
|
||||
pair_encoding,
|
||||
MODEL,
|
||||
{
|
||||
"system": "Be terse.",
|
||||
"messages": [
|
||||
{"role": "user", "content": "Weather in Paris?"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [{"type": "tool_use", "id": "c1", "name": "get_weather", "input": {"city": "Paris"}}],
|
||||
},
|
||||
{"role": "user", "content": [{"type": "tool_result", "tool_use_id": "c1", "content": "Sun"}]},
|
||||
],
|
||||
"tools": [
|
||||
{
|
||||
"name": "get_weather",
|
||||
"description": "Get the weather in °C",
|
||||
"input_schema": _WEATHER["function"]["parameters"],
|
||||
}
|
||||
],
|
||||
},
|
||||
)
|
||||
openai: Final = kimi_k3_tokens(
|
||||
pair_encoding,
|
||||
MODEL,
|
||||
[
|
||||
{"role": "system", "content": "Be terse."},
|
||||
{"role": "user", "content": "Weather in Paris?"},
|
||||
{"role": "assistant", "content": None, "tool_calls": [_call("c1", {"city": "Paris"})]},
|
||||
{"role": "tool", "tool_call_id": "c1", "content": "Sun"},
|
||||
],
|
||||
[_WEATHER],
|
||||
)
|
||||
|
||||
assert openai is not None
|
||||
assert anthropic == openai
|
||||
|
||||
|
||||
@pytest.mark.parametrize("failure", ["checksum", "missing_file"])
|
||||
def test_tokenizer_that_fails_its_checksum_or_cannot_be_read_is_not_loaded(
|
||||
byte_source: KimiK3TokenizerSource, tmp_path: Path, caplog: pytest.LogCaptureFixture, failure: str
|
||||
):
|
||||
source: Final = (
|
||||
KimiK3TokenizerSource(ranks_path=byte_source.ranks_path, sha256=hashlib.sha256(b"other").hexdigest())
|
||||
if failure == "checksum"
|
||||
else KimiK3TokenizerSource(ranks_path=lambda: str(tmp_path / "missing"), sha256=byte_source.sha256)
|
||||
)
|
||||
verbose_logger.addHandler(caplog.handler)
|
||||
try:
|
||||
with caplog.at_level(logging.WARNING, logger=verbose_logger.name):
|
||||
loaded: Final = load_kimi_k3_encoding(source)
|
||||
finally:
|
||||
verbose_logger.removeHandler(caplog.handler)
|
||||
|
||||
assert loaded is None
|
||||
assert any("Kimi K3 tokenizer unavailable" in record.getMessage() for record in caplog.records)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_slow_tokenizer_load_is_not_waited_on_past_the_wait_and_is_shared_once_done(
|
||||
byte_source: KimiK3TokenizerSource,
|
||||
):
|
||||
released: Final = threading.Event()
|
||||
|
||||
def slow_ranks_path() -> str:
|
||||
released.wait(10)
|
||||
return byte_source.ranks_path()
|
||||
|
||||
encodings: Final = KimiK3Encodings(
|
||||
source=KimiK3TokenizerSource(ranks_path=slow_ranks_path, sha256=byte_source.sha256),
|
||||
retry_after_seconds=60,
|
||||
clock=lambda: 0.0,
|
||||
)
|
||||
|
||||
assert await encodings.encoding(wait_seconds=0) is None
|
||||
released.set()
|
||||
loaded: Final = await encodings.encoding(wait_seconds=5)
|
||||
assert loaded is not None
|
||||
assert await encodings.encoding(wait_seconds=0) is loaded
|
||||
|
||||
|
||||
class _Clock:
|
||||
def __init__(self) -> None:
|
||||
self.now = 0.0
|
||||
|
||||
def __call__(self) -> float:
|
||||
return self.now
|
||||
|
||||
|
||||
async def _loaded_eventually(encodings: KimiK3Encodings) -> tiktoken.Encoding | None:
|
||||
for _ in range(500):
|
||||
if (loaded := await encodings.encoding(wait_seconds=0)) is not None:
|
||||
return loaded
|
||||
await asyncio.sleep(0.01)
|
||||
return None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_failed_tokenizer_load_is_retried_in_the_background_only_after_the_retry_window(
|
||||
byte_source: KimiK3TokenizerSource, tmp_path: Path
|
||||
):
|
||||
clock: Final = _Clock()
|
||||
reads: Final[list[float]] = []
|
||||
released: Final = threading.Event()
|
||||
|
||||
def ranks_path() -> str:
|
||||
reads.append(clock.now)
|
||||
if len(reads) == 1:
|
||||
return str(tmp_path / "missing")
|
||||
released.wait(10)
|
||||
return byte_source.ranks_path()
|
||||
|
||||
encodings: Final = KimiK3Encodings(
|
||||
source=KimiK3TokenizerSource(ranks_path=ranks_path, sha256=byte_source.sha256),
|
||||
retry_after_seconds=60,
|
||||
clock=clock,
|
||||
)
|
||||
|
||||
assert await encodings.encoding(wait_seconds=5) is None
|
||||
clock.now = 59.0
|
||||
assert await encodings.encoding(wait_seconds=5) is None
|
||||
clock.now = 60.0
|
||||
assert await asyncio.wait_for(encodings.encoding(wait_seconds=5), timeout=1) is None
|
||||
released.set()
|
||||
assert await _loaded_eventually(encodings) is not None
|
||||
assert reads == [0.0, 60.0]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_tokenizer_load_still_running_after_the_wait_is_logged(
|
||||
byte_source: KimiK3TokenizerSource, caplog: pytest.LogCaptureFixture
|
||||
):
|
||||
released: Final = threading.Event()
|
||||
|
||||
def hung_ranks_path() -> str:
|
||||
released.wait(10)
|
||||
return byte_source.ranks_path()
|
||||
|
||||
encodings: Final = KimiK3Encodings(
|
||||
source=KimiK3TokenizerSource(ranks_path=hung_ranks_path, sha256=byte_source.sha256), retry_after_seconds=60
|
||||
)
|
||||
verbose_logger.addHandler(caplog.handler)
|
||||
try:
|
||||
with caplog.at_level(logging.WARNING, logger=verbose_logger.name):
|
||||
loaded: Final = await encodings.encoding(wait_seconds=0.05)
|
||||
finally:
|
||||
verbose_logger.removeHandler(caplog.handler)
|
||||
released.set()
|
||||
|
||||
assert loaded is None
|
||||
assert any("Kimi K3 tokenizer still loading" in record.getMessage() for record in caplog.records)
|
||||
|
|
@ -36,6 +36,8 @@ from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_tran
|
|||
AmazonAnthropicClaudeMessagesConfig,
|
||||
AmazonAnthropicClaudeMessagesStreamDecoder,
|
||||
)
|
||||
from litellm.llms.bedrock.messages.mantle_transformation import AmazonMantleMessagesConfig
|
||||
from litellm.llms.bedrock_mantle.messages.transformation import BedrockMantleAnthropicMessagesConfig
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -3529,6 +3531,94 @@ def test_bedrock_clear_thinking_preserves_display_updates() -> None:
|
|||
assert ANTHROPIC_THINKING_DISPLAY_UPDATES_BETA_HEADER in result.get("anthropic_beta", [])
|
||||
|
||||
|
||||
BLOCK_BINDING: Final = {"prefix_mismatch_behavior": "error"}
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("local_model_cost_map", "local_beta_headers_config")
|
||||
@pytest.mark.parametrize(
|
||||
("model", "requested_thinking", "extra_params"),
|
||||
(
|
||||
(
|
||||
"eu.anthropic.claude-opus-5",
|
||||
{"type": "adaptive", "display": "updates", "block_binding": BLOCK_BINDING, "future_key": 1},
|
||||
{},
|
||||
),
|
||||
(
|
||||
"eu.anthropic.claude-opus-5",
|
||||
{
|
||||
"type": "enabled",
|
||||
"budget_tokens": 24000,
|
||||
"display": "updates",
|
||||
"block_binding": BLOCK_BINDING,
|
||||
"future_key": 1,
|
||||
},
|
||||
{},
|
||||
),
|
||||
(
|
||||
"us.anthropic.claude-opus-4-6",
|
||||
{
|
||||
"type": "enabled",
|
||||
"budget_tokens": 2048,
|
||||
"display": "updates",
|
||||
"block_binding": BLOCK_BINDING,
|
||||
"future_key": 1,
|
||||
},
|
||||
{"context_management": {"edits": [{"type": "clear_thinking_20251015"}]}},
|
||||
),
|
||||
),
|
||||
ids=("adaptive", "legacy", "legacy_clear_thinking"),
|
||||
)
|
||||
def test_bedrock_messages_keeps_block_binding_with_display_updates(
|
||||
model: str, requested_thinking: dict[str, object], extra_params: dict[str, object]
|
||||
) -> None:
|
||||
from litellm.types.router import GenericLiteLLMParams
|
||||
|
||||
result: Final = AmazonAnthropicClaudeMessagesConfig().transform_anthropic_messages_request(
|
||||
model=model,
|
||||
messages=[{"role": "user", "content": "Reply with OK"}],
|
||||
anthropic_messages_optional_request_params={"max_tokens": 512, "thinking": requested_thinking, **extra_params},
|
||||
litellm_params=GenericLiteLLMParams(),
|
||||
headers={},
|
||||
)
|
||||
|
||||
assert result.get("thinking") == {"type": "adaptive", "display": "updates", "block_binding": BLOCK_BINDING}
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("local_model_cost_map", "local_beta_headers_config")
|
||||
@pytest.mark.parametrize(
|
||||
"config_class",
|
||||
(AmazonMantleMessagesConfig, BedrockMantleAnthropicMessagesConfig),
|
||||
ids=("bedrock_mantle_route", "bedrock_mantle_provider"),
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
("requested_thinking", "extra_params"),
|
||||
(
|
||||
({"type": "adaptive", "display": "updates", "block_binding": BLOCK_BINDING, "future_key": 1}, {}),
|
||||
(
|
||||
{"type": "enabled", "budget_tokens": 2048, "display": "updates", "block_binding": BLOCK_BINDING},
|
||||
{"context_management": {"edits": [{"type": "clear_thinking_20251015"}]}},
|
||||
),
|
||||
),
|
||||
ids=("adaptive", "legacy_clear_thinking"),
|
||||
)
|
||||
def test_bedrock_mantle_messages_rebuilds_display_updates_without_block_binding(
|
||||
config_class: type[AmazonMantleMessagesConfig],
|
||||
requested_thinking: dict[str, object],
|
||||
extra_params: dict[str, object],
|
||||
) -> None:
|
||||
from litellm.types.router import GenericLiteLLMParams
|
||||
|
||||
result: Final = config_class().transform_anthropic_messages_request(
|
||||
model="anthropic.claude-opus-4-8",
|
||||
messages=[{"role": "user", "content": "Reply with OK"}],
|
||||
anthropic_messages_optional_request_params={"max_tokens": 512, "thinking": requested_thinking, **extra_params},
|
||||
litellm_params=GenericLiteLLMParams(),
|
||||
headers={},
|
||||
)
|
||||
|
||||
assert result.get("thinking") == {"type": "adaptive", "display": "updates"}
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("local_model_cost_map", "local_beta_headers_config")
|
||||
@pytest.mark.parametrize("action", (None, "tool_addition", "tool_removal"))
|
||||
@pytest.mark.parametrize("explicit_beta", (False, True))
|
||||
|
|
|
|||
|
|
@ -1804,6 +1804,71 @@ async def test_anthropic_post_retry_reserializes_mutated_body():
|
|||
assert "MUTATED" in second_sent # ... the mutated body
|
||||
|
||||
|
||||
_BOUND_THINKING_ERROR: Final = (
|
||||
'{"type":"error","error":{"type":"invalid_request_error","message":"messages.1.content.0: '
|
||||
'Invalid `signature` in `thinking` block. The block is bound to a different conversation."}}'
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("prefix_mismatch_behavior", "expected_posts"),
|
||||
(("error", 1), ("drop_block", 2)),
|
||||
)
|
||||
async def test_anthropic_post_skips_the_thinking_strip_retry_when_block_binding_errors(
|
||||
prefix_mismatch_behavior: str, expected_posts: int
|
||||
) -> None:
|
||||
"""A caller that pinned ``prefix_mismatch_behavior: error`` asked for the 400,
|
||||
so the handler must surface it instead of stripping the bound blocks and retrying."""
|
||||
from litellm.llms.anthropic.pass_through.messages.transformation import AnthropicMessagesConfig
|
||||
|
||||
handler = BaseLLMHTTPHandler()
|
||||
thinking_block = {"type": "thinking", "thinking": "", "signature": "sig"}
|
||||
request_body = {
|
||||
"model": "claude-opus-5-5",
|
||||
"messages": [
|
||||
{"role": "user", "content": "hi"},
|
||||
{"role": "assistant", "content": [thinking_block, {"type": "text", "text": "hello"}]},
|
||||
],
|
||||
"thinking": {"type": "adaptive", "block_binding": {"prefix_mismatch_behavior": prefix_mismatch_behavior}},
|
||||
}
|
||||
request = httpx.Request("POST", "https://api.anthropic.com/v1/messages")
|
||||
http_client = Mock()
|
||||
http_client.post = AsyncMock(
|
||||
side_effect=[
|
||||
httpx.Response(400, request=request, text=_BOUND_THINKING_ERROR),
|
||||
httpx.Response(200, request=request, json={"type": "message", "content": []}),
|
||||
]
|
||||
)
|
||||
logging_obj: Final = Mock(baseline_cache_context=None)
|
||||
logging_obj.model_call_details = {}
|
||||
|
||||
async def _post() -> httpx.Response:
|
||||
return await handler._async_post_anthropic_messages_with_http_error_retry(
|
||||
async_httpx_client=http_client,
|
||||
request_url="https://api.anthropic.com/v1/messages",
|
||||
headers={},
|
||||
signed_json_body=None,
|
||||
request_body=request_body,
|
||||
stream=False,
|
||||
logging_obj=logging_obj,
|
||||
provider_config=AnthropicMessagesConfig(),
|
||||
litellm_params=GenericLiteLLMParams(),
|
||||
api_key="k",
|
||||
model="claude-opus-5-5",
|
||||
)
|
||||
|
||||
if expected_posts == 1:
|
||||
with pytest.raises(BaseLLMException) as raised:
|
||||
await _post()
|
||||
assert raised.value.status_code == 400
|
||||
assert "bound to a different conversation" in raised.value.message
|
||||
assert request_body["messages"][1]["content"][0] == thinking_block
|
||||
else:
|
||||
assert (await _post()).status_code == 200
|
||||
assert http_client.post.await_count == expected_posts
|
||||
|
||||
|
||||
def test_base_responses_config_sign_request_is_noop_by_default():
|
||||
"""Default responses sign_request must be a no-op: unchanged headers, no signed body.
|
||||
|
||||
|
|
|
|||
|
|
@ -1,8 +1,200 @@
|
|||
from typing import Final
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
|
||||
import litellm
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
|
||||
from litellm.llms.gemini.common_utils import GeminiModelInfo, GoogleAIStudioTokenCounter
|
||||
from litellm.llms.gemini.count_tokens import handler as gemini_count_tokens_handler
|
||||
from litellm.llms.gemini.count_tokens.handler import GoogleAIStudioTokenCounter as GeminiCountTokensHandler
|
||||
|
||||
_JSON_BODY: Final = TypeAdapter(dict[str, JsonValue])
|
||||
|
||||
|
||||
class _GeminiCountTokensPeer:
|
||||
def __init__(self) -> None:
|
||||
self.request_body: dict[str, JsonValue] | None = None
|
||||
|
||||
def __call__(self, request: httpx.Request) -> httpx.Response:
|
||||
self.request_body = _JSON_BODY.validate_json(request.content)
|
||||
return httpx.Response(200, json={"totalTokens": 11}, request=request)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_google_ai_studio_count_tokens_forwards_system_and_tools() -> None:
|
||||
peer: Final = _GeminiCountTokensPeer()
|
||||
client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(peer))
|
||||
|
||||
with patch.object(gemini_count_tokens_handler, "get_async_httpx_client", return_value=client):
|
||||
result: Final = await GoogleAIStudioTokenCounter().count_tokens(
|
||||
model_to_use="gemini-2.5-flash",
|
||||
messages=None,
|
||||
contents=[{"role": "user", "parts": [{"text": "hello"}]}],
|
||||
deployment={
|
||||
"litellm_params": {
|
||||
"api_key": "test-key",
|
||||
"api_base": "https://count-tokens.test",
|
||||
"system_instruction": {"parts": [{"text": "deployment-level"}]},
|
||||
}
|
||||
},
|
||||
request_model="gemini/gemini-2.5-flash",
|
||||
tools=[
|
||||
{
|
||||
"name": "lookup",
|
||||
"description": "Find a record",
|
||||
"input_schema": {"type": "object", "properties": {}},
|
||||
}
|
||||
],
|
||||
system="You are terse.",
|
||||
)
|
||||
|
||||
assert result is not None
|
||||
assert result.error is False
|
||||
assert result.total_tokens == 11
|
||||
assert peer.request_body == {
|
||||
"generateContentRequest": {
|
||||
"model": "models/gemini-2.5-flash",
|
||||
"contents": [{"role": "user", "parts": [{"text": "hello"}]}],
|
||||
"systemInstruction": {"parts": [{"text": "You are terse."}]},
|
||||
"tools": [
|
||||
{
|
||||
"functionDeclarations": [
|
||||
{
|
||||
"name": "lookup",
|
||||
"description": "Find a record",
|
||||
"parameters": {"type": "object"},
|
||||
}
|
||||
]
|
||||
}
|
||||
],
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_google_ai_studio_count_tokens_rejects_image_system_without_request() -> None:
|
||||
peer: Final = _GeminiCountTokensPeer()
|
||||
client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(peer))
|
||||
|
||||
with patch.object(gemini_count_tokens_handler, "get_async_httpx_client", return_value=client):
|
||||
result: Final = await GoogleAIStudioTokenCounter().count_tokens(
|
||||
model_to_use="gemini-2.5-flash",
|
||||
messages=None,
|
||||
contents=[{"role": "user", "parts": [{"text": "hello"}]}],
|
||||
deployment={"litellm_params": {"api_key": "test-key", "api_base": "https://count-tokens.test"}},
|
||||
request_model="gemini/gemini-2.5-flash",
|
||||
system=[
|
||||
{"type": "text", "text": "You are terse."},
|
||||
{
|
||||
"type": "image",
|
||||
"source": {"type": "base64", "media_type": "image/png", "data": "iVBORw0KGgo="},
|
||||
},
|
||||
],
|
||||
)
|
||||
|
||||
assert result is not None
|
||||
assert result.error is True
|
||||
assert result.status_code == 400
|
||||
assert peer.request_body is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_google_ai_studio_count_tokens_rejects_an_image_url_it_refuses_to_fetch_without_request(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "user_url_validation", True)
|
||||
peer: Final = _GeminiCountTokensPeer()
|
||||
client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(peer))
|
||||
|
||||
with patch.object(gemini_count_tokens_handler, "get_async_httpx_client", return_value=client):
|
||||
result: Final = await GoogleAIStudioTokenCounter().count_tokens(
|
||||
model_to_use="gemini-2.5-flash",
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [{"type": "image", "source": {"type": "url", "url": "http://127.0.0.1:9/image.png"}}],
|
||||
}
|
||||
],
|
||||
contents=None,
|
||||
deployment={"litellm_params": {"api_key": "test-key", "api_base": "https://count-tokens.test"}},
|
||||
request_model="gemini/gemini-2.5-flash",
|
||||
messages_format="anthropic",
|
||||
)
|
||||
|
||||
assert result is not None
|
||||
assert result.error is True
|
||||
assert result.status_code == 400
|
||||
assert "http://127.0.0.1:9/image.png" in (result.error_message or "")
|
||||
assert peer.request_body is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_google_ai_studio_count_tokens_answers_with_the_status_google_rejected_it_with() -> None:
|
||||
rejection: Final = '{"error": {"code": 400, "message": "Invalid function name", "status": "INVALID_ARGUMENT"}}'
|
||||
client: Final = AsyncHTTPHandler(
|
||||
transport=httpx.MockTransport(lambda request: httpx.Response(400, text=rejection, request=request))
|
||||
)
|
||||
|
||||
with patch.object(gemini_count_tokens_handler, "get_async_httpx_client", return_value=client):
|
||||
result: Final = await GoogleAIStudioTokenCounter().count_tokens(
|
||||
model_to_use="gemini-2.5-flash",
|
||||
messages=None,
|
||||
contents=[{"role": "user", "parts": [{"text": "hello"}]}],
|
||||
deployment={"litellm_params": {"api_key": "test-key", "api_base": "https://count-tokens.test"}},
|
||||
request_model="gemini/gemini-2.5-flash",
|
||||
)
|
||||
|
||||
assert result is not None
|
||||
assert result.error is True
|
||||
assert result.status_code == 400
|
||||
assert "Invalid function name" in (result.error_message or "")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("system_instruction", "tools", "expected_body"),
|
||||
(
|
||||
(
|
||||
{"parts": [{"text": "Follow the system rule."}]},
|
||||
[{"functionDeclarations": [{"name": "lookup", "parameters": {"type": "OBJECT"}}]}],
|
||||
{
|
||||
"generateContentRequest": {
|
||||
"model": "models/gemini-4-argon",
|
||||
"contents": [{"role": "user", "parts": [{"text": "hello"}]}],
|
||||
"systemInstruction": {"parts": [{"text": "Follow the system rule."}]},
|
||||
"tools": [{"functionDeclarations": [{"name": "lookup", "parameters": {"type": "OBJECT"}}]}],
|
||||
}
|
||||
},
|
||||
),
|
||||
(
|
||||
None,
|
||||
[],
|
||||
{"contents": [{"role": "user", "parts": [{"text": "hello"}]}]},
|
||||
),
|
||||
),
|
||||
)
|
||||
async def test_gemini_count_tokens_sends_optional_request_fields(
|
||||
system_instruction: dict[str, JsonValue] | None,
|
||||
tools: list[dict[str, JsonValue]],
|
||||
expected_body: dict[str, JsonValue],
|
||||
) -> None:
|
||||
peer: Final = _GeminiCountTokensPeer()
|
||||
token_counter: Final = GeminiCountTokensHandler(client=AsyncHTTPHandler(transport=httpx.MockTransport(peer)))
|
||||
|
||||
result: Final = await token_counter.acount_tokens(
|
||||
contents=[{"role": "user", "parts": [{"text": "hello"}]}],
|
||||
model="gemini-4-argon",
|
||||
api_key="test-key",
|
||||
api_base="https://count-tokens.test",
|
||||
system_instruction=system_instruction,
|
||||
tools=tools,
|
||||
)
|
||||
|
||||
assert result == {"totalTokens": 11}
|
||||
assert peer.request_body == expected_body
|
||||
|
||||
|
||||
class TestGeminiModelInfo:
|
||||
|
|
@ -36,18 +228,10 @@ class TestGeminiModelInfo:
|
|||
# Test edge cases where model names end with characters from "models/"
|
||||
# These would be incorrectly processed if using strip("models/") instead of replace("models/", "")
|
||||
models = [
|
||||
{
|
||||
"name": "models/gemini-1.5-pro"
|
||||
}, # ends with 'o' - would become "gemini-1.5-pr" with strip()
|
||||
{
|
||||
"name": "models/test-model"
|
||||
}, # ends with 'l' - would become "gemini/test-mode" with strip()
|
||||
{
|
||||
"name": "models/custom-models"
|
||||
}, # ends with 's' - would become "gemini/custom-model" with strip()
|
||||
{
|
||||
"name": "models/demo"
|
||||
}, # ends with 'o' - would become "gemini/dem" with strip()
|
||||
{"name": "models/gemini-1.5-pro"}, # ends with 'o' - would become "gemini-1.5-pr" with strip()
|
||||
{"name": "models/test-model"}, # ends with 'l' - would become "gemini/test-mode" with strip()
|
||||
{"name": "models/custom-models"}, # ends with 's' - would become "gemini/custom-model" with strip()
|
||||
{"name": "models/demo"}, # ends with 'o' - would become "gemini/dem" with strip()
|
||||
]
|
||||
|
||||
result = gemini_model_info.process_model_name(models)
|
||||
|
|
@ -98,16 +282,10 @@ class TestGoogleAIStudioTokenCounter:
|
|||
token_counter = GoogleAIStudioTokenCounter()
|
||||
|
||||
# Test with gemini provider - should return True
|
||||
assert (
|
||||
token_counter.should_use_token_counting_api(LlmProviders.GEMINI.value)
|
||||
is True
|
||||
)
|
||||
assert token_counter.should_use_token_counting_api(LlmProviders.GEMINI.value) is True
|
||||
|
||||
# Test with other providers - should return False
|
||||
assert (
|
||||
token_counter.should_use_token_counting_api(LlmProviders.OPENAI.value)
|
||||
is False
|
||||
)
|
||||
assert token_counter.should_use_token_counting_api(LlmProviders.OPENAI.value) is False
|
||||
assert token_counter.should_use_token_counting_api("anthropic") is False
|
||||
assert token_counter.should_use_token_counting_api("vertex_ai") is False
|
||||
|
||||
|
|
@ -158,7 +336,10 @@ class TestGoogleAIStudioTokenCounter:
|
|||
|
||||
# Verify the mock was called correctly
|
||||
mock_acount_tokens.assert_called_once_with(
|
||||
model=model_to_use, contents=contents
|
||||
model=model_to_use,
|
||||
contents=contents,
|
||||
system_instruction=None,
|
||||
tools=None,
|
||||
)
|
||||
|
||||
def test_clean_contents_for_gemini_api_removes_id_field(self):
|
||||
|
|
@ -176,9 +357,7 @@ class TestGoogleAIStudioTokenCounter:
|
|||
"functionResponse": {
|
||||
"id": "read_many_files-1757526647518-730a691aac11c", # This should be removed
|
||||
"name": "read_many_files",
|
||||
"response": {
|
||||
"output": "No files matching the criteria were found or all were skipped."
|
||||
},
|
||||
"response": {"output": "No files matching the criteria were found or all were skipped."},
|
||||
}
|
||||
}
|
||||
],
|
||||
|
|
@ -187,9 +366,7 @@ class TestGoogleAIStudioTokenCounter:
|
|||
]
|
||||
|
||||
# Clean the contents
|
||||
cleaned_contents = token_counter._clean_contents_for_gemini_api(
|
||||
contents_with_id
|
||||
)
|
||||
cleaned_contents = token_counter._clean_contents_for_gemini_api(contents_with_id)
|
||||
|
||||
# Verify the 'id' field was removed
|
||||
function_response = cleaned_contents[1]["parts"][0]["functionResponse"]
|
||||
|
|
@ -198,8 +375,7 @@ class TestGoogleAIStudioTokenCounter:
|
|||
assert "response" in function_response
|
||||
assert function_response["name"] == "read_many_files"
|
||||
assert (
|
||||
function_response["response"]["output"]
|
||||
== "No files matching the criteria were found or all were skipped."
|
||||
function_response["response"]["output"] == "No files matching the criteria were found or all were skipped."
|
||||
)
|
||||
|
||||
def test_clean_contents_for_gemini_api_preserves_other_fields(self):
|
||||
|
|
@ -215,9 +391,7 @@ class TestGoogleAIStudioTokenCounter:
|
|||
]
|
||||
|
||||
# Clean the contents
|
||||
cleaned_contents = token_counter._clean_contents_for_gemini_api(
|
||||
contents_without_function_response
|
||||
)
|
||||
cleaned_contents = token_counter._clean_contents_for_gemini_api(contents_without_function_response)
|
||||
|
||||
# Verify the contents are unchanged
|
||||
assert cleaned_contents == contents_without_function_response
|
||||
|
|
|
|||
|
|
@ -0,0 +1,436 @@
|
|||
import json
|
||||
import operator
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import respx
|
||||
from pydantic import JsonValue, ValidationError
|
||||
|
||||
import litellm
|
||||
from litellm.llms.openai.responses.count_tokens.token_counter import OpenAITokenCounter
|
||||
from litellm.responses.utils import ResponsesAPIRequestUtils
|
||||
from litellm.types.utils import TokenCountMessagesFormat, TokenCountResponse
|
||||
|
||||
COUNT_URL: Final = "https://api.openai.com/v1/responses/input_tokens"
|
||||
|
||||
|
||||
async def _count(monkeypatch: pytest.MonkeyPatch, response: httpx.Response) -> TokenCountResponse | None:
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
with respx.mock:
|
||||
respx.post(COUNT_URL).mock(return_value=response)
|
||||
return await OpenAITokenCounter().count_tokens(
|
||||
model_to_use="gpt-4o",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
contents=None,
|
||||
deployment={"litellm_params": {"api_key": "sk-test"}},
|
||||
request_model="openai/gpt-4o",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("payload", "expected_total"),
|
||||
[
|
||||
({"input_tokens": 15}, 15),
|
||||
({"input_tokens": "15"}, 15),
|
||||
({"input_tokens": 15.0}, 15),
|
||||
({"input_tokens": True}, 1),
|
||||
({"input_tokens": -3}, -3),
|
||||
({}, 0),
|
||||
({"object": "response.input_tokens", "error": {"message": "429 rate limit"}}, 0),
|
||||
({"input_tokens": 7, "object": "response.input_tokens", "nested": {"a": [1, None]}}, 7),
|
||||
],
|
||||
)
|
||||
async def test_count_reports_upstream_input_tokens_and_keeps_the_raw_payload(
|
||||
monkeypatch: pytest.MonkeyPatch, payload: dict[str, object], expected_total: int
|
||||
) -> None:
|
||||
result: Final = await _count(monkeypatch, httpx.Response(200, json=payload))
|
||||
|
||||
assert result is not None
|
||||
assert repr(result.model_dump()) == repr(
|
||||
{
|
||||
"total_tokens": expected_total,
|
||||
"request_model": "openai/gpt-4o",
|
||||
"model_used": "gpt-4o",
|
||||
"tokenizer_type": "openai_api",
|
||||
"count_source": None,
|
||||
"original_response": payload,
|
||||
"error": False,
|
||||
"error_message": None,
|
||||
"status_code": None,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"rejected_count",
|
||||
[
|
||||
None,
|
||||
15.5,
|
||||
"Request timed out",
|
||||
"429 Too Many Requests",
|
||||
[15],
|
||||
[*range(403), "Request timed out"],
|
||||
{"value": 429},
|
||||
],
|
||||
)
|
||||
async def test_non_integer_input_tokens_becomes_an_error_response_carrying_the_validation_message(
|
||||
monkeypatch: pytest.MonkeyPatch, rejected_count: object
|
||||
) -> None:
|
||||
with pytest.raises(ValidationError) as rejection:
|
||||
TokenCountResponse.model_validate(
|
||||
{
|
||||
"total_tokens": rejected_count,
|
||||
"request_model": "openai/gpt-4o",
|
||||
"model_used": "gpt-4o",
|
||||
"tokenizer_type": "openai_api",
|
||||
}
|
||||
)
|
||||
|
||||
result: Final = await _count(monkeypatch, httpx.Response(200, json={"input_tokens": rejected_count}))
|
||||
|
||||
assert result is not None
|
||||
assert [error["loc"] for error in rejection.value.errors()] == [("total_tokens",)]
|
||||
assert repr(result.model_dump()) == repr(
|
||||
{
|
||||
"total_tokens": 0,
|
||||
"request_model": "openai/gpt-4o",
|
||||
"model_used": "gpt-4o",
|
||||
"tokenizer_type": "openai_api",
|
||||
"count_source": None,
|
||||
"original_response": None,
|
||||
"error": True,
|
||||
"error_message": str(rejection.value),
|
||||
"status_code": 500,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"payload",
|
||||
[[*({"input_tokens": position} for position in range(403)), "429"], "input_tokens", 15, 0, False],
|
||||
)
|
||||
async def test_count_payload_that_is_not_an_object_becomes_an_error_response_naming_the_failed_lookup(
|
||||
monkeypatch: pytest.MonkeyPatch, payload: object
|
||||
) -> None:
|
||||
with pytest.raises(AttributeError) as failed_lookup:
|
||||
operator.attrgetter("get")(payload)
|
||||
|
||||
result: Final = await _count(monkeypatch, httpx.Response(200, json=payload))
|
||||
|
||||
assert result is not None
|
||||
assert repr(result.model_dump()) == repr(
|
||||
{
|
||||
"total_tokens": 0,
|
||||
"request_model": "openai/gpt-4o",
|
||||
"model_used": "gpt-4o",
|
||||
"tokenizer_type": "openai_api",
|
||||
"count_source": None,
|
||||
"original_response": None,
|
||||
"error": True,
|
||||
"error_message": str(failed_lookup.value),
|
||||
"status_code": 500,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_null_count_payload_falls_back_to_local_counting(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
assert await _count(monkeypatch, httpx.Response(200, content=b"null")) is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("upstream_status", "upstream_text"),
|
||||
[(429, "rate limit exceeded"), (403, "forbidden"), (408, "Request timed out")],
|
||||
)
|
||||
async def test_upstream_error_status_becomes_an_error_response_with_that_status_and_text(
|
||||
monkeypatch: pytest.MonkeyPatch, upstream_status: int, upstream_text: str
|
||||
) -> None:
|
||||
result: Final = await _count(monkeypatch, httpx.Response(upstream_status, text=upstream_text))
|
||||
|
||||
assert result is not None
|
||||
assert repr(result.model_dump()) == repr(
|
||||
{
|
||||
"total_tokens": 0,
|
||||
"request_model": "openai/gpt-4o",
|
||||
"model_used": "gpt-4o",
|
||||
"tokenizer_type": "openai_api",
|
||||
"count_source": None,
|
||||
"original_response": None,
|
||||
"error": True,
|
||||
"error_message": upstream_text,
|
||||
"status_code": upstream_status,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
async def _sent_count_body(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
messages: list[dict[str, object]],
|
||||
system: object,
|
||||
messages_format: TokenCountMessagesFormat,
|
||||
output_format: JsonValue = None,
|
||||
) -> object:
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
with respx.mock:
|
||||
route: Final = respx.post(COUNT_URL).mock(return_value=httpx.Response(200, json={"input_tokens": 1}))
|
||||
await OpenAITokenCounter().count_tokens(
|
||||
model_to_use="gpt-4o",
|
||||
messages=messages,
|
||||
contents=None,
|
||||
deployment={"litellm_params": {"api_key": "sk-test"}},
|
||||
request_model="openai/gpt-4o",
|
||||
system=system,
|
||||
messages_format=messages_format,
|
||||
output_format=output_format,
|
||||
)
|
||||
return json.loads(route.calls.last.request.content)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_anthropic_format_sends_the_document_image_and_tool_turns_to_the_count_endpoint(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
messages: Final[list[dict[str, object]]] = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "Summarize."},
|
||||
{"type": "document", "source": {"type": "base64", "media_type": "application/pdf", "data": "JVBE"}},
|
||||
{"type": "image", "source": {"type": "url", "url": "https://example.com/page.png"}},
|
||||
],
|
||||
},
|
||||
{"role": "assistant", "content": [{"type": "tool_use", "id": "toolu_01", "name": "lookup", "input": {}}]},
|
||||
{"role": "user", "content": [{"type": "tool_result", "tool_use_id": "toolu_01", "content": "found"}]},
|
||||
]
|
||||
|
||||
sent: Final = await _sent_count_body(
|
||||
monkeypatch, messages, [{"type": "text", "text": "Be brief."}], messages_format="anthropic"
|
||||
)
|
||||
|
||||
assert sent == {
|
||||
"model": "gpt-4o",
|
||||
"input": [
|
||||
{
|
||||
"type": "message",
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "input_text", "text": "Summarize."},
|
||||
{"type": "input_file", "filename": "document.pdf", "file_data": "data:application/pdf;base64,JVBE"},
|
||||
{"type": "input_image", "image_url": "https://example.com/page.png"},
|
||||
],
|
||||
},
|
||||
{"type": "function_call", "call_id": "toolu_01", "name": "lookup", "arguments": "{}"},
|
||||
{"type": "function_call_output", "call_id": "toolu_01", "output": "found"},
|
||||
],
|
||||
"instructions": "Be brief.",
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_openai_format_keeps_the_chat_conversion(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
messages: Final[list[dict[str, object]]] = [
|
||||
{"role": "system", "content": "Be brief."},
|
||||
{"role": "user", "content": [{"type": "text", "text": "Summarize."}]},
|
||||
{"role": "tool", "tool_call_id": "call_1", "content": "found"},
|
||||
]
|
||||
|
||||
sent: Final = await _sent_count_body(monkeypatch, messages, system=None, messages_format="openai")
|
||||
|
||||
assert sent == {
|
||||
"model": "gpt-4o",
|
||||
"input": [
|
||||
{"role": "user", "content": "Summarize."},
|
||||
{"type": "function_call_output", "call_id": "call_1", "output": "found"},
|
||||
],
|
||||
"instructions": "Be brief.",
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_anthropic_format_sends_affinity_wrapped_reasoning_unwrapped_as_inference_does(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
wrapped: Final = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id("gAAAAoriginal", "deployment-1") # pyright: ignore[reportPrivateUsage] # the wrap inference applies
|
||||
messages: Final[list[dict[str, object]]] = [
|
||||
{"role": "user", "content": "hi"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{"type": "redacted_thinking", "data": f"litellm_encrypted_reasoning:{wrapped}"},
|
||||
{"type": "text", "text": "OK"},
|
||||
],
|
||||
},
|
||||
{"role": "user", "content": "again"},
|
||||
]
|
||||
|
||||
sent: Final = await _sent_count_body(monkeypatch, messages, system=None, messages_format="anthropic")
|
||||
|
||||
assert sent == {
|
||||
"model": "gpt-4o",
|
||||
"input": [
|
||||
{"type": "message", "role": "user", "content": [{"type": "input_text", "text": "hi"}]},
|
||||
{"type": "reasoning", "summary": [], "encrypted_content": "gAAAAoriginal"},
|
||||
{"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "OK"}]},
|
||||
{"type": "message", "role": "user", "content": [{"type": "input_text", "text": "again"}]},
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_anthropic_format_counts_the_structured_output_schema(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
schema: Final[dict[str, JsonValue]] = {"type": "object", "properties": {"answer": {"type": "string"}}}
|
||||
|
||||
sent: Final = await _sent_count_body(
|
||||
monkeypatch,
|
||||
[{"role": "user", "content": "Answer in JSON."}],
|
||||
system=None,
|
||||
messages_format="anthropic",
|
||||
output_format={"type": "json_schema", "schema": schema},
|
||||
)
|
||||
|
||||
assert sent == {
|
||||
"model": "gpt-4o",
|
||||
"input": [{"type": "message", "role": "user", "content": [{"type": "input_text", "text": "Answer in JSON."}]}],
|
||||
"text": {"format": {"type": "json_schema", "name": "structured_output", "schema": schema, "strict": False}},
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_anthropic_format_sends_a_replayed_web_search_result_as_its_text(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
snippet: Final = "LiteLLM ships a weekly release candidate and a stable line."
|
||||
messages: Final[list[dict[str, object]]] = [
|
||||
{"role": "user", "content": "What is new?"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{"type": "server_tool_use", "id": "srvtoolu_01", "name": "web_search", "input": {"query": "news"}},
|
||||
{
|
||||
"type": "web_search_tool_result",
|
||||
"tool_use_id": "srvtoolu_01",
|
||||
"content": [
|
||||
{
|
||||
"type": "web_search_result",
|
||||
"url": "https://docs.litellm.ai/release_notes",
|
||||
"title": "Release notes",
|
||||
"encrypted_content": "",
|
||||
"snippet": snippet,
|
||||
}
|
||||
],
|
||||
},
|
||||
],
|
||||
},
|
||||
{"role": "user", "content": "Thanks."},
|
||||
]
|
||||
|
||||
sent: Final = json.dumps(await _sent_count_body(monkeypatch, messages, system=None, messages_format="anthropic"))
|
||||
|
||||
assert snippet in sent
|
||||
assert "https://docs.litellm.ai/release_notes" in sent
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_anthropic_format_message_without_a_role_answers_400_without_calling_openai(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
with respx.mock:
|
||||
route: Final = respx.post(COUNT_URL).mock(return_value=httpx.Response(200, json={"input_tokens": 1}))
|
||||
result: Final = await OpenAITokenCounter().count_tokens(
|
||||
model_to_use="gpt-4o",
|
||||
messages=[{"role": "user", "content": "hi"}, {"content": "no role"}],
|
||||
contents=None,
|
||||
deployment={"litellm_params": {"api_key": "sk-test"}},
|
||||
request_model="openai/gpt-4o",
|
||||
messages_format="anthropic",
|
||||
)
|
||||
|
||||
assert route.called is False
|
||||
assert result is not None
|
||||
assert (result.error, result.status_code, result.error_message) == (True, 400, "messages.1.role: Field required")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_openai_format_sends_a_list_system_as_its_text(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
sent: Final = await _sent_count_body(
|
||||
monkeypatch,
|
||||
[{"role": "user", "content": "hi"}],
|
||||
system=[{"type": "text", "text": "Be brief."}, {"type": "text", "text": "Answer in English."}],
|
||||
messages_format="openai",
|
||||
)
|
||||
|
||||
assert sent == {
|
||||
"model": "gpt-4o",
|
||||
"input": [{"role": "user", "content": "hi"}],
|
||||
"instructions": "Be brief.\nAnswer in English.",
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_undecodable_count_body_becomes_a_500_error_response_carrying_the_decoder_message(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
with pytest.raises(json.JSONDecodeError) as undecodable:
|
||||
json.loads("not json")
|
||||
|
||||
result: Final = await _count(monkeypatch, httpx.Response(200, text="not json"))
|
||||
|
||||
assert result is not None
|
||||
assert repr(result.model_dump()) == repr(
|
||||
{
|
||||
"total_tokens": 0,
|
||||
"request_model": "openai/gpt-4o",
|
||||
"model_used": "gpt-4o",
|
||||
"tokenizer_type": "openai_api",
|
||||
"count_source": None,
|
||||
"original_response": None,
|
||||
"error": True,
|
||||
"error_message": f"CountTokens processing error: {undecodable.value}",
|
||||
"status_code": 500,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_count_uses_the_global_openai_key_chat_would_use(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
monkeypatch.setattr(litellm, "openai_key", "sk-from-litellm-settings")
|
||||
monkeypatch.delenv("OPENAI_API_KEY", raising=False)
|
||||
with respx.mock:
|
||||
route: Final = respx.post(COUNT_URL).mock(return_value=httpx.Response(200, json={"input_tokens": 9}))
|
||||
result: Final = await OpenAITokenCounter().count_tokens(
|
||||
model_to_use="gpt-4o",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
contents=None,
|
||||
deployment={"litellm_params": {}},
|
||||
request_model="openai/gpt-4o",
|
||||
)
|
||||
assert result is not None
|
||||
assert result.total_tokens == 9
|
||||
assert route.calls.last.request.headers["authorization"] == "Bearer sk-from-litellm-settings"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_count_sends_the_global_key_only_to_the_global_api_base(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
monkeypatch.setattr(litellm, "openai_key", "sk-private-gateway")
|
||||
monkeypatch.setattr(litellm, "api_base", "https://gateway.internal.example/v1")
|
||||
monkeypatch.delenv("OPENAI_API_KEY", raising=False)
|
||||
with respx.mock:
|
||||
route: Final = respx.post("https://gateway.internal.example/v1/responses/input_tokens").mock(
|
||||
return_value=httpx.Response(200, json={"input_tokens": 4})
|
||||
)
|
||||
result: Final = await OpenAITokenCounter().count_tokens(
|
||||
model_to_use="gpt-4o",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
contents=None,
|
||||
deployment={"litellm_params": {}},
|
||||
request_model="openai/gpt-4o",
|
||||
)
|
||||
assert result is not None
|
||||
assert result.total_tokens == 4
|
||||
assert route.calls.last.request.headers["authorization"] == "Bearer sk-private-gateway"
|
||||
|
|
@ -1,6 +1,10 @@
|
|||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.llms.anthropic.pass_through.responses_adapters.transformation import (
|
||||
LiteLLMAnthropicToResponsesAPIAdapter,
|
||||
)
|
||||
from litellm.llms.openai.responses.count_tokens.transformation import (
|
||||
OpenAICountTokensConfig,
|
||||
)
|
||||
|
|
@ -446,3 +450,108 @@ def test_get_required_headers():
|
|||
|
||||
assert headers["Authorization"] == "Bearer sk-test-key"
|
||||
assert headers["Content-Type"] == "application/json"
|
||||
|
||||
|
||||
PDF_BASE64: Final = "JVBERi0xLjcK"
|
||||
PNG_BASE64: Final = "iVBORw0KGgo="
|
||||
|
||||
ANTHROPIC_AGENT_TURN: Final[list[dict[str, object]]] = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "Summarize the indemnity terms."},
|
||||
{
|
||||
"type": "document",
|
||||
"title": "contract.pdf",
|
||||
"source": {"type": "base64", "media_type": "application/pdf", "data": PDF_BASE64},
|
||||
},
|
||||
{"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": PNG_BASE64}},
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{"type": "text", "text": "Let me pull clause 14."},
|
||||
{"type": "tool_use", "id": "toolu_01", "name": "lookup_clause", "input": {"id": "14"}},
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [{"type": "tool_result", "tool_use_id": "toolu_01", "content": "14. Indemnification."}],
|
||||
},
|
||||
]
|
||||
|
||||
ANTHROPIC_TOOLS: Final[list[dict[str, object]]] = [
|
||||
{
|
||||
"name": "lookup_clause",
|
||||
"description": "Fetch a clause by id",
|
||||
"input_schema": {"type": "object", "properties": {"id": {"type": "string"}}},
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"system",
|
||||
[None, "You review contracts.", [{"type": "text", "text": "You review contracts."}]],
|
||||
)
|
||||
def test_anthropic_messages_convert_exactly_as_v1_messages_inference_converts_them(system: object) -> None:
|
||||
inference_request: Final[dict[str, object]] = {
|
||||
"model": "gpt-4o",
|
||||
"messages": ANTHROPIC_AGENT_TURN,
|
||||
"tools": ANTHROPIC_TOOLS,
|
||||
**({"system": system} if system is not None else {}),
|
||||
}
|
||||
inference: Final = LiteLLMAnthropicToResponsesAPIAdapter().translate_request(
|
||||
inference_request, include_encrypted_reasoning=False
|
||||
)
|
||||
|
||||
counted: Final = OpenAICountTokensConfig.anthropic_messages_to_responses_input(
|
||||
model="gpt-4o", messages=ANTHROPIC_AGENT_TURN, system=system, tools=ANTHROPIC_TOOLS
|
||||
)
|
||||
|
||||
assert counted == (
|
||||
tuple(inference["input"]),
|
||||
inference.get("instructions"),
|
||||
tuple(inference["tools"]),
|
||||
inference.get("text"),
|
||||
)
|
||||
|
||||
|
||||
def test_anthropic_messages_keep_every_block_the_chat_converter_drops() -> None:
|
||||
input_items, _, _, _ = OpenAICountTokensConfig.anthropic_messages_to_responses_input(
|
||||
model="gpt-4o", messages=ANTHROPIC_AGENT_TURN, system=None, tools=None
|
||||
)
|
||||
|
||||
assert input_items == (
|
||||
{
|
||||
"type": "message",
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "input_text", "text": "Summarize the indemnity terms."},
|
||||
{
|
||||
"type": "input_file",
|
||||
"filename": "contract.pdf",
|
||||
"file_data": f"data:application/pdf;base64,{PDF_BASE64}",
|
||||
},
|
||||
{"type": "input_image", "image_url": f"data:image/png;base64,{PNG_BASE64}"},
|
||||
],
|
||||
},
|
||||
{"type": "function_call", "call_id": "toolu_01", "name": "lookup_clause", "arguments": '{"id": "14"}'},
|
||||
{
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"content": [{"type": "output_text", "text": "Let me pull clause 14."}],
|
||||
},
|
||||
{"type": "function_call_output", "call_id": "toolu_01", "output": "14. Indemnification."},
|
||||
)
|
||||
|
||||
|
||||
def test_anthropic_list_system_becomes_joined_instruction_text() -> None:
|
||||
_, instructions, _, _ = OpenAICountTokensConfig.anthropic_messages_to_responses_input(
|
||||
model="gpt-4o",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
system=[{"type": "text", "text": "First rule."}, {"type": "text", "text": "Second rule."}],
|
||||
tools=None,
|
||||
)
|
||||
|
||||
assert instructions == "First rule.\nSecond rule."
|
||||
|
|
|
|||
|
|
@ -630,3 +630,56 @@ class TestClientsideBaseOverride:
|
|||
)
|
||||
is None
|
||||
)
|
||||
|
||||
|
||||
class TestCountTokens:
|
||||
@pytest.mark.asyncio
|
||||
@respx.mock
|
||||
async def test_count_authenticates_with_the_workload_identity_bearer(
|
||||
self, wif_env: OpenAIWorkloadIdentityConfig, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
from litellm.llms.openai.responses.count_tokens.token_counter import OpenAITokenCounter
|
||||
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
monkeypatch.setattr(litellm, "api_key", None)
|
||||
monkeypatch.setattr(litellm, "openai_key", None)
|
||||
mock_token_exchange()
|
||||
count_route: Final = respx.post("https://api.openai.com/v1/responses/input_tokens").mock(
|
||||
return_value=httpx.Response(200, json={"input_tokens": 12})
|
||||
)
|
||||
result: Final = await OpenAITokenCounter().count_tokens(
|
||||
model_to_use="gpt-4o",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
contents=None,
|
||||
deployment={"litellm_params": {}},
|
||||
request_model="openai/gpt-4o",
|
||||
)
|
||||
assert result is not None
|
||||
assert result.error is False
|
||||
assert result.total_tokens == 12
|
||||
assert count_route.calls.last.request.headers["authorization"] == "Bearer exchanged-bearer-token"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@respx.mock
|
||||
async def test_failed_mint_becomes_a_count_error_response(
|
||||
self, wif_env: OpenAIWorkloadIdentityConfig, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
from litellm.llms.openai.responses.count_tokens.token_counter import OpenAITokenCounter
|
||||
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
monkeypatch.setattr(litellm, "api_key", None)
|
||||
monkeypatch.setattr(litellm, "openai_key", None)
|
||||
respx.post(TOKEN_EXCHANGE_URL).mock(return_value=httpx.Response(400, json={"error": "invalid_grant"}))
|
||||
count_route: Final = respx.post("https://api.openai.com/v1/responses/input_tokens")
|
||||
result: Final = await OpenAITokenCounter().count_tokens(
|
||||
model_to_use="gpt-4o",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
contents=None,
|
||||
deployment={"litellm_params": {}},
|
||||
request_model="openai/gpt-4o",
|
||||
)
|
||||
assert result is not None
|
||||
assert result.error is True
|
||||
assert result.status_code == 401
|
||||
assert result.total_tokens == 0
|
||||
assert not count_route.called
|
||||
|
|
|
|||
0
tests/unit/llms/vertex_ai/count_tokens/__init__.py
Normal file
0
tests/unit/llms/vertex_ai/count_tokens/__init__.py
Normal file
507
tests/unit/llms/vertex_ai/count_tokens/test_context.py
Normal file
507
tests/unit/llms/vertex_ai/count_tokens/test_context.py
Normal file
|
|
@ -0,0 +1,507 @@
|
|||
import asyncio
|
||||
import base64
|
||||
import time
|
||||
import uuid
|
||||
from collections.abc import Generator
|
||||
from contextlib import contextmanager
|
||||
from copy import deepcopy
|
||||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||||
from threading import Thread
|
||||
from typing import Final, Literal
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
|
||||
import litellm
|
||||
from litellm.llms.anthropic.pass_through.adapters.transformation import (
|
||||
LiteLLMAnthropicMessagesAdapter,
|
||||
truncate_tool_name,
|
||||
)
|
||||
from litellm.llms.vertex_ai.count_tokens.context import (
|
||||
VertexGeminiCountContext,
|
||||
VertexGeminiCountContextError,
|
||||
gemini_count_contents,
|
||||
inference_request_view,
|
||||
vertex_gemini_count_context,
|
||||
)
|
||||
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexGeminiConfig
|
||||
|
||||
_SCHEMA: Final[JsonValue] = {
|
||||
"type": "object",
|
||||
"properties": {},
|
||||
"additionalProperties": False,
|
||||
}
|
||||
_FUNCTION_DECLARATION: Final = {
|
||||
"functionDeclarations": [
|
||||
{
|
||||
"name": "lookup",
|
||||
"description": "Look up a value",
|
||||
"parameters": {"type": "object"},
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"tool",
|
||||
(
|
||||
{
|
||||
"name": "lookup",
|
||||
"description": "Look up a value",
|
||||
"input_schema": _SCHEMA,
|
||||
},
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "lookup",
|
||||
"description": "Look up a value",
|
||||
"parameters": _SCHEMA,
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": "function",
|
||||
"name": "lookup",
|
||||
"description": "Look up a value",
|
||||
"parameters": _SCHEMA,
|
||||
},
|
||||
),
|
||||
ids=("anthropic", "openai-chat", "responses"),
|
||||
)
|
||||
def test_function_tool_shapes_map_to_vertex_declarations(tool: JsonValue) -> None:
|
||||
context: Final = vertex_gemini_count_context(system=None, tools=[tool])
|
||||
|
||||
assert isinstance(context, VertexGeminiCountContext)
|
||||
assert context.tools == (_FUNCTION_DECLARATION,)
|
||||
assert _SCHEMA["additionalProperties"] is False
|
||||
|
||||
|
||||
def test_long_anthropic_tool_name_is_truncated_for_vertex() -> None:
|
||||
name: Final = "mcp__docs__" + "a" * 70
|
||||
context: Final = vertex_gemini_count_context(
|
||||
system=None,
|
||||
tools=[{"name": name, "description": "Look up a value", "input_schema": _SCHEMA}],
|
||||
)
|
||||
|
||||
assert isinstance(context, VertexGeminiCountContext)
|
||||
assert context.tools == (
|
||||
{
|
||||
"functionDeclarations": [
|
||||
{
|
||||
"name": truncate_tool_name(name),
|
||||
"description": "Look up a value",
|
||||
"parameters": {"type": "object"},
|
||||
}
|
||||
]
|
||||
},
|
||||
)
|
||||
assert len(truncate_tool_name(name)) == 64
|
||||
|
||||
|
||||
def test_64_character_anthropic_tool_name_is_unchanged_for_vertex() -> None:
|
||||
name: Final = "mcp__docs__" + "a" * 53
|
||||
context: Final = vertex_gemini_count_context(
|
||||
system=None,
|
||||
tools=[{"name": name, "input_schema": _SCHEMA}],
|
||||
)
|
||||
|
||||
assert isinstance(context, VertexGeminiCountContext)
|
||||
assert context.tools == ({"functionDeclarations": [{"name": name, "parameters": {"type": "object"}}]},)
|
||||
|
||||
|
||||
def test_long_openai_chat_tool_name_is_unchanged_for_vertex() -> None:
|
||||
name: Final = "mcp__docs__" + "a" * 70
|
||||
context: Final = vertex_gemini_count_context(
|
||||
system=None,
|
||||
tools=[{"type": "function", "function": {"name": name, "parameters": _SCHEMA}}],
|
||||
)
|
||||
|
||||
assert isinstance(context, VertexGeminiCountContext)
|
||||
assert context.tools == ({"functionDeclarations": [{"name": name, "parameters": {"type": "object"}}]},)
|
||||
|
||||
|
||||
def test_unknown_tools_are_omitted() -> None:
|
||||
context: Final = vertex_gemini_count_context(
|
||||
system=None,
|
||||
tools=[{"type": "web_search_preview"}, {"type": "computer_use"}],
|
||||
)
|
||||
|
||||
assert isinstance(context, VertexGeminiCountContext)
|
||||
assert context.tools is None
|
||||
|
||||
|
||||
def test_text_block_system_instruction_keeps_one_part_per_block() -> None:
|
||||
context: Final = vertex_gemini_count_context(
|
||||
system=[{"type": "text", "text": "First instruction."}, {"type": "text", "text": "Second instruction."}],
|
||||
tools=None,
|
||||
)
|
||||
|
||||
assert isinstance(context, VertexGeminiCountContext)
|
||||
assert context.system_instruction == {"parts": ({"text": "First instruction."}, {"text": "Second instruction."})}
|
||||
|
||||
|
||||
def test_string_system_instruction_maps_to_one_part() -> None:
|
||||
context: Final = vertex_gemini_count_context(system="Follow this instruction.", tools=None)
|
||||
|
||||
assert isinstance(context, VertexGeminiCountContext)
|
||||
assert context.system_instruction == {"parts": ({"text": "Follow this instruction."},)}
|
||||
|
||||
|
||||
def test_non_text_system_block_returns_an_error_value() -> None:
|
||||
context: Final = vertex_gemini_count_context(
|
||||
system=[{"type": "image", "source": {"type": "base64", "data": "AAAA"}}],
|
||||
tools=None,
|
||||
)
|
||||
|
||||
assert isinstance(context, VertexGeminiCountContextError)
|
||||
assert "text blocks" in context.error_message
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("system", "tools"), ((None, None), ("", None), ([], [])))
|
||||
def test_empty_context_is_omitted(system: JsonValue | None, tools: JsonValue | None) -> None:
|
||||
context: Final = vertex_gemini_count_context(system=system, tools=tools)
|
||||
|
||||
assert isinstance(context, VertexGeminiCountContext)
|
||||
assert context.system_instruction is None
|
||||
assert context.tools is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"tool",
|
||||
(
|
||||
{"description": "Look up a value", "input_schema": _SCHEMA},
|
||||
{"type": "bash_20250124", "name": "bash"},
|
||||
),
|
||||
ids=("unnamed", "anthropic-hosted"),
|
||||
)
|
||||
def test_anthropic_tools_count_the_same_declarations_inference_sends(tool: dict[str, JsonValue]) -> None:
|
||||
translated_tools, _ = LiteLLMAnthropicMessagesAdapter().translate_anthropic_tools_to_openai([deepcopy(tool)]) # pyright: ignore[reportArgumentType] # raw Anthropic tool dict
|
||||
inference_tools: Final = VertexGeminiConfig()._map_function( # pyright: ignore[reportPrivateUsage] # the tool mapping inference uses
|
||||
value=[dict(translated) for translated in translated_tools],
|
||||
optional_params={},
|
||||
)
|
||||
|
||||
context: Final = vertex_gemini_count_context(system=None, tools=[tool])
|
||||
|
||||
assert isinstance(context, VertexGeminiCountContext)
|
||||
assert context.tools is not None
|
||||
assert [tool_value["functionDeclarations"] for tool_value in context.tools] == [
|
||||
inference_tool["function_declarations"] for inference_tool in inference_tools
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"tools",
|
||||
(
|
||||
"lookup",
|
||||
[1],
|
||||
[{"name": "lookup", "input_schema": "not-a-schema"}],
|
||||
[{"type": "function", "function": {"name": "lookup", "parameters": [1]}}],
|
||||
[{"name": "lookup", "input_schema": {"type": "object", "properties": {"query": "not-a-schema"}}}],
|
||||
),
|
||||
ids=("not-a-list", "non-object-tool", "string-schema", "list-parameters", "malformed-property"),
|
||||
)
|
||||
def test_invalid_tools_return_an_error_value(tools: JsonValue) -> None:
|
||||
context: Final = vertex_gemini_count_context(system=None, tools=tools)
|
||||
|
||||
assert isinstance(context, VertexGeminiCountContextError)
|
||||
assert "tools" in context.error_message
|
||||
|
||||
|
||||
_ANTHROPIC_CONVERSATION: Final = (
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "Summarize both"},
|
||||
{"type": "document", "source": {"type": "base64", "media_type": "application/pdf", "data": "JVBERi0="}},
|
||||
{"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "aW1hZ2U="}},
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [{"type": "tool_use", "id": "toolu_01", "name": "lookup", "input": {"query": "x"}}],
|
||||
},
|
||||
{"role": "user", "content": [{"type": "tool_result", "tool_use_id": "toolu_01", "content": "42"}]},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("custom_llm_provider", ("vertex_ai", "gemini"))
|
||||
async def test_anthropic_documents_images_and_tool_turns_become_gemini_parts(
|
||||
custom_llm_provider: Literal["vertex_ai", "gemini"],
|
||||
) -> None:
|
||||
contents: Final = await gemini_count_contents(
|
||||
_ANTHROPIC_CONVERSATION, "anthropic", "gemini-2.5-flash", custom_llm_provider, {}
|
||||
)
|
||||
|
||||
assert contents == [
|
||||
{
|
||||
"role": "user",
|
||||
"parts": [
|
||||
{"text": "Summarize both"},
|
||||
{"inline_data": {"data": "JVBERi0=", "mime_type": "application/pdf"}},
|
||||
{"inline_data": {"data": "aW1hZ2U=", "mime_type": "image/png"}},
|
||||
],
|
||||
},
|
||||
{"role": "model", "parts": [{"function_call": {"name": "lookup", "args": {"query": "x"}}}]},
|
||||
{"role": "user", "parts": [{"function_response": {"name": "lookup", "response": {"content": "42"}}}]},
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("custom_llm_provider", ("vertex_ai", "gemini"))
|
||||
async def test_openai_messages_keep_their_chat_conversion(custom_llm_provider: Literal["vertex_ai", "gemini"]) -> None:
|
||||
contents: Final = await gemini_count_contents(
|
||||
({"role": "user", "content": "hello"}, {"role": "assistant", "content": "hi"}),
|
||||
"openai",
|
||||
"gemini-2.5-flash",
|
||||
custom_llm_provider,
|
||||
{},
|
||||
)
|
||||
|
||||
assert contents == [
|
||||
{"role": "user", "parts": [{"text": "hello"}]},
|
||||
{"role": "model", "parts": [{"text": "hi"}]},
|
||||
]
|
||||
|
||||
|
||||
_PNG_BYTES: Final = base64.b64decode(
|
||||
"iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg=="
|
||||
)
|
||||
_DOWNLOAD_SECONDS: Final = 0.5
|
||||
|
||||
|
||||
class _SlowPngHandler(BaseHTTPRequestHandler):
|
||||
def do_GET(self) -> None:
|
||||
time.sleep(_DOWNLOAD_SECONDS)
|
||||
self.send_response(200)
|
||||
self.send_header("Content-Type", "image/png")
|
||||
self.send_header("Content-Length", str(len(_PNG_BYTES)))
|
||||
self.end_headers()
|
||||
self.wfile.write(_PNG_BYTES)
|
||||
|
||||
def log_message(self, format: str, *args: object) -> None:
|
||||
return None
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _slow_png_server() -> Generator[str, None, None]:
|
||||
with ThreadingHTTPServer(("127.0.0.1", 0), _SlowPngHandler) as server:
|
||||
thread: Final = Thread(target=server.serve_forever, daemon=True)
|
||||
thread.start()
|
||||
try:
|
||||
yield f"http://127.0.0.1:{server.server_address[1]}"
|
||||
finally:
|
||||
server.shutdown()
|
||||
|
||||
|
||||
async def _event_loop_wake_delay(nap_seconds: float) -> float:
|
||||
started: Final = time.monotonic()
|
||||
await asyncio.sleep(nap_seconds)
|
||||
return time.monotonic() - started
|
||||
|
||||
|
||||
@pytest.mark.parametrize("custom_llm_provider", ("vertex_ai", "gemini"))
|
||||
async def test_remote_image_downloads_without_blocking_the_event_loop(
|
||||
custom_llm_provider: Literal["vertex_ai", "gemini"], monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "user_url_validation", False)
|
||||
with _slow_png_server() as base_url:
|
||||
image_url: Final = f"{base_url}/{uuid.uuid4().hex}.png"
|
||||
wake_delay, contents = await asyncio.gather(
|
||||
_event_loop_wake_delay(0.05),
|
||||
gemini_count_contents(
|
||||
(
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "Describe it"},
|
||||
{"type": "image", "source": {"type": "url", "url": image_url}},
|
||||
],
|
||||
},
|
||||
),
|
||||
"anthropic",
|
||||
"gemini-2.5-flash",
|
||||
custom_llm_provider,
|
||||
{},
|
||||
),
|
||||
)
|
||||
|
||||
assert contents == [
|
||||
{
|
||||
"role": "user",
|
||||
"parts": [
|
||||
{"text": "Describe it"},
|
||||
{"inline_data": {"data": base64.b64encode(_PNG_BYTES).decode(), "mime_type": "image/png"}},
|
||||
],
|
||||
}
|
||||
]
|
||||
assert wake_delay < _DOWNLOAD_SECONDS / 2
|
||||
|
||||
|
||||
class _MissingImageHandler(BaseHTTPRequestHandler):
|
||||
def do_GET(self) -> None:
|
||||
self.send_response(404)
|
||||
self.send_header("Content-Length", "0")
|
||||
self.end_headers()
|
||||
|
||||
def log_message(self, format: str, *args: object) -> None:
|
||||
return None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("custom_llm_provider", ("vertex_ai", "gemini"))
|
||||
async def test_image_that_cannot_be_fetched_returns_an_error_value(
|
||||
custom_llm_provider: Literal["vertex_ai", "gemini"], monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "user_url_validation", False)
|
||||
with ThreadingHTTPServer(("127.0.0.1", 0), _MissingImageHandler) as server:
|
||||
Thread(target=server.serve_forever, daemon=True).start()
|
||||
try:
|
||||
contents = await gemini_count_contents(
|
||||
(
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "image",
|
||||
"source": {
|
||||
"type": "url",
|
||||
"url": f"http://127.0.0.1:{server.server_address[1]}/{uuid.uuid4().hex}.png",
|
||||
},
|
||||
}
|
||||
],
|
||||
},
|
||||
),
|
||||
"anthropic",
|
||||
"gemini-2.5-flash",
|
||||
custom_llm_provider,
|
||||
{},
|
||||
)
|
||||
finally:
|
||||
server.shutdown()
|
||||
|
||||
assert isinstance(contents, VertexGeminiCountContextError)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"system", ("Be brief.", [{"type": "text", "text": "Be brief."}]), ids=("string-system", "text-block-system")
|
||||
)
|
||||
def test_anthropic_view_keeps_only_the_turns_after_a_client_compaction_block(system: JsonValue) -> None:
|
||||
view: Final = inference_request_view(
|
||||
(
|
||||
{"role": "user", "content": "first question " * 50},
|
||||
{"role": "assistant", "content": [{"type": "compaction", "content": "Summary of the earlier turns."}]},
|
||||
{"role": "user", "content": "latest question"},
|
||||
),
|
||||
system,
|
||||
"anthropic",
|
||||
)
|
||||
|
||||
assert view.messages == [{"role": "user", "content": "latest question"}]
|
||||
assert "Summary of the earlier turns." in str(view.system)
|
||||
assert "Be brief." in str(view.system)
|
||||
|
||||
|
||||
def test_anthropic_view_flattens_a_replayed_web_search_result_into_text() -> None:
|
||||
snippet: Final = "LiteLLM ships a weekly release candidate and a stable line."
|
||||
view: Final = inference_request_view(
|
||||
(
|
||||
{"role": "user", "content": "What is new?"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{"type": "server_tool_use", "id": "srvtoolu_01", "name": "web_search", "input": {"query": "news"}},
|
||||
{
|
||||
"type": "web_search_tool_result",
|
||||
"tool_use_id": "srvtoolu_01",
|
||||
"content": [
|
||||
{
|
||||
"type": "web_search_result",
|
||||
"url": "https://docs.litellm.ai/release_notes",
|
||||
"title": "Release notes",
|
||||
"encrypted_content": "",
|
||||
"snippet": snippet,
|
||||
}
|
||||
],
|
||||
},
|
||||
],
|
||||
},
|
||||
),
|
||||
None,
|
||||
"anthropic",
|
||||
)
|
||||
|
||||
assistant_blocks: Final = TypeAdapter(list[dict[str, str]]).validate_python(view.messages[1]["content"])
|
||||
assert [block["type"] for block in assistant_blocks] == ["text"]
|
||||
assert snippet in assistant_blocks[0]["text"]
|
||||
|
||||
|
||||
def test_openai_view_is_the_request_unchanged() -> None:
|
||||
messages: Final = ({"role": "user", "content": [{"type": "compaction", "content": "kept"}]},)
|
||||
|
||||
view: Final = inference_request_view(messages, "Be brief.", "openai")
|
||||
|
||||
assert view.messages == messages
|
||||
assert view.system == "Be brief."
|
||||
|
||||
|
||||
async def test_gcs_media_type_lookup_uses_the_deployment_vertex_project() -> None:
|
||||
with patch(
|
||||
"litellm.llms.vertex_ai.gemini.transformation._get_gcs_object_content_type",
|
||||
return_value="image/png",
|
||||
) as lookup:
|
||||
contents: Final = await gemini_count_contents(
|
||||
(
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "image", "source": {"type": "url", "url": "gs://bucket/image-without-extension"}}
|
||||
],
|
||||
},
|
||||
),
|
||||
"anthropic",
|
||||
"gemini-2.5-flash",
|
||||
"vertex_ai",
|
||||
{"vertex_project": "deployment-project"},
|
||||
)
|
||||
|
||||
assert lookup.call_args.kwargs.get("vertex_project") == "deployment-project"
|
||||
assert not isinstance(contents, VertexGeminiCountContextError)
|
||||
assert contents[0]["parts"][0] == {
|
||||
"file_data": {"mime_type": "image/png", "file_uri": "gs://bucket/image-without-extension"}
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("custom_llm_provider", ("vertex_ai", "gemini"))
|
||||
def test_anthropic_replayed_web_search_result_is_counted_as_its_text(
|
||||
custom_llm_provider: Literal["vertex_ai", "gemini"],
|
||||
) -> None:
|
||||
snippet: Final = "LiteLLM ships a weekly release candidate and a stable line."
|
||||
messages: Final = (
|
||||
{"role": "user", "content": "What is new?"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{"type": "server_tool_use", "id": "srvtoolu_01", "name": "web_search", "input": {"query": "news"}},
|
||||
{
|
||||
"type": "web_search_tool_result",
|
||||
"tool_use_id": "srvtoolu_01",
|
||||
"content": [
|
||||
{
|
||||
"type": "web_search_result",
|
||||
"url": "https://docs.litellm.ai/release_notes",
|
||||
"title": "Release notes",
|
||||
"encrypted_content": "",
|
||||
"snippet": snippet,
|
||||
}
|
||||
],
|
||||
},
|
||||
],
|
||||
},
|
||||
{"role": "user", "content": "Thanks."},
|
||||
)
|
||||
|
||||
view: Final = inference_request_view(messages, None, "anthropic")
|
||||
contents: Final = asyncio.run(
|
||||
gemini_count_contents(view.messages, "anthropic", "gemini-2.5-flash", custom_llm_provider, {})
|
||||
)
|
||||
|
||||
assert snippet in str(contents)
|
||||
|
|
@ -1,9 +1,18 @@
|
|||
from unittest.mock import patch
|
||||
from collections.abc import Iterator
|
||||
from contextlib import contextmanager
|
||||
from typing import Final
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
|
||||
import litellm
|
||||
from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
|
||||
from litellm.llms.gemini.count_tokens import handler as gemini_count_tokens_handler
|
||||
from litellm.llms.vertex_ai.common_utils import (
|
||||
VertexAITokenCounter,
|
||||
convert_anyof_null_to_nullable,
|
||||
get_vertex_location_from_url,
|
||||
get_vertex_project_id_from_url,
|
||||
|
|
@ -14,6 +23,110 @@ from litellm.llms.vertex_ai.common_utils import (
|
|||
validate_vertex_location,
|
||||
vertex_request_labels_from_litellm_params,
|
||||
)
|
||||
from litellm.llms.vertex_ai.count_tokens import handler as vertex_count_tokens_handler
|
||||
from litellm.types.utils import TokenCountMessagesFormat
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("messages_format", ("openai", "anthropic"))
|
||||
async def test_vertex_gemini_rejects_non_text_system_blocks_before_request(
|
||||
messages_format: TokenCountMessagesFormat,
|
||||
) -> None:
|
||||
result: Final = await VertexAITokenCounter().count_tokens(
|
||||
model_to_use="gemini-4-argon",
|
||||
messages=[{"role": "user", "content": "Count this"}],
|
||||
contents=None,
|
||||
request_model="vertex_ai/gemini-4-argon",
|
||||
system=[{"type": "image", "source": {"type": "base64", "data": "AAAA"}}],
|
||||
messages_format=messages_format,
|
||||
)
|
||||
|
||||
assert result is not None
|
||||
assert result.error is True
|
||||
assert result.status_code == 400
|
||||
assert "text blocks" in (result.error_message or "")
|
||||
|
||||
|
||||
class _VertexCountTokensPeer:
|
||||
def __init__(self, status_code: int) -> None:
|
||||
self.status_code: Final = status_code
|
||||
self.requests: tuple[httpx.Request, ...] = ()
|
||||
|
||||
def __call__(self, request: httpx.Request) -> httpx.Response:
|
||||
self.requests = (*self.requests, request)
|
||||
if self.status_code == 200:
|
||||
return httpx.Response(200, json={"totalTokens": 37}, request=request)
|
||||
return httpx.Response(
|
||||
self.status_code,
|
||||
json={"error": {"code": self.status_code, "message": "Invalid function name", "status": "INVALID_ARGUMENT"}},
|
||||
request=request,
|
||||
)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _vertex_count_tokens_peer(status_code: int) -> Iterator[_VertexCountTokensPeer]:
|
||||
peer: Final = _VertexCountTokensPeer(status_code)
|
||||
with (
|
||||
patch.object(
|
||||
gemini_count_tokens_handler,
|
||||
"get_async_httpx_client",
|
||||
return_value=AsyncHTTPHandler(transport=httpx.MockTransport(peer)),
|
||||
),
|
||||
patch.object(
|
||||
vertex_count_tokens_handler.VertexAITokenCounter,
|
||||
"_ensure_access_token_async",
|
||||
new=AsyncMock(return_value=("fake-token", "test-project")),
|
||||
),
|
||||
):
|
||||
yield peer
|
||||
|
||||
|
||||
_VERTEX_DEPLOYMENT: Final = {"litellm_params": {"vertex_project": "test-project", "vertex_location": "us-central1"}}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_vertex_gemini_count_answers_with_the_status_google_rejected_it_with() -> None:
|
||||
with _vertex_count_tokens_peer(400) as peer:
|
||||
result: Final = await VertexAITokenCounter().count_tokens(
|
||||
model_to_use="gemini-2.5-flash",
|
||||
messages=[{"role": "user", "content": "Hello"}],
|
||||
contents=None,
|
||||
deployment=_VERTEX_DEPLOYMENT,
|
||||
request_model="vertex_ai/gemini-2.5-flash",
|
||||
)
|
||||
|
||||
assert result is not None
|
||||
assert result.error is True
|
||||
assert result.status_code == 400
|
||||
assert "Invalid function name" in (result.error_message or "")
|
||||
assert len(peer.requests) == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_vertex_gemini_count_rejects_an_image_url_it_refuses_to_fetch_without_request(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "user_url_validation", True)
|
||||
with _vertex_count_tokens_peer(200) as peer:
|
||||
result: Final = await VertexAITokenCounter().count_tokens(
|
||||
model_to_use="gemini-2.5-flash",
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [{"type": "image", "source": {"type": "url", "url": "http://127.0.0.1:9/image.png"}}],
|
||||
}
|
||||
],
|
||||
contents=None,
|
||||
deployment=_VERTEX_DEPLOYMENT,
|
||||
request_model="vertex_ai/gemini-2.5-flash",
|
||||
messages_format="anthropic",
|
||||
)
|
||||
|
||||
assert result is not None
|
||||
assert result.error is True
|
||||
assert result.status_code == 400
|
||||
assert "http://127.0.0.1:9/image.png" in (result.error_message or "")
|
||||
assert peer.requests == ()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("location", ["us", "eu", "us-central1", "europe-west1", "global"])
|
||||
|
|
@ -147,15 +260,11 @@ async def test_get_supports_system_message():
|
|||
from litellm.llms.vertex_ai.common_utils import get_supports_system_message
|
||||
|
||||
# fine-tuned vertex gemini models will specifiy they are in the /gemini spec format
|
||||
result = get_supports_system_message(
|
||||
model="gemini/1234567890", custom_llm_provider="vertex_ai"
|
||||
)
|
||||
result = get_supports_system_message(model="gemini/1234567890", custom_llm_provider="vertex_ai")
|
||||
assert result == True
|
||||
|
||||
# non-fine-tuned vertex gemini models will not specifiy they are in the /gemini spec format
|
||||
result = get_supports_system_message(
|
||||
model="random-model-name", custom_llm_provider="vertex_ai"
|
||||
)
|
||||
result = get_supports_system_message(model="random-model-name", custom_llm_provider="vertex_ai")
|
||||
assert result == False
|
||||
|
||||
|
||||
|
|
@ -228,13 +337,9 @@ def test_build_vertex_schema():
|
|||
"properties": {
|
||||
"tags": {"items": {"type": "string"}, "type": "array"},
|
||||
"metadata": {"type": "object"},
|
||||
"callbacks": {
|
||||
"anyOf": [{"items": {}, "type": "array"}, {}, {"type": "null"}]
|
||||
},
|
||||
"callbacks": {"anyOf": [{"items": {}, "type": "array"}, {}, {"type": "null"}]},
|
||||
"run_name": {"type": "string"},
|
||||
"max_concurrency": {
|
||||
"anyOf": [{"type": "integer"}, {"type": "null"}]
|
||||
},
|
||||
"max_concurrency": {"anyOf": [{"type": "integer"}, {"type": "null"}]},
|
||||
"recursion_limit": {"type": "integer"},
|
||||
"configurable": {"type": "object"},
|
||||
"run_id": {
|
||||
|
|
@ -278,9 +383,7 @@ def test_build_vertex_schema():
|
|||
]
|
||||
},
|
||||
"run_name": {"type": "string"},
|
||||
"max_concurrency": {
|
||||
"anyOf": [{"type": "integer", "nullable": True}]
|
||||
},
|
||||
"max_concurrency": {"anyOf": [{"type": "integer", "nullable": True}]},
|
||||
"recursion_limit": {"type": "integer"},
|
||||
"configurable": {"type": "object"},
|
||||
"run_id": {"anyOf": [{"type": "string", "nullable": True}]},
|
||||
|
|
@ -685,9 +788,7 @@ def test_get_vertex_url_global_region(stream, expected_endpoint_suffix):
|
|||
({}, "us-east1", "us-east1"),
|
||||
],
|
||||
)
|
||||
def test_get_vertex_region_global_only_model(
|
||||
model_cost_entry, vertex_region, expected_region
|
||||
):
|
||||
def test_get_vertex_region_global_only_model(model_cost_entry, vertex_region, expected_region):
|
||||
"""Test get_vertex_region resolves region from model_cost supported_regions"""
|
||||
import litellm
|
||||
from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
|
||||
|
|
@ -699,9 +800,7 @@ def test_get_vertex_region_global_only_model(
|
|||
{"vertex_ai/test-model": model_cost_entry},
|
||||
clear=False,
|
||||
):
|
||||
result = vertex_base.get_vertex_region(
|
||||
vertex_region=vertex_region, model="test-model"
|
||||
)
|
||||
result = vertex_base.get_vertex_region(vertex_region=vertex_region, model="test-model")
|
||||
|
||||
assert result == expected_region
|
||||
|
||||
|
|
@ -819,9 +918,7 @@ def test_convert_schema_types_type_array_conversion():
|
|||
assert anyof_types[1]["type"] == "number"
|
||||
|
||||
# 4. Other properties preserved
|
||||
assert (
|
||||
input_schema["properties"]["studio"]["description"] == "The studio ID or name"
|
||||
)
|
||||
assert input_schema["properties"]["studio"]["description"] == "The studio ID or name"
|
||||
assert input_schema["required"] == ["studio"]
|
||||
|
||||
|
||||
|
|
@ -988,7 +1085,9 @@ def test_construct_target_url_with_version_prefix():
|
|||
),
|
||||
],
|
||||
)
|
||||
def test_construct_target_url_versionless_project_route_gets_api_version(requested_route: str, expected_url: str) -> None:
|
||||
def test_construct_target_url_versionless_project_route_gets_api_version(
|
||||
requested_route: str, expected_url: str
|
||||
) -> None:
|
||||
from litellm.llms.vertex_ai.common_utils import construct_target_url
|
||||
|
||||
target_url = construct_target_url(
|
||||
|
|
@ -1121,10 +1220,7 @@ def test_fix_enum_types():
|
|||
# 2. Non-string enums are removed
|
||||
assert "enum" not in input_schema["properties"]["maxLength"]
|
||||
assert "enum" not in input_schema["properties"]["enabled"]
|
||||
assert (
|
||||
"enum"
|
||||
not in input_schema["properties"]["nested"]["properties"]["innerNonStringEnum"]
|
||||
)
|
||||
assert "enum" not in input_schema["properties"]["nested"]["properties"]["innerNonStringEnum"]
|
||||
|
||||
# 3. anyOf with string type keeps enum, non-string removes it
|
||||
assert "enum" in input_schema["properties"]["anyOfField"]["anyOf"][0]
|
||||
|
|
@ -1178,8 +1274,6 @@ def test_get_token_url():
|
|||
assert "v1beta1" not in url
|
||||
assert "/v1/" in url
|
||||
|
||||
pass
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_vertex_ai_token_counter_routes_partner_models():
|
||||
|
|
@ -1457,6 +1551,49 @@ async def test_vertex_ai_token_counter_routes_gemini_models():
|
|||
assert result.total_tokens == 50
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_vertex_gemini_count_tokens_uses_request_system_over_deployment_system_instruction() -> None:
|
||||
class VertexCountTokensPeer:
|
||||
def __init__(self) -> None:
|
||||
self.request_body: dict[str, JsonValue] | None = None
|
||||
|
||||
def __call__(self, request: httpx.Request) -> httpx.Response:
|
||||
self.request_body = TypeAdapter(dict[str, JsonValue]).validate_json(request.content)
|
||||
return httpx.Response(200, json={"totalTokens": 37}, request=request)
|
||||
|
||||
peer: Final = VertexCountTokensPeer()
|
||||
client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(peer))
|
||||
|
||||
with (
|
||||
patch.object(gemini_count_tokens_handler, "get_async_httpx_client", return_value=client),
|
||||
patch.object(
|
||||
vertex_count_tokens_handler.VertexAITokenCounter,
|
||||
"_ensure_access_token_async",
|
||||
new=AsyncMock(return_value=("fake-token", "test-project")),
|
||||
),
|
||||
):
|
||||
result: Final = await VertexAITokenCounter().count_tokens(
|
||||
model_to_use="gemini-2.5-flash",
|
||||
messages=[{"role": "user", "content": "Hello"}],
|
||||
contents=None,
|
||||
deployment={
|
||||
"litellm_params": {
|
||||
"vertex_project": "test-project",
|
||||
"vertex_location": "us-central1",
|
||||
"system_instruction": {"parts": [{"text": "deployment-level"}]},
|
||||
}
|
||||
},
|
||||
request_model="vertex_ai/gemini-2.5-flash",
|
||||
system="You are terse.",
|
||||
)
|
||||
|
||||
assert result is not None
|
||||
assert result.error is False
|
||||
assert result.total_tokens == 37
|
||||
assert peer.request_body is not None
|
||||
assert peer.request_body["systemInstruction"] == {"parts": [{"text": "You are terse."}]}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_vertex_ai_token_counter_converts_messages_to_contents_for_gemini():
|
||||
"""
|
||||
|
|
@ -1471,9 +1608,7 @@ async def test_vertex_ai_token_counter_converts_messages_to_contents_for_gemini(
|
|||
|
||||
token_counter = VertexAITokenCounter()
|
||||
|
||||
with patch(
|
||||
"litellm.llms.vertex_ai.count_tokens.handler.VertexAITokenCounter.acount_tokens"
|
||||
) as mock_acount_tokens:
|
||||
with patch("litellm.llms.vertex_ai.count_tokens.handler.VertexAITokenCounter.acount_tokens") as mock_acount_tokens:
|
||||
mock_acount_tokens.return_value = {
|
||||
"totalTokens": 42,
|
||||
"tokenizer_used": "gemini",
|
||||
|
|
@ -1515,9 +1650,7 @@ async def test_vertex_ai_token_counter_returns_none_when_api_omits_total_tokens(
|
|||
|
||||
token_counter = VertexAITokenCounter()
|
||||
|
||||
with patch(
|
||||
"litellm.llms.vertex_ai.count_tokens.handler.VertexAITokenCounter.acount_tokens"
|
||||
) as mock_acount_tokens:
|
||||
with patch("litellm.llms.vertex_ai.count_tokens.handler.VertexAITokenCounter.acount_tokens") as mock_acount_tokens:
|
||||
mock_acount_tokens.return_value = {"tokenizer_used": "gemini"}
|
||||
|
||||
result = await token_counter.count_tokens(
|
||||
|
|
@ -1560,9 +1693,7 @@ async def test_vertex_ai_partner_model_detection():
|
|||
# Test Minimax models
|
||||
assert VertexAIPartnerModels.is_vertex_partner_model("minimaxai/minimax-m2-maas")
|
||||
# Test Moonshot models
|
||||
assert VertexAIPartnerModels.is_vertex_partner_model(
|
||||
"moonshotai/kimi-k2-thinking-maas"
|
||||
)
|
||||
assert VertexAIPartnerModels.is_vertex_partner_model("moonshotai/kimi-k2-thinking-maas")
|
||||
|
||||
# Test Gemini models (should NOT be detected as partner model)
|
||||
assert not VertexAIPartnerModels.is_vertex_partner_model("gemini-1.5-pro")
|
||||
|
|
@ -1593,9 +1724,7 @@ def test_vertex_ai_moonshot_uses_openai_handler():
|
|||
VertexAIPartnerModels,
|
||||
)
|
||||
|
||||
assert VertexAIPartnerModels.should_use_openai_handler(
|
||||
"moonshotai/kimi-k2-thinking-maas"
|
||||
)
|
||||
assert VertexAIPartnerModels.should_use_openai_handler("moonshotai/kimi-k2-thinking-maas")
|
||||
|
||||
|
||||
def test_vertex_ai_zai_uses_openai_handler():
|
||||
|
|
@ -1630,9 +1759,7 @@ def test_vertex_ai_gemma_maas_is_partner_model():
|
|||
VertexAIPartnerModels,
|
||||
)
|
||||
|
||||
assert VertexAIPartnerModels.is_vertex_partner_model(
|
||||
"google/gemma-4-26b-a4b-it-maas"
|
||||
)
|
||||
assert VertexAIPartnerModels.is_vertex_partner_model("google/gemma-4-26b-a4b-it-maas")
|
||||
|
||||
|
||||
def test_vertex_ai_gemma_maas_uses_openai_handler():
|
||||
|
|
@ -1643,9 +1770,7 @@ def test_vertex_ai_gemma_maas_uses_openai_handler():
|
|||
VertexAIPartnerModels,
|
||||
)
|
||||
|
||||
assert VertexAIPartnerModels.should_use_openai_handler(
|
||||
"google/gemma-4-26b-a4b-it-maas"
|
||||
)
|
||||
assert VertexAIPartnerModels.should_use_openai_handler("google/gemma-4-26b-a4b-it-maas")
|
||||
|
||||
|
||||
def test_vertex_ai_gemma_maas_routes_to_partner_models():
|
||||
|
|
@ -1835,12 +1960,8 @@ def test_pop_vertex_request_labels_prefers_explicit_labels_then_metadata():
|
|||
|
||||
def test_pop_vertex_request_labels_uses_litellm_metadata_when_metadata_absent():
|
||||
optional: dict = {}
|
||||
litellm_params = {
|
||||
"litellm_metadata": {"requester_metadata": {"team": "from_litellm_meta"}}
|
||||
}
|
||||
assert pop_vertex_request_labels(optional, litellm_params) == {
|
||||
"team": "from_litellm_meta"
|
||||
}
|
||||
litellm_params = {"litellm_metadata": {"requester_metadata": {"team": "from_litellm_meta"}}}
|
||||
assert pop_vertex_request_labels(optional, litellm_params) == {"team": "from_litellm_meta"}
|
||||
|
||||
|
||||
def test_vertex_text_embedding_request_includes_labels_from_metadata():
|
||||
|
|
@ -1850,9 +1971,7 @@ def test_vertex_text_embedding_request_includes_labels_from_metadata():
|
|||
input="hi",
|
||||
optional_params={},
|
||||
model="text-embedding-004",
|
||||
litellm_params={
|
||||
"metadata": {"requester_metadata": {"project_id": "cost-center-1"}}
|
||||
},
|
||||
litellm_params={"metadata": {"requester_metadata": {"project_id": "cost-center-1"}}},
|
||||
)
|
||||
assert req.get("labels") == {"project_id": "cost-center-1"}
|
||||
|
||||
|
|
|
|||
|
|
@ -2994,11 +2994,7 @@ async def test_call_mcp_tool_uses_manager_permission_lookup():
|
|||
"get_mcp_server_by_id",
|
||||
return_value=mock_server,
|
||||
),
|
||||
patch.object(
|
||||
global_mcp_server_manager,
|
||||
"get_mcp_server_from_tool_name",
|
||||
return_value=mock_server,
|
||||
) as mock_get_server,
|
||||
patch.object(global_mcp_server_manager, "server_exposes_tool", return_value=True),
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.operations.global_mcp_tool_registry"
|
||||
) as mock_tool_registry,
|
||||
|
|
@ -3023,24 +3019,18 @@ async def test_call_mcp_tool_uses_manager_permission_lookup():
|
|||
|
||||
assert result == expected_response
|
||||
mock_get_allowed.assert_awaited_once()
|
||||
# We call `_get_mcp_server_from_tool_name` multiple times:
|
||||
# - for logging/metadata
|
||||
# - for resolving the server during dispatch
|
||||
# - and inside `call_tool` for guardrails/hooks
|
||||
# The exact count isn't important, only that it is used.
|
||||
assert mock_get_server.call_count >= 2
|
||||
# First call should use the prefixed tool name
|
||||
assert (
|
||||
mock_get_server.call_args_list[0][0][0]
|
||||
== f"{mock_server.name}/gmail_send_email"
|
||||
)
|
||||
# The request names exactly one server, so that server is the dispatch target
|
||||
# and the tool name reaches it untouched (the "/" is not a LiteLLM prefix).
|
||||
mock_handle_managed.assert_awaited_once()
|
||||
assert mock_handle_managed.await_args.kwargs["server_name"] == mock_server.name
|
||||
assert mock_handle_managed.await_args.kwargs["name"] == f"{mock_server.name}/gmail_send_email"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_mcp_tool_resolves_unprefixed_tool_name_and_checks_permissions():
|
||||
"""
|
||||
Ensure `call_mcp_tool` correctly resolves the MCP server for an unprefixed tool
|
||||
name and enforces server-level permissions using that resolved server.
|
||||
Ensure `call_mcp_tool` resolves an unprefixed tool name on a request scoped to
|
||||
one server against that server and enforces server-level permissions with it.
|
||||
"""
|
||||
from litellm.proxy._experimental.mcp_server.server import (
|
||||
call_mcp_tool,
|
||||
|
|
@ -3070,11 +3060,7 @@ async def test_call_mcp_tool_resolves_unprefixed_tool_name_and_checks_permission
|
|||
"get_mcp_server_by_id",
|
||||
return_value=mock_server,
|
||||
),
|
||||
patch.object(
|
||||
global_mcp_server_manager,
|
||||
"get_mcp_server_from_tool_name",
|
||||
return_value=mock_server,
|
||||
) as mock_get_server,
|
||||
patch.object(global_mcp_server_manager, "server_exposes_tool", return_value=True),
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.operations.global_mcp_tool_registry"
|
||||
) as mock_tool_registry,
|
||||
|
|
@ -3091,7 +3077,7 @@ async def test_call_mcp_tool_resolves_unprefixed_tool_name_and_checks_permission
|
|||
mock_tool_registry.get_tool.return_value = None
|
||||
mock_handle_managed.return_value = expected_response
|
||||
|
||||
# Call with UNPREFIXED tool name; server should be resolved via mapping
|
||||
# Call with UNPREFIXED tool name; the single server the request names is the target
|
||||
result = await call_mcp_tool(
|
||||
name="gmail_send_email",
|
||||
arguments={"body": "hello"},
|
||||
|
|
@ -3100,8 +3086,9 @@ async def test_call_mcp_tool_resolves_unprefixed_tool_name_and_checks_permission
|
|||
|
||||
assert result == expected_response
|
||||
mock_get_allowed.assert_awaited_once()
|
||||
# We should resolve the server at least once using the unprefixed name
|
||||
assert mock_get_server.call_count >= 1
|
||||
assert mock_get_server.call_args_list[0][0][0] == "gmail_send_email"
|
||||
mock_handle_managed.assert_awaited_once()
|
||||
assert mock_handle_managed.await_args.kwargs["server_name"] == mock_server.name
|
||||
assert mock_handle_managed.await_args.kwargs["name"] == "gmail_send_email"
|
||||
# Permissions check should be invoked with the resolved server name
|
||||
mock_is_allowed.assert_called_once()
|
||||
assert mock_is_allowed.call_args.kwargs["server_name"] == mock_server.name
|
||||
|
|
|
|||
|
|
@ -565,6 +565,88 @@ async def test_tools_call_warmup_does_not_consume_mcp_server_rpm() -> None:
|
|||
assert [tool.name for tool in listing.tools] == ["echo"]
|
||||
|
||||
|
||||
|
||||
def _scoped_rig(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> tuple[MCPServer, MCPServer, AsyncMock, AsyncMock]:
|
||||
"""Two config servers that both expose a bare ``list_rules`` on a manager that has listed neither."""
|
||||
from mcp.types import CallToolResult, TextContent
|
||||
|
||||
from litellm.proxy import proxy_server
|
||||
from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServerManager
|
||||
|
||||
manager: Final = MCPServerManager()
|
||||
memory: Final = MCPServer(
|
||||
server_id="id-memory", name="memory_server", server_name="memory_server", transport=MCPTransport.http
|
||||
)
|
||||
other: Final = MCPServer(
|
||||
server_id="id-other", name="other_server", server_name="other_server", transport=MCPTransport.http
|
||||
)
|
||||
manager.config_mcp_servers = {server.server_id: server for server in (memory, other)}
|
||||
monkeypatch.setattr(proxy_server, "prisma_client", None)
|
||||
monkeypatch.setattr(operations, "global_mcp_server_manager", manager)
|
||||
monkeypatch.setattr(
|
||||
manager, "get_allowed_mcp_servers", AsyncMock(return_value=[memory.server_id, other.server_id])
|
||||
)
|
||||
monkeypatch.setattr(global_mcp_tool_registry, "published_tools", {})
|
||||
listing: Final = AsyncMock()
|
||||
monkeypatch.setattr(operations, "_get_tools_from_mcp_servers", listing)
|
||||
handled: Final = AsyncMock(
|
||||
return_value=CallToolResult(content=[TextContent(type="text", text="ok")], isError=False)
|
||||
)
|
||||
monkeypatch.setattr(operations, "_handle_managed_mcp_tool", handled)
|
||||
return memory, other, listing, handled
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("scope", ["memory_server", "other_server"])
|
||||
async def test_bare_tool_name_on_a_single_server_scope_routes_to_that_server_cold(
|
||||
monkeypatch: pytest.MonkeyPatch, scope: str
|
||||
) -> None:
|
||||
"""``/{server_name}/mcp`` (or ``x-mcp-servers`` naming one server) with the upstream's bare tool name
|
||||
must list and call the scoped server on a process whose tool-name map has never seen it, and when two
|
||||
servers share the bare name the scope decides (BerriAI/litellm#44831)."""
|
||||
memory, other, listing, handled = _scoped_rig(monkeypatch)
|
||||
expected: Final = memory if scope == "memory_server" else other
|
||||
|
||||
async with operations.global_mcp_server_manager.catalog.operation():
|
||||
result: Final = await operations.call_mcp_tool(
|
||||
name="list_rules",
|
||||
arguments={},
|
||||
user_api_key_auth=UserAPIKeyAuth(api_key="sk-scoped"),
|
||||
mcp_servers=[scope],
|
||||
)
|
||||
|
||||
assert result.is_error is False
|
||||
assert listing.await_args.kwargs["mcp_servers"] == [expected.server_id]
|
||||
assert handled.await_args.kwargs["server_name"] == expected.name
|
||||
assert handled.await_args.kwargs["name"] == "list_rules"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("mcp_servers", [None, ["memory_server", "other_server"]])
|
||||
async def test_bare_tool_name_without_a_single_server_scope_is_still_unresolved_cold(
|
||||
monkeypatch: pytest.MonkeyPatch, mcp_servers: list[str] | None
|
||||
) -> None:
|
||||
"""The scope shortcut only applies when the request names exactly one server; an ambiguous bare name
|
||||
keeps today's answer rather than being guessed onto the first allowed server."""
|
||||
from fastapi import HTTPException
|
||||
|
||||
_memory, _other, listing, handled = _scoped_rig(monkeypatch)
|
||||
|
||||
async with operations.global_mcp_server_manager.catalog.operation():
|
||||
with pytest.raises(HTTPException) as unresolved:
|
||||
await operations.call_mcp_tool(
|
||||
name="list_rules",
|
||||
arguments={},
|
||||
user_api_key_auth=UserAPIKeyAuth(api_key="sk-scoped"),
|
||||
mcp_servers=mcp_servers,
|
||||
)
|
||||
|
||||
assert unresolved.value.status_code == 404
|
||||
listing.assert_not_awaited()
|
||||
handled.assert_not_awaited()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_tools_call_pre_call_check_enforces_mcp_server_rpm() -> None:
|
||||
from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServerManager
|
||||
|
|
|
|||
|
|
@ -284,7 +284,9 @@ class TestHttpExceptionDictDetail:
|
|||
"base_process_llm_request",
|
||||
new=AsyncMock(side_effect=exc),
|
||||
),
|
||||
patch.object(proxy_server, "proxy_logging_obj") as mock_logging, # test-quality-ok: module global imported at call time; no injection seam
|
||||
patch.object(
|
||||
proxy_server, "proxy_logging_obj"
|
||||
) as mock_logging, # test-quality-ok: module global imported at call time; no injection seam
|
||||
):
|
||||
mock_logging.post_call_failure_hook = AsyncMock()
|
||||
response = await ep.anthropic_response(
|
||||
|
|
@ -424,8 +426,6 @@ class TestErrorLogCarriesCallId:
|
|||
|
||||
@pytest.mark.asyncio
|
||||
async def test_count_tokens_failure_log_carries_callers_call_id(self, caplog: pytest.LogCaptureFixture):
|
||||
from fastapi import HTTPException
|
||||
|
||||
import litellm.proxy.anthropic_endpoints.endpoints as ep
|
||||
import litellm.proxy.proxy_server as proxy_server
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
|
@ -440,13 +440,14 @@ class TestErrorLogCarriesCallId:
|
|||
"read_request_body",
|
||||
new=AsyncMock(return_value={"model": "claude-sonnet", "messages": [{"role": "user", "content": "hi"}]}),
|
||||
),
|
||||
patch.object(proxy_server, "token_counter", new=AsyncMock(side_effect=RuntimeError("tokenizer down"))), # test-quality-ok: module global imported at call time; the test targets the endpoint's except block
|
||||
patch.object(
|
||||
proxy_server, "count_request_tokens", new=AsyncMock(side_effect=RuntimeError("tokenizer down"))
|
||||
), # test-quality-ok: module global imported at call time; the test targets the endpoint's except block
|
||||
caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"),
|
||||
pytest.raises(HTTPException) as raised,
|
||||
):
|
||||
await ep.count_tokens(request=request, user_api_key_dict=UserAPIKeyAuth())
|
||||
response = await ep.count_tokens(request=request, user_api_key_dict=UserAPIKeyAuth())
|
||||
|
||||
assert raised.value.status_code == 500
|
||||
assert response.status_code == 500
|
||||
record = self._error_record(caplog)
|
||||
assert record.litellm_call_id == call_id
|
||||
assert call_id in record.getMessage()
|
||||
|
|
|
|||
|
|
@ -1354,6 +1354,124 @@ async def test_managed_batch_routes_pass_team_model_access_check(route, request_
|
|||
)
|
||||
|
||||
|
||||
_TEAM_DEPLOYMENT_ID: Final = "team-a-deployment-id"
|
||||
_TEAM_INTERNAL_MODEL_NAME: Final = "model_name_team-a_6d1f3c2e-0000-4000-8000-000000000000"
|
||||
_TEAM_PUBLIC_MODEL_NAME: Final = "team-batch-model"
|
||||
_TEAM_MANAGED_BATCH_ID: Final = _encode_managed_id(
|
||||
f"litellm_proxy;model_id:{_TEAM_DEPLOYMENT_ID};llm_batch_id:provider-batch-789"
|
||||
)
|
||||
|
||||
|
||||
def _team_scoped_router(public_model_name: str | None = _TEAM_PUBLIC_MODEL_NAME) -> Router:
|
||||
return Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": _TEAM_INTERNAL_MODEL_NAME,
|
||||
"litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "test-key"},
|
||||
"model_info": {
|
||||
"id": _TEAM_DEPLOYMENT_ID,
|
||||
"team_id": "team-a",
|
||||
**({"team_public_model_name": public_model_name} if public_model_name else {}),
|
||||
},
|
||||
},
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
def _team_scoped_character_id() -> str:
|
||||
from litellm.types.videos.utils import encode_character_id_with_provider
|
||||
|
||||
return encode_character_id_with_provider(character_id="char_team", provider="openai", model_id=_TEAM_DEPLOYMENT_ID)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"route, request_data",
|
||||
[
|
||||
pytest.param("/v1/batches/{batch_id}", {"batch_id": _TEAM_MANAGED_BATCH_ID}, id="batch-retrieve"),
|
||||
pytest.param("/v1/batches/{batch_id}/cancel", {"batch_id": _TEAM_MANAGED_BATCH_ID}, id="batch-cancel"),
|
||||
pytest.param(
|
||||
"/v1/fine_tuning/jobs/{fine_tuning_job_id}",
|
||||
{
|
||||
"fine_tuning_job_id": _encode_managed_id(
|
||||
f"litellm_proxy;model_id:{_TEAM_DEPLOYMENT_ID};generic_response_id:ftjob-1"
|
||||
)
|
||||
},
|
||||
id="fine-tuning-job",
|
||||
),
|
||||
pytest.param(
|
||||
"/v1/vector_stores/{vector_store_id}",
|
||||
{
|
||||
"vector_store_id": _encode_managed_id(
|
||||
"litellm_proxy:vector_store;unified_id,vs-unified;"
|
||||
f"target_model_names,{_TEAM_PUBLIC_MODEL_NAME};resource_id,vs_1;model_id,{_TEAM_DEPLOYMENT_ID}"
|
||||
)
|
||||
},
|
||||
id="vector-store",
|
||||
),
|
||||
pytest.param(
|
||||
"/v1/videos/characters/{character_id}",
|
||||
{"character_id": _team_scoped_character_id()},
|
||||
id="character",
|
||||
),
|
||||
],
|
||||
)
|
||||
async def test_team_scoped_managed_ids_authorize_only_the_owning_team_by_public_name(
|
||||
route: str, request_data: dict[str, str]
|
||||
) -> None:
|
||||
from litellm.proxy._types import ProxyException
|
||||
from litellm.proxy.auth.auth_checks import can_key_call_resolved_model
|
||||
|
||||
llm_router: Final = _team_scoped_router()
|
||||
|
||||
async def check(team_id: str) -> None:
|
||||
await can_key_call_resolved_model(
|
||||
model=get_model_from_request(
|
||||
request_data=request_data, route=route, llm_router=llm_router, team_id=team_id
|
||||
),
|
||||
llm_model_list=None,
|
||||
valid_token=UserAPIKeyAuth(api_key="sk-team", team_id=team_id, team_models=[_TEAM_PUBLIC_MODEL_NAME]),
|
||||
llm_router=llm_router,
|
||||
)
|
||||
|
||||
await check("team-a")
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
await check("team-b")
|
||||
assert exc_info.value.code == "403"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_legacy_team_alias_authorizes_a_team_batch_only_for_keys_granted_the_alias() -> None:
|
||||
from litellm.proxy._types import ProxyException
|
||||
from litellm.proxy.auth.auth_checks import can_key_call_resolved_model
|
||||
|
||||
llm_router: Final = _team_scoped_router(public_model_name=None)
|
||||
|
||||
async def check(key_models: list[str]) -> None:
|
||||
await can_key_call_resolved_model(
|
||||
model=get_model_from_request(
|
||||
request_data={"batch_id": _TEAM_MANAGED_BATCH_ID},
|
||||
route="/v1/batches/{batch_id}",
|
||||
llm_router=llm_router,
|
||||
team_id="team-a",
|
||||
),
|
||||
llm_model_list=None,
|
||||
valid_token=UserAPIKeyAuth(
|
||||
api_key="sk-team",
|
||||
team_id="team-a",
|
||||
models=key_models,
|
||||
team_models=["legacy-team-model"],
|
||||
team_model_aliases={"legacy-team-model": _TEAM_INTERNAL_MODEL_NAME},
|
||||
),
|
||||
llm_router=llm_router,
|
||||
)
|
||||
|
||||
await check(["legacy-team-model"])
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
await check(["some-other-model"])
|
||||
assert exc_info.value.code == "403"
|
||||
|
||||
|
||||
def test_get_model_from_request_resolves_character_id_model_with_router():
|
||||
from litellm.types.videos.utils import encode_character_id_with_provider
|
||||
|
||||
|
|
|
|||
|
|
@ -35,7 +35,7 @@ import logging
|
|||
from contextlib import ExitStack
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import Any, Dict, Optional
|
||||
from typing import Any, Dict, Final, Optional
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import httpx
|
||||
|
|
@ -3259,6 +3259,129 @@ async def test_cancel__unified_batch_id_rejects_key_without_model_grant(cancel_h
|
|||
cancel_harness.router_acancel.assert_not_called()
|
||||
|
||||
|
||||
TEAM_INTERNAL_MODEL_NAME: Final = "model_name_team-a_00000000-0000-0000-0000-000000000000"
|
||||
TEAM_PUBLIC_MODEL_NAME: Final = "team-gpt4o-mini"
|
||||
|
||||
|
||||
def _team_scoped_deployment() -> Deployment:
|
||||
return Deployment(
|
||||
model_name=TEAM_INTERNAL_MODEL_NAME,
|
||||
litellm_params=LiteLLM_Params(model="gpt-4o-mini"),
|
||||
model_info=ModelInfo(id="dep-team-a-1", team_id="team-a", team_public_model_name=TEAM_PUBLIC_MODEL_NAME),
|
||||
)
|
||||
|
||||
|
||||
def _route_unified_batch_to_team_deployment(
|
||||
harness: RetrieveHarness | CancelHarness, deployment: Deployment | None = None
|
||||
) -> None:
|
||||
target: Final = deployment or _team_scoped_deployment()
|
||||
harness.router.resolve_model_name_from_model_id = MagicMock(return_value=target.model_name)
|
||||
harness.router.get_deployment_by_model_group_name = MagicMock(return_value=target)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retrieve__team_scoped_deployment_authorizes_owning_team_by_public_name(
|
||||
retrieve_harness: RetrieveHarness,
|
||||
) -> None:
|
||||
_route_unified_batch_to_team_deployment(retrieve_harness)
|
||||
|
||||
resp: Final = await call_retrieve(
|
||||
retrieve_harness, UNIFIED_BATCH_ID_FOR_GPT4O_MINI, user=_key_restricted_to(TEAM_PUBLIC_MODEL_NAME)
|
||||
)
|
||||
|
||||
assert retrieve_harness.router_aretrieve.call_count == 1
|
||||
assert resp._hidden_params["model_id"] == "gpt-4o-mini"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancel__team_scoped_deployment_authorizes_owning_team_by_public_name(
|
||||
cancel_harness: CancelHarness,
|
||||
) -> None:
|
||||
_route_unified_batch_to_team_deployment(cancel_harness)
|
||||
|
||||
await call_cancel(cancel_harness, UNIFIED_BATCH_ID_FOR_GPT4O_MINI, user=_key_restricted_to(TEAM_PUBLIC_MODEL_NAME))
|
||||
|
||||
assert cancel_harness.router_acancel.call_count == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"deployment, user",
|
||||
[
|
||||
pytest.param(
|
||||
_team_scoped_deployment(), _key_restricted_to("some-other-model"), id="owning-team-key-without-the-model"
|
||||
),
|
||||
pytest.param(
|
||||
_team_scoped_deployment(),
|
||||
UserAPIKeyAuth(
|
||||
api_key="sk-team-b",
|
||||
team_id="team-b",
|
||||
team_models=[TEAM_PUBLIC_MODEL_NAME],
|
||||
models=[TEAM_PUBLIC_MODEL_NAME],
|
||||
),
|
||||
id="other-team-key-granted-the-same-public-name",
|
||||
),
|
||||
pytest.param(
|
||||
Deployment(
|
||||
model_name="model_name_team-a_ml_00000000-0000-0000-0000-000000000000",
|
||||
litellm_params=LiteLLM_Params(model="gpt-4o-mini"),
|
||||
model_info=ModelInfo(
|
||||
id="dep-team-a-ml-1", team_id="team-a_ml", team_public_model_name=TEAM_PUBLIC_MODEL_NAME
|
||||
),
|
||||
),
|
||||
_key_restricted_to(TEAM_PUBLIC_MODEL_NAME),
|
||||
id="team-whose-id-prefixes-the-owning-team-id",
|
||||
),
|
||||
],
|
||||
)
|
||||
async def test_retrieve__team_scoped_deployment_rejects_keys_outside_its_grant(
|
||||
retrieve_harness: RetrieveHarness, deployment: Deployment, user: UserAPIKeyAuth
|
||||
) -> None:
|
||||
_route_unified_batch_to_team_deployment(retrieve_harness, deployment)
|
||||
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
await call_retrieve(retrieve_harness, UNIFIED_BATCH_ID_FOR_GPT4O_MINI, user=user)
|
||||
|
||||
assert exc_info.value.code == "403"
|
||||
retrieve_harness.router_aretrieve.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"deployment, user",
|
||||
[
|
||||
pytest.param(
|
||||
_team_scoped_deployment(),
|
||||
_key_restricted_to(TEAM_INTERNAL_MODEL_NAME),
|
||||
id="owning-team-key-granted-the-internal-name",
|
||||
),
|
||||
pytest.param(
|
||||
Deployment(
|
||||
model_name=TEAM_INTERNAL_MODEL_NAME,
|
||||
litellm_params=LiteLLM_Params(model="gpt-4o-mini"),
|
||||
model_info=ModelInfo(id="dep-team-a-1", team_id="team-a"),
|
||||
),
|
||||
UserAPIKeyAuth(
|
||||
api_key="sk-legacy",
|
||||
team_id="team-a",
|
||||
team_models=["legacy-team-model"],
|
||||
models=["legacy-team-model"],
|
||||
team_model_aliases={"legacy-team-model": TEAM_INTERNAL_MODEL_NAME},
|
||||
),
|
||||
id="legacy-team-model-reached-through-a-team-alias",
|
||||
),
|
||||
],
|
||||
)
|
||||
async def test_retrieve__team_scoped_deployment_authorizes_internal_name_and_legacy_alias_grants(
|
||||
retrieve_harness: RetrieveHarness, deployment: Deployment, user: UserAPIKeyAuth
|
||||
) -> None:
|
||||
_route_unified_batch_to_team_deployment(retrieve_harness, deployment)
|
||||
|
||||
await call_retrieve(retrieve_harness, UNIFIED_BATCH_ID_FOR_GPT4O_MINI, user=user)
|
||||
|
||||
assert retrieve_harness.router_aretrieve.call_count == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancel__executed_batch_rejects_key_without_model_grant(cancel_harness, executed_runner):
|
||||
runner, factory = executed_runner
|
||||
|
|
|
|||
190
tests/unit/proxy/common_utils/test_token_count_settings.py
Normal file
190
tests/unit/proxy/common_utils/test_token_count_settings.py
Normal file
|
|
@ -0,0 +1,190 @@
|
|||
import json
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.proxy.common_utils.token_count_settings import (
|
||||
ProviderCountFailure,
|
||||
TokenCountModel,
|
||||
TokenCountSettings,
|
||||
count_source_fields,
|
||||
resolve_token_count_settings,
|
||||
)
|
||||
from litellm.types.utils import TokenCountResponse
|
||||
|
||||
_OPUS_US_EAST_1: Final = TokenCountModel(model="anthropic.claude-opus-5-5", region="us-east-1")
|
||||
_DEFAULT: Final = {
|
||||
"token_count_framing": 97,
|
||||
"token_count_model": {"model": "anthropic.claude-opus-5-5", "region": "us-east-1"},
|
||||
}
|
||||
|
||||
|
||||
def _in_us_west_2() -> str:
|
||||
return "us-west-2"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("configured", "expected"),
|
||||
[
|
||||
({}, TokenCountSettings(framing=97, count_model=_OPUS_US_EAST_1)),
|
||||
(
|
||||
{"token_count_framing": 5, "token_count_model": {"model": "custom-model", "region": "ap-southeast-4"}},
|
||||
TokenCountSettings(framing=5, count_model=TokenCountModel(model="custom-model", region="ap-southeast-4")),
|
||||
),
|
||||
(
|
||||
{"token_count_model": {"model": "custom-model"}},
|
||||
TokenCountSettings(framing=97, count_model=TokenCountModel(model="custom-model", region=None)),
|
||||
),
|
||||
(
|
||||
{"token_count_model": {"model": "custom-model", "region": " "}},
|
||||
TokenCountSettings(framing=97, count_model=TokenCountModel(model="custom-model", region=None)),
|
||||
),
|
||||
({"token_count_framing": 0}, TokenCountSettings(framing=0, count_model=_OPUS_US_EAST_1)),
|
||||
(
|
||||
{"token_count_framing": None, "token_count_model": None},
|
||||
TokenCountSettings(framing=97, count_model=_OPUS_US_EAST_1),
|
||||
),
|
||||
(
|
||||
{"token_count_framing": True, "token_count_model": "custom-model"},
|
||||
TokenCountSettings(framing=97, count_model=_OPUS_US_EAST_1),
|
||||
),
|
||||
({"token_count_framing": 78.0}, TokenCountSettings(framing=78, count_model=_OPUS_US_EAST_1)),
|
||||
({"token_count_framing": 1.5}, TokenCountSettings(framing=97, count_model=_OPUS_US_EAST_1)),
|
||||
({"token_count_framing": -1}, TokenCountSettings(framing=97, count_model=_OPUS_US_EAST_1)),
|
||||
({"token_count_framing": "78"}, TokenCountSettings(framing=97, count_model=_OPUS_US_EAST_1)),
|
||||
({"token_count_model": {}}, TokenCountSettings(framing=97, count_model=None)),
|
||||
({"token_count_model": ""}, TokenCountSettings(framing=97, count_model=_OPUS_US_EAST_1)),
|
||||
(
|
||||
{"token_count_model": {"model": " ", "region": "us-east-1"}},
|
||||
TokenCountSettings(framing=97, count_model=None),
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_config_overrides_the_cost_map_default_key_by_key(
|
||||
configured: dict[str, object], expected: TokenCountSettings
|
||||
) -> None:
|
||||
assert (
|
||||
resolve_token_count_settings(configured=configured, default=_DEFAULT, deployment_region=_in_us_west_2)
|
||||
== expected
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("deployment_region", "configured", "expected_region"),
|
||||
[
|
||||
("us-west-2", {}, "us-east-1"),
|
||||
("eu-central-1", {}, None),
|
||||
("ca-central-1", {}, None),
|
||||
("us-gov-west-1", {}, None),
|
||||
(
|
||||
"eu-central-1",
|
||||
{"token_count_model": {"model": "anthropic.claude-opus-5-5", "region": "us-east-1"}},
|
||||
"us-east-1",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_shipped_count_region_never_takes_the_prompt_out_of_the_deployment_geography(
|
||||
deployment_region: str, configured: dict[str, object], expected_region: str | None
|
||||
) -> None:
|
||||
settings: Final = resolve_token_count_settings(
|
||||
configured=configured, default=_DEFAULT, deployment_region=lambda: deployment_region
|
||||
)
|
||||
assert settings.count_model == TokenCountModel(model="anthropic.claude-opus-5-5", region=expected_region)
|
||||
|
||||
|
||||
def test_shipped_global_count_model_keeps_its_region_for_a_deployment_in_any_geography() -> None:
|
||||
settings: Final = resolve_token_count_settings(
|
||||
configured={},
|
||||
default={"token_count_model": {"model": "global.anthropic.claude-opus-5-5", "region": "us-east-1"}},
|
||||
deployment_region=lambda: "eu-west-1",
|
||||
)
|
||||
assert settings.count_model == TokenCountModel(model="global.anthropic.claude-opus-5-5", region="us-east-1")
|
||||
|
||||
|
||||
def test_shipped_australian_count_region_stays_with_an_australian_deployment() -> None:
|
||||
settings: Final = resolve_token_count_settings(
|
||||
configured={},
|
||||
default={"token_count_model": {"model": "au.anthropic.claude-opus-5-5", "region": "ap-southeast-4"}},
|
||||
deployment_region=lambda: "ap-southeast-2",
|
||||
)
|
||||
assert settings.count_model == TokenCountModel(model="au.anthropic.claude-opus-5-5", region="ap-southeast-4")
|
||||
|
||||
|
||||
def test_no_setting_anywhere_means_the_model_counts_itself_without_framing() -> None:
|
||||
assert resolve_token_count_settings(
|
||||
configured={}, default={}, deployment_region=_in_us_west_2
|
||||
) == TokenCountSettings(framing=None, count_model=None)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("value", "expected"),
|
||||
[("kimi_k3", "kimi_k3"), ("cl100k_base", None), (None, None), (True, None)],
|
||||
)
|
||||
def test_only_a_tokenizer_litellm_ships_is_taken_from_the_cost_map(value: object, expected: str | None) -> None:
|
||||
settings: Final = resolve_token_count_settings(
|
||||
configured={}, default={"token_count_tokenizer": value}, deployment_region=_in_us_west_2
|
||||
)
|
||||
assert settings.tokenizer == expected
|
||||
|
||||
|
||||
def _counted(source: str | None, model_used: str = "us.anthropic.claude-opus-5-5") -> TokenCountResponse:
|
||||
return TokenCountResponse(
|
||||
total_tokens=7,
|
||||
request_model="sonnet",
|
||||
model_used=model_used,
|
||||
tokenizer_type="bedrock_api",
|
||||
count_source=source,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("source", "expected"),
|
||||
[
|
||||
(None, {}),
|
||||
("provider", {"litellm_count_source": "provider"}),
|
||||
("local_estimate", {"litellm_count_source": "local_estimate"}),
|
||||
(
|
||||
"token_count_model",
|
||||
{"litellm_count_source": "token_count_model", "litellm_count_model": "us.anthropic.claude-opus-5-5"},
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_count_source_fields_name_the_counting_model_only_when_token_count_model_counted(
|
||||
source: str | None, expected: dict[str, str]
|
||||
) -> None:
|
||||
assert count_source_fields(_counted(source)) == expected
|
||||
|
||||
|
||||
_UPSTREAM_404: Final = {
|
||||
"type": "error",
|
||||
"request_id": "req_upstream",
|
||||
"error": {"type": "not_found_error", "message": "The model 'anthropic.claude-opus-5-5' does not exist"},
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("message", "expected"),
|
||||
[
|
||||
(
|
||||
json.dumps(_UPSTREAM_404),
|
||||
{
|
||||
**_UPSTREAM_404,
|
||||
"error": {
|
||||
"type": "not_found_error",
|
||||
"message": "The model 'anthropic.claude-opus-5-5' does not exist. Enable eu-west-1",
|
||||
},
|
||||
},
|
||||
),
|
||||
(json.dumps({"message": "Too many tokens"}), "Too many tokens. Enable eu-west-1"),
|
||||
("scripted 404", "scripted 404. Enable eu-west-1"),
|
||||
(
|
||||
"Credential should be scoped to a valid region. ",
|
||||
"Credential should be scoped to a valid region. Enable eu-west-1",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_hint_lands_in_the_upstream_error_message_and_keeps_its_envelope(message: str, expected: object) -> None:
|
||||
hinted: Final = ProviderCountFailure(status_code=404, message=message).with_hint("Enable eu-west-1")
|
||||
|
||||
assert (json.loads(hinted.message) if isinstance(expected, dict) else hinted.message) == expected
|
||||
assert hinted.status_code == 404
|
||||
|
|
@ -4,10 +4,11 @@ Test to verify the Google GenAI proxy API endpoints
|
|||
"""
|
||||
|
||||
from typing import Final
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.types.utils import TokenCountResponse
|
||||
|
||||
|
||||
def _build_test_client():
|
||||
|
|
@ -88,9 +89,7 @@ def test_google_stream_generate_content_endpoint():
|
|||
assert init_kwargs["data"]["_litellm_raw_sse_stream"] is True
|
||||
assert init_kwargs["data"]["_litellm_skip_openai_stream_done"] is True
|
||||
assert init_kwargs["data"]["model"] == "test-model"
|
||||
assert init_kwargs["data"]["contents"] == [
|
||||
{"role": "user", "parts": [{"text": "Hello"}]}
|
||||
]
|
||||
assert init_kwargs["data"]["contents"] == [{"role": "user", "parts": [{"text": "Hello"}]}]
|
||||
|
||||
|
||||
def test_google_generate_content_data_flows_through_processor():
|
||||
|
|
@ -123,9 +122,7 @@ def test_google_generate_content_data_flows_through_processor():
|
|||
data = mock_init.call_args.kwargs["data"]
|
||||
assert data["model"] == "test-model"
|
||||
assert data["contents"] == [{"role": "user", "parts": [{"text": "Hello"}]}]
|
||||
assert data["systemInstruction"] == {
|
||||
"parts": [{"text": "Your name is Doodle."}]
|
||||
}
|
||||
assert data["systemInstruction"] == {"parts": [{"text": "Your name is Doodle."}]}
|
||||
# generationConfig arrives intact here; the rename to `config` is
|
||||
# done downstream in route_request (see test_route_llm_request).
|
||||
assert data["generationConfig"]["responseModalities"] == ["TEXT", "IMAGE"]
|
||||
|
|
@ -160,12 +157,17 @@ def test_google_count_tokens_unchanged():
|
|||
except ImportError as e:
|
||||
pytest.skip(f"Skipping test due to missing dependency: {e}")
|
||||
|
||||
fake_response = MagicMock()
|
||||
fake_response.original_response = {
|
||||
"totalTokens": 7,
|
||||
"promptTokensDetails": [{"modality": "TEXT", "tokenCount": 7}],
|
||||
}
|
||||
fake_response.total_tokens = 7
|
||||
fake_response = TokenCountResponse(
|
||||
total_tokens=7,
|
||||
request_model="test-model",
|
||||
model_used="test-model",
|
||||
tokenizer_type="gemini_api",
|
||||
original_response={
|
||||
"totalTokens": 7,
|
||||
"promptTokensDetails": [{"modality": "TEXT", "tokenCount": 7}],
|
||||
},
|
||||
count_source="provider",
|
||||
)
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.proxy_server.token_counter",
|
||||
|
|
@ -180,6 +182,30 @@ def test_google_count_tokens_unchanged():
|
|||
assert response.status_code == 200
|
||||
body = response.json()
|
||||
assert body["totalTokens"] == 7
|
||||
assert body["litellm_count_source"] == "provider"
|
||||
|
||||
|
||||
def test_google_count_tokens_unexpected_failure_uses_the_google_error_envelope():
|
||||
try:
|
||||
client = _build_test_client()
|
||||
except ImportError as e:
|
||||
pytest.skip(f"Skipping test due to missing dependency: {e}")
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.proxy_server.token_counter",
|
||||
new_callable=AsyncMock,
|
||||
side_effect=RuntimeError("tokenizer unavailable"),
|
||||
):
|
||||
response = client.post(
|
||||
"/v1beta/models/test-model:countTokens",
|
||||
json={"contents": [{"role": "user", "parts": [{"text": "Hello"}]}]},
|
||||
)
|
||||
|
||||
assert response.status_code == 500
|
||||
error = response.json()["error"]
|
||||
assert error["code"] == 500
|
||||
assert error["status"] == "INTERNAL"
|
||||
assert "tokenizer unavailable" not in error["message"]
|
||||
|
||||
|
||||
def test_google_count_tokens_uses_normalized_total_when_provider_shape_differs():
|
||||
|
|
@ -189,9 +215,14 @@ def test_google_count_tokens_uses_normalized_total_when_provider_shape_differs()
|
|||
except ImportError as e:
|
||||
pytest.skip(f"Skipping test due to missing dependency: {e}")
|
||||
|
||||
fake_response: Final = MagicMock()
|
||||
fake_response.original_response = {"input_tokens": 2167}
|
||||
fake_response.total_tokens = 2167
|
||||
fake_response: Final = TokenCountResponse(
|
||||
total_tokens=2167,
|
||||
request_model="claude-opus-4-8",
|
||||
model_used="us.anthropic.claude-opus-5-5",
|
||||
tokenizer_type="bedrock_api",
|
||||
original_response={"input_tokens": 2167},
|
||||
count_source="token_count_model",
|
||||
)
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.proxy_server.token_counter",
|
||||
|
|
@ -207,3 +238,5 @@ def test_google_count_tokens_uses_normalized_total_when_provider_shape_differs()
|
|||
body: Final = response.json()
|
||||
assert body["totalTokens"] == 2167
|
||||
assert body["promptTokensDetails"] == []
|
||||
assert body["litellm_count_source"] == "token_count_model"
|
||||
assert body["litellm_count_model"] == "us.anthropic.claude-opus-5-5"
|
||||
|
|
|
|||
|
|
@ -9,11 +9,12 @@ The request payload is correctly processed and forwarded to the httpx client.
|
|||
|
||||
import json
|
||||
import unittest.mock
|
||||
from typing import Optional
|
||||
from typing import Final, Optional
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from pydantic import JsonValue, ValidationError
|
||||
|
||||
# Add the parent directory to the system path
|
||||
|
||||
|
|
@ -493,3 +494,39 @@ async def test_gemini_proxy_config_with_custom_api_base():
|
|||
if __name__ == "__main__":
|
||||
# Run the tests
|
||||
pytest.main([__file__, "-v"])
|
||||
|
||||
|
||||
@pytest.mark.parametrize("system_field", ("systemInstruction", "system_instruction"))
|
||||
@pytest.mark.parametrize("parts", ({"text": "Follow these rules"}, [{"text": "Follow these rules"}]))
|
||||
@pytest.mark.parametrize("tools", (None, []))
|
||||
def test_google_count_tokens_accepts_equivalent_system_parts_and_no_tools(
|
||||
system_field: str, parts: JsonValue, tools: JsonValue
|
||||
) -> None:
|
||||
from litellm.proxy.google_endpoints.endpoints import _google_count_tokens_request
|
||||
|
||||
contents: Final = [{"role": "user", "parts": [{"text": "Count this request"}]}]
|
||||
request: Final = _google_count_tokens_request(
|
||||
"counted", {"contents": contents, system_field: {"parts": parts}, "tools": tools}
|
||||
)
|
||||
|
||||
assert request.messages == [{"role": "user", "content": "Count this request"}]
|
||||
assert request.system == [{"type": "text", "text": "Follow these rules"}]
|
||||
assert request.tools is None
|
||||
assert request.contents == contents
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"parts",
|
||||
(
|
||||
{"inlineData": {"mimeType": "image/png", "data": "aGVsbG8="}},
|
||||
[{"inlineData": {"mimeType": "image/png", "data": "aGVsbG8="}}],
|
||||
),
|
||||
)
|
||||
def test_google_count_tokens_rejects_media_only_system_parts(parts: JsonValue) -> None:
|
||||
from litellm.proxy.google_endpoints.endpoints import _google_count_tokens_request
|
||||
|
||||
with pytest.raises(ValidationError):
|
||||
_google_count_tokens_request(
|
||||
"counted",
|
||||
{"contents": [{"role": "user", "parts": [{"text": "Hello"}]}], "systemInstruction": {"parts": parts}},
|
||||
)
|
||||
|
|
|
|||
|
|
@ -4147,6 +4147,54 @@ class TestModelInfoCostMapEchoFilter:
|
|||
assert "litellm_provider" not in info
|
||||
assert "supported_openai_params" not in info
|
||||
|
||||
@pytest.fixture
|
||||
def bedrock_sonnet_5_5_mapped_as_on_main(self) -> Iterator[None]:
|
||||
import litellm
|
||||
|
||||
entry: Final = {
|
||||
"litellm_provider": "bedrock_converse",
|
||||
"mode": "chat",
|
||||
"token_count_model": {"model": "us.anthropic.claude-opus-5-5", "region": "us-east-1"},
|
||||
}
|
||||
with pytest.MonkeyPatch.context() as patched:
|
||||
patched.setattr(litellm, "model_cost", {**litellm.model_cost, "us.anthropic.claude-sonnet-5-5": entry})
|
||||
litellm.utils._invalidate_model_cost_lowercase_map()
|
||||
yield
|
||||
litellm.utils._invalidate_model_cost_lowercase_map()
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("model", "field", "edited"),
|
||||
[
|
||||
("bedrock/us.anthropic.claude-sonnet-5-5", "token_count_model", {}),
|
||||
("bedrock/us.moonshotai.kimi-k3", "token_count_framing", 0),
|
||||
],
|
||||
)
|
||||
def test_token_count_setting_is_stored_only_when_the_operator_changed_it(
|
||||
self, model: str, field: str, edited: object, bedrock_sonnet_5_5_mapped_as_on_main: None
|
||||
) -> None:
|
||||
import litellm
|
||||
from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model
|
||||
|
||||
entry: Final = litellm.get_model_info(model)
|
||||
assert entry[field] not in (None, edited), entry
|
||||
db_model: Final = Deployment(
|
||||
model_name="count-settings",
|
||||
litellm_params=LiteLLM_Params(model=model),
|
||||
model_info=ModelInfo(id="dep-echo-count"),
|
||||
)
|
||||
|
||||
echoed: Final = update_db_model(
|
||||
db_model=db_model,
|
||||
updated_patch=updateDeployment(model_info=ModelInfo(**{**entry, "id": "dep-echo-count"})),
|
||||
)
|
||||
changed: Final = update_db_model(
|
||||
db_model=db_model,
|
||||
updated_patch=updateDeployment(model_info=ModelInfo(**{**entry, "id": "dep-echo-count", field: edited})),
|
||||
)
|
||||
|
||||
assert field not in json.loads(echoed["model_info"])
|
||||
assert json.loads(changed["model_info"])[field] == edited
|
||||
|
||||
def test_metadata_without_a_cost_map_key_is_persisted(self):
|
||||
import litellm
|
||||
|
||||
|
|
|
|||
|
|
@ -54,6 +54,7 @@ def test_token_counter_happy_path(client, auth_as, patched_token_counter):
|
|||
"request_model": "gpt-4",
|
||||
"model_used": "gpt-4",
|
||||
"tokenizer_type": "openai_tokenizer",
|
||||
"count_source": "local_estimate",
|
||||
"original_response": None,
|
||||
"error": False,
|
||||
"error_message": None,
|
||||
|
|
@ -390,8 +391,7 @@ def test_token_counter_fallback_counts_tools_system_and_anthropic_blocks(client,
|
|||
assert full > bare
|
||||
|
||||
|
||||
def test_token_counter_fallback_prompt_with_tools_does_not_500(client, auth_as, monkeypatch):
|
||||
"""Regression: a ``prompt`` request carrying ``tools`` but no ``messages`` still counts, because the fallback attaches tools only when counting messages (``token_counter`` rejects tools on the text path)."""
|
||||
def test_token_counter_fallback_counts_a_prompt_with_tools_like_its_messages_form(client, auth_as, monkeypatch):
|
||||
monkeypatch.setattr(proxy_server, "llm_router", None)
|
||||
monkeypatch.setattr(litellm, "disable_token_counter", False, raising=False)
|
||||
prompt = "count the tokens in this sentence please"
|
||||
|
|
@ -412,4 +412,6 @@ def test_token_counter_fallback_prompt_with_tools_does_not_500(client, auth_as,
|
|||
)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.json()["total_tokens"] == litellm.token_counter(model="claude-fable-5", text=prompt)
|
||||
assert response.json()["total_tokens"] == litellm.token_counter(
|
||||
model="claude-fable-5", messages=[{"role": "user", "content": prompt}], tools=tools
|
||||
)
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
|
|
@ -6,6 +6,7 @@ import pytest
|
|||
from fastapi import HTTPException, Request
|
||||
|
||||
import litellm
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook import (
|
||||
LiteLLM_ManagedVectorStore,
|
||||
)
|
||||
|
|
@ -3256,3 +3257,159 @@ class TestConfigOwnedVectorStores:
|
|||
assert response["status"] == "success", response
|
||||
prisma.db.litellm_managedvectorstorestable.delete.assert_awaited_once_with(where={"vector_store_id": self.DB_ID})
|
||||
assert registry.get_litellm_managed_vector_store_from_registry(self.DB_ID) is None
|
||||
|
||||
|
||||
CALLERS_METADATA = {"team": "llmproxy", "endpoint": "ingest", "trace_id": "callers-trace"}
|
||||
|
||||
|
||||
class _PreCallBodies(CustomLogger):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.logged_bodies = []
|
||||
|
||||
def log_pre_api_call(self, model, messages, kwargs):
|
||||
self.logged_bodies.append(kwargs["litellm_params"]["proxy_server_request"]["body"])
|
||||
|
||||
|
||||
@pytest.mark.parametrize("operation", ["create", "update"])
|
||||
@pytest.mark.parametrize("model", [None, "gpt-5"])
|
||||
@pytest.mark.parametrize(
|
||||
("metadata_field", "expected_metadata"),
|
||||
[
|
||||
({"metadata": CALLERS_METADATA}, CALLERS_METADATA),
|
||||
({"metadata": {}}, {}),
|
||||
({"metadata": json.dumps(CALLERS_METADATA)}, CALLERS_METADATA),
|
||||
({"metadata": None}, None),
|
||||
({}, None),
|
||||
],
|
||||
)
|
||||
def test_proxy_vector_store_writes_send_the_provider_exactly_the_callers_metadata(
|
||||
monkeypatch, operation, model, metadata_field, expected_metadata
|
||||
):
|
||||
import httpx
|
||||
import respx
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from litellm.constants import VECTOR_STORE_CALLER_METADATA_KWARG
|
||||
from litellm.proxy import proxy_server
|
||||
from litellm.proxy.proxy_server import app
|
||||
|
||||
monkeypatch.setenv("OPENAI_API_KEY", "sk-test")
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
litellm.in_memory_llm_clients_cache.flush_cache()
|
||||
monkeypatch.setattr(proxy_server, "master_key", "sk-master-test")
|
||||
monkeypatch.setattr(
|
||||
proxy_server,
|
||||
"llm_router",
|
||||
litellm.Router(
|
||||
model_list=[{"model_name": "gpt-5", "litellm_params": {"model": "openai/gpt-5", "api_key": "sk-test"}}]
|
||||
),
|
||||
)
|
||||
pre_call_bodies = _PreCallBodies()
|
||||
monkeypatch.setattr(litellm, "callbacks", [pre_call_bodies])
|
||||
sent_bodies = []
|
||||
|
||||
def echo(request: httpx.Request) -> httpx.Response:
|
||||
body = json.loads(request.content)
|
||||
sent_bodies.append(body)
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"id": "vs_1",
|
||||
"object": "vector_store",
|
||||
"created_at": 0,
|
||||
"name": "x",
|
||||
"usage_bytes": 0,
|
||||
"file_counts": {"in_progress": 0, "completed": 0, "failed": 0, "cancelled": 0, "total": 0},
|
||||
"status": "completed",
|
||||
"metadata": body.get("metadata"),
|
||||
},
|
||||
)
|
||||
|
||||
path = "/v1/vector_stores" if operation == "create" else "/v1/vector_stores/vs_1"
|
||||
routing = {} if model is None else {"model": model}
|
||||
with respx.mock(assert_all_called=False) as upstream:
|
||||
upstream.post(url__regex=r"https://api\.openai\.com/v1/vector_stores.*").mock(side_effect=echo)
|
||||
response = TestClient(app).post(
|
||||
path,
|
||||
json={"name": "x", **routing, **metadata_field},
|
||||
headers={
|
||||
"Authorization": "Bearer sk-master-test",
|
||||
"x-litellm-tags": "header-tag",
|
||||
"x-litellm-trace-id": "header-trace",
|
||||
"x-litellm-session-id": "header-session",
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert len(sent_bodies) == 1
|
||||
assert sent_bodies[0].get("metadata") == expected_metadata
|
||||
assert len(pre_call_bodies.logged_bodies) == 1
|
||||
assert "name" in pre_call_bodies.logged_bodies[0]
|
||||
assert VECTOR_STORE_CALLER_METADATA_KWARG not in pre_call_bodies.logged_bodies[0]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("operation", ["create", "update"])
|
||||
@pytest.mark.parametrize(
|
||||
"request_body",
|
||||
[
|
||||
{"data": {"name": "x", "metadata": json.dumps(CALLERS_METADATA)}},
|
||||
{
|
||||
"content": json.dumps({"name": "x", "metadata": CALLERS_METADATA})[:-1]
|
||||
+ ', "nested": '
|
||||
+ "[" * 300
|
||||
+ "]" * 300
|
||||
+ "}",
|
||||
"headers": {"Content-Type": "application/json"},
|
||||
},
|
||||
],
|
||||
ids=["form", "json-too-deep-to-reparse"],
|
||||
)
|
||||
def test_proxy_vector_store_writes_send_only_the_callers_metadata_when_the_raw_body_cannot_be_reparsed(
|
||||
monkeypatch, operation, request_body
|
||||
):
|
||||
import httpx
|
||||
import respx
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from litellm.proxy import proxy_server
|
||||
from litellm.proxy.proxy_server import app
|
||||
|
||||
monkeypatch.setenv("OPENAI_API_KEY", "sk-test")
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
litellm.in_memory_llm_clients_cache.flush_cache()
|
||||
monkeypatch.setattr(proxy_server, "master_key", "sk-master-test")
|
||||
monkeypatch.setattr(proxy_server, "llm_router", litellm.Router(model_list=[]))
|
||||
sent_bodies = []
|
||||
|
||||
def echo(request: httpx.Request) -> httpx.Response:
|
||||
body = json.loads(request.content)
|
||||
sent_bodies.append(body)
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"id": "vs_1",
|
||||
"object": "vector_store",
|
||||
"created_at": 0,
|
||||
"name": "x",
|
||||
"usage_bytes": 0,
|
||||
"file_counts": {"in_progress": 0, "completed": 0, "failed": 0, "cancelled": 0, "total": 0},
|
||||
"status": "completed",
|
||||
"metadata": body.get("metadata"),
|
||||
},
|
||||
)
|
||||
|
||||
path = "/v1/vector_stores" if operation == "create" else "/v1/vector_stores/vs_1"
|
||||
headers = {
|
||||
"Authorization": "Bearer sk-master-test",
|
||||
"x-litellm-trace-id": "header-trace",
|
||||
**request_body.get("headers", {}),
|
||||
}
|
||||
body_kwargs = {key: value for key, value in request_body.items() if key != "headers"}
|
||||
with respx.mock(assert_all_called=False) as upstream:
|
||||
upstream.post(url__regex=r"https://api\.openai\.com/v1/vector_stores.*").mock(side_effect=echo)
|
||||
response = TestClient(app).post(path, headers=headers, **body_kwargs)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert len(sent_bodies) == 1
|
||||
assert sent_bodies[0].get("metadata") == CALLERS_METADATA
|
||||
|
|
|
|||
|
|
@ -97,7 +97,7 @@ def test_acount_tokens_with_tools():
|
|||
assert result.total_tokens == 30
|
||||
mock_handler.assert_called_once()
|
||||
call_kwargs = mock_handler.call_args
|
||||
assert call_kwargs.kwargs.get("tools") == tools
|
||||
assert list(call_kwargs.kwargs.get("tools") or ()) == tools
|
||||
|
||||
|
||||
def test_acount_tokens_with_system():
|
||||
|
|
|
|||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Reference in a new issue