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