Python: Azure AI Search Citation Extraction - V2

python/samples/getting_started/agents/azure_ai/azure_ai_with_azure_ai_search.py

@@ -1,7 +1,9 @@
 # Copyright (c) Microsoft. All rights reserved.
 import asyncio
 import os
+from typing import Any
 
+from agent_framework._types import AgentRunResponse, CitationAnnotation, TextContent
 from agent_framework.azure import AzureAIClient
 from azure.identity.aio import AzureCliCredential
 
@@ -18,6 +20,38 @@
 """
 
 
+def extract_citations_from_response(response: AgentRunResponse) -> list[dict[str, Any]]:
+    """Extract citation information from an AgentRunResponse."""
+    citations: list[dict[str, Any]] = []
+
+    if hasattr(response, "messages") and response.messages:
+        for message in response.messages:
+            if hasattr(message, "contents") and message.contents:
+                for content in message.contents:
+                    if isinstance(content, TextContent) and hasattr(content, "annotations") and content.annotations:
+                        for annotation in content.annotations:
+                            if isinstance(annotation, CitationAnnotation):
+                                citation_info = {
+                                    "url": annotation.url,
+                                    "title": getattr(annotation, "title", None),
+                                    "file_id": getattr(annotation, "file_id", None),
+                                }
+
+                                # Extract text position information
+                                if hasattr(annotation, "annotated_regions") and annotation.annotated_regions:
+                                    citation_info["positions"] = [
+                                        {
+                                            "start_index": region.start_index,
+                                            "end_index": region.end_index,
+                                        }
+                                        for region in annotation.annotated_regions
+                                    ]
+
+                                citations.append(citation_info)
+
+    return citations
+
+
 async def main() -> None:
     async with (
         AzureCliCredential() as credential,
@@ -33,17 +67,33 @@ async def main() -> None:
                             "project_connection_id": os.environ["AI_SEARCH_PROJECT_CONNECTION_ID"],
                             "index_name": os.environ["AI_SEARCH_INDEX_NAME"],
                             # For query_type=vector, ensure your index has a field with vectorized data.
-                            "query_type": "simple",
+                            "query_type": "vector",
                         }
                     ]
                 },
             },
         ) as agent,
     ):
-        query = "Tell me about insurance options"
+        query = (
+            "Use Azure AI search knowledge tool to find detailed information about a winter hotel."
+            " Use the search tool and index."
+        )
         print(f"User: {query}")
-        result = await agent.run(query)
-        print(f"Result: {result}\n")
+
+        # Get the response
+        response = await agent.run(query)
+        print(f"Agent: {response}\n")
+
+        # Extract citation data
+        citations = extract_citations_from_response(response)
+
+        print("=== CITATION INFORMATION ===")
+        if citations:
+            for i, citation in enumerate(citations, 1):
+                print(f"Citation {i}:")
+                print(f"  URL: {citation['url']}")
+        else:
+            print("No citations found in the response.")
 
 
 if __name__ == "__main__":