Python: Enhance Azure AI Search Citations with Document URLs in Foundry V2

[giles17]

Summary

Enriches Azure AI Search url_citation annotations with per-document REST API URLs (get_url) for the Foundry V2 (Responses API) path. Previously, url_citation annotations only contained the search service base URL, making it hard for users to identify which specific document was referenced.

This is the V2 counterpart to PR #2066 which solved the same problem for Foundry V1 (Assistants API).

Resolves #2496

Problem

When using Azure AI Search as a tool with the Responses API, the url_citation annotation in the assistant's response only contains the search service base URL (e.g., https://search.example.com/), not the document-specific URL. The actual per-document URLs exist in the azure_ai_search_call_output response items under output.get_urls[].

Solution

Override _inner_get_response in RawAzureAIClient to post-process both streaming and non-streaming responses:

Non-streaming: Wraps the awaitable to extract get_urls from raw_representation.output after the base class parses the response, then enriches citation annotations.

Streaming: Registers a transform hook on ResponseStream that:

1. Captures get_urls from azure_ai_search_call_output items (via response.output_item.added and response.output_item.done events — the data is only fully populated in the done event)
2. Handles url_citation annotations (which the base class doesn't handle in streaming) by creating proper Content objects with enriched Annotation entries
3. Uses closure-local state instead of instance state, so concurrent streams don't interfere

How users access citations

# Non-streaming
result = await agent.run(query)
for msg in result.messages:
    for content in msg.contents:
        for ann in (content.annotations or []):
            if ann.get("type") == "citation":
                doc_url = ann.get("additional_properties", {}).get("get_url")
                print(f"Document URL: {doc_url}")

# Streaming
async for chunk in agent.run(query, stream=True):
    for content in (chunk.contents or []):
        for ann in (content.annotations or []):
            if ann.get("type") == "citation":
                doc_url = ann.get("additional_properties", {}).get("get_url")
                print(f"Document URL: {doc_url}")

Changes

• packages/azure-ai/agent_framework_azure_ai/_client.py — Added _inner_get_response override, helper methods (_extract_azure_search_urls, _get_search_doc_url, _enrich_annotations_with_search_urls, _build_url_citation_content)
• packages/azure-ai/tests/test_azure_ai_client.py — Added 14 unit tests covering helpers, non-streaming enrichment, streaming hook registration, streaming URL capture and annotation enrichment
• samples/02-agents/providers/azure_ai/azure_ai_with_azure_ai_search.py — Updated sample demonstrating citation extraction for both streaming and non-streaming

[markwallace-microsoft]

Python Test Coverage Report •

File	Stmts	Miss	Cover	Missing
packages/azure-ai/agent_framework_azure_ai
_client.py	380	42	88%	386, 388, 431, 439–451, 464, 524, 539–544, 587, 622, 624, 689, 692, 694, 795, 830, 870, 1078, 1081, 1084–1085, 1087–1090, 1133
TOTAL	21336	3308	84%

Python Unit Test Overview

Tests	Skipped	Failures	Errors	Time
4199	239 💤	0 ❌	0 🔥	1m 19s ⏱️

[Copilot]

Pull request overview

Enhances the Azure AI (Foundry V2 / Responses API) integration to surface Azure AI Search per-document REST URLs (get_url) on citation annotations, for both non-streaming and streaming responses.

Changes:

• Added extraction/mapping/enrichment utilities in the Azure AI Responses client to attach get_url to citation annotations.
• Overrode non-streaming and streaming parsing to capture Azure AI Search get_urls and enrich citations.
• Updated the Azure AI Search sample and added a dedicated set of unit tests for this behavior.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 3 comments.

File	Description
python/samples/02-agents/providers/azure_ai/azure_ai_with_azure_ai_search.py	Demonstrates reading enriched citation additional_properties.get_url for non-streaming and streaming runs.
python/packages/azure-ai/agent_framework_azure_ai/_client.py	Implements extraction of get_urls from azure_ai_search_call_output items and citation enrichment in response parsing (including streaming).
python/packages/azure-ai/tests/test_azure_ai_client.py	Adds unit tests for URL extraction, mapping, enrichment, and streaming behavior/state handling.