diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260305000000_add_agent_headers/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260305000000_add_agent_headers/migration.sql new file mode 100644 index 00000000000..acb35baba96 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260305000000_add_agent_headers/migration.sql @@ -0,0 +1,5 @@ +-- Add static_headers and extra_headers to LiteLLM_AgentsTable + +ALTER TABLE "LiteLLM_AgentsTable" + ADD COLUMN IF NOT EXISTS "static_headers" JSONB DEFAULT '{}', + ADD COLUMN IF NOT EXISTS "extra_headers" TEXT[] DEFAULT ARRAY[]::TEXT[]; diff --git a/litellm/a2a_protocol/main.py b/litellm/a2a_protocol/main.py index 485b57e311b..b9eb26f76c2 100644 --- a/litellm/a2a_protocol/main.py +++ b/litellm/a2a_protocol/main.py @@ -162,6 +162,31 @@ async def _send_message_via_completion_bridge( return LiteLLMSendMessageResponse.from_dict(response_dict) +async def _create_a2a_client_with_headers( + api_base: Optional[str], + trace_id: Optional[str], + agent_id: Optional[str], +) -> tuple: + """Create an A2A client with LiteLLM trace headers.""" + if api_base is None: + raise ValueError("Either a2a_client or api_base is required for standard A2A flow") + trace_id = trace_id or str(uuid.uuid4()) + extra_headers: Dict[str, str] = {"X-LiteLLM-Trace-Id": trace_id} + if agent_id: + extra_headers["X-LiteLLM-Agent-Id"] = agent_id + return await create_a2a_client(base_url=api_base, extra_headers=extra_headers), trace_id + + +def _set_message_context_id(message: Any, context_id: str) -> None: + """Set context_id on an A2A message if not already set.""" + if isinstance(message, dict): + if message.get("context_id") is None: + message["context_id"] = context_id + else: + if getattr(message, "context_id", None) is None: + message.context_id = context_id + + @client async def asend_message( a2a_client: Optional["A2AClientType"] = None, @@ -245,16 +270,8 @@ async def asend_message( # Create A2A client if not provided but api_base is available if a2a_client is None: - if api_base is None: - raise ValueError( - "Either a2a_client or api_base is required for standard A2A flow" - ) - trace_id = trace_id or str(uuid.uuid4()) - extra_headers = {"X-LiteLLM-Trace-Id": trace_id} - if agent_id: - extra_headers["X-LiteLLM-Agent-Id"] = agent_id - a2a_client = await create_a2a_client( - base_url=api_base, extra_headers=extra_headers + a2a_client, trace_id = await _create_a2a_client_with_headers( + api_base=api_base, trace_id=trace_id, agent_id=agent_id ) # Type assertion: a2a_client is guaranteed to be non-None here @@ -271,13 +288,7 @@ async def asend_message( card_url = getattr(agent_card, "url", None) if agent_card else None context_id = trace_id or str(uuid.uuid4()) - message = request.params.message - if isinstance(message, dict): - if message.get("context_id") is None: - message["context_id"] = context_id - else: - if getattr(message, "context_id", None) is None: - message.context_id = context_id + _set_message_context_id(request.params.message, context_id) # Retry loop: if connection fails due to localhost URL in agent card, retry with fixed URL a2a_response = None diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 80600585362..6f4ef0c24b6 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -63,6 +63,8 @@ model LiteLLM_AgentsTable { agent_name String @unique litellm_params Json? agent_card_params Json + static_headers Json? @default("{}") + extra_headers String[] @default([]) agent_access_groups String[] @default([]) object_permission_id String? object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id]) diff --git a/schema.prisma b/schema.prisma index 80600585362..6f4ef0c24b6 100644 --- a/schema.prisma +++ b/schema.prisma @@ -63,6 +63,8 @@ model LiteLLM_AgentsTable { agent_name String @unique litellm_params Json? agent_card_params Json + static_headers Json? @default("{}") + extra_headers String[] @default([]) agent_access_groups String[] @default([]) object_permission_id String? object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id]) diff --git a/tests/test_litellm/llms/azure/chat/test_azure_gpt5_transformation.py b/tests/test_litellm/llms/azure/chat/test_azure_gpt5_transformation.py index 25f3d1364f6..0cd3c4682ba 100644 --- a/tests/test_litellm/llms/azure/chat/test_azure_gpt5_transformation.py +++ b/tests/test_litellm/llms/azure/chat/test_azure_gpt5_transformation.py @@ -181,15 +181,15 @@ def test_azure_gpt5_1_temperature_with_reasoning_effort_other_values(config: Azu def test_azure_gpt5_1_series_temperature_handling(config: AzureOpenAIGPT5Config): - """Test that Azure GPT-5.1 with gpt5_series prefix supports temperature with reasoning_effort='none'.""" + """Test that Azure GPT-5.1 with gpt5_series prefix passes through temperature=1.""" params = config.map_openai_params( - non_default_params={"temperature": 0.6}, + non_default_params={"temperature": 1}, optional_params={}, model="gpt5_series/gpt-5.1", drop_params=False, api_version="2024-05-01-preview", ) - assert params["temperature"] == 0.6 + assert params["temperature"] == 1 def test_azure_gpt5_reasoning_effort_none_error(config: AzureOpenAIGPT5Config): diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 238167713c7..51b6d474308 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -795,6 +795,8 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): }, }, "supports_native_streaming": {"type": "boolean"}, + "supports_none_reasoning_effort": {"type": "boolean"}, + "supports_xhigh_reasoning_effort": {"type": "boolean"}, "tiered_pricing": { "type": "array", "items": {