fixes to azure ai search init, samples (#5021)

This commit is contained in:
Eduard van Valkenburg
2026-04-01 11:59:52 +02:00
committed by GitHub
Unverified
parent e43fc8ccec
commit 2a8c3e2dcf
9 changed files with 333 additions and 53 deletions
@@ -11,11 +11,21 @@ from __future__ import annotations
import logging
import sys
from collections.abc import Awaitable, Callable
from typing import TYPE_CHECKING, Any, ClassVar, Literal, TypedDict
from typing import TYPE_CHECKING, Any, ClassVar, Literal, TypedDict, overload
from agent_framework import AGENT_FRAMEWORK_USER_AGENT, Annotation, Content, Message, SupportsGetEmbeddings
from agent_framework._sessions import AgentSession, BaseContextProvider, SessionContext
from agent_framework._settings import SecretString, load_settings
from agent_framework import (
AGENT_FRAMEWORK_USER_AGENT,
AgentSession,
Annotation,
BaseContextProvider,
Content,
Message,
SecretString,
SessionContext,
SupportsGetEmbeddings,
load_settings,
)
from agent_framework.exceptions import SettingNotFoundError
from azure.core.credentials import AzureKeyCredential, TokenCredential
from azure.core.credentials_async import AsyncTokenCredential
from azure.core.exceptions import ResourceNotFoundError
@@ -111,6 +121,9 @@ except ImportError:
_agentic_retrieval_available = False
AzureCredentialTypes = TokenCredential | AsyncTokenCredential
EmbeddingFunction = Callable[[str], Awaitable[list[float]]] | SupportsGetEmbeddings[str, list[float], Any]
KnowledgeBaseOutputModeLiteral = Literal["extractive_data", "answer_synthesis"]
RetrievalReasoningEffortLiteral = Literal["minimal", "medium", "low"]
logger = logging.getLogger("agent_framework.azure_ai_search")
@@ -151,6 +164,230 @@ class AzureAISearchContextProvider(BaseContextProvider):
_DEFAULT_SEARCH_CONTEXT_PROMPT: ClassVar[str] = "Use the following context to answer the question:"
DEFAULT_SOURCE_ID: ClassVar[str] = "azure_ai_search"
@overload
def __init__(
self,
source_id: str = DEFAULT_SOURCE_ID,
endpoint: str | None = None,
index_name: str | None = None,
api_key: str | AzureKeyCredential | None = None,
credential: AzureCredentialTypes | None = None,
*,
mode: Literal["semantic"] = "semantic",
top_k: int = 5,
semantic_configuration_name: str | None = None,
vector_field_name: str | None = None,
embedding_function: EmbeddingFunction | None = None,
context_prompt: str | None = None,
azure_openai_resource_url: str | None = None,
model_deployment_name: str | None = None,
model_name: str | None = None,
knowledge_base_name: None = None,
retrieval_instructions: str | None = None,
azure_openai_api_key: str | None = None,
knowledge_base_output_mode: KnowledgeBaseOutputModeLiteral = "extractive_data",
retrieval_reasoning_effort: RetrievalReasoningEffortLiteral = "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 a semantic Azure AI Search context provider.
Keyword Args:
source_id: Unique identifier for this provider instance.
endpoint: Azure AI Search endpoint URL.
index_name: Name of the search index to query.
api_key: API key for authentication.
credential: Azure credential for managed identity authentication.
mode: Must be ``"semantic"`` for this overload.
top_k: Maximum number of documents to retrieve.
semantic_configuration_name: Name of the semantic configuration in the index.
vector_field_name: Name of the vector field in the index.
embedding_function: Embedding provider used for vector search.
context_prompt: Custom prompt to prepend to retrieved context.
azure_openai_resource_url: Unused in semantic mode.
model_deployment_name: Unused in semantic mode.
model_name: Unused in semantic mode.
knowledge_base_name: Must be ``None`` for this overload.
retrieval_instructions: Unused in semantic mode.
azure_openai_api_key: Unused in semantic mode.
knowledge_base_output_mode: Unused in semantic mode.
retrieval_reasoning_effort: Unused in semantic mode.
agentic_message_history_count: Unused in semantic mode.
env_file_path: Optional ``.env`` file checked before process environment variables.
env_file_encoding: Encoding for the ``.env`` file.
"""
...
@overload
def __init__(
self,
source_id: str = DEFAULT_SOURCE_ID,
endpoint: str | None = None,
index_name: str | None = None,
api_key: str | AzureKeyCredential | None = None,
credential: AzureCredentialTypes | None = None,
*,
mode: Literal["agentic"],
top_k: int = 5,
semantic_configuration_name: str | None = None,
vector_field_name: str | None = None,
embedding_function: EmbeddingFunction | None = None,
context_prompt: str | None = None,
azure_openai_resource_url: str,
model_deployment_name: str,
model_name: str | None = None,
knowledge_base_name: None = None,
retrieval_instructions: str | None = None,
azure_openai_api_key: str | None = None,
knowledge_base_output_mode: KnowledgeBaseOutputModeLiteral = "extractive_data",
retrieval_reasoning_effort: RetrievalReasoningEffortLiteral = "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 an agentic provider that creates a Knowledge Base from an index.
Keyword Args:
source_id: Unique identifier for this provider instance.
endpoint: Azure AI Search endpoint URL.
index_name: Name of the search index used to create the Knowledge Base.
api_key: API key for authentication.
credential: Azure credential for managed identity authentication.
mode: Must be ``"agentic"`` for this overload.
top_k: Maximum number of documents to retrieve.
semantic_configuration_name: Semantic configuration name used by hybrid search operations.
vector_field_name: Vector field name used by hybrid search operations.
embedding_function: Embedding provider used for vector search.
context_prompt: Custom prompt to prepend to retrieved context.
azure_openai_resource_url: Azure OpenAI resource URL for Knowledge Base creation.
model_deployment_name: Azure OpenAI deployment used by the generated Knowledge Base.
model_name: Underlying model name for the Knowledge Base model configuration.
knowledge_base_name: Must be ``None`` for this overload.
retrieval_instructions: Custom instructions for Knowledge Base retrieval.
azure_openai_api_key: Optional Azure OpenAI API key for Knowledge Base creation.
knowledge_base_output_mode: Output mode for Knowledge Base retrieval.
retrieval_reasoning_effort: Reasoning effort for query planning.
agentic_message_history_count: Number of recent messages included in retrieval.
env_file_path: Optional ``.env`` file checked before process environment variables.
env_file_encoding: Encoding for the ``.env`` file.
"""
...
@overload
def __init__(
self,
source_id: str = DEFAULT_SOURCE_ID,
endpoint: str | None = None,
index_name: None = None,
api_key: str | AzureKeyCredential | None = None,
credential: AzureCredentialTypes | None = None,
*,
mode: Literal["agentic"],
top_k: int = 5,
semantic_configuration_name: str | None = None,
vector_field_name: str | None = None,
embedding_function: EmbeddingFunction | None = None,
context_prompt: 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,
retrieval_instructions: str | None = None,
azure_openai_api_key: str | None = None,
knowledge_base_output_mode: KnowledgeBaseOutputModeLiteral = "extractive_data",
retrieval_reasoning_effort: RetrievalReasoningEffortLiteral = "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 an agentic provider that connects to an existing Knowledge Base.
Keyword Args:
source_id: Unique identifier for this provider instance.
endpoint: Azure AI Search endpoint URL.
index_name: Must be ``None`` for this overload.
knowledge_base_name: Name of the existing Knowledge Base to use.
api_key: API key for authentication.
credential: Azure credential for managed identity authentication.
mode: Must be ``"agentic"`` for this overload.
top_k: Maximum number of documents to retrieve.
semantic_configuration_name: Semantic configuration name used by hybrid search operations.
vector_field_name: Vector field name used by hybrid search operations.
embedding_function: Embedding provider used for vector search.
context_prompt: Custom prompt to prepend to retrieved context.
azure_openai_resource_url: Unused when connecting to an existing Knowledge Base.
model_deployment_name: Unused when connecting to an existing Knowledge Base.
model_name: Unused when connecting to an existing Knowledge Base.
retrieval_instructions: Custom instructions for Knowledge Base retrieval.
azure_openai_api_key: Unused when connecting to an existing Knowledge Base.
knowledge_base_output_mode: Output mode for Knowledge Base retrieval.
retrieval_reasoning_effort: Reasoning effort for query planning.
agentic_message_history_count: Number of recent messages included in retrieval.
env_file_path: Optional ``.env`` file checked before process environment variables.
env_file_encoding: Encoding for the ``.env`` file.
"""
...
@overload
def __init__(
self,
source_id: str = DEFAULT_SOURCE_ID,
endpoint: str | None = None,
index_name: None = None,
api_key: str | AzureKeyCredential | None = None,
credential: AzureCredentialTypes | None = None,
*,
mode: Literal["agentic"],
top_k: int = 5,
semantic_configuration_name: str | None = None,
vector_field_name: str | None = None,
embedding_function: EmbeddingFunction | None = None,
context_prompt: str | None = None,
azure_openai_resource_url: str | None = None,
model_deployment_name: str | None = None,
model_name: str | None = None,
knowledge_base_name: None = None,
retrieval_instructions: str | None = None,
azure_openai_api_key: str | None = None,
knowledge_base_output_mode: KnowledgeBaseOutputModeLiteral = "extractive_data",
retrieval_reasoning_effort: RetrievalReasoningEffortLiteral = "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 an agentic provider using environment-resolved setup.
This overload is for agentic initialization where ``index_name`` or
``knowledge_base_name`` is supplied by ``env_file_path`` or the
``AZURE_SEARCH_*`` environment variables.
Keyword Args:
source_id: Unique identifier for this provider instance.
endpoint: Azure AI Search endpoint URL.
index_name: Resolved from ``env_file_path`` or ``AZURE_SEARCH_INDEX_NAME``.
api_key: API key for authentication.
credential: Azure credential for managed identity authentication.
mode: Must be ``"agentic"`` for this overload.
top_k: Maximum number of documents to retrieve.
semantic_configuration_name: Semantic configuration name used by hybrid search operations.
vector_field_name: Vector field name used by hybrid search operations.
embedding_function: Embedding provider used for vector search.
context_prompt: Custom prompt to prepend to retrieved context.
azure_openai_resource_url: Azure OpenAI resource URL when creating a Knowledge Base from an index.
model_deployment_name: Azure OpenAI deployment when creating a Knowledge Base from an index.
model_name: Underlying model name for Knowledge Base model configuration.
knowledge_base_name: Resolved from ``env_file_path`` or ``AZURE_SEARCH_KNOWLEDGE_BASE_NAME``.
retrieval_instructions: Custom instructions for Knowledge Base retrieval.
azure_openai_api_key: Optional Azure OpenAI API key for Knowledge Base creation.
knowledge_base_output_mode: Output mode for Knowledge Base retrieval.
retrieval_reasoning_effort: Reasoning effort for query planning.
agentic_message_history_count: Number of recent messages included in retrieval.
env_file_path: Optional ``.env`` file checked before process environment variables.
env_file_encoding: Encoding for the ``.env`` file.
"""
...
def __init__(
self,
source_id: str = DEFAULT_SOURCE_ID,
@@ -163,9 +400,7 @@ class AzureAISearchContextProvider(BaseContextProvider):
top_k: int = 5,
semantic_configuration_name: str | None = None,
vector_field_name: str | None = None,
embedding_function: Callable[[str], Awaitable[list[float]]]
| SupportsGetEmbeddings[str, list[float], Any]
| None = None,
embedding_function: EmbeddingFunction | None = None,
context_prompt: str | None = None,
azure_openai_resource_url: str | None = None,
model_deployment_name: str | None = None,
@@ -173,8 +408,8 @@ class AzureAISearchContextProvider(BaseContextProvider):
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",
knowledge_base_output_mode: KnowledgeBaseOutputModeLiteral = "extractive_data",
retrieval_reasoning_effort: RetrievalReasoningEffortLiteral = "minimal",
agentic_message_history_count: int = _DEFAULT_AGENTIC_MESSAGE_HISTORY_COUNT,
env_file_path: str | None = None,
env_file_encoding: str | None = None,
@@ -184,7 +419,9 @@ class AzureAISearchContextProvider(BaseContextProvider):
Args:
source_id: Unique identifier for this provider instance.
endpoint: Azure AI Search endpoint URL.
index_name: Name of the search index to query.
index_name: Name of the search index to query. In agentic mode, providing this
explicitly selects the index-backed setup and ignores any environment-provided
knowledge base name.
api_key: API key for authentication.
credential: Azure credential for managed identity authentication.
Accepts a TokenCredential, AsyncTokenCredential, or a callable token provider.
@@ -197,7 +434,9 @@ class AzureAISearchContextProvider(BaseContextProvider):
azure_openai_resource_url: Azure OpenAI resource URL for Knowledge Base.
model_deployment_name: Model deployment name in Azure OpenAI.
model_name: The underlying model name.
knowledge_base_name: Name of an existing Knowledge Base to use.
knowledge_base_name: Name of an existing Knowledge Base to use. In agentic mode,
providing this explicitly selects the Knowledge Base-backed setup and ignores any
environment-provided index name.
retrieval_instructions: Custom instructions for Knowledge Base retrieval.
azure_openai_api_key: Azure OpenAI API key.
knowledge_base_output_mode: Output mode for Knowledge Base retrieval.
@@ -208,12 +447,26 @@ class AzureAISearchContextProvider(BaseContextProvider):
"""
super().__init__(source_id)
# Determine which fields are required based on mode
required: list[str | tuple[str, ...]] = ["endpoint"]
required: list[str | tuple[str, ...]]
ignored_agentic_field: Literal["index_name", "knowledge_base_name"] | None = None
explicit_index_name = index_name is not None
explicit_knowledge_base_name = knowledge_base_name is not None
if mode == "semantic":
required.append("index_name")
elif mode == "agentic":
required.append(("index_name", "knowledge_base_name"))
required = ["endpoint", "index_name"]
elif explicit_index_name and explicit_knowledge_base_name:
raise SettingNotFoundError(
"Only one of 'index_name', 'knowledge_base_name' may be provided, "
"but multiple were set: 'index_name', 'knowledge_base_name'."
)
elif explicit_index_name:
required = ["endpoint", "index_name"]
ignored_agentic_field = "knowledge_base_name"
elif explicit_knowledge_base_name:
required = ["endpoint", "knowledge_base_name"]
ignored_agentic_field = "index_name"
else:
required = ["endpoint", ("index_name", "knowledge_base_name")]
# Load settings from environment/file
settings = load_settings(
@@ -227,6 +480,8 @@ class AzureAISearchContextProvider(BaseContextProvider):
env_file_path=env_file_path,
env_file_encoding=env_file_encoding,
)
if ignored_agentic_field is not None:
settings[ignored_agentic_field] = None
if mode == "agentic" and settings.get("index_name") and not model_deployment_name:
raise ValueError(
@@ -283,6 +283,37 @@ class TestInitAgenticValidation:
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_deployment_name="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
# -- __aenter__ / __aexit__ ---------------------------------------------------
@@ -808,22 +808,21 @@ class SupportsFileSearchTool(Protocol):
# region SupportsGetEmbeddings Protocol
# Contravariant TypeVars for the Protocol
# TypeVars for the Protocol
EmbeddingInputContraT = TypeVar(
"EmbeddingInputContraT",
default="str",
contravariant=True,
)
EmbeddingOptionsContraT = TypeVar(
"EmbeddingOptionsContraT",
EmbeddingProtocolOptionsT = TypeVar(
"EmbeddingProtocolOptionsT",
bound=TypedDict, # type: ignore[valid-type]
default="EmbeddingGenerationOptions",
contravariant=True,
)
@runtime_checkable
class SupportsGetEmbeddings(Protocol[EmbeddingInputContraT, EmbeddingT, EmbeddingOptionsContraT]):
class SupportsGetEmbeddings(Protocol[EmbeddingInputContraT, EmbeddingT, EmbeddingProtocolOptionsT]):
"""Protocol for an embedding client that can generate embeddings.
This protocol enables duck-typing for embedding generation. Any class that
@@ -850,8 +849,8 @@ class SupportsGetEmbeddings(Protocol[EmbeddingInputContraT, EmbeddingT, Embeddin
self,
values: Sequence[EmbeddingInputContraT],
*,
options: EmbeddingOptionsContraT | None = None,
) -> Awaitable[GeneratedEmbeddings[EmbeddingT]]:
options: EmbeddingProtocolOptionsT | None = None,
) -> Awaitable[GeneratedEmbeddings[EmbeddingT, EmbeddingProtocolOptionsT]]:
"""Generate embeddings for the given values.
Args:
@@ -96,6 +96,7 @@ class ConversationSplitter(Protocol):
# Fallback: split at last user message
return EvalItem._split_last_turn_static(conversation)
item.split_messages(split=split_before_memory)
"""
@@ -468,10 +469,7 @@ class EvalResults:
"""
if not self.all_passed:
errored = (self.result_counts or {}).get("errored", 0)
detail = msg or (
f"Eval run {self.run_id} {self.status}: "
f"{self.passed} passed, {self.failed} failed."
)
detail = msg or (f"Eval run {self.run_id} {self.status}: {self.passed} passed, {self.failed} failed.")
if errored:
detail += f" {errored} errored."
if self.report_url:
@@ -1188,8 +1186,7 @@ def _coerce_result(value: Any, check_name: str) -> CheckResult:
score = float(d["score"])
except (TypeError, ValueError) as exc:
raise TypeError(
f"Function evaluator '{check_name}' returned dict with non-numeric 'score' value:"
f" {d['score']!r}"
f"Function evaluator '{check_name}' returned dict with non-numeric 'score' value: {d['score']!r}"
) from exc
# Honour an explicit 'passed' override; otherwise threshold-based.
passed = bool(d["passed"]) if "passed" in d else score >= float(d.get("threshold", 0.5))
@@ -7,6 +7,7 @@ import os
from unittest.mock import AsyncMock, MagicMock
import pytest
from agent_framework import SupportsGetEmbeddings
from agent_framework.exceptions import SettingNotFoundError
from openai.types import CreateEmbeddingResponse
from openai.types import Embedding as OpenAIEmbedding
@@ -44,6 +45,7 @@ def test_openai_construction_with_explicit_params() -> None:
api_key="test-key",
)
assert client.model == "text-embedding-3-small"
assert isinstance(client, SupportsGetEmbeddings)
def test_raw_openai_embedding_client_init_uses_explicit_parameters() -> None:
@@ -22,31 +22,29 @@ This sample demonstrates how to run the same prompt flow against different built
chat clients using a single `get_client` factory.
Select one of these client names:
- openai_responses
- openai_chat
- openai_chat_completion
- anthropic
- ollama
- bedrock
- azure_openai_responses
- azure_openai_chat
- azure_openai_chat_completion
- foundry_chat
"""
ClientName = Literal[
"openai_responses",
"openai_chat",
"openai_chat_completion",
"anthropic",
"ollama",
"bedrock",
"azure_openai_responses",
"azure_openai_chat",
"azure_openai_chat_completion",
"foundry_chat",
]
# NOTE: approval_mode="never_require" is for sample brevity.
# Use "always_require" in production; see samples/02-agents/tools/function_tool_with_approval.py
# and samples/02-agents/tools/function_tool_with_approval_and_sessions.py.
@tool(approval_mode="never_require")
def get_weather(
location: Annotated[str, Field(description="The location to get the weather for.")],
@@ -62,7 +60,7 @@ def get_client(client_name: ClientName) -> SupportsChatGetResponse[Any]:
from agent_framework.anthropic import AnthropicClient
from agent_framework.ollama import OllamaChatClient
if client_name == "openai_responses":
if client_name == "openai_chat":
return OpenAIChatClient()
if client_name == "openai_chat_completion":
return OpenAIChatCompletionClient()
@@ -72,7 +70,7 @@ def get_client(client_name: ClientName) -> SupportsChatGetResponse[Any]:
return OllamaChatClient()
if client_name == "bedrock":
return BedrockChatClient()
if client_name == "azure_openai_responses":
if client_name == "azure_openai_chat":
return OpenAIChatClient(credential=AzureCliCredential())
if client_name == "azure_openai_chat_completion":
return OpenAIChatCompletionClient(credential=AzureCliCredential())
@@ -86,7 +84,7 @@ def get_client(client_name: ClientName) -> SupportsChatGetResponse[Any]:
raise ValueError(f"Unsupported client name: {client_name}")
async def main(client_name: ClientName = "openai_responses") -> None:
async def main(client_name: ClientName = "openai_chat") -> None:
"""Run a basic prompt using a selected built-in client."""
client = get_client(client_name)
@@ -122,7 +120,7 @@ async def main(client_name: ClientName = "openai_responses") -> None:
if __name__ == "__main__":
asyncio.run(main("openai_responses"))
asyncio.run(main("openai_chat"))
"""
@@ -11,7 +11,7 @@
import asyncio
from typing import Any
import tiktoken
import tiktoken # type: ignore
from agent_framework import (
Message,
TokenizerProtocol,
@@ -100,7 +100,7 @@ async def main() -> None:
credential=AzureCliCredential() if not search_key else None,
mode="agentic",
azure_openai_resource_url=azure_openai_resource_url,
model_model=model_deployment,
model_deployment_name=model_deployment,
# Optional: Configure retrieval behavior
knowledge_base_output_mode="extractive_data", # or "answer_synthesis"
retrieval_reasoning_effort="minimal", # or "medium", "low"
@@ -110,13 +110,12 @@ async def main() -> None:
# Create agent with search context provider
async with (
search_provider,
FoundryChatClient(
project_endpoint=project_endpoint,
model_model=model_deployment,
credential=AzureCliCredential(),
) as client,
Agent(
client=client,
client=FoundryChatClient(
project_endpoint=project_endpoint,
model=model_deployment,
credential=AzureCliCredential(),
),
name="SearchAgent",
instructions=(
"You are a helpful assistant with advanced reasoning capabilities. "
@@ -85,13 +85,12 @@ async def main() -> None:
# Create agent with search context provider
async with (
search_provider,
FoundryChatClient(
project_endpoint=project_endpoint,
model_model=model_deployment,
credential=credential,
) as client,
Agent(
client=client,
client=FoundryChatClient(
project_endpoint=project_endpoint,
model=model_deployment,
credential=credential,
),
name="SearchAgent",
instructions=(
"You are a helpful assistant. Use the provided context from the "