From ed8af313797e6cb8aae26bf997dfd9f92396ee66 Mon Sep 17 00:00:00 2001 From: Farzad Sunavala Date: Tue, 28 Oct 2025 23:05:30 -0500 Subject: [PATCH 01/13] Python: Fix pyright errors and move search provider to core (#1546) --- python/.env.example | 8 + .../core/agent_framework/azure/__init__.py | 1 + .../agent_framework/azure/_search_provider.py | 538 ++++++++++++++++++ python/packages/core/pyproject.toml | 1 + .../core/tests/azure/test_search_provider.py | 338 +++++++++++ .../azure_ai/azure_ai_with_search_context.py | 137 +++++ python/uv.lock | 26 + 7 files changed, 1049 insertions(+) create mode 100644 python/packages/core/agent_framework/azure/_search_provider.py create mode 100644 python/packages/core/tests/azure/test_search_provider.py create mode 100644 python/samples/getting_started/agents/azure_ai/azure_ai_with_search_context.py 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/core/agent_framework/azure/__init__.py b/python/packages/core/agent_framework/azure/__init__.py index 5dfab603cb..39dbf11855 100644 --- a/python/packages/core/agent_framework/azure/__init__.py +++ b/python/packages/core/agent_framework/azure/__init__.py @@ -6,6 +6,7 @@ _IMPORTS: dict[str, tuple[str, str]] = { "AzureAIAgentClient": ("agent_framework_azure_ai", "azure-ai"), + "AzureAISearchContextProvider": ("agent_framework.azure._search_provider", "core"), "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/agent_framework/azure/_search_provider.py b/python/packages/core/agent_framework/azure/_search_provider.py new file mode 100644 index 0000000000..2a95592d26 --- /dev/null +++ b/python/packages/core/agent_framework/azure/_search_provider.py @@ -0,0 +1,538 @@ +# 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: +- Semantic: Fast hybrid search (vector + keyword) with semantic ranker +- Agentic: Slower multi-hop reasoning using Knowledge Bases for complex queries + +Use semantic mode for most cases. Use agentic mode only when you need multi-hop +reasoning across documents with Knowledge Bases. +""" + +import sys +from collections.abc import MutableSequence +from typing import TYPE_CHECKING, Any, Literal, cast + +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, + KnowledgeAgent, + KnowledgeAgentAzureOpenAIModel, + KnowledgeAgentOutputConfiguration, + KnowledgeAgentOutputConfigurationModality, + KnowledgeAgentRequestLimits, + KnowledgeSourceReference, + SearchIndexKnowledgeSource, + SearchIndexKnowledgeSourceParameters, +) +from azure.search.documents.models import ( + QueryCaptionType, + QueryType, + VectorizedQuery, +) + +from agent_framework import ChatMessage, Context, ContextProvider + +# Type checking imports for optional agentic mode dependencies +if TYPE_CHECKING: + from azure.search.documents.agent.aio import KnowledgeAgentRetrievalClient + from azure.search.documents.agent.models import ( + KnowledgeAgentMessage, + KnowledgeAgentMessageTextContent, + KnowledgeAgentRetrievalRequest, + ) + +# Runtime imports for agentic mode (optional dependency) +try: + from azure.search.documents.agent.aio import KnowledgeAgentRetrievalClient + from azure.search.documents.agent.models import ( + KnowledgeAgentMessage, + KnowledgeAgentMessageTextContent, + KnowledgeAgentRetrievalRequest, + ) + + _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 + + +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: + + - **semantic** (default): Fast hybrid search combining vector and keyword search + with semantic reranking. Suitable for most RAG use cases. + - **agentic**: Slower multi-hop reasoning across documents using Knowledge Bases. + Use only for complex queries requiring cross-document reasoning. + + Examples: + Semantic hybrid search (recommended for most cases): + + .. code-block:: python + + from agent_framework import ChatAgent + from agent_framework_azure_ai import AzureAIAgentClient, AzureAISearchContextProvider + from azure.identity.aio import DefaultAzureCredential + + # Create context provider with semantic hybrid search + search_provider = AzureAISearchContextProvider( + endpoint="https://mysearch.search.windows.net", + index_name="my-index", + credential=DefaultAzureCredential(), + mode="semantic", # Fast hybrid + semantic ranker (default) + ) + + # Use with agent + async with ( + AzureAIAgentClient() as client, + ChatAgent( + chat_client=client, + context_providers=[search_provider], + ) as agent, + ): + response = await agent.run("What is in the documents?") + + Agentic retrieval for complex queries: + + .. code-block:: python + + # Use agentic mode for multi-hop reasoning (slower) + search_provider = AzureAISearchContextProvider( + endpoint="https://mysearch.search.windows.net", + index_name="my-index", + credential=DefaultAzureCredential(), + mode="agentic", # Multi-hop reasoning + azure_ai_project_endpoint="https://myproject.services.ai.azure.com", + model_deployment_name="gpt-4o", + knowledge_base_name="my-knowledge-base", # Required for agentic mode + ) + """ + + def __init__( + self, + endpoint: str, + index_name: str, + credential: AzureKeyCredential | AsyncTokenCredential, + mode: Literal["semantic", "agentic"] = "semantic", + top_k: int = 5, + semantic_configuration_name: str | None = None, + vector_field_name: str | None = None, + embedding_function: Any | None = None, + context_prompt: str | None = None, + # Agentic mode parameters (Knowledge Base) + azure_ai_project_endpoint: 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, + azure_openai_resource_url: str | None = None, + # Deprecated parameters (for backwards compatibility) + azure_openai_endpoint: str | None = None, + azure_openai_deployment_name: str | None = None, + azure_openai_api_version: str | None = None, + ) -> None: + """Initialize Azure AI Search Context Provider. + + Args: + endpoint: Azure AI Search endpoint URL. + index_name: Name of the search index to query. + credential: Azure credential (API key or DefaultAzureCredential). + 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. 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. + context_prompt: Custom prompt to prepend to retrieved context. + Default: Uses DEFAULT_CONTEXT_PROMPT. + azure_ai_project_endpoint: Azure AI Foundry project endpoint URL. + Required for agentic mode. Example: "https://myproject.services.ai.azure.com" + model_deployment_name: Model deployment name in the Azure AI project. + 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. + azure_openai_resource_url: Azure OpenAI resource URL for Knowledge Base model calls. + Required for agentic mode. Example: "https://myresource.openai.azure.com" + This is different from azure_ai_project_endpoint (which is Foundry-specific). + azure_openai_endpoint: (Deprecated) Use azure_ai_project_endpoint instead. + azure_openai_deployment_name: (Deprecated) Use model_deployment_name instead. + azure_openai_api_version: (Deprecated) No longer used. + """ + self.endpoint = endpoint + self.index_name = index_name + self.credential = 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_CONTEXT_PROMPT + + # Agentic mode parameters (Knowledge Base) + # azure_openai_resource_url: The actual Azure OpenAI endpoint for model calls + # azure_openai_endpoint (deprecated): Fall back to this if resource_url not provided + self.azure_openai_resource_url = azure_openai_resource_url or azure_openai_endpoint + + self.azure_openai_deployment_name = model_deployment_name or azure_openai_deployment_name + # If model_name not provided, default to deployment name for backwards compatibility + self.model_name = model_name or self.azure_openai_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 + + # 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 (or deprecated azure_openai_endpoint) 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 (or deprecated azure_openai_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=endpoint, + index_name=index_name, + credential=credential, + ) + + # Create index client for agentic mode (Knowledge Base) + # Note: Retrieval client is created fresh for each query to avoid transport issues + self._index_client: SearchIndexClient | None = None + if mode == "agentic": + self._index_client = SearchIndexClient( + endpoint=endpoint, + credential=credential, + ) + + 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 handled by client destructors. + + 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. + + Note: + We don't explicitly close the Azure SDK clients here because doing so + can cause "transport already closed" errors on subsequent uses within + the same context. The clients will clean up their resources when they + are garbage collected. + """ + pass + + @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. + """ + # Extract query from messages + messages_list = [messages] if isinstance(messages, ChatMessage) else list(messages) + query = "\n".join(msg.text for msg in messages_list if msg and msg.text and msg.text.strip()) + + if not query: + return Context() + + # Perform search based on mode + if self.mode == "semantic": + search_results = await self._semantic_search(query) + else: # agentic + search_results = await self._agentic_search(query) + + # Format results as context + if not search_results: + return Context() + + context_text = f"{self.context_prompt}\n\n{search_results}" + + return Context(messages=[ChatMessage(role="system", text=context_text)]) + + async def _semantic_search(self, query: str) -> 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: + Formatted search results as string. + """ + vector_queries = [] + + # Generate vector query if embedding function provided + if self.embedding_function and self.vector_field_name: + query_vector = await self.embedding_function(query) + vector_queries = [ + VectorizedQuery( + vector=query_vector, + k_nearest_neighbors=self.top_k, + fields=self.vector_field_name, + ) + ] + + # 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 + formatted_results: list[str] = [] + async for doc in results: # type: ignore[reportUnknownVariableType] + # Extract semantic captions if available + caption: str | None = None + if hasattr(doc, "@search.captions"): # type: ignore[reportUnknownArgumentType] + captions: Any = doc.get("@search.captions", []) # type: ignore[reportUnknownVariableType] + if captions: + caption = captions[0].text if hasattr(captions[0], "text") else str(captions[0]) # type: ignore[reportUnknownArgumentType, reportUnknownMemberType] + + # Build document text + doc_text: str = caption if caption else self._extract_document_text(doc) # type: ignore[reportUnknownArgumentType] + if doc_text: + formatted_results.append(doc_text) # type: ignore[reportUnknownArgumentType] + + return "\n\n".join(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 + + # Type narrowing: these are validated as non-None in __init__ for agentic mode + # Using cast() for type checker - actual validation happens in __init__ + knowledge_base_name = cast(str, self.knowledge_base_name) + azure_openai_resource_url = cast(str, self.azure_openai_resource_url) + azure_openai_deployment_name = cast(str, self.azure_openai_deployment_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 get Knowledge Base (using KnowledgeAgent SDK class) + try: + # Try to get existing Knowledge Base + await self._index_client.get_agent(knowledge_base_name) + except ResourceNotFoundError: + # Create new Knowledge Base if it doesn't exist + aoai_params = AzureOpenAIVectorizerParameters( + resource_url=azure_openai_resource_url, + deployment_name=azure_openai_deployment_name, + model_name=self.model_name, # Underlying model name (e.g., "gpt-4o") + api_key=self.azure_openai_api_key, # Optional: for API key auth instead of managed identity + ) + + # Note: SDK uses KnowledgeAgent class name, but this represents a Knowledge Base + knowledge_base = KnowledgeAgent( + name=knowledge_base_name, + description=f"Knowledge Base for multi-hop retrieval across {self.index_name}", + models=[KnowledgeAgentAzureOpenAIModel(azure_open_ai_parameters=aoai_params)], + knowledge_sources=[ + KnowledgeSourceReference( + name=knowledge_source_name, + include_references=True, + include_reference_source_data=True, + ) + ], + output_configuration=KnowledgeAgentOutputConfiguration( + modality=KnowledgeAgentOutputConfigurationModality.ANSWER_SYNTHESIS, + attempt_fast_path=True, + ), + request_limits=KnowledgeAgentRequestLimits( + max_output_size=10000, + max_runtime_in_seconds=60, + ), + retrieval_instructions=self.retrieval_instructions, + ) + await self._index_client.create_agent(knowledge_base) + + self._knowledge_base_initialized = True + + async def _agentic_search(self, query: str) -> str: + """Perform agentic retrieval with multi-hop reasoning using Knowledge Bases. + + NOTE: This mode is significantly slower than semantic search and should + only be used for complex queries requiring cross-document reasoning. + + 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: + query: Search query text. + + Returns: + Synthesized answer from the Knowledge Base. + """ + # Ensure Knowledge Base is initialized + await self._ensure_knowledge_base() + + # Type narrowing: knowledge_base_name is validated in __init__ for agentic mode + # Using cast() for type checker - actual validation happens in __init__ + knowledge_base_name = cast(str, self.knowledge_base_name) + + # Create retrieval request with query as a conversation message + # Note: SDK uses KnowledgeAgent class names, but represents Knowledge Base operations + retrieval_request = KnowledgeAgentRetrievalRequest( + messages=[ + KnowledgeAgentMessage( + role="user", + content=[KnowledgeAgentMessageTextContent(text=query)], + ) + ] + ) + + # Create a fresh retrieval client for each query to avoid transport closure issues + if not _agentic_retrieval_available: + raise ImportError("KnowledgeAgentRetrievalClient not available") + + retrieval_client = KnowledgeAgentRetrievalClient( + endpoint=self.endpoint, + agent_name=knowledge_base_name, + credential=self.credential, + ) + + try: + # Perform retrieval via Knowledge Base + retrieval_result = await retrieval_client.retrieve(retrieval_request=retrieval_request) + finally: + # Ensure client is closed after use + await retrieval_client.close() + + # Extract synthesized answer 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: + # Combine all text content + answer_parts: list[str] = [] + for content_item in assistant_message.content: + # Check if this is a text content item + if isinstance(content_item, KnowledgeAgentMessageTextContent) and content_item.text: + answer_parts.append(content_item.text) + + if answer_parts: + return "\n".join(answer_parts) + + # Fallback if no answer generated + return "No results found from Knowledge Base." + + def _extract_document_text(self, doc: dict[str, Any]) -> str: + """Extract readable text from a search document. + + Args: + doc: Search result document. + + Returns: + Formatted document text. + """ + # Try common text field names + for field in ["content", "text", "description", "body", "chunk"]: + if doc.get(field): + return str(doc[field])[:500] # Limit to 500 chars + + # Fallback: concatenate all string fields + 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}") + + return " | ".join(text_parts)[:500] if text_parts else "" diff --git a/python/packages/core/pyproject.toml b/python/packages/core/pyproject.toml index f5b2c23e7d..80355fa6c0 100644 --- a/python/packages/core/pyproject.toml +++ b/python/packages/core/pyproject.toml @@ -35,6 +35,7 @@ dependencies = [ "openai>=1.99.0,<2", "azure-identity>=1,<2", "mcp[ws]>=1.13", + "azure-search-documents>=11.7.0b1", ] [project.optional-dependencies] diff --git a/python/packages/core/tests/azure/test_search_provider.py b/python/packages/core/tests/azure/test_search_provider.py new file mode 100644 index 0000000000..149ad12524 --- /dev/null +++ b/python/packages/core/tests/azure/test_search_provider.py @@ -0,0 +1,338 @@ +# Copyright (c) Microsoft. All rights reserved. +# pyright: reportPrivateUsage=false + +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from azure.core.credentials import AzureKeyCredential +from azure.core.exceptions import ResourceNotFoundError + +from agent_framework import ChatMessage, Context, Role +from agent_framework.azure import AzureAISearchContextProvider + + +@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 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", + credential=AzureKeyCredential("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", + credential=AzureKeyCredential("test-key"), + mode="semantic", + vector_field_name="embedding", + ) + + def test_init_agentic_mode_requires_parameters(self) -> None: + """Test that agentic mode requires additional parameters.""" + with pytest.raises(ValueError, match="azure_openai_resource_url"): + AzureAISearchContextProvider( + endpoint="https://test.search.windows.net", + index_name="test-index", + credential=AzureKeyCredential("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", + credential=AzureKeyCredential("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", + credential=AzureKeyCredential("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", + credential=AzureKeyCredential("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", + credential=AzureKeyCredential("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", + credential=AzureKeyCredential("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", + credential=AzureKeyCredential("test-key"), + mode="semantic", + ) + assert provider.context_prompt == provider.DEFAULT_CONTEXT_PROMPT + + +class TestSemanticSearch: + """Test semantic search functionality.""" + + @pytest.mark.asyncio + @patch("agent_framework.azure._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", + credential=AzureKeyCredential("test-key"), + mode="semantic", + ) + + context = await provider.invoking(sample_messages) + + assert isinstance(context, Context) + assert len(context.messages) > 0 + assert "Test document content" in context.messages[0].text + + @pytest.mark.asyncio + @patch("agent_framework.azure._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", + credential=AzureKeyCredential("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.azure._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", + credential=AzureKeyCredential("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.azure._search_provider.SearchIndexClient") + @patch("agent_framework.azure._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_agent.side_effect = ResourceNotFoundError("Not found") + mock_index_client.create_agent = 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", + credential=AzureKeyCredential("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 agent (Knowledge Base) was created + mock_index_client.create_agent.assert_called_once() + + @pytest.mark.asyncio + @patch("agent_framework.azure._search_provider.SearchIndexClient") + @patch("agent_framework.azure._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_agent.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", + credential=AzureKeyCredential("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.azure._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", + credential=AzureKeyCredential("test-key"), + mode="semantic", + ) as provider: + assert provider is not None + assert isinstance(provider, AzureAISearchContextProvider) diff --git a/python/samples/getting_started/agents/azure_ai/azure_ai_with_search_context.py b/python/samples/getting_started/agents/azure_ai/azure_ai_with_search_context.py new file mode 100644 index 0000000000..2f7f308944 --- /dev/null +++ b/python/samples/getting_started/agents/azure_ai/azure_ai_with_search_context.py @@ -0,0 +1,137 @@ +# Copyright (c) Microsoft. All rights reserved. + +import asyncio +import os + +from agent_framework import ChatAgent +from agent_framework.azure import AzureAISearchContextProvider +from agent_framework_azure_ai import AzureAIAgentClient +from azure.core.credentials import AzureKeyCredential +from azure.identity.aio import DefaultAzureCredential + +""" +The following sample demonstrates how to use Azure AI Search as a context provider +for RAG (Retrieval Augmented Generation) with Azure AI agents. + +AzureAISearchContextProvider supports two modes: + +1. **Semantic mode** (default, recommended): + - 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 + +2. **Agentic mode** (slower, advanced): + - Uses Knowledge Bases in Azure AI Search + - Performs multi-hop reasoning across documents + - Uses an LLM to synthesize answers + - Best for complex queries requiring cross-document reasoning + - Significantly slower (order of magnitude) + +Prerequisites: +1. An Azure AI Search service with a search index +2. An Azure AI Foundry project with a model deployment +3. Ensure the model deployment name exists in your Azure AI Foundry project +4. Set the following environment variables: + + For both modes: + - AZURE_SEARCH_ENDPOINT: Your Azure AI Search endpoint + - AZURE_SEARCH_API_KEY: Your search API key (or use Azure AD) + - 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") + + Additional for agentic mode (Knowledge Bases): + - USE_AGENTIC_MODE: Set to "true" to use agentic retrieval + - AZURE_SEARCH_KNOWLEDGE_BASE_NAME: Your Knowledge Base name + - AZURE_OPENAI_ENDPOINT: Your Azure OpenAI endpoint (e.g., "https://myresource.openai.azure.com") + (This is different from AZURE_AI_PROJECT_ENDPOINT - Knowledge Base needs the OpenAI endpoint for model calls) +""" + +# 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 context provider.""" + + # 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") + + # Check if agentic mode is requested + use_agentic = os.environ.get("USE_AGENTIC_MODE", "false").lower() == "true" + + # Create credential + search_credential = AzureKeyCredential(search_key) if search_key else DefaultAzureCredential() + + # Create Azure AI Search context provider + if use_agentic: + # Agentic mode: Multi-hop reasoning with Knowledge Bases (slower) + print("Using AGENTIC mode (Knowledge Bases with multi-hop reasoning, slower)\n") + knowledge_base_name = os.environ["AZURE_SEARCH_KNOWLEDGE_BASE_NAME"] + azure_openai_endpoint = os.environ["AZURE_OPENAI_ENDPOINT"] + search_provider = AzureAISearchContextProvider( + endpoint=search_endpoint, + index_name=index_name, + credential=search_credential, + mode="agentic", + # Agentic mode uses Azure AI Foundry project for model inference + azure_ai_project_endpoint=project_endpoint, + model_deployment_name=model_deployment, + knowledge_base_name=knowledge_base_name, + azure_openai_resource_url=azure_openai_endpoint, + top_k=3, + ) + else: + # Semantic mode: Fast hybrid search + semantic ranking (recommended) + print("Using SEMANTIC mode (hybrid search + semantic ranking, fast)\n") + search_provider = AzureAISearchContextProvider( + endpoint=search_endpoint, + index_name=index_name, + credential=search_credential, + mode="semantic", + 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 ===\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 0e11d3bd11..90e3404f57 100644 --- a/python/uv.lock +++ b/python/uv.lock @@ -208,6 +208,7 @@ version = "1.0.0b251028" source = { editable = "packages/core" } dependencies = [ { name = "azure-identity", 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'" }, { name = "mcp", extra = ["ws"], marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, { name = "openai", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, { name = "opentelemetry-api", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, @@ -238,6 +239,7 @@ requires-dist = [ { name = "agent-framework-mem0", marker = "extra == 'all'", editable = "packages/mem0" }, { name = "agent-framework-redis", marker = "extra == 'all'", editable = "packages/redis" }, { name = "azure-identity", specifier = ">=1,<2" }, + { name = "azure-search-documents", specifier = ">=11.7.0b1" }, { name = "mcp", extras = ["ws"], specifier = ">=1.13" }, { name = "openai", specifier = ">=1.99.0,<2" }, { name = "opentelemetry-api", specifier = ">=1.24" }, @@ -780,6 +782,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/db/10/8b7bd070e3cc804343dab124ce66a3b7999a72d5be0e49232cbcd1d36e18/azure_ai_projects-1.1.0b4-py3-none-any.whl", hash = "sha256:d8aab84fd7cd7c5937e78141e37ca4473dc5ed6cce2c0490c634418abe14afea", size = 126670, upload-time = "2025-09-12T17:35:10.039Z" }, ] +[[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" @@ -809,6 +820,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.0b1" +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/bc/9a/819cb9375e431eac8f337777526a88e5f8e10cbf2efd877d39c470f49aa2/azure_search_documents-11.7.0b1.tar.gz", hash = "sha256:0324bd6732dd79c2bc4b6f6429d5a1b96129077d4706620af6f74af4b952fe0c", size = 395057, upload-time = "2025-09-04T23:30:32.571Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fb/14/8e9bc312fffc21e5220c8ce87f064009e7f577fef69bade2a9713032b646/azure_search_documents-11.7.0b1-py3-none-any.whl", hash = "sha256:aae83b501246438eff77aa3930de9fdd3532c0d219018342fd9d85c682814376", size = 413982, upload-time = "2025-09-04T23:30:34.252Z" }, +] + [[package]] name = "azure-storage-blob" version = "12.27.1" From 8de2673645bb8c21ea4839361c77d5e1074a53ba Mon Sep 17 00:00:00 2001 From: Farzad Sunavala Date: Sun, 2 Nov 2025 18:04:38 +0000 Subject: [PATCH 02/13] address pablo coments --- .../agent_framework/azure/_search_provider.py | 235 ++++++++++++------ .../core/tests/azure/test_search_provider.py | 206 +++++++++++++++ .../getting_started/agents/azure_ai/README.md | 2 +- .../azure_ai/azure_ai_with_search_context.py | 2 +- 4 files changed, 373 insertions(+), 72 deletions(-) diff --git a/python/packages/core/agent_framework/azure/_search_provider.py b/python/packages/core/agent_framework/azure/_search_provider.py index 2a95592d26..267002b6fa 100644 --- a/python/packages/core/agent_framework/azure/_search_provider.py +++ b/python/packages/core/agent_framework/azure/_search_provider.py @@ -12,7 +12,7 @@ import sys from collections.abc import MutableSequence -from typing import TYPE_CHECKING, Any, Literal, cast +from typing import TYPE_CHECKING, Any, Literal from azure.core.credentials import AzureKeyCredential from azure.core.credentials_async import AsyncTokenCredential @@ -125,11 +125,14 @@ class AzureAISearchContextProvider(ContextProvider): ) """ + DEFAULT_CONTEXT_PROMPT = "Use the following context to answer the question:" + def __init__( self, endpoint: str, index_name: str, credential: AzureKeyCredential | AsyncTokenCredential, + *, mode: Literal["semantic", "agentic"] = "semantic", top_k: int = 5, semantic_configuration_name: str | None = None, @@ -208,6 +211,13 @@ def __init__( self.azure_openai_api_key = azure_openai_api_key self.azure_ai_project_endpoint = azure_ai_project_endpoint + # Auto-discover vector field if not specified + self._auto_discovered_vector_field = False + 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") @@ -237,14 +247,15 @@ def __init__( credential=credential, ) - # Create index client for agentic mode (Knowledge Base) - # Note: Retrieval client is created fresh for each query to avoid transport issues + # Create index client and retrieval client for agentic mode (Knowledge Base) self._index_client: SearchIndexClient | None = None + self._retrieval_client: KnowledgeAgentRetrievalClient | None = None if mode == "agentic": self._index_client = SearchIndexClient( endpoint=endpoint, credential=credential, ) + # Retrieval client will be created after Knowledge Base initialization self._knowledge_base_initialized = False @@ -258,20 +269,17 @@ async def __aexit__( exc_val: BaseException | None, exc_tb: Any, ) -> None: - """Async context manager exit - cleanup handled by client destructors. + """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. - - Note: - We don't explicitly close the Azure SDK clients here because doing so - can cause "transport already closed" errors on subsequent uses within - the same context. The clients will clean up their resources when they - are garbage collected. """ - pass + # 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( @@ -288,18 +296,28 @@ async def invoking( Returns: Context object with retrieved documents as messages. """ - # Extract query from messages + # Convert to list and filter to USER/ASSISTANT messages with text only messages_list = [messages] if isinstance(messages, ChatMessage) else list(messages) - query = "\n".join(msg.text for msg in messages_list if msg and msg.text and msg.text.strip()) + from agent_framework import Role - if not query: + 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_results = await self._semantic_search(query) else: # agentic - search_results = await self._agentic_search(query) + # Agentic mode: pass last 10 messages as conversation history + recent_messages = filtered_messages[-10:] + search_results = await self._agentic_search(recent_messages) # Format results as context if not search_results: @@ -309,6 +327,72 @@ async def invoking( return Context(messages=[ChatMessage(role="system", text=context_text)]) + async def _auto_discover_vector_field(self) -> None: + """Auto-discover vector field from index schema. + + Attempts to find vector fields in the index and use the vectorizer configuration. + If successful, sets self.vector_field_name. If embedding function is not provided, + logs a warning that vector search is recommended for RAG. + """ + if self._auto_discovered_vector_field or self.vector_field_name: + return # Already discovered or manually specified + + try: + # Need index client to get schema + if not self._index_client: + from azure.search.documents.indexes.aio import SearchIndexClient + + index_client = SearchIndexClient(endpoint=self.endpoint, credential=self.credential) + else: + index_client = self._index_client + + # Get index schema + index = await index_client.get_index(self.index_name) + + # Find vector fields - must have vector_search_dimensions set (not None) + vector_fields = [ + field + for field in index.fields + if field.vector_search_dimensions is not None and field.vector_search_dimensions > 0 + ] + + if len(vector_fields) == 1: + # Exactly one vector field found - auto-select it + self.vector_field_name = vector_fields[0].name + self._auto_discovered_vector_field = True + + # Warn if no embedding function provided + if not self.embedding_function: + import logging + + logging.warning( + f"Auto-discovered vector field '{self.vector_field_name}' but no embedding_function provided. " + "Vector search is recommended for RAG use cases. Falling back to keyword-only search." + ) + self.vector_field_name = None # Clear it since we can't use it + elif len(vector_fields) > 1: + # Multiple vector fields - warn and continue with keyword search + import logging + + logging.warning( + f"Multiple vector fields found in index '{self.index_name}': " + f"{[f.name for f in vector_fields]}. " + "Please specify vector_field_name explicitly. Using keyword-only search." + ) + # If no vector fields found, silently continue with keyword search + + # Close index client if we created it + if not self._index_client: + await index_client.close() + + except Exception as e: + # Log warning but continue with keyword search + import logging + + logging.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) -> str: """Perform semantic hybrid search with semantic ranking. @@ -323,15 +407,20 @@ async def _semantic_search(self, query: str) -> str: Returns: Formatted search results as string. """ + # Auto-discover vector field if not already done + await self._auto_discover_vector_field() + vector_queries = [] # Generate vector query if embedding function provided if self.embedding_function and self.vector_field_name: query_vector = await self.embedding_function(query) + # 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 vector_queries = [ VectorizedQuery( vector=query_vector, - k_nearest_neighbors=self.top_k, + k_nearest_neighbors=vector_k, fields=self.vector_field_name, ) ] @@ -354,18 +443,14 @@ async def _semantic_search(self, query: str) -> str: # Execute search results = await self._search_client.search(**search_params) # type: ignore[reportUnknownVariableType] - # Format results + # Format results with citations formatted_results: list[str] = [] async for doc in results: # type: ignore[reportUnknownVariableType] - # Extract semantic captions if available - caption: str | None = None - if hasattr(doc, "@search.captions"): # type: ignore[reportUnknownArgumentType] - captions: Any = doc.get("@search.captions", []) # type: ignore[reportUnknownVariableType] - if captions: - caption = captions[0].text if hasattr(captions[0], "text") else str(captions[0]) # type: ignore[reportUnknownArgumentType, reportUnknownMemberType] - - # Build document text - doc_text: str = caption if caption else self._extract_document_text(doc) # type: ignore[reportUnknownArgumentType] + # 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] @@ -382,11 +467,17 @@ async def _ensure_knowledge_base(self) -> None: if self._knowledge_base_initialized or not self._index_client: return - # Type narrowing: these are validated as non-None in __init__ for agentic mode - # Using cast() for type checker - actual validation happens in __init__ - knowledge_base_name = cast(str, self.knowledge_base_name) - azure_openai_resource_url = cast(str, self.azure_openai_resource_url) - azure_openai_deployment_name = cast(str, self.azure_openai_deployment_name) + # 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 + azure_openai_resource_url = self.azure_openai_resource_url + azure_openai_deployment_name = self.azure_openai_deployment_name # Step 1: Create or get knowledge source knowledge_source_name = f"{self.index_name}-source" @@ -444,7 +535,15 @@ async def _ensure_knowledge_base(self) -> None: self._knowledge_base_initialized = True - async def _agentic_search(self, query: str) -> str: + # Create retrieval client now that Knowledge Base is initialized + if _agentic_retrieval_available and self._retrieval_client is None: + self._retrieval_client = KnowledgeAgentRetrievalClient( + endpoint=self.endpoint, + agent_name=knowledge_base_name, + credential=self.credential, + ) + + async def _agentic_search(self, messages: list[ChatMessage]) -> str: """Perform agentic retrieval with multi-hop reasoning using Knowledge Bases. NOTE: This mode is significantly slower than semantic search and should @@ -457,7 +556,7 @@ async def _agentic_search(self, query: str) -> str: 4. Synthesize a comprehensive answer with references Args: - query: Search query text. + messages: Conversation history (last 10 messages) to use for retrieval context. Returns: Synthesized answer from the Knowledge Base. @@ -465,37 +564,25 @@ async def _agentic_search(self, query: str) -> str: # Ensure Knowledge Base is initialized await self._ensure_knowledge_base() - # Type narrowing: knowledge_base_name is validated in __init__ for agentic mode - # Using cast() for type checker - actual validation happens in __init__ - knowledge_base_name = cast(str, self.knowledge_base_name) - - # Create retrieval request with query as a conversation message + # Convert ChatMessage list to KnowledgeAgent message format # Note: SDK uses KnowledgeAgent class names, but represents Knowledge Base operations - retrieval_request = KnowledgeAgentRetrievalRequest( - messages=[ - KnowledgeAgentMessage( - role="user", - content=[KnowledgeAgentMessageTextContent(text=query)], - ) - ] - ) + kb_messages = [ + KnowledgeAgentMessage( + role=msg.role.value if hasattr(msg.role, "value") else str(msg.role), + content=[KnowledgeAgentMessageTextContent(text=msg.text)], + ) + for msg in messages + if msg.text + ] - # Create a fresh retrieval client for each query to avoid transport closure issues - if not _agentic_retrieval_available: - raise ImportError("KnowledgeAgentRetrievalClient not available") + retrieval_request = KnowledgeAgentRetrievalRequest(messages=kb_messages) - retrieval_client = KnowledgeAgentRetrievalClient( - endpoint=self.endpoint, - agent_name=knowledge_base_name, - credential=self.credential, - ) + # Use reusable retrieval client + if not self._retrieval_client: + raise RuntimeError("Retrieval client not initialized. Ensure Knowledge Base is set up correctly.") - try: - # Perform retrieval via Knowledge Base - retrieval_result = await retrieval_client.retrieve(retrieval_request=retrieval_request) - finally: - # Ensure client is closed after use - await retrieval_client.close() + # Perform retrieval via Knowledge Base + retrieval_result = await self._retrieval_client.retrieve(retrieval_request=retrieval_request) # Extract synthesized answer from response if retrieval_result.response and len(retrieval_result.response) > 0: @@ -515,24 +602,32 @@ async def _agentic_search(self, query: str) -> str: # Fallback if no answer generated return "No results found from Knowledge Base." - def _extract_document_text(self, doc: dict[str, Any]) -> str: - """Extract readable text from a search document. + 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. + 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): - return str(doc[field])[:500] # Limit to 500 chars + text = str(doc[field]) + break # Fallback: concatenate all string fields - 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}") - - return " | ".join(text_parts)[:500] if text_parts else "" + 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/core/tests/azure/test_search_provider.py b/python/packages/core/tests/azure/test_search_provider.py index 149ad12524..29e86f9352 100644 --- a/python/packages/core/tests/azure/test_search_provider.py +++ b/python/packages/core/tests/azure/test_search_provider.py @@ -336,3 +336,209 @@ async def test_context_manager(self, mock_search_class: MagicMock) -> None: ) as provider: assert provider is not None assert isinstance(provider, AzureAISearchContextProvider) + + +class TestMessageFiltering: + """Test message filtering functionality.""" + + @pytest.mark.asyncio + @patch("agent_framework.azure._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", + credential=AzureKeyCredential("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.azure._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", + credential=AzureKeyCredential("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.azure._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", + credential=AzureKeyCredential("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) > 0 + assert "[Source: doc123]" in context.messages[0].text + assert "Test document content" in context.messages[0].text + + +class TestVectorFieldAutoDiscovery: + """Test vector field auto-discovery functionality.""" + + @pytest.mark.asyncio + @patch("agent_framework.azure._search_provider.SearchIndexClient") + @patch("agent_framework.azure._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", + credential=AzureKeyCredential("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.azure._search_provider.SearchIndexClient") + @patch("agent_framework.azure._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", + credential=AzureKeyCredential("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 diff --git a/python/samples/getting_started/agents/azure_ai/README.md b/python/samples/getting_started/agents/azure_ai/README.md index 77faa26192..a548f74c96 100644 --- a/python/samples/getting_started/agents/azure_ai/README.md +++ b/python/samples/getting_started/agents/azure_ai/README.md @@ -12,7 +12,7 @@ This folder contains examples demonstrating different ways to create and use age | [`azure_ai_with_existing_agent.py`](azure_ai_with_existing_agent.py) | Shows how to work with a pre-existing agent by providing the agent ID to the Azure AI chat client. This example also demonstrates proper cleanup of manually created agents. | | [`azure_ai_with_existing_thread.py`](azure_ai_with_existing_thread.py) | Shows how to work with a pre-existing thread by providing the thread ID to the Azure AI chat client. This example also demonstrates proper cleanup of manually created threads. | | [`azure_ai_with_explicit_settings.py`](azure_ai_with_explicit_settings.py) | Shows how to create an agent with explicitly configured `AzureAIAgentClient` settings, including project endpoint, model deployment, credentials, and agent name. | -| [`azure_ai_with_azure_ai_search.py`](azure_ai_with_azure_ai_search.py) | Demonstrates how to use Azure AI Search with Azure AI agents to search through indexed data. Shows how to configure search parameters, query types, and integrate with existing search indexes. | +| [`azure_ai_with_search_context.py`](azure_ai_with_search_context.py) | Shows how to use AzureAISearchContextProvider for RAG with hybrid search and semantic ranking. Demonstrates semantic (fast) and agentic (slower, uses Knowledge Bases for complex multi-hop reasoning) retrieval modes. Provides context from Azure AI Search to enhance agent responses. | | [`azure_ai_with_file_search.py`](azure_ai_with_file_search.py) | Demonstrates how to use the HostedFileSearchTool with Azure AI agents to search through uploaded documents. Shows file upload, vector store creation, and querying document content. Includes both streaming and non-streaming examples. | | [`azure_ai_with_function_tools.py`](azure_ai_with_function_tools.py) | Demonstrates how to use function tools with agents. Shows both agent-level tools (defined when creating the agent) and query-level tools (provided with specific queries). | | [`azure_ai_with_hosted_mcp.py`](azure_ai_with_hosted_mcp.py) | Shows how to integrate Azure AI agents with hosted Model Context Protocol (MCP) servers for enhanced functionality and tool integration. Demonstrates remote MCP server connections and tool discovery. | diff --git a/python/samples/getting_started/agents/azure_ai/azure_ai_with_search_context.py b/python/samples/getting_started/agents/azure_ai/azure_ai_with_search_context.py index 2f7f308944..c9f05dba2b 100644 --- a/python/samples/getting_started/agents/azure_ai/azure_ai_with_search_context.py +++ b/python/samples/getting_started/agents/azure_ai/azure_ai_with_search_context.py @@ -36,7 +36,7 @@ For both modes: - AZURE_SEARCH_ENDPOINT: Your Azure AI Search endpoint - - AZURE_SEARCH_API_KEY: Your search API key (or use Azure AD) + - AZURE_SEARCH_API_KEY: (Optional) Your search API key - if not provided, uses DefaultAzureCredential for Entra ID authentication - 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") From 0cd48ab5ba62937826329332fed350bc40ebaa50 Mon Sep 17 00:00:00 2001 From: farzad528 Date: Mon, 17 Nov 2025 16:31:02 +0000 Subject: [PATCH 03/13] update azure ai search pypi version to latest prev --- .../agent_framework/azure/_search_provider.py | 259 +++++++++++------- python/packages/core/pyproject.toml | 2 +- .../core/tests/azure/test_search_provider.py | 10 +- python/uv.lock | 18 +- 4 files changed, 186 insertions(+), 103 deletions(-) diff --git a/python/packages/core/agent_framework/azure/_search_provider.py b/python/packages/core/agent_framework/azure/_search_provider.py index 267002b6fa..8852e65839 100644 --- a/python/packages/core/agent_framework/azure/_search_provider.py +++ b/python/packages/core/agent_framework/azure/_search_provider.py @@ -21,11 +21,9 @@ from azure.search.documents.indexes.aio import SearchIndexClient from azure.search.documents.indexes.models import ( AzureOpenAIVectorizerParameters, - KnowledgeAgent, - KnowledgeAgentAzureOpenAIModel, - KnowledgeAgentOutputConfiguration, - KnowledgeAgentOutputConfigurationModality, - KnowledgeAgentRequestLimits, + KnowledgeBase, + KnowledgeBaseAzureOpenAIModel, + KnowledgeRetrievalOutputMode, KnowledgeSourceReference, SearchIndexKnowledgeSource, SearchIndexKnowledgeSourceParameters, @@ -33,6 +31,7 @@ from azure.search.documents.models import ( QueryCaptionType, QueryType, + VectorizableTextQuery, VectorizedQuery, ) @@ -40,20 +39,20 @@ # Type checking imports for optional agentic mode dependencies if TYPE_CHECKING: - from azure.search.documents.agent.aio import KnowledgeAgentRetrievalClient - from azure.search.documents.agent.models import ( - KnowledgeAgentMessage, - KnowledgeAgentMessageTextContent, - KnowledgeAgentRetrievalRequest, + from azure.search.documents.knowledgebases.aio import KnowledgeBaseRetrievalClient + from azure.search.documents.knowledgebases.models import ( + KnowledgeBaseMessage, + KnowledgeBaseMessageTextContent, + KnowledgeBaseRetrievalRequest, ) # Runtime imports for agentic mode (optional dependency) try: - from azure.search.documents.agent.aio import KnowledgeAgentRetrievalClient - from azure.search.documents.agent.models import ( - KnowledgeAgentMessage, - KnowledgeAgentMessageTextContent, - KnowledgeAgentRetrievalRequest, + from azure.search.documents.knowledgebases.aio import KnowledgeBaseRetrievalClient + from azure.search.documents.knowledgebases.models import ( + KnowledgeBaseMessage, + KnowledgeBaseMessageTextContent, + KnowledgeBaseRetrievalRequest, ) _agentic_retrieval_available = True @@ -213,6 +212,7 @@ def __init__( # 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 @@ -249,7 +249,7 @@ def __init__( # Create index client and retrieval client for agentic mode (Knowledge Base) self._index_client: SearchIndexClient | None = None - self._retrieval_client: KnowledgeAgentRetrievalClient | None = None + self._retrieval_client: KnowledgeBaseRetrievalClient | None = None if mode == "agentic": self._index_client = SearchIndexClient( endpoint=endpoint, @@ -327,12 +327,60 @@ async def invoking( return Context(messages=[ChatMessage(role="system", text=context_text)]) + 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 = [] + + # 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 use the vectorizer configuration. - If successful, sets self.vector_field_name. If embedding function is not provided, - logs a warning that vector search is recommended for RAG. + 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 @@ -349,37 +397,65 @@ async def _auto_discover_vector_field(self) -> None: # Get index schema index = await index_client.get_index(self.index_name) - # Find vector fields - must have vector_search_dimensions set (not None) - vector_fields = [ - field - for field in index.fields - if field.vector_search_dimensions is not None and field.vector_search_dimensions > 0 - ] + # Step 1: Find all vector fields + vector_fields = self._find_vector_fields(index) + + if not vector_fields: + # No vector fields found - keyword search only + import logging + + logging.info(f"No vector fields found in index '{self.index_name}'. Using keyword-only search.") + self._auto_discovered_vector_field = True + if not self._index_client: + await index_client.close() + 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 + import logging + + logging.info( + f"Auto-discovered vectorizable field '{self.vector_field_name}' " + f"with server-side vectorization. No embedding_function needed." + ) + else: + # Multiple vectorizable fields + import logging - if len(vector_fields) == 1: - # Exactly one vector field found - auto-select it - self.vector_field_name = vector_fields[0].name + logging.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 - # Warn if no embedding function provided if not self.embedding_function: import logging logging.warning( - f"Auto-discovered vector field '{self.vector_field_name}' but no embedding_function provided. " - "Vector search is recommended for RAG use cases. Falling back to keyword-only search." + 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 # Clear it since we can't use it - elif len(vector_fields) > 1: - # Multiple vector fields - warn and continue with keyword search + self.vector_field_name = None + else: + # Multiple vector fields without vectorizers import logging logging.warning( - f"Multiple vector fields found in index '{self.index_name}': " - f"{[f.name for f in vector_fields]}. " - "Please specify vector_field_name explicitly. Using keyword-only search." + f"Multiple vector fields found: {vector_fields}. " + f"Please specify vector_field_name explicitly. Using keyword-only search." ) - # If no vector fields found, silently continue with keyword search # Close index client if we created it if not self._index_client: @@ -412,18 +488,31 @@ async def _semantic_search(self, query: str) -> str: vector_queries = [] - # Generate vector query if embedding function provided - if self.embedding_function and self.vector_field_name: - query_vector = await self.embedding_function(query) + # 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 - vector_queries = [ - VectorizedQuery( - vector=query_vector, - k_nearest_neighbors=vector_k, - fields=self.vector_field_name, - ) - ] + + 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] = { @@ -476,8 +565,6 @@ async def _ensure_knowledge_base(self) -> None: raise ValueError("model_deployment_name is required for agentic mode") knowledge_base_name = self.knowledge_base_name - azure_openai_resource_url = self.azure_openai_resource_url - azure_openai_deployment_name = self.azure_openai_deployment_name # Step 1: Create or get knowledge source knowledge_source_name = f"{self.index_name}-source" @@ -496,50 +583,37 @@ async def _ensure_knowledge_base(self) -> None: ) await self._index_client.create_knowledge_source(knowledge_source) - # Step 2: Create or get Knowledge Base (using KnowledgeAgent SDK class) - try: - # Try to get existing Knowledge Base - await self._index_client.get_agent(knowledge_base_name) - except ResourceNotFoundError: - # Create new Knowledge Base if it doesn't exist - aoai_params = AzureOpenAIVectorizerParameters( - resource_url=azure_openai_resource_url, - deployment_name=azure_openai_deployment_name, - model_name=self.model_name, # Underlying model name (e.g., "gpt-4o") - api_key=self.azure_openai_api_key, # Optional: for API key auth instead of managed identity - ) + # Step 2: Create or update Knowledge Base + # Always create/update to ensure configuration is current + # Note: EXTRACTIVE_DATA mode returns raw chunks without synthesis + # Model is still needed for query planning and multi-hop reasoning + 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, + ) - # Note: SDK uses KnowledgeAgent class name, but this represents a Knowledge Base - knowledge_base = KnowledgeAgent( - name=knowledge_base_name, - description=f"Knowledge Base for multi-hop retrieval across {self.index_name}", - models=[KnowledgeAgentAzureOpenAIModel(azure_open_ai_parameters=aoai_params)], - knowledge_sources=[ - KnowledgeSourceReference( - name=knowledge_source_name, - include_references=True, - include_reference_source_data=True, - ) - ], - output_configuration=KnowledgeAgentOutputConfiguration( - modality=KnowledgeAgentOutputConfigurationModality.ANSWER_SYNTHESIS, - attempt_fast_path=True, - ), - request_limits=KnowledgeAgentRequestLimits( - max_output_size=10000, - max_runtime_in_seconds=60, - ), - retrieval_instructions=self.retrieval_instructions, - ) - await self._index_client.create_agent(knowledge_base) + 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=KnowledgeRetrievalOutputMode.EXTRACTIVE_DATA, + ) + 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 = KnowledgeAgentRetrievalClient( + self._retrieval_client = KnowledgeBaseRetrievalClient( endpoint=self.endpoint, - agent_name=knowledge_base_name, + knowledge_base_name=knowledge_base_name, credential=self.credential, ) @@ -564,18 +638,17 @@ async def _agentic_search(self, messages: list[ChatMessage]) -> str: # Ensure Knowledge Base is initialized await self._ensure_knowledge_base() - # Convert ChatMessage list to KnowledgeAgent message format - # Note: SDK uses KnowledgeAgent class names, but represents Knowledge Base operations + # Convert ChatMessage list to KnowledgeBase message format kb_messages = [ - KnowledgeAgentMessage( + KnowledgeBaseMessage( role=msg.role.value if hasattr(msg.role, "value") else str(msg.role), - content=[KnowledgeAgentMessageTextContent(text=msg.text)], + content=[KnowledgeBaseMessageTextContent(text=msg.text)], ) for msg in messages if msg.text ] - retrieval_request = KnowledgeAgentRetrievalRequest(messages=kb_messages) + retrieval_request = KnowledgeBaseRetrievalRequest(messages=kb_messages) # Use reusable retrieval client if not self._retrieval_client: @@ -593,7 +666,7 @@ async def _agentic_search(self, messages: list[ChatMessage]) -> str: answer_parts: list[str] = [] for content_item in assistant_message.content: # Check if this is a text content item - if isinstance(content_item, KnowledgeAgentMessageTextContent) and content_item.text: + if isinstance(content_item, KnowledgeBaseMessageTextContent) and content_item.text: answer_parts.append(content_item.text) if answer_parts: diff --git a/python/packages/core/pyproject.toml b/python/packages/core/pyproject.toml index 80355fa6c0..dc6de209e6 100644 --- a/python/packages/core/pyproject.toml +++ b/python/packages/core/pyproject.toml @@ -35,7 +35,7 @@ dependencies = [ "openai>=1.99.0,<2", "azure-identity>=1,<2", "mcp[ws]>=1.13", - "azure-search-documents>=11.7.0b1", + "azure-search-documents==11.7.0b2", ] [project.optional-dependencies] diff --git a/python/packages/core/tests/azure/test_search_provider.py b/python/packages/core/tests/azure/test_search_provider.py index 29e86f9352..2db0968ac2 100644 --- a/python/packages/core/tests/azure/test_search_provider.py +++ b/python/packages/core/tests/azure/test_search_provider.py @@ -258,8 +258,8 @@ async def test_ensure_knowledge_base_creates_when_not_exists( 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_agent.side_effect = ResourceNotFoundError("Not found") - mock_index_client.create_agent = 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() @@ -281,8 +281,8 @@ async def test_ensure_knowledge_base_creates_when_not_exists( # Verify knowledge source was created mock_index_client.create_knowledge_source.assert_called_once() - # Verify agent (Knowledge Base) was created - mock_index_client.create_agent.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.azure._search_provider.SearchIndexClient") @@ -294,7 +294,7 @@ async def test_ensure_knowledge_base_skips_when_exists( # Setup mocks mock_index_client = AsyncMock() mock_index_client.get_knowledge_source.return_value = MagicMock() # Exists - mock_index_client.get_agent.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() diff --git a/python/uv.lock b/python/uv.lock index 90e3404f57..33db600462 100644 --- a/python/uv.lock +++ b/python/uv.lock @@ -239,7 +239,7 @@ requires-dist = [ { name = "agent-framework-mem0", marker = "extra == 'all'", editable = "packages/mem0" }, { name = "agent-framework-redis", marker = "extra == 'all'", editable = "packages/redis" }, { name = "azure-identity", specifier = ">=1,<2" }, - { name = "azure-search-documents", specifier = ">=11.7.0b1" }, + { name = "azure-search-documents", specifier = "==11.7.0b2" }, { name = "mcp", extras = ["ws"], specifier = ">=1.13" }, { name = "openai", specifier = ">=1.99.0,<2" }, { name = "opentelemetry-api", specifier = ">=1.24" }, @@ -822,7 +822,7 @@ wheels = [ [[package]] name = "azure-search-documents" -version = "11.7.0b1" +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'" }, @@ -830,9 +830,9 @@ dependencies = [ { 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/bc/9a/819cb9375e431eac8f337777526a88e5f8e10cbf2efd877d39c470f49aa2/azure_search_documents-11.7.0b1.tar.gz", hash = "sha256:0324bd6732dd79c2bc4b6f6429d5a1b96129077d4706620af6f74af4b952fe0c", size = 395057, upload-time = "2025-09-04T23:30:32.571Z" } +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/fb/14/8e9bc312fffc21e5220c8ce87f064009e7f577fef69bade2a9713032b646/azure_search_documents-11.7.0b1-py3-none-any.whl", hash = "sha256:aae83b501246438eff77aa3930de9fdd3532c0d219018342fd9d85c682814376", size = 413982, upload-time = "2025-09-04T23:30:34.252Z" }, + { 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]] @@ -2024,6 +2024,8 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/7f/91/ae2eb6b7979e2f9b035a9f612cf70f1bf54aad4e1d125129bef1eae96f19/greenlet-3.2.4-cp310-cp310-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:c2ca18a03a8cfb5b25bc1cbe20f3d9a4c80d8c3b13ba3df49ac3961af0b1018d", size = 584358, upload-time = "2025-08-07T13:18:23.708Z" }, { url = "https://files.pythonhosted.org/packages/f7/85/433de0c9c0252b22b16d413c9407e6cb3b41df7389afc366ca204dbc1393/greenlet-3.2.4-cp310-cp310-musllinux_1_1_aarch64.whl", hash = "sha256:9fe0a28a7b952a21e2c062cd5756d34354117796c6d9215a87f55e38d15402c5", size = 1113550, upload-time = "2025-08-07T13:42:37.467Z" }, { url = "https://files.pythonhosted.org/packages/a1/8d/88f3ebd2bc96bf7747093696f4335a0a8a4c5acfcf1b757717c0d2474ba3/greenlet-3.2.4-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:8854167e06950ca75b898b104b63cc646573aa5fef1353d4508ecdd1ee76254f", size = 1137126, upload-time = "2025-08-07T13:18:20.239Z" }, + { url = "https://files.pythonhosted.org/packages/f1/29/74242b7d72385e29bcc5563fba67dad94943d7cd03552bac320d597f29b2/greenlet-3.2.4-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:f47617f698838ba98f4ff4189aef02e7343952df3a615f847bb575c3feb177a7", size = 1544904, upload-time = "2025-11-04T12:42:04.763Z" }, + { url = "https://files.pythonhosted.org/packages/c8/e2/1572b8eeab0f77df5f6729d6ab6b141e4a84ee8eb9bc8c1e7918f94eda6d/greenlet-3.2.4-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:af41be48a4f60429d5cad9d22175217805098a9ef7c40bfef44f7669fb9d74d8", size = 1611228, upload-time = "2025-11-04T12:42:08.423Z" }, { url = "https://files.pythonhosted.org/packages/d6/6f/b60b0291d9623c496638c582297ead61f43c4b72eef5e9c926ef4565ec13/greenlet-3.2.4-cp310-cp310-win_amd64.whl", hash = "sha256:73f49b5368b5359d04e18d15828eecc1806033db5233397748f4ca813ff1056c", size = 298654, upload-time = "2025-08-07T13:50:00.469Z" }, { url = "https://files.pythonhosted.org/packages/a4/de/f28ced0a67749cac23fecb02b694f6473f47686dff6afaa211d186e2ef9c/greenlet-3.2.4-cp311-cp311-macosx_11_0_universal2.whl", hash = "sha256:96378df1de302bc38e99c3a9aa311967b7dc80ced1dcc6f171e99842987882a2", size = 272305, upload-time = "2025-08-07T13:15:41.288Z" }, { url = "https://files.pythonhosted.org/packages/09/16/2c3792cba130000bf2a31c5272999113f4764fd9d874fb257ff588ac779a/greenlet-3.2.4-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:1ee8fae0519a337f2329cb78bd7a8e128ec0f881073d43f023c7b8d4831d5246", size = 632472, upload-time = "2025-08-07T13:42:55.044Z" }, @@ -2033,6 +2035,8 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/1f/8e/abdd3f14d735b2929290a018ecf133c901be4874b858dd1c604b9319f064/greenlet-3.2.4-cp311-cp311-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:2523e5246274f54fdadbce8494458a2ebdcdbc7b802318466ac5606d3cded1f8", size = 587684, upload-time = "2025-08-07T13:18:25.164Z" }, { url = "https://files.pythonhosted.org/packages/5d/65/deb2a69c3e5996439b0176f6651e0052542bb6c8f8ec2e3fba97c9768805/greenlet-3.2.4-cp311-cp311-musllinux_1_1_aarch64.whl", hash = "sha256:1987de92fec508535687fb807a5cea1560f6196285a4cde35c100b8cd632cc52", size = 1116647, upload-time = "2025-08-07T13:42:38.655Z" }, { url = "https://files.pythonhosted.org/packages/3f/cc/b07000438a29ac5cfb2194bfc128151d52f333cee74dd7dfe3fb733fc16c/greenlet-3.2.4-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:55e9c5affaa6775e2c6b67659f3a71684de4c549b3dd9afca3bc773533d284fa", size = 1142073, upload-time = "2025-08-07T13:18:21.737Z" }, + { url = "https://files.pythonhosted.org/packages/67/24/28a5b2fa42d12b3d7e5614145f0bd89714c34c08be6aabe39c14dd52db34/greenlet-3.2.4-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:c9c6de1940a7d828635fbd254d69db79e54619f165ee7ce32fda763a9cb6a58c", size = 1548385, upload-time = "2025-11-04T12:42:11.067Z" }, + { url = "https://files.pythonhosted.org/packages/6a/05/03f2f0bdd0b0ff9a4f7b99333d57b53a7709c27723ec8123056b084e69cd/greenlet-3.2.4-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:03c5136e7be905045160b1b9fdca93dd6727b180feeafda6818e6496434ed8c5", size = 1613329, upload-time = "2025-11-04T12:42:12.928Z" }, { url = "https://files.pythonhosted.org/packages/d8/0f/30aef242fcab550b0b3520b8e3561156857c94288f0332a79928c31a52cf/greenlet-3.2.4-cp311-cp311-win_amd64.whl", hash = "sha256:9c40adce87eaa9ddb593ccb0fa6a07caf34015a29bf8d344811665b573138db9", size = 299100, upload-time = "2025-08-07T13:44:12.287Z" }, { url = "https://files.pythonhosted.org/packages/44/69/9b804adb5fd0671f367781560eb5eb586c4d495277c93bde4307b9e28068/greenlet-3.2.4-cp312-cp312-macosx_11_0_universal2.whl", hash = "sha256:3b67ca49f54cede0186854a008109d6ee71f66bd57bb36abd6d0a0267b540cdd", size = 274079, upload-time = "2025-08-07T13:15:45.033Z" }, { url = "https://files.pythonhosted.org/packages/46/e9/d2a80c99f19a153eff70bc451ab78615583b8dac0754cfb942223d2c1a0d/greenlet-3.2.4-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:ddf9164e7a5b08e9d22511526865780a576f19ddd00d62f8a665949327fde8bb", size = 640997, upload-time = "2025-08-07T13:42:56.234Z" }, @@ -2042,6 +2046,8 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/19/0d/6660d55f7373b2ff8152401a83e02084956da23ae58cddbfb0b330978fe9/greenlet-3.2.4-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:3b3812d8d0c9579967815af437d96623f45c0f2ae5f04e366de62a12d83a8fb0", size = 607586, upload-time = "2025-08-07T13:18:28.544Z" }, { url = "https://files.pythonhosted.org/packages/8e/1a/c953fdedd22d81ee4629afbb38d2f9d71e37d23caace44775a3a969147d4/greenlet-3.2.4-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:abbf57b5a870d30c4675928c37278493044d7c14378350b3aa5d484fa65575f0", size = 1123281, upload-time = "2025-08-07T13:42:39.858Z" }, { url = "https://files.pythonhosted.org/packages/3f/c7/12381b18e21aef2c6bd3a636da1088b888b97b7a0362fac2e4de92405f97/greenlet-3.2.4-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:20fb936b4652b6e307b8f347665e2c615540d4b42b3b4c8a321d8286da7e520f", size = 1151142, upload-time = "2025-08-07T13:18:22.981Z" }, + { url = "https://files.pythonhosted.org/packages/27/45/80935968b53cfd3f33cf99ea5f08227f2646e044568c9b1555b58ffd61c2/greenlet-3.2.4-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:ee7a6ec486883397d70eec05059353b8e83eca9168b9f3f9a361971e77e0bcd0", size = 1564846, upload-time = "2025-11-04T12:42:15.191Z" }, + { url = "https://files.pythonhosted.org/packages/69/02/b7c30e5e04752cb4db6202a3858b149c0710e5453b71a3b2aec5d78a1aab/greenlet-3.2.4-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:326d234cbf337c9c3def0676412eb7040a35a768efc92504b947b3e9cfc7543d", size = 1633814, upload-time = "2025-11-04T12:42:17.175Z" }, { url = "https://files.pythonhosted.org/packages/e9/08/b0814846b79399e585f974bbeebf5580fbe59e258ea7be64d9dfb253c84f/greenlet-3.2.4-cp312-cp312-win_amd64.whl", hash = "sha256:a7d4e128405eea3814a12cc2605e0e6aedb4035bf32697f72deca74de4105e02", size = 299899, upload-time = "2025-08-07T13:38:53.448Z" }, { url = "https://files.pythonhosted.org/packages/49/e8/58c7f85958bda41dafea50497cbd59738c5c43dbbea5ee83d651234398f4/greenlet-3.2.4-cp313-cp313-macosx_11_0_universal2.whl", hash = "sha256:1a921e542453fe531144e91e1feedf12e07351b1cf6c9e8a3325ea600a715a31", size = 272814, upload-time = "2025-08-07T13:15:50.011Z" }, { url = "https://files.pythonhosted.org/packages/62/dd/b9f59862e9e257a16e4e610480cfffd29e3fae018a68c2332090b53aac3d/greenlet-3.2.4-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:cd3c8e693bff0fff6ba55f140bf390fa92c994083f838fece0f63be121334945", size = 641073, upload-time = "2025-08-07T13:42:57.23Z" }, @@ -2051,6 +2057,8 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ee/43/3cecdc0349359e1a527cbf2e3e28e5f8f06d3343aaf82ca13437a9aa290f/greenlet-3.2.4-cp313-cp313-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:23768528f2911bcd7e475210822ffb5254ed10d71f4028387e5a99b4c6699671", size = 610497, upload-time = "2025-08-07T13:18:31.636Z" }, { url = "https://files.pythonhosted.org/packages/b8/19/06b6cf5d604e2c382a6f31cafafd6f33d5dea706f4db7bdab184bad2b21d/greenlet-3.2.4-cp313-cp313-musllinux_1_1_aarch64.whl", hash = "sha256:00fadb3fedccc447f517ee0d3fd8fe49eae949e1cd0f6a611818f4f6fb7dc83b", size = 1121662, upload-time = "2025-08-07T13:42:41.117Z" }, { url = "https://files.pythonhosted.org/packages/a2/15/0d5e4e1a66fab130d98168fe984c509249c833c1a3c16806b90f253ce7b9/greenlet-3.2.4-cp313-cp313-musllinux_1_1_x86_64.whl", hash = "sha256:d25c5091190f2dc0eaa3f950252122edbbadbb682aa7b1ef2f8af0f8c0afefae", size = 1149210, upload-time = "2025-08-07T13:18:24.072Z" }, + { url = "https://files.pythonhosted.org/packages/1c/53/f9c440463b3057485b8594d7a638bed53ba531165ef0ca0e6c364b5cc807/greenlet-3.2.4-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:6e343822feb58ac4d0a1211bd9399de2b3a04963ddeec21530fc426cc121f19b", size = 1564759, upload-time = "2025-11-04T12:42:19.395Z" }, + { url = "https://files.pythonhosted.org/packages/47/e4/3bb4240abdd0a8d23f4f88adec746a3099f0d86bfedb623f063b2e3b4df0/greenlet-3.2.4-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:ca7f6f1f2649b89ce02f6f229d7c19f680a6238af656f61e0115b24857917929", size = 1634288, upload-time = "2025-11-04T12:42:21.174Z" }, { url = "https://files.pythonhosted.org/packages/0b/55/2321e43595e6801e105fcfdee02b34c0f996eb71e6ddffca6b10b7e1d771/greenlet-3.2.4-cp313-cp313-win_amd64.whl", hash = "sha256:554b03b6e73aaabec3745364d6239e9e012d64c68ccd0b8430c64ccc14939a8b", size = 299685, upload-time = "2025-08-07T13:24:38.824Z" }, { url = "https://files.pythonhosted.org/packages/22/5c/85273fd7cc388285632b0498dbbab97596e04b154933dfe0f3e68156c68c/greenlet-3.2.4-cp314-cp314-macosx_11_0_universal2.whl", hash = "sha256:49a30d5fda2507ae77be16479bdb62a660fa51b1eb4928b524975b3bde77b3c0", size = 273586, upload-time = "2025-08-07T13:16:08.004Z" }, { url = "https://files.pythonhosted.org/packages/d1/75/10aeeaa3da9332c2e761e4c50d4c3556c21113ee3f0afa2cf5769946f7a3/greenlet-3.2.4-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:299fd615cd8fc86267b47597123e3f43ad79c9d8a22bebdce535e53550763e2f", size = 686346, upload-time = "2025-08-07T13:42:59.944Z" }, @@ -2058,6 +2066,8 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/dc/8b/29aae55436521f1d6f8ff4e12fb676f3400de7fcf27fccd1d4d17fd8fecd/greenlet-3.2.4-cp314-cp314-manylinux2014_s390x.manylinux_2_17_s390x.whl", hash = "sha256:b4a1870c51720687af7fa3e7cda6d08d801dae660f75a76f3845b642b4da6ee1", size = 694659, upload-time = "2025-08-07T13:53:17.759Z" }, { url = "https://files.pythonhosted.org/packages/92/2e/ea25914b1ebfde93b6fc4ff46d6864564fba59024e928bdc7de475affc25/greenlet-3.2.4-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:061dc4cf2c34852b052a8620d40f36324554bc192be474b9e9770e8c042fd735", size = 695355, upload-time = "2025-08-07T13:18:34.517Z" }, { url = "https://files.pythonhosted.org/packages/72/60/fc56c62046ec17f6b0d3060564562c64c862948c9d4bc8aa807cf5bd74f4/greenlet-3.2.4-cp314-cp314-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:44358b9bf66c8576a9f57a590d5f5d6e72fa4228b763d0e43fee6d3b06d3a337", size = 657512, upload-time = "2025-08-07T13:18:33.969Z" }, + { url = "https://files.pythonhosted.org/packages/23/6e/74407aed965a4ab6ddd93a7ded3180b730d281c77b765788419484cdfeef/greenlet-3.2.4-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:2917bdf657f5859fbf3386b12d68ede4cf1f04c90c3a6bc1f013dd68a22e2269", size = 1612508, upload-time = "2025-11-04T12:42:23.427Z" }, + { url = "https://files.pythonhosted.org/packages/0d/da/343cd760ab2f92bac1845ca07ee3faea9fe52bee65f7bcb19f16ad7de08b/greenlet-3.2.4-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:015d48959d4add5d6c9f6c5210ee3803a830dce46356e3bc326d6776bde54681", size = 1680760, upload-time = "2025-11-04T12:42:25.341Z" }, { url = "https://files.pythonhosted.org/packages/e3/a5/6ddab2b4c112be95601c13428db1d8b6608a8b6039816f2ba09c346c08fc/greenlet-3.2.4-cp314-cp314-win_amd64.whl", hash = "sha256:e37ab26028f12dbb0ff65f29a8d3d44a765c61e729647bf2ddfbbed621726f01", size = 303425, upload-time = "2025-08-07T13:32:27.59Z" }, ] From a9244c239a403d0452163fb52d1cbb3c02869196 Mon Sep 17 00:00:00 2001 From: farzad528 Date: Mon, 17 Nov 2025 17:56:21 +0000 Subject: [PATCH 04/13] init update --- python/packages/core/agent_framework/azure/__init__.py | 1 + 1 file changed, 1 insertion(+) diff --git a/python/packages/core/agent_framework/azure/__init__.py b/python/packages/core/agent_framework/azure/__init__.py index e28acf6a7f..80c6ece5a4 100644 --- a/python/packages/core/agent_framework/azure/__init__.py +++ b/python/packages/core/agent_framework/azure/__init__.py @@ -9,6 +9,7 @@ "AgentFunctionApp": ("agent_framework_azurefunctions", "azurefunctions"), "AgentResponseCallbackProtocol": ("agent_framework_azurefunctions", "azurefunctions"), "AzureAIAgentClient": ("agent_framework_azure_ai", "azure-ai"), + "AzureAIClient": ("agent_framework_azure_ai", "azure-ai"), "AzureAISearchContextProvider": ("agent_framework.azure._search_provider", "core"), "AzureOpenAIAssistantsClient": ("agent_framework.azure._assistants_client", "core"), "AzureOpenAIChatClient": ("agent_framework.azure._chat_client", "core"), From c7d1e79aa1b77343177058d4e08f42e9934361c3 Mon Sep 17 00:00:00 2001 From: farzad528 Date: Mon, 17 Nov 2025 19:25:08 +0000 Subject: [PATCH 05/13] Fix MyPy type annotation errors in search provider - Add type annotation to DEFAULT_CONTEXT_PROMPT - Add type annotation to vectorizable_fields - Add union type annotation to vector_queries --- .../packages/core/agent_framework/azure/_search_provider.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/python/packages/core/agent_framework/azure/_search_provider.py b/python/packages/core/agent_framework/azure/_search_provider.py index 8852e65839..2cac7d48f9 100644 --- a/python/packages/core/agent_framework/azure/_search_provider.py +++ b/python/packages/core/agent_framework/azure/_search_provider.py @@ -124,7 +124,7 @@ class AzureAISearchContextProvider(ContextProvider): ) """ - DEFAULT_CONTEXT_PROMPT = "Use the following context to answer the question:" + DEFAULT_CONTEXT_PROMPT: str = "Use the following context to answer the question:" def __init__( self, @@ -355,7 +355,7 @@ def _find_vectorizable_fields(self, index: Any, vector_fields: list[str]) -> lis Returns: List of vectorizable field names (subset of vector_fields). """ - vectorizable_fields = [] + vectorizable_fields: list[str] = [] # Check if index has vector search configuration if not index.vector_search or not index.vector_search.profiles: @@ -486,7 +486,7 @@ async def _semantic_search(self, query: str) -> str: # Auto-discover vector field if not already done await self._auto_discover_vector_field() - vector_queries = [] + vector_queries: list[VectorizableTextQuery | VectorizedQuery] = [] # Build vector query based on server-side vectorization or client-side embedding if self.vector_field_name: From df2aaac5bc938534167ba582618941d4864669e6 Mon Sep 17 00:00:00 2001 From: farzad528 Date: Mon, 17 Nov 2025 19:32:08 +0000 Subject: [PATCH 06/13] Fix DEFAULT_CONTEXT_PROMPT MyPy error and update test - Rename DEFAULT_CONTEXT_PROMPT to _DEFAULT_SEARCH_CONTEXT_PROMPT to avoid conflict with base class Final variable - Update test to use new constant name - All core package tests passing (1123 passed) --- .../packages/core/agent_framework/azure/_search_provider.py | 6 +++--- python/packages/core/tests/azure/test_search_provider.py | 2 +- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/python/packages/core/agent_framework/azure/_search_provider.py b/python/packages/core/agent_framework/azure/_search_provider.py index 2cac7d48f9..0c066bb08e 100644 --- a/python/packages/core/agent_framework/azure/_search_provider.py +++ b/python/packages/core/agent_framework/azure/_search_provider.py @@ -124,7 +124,7 @@ class AzureAISearchContextProvider(ContextProvider): ) """ - DEFAULT_CONTEXT_PROMPT: str = "Use the following context to answer the question:" + _DEFAULT_SEARCH_CONTEXT_PROMPT = "Use the following context to answer the question:" def __init__( self, @@ -168,7 +168,7 @@ def __init__( Signature: async def embed(text: str) -> list[float] Required if vector_field_name is specified. context_prompt: Custom prompt to prepend to retrieved context. - Default: Uses DEFAULT_CONTEXT_PROMPT. + Default: "Use the following context to answer the question:" azure_ai_project_endpoint: Azure AI Foundry project endpoint URL. Required for agentic mode. Example: "https://myproject.services.ai.azure.com" model_deployment_name: Model deployment name in the Azure AI project. @@ -195,7 +195,7 @@ def __init__( 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_CONTEXT_PROMPT + self.context_prompt = context_prompt or self._DEFAULT_SEARCH_CONTEXT_PROMPT # Agentic mode parameters (Knowledge Base) # azure_openai_resource_url: The actual Azure OpenAI endpoint for model calls diff --git a/python/packages/core/tests/azure/test_search_provider.py b/python/packages/core/tests/azure/test_search_provider.py index 2db0968ac2..e32d3e9955 100644 --- a/python/packages/core/tests/azure/test_search_provider.py +++ b/python/packages/core/tests/azure/test_search_provider.py @@ -158,7 +158,7 @@ def test_init_uses_default_context_prompt(self) -> None: credential=AzureKeyCredential("test-key"), mode="semantic", ) - assert provider.context_prompt == provider.DEFAULT_CONTEXT_PROMPT + assert provider.context_prompt == provider._DEFAULT_SEARCH_CONTEXT_PROMPT class TestSemanticSearch: From a02b1c7b944bc08c1f8a8836641014d72f17d919 Mon Sep 17 00:00:00 2001 From: farzad528 Date: Wed, 19 Nov 2025 14:09:31 +0000 Subject: [PATCH 07/13] Python: Move Azure AI Search to separate package per PR feedback MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Addresses reviewer feedback from PR #1546 by isolating the beta dependency (azure-search-documents==11.7.0b2) into a new agent-framework-aisearch package. Changes: - Created new agent-framework-aisearch package with complete structure - Moved AzureAISearchContextProvider from core to aisearch package - Added AzureAISearchSettings class for environment variable auto-loading - Added support for direct API key string (auto-converts to AzureKeyCredential) - Added azure_openai_api_key parameter for Knowledge Base authentication - Updated embedding_function type to Callable[[str], Awaitable[list[float]]] - Moved Role import to top-level imports - Maintained lazy loading through agent_framework.azure module - Removed beta dependency from core package - Updated all tests to use new package location - All quality checks pass: ruff format/lint, pyright, mypy (0 errors) - All 21 unit tests pass with 59% coverage Semantic search mode verified working with both API key and managed identity authentication. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude --- python/packages/aisearch/LICENSE | 21 +++ python/packages/aisearch/README.md | 23 ++++ .../agent_framework_aisearch/__init__.py | 16 +++ .../_search_provider.py | 127 ++++++++++++++---- python/packages/aisearch/pyproject.toml | 92 +++++++++++++ python/packages/aisearch/tests/__init__.py | 1 + .../tests}/test_search_provider.py | 35 +++-- .../core/agent_framework/azure/__init__.py | 3 +- python/packages/core/pyproject.toml | 1 - python/uv.lock | 22 ++- 10 files changed, 293 insertions(+), 48 deletions(-) create mode 100644 python/packages/aisearch/LICENSE create mode 100644 python/packages/aisearch/README.md create mode 100644 python/packages/aisearch/agent_framework_aisearch/__init__.py rename python/packages/{core/agent_framework/azure => aisearch/agent_framework_aisearch}/_search_provider.py (84%) create mode 100644 python/packages/aisearch/pyproject.toml create mode 100644 python/packages/aisearch/tests/__init__.py rename python/packages/{core/tests/azure => aisearch/tests}/test_search_provider.py (95%) 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/core/agent_framework/azure/_search_provider.py b/python/packages/aisearch/agent_framework_aisearch/_search_provider.py similarity index 84% rename from python/packages/core/agent_framework/azure/_search_provider.py rename to python/packages/aisearch/agent_framework_aisearch/_search_provider.py index 0c066bb08e..6648490864 100644 --- a/python/packages/core/agent_framework/azure/_search_provider.py +++ b/python/packages/aisearch/agent_framework_aisearch/_search_provider.py @@ -10,10 +10,12 @@ reasoning across documents with Knowledge Bases. """ +import os import sys -from collections.abc import MutableSequence +from collections.abc import Awaitable, Callable, MutableSequence from typing import TYPE_CHECKING, Any, Literal +from agent_framework import ChatMessage, Context, ContextProvider, Role from azure.core.credentials import AzureKeyCredential from azure.core.credentials_async import AsyncTokenCredential from azure.core.exceptions import ResourceNotFoundError @@ -35,8 +37,6 @@ VectorizedQuery, ) -from agent_framework import ChatMessage, Context, ContextProvider - # Type checking imports for optional agentic mode dependencies if TYPE_CHECKING: from azure.search.documents.knowledgebases.aio import KnowledgeBaseRetrievalClient @@ -70,6 +70,57 @@ from typing_extensions import override # type: ignore[import] # pragma: no cover +class AzureAISearchSettings: + """Settings for Azure AI Search Context Provider with auto-loading from environment. + + Environment variables: + AZURE_SEARCH_ENDPOINT: Azure AI Search endpoint URL + AZURE_SEARCH_INDEX_NAME: Name of the search index + AZURE_SEARCH_API_KEY: API key for authentication + AZURE_AI_PROJECT_ENDPOINT: Azure AI Foundry project endpoint (for agentic mode) + AZURE_OPENAI_RESOURCE_URL: Azure OpenAI resource URL (for agentic mode) + AZURE_OPENAI_DEPLOYMENT_NAME: Model deployment name (for agentic mode) + AZURE_OPENAI_API_KEY: Azure OpenAI API key (for agentic mode with API key auth) + """ + + def __init__( + self, + *, + endpoint: str | None = None, + index_name: str | None = None, + api_key: str | None = None, + azure_ai_project_endpoint: str | None = None, + azure_openai_resource_url: str | None = None, + model_deployment_name: str | None = None, + azure_openai_api_key: str | None = None, + ) -> None: + """Initialize settings with auto-loading from environment variables. + + Args: + endpoint: Azure AI Search endpoint URL (or AZURE_SEARCH_ENDPOINT env var) + index_name: Search index name (or AZURE_SEARCH_INDEX_NAME env var) + api_key: API key (or AZURE_SEARCH_API_KEY env var) + azure_ai_project_endpoint: Azure AI Foundry project endpoint (or AZURE_AI_PROJECT_ENDPOINT env var) + azure_openai_resource_url: Azure OpenAI resource URL (or AZURE_OPENAI_RESOURCE_URL env var) + model_deployment_name: Model deployment name (or AZURE_OPENAI_DEPLOYMENT_NAME env var) + azure_openai_api_key: Azure OpenAI API key (or AZURE_OPENAI_API_KEY env var) + """ + self.endpoint = endpoint or os.getenv("AZURE_SEARCH_ENDPOINT") + self.index_name = index_name or os.getenv("AZURE_SEARCH_INDEX_NAME") + self.api_key = api_key or os.getenv("AZURE_SEARCH_API_KEY") + self.azure_ai_project_endpoint = azure_ai_project_endpoint or os.getenv("AZURE_AI_PROJECT_ENDPOINT") + self.azure_openai_resource_url = azure_openai_resource_url or os.getenv("AZURE_OPENAI_RESOURCE_URL") + self.model_deployment_name = model_deployment_name or os.getenv("AZURE_OPENAI_DEPLOYMENT_NAME") + self.azure_openai_api_key = azure_openai_api_key or os.getenv("AZURE_OPENAI_API_KEY") + + def validate(self) -> None: + """Validate required settings are present.""" + if not self.endpoint: + raise ValueError("endpoint is required (or set AZURE_SEARCH_ENDPOINT)") + if not self.index_name: + raise ValueError("index_name is required (or set AZURE_SEARCH_INDEX_NAME)") + + class AzureAISearchContextProvider(ContextProvider): """Azure AI Search Context Provider with hybrid search and semantic ranking. @@ -82,32 +133,31 @@ class AzureAISearchContextProvider(ContextProvider): Use only for complex queries requiring cross-document reasoning. Examples: - Semantic hybrid search (recommended for most cases): + Using Settings class with environment variables: .. code-block:: python from agent_framework import ChatAgent - from agent_framework_azure_ai import AzureAIAgentClient, AzureAISearchContextProvider + from agent_framework_azure_ai import AzureAIAgentClient + from agent_framework.azure import AzureAISearchContextProvider, AzureAISearchSettings from azure.identity.aio import DefaultAzureCredential - # Create context provider with semantic hybrid search + # Load from environment variables + settings = AzureAISearchSettings() + search_provider = AzureAISearchContextProvider(settings=settings, credential=DefaultAzureCredential()) + + Semantic hybrid search with API key: + + .. code-block:: python + + # Direct API key string (converted to AzureKeyCredential automatically) search_provider = AzureAISearchContextProvider( endpoint="https://mysearch.search.windows.net", index_name="my-index", - credential=DefaultAzureCredential(), - mode="semantic", # Fast hybrid + semantic ranker (default) + credential="my-api-key", # String converted to AzureKeyCredential + mode="semantic", ) - # Use with agent - async with ( - AzureAIAgentClient() as client, - ChatAgent( - chat_client=client, - context_providers=[search_provider], - ) as agent, - ): - response = await agent.run("What is in the documents?") - Agentic retrieval for complex queries: .. code-block:: python @@ -128,15 +178,16 @@ class AzureAISearchContextProvider(ContextProvider): def __init__( self, - endpoint: str, - index_name: str, - credential: AzureKeyCredential | AsyncTokenCredential, + endpoint: str | None = None, + index_name: str | None = None, + credential: str | AzureKeyCredential | AsyncTokenCredential | None = None, *, + settings: AzureAISearchSettings | 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: Any | 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, @@ -156,7 +207,10 @@ def __init__( Args: endpoint: Azure AI Search endpoint URL. index_name: Name of the search index to query. - credential: Azure credential (API key or DefaultAzureCredential). + credential: Azure credential (API key string, AzureKeyCredential, or AsyncTokenCredential). + If a string is provided, it will be converted to AzureKeyCredential. + settings: Settings object with auto-loaded configuration from environment. + If provided, endpoint/index_name/api_key from settings will be used as defaults. 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. Default: 5. @@ -166,7 +220,7 @@ def __init__( 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. + 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. @@ -187,6 +241,30 @@ def __init__( azure_openai_deployment_name: (Deprecated) Use model_deployment_name instead. azure_openai_api_version: (Deprecated) No longer used. """ + # Load from settings if provided + if settings: + settings.validate() + endpoint = endpoint or settings.endpoint + index_name = index_name or settings.index_name + if not credential and settings.api_key: + credential = settings.api_key + azure_ai_project_endpoint = azure_ai_project_endpoint or settings.azure_ai_project_endpoint + azure_openai_resource_url = azure_openai_resource_url or settings.azure_openai_resource_url + model_deployment_name = model_deployment_name or settings.model_deployment_name + azure_openai_api_key = azure_openai_api_key or settings.azure_openai_api_key + + # Validate required parameters + if not endpoint: + raise ValueError("endpoint is required") + if not index_name: + raise ValueError("index_name is required") + if not credential: + raise ValueError("credential is required") + + # Convert string credential to AzureKeyCredential + if isinstance(credential, str): + credential = AzureKeyCredential(credential) + self.endpoint = endpoint self.index_name = index_name self.credential = credential @@ -298,7 +376,6 @@ async def invoking( """ # Convert to list and filter to USER/ASSISTANT messages with text only messages_list = [messages] if isinstance(messages, ChatMessage) else list(messages) - from agent_framework import Role filtered_messages = [ msg diff --git a/python/packages/aisearch/pyproject.toml b/python/packages/aisearch/pyproject.toml new file mode 100644 index 0000000000..07720b4f6d --- /dev/null +++ b/python/packages/aisearch/pyproject.toml @@ -0,0 +1,92 @@ +[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-identity>=1,<2", + "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/__init__.py b/python/packages/aisearch/tests/__init__.py new file mode 100644 index 0000000000..2a50eae894 --- /dev/null +++ b/python/packages/aisearch/tests/__init__.py @@ -0,0 +1 @@ +# Copyright (c) Microsoft. All rights reserved. diff --git a/python/packages/core/tests/azure/test_search_provider.py b/python/packages/aisearch/tests/test_search_provider.py similarity index 95% rename from python/packages/core/tests/azure/test_search_provider.py rename to python/packages/aisearch/tests/test_search_provider.py index e32d3e9955..e17ec89db8 100644 --- a/python/packages/core/tests/azure/test_search_provider.py +++ b/python/packages/aisearch/tests/test_search_provider.py @@ -4,11 +4,10 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest -from azure.core.credentials import AzureKeyCredential -from azure.core.exceptions import ResourceNotFoundError - from agent_framework import ChatMessage, Context, Role from agent_framework.azure import AzureAISearchContextProvider +from azure.core.credentials import AzureKeyCredential +from azure.core.exceptions import ResourceNotFoundError @pytest.fixture @@ -165,7 +164,7 @@ class TestSemanticSearch: """Test semantic search functionality.""" @pytest.mark.asyncio - @patch("agent_framework.azure._search_provider.SearchClient") + @patch("agent_framework_aisearch._search_provider.SearchClient") async def test_semantic_search_basic( self, mock_search_class: MagicMock, sample_messages: list[ChatMessage] ) -> None: @@ -191,7 +190,7 @@ async def test_semantic_search_basic( assert "Test document content" in context.messages[0].text @pytest.mark.asyncio - @patch("agent_framework.azure._search_provider.SearchClient") + @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() @@ -211,7 +210,7 @@ async def test_semantic_search_empty_query(self, mock_search_class: MagicMock) - assert len(context.messages) == 0 @pytest.mark.asyncio - @patch("agent_framework.azure._search_provider.SearchClient") + @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: @@ -248,8 +247,8 @@ class TestKnowledgeBaseSetup: """Test Knowledge Base setup for agentic mode.""" @pytest.mark.asyncio - @patch("agent_framework.azure._search_provider.SearchIndexClient") - @patch("agent_framework.azure._search_provider.SearchClient") + @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: @@ -285,8 +284,8 @@ async def test_ensure_knowledge_base_creates_when_not_exists( mock_index_client.create_or_update_knowledge_base.assert_called_once() @pytest.mark.asyncio - @patch("agent_framework.azure._search_provider.SearchIndexClient") - @patch("agent_framework.azure._search_provider.SearchClient") + @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: @@ -322,7 +321,7 @@ class TestContextProviderLifecycle: """Test context provider lifecycle methods.""" @pytest.mark.asyncio - @patch("agent_framework.azure._search_provider.SearchClient") + @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() @@ -342,7 +341,7 @@ class TestMessageFiltering: """Test message filtering functionality.""" @pytest.mark.asyncio - @patch("agent_framework.azure._search_provider.SearchClient") + @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 @@ -374,7 +373,7 @@ async def test_filters_non_user_assistant_messages(self, mock_search_class: Magi mock_search_client.search.assert_called_once() @pytest.mark.asyncio - @patch("agent_framework.azure._search_provider.SearchClient") + @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() @@ -404,7 +403,7 @@ class TestCitations: """Test citation functionality.""" @pytest.mark.asyncio - @patch("agent_framework.azure._search_provider.SearchClient") + @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 @@ -435,8 +434,8 @@ class TestVectorFieldAutoDiscovery: """Test vector field auto-discovery functionality.""" @pytest.mark.asyncio - @patch("agent_framework.azure._search_provider.SearchIndexClient") - @patch("agent_framework.azure._search_provider.SearchClient") + @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: @@ -500,8 +499,8 @@ async def test_vector_detection_accuracy(self) -> None: assert is_vector_3 is False @pytest.mark.asyncio - @patch("agent_framework.azure._search_provider.SearchIndexClient") - @patch("agent_framework.azure._search_provider.SearchClient") + @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: diff --git a/python/packages/core/agent_framework/azure/__init__.py b/python/packages/core/agent_framework/azure/__init__.py index 80c6ece5a4..09670188ee 100644 --- a/python/packages/core/agent_framework/azure/__init__.py +++ b/python/packages/core/agent_framework/azure/__init__.py @@ -10,7 +10,8 @@ "AgentResponseCallbackProtocol": ("agent_framework_azurefunctions", "azurefunctions"), "AzureAIAgentClient": ("agent_framework_azure_ai", "azure-ai"), "AzureAIClient": ("agent_framework_azure_ai", "azure-ai"), - "AzureAISearchContextProvider": ("agent_framework.azure._search_provider", "core"), + "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 7d357e237a..8f19719c59 100644 --- a/python/packages/core/pyproject.toml +++ b/python/packages/core/pyproject.toml @@ -36,7 +36,6 @@ dependencies = [ "openai>=1.99.0", "azure-identity>=1,<2", "mcp[ws]>=1.13", - "azure-search-documents==11.7.0b2", "packaging>=24.1", ] diff --git a/python/uv.lock b/python/uv.lock index 0ce58797a0..c62f586d30 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", @@ -190,7 +191,7 @@ requires-dist = [ [[package]] name = "agent-framework-ag-ui" -version = "1.0.0b251114" +version = "1.0.0b251117" source = { editable = "packages/ag-ui" } dependencies = [ { name = "ag-ui-protocol", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, @@ -218,6 +219,23 @@ 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-identity", 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-identity", specifier = ">=1,<2" }, + { name = "azure-search-documents", specifier = "==11.7.0b2" }, +] + [[package]] name = "agent-framework-anthropic" version = "1.0.0b251114" @@ -305,7 +323,6 @@ version = "1.0.0b251114" source = { editable = "packages/core" } dependencies = [ { name = "azure-identity", 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'" }, { name = "mcp", extra = ["ws"], marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, { name = "openai", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, { name = "opentelemetry-api", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, @@ -343,7 +360,6 @@ requires-dist = [ { name = "agent-framework-purview", marker = "extra == 'all'", editable = "packages/purview" }, { name = "agent-framework-redis", marker = "extra == 'all'", editable = "packages/redis" }, { name = "azure-identity", specifier = ">=1,<2" }, - { name = "azure-search-documents", specifier = "==11.7.0b2" }, { name = "mcp", extras = ["ws"], specifier = ">=1.13" }, { name = "openai", specifier = ">=1.99.0" }, { name = "opentelemetry-api", specifier = ">=1.24" }, From be6b61f274d49d70373408ca84bdad8b86ecbe34 Mon Sep 17 00:00:00 2001 From: farzad528 Date: Wed, 19 Nov 2025 14:38:45 +0000 Subject: [PATCH 08/13] Python: Clarify top_k parameter only applies to semantic mode MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Updated documentation to clarify that the top_k parameter only affects semantic search mode. In agentic mode, the server-side Knowledge Base determines retrieval based on query complexity and reasoning effort. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude --- .../aisearch/agent_framework_aisearch/_search_provider.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/python/packages/aisearch/agent_framework_aisearch/_search_provider.py b/python/packages/aisearch/agent_framework_aisearch/_search_provider.py index 6648490864..b2d792a2f8 100644 --- a/python/packages/aisearch/agent_framework_aisearch/_search_provider.py +++ b/python/packages/aisearch/agent_framework_aisearch/_search_provider.py @@ -213,7 +213,9 @@ def __init__( If provided, endpoint/index_name/api_key from settings will be used as defaults. 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. Default: 5. + 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. From f743aaa0e93fa9b2ca93561370a8887be15bbbf1 Mon Sep 17 00:00:00 2001 From: farzad528 Date: Wed, 19 Nov 2025 15:03:39 +0000 Subject: [PATCH 09/13] Python: Add Knowledge Base output mode and retrieval reasoning effort parameters MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Added support for configurable Knowledge Base behavior in agentic mode: - knowledge_base_output_mode: "extractive_data" (default) or "answer_synthesis" Some knowledge sources require answer_synthesis mode for proper functionality. - retrieval_reasoning_effort: "minimal" (default), "medium", or "low" Controls query planning complexity and multi-hop reasoning depth. These parameters give users fine-grained control over Knowledge Base behavior and enable support for knowledge sources that require answer synthesis. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude --- .../_search_provider.py | 37 +++++++++++++++++-- 1 file changed, 34 insertions(+), 3 deletions(-) diff --git a/python/packages/aisearch/agent_framework_aisearch/_search_provider.py b/python/packages/aisearch/agent_framework_aisearch/_search_provider.py index b2d792a2f8..8116715805 100644 --- a/python/packages/aisearch/agent_framework_aisearch/_search_provider.py +++ b/python/packages/aisearch/agent_framework_aisearch/_search_provider.py @@ -25,7 +25,11 @@ AzureOpenAIVectorizerParameters, KnowledgeBase, KnowledgeBaseAzureOpenAIModel, + KnowledgeRetrievalLowReasoningEffort, + KnowledgeRetrievalMediumReasoningEffort, + KnowledgeRetrievalMinimalReasoningEffort, KnowledgeRetrievalOutputMode, + KnowledgeRetrievalReasoningEffort, KnowledgeSourceReference, SearchIndexKnowledgeSource, SearchIndexKnowledgeSourceParameters, @@ -197,6 +201,8 @@ def __init__( retrieval_instructions: str | None = None, azure_openai_api_key: str | None = None, azure_openai_resource_url: str | None = None, + knowledge_base_output_mode: Literal["extractive_data", "answer_synthesis"] = "extractive_data", + retrieval_reasoning_effort: Literal["minimal", "medium", "low"] = "minimal", # Deprecated parameters (for backwards compatibility) azure_openai_endpoint: str | None = None, azure_openai_deployment_name: str | None = None, @@ -239,6 +245,15 @@ def __init__( azure_openai_resource_url: Azure OpenAI resource URL for Knowledge Base model calls. Required for agentic mode. Example: "https://myresource.openai.azure.com" This is different from azure_ai_project_endpoint (which is Foundry-specific). + 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". azure_openai_endpoint: (Deprecated) Use azure_ai_project_endpoint instead. azure_openai_deployment_name: (Deprecated) Use model_deployment_name instead. azure_openai_api_version: (Deprecated) No longer used. @@ -289,6 +304,8 @@ def __init__( 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 # Auto-discover vector field if not specified self._auto_discovered_vector_field = False @@ -664,8 +681,6 @@ async def _ensure_knowledge_base(self) -> None: # Step 2: Create or update Knowledge Base # Always create/update to ensure configuration is current - # Note: EXTRACTIVE_DATA mode returns raw chunks without synthesis - # Model is still needed for query planning and multi-hop reasoning aoai_params = AzureOpenAIVectorizerParameters( resource_url=self.azure_openai_resource_url, deployment_name=self.azure_openai_deployment_name, @@ -673,6 +688,21 @@ async def _ensure_knowledge_base(self) -> None: 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}", @@ -682,7 +712,8 @@ async def _ensure_knowledge_base(self) -> None: ) ], models=[KnowledgeBaseAzureOpenAIModel(azure_open_ai_parameters=aoai_params)], - output_mode=KnowledgeRetrievalOutputMode.EXTRACTIVE_DATA, + output_mode=output_mode, + retrieval_reasoning_effort=reasoning_effort, ) await self._index_client.create_or_update_knowledge_base(knowledge_base) From 004a55cb9e6c1d66199b27f71164268f5438d269 Mon Sep 17 00:00:00 2001 From: farzad528 Date: Wed, 19 Nov 2025 16:16:46 +0000 Subject: [PATCH 10/13] effort and outputmode query params --- .../_search_provider.py | 85 ++++++++++++++++--- 1 file changed, 75 insertions(+), 10 deletions(-) diff --git a/python/packages/aisearch/agent_framework_aisearch/_search_provider.py b/python/packages/aisearch/agent_framework_aisearch/_search_provider.py index 8116715805..02efdb6e81 100644 --- a/python/packages/aisearch/agent_framework_aisearch/_search_provider.py +++ b/python/packages/aisearch/agent_framework_aisearch/_search_provider.py @@ -48,6 +48,23 @@ 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) @@ -57,6 +74,23 @@ 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 @@ -748,17 +782,48 @@ async def _agentic_search(self, messages: list[ChatMessage]) -> str: # Ensure Knowledge Base is initialized await self._ensure_knowledge_base() - # Convert ChatMessage list to KnowledgeBase message format - 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 - ] + # 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] - retrieval_request = KnowledgeBaseRetrievalRequest(messages=kb_messages) + # 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: From 5d968453452c1666c3e61828e9f2383b0f856b16 Mon Sep 17 00:00:00 2001 From: farzad528 Date: Wed, 19 Nov 2025 18:04:40 +0000 Subject: [PATCH 11/13] Address PR review feedback for Azure AI Search context provider --- .../_search_provider.py | 191 ++++++++---------- python/packages/aisearch/pyproject.toml | 1 - .../aisearch/tests/test_search_provider.py | 14 +- python/packages/core/pyproject.toml | 1 + python/pyproject.toml | 8 +- .../getting_started/agents/azure_ai/README.md | 4 +- .../azure_ai_with_search_context_agentic.py | 117 +++++++++++ .../azure_ai_with_search_context_semantic.py | 97 +++++++++ python/uv.lock | 8 +- 9 files changed, 317 insertions(+), 124 deletions(-) create mode 100644 python/samples/getting_started/agents/azure_ai/azure_ai_with_search_context_agentic.py create mode 100644 python/samples/getting_started/agents/azure_ai/azure_ai_with_search_context_semantic.py diff --git a/python/packages/aisearch/agent_framework_aisearch/_search_provider.py b/python/packages/aisearch/agent_framework_aisearch/_search_provider.py index 02efdb6e81..c633483c00 100644 --- a/python/packages/aisearch/agent_framework_aisearch/_search_provider.py +++ b/python/packages/aisearch/agent_framework_aisearch/_search_provider.py @@ -10,12 +10,13 @@ reasoning across documents with Knowledge Bases. """ -import os import sys from collections.abc import Awaitable, Callable, MutableSequence -from typing import TYPE_CHECKING, Any, Literal +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 azure.core.credentials import AzureKeyCredential from azure.core.credentials_async import AsyncTokenCredential from azure.core.exceptions import ResourceNotFoundError @@ -107,8 +108,12 @@ 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: + +class AzureAISearchSettings(AFBaseSettings): """Settings for Azure AI Search Context Provider with auto-loading from environment. Environment variables: @@ -121,43 +126,35 @@ class AzureAISearchSettings: AZURE_OPENAI_API_KEY: Azure OpenAI API key (for agentic mode with API key auth) """ - def __init__( - self, - *, - endpoint: str | None = None, - index_name: str | None = None, - api_key: str | None = None, - azure_ai_project_endpoint: str | None = None, - azure_openai_resource_url: str | None = None, - model_deployment_name: str | None = None, - azure_openai_api_key: str | None = None, - ) -> None: - """Initialize settings with auto-loading from environment variables. + env_prefix: ClassVar[str] = "AZURE_" - Args: - endpoint: Azure AI Search endpoint URL (or AZURE_SEARCH_ENDPOINT env var) - index_name: Search index name (or AZURE_SEARCH_INDEX_NAME env var) - api_key: API key (or AZURE_SEARCH_API_KEY env var) - azure_ai_project_endpoint: Azure AI Foundry project endpoint (or AZURE_AI_PROJECT_ENDPOINT env var) - azure_openai_resource_url: Azure OpenAI resource URL (or AZURE_OPENAI_RESOURCE_URL env var) - model_deployment_name: Model deployment name (or AZURE_OPENAI_DEPLOYMENT_NAME env var) - azure_openai_api_key: Azure OpenAI API key (or AZURE_OPENAI_API_KEY env var) - """ - self.endpoint = endpoint or os.getenv("AZURE_SEARCH_ENDPOINT") - self.index_name = index_name or os.getenv("AZURE_SEARCH_INDEX_NAME") - self.api_key = api_key or os.getenv("AZURE_SEARCH_API_KEY") - self.azure_ai_project_endpoint = azure_ai_project_endpoint or os.getenv("AZURE_AI_PROJECT_ENDPOINT") - self.azure_openai_resource_url = azure_openai_resource_url or os.getenv("AZURE_OPENAI_RESOURCE_URL") - self.model_deployment_name = model_deployment_name or os.getenv("AZURE_OPENAI_DEPLOYMENT_NAME") - self.azure_openai_api_key = azure_openai_api_key or os.getenv("AZURE_OPENAI_API_KEY") - - def validate(self) -> None: + endpoint: str | None = None + search_endpoint: str | None = None # Alias for endpoint + index_name: str | None = None + search_index_name: str | None = None # Alias for index_name + api_key: str | None = None + search_api_key: str | None = None # Alias for api_key + ai_project_endpoint: str | None = None + openai_resource_url: str | None = None + openai_deployment_name: str | None = None + openai_api_key: str | None = None + + def validate_settings(self) -> None: """Validate required settings are present.""" - if not self.endpoint: + # Use aliases if main fields not set + endpoint = self.endpoint or self.search_endpoint + index_name = self.index_name or self.search_index_name + + if not endpoint: raise ValueError("endpoint is required (or set AZURE_SEARCH_ENDPOINT)") - if not self.index_name: + if not index_name: raise ValueError("index_name is required (or set AZURE_SEARCH_INDEX_NAME)") + # Update main fields from aliases + self.endpoint = endpoint + self.index_name = index_name + self.api_key = self.api_key or self.search_api_key + class AzureAISearchContextProvider(ContextProvider): """Azure AI Search Context Provider with hybrid search and semantic ranking. @@ -229,18 +226,15 @@ def __init__( 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, - azure_openai_resource_url: str | None = None, knowledge_base_output_mode: Literal["extractive_data", "answer_synthesis"] = "extractive_data", retrieval_reasoning_effort: Literal["minimal", "medium", "low"] = "minimal", - # Deprecated parameters (for backwards compatibility) - azure_openai_endpoint: str | None = None, - azure_openai_deployment_name: str | None = None, - azure_openai_api_version: str | None = None, + agentic_message_history_count: int = _DEFAULT_AGENTIC_MESSAGE_HISTORY_COUNT, ) -> None: """Initialize Azure AI Search Context Provider. @@ -250,7 +244,7 @@ def __init__( credential: Azure credential (API key string, AzureKeyCredential, or AsyncTokenCredential). If a string is provided, it will be converted to AzureKeyCredential. settings: Settings object with auto-loaded configuration from environment. - If provided, endpoint/index_name/api_key from settings will be used as defaults. + If provided, endpoint/index_name/credential from settings will be used as defaults. 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. @@ -274,11 +268,10 @@ def __init__( 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. azure_openai_resource_url: Azure OpenAI resource URL for Knowledge Base model calls. Required for agentic mode. Example: "https://myresource.openai.azure.com" - This is different from azure_ai_project_endpoint (which is Foundry-specific). + 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. @@ -288,21 +281,20 @@ def __init__( "medium": Moderate reasoning with some query decomposition. "low": Lower reasoning effort than medium. Default: "minimal". - azure_openai_endpoint: (Deprecated) Use azure_ai_project_endpoint instead. - azure_openai_deployment_name: (Deprecated) Use model_deployment_name instead. - azure_openai_api_version: (Deprecated) No longer used. + agentic_message_history_count: Number of recent messages to send to the Knowledge Base for context + in agentic mode. Default: 10. """ # Load from settings if provided if settings: - settings.validate() + settings.validate_settings() endpoint = endpoint or settings.endpoint index_name = index_name or settings.index_name if not credential and settings.api_key: credential = settings.api_key - azure_ai_project_endpoint = azure_ai_project_endpoint or settings.azure_ai_project_endpoint - azure_openai_resource_url = azure_openai_resource_url or settings.azure_openai_resource_url - model_deployment_name = model_deployment_name or settings.model_deployment_name - azure_openai_api_key = azure_openai_api_key or settings.azure_openai_api_key + azure_ai_project_endpoint = azure_ai_project_endpoint or settings.ai_project_endpoint + azure_openai_resource_url = azure_openai_resource_url or settings.openai_resource_url + model_deployment_name = model_deployment_name or settings.openai_deployment_name + azure_openai_api_key = azure_openai_api_key or settings.openai_api_key # Validate required parameters if not endpoint: @@ -327,19 +319,17 @@ def __init__( self.context_prompt = context_prompt or self._DEFAULT_SEARCH_CONTEXT_PROMPT # Agentic mode parameters (Knowledge Base) - # azure_openai_resource_url: The actual Azure OpenAI endpoint for model calls - # azure_openai_endpoint (deprecated): Fall back to this if resource_url not provided - self.azure_openai_resource_url = azure_openai_resource_url or azure_openai_endpoint - - self.azure_openai_deployment_name = model_deployment_name or azure_openai_deployment_name - # If model_name not provided, default to deployment name for backwards compatibility - self.model_name = model_name or self.azure_openai_deployment_name + 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 @@ -361,13 +351,11 @@ def __init__( ) if not self.azure_openai_resource_url: raise ValueError( - "azure_openai_resource_url (or deprecated azure_openai_endpoint) is required for agentic mode. " + "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 (or deprecated azure_openai_deployment_name) is required for agentic mode" - ) + 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") @@ -443,19 +431,21 @@ async def invoking( if self.mode == "semantic": # Semantic mode: flatten messages to single query query = "\n".join(msg.text for msg in filtered_messages) - search_results = await self._semantic_search(query) + search_result_parts = await self._semantic_search(query) else: # agentic - # Agentic mode: pass last 10 messages as conversation history - recent_messages = filtered_messages[-10:] - search_results = await self._agentic_search(recent_messages) + # 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 - if not search_results: + # Format results as context - return multiple messages for each result part + if not search_result_parts: return Context() - context_text = f"{self.context_prompt}\n\n{search_results}" + # 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=[ChatMessage(role="system", text=context_text)]) + return Context(messages=context_messages) def _find_vector_fields(self, index: Any) -> list[str]: """Find all fields that can store vectors (have dimensions defined). @@ -516,13 +506,10 @@ async def _auto_discover_vector_field(self) -> None: return # Already discovered or manually specified try: - # Need index client to get schema + # Use existing index client or create temporary one if not self._index_client: - from azure.search.documents.indexes.aio import SearchIndexClient - - index_client = SearchIndexClient(endpoint=self.endpoint, credential=self.credential) - else: - index_client = 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) @@ -532,12 +519,8 @@ async def _auto_discover_vector_field(self) -> None: if not vector_fields: # No vector fields found - keyword search only - import logging - - logging.info(f"No vector fields found in index '{self.index_name}'. Using keyword-only search.") + logger.info(f"No vector fields found in index '{self.index_name}'. Using keyword-only search.") self._auto_discovered_vector_field = True - if not self._index_client: - await index_client.close() return # Step 2: Find which vector fields have server-side vectorization @@ -550,17 +533,13 @@ async def _auto_discover_vector_field(self) -> None: self.vector_field_name = vectorizable_fields[0] self._auto_discovered_vector_field = True self._use_vectorizable_query = True # Use VectorizableTextQuery - import logging - - logging.info( + logger.info( f"Auto-discovered vectorizable field '{self.vector_field_name}' " f"with server-side vectorization. No embedding_function needed." ) else: # Multiple vectorizable fields - import logging - - logging.warning( + logger.warning( f"Multiple vectorizable fields found: {vectorizable_fields}. " f"Please specify vector_field_name explicitly. Using keyword-only search." ) @@ -571,35 +550,25 @@ async def _auto_discover_vector_field(self) -> None: self._use_vectorizable_query = False if not self.embedding_function: - import logging - - logging.warning( + 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 - import logging - - logging.warning( + logger.warning( f"Multiple vector fields found: {vector_fields}. " f"Please specify vector_field_name explicitly. Using keyword-only search." ) - # Close index client if we created it - if not self._index_client: - await index_client.close() - except Exception as e: # Log warning but continue with keyword search - import logging - - logging.warning(f"Failed to auto-discover vector field: {e}. Using keyword-only 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) -> str: + 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: @@ -611,7 +580,7 @@ async def _semantic_search(self, query: str) -> str: query: Search query text. Returns: - Formatted search results as string. + List of formatted search result strings, one per document. """ # Auto-discover vector field if not already done await self._auto_discover_vector_field() @@ -673,7 +642,7 @@ async def _semantic_search(self, query: str) -> str: if doc_text: formatted_results.append(doc_text) # type: ignore[reportUnknownArgumentType] - return "\n\n".join(formatted_results) + return formatted_results async def _ensure_knowledge_base(self) -> None: """Ensure Knowledge Base and knowledge source are created. @@ -761,7 +730,7 @@ async def _ensure_knowledge_base(self) -> None: credential=self.credential, ) - async def _agentic_search(self, messages: list[ChatMessage]) -> str: + async def _agentic_search(self, messages: list[ChatMessage]) -> list[str]: """Perform agentic retrieval with multi-hop reasoning using Knowledge Bases. NOTE: This mode is significantly slower than semantic search and should @@ -774,10 +743,10 @@ async def _agentic_search(self, messages: list[ChatMessage]) -> str: 4. Synthesize a comprehensive answer with references Args: - messages: Conversation history (last 10 messages) to use for retrieval context. + messages: Conversation history to use for retrieval context. Returns: - Synthesized answer from the Knowledge Base. + List of answer parts from the Knowledge Base, one per content item. """ # Ensure Knowledge Base is initialized await self._ensure_knowledge_base() @@ -832,12 +801,12 @@ async def _agentic_search(self, messages: list[ChatMessage]) -> str: # Perform retrieval via Knowledge Base retrieval_result = await self._retrieval_client.retrieve(retrieval_request=retrieval_request) - # Extract synthesized answer from response + # 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: - # Combine all text 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 @@ -845,10 +814,10 @@ async def _agentic_search(self, messages: list[ChatMessage]) -> str: answer_parts.append(content_item.text) if answer_parts: - return "\n".join(answer_parts) + return answer_parts # Fallback if no answer generated - return "No results found from Knowledge Base." + 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. diff --git a/python/packages/aisearch/pyproject.toml b/python/packages/aisearch/pyproject.toml index 07720b4f6d..5d64b398aa 100644 --- a/python/packages/aisearch/pyproject.toml +++ b/python/packages/aisearch/pyproject.toml @@ -24,7 +24,6 @@ classifiers = [ ] dependencies = [ "agent-framework-core", - "azure-identity>=1,<2", "azure-search-documents==11.7.0b2", ] diff --git a/python/packages/aisearch/tests/test_search_provider.py b/python/packages/aisearch/tests/test_search_provider.py index e17ec89db8..0b074f9d44 100644 --- a/python/packages/aisearch/tests/test_search_provider.py +++ b/python/packages/aisearch/tests/test_search_provider.py @@ -186,8 +186,11 @@ async def test_semantic_search_basic( context = await provider.invoking(sample_messages) assert isinstance(context, Context) - assert len(context.messages) > 0 - assert "Test document content" in context.messages[0].text + 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") @@ -425,9 +428,10 @@ async def test_citations_included_in_semantic_search(self, mock_search_class: Ma # Check that citation is included assert isinstance(context, Context) - assert len(context.messages) > 0 - assert "[Source: doc123]" in context.messages[0].text - assert "Test document content" in context.messages[0].text + 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 TestVectorFieldAutoDiscovery: diff --git a/python/packages/core/pyproject.toml b/python/packages/core/pyproject.toml index 8f19719c59..5e4d16e039 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-copilotstudio", diff --git a/python/pyproject.toml b/python/pyproject.toml index 92a7d525ed..7dd43dbeff 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -26,6 +26,7 @@ dependencies = [ "agent-framework-core", "agent-framework-a2a", "agent-framework-ag-ui", + "agent-framework-aisearch", "agent-framework-anthropic", "agent-framework-azure-ai", "agent-framework-azurefunctions", @@ -93,16 +94,17 @@ 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 } agent-framework-chatkit = { workspace = true } agent-framework-copilotstudio = { workspace = true } +agent-framework-devui = { workspace = true } agent-framework-lab = { workspace = true } agent-framework-mem0 = { workspace = true } -agent-framework-redis = { workspace = true } -agent-framework-devui = { workspace = true } agent-framework-purview = { workspace = true } -agent-framework-anthropic = { workspace = true } +agent-framework-redis = { workspace = true } [tool.ruff] line-length = 120 diff --git a/python/samples/getting_started/agents/azure_ai/README.md b/python/samples/getting_started/agents/azure_ai/README.md index 17d3e969ab..1d512149d1 100644 --- a/python/samples/getting_started/agents/azure_ai/README.md +++ b/python/samples/getting_started/agents/azure_ai/README.md @@ -20,7 +20,9 @@ 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.py`](azure_ai_with_search_context.py) | Shows how to use AzureAISearchContextProvider for RAG with hybrid search and semantic ranking. Demonstrates semantic (fast) and agentic (slower, uses Knowledge Bases for complex multi-hop reasoning) retrieval modes. Provides context from Azure AI Search to enhance agent responses. | +| [`azure_ai_with_search_context.py`](azure_ai_with_search_context.py) | Shows how to use AzureAISearchContextProvider for RAG with hybrid search and semantic ranking. Demonstrates both semantic and agentic retrieval modes using an environment variable. | +| [`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 most use cases. | +| [`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. Best for complex multi-hop reasoning queries. | | [`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..aecd1f5b06 --- /dev/null +++ b/python/samples/getting_started/agents/azure_ai/azure_ai_with_search_context_agentic.py @@ -0,0 +1,117 @@ +# Copyright (c) Microsoft. All rights reserved. + +import asyncio +import os + +from agent_framework import ChatAgent +from agent_framework_aisearch import AzureAISearchContextProvider +from agent_framework_azure_ai import AzureAIAgentClient +from azure.core.credentials import AzureKeyCredential +from azure.identity.aio import DefaultAzureCredential + +""" +This sample demonstrates how to use Azure AI Search with agentic mode for advanced RAG +(Retrieval Augmented Generation) with Azure AI agents. + +**Agentic mode** is an advanced mode for complex scenarios: +- Uses Knowledge Bases in Azure AI Search +- Performs multi-hop reasoning across documents +- Uses an LLM to synthesize answers +- Best for complex queries requiring cross-document reasoning +- **Significantly slower** (order of magnitude) than semantic mode + +⚠️ Only use agentic mode when you need multi-hop reasoning across documents. + For most RAG use cases, 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 credential + search_credential = AzureKeyCredential(search_key) if search_key else DefaultAzureCredential() + + # Create Azure AI Search context provider with agentic mode (slower, multi-hop reasoning) + print("Using AGENTIC mode (Knowledge Bases with multi-hop reasoning, slower)\n") + print("⚠️ This mode is significantly slower than semantic mode.\n") + search_provider = AzureAISearchContextProvider( + endpoint=search_endpoint, + index_name=index_name, + credential=search_credential, + 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..80ecdb1718 --- /dev/null +++ b/python/samples/getting_started/agents/azure_ai/azure_ai_with_search_context_semantic.py @@ -0,0 +1,97 @@ +# Copyright (c) Microsoft. All rights reserved. + +import asyncio +import os + +from agent_framework import ChatAgent +from agent_framework_aisearch import AzureAISearchContextProvider +from agent_framework_azure_ai import AzureAIAgentClient +from azure.core.credentials import AzureKeyCredential +from azure.identity.aio import DefaultAzureCredential + +""" +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 credential + search_credential = AzureKeyCredential(search_key) if search_key else DefaultAzureCredential() + + # 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, + credential=search_credential, + 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 862b585122..eeff2584b7 100644 --- a/python/uv.lock +++ b/python/uv.lock @@ -88,6 +88,7 @@ source = { virtual = "." } dependencies = [ { 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'" }, @@ -133,6 +134,7 @@ docs = [ requires-dist = [ { name = "agent-framework-a2a", editable = "packages/a2a" }, { name = "agent-framework-ag-ui", editable = "packages/ag-ui" }, + { name = "agent-framework-aisearch", editable = "packages/aisearch" }, { name = "agent-framework-anthropic", editable = "packages/anthropic" }, { name = "agent-framework-azure-ai", editable = "packages/azure-ai" }, { name = "agent-framework-azurefunctions", editable = "packages/azurefunctions" }, @@ -225,14 +227,12 @@ 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-identity", 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-identity", specifier = ">=1,<2" }, { name = "azure-search-documents", specifier = "==11.7.0b2" }, ] @@ -339,6 +339,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-copilotstudio", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, @@ -352,6 +353,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-copilotstudio", marker = "extra == 'all'", editable = "packages/copilotstudio" }, @@ -1771,7 +1773,7 @@ name = "exceptiongroup" version = "1.3.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "typing-extensions", marker = "(python_full_version < '3.13' and sys_platform == 'darwin') or (python_full_version < '3.13' and sys_platform == 'linux') or (python_full_version < '3.13' and sys_platform == 'win32')" }, + { name = "typing-extensions", marker = "(python_full_version < '3.11' and sys_platform == 'darwin') or (python_full_version < '3.11' and sys_platform == 'linux') or (python_full_version < '3.11' and sys_platform == 'win32')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/0b/9f/a65090624ecf468cdca03533906e7c69ed7588582240cfe7cc9e770b50eb/exceptiongroup-1.3.0.tar.gz", hash = "sha256:b241f5885f560bc56a59ee63ca4c6a8bfa46ae4ad651af316d4e81817bb9fd88", size = 29749, upload-time = "2025-05-10T17:42:51.123Z" } wheels = [ From a2db26fd2888924df32137c69ef030626c94c7bf Mon Sep 17 00:00:00 2001 From: farzad528 Date: Thu, 20 Nov 2025 12:22:43 +0000 Subject: [PATCH 12/13] comments eduward --- .../_search_provider.py | 213 ++++++--- .../aisearch/tests/test_search_provider.py | 449 +++++++++++++++++- .../getting_started/agents/azure_ai/README.md | 5 +- .../azure_ai/azure_ai_with_search_context.py | 137 ------ .../azure_ai_with_search_context_agentic.py | 26 +- .../azure_ai_with_search_context_semantic.py | 5 + 6 files changed, 606 insertions(+), 229 deletions(-) delete mode 100644 python/samples/getting_started/agents/azure_ai/azure_ai_with_search_context.py diff --git a/python/packages/aisearch/agent_framework_aisearch/_search_provider.py b/python/packages/aisearch/agent_framework_aisearch/_search_provider.py index c633483c00..0013f34404 100644 --- a/python/packages/aisearch/agent_framework_aisearch/_search_provider.py +++ b/python/packages/aisearch/agent_framework_aisearch/_search_provider.py @@ -3,11 +3,12 @@ """Azure AI Search Context Provider for Agent Framework. This module provides context providers for Azure AI Search integration with two modes: -- Semantic: Fast hybrid search (vector + keyword) with semantic ranker -- Agentic: Slower multi-hop reasoning using Knowledge Bases for complex queries +- 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. -Use semantic mode for most cases. Use agentic mode only when you need multi-hop -reasoning across documents with Knowledge Bases. +See: https://techcommunity.microsoft.com/blog/azure-ai-foundry-blog/foundry-iq-boost-response-relevance-by-36-with-agentic-retrieval/4470720 """ import sys @@ -17,6 +18,7 @@ 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 @@ -41,6 +43,7 @@ VectorizableTextQuery, VectorizedQuery, ) +from pydantic import SecretStr, ValidationError # Type checking imports for optional agentic mode dependencies if TYPE_CHECKING: @@ -116,44 +119,44 @@ class AzureAISearchSettings(AFBaseSettings): """Settings for Azure AI Search Context Provider with auto-loading from environment. - Environment variables: - AZURE_SEARCH_ENDPOINT: Azure AI Search endpoint URL - AZURE_SEARCH_INDEX_NAME: Name of the search index - AZURE_SEARCH_API_KEY: API key for authentication - AZURE_AI_PROJECT_ENDPOINT: Azure AI Foundry project endpoint (for agentic mode) - AZURE_OPENAI_RESOURCE_URL: Azure OpenAI resource URL (for agentic mode) - AZURE_OPENAI_DEPLOYMENT_NAME: Model deployment name (for agentic mode) - AZURE_OPENAI_API_KEY: Azure OpenAI API key (for agentic mode with API key auth) + 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_" + env_prefix: ClassVar[str] = "AZURE_SEARCH_" endpoint: str | None = None - search_endpoint: str | None = None # Alias for endpoint index_name: str | None = None - search_index_name: str | None = None # Alias for index_name - api_key: str | None = None - search_api_key: str | None = None # Alias for api_key - ai_project_endpoint: str | None = None - openai_resource_url: str | None = None - openai_deployment_name: str | None = None - openai_api_key: str | None = None - - def validate_settings(self) -> None: - """Validate required settings are present.""" - # Use aliases if main fields not set - endpoint = self.endpoint or self.search_endpoint - index_name = self.index_name or self.search_index_name - - if not endpoint: - raise ValueError("endpoint is required (or set AZURE_SEARCH_ENDPOINT)") - if not index_name: - raise ValueError("index_name is required (or set AZURE_SEARCH_INDEX_NAME)") - - # Update main fields from aliases - self.endpoint = endpoint - self.index_name = index_name - self.api_key = self.api_key or self.search_api_key + api_key: SecretStr | None = None class AzureAISearchContextProvider(ContextProvider): @@ -162,24 +165,22 @@ class AzureAISearchContextProvider(ContextProvider): 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. Suitable for most RAG use cases. - - **agentic**: Slower multi-hop reasoning across documents using Knowledge Bases. - Use only for complex queries requiring cross-document reasoning. + with semantic reranking. Best for simple queries where speed is critical. Examples: - Using Settings class with environment variables: + Using environment variables (recommended): .. code-block:: python - from agent_framework import ChatAgent - from agent_framework_azure_ai import AzureAIAgentClient - from agent_framework.azure import AzureAISearchContextProvider, AzureAISearchSettings + from agent_framework_aisearch import AzureAISearchContextProvider from azure.identity.aio import DefaultAzureCredential - # Load from environment variables - settings = AzureAISearchSettings() - search_provider = AzureAISearchContextProvider(settings=settings, credential=DefaultAzureCredential()) + # Set AZURE_SEARCH_ENDPOINT and AZURE_SEARCH_INDEX_NAME in environment + search_provider = AzureAISearchContextProvider(credential=DefaultAzureCredential()) Semantic hybrid search with API key: @@ -193,19 +194,30 @@ class AzureAISearchContextProvider(ContextProvider): 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 (slower) + # 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", # Multi-hop reasoning - azure_ai_project_endpoint="https://myproject.services.ai.azure.com", + mode="agentic", + azure_openai_resource_url="https://myresource.openai.azure.com", model_deployment_name="gpt-4o", - knowledge_base_name="my-knowledge-base", # Required for agentic mode + knowledge_base_name="my-knowledge-base", ) """ @@ -217,7 +229,6 @@ def __init__( index_name: str | None = None, credential: str | AzureKeyCredential | AsyncTokenCredential | None = None, *, - settings: AzureAISearchSettings | None = None, mode: Literal["semantic", "agentic"] = "semantic", top_k: int = 5, semantic_configuration_name: str | None = None, @@ -235,16 +246,19 @@ def __init__( 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. credential: Azure credential (API key string, AzureKeyCredential, or AsyncTokenCredential). If a string is provided, it will be converted to AzureKeyCredential. - settings: Settings object with auto-loaded configuration from environment. - If provided, endpoint/index_name/credential from settings will be used as defaults. + Can also be set via environment variable AZURE_SEARCH_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. @@ -260,16 +274,23 @@ def __init__( 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. - Required for agentic mode. Example: "https://myproject.services.ai.azure.com" - model_deployment_name: Model deployment name in the Azure AI project. + 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_resource_url: Azure OpenAI resource URL for Knowledge Base model calls. - Required for agentic mode. Example: "https://myresource.openai.azure.com" 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. @@ -281,28 +302,68 @@ def __init__( "medium": Moderate reasoning with some query decomposition. "low": Lower reasoning effort than medium. Default: "minimal". - agentic_message_history_count: Number of recent messages to send to the Knowledge Base for context - in agentic mode. Default: 10. + 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 from settings if provided - if settings: - settings.validate_settings() - endpoint = endpoint or settings.endpoint - index_name = index_name or settings.index_name - if not credential and settings.api_key: - credential = settings.api_key - azure_ai_project_endpoint = azure_ai_project_endpoint or settings.ai_project_endpoint - azure_openai_resource_url = azure_openai_resource_url or settings.openai_resource_url - model_deployment_name = model_deployment_name or settings.openai_deployment_name - azure_openai_api_key = azure_openai_api_key or settings.openai_api_key + # Load settings from environment/file + try: + settings = AzureAISearchSettings( + endpoint=endpoint, + index_name=index_name, + 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 + + # Use settings values, with explicit parameters taking precedence + endpoint = endpoint or settings.endpoint + index_name = index_name or settings.index_name + if not credential and settings.api_key: + credential = settings.api_key.get_secret_value() # Validate required parameters if not endpoint: - raise ValueError("endpoint is required") + raise ServiceInitializationError( + "Azure AI Search endpoint is required. Set via 'endpoint' parameter " + "or 'AZURE_SEARCH_ENDPOINT' environment variable." + ) if not index_name: - raise ValueError("index_name is required") + raise ServiceInitializationError( + "Azure AI Search index name is required. Set via 'index_name' parameter " + "or 'AZURE_SEARCH_INDEX_NAME' environment variable." + ) if not credential: - raise ValueError("credential is required") + raise ServiceInitializationError( + "Azure credential is required. Provide 'credential' parameter " + "or set 'AZURE_SEARCH_API_KEY' environment variable." + ) # Convert string credential to AzureKeyCredential if isinstance(credential, str): @@ -733,8 +794,8 @@ async def _ensure_knowledge_base(self) -> None: async def _agentic_search(self, messages: list[ChatMessage]) -> list[str]: """Perform agentic retrieval with multi-hop reasoning using Knowledge Bases. - NOTE: This mode is significantly slower than semantic search and should - only be used for complex queries requiring cross-document reasoning. + 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 diff --git a/python/packages/aisearch/tests/test_search_provider.py b/python/packages/aisearch/tests/test_search_provider.py index 0b074f9d44..a7fc4991d5 100644 --- a/python/packages/aisearch/tests/test_search_provider.py +++ b/python/packages/aisearch/tests/test_search_provider.py @@ -1,14 +1,18 @@ # 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: @@ -41,6 +45,84 @@ def sample_messages() -> list[ChatMessage]: ] +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", + credential=AzureKeyCredential("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", + credential=AzureKeyCredential("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", + credential=AzureKeyCredential("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.""" @@ -68,8 +150,8 @@ def test_init_semantic_mode_with_vector_field_requires_embedding_function(self) vector_field_name="embedding", ) - def test_init_agentic_mode_requires_parameters(self) -> None: - """Test that agentic mode requires additional parameters.""" + 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", @@ -339,6 +421,50 @@ async def test_context_manager(self, mock_search_class: MagicMock) -> None: 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", + credential=AzureKeyCredential("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_credential_conversion(self) -> None: + """Test that string credential is converted to AzureKeyCredential.""" + provider = AzureAISearchContextProvider( + endpoint="https://test.search.windows.net", + index_name="test-index", + credential="my-api-key", # String instead of AzureKeyCredential + mode="semantic", + ) + assert isinstance(provider.credential, AzureKeyCredential) + class TestMessageFiltering: """Test message filtering functionality.""" @@ -434,6 +560,172 @@ async def test_citations_included_in_semantic_search(self, mock_search_class: Ma 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", + credential=AzureKeyCredential("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", + credential=AzureKeyCredential("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", + credential=AzureKeyCredential("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.""" @@ -545,3 +837,156 @@ async def test_no_false_positives_on_string_fields( # 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", + credential=AzureKeyCredential("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", + credential=AzureKeyCredential("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", + credential=AzureKeyCredential("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/samples/getting_started/agents/azure_ai/README.md b/python/samples/getting_started/agents/azure_ai/README.md index 1d512149d1..7a3b2d4520 100644 --- a/python/samples/getting_started/agents/azure_ai/README.md +++ b/python/samples/getting_started/agents/azure_ai/README.md @@ -20,9 +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.py`](azure_ai_with_search_context.py) | Shows how to use AzureAISearchContextProvider for RAG with hybrid search and semantic ranking. Demonstrates both semantic and agentic retrieval modes using an environment variable. | -| [`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 most use cases. | -| [`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. Best for complex multi-hop reasoning queries. | +| [`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.py b/python/samples/getting_started/agents/azure_ai/azure_ai_with_search_context.py deleted file mode 100644 index c9f05dba2b..0000000000 --- a/python/samples/getting_started/agents/azure_ai/azure_ai_with_search_context.py +++ /dev/null @@ -1,137 +0,0 @@ -# Copyright (c) Microsoft. All rights reserved. - -import asyncio -import os - -from agent_framework import ChatAgent -from agent_framework.azure import AzureAISearchContextProvider -from agent_framework_azure_ai import AzureAIAgentClient -from azure.core.credentials import AzureKeyCredential -from azure.identity.aio import DefaultAzureCredential - -""" -The following sample demonstrates how to use Azure AI Search as a context provider -for RAG (Retrieval Augmented Generation) with Azure AI agents. - -AzureAISearchContextProvider supports two modes: - -1. **Semantic mode** (default, recommended): - - 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 - -2. **Agentic mode** (slower, advanced): - - Uses Knowledge Bases in Azure AI Search - - Performs multi-hop reasoning across documents - - Uses an LLM to synthesize answers - - Best for complex queries requiring cross-document reasoning - - Significantly slower (order of magnitude) - -Prerequisites: -1. An Azure AI Search service with a search index -2. An Azure AI Foundry project with a model deployment -3. Ensure the model deployment name exists in your Azure AI Foundry project -4. Set the following environment variables: - - For both modes: - - 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 authentication - - 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") - - Additional for agentic mode (Knowledge Bases): - - USE_AGENTIC_MODE: Set to "true" to use agentic retrieval - - AZURE_SEARCH_KNOWLEDGE_BASE_NAME: Your Knowledge Base name - - AZURE_OPENAI_ENDPOINT: Your Azure OpenAI endpoint (e.g., "https://myresource.openai.azure.com") - (This is different from AZURE_AI_PROJECT_ENDPOINT - Knowledge Base needs the OpenAI endpoint for model calls) -""" - -# 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 context provider.""" - - # 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") - - # Check if agentic mode is requested - use_agentic = os.environ.get("USE_AGENTIC_MODE", "false").lower() == "true" - - # Create credential - search_credential = AzureKeyCredential(search_key) if search_key else DefaultAzureCredential() - - # Create Azure AI Search context provider - if use_agentic: - # Agentic mode: Multi-hop reasoning with Knowledge Bases (slower) - print("Using AGENTIC mode (Knowledge Bases with multi-hop reasoning, slower)\n") - knowledge_base_name = os.environ["AZURE_SEARCH_KNOWLEDGE_BASE_NAME"] - azure_openai_endpoint = os.environ["AZURE_OPENAI_ENDPOINT"] - search_provider = AzureAISearchContextProvider( - endpoint=search_endpoint, - index_name=index_name, - credential=search_credential, - mode="agentic", - # Agentic mode uses Azure AI Foundry project for model inference - azure_ai_project_endpoint=project_endpoint, - model_deployment_name=model_deployment, - knowledge_base_name=knowledge_base_name, - azure_openai_resource_url=azure_openai_endpoint, - top_k=3, - ) - else: - # Semantic mode: Fast hybrid search + semantic ranking (recommended) - print("Using SEMANTIC mode (hybrid search + semantic ranking, fast)\n") - search_provider = AzureAISearchContextProvider( - endpoint=search_endpoint, - index_name=index_name, - credential=search_credential, - mode="semantic", - 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 ===\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_agentic.py b/python/samples/getting_started/agents/azure_ai/azure_ai_with_search_context_agentic.py index aecd1f5b06..f84b0fb75f 100644 --- 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 @@ -3,25 +3,29 @@ 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.core.credentials import AzureKeyCredential 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 advanced RAG +This sample demonstrates how to use Azure AI Search with agentic mode for RAG (Retrieval Augmented Generation) with Azure AI agents. -**Agentic mode** is an advanced mode for complex scenarios: -- Uses Knowledge Bases in Azure AI Search +**Agentic mode** is recommended for most scenarios: +- Uses Knowledge Bases in Azure AI Search for query planning - Performs multi-hop reasoning across documents -- Uses an LLM to synthesize answers -- Best for complex queries requiring cross-document reasoning -- **Significantly slower** (order of magnitude) than semantic mode +- 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 -⚠️ Only use agentic mode when you need multi-hop reasoning across documents. - For most RAG use cases, use semantic mode instead (see azure_ai_with_search_context_semantic.py). +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 @@ -61,9 +65,9 @@ async def main() -> None: # Create credential search_credential = AzureKeyCredential(search_key) if search_key else DefaultAzureCredential() - # Create Azure AI Search context provider with agentic mode (slower, multi-hop reasoning) - print("Using AGENTIC mode (Knowledge Bases with multi-hop reasoning, slower)\n") - print("⚠️ This mode is significantly slower than semantic mode.\n") + # 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, 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 index 80ecdb1718..4eafbb9f6a 100644 --- 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 @@ -3,12 +3,17 @@ 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.core.credentials import AzureKeyCredential 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. From 97facd7b65428c52c563f64d4dfce5bc70f6fbe0 Mon Sep 17 00:00:00 2001 From: farzad528 Date: Thu, 20 Nov 2025 14:03:45 +0000 Subject: [PATCH 13/13] ed latest comments --- .../_search_provider.py | 57 ++++++++-------- python/packages/aisearch/tests/__init__.py | 1 - .../aisearch/tests/test_search_provider.py | 66 +++++++++---------- .../azure_ai_with_search_context_agentic.py | 7 +- .../azure_ai_with_search_context_semantic.py | 7 +- 5 files changed, 67 insertions(+), 71 deletions(-) delete mode 100644 python/packages/aisearch/tests/__init__.py diff --git a/python/packages/aisearch/agent_framework_aisearch/_search_provider.py b/python/packages/aisearch/agent_framework_aisearch/_search_provider.py index 0013f34404..23c8c3f309 100644 --- a/python/packages/aisearch/agent_framework_aisearch/_search_provider.py +++ b/python/packages/aisearch/agent_framework_aisearch/_search_provider.py @@ -186,11 +186,11 @@ class AzureAISearchContextProvider(ContextProvider): .. code-block:: python - # Direct API key string (converted to AzureKeyCredential automatically) + # Direct API key string search_provider = AzureAISearchContextProvider( endpoint="https://mysearch.search.windows.net", index_name="my-index", - credential="my-api-key", # String converted to AzureKeyCredential + api_key="my-api-key", mode="semantic", ) @@ -227,7 +227,8 @@ def __init__( self, endpoint: str | None = None, index_name: str | None = None, - credential: str | AzureKeyCredential | AsyncTokenCredential | None = None, + api_key: str | AzureKeyCredential | None = None, + credential: AsyncTokenCredential | None = None, *, mode: Literal["semantic", "agentic"] = "semantic", top_k: int = 5, @@ -256,9 +257,10 @@ def __init__( 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. - credential: Azure credential (API key string, AzureKeyCredential, or AsyncTokenCredential). - If a string is provided, it will be converted to AzureKeyCredential. + 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. @@ -336,42 +338,43 @@ def __init__( 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 - # Use settings values, with explicit parameters taking precedence - endpoint = endpoint or settings.endpoint - index_name = index_name or settings.index_name - if not credential and settings.api_key: - credential = settings.api_key.get_secret_value() - # Validate required parameters - if not endpoint: + if not settings.endpoint: raise ServiceInitializationError( "Azure AI Search endpoint is required. Set via 'endpoint' parameter " "or 'AZURE_SEARCH_ENDPOINT' environment variable." ) - if not index_name: + 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." ) - if not credential: + + # 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 'credential' parameter " + "Azure credential is required. Provide 'api_key' or 'credential' parameter " "or set 'AZURE_SEARCH_API_KEY' environment variable." ) - # Convert string credential to AzureKeyCredential - if isinstance(credential, str): - credential = AzureKeyCredential(credential) - - self.endpoint = endpoint - self.index_name = index_name - self.credential = credential + 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 @@ -422,9 +425,9 @@ def __init__( # Create search client for semantic mode self._search_client = SearchClient( - endpoint=endpoint, - index_name=index_name, - credential=credential, + endpoint=self.endpoint, + index_name=self.index_name, + credential=self.credential, ) # Create index client and retrieval client for agentic mode (Knowledge Base) @@ -432,8 +435,8 @@ def __init__( self._retrieval_client: KnowledgeBaseRetrievalClient | None = None if mode == "agentic": self._index_client = SearchIndexClient( - endpoint=endpoint, - credential=credential, + endpoint=self.endpoint, + credential=self.credential, ) # Retrieval client will be created after Knowledge Base initialization diff --git a/python/packages/aisearch/tests/__init__.py b/python/packages/aisearch/tests/__init__.py deleted file mode 100644 index 2a50eae894..0000000000 --- a/python/packages/aisearch/tests/__init__.py +++ /dev/null @@ -1 +0,0 @@ -# Copyright (c) Microsoft. All rights reserved. diff --git a/python/packages/aisearch/tests/test_search_provider.py b/python/packages/aisearch/tests/test_search_provider.py index a7fc4991d5..6813c3d16a 100644 --- a/python/packages/aisearch/tests/test_search_provider.py +++ b/python/packages/aisearch/tests/test_search_provider.py @@ -75,7 +75,7 @@ def test_provider_uses_settings_from_env(self) -> None: provider = AzureAISearchContextProvider( endpoint="https://test.search.windows.net", index_name="test-index", - credential=AzureKeyCredential("test-key"), + api_key="test-key", ) assert provider.endpoint == "https://test.search.windows.net" assert provider.index_name == "test-index" @@ -90,7 +90,7 @@ def test_provider_missing_endpoint_raises_error(self) -> None: ): AzureAISearchContextProvider( index_name="test-index", - credential=AzureKeyCredential("test-key"), + api_key="test-key", env_file_path="", # Disable .env file loading ) @@ -104,7 +104,7 @@ def test_provider_missing_index_name_raises_error(self) -> None: ): AzureAISearchContextProvider( endpoint="https://test.search.windows.net", - credential=AzureKeyCredential("test-key"), + api_key="test-key", env_file_path="", # Disable .env file loading ) @@ -131,7 +131,7 @@ def test_init_semantic_mode_minimal(self) -> None: provider = AzureAISearchContextProvider( endpoint="https://test.search.windows.net", index_name="test-index", - credential=AzureKeyCredential("test-key"), + api_key="test-key", mode="semantic", ) assert provider.endpoint == "https://test.search.windows.net" @@ -145,7 +145,7 @@ def test_init_semantic_mode_with_vector_field_requires_embedding_function(self) AzureAISearchContextProvider( endpoint="https://test.search.windows.net", index_name="test-index", - credential=AzureKeyCredential("test-key"), + api_key="test-key", mode="semantic", vector_field_name="embedding", ) @@ -156,7 +156,7 @@ def test_init_agentic_mode_requires_azure_openai_resource_url(self) -> None: AzureAISearchContextProvider( endpoint="https://test.search.windows.net", index_name="test-index", - credential=AzureKeyCredential("test-key"), + api_key="test-key", mode="agentic", ) @@ -166,7 +166,7 @@ def test_init_agentic_mode_requires_model_deployment_name(self) -> None: AzureAISearchContextProvider( endpoint="https://test.search.windows.net", index_name="test-index", - credential=AzureKeyCredential("test-key"), + 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", @@ -178,7 +178,7 @@ def test_init_agentic_mode_requires_knowledge_base_name(self) -> None: AzureAISearchContextProvider( endpoint="https://test.search.windows.net", index_name="test-index", - credential=AzureKeyCredential("test-key"), + api_key="test-key", mode="agentic", azure_ai_project_endpoint="https://test.services.ai.azure.com", model_deployment_name="gpt-4o", @@ -190,7 +190,7 @@ def test_init_agentic_mode_with_all_params(self) -> None: provider = AzureAISearchContextProvider( endpoint="https://test.search.windows.net", index_name="test-index", - credential=AzureKeyCredential("test-key"), + api_key="test-key", mode="agentic", azure_ai_project_endpoint="https://test.services.ai.azure.com", model_deployment_name="my-gpt-4o-deployment", @@ -210,7 +210,7 @@ def test_init_model_name_defaults_to_deployment_name(self) -> None: provider = AzureAISearchContextProvider( endpoint="https://test.search.windows.net", index_name="test-index", - credential=AzureKeyCredential("test-key"), + api_key="test-key", mode="agentic", azure_ai_project_endpoint="https://test.services.ai.azure.com", model_deployment_name="gpt-4o", @@ -225,7 +225,7 @@ def test_init_with_custom_context_prompt(self) -> None: provider = AzureAISearchContextProvider( endpoint="https://test.search.windows.net", index_name="test-index", - credential=AzureKeyCredential("test-key"), + api_key="test-key", mode="semantic", context_prompt=custom_prompt, ) @@ -236,7 +236,7 @@ def test_init_uses_default_context_prompt(self) -> None: provider = AzureAISearchContextProvider( endpoint="https://test.search.windows.net", index_name="test-index", - credential=AzureKeyCredential("test-key"), + api_key="test-key", mode="semantic", ) assert provider.context_prompt == provider._DEFAULT_SEARCH_CONTEXT_PROMPT @@ -261,7 +261,7 @@ async def test_semantic_search_basic( provider = AzureAISearchContextProvider( endpoint="https://test.search.windows.net", index_name="test-index", - credential=AzureKeyCredential("test-key"), + api_key="test-key", mode="semantic", ) @@ -284,7 +284,7 @@ async def test_semantic_search_empty_query(self, mock_search_class: MagicMock) - provider = AzureAISearchContextProvider( endpoint="https://test.search.windows.net", index_name="test-index", - credential=AzureKeyCredential("test-key"), + api_key="test-key", mode="semantic", ) @@ -314,7 +314,7 @@ async def mock_embed(text: str) -> list[float]: provider = AzureAISearchContextProvider( endpoint="https://test.search.windows.net", index_name="test-index", - credential=AzureKeyCredential("test-key"), + api_key="test-key", mode="semantic", vector_field_name="embedding", embedding_function=mock_embed, @@ -352,7 +352,7 @@ async def test_ensure_knowledge_base_creates_when_not_exists( provider = AzureAISearchContextProvider( endpoint="https://test.search.windows.net", index_name="test-index", - credential=AzureKeyCredential("test-key"), + api_key="test-key", mode="agentic", azure_ai_project_endpoint="https://test.services.ai.azure.com", model_deployment_name="gpt-4o", @@ -387,7 +387,7 @@ async def test_ensure_knowledge_base_skips_when_exists( provider = AzureAISearchContextProvider( endpoint="https://test.search.windows.net", index_name="test-index", - credential=AzureKeyCredential("test-key"), + api_key="test-key", mode="agentic", azure_ai_project_endpoint="https://test.services.ai.azure.com", model_deployment_name="gpt-4o", @@ -415,7 +415,7 @@ async def test_context_manager(self, mock_search_class: MagicMock) -> None: async with AzureAISearchContextProvider( endpoint="https://test.search.windows.net", index_name="test-index", - credential=AzureKeyCredential("test-key"), + api_key="test-key", mode="semantic", ) as provider: assert provider is not None @@ -442,7 +442,7 @@ async def test_context_manager_agentic_cleanup( async with AzureAISearchContextProvider( endpoint="https://test.search.windows.net", index_name="test-index", - credential=AzureKeyCredential("test-key"), + api_key="test-key", mode="agentic", azure_ai_project_endpoint="https://test.services.ai.azure.com", model_deployment_name="gpt-4o", @@ -455,12 +455,12 @@ async def test_context_manager_agentic_cleanup( # Verify cleanup was called mock_retrieval_client.close.assert_called_once() - def test_string_credential_conversion(self) -> None: - """Test that string credential is converted to AzureKeyCredential.""" + 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", - credential="my-api-key", # String instead of AzureKeyCredential + api_key="my-api-key", # String api_key mode="semantic", ) assert isinstance(provider.credential, AzureKeyCredential) @@ -483,7 +483,7 @@ async def test_filters_non_user_assistant_messages(self, mock_search_class: Magi provider = AzureAISearchContextProvider( endpoint="https://test.search.windows.net", index_name="test-index", - credential=AzureKeyCredential("test-key"), + api_key="test-key", mode="semantic", ) @@ -511,7 +511,7 @@ async def test_filters_empty_messages(self, mock_search_class: MagicMock) -> Non provider = AzureAISearchContextProvider( endpoint="https://test.search.windows.net", index_name="test-index", - credential=AzureKeyCredential("test-key"), + api_key="test-key", mode="semantic", ) @@ -546,7 +546,7 @@ async def test_citations_included_in_semantic_search(self, mock_search_class: Ma provider = AzureAISearchContextProvider( endpoint="https://test.search.windows.net", index_name="test-index", - credential=AzureKeyCredential("test-key"), + api_key="test-key", mode="semantic", ) @@ -608,7 +608,7 @@ async def test_agentic_search_basic( provider = AzureAISearchContextProvider( endpoint="https://test.search.windows.net", index_name="test-index", - credential=AzureKeyCredential("test-key"), + api_key="test-key", mode="agentic", azure_ai_project_endpoint="https://test.services.ai.azure.com", model_deployment_name="gpt-4o", @@ -655,7 +655,7 @@ async def test_agentic_search_no_results( provider = AzureAISearchContextProvider( endpoint="https://test.search.windows.net", index_name="test-index", - credential=AzureKeyCredential("test-key"), + api_key="test-key", mode="agentic", azure_ai_project_endpoint="https://test.services.ai.azure.com", model_deployment_name="gpt-4o", @@ -711,7 +711,7 @@ async def test_agentic_search_with_medium_reasoning( provider = AzureAISearchContextProvider( endpoint="https://test.search.windows.net", index_name="test-index", - credential=AzureKeyCredential("test-key"), + api_key="test-key", mode="agentic", azure_ai_project_endpoint="https://test.services.ai.azure.com", model_deployment_name="gpt-4o", @@ -758,7 +758,7 @@ async def test_auto_discovers_single_vector_field( provider = AzureAISearchContextProvider( endpoint="https://test.search.windows.net", index_name="test-index", - credential=AzureKeyCredential("test-key"), + api_key="test-key", mode="semantic", ) @@ -827,7 +827,7 @@ async def test_no_false_positives_on_string_fields( provider = AzureAISearchContextProvider( endpoint="https://test.search.windows.net", index_name="test-index", - credential=AzureKeyCredential("test-key"), + api_key="test-key", mode="semantic", ) @@ -872,7 +872,7 @@ async def test_multiple_vector_fields_without_vectorizer( provider = AzureAISearchContextProvider( endpoint="https://test.search.windows.net", index_name="test-index", - credential=AzureKeyCredential("test-key"), + api_key="test-key", mode="semantic", ) @@ -929,7 +929,7 @@ async def test_multiple_vectorizable_fields( provider = AzureAISearchContextProvider( endpoint="https://test.search.windows.net", index_name="test-index", - credential=AzureKeyCredential("test-key"), + api_key="test-key", mode="semantic", ) @@ -979,7 +979,7 @@ async def test_single_vectorizable_field_detected( provider = AzureAISearchContextProvider( endpoint="https://test.search.windows.net", index_name="test-index", - credential=AzureKeyCredential("test-key"), + api_key="test-key", mode="semantic", ) 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 index f84b0fb75f..d8f6a9dc30 100644 --- 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 @@ -8,7 +8,6 @@ from agent_framework import ChatAgent from agent_framework_aisearch import AzureAISearchContextProvider from agent_framework_azure_ai import AzureAIAgentClient -from azure.core.credentials import AzureKeyCredential from azure.identity.aio import DefaultAzureCredential # Load environment variables from .env file @@ -62,16 +61,14 @@ async def main() -> None: knowledge_base_name = os.environ["AZURE_SEARCH_KNOWLEDGE_BASE_NAME"] azure_openai_resource_url = os.environ["AZURE_OPENAI_RESOURCE_URL"] - # Create credential - search_credential = AzureKeyCredential(search_key) if search_key else DefaultAzureCredential() - # 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, - credential=search_credential, + 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, 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 index 4eafbb9f6a..30504354f7 100644 --- 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 @@ -8,7 +8,6 @@ from agent_framework import ChatAgent from agent_framework_aisearch import AzureAISearchContextProvider from agent_framework_azure_ai import AzureAIAgentClient -from azure.core.credentials import AzureKeyCredential from azure.identity.aio import DefaultAzureCredential # Load environment variables from .env file @@ -53,15 +52,13 @@ async def main() -> None: project_endpoint = os.environ["AZURE_AI_PROJECT_ENDPOINT"] model_deployment = os.environ.get("AZURE_AI_MODEL_DEPLOYMENT_NAME", "gpt-4o") - # Create credential - search_credential = AzureKeyCredential(search_key) if search_key else DefaultAzureCredential() - # 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, - credential=search_credential, + 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 )