Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions python/.cspell.json
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@
"aiplatform",
"azuredocindex",
"azuredocs",
"azurefunctions",
"boto",
"contentvector",
"contoso",
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
Please install this package via pip:

```bash
pip install agent-framework-aisearch --pre
pip install agent-framework-azure-ai-search --pre
```

## Azure AI Search Integration
Expand Down
Original file line number Diff line number Diff line change
@@ -1,15 +1,5 @@
# 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
Expand Down Expand Up @@ -111,6 +101,18 @@
else:
from typing_extensions import override # type: ignore[import] # pragma: no cover

"""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
"""


# Module-level constants
logger = get_logger("agent_framework.azure")
_DEFAULT_AGENTIC_MESSAGE_HISTORY_COUNT = 10
Expand Down
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
[project]
name = "agent-framework-aisearch"
name = "agent-framework-azure-ai-search"
description = "Azure AI Search integration for Microsoft Agent Framework."
authors = [{ name = "Microsoft", email = "af-support@microsoft.com"}]
readme = "README.md"
Expand Down Expand Up @@ -76,15 +76,15 @@ disallow_incomplete_defs = true
disallow_untyped_decorators = true

[tool.bandit]
targets = ["agent_framework_aisearch"]
targets = ["agent_framework_azure_ai_search"]
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"
mypy = "mypy --config-file $POE_ROOT/pyproject.toml agent_framework_azure_ai_search"
test = "pytest --cov=agent_framework_azure_ai_search --cov-report=term-missing:skip-covered tests"

[build-system]
requires = ["flit-core >= 3.11,<4.0"]
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -6,13 +6,11 @@

import pytest
from agent_framework import ChatMessage, Context, Role
from agent_framework.azure import AzureAISearchContextProvider
from agent_framework.azure import AzureAISearchContextProvider, AzureAISearchSettings
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:
Expand Down Expand Up @@ -246,7 +244,7 @@ class TestSemanticSearch:
"""Test semantic search functionality."""

@pytest.mark.asyncio
@patch("agent_framework_aisearch._search_provider.SearchClient")
@patch("agent_framework_azure_ai_search._search_provider.SearchClient")
async def test_semantic_search_basic(
self, mock_search_class: MagicMock, sample_messages: list[ChatMessage]
) -> None:
Expand Down Expand Up @@ -275,7 +273,7 @@ async def test_semantic_search_basic(
assert "Test document content" in context.messages[1].text

@pytest.mark.asyncio
@patch("agent_framework_aisearch._search_provider.SearchClient")
@patch("agent_framework_azure_ai_search._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()
Expand All @@ -295,7 +293,7 @@ async def test_semantic_search_empty_query(self, mock_search_class: MagicMock) -
assert len(context.messages) == 0

@pytest.mark.asyncio
@patch("agent_framework_aisearch._search_provider.SearchClient")
@patch("agent_framework_azure_ai_search._search_provider.SearchClient")
async def test_semantic_search_with_vector_query(
self, mock_search_class: MagicMock, sample_messages: list[ChatMessage]
) -> None:
Expand Down Expand Up @@ -332,8 +330,8 @@ 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")
@patch("agent_framework_azure_ai_search._search_provider.SearchIndexClient")
@patch("agent_framework_azure_ai_search._search_provider.SearchClient")
async def test_ensure_knowledge_base_creates_when_not_exists(
self, mock_search_class: MagicMock, mock_index_class: MagicMock
) -> None:
Expand Down Expand Up @@ -369,8 +367,8 @@ async def test_ensure_knowledge_base_creates_when_not_exists(
mock_index_client.create_or_update_knowledge_base.assert_called_once()

@pytest.mark.asyncio
@patch("agent_framework_aisearch._search_provider.SearchIndexClient")
@patch("agent_framework_aisearch._search_provider.SearchClient")
@patch("agent_framework_azure_ai_search._search_provider.SearchIndexClient")
@patch("agent_framework_azure_ai_search._search_provider.SearchClient")
async def test_ensure_knowledge_base_skips_when_exists(
self, mock_search_class: MagicMock, mock_index_class: MagicMock
) -> None:
Expand Down Expand Up @@ -406,7 +404,7 @@ class TestContextProviderLifecycle:
"""Test context provider lifecycle methods."""

@pytest.mark.asyncio
@patch("agent_framework_aisearch._search_provider.SearchClient")
@patch("agent_framework_azure_ai_search._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()
Expand All @@ -422,9 +420,9 @@ async def test_context_manager(self, mock_search_class: MagicMock) -> 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")
@patch("agent_framework_azure_ai_search._search_provider.KnowledgeBaseRetrievalClient")
@patch("agent_framework_azure_ai_search._search_provider.SearchIndexClient")
@patch("agent_framework_azure_ai_search._search_provider.SearchClient")
async def test_context_manager_agentic_cleanup(
self, mock_search_class: MagicMock, mock_index_class: MagicMock, mock_retrieval_class: MagicMock
) -> None:
Expand Down Expand Up @@ -470,7 +468,7 @@ class TestMessageFiltering:
"""Test message filtering functionality."""

@pytest.mark.asyncio
@patch("agent_framework_aisearch._search_provider.SearchClient")
@patch("agent_framework_azure_ai_search._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
Expand Down Expand Up @@ -502,7 +500,7 @@ async def test_filters_non_user_assistant_messages(self, mock_search_class: Magi
mock_search_client.search.assert_called_once()

@pytest.mark.asyncio
@patch("agent_framework_aisearch._search_provider.SearchClient")
@patch("agent_framework_azure_ai_search._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()
Expand Down Expand Up @@ -532,7 +530,7 @@ class TestCitations:
"""Test citation functionality."""

@pytest.mark.asyncio
@patch("agent_framework_aisearch._search_provider.SearchClient")
@patch("agent_framework_azure_ai_search._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
Expand Down Expand Up @@ -564,9 +562,9 @@ 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")
@patch("agent_framework_azure_ai_search._search_provider.KnowledgeBaseRetrievalClient")
@patch("agent_framework_azure_ai_search._search_provider.SearchIndexClient")
@patch("agent_framework_azure_ai_search._search_provider.SearchClient")
async def test_agentic_search_basic(
self,
mock_search_class: MagicMock,
Expand All @@ -593,7 +591,7 @@ async def test_agentic_search_basic(
mock_content = MagicMock()
mock_content.text = "Agentic search result"
# Make it pass isinstance check
from agent_framework_aisearch._search_provider import _agentic_retrieval_available
from agent_framework_azure_ai_search._search_provider import _agentic_retrieval_available

if _agentic_retrieval_available:
from azure.search.documents.knowledgebases.models import KnowledgeBaseMessageTextContent
Expand Down Expand Up @@ -623,9 +621,9 @@ async def test_agentic_search_basic(
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")
@patch("agent_framework_azure_ai_search._search_provider.KnowledgeBaseRetrievalClient")
@patch("agent_framework_azure_ai_search._search_provider.SearchIndexClient")
@patch("agent_framework_azure_ai_search._search_provider.SearchClient")
async def test_agentic_search_no_results(
self,
mock_search_class: MagicMock,
Expand Down Expand Up @@ -670,9 +668,9 @@ async def test_agentic_search_no_results(
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")
@patch("agent_framework_azure_ai_search._search_provider.KnowledgeBaseRetrievalClient")
@patch("agent_framework_azure_ai_search._search_provider.SearchIndexClient")
@patch("agent_framework_azure_ai_search._search_provider.SearchClient")
async def test_agentic_search_with_medium_reasoning(
self,
mock_search_class: MagicMock,
Expand All @@ -696,7 +694,7 @@ async def test_agentic_search_with_medium_reasoning(
mock_message = MagicMock()
mock_content = MagicMock()
mock_content.text = "Medium reasoning result"
from agent_framework_aisearch._search_provider import _agentic_retrieval_available
from agent_framework_azure_ai_search._search_provider import _agentic_retrieval_available

if _agentic_retrieval_available:
from azure.search.documents.knowledgebases.models import KnowledgeBaseMessageTextContent
Expand Down Expand Up @@ -730,8 +728,8 @@ 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")
@patch("agent_framework_azure_ai_search._search_provider.SearchIndexClient")
@patch("agent_framework_azure_ai_search._search_provider.SearchClient")
async def test_auto_discovers_single_vector_field(
self, mock_search_class: MagicMock, mock_index_class: MagicMock
) -> None:
Expand Down Expand Up @@ -795,8 +793,8 @@ async def test_vector_detection_accuracy(self) -> None:
assert is_vector_3 is False

@pytest.mark.asyncio
@patch("agent_framework_aisearch._search_provider.SearchIndexClient")
@patch("agent_framework_aisearch._search_provider.SearchClient")
@patch("agent_framework_azure_ai_search._search_provider.SearchIndexClient")
@patch("agent_framework_azure_ai_search._search_provider.SearchClient")
async def test_no_false_positives_on_string_fields(
self, mock_search_class: MagicMock, mock_index_class: MagicMock
) -> None:
Expand Down Expand Up @@ -839,8 +837,8 @@ async def test_no_false_positives_on_string_fields(
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")
@patch("agent_framework_azure_ai_search._search_provider.SearchIndexClient")
@patch("agent_framework_azure_ai_search._search_provider.SearchClient")
async def test_multiple_vector_fields_without_vectorizer(
self, mock_search_class: MagicMock, mock_index_class: MagicMock
) -> None:
Expand Down Expand Up @@ -884,8 +882,8 @@ async def test_multiple_vector_fields_without_vectorizer(
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")
@patch("agent_framework_azure_ai_search._search_provider.SearchIndexClient")
@patch("agent_framework_azure_ai_search._search_provider.SearchClient")
async def test_multiple_vectorizable_fields(
self, mock_search_class: MagicMock, mock_index_class: MagicMock
) -> None:
Expand Down Expand Up @@ -941,8 +939,8 @@ async def test_multiple_vectorizable_fields(
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")
@patch("agent_framework_azure_ai_search._search_provider.SearchIndexClient")
@patch("agent_framework_azure_ai_search._search_provider.SearchClient")
async def test_single_vectorizable_field_detected(
self, mock_search_class: MagicMock, mock_index_class: MagicMock
) -> None:
Expand Down
Loading
Loading