# Copyright (c) Microsoft. All rights reserved. # pyright: reportPrivateUsage=false import os from collections.abc import Iterator from contextlib import contextmanager from types import SimpleNamespace from typing import Any, cast from unittest.mock import AsyncMock, Mock, patch import pytest from agent_framework import Content, Message from agent_framework._sessions import AgentSession, SessionContext from agent_framework.exceptions import SettingNotFoundError from azure.core.credentials import AzureKeyCredential from agent_framework_azure_ai_search import _context_provider from agent_framework_azure_ai_search._context_provider import ( AzureAISearchContextProvider, KnowledgeBaseOutputModeLiteral, RetrievalReasoningEffortLiteral, ) # -- Helpers ------------------------------------------------------------------- @pytest.fixture(autouse=True) def clear_azure_search_env(monkeypatch: pytest.MonkeyPatch) -> None: """Keep tests isolated from ambient Azure Search environment variables.""" for key in ( "AZURE_SEARCH_ENDPOINT", "AZURE_SEARCH_INDEX_NAME", "AZURE_SEARCH_KNOWLEDGE_BASE_NAME", "AZURE_SEARCH_API_KEY", ): monkeypatch.delenv(key, raising=False) class MockSearchResults: """Async-iterable mock for Azure SearchClient.search() results.""" def __init__(self, docs: list[dict]): self._docs = docs self._index = 0 def __aiter__(self): return self async def __anext__(self): if self._index >= len(self._docs): raise StopAsyncIteration doc = self._docs[self._index] self._index += 1 return doc def _make_mock_index( fields: list[SimpleNamespace] | None = None, profiles: list[SimpleNamespace] | None = None, has_vector_search: bool = True, ) -> SimpleNamespace: """Create a mock search index with the given fields and vector search profiles.""" vector_search = None if has_vector_search: vector_search = SimpleNamespace(profiles=profiles or []) return SimpleNamespace(fields=fields or [], vector_search=vector_search) @pytest.fixture def mock_search_client() -> AsyncMock: """Create a mock SearchClient that returns one document.""" client = AsyncMock() async def _search(**kwargs): return MockSearchResults([{"id": "doc1", "content": "test document"}]) client.search = AsyncMock(side_effect=_search) return client @pytest.fixture def mock_search_client_empty() -> AsyncMock: """Create a mock SearchClient that returns no results.""" client = AsyncMock() async def _search(**kwargs): return MockSearchResults([]) client.search = AsyncMock(side_effect=_search) return client def _make_provider(**overrides: Any) -> AzureAISearchContextProvider: """Create a semantic-mode provider with mocked internals (skips auto-discovery).""" defaults: dict[str, Any] = { "source_id": AzureAISearchContextProvider.DEFAULT_SOURCE_ID, "endpoint": "https://test.search.windows.net", "index_name": "test-index", "api_key": "test-key", } defaults.update(overrides) provider = AzureAISearchContextProvider(**defaults) provider._auto_discovered_vector_field = True # skip auto-discovery return provider # -- Preview-feature stubs ---------------------------------------------------- # The stable/GA azure-search-documents SDK (api-version 2026-04-01) does not ship the # preview-only agentic symbols (output mode, low/medium reasoning effort, image content, # messages-based retrieval request). These stubs let the preview code paths be exercised # deterministically regardless of which SDK is installed. class _StubReasoningEffort: def __init__(self, *args: object, **kwargs: object) -> None: ... class _StubOutputMode: EXTRACTIVE_DATA = "extractiveData" ANSWER_SYNTHESIS = "answerSynthesis" class _StubRetrievalRequest: """Lenient stand-in for the preview KnowledgeBaseRetrievalRequest that accepts any kwargs.""" def __init__(self, **kwargs: object) -> None: self.__dict__.update(kwargs) @contextmanager def force_preview_features() -> Iterator[None]: """Force preview-only agentic features on (with lightweight stubs).""" with patch.multiple( _context_provider, _preview_agentic_features_available=True, KBRetrievalMinimalReasoningEffort=_StubReasoningEffort, KBRetrievalMediumReasoningEffort=_StubReasoningEffort, KBRetrievalLowReasoningEffort=_StubReasoningEffort, KBRetrievalOutputMode=_StubOutputMode, KnowledgeBaseRetrievalRequest=_StubRetrievalRequest, ): yield # -- Initialization: semantic mode --------------------------------------------- class TestInitSemantic: """Initialization tests for semantic mode.""" def test_valid_init(self) -> None: provider = _make_provider() assert provider.source_id == AzureAISearchContextProvider.DEFAULT_SOURCE_ID assert provider.endpoint == "https://test.search.windows.net" assert provider.index_name == "test-index" assert provider.mode == "semantic" def test_source_id_set(self) -> None: provider = _make_provider(source_id="my-source") assert provider.source_id == "my-source" def test_missing_endpoint_raises(self) -> None: with patch.dict(os.environ, {}, clear=True), pytest.raises(SettingNotFoundError, match="endpoint"): AzureAISearchContextProvider( source_id="s", endpoint=None, index_name="idx", api_key="key", ) def test_missing_index_name_semantic_raises(self) -> None: with pytest.raises(SettingNotFoundError, match="index_name"): AzureAISearchContextProvider( source_id="s", endpoint="https://test.search.windows.net", index_name=None, api_key="key", ) def test_env_variable_fallback(self) -> None: env = { "AZURE_SEARCH_ENDPOINT": "https://env.search.windows.net", "AZURE_SEARCH_INDEX_NAME": "env-index", "AZURE_SEARCH_API_KEY": "env-key", } with patch.dict(os.environ, env, clear=False): provider = AzureAISearchContextProvider(source_id="env-test") assert provider.endpoint == "https://env.search.windows.net" assert provider.index_name == "env-index" def test_top_k_and_semantic_config(self) -> None: provider = _make_provider(top_k=10, semantic_configuration_name="my-config") assert provider.top_k == 10 assert provider.semantic_configuration_name == "my-config" def test_default_context_prompt(self) -> None: provider = _make_provider() assert provider.context_prompt == AzureAISearchContextProvider._DEFAULT_SEARCH_CONTEXT_PROMPT def test_custom_context_prompt(self) -> None: provider = _make_provider(context_prompt="Custom prompt:") assert provider.context_prompt == "Custom prompt:" def test_model_is_stored(self) -> None: """Model is stored on the provider for Azure OpenAI vectorization.""" provider = _make_provider(model="my-deploy") assert provider.azure_openai_model == "my-deploy" def test_model_explicit(self) -> None: provider = _make_provider(model="gpt-4") assert provider.azure_openai_model == "gpt-4" # -- Initialization: credential resolution ------------------------------------ class TestInitCredentialResolution: """Tests for credential resolution paths.""" def test_token_credential_used(self) -> None: mock_cred = AsyncMock() provider = AzureAISearchContextProvider( endpoint="https://test.search.windows.net", index_name="idx", credential=mock_cred, ) provider._auto_discovered_vector_field = True assert provider.credential is mock_cred def test_azure_key_credential_passed_through(self) -> None: akc = AzureKeyCredential("my-key") provider = AzureAISearchContextProvider( endpoint="https://test.search.windows.net", index_name="idx", api_key=akc, ) provider._auto_discovered_vector_field = True assert provider.credential is akc def test_no_credential_raises(self) -> None: with pytest.raises(ValueError, match="Azure credential is required"): AzureAISearchContextProvider( endpoint="https://test.search.windows.net", index_name="idx", ) # -- Initialization: agentic mode validation ----------------------------------- class TestInitAgenticValidation: """Initialization validation tests for agentic mode.""" def test_both_index_and_kb_raises(self) -> None: with pytest.raises(SettingNotFoundError, match="multiple were set"): cast(Any, AzureAISearchContextProvider)( source_id="s", endpoint="https://test.search.windows.net", index_name="idx", knowledge_base_name="kb", api_key="key", mode="agentic", model="deploy", azure_openai_resource_url="https://aoai.openai.azure.com", ) def test_neither_index_nor_kb_raises(self) -> None: with pytest.raises(SettingNotFoundError, match="none was set"): AzureAISearchContextProvider( source_id="s", endpoint="https://test.search.windows.net", api_key="key", mode="agentic", ) def test_missing_model_raises(self) -> None: with pytest.raises(ValueError, match="model"): cast(Any, AzureAISearchContextProvider)( source_id="s", endpoint="https://test.search.windows.net", index_name="idx", api_key="key", mode="agentic", azure_openai_resource_url="https://aoai.openai.azure.com", ) def test_vector_field_without_embedding_raises(self) -> None: with pytest.raises(ValueError, match="embedding_function"): AzureAISearchContextProvider( source_id="s", endpoint="https://test.search.windows.net", index_name="idx", api_key="key", vector_field_name="embedding", ) def test_agentic_missing_aoai_url_with_index_raises(self) -> None: with pytest.raises(ValueError, match="azure_openai_resource_url"): cast(Any, AzureAISearchContextProvider)( source_id="s", endpoint="https://test.search.windows.net", index_name="idx", api_key="key", mode="agentic", model="deploy", ) def test_agentic_with_kb_name_sets_use_existing(self) -> None: provider = AzureAISearchContextProvider( source_id="s", endpoint="https://test.search.windows.net", knowledge_base_name="my-kb", api_key="key", mode="agentic", ) assert provider._use_existing_knowledge_base is True assert provider.knowledge_base_name == "my-kb" def test_agentic_with_index_generates_kb_name(self) -> None: provider = AzureAISearchContextProvider( source_id="s", endpoint="https://test.search.windows.net", index_name="idx", api_key="key", mode="agentic", model="deploy", azure_openai_resource_url="https://aoai.openai.azure.com", ) assert provider._use_existing_knowledge_base is False assert provider.knowledge_base_name == "idx-kb" def test_agentic_explicit_kb_ignores_env_index_name(self) -> None: with patch.dict(os.environ, {"AZURE_SEARCH_INDEX_NAME": "env-index"}, clear=False): provider = AzureAISearchContextProvider( source_id="s", endpoint="https://test.search.windows.net", knowledge_base_name="my-kb", api_key="key", mode="agentic", ) assert provider.index_name is None assert provider.knowledge_base_name == "my-kb" assert provider._use_existing_knowledge_base is True assert provider._search_client is None def test_agentic_explicit_index_ignores_env_kb_name(self) -> None: with patch.dict(os.environ, {"AZURE_SEARCH_KNOWLEDGE_BASE_NAME": "env-kb"}, clear=False): provider = AzureAISearchContextProvider( source_id="s", endpoint="https://test.search.windows.net", index_name="idx", api_key="key", mode="agentic", model="deploy", azure_openai_resource_url="https://aoai.openai.azure.com", ) assert provider.index_name == "idx" assert provider.knowledge_base_name == "idx-kb" assert provider._use_existing_knowledge_base is False # -- client construction + stable/preview feature gating ---------------------- class TestClientConstruction: """The provider must not pin an api-version; the installed SDK picks its own default.""" def test_common_client_kwargs_has_no_api_version(self) -> None: provider = _make_provider() kwargs = provider._common_client_kwargs() assert "user_agent" in kwargs assert "api_version" not in kwargs def test_search_client_built_without_api_version(self) -> None: with patch("agent_framework_azure_ai_search._context_provider.SearchClient") as mock_sc: _make_provider() _, kwargs = mock_sc.call_args assert "api_version" not in kwargs def test_index_client_built_without_api_version_agentic(self) -> None: with patch("agent_framework_azure_ai_search._context_provider.SearchIndexClient") as mock_ic: AzureAISearchContextProvider( endpoint="https://test.search.windows.net", knowledge_base_name="kb", api_key="key", mode="agentic", ) _, kwargs = mock_ic.call_args assert "api_version" not in kwargs class TestPreviewFeatureGating: """Auto-detect gating: preview-only agentic options require the preview (prerelease) build.""" def _agentic( self, *, knowledge_base_output_mode: KnowledgeBaseOutputModeLiteral = "extractive_data", retrieval_reasoning_effort: RetrievalReasoningEffortLiteral = "minimal", ) -> AzureAISearchContextProvider: return AzureAISearchContextProvider( endpoint="https://test.search.windows.net", knowledge_base_name="kb", api_key="key", mode="agentic", knowledge_base_output_mode=knowledge_base_output_mode, retrieval_reasoning_effort=retrieval_reasoning_effort, ) def test_answer_synthesis_rejected_without_preview_sdk(self) -> None: with ( patch.object(_context_provider, "_preview_agentic_features_available", False), pytest.raises(ValueError, match="answer_synthesis"), ): self._agentic(knowledge_base_output_mode="answer_synthesis") def test_medium_effort_rejected_without_preview_sdk(self) -> None: with ( patch.object(_context_provider, "_preview_agentic_features_available", False), pytest.raises(ValueError, match="reasoning_effort"), ): self._agentic(retrieval_reasoning_effort="medium") def test_low_effort_rejected_without_preview_sdk(self) -> None: with ( patch.object(_context_provider, "_preview_agentic_features_available", False), pytest.raises(ValueError, match="reasoning_effort"), ): self._agentic(retrieval_reasoning_effort="low") def test_defaults_allowed_without_preview_sdk(self) -> None: with patch.object(_context_provider, "_preview_agentic_features_available", False): provider = self._agentic() assert provider.knowledge_base_output_mode == "extractive_data" assert provider.retrieval_reasoning_effort == "minimal" def test_preview_options_allowed_with_preview_sdk(self) -> None: with patch.object(_context_provider, "_preview_agentic_features_available", True): provider = self._agentic( knowledge_base_output_mode="answer_synthesis", retrieval_reasoning_effort="medium", ) assert provider.knowledge_base_output_mode == "answer_synthesis" assert provider.retrieval_reasoning_effort == "medium" async def test_kb_creation_omits_preview_fields_on_stable_sdk(self) -> None: provider = _make_provider() provider._knowledge_base_initialized = False provider._use_existing_knowledge_base = False provider.knowledge_base_name = "test-kb" provider.azure_openai_resource_url = "https://aoai.openai.azure.com" provider.azure_openai_model = "gpt-4" provider.index_name = "test-index" captured: dict[str, object] = {} async def _capture(kb: object) -> None: captured["kb"] = kb mock_index_client = AsyncMock() mock_index_client.get_knowledge_source.return_value = Mock() mock_index_client.create_or_update_knowledge_base = AsyncMock(side_effect=_capture) provider._index_client = mock_index_client with ( patch.object(_context_provider, "_preview_agentic_features_available", False), patch("agent_framework_azure_ai_search._context_provider.KnowledgeBaseRetrievalClient") as mock_cls, ): mock_cls.return_value = AsyncMock() await provider._ensure_knowledge_base() kb = captured["kb"] assert getattr(kb, "output_mode", None) is None assert getattr(kb, "retrieval_reasoning_effort", None) is None # -- __aenter__ / __aexit__ --------------------------------------------------- class TestAsyncContextManager: """Tests for async context manager.""" async def test_aenter_returns_self(self) -> None: provider = _make_provider() result = await provider.__aenter__() assert result is provider async def test_closes_retrieval_client(self) -> None: provider = _make_provider() mock_retrieval = AsyncMock() provider._retrieval_client = mock_retrieval await provider.__aexit__(None, None, None) mock_retrieval.close.assert_awaited_once() assert provider._retrieval_client is None async def test_no_retrieval_client_no_error(self) -> None: provider = _make_provider() assert provider._retrieval_client is None await provider.__aexit__(None, None, None) # should not raise # -- before_run: semantic mode ------------------------------------------------- class TestBeforeRunSemantic: """Tests for before_run in semantic mode.""" async def test_results_added_to_context(self, mock_search_client: AsyncMock) -> None: provider = _make_provider() provider._search_client = mock_search_client session = AgentSession(session_id="test-session") ctx = SessionContext( input_messages=[Message(role="user", contents=["test query"])], session_id="s1", ) await provider.before_run( agent=cast(Any, None), session=session, context=ctx, state=session.state.setdefault(provider.source_id, {}), ) # type: ignore[arg-type] mock_search_client.search.assert_awaited_once() msgs = ctx.context_messages.get(provider.source_id, []) assert len(msgs) >= 2 # context_prompt + at least one result assert msgs[0].text == provider.context_prompt async def test_empty_input_no_search(self, mock_search_client: AsyncMock) -> None: provider = _make_provider() provider._search_client = mock_search_client session = AgentSession(session_id="test-session") ctx = SessionContext(input_messages=[], session_id="s1") await provider.before_run( agent=cast(Any, None), session=session, context=ctx, state=session.state.setdefault(provider.source_id, {}), ) # type: ignore[arg-type] mock_search_client.search.assert_not_awaited() assert ctx.context_messages.get(provider.source_id) is None async def test_no_results_no_messages(self, mock_search_client_empty: AsyncMock) -> None: provider = _make_provider() provider._search_client = mock_search_client_empty session = AgentSession(session_id="test-session") ctx = SessionContext( input_messages=[Message(role="user", contents=["test query"])], session_id="s1", ) await provider.before_run( agent=cast(Any, None), session=session, context=ctx, state=session.state.setdefault(provider.source_id, {}), ) # type: ignore[arg-type] mock_search_client_empty.search.assert_awaited_once() assert ctx.context_messages.get(provider.source_id) is None async def test_context_prompt_prepended(self, mock_search_client: AsyncMock) -> None: custom_prompt = "Custom search context:" provider = _make_provider(context_prompt=custom_prompt) provider._search_client = mock_search_client session = AgentSession(session_id="test-session") ctx = SessionContext( input_messages=[Message(role="user", contents=["test query"])], session_id="s1", ) await provider.before_run( agent=cast(Any, None), session=session, context=ctx, state=session.state.setdefault(provider.source_id, {}), ) # type: ignore[arg-type] msgs = ctx.context_messages[provider.source_id] assert msgs[0].text == custom_prompt # -- before_run: message filtering --------------------------------------------- class TestBeforeRunFiltering: """Tests that only user/assistant messages are used for search.""" async def test_filters_non_user_assistant(self, mock_search_client: AsyncMock) -> None: provider = _make_provider() provider._search_client = mock_search_client session = AgentSession(session_id="test-session") ctx = SessionContext( input_messages=[ Message(role="system", contents=["system prompt"]), Message(role="user", contents=["actual question"]), ], session_id="s1", ) await provider.before_run( agent=cast(Any, None), session=session, context=ctx, state=session.state.setdefault(provider.source_id, {}), ) # type: ignore[arg-type] mock_search_client.search.assert_awaited_once() call_kwargs = mock_search_client.search.call_args[1] # The search text should contain only the user message, not the system message assert "actual question" in call_kwargs["search_text"] assert "system prompt" not in call_kwargs["search_text"] async def test_only_system_messages_no_search(self, mock_search_client: AsyncMock) -> None: provider = _make_provider() provider._search_client = mock_search_client session = AgentSession(session_id="test-session") ctx = SessionContext( input_messages=[Message(role="system", contents=["system prompt"])], session_id="s1", ) await provider.before_run( agent=cast(Any, None), session=session, context=ctx, state=session.state.setdefault(provider.source_id, {}), ) # type: ignore[arg-type] mock_search_client.search.assert_not_awaited() async def test_whitespace_only_messages_filtered(self, mock_search_client: AsyncMock) -> None: provider = _make_provider() provider._search_client = mock_search_client session = AgentSession(session_id="test-session") ctx = SessionContext( input_messages=[Message(role="user", contents=[" "])], session_id="s1", ) await provider.before_run( agent=cast(Any, None), session=session, context=ctx, state=session.state.setdefault(provider.source_id, {}), ) # type: ignore[arg-type] mock_search_client.search.assert_not_awaited() async def test_assistant_messages_included(self, mock_search_client: AsyncMock) -> None: provider = _make_provider() provider._search_client = mock_search_client session = AgentSession(session_id="test-session") ctx = SessionContext( input_messages=[ Message(role="user", contents=["first question"]), Message(role="assistant", contents=["first answer"]), Message(role="user", contents=["follow up"]), ], session_id="s1", ) await provider.before_run( agent=cast(Any, None), session=session, context=ctx, state=session.state.setdefault(provider.source_id, {}), ) # type: ignore[arg-type] call_kwargs = mock_search_client.search.call_args[1] assert "first question" in call_kwargs["search_text"] assert "first answer" in call_kwargs["search_text"] assert "follow up" in call_kwargs["search_text"] # -- _find_vector_fields ------------------------------------------------------- class TestFindVectorFields: """Tests for _find_vector_fields helper.""" def test_finds_fields_with_dimensions(self) -> None: provider = _make_provider() index = _make_mock_index( fields=[ SimpleNamespace(name="embedding", vector_search_dimensions=1536), SimpleNamespace(name="content", vector_search_dimensions=None), SimpleNamespace(name="title", vector_search_dimensions=0), ] ) result = provider._find_vector_fields(index) assert result == ["embedding"] def test_returns_empty_for_no_vector_fields(self) -> None: provider = _make_provider() index = _make_mock_index( fields=[ SimpleNamespace(name="content", vector_search_dimensions=None), SimpleNamespace(name="title", vector_search_dimensions=0), ] ) result = provider._find_vector_fields(index) assert result == [] def test_multiple_vector_fields(self) -> None: provider = _make_provider() index = _make_mock_index( fields=[ SimpleNamespace(name="emb1", vector_search_dimensions=768), SimpleNamespace(name="emb2", vector_search_dimensions=1536), ] ) result = provider._find_vector_fields(index) assert result == ["emb1", "emb2"] # -- _find_vectorizable_fields ------------------------------------------------ class TestFindVectorizableFields: """Tests for _find_vectorizable_fields helper.""" def test_finds_vectorizable_fields(self) -> None: provider = _make_provider() profiles = [SimpleNamespace(name="profile1", vectorizer_name="my-vectorizer")] fields = [ SimpleNamespace(name="embedding", vector_search_dimensions=1536, vector_search_profile_name="profile1"), ] index = _make_mock_index(fields=fields, profiles=profiles) result = provider._find_vectorizable_fields(index, ["embedding"]) assert result == ["embedding"] def test_returns_empty_when_no_vector_search(self) -> None: provider = _make_provider() index = _make_mock_index(has_vector_search=False) result = provider._find_vectorizable_fields(index, ["embedding"]) assert result == [] def test_returns_empty_when_no_profiles(self) -> None: provider = _make_provider() index = _make_mock_index(profiles=None) index.vector_search = SimpleNamespace(profiles=None) result = provider._find_vectorizable_fields(index, ["embedding"]) assert result == [] def test_field_not_in_vector_fields_excluded(self) -> None: provider = _make_provider() profiles = [SimpleNamespace(name="profile1", vectorizer_name="my-vectorizer")] fields = [ SimpleNamespace(name="other_field", vector_search_dimensions=1536, vector_search_profile_name="profile1"), ] index = _make_mock_index(fields=fields, profiles=profiles) result = provider._find_vectorizable_fields(index, ["embedding"]) assert result == [] def test_profile_without_vectorizer_not_included(self) -> None: provider = _make_provider() profiles = [SimpleNamespace(name="profile1", vectorizer_name=None)] fields = [ SimpleNamespace(name="embedding", vector_search_dimensions=1536, vector_search_profile_name="profile1"), ] index = _make_mock_index(fields=fields, profiles=profiles) result = provider._find_vectorizable_fields(index, ["embedding"]) assert result == [] def test_field_without_profile_name_excluded(self) -> None: provider = _make_provider() profiles = [SimpleNamespace(name="profile1", vectorizer_name="my-vectorizer")] fields = [ SimpleNamespace(name="embedding", vector_search_dimensions=1536, vector_search_profile_name=None), ] index = _make_mock_index(fields=fields, profiles=profiles) result = provider._find_vectorizable_fields(index, ["embedding"]) assert result == [] # -- _auto_discover_vector_field ----------------------------------------------- class TestAutoDiscoverVectorField: """Tests for _auto_discover_vector_field.""" async def test_skip_if_already_discovered(self) -> None: provider = _make_provider() provider._auto_discovered_vector_field = True await provider._auto_discover_vector_field() # No error, no side effects async def test_skip_if_vector_field_set(self) -> None: provider = _make_provider() provider._auto_discovered_vector_field = False provider.vector_field_name = "my_field" await provider._auto_discover_vector_field() # Should return immediately async def test_no_index_name_warns(self) -> None: provider = _make_provider() provider._auto_discovered_vector_field = False provider.index_name = None provider._index_client = AsyncMock() await provider._auto_discover_vector_field() assert provider._auto_discovered_vector_field is True async def test_no_vector_fields_sets_flag(self) -> None: provider = _make_provider() provider._auto_discovered_vector_field = False mock_index_client = AsyncMock() mock_index_client.get_index.return_value = _make_mock_index( fields=[SimpleNamespace(name="content", vector_search_dimensions=None)] ) provider._index_client = mock_index_client await provider._auto_discover_vector_field() assert provider._auto_discovered_vector_field is True assert provider.vector_field_name is None async def test_single_vectorizable_field_discovered(self) -> None: provider = _make_provider() provider._auto_discovered_vector_field = False profiles = [SimpleNamespace(name="profile1", vectorizer_name="my-vectorizer")] fields = [ SimpleNamespace(name="embedding", vector_search_dimensions=1536, vector_search_profile_name="profile1"), ] mock_index_client = AsyncMock() mock_index_client.get_index.return_value = _make_mock_index(fields=fields, profiles=profiles) provider._index_client = mock_index_client await provider._auto_discover_vector_field() assert provider.vector_field_name == "embedding" assert provider._use_vectorizable_query is True assert provider._auto_discovered_vector_field is True async def test_multiple_vectorizable_fields_warns(self) -> None: provider = _make_provider() provider._auto_discovered_vector_field = False profiles = [ SimpleNamespace(name="profile1", vectorizer_name="v1"), SimpleNamespace(name="profile2", vectorizer_name="v2"), ] fields = [ SimpleNamespace(name="emb1", vector_search_dimensions=768, vector_search_profile_name="profile1"), SimpleNamespace(name="emb2", vector_search_dimensions=1536, vector_search_profile_name="profile2"), ] mock_index_client = AsyncMock() mock_index_client.get_index.return_value = _make_mock_index(fields=fields, profiles=profiles) provider._index_client = mock_index_client await provider._auto_discover_vector_field() assert provider._auto_discovered_vector_field is True # vector_field_name should not be set when multiple found assert provider.vector_field_name is None async def test_single_vector_field_without_embedding_clears_field(self) -> None: provider = _make_provider() provider._auto_discovered_vector_field = False provider.embedding_function = None fields = [ SimpleNamespace(name="embedding", vector_search_dimensions=1536, vector_search_profile_name=None), ] mock_index_client = AsyncMock() mock_index_client.get_index.return_value = _make_mock_index(fields=fields, profiles=[]) provider._index_client = mock_index_client await provider._auto_discover_vector_field() assert provider._auto_discovered_vector_field is True assert provider.vector_field_name is None async def test_single_vector_field_with_embedding_function(self) -> None: provider = _make_provider() provider._auto_discovered_vector_field = False provider.embedding_function = AsyncMock(return_value=[0.1] * 1536) fields = [ SimpleNamespace(name="embedding", vector_search_dimensions=1536, vector_search_profile_name=None), ] mock_index_client = AsyncMock() mock_index_client.get_index.return_value = _make_mock_index(fields=fields, profiles=[]) provider._index_client = mock_index_client await provider._auto_discover_vector_field() assert provider.vector_field_name == "embedding" assert provider._use_vectorizable_query is False async def test_multiple_vector_fields_no_vectorizable_warns(self) -> None: provider = _make_provider() provider._auto_discovered_vector_field = False fields = [ SimpleNamespace(name="emb1", vector_search_dimensions=768, vector_search_profile_name=None), SimpleNamespace(name="emb2", vector_search_dimensions=1536, vector_search_profile_name=None), ] mock_index_client = AsyncMock() mock_index_client.get_index.return_value = _make_mock_index(fields=fields, profiles=[]) provider._index_client = mock_index_client await provider._auto_discover_vector_field() assert provider._auto_discovered_vector_field is True assert provider.vector_field_name is None async def test_exception_falls_back_to_keyword_search(self) -> None: provider = _make_provider() provider._auto_discovered_vector_field = False mock_index_client = AsyncMock() mock_index_client.get_index.side_effect = Exception("network error") provider._index_client = mock_index_client await provider._auto_discover_vector_field() assert provider._auto_discovered_vector_field is True async def test_creates_index_client_if_none(self) -> None: provider = _make_provider() provider._auto_discovered_vector_field = False provider._index_client = None with patch("agent_framework_azure_ai_search._context_provider.SearchIndexClient") as mock_cls: mock_client = AsyncMock() mock_client.get_index.return_value = _make_mock_index( fields=[SimpleNamespace(name="content", vector_search_dimensions=None)] ) mock_cls.return_value = mock_client await provider._auto_discover_vector_field() mock_cls.assert_called_once() assert provider._auto_discovered_vector_field is True # -- _semantic_search ---------------------------------------------------------- class TestSemanticSearch: """Tests for _semantic_search method.""" async def test_basic_keyword_search(self) -> None: provider = _make_provider() mock_client = AsyncMock() async def _search(**kwargs): return MockSearchResults([{"id": "d1", "content": "result text"}]) mock_client.search = AsyncMock(side_effect=_search) provider._search_client = mock_client results = await provider._semantic_search("test query") assert len(results) == 1 assert "result text" in results[0].text call_kwargs = mock_client.search.call_args[1] assert call_kwargs["search_text"] == "test query" async def test_vectorizable_text_query(self) -> None: provider = _make_provider() provider._use_vectorizable_query = True provider.vector_field_name = "embedding" mock_client = AsyncMock() async def _search(**kwargs): return MockSearchResults([{"id": "d1", "content": "vector result"}]) mock_client.search = AsyncMock(side_effect=_search) provider._search_client = mock_client results = await provider._semantic_search("vector query") assert len(results) == 1 call_kwargs = mock_client.search.call_args[1] assert "vector_queries" in call_kwargs assert len(call_kwargs["vector_queries"]) == 1 async def test_vectorized_query_with_embedding_function(self) -> None: provider = _make_provider() provider._use_vectorizable_query = False provider.vector_field_name = "embedding" async def _embed(query: str) -> list[float]: return [0.1, 0.2, 0.3] provider.embedding_function = _embed mock_client = AsyncMock() async def _search(**kwargs): return MockSearchResults([{"id": "d1", "content": "embed result"}]) mock_client.search = AsyncMock(side_effect=_search) provider._search_client = mock_client results = await provider._semantic_search("embed query") assert len(results) == 1 call_kwargs = mock_client.search.call_args[1] assert "vector_queries" in call_kwargs async def test_semantic_configuration_params(self) -> None: provider = _make_provider(semantic_configuration_name="my-semantic-config") mock_client = AsyncMock() async def _search(**kwargs): return MockSearchResults([{"id": "d1", "content": "semantic result"}]) mock_client.search = AsyncMock(side_effect=_search) provider._search_client = mock_client await provider._semantic_search("sem query") call_kwargs = mock_client.search.call_args[1] # Must be plain strings: in azure-search-documents 12.x these are non-str enums whose # str() form ("querytype.semantic") is rejected by the service. assert call_kwargs["query_type"] == "semantic" assert isinstance(call_kwargs["query_type"], str) assert call_kwargs["semantic_configuration_name"] == "my-semantic-config" assert call_kwargs["query_caption"] == "extractive" assert isinstance(call_kwargs["query_caption"], str) async def test_vector_k_with_semantic_config(self) -> None: provider = _make_provider(semantic_configuration_name="sc", top_k=3) provider._use_vectorizable_query = True provider.vector_field_name = "embedding" mock_client = AsyncMock() async def _search(**kwargs): return MockSearchResults([]) mock_client.search = AsyncMock(side_effect=_search) provider._search_client = mock_client await provider._semantic_search("query") call_kwargs = mock_client.search.call_args[1] assert "vector_queries" in call_kwargs assert len(call_kwargs["vector_queries"]) == 1 async def test_no_search_client_raises(self) -> None: provider = _make_provider() provider._search_client = None with pytest.raises(RuntimeError, match="Search client is not initialized"): await provider._semantic_search("query") async def test_empty_results_returns_empty_list(self) -> None: provider = _make_provider() mock_client = AsyncMock() async def _search(**kwargs): return MockSearchResults([]) mock_client.search = AsyncMock(side_effect=_search) provider._search_client = mock_client results = await provider._semantic_search("query") assert results == [] async def test_doc_without_text_excluded(self) -> None: provider = _make_provider() mock_client = AsyncMock() async def _search(**kwargs): # doc with only @search metadata and id - no extractable text return MockSearchResults([{"id": "d1", "@search.score": 0.9}]) mock_client.search = AsyncMock(side_effect=_search) provider._search_client = mock_client results = await provider._semantic_search("query") assert results == [] # -- _extract_document_text ---------------------------------------------------- class TestExtractDocumentText: """Tests for _extract_document_text.""" def test_content_field_extracted(self) -> None: provider = _make_provider() result = provider._extract_document_text({"content": "Hello world"}, doc_id="d1") assert result == "[Source: d1] Hello world" def test_text_field_extracted(self) -> None: provider = _make_provider() result = provider._extract_document_text({"text": "Some text"}, doc_id="d1") assert result == "[Source: d1] Some text" def test_description_field_extracted(self) -> None: provider = _make_provider() result = provider._extract_document_text({"description": "A description"}, doc_id="d1") assert result == "[Source: d1] A description" def test_body_field_extracted(self) -> None: provider = _make_provider() result = provider._extract_document_text({"body": "Body content"}, doc_id="d1") assert result == "[Source: d1] Body content" def test_chunk_field_extracted(self) -> None: provider = _make_provider() result = provider._extract_document_text({"chunk": "Chunk data"}, doc_id="d1") assert result == "[Source: d1] Chunk data" def test_content_field_priority(self) -> None: provider = _make_provider() result = provider._extract_document_text( {"content": "Primary", "text": "Secondary", "description": "Tertiary"}, doc_id="d1" ) assert result == "[Source: d1] Primary" def test_fallback_to_string_fields(self) -> None: provider = _make_provider() result = provider._extract_document_text( {"title": "My Title", "summary": "My Summary", "id": "skip-this", "@search.score": "skip-meta"}, doc_id="d1", ) assert "title: My Title" in result assert "summary: My Summary" in result assert "id" not in result.split("] ")[1] # id should be excluded from fallback assert "@search.score" not in result def test_empty_doc_returns_empty(self) -> None: provider = _make_provider() result = provider._extract_document_text({}) assert result == "" def test_no_doc_id_returns_text_only(self) -> None: provider = _make_provider() result = provider._extract_document_text({"content": "Hello"}, doc_id=None) assert result == "Hello" def test_search_id_fallback(self) -> None: """Test that doc results using @search.id work too (via before_run path).""" provider = _make_provider() result = provider._extract_document_text({"content": "data"}, doc_id="alt-id") assert result == "[Source: alt-id] data" def test_only_id_and_metadata_returns_empty(self) -> None: provider = _make_provider() result = provider._extract_document_text({"id": "d1", "@search.score": 0.9}) assert result == "" def test_non_string_values_excluded_from_fallback(self) -> None: provider = _make_provider() result = provider._extract_document_text({"count": 42, "tags": ["a", "b"]}, doc_id="d1") # Non-string values should not appear in fallback assert result == "" # -- _ensure_knowledge_base --------------------------------------------------- class TestEnsureKnowledgeBase: """Tests for _ensure_knowledge_base.""" async def test_already_initialized_returns_early(self) -> None: provider = _make_provider() provider._knowledge_base_initialized = True await provider._ensure_knowledge_base() # should not raise async def test_missing_kb_name_raises(self) -> None: provider = _make_provider() provider._knowledge_base_initialized = False provider.knowledge_base_name = None with pytest.raises(ValueError, match="knowledge_base_name is required"): await provider._ensure_knowledge_base() async def test_existing_kb_sets_initialized(self) -> None: provider = _make_provider() provider._knowledge_base_initialized = False provider._use_existing_knowledge_base = True provider.knowledge_base_name = "existing-kb" with patch("agent_framework_azure_ai_search._context_provider.KnowledgeBaseRetrievalClient") as mock_cls: mock_cls.return_value = AsyncMock() await provider._ensure_knowledge_base() assert provider._knowledge_base_initialized is True async def test_missing_index_client_raises(self) -> None: provider = _make_provider() provider._knowledge_base_initialized = False provider._use_existing_knowledge_base = False provider.knowledge_base_name = "test-kb" provider._index_client = None with pytest.raises(ValueError, match="Index client is required"): await provider._ensure_knowledge_base() async def test_missing_aoai_url_raises(self) -> None: provider = _make_provider() provider._knowledge_base_initialized = False provider._use_existing_knowledge_base = False provider.knowledge_base_name = "test-kb" provider._index_client = AsyncMock() provider.azure_openai_resource_url = None with pytest.raises(ValueError, match="azure_openai_resource_url is required"): await provider._ensure_knowledge_base() async def test_missing_deployment_name_raises(self) -> None: provider = _make_provider() provider._knowledge_base_initialized = False provider._use_existing_knowledge_base = False provider.knowledge_base_name = "test-kb" provider._index_client = AsyncMock() provider.azure_openai_resource_url = "https://aoai.openai.azure.com" provider.azure_openai_model = None with pytest.raises(ValueError, match="model is required"): await provider._ensure_knowledge_base() async def test_missing_index_name_raises(self) -> None: provider = _make_provider() provider._knowledge_base_initialized = False provider._use_existing_knowledge_base = False provider.knowledge_base_name = "test-kb" provider._index_client = AsyncMock() provider.azure_openai_resource_url = "https://aoai.openai.azure.com" provider.azure_openai_model = "deploy" provider.index_name = None with pytest.raises(ValueError, match="index_name is required"): await provider._ensure_knowledge_base() async def test_creates_knowledge_source_when_not_found(self) -> None: from azure.core.exceptions import ResourceNotFoundError provider = _make_provider() provider._knowledge_base_initialized = False provider._use_existing_knowledge_base = False provider.knowledge_base_name = "test-kb" provider.azure_openai_resource_url = "https://aoai.openai.azure.com" provider.azure_openai_model = "gpt-4" provider.index_name = "test-index" 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() provider._index_client = mock_index_client with patch("agent_framework_azure_ai_search._context_provider.KnowledgeBaseRetrievalClient") as mock_cls: mock_cls.return_value = AsyncMock() await provider._ensure_knowledge_base() mock_index_client.create_knowledge_source.assert_awaited_once() mock_index_client.create_or_update_knowledge_base.assert_awaited_once() assert provider._knowledge_base_initialized is True async def test_uses_existing_knowledge_source(self) -> None: provider = _make_provider() provider._knowledge_base_initialized = False provider._use_existing_knowledge_base = False provider.knowledge_base_name = "test-kb" provider.azure_openai_resource_url = "https://aoai.openai.azure.com" provider.azure_openai_model = "gpt-4" provider.index_name = "test-index" mock_index_client = AsyncMock() mock_index_client.get_knowledge_source.return_value = Mock() # source already exists mock_index_client.create_or_update_knowledge_base = AsyncMock() provider._index_client = mock_index_client with patch("agent_framework_azure_ai_search._context_provider.KnowledgeBaseRetrievalClient") as mock_cls: mock_cls.return_value = AsyncMock() await provider._ensure_knowledge_base() mock_index_client.create_knowledge_source.assert_not_awaited() mock_index_client.create_or_update_knowledge_base.assert_awaited_once() async def test_answer_synthesis_output_mode(self) -> None: provider = _make_provider() provider._knowledge_base_initialized = False provider._use_existing_knowledge_base = False provider.knowledge_base_name = "test-kb" provider.azure_openai_resource_url = "https://aoai.openai.azure.com" provider.azure_openai_model = "gpt-4" provider.index_name = "test-index" provider.knowledge_base_output_mode = "answer_synthesis" mock_index_client = AsyncMock() mock_index_client.get_knowledge_source.return_value = Mock() mock_index_client.create_or_update_knowledge_base = AsyncMock() provider._index_client = mock_index_client with patch("agent_framework_azure_ai_search._context_provider.KnowledgeBaseRetrievalClient") as mock_cls: mock_cls.return_value = AsyncMock() await provider._ensure_knowledge_base() assert provider._knowledge_base_initialized is True async def test_medium_reasoning_effort(self) -> None: provider = _make_provider() provider._knowledge_base_initialized = False provider._use_existing_knowledge_base = False provider.knowledge_base_name = "test-kb" provider.azure_openai_resource_url = "https://aoai.openai.azure.com" provider.azure_openai_model = "gpt-4" provider.index_name = "test-index" provider.retrieval_reasoning_effort = "medium" mock_index_client = AsyncMock() mock_index_client.get_knowledge_source.return_value = Mock() mock_index_client.create_or_update_knowledge_base = AsyncMock() provider._index_client = mock_index_client with patch("agent_framework_azure_ai_search._context_provider.KnowledgeBaseRetrievalClient") as mock_cls: mock_cls.return_value = AsyncMock() await provider._ensure_knowledge_base() assert provider._knowledge_base_initialized is True # -- _agentic_search ---------------------------------------------------------- class TestAgenticSearch: """Tests for _agentic_search.""" async def test_no_retrieval_client_raises(self) -> None: provider = _make_provider() provider._knowledge_base_initialized = True provider.knowledge_base_name = "kb" provider._retrieval_client = None with pytest.raises(RuntimeError, match="Retrieval client not initialized"): await provider._agentic_search([Message(role="user", contents=["query"])]) async def test_minimal_reasoning_returns_results(self) -> None: provider = _make_provider() provider._knowledge_base_initialized = True provider.knowledge_base_name = "kb" provider.retrieval_reasoning_effort = "minimal" mock_content = Mock() mock_content.text = "Answer text" mock_message = Mock() mock_message.role = "assistant" mock_message.content = [mock_content] mock_result = Mock() mock_result.response = [mock_message] mock_result.references = None mock_retrieval = AsyncMock() mock_retrieval.retrieve = AsyncMock(return_value=mock_result) provider._retrieval_client = mock_retrieval # Patch isinstance check for KnowledgeBaseMessageTextContent with patch( "agent_framework_azure_ai_search._context_provider.KnowledgeBaseMessageTextContent", type(mock_content), ): results = await provider._agentic_search([Message(role="user", contents=["test query"])]) assert len(results) == 1 assert results[0].text == "Answer text" assert results[0].role == "assistant" async def test_non_minimal_reasoning_uses_messages(self) -> None: provider = _make_provider() provider._knowledge_base_initialized = True provider.knowledge_base_name = "kb" provider.retrieval_reasoning_effort = "medium" mock_content = Mock() mock_content.text = "Medium answer" mock_message = Mock() mock_message.role = "assistant" mock_message.content = [mock_content] mock_result = Mock() mock_result.response = [mock_message] mock_result.references = None mock_retrieval = AsyncMock() mock_retrieval.retrieve = AsyncMock(return_value=mock_result) provider._retrieval_client = mock_retrieval with ( force_preview_features(), patch( "agent_framework_azure_ai_search._context_provider.KnowledgeBaseMessageTextContent", type(mock_content), ), ): results = await provider._agentic_search([ Message(role="user", contents=["question"]), Message(role="assistant", contents=["answer"]), ]) assert len(results) == 1 assert results[0].text == "Medium answer" mock_retrieval.retrieve.assert_awaited_once() async def test_no_response_returns_default_message(self) -> None: provider = _make_provider() provider._knowledge_base_initialized = True provider.knowledge_base_name = "kb" provider.retrieval_reasoning_effort = "minimal" mock_result = Mock() mock_result.response = [] mock_result.references = None mock_retrieval = AsyncMock() mock_retrieval.retrieve = AsyncMock(return_value=mock_result) provider._retrieval_client = mock_retrieval results = await provider._agentic_search([Message(role="user", contents=["query"])]) assert len(results) == 1 assert results[0].text == "No results found from Knowledge Base." async def test_empty_content_returns_default_message(self) -> None: provider = _make_provider() provider._knowledge_base_initialized = True provider.knowledge_base_name = "kb" provider.retrieval_reasoning_effort = "minimal" mock_message = Mock() mock_message.content = None mock_result = Mock() mock_result.response = [mock_message] mock_result.references = None mock_retrieval = AsyncMock() mock_retrieval.retrieve = AsyncMock(return_value=mock_result) provider._retrieval_client = mock_retrieval results = await provider._agentic_search([Message(role="user", contents=["query"])]) assert len(results) == 1 assert results[0].text == "No results found from Knowledge Base." async def test_answer_synthesis_output_mode(self) -> None: provider = _make_provider() provider._knowledge_base_initialized = True provider.knowledge_base_name = "kb" provider.retrieval_reasoning_effort = "low" provider.knowledge_base_output_mode = "answer_synthesis" mock_content = Mock() mock_content.text = "Synthesized answer" mock_message = Mock() mock_message.role = "assistant" mock_message.content = [mock_content] mock_result = Mock() mock_result.response = [mock_message] mock_result.references = None mock_retrieval = AsyncMock() mock_retrieval.retrieve = AsyncMock(return_value=mock_result) provider._retrieval_client = mock_retrieval with ( force_preview_features(), patch( "agent_framework_azure_ai_search._context_provider.KnowledgeBaseMessageTextContent", type(mock_content), ), ): results = await provider._agentic_search([Message(role="user", contents=["query"])]) assert len(results) == 1 assert results[0].text == "Synthesized answer" async def test_content_without_text_excluded(self) -> None: provider = _make_provider() provider._knowledge_base_initialized = True provider.knowledge_base_name = "kb" provider.retrieval_reasoning_effort = "minimal" mock_content_with_text = Mock() mock_content_with_text.text = "Good content" mock_content_no_text = Mock() mock_content_no_text.text = None mock_message = Mock() mock_message.role = "assistant" mock_message.content = [mock_content_no_text, mock_content_with_text] mock_result = Mock() mock_result.response = [mock_message] mock_result.references = None mock_retrieval = AsyncMock() mock_retrieval.retrieve = AsyncMock(return_value=mock_result) provider._retrieval_client = mock_retrieval with patch( "agent_framework_azure_ai_search._context_provider.KnowledgeBaseMessageTextContent", type(mock_content_with_text), ): results = await provider._agentic_search([Message(role="user", contents=["query"])]) assert len(results) == 1 assert results[0].text == "Good content" async def test_none_response_returns_default_message(self) -> None: provider = _make_provider() provider._knowledge_base_initialized = True provider.knowledge_base_name = "kb" provider.retrieval_reasoning_effort = "minimal" mock_result = Mock() mock_result.response = None mock_result.references = None mock_retrieval = AsyncMock() mock_retrieval.retrieve = AsyncMock(return_value=mock_result) provider._retrieval_client = mock_retrieval results = await provider._agentic_search([Message(role="user", contents=["query"])]) assert len(results) == 1 assert results[0].text == "No results found from Knowledge Base." async def test_requests_reference_source_data_per_source(self) -> None: # Regression for #5095: the retrieval request must carry # knowledge_source_params with include_reference_source_data=True for # each resolved knowledge source, otherwise ref.source_data is always # None even when the source has source_data_fields configured. provider = _make_provider() provider._knowledge_base_initialized = True provider.knowledge_base_name = "kb" provider.retrieval_reasoning_effort = "minimal" provider._knowledge_source_names = ["test-index-source"] mock_result = Mock() mock_result.response = [] mock_result.references = None mock_retrieval = AsyncMock() mock_retrieval.retrieve = AsyncMock(return_value=mock_result) provider._retrieval_client = mock_retrieval await provider._agentic_search([Message(role="user", contents=["query"])]) mock_retrieval.retrieve.assert_awaited_once() request = mock_retrieval.retrieve.call_args.kwargs["retrieval_request"] params = request.knowledge_source_params assert params is not None assert len(params) == 1 assert params[0].knowledge_source_name == "test-index-source" assert params[0].include_reference_source_data is True async def test_no_source_params_when_no_sources(self) -> None: # When no knowledge source names are resolved, the request must not # carry an empty/placeholder knowledge_source_params list. provider = _make_provider() provider._knowledge_base_initialized = True provider.knowledge_base_name = "kb" provider.retrieval_reasoning_effort = "minimal" provider._knowledge_source_names = [] mock_result = Mock() mock_result.response = [] mock_result.references = None mock_retrieval = AsyncMock() mock_retrieval.retrieve = AsyncMock(return_value=mock_result) provider._retrieval_client = mock_retrieval await provider._agentic_search([Message(role="user", contents=["query"])]) request = mock_retrieval.retrieve.call_args.kwargs["retrieval_request"] assert request.knowledge_source_params is None # -- before_run: agentic mode -------------------------------------------------- # -- _prepare_messages_for_kb_search / _parse_content_from_kb_response -------- class TestPrepareMessagesForKbSearch: """Tests for _prepare_messages_for_kb_search.""" def test_text_only_messages(self) -> None: messages = [ Message(role="user", contents=["hello"]), Message(role="assistant", contents=["world"]), ] result = AzureAISearchContextProvider._prepare_messages_for_kb_search(messages) assert len(result) == 2 assert result[0].role == "user" assert result[1].role == "assistant" # Verify content is KnowledgeBaseMessageTextContent from azure.search.documents.knowledgebases.models import KnowledgeBaseMessageTextContent assert isinstance(result[0].content[0], KnowledgeBaseMessageTextContent) assert result[0].content[0].text == "hello" def test_image_uri_content(self) -> None: img = Content.from_uri(uri="https://example.com/photo.png", media_type="image/png") messages = [Message(role="user", contents=[img])] result = AzureAISearchContextProvider._prepare_messages_for_kb_search(messages) assert len(result) == 1 from azure.search.documents.knowledgebases.models import KnowledgeBaseMessageImageContent assert isinstance(result[0].content[0], KnowledgeBaseMessageImageContent) assert result[0].content[0].image.url == "https://example.com/photo.png" def test_mixed_text_and_image_content(self) -> None: text = Content.from_text("describe this image") img = Content.from_uri(uri="https://example.com/img.jpg", media_type="image/jpeg") messages = [Message(role="user", contents=[text, img])] result = AzureAISearchContextProvider._prepare_messages_for_kb_search(messages) assert len(result) == 1 assert len(result[0].content) == 2 def test_skips_non_text_non_image_content(self) -> None: error = Content.from_error(message="oops") messages = [Message(role="user", contents=[error])] result = AzureAISearchContextProvider._prepare_messages_for_kb_search(messages) assert len(result) == 0 # message had no usable content def test_skips_empty_text(self) -> None: empty = Content.from_text("") messages = [Message(role="user", contents=[empty])] result = AzureAISearchContextProvider._prepare_messages_for_kb_search(messages) assert len(result) == 0 def test_fallback_to_msg_text_when_no_contents(self) -> None: msg = Message(role="user", contents=["fallback text"]) result = AzureAISearchContextProvider._prepare_messages_for_kb_search([msg]) assert len(result) == 1 from azure.search.documents.knowledgebases.models import KnowledgeBaseMessageTextContent assert isinstance(result[0].content[0], KnowledgeBaseMessageTextContent) assert result[0].content[0].text == "fallback text" def test_data_uri_image(self) -> None: img = Content.from_data(data=b"\x89PNG", media_type="image/png") messages = [Message(role="user", contents=[img])] result = AzureAISearchContextProvider._prepare_messages_for_kb_search(messages) assert len(result) == 1 from azure.search.documents.knowledgebases.models import KnowledgeBaseMessageImageContent assert isinstance(result[0].content[0], KnowledgeBaseMessageImageContent) def test_non_image_uri_skipped(self) -> None: pdf = Content.from_uri(uri="https://example.com/doc.pdf", media_type="application/pdf") messages = [Message(role="user", contents=[pdf])] result = AzureAISearchContextProvider._prepare_messages_for_kb_search(messages) assert len(result) == 0 class TestParseReferencesToAnnotations: """Tests for _parse_references_to_annotations.""" def test_none_references(self) -> None: result = AzureAISearchContextProvider._parse_references_to_annotations(None) assert result == [] def test_empty_references(self) -> None: result = AzureAISearchContextProvider._parse_references_to_annotations([]) assert result == [] def test_search_index_reference_captures_doc_key(self) -> None: from azure.search.documents.knowledgebases.models import KnowledgeBaseSearchIndexReference ref = KnowledgeBaseSearchIndexReference(id="ref-1", activity_source=0, doc_key="doc-1") result = AzureAISearchContextProvider._parse_references_to_annotations([ref]) assert len(result) == 1 assert result[0]["type"] == "citation" assert result[0]["title"] == "ref-1" extra = result[0]["additional_properties"] assert extra["reference_id"] == "ref-1" assert extra["reference_type"] == "searchIndex" assert extra["activity_source"] == 0 assert extra["doc_key"] == "doc-1" def test_web_reference_with_url_and_title(self) -> None: from azure.search.documents.knowledgebases.models import KnowledgeBaseWebReference ref = KnowledgeBaseWebReference( id="ref-2", activity_source=0, url="https://example.com/page", title="Example Page" ) result = AzureAISearchContextProvider._parse_references_to_annotations([ref]) assert len(result) == 1 assert result[0]["url"] == "https://example.com/page" assert result[0]["title"] == "Example Page" assert result[0]["additional_properties"]["reference_type"] == "web" def test_blob_reference_extracts_blob_url(self) -> None: from azure.search.documents.knowledgebases.models import KnowledgeBaseAzureBlobReference ref = KnowledgeBaseAzureBlobReference( id="ref-3", activity_source=0, blob_url="https://storage.blob.core.windows.net/doc.pdf" ) result = AzureAISearchContextProvider._parse_references_to_annotations([ref]) assert result[0]["url"] == "https://storage.blob.core.windows.net/doc.pdf" assert result[0]["additional_properties"]["reference_type"] == "azureBlob" def test_source_data_and_reranker_score(self) -> None: from azure.search.documents.knowledgebases.models import KnowledgeBaseSearchIndexReference ref = KnowledgeBaseSearchIndexReference( id="ref-4", activity_source=0, source_data={"chunk": "some text"}, reranker_score=0.95 ) result = AzureAISearchContextProvider._parse_references_to_annotations([ref]) extra = result[0]["additional_properties"] assert extra["source_data"] == {"chunk": "some text"} assert extra["reranker_score"] == 0.95 def test_raw_representation_stores_original_ref(self) -> None: from azure.search.documents.knowledgebases.models import KnowledgeBaseSearchIndexReference ref = KnowledgeBaseSearchIndexReference(id="ref-5", activity_source=0) result = AzureAISearchContextProvider._parse_references_to_annotations([ref]) assert result[0]["raw_representation"] is ref def test_remote_sharepoint_captures_sensitivity_label(self) -> None: # KnowledgeBaseRemoteSharePointReference is preview-only; a SimpleNamespace fake keeps this # test runnable on the stable/GA SDK while exercising the same parsing branches. ref = SimpleNamespace( id="ref-6", activity_source=0, reranker_score=None, source_data=None, web_url="https://sp.example.com/doc", search_sensitivity_label_info=SimpleNamespace( display_name="Confidential", sensitivity_label_id="lbl-1", is_encrypted=True ), ) result = AzureAISearchContextProvider._parse_references_to_annotations(cast(Any, [ref])) assert result[0]["url"] == "https://sp.example.com/doc" sl = result[0]["additional_properties"]["sensitivity_label"] assert sl["display_name"] == "Confidential" assert sl["sensitivity_label_id"] == "lbl-1" assert sl["is_encrypted"] is True def test_multiple_references(self) -> None: from azure.search.documents.knowledgebases.models import ( KnowledgeBaseSearchIndexReference, KnowledgeBaseWebReference, ) refs = [ KnowledgeBaseSearchIndexReference(id="ref-a", activity_source=0), KnowledgeBaseWebReference(id="ref-b", activity_source=1, url="https://example.com"), ] result = AzureAISearchContextProvider._parse_references_to_annotations(refs) # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type] assert len(result) == 2 assert result[0]["additional_properties"]["activity_source"] == 0 assert result[1]["additional_properties"]["activity_source"] == 1 class TestParseMessagesFromKbResponse: """Tests for _parse_messages_from_kb_response.""" def test_converts_all_messages(self) -> None: from azure.search.documents.knowledgebases.models import ( KnowledgeBaseMessage, KnowledgeBaseMessageTextContent, KnowledgeBaseRetrievalResponse, ) response = KnowledgeBaseRetrievalResponse( response=[ KnowledgeBaseMessage(role="user", content=[KnowledgeBaseMessageTextContent(text="q")]), KnowledgeBaseMessage(role="assistant", content=[KnowledgeBaseMessageTextContent(text="answer")]), ], references=None, ) result = AzureAISearchContextProvider._parse_messages_from_kb_response(response) assert len(result) == 2 assert result[0].role == "user" assert result[0].text == "q" assert result[1].role == "assistant" assert result[1].text == "answer" def test_none_response_returns_default(self) -> None: from azure.search.documents.knowledgebases.models import KnowledgeBaseRetrievalResponse response = KnowledgeBaseRetrievalResponse(response=None, references=None) result = AzureAISearchContextProvider._parse_messages_from_kb_response(response) assert len(result) == 1 assert result[0].text == "No results found from Knowledge Base." def test_empty_response_returns_default(self) -> None: from azure.search.documents.knowledgebases.models import KnowledgeBaseRetrievalResponse response = KnowledgeBaseRetrievalResponse(response=[], references=None) result = AzureAISearchContextProvider._parse_messages_from_kb_response(response) assert len(result) == 1 assert result[0].text == "No results found from Knowledge Base." def test_image_content(self) -> None: from azure.search.documents.knowledgebases.models import ( KnowledgeBaseImageContent, KnowledgeBaseMessage, KnowledgeBaseMessageImageContent, KnowledgeBaseRetrievalResponse, ) # KnowledgeBaseImageContent is available on both the stable and preview SDKs and # exposes ``url``, so the parsing assertion stays SDK-agnostic. response = KnowledgeBaseRetrievalResponse( response=[ KnowledgeBaseMessage( role="assistant", content=[ KnowledgeBaseMessageImageContent( image=KnowledgeBaseImageContent(url="https://img.example.com/a.png") ) ], ), ], references=None, ) result = AzureAISearchContextProvider._parse_messages_from_kb_response(response) assert len(result) == 1 assert result[0].contents[0].type == "uri" assert result[0].contents[0].uri == "https://img.example.com/a.png" def test_mixed_text_and_image_content(self) -> None: from azure.search.documents.knowledgebases.models import ( KnowledgeBaseImageContent, KnowledgeBaseMessage, KnowledgeBaseMessageImageContent, KnowledgeBaseMessageTextContent, KnowledgeBaseRetrievalResponse, ) response = KnowledgeBaseRetrievalResponse( response=[ KnowledgeBaseMessage( role="assistant", content=[ KnowledgeBaseMessageTextContent(text="description"), KnowledgeBaseMessageImageContent( image=KnowledgeBaseImageContent(url="https://img.example.com/b.png") ), ], ), ], references=None, ) result = AzureAISearchContextProvider._parse_messages_from_kb_response(response) assert len(result) == 1 assert len(result[0].contents) == 2 assert result[0].contents[0].type == "text" assert result[0].contents[1].type == "uri" def test_references_become_annotations(self) -> None: from azure.search.documents.knowledgebases.models import ( KnowledgeBaseMessage, KnowledgeBaseMessageTextContent, KnowledgeBaseRetrievalResponse, KnowledgeBaseWebReference, ) response = KnowledgeBaseRetrievalResponse( response=[ KnowledgeBaseMessage(role="assistant", content=[KnowledgeBaseMessageTextContent(text="answer")]), ], references=[ KnowledgeBaseWebReference(id="ref-1", activity_source=0, url="https://example.com", title="Example"), ], ) result = AzureAISearchContextProvider._parse_messages_from_kb_response(response) assert len(result) == 1 annotations = result[0].contents[0].annotations assert annotations is not None assert len(annotations) == 1 assert annotations[0]["type"] == "citation" assert annotations[0]["url"] == "https://example.com" assert annotations[0]["title"] == "Example" def test_multiple_messages_with_references(self) -> None: from azure.search.documents.knowledgebases.models import ( KnowledgeBaseMessage, KnowledgeBaseMessageTextContent, KnowledgeBaseRetrievalResponse, KnowledgeBaseSearchIndexReference, ) response = KnowledgeBaseRetrievalResponse( response=[ KnowledgeBaseMessage(role="user", content=[KnowledgeBaseMessageTextContent(text="q")]), KnowledgeBaseMessage( role="assistant", content=[ KnowledgeBaseMessageTextContent(text="part1"), KnowledgeBaseMessageTextContent(text="part2"), ], ), ], references=[KnowledgeBaseSearchIndexReference(id="doc-1", activity_source=0)], ) result = AzureAISearchContextProvider._parse_messages_from_kb_response(response) assert len(result) == 2 # All content items get annotations for msg in result: for c in msg.contents: assert c.annotations is not None assert len(c.annotations) == 1 # -- before_run: agentic mode -------------------------------------------------- class TestBeforeRunAgentic: """Tests for before_run in agentic mode.""" async def test_agentic_mode_calls_agentic_search(self) -> None: provider = _make_provider() provider.mode = "agentic" provider.agentic_message_history_count = 5 provider._knowledge_base_initialized = True provider.knowledge_base_name = "kb" mock_content = Mock() mock_content.text = "agentic result" mock_message = Mock() mock_message.role = "assistant" mock_message.content = [mock_content] mock_result = Mock() mock_result.response = [mock_message] mock_result.references = None mock_retrieval = AsyncMock() mock_retrieval.retrieve = AsyncMock(return_value=mock_result) provider._retrieval_client = mock_retrieval session = AgentSession(session_id="test-session") ctx = SessionContext( input_messages=[Message(role="user", contents=["agentic question"])], session_id="s1", ) with patch( "agent_framework_azure_ai_search._context_provider.KnowledgeBaseMessageTextContent", type(mock_content), ): await provider.before_run( agent=cast(Any, None), session=session, context=ctx, state=session.state.setdefault(provider.source_id, {}), ) # type: ignore[arg-type] msgs = ctx.context_messages.get(provider.source_id, []) assert len(msgs) >= 2 assert msgs[0].text == provider.context_prompt assert msgs[1].text == "agentic result"