mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-16 04:29:31 +00:00
fix(count_tokens): count Anthropic-format requests for Gemini the way /v1/messages inference converts them (#46767)
* fix(count_tokens): convert Anthropic-format messages for the OpenAI counter the way inference does * test(count_tokens): check /v1/messages/count_tokens sends OpenAI the input /v1/messages inference sends * fix(count_tokens): build the Anthropic count request without a cast and sync dashboard API types * fix(count_tokens): count Anthropic-format requests for Gemini the way /v1/messages inference converts them * test(count_tokens): compare tools by value in the public api test * fix(count_tokens): count Anthropic bodies the way /v1/messages inference sends them to OpenAI Unwrap affinity-encoded reasoning, count the structured-output schema, flatten replayed web search results, join a list system prompt, and answer a message without a role with a 400 before calling OpenAI. messages_format is now internal to the Anthropic count route instead of a public /utils/token_counter field. * test(count_tokens): expect the image detail the adapter now sends after merging main * fix(count_tokens): fetch Gemini remote media without blocking the event loop * fix(count_tokens): type normalize_native_messages without the unimported Any On Python 3.10 to 3.13 the annotation is evaluated at import, so the proxy failed to start with NameError * chore(count_tokens): give the Gemini count helpers' list returns a reason * fix(count_tokens): normalize Anthropic messages for Gemini counts the way /v1/messages does * fix(count_tokens): count Gemini bodies after inference's history rewrites and answer converter errors with 400 * test(count_tokens): cover the Gemini count error paths and compaction with a text-block system --------- Co-authored-by: Mrinal Chanshetty <mrinal@berri.ai>
This commit is contained in:
Mateo Wang
GitHub
parent
dfb949d437
commit
dba86bec28
9 changed files with 703 additions and 63 deletions
|
|
@ -465,7 +465,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(
|
||||
|
|
|
|||
|
|
@ -504,25 +504,41 @@ class GoogleAIStudioTokenCounter(BaseTokenCounter):
|
|||
from litellm.llms.vertex_ai.count_tokens.context import (
|
||||
VertexGeminiCountContextError,
|
||||
count_context_error_response,
|
||||
gemini_count_contents,
|
||||
inference_request_view,
|
||||
provider_error_response,
|
||||
vertex_gemini_count_context,
|
||||
)
|
||||
|
||||
count_context: Final = vertex_gemini_count_context(system=system, tools=tools)
|
||||
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)
|
||||
|
||||
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(
|
||||
|
|
|
|||
|
|
@ -1359,24 +1359,27 @@ class VertexAITokenCounter(BaseTokenCounter):
|
|||
from litellm.llms.vertex_ai.count_tokens.context import (
|
||||
VertexGeminiCountContextError,
|
||||
count_context_error_response,
|
||||
gemini_count_contents,
|
||||
inference_request_view,
|
||||
provider_error_response,
|
||||
vertex_gemini_count_context,
|
||||
)
|
||||
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
|
||||
)
|
||||
|
||||
count_context: Final = vertex_gemini_count_context(system=system, tools=tools)
|
||||
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)
|
||||
|
||||
resolved_contents: Final = (
|
||||
contents
|
||||
if contents is not None
|
||||
else gemini_convert_messages_with_history(
|
||||
messages=messages or [] # pyright: ignore[reportArgumentType] # preserve direct caller message handling
|
||||
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,
|
||||
|
|
@ -1390,15 +1393,7 @@ class VertexAITokenCounter(BaseTokenCounter):
|
|||
**count_tokens_params_request,
|
||||
)
|
||||
except litellm.APIError as error:
|
||||
return TokenCountResponse(
|
||||
total_tokens=0,
|
||||
request_model=request_model,
|
||||
model_used=model_to_use,
|
||||
tokenizer_type="gemini",
|
||||
error=True,
|
||||
error_message=error.message,
|
||||
status_code=error.status_code,
|
||||
)
|
||||
return provider_error_response(error, request_model, model_to_use)
|
||||
|
||||
if "totalTokens" in result:
|
||||
return TokenCountResponse(
|
||||
|
|
|
|||
|
|
@ -1,12 +1,20 @@
|
|||
from collections.abc import Mapping
|
||||
import copy
|
||||
from collections.abc import Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
from typing import TYPE_CHECKING, Final, Literal
|
||||
|
||||
from pydantic import JsonValue, TypeAdapter, ValidationError
|
||||
from typing_extensions import assert_never
|
||||
|
||||
from litellm.types.llms.anthropic import AnthropicMessagesTextParam
|
||||
from litellm.types.utils import TokenCountResponse
|
||||
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)
|
||||
|
|
@ -142,3 +150,103 @@ def count_context_error_response(
|
|||
error_message=error.error_message,
|
||||
status_code=400,
|
||||
)
|
||||
|
||||
|
||||
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)(
|
||||
messages=chat_messages, model=model, litellm_params=dict(litellm_params)
|
||||
)
|
||||
except BadRequestError as error:
|
||||
return VertexGeminiCountContextError(error_message=error.message)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
@ -1318,6 +1318,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],
|
||||
|
|
@ -1360,7 +1364,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):
|
||||
|
|
|
|||
|
|
@ -167,6 +167,93 @@ def test_messages_and_count_tokens_send_matching_gemini_context(
|
|||
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) -> str:
|
||||
return scenario.model(model=f"gemini/{_MODEL}", api_base=api_base, api_key="scripted-gemini-key")
|
||||
|
||||
|
||||
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"]
|
||||
|
||||
|
||||
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:
|
||||
|
|
|
|||
|
|
@ -5,6 +5,7 @@ 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
|
||||
|
|
@ -100,6 +101,58 @@ async def test_google_ai_studio_count_tokens_rejects_image_system_without_reques
|
|||
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"),
|
||||
|
|
@ -129,9 +182,7 @@ async def test_gemini_count_tokens_sends_optional_request_fields(
|
|||
expected_body: dict[str, JsonValue],
|
||||
) -> None:
|
||||
peer: Final = _GeminiCountTokensPeer()
|
||||
token_counter: Final = GeminiCountTokensHandler(
|
||||
client=AsyncHTTPHandler(transport=httpx.MockTransport(peer))
|
||||
)
|
||||
token_counter: Final = GeminiCountTokensHandler(client=AsyncHTTPHandler(transport=httpx.MockTransport(peer)))
|
||||
|
||||
result: Final = await token_counter.acount_tokens(
|
||||
contents=[{"role": "user", "parts": [{"text": "hello"}]}],
|
||||
|
|
@ -177,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)
|
||||
|
|
@ -239,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
|
||||
|
||||
|
|
@ -320,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."},
|
||||
}
|
||||
}
|
||||
],
|
||||
|
|
@ -331,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"]
|
||||
|
|
@ -342,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):
|
||||
|
|
@ -359,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
|
||||
|
|
|
|||
|
|
@ -1,8 +1,18 @@
|
|||
from typing import Final
|
||||
import asyncio
|
||||
import base64
|
||||
import time
|
||||
import uuid
|
||||
from collections.abc import Generator
|
||||
from contextlib import contextmanager
|
||||
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
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
|
||||
import litellm
|
||||
from litellm.llms.anthropic.pass_through.adapters.transformation import (
|
||||
LiteLLMAnthropicMessagesAdapter,
|
||||
truncate_tool_name,
|
||||
|
|
@ -10,6 +20,8 @@ from litellm.llms.anthropic.pass_through.adapters.transformation import (
|
|||
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
|
||||
|
|
@ -193,3 +205,302 @@ def test_invalid_tools_return_an_error_value(tools: JsonValue) -> None:
|
|||
|
||||
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,5 +1,6 @@
|
|||
import time
|
||||
from collections.abc import Iterator
|
||||
from contextlib import contextmanager
|
||||
from typing import Final
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
|
|
@ -7,6 +8,7 @@ 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
|
||||
|
|
@ -24,16 +26,21 @@ from litellm.llms.vertex_ai.common_utils import (
|
|||
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
|
||||
async def test_vertex_gemini_rejects_non_text_system_blocks_before_request() -> None:
|
||||
@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
|
||||
|
|
@ -42,6 +49,88 @@ async def test_vertex_gemini_rejects_non_text_system_blocks_before_request() ->
|
|||
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"])
|
||||
def test_validate_vertex_location_accepts_valid(location):
|
||||
assert validate_vertex_location(location) == location
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue