Skip to content

Commit 0d9b574

Browse files
jay-thakurekzhu
andauthored
Add Azure AI Search tool implementation (#5844)
# Azure AI Search Tool Implementation This PR adds a new tool for Azure AI Search integration to autogen-ext, enabling agents to search and retrieve information from Azure AI Search indexes. ## Why Are These Changes Needed? AutoGen currently lacks native integration with Azure AI Search, which is a powerful enterprise search service that supports semantic, vector, and hybrid search capabilities. This integration enables agents to: 1. Retrieve relevant information from large document collections 2. Perform semantic search with AI-powered ranking 3. Execute vector similarity search using embeddings 4. Combine text and vector approaches for optimal results This tool complements existing retrieval capabilities and provides a seamless way to integrate with Azure's search infrastructure. ## Features - **Multiple Search Types**: Support for text, semantic, vector, and hybrid search - **Flexible Configuration**: Customizable search parameters and fields - **Robust Error Handling**: User-friendly error messages with actionable guidance - **Performance Optimizations**: Configurable caching and retry mechanisms - **Vector Search Support**: Built-in embedding generation with extensibility ## Usage Example ```python from autogen_ext.tools.azure import AzureAISearchTool from azure.core.credentials import AzureKeyCredential from autogen import AssistantAgent, UserProxyAgent # Create the search tool search_tool = AzureAISearchTool.load_component({ "provider": "autogen_ext.tools.azure.AzureAISearchTool", "config": { "name": "DocumentSearch", "description": "Search for information in the knowledge base", "endpoint": "https://your-service.search.windows.net", "index_name": "your-index", "credential": {"api_key": "your-api-key"}, "query_type": "semantic", "semantic_config_name": "default" } }) # Create an agent with the search tool assistant = AssistantAgent( "assistant", llm_config={"tools": [search_tool]} ) # Create a user proxy agent user_proxy = UserProxyAgent( "user_proxy", human_input_mode="TERMINATE", max_consecutive_auto_reply=10, code_execution_config={"work_dir": "coding"} ) # Start the conversation user_proxy.initiate_chat( assistant, message="What information do we have about quantum computing in our knowledge base?" ) ``` ## Testing - Added unit tests for all search types (text, semantic, vector, hybrid) - Added tests for error handling and cancellation - All tests pass locally ## Documentation - Added comprehensive docstrings with examples - Included warnings about placeholder embedding implementation - Added links to Azure AI Search documentation ## Related issue number Closes #5419 ## Checks - [x] I've included any doc changes needed for <https://microsoft.github.io/autogen/>. See <https://github.com/microsoft/autogen/blob/main/CONTRIBUTING.md> to build and test documentation locally. - [x] I've added tests (if relevant) corresponding to the changes introduced in this PR. - [x] I've made sure all auto checks have passed. --------- Co-authored-by: Eric Zhu <[email protected]>
1 parent d7f2b56 commit 0d9b574

File tree

9 files changed

+2684
-0
lines changed

9 files changed

+2684
-0
lines changed

python/packages/autogen-core/docs/src/reference/index.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,7 @@ python/autogen_ext.models.anthropic
5555
python/autogen_ext.models.semantic_kernel
5656
python/autogen_ext.models.ollama
5757
python/autogen_ext.models.llama_cpp
58+
python/autogen_ext.tools.azure
5859
python/autogen_ext.tools.code_execution
5960
python/autogen_ext.tools.graphrag
6061
python/autogen_ext.tools.http
Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
autogen\_ext.tools.azure
2+
========================
3+
4+
.. automodule:: autogen_ext.tools.azure
5+
:members:
6+
:undoc-members:
7+
:show-inheritance:
8+

python/packages/autogen-ext/pyproject.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@ azure = [
2525
"azure-ai-inference>=1.0.0b7",
2626
"azure-core",
2727
"azure-identity",
28+
"azure-search-documents>=11.4.0",
2829
]
2930
docker = ["docker~=7.0", "asyncio_atexit>=1.0.1"]
3031
ollama = ["ollama>=0.4.7", "tiktoken>=0.8.0"]
Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
from ._ai_search import (
2+
AzureAISearchTool,
3+
BaseAzureAISearchTool,
4+
SearchQuery,
5+
SearchResult,
6+
SearchResults,
7+
)
8+
from ._config import AzureAISearchConfig
9+
10+
__all__ = [
11+
"AzureAISearchTool",
12+
"BaseAzureAISearchTool",
13+
"SearchQuery",
14+
"SearchResult",
15+
"SearchResults",
16+
"AzureAISearchConfig",
17+
]

0 commit comments

Comments
 (0)