tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-12 20:22:46 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #23939 from Sameerlite/Sameerlite/azure-ai-annotations

Browse Source 
fix(azure-ai-agents): preserve annotations in Bing Search grounding responses
This commit is contained in:
Sameer Kankute
2026-03-20 23:33:08 +05:30
committed by GitHub
parent 8d843fd93b 6514446dcb
commit 0673c57fef
2 changed files with 371 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/llms/azure_ai/agents/handler.py
+94 -18
View File
@@ -97,14 +97,63 @@ class AzureAIAgentsHandler:
# -------------------------------------------------------------------------
# Response Helpers
# -------------------------------------------------------------------------
def _extract_content_from_messages(self, messages_data: dict) -> str:
"""Extract assistant content from the messages response."""
def _extract_content_from_messages(
self, messages_data: dict
) -> Tuple[str, Optional[List[Dict[str, Any]]]]:
"""Extract assistant content and annotations from the messages response.
Returns (content, annotations) where annotations is a list of
OpenAI-compatible ChatCompletionAnnotation dicts, or None.
"""
for msg in messages_data.get("data", []):
if msg.get("role") == "assistant":
for content_item in msg.get("content", []):
if content_item.get("type") == "text":
return content_item.get("text", {}).get("value", "")
return ""
text_obj = content_item.get("text", {})
content = text_obj.get("value", "")
raw_annotations = text_obj.get("annotations")
annotations = self._transform_annotations(
raw_annotations
)
return content, annotations
return "", None
def _transform_annotations(
self,
raw_annotations: Optional[List[Dict[str, Any]]],
) -> Optional[List[Dict[str, Any]]]:
"""Transform Azure AI Foundry annotations to OpenAI-compatible format.
Azure AI returns annotations like:
{"type": "url_citation", "text": "[1]", "start_index": 10,
"end_index": 13, "url_citation": {"url": "...", "title": "..."}}
OpenAI expects:
{"type": "url_citation", "url_citation": {"url": "...", "title": "...",
"start_index": 10, "end_index": 13}}
"""
if not raw_annotations:
return None
result: List[Dict[str, Any]] = []
for ann in raw_annotations:
ann_type = ann.get("type")
if ann_type == "url_citation":
url_citation = dict(ann.get("url_citation", {}))
# Azure puts start/end_index at annotation level; OpenAI
# expects them inside url_citation
if "start_index" in ann and "start_index" not in url_citation:
url_citation["start_index"] = ann["start_index"]
if "end_index" in ann and "end_index" not in url_citation:
url_citation["end_index"] = ann["end_index"]
result.append(
{"type": "url_citation", "url_citation": url_citation}
)
else:
# Pass through unknown annotation types as-is
result.append(ann)
return result if result else None
def _build_model_response(
self,
@@ -113,15 +162,23 @@ class AzureAIAgentsHandler:
model_response: ModelResponse,
thread_id: str,
messages: List[Dict[str, Any]],
annotations: Optional[List[Dict[str, Any]]] = None,
) -> ModelResponse:
"""Build the ModelResponse from agent output."""
from litellm.types.utils import Choices, Message, Usage
message_kwargs: Dict[str, Any] = {
"content": content,
"role": "assistant",
}
if annotations:
message_kwargs["annotations"] = annotations
model_response.choices = [
Choices(
finish_reason="stop",
index=0,
message=Message(content=content, role="assistant"),
message=Message(**message_kwargs),
)
]
model_response.model = model
@@ -250,7 +307,7 @@ class AzureAIAgentsHandler:
)
# Execute the agent flow
thread_id, content = self._execute_agent_flow_sync(
thread_id, content, annotations = self._execute_agent_flow_sync(
make_request=make_request,
api_base=api_base,
api_version=api_version,
@@ -261,7 +318,7 @@ class AzureAIAgentsHandler:
)
return self._build_model_response(
model, content, model_response, thread_id, messages
model, content, model_response, thread_id, messages, annotations
)
def _execute_agent_flow_sync(
@@ -273,8 +330,8 @@ class AzureAIAgentsHandler:
thread_id: Optional[str],
messages: List[Dict[str, Any]],
optional_params: dict,
) -> Tuple[str, str]:
"""Execute the agent flow synchronously. Returns (thread_id, content)."""
) -> Tuple[str, str, Optional[List[Dict[str, Any]]]]:
"""Execute the agent flow synchronously. Returns (thread_id, content, annotations)."""
# Step 1: Create thread if not provided
if not thread_id:
@@ -347,8 +404,8 @@ class AzureAIAgentsHandler:
)
self._check_response(response, [200], "Failed to get messages")
content = self._extract_content_from_messages(response.json())
return thread_id, content
content, annotations = self._extract_content_from_messages(response.json())
return thread_id, content, annotations
# -------------------------------------------------------------------------
# Async Completion
@@ -399,7 +456,7 @@ class AzureAIAgentsHandler:
)
# Execute the agent flow
thread_id, content = await self._execute_agent_flow_async(
thread_id, content, annotations = await self._execute_agent_flow_async(
make_request=make_request,
api_base=api_base,
api_version=api_version,
@@ -410,7 +467,7 @@ class AzureAIAgentsHandler:
)
return self._build_model_response(
model, content, model_response, thread_id, messages
model, content, model_response, thread_id, messages, annotations
)
async def _execute_agent_flow_async(
@@ -422,8 +479,8 @@ class AzureAIAgentsHandler:
thread_id: Optional[str],
messages: List[Dict[str, Any]],
optional_params: dict,
) -> Tuple[str, str]:
"""Execute the agent flow asynchronously. Returns (thread_id, content)."""
) -> Tuple[str, str, Optional[List[Dict[str, Any]]]]:
"""Execute the agent flow asynchronously. Returns (thread_id, content, annotations)."""
# Step 1: Create thread if not provided
if not thread_id:
@@ -496,8 +553,8 @@ class AzureAIAgentsHandler:
)
self._check_response(response, [200], "Failed to get messages")
content = self._extract_content_from_messages(response.json())
return thread_id, content
content, annotations = self._extract_content_from_messages(response.json())
return thread_id, content, annotations
# -------------------------------------------------------------------------
# Streaming Completion (Native SSE)
@@ -585,6 +642,7 @@ class AzureAIAgentsHandler:
response_id = f"chatcmpl-{uuid.uuid4().hex[:8]}"
created = int(time.time())
thread_id = None
collected_annotations: Optional[List[Dict[str, Any]]] = None
current_event = None
@@ -600,6 +658,9 @@ class AzureAIAgentsHandler:
if data_str == "[DONE]":
# Send final chunk with finish_reason
final_delta_kwargs: Dict[str, Any] = {"content": None}
if collected_annotations:
final_delta_kwargs["annotations"] = collected_annotations
final_chunk = ModelResponseStream(
id=response_id,
created=created,
@@ -609,7 +670,7 @@ class AzureAIAgentsHandler:
StreamingChoices(
finish_reason="stop",
index=0,
delta=Delta(content=None),
delta=Delta(**final_delta_kwargs),
)
],
)
@@ -628,6 +689,21 @@ class AzureAIAgentsHandler:
thread_id = data["id"]
verbose_logger.debug(f"Stream created thread: {thread_id}")
# Extract annotations from completed message
if current_event == "thread.message.completed":
for content_item in data.get("content", []):
if content_item.get("type") == "text":
raw_annotations = content_item.get("text", {}).get(
"annotations"
)
transformed = self._transform_annotations(
raw_annotations
)
if transformed:
if collected_annotations is None:
collected_annotations = []
collected_annotations.extend(transformed)
# Process message deltas - this is where the actual content comes
if current_event == "thread.message.delta":
delta_content = data.get("delta", {}).get("content", [])
tests/llm_translation/test_azure_agents.py
+277 -2
View File
@@ -23,12 +23,14 @@ Example environment variables:
See: https://learn.microsoft.com/en-us/azure/ai-foundry/agents/quickstart
"""
import json
import os
import sys
sys.path.insert(0, os.path.abspath("../.."))
import pytest
from unittest.mock import MagicMock
import litellm
@@ -343,13 +345,286 @@ def test_azure_ai_agents_extract_content_from_messages():
]
}
content = handler._extract_content_from_messages(messages_data)
content, annotations = handler._extract_content_from_messages(messages_data)
assert content == "The answer is 100."
assert annotations is None
# Test empty response
empty_data = {"data": []}
content = handler._extract_content_from_messages(empty_data)
content, annotations = handler._extract_content_from_messages(empty_data)
assert content == ""
assert annotations is None
def test_azure_ai_agents_extract_content_with_annotations():
"""
Test that annotations (e.g., Bing Search citations) are extracted from
Azure Agents message responses and transformed to OpenAI-compatible format.
Ref: https://github.com/BerriAI/litellm/issues/19126
"""
from litellm.llms.azure_ai.agents.handler import AzureAIAgentsHandler
handler = AzureAIAgentsHandler()
messages_data = {
"data": [
{
"id": "msg_abc",
"role": "assistant",
"content": [
{
"type": "text",
"text": {
"value": "According to sources [1], the answer is yes.",
"annotations": [
{
"type": "url_citation",
"text": "[1]",
"start_index": 22,
"end_index": 25,
"url_citation": {
"url": "https://example.com/source",
"title": "Example Source"
}
}
]
}
}
]
}
]
}
content, annotations = handler._extract_content_from_messages(messages_data)
assert content == "According to sources [1], the answer is yes."
assert annotations is not None
assert len(annotations) == 1
assert annotations[0]["type"] == "url_citation"
assert annotations[0]["url_citation"]["url"] == "https://example.com/source"
assert annotations[0]["url_citation"]["title"] == "Example Source"
# start/end_index should be moved into url_citation for OpenAI compatibility
assert annotations[0]["url_citation"]["start_index"] == 22
assert annotations[0]["url_citation"]["end_index"] == 25
def test_azure_ai_agents_build_model_response_with_annotations():
"""
Test that _build_model_response includes annotations in the Message object.
"""
from litellm.llms.azure_ai.agents.handler import AzureAIAgentsHandler
from litellm.types.utils import ModelResponse
handler = AzureAIAgentsHandler()
model_response = ModelResponse()
annotations = [
{
"type": "url_citation",
"url_citation": {
"url": "https://example.com",
"title": "Example",
"start_index": 0,
"end_index": 5,
},
}
]
result = handler._build_model_response(
model="azure_ai/agents/asst_123",
content="Hello [1]",
model_response=model_response,
thread_id="thread_abc",
messages=[{"role": "user", "content": "test"}],
annotations=annotations,
)
assert result.choices[0].message.content == "Hello [1]"
assert result.choices[0].message.annotations is not None
assert len(result.choices[0].message.annotations) == 1
assert result.choices[0].message.annotations[0]["type"] == "url_citation"
def test_azure_ai_agents_build_model_response_without_annotations():
"""
Test that _build_model_response works correctly without annotations.
"""
from litellm.llms.azure_ai.agents.handler import AzureAIAgentsHandler
from litellm.types.utils import ModelResponse
handler = AzureAIAgentsHandler()
model_response = ModelResponse()
result = handler._build_model_response(
model="azure_ai/agents/asst_123",
content="Hello",
model_response=model_response,
thread_id="thread_abc",
messages=[{"role": "user", "content": "test"}],
)
assert result.choices[0].message.content == "Hello"
assert getattr(result.choices[0].message, "annotations", None) is None
@pytest.mark.asyncio
async def test_azure_ai_agents_streaming_annotations_from_completed_message():
"""
Test that annotations from thread.message.completed SSE events are collected
and attached to the final chunk's delta.
Ref: https://github.com/BerriAI/litellm/issues/19126
"""
from litellm.llms.azure_ai.agents.handler import AzureAIAgentsHandler
handler = AzureAIAgentsHandler()
# SSE lines simulating a stream with annotations in thread.message.completed
completed_data = {
"content": [
{
"type": "text",
"text": {
"value": "According to [1], the answer is 42.",
"annotations": [
{
"type": "url_citation",
"text": "[1]",
"start_index": 12,
"end_index": 15,
"url_citation": {
"url": "https://example.com/citation",
"title": "Citation Source",
},
}
],
},
}
]
}
sse_lines = [
"event: thread.created",
"",
'data: {"id": "thread_stream_123"}',
"",
"event: thread.message.delta",
"",
'data: {"delta": {"content": [{"type": "text", "text": {"value": "According to [1], the answer is 42."}}]}}',
"",
"event: thread.message.completed",
"",
f"data: {json.dumps(completed_data)}",
"",
"data: [DONE]",
]
async def mock_aiter_lines():
for line in sse_lines:
yield line
mock_response = MagicMock()
mock_response.aiter_lines = MagicMock(return_value=mock_aiter_lines())
chunks = []
async for chunk in handler._process_sse_stream(mock_response, "azure_ai/agents/asst_123"):
chunks.append(chunk)
# Should have content chunks + final [DONE] chunk
assert len(chunks) >= 1
final_chunk = chunks[-1]
assert final_chunk.choices[0].finish_reason == "stop"
assert final_chunk.choices[0].delta.annotations is not None
assert len(final_chunk.choices[0].delta.annotations) == 1
ann = final_chunk.choices[0].delta.annotations[0]
assert ann["type"] == "url_citation"
assert ann["url_citation"]["url"] == "https://example.com/citation"
assert ann["url_citation"]["title"] == "Citation Source"
@pytest.mark.asyncio
async def test_azure_ai_agents_streaming_accumulates_annotations_from_multiple_text_items():
"""
Test that annotations from multiple text content items in thread.message.completed
are accumulated (not overwritten).
Ref: Greptile review on PR #23849
"""
from litellm.llms.azure_ai.agents.handler import AzureAIAgentsHandler
handler = AzureAIAgentsHandler()
# Two text blocks, each with distinct citations
completed_data = {
"content": [
{
"type": "text",
"text": {
"value": "First source [1].",
"annotations": [
{
"type": "url_citation",
"text": "[1]",
"start_index": 12,
"end_index": 15,
"url_citation": {
"url": "https://example.com/first",
"title": "First",
},
}
],
},
},
{
"type": "text",
"text": {
"value": "Second source [2].",
"annotations": [
{
"type": "url_citation",
"text": "[2]",
"start_index": 13,
"end_index": 16,
"url_citation": {
"url": "https://example.com/second",
"title": "Second",
},
}
],
},
},
]
}
sse_lines = [
"event: thread.created",
"",
'data: {"id": "thread_multi"}',
"",
"event: thread.message.completed",
"",
f"data: {json.dumps(completed_data)}",
"",
"data: [DONE]",
]
async def mock_aiter_lines():
for line in sse_lines:
yield line
mock_response = MagicMock()
mock_response.aiter_lines = MagicMock(return_value=mock_aiter_lines())
chunks = []
async for chunk in handler._process_sse_stream(mock_response, "azure_ai/agents/asst_123"):
chunks.append(chunk)
final_chunk = chunks[-1]
assert final_chunk.choices[0].delta.annotations is not None
assert len(final_chunk.choices[0].delta.annotations) == 2
urls = [a["url_citation"]["url"] for a in final_chunk.choices[0].delta.annotations]
assert "https://example.com/first" in urls
assert "https://example.com/second" in urls
@pytest.mark.asyncio