Fix Azure AI Search citation enrichment follow-ups

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: eavanvalkenburg

python/packages/foundry/tests/foundry/test_foundry_agent.py

@@ -57,7 +57,6 @@ def _get_foundry_azure_ai_search_model() -> str | None:
 
 skip_if_foundry_azure_ai_search_integration_tests_disabled = pytest.mark.skipif(
     os.getenv("FOUNDRY_PROJECT_ENDPOINT", "") in ("", "https://test-project.services.ai.azure.com/")
-    or os.getenv("AZURE_SEARCH_ENDPOINT", "") == ""
     or os.getenv("AZURE_SEARCH_INDEX_NAME", "") == ""
     or _get_foundry_azure_ai_search_model() is None,
     reason="No live Foundry project, Azure Search index, or model provided for Azure AI Search integration tests.",

python/packages/openai/agent_framework_openai/_chat_client.py

@@ -122,6 +122,7 @@
 OPENAI_SHELL_OUTPUT_TYPE_LOCAL_SHELL_CALL = "local_shell_call_output"
 _AZURE_AI_SEARCH_CALL_OUTPUT_TYPE = "azure_ai_search_call_output"
 _AZURE_AI_SEARCH_OUTPUT_EVENT_TYPES = {"response.output_item.added", "response.output_item.done"}
+_AZURE_AI_SEARCH_OUTPUT_EVENT_PREFIX = "response.azure_ai_search_call_output."
 
 # Internal marker emitted by `_prepare_content_for_openai` for an
 # `mcp_server_tool_result` Content. The Responses API expects an `mcp_call`
@@ -1978,16 +1979,27 @@ def _parse_azure_ai_search_output_payload(output: Any) -> Mapping[str, Any] | No
         return None
 
     @staticmethod
-    def _extract_azure_ai_search_get_urls(event: Any) -> list[str]:
-        """Extract per-document Azure AI Search REST URLs from a streamed Responses event."""
-        if getattr(event, "type", None) not in _AZURE_AI_SEARCH_OUTPUT_EVENT_TYPES:
-            return []
+    def _extract_azure_ai_search_output_payload(event: Any) -> Mapping[str, Any] | None:
+        """Return Azure AI Search output payload from either a top-level event or its nested item."""
+        payload = RawOpenAIChatClient._parse_azure_ai_search_output_payload(getattr(event, "output", None))
+        if payload is not None:
+            return payload
 
         item = getattr(event, "item", None)
-        if getattr(item, "type", None) != _AZURE_AI_SEARCH_CALL_OUTPUT_TYPE:
+        if getattr(item, "type", None) == _AZURE_AI_SEARCH_CALL_OUTPUT_TYPE:
+            return RawOpenAIChatClient._parse_azure_ai_search_output_payload(getattr(item, "output", None))
+        return None
+
+    @staticmethod
+    def _extract_azure_ai_search_get_urls(event: Any) -> list[str]:
+        """Extract per-document Azure AI Search REST URLs from a streamed Responses event."""
+        event_type = getattr(event, "type", None)
+        if event_type not in _AZURE_AI_SEARCH_OUTPUT_EVENT_TYPES and not (
+            isinstance(event_type, str) and event_type.startswith(_AZURE_AI_SEARCH_OUTPUT_EVENT_PREFIX)
+        ):
             return []
 
-        payload = RawOpenAIChatClient._parse_azure_ai_search_output_payload(getattr(item, "output", None))
+        payload = RawOpenAIChatClient._extract_azure_ai_search_output_payload(event)
         if payload is None:
             return []
 
@@ -2028,12 +2040,18 @@ def _enrich_streamed_azure_ai_search_citations(cls, updates: Sequence[ChatRespon
                 for annotation in content.annotations:
                     if annotation.get("type") != "citation" or annotation.get("file_id"):
                         continue
-                    additional_properties = annotation.setdefault("additional_properties", {})
-                    if "get_url" in additional_properties:
-                        continue
+
                     doc_index = cls._azure_ai_search_doc_index(annotation)
                     if doc_index is None or doc_index >= len(get_urls):
                         continue
+
+                    additional_properties = annotation.get("additional_properties")
+                    if not isinstance(additional_properties, dict):
+                        additional_properties = {}
+                        annotation["additional_properties"] = additional_properties
+                    if "get_url" in additional_properties:
+                        continue
+
                     additional_properties["get_url"] = get_urls[doc_index]
 
     @staticmethod
@@ -3136,7 +3154,8 @@ def _get_ann_value(key: str) -> Any:
                 elif getattr(done_item, "type", None) == _AZURE_AI_SEARCH_CALL_OUTPUT_TYPE:
                     pass
             case _:
-                logger.debug("Unparsed event of type: %s: %s", event.type, event)
+                if not isinstance(event.type, str) or not event.type.startswith(_AZURE_AI_SEARCH_OUTPUT_EVENT_PREFIX):
+                    logger.debug("Unparsed event of type: %s: %s", event.type, event)
 
         return ChatResponseUpdate(
             contents=contents,

python/packages/openai/tests/openai/test_openai_chat_client.py

@@ -13,6 +13,7 @@
 import pytest
 from agent_framework import (
     Agent,
+    Annotation,
     ChatOptions,
     ChatResponse,
     ChatResponseUpdate,
@@ -3574,9 +3575,17 @@ def _make_mcp_call_done_event(output: str) -> MagicMock:
     return event
 
 
-def _make_azure_ai_search_output_event(output: Any, *, event_type: str = "response.output_item.done") -> MagicMock:
+def _make_azure_ai_search_output_event(
+    output: Any,
+    *,
+    event_type: str = "response.output_item.done",
+    top_level_output: bool = False,
+) -> MagicMock:
     event = MagicMock()
     event.type = event_type
+    if top_level_output:
+        event.output = output
+        return event
     event.item = MagicMock()
     event.item.type = "azure_ai_search_call_output"
     event.item.output = output
@@ -3675,6 +3684,81 @@ def test_streaming_azure_ai_search_output_does_not_overwrite_existing_get_url()
     assert annotation["additional_properties"]["get_url"] == existing_get_url
 
 
+def test_streaming_azure_ai_search_output_normalizes_non_dict_additional_properties() -> None:
+    """Existing non-dict additional_properties should be normalized before enriching get_url."""
+    client = OpenAIChatClient(model="test-model", api_key="test-key")
+    chat_options = ChatOptions()
+    function_call_ids: dict[int, tuple[str, str]] = {}
+    get_url = "https://example.search.windows.net/indexes/my-index/docs/doc-123?api-version=2024-07-01"
+
+    citation_update = client._parse_chunk_from_openai(
+        _make_url_citation_event(title="doc_0"),
+        chat_options,
+        function_call_ids,
+    )
+    citation_update.contents[0].annotations[0]["additional_properties"] = None
+    search_update = client._parse_chunk_from_openai(
+        _make_azure_ai_search_output_event(json.dumps({"get_urls": [get_url]})),
+        chat_options,
+        function_call_ids,
+    )
+
+    response = client._finalize_response_updates([citation_update, search_update])
+
+    annotation = response.messages[0].contents[0].annotations[0]
+    assert annotation["additional_properties"] == {"get_url": get_url}
+
+
+def test_streaming_azure_ai_search_output_does_not_create_additional_properties_for_unusable_citation() -> None:
+    """Unenrichable Azure AI Search citations should keep their original annotation shape."""
+    update = ChatResponseUpdate(
+        contents=[
+            Content.from_text(
+                text="hello",
+                annotations=[Annotation(type="citation", title="source_0", url="https://example.invalid")],
+            )
+        ],
+        raw_representation=_make_azure_ai_search_output_event(
+            json.dumps({"get_urls": ["https://example.search.windows.net/indexes/my-index/docs/doc-0"]})
+        ),
+    )
+
+    RawOpenAIChatClient._enrich_streamed_azure_ai_search_citations([update])
+
+    annotation = update.contents[0].annotations[0]
+    assert annotation.get("additional_properties") is None
+
+
+def test_extract_azure_ai_search_get_urls_accepts_dedicated_output_event() -> None:
+    """Dedicated response.azure_ai_search_call_output.* events should yield get_urls too."""
+    get_url = "https://example.search.windows.net/indexes/my-index/docs/doc-123?api-version=2024-07-01"
+    event = _make_azure_ai_search_output_event(
+        json.dumps({"get_urls": [get_url]}),
+        event_type="response.azure_ai_search_call_output.done",
+        top_level_output=True,
+    )
+
+    assert RawOpenAIChatClient._extract_azure_ai_search_get_urls(event) == [get_url]
+
+
+def test_parse_chunk_from_openai_ignores_dedicated_azure_ai_search_events() -> None:
+    """Dedicated Azure AI Search events should be treated as intentional no-op updates."""
+    client = OpenAIChatClient(model="test-model", api_key="test-key")
+    chat_options = ChatOptions()
+    function_call_ids: dict[int, tuple[str, str]] = {}
+    event = _make_azure_ai_search_output_event(
+        json.dumps({"get_urls": ["https://example.search.windows.net/indexes/my-index/docs/doc-0"]}),
+        event_type="response.azure_ai_search_call_output.done",
+        top_level_output=True,
+    )
+
+    with patch("agent_framework_openai._chat_client.logger.debug") as mock_debug:
+        update = client._parse_chunk_from_openai(event, chat_options, function_call_ids)
+
+    assert update.contents == []
+    mock_debug.assert_not_called()
+
+
 @pytest.mark.parametrize(
     ("title", "output"),
     [