fix: a2a gateway - azure foundry invoke

This commit is contained in:
Ishaan Jaffer 2025-12-14 14:06:35 -08:00
parent 71ab4a3108
commit d65d97d609

View file

@ -55,7 +55,6 @@ class A2ACompletionBridgeHandler:
# Get completion params
custom_llm_provider = litellm_params.get("custom_llm_provider")
model = litellm_params.get("model", "agent")
api_key = litellm_params.get("api_key")
# Build full model string if provider specified
# Skip prepending if model already starts with the provider prefix
@ -68,14 +67,22 @@ class A2ACompletionBridgeHandler:
f"A2A completion bridge: model={full_model}, api_base={api_base}"
)
# Build completion params dict
completion_params = {
"model": full_model,
"messages": openai_messages,
"api_base": api_base,
"stream": False,
}
# Add litellm_params (contains api_key, client_id, client_secret, tenant_id, etc.)
litellm_params_to_add = {
k: v for k, v in litellm_params.items()
if k not in ("model", "custom_llm_provider")
}
completion_params.update(litellm_params_to_add)
# Call litellm.acompletion
response = await litellm.acompletion(
model=full_model,
messages=openai_messages,
api_base=api_base,
api_key=api_key,
stream=False,
)
response = await litellm.acompletion(**completion_params)
# Transform response to A2A format
a2a_response = A2ACompletionBridgeTransformation.openai_response_to_a2a_response(
@ -129,7 +136,6 @@ class A2ACompletionBridgeHandler:
# Get completion params
custom_llm_provider = litellm_params.get("custom_llm_provider")
model = litellm_params.get("model", "agent")
api_key = litellm_params.get("api_key")
# Build full model string if provider specified
# Skip prepending if model already starts with the provider prefix
@ -142,6 +148,20 @@ class A2ACompletionBridgeHandler:
f"A2A completion bridge streaming: model={full_model}, api_base={api_base}"
)
# Build completion params dict
completion_params = {
"model": full_model,
"messages": openai_messages,
"api_base": api_base,
"stream": True,
}
# Add litellm_params (contains api_key, client_id, client_secret, tenant_id, etc.)
litellm_params_to_add = {
k: v for k, v in litellm_params.items()
if k not in ("model", "custom_llm_provider")
}
completion_params.update(litellm_params_to_add)
# 1. Emit initial task event (kind: "task", status: "submitted")
task_event = A2ACompletionBridgeTransformation.create_task_event(ctx)
yield task_event
@ -156,13 +176,7 @@ class A2ACompletionBridgeHandler:
yield working_event
# Call litellm.acompletion with streaming
response = await litellm.acompletion(
model=full_model,
messages=openai_messages,
api_base=api_base,
api_key=api_key,
stream=True,
)
response = await litellm.acompletion(**completion_params)
# 3. Accumulate content and emit artifact update
accumulated_text = ""