diff --git a/python/.env.example b/python/.env.example index 82458a3fda..f864f18f72 100644 --- a/python/.env.example +++ b/python/.env.example @@ -3,6 +3,14 @@ AZURE_AI_PROJECT_ENDPOINT="" AZURE_AI_MODEL_DEPLOYMENT_NAME="" # Bing connection for web search (optional, used by samples with web search) BING_CONNECTION_ID="" +# Azure AI Search (optional, used by AzureAISearchContextProvider samples) +AZURE_SEARCH_ENDPOINT="" +AZURE_SEARCH_API_KEY="" +AZURE_SEARCH_INDEX_NAME="" +AZURE_SEARCH_SEMANTIC_CONFIG="" +AZURE_SEARCH_KNOWLEDGE_BASE_NAME="" +# Note: For agentic mode Knowledge Bases, also set AZURE_OPENAI_ENDPOINT below +# (different from AZURE_AI_PROJECT_ENDPOINT - Knowledge Base needs OpenAI endpoint for model calls) # OpenAI OPENAI_API_KEY="" OPENAI_CHAT_MODEL_ID="" diff --git a/python/packages/aisearch/LICENSE b/python/packages/aisearch/LICENSE new file mode 100644 index 0000000000..9e841e7a26 --- /dev/null +++ b/python/packages/aisearch/LICENSE @@ -0,0 +1,21 @@ + MIT License + + Copyright (c) Microsoft Corporation. + + Permission is hereby granted, free of charge, to any person obtaining a copy + of this software and associated documentation files (the "Software"), to deal + in the Software without restriction, including without limitation the rights + to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + copies of the Software, and to permit persons to whom the Software is + furnished to do so, subject to the following conditions: + + The above copyright notice and this permission notice shall be included in all + copies or substantial portions of the Software. + + THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE + SOFTWARE diff --git a/python/packages/aisearch/README.md b/python/packages/aisearch/README.md new file mode 100644 index 0000000000..6631a2c863 --- /dev/null +++ b/python/packages/aisearch/README.md @@ -0,0 +1,23 @@ +# Get Started with Microsoft Agent Framework Azure AI Search + +Please install this package via pip: + +```bash +pip install agent-framework-aisearch --pre +``` + +## Azure AI Search Integration + +The Azure AI Search integration provides context providers for RAG (Retrieval Augmented Generation) capabilities with two modes: + +- **Semantic Mode**: Fast hybrid search (vector + keyword) with semantic ranking +- **Agentic Mode**: Multi-hop reasoning using Knowledge Bases for complex queries + +### Basic Usage Example + +See the [Azure AI Search context provider examples](https://github.com/microsoft/agent-framework/tree/main/python/samples/getting_started/agents/azure_ai/) which demonstrate: + +- Semantic search with hybrid (vector + keyword) queries +- Agentic mode with Knowledge Bases for complex multi-hop reasoning +- Environment variable configuration with Settings class +- API key and managed identity authentication diff --git a/python/packages/aisearch/agent_framework_aisearch/__init__.py b/python/packages/aisearch/agent_framework_aisearch/__init__.py new file mode 100644 index 0000000000..fedfb05bcd --- /dev/null +++ b/python/packages/aisearch/agent_framework_aisearch/__init__.py @@ -0,0 +1,16 @@ +# Copyright (c) Microsoft. All rights reserved. + +import importlib.metadata + +from ._search_provider import AzureAISearchContextProvider, AzureAISearchSettings + +try: + __version__ = importlib.metadata.version(__name__) +except importlib.metadata.PackageNotFoundError: + __version__ = "0.0.0" # Fallback for development mode + +__all__ = [ + "AzureAISearchContextProvider", + "AzureAISearchSettings", + "__version__", +] diff --git a/python/packages/aisearch/agent_framework_aisearch/_search_provider.py b/python/packages/aisearch/agent_framework_aisearch/_search_provider.py new file mode 100644 index 0000000000..23c8c3f309 --- /dev/null +++ b/python/packages/aisearch/agent_framework_aisearch/_search_provider.py @@ -0,0 +1,914 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Azure AI Search Context Provider for Agent Framework. + +This module provides context providers for Azure AI Search integration with two modes: +- Agentic: Recommended for most scenarios. Uses Knowledge Bases for query planning and + multi-hop reasoning. Slightly slower with more token consumption, but more accurate. +- Semantic: Fast hybrid search (vector + keyword) with semantic ranker. Best for simple + queries where speed is critical. + +See: https://techcommunity.microsoft.com/blog/azure-ai-foundry-blog/foundry-iq-boost-response-relevance-by-36-with-agentic-retrieval/4470720 +""" + +import sys +from collections.abc import Awaitable, Callable, MutableSequence +from typing import TYPE_CHECKING, Any, ClassVar, Literal + +from agent_framework import ChatMessage, Context, ContextProvider, Role +from agent_framework._logging import get_logger +from agent_framework._pydantic import AFBaseSettings +from agent_framework.exceptions import ServiceInitializationError +from azure.core.credentials import AzureKeyCredential +from azure.core.credentials_async import AsyncTokenCredential +from azure.core.exceptions import ResourceNotFoundError +from azure.search.documents.aio import SearchClient +from azure.search.documents.indexes.aio import SearchIndexClient +from azure.search.documents.indexes.models import ( + AzureOpenAIVectorizerParameters, + KnowledgeBase, + KnowledgeBaseAzureOpenAIModel, + KnowledgeRetrievalLowReasoningEffort, + KnowledgeRetrievalMediumReasoningEffort, + KnowledgeRetrievalMinimalReasoningEffort, + KnowledgeRetrievalOutputMode, + KnowledgeRetrievalReasoningEffort, + KnowledgeSourceReference, + SearchIndexKnowledgeSource, + SearchIndexKnowledgeSourceParameters, +) +from azure.search.documents.models import ( + QueryCaptionType, + QueryType, + VectorizableTextQuery, + VectorizedQuery, +) +from pydantic import SecretStr, ValidationError + +# Type checking imports for optional agentic mode dependencies +if TYPE_CHECKING: + from azure.search.documents.knowledgebases.aio import KnowledgeBaseRetrievalClient + from azure.search.documents.knowledgebases.models import ( + KnowledgeBaseMessage, + KnowledgeBaseMessageTextContent, + KnowledgeBaseRetrievalRequest, + KnowledgeRetrievalIntent, + KnowledgeRetrievalSemanticIntent, + ) + from azure.search.documents.knowledgebases.models import ( + KnowledgeRetrievalLowReasoningEffort as KBRetrievalLowReasoningEffort, + ) + from azure.search.documents.knowledgebases.models import ( + KnowledgeRetrievalMediumReasoningEffort as KBRetrievalMediumReasoningEffort, + ) + from azure.search.documents.knowledgebases.models import ( + KnowledgeRetrievalMinimalReasoningEffort as KBRetrievalMinimalReasoningEffort, + ) + from azure.search.documents.knowledgebases.models import ( + KnowledgeRetrievalOutputMode as KBRetrievalOutputMode, + ) + from azure.search.documents.knowledgebases.models import ( + KnowledgeRetrievalReasoningEffort as KBRetrievalReasoningEffort, + ) + +# Runtime imports for agentic mode (optional dependency) +try: + from azure.search.documents.knowledgebases.aio import KnowledgeBaseRetrievalClient + from azure.search.documents.knowledgebases.models import ( + KnowledgeBaseMessage, + KnowledgeBaseMessageTextContent, + KnowledgeBaseRetrievalRequest, + KnowledgeRetrievalIntent, + KnowledgeRetrievalSemanticIntent, + ) + from azure.search.documents.knowledgebases.models import ( + KnowledgeRetrievalLowReasoningEffort as KBRetrievalLowReasoningEffort, + ) + from azure.search.documents.knowledgebases.models import ( + KnowledgeRetrievalMediumReasoningEffort as KBRetrievalMediumReasoningEffort, + ) + from azure.search.documents.knowledgebases.models import ( + KnowledgeRetrievalMinimalReasoningEffort as KBRetrievalMinimalReasoningEffort, + ) + from azure.search.documents.knowledgebases.models import ( + KnowledgeRetrievalOutputMode as KBRetrievalOutputMode, + ) + from azure.search.documents.knowledgebases.models import ( + KnowledgeRetrievalReasoningEffort as KBRetrievalReasoningEffort, + ) + + _agentic_retrieval_available = True +except ImportError: + _agentic_retrieval_available = False + +if sys.version_info >= (3, 11): + from typing import Self # pragma: no cover +else: + from typing_extensions import Self # pragma: no cover + +if sys.version_info >= (3, 12): + from typing import override # type: ignore # pragma: no cover +else: + from typing_extensions import override # type: ignore[import] # pragma: no cover + +# Module-level constants +logger = get_logger("agent_framework.azure") +_DEFAULT_AGENTIC_MESSAGE_HISTORY_COUNT = 10 + + +class AzureAISearchSettings(AFBaseSettings): + """Settings for Azure AI Search Context Provider with auto-loading from environment. + + The settings are first loaded from environment variables with the prefix 'AZURE_SEARCH_'. + If the environment variables are not found, the settings can be loaded from a .env file. + + Keyword Args: + endpoint: Azure AI Search endpoint URL. + Can be set via environment variable AZURE_SEARCH_ENDPOINT. + index_name: Name of the search index. + Can be set via environment variable AZURE_SEARCH_INDEX_NAME. + api_key: API key for authentication (optional, use managed identity if not provided). + Can be set via environment variable AZURE_SEARCH_API_KEY. + env_file_path: If provided, the .env settings are read from this file path location. + env_file_encoding: The encoding of the .env file, defaults to 'utf-8'. + + Examples: + .. code-block:: python + + from agent_framework_aisearch import AzureAISearchSettings + + # Using environment variables + # Set AZURE_SEARCH_ENDPOINT=https://mysearch.search.windows.net + # Set AZURE_SEARCH_INDEX_NAME=my-index + settings = AzureAISearchSettings() + + # Or passing parameters directly + settings = AzureAISearchSettings( + endpoint="https://mysearch.search.windows.net", + index_name="my-index", + ) + + # Or loading from a .env file + settings = AzureAISearchSettings(env_file_path="path/to/.env") + """ + + env_prefix: ClassVar[str] = "AZURE_SEARCH_" + + endpoint: str | None = None + index_name: str | None = None + api_key: SecretStr | None = None + + +class AzureAISearchContextProvider(ContextProvider): + """Azure AI Search Context Provider with hybrid search and semantic ranking. + + This provider retrieves relevant documents from Azure AI Search to provide context + to the AI agent. It supports two modes: + + - **agentic**: Recommended for most scenarios. Uses Knowledge Bases for query planning + and multi-hop reasoning. Slightly slower with more token consumption, but provides + more accurate results (up to 36% improvement in response relevance). + - **semantic** (default): Fast hybrid search combining vector and keyword search + with semantic reranking. Best for simple queries where speed is critical. + + Examples: + Using environment variables (recommended): + + .. code-block:: python + + from agent_framework_aisearch import AzureAISearchContextProvider + from azure.identity.aio import DefaultAzureCredential + + # Set AZURE_SEARCH_ENDPOINT and AZURE_SEARCH_INDEX_NAME in environment + search_provider = AzureAISearchContextProvider(credential=DefaultAzureCredential()) + + Semantic hybrid search with API key: + + .. code-block:: python + + # Direct API key string + search_provider = AzureAISearchContextProvider( + endpoint="https://mysearch.search.windows.net", + index_name="my-index", + api_key="my-api-key", + mode="semantic", + ) + + Loading from .env file: + + .. code-block:: python + + # Load settings from a .env file + search_provider = AzureAISearchContextProvider( + credential=DefaultAzureCredential(), env_file_path="path/to/.env" + ) + + Agentic retrieval for complex queries: + + .. code-block:: python + + # Use agentic mode for multi-hop reasoning + # Note: azure_openai_resource_url is the OpenAI endpoint for Knowledge Base model calls, + # which is different from azure_ai_project_endpoint (the AI Foundry project endpoint) + search_provider = AzureAISearchContextProvider( + endpoint="https://mysearch.search.windows.net", + index_name="my-index", + credential=DefaultAzureCredential(), + mode="agentic", + azure_openai_resource_url="https://myresource.openai.azure.com", + model_deployment_name="gpt-4o", + knowledge_base_name="my-knowledge-base", + ) + """ + + _DEFAULT_SEARCH_CONTEXT_PROMPT = "Use the following context to answer the question:" + + def __init__( + self, + endpoint: str | None = None, + index_name: str | None = None, + api_key: str | AzureKeyCredential | None = None, + credential: AsyncTokenCredential | None = None, + *, + mode: Literal["semantic", "agentic"] = "semantic", + top_k: int = 5, + semantic_configuration_name: str | None = None, + vector_field_name: str | None = None, + embedding_function: Callable[[str], Awaitable[list[float]]] | None = None, + context_prompt: str | None = None, + # Agentic mode parameters (Knowledge Base) + azure_ai_project_endpoint: str | None = None, + azure_openai_resource_url: str | None = None, + model_deployment_name: str | None = None, + model_name: str | None = None, + knowledge_base_name: str | None = None, + retrieval_instructions: str | None = None, + azure_openai_api_key: str | None = None, + knowledge_base_output_mode: Literal["extractive_data", "answer_synthesis"] = "extractive_data", + retrieval_reasoning_effort: Literal["minimal", "medium", "low"] = "minimal", + agentic_message_history_count: int = _DEFAULT_AGENTIC_MESSAGE_HISTORY_COUNT, + env_file_path: str | None = None, + env_file_encoding: str | None = None, + ) -> None: + """Initialize Azure AI Search Context Provider. + + Args: + endpoint: Azure AI Search endpoint URL. + Can also be set via environment variable AZURE_SEARCH_ENDPOINT. + index_name: Name of the search index to query. + Can also be set via environment variable AZURE_SEARCH_INDEX_NAME. + api_key: API key for authentication (string or AzureKeyCredential). + Can also be set via environment variable AZURE_SEARCH_API_KEY. + credential: AsyncTokenCredential for managed identity authentication. + Use this for Entra ID authentication instead of api_key. + mode: Search mode - "semantic" for hybrid search with semantic ranking (fast) + or "agentic" for multi-hop reasoning (slower). Default: "semantic". + top_k: Maximum number of documents to retrieve. Only applies to semantic mode. + In agentic mode, the server-side Knowledge Base determines retrieval based on + query complexity and reasoning effort. Default: 5. + semantic_configuration_name: Name of semantic configuration in the index. + Required for semantic ranking. If None, uses index default. + vector_field_name: Name of the vector field in the index for hybrid search. + Required if using vector search. Default: None (keyword search only). + embedding_function: Async function to generate embeddings for vector search. + Signature: async def embed(text: str) -> list[float] + Required if vector_field_name is specified and no server-side vectorization. + context_prompt: Custom prompt to prepend to retrieved context. + Default: "Use the following context to answer the question:" + azure_ai_project_endpoint: Azure AI Foundry project endpoint URL. + This is NOT the same as azure_openai_resource_url - the project endpoint is used + for Azure AI Foundry services, while the OpenAI endpoint is used by the Knowledge + Base to call the model for query planning. Required for agentic mode. + Example: "https://myproject.services.ai.azure.com/api/projects/myproject" + azure_openai_resource_url: Azure OpenAI resource URL for Knowledge Base model calls. + This is the OpenAI endpoint used by the Knowledge Base to call the LLM for + query planning and reasoning. This is separate from the project endpoint because + the Knowledge Base directly calls Azure OpenAI for its internal operations. + Required for agentic mode. Example: "https://myresource.openai.azure.com" + model_deployment_name: Model deployment name in Azure OpenAI for Knowledge Base. + This is the deployment name the Knowledge Base uses to call the LLM. + Required for agentic mode. + model_name: The underlying model name (e.g., "gpt-4o", "gpt-4o-mini"). + If not provided, defaults to model_deployment_name. Used for Knowledge Base configuration. + knowledge_base_name: Name for the Knowledge Base. Required for agentic mode. + retrieval_instructions: Custom instructions for the Knowledge Base's + retrieval planning. Only used in agentic mode. + azure_openai_api_key: Azure OpenAI API key for Knowledge Base to call the model. + Only needed when using API key authentication instead of managed identity. + knowledge_base_output_mode: Output mode for Knowledge Base retrieval. Only used in agentic mode. + "extractive_data": Returns raw chunks without synthesis (default, recommended for agent integration). + "answer_synthesis": Returns synthesized answer from the LLM. + Some knowledge sources require answer_synthesis mode. Default: "extractive_data". + retrieval_reasoning_effort: Reasoning effort for Knowledge Base query planning. Only used in agentic mode. + "minimal": Fastest, basic query planning. + "medium": Moderate reasoning with some query decomposition. + "low": Lower reasoning effort than medium. + Default: "minimal". + agentic_message_history_count: Number of recent messages from conversation history to send to + the Knowledge Base. This context helps with query planning in agentic mode, allowing the + Knowledge Base to understand the conversation flow and generate better retrieval queries. + There is no technical limit - adjust based on your use case. Default: 10. + env_file_path: Path to environment file for loading settings. + env_file_encoding: Encoding of the environment file. + + Examples: + .. code-block:: python + + from agent_framework_aisearch import AzureAISearchContextProvider + from azure.identity.aio import DefaultAzureCredential + + # Using environment variables + # Set AZURE_SEARCH_ENDPOINT=https://mysearch.search.windows.net + # Set AZURE_SEARCH_INDEX_NAME=my-index + credential = DefaultAzureCredential() + provider = AzureAISearchContextProvider(credential=credential) + + # Or passing parameters directly + provider = AzureAISearchContextProvider( + endpoint="https://mysearch.search.windows.net", + index_name="my-index", + credential=credential, + ) + + # Or loading from a .env file + provider = AzureAISearchContextProvider(credential=credential, env_file_path="path/to/.env") + """ + # Load settings from environment/file + try: + settings = AzureAISearchSettings( + endpoint=endpoint, + index_name=index_name, + api_key=api_key if isinstance(api_key, str) else None, + env_file_path=env_file_path, + env_file_encoding=env_file_encoding, + ) + except ValidationError as ex: + raise ServiceInitializationError("Failed to create Azure AI Search settings.", ex) from ex + + # Validate required parameters + if not settings.endpoint: + raise ServiceInitializationError( + "Azure AI Search endpoint is required. Set via 'endpoint' parameter " + "or 'AZURE_SEARCH_ENDPOINT' environment variable." + ) + if not settings.index_name: + raise ServiceInitializationError( + "Azure AI Search index name is required. Set via 'index_name' parameter " + "or 'AZURE_SEARCH_INDEX_NAME' environment variable." + ) + + # Determine the credential to use + resolved_credential: AzureKeyCredential | AsyncTokenCredential + if credential: + # AsyncTokenCredential takes precedence + resolved_credential = credential + elif isinstance(api_key, AzureKeyCredential): + resolved_credential = api_key + elif settings.api_key: + resolved_credential = AzureKeyCredential(settings.api_key.get_secret_value()) + else: + raise ServiceInitializationError( + "Azure credential is required. Provide 'api_key' or 'credential' parameter " + "or set 'AZURE_SEARCH_API_KEY' environment variable." + ) + + self.endpoint = settings.endpoint + self.index_name = settings.index_name + self.credential = resolved_credential + self.mode = mode + self.top_k = top_k + self.semantic_configuration_name = semantic_configuration_name + self.vector_field_name = vector_field_name + self.embedding_function = embedding_function + self.context_prompt = context_prompt or self._DEFAULT_SEARCH_CONTEXT_PROMPT + + # Agentic mode parameters (Knowledge Base) + self.azure_openai_resource_url = azure_openai_resource_url + self.azure_openai_deployment_name = model_deployment_name + # If model_name not provided, default to deployment name + self.model_name = model_name or model_deployment_name + self.knowledge_base_name = knowledge_base_name + self.retrieval_instructions = retrieval_instructions + self.azure_openai_api_key = azure_openai_api_key + self.azure_ai_project_endpoint = azure_ai_project_endpoint + self.knowledge_base_output_mode = knowledge_base_output_mode + self.retrieval_reasoning_effort = retrieval_reasoning_effort + self.agentic_message_history_count = agentic_message_history_count + + # Auto-discover vector field if not specified + self._auto_discovered_vector_field = False + self._use_vectorizable_query = False # Will be set to True if server-side vectorization detected + if not vector_field_name and mode == "semantic": + # Attempt to auto-discover vector field from index schema + # This will be done lazily on first search to avoid blocking initialization + pass + + # Validation + if vector_field_name and not embedding_function: + raise ValueError("embedding_function is required when vector_field_name is specified") + + if mode == "agentic": + if not _agentic_retrieval_available: + raise ImportError( + "Agentic retrieval requires azure-search-documents >= 11.7.0b1 with Knowledge Base support. " + "Please upgrade: pip install azure-search-documents>=11.7.0b1" + ) + if not self.azure_openai_resource_url: + raise ValueError( + "azure_openai_resource_url is required for agentic mode. " + "This should be your Azure OpenAI endpoint (e.g., 'https://myresource.openai.azure.com')" + ) + if not self.azure_openai_deployment_name: + raise ValueError("model_deployment_name is required for agentic mode") + if not knowledge_base_name: + raise ValueError("knowledge_base_name is required for agentic mode") + + # Create search client for semantic mode + self._search_client = SearchClient( + endpoint=self.endpoint, + index_name=self.index_name, + credential=self.credential, + ) + + # Create index client and retrieval client for agentic mode (Knowledge Base) + self._index_client: SearchIndexClient | None = None + self._retrieval_client: KnowledgeBaseRetrievalClient | None = None + if mode == "agentic": + self._index_client = SearchIndexClient( + endpoint=self.endpoint, + credential=self.credential, + ) + # Retrieval client will be created after Knowledge Base initialization + + self._knowledge_base_initialized = False + + async def __aenter__(self) -> Self: + """Async context manager entry.""" + return self + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc_val: BaseException | None, + exc_tb: Any, + ) -> None: + """Async context manager exit - cleanup clients. + + Args: + exc_type: Exception type if an error occurred. + exc_val: Exception value if an error occurred. + exc_tb: Exception traceback if an error occurred. + """ + # Close retrieval client if it was created + if self._retrieval_client is not None: + await self._retrieval_client.close() + self._retrieval_client = None + + @override + async def invoking( + self, + messages: ChatMessage | MutableSequence[ChatMessage], + **kwargs: Any, + ) -> Context: + """Retrieve relevant context from Azure AI Search before model invocation. + + Args: + messages: User messages to use for context retrieval. + **kwargs: Additional arguments (unused). + + Returns: + Context object with retrieved documents as messages. + """ + # Convert to list and filter to USER/ASSISTANT messages with text only + messages_list = [messages] if isinstance(messages, ChatMessage) else list(messages) + + filtered_messages = [ + msg + for msg in messages_list + if msg and msg.text and msg.text.strip() and msg.role in [Role.USER, Role.ASSISTANT] + ] + + if not filtered_messages: + return Context() + + # Perform search based on mode + if self.mode == "semantic": + # Semantic mode: flatten messages to single query + query = "\n".join(msg.text for msg in filtered_messages) + search_result_parts = await self._semantic_search(query) + else: # agentic + # Agentic mode: pass recent messages as conversation history + recent_messages = filtered_messages[-self.agentic_message_history_count :] + search_result_parts = await self._agentic_search(recent_messages) + + # Format results as context - return multiple messages for each result part + if not search_result_parts: + return Context() + + # Create context messages: first message with prompt, then one message per result part + context_messages = [ChatMessage(role=Role.USER, text=self.context_prompt)] + context_messages.extend([ChatMessage(role=Role.USER, text=part) for part in search_result_parts]) + + return Context(messages=context_messages) + + def _find_vector_fields(self, index: Any) -> list[str]: + """Find all fields that can store vectors (have dimensions defined). + + Args: + index: SearchIndex object from Azure Search. + + Returns: + List of vector field names. + """ + return [ + field.name + for field in index.fields + if field.vector_search_dimensions is not None and field.vector_search_dimensions > 0 + ] + + def _find_vectorizable_fields(self, index: Any, vector_fields: list[str]) -> list[str]: + """Find vector fields that have auto-vectorization configured. + + These are fields that have a vectorizer in their profile, meaning the index + can automatically vectorize text queries without needing a client-side embedding function. + + Args: + index: SearchIndex object from Azure Search. + vector_fields: List of vector field names. + + Returns: + List of vectorizable field names (subset of vector_fields). + """ + vectorizable_fields: list[str] = [] + + # Check if index has vector search configuration + if not index.vector_search or not index.vector_search.profiles: + return vectorizable_fields + + # For each vector field, check if it has a vectorizer configured + for field in index.fields: + if field.name in vector_fields and field.vector_search_profile_name: + # Find the profile for this field + profile = next( + (p for p in index.vector_search.profiles if p.name == field.vector_search_profile_name), None + ) + + if profile and hasattr(profile, "vectorizer_name") and profile.vectorizer_name: + # This field has server-side vectorization configured + vectorizable_fields.append(field.name) + + return vectorizable_fields + + async def _auto_discover_vector_field(self) -> None: + """Auto-discover vector field from index schema. + + Attempts to find vector fields in the index and detect which have server-side + vectorization configured. Prioritizes vectorizable fields (which can auto-embed text) + over regular vector fields (which require client-side embedding). + """ + if self._auto_discovered_vector_field or self.vector_field_name: + return # Already discovered or manually specified + + try: + # Use existing index client or create temporary one + if not self._index_client: + self._index_client = SearchIndexClient(endpoint=self.endpoint, credential=self.credential) + index_client = self._index_client + + # Get index schema + index = await index_client.get_index(self.index_name) + + # Step 1: Find all vector fields + vector_fields = self._find_vector_fields(index) + + if not vector_fields: + # No vector fields found - keyword search only + logger.info(f"No vector fields found in index '{self.index_name}'. Using keyword-only search.") + self._auto_discovered_vector_field = True + return + + # Step 2: Find which vector fields have server-side vectorization + vectorizable_fields = self._find_vectorizable_fields(index, vector_fields) + + # Step 3: Decide which field to use + if vectorizable_fields: + # Prefer vectorizable fields (server-side embedding) + if len(vectorizable_fields) == 1: + self.vector_field_name = vectorizable_fields[0] + self._auto_discovered_vector_field = True + self._use_vectorizable_query = True # Use VectorizableTextQuery + logger.info( + f"Auto-discovered vectorizable field '{self.vector_field_name}' " + f"with server-side vectorization. No embedding_function needed." + ) + else: + # Multiple vectorizable fields + logger.warning( + f"Multiple vectorizable fields found: {vectorizable_fields}. " + f"Please specify vector_field_name explicitly. Using keyword-only search." + ) + elif len(vector_fields) == 1: + # Single vector field without vectorizer - needs client-side embedding + self.vector_field_name = vector_fields[0] + self._auto_discovered_vector_field = True + self._use_vectorizable_query = False + + if not self.embedding_function: + logger.warning( + f"Auto-discovered vector field '{self.vector_field_name}' without server-side vectorization. " + f"Provide embedding_function for vector search, or it will fall back to keyword-only search." + ) + self.vector_field_name = None + else: + # Multiple vector fields without vectorizers + logger.warning( + f"Multiple vector fields found: {vector_fields}. " + f"Please specify vector_field_name explicitly. Using keyword-only search." + ) + + except Exception as e: + # Log warning but continue with keyword search + logger.warning(f"Failed to auto-discover vector field: {e}. Using keyword-only search.") + + self._auto_discovered_vector_field = True # Mark as attempted + + async def _semantic_search(self, query: str) -> list[str]: + """Perform semantic hybrid search with semantic ranking. + + This is the recommended mode for most use cases. It combines: + - Vector search (if embedding_function provided) + - Keyword search (BM25) + - Semantic reranking (if semantic_configuration_name provided) + + Args: + query: Search query text. + + Returns: + List of formatted search result strings, one per document. + """ + # Auto-discover vector field if not already done + await self._auto_discover_vector_field() + + vector_queries: list[VectorizableTextQuery | VectorizedQuery] = [] + + # Build vector query based on server-side vectorization or client-side embedding + if self.vector_field_name: + # Use larger k for vector query when semantic reranker is enabled for better ranking quality + vector_k = max(self.top_k, 50) if self.semantic_configuration_name else self.top_k + + if self._use_vectorizable_query: + # Server-side vectorization: Index will auto-embed the text query + vector_queries = [ + VectorizableTextQuery( + text=query, + k_nearest_neighbors=vector_k, + fields=self.vector_field_name, + ) + ] + elif self.embedding_function: + # Client-side embedding: We provide the vector + query_vector = await self.embedding_function(query) + vector_queries = [ + VectorizedQuery( + vector=query_vector, + k_nearest_neighbors=vector_k, + fields=self.vector_field_name, + ) + ] + # else: vector_field_name is set but no vectorization available - skip vector search + + # Build search parameters + search_params: dict[str, Any] = { + "search_text": query, + "top": self.top_k, + } + + if vector_queries: + search_params["vector_queries"] = vector_queries + + # Add semantic ranking if configured + if self.semantic_configuration_name: + search_params["query_type"] = QueryType.SEMANTIC + search_params["semantic_configuration_name"] = self.semantic_configuration_name + search_params["query_caption"] = QueryCaptionType.EXTRACTIVE + + # Execute search + results = await self._search_client.search(**search_params) # type: ignore[reportUnknownVariableType] + + # Format results with citations + formatted_results: list[str] = [] + async for doc in results: # type: ignore[reportUnknownVariableType] + # Extract document ID for citation + doc_id = doc.get("id") or doc.get("@search.id") # type: ignore[reportUnknownVariableType] + + # Use full document chunks with citation + doc_text: str = self._extract_document_text(doc, doc_id=doc_id) # type: ignore[reportUnknownArgumentType] + if doc_text: + formatted_results.append(doc_text) # type: ignore[reportUnknownArgumentType] + + return formatted_results + + async def _ensure_knowledge_base(self) -> None: + """Ensure Knowledge Base and knowledge source are created. + + This method is idempotent - it will only create resources if they don't exist. + + Note: Azure SDK uses KnowledgeAgent classes internally, but the feature + is marketed as "Knowledge Bases" in Azure AI Search. + """ + if self._knowledge_base_initialized or not self._index_client: + return + + # Runtime validation for agentic mode parameters + if not self.knowledge_base_name: + raise ValueError("knowledge_base_name is required for agentic mode") + if not self.azure_openai_resource_url: + raise ValueError("azure_openai_resource_url is required for agentic mode") + if not self.azure_openai_deployment_name: + raise ValueError("model_deployment_name is required for agentic mode") + + knowledge_base_name = self.knowledge_base_name + + # Step 1: Create or get knowledge source + knowledge_source_name = f"{self.index_name}-source" + + try: + # Try to get existing knowledge source + await self._index_client.get_knowledge_source(knowledge_source_name) + except ResourceNotFoundError: + # Create new knowledge source if it doesn't exist + knowledge_source = SearchIndexKnowledgeSource( + name=knowledge_source_name, + description=f"Knowledge source for {self.index_name} search index", + search_index_parameters=SearchIndexKnowledgeSourceParameters( + search_index_name=self.index_name, + ), + ) + await self._index_client.create_knowledge_source(knowledge_source) + + # Step 2: Create or update Knowledge Base + # Always create/update to ensure configuration is current + aoai_params = AzureOpenAIVectorizerParameters( + resource_url=self.azure_openai_resource_url, + deployment_name=self.azure_openai_deployment_name, + model_name=self.model_name, + api_key=self.azure_openai_api_key, + ) + + # Map output mode string to SDK enum + output_mode = ( + KnowledgeRetrievalOutputMode.EXTRACTIVE_DATA + if self.knowledge_base_output_mode == "extractive_data" + else KnowledgeRetrievalOutputMode.ANSWER_SYNTHESIS + ) + + # Map reasoning effort string to SDK class + reasoning_effort_map: dict[str, KnowledgeRetrievalReasoningEffort] = { + "minimal": KnowledgeRetrievalMinimalReasoningEffort(), + "medium": KnowledgeRetrievalMediumReasoningEffort(), + "low": KnowledgeRetrievalLowReasoningEffort(), + } + reasoning_effort = reasoning_effort_map[self.retrieval_reasoning_effort] + + knowledge_base = KnowledgeBase( + name=knowledge_base_name, + description=f"Knowledge Base for multi-hop retrieval across {self.index_name}", + knowledge_sources=[ + KnowledgeSourceReference( + name=knowledge_source_name, + ) + ], + models=[KnowledgeBaseAzureOpenAIModel(azure_open_ai_parameters=aoai_params)], + output_mode=output_mode, + retrieval_reasoning_effort=reasoning_effort, + ) + await self._index_client.create_or_update_knowledge_base(knowledge_base) + + self._knowledge_base_initialized = True + + # Create retrieval client now that Knowledge Base is initialized + if _agentic_retrieval_available and self._retrieval_client is None: + self._retrieval_client = KnowledgeBaseRetrievalClient( + endpoint=self.endpoint, + knowledge_base_name=knowledge_base_name, + credential=self.credential, + ) + + async def _agentic_search(self, messages: list[ChatMessage]) -> list[str]: + """Perform agentic retrieval with multi-hop reasoning using Knowledge Bases. + + This mode uses query planning and is slightly slower than semantic search, + but provides more accurate results through intelligent retrieval. + + This method uses Azure AI Search Knowledge Bases which: + 1. Analyze the query and plan sub-queries + 2. Retrieve relevant documents across multiple sources + 3. Perform multi-hop reasoning with an LLM + 4. Synthesize a comprehensive answer with references + + Args: + messages: Conversation history to use for retrieval context. + + Returns: + List of answer parts from the Knowledge Base, one per content item. + """ + # Ensure Knowledge Base is initialized + await self._ensure_knowledge_base() + + # Map reasoning effort string to SDK class (for retrieval requests) + reasoning_effort_map: dict[str, KBRetrievalReasoningEffort] = { + "minimal": KBRetrievalMinimalReasoningEffort(), + "medium": KBRetrievalMediumReasoningEffort(), + "low": KBRetrievalLowReasoningEffort(), + } + reasoning_effort = reasoning_effort_map[self.retrieval_reasoning_effort] + + # Map output mode string to SDK enum (for retrieval requests) + output_mode = ( + KBRetrievalOutputMode.EXTRACTIVE_DATA + if self.knowledge_base_output_mode == "extractive_data" + else KBRetrievalOutputMode.ANSWER_SYNTHESIS + ) + + # For minimal reasoning, use intents API; for medium/low, use messages API + if self.retrieval_reasoning_effort == "minimal": + # Minimal reasoning uses intents with a single search query + query = "\n".join(msg.text for msg in messages if msg.text) + intents: list[KnowledgeRetrievalIntent] = [KnowledgeRetrievalSemanticIntent(search=query)] + retrieval_request = KnowledgeBaseRetrievalRequest( + intents=intents, + retrieval_reasoning_effort=reasoning_effort, + output_mode=output_mode, + include_activity=True, + ) + else: + # Medium/low reasoning uses messages with conversation history + kb_messages = [ + KnowledgeBaseMessage( + role=msg.role.value if hasattr(msg.role, "value") else str(msg.role), + content=[KnowledgeBaseMessageTextContent(text=msg.text)], + ) + for msg in messages + if msg.text + ] + retrieval_request = KnowledgeBaseRetrievalRequest( + messages=kb_messages, + retrieval_reasoning_effort=reasoning_effort, + output_mode=output_mode, + include_activity=True, + ) + + # Use reusable retrieval client + if not self._retrieval_client: + raise RuntimeError("Retrieval client not initialized. Ensure Knowledge Base is set up correctly.") + + # Perform retrieval via Knowledge Base + retrieval_result = await self._retrieval_client.retrieve(retrieval_request=retrieval_request) + + # Extract answer parts from response + if retrieval_result.response and len(retrieval_result.response) > 0: + # Get the assistant's response (last message) + assistant_message = retrieval_result.response[-1] + if assistant_message.content: + # Extract all text content items as separate parts + answer_parts: list[str] = [] + for content_item in assistant_message.content: + # Check if this is a text content item + if isinstance(content_item, KnowledgeBaseMessageTextContent) and content_item.text: + answer_parts.append(content_item.text) + + if answer_parts: + return answer_parts + + # Fallback if no answer generated + return ["No results found from Knowledge Base."] + + def _extract_document_text(self, doc: dict[str, Any], doc_id: str | None = None) -> str: + """Extract readable text from a search document with optional citation. + + Args: + doc: Search result document. + doc_id: Optional document ID for citation. + + Returns: + Formatted document text with citation if doc_id provided. + """ + # Try common text field names + text = "" + for field in ["content", "text", "description", "body", "chunk"]: + if doc.get(field): + text = str(doc[field]) + break + + # Fallback: concatenate all string fields + if not text: + text_parts: list[str] = [] + for key, value in doc.items(): + if isinstance(value, str) and not key.startswith("@") and key != "id": + text_parts.append(f"{key}: {value}") + text = " | ".join(text_parts) if text_parts else "" + + # Add citation if document ID provided + if doc_id and text: + return f"[Source: {doc_id}] {text}" + return text diff --git a/python/packages/aisearch/pyproject.toml b/python/packages/aisearch/pyproject.toml new file mode 100644 index 0000000000..5d64b398aa --- /dev/null +++ b/python/packages/aisearch/pyproject.toml @@ -0,0 +1,91 @@ +[project] +name = "agent-framework-aisearch" +description = "Azure AI Search integration for Microsoft Agent Framework." +authors = [{ name = "Microsoft", email = "af-support@microsoft.com"}] +readme = "README.md" +requires-python = ">=3.10" +version = "1.0.0b251118" +license-files = ["LICENSE"] +urls.homepage = "https://aka.ms/agent-framework" +urls.source = "https://github.com/microsoft/agent-framework/tree/main/python" +urls.release_notes = "https://github.com/microsoft/agent-framework/releases?q=tag%3Apython-1&expanded=true" +urls.issues = "https://github.com/microsoft/agent-framework/issues" +classifiers = [ + "License :: OSI Approved :: MIT License", + "Development Status :: 4 - Beta", + "Intended Audience :: Developers", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", + "Programming Language :: Python :: 3.14", + "Typing :: Typed", +] +dependencies = [ + "agent-framework-core", + "azure-search-documents==11.7.0b2", +] + +[tool.uv] +prerelease = "if-necessary-or-explicit" +environments = [ + "sys_platform == 'darwin'", + "sys_platform == 'linux'", + "sys_platform == 'win32'" +] + +[tool.uv-dynamic-versioning] +fallback-version = "0.0.0" + +[tool.pytest.ini_options] +testpaths = 'tests' +addopts = "-ra -q -r fEX" +asyncio_mode = "auto" +asyncio_default_fixture_loop_scope = "function" +filterwarnings = [ + "ignore:Support for class-based `config` is deprecated:DeprecationWarning:pydantic.*" +] +timeout = 120 + +[tool.ruff] +extend = "../../pyproject.toml" +exclude = ["examples"] + +[tool.coverage.run] +omit = [ + "**/__init__.py" +] + +[tool.pyright] +extends = "../../pyproject.toml" +exclude = ['tests'] + +[tool.mypy] +plugins = ['pydantic.mypy'] +strict = true +python_version = "3.10" +ignore_missing_imports = true +disallow_untyped_defs = true +no_implicit_optional = true +check_untyped_defs = true +warn_return_any = true +show_error_codes = true +warn_unused_ignores = false +disallow_incomplete_defs = true +disallow_untyped_decorators = true + +[tool.bandit] +targets = ["agent_framework_aisearch"] +exclude_dirs = ["tests"] + +[tool.poe] +executor.type = "uv" +include = "../../shared_tasks.toml" +[tool.poe.tasks] +mypy = "mypy --config-file $POE_ROOT/pyproject.toml agent_framework_aisearch" +test = "pytest --cov=agent_framework_aisearch --cov-report=term-missing:skip-covered tests" + +[build-system] +requires = ["flit-core >= 3.11,<4.0"] +build-backend = "flit_core.buildapi" diff --git a/python/packages/aisearch/tests/test_search_provider.py b/python/packages/aisearch/tests/test_search_provider.py new file mode 100644 index 0000000000..6813c3d16a --- /dev/null +++ b/python/packages/aisearch/tests/test_search_provider.py @@ -0,0 +1,992 @@ +# Copyright (c) Microsoft. All rights reserved. +# pyright: reportPrivateUsage=false + +import os +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from agent_framework import ChatMessage, Context, Role +from agent_framework.azure import AzureAISearchContextProvider +from agent_framework.exceptions import ServiceInitializationError +from azure.core.credentials import AzureKeyCredential +from azure.core.exceptions import ResourceNotFoundError + +from agent_framework_aisearch import AzureAISearchSettings + + +@pytest.fixture +def mock_search_client() -> AsyncMock: + """Create a mock SearchClient.""" + mock_client = AsyncMock() + mock_client.search = AsyncMock() + mock_client.__aenter__ = AsyncMock(return_value=mock_client) + mock_client.__aexit__ = AsyncMock() + return mock_client + + +@pytest.fixture +def mock_index_client() -> AsyncMock: + """Create a mock SearchIndexClient.""" + mock_client = AsyncMock() + mock_client.get_knowledge_source = AsyncMock() + mock_client.create_knowledge_source = AsyncMock() + mock_client.get_agent = AsyncMock() + mock_client.create_agent = AsyncMock() + mock_client.__aenter__ = AsyncMock(return_value=mock_client) + mock_client.__aexit__ = AsyncMock() + return mock_client + + +@pytest.fixture +def sample_messages() -> list[ChatMessage]: + """Create sample chat messages for testing.""" + return [ + ChatMessage(role=Role.USER, text="What is in the documents?"), + ] + + +class TestAzureAISearchSettings: + """Test AzureAISearchSettings configuration.""" + + def test_settings_with_direct_values(self) -> None: + """Test settings with direct values.""" + settings = AzureAISearchSettings( + endpoint="https://test.search.windows.net", + index_name="test-index", + api_key="test-key", + ) + assert settings.endpoint == "https://test.search.windows.net" + assert settings.index_name == "test-index" + # api_key is now SecretStr + assert settings.api_key.get_secret_value() == "test-key" + + def test_settings_with_env_file_path(self) -> None: + """Test settings with env_file_path parameter.""" + settings = AzureAISearchSettings( + endpoint="https://test.search.windows.net", + index_name="test-index", + env_file_path="test.env", + ) + assert settings.endpoint == "https://test.search.windows.net" + assert settings.index_name == "test-index" + + def test_provider_uses_settings_from_env(self) -> None: + """Test that provider creates settings internally from env.""" + provider = AzureAISearchContextProvider( + endpoint="https://test.search.windows.net", + index_name="test-index", + api_key="test-key", + ) + assert provider.endpoint == "https://test.search.windows.net" + assert provider.index_name == "test-index" + + def test_provider_missing_endpoint_raises_error(self) -> None: + """Test that provider raises ServiceInitializationError without endpoint.""" + # Use patch.dict to clear environment and pass env_file_path="" to prevent .env file loading + clean_env = {k: v for k, v in os.environ.items() if not k.startswith("AZURE_SEARCH_")} + with ( + patch.dict(os.environ, clean_env, clear=True), + pytest.raises(ServiceInitializationError, match="endpoint is required"), + ): + AzureAISearchContextProvider( + index_name="test-index", + api_key="test-key", + env_file_path="", # Disable .env file loading + ) + + def test_provider_missing_index_name_raises_error(self) -> None: + """Test that provider raises ServiceInitializationError without index_name.""" + # Use patch.dict to clear environment and pass env_file_path="" to prevent .env file loading + clean_env = {k: v for k, v in os.environ.items() if not k.startswith("AZURE_SEARCH_")} + with ( + patch.dict(os.environ, clean_env, clear=True), + pytest.raises(ServiceInitializationError, match="index name is required"), + ): + AzureAISearchContextProvider( + endpoint="https://test.search.windows.net", + api_key="test-key", + env_file_path="", # Disable .env file loading + ) + + def test_provider_missing_credential_raises_error(self) -> None: + """Test that provider raises ServiceInitializationError without credential.""" + # Use patch.dict to clear environment and pass env_file_path="" to prevent .env file loading + clean_env = {k: v for k, v in os.environ.items() if not k.startswith("AZURE_SEARCH_")} + with ( + patch.dict(os.environ, clean_env, clear=True), + pytest.raises(ServiceInitializationError, match="credential is required"), + ): + AzureAISearchContextProvider( + endpoint="https://test.search.windows.net", + index_name="test-index", + env_file_path="", # Disable .env file loading + ) + + +class TestSearchProviderInitialization: + """Test initialization and configuration of AzureAISearchContextProvider.""" + + def test_init_semantic_mode_minimal(self) -> None: + """Test initialization with minimal semantic mode parameters.""" + provider = AzureAISearchContextProvider( + endpoint="https://test.search.windows.net", + index_name="test-index", + api_key="test-key", + mode="semantic", + ) + assert provider.endpoint == "https://test.search.windows.net" + assert provider.index_name == "test-index" + assert provider.mode == "semantic" + assert provider.top_k == 5 + + def test_init_semantic_mode_with_vector_field_requires_embedding_function(self) -> None: + """Test that vector_field_name requires embedding_function.""" + with pytest.raises(ValueError, match="embedding_function is required"): + AzureAISearchContextProvider( + endpoint="https://test.search.windows.net", + index_name="test-index", + api_key="test-key", + mode="semantic", + vector_field_name="embedding", + ) + + def test_init_agentic_mode_requires_azure_openai_resource_url(self) -> None: + """Test that agentic mode requires azure_openai_resource_url.""" + with pytest.raises(ValueError, match="azure_openai_resource_url"): + AzureAISearchContextProvider( + endpoint="https://test.search.windows.net", + index_name="test-index", + api_key="test-key", + mode="agentic", + ) + + def test_init_agentic_mode_requires_model_deployment_name(self) -> None: + """Test that agentic mode requires model_deployment_name.""" + with pytest.raises(ValueError, match="model_deployment_name"): + AzureAISearchContextProvider( + endpoint="https://test.search.windows.net", + index_name="test-index", + api_key="test-key", + mode="agentic", + azure_ai_project_endpoint="https://test.services.ai.azure.com", + azure_openai_resource_url="https://test.openai.azure.com", + ) + + def test_init_agentic_mode_requires_knowledge_base_name(self) -> None: + """Test that agentic mode requires knowledge_base_name.""" + with pytest.raises(ValueError, match="knowledge_base_name"): + AzureAISearchContextProvider( + endpoint="https://test.search.windows.net", + index_name="test-index", + api_key="test-key", + mode="agentic", + azure_ai_project_endpoint="https://test.services.ai.azure.com", + model_deployment_name="gpt-4o", + azure_openai_resource_url="https://test.openai.azure.com", + ) + + def test_init_agentic_mode_with_all_params(self) -> None: + """Test initialization with all agentic mode parameters.""" + provider = AzureAISearchContextProvider( + endpoint="https://test.search.windows.net", + index_name="test-index", + api_key="test-key", + mode="agentic", + azure_ai_project_endpoint="https://test.services.ai.azure.com", + model_deployment_name="my-gpt-4o-deployment", + model_name="gpt-4o", + knowledge_base_name="test-kb", + azure_openai_resource_url="https://test.openai.azure.com", + ) + assert provider.mode == "agentic" + assert provider.azure_ai_project_endpoint == "https://test.services.ai.azure.com" + assert provider.azure_openai_resource_url == "https://test.openai.azure.com" + assert provider.azure_openai_deployment_name == "my-gpt-4o-deployment" + assert provider.model_name == "gpt-4o" + assert provider.knowledge_base_name == "test-kb" + + def test_init_model_name_defaults_to_deployment_name(self) -> None: + """Test that model_name defaults to deployment_name if not provided.""" + provider = AzureAISearchContextProvider( + endpoint="https://test.search.windows.net", + index_name="test-index", + api_key="test-key", + mode="agentic", + azure_ai_project_endpoint="https://test.services.ai.azure.com", + model_deployment_name="gpt-4o", + knowledge_base_name="test-kb", + azure_openai_resource_url="https://test.openai.azure.com", + ) + assert provider.model_name == "gpt-4o" + + def test_init_with_custom_context_prompt(self) -> None: + """Test initialization with custom context prompt.""" + custom_prompt = "Use the following information:" + provider = AzureAISearchContextProvider( + endpoint="https://test.search.windows.net", + index_name="test-index", + api_key="test-key", + mode="semantic", + context_prompt=custom_prompt, + ) + assert provider.context_prompt == custom_prompt + + def test_init_uses_default_context_prompt(self) -> None: + """Test that default context prompt is used when not provided.""" + provider = AzureAISearchContextProvider( + endpoint="https://test.search.windows.net", + index_name="test-index", + api_key="test-key", + mode="semantic", + ) + assert provider.context_prompt == provider._DEFAULT_SEARCH_CONTEXT_PROMPT + + +class TestSemanticSearch: + """Test semantic search functionality.""" + + @pytest.mark.asyncio + @patch("agent_framework_aisearch._search_provider.SearchClient") + async def test_semantic_search_basic( + self, mock_search_class: MagicMock, sample_messages: list[ChatMessage] + ) -> None: + """Test basic semantic search without vector search.""" + # Setup mock + mock_search_client = AsyncMock() + mock_results = AsyncMock() + mock_results.__aiter__.return_value = iter([{"content": "Test document content"}]) + mock_search_client.search.return_value = mock_results + mock_search_class.return_value = mock_search_client + + provider = AzureAISearchContextProvider( + endpoint="https://test.search.windows.net", + index_name="test-index", + api_key="test-key", + mode="semantic", + ) + + context = await provider.invoking(sample_messages) + + assert isinstance(context, Context) + assert len(context.messages) > 1 # First message is prompt, rest are results + # First message should be the context prompt + assert "Use the following context" in context.messages[0].text + # Second message should contain the search result + assert "Test document content" in context.messages[1].text + + @pytest.mark.asyncio + @patch("agent_framework_aisearch._search_provider.SearchClient") + async def test_semantic_search_empty_query(self, mock_search_class: MagicMock) -> None: + """Test that empty queries return empty context.""" + mock_search_client = AsyncMock() + mock_search_class.return_value = mock_search_client + + provider = AzureAISearchContextProvider( + endpoint="https://test.search.windows.net", + index_name="test-index", + api_key="test-key", + mode="semantic", + ) + + # Empty message + context = await provider.invoking([ChatMessage(role=Role.USER, text="")]) + + assert isinstance(context, Context) + assert len(context.messages) == 0 + + @pytest.mark.asyncio + @patch("agent_framework_aisearch._search_provider.SearchClient") + async def test_semantic_search_with_vector_query( + self, mock_search_class: MagicMock, sample_messages: list[ChatMessage] + ) -> None: + """Test semantic search with vector query.""" + # Setup mock + mock_search_client = AsyncMock() + mock_results = AsyncMock() + mock_results.__aiter__.return_value = iter([{"content": "Vector search result"}]) + mock_search_client.search.return_value = mock_results + mock_search_class.return_value = mock_search_client + + # Mock embedding function + async def mock_embed(text: str) -> list[float]: + return [0.1, 0.2, 0.3] + + provider = AzureAISearchContextProvider( + endpoint="https://test.search.windows.net", + index_name="test-index", + api_key="test-key", + mode="semantic", + vector_field_name="embedding", + embedding_function=mock_embed, + ) + + context = await provider.invoking(sample_messages) + + assert isinstance(context, Context) + assert len(context.messages) > 0 + # Verify that search was called + mock_search_client.search.assert_called_once() + + +class TestKnowledgeBaseSetup: + """Test Knowledge Base setup for agentic mode.""" + + @pytest.mark.asyncio + @patch("agent_framework_aisearch._search_provider.SearchIndexClient") + @patch("agent_framework_aisearch._search_provider.SearchClient") + async def test_ensure_knowledge_base_creates_when_not_exists( + self, mock_search_class: MagicMock, mock_index_class: MagicMock + ) -> None: + """Test that Knowledge Base is created when it doesn't exist.""" + # Setup mocks + mock_index_client = AsyncMock() + mock_index_client.get_knowledge_source.side_effect = ResourceNotFoundError("Not found") + mock_index_client.create_knowledge_source = AsyncMock() + mock_index_client.get_knowledge_base.side_effect = ResourceNotFoundError("Not found") + mock_index_client.create_or_update_knowledge_base = AsyncMock() + mock_index_class.return_value = mock_index_client + + mock_search_client = AsyncMock() + mock_search_class.return_value = mock_search_client + + provider = AzureAISearchContextProvider( + endpoint="https://test.search.windows.net", + index_name="test-index", + api_key="test-key", + mode="agentic", + azure_ai_project_endpoint="https://test.services.ai.azure.com", + model_deployment_name="gpt-4o", + model_name="gpt-4o", + knowledge_base_name="test-kb", + azure_openai_resource_url="https://test.openai.azure.com", + ) + + await provider._ensure_knowledge_base() + + # Verify knowledge source was created + mock_index_client.create_knowledge_source.assert_called_once() + # Verify Knowledge Base was created + mock_index_client.create_or_update_knowledge_base.assert_called_once() + + @pytest.mark.asyncio + @patch("agent_framework_aisearch._search_provider.SearchIndexClient") + @patch("agent_framework_aisearch._search_provider.SearchClient") + async def test_ensure_knowledge_base_skips_when_exists( + self, mock_search_class: MagicMock, mock_index_class: MagicMock + ) -> None: + """Test that Knowledge Base setup is skipped when already exists.""" + # Setup mocks + mock_index_client = AsyncMock() + mock_index_client.get_knowledge_source.return_value = MagicMock() # Exists + mock_index_client.get_knowledge_base.return_value = MagicMock() # Exists + mock_index_class.return_value = mock_index_client + + mock_search_client = AsyncMock() + mock_search_class.return_value = mock_search_client + + provider = AzureAISearchContextProvider( + endpoint="https://test.search.windows.net", + index_name="test-index", + api_key="test-key", + mode="agentic", + azure_ai_project_endpoint="https://test.services.ai.azure.com", + model_deployment_name="gpt-4o", + knowledge_base_name="test-kb", + azure_openai_resource_url="https://test.openai.azure.com", + ) + + await provider._ensure_knowledge_base() + + # Verify nothing was created + mock_index_client.create_knowledge_source.assert_not_called() + mock_index_client.create_agent.assert_not_called() + + +class TestContextProviderLifecycle: + """Test context provider lifecycle methods.""" + + @pytest.mark.asyncio + @patch("agent_framework_aisearch._search_provider.SearchClient") + async def test_context_manager(self, mock_search_class: MagicMock) -> None: + """Test that provider can be used as async context manager.""" + mock_search_client = AsyncMock() + mock_search_class.return_value = mock_search_client + + async with AzureAISearchContextProvider( + endpoint="https://test.search.windows.net", + index_name="test-index", + api_key="test-key", + mode="semantic", + ) as provider: + assert provider is not None + assert isinstance(provider, AzureAISearchContextProvider) + + @pytest.mark.asyncio + @patch("agent_framework_aisearch._search_provider.KnowledgeBaseRetrievalClient") + @patch("agent_framework_aisearch._search_provider.SearchIndexClient") + @patch("agent_framework_aisearch._search_provider.SearchClient") + async def test_context_manager_agentic_cleanup( + self, mock_search_class: MagicMock, mock_index_class: MagicMock, mock_retrieval_class: MagicMock + ) -> None: + """Test that agentic mode provider cleans up retrieval client.""" + mock_search_client = AsyncMock() + mock_search_class.return_value = mock_search_client + + mock_index_client = AsyncMock() + mock_index_class.return_value = mock_index_client + + mock_retrieval_client = AsyncMock() + mock_retrieval_client.close = AsyncMock() + mock_retrieval_class.return_value = mock_retrieval_client + + async with AzureAISearchContextProvider( + endpoint="https://test.search.windows.net", + index_name="test-index", + api_key="test-key", + mode="agentic", + azure_ai_project_endpoint="https://test.services.ai.azure.com", + model_deployment_name="gpt-4o", + knowledge_base_name="test-kb", + azure_openai_resource_url="https://test.openai.azure.com", + ) as provider: + # Simulate retrieval client being created + provider._retrieval_client = mock_retrieval_client + + # Verify cleanup was called + mock_retrieval_client.close.assert_called_once() + + def test_string_api_key_conversion(self) -> None: + """Test that string api_key is converted to AzureKeyCredential.""" + provider = AzureAISearchContextProvider( + endpoint="https://test.search.windows.net", + index_name="test-index", + api_key="my-api-key", # String api_key + mode="semantic", + ) + assert isinstance(provider.credential, AzureKeyCredential) + + +class TestMessageFiltering: + """Test message filtering functionality.""" + + @pytest.mark.asyncio + @patch("agent_framework_aisearch._search_provider.SearchClient") + async def test_filters_non_user_assistant_messages(self, mock_search_class: MagicMock) -> None: + """Test that only USER and ASSISTANT messages are processed.""" + # Setup mock + mock_search_client = AsyncMock() + mock_results = AsyncMock() + mock_results.__aiter__.return_value = iter([{"content": "Test result"}]) + mock_search_client.search.return_value = mock_results + mock_search_class.return_value = mock_search_client + + provider = AzureAISearchContextProvider( + endpoint="https://test.search.windows.net", + index_name="test-index", + api_key="test-key", + mode="semantic", + ) + + # Mix of message types + messages = [ + ChatMessage(role=Role.SYSTEM, text="System message"), + ChatMessage(role=Role.USER, text="User message"), + ChatMessage(role=Role.ASSISTANT, text="Assistant message"), + ChatMessage(role=Role.TOOL, text="Tool message"), + ] + + context = await provider.invoking(messages) + + # Should have processed only USER and ASSISTANT messages + assert isinstance(context, Context) + mock_search_client.search.assert_called_once() + + @pytest.mark.asyncio + @patch("agent_framework_aisearch._search_provider.SearchClient") + async def test_filters_empty_messages(self, mock_search_class: MagicMock) -> None: + """Test that empty/whitespace messages are filtered out.""" + mock_search_client = AsyncMock() + mock_search_class.return_value = mock_search_client + + provider = AzureAISearchContextProvider( + endpoint="https://test.search.windows.net", + index_name="test-index", + api_key="test-key", + mode="semantic", + ) + + # Messages with empty/whitespace text + messages = [ + ChatMessage(role=Role.USER, text=""), + ChatMessage(role=Role.USER, text=" "), + ChatMessage(role=Role.USER, text=None), + ] + + context = await provider.invoking(messages) + + # Should return empty context + assert len(context.messages) == 0 + + +class TestCitations: + """Test citation functionality.""" + + @pytest.mark.asyncio + @patch("agent_framework_aisearch._search_provider.SearchClient") + async def test_citations_included_in_semantic_search(self, mock_search_class: MagicMock) -> None: + """Test that citations are included in semantic search results.""" + # Setup mock with document ID + mock_search_client = AsyncMock() + mock_results = AsyncMock() + mock_doc = {"id": "doc123", "content": "Test document content"} + mock_results.__aiter__.return_value = iter([mock_doc]) + mock_search_client.search.return_value = mock_results + mock_search_class.return_value = mock_search_client + + provider = AzureAISearchContextProvider( + endpoint="https://test.search.windows.net", + index_name="test-index", + api_key="test-key", + mode="semantic", + ) + + context = await provider.invoking([ChatMessage(role=Role.USER, text="test query")]) + + # Check that citation is included + assert isinstance(context, Context) + assert len(context.messages) > 1 # First message is prompt, rest are results + # Citation should be in the result message (second message) + assert "[Source: doc123]" in context.messages[1].text + assert "Test document content" in context.messages[1].text + + +class TestAgenticSearch: + """Test agentic search functionality.""" + + @pytest.mark.asyncio + @patch("agent_framework_aisearch._search_provider.KnowledgeBaseRetrievalClient") + @patch("agent_framework_aisearch._search_provider.SearchIndexClient") + @patch("agent_framework_aisearch._search_provider.SearchClient") + async def test_agentic_search_basic( + self, + mock_search_class: MagicMock, + mock_index_class: MagicMock, + mock_retrieval_class: MagicMock, + sample_messages: list[ChatMessage], + ) -> None: + """Test basic agentic search with Knowledge Base retrieval.""" + # Setup search client mock + mock_search_client = AsyncMock() + mock_search_class.return_value = mock_search_client + + # Setup index client mock + mock_index_client = AsyncMock() + mock_index_client.get_knowledge_source.side_effect = ResourceNotFoundError("Not found") + mock_index_client.create_knowledge_source = AsyncMock() + mock_index_client.create_or_update_knowledge_base = AsyncMock() + mock_index_class.return_value = mock_index_client + + # Setup retrieval client mock with response + mock_retrieval_client = AsyncMock() + mock_response = MagicMock() + mock_message = MagicMock() + mock_content = MagicMock() + mock_content.text = "Agentic search result" + # Make it pass isinstance check + from agent_framework_aisearch._search_provider import _agentic_retrieval_available + + if _agentic_retrieval_available: + from azure.search.documents.knowledgebases.models import KnowledgeBaseMessageTextContent + + mock_content.__class__ = KnowledgeBaseMessageTextContent + mock_message.content = [mock_content] + mock_response.response = [mock_message] + mock_retrieval_client.retrieve.return_value = mock_response + mock_retrieval_client.close = AsyncMock() + mock_retrieval_class.return_value = mock_retrieval_client + + provider = AzureAISearchContextProvider( + endpoint="https://test.search.windows.net", + index_name="test-index", + api_key="test-key", + mode="agentic", + azure_ai_project_endpoint="https://test.services.ai.azure.com", + model_deployment_name="gpt-4o", + knowledge_base_name="test-kb", + azure_openai_resource_url="https://test.openai.azure.com", + ) + + context = await provider.invoking(sample_messages) + + assert isinstance(context, Context) + # Should have at least the prompt message + assert len(context.messages) >= 1 + + @pytest.mark.asyncio + @patch("agent_framework_aisearch._search_provider.KnowledgeBaseRetrievalClient") + @patch("agent_framework_aisearch._search_provider.SearchIndexClient") + @patch("agent_framework_aisearch._search_provider.SearchClient") + async def test_agentic_search_no_results( + self, + mock_search_class: MagicMock, + mock_index_class: MagicMock, + mock_retrieval_class: MagicMock, + sample_messages: list[ChatMessage], + ) -> None: + """Test agentic search when no results are returned.""" + # Setup mocks + mock_search_client = AsyncMock() + mock_search_class.return_value = mock_search_client + + mock_index_client = AsyncMock() + mock_index_client.get_knowledge_source.side_effect = ResourceNotFoundError("Not found") + mock_index_client.create_knowledge_source = AsyncMock() + mock_index_client.create_or_update_knowledge_base = AsyncMock() + mock_index_class.return_value = mock_index_client + + # Empty response + mock_retrieval_client = AsyncMock() + mock_response = MagicMock() + mock_response.response = [] + mock_retrieval_client.retrieve.return_value = mock_response + mock_retrieval_client.close = AsyncMock() + mock_retrieval_class.return_value = mock_retrieval_client + + provider = AzureAISearchContextProvider( + endpoint="https://test.search.windows.net", + index_name="test-index", + api_key="test-key", + mode="agentic", + azure_ai_project_endpoint="https://test.services.ai.azure.com", + model_deployment_name="gpt-4o", + knowledge_base_name="test-kb", + azure_openai_resource_url="https://test.openai.azure.com", + ) + + context = await provider.invoking(sample_messages) + + assert isinstance(context, Context) + # Should have fallback message + assert len(context.messages) >= 1 + + @pytest.mark.asyncio + @patch("agent_framework_aisearch._search_provider.KnowledgeBaseRetrievalClient") + @patch("agent_framework_aisearch._search_provider.SearchIndexClient") + @patch("agent_framework_aisearch._search_provider.SearchClient") + async def test_agentic_search_with_medium_reasoning( + self, + mock_search_class: MagicMock, + mock_index_class: MagicMock, + mock_retrieval_class: MagicMock, + sample_messages: list[ChatMessage], + ) -> None: + """Test agentic search with medium reasoning effort.""" + # Setup mocks + mock_search_client = AsyncMock() + mock_search_class.return_value = mock_search_client + + mock_index_client = AsyncMock() + mock_index_client.get_knowledge_source.side_effect = ResourceNotFoundError("Not found") + mock_index_client.create_knowledge_source = AsyncMock() + mock_index_client.create_or_update_knowledge_base = AsyncMock() + mock_index_class.return_value = mock_index_client + + mock_retrieval_client = AsyncMock() + mock_response = MagicMock() + mock_message = MagicMock() + mock_content = MagicMock() + mock_content.text = "Medium reasoning result" + from agent_framework_aisearch._search_provider import _agentic_retrieval_available + + if _agentic_retrieval_available: + from azure.search.documents.knowledgebases.models import KnowledgeBaseMessageTextContent + + mock_content.__class__ = KnowledgeBaseMessageTextContent + mock_message.content = [mock_content] + mock_response.response = [mock_message] + mock_retrieval_client.retrieve.return_value = mock_response + mock_retrieval_client.close = AsyncMock() + mock_retrieval_class.return_value = mock_retrieval_client + + provider = AzureAISearchContextProvider( + endpoint="https://test.search.windows.net", + index_name="test-index", + api_key="test-key", + mode="agentic", + azure_ai_project_endpoint="https://test.services.ai.azure.com", + model_deployment_name="gpt-4o", + knowledge_base_name="test-kb", + azure_openai_resource_url="https://test.openai.azure.com", + retrieval_reasoning_effort="medium", # Test medium reasoning + ) + + context = await provider.invoking(sample_messages) + + assert isinstance(context, Context) + assert len(context.messages) >= 1 + + +class TestVectorFieldAutoDiscovery: + """Test vector field auto-discovery functionality.""" + + @pytest.mark.asyncio + @patch("agent_framework_aisearch._search_provider.SearchIndexClient") + @patch("agent_framework_aisearch._search_provider.SearchClient") + async def test_auto_discovers_single_vector_field( + self, mock_search_class: MagicMock, mock_index_class: MagicMock + ) -> None: + """Test that single vector field is auto-discovered.""" + # Setup search client mock + mock_search_client = AsyncMock() + mock_search_class.return_value = mock_search_client + + # Setup index client mock + mock_index_client = AsyncMock() + mock_index = MagicMock() + + # Create mock field with vector_search_dimensions attribute + mock_vector_field = MagicMock() + mock_vector_field.name = "embedding_vector" + mock_vector_field.vector_search_dimensions = 1536 + + mock_index.fields = [mock_vector_field] + mock_index_client.get_index.return_value = mock_index + mock_index_client.close = AsyncMock() + mock_index_class.return_value = mock_index_client + + # Create provider without specifying vector_field_name + provider = AzureAISearchContextProvider( + endpoint="https://test.search.windows.net", + index_name="test-index", + api_key="test-key", + mode="semantic", + ) + + # Trigger auto-discovery + await provider._auto_discover_vector_field() + + # Vector field should be auto-discovered but not used without embedding function + assert provider._auto_discovered_vector_field is True + # Should be cleared since no embedding function + assert provider.vector_field_name is None + + @pytest.mark.asyncio + async def test_vector_detection_accuracy(self) -> None: + """Test that vector field detection logic correctly identifies vector fields.""" + from azure.search.documents.indexes.models import SearchField + + # Create real SearchField objects to test the detection logic + vector_field = SearchField( + name="embedding_vector", type="Collection(Edm.Single)", vector_search_dimensions=1536, searchable=True + ) + + string_field = SearchField(name="content", type="Edm.String", searchable=True) + + number_field = SearchField(name="price", type="Edm.Double", filterable=True) + + # Test detection logic directly + is_vector_1 = vector_field.vector_search_dimensions is not None and vector_field.vector_search_dimensions > 0 + is_vector_2 = string_field.vector_search_dimensions is not None and string_field.vector_search_dimensions > 0 + is_vector_3 = number_field.vector_search_dimensions is not None and number_field.vector_search_dimensions > 0 + + # Only the vector field should be detected + assert is_vector_1 is True + assert is_vector_2 is False + assert is_vector_3 is False + + @pytest.mark.asyncio + @patch("agent_framework_aisearch._search_provider.SearchIndexClient") + @patch("agent_framework_aisearch._search_provider.SearchClient") + async def test_no_false_positives_on_string_fields( + self, mock_search_class: MagicMock, mock_index_class: MagicMock + ) -> None: + """Test that regular string fields are not detected as vector fields.""" + # Setup search client mock + mock_search_client = AsyncMock() + mock_search_class.return_value = mock_search_client + + # Setup index with only string fields (no vectors) + mock_index_client = AsyncMock() + mock_index = MagicMock() + + # All fields have vector_search_dimensions = None + mock_fields = [] + for name in ["id", "title", "content", "category"]: + field = MagicMock() + field.name = name + field.vector_search_dimensions = None + field.vector_search_profile_name = None + mock_fields.append(field) + + mock_index.fields = mock_fields + mock_index_client.get_index.return_value = mock_index + mock_index_client.close = AsyncMock() + mock_index_class.return_value = mock_index_client + + # Create provider + provider = AzureAISearchContextProvider( + endpoint="https://test.search.windows.net", + index_name="test-index", + api_key="test-key", + mode="semantic", + ) + + # Trigger auto-discovery + await provider._auto_discover_vector_field() + + # Should NOT detect any vector fields + assert provider.vector_field_name is None + assert provider._auto_discovered_vector_field is True + + @pytest.mark.asyncio + @patch("agent_framework_aisearch._search_provider.SearchIndexClient") + @patch("agent_framework_aisearch._search_provider.SearchClient") + async def test_multiple_vector_fields_without_vectorizer( + self, mock_search_class: MagicMock, mock_index_class: MagicMock + ) -> None: + """Test that multiple vector fields without vectorizer logs warning and uses keyword search.""" + # Setup search client mock + mock_search_client = AsyncMock() + mock_search_class.return_value = mock_search_client + + # Setup index with multiple vector fields (no vectorizers) + mock_index_client = AsyncMock() + mock_index = MagicMock() + + # Multiple vector fields + mock_fields = [] + for name in ["embedding1", "embedding2"]: + field = MagicMock() + field.name = name + field.vector_search_dimensions = 1536 + field.vector_search_profile_name = None # No vectorizer + mock_fields.append(field) + + mock_index.fields = mock_fields + mock_index.vector_search = None # No vector search config + mock_index_client.get_index.return_value = mock_index + mock_index_client.close = AsyncMock() + mock_index_class.return_value = mock_index_client + + # Create provider + provider = AzureAISearchContextProvider( + endpoint="https://test.search.windows.net", + index_name="test-index", + api_key="test-key", + mode="semantic", + ) + + # Trigger auto-discovery + await provider._auto_discover_vector_field() + + # Should NOT use any vector field (multiple fields, can't choose) + assert provider.vector_field_name is None + assert provider._auto_discovered_vector_field is True + + @pytest.mark.asyncio + @patch("agent_framework_aisearch._search_provider.SearchIndexClient") + @patch("agent_framework_aisearch._search_provider.SearchClient") + async def test_multiple_vectorizable_fields( + self, mock_search_class: MagicMock, mock_index_class: MagicMock + ) -> None: + """Test that multiple vectorizable fields logs warning and uses keyword search.""" + # Setup search client mock + mock_search_client = AsyncMock() + mock_search_class.return_value = mock_search_client + + # Setup index with multiple vectorizable fields + mock_index_client = AsyncMock() + mock_index = MagicMock() + + # Multiple vector fields with vectorizers + mock_fields = [] + for name in ["embedding1", "embedding2"]: + field = MagicMock() + field.name = name + field.vector_search_dimensions = 1536 + field.vector_search_profile_name = f"{name}-profile" + mock_fields.append(field) + + mock_index.fields = mock_fields + + # Setup vector search config with profiles that have vectorizers + mock_profile1 = MagicMock() + mock_profile1.name = "embedding1-profile" + mock_profile1.vectorizer_name = "vectorizer1" + + mock_profile2 = MagicMock() + mock_profile2.name = "embedding2-profile" + mock_profile2.vectorizer_name = "vectorizer2" + + mock_index.vector_search = MagicMock() + mock_index.vector_search.profiles = [mock_profile1, mock_profile2] + + mock_index_client.get_index.return_value = mock_index + mock_index_client.close = AsyncMock() + mock_index_class.return_value = mock_index_client + + # Create provider + provider = AzureAISearchContextProvider( + endpoint="https://test.search.windows.net", + index_name="test-index", + api_key="test-key", + mode="semantic", + ) + + # Trigger auto-discovery + await provider._auto_discover_vector_field() + + # Should NOT use any vector field (multiple vectorizable fields, can't choose) + assert provider.vector_field_name is None + assert provider._auto_discovered_vector_field is True + + @pytest.mark.asyncio + @patch("agent_framework_aisearch._search_provider.SearchIndexClient") + @patch("agent_framework_aisearch._search_provider.SearchClient") + async def test_single_vectorizable_field_detected( + self, mock_search_class: MagicMock, mock_index_class: MagicMock + ) -> None: + """Test that single vectorizable field is auto-detected for server-side vectorization.""" + # Setup search client mock + mock_search_client = AsyncMock() + mock_search_class.return_value = mock_search_client + + # Setup index with single vectorizable field + mock_index_client = AsyncMock() + mock_index = MagicMock() + + # Single vector field with vectorizer + mock_field = MagicMock() + mock_field.name = "embedding" + mock_field.vector_search_dimensions = 1536 + mock_field.vector_search_profile_name = "embedding-profile" + + mock_index.fields = [mock_field] + + # Setup vector search config with profile that has vectorizer + mock_profile = MagicMock() + mock_profile.name = "embedding-profile" + mock_profile.vectorizer_name = "openai-vectorizer" + + mock_index.vector_search = MagicMock() + mock_index.vector_search.profiles = [mock_profile] + + mock_index_client.get_index.return_value = mock_index + mock_index_client.close = AsyncMock() + mock_index_class.return_value = mock_index_client + + # Create provider + provider = AzureAISearchContextProvider( + endpoint="https://test.search.windows.net", + index_name="test-index", + api_key="test-key", + mode="semantic", + ) + + # Trigger auto-discovery + await provider._auto_discover_vector_field() + + # Should detect the vectorizable field + assert provider.vector_field_name == "embedding" + assert provider._auto_discovered_vector_field is True + assert provider._use_vectorizable_query is True # Server-side vectorization diff --git a/python/packages/core/agent_framework/azure/__init__.py b/python/packages/core/agent_framework/azure/__init__.py index c1b64f2117..09670188ee 100644 --- a/python/packages/core/agent_framework/azure/__init__.py +++ b/python/packages/core/agent_framework/azure/__init__.py @@ -10,6 +10,8 @@ "AgentResponseCallbackProtocol": ("agent_framework_azurefunctions", "azurefunctions"), "AzureAIAgentClient": ("agent_framework_azure_ai", "azure-ai"), "AzureAIClient": ("agent_framework_azure_ai", "azure-ai"), + "AzureAISearchContextProvider": ("agent_framework_aisearch", "aisearch"), + "AzureAISearchSettings": ("agent_framework_aisearch", "aisearch"), "AzureOpenAIAssistantsClient": ("agent_framework.azure._assistants_client", "core"), "AzureOpenAIChatClient": ("agent_framework.azure._chat_client", "core"), "AzureAISettings": ("agent_framework_azure_ai", "azure-ai"), diff --git a/python/packages/core/pyproject.toml b/python/packages/core/pyproject.toml index 8d19afa0c5..0559208dfa 100644 --- a/python/packages/core/pyproject.toml +++ b/python/packages/core/pyproject.toml @@ -43,6 +43,7 @@ dependencies = [ all = [ "agent-framework-a2a", "agent-framework-ag-ui", + "agent-framework-aisearch", "agent-framework-anthropic", "agent-framework-azure-ai", "agent-framework-azurefunctions", diff --git a/python/pyproject.toml b/python/pyproject.toml index cbe69d4634..4fcfe56843 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -81,6 +81,7 @@ agent-framework = { workspace = true } agent-framework-core = { workspace = true } agent-framework-a2a = { workspace = true } agent-framework-ag-ui = { workspace = true } +agent-framework-aisearch = { workspace = true } agent-framework-anthropic = { workspace = true } agent-framework-azure-ai = { workspace = true } agent-framework-azurefunctions = { workspace = true } diff --git a/python/samples/getting_started/agents/azure_ai/README.md b/python/samples/getting_started/agents/azure_ai/README.md index 8fd883b0cb..7a3b2d4520 100644 --- a/python/samples/getting_started/agents/azure_ai/README.md +++ b/python/samples/getting_started/agents/azure_ai/README.md @@ -20,6 +20,8 @@ This folder contains examples demonstrating different ways to create and use age | [`azure_ai_with_file_search.py`](azure_ai_with_file_search.py) | Shows how to use the `HostedFileSearchTool` with Azure AI agents to upload files, create vector stores, and enable agents to search through uploaded documents to answer user questions. | | [`azure_ai_with_hosted_mcp.py`](azure_ai_with_hosted_mcp.py) | Shows how to integrate hosted Model Context Protocol (MCP) tools with Azure AI Agent. | | [`azure_ai_with_response_format.py`](azure_ai_with_response_format.py) | Shows how to use structured outputs (response format) with Azure AI agents using Pydantic models to enforce specific response schemas. | +| [`azure_ai_with_search_context_agentic.py`](azure_ai_with_search_context_agentic.py) | Shows how to use AzureAISearchContextProvider with agentic mode. Uses Knowledge Bases for multi-hop reasoning across documents with query planning. Recommended for most scenarios - slightly slower with more token consumption for query planning, but more accurate results. | +| [`azure_ai_with_search_context_semantic.py`](azure_ai_with_search_context_semantic.py) | Shows how to use AzureAISearchContextProvider with semantic mode. Fast hybrid search with vector + keyword search and semantic ranking for RAG. Best for simple queries where speed is critical. | | [`azure_ai_with_sharepoint.py`](azure_ai_with_sharepoint.py) | Shows how to use SharePoint grounding with Azure AI agents to search through SharePoint content and answer user questions with proper citations. Requires a SharePoint connection configured in your Azure AI project. | | [`azure_ai_with_thread.py`](azure_ai_with_thread.py) | Demonstrates thread management with Azure AI agents, including automatic thread creation for stateless conversations and explicit thread management for maintaining conversation context across multiple interactions. | | [`azure_ai_with_image_generation.py`](azure_ai_with_image_generation.py) | Shows how to use the `ImageGenTool` with Azure AI agents to generate images based on text prompts. | diff --git a/python/samples/getting_started/agents/azure_ai/azure_ai_with_search_context_agentic.py b/python/samples/getting_started/agents/azure_ai/azure_ai_with_search_context_agentic.py new file mode 100644 index 0000000000..d8f6a9dc30 --- /dev/null +++ b/python/samples/getting_started/agents/azure_ai/azure_ai_with_search_context_agentic.py @@ -0,0 +1,118 @@ +# Copyright (c) Microsoft. All rights reserved. + +import asyncio +import os + +from dotenv import load_dotenv + +from agent_framework import ChatAgent +from agent_framework_aisearch import AzureAISearchContextProvider +from agent_framework_azure_ai import AzureAIAgentClient +from azure.identity.aio import DefaultAzureCredential + +# Load environment variables from .env file +load_dotenv() + +""" +This sample demonstrates how to use Azure AI Search with agentic mode for RAG +(Retrieval Augmented Generation) with Azure AI agents. + +**Agentic mode** is recommended for most scenarios: +- Uses Knowledge Bases in Azure AI Search for query planning +- Performs multi-hop reasoning across documents +- Provides more accurate results through intelligent retrieval +- Slightly slower with more token consumption for query planning +- See: https://techcommunity.microsoft.com/blog/azure-ai-foundry-blog/foundry-iq-boost-response-relevance-by-36-with-agentic-retrieval/4470720 + +For simple queries where speed is critical, use semantic mode instead (see azure_ai_with_search_context_semantic.py). + +Prerequisites: +1. An Azure AI Search service with a search index +2. An Azure AI Foundry project with a model deployment +3. An Azure OpenAI resource (for Knowledge Base model calls) +4. Set the following environment variables: + - AZURE_SEARCH_ENDPOINT: Your Azure AI Search endpoint + - AZURE_SEARCH_API_KEY: (Optional) Your search API key - if not provided, uses DefaultAzureCredential for Entra ID + - AZURE_SEARCH_INDEX_NAME: Your search index name + - AZURE_AI_PROJECT_ENDPOINT: Your Azure AI Foundry project endpoint + - AZURE_AI_MODEL_DEPLOYMENT_NAME: Your model deployment name (e.g., "gpt-4o") + - AZURE_SEARCH_KNOWLEDGE_BASE_NAME: Your Knowledge Base name + - AZURE_OPENAI_RESOURCE_URL: Your Azure OpenAI resource URL (e.g., "https://myresource.openai.azure.com") + Note: This is different from AZURE_AI_PROJECT_ENDPOINT - Knowledge Base needs the OpenAI endpoint for model calls +""" + +# Sample queries to demonstrate agentic RAG +USER_INPUTS = [ + "What information is available in the knowledge base?", + "Analyze and compare the main topics from different documents", + "What connections can you find across different sections?", +] + + +async def main() -> None: + """Main function demonstrating Azure AI Search agentic mode.""" + + # Get configuration from environment + search_endpoint = os.environ["AZURE_SEARCH_ENDPOINT"] + search_key = os.environ.get("AZURE_SEARCH_API_KEY") + index_name = os.environ["AZURE_SEARCH_INDEX_NAME"] + project_endpoint = os.environ["AZURE_AI_PROJECT_ENDPOINT"] + model_deployment = os.environ.get("AZURE_AI_MODEL_DEPLOYMENT_NAME", "gpt-4o") + knowledge_base_name = os.environ["AZURE_SEARCH_KNOWLEDGE_BASE_NAME"] + azure_openai_resource_url = os.environ["AZURE_OPENAI_RESOURCE_URL"] + + # Create Azure AI Search context provider with agentic mode (recommended for accuracy) + print("Using AGENTIC mode (Knowledge Bases with query planning, recommended)\n") + print("ℹ️ This mode is slightly slower but provides more accurate results.\n") + search_provider = AzureAISearchContextProvider( + endpoint=search_endpoint, + index_name=index_name, + api_key=search_key, # Use api_key for API key auth, or credential for managed identity + credential=DefaultAzureCredential() if not search_key else None, + mode="agentic", # Advanced mode for multi-hop reasoning + # Agentic mode configuration + azure_ai_project_endpoint=project_endpoint, + azure_openai_resource_url=azure_openai_resource_url, + model_deployment_name=model_deployment, + knowledge_base_name=knowledge_base_name, + # Optional: Configure retrieval behavior + knowledge_base_output_mode="extractive_data", # or "answer_synthesis" + retrieval_reasoning_effort="minimal", # or "medium", "low" + top_k=3, # Note: In agentic mode, the server-side Knowledge Base determines final retrieval + ) + + # Create agent with search context provider + async with ( + search_provider, + AzureAIAgentClient( + project_endpoint=project_endpoint, + model_deployment_name=model_deployment, + async_credential=DefaultAzureCredential(), + ) as client, + ChatAgent( + chat_client=client, + name="SearchAgent", + instructions=( + "You are a helpful assistant with advanced reasoning capabilities. " + "Use the provided context from the knowledge base to answer complex " + "questions that may require synthesizing information from multiple sources." + ), + context_providers=[search_provider], + ) as agent, + ): + print("=== Azure AI Agent with Search Context (Agentic Mode) ===\n") + + for user_input in USER_INPUTS: + print(f"User: {user_input}") + print("Agent: ", end="", flush=True) + + # Stream response + async for chunk in agent.run_stream(user_input): + if chunk.text: + print(chunk.text, end="", flush=True) + + print("\n") + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/python/samples/getting_started/agents/azure_ai/azure_ai_with_search_context_semantic.py b/python/samples/getting_started/agents/azure_ai/azure_ai_with_search_context_semantic.py new file mode 100644 index 0000000000..30504354f7 --- /dev/null +++ b/python/samples/getting_started/agents/azure_ai/azure_ai_with_search_context_semantic.py @@ -0,0 +1,99 @@ +# Copyright (c) Microsoft. All rights reserved. + +import asyncio +import os + +from dotenv import load_dotenv + +from agent_framework import ChatAgent +from agent_framework_aisearch import AzureAISearchContextProvider +from agent_framework_azure_ai import AzureAIAgentClient +from azure.identity.aio import DefaultAzureCredential + +# Load environment variables from .env file +load_dotenv() + +""" +This sample demonstrates how to use Azure AI Search with semantic mode for RAG +(Retrieval Augmented Generation) with Azure AI agents. + +**Semantic mode** is the recommended default mode: +- Fast hybrid search combining vector and keyword search +- Uses semantic ranking for improved relevance +- Returns raw search results as context +- Best for most RAG use cases + +Prerequisites: +1. An Azure AI Search service with a search index +2. An Azure AI Foundry project with a model deployment +3. Set the following environment variables: + - AZURE_SEARCH_ENDPOINT: Your Azure AI Search endpoint + - AZURE_SEARCH_API_KEY: (Optional) Your search API key - if not provided, uses DefaultAzureCredential for Entra ID + - AZURE_SEARCH_INDEX_NAME: Your search index name + - AZURE_AI_PROJECT_ENDPOINT: Your Azure AI Foundry project endpoint + - AZURE_AI_MODEL_DEPLOYMENT_NAME: Your model deployment name (e.g., "gpt-4o") +""" + +# Sample queries to demonstrate RAG +USER_INPUTS = [ + "What information is available in the knowledge base?", + "Summarize the main topics from the documents", + "Find specific details about the content", +] + + +async def main() -> None: + """Main function demonstrating Azure AI Search semantic mode.""" + + # Get configuration from environment + search_endpoint = os.environ["AZURE_SEARCH_ENDPOINT"] + search_key = os.environ.get("AZURE_SEARCH_API_KEY") + index_name = os.environ["AZURE_SEARCH_INDEX_NAME"] + project_endpoint = os.environ["AZURE_AI_PROJECT_ENDPOINT"] + model_deployment = os.environ.get("AZURE_AI_MODEL_DEPLOYMENT_NAME", "gpt-4o") + + # Create Azure AI Search context provider with semantic mode (recommended, fast) + print("Using SEMANTIC mode (hybrid search + semantic ranking, fast)\n") + search_provider = AzureAISearchContextProvider( + endpoint=search_endpoint, + index_name=index_name, + api_key=search_key, # Use api_key for API key auth, or credential for managed identity + credential=DefaultAzureCredential() if not search_key else None, + mode="semantic", # Default mode + top_k=3, # Retrieve top 3 most relevant documents + ) + + # Create agent with search context provider + async with ( + search_provider, + AzureAIAgentClient( + project_endpoint=project_endpoint, + model_deployment_name=model_deployment, + async_credential=DefaultAzureCredential(), + ) as client, + ChatAgent( + chat_client=client, + name="SearchAgent", + instructions=( + "You are a helpful assistant. Use the provided context from the " + "knowledge base to answer questions accurately." + ), + context_providers=[search_provider], + ) as agent, + ): + print("=== Azure AI Agent with Search Context (Semantic Mode) ===\n") + + for user_input in USER_INPUTS: + print(f"User: {user_input}") + print("Agent: ", end="", flush=True) + + # Stream response + async for chunk in agent.run_stream(user_input): + if chunk.text: + print(chunk.text, end="", flush=True) + + print("\n") + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/python/uv.lock b/python/uv.lock index cbbd28046a..eda7decab3 100644 --- a/python/uv.lock +++ b/python/uv.lock @@ -26,6 +26,7 @@ members = [ "agent-framework", "agent-framework-a2a", "agent-framework-ag-ui", + "agent-framework-aisearch", "agent-framework-anthropic", "agent-framework-azure-ai", "agent-framework-azurefunctions", @@ -193,6 +194,21 @@ requires-dist = [ ] provides-extras = ["dev"] +[[package]] +name = "agent-framework-aisearch" +version = "1.0.0b251118" +source = { editable = "packages/aisearch" } +dependencies = [ + { name = "agent-framework-core", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, + { name = "azure-search-documents", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, +] + +[package.metadata] +requires-dist = [ + { name = "agent-framework-core", editable = "packages/core" }, + { name = "azure-search-documents", specifier = "==11.7.0b2" }, +] + [[package]] name = "agent-framework-anthropic" version = "1.0.0b251114" @@ -304,6 +320,7 @@ dependencies = [ all = [ { name = "agent-framework-a2a", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, { name = "agent-framework-ag-ui", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, + { name = "agent-framework-aisearch", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, { name = "agent-framework-anthropic", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, { name = "agent-framework-azure-ai", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, { name = "agent-framework-azurefunctions", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, @@ -321,6 +338,7 @@ all = [ requires-dist = [ { name = "agent-framework-a2a", marker = "extra == 'all'", editable = "packages/a2a" }, { name = "agent-framework-ag-ui", marker = "extra == 'all'", editable = "packages/ag-ui" }, + { name = "agent-framework-aisearch", marker = "extra == 'all'", editable = "packages/aisearch" }, { name = "agent-framework-anthropic", marker = "extra == 'all'", editable = "packages/anthropic" }, { name = "agent-framework-azure-ai", marker = "extra == 'all'", editable = "packages/azure-ai" }, { name = "agent-framework-azurefunctions", marker = "extra == 'all'", editable = "packages/azurefunctions" }, @@ -928,6 +946,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/a9/41/d9a2b3eb33b4ffd9acfaa115cfd456e32d0c754227d6d78ec5d039ff75c2/azure_ai_projects-2.0.0b2-py3-none-any.whl", hash = "sha256:642496fdf9846c91f3557d39899d3893f0ce8f910334320686fc8f617492351d", size = 234023, upload-time = "2025-11-15T06:17:48.141Z" }, ] +[[package]] +name = "azure-common" +version = "1.1.28" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/3e/71/f6f71a276e2e69264a97ad39ef850dca0a04fce67b12570730cb38d0ccac/azure-common-1.1.28.zip", hash = "sha256:4ac0cd3214e36b6a1b6a442686722a5d8cc449603aa833f3f0f40bda836704a3", size = 20914, upload-time = "2022-02-03T19:39:44.373Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/62/55/7f118b9c1b23ec15ca05d15a578d8207aa1706bc6f7c87218efffbbf875d/azure_common-1.1.28-py2.py3-none-any.whl", hash = "sha256:5c12d3dcf4ec20599ca6b0d3e09e86e146353d443e7fcc050c9a19c1f9df20ad", size = 14462, upload-time = "2022-02-03T19:39:42.417Z" }, +] + [[package]] name = "azure-core" version = "1.36.0" @@ -987,6 +1014,21 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/83/7b/5652771e24fff12da9dde4c20ecf4682e606b104f26419d139758cc935a6/azure_identity-1.25.1-py3-none-any.whl", hash = "sha256:e9edd720af03dff020223cd269fa3a61e8f345ea75443858273bcb44844ab651", size = 191317, upload-time = "2025-10-06T20:30:04.251Z" }, ] +[[package]] +name = "azure-search-documents" +version = "11.7.0b2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "azure-common", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, + { name = "azure-core", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, + { name = "isodate", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, + { name = "typing-extensions", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/f9/ba/bde0f03e0a742ba3bbcc929f91ed2f3b1420c2bb84c9a7f878f3b87ebfce/azure_search_documents-11.7.0b2.tar.gz", hash = "sha256:b6e039f8038ff2210d2057e704e867c6e29bb46bfcd400da4383e45e4b8bb189", size = 423956, upload-time = "2025-11-14T20:09:32.876Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e5/26/ed4498374f9088818278ac225f2bea688b4ec979d81bf83a5355c8c366af/azure_search_documents-11.7.0b2-py3-none-any.whl", hash = "sha256:f82117b321344a84474269ed26df194c24cca619adc024d981b1b86aee3c6f05", size = 432037, upload-time = "2025-11-14T20:09:34.347Z" }, +] + [[package]] name = "azure-storage-blob" version = "12.27.1"