mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-29 06:33:28 +00:00
ea6dbd2109
* feat(proxy): log one line per token count call Each count call on /v1/messages/count_tokens, /utils/token_counter, /v1/responses/input_tokens and the Gemini countTokens routes now logs one INFO line with the litellm call id, the requested model, the model the count ran against, the count, the count source and the status code, so it joins the inference call that sends the same x-litellm-call-id * test(google_endpoints): give the count test's fake request a headers list * fix(proxy): keep token count log lines one per call and trusted in JSON logs, and log Bedrock native counts Quote log values that carry whitespace, quotes, equals, or control characters so a model name can't split or forge a line. Let JsonFormatter record extras win over a dict parsed from the message. Log the Bedrock native count-tokens route through the same line. * fix(proxy): escape braces in token count log values, log the Bedrock model AWS counts, and skip the line when there is nothing to count * fix(proxy): log the deployment id the count routed to on every token count line * fix(proxy): log count requests rejected before counting with the status they answer A count route that rejects a request before the shared count function runs now logs the same token count line with its 400. An unsupported content block on the Messages count route becomes its 400 inside the logged wrapper, so its line says 400 instead of 500. A count that ends in a BaseException other than cancellation logs a 500. * fix(proxy): compute the Bedrock counted model id only after a successful count * revert(logging): keep the JSON log formatter's field order as on main The count line already escapes braces in its values, so the formatter never parses a dict from it. Changing the order for every logger is a separate decision.
107 lines
3.3 KiB
Python
107 lines
3.3 KiB
Python
"""
|
|
Test for google_endpoints/endpoints.py
|
|
"""
|
|
|
|
from typing import Final
|
|
|
|
import pytest
|
|
from dotenv import load_dotenv
|
|
from pydantic import JsonValue
|
|
from starlette.requests import Request
|
|
|
|
from litellm.proxy.google_endpoints.endpoints import _google_count_tokens_request, google_count_tokens
|
|
from litellm.types.llms.vertex_ai import TokenCountDetailsResponse
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("system_key", "schema_key"),
|
|
(("systemInstruction", "parameters"), ("system_instruction", "parametersJsonSchema")),
|
|
)
|
|
def test_google_count_tokens_translates_system_and_function_declarations(
|
|
system_key: str, schema_key: str
|
|
) -> None:
|
|
schema: Final = {"type": "object", "properties": {"query": {"type": "string"}}}
|
|
request_body: Final[dict[str, JsonValue]] = {
|
|
"contents": [{"role": "user", "parts": [{"text": "Look this up"}]}],
|
|
system_key: {"parts": [{"text": "Use the provided tool."}]},
|
|
"tools": [
|
|
{
|
|
"functionDeclarations": [
|
|
{
|
|
"name": "lookup",
|
|
"description": "Find a record",
|
|
schema_key: schema,
|
|
}
|
|
]
|
|
}
|
|
],
|
|
}
|
|
request: Final = _google_count_tokens_request("claude-sonnet", request_body)
|
|
|
|
assert request.system == [{"type": "text", "text": "Use the provided tool."}]
|
|
assert request.tools == [
|
|
{
|
|
"name": "lookup",
|
|
"description": "Find a record",
|
|
"input_schema": schema,
|
|
}
|
|
]
|
|
assert request.messages == [{"role": "user", "content": "Look this up"}]
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("schema_key", "expected_schema"),
|
|
(
|
|
("parameters", {"type": "object", "properties": {"query": {"type": "string"}}}),
|
|
("parametersJsonSchema", {"type": "OBJECT", "properties": {"query": {"type": "STRING"}}}),
|
|
),
|
|
)
|
|
def test_google_count_tokens_normalizes_parameters_only(
|
|
schema_key: str, expected_schema: dict[str, JsonValue]
|
|
) -> None:
|
|
request: Final = _google_count_tokens_request(
|
|
"claude-sonnet",
|
|
{
|
|
"contents": [{"role": "user", "parts": [{"text": "Look this up"}]}],
|
|
"tools": [
|
|
{
|
|
"functionDeclarations": [
|
|
{
|
|
"name": "lookup",
|
|
schema_key: {
|
|
"type": "OBJECT",
|
|
"properties": {"query": {"type": "STRING"}},
|
|
},
|
|
}
|
|
]
|
|
}
|
|
],
|
|
},
|
|
)
|
|
|
|
assert request.tools == [{"name": "lookup", "input_schema": expected_schema}]
|
|
|
|
load_dotenv()
|
|
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_proxy_gemini_to_openai_like_model_token_counting():
|
|
"""
|
|
Test the token counting endpoint for proxing gemini to openai-like models.
|
|
"""
|
|
response: TokenCountDetailsResponse = await google_count_tokens(
|
|
request=Request(
|
|
scope={
|
|
"type": "http",
|
|
"headers": [],
|
|
"parsed_body": (
|
|
["contents"],
|
|
{"contents": [{"parts": [{"text": "Hello, how are you?"}]}]},
|
|
),
|
|
}
|
|
),
|
|
model_name="volcengine/foo",
|
|
)
|
|
|
|
assert response.get("totalTokens") > 0
|