mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-29 06:33:28 +00:00
464932b1d1
* refactor(typing): narrow Any sources across providers, proxy and guardrails (wave 1) * refactor(typing): narrow Any sources across hubs, providers and proxy (wave 2) * refactor(typing): validate RunwayML task JSON and let dotprompt loaders accept exported prompts * fix: restore pass-through behavior on paths the typing pass tightened MCP client IP extraction trusts X-Forwarded-For only when use_x_forwarded_for is True, dotprompt skips a malformed JSON prompt entry instead of raising, RunwayML task payloads pass through unvalidated, and a malformed provider body raises the same AttributeError as before at the CompactifAI, Gemini, and mock transport parse sites * fix(runwayml): keep a malformed 2xx task body on the 500 path * fix(edenai,cli): keep a non-object body in the error and fail on a malformed migration report EdenAI's embedding and image generation parse dropped a non-object body from the error they raised, so the message no longer named what the provider sent. A helper now raises with the body in the message, and both paths have a regression test. `lite encryption migrate` exited 0 with no message when the proxy's report was not an object or its residual_legacy was not a number, where base crashed. It now fails with a message naming the value. The CompactifAI and Gemini image-edit non-object body tests accept TypeError or AttributeError, which both base and this branch raise, so they guard against a ValidationError instead of pinning one class. RunwayML's async parse gets the same non-object body test as the sync one. * fix: keep base's error classes, duck-typed checks and fail-closed guardrails the typing pass changed * test: annotate the new regression tests' locals with Final * test: annotate the remaining new test locals with Final
55 lines
1.9 KiB
Python
55 lines
1.9 KiB
Python
"""Tests for litellm/a2a_protocol/utils.py token/usage extraction."""
|
|
|
|
from types import SimpleNamespace
|
|
from typing import Final
|
|
|
|
import pytest
|
|
|
|
pytest.importorskip("a2a.compat.v0_3.types")
|
|
|
|
from a2a.compat.v0_3.types import MessageSendParams, SendMessageRequest
|
|
|
|
from litellm.a2a_protocol.utils import A2ARequestUtils
|
|
from tests.unit.litellm_core_utils.attribute_forwarder import AttributeForwarder
|
|
|
|
|
|
def _request(user_text: str) -> SendMessageRequest:
|
|
return SendMessageRequest(
|
|
id="r1",
|
|
params=MessageSendParams(
|
|
message={
|
|
"messageId": "m1",
|
|
"role": "user",
|
|
"parts": [{"kind": "text", "text": user_text}],
|
|
}
|
|
),
|
|
)
|
|
|
|
|
|
def test_calculate_usage_counts_input_tokens_from_request_object():
|
|
"""Regression: request-side Part is a RootModel; input tokens must be counted."""
|
|
request = _request("count these input tokens please")
|
|
response_dict = {
|
|
"result": {
|
|
"kind": "message",
|
|
"parts": [{"kind": "text", "text": "ok"}],
|
|
}
|
|
}
|
|
|
|
prompt_tokens, completion_tokens, total_tokens = A2ARequestUtils.calculate_usage_from_request_response(
|
|
request=request, response_dict=response_dict
|
|
)
|
|
|
|
assert prompt_tokens > 0
|
|
assert completion_tokens > 0
|
|
assert total_tokens == prompt_tokens + completion_tokens
|
|
|
|
|
|
def test_calculate_usage_counts_input_tokens_from_a_message_that_forwards_its_fields():
|
|
request: Final = _request("count these input tokens please")
|
|
forwarded_request: Final = SimpleNamespace(params=SimpleNamespace(message=AttributeForwarder(request.params.message)))
|
|
response_dict: Final = {"result": {"kind": "message", "parts": [{"kind": "text", "text": "ok"}]}}
|
|
|
|
assert A2ARequestUtils.calculate_usage_from_request_response(
|
|
request=forwarded_request, response_dict=response_dict
|
|
) == A2ARequestUtils.calculate_usage_from_request_response(request=request, response_dict=response_dict)
|