diff --git a/litellm/llms/azure_ai/agents/handler.py b/litellm/llms/azure_ai/agents/handler.py index 9eeec7f4e3..3fbda13d3c 100644 --- a/litellm/llms/azure_ai/agents/handler.py +++ b/litellm/llms/azure_ai/agents/handler.py @@ -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", []) diff --git a/tests/llm_translation/test_azure_agents.py b/tests/llm_translation/test_azure_agents.py index 66a46d5338..3e6b1e00a7 100644 --- a/tests/llm_translation/test_azure_agents.py +++ b/tests/llm_translation/test_azure_agents.py @@ -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