Python: Azure AI Search provider improvements - EmbeddingGenerator, async context manager, KB message handling

[eavanvalkenburg]

Summary

This PR improves the Azure AI Search context provider with several fixes and enhancements:

Key Changes

• EmbeddingGenerator support for semantic mode — The provider now accepts SupportsGetEmbeddings instances in addition to plain async callables for the embedding_function parameter, enabling richer embedding workflows.
• Async context manager improvements — Added a dedicated close() method that properly cleans up all clients (search, index, and retrieval), and refactored __aexit__ to delegate to it.
• Streamlined KB message conversion — Extracted _prepare_messages_for_kb_search and _parse_messages_from_kb_response helpers that handle both text and image content types between framework Message objects and KB SDK types.
• Rich reference-to-annotation mapping — KB retrieval references are now fully captured as Annotation objects with all available metadata: doc_key, reranker_score, source_data, reference_type, sensitivity_label, and the raw reference object.
• Both search paths return Message objects — _semantic_search and _agentic_search now both return list[Message], simplifying the before_run hook.
• Sample updates

Testing

• All existing tests updated to reflect the new return types
• Added comprehensive tests for _prepare_messages_for_kb_search, _parse_messages_from_kb_response, and _parse_references_to_annotations

Fixes #3001
Supersedes #3333

[Copilot]

Pull request overview

This PR enhances the Azure AI Search context provider with several improvements to support richer embedding workflows, proper resource cleanup, and better knowledge base message handling.

Changes:

• Added support for SupportsGetEmbeddings instances in addition to plain async callables for the embedding_function parameter
• Implemented proper async context manager cleanup with a dedicated close() method that resets internal state flags
• Refactored KB message conversion with dedicated helper methods (_prepare_messages_for_kb_search, _parse_messages_from_kb_response, _parse_references_to_annotations)
• Changed return types from list[str] to list[Message] for both search paths, with rich annotation support for KB references

Reviewed changes

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

File	Description
python/packages/azure-ai-search/agent_framework_azure_ai_search/_context_provider.py	Core changes: Added SupportsGetEmbeddings support, close() method, message conversion helpers, fixed k_nearest_neighbors→k parameter, and unified return types to Message objects
python/packages/azure-ai-search/tests/test_aisearch_context_provider.py	Updated all test assertions for new list[Message] return type, added comprehensive test classes for the three new helper methods
python/samples/02-agents/context_providers/azure_ai_search/azure_ai_with_search_context_semantic.py	Added optional embedding client configuration for hybrid search demonstration
python/samples/02-agents/context_providers/azure_ai_search/azure_ai_with_search_context_agentic.py	Updated default retrieval settings and added annotation display in streaming output

[markwallace-microsoft]

Python Test Coverage Report •

File	Stmts	Miss	Cover	Missing
packages/azure-ai-search/agent_framework_azure_ai_search
_context_provider.py	334	10	97%	111–112, 283, 334–335, 458–459, 635–636, 683
TOTAL	21633	3100	85%

Python Unit Test Overview

Tests	Skipped	Failures	Errors	Time
4346	246 💤	0 ❌	0 🔥	1m 19s ⏱️