mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
1eb7122465
Guard the per-request CPU cost of the chat completion, MCP tool and A2A message transforms against regressions on every commit. All benchmarks are pure in-process work with no network I/O so they stay deterministic under CodSpeed's simulation mode, and they import under the base dependency set the benchmark job installs. Inference covers the full SDK overhead via mock_response (simple, multi-turn, tools, streaming) plus convert_to_model_response_object as a deterministic anchor. MCP covers the client-side tool translation and the proxy server-side tool-name prefix round-trip. A2A covers the client request/response transforms and the proxy server-ingress message conversion. Adds the mcp and a2a-sdk packages to the benchmark run since those transform modules need them, and broadens the workflow triggers to litellm_internal_staging so the internal branch flow is benchmarked too.
76 lines
2.5 KiB
Python
76 lines
2.5 KiB
Python
"""
|
|
Performance benchmarks for the A2A (agent-to-agent) message-translation hot path.
|
|
|
|
Both directions are covered: the client direction (litellm.completion talking to
|
|
an upstream A2A agent) converts OpenAI messages into a prompt and extracts text
|
|
from the A2A response, and the proxy server-ingress direction converts an inbound
|
|
A2A message into OpenAI messages before bridging to a completion. All are pure-CPU
|
|
per-request transforms.
|
|
"""
|
|
|
|
import pytest
|
|
|
|
from litellm.a2a_protocol.litellm_completion_bridge.transformation import (
|
|
A2ACompletionBridgeTransformation,
|
|
)
|
|
from litellm.llms.a2a.common_utils import (
|
|
convert_messages_to_prompt,
|
|
extract_text_from_a2a_response,
|
|
)
|
|
|
|
MESSAGES = [
|
|
{"role": "system", "content": "You are a helpful research assistant."},
|
|
{"role": "user", "content": "What is the capital of France?"},
|
|
{"role": "assistant", "content": "The capital of France is Paris."},
|
|
{"role": "user", "content": "And what is its population?"},
|
|
]
|
|
|
|
MESSAGE_RESPONSE = {
|
|
"result": {
|
|
"kind": "message",
|
|
"parts": [
|
|
{"kind": "text", "text": "The population of Paris is about 2.1 million."},
|
|
{"kind": "text", "text": "The metro area has over 12 million people."},
|
|
],
|
|
}
|
|
}
|
|
|
|
TASK_RESPONSE = {
|
|
"result": {
|
|
"kind": "task",
|
|
"artifacts": [{"parts": [{"kind": "text", "text": "Paris has a population of about 2.1 million."}]}],
|
|
}
|
|
}
|
|
|
|
A2A_INBOUND_MESSAGE = {
|
|
"role": "user",
|
|
"parts": [
|
|
{"kind": "text", "text": "Summarize the latest quarterly report."},
|
|
{"kind": "text", "text": "Focus on revenue and margins."},
|
|
],
|
|
"messageId": "msg-1",
|
|
}
|
|
|
|
|
|
@pytest.mark.benchmark
|
|
def test_convert_messages_to_a2a_prompt():
|
|
"""Benchmark converting OpenAI messages into an A2A prompt string."""
|
|
convert_messages_to_prompt(messages=MESSAGES)
|
|
|
|
|
|
@pytest.mark.benchmark
|
|
def test_extract_text_from_a2a_message_response():
|
|
"""Benchmark extracting text from a direct-message A2A response."""
|
|
extract_text_from_a2a_response(response_dict=MESSAGE_RESPONSE)
|
|
|
|
|
|
@pytest.mark.benchmark
|
|
def test_extract_text_from_a2a_task_response():
|
|
"""Benchmark extracting text from a task-with-artifacts A2A response."""
|
|
extract_text_from_a2a_response(response_dict=TASK_RESPONSE)
|
|
|
|
|
|
@pytest.mark.benchmark
|
|
def test_a2a_inbound_message_to_openai_messages():
|
|
"""Benchmark the proxy converting an inbound A2A message into OpenAI messages."""
|
|
A2ACompletionBridgeTransformation.a2a_message_to_openai_messages(A2A_INBOUND_MESSAGE)
|