agent/agent-framework
1
0
Fork 0
mirror of https://github.com/microsoft/agent-framework.git synced 2026-06-16 21:04:09 +08:00
Code Issues Packages Projects Releases Wiki Activity

Python: Added URL Citation Support to Azure AI Agent (#1397)

Browse Source 
* added url citation to agentrunresponse

* small fix

* fix

* Update python/packages/azure-ai/agent_framework_azure_ai/_chat_client.py

Co-authored-by: Tao Chen <taochen@microsoft.com>

---------

Co-authored-by: Tao Chen <taochen@microsoft.com>
This commit is contained in:
Giles Odigwe
2025-10-13 12:04:08 -07:00
committed by GitHub
Unverified
parent 166aa8fd54
commit baf59ca1ed
2 changed files with 187 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/packages/azure-ai/agent_framework_azure_ai/_chat_client.py
+57 -8
View File
@@ -14,6 +14,7 @@ from agent_framework import (
ChatOptions,
ChatResponse,
ChatResponseUpdate,
CitationAnnotation,
Contents,
DataContent,
FunctionApprovalRequestContent,
@@ -28,6 +29,7 @@ from agent_framework import (
HostedWebSearchTool,
Role,
TextContent,
TextSpanRegion,
ToolMode,
ToolProtocol,
UriContent,
@@ -60,6 +62,8 @@ from azure.ai.agents.models import (
ListSortOrder,
McpTool,
MessageDeltaChunk,
MessageDeltaTextContent,
MessageDeltaTextUrlCitationAnnotation,
MessageImageUrlParam,
MessageInputContentBlock,
MessageInputImageUrlBlock,
@@ -480,6 +484,37 @@ class AzureAIAgentClient(BaseChatClient):
# and remove until here.
return thread_id
def _extract_url_citations(self, message_delta_chunk: MessageDeltaChunk) -> list[CitationAnnotation]:
"""Extract URL citations from MessageDeltaChunk."""
url_citations: list[CitationAnnotation] = []
# Process each content item in the delta to find citations
for content in message_delta_chunk.delta.content:
if isinstance(content, MessageDeltaTextContent) and content.text and content.text.annotations:
for annotation in content.text.annotations:
if isinstance(annotation, MessageDeltaTextUrlCitationAnnotation):
# Create annotated regions only if both start and end indices are available
annotated_regions = []
if annotation.start_index and annotation.end_index:
annotated_regions = [
TextSpanRegion(
start_index=annotation.start_index,
end_index=annotation.end_index,
)
]
# Create CitationAnnotation from AzureAI annotation
citation = CitationAnnotation(
title=getattr(annotation.url_citation, "title", None),
url=annotation.url_citation.url,
snippet=None,
annotated_regions=annotated_regions,
raw_representation=annotation,
)
url_citations.append(citation)
return url_citations
async def _process_stream(
self, stream: AsyncAgentRunStream[AsyncAgentEventHandler[Any]] | AsyncAgentEventHandler[Any], thread_id: str
) -> AsyncIterable[ChatResponseUpdate]:
@@ -492,9 +527,21 @@ class AzureAIAgentClient(BaseChatClient):
case MessageDeltaChunk():
# only one event_type: AgentStreamEvent.THREAD_MESSAGE_DELTA
role = Role.USER if event_data.delta.role == MessageRole.USER else Role.ASSISTANT
# Extract URL citations from the delta chunk
url_citations = self._extract_url_citations(event_data)
# Create contents with citations if any exist
citation_content: list[Contents] = []
if event_data.text or url_citations:
text_content_obj = TextContent(text=event_data.text or "")
if url_citations:
text_content_obj.annotations = url_citations
citation_content.append(text_content_obj)
yield ChatResponseUpdate(
role=role,
text=event_data.text,
contents=citation_content if citation_content else None,
conversation_id=thread_id,
message_id=response_id,
raw_representation=event_data,
@@ -518,11 +565,13 @@ class AzureAIAgentClient(BaseChatClient):
"submit_tool_outputs",
"submit_tool_approval",
]:
contents = self._create_function_call_contents(event_data, response_id)
if contents:
function_call_contents = self._create_function_call_contents(
event_data, response_id
)
if function_call_contents:
yield ChatResponseUpdate(
role=Role.ASSISTANT,
contents=contents,
contents=function_call_contents,
conversation_id=thread_id,
message_id=response_id,
raw_representation=event_data,
@@ -589,22 +638,22 @@ class AzureAIAgentClient(BaseChatClient):
tool_call.code_interpreter,
RunStepDeltaCodeInterpreterDetailItemObject,
):
contents = []
code_contents: list[Contents] = []
if tool_call.code_interpreter.input is not None:
logger.debug(f"Code Interpreter Input: {tool_call.code_interpreter.input}")
if tool_call.code_interpreter.outputs is not None:
for output in tool_call.code_interpreter.outputs:
if isinstance(output, RunStepDeltaCodeInterpreterLogOutput) and output.logs:
contents.append(TextContent(text=output.logs))
code_contents.append(TextContent(text=output.logs))
if (
isinstance(output, RunStepDeltaCodeInterpreterImageOutput)
and output.image is not None
and output.image.file_id is not None
):
contents.append(HostedFileContent(file_id=output.image.file_id))
code_contents.append(HostedFileContent(file_id=output.image.file_id))
yield ChatResponseUpdate(
role=Role.ASSISTANT,
contents=contents,
contents=code_contents,
conversation_id=thread_id,
message_id=response_id,
raw_representation=tool_call.code_interpreter,
python/packages/azure-ai/tests/test_azure_ai_agent_client.py
+130
View File
@@ -18,6 +18,7 @@ from agent_framework import (
ChatOptions,
ChatResponse,
ChatResponseUpdate,
CitationAnnotation,
FunctionApprovalRequestContent,
FunctionApprovalResponseContent,
FunctionCallContent,
@@ -36,6 +37,9 @@ from agent_framework.exceptions import ServiceInitializationError
from azure.ai.agents.models import (
CodeInterpreterToolDefinition,
FileInfo,
MessageDeltaChunk,
MessageDeltaTextContent,
MessageDeltaTextUrlCitationAnnotation,
RequiredFunctionToolCall,
RequiredMcpToolCall,
RunStatus,
@@ -1439,6 +1443,132 @@ async def test_azure_ai_chat_client_create_agent_stream_submit_tool_outputs(
assert final_thread_id == "test-thread"
def test_azure_ai_chat_client_extract_url_citations_with_citations(mock_ai_project_client: MagicMock) -> None:
"""Test _extract_url_citations with MessageDeltaChunk containing URL citations."""
chat_client = create_test_azure_ai_chat_client(mock_ai_project_client, agent_id="test-agent")
# Create mock URL citation annotation
mock_url_citation = MagicMock()
mock_url_citation.url = "https://example.com/test"
mock_url_citation.title = "Test Title"
mock_annotation = MagicMock(spec=MessageDeltaTextUrlCitationAnnotation)
mock_annotation.url_citation = mock_url_citation
mock_annotation.start_index = 10
mock_annotation.end_index = 20
# Create mock text content with annotations
mock_text = MagicMock()
mock_text.annotations = [mock_annotation]
mock_text_content = MagicMock(spec=MessageDeltaTextContent)
mock_text_content.text = mock_text
# Create mock delta
mock_delta = MagicMock()
mock_delta.content = [mock_text_content]
# Create mock MessageDeltaChunk
mock_chunk = MagicMock(spec=MessageDeltaChunk)
mock_chunk.delta = mock_delta
# Call the method
citations = chat_client._extract_url_citations(mock_chunk) # type: ignore
# Verify results
assert len(citations) == 1
citation = citations[0]
assert isinstance(citation, CitationAnnotation)
assert citation.url == "https://example.com/test"
assert citation.title == "Test Title"
assert citation.snippet is None
assert citation.annotated_regions is not None
assert len(citation.annotated_regions) == 1
assert citation.annotated_regions[0].start_index == 10
assert citation.annotated_regions[0].end_index == 20
def test_azure_ai_chat_client_extract_url_citations_no_citations(mock_ai_project_client: MagicMock) -> None:
"""Test _extract_url_citations with MessageDeltaChunk containing no citations."""
chat_client = create_test_azure_ai_chat_client(mock_ai_project_client, agent_id="test-agent")
# Create mock text content without annotations
mock_text_content = MagicMock(spec=MessageDeltaTextContent)
mock_text_content.text = None # No text, so no annotations
# Create mock delta
mock_delta = MagicMock()
mock_delta.content = [mock_text_content]
# Create mock MessageDeltaChunk
mock_chunk = MagicMock(spec=MessageDeltaChunk)
mock_chunk.delta = mock_delta
# Call the method
citations = chat_client._extract_url_citations(mock_chunk) # type: ignore
# Verify no citations returned
assert len(citations) == 0
def test_azure_ai_chat_client_extract_url_citations_empty_delta(mock_ai_project_client: MagicMock) -> None:
"""Test _extract_url_citations with empty delta content."""
chat_client = create_test_azure_ai_chat_client(mock_ai_project_client, agent_id="test-agent")
# Create mock delta with empty content
mock_delta = MagicMock()
mock_delta.content = []
# Create mock MessageDeltaChunk
mock_chunk = MagicMock(spec=MessageDeltaChunk)
mock_chunk.delta = mock_delta
# Call the method
citations = chat_client._extract_url_citations(mock_chunk) # type: ignore
# Verify no citations returned
assert len(citations) == 0
def test_azure_ai_chat_client_extract_url_citations_without_indices(mock_ai_project_client: MagicMock) -> None:
"""Test _extract_url_citations with URL citations that don't have start/end indices."""
chat_client = create_test_azure_ai_chat_client(mock_ai_project_client, agent_id="test-agent")
# Create mock URL citation annotation without indices
mock_url_citation = MagicMock()
mock_url_citation.url = "https://example.com/no-indices"
mock_annotation = MagicMock(spec=MessageDeltaTextUrlCitationAnnotation)
mock_annotation.url_citation = mock_url_citation
mock_annotation.start_index = None
mock_annotation.end_index = None
# Create mock text content with annotations
mock_text = MagicMock()
mock_text.annotations = [mock_annotation]
mock_text_content = MagicMock(spec=MessageDeltaTextContent)
mock_text_content.text = mock_text
# Create mock delta
mock_delta = MagicMock()
mock_delta.content = [mock_text_content]
# Create mock MessageDeltaChunk
mock_chunk = MagicMock(spec=MessageDeltaChunk)
mock_chunk.delta = mock_delta
# Call the method
citations = chat_client._extract_url_citations(mock_chunk) # type: ignore
# Verify results
assert len(citations) == 1
citation = citations[0]
assert citation.url == "https://example.com/no-indices"
assert citation.annotated_regions is not None
assert len(citation.annotated_regions) == 0 # No regions when indices are None
async def test_azure_ai_chat_client_setup_azure_ai_observability_resource_not_found(
mock_ai_project_client: MagicMock,
) -> None: