From 24e18e1549c5b3bebc28ded0341d0dd7f055475b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 19:57:25 -0700 Subject: [PATCH] fix(azure_ai): keep system prompts and project-endpoint auth on Foundry agent calls and follow-ups --- .../llms/azure_ai/responses/transformation.py | 56 ++++++++---- .../public_endpoints/agent_create_fields.json | 4 +- .../test_azure_ai_responses_transformation.py | 86 +++++++++++++++++-- 3 files changed, 117 insertions(+), 29 deletions(-) diff --git a/litellm/llms/azure_ai/responses/transformation.py b/litellm/llms/azure_ai/responses/transformation.py index 3f47ad30c11..b7c69e4b507 100644 --- a/litellm/llms/azure_ai/responses/transformation.py +++ b/litellm/llms/azure_ai/responses/transformation.py @@ -1,7 +1,8 @@ -from collections.abc import Collection, Mapping +from collections.abc import Collection, Mapping, Sequence from typing import Final import httpx +from pydantic import TypeAdapter from litellm.llms.azure.responses.transformation import AzureOpenAIResponsesAPIConfig from litellm.llms.azure_ai.chat.transformation import AzureAIGPT5Config @@ -18,14 +19,17 @@ from litellm.types.utils import LlmProviders _PROJECT_PATH_PREFIX: Final = ("api", "projects") _RESPONSES_PATH: Final = ("openai", "v1", "responses") -_AGENT_OWNED_PARAMS: Final = frozenset({"instructions", "temperature", "top_p", "tools", "text", "reasoning"}) +_AGENT_OWNED_PARAMS: Final = frozenset({"temperature", "top_p", "tools", "text", "reasoning"}) + + +def _project_root(base_url: httpx.URL) -> tuple[str, ...]: + segments: Final = tuple(segment for segment in base_url.path.split("/") if segment) + return segments[:3] if segments[:2] == _PROJECT_PATH_PREFIX else () def _responses_url(api_base: str) -> str: base_url: Final = httpx.URL(api_base) - segments: Final = tuple(segment for segment in base_url.path.split("/") if segment) - project_root: Final = segments[:3] if segments[:2] == _PROJECT_PATH_PREFIX else () - return str(base_url.copy_with(path="/" + "/".join((*project_root, *_RESPONSES_PATH)), query=None)) + return str(base_url.copy_with(path="/" + "/".join((*_project_root(base_url), *_RESPONSES_PATH)), query=None)) class AzureAIResponsesAPIConfig(AzureOpenAIResponsesAPIConfig): @@ -47,10 +51,17 @@ class AzureAIResponsesAPIConfig(AzureOpenAIResponsesAPIConfig): def validate_environment(self, headers: dict, model: str, litellm_params: GenericLiteLLMParams | None) -> dict: params: Final = litellm_params or GenericLiteLLMParams() + api_base: Final = AzureFoundryModelInfo.get_api_base(params.api_base) + on_project: Final = api_base is not None and bool(_project_root(httpx.URL(api_base))) + auth_params: Final = ( + params.model_copy(update={"azure_scope": params.azure_scope or AZURE_AI_AGENTS_SCOPE}) + if on_project + else params + ) auth_headers: Final = get_azure_ai_auth_headers( api_key=AzureFoundryModelInfo.get_api_key(params.api_key), - litellm_params=params.model_dump(), - api_key_header=api_key_header_for_base(AzureFoundryModelInfo.get_api_base(params.api_base)), + litellm_params=auth_params.model_dump(), + api_key_header="Authorization" if on_project else api_key_header_for_base(api_base), ) return { **headers, @@ -78,18 +89,19 @@ def _agent_reference(model: str) -> FoundryAgentReference: return FoundryAgentReference(type="agent_reference", name=name) -class AzureAIAgentResponsesAPIConfig(AzureAIResponsesAPIConfig): - def validate_environment( - self, - headers: dict[str, str], # mutable-ok: signature fixed by the override contract - model: str, - litellm_params: GenericLiteLLMParams | None, - ) -> dict[str, str]: # mutable-ok: same - params: Final = litellm_params or GenericLiteLLMParams() - return super().validate_environment( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # parent returns a bare dict - headers, model, params.model_copy(update={"azure_scope": params.azure_scope or AZURE_AI_AGENTS_SCOPE}) - ) +_INPUT_ITEMS: Final = TypeAdapter(list[object]) + +def _with_system_prompt(instructions: str, request_input: object) -> Sequence[object]: + conversation: Final = ( + [{"type": "message", "role": "user", "content": request_input}] + if isinstance(request_input, str) + else _INPUT_ITEMS.validate_python(request_input) + ) + return [{"type": "message", "role": "system", "content": instructions}, *conversation] + + +class AzureAIAgentResponsesAPIConfig(AzureAIResponsesAPIConfig): def should_fake_stream( self, model: str | None, stream: bool | None, custom_llm_provider: str | None = None ) -> bool: @@ -127,8 +139,14 @@ class AzureAIAgentResponsesAPIConfig(AzureAIResponsesAPIConfig): litellm_params=litellm_params, headers=headers, ) + instructions: Final = request.get("instructions") return { - **{key: value for key, value in request.items() if key != "model"}, + **{key: value for key, value in request.items() if key not in ("model", "instructions")}, + **( + {"input": _with_system_prompt(instructions, request.get("input"))} + if isinstance(instructions, str) and instructions + else {} + ), "agent_reference": _agent_reference(model), } diff --git a/litellm/proxy/public_endpoints/agent_create_fields.json b/litellm/proxy/public_endpoints/agent_create_fields.json index f289a6c1805..53b32d249c7 100644 --- a/litellm/proxy/public_endpoints/agent_create_fields.json +++ b/litellm/proxy/public_endpoints/agent_create_fields.json @@ -127,8 +127,8 @@ { "key": "agent_id", "label": "Agent ID", - "placeholder": "asst_abc123", - "tooltip": "The agent/assistant ID from your Azure AI Foundry project (e.g., asst_abc123)", + "placeholder": "my-agent", + "tooltip": "The agent name from your Azure AI Foundry project, optionally with a version (e.g., my-agent or my-agent:2), or a classic assistant ID (e.g., asst_abc123)", "required": true, "field_type": "text", "default_value": null, diff --git a/tests/unit/llms/azure_ai/responses/test_azure_ai_responses_transformation.py b/tests/unit/llms/azure_ai/responses/test_azure_ai_responses_transformation.py index 40459856e1d..d7b285453a6 100644 --- a/tests/unit/llms/azure_ai/responses/test_azure_ai_responses_transformation.py +++ b/tests/unit/llms/azure_ai/responses/test_azure_ai_responses_transformation.py @@ -163,19 +163,26 @@ def test_native_websocket_stays_off(): assert AzureAIResponsesAPIConfig().supports_native_websocket() is False -def test_validate_environment_sends_api_key_header(): +@pytest.mark.parametrize( + "api_base,expected_auth", + [ + ("https://res.services.ai.azure.com", {"api-key": "secret"}), + (FOUNDRY_PROJECT_BASE, {"Authorization": "Bearer secret"}), + ], +) +def test_validate_environment_sends_the_api_key_the_way_the_endpoint_takes_it(api_base, expected_auth): headers = AzureAIResponsesAPIConfig().validate_environment( headers={"x-custom": "1"}, model="gpt-5.6-luna", - litellm_params=GenericLiteLLMParams(api_key="secret", api_base=FOUNDRY_PROJECT_BASE), + litellm_params=GenericLiteLLMParams(api_key="secret", api_base=api_base), ) - assert headers == {"x-custom": "1", "api-key": "secret", "Content-Type": "application/json"} + assert headers == {"x-custom": "1", **expected_auth, "Content-Type": "application/json"} def test_validate_environment_reads_api_key_from_env(monkeypatch): monkeypatch.setenv("AZURE_AI_API_KEY", "env-secret") headers = AzureAIResponsesAPIConfig().validate_environment( - headers={}, model="gpt-5.6-luna", litellm_params=GenericLiteLLMParams(api_base=FOUNDRY_PROJECT_BASE) + headers={}, model="gpt-5.6-luna", litellm_params=GenericLiteLLMParams(api_base="https://res.services.ai.azure.com") ) assert headers["api-key"] == "env-secret" @@ -230,7 +237,10 @@ def _assert_native_responses_request(route, expected_url, expected_model): request = route.calls.last.request body = json.loads(request.content) assert f"{request.url.scheme}://{request.url.host}{request.url.path}" == expected_url - assert request.headers["api-key"] == "fake-key" + on_project: Final = expected_url == FOUNDRY_RESPONSES_URL + assert request.headers.get("authorization" if on_project else "api-key") == ( + "Bearer fake-key" if on_project else "fake-key" + ) assert body["model"] == expected_model assert body["input"] == "What is the weather in SF?" assert "messages" not in body @@ -447,7 +457,6 @@ async def test_router_aresponses_keeps_the_deployment_agent_over_extra_body(): AGENT_OWNED_PARAMS = { - "instructions": "Answer in French", "temperature": 0.2, "top_p": 0.5, "tools": [WEATHER_TOOL], @@ -517,6 +526,39 @@ async def test_aresponses_rejects_agent_owned_params_without_drop_params(param): assert not route.called +SYSTEM_PROMPT = {"type": "message", "role": "system", "content": "Answer in French"} + + +@pytest.mark.asyncio +@respx.mock +@pytest.mark.parametrize( + "request_input,expected_input", + [ + ("What is 2+2?", [SYSTEM_PROMPT, {"type": "message", "role": "user", "content": "What is 2+2?"}]), + ( + [{"type": "message", "role": "user", "content": "What is 2+2?"}], + [SYSTEM_PROMPT, {"type": "message", "role": "user", "content": "What is 2+2?"}], + ), + ], +) +async def test_aresponses_sends_instructions_to_the_agent_as_a_leading_system_message(request_input, expected_input): + route = respx.post(url__regex=r".*/openai/v1/responses(\?.*)?$").mock( + return_value=httpx.Response(200, json=_responses_payload("gpt-4.1")) + ) + + await litellm.aresponses( + model="azure_ai/agents/support-agent", + input=request_input, + instructions="Answer in French", + api_base=FOUNDRY_PROJECT_BASE, + api_key="fake-key", + ) + + body = json.loads(route.calls.last.request.content) + assert "instructions" not in body + assert body["input"] == expected_input + + @pytest.mark.asyncio @respx.mock async def test_aresponses_streams_agent_answers_from_foundry(): @@ -553,7 +595,7 @@ async def test_aresponses_streams_agent_answers_from_foundry(): "configured_scope,expected_scope", [(None, AZURE_AI_AGENTS_SCOPE), ("api://custom-app/.default", "api://custom-app/.default")], ) -async def test_aresponses_mints_the_service_principal_token_for_the_agents_scope(configured_scope, expected_scope): +async def test_aresponses_mints_the_service_principal_token_for_the_project_scope(configured_scope, expected_scope): route = respx.post(url__regex=r".*/openai/v1/responses(\?.*)?$").mock( return_value=httpx.Response(200, json=_responses_payload("gpt-4.1")) ) @@ -576,6 +618,32 @@ async def test_aresponses_mints_the_service_principal_token_for_the_agents_scope assert route.calls.last.request.headers["Authorization"] == "Bearer agent-sp-token" +@pytest.mark.asyncio +@respx.mock +async def test_aget_responses_mints_the_service_principal_token_for_the_project_scope(): + route = respx.get(url__regex=r".*/api/projects/proj/openai/v1/responses/resp_123(\?.*)?$").mock( + return_value=httpx.Response(200, json=_responses_payload("gpt-4.1")) + ) + + with patch("litellm.llms.azure.common_utils.get_azure_ad_token_from_entra_id") as mock_entra_id: # test-quality-ok: stubs the Entra token fetch to assert the scope a follow-up on a project endpoint requests; live service principal path proven by the PR's Foundry e2e QA + mock_entra_id.return_value = lambda: "project-sp-token" + response = await litellm.aget_responses( + response_id="resp_123", + custom_llm_provider="azure_ai", + model="azure_ai/agents/support-agent", + api_base=FOUNDRY_PROJECT_BASE, + tenant_id="tenant", + client_id="client", + client_secret="secret", + ) + + mock_entra_id.assert_called_once_with( + tenant_id="tenant", client_id="client", client_secret="secret", scope=AZURE_AI_AGENTS_SCOPE + ) + assert route.calls.last.request.headers["Authorization"] == "Bearer project-sp-token" + assert response.status == "completed" + + def _agent_answer_payload(text: str) -> dict: return { **_responses_payload("gpt-4.1"), @@ -600,7 +668,7 @@ async def test_acompletion_reaches_the_agent_through_the_project_responses_api() response = await litellm.acompletion( model="azure_ai/agents/support-agent", - messages=[{"role": "user", "content": "What is 2+2?"}], + messages=[{"role": "system", "content": "Answer in French"}, {"role": "user", "content": "What is 2+2?"}], api_base=FOUNDRY_PROJECT_BASE, api_key="fake-key", ) @@ -608,4 +676,6 @@ async def test_acompletion_reaches_the_agent_through_the_project_responses_api() body = json.loads(route.calls.last.request.content) assert body["agent_reference"] == {"type": "agent_reference", "name": "support-agent"} assert "model" not in body + assert "instructions" not in body + assert body["input"][0] == SYSTEM_PROMPT assert response.choices[0].message.content == "4"