[BREAKING] Python: fix OpenAI Azure routing and provider samples (#4925)

* Python: fix OpenAI Azure routing and provider samples

Prefer OpenAI when OPENAI_API_KEY is present unless Azure is explicitly requested. Clarify constructor docs, keep deprecated Azure wrappers compatible with stricter settings validation, and refresh the provider samples and tests to use the current client patterns.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* fix bandit

* Python: align OpenAI embedding Azure routing

Extend the shared OpenAI-vs-Azure routing and credential behavior to the embedding client, add Azure embedding regression coverage, and refresh the embedding samples to use the generic client path.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Python: fix embedding client pyright check

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Python: thin OpenAI embedding wrapper

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Python: document embedding overload routing

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Python: fix callable OpenAI key routing

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Python: fix Azure credential routing tests

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Python: address OpenAI review feedback

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Python: narrow Azure routing markers

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Python: refine OpenAI model fallback order

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Python: narrow Azure deployment docs

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Python: remove embedding routing wording

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Python: run embedding Azure integration tests

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* changed variable name

* Python: expand OpenAI package README

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* clarified readme

* Python: fix Azure OpenAI integration setup

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Python: correct Azure integration env mapping

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* updated code to fix int tests

* test updates

* test fix

* fix test setup

* updates to tests and setup

* remove openai assistants int tests

* improvements in int tests

* fix env var

* fix env vars

* fix azure responses test

* trigger actions

---------

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This commit is contained in:
Eduard van Valkenburg
2026-03-27 14:33:39 +01:00
committed by GitHub
Unverified
parent 3611be82cf
commit cc0cfaaac8
103 changed files with 5451 additions and 4216 deletions
@@ -2,10 +2,9 @@
import importlib.metadata
from ._foundry_agent import FoundryAgent, RawFoundryAgent
from ._foundry_agent_client import RawFoundryAgentChatClient
from ._foundry_chat_client import FoundryChatClient, FoundryChatOptions, RawFoundryChatClient
from ._foundry_memory_provider import FoundryMemoryProvider
from ._agent import FoundryAgent, RawFoundryAgent, RawFoundryAgentChatClient
from ._chat_client import FoundryChatClient, FoundryChatOptions, RawFoundryChatClient
from ._memory_provider import FoundryMemoryProvider
try:
__version__ = importlib.metadata.version(__name__)
@@ -1,30 +1,37 @@
# Copyright (c) Microsoft. All rights reserved.
"""Microsoft Foundry Agent client for connecting to pre-configured agents in Foundry.
"""Microsoft Foundry Agent for connecting to pre-configured agents in Foundry.
This module provides ``RawFoundryAgentClient`` and ``FoundryAgentClient`` for
communicating with PromptAgents and HostedAgents via the Responses API.
This module provides ``RawFoundryAgent`` and ``FoundryAgent`` Agent subclasses
that connect to existing PromptAgents or HostedAgents in Foundry. Use
``FoundryAgent`` for the recommended experience with full middleware and telemetry.
"""
from __future__ import annotations
import logging
import sys
from collections.abc import Callable, Mapping, MutableMapping, Sequence
from collections.abc import Awaitable, Callable, Mapping, MutableMapping, Sequence
from typing import TYPE_CHECKING, Any, ClassVar, Generic, cast
from agent_framework._middleware import ChatMiddlewareLayer
from agent_framework._settings import load_settings
from agent_framework._telemetry import AGENT_FRAMEWORK_USER_AGENT
from agent_framework._tools import FunctionInvocationConfiguration, FunctionInvocationLayer, FunctionTool
from agent_framework._types import Message
from agent_framework.observability import ChatTelemetryLayer
from agent_framework import (
AGENT_FRAMEWORK_USER_AGENT,
AgentMiddlewareLayer,
BaseContextProvider,
ChatAndFunctionMiddlewareTypes,
ChatMiddlewareLayer,
FunctionInvocationConfiguration,
FunctionInvocationLayer,
FunctionTool,
Message,
RawAgent,
load_settings,
)
from agent_framework.observability import AgentTelemetryLayer, ChatTelemetryLayer
from agent_framework_openai._chat_client import OpenAIChatOptions, RawOpenAIChatClient
from azure.ai.projects.aio import AIProjectClient
from ._entra_id_authentication import AzureCredentialTypes
logger: logging.Logger = logging.getLogger(__name__)
from azure.core.credentials import TokenCredential
from azure.core.credentials_async import AsyncTokenCredential
if sys.version_info >= (3, 13):
from typing import TypeVar # type: ignore # pragma: no cover
@@ -33,22 +40,25 @@ else:
if sys.version_info >= (3, 12):
from typing import override # type: ignore # pragma: no cover
else:
from typing_extensions import override # type: ignore # pragma: no cover
from typing_extensions import override # type: ignore[import] # pragma: no cover
if sys.version_info >= (3, 11):
from typing import TypedDict # type: ignore # pragma: no cover
else:
from typing_extensions import TypedDict # type: ignore # pragma: no cover
if TYPE_CHECKING:
from agent_framework import Agent, BaseContextProvider
from agent_framework._middleware import (
ChatMiddleware,
ChatMiddlewareCallable,
FunctionMiddleware,
FunctionMiddlewareCallable,
from agent_framework import (
Agent,
BaseContextProvider,
ChatAndFunctionMiddlewareTypes,
MiddlewareTypes,
ToolTypes,
)
from agent_framework._tools import ToolTypes
logger: logging.Logger = logging.getLogger("agent_framework.foundry")
AzureTokenProvider = Callable[[], str | Awaitable[str]]
AzureCredentialTypes = TokenCredential | AsyncTokenCredential
class FoundryAgentSettings(TypedDict, total=False):
@@ -203,8 +213,6 @@ class RawFoundryAgentChatClient( # type: ignore[misc]
**kwargs: Any,
) -> Agent[FoundryAgentOptionsT]:
"""Create a FoundryAgent that reuses this client's Foundry configuration."""
from ._foundry_agent import FoundryAgent
function_tools = cast(
FunctionTool | Callable[..., Any] | Sequence[FunctionTool | Callable[..., Any]] | None,
tools,
@@ -359,9 +367,7 @@ class _FoundryAgentChatClient( # type: ignore[misc]
allow_preview: bool | None = None,
env_file_path: str | None = None,
env_file_encoding: str | None = None,
middleware: (
Sequence[ChatMiddleware | ChatMiddlewareCallable | FunctionMiddleware | FunctionMiddlewareCallable] | None
) = None,
middleware: (Sequence[ChatAndFunctionMiddlewareTypes] | None) = None,
function_invocation_configuration: FunctionInvocationConfiguration | None = None,
**kwargs: Any,
) -> None:
@@ -393,3 +399,236 @@ class _FoundryAgentChatClient( # type: ignore[misc]
function_invocation_configuration=function_invocation_configuration,
**kwargs,
)
class RawFoundryAgent( # type: ignore[misc]
RawAgent[FoundryAgentOptionsT],
):
"""Raw Microsoft Foundry Agent without agent-level middleware or telemetry.
Connects to an existing PromptAgent or HostedAgent in Foundry.
For full middleware and telemetry support, use :class:`FoundryAgent`.
Examples:
.. code-block:: python
from agent_framework.foundry import RawFoundryAgent
from azure.identity import AzureCliCredential
agent = RawFoundryAgent(
project_endpoint="https://your-project.services.ai.azure.com",
agent_name="my-prompt-agent",
agent_version="1.0",
credential=AzureCliCredential(),
)
result = await agent.run("Hello!")
"""
def __init__(
self,
*,
project_endpoint: str | None = None,
agent_name: str | None = None,
agent_version: str | None = None,
credential: AzureCredentialTypes | None = None,
project_client: AIProjectClient | None = None,
allow_preview: bool | None = None,
tools: FunctionTool | Callable[..., Any] | Sequence[FunctionTool | Callable[..., Any]] | None = None,
context_providers: Sequence[BaseContextProvider] | None = None,
client_type: type[RawFoundryAgentChatClient] | None = None,
env_file_path: str | None = None,
env_file_encoding: str | None = None,
**kwargs: Any,
) -> None:
"""Initialize a Foundry Agent.
Keyword Args:
project_endpoint: The Foundry project endpoint URL.
Can also be set via environment variable FOUNDRY_PROJECT_ENDPOINT.
agent_name: The name of the Foundry agent to connect to.
Can also be set via environment variable FOUNDRY_AGENT_NAME.
agent_version: The version of the agent (required for PromptAgents, optional for HostedAgents).
Can also be set via environment variable FOUNDRY_AGENT_VERSION.
credential: Azure credential for authentication.
project_client: An existing AIProjectClient to use.
allow_preview: Enables preview opt-in on internally-created AIProjectClient.
tools: Function tools to provide to the agent. Only ``FunctionTool`` objects are accepted.
context_providers: Optional context providers for injecting dynamic context.
client_type: Custom client class to use (must be a subclass of ``RawFoundryAgentChatClient``).
Defaults to ``_FoundryAgentChatClient`` (full client middleware).
env_file_path: Path to .env file for settings.
env_file_encoding: Encoding for .env file.
kwargs: Additional keyword arguments passed to the Agent base class.
"""
# Create the client
actual_client_type = client_type or _FoundryAgentChatClient
if not issubclass(actual_client_type, RawFoundryAgentChatClient):
raise TypeError(
f"client_type must be a subclass of RawFoundryAgentChatClient, got {actual_client_type.__name__}"
)
client = actual_client_type(
project_endpoint=project_endpoint,
agent_name=agent_name,
agent_version=agent_version,
credential=credential,
project_client=project_client,
allow_preview=allow_preview,
env_file_path=env_file_path,
env_file_encoding=env_file_encoding,
)
super().__init__(
client=client, # type: ignore[arg-type]
tools=tools, # type: ignore[arg-type]
context_providers=context_providers,
**kwargs,
)
async def configure_azure_monitor(
self,
enable_sensitive_data: bool = False,
**kwargs: Any,
) -> None:
"""Setup observability with Azure Monitor (Microsoft Foundry integration).
This method configures Azure Monitor for telemetry collection using the
connection string from the Foundry project client (accessed via the internal client).
Args:
enable_sensitive_data: Enable sensitive data logging (prompts, responses).
Should only be enabled in development/test environments. Default is False.
**kwargs: Additional arguments passed to configure_azure_monitor().
Raises:
ImportError: If azure-monitor-opentelemetry-exporter is not installed.
"""
from azure.core.exceptions import ResourceNotFoundError
client = self.client
if not isinstance(client, RawFoundryAgentChatClient):
raise TypeError("configure_azure_monitor requires a RawFoundryAgentChatClient-based client.")
try:
conn_string = await client.project_client.telemetry.get_application_insights_connection_string()
except ResourceNotFoundError:
logger.warning(
"No Application Insights connection string found for the Foundry project. "
"Please ensure Application Insights is configured in your project, "
"or call configure_otel_providers() manually with custom exporters."
)
return
try:
from azure.monitor.opentelemetry import configure_azure_monitor # type: ignore[import]
except ImportError as exc:
raise ImportError(
"azure-monitor-opentelemetry is required for Azure Monitor integration. "
"Install it with: pip install azure-monitor-opentelemetry"
) from exc
from agent_framework.observability import create_metric_views, create_resource, enable_instrumentation
if "resource" not in kwargs:
kwargs["resource"] = create_resource()
configure_azure_monitor(
connection_string=conn_string,
views=create_metric_views(),
**kwargs,
)
enable_instrumentation(enable_sensitive_data=enable_sensitive_data)
class FoundryAgent( # type: ignore[misc]
AgentMiddlewareLayer,
AgentTelemetryLayer,
RawFoundryAgent[FoundryAgentOptionsT],
):
"""Microsoft Foundry Agent with full middleware and telemetry support.
Connects to an existing PromptAgent or HostedAgent in Foundry.
This is the recommended class for production use.
Examples:
.. code-block:: python
from agent_framework.foundry import FoundryAgent
from azure.identity import AzureCliCredential
# Connect to a PromptAgent
agent = FoundryAgent(
project_endpoint="https://your-project.services.ai.azure.com",
agent_name="my-prompt-agent",
agent_version="1.0",
credential=AzureCliCredential(),
tools=[my_function_tool],
)
result = await agent.run("Hello!")
# Connect to a HostedAgent (no version needed)
agent = FoundryAgent(
project_endpoint="https://your-project.services.ai.azure.com",
agent_name="my-hosted-agent",
credential=AzureCliCredential(),
)
# Custom client (e.g., raw client without client middleware)
agent = FoundryAgent(
project_endpoint="https://your-project.services.ai.azure.com",
agent_name="my-agent",
credential=AzureCliCredential(),
client_type=RawFoundryAgentChatClient,
)
"""
def __init__(
self,
*,
project_endpoint: str | None = None,
agent_name: str | None = None,
agent_version: str | None = None,
credential: AzureCredentialTypes | None = None,
project_client: AIProjectClient | None = None,
allow_preview: bool | None = None,
tools: FunctionTool | Callable[..., Any] | Sequence[FunctionTool | Callable[..., Any]] | None = None,
context_providers: Sequence[BaseContextProvider] | None = None,
middleware: Sequence[MiddlewareTypes] | None = None,
client_type: type[RawFoundryAgentChatClient] | None = None,
env_file_path: str | None = None,
env_file_encoding: str | None = None,
**kwargs: Any,
) -> None:
"""Initialize a Foundry Agent with full middleware and telemetry.
Keyword Args:
project_endpoint: The Foundry project endpoint URL.
agent_name: The name of the Foundry agent to connect to.
agent_version: The version of the agent (for PromptAgents).
credential: Azure credential for authentication.
project_client: An existing AIProjectClient to use.
allow_preview: Enables preview opt-in on internally-created AIProjectClient.
tools: Function tools to provide to the agent. Only ``FunctionTool`` objects are accepted.
context_providers: Optional context providers.
middleware: Optional agent-level middleware.
client_type: Custom client class (must subclass ``RawFoundryAgentChatClient``).
env_file_path: Path to .env file for settings.
env_file_encoding: Encoding for .env file.
kwargs: Additional keyword arguments.
"""
super().__init__(
project_endpoint=project_endpoint,
agent_name=agent_name,
agent_version=agent_version,
credential=credential,
project_client=project_client,
allow_preview=allow_preview,
tools=tools,
context_providers=context_providers,
middleware=middleware,
client_type=client_type,
env_file_path=env_file_path,
env_file_encoding=env_file_encoding,
**kwargs,
)
@@ -4,14 +4,17 @@ from __future__ import annotations
import logging
import sys
from collections.abc import Sequence
from collections.abc import Awaitable, Callable, Sequence
from typing import TYPE_CHECKING, Any, ClassVar, Generic, Literal
from agent_framework._middleware import ChatMiddlewareLayer
from agent_framework._settings import load_settings
from agent_framework._telemetry import AGENT_FRAMEWORK_USER_AGENT
from agent_framework._tools import FunctionInvocationConfiguration, FunctionInvocationLayer
from agent_framework._types import Content
from agent_framework import (
AGENT_FRAMEWORK_USER_AGENT,
ChatMiddlewareLayer,
Content,
FunctionInvocationConfiguration,
FunctionInvocationLayer,
load_settings,
)
from agent_framework.observability import ChatTelemetryLayer
from agent_framework_openai._chat_client import OpenAIChatOptions, RawOpenAIChatClient
from azure.ai.projects.aio import AIProjectClient
@@ -25,9 +28,8 @@ from azure.ai.projects.models import (
)
from azure.ai.projects.models import FileSearchTool as ProjectsFileSearchTool
from azure.ai.projects.models import MCPTool as FoundryMCPTool
from ._entra_id_authentication import AzureCredentialTypes, AzureTokenProvider
from ._shared import resolve_file_ids
from azure.core.credentials import TokenCredential
from azure.core.credentials_async import AsyncTokenCredential
if sys.version_info >= (3, 13):
from typing import TypeVar # type: ignore # pragma: no cover
@@ -43,15 +45,13 @@ else:
from typing_extensions import TypedDict # type: ignore # pragma: no cover
if TYPE_CHECKING:
from agent_framework._middleware import (
ChatMiddleware,
ChatMiddlewareCallable,
FunctionMiddleware,
FunctionMiddlewareCallable,
)
from agent_framework import ChatAndFunctionMiddlewareTypes
logger: logging.Logger = logging.getLogger("agent_framework.foundry")
AzureTokenProvider = Callable[[], str | Awaitable[str]]
AzureCredentialTypes = TokenCredential | AsyncTokenCredential
class FoundrySettings(TypedDict, total=False):
"""Settings for Microsoft FoundryChatClient resolved from args and environment.
@@ -67,6 +67,33 @@ class FoundrySettings(TypedDict, total=False):
project_endpoint: str | None
def resolve_file_ids(file_ids: Sequence[str | Content] | None) -> list[str] | None:
"""Resolve file IDs from strings or hosted-file Content objects."""
if not file_ids:
return None
resolved: list[str] = []
for item in file_ids:
if isinstance(item, str):
if not item:
raise ValueError("file_ids must not contain empty strings.")
resolved.append(item)
elif isinstance(item, Content):
if item.type != "hosted_file":
raise ValueError(
f"Unsupported Content type {item.type!r} for code interpreter file_ids. "
"Only Content.from_hosted_file() is supported."
)
if item.file_id is None:
raise ValueError(
"Content.from_hosted_file() item is missing a file_id. "
"Ensure the Content object has a valid file_id before using it in file_ids."
)
resolved.append(item.file_id)
return resolved if resolved else None
FoundryChatOptionsT = TypeVar(
"FoundryChatOptionsT",
bound=TypedDict, # type: ignore[valid-type]
@@ -492,9 +519,7 @@ class FoundryChatClient( # type: ignore[misc]
env_file_path: str | None = None,
env_file_encoding: str | None = None,
instruction_role: str | None = None,
middleware: (
Sequence[ChatMiddleware | ChatMiddlewareCallable | FunctionMiddleware | FunctionMiddlewareCallable] | None
) = None,
middleware: (Sequence[ChatAndFunctionMiddlewareTypes] | None) = None,
function_invocation_configuration: FunctionInvocationConfiguration | None = None,
**kwargs: Any,
) -> None:
@@ -1,67 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
from __future__ import annotations
import logging
from collections.abc import Awaitable, Callable
from typing import Union
from agent_framework.exceptions import ChatClientInvalidAuthException
from azure.core.credentials import TokenCredential
from azure.core.credentials_async import AsyncTokenCredential
logger: logging.Logger = logging.getLogger(__name__)
AzureTokenProvider = Callable[[], Union[str, Awaitable[str]]]
"""A callable that returns a bearer token string, either synchronously or asynchronously."""
AzureCredentialTypes = Union[TokenCredential, AsyncTokenCredential]
"""Union of Azure credential types.
Accepts:
- ``TokenCredential`` — synchronous Azure credential (e.g. ``DefaultAzureCredential()``)
- ``AsyncTokenCredential`` — asynchronous Azure credential (e.g. ``azure.identity.aio.DefaultAzureCredential()``)
"""
def resolve_credential_to_token_provider(
credential: AzureCredentialTypes | AzureTokenProvider,
token_endpoint: str | None,
) -> AzureTokenProvider:
"""Convert an Azure credential or token provider into an ``ad_token_provider`` callable.
If the credential is already a callable token provider, it is returned as-is
(``token_endpoint`` is not required in this case).
If it is a ``TokenCredential`` or ``AsyncTokenCredential``, it is wrapped using
``azure.identity.get_bearer_token_provider`` (sync or async variant) which
handles token caching and automatic refresh.
Args:
credential: An Azure credential or token provider callable.
token_endpoint: The token scope/endpoint
(e.g. ``"https://cognitiveservices.azure.com/.default"``).
Required when ``credential`` is a ``TokenCredential`` or ``AsyncTokenCredential``.
Returns:
A callable that returns a bearer token string (sync or async).
Raises:
ServiceInvalidAuthError: If the token endpoint is empty when needed for credential wrapping.
"""
# Already a token provider callable (not a credential object) — use directly
if callable(credential) and not isinstance(credential, (TokenCredential, AsyncTokenCredential)):
return credential
if not token_endpoint:
raise ChatClientInvalidAuthException(
"A token endpoint must be provided either in settings, as an environment variable, or as an argument."
)
if isinstance(credential, AsyncTokenCredential):
from azure.identity.aio import get_bearer_token_provider as get_async_bearer_token_provider
return get_async_bearer_token_provider(credential, token_endpoint)
from azure.identity import get_bearer_token_provider
return get_bearer_token_provider(credential, token_endpoint) # type: ignore[arg-type]
@@ -1,287 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
"""Microsoft Foundry Agent for connecting to pre-configured agents in Foundry.
This module provides ``RawFoundryAgent`` and ``FoundryAgent`` — Agent subclasses
that connect to existing PromptAgents or HostedAgents in Foundry. Use
``FoundryAgent`` for the recommended experience with full middleware and telemetry.
"""
from __future__ import annotations
import logging
import sys
from collections.abc import Callable, Sequence
from typing import TYPE_CHECKING, Any
from agent_framework import (
AgentMiddlewareLayer,
BaseContextProvider,
RawAgent,
)
from agent_framework.observability import AgentTelemetryLayer
from azure.ai.projects.aio import AIProjectClient
from ._entra_id_authentication import AzureCredentialTypes
from ._foundry_agent_client import (
RawFoundryAgentChatClient,
_FoundryAgentChatClient, # pyright: ignore[reportPrivateUsage]
)
if sys.version_info >= (3, 13):
from typing import TypeVar # type: ignore # pragma: no cover
else:
from typing_extensions import TypeVar # type: ignore # pragma: no cover
if sys.version_info >= (3, 11):
from typing import TypedDict # type: ignore # pragma: no cover
else:
from typing_extensions import TypedDict # type: ignore # pragma: no cover
if TYPE_CHECKING:
from agent_framework._middleware import MiddlewareTypes
from agent_framework._tools import FunctionTool
from agent_framework_openai._chat_client import OpenAIChatOptions
logger: logging.Logger = logging.getLogger("agent_framework.foundry")
FoundryAgentOptionsT = TypeVar(
"FoundryAgentOptionsT",
bound=TypedDict, # type: ignore[valid-type]
default="OpenAIChatOptions",
covariant=True,
)
class RawFoundryAgent( # type: ignore[misc]
RawAgent[FoundryAgentOptionsT],
):
"""Raw Microsoft Foundry Agent without agent-level middleware or telemetry.
Connects to an existing PromptAgent or HostedAgent in Foundry.
For full middleware and telemetry support, use :class:`FoundryAgent`.
Examples:
.. code-block:: python
from agent_framework.foundry import RawFoundryAgent
from azure.identity import AzureCliCredential
agent = RawFoundryAgent(
project_endpoint="https://your-project.services.ai.azure.com",
agent_name="my-prompt-agent",
agent_version="1.0",
credential=AzureCliCredential(),
)
result = await agent.run("Hello!")
"""
def __init__(
self,
*,
project_endpoint: str | None = None,
agent_name: str | None = None,
agent_version: str | None = None,
credential: AzureCredentialTypes | None = None,
project_client: AIProjectClient | None = None,
allow_preview: bool | None = None,
tools: FunctionTool | Callable[..., Any] | Sequence[FunctionTool | Callable[..., Any]] | None = None,
context_providers: Sequence[BaseContextProvider] | None = None,
client_type: type[RawFoundryAgentChatClient] | None = None,
env_file_path: str | None = None,
env_file_encoding: str | None = None,
**kwargs: Any,
) -> None:
"""Initialize a Foundry Agent.
Keyword Args:
project_endpoint: The Foundry project endpoint URL.
Can also be set via environment variable FOUNDRY_PROJECT_ENDPOINT.
agent_name: The name of the Foundry agent to connect to.
Can also be set via environment variable FOUNDRY_AGENT_NAME.
agent_version: The version of the agent (required for PromptAgents, optional for HostedAgents).
Can also be set via environment variable FOUNDRY_AGENT_VERSION.
credential: Azure credential for authentication.
project_client: An existing AIProjectClient to use.
allow_preview: Enables preview opt-in on internally-created AIProjectClient.
tools: Function tools to provide to the agent. Only ``FunctionTool`` objects are accepted.
context_providers: Optional context providers for injecting dynamic context.
client_type: Custom client class to use (must be a subclass of ``RawFoundryAgentChatClient``).
Defaults to ``_FoundryAgentChatClient`` (full client middleware).
env_file_path: Path to .env file for settings.
env_file_encoding: Encoding for .env file.
kwargs: Additional keyword arguments passed to the Agent base class.
"""
# Create the client
actual_client_type = client_type or _FoundryAgentChatClient
if not issubclass(actual_client_type, RawFoundryAgentChatClient):
raise TypeError(
f"client_type must be a subclass of RawFoundryAgentChatClient, got {actual_client_type.__name__}"
)
client = actual_client_type(
project_endpoint=project_endpoint,
agent_name=agent_name,
agent_version=agent_version,
credential=credential,
project_client=project_client,
allow_preview=allow_preview,
env_file_path=env_file_path,
env_file_encoding=env_file_encoding,
)
super().__init__(
client=client, # type: ignore[arg-type]
tools=tools, # type: ignore[arg-type]
context_providers=context_providers,
**kwargs,
)
async def configure_azure_monitor(
self,
enable_sensitive_data: bool = False,
**kwargs: Any,
) -> None:
"""Setup observability with Azure Monitor (Microsoft Foundry integration).
This method configures Azure Monitor for telemetry collection using the
connection string from the Foundry project client (accessed via the internal client).
Args:
enable_sensitive_data: Enable sensitive data logging (prompts, responses).
Should only be enabled in development/test environments. Default is False.
**kwargs: Additional arguments passed to configure_azure_monitor().
Raises:
ImportError: If azure-monitor-opentelemetry-exporter is not installed.
"""
from azure.core.exceptions import ResourceNotFoundError
from ._foundry_agent_client import RawFoundryAgentChatClient
client = self.client
if not isinstance(client, RawFoundryAgentChatClient):
raise TypeError("configure_azure_monitor requires a RawFoundryAgentChatClient-based client.")
try:
conn_string = await client.project_client.telemetry.get_application_insights_connection_string()
except ResourceNotFoundError:
logger.warning(
"No Application Insights connection string found for the Foundry project. "
"Please ensure Application Insights is configured in your project, "
"or call configure_otel_providers() manually with custom exporters."
)
return
try:
from azure.monitor.opentelemetry import configure_azure_monitor # type: ignore[import]
except ImportError as exc:
raise ImportError(
"azure-monitor-opentelemetry is required for Azure Monitor integration. "
"Install it with: pip install azure-monitor-opentelemetry"
) from exc
from agent_framework.observability import create_metric_views, create_resource, enable_instrumentation
if "resource" not in kwargs:
kwargs["resource"] = create_resource()
configure_azure_monitor(
connection_string=conn_string,
views=create_metric_views(),
**kwargs,
)
enable_instrumentation(enable_sensitive_data=enable_sensitive_data)
class FoundryAgent( # type: ignore[misc]
AgentMiddlewareLayer,
AgentTelemetryLayer,
RawFoundryAgent[FoundryAgentOptionsT],
):
"""Microsoft Foundry Agent with full middleware and telemetry support.
Connects to an existing PromptAgent or HostedAgent in Foundry.
This is the recommended class for production use.
Examples:
.. code-block:: python
from agent_framework.foundry import FoundryAgent
from azure.identity import AzureCliCredential
# Connect to a PromptAgent
agent = FoundryAgent(
project_endpoint="https://your-project.services.ai.azure.com",
agent_name="my-prompt-agent",
agent_version="1.0",
credential=AzureCliCredential(),
tools=[my_function_tool],
)
result = await agent.run("Hello!")
# Connect to a HostedAgent (no version needed)
agent = FoundryAgent(
project_endpoint="https://your-project.services.ai.azure.com",
agent_name="my-hosted-agent",
credential=AzureCliCredential(),
)
# Custom client (e.g., raw client without client middleware)
agent = FoundryAgent(
project_endpoint="https://your-project.services.ai.azure.com",
agent_name="my-agent",
credential=AzureCliCredential(),
client_type=RawFoundryAgentChatClient,
)
"""
def __init__(
self,
*,
project_endpoint: str | None = None,
agent_name: str | None = None,
agent_version: str | None = None,
credential: AzureCredentialTypes | None = None,
project_client: AIProjectClient | None = None,
allow_preview: bool | None = None,
tools: FunctionTool | Callable[..., Any] | Sequence[FunctionTool | Callable[..., Any]] | None = None,
context_providers: Sequence[BaseContextProvider] | None = None,
middleware: Sequence[MiddlewareTypes] | None = None,
client_type: type[RawFoundryAgentChatClient] | None = None,
env_file_path: str | None = None,
env_file_encoding: str | None = None,
**kwargs: Any,
) -> None:
"""Initialize a Foundry Agent with full middleware and telemetry.
Keyword Args:
project_endpoint: The Foundry project endpoint URL.
agent_name: The name of the Foundry agent to connect to.
agent_version: The version of the agent (for PromptAgents).
credential: Azure credential for authentication.
project_client: An existing AIProjectClient to use.
allow_preview: Enables preview opt-in on internally-created AIProjectClient.
tools: Function tools to provide to the agent. Only ``FunctionTool`` objects are accepted.
context_providers: Optional context providers.
middleware: Optional agent-level middleware.
client_type: Custom client class (must subclass ``RawFoundryAgentChatClient``).
env_file_path: Path to .env file for settings.
env_file_encoding: Encoding for .env file.
kwargs: Additional keyword arguments.
"""
super().__init__(
project_endpoint=project_endpoint,
agent_name=agent_name,
agent_version=agent_version,
credential=credential,
project_client=project_client,
allow_preview=allow_preview,
tools=tools,
context_providers=context_providers,
middleware=middleware,
client_type=client_type,
env_file_path=env_file_path,
env_file_encoding=env_file_encoding,
**kwargs,
)
@@ -13,25 +13,38 @@ import sys
from contextlib import AbstractAsyncContextManager
from typing import TYPE_CHECKING, Any, ClassVar
from agent_framework import AGENT_FRAMEWORK_USER_AGENT, Message
from agent_framework._sessions import AgentSession, BaseContextProvider, SessionContext
from agent_framework._settings import load_settings
from agent_framework import (
AGENT_FRAMEWORK_USER_AGENT,
AgentSession,
BaseContextProvider,
Message,
SessionContext,
load_settings,
)
from azure.ai.projects.aio import AIProjectClient
from azure.core.credentials import TokenCredential
from azure.core.credentials_async import AsyncTokenCredential
from openai.types.responses import ResponseInputItemParam
from ._entra_id_authentication import AzureCredentialTypes
from ._shared import FoundryProjectSettings
if sys.version_info >= (3, 11):
from typing import Self # pragma: no cover
from typing import Self, TypedDict # pragma: no cover
else:
from typing_extensions import Self # pragma: no cover
from typing_extensions import Self, TypedDict # pragma: no cover
if TYPE_CHECKING:
from agent_framework._agents import SupportsAgentRun
from agent_framework import SupportsAgentRun
logger = logging.getLogger(__name__)
AzureCredentialTypes = TokenCredential | AsyncTokenCredential
class FoundryProjectSettings(TypedDict, total=False):
"""Foundry project settings loaded from FOUNDRY_ environment variables."""
project_endpoint: str | None
class FoundryMemoryProvider(BaseContextProvider):
"""Foundry Memory context provider using the new BaseContextProvider hooks pattern.
@@ -1,49 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
from __future__ import annotations
import logging
import sys
from collections.abc import Sequence
from agent_framework import Content
if sys.version_info >= (3, 11):
from typing import TypedDict # pragma: no cover
else:
from typing_extensions import TypedDict # type: ignore # pragma: no cover
logger = logging.getLogger("agent_framework.foundry")
class FoundryProjectSettings(TypedDict, total=False):
"""Foundry project settings loaded from FOUNDRY_ environment variables."""
project_endpoint: str | None
def resolve_file_ids(file_ids: Sequence[str | Content] | None) -> list[str] | None:
"""Resolve file IDs from strings or hosted-file Content objects."""
if not file_ids:
return None
resolved: list[str] = []
for item in file_ids:
if isinstance(item, str):
if not item:
raise ValueError("file_ids must not contain empty strings.")
resolved.append(item)
elif isinstance(item, Content):
if item.type != "hosted_file":
raise ValueError(
f"Unsupported Content type {item.type!r} for code interpreter file_ids. "
"Only Content.from_hosted_file() is supported."
)
if item.file_id is None:
raise ValueError(
"Content.from_hosted_file() item is missing a file_id. "
"Ensure the Content object has a valid file_id before using it in file_ids."
)
resolved.append(item.file_id)
return resolved if resolved else None
Binary file not shown.

After

Width:  |  Height:  |  Size: 178 KiB

@@ -0,0 +1,413 @@
# Copyright (c) Microsoft. All rights reserved.
from __future__ import annotations
import os
import sys
from typing import Any
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from agent_framework import AgentResponse, ChatContext, ChatMiddleware, Message, tool
from azure.core.exceptions import ResourceNotFoundError
from azure.identity import AzureCliCredential
from agent_framework_foundry._agent import (
FoundryAgent,
RawFoundryAgent,
RawFoundryAgentChatClient,
_FoundryAgentChatClient,
)
skip_if_foundry_agent_integration_tests_disabled = pytest.mark.skipif(
os.getenv("FOUNDRY_PROJECT_ENDPOINT", "") in ("", "https://test-project.services.ai.azure.com/")
or os.getenv("FOUNDRY_AGENT_NAME", "") == "",
reason="No real FOUNDRY_PROJECT_ENDPOINT or FOUNDRY_AGENT_NAME provided; skipping integration tests.",
)
_FOUNDRY_AGENT_ENV_VARS = (
"FOUNDRY_PROJECT_ENDPOINT",
"FOUNDRY_AGENT_NAME",
"FOUNDRY_AGENT_VERSION",
)
@pytest.fixture(autouse=True)
def clear_foundry_agent_settings_env(monkeypatch: pytest.MonkeyPatch, request: pytest.FixtureRequest) -> None:
"""Prevent unit tests from inheriting Foundry agent settings from the shell."""
if request.node.get_closest_marker("integration") is not None:
return
for env_var in _FOUNDRY_AGENT_ENV_VARS:
monkeypatch.delenv(env_var, raising=False)
def test_raw_foundry_agent_chat_client_init_requires_agent_name() -> None:
"""Test that agent_name is required."""
with pytest.raises(ValueError, match="Agent name is required"):
RawFoundryAgentChatClient(
project_client=MagicMock(),
)
def test_raw_foundry_agent_chat_client_init_with_agent_name() -> None:
"""Test construction with agent_name and project_client."""
mock_project = MagicMock()
mock_project.get_openai_client.return_value = MagicMock()
client = RawFoundryAgentChatClient(
project_client=mock_project,
agent_name="test-agent",
agent_version="1.0",
)
assert client.agent_name == "test-agent"
assert client.agent_version == "1.0"
def test_raw_foundry_agent_chat_client_get_agent_reference_with_version() -> None:
"""Test agent reference includes version when provided."""
mock_project = MagicMock()
mock_project.get_openai_client.return_value = MagicMock()
client = RawFoundryAgentChatClient(
project_client=mock_project,
agent_name="my-agent",
agent_version="2.0",
)
ref = client._get_agent_reference()
assert ref == {"name": "my-agent", "version": "2.0", "type": "agent_reference"}
def test_raw_foundry_agent_chat_client_get_agent_reference_without_version() -> None:
"""Test agent reference omits version for HostedAgents."""
mock_project = MagicMock()
mock_project.get_openai_client.return_value = MagicMock()
client = RawFoundryAgentChatClient(
project_client=mock_project,
agent_name="hosted-agent",
)
ref = client._get_agent_reference()
assert ref == {"name": "hosted-agent", "type": "agent_reference"}
assert "version" not in ref
def test_raw_foundry_agent_chat_client_as_agent_preserves_client_type() -> None:
"""Test that as_agent() wraps the client in FoundryAgent using the same client class."""
class CustomClient(RawFoundryAgentChatClient):
pass
mock_project = MagicMock()
mock_project.get_openai_client.return_value = MagicMock()
client = CustomClient(
project_client=mock_project,
agent_name="test-agent",
agent_version="1.0",
)
agent = client.as_agent(instructions="You are helpful.")
assert isinstance(agent, FoundryAgent)
assert agent.name == "test-agent"
assert isinstance(agent.client, CustomClient)
assert agent.client.project_client is mock_project
assert agent.client.agent_name == "test-agent"
assert agent.client.agent_version == "1.0"
named_agent = client.as_agent(name="display-name", instructions="You are helpful.")
assert named_agent.name == "display-name"
assert named_agent.client.agent_name == "test-agent"
async def test_raw_foundry_agent_chat_client_prepare_options_validates_tools() -> None:
"""Test that _prepare_options rejects non-FunctionTool objects."""
mock_project = MagicMock()
mock_project.get_openai_client.return_value = MagicMock()
client = RawFoundryAgentChatClient(
project_client=mock_project,
agent_name="test-agent",
)
with pytest.raises(TypeError, match="Only FunctionTool objects are accepted"):
await client._prepare_options(
messages=[Message(role="user", contents="hi")],
options={"tools": [{"type": "function", "function": {"name": "bad"}}]},
)
async def test_raw_foundry_agent_chat_client_prepare_options_accepts_function_tools() -> None:
"""Test that _prepare_options accepts FunctionTool objects."""
mock_project = MagicMock()
mock_openai = MagicMock()
mock_project.get_openai_client.return_value = mock_openai
client = RawFoundryAgentChatClient(
project_client=mock_project,
agent_name="test-agent",
)
@tool(approval_mode="never_require")
def my_func() -> str:
"""A test function."""
return "ok"
with patch(
"agent_framework_openai._chat_client.RawOpenAIChatClient._prepare_options",
new_callable=AsyncMock,
return_value={},
):
result = await client._prepare_options(
messages=[Message(role="user", contents="hi")],
options={"tools": [my_func]},
)
assert "extra_body" in result
assert result["extra_body"]["agent_reference"]["name"] == "test-agent"
def test_raw_foundry_agent_chat_client_check_model_presence_is_noop() -> None:
"""Test that _check_model_presence does nothing (model is on service)."""
mock_project = MagicMock()
mock_project.get_openai_client.return_value = MagicMock()
client = RawFoundryAgentChatClient(
project_client=mock_project,
agent_name="test-agent",
)
options: dict[str, Any] = {}
client._check_model_presence(options)
assert "model" not in options
def test_foundry_agent_chat_client_init() -> None:
"""Test construction of the full-middleware client."""
mock_project = MagicMock()
mock_project.get_openai_client.return_value = MagicMock()
client = _FoundryAgentChatClient(
project_client=mock_project,
agent_name="test-agent",
agent_version="1.0",
)
assert client.agent_name == "test-agent"
def test_raw_foundry_agent_init_creates_client() -> None:
"""Test that RawFoundryAgent creates a client internally."""
mock_project = MagicMock()
mock_project.get_openai_client.return_value = MagicMock()
agent = RawFoundryAgent(
project_client=mock_project,
agent_name="test-agent",
agent_version="1.0",
)
assert agent.client is not None
assert agent.client.agent_name == "test-agent"
def test_raw_foundry_agent_init_with_custom_client_type() -> None:
"""Test that client_type parameter is respected."""
mock_project = MagicMock()
mock_project.get_openai_client.return_value = MagicMock()
agent = RawFoundryAgent(
project_client=mock_project,
agent_name="test-agent",
client_type=RawFoundryAgentChatClient,
)
assert isinstance(agent.client, RawFoundryAgentChatClient)
def test_raw_foundry_agent_init_rejects_invalid_client_type() -> None:
"""Test that invalid client_type raises TypeError."""
with pytest.raises(TypeError, match="must be a subclass of RawFoundryAgentChatClient"):
RawFoundryAgent(
project_client=MagicMock(),
agent_name="test-agent",
client_type=object, # type: ignore[arg-type]
)
def test_raw_foundry_agent_init_with_function_tools() -> None:
"""Test that FunctionTool and callables are accepted."""
mock_project = MagicMock()
mock_project.get_openai_client.return_value = MagicMock()
@tool(approval_mode="never_require")
def my_func() -> str:
"""A test function."""
return "ok"
agent = RawFoundryAgent(
project_client=mock_project,
agent_name="test-agent",
tools=[my_func],
)
assert agent.default_options.get("tools") is not None
def test_foundry_agent_init() -> None:
"""Test construction of the full-middleware agent."""
mock_project = MagicMock()
mock_project.get_openai_client.return_value = MagicMock()
agent = FoundryAgent(
project_client=mock_project,
agent_name="test-agent",
agent_version="1.0",
)
assert agent.client is not None
assert agent.client.agent_name == "test-agent"
def test_foundry_agent_init_with_middleware() -> None:
"""Test that agent-level middleware is accepted."""
mock_project = MagicMock()
mock_project.get_openai_client.return_value = MagicMock()
class MyMiddleware(ChatMiddleware):
async def process(self, context: ChatContext) -> None:
pass
agent = FoundryAgent(
project_client=mock_project,
agent_name="test-agent",
middleware=[MyMiddleware()],
)
assert agent.client is not None
async def test_foundry_agent_configure_azure_monitor() -> None:
"""Test configure_azure_monitor delegates through the underlying client."""
mock_project = MagicMock()
mock_project.get_openai_client.return_value = MagicMock()
mock_project.telemetry.get_application_insights_connection_string = AsyncMock(
return_value="InstrumentationKey=test-key;IngestionEndpoint=https://test.endpoint"
)
agent = FoundryAgent(project_client=mock_project, agent_name="test-agent")
mock_configure = MagicMock()
mock_views = MagicMock(return_value=[])
mock_resource = MagicMock()
mock_enable = MagicMock()
with (
patch.dict(
"sys.modules",
{"azure.monitor.opentelemetry": MagicMock(configure_azure_monitor=mock_configure)},
),
patch("agent_framework.observability.create_metric_views", mock_views),
patch("agent_framework.observability.create_resource", return_value=mock_resource),
patch("agent_framework.observability.enable_instrumentation", mock_enable),
):
await agent.configure_azure_monitor(enable_sensitive_data=True)
mock_project.telemetry.get_application_insights_connection_string.assert_called_once()
call_kwargs = mock_configure.call_args.kwargs
assert call_kwargs["connection_string"] == "InstrumentationKey=test-key;IngestionEndpoint=https://test.endpoint"
assert call_kwargs["views"] == []
assert call_kwargs["resource"] is mock_resource
mock_enable.assert_called_once_with(enable_sensitive_data=True)
async def test_foundry_agent_configure_azure_monitor_resource_not_found() -> None:
"""Test configure_azure_monitor handles ResourceNotFoundError gracefully."""
mock_project = MagicMock()
mock_project.get_openai_client.return_value = MagicMock()
mock_project.telemetry.get_application_insights_connection_string = AsyncMock(
side_effect=ResourceNotFoundError("No Application Insights found")
)
agent = FoundryAgent(project_client=mock_project, agent_name="test-agent")
await agent.configure_azure_monitor()
mock_project.telemetry.get_application_insights_connection_string.assert_called_once()
async def test_foundry_agent_configure_azure_monitor_import_error() -> None:
"""Test configure_azure_monitor raises ImportError when Azure Monitor is unavailable."""
mock_project = MagicMock()
mock_project.get_openai_client.return_value = MagicMock()
mock_project.telemetry.get_application_insights_connection_string = AsyncMock(
return_value="InstrumentationKey=test-key"
)
agent = FoundryAgent(project_client=mock_project, agent_name="test-agent")
original_import = __import__
def _import_with_missing_azure_monitor(
name: str,
globals: dict[str, Any] | None = None,
locals: dict[str, Any] | None = None,
fromlist: tuple[str, ...] = (),
level: int = 0,
) -> Any:
if name == "azure.monitor.opentelemetry":
raise ImportError("No module named 'azure.monitor.opentelemetry'")
return original_import(name, globals, locals, fromlist, level)
with (
patch.dict(sys.modules, {"azure.monitor.opentelemetry": None}),
patch("builtins.__import__", side_effect=_import_with_missing_azure_monitor),
pytest.raises(ImportError, match="azure-monitor-opentelemetry is required"),
):
await agent.configure_azure_monitor()
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_foundry_agent_integration_tests_disabled
async def test_foundry_agent_basic_run() -> None:
"""Smoke-test FoundryAgent against a real configured agent."""
async with FoundryAgent(credential=AzureCliCredential()) as agent:
response = await agent.run("Please respond with exactly: 'This is a response test.'")
assert isinstance(response, AgentResponse)
assert response.text is not None
assert "response test" in response.text.lower()
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_foundry_agent_integration_tests_disabled
async def test_foundry_agent_custom_client_run() -> None:
"""Smoke-test FoundryAgent against a real configured agent."""
async with FoundryAgent(credential=AzureCliCredential(), client_type=RawFoundryAgentChatClient) as agent:
response = await agent.run("Please respond with exactly: 'This is a response test.'")
assert isinstance(response, AgentResponse)
assert response.text is not None
assert "response test" in response.text.lower()
@@ -0,0 +1,751 @@
# Copyright (c) Microsoft. All rights reserved.
from __future__ import annotations
import json
import os
import sys
from functools import wraps
from pathlib import Path
from typing import Annotated, Any
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from agent_framework import ChatResponse, Content, Message, SupportsChatGetResponse, tool
from agent_framework._telemetry import AGENT_FRAMEWORK_USER_AGENT
from agent_framework.exceptions import ChatClientException, ChatClientInvalidRequestException
from agent_framework_openai import OpenAIContentFilterException
from azure.core.exceptions import ResourceNotFoundError
from azure.identity import AzureCliCredential
from openai import BadRequestError
from pydantic import BaseModel
from pytest import param
from agent_framework_foundry import FoundryChatClient, RawFoundryChatClient
class OutputStruct(BaseModel):
"""A structured output for testing purposes."""
location: str
weather: str | None = None
@tool(approval_mode="never_require")
async def get_weather(location: Annotated[str, "The location as a city name"]) -> str:
"""Get the current weather in a given location."""
return f"The current weather in {location} is sunny."
skip_if_foundry_integration_tests_disabled = pytest.mark.skipif(
os.getenv("FOUNDRY_PROJECT_ENDPOINT", "") in ("", "https://test-project.services.ai.azure.com/")
or os.getenv("FOUNDRY_MODEL", "") == "",
reason="No real FOUNDRY_PROJECT_ENDPOINT or FOUNDRY_MODEL provided; skipping integration tests.",
)
_TEST_FOUNDRY_PROJECT_ENDPOINT = "https://test-project.services.ai.azure.com/"
_TEST_FOUNDRY_MODEL = "test-gpt-4o"
_FOUNDRY_CHAT_ENV_VARS = ("FOUNDRY_PROJECT_ENDPOINT", "FOUNDRY_MODEL")
@pytest.fixture(autouse=True)
def clear_foundry_chat_settings_env(monkeypatch: pytest.MonkeyPatch, request: pytest.FixtureRequest) -> None:
"""Prevent unit tests from inheriting Foundry chat settings from the shell."""
if request.node.get_closest_marker("integration") is not None:
return
for env_var in _FOUNDRY_CHAT_ENV_VARS:
monkeypatch.delenv(env_var, raising=False)
def _with_foundry_debug() -> Any:
def decorator(func: Any) -> Any:
@wraps(func)
async def wrapper(*args: Any, **kwargs: Any) -> Any:
try:
return await func(*args, **kwargs)
except Exception as exc:
debug_message = (
"Foundry debug: "
f"project_endpoint={os.getenv('FOUNDRY_PROJECT_ENDPOINT', '<unset>')}, "
f"model={os.getenv('FOUNDRY_MODEL', '<unset>')}"
)
if hasattr(exc, "add_note"):
exc.add_note(debug_message)
elif exc.args:
exc.args = (f"{exc.args[0]}\n{debug_message}", *exc.args[1:])
else:
exc.args = (debug_message,)
raise
return wrapper
return decorator
def _make_mock_openai_client() -> MagicMock:
client = MagicMock()
client.default_headers = {}
client.responses = MagicMock()
client.responses.create = AsyncMock()
client.responses.parse = AsyncMock()
client.files = MagicMock()
client.files.create = AsyncMock()
client.files.delete = AsyncMock()
client.vector_stores = MagicMock()
client.vector_stores.create = AsyncMock()
client.vector_stores.delete = AsyncMock()
client.vector_stores.files = MagicMock()
client.vector_stores.files.create_and_poll = AsyncMock()
return client
async def create_vector_store(client: FoundryChatClient) -> tuple[str, Content]:
"""Create a vector store with sample documents for testing."""
file = await client.client.files.create(
file=("todays_weather.txt", b"The weather today is sunny with a high of 75F."),
purpose="user_data",
)
vector_store = await client.client.vector_stores.create(
name="knowledge_base",
expires_after={"anchor": "last_active_at", "days": 1},
)
result = await client.client.vector_stores.files.create_and_poll(
vector_store_id=vector_store.id,
file_id=file.id,
poll_interval_ms=1000,
)
if result.last_error is not None:
raise RuntimeError(f"Vector store file processing failed with status: {result.last_error.message}")
return file.id, Content.from_hosted_vector_store(vector_store_id=vector_store.id)
async def delete_vector_store(client: FoundryChatClient, file_id: str, vector_store_id: str) -> None:
"""Delete the vector store after tests."""
await client.client.vector_stores.delete(vector_store_id=vector_store_id)
await client.client.files.delete(file_id=file_id)
def test_init() -> None:
mock_openai_client = _make_mock_openai_client()
mock_project_client = MagicMock()
mock_project_client.get_openai_client.return_value = mock_openai_client
client = FoundryChatClient(project_client=mock_project_client, model=_TEST_FOUNDRY_MODEL)
assert client.model == _TEST_FOUNDRY_MODEL
assert isinstance(client, SupportsChatGetResponse)
assert client.project_client is mock_project_client
def test_init_with_default_header() -> None:
default_headers = {"X-Unit-Test": "test-guid"}
mock_openai_client = _make_mock_openai_client()
project_client = MagicMock()
project_client.get_openai_client.return_value = mock_openai_client
client = FoundryChatClient(
project_client=project_client,
model=_TEST_FOUNDRY_MODEL,
default_headers=default_headers,
)
assert client.model == _TEST_FOUNDRY_MODEL
for key, value in default_headers.items():
assert client.default_headers is not None
assert key in client.default_headers
assert client.default_headers[key] == value
def test_init_with_project_endpoint_creates_project_client() -> None:
credential = MagicMock()
mock_openai_client = _make_mock_openai_client()
project_client = MagicMock()
project_client.get_openai_client.return_value = mock_openai_client
with patch("agent_framework_foundry._chat_client.AIProjectClient", return_value=project_client) as factory:
client = FoundryChatClient(
project_endpoint=_TEST_FOUNDRY_PROJECT_ENDPOINT,
model=_TEST_FOUNDRY_MODEL,
credential=credential,
allow_preview=True,
)
assert client.project_client is project_client
assert client.model == _TEST_FOUNDRY_MODEL
assert factory.call_args.kwargs["endpoint"] == _TEST_FOUNDRY_PROJECT_ENDPOINT
assert factory.call_args.kwargs["credential"] is credential
assert factory.call_args.kwargs["allow_preview"] is True
assert factory.call_args.kwargs["user_agent"] == AGENT_FRAMEWORK_USER_AGENT
def test_init_with_empty_model_raises(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.delenv("FOUNDRY_MODEL", raising=False)
mock_openai_client = _make_mock_openai_client()
mock_project_client = MagicMock()
mock_project_client.get_openai_client.return_value = mock_openai_client
with pytest.raises(ValueError, match="Model is required"):
FoundryChatClient(project_client=mock_project_client)
def test_init_with_empty_project_source_raises(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.delenv("FOUNDRY_PROJECT_ENDPOINT", raising=False)
with pytest.raises(ValueError, match="Either 'project_endpoint' or 'project_client' is required"):
FoundryChatClient(model=_TEST_FOUNDRY_MODEL)
def test_init_with_project_endpoint_requires_credential() -> None:
with pytest.raises(ValueError, match="Azure credential is required"):
FoundryChatClient(
project_endpoint=_TEST_FOUNDRY_PROJECT_ENDPOINT,
model=_TEST_FOUNDRY_MODEL,
)
async def test_configure_azure_monitor() -> None:
mock_openai_client = _make_mock_openai_client()
project_client = MagicMock()
project_client.get_openai_client.return_value = mock_openai_client
project_client.telemetry.get_application_insights_connection_string = AsyncMock(
return_value="InstrumentationKey=test-key;IngestionEndpoint=https://test.endpoint"
)
client = FoundryChatClient(project_client=project_client, model=_TEST_FOUNDRY_MODEL)
mock_configure = MagicMock()
mock_views = MagicMock(return_value=[])
mock_resource = MagicMock()
mock_enable = MagicMock()
with (
patch.dict(
"sys.modules",
{"azure.monitor.opentelemetry": MagicMock(configure_azure_monitor=mock_configure)},
),
patch("agent_framework.observability.create_metric_views", mock_views),
patch("agent_framework.observability.create_resource", return_value=mock_resource),
patch("agent_framework.observability.enable_instrumentation", mock_enable),
):
await client.configure_azure_monitor(enable_sensitive_data=True)
project_client.telemetry.get_application_insights_connection_string.assert_called_once()
mock_configure.assert_called_once()
call_kwargs = mock_configure.call_args.kwargs
assert call_kwargs["connection_string"] == "InstrumentationKey=test-key;IngestionEndpoint=https://test.endpoint"
assert call_kwargs["views"] == []
assert call_kwargs["resource"] is mock_resource
mock_enable.assert_called_once_with(enable_sensitive_data=True)
async def test_configure_azure_monitor_resource_not_found() -> None:
mock_openai_client = _make_mock_openai_client()
project_client = MagicMock()
project_client.get_openai_client.return_value = mock_openai_client
project_client.telemetry.get_application_insights_connection_string = AsyncMock(
side_effect=ResourceNotFoundError("No Application Insights found")
)
client = FoundryChatClient(project_client=project_client, model=_TEST_FOUNDRY_MODEL)
await client.configure_azure_monitor()
project_client.telemetry.get_application_insights_connection_string.assert_called_once()
async def test_configure_azure_monitor_import_error() -> None:
mock_openai_client = _make_mock_openai_client()
project_client = MagicMock()
project_client.get_openai_client.return_value = mock_openai_client
project_client.telemetry.get_application_insights_connection_string = AsyncMock(
return_value="InstrumentationKey=test-key"
)
client = FoundryChatClient(project_client=project_client, model=_TEST_FOUNDRY_MODEL)
original_import = __import__
def _import_with_missing_azure_monitor(
name: str,
globals: dict[str, Any] | None = None,
locals: dict[str, Any] | None = None,
fromlist: tuple[str, ...] = (),
level: int = 0,
) -> Any:
if name == "azure.monitor.opentelemetry":
raise ImportError("No module named 'azure.monitor.opentelemetry'")
return original_import(name, globals, locals, fromlist, level)
with (
patch.dict(sys.modules, {"azure.monitor.opentelemetry": None}),
patch("builtins.__import__", side_effect=_import_with_missing_azure_monitor),
pytest.raises(ImportError, match="azure-monitor-opentelemetry is required"),
):
await client.configure_azure_monitor()
async def test_configure_azure_monitor_with_custom_resource() -> None:
mock_openai_client = _make_mock_openai_client()
project_client = MagicMock()
project_client.get_openai_client.return_value = mock_openai_client
project_client.telemetry.get_application_insights_connection_string = AsyncMock(
return_value="InstrumentationKey=test-key"
)
client = FoundryChatClient(project_client=project_client, model=_TEST_FOUNDRY_MODEL)
custom_resource = MagicMock()
mock_configure = MagicMock()
with (
patch.dict(
"sys.modules",
{"azure.monitor.opentelemetry": MagicMock(configure_azure_monitor=mock_configure)},
),
patch("agent_framework.observability.create_metric_views", return_value=[]),
patch("agent_framework.observability.create_resource") as mock_create_resource,
patch("agent_framework.observability.enable_instrumentation"),
):
await client.configure_azure_monitor(resource=custom_resource)
mock_create_resource.assert_not_called()
call_kwargs = mock_configure.call_args.kwargs
assert call_kwargs["resource"] is custom_resource
async def test_get_response_with_invalid_input() -> None:
mock_openai_client = _make_mock_openai_client()
project_client = MagicMock()
project_client.get_openai_client.return_value = mock_openai_client
client = FoundryChatClient(project_client=project_client, model="test-model")
with pytest.raises(ChatClientInvalidRequestException, match="Messages are required"):
await client.get_response(messages=[])
async def test_web_search_tool_with_location() -> None:
mock_openai_client = _make_mock_openai_client()
project_client = MagicMock()
project_client.get_openai_client.return_value = mock_openai_client
client = FoundryChatClient(project_client=project_client, model="test-model")
web_search_tool = FoundryChatClient.get_web_search_tool(
user_location={
"city": "Seattle",
"country": "US",
"region": "WA",
"timezone": "America/Los_Angeles",
}
)
assert web_search_tool.user_location.city == "Seattle"
assert web_search_tool.user_location.country == "US"
_, run_options, _ = await client._prepare_request(
messages=[Message(role="user", text="What's the weather?")],
options={"tools": [web_search_tool], "tool_choice": "auto"},
)
assert run_options["tools"] == [web_search_tool]
assert run_options["tool_choice"] == "auto"
async def test_code_interpreter_tool_variations() -> None:
mock_openai_client = _make_mock_openai_client()
project_client = MagicMock()
project_client.get_openai_client.return_value = mock_openai_client
client = FoundryChatClient(project_client=project_client, model="test-model")
code_tool = FoundryChatClient.get_code_interpreter_tool()
assert code_tool.container["type"] == "auto"
_, run_options, _ = await client._prepare_request(
messages=[Message("user", ["Run some code"])],
options={"tools": [code_tool]},
)
assert run_options["tools"] == [code_tool]
code_tool_with_files = FoundryChatClient.get_code_interpreter_tool(file_ids=["file1", "file2"])
assert code_tool_with_files.container.file_ids == ["file1", "file2"]
_, run_options, _ = await client._prepare_request(
messages=[Message(role="user", text="Process these files")],
options={"tools": [code_tool_with_files]},
)
assert run_options["tools"] == [code_tool_with_files]
async def test_hosted_file_search_tool_validation() -> None:
mock_openai_client = _make_mock_openai_client()
project_client = MagicMock()
project_client.get_openai_client.return_value = mock_openai_client
client = FoundryChatClient(project_client=project_client, model="test-model")
with pytest.raises(ValueError, match="vector_store_ids"):
FoundryChatClient.get_file_search_tool(vector_store_ids=[])
file_search_tool = FoundryChatClient.get_file_search_tool(vector_store_ids=["vs_123"])
assert file_search_tool.vector_store_ids == ["vs_123"]
_, run_options, _ = await client._prepare_request(
messages=[Message("user", ["Test"])],
options={"tools": [file_search_tool]},
)
assert run_options["tools"] == [file_search_tool]
async def test_chat_message_parsing_with_function_calls() -> None:
mock_openai_client = _make_mock_openai_client()
project_client = MagicMock()
project_client.get_openai_client.return_value = mock_openai_client
client = FoundryChatClient(project_client=project_client, model="test-model")
function_call = Content.from_function_call(
call_id="test-call-id",
name="test_function",
arguments='{"param": "value"}',
additional_properties={"fc_id": "test-fc-id"},
)
function_result = Content.from_function_result(call_id="test-call-id", result="Function executed successfully")
messages = [
Message(role="user", text="Call a function"),
Message(role="assistant", contents=[function_call]),
Message(role="tool", contents=[function_result]),
]
prepared_messages = client._prepare_messages_for_openai(messages)
assert prepared_messages == [
{
"type": "message",
"role": "user",
"content": [{"type": "input_text", "text": "Call a function"}],
},
{
"call_id": "test-call-id",
"id": "fc_test-fc-id",
"type": "function_call",
"name": "test_function",
"arguments": '{"param": "value"}',
},
{
"call_id": "test-call-id",
"type": "function_call_output",
"output": "Function executed successfully",
},
]
async def test_content_filter_exception() -> None:
mock_openai_client = _make_mock_openai_client()
project_client = MagicMock()
project_client.get_openai_client.return_value = mock_openai_client
client = FoundryChatClient(project_client=project_client, model="test-model")
mock_error = BadRequestError(
message="Content filter error",
response=MagicMock(),
body={"error": {"code": "content_filter", "message": "Content filter error"}},
)
mock_error.code = "content_filter"
client.client.responses.create.side_effect = mock_error
with pytest.raises(OpenAIContentFilterException) as exc_info:
await client.get_response(messages=[Message(role="user", text="Test message")])
assert "content error" in str(exc_info.value)
async def test_response_format_parse_path() -> None:
mock_openai_client = _make_mock_openai_client()
project_client = MagicMock()
project_client.get_openai_client.return_value = mock_openai_client
client = FoundryChatClient(project_client=project_client, model="test-model")
mock_parsed_response = MagicMock()
mock_parsed_response.id = "parsed_response_123"
mock_parsed_response.text = "Parsed response"
mock_parsed_response.model = "test-model"
mock_parsed_response.created_at = 1000000000
mock_parsed_response.metadata = {}
mock_parsed_response.output_parsed = None
mock_parsed_response.usage = None
mock_parsed_response.finish_reason = None
mock_parsed_response.conversation = None
client.client.responses.parse = AsyncMock(return_value=mock_parsed_response)
response = await client.get_response(
messages=[Message(role="user", text="Test message")],
options={"response_format": OutputStruct, "store": True},
)
assert response.response_id == "parsed_response_123"
assert response.conversation_id == "parsed_response_123"
assert response.model == "test-model"
async def test_response_format_parse_path_with_conversation_id() -> None:
mock_openai_client = _make_mock_openai_client()
project_client = MagicMock()
project_client.get_openai_client.return_value = mock_openai_client
client = FoundryChatClient(project_client=project_client, model="test-model")
mock_parsed_response = MagicMock()
mock_parsed_response.id = "parsed_response_123"
mock_parsed_response.text = "Parsed response"
mock_parsed_response.model = "test-model"
mock_parsed_response.created_at = 1000000000
mock_parsed_response.metadata = {}
mock_parsed_response.output_parsed = None
mock_parsed_response.usage = None
mock_parsed_response.finish_reason = None
mock_parsed_response.conversation = MagicMock()
mock_parsed_response.conversation.id = "conversation_456"
client.client.responses.parse = AsyncMock(return_value=mock_parsed_response)
response = await client.get_response(
messages=[Message(role="user", text="Test message")],
options={"response_format": OutputStruct, "store": True},
)
assert response.response_id == "parsed_response_123"
assert response.conversation_id == "conversation_456"
assert response.model == "test-model"
async def test_bad_request_error_non_content_filter() -> None:
mock_openai_client = _make_mock_openai_client()
project_client = MagicMock()
project_client.get_openai_client.return_value = mock_openai_client
client = FoundryChatClient(project_client=project_client, model="test-model")
mock_error = BadRequestError(
message="Invalid request",
response=MagicMock(),
body={"error": {"code": "invalid_request", "message": "Invalid request"}},
)
mock_error.code = "invalid_request"
client.client.responses.parse = AsyncMock(side_effect=mock_error)
with pytest.raises(ChatClientException) as exc_info:
await client.get_response(
messages=[Message(role="user", text="Test message")],
options={"response_format": OutputStruct},
)
assert "failed to complete the prompt" in str(exc_info.value)
def test_get_mcp_tool_with_project_connection_id() -> None:
tool_config = FoundryChatClient.get_mcp_tool(
name="Docs MCP",
project_connection_id="conn-123",
allowed_tools=["search_docs"],
)
assert tool_config["project_connection_id"] == "conn-123"
assert tool_config["allowed_tools"] == ["search_docs"]
assert tool_config["server_label"] == "Docs_MCP"
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_foundry_integration_tests_disabled
@pytest.mark.parametrize(
"option_name,option_value,needs_validation",
[
param("max_tokens", 500, False, id="max_tokens"),
param("seed", 123, False, id="seed"),
param("user", "test-user-id", False, id="user"),
param("metadata", {"test_key": "test_value"}, False, id="metadata"),
param("tool_choice", "none", True, id="tool_choice_none"),
param("tools", [get_weather], True, id="tools_function"),
param("tool_choice", "auto", True, id="tool_choice_auto"),
param("response_format", OutputStruct, True, id="response_format_pydantic"),
param(
"response_format",
{
"type": "json_schema",
"json_schema": {
"name": "WeatherDigest",
"strict": True,
"schema": {
"title": "WeatherDigest",
"type": "object",
"properties": {
"location": {"type": "string"},
"conditions": {"type": "string"},
},
"required": ["location", "conditions"],
"additionalProperties": False,
},
},
},
True,
id="response_format_runtime_json_schema",
),
],
)
@_with_foundry_debug()
async def test_integration_options(
option_name: str,
option_value: Any,
needs_validation: bool,
) -> None:
client = FoundryChatClient(credential=AzureCliCredential())
client.function_invocation_configuration["max_iterations"] = 2
if option_name.startswith("tools") or option_name.startswith("tool_choice"):
messages = [Message(role="user", text="What is the weather in Seattle?")]
elif option_name.startswith("response_format"):
messages = [Message(role="user", text="The weather in Seattle is sunny")]
messages.append(Message(role="user", text="What is the weather in Seattle?"))
else:
messages = [Message(role="user", text="Say 'Hello World' briefly.")]
options: dict[str, Any] = {option_name: option_value}
if option_name.startswith("tool_choice"):
options["tools"] = [get_weather]
response = await client.get_response(messages=messages, options=options, stream=True).get_final_response()
assert isinstance(response, ChatResponse)
assert response.text is not None
assert len(response.text) > 0
if needs_validation:
if option_name.startswith("tools") or option_name.startswith("tool_choice"):
text = response.text.lower()
assert "sunny" in text or "seattle" in text
elif option_name.startswith("response_format"):
if option_value == OutputStruct:
assert response.value is not None
assert isinstance(response.value, OutputStruct)
assert "seattle" in response.value.location.lower()
else:
assert response.value is None
response_value = json.loads(response.text)
assert isinstance(response_value, dict)
assert "location" in response_value
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_foundry_integration_tests_disabled
@_with_foundry_debug()
async def test_integration_web_search() -> None:
client = FoundryChatClient(credential=AzureCliCredential())
web_search_tool = FoundryChatClient.get_web_search_tool()
content = {
"messages": [
Message(
role="user",
text="Who are the main characters of Kpop Demon Hunters? Do a web search to find the answer.",
)
],
"options": {"tool_choice": "auto", "tools": [web_search_tool]},
}
response = await client.get_response(stream=True, **content).get_final_response()
assert isinstance(response, ChatResponse)
assert "Rumi" in response.text
assert "Mira" in response.text
assert "Zoey" in response.text
@pytest.mark.flaky
@pytest.mark.integration
@skip_if_foundry_integration_tests_disabled
@_with_foundry_debug()
async def test_integration_tool_rich_content_image() -> None:
image_path = Path(__file__).parent.parent / "assets" / "sample_image.jpg"
image_bytes = image_path.read_bytes()
@tool(approval_mode="never_require")
def get_test_image() -> Content:
return Content.from_data(data=image_bytes, media_type="image/jpeg")
client = FoundryChatClient(credential=AzureCliCredential())
client.function_invocation_configuration["max_iterations"] = 2
messages = [Message(role="user", text="Call the get_test_image tool and describe what you see.")]
options: dict[str, Any] = {"tools": [get_test_image], "tool_choice": "auto"}
response = await client.get_response(messages=messages, options=options, stream=True).get_final_response()
assert isinstance(response, ChatResponse)
assert response.text is not None
assert len(response.text) > 0
assert "house" in response.text.lower(), f"Model did not describe the house image. Response: {response.text}"
def test_get_code_interpreter_tool() -> None:
"""Test code interpreter tool creation."""
tool_obj = RawFoundryChatClient.get_code_interpreter_tool()
assert tool_obj is not None
def test_get_code_interpreter_tool_with_file_ids() -> None:
"""Test code interpreter tool with file IDs."""
tool_obj = RawFoundryChatClient.get_code_interpreter_tool(file_ids=["file-abc123"])
assert tool_obj is not None
def test_get_file_search_tool() -> None:
"""Test file search tool creation."""
tool_obj = RawFoundryChatClient.get_file_search_tool(vector_store_ids=["vs_abc123"])
assert tool_obj is not None
def test_get_file_search_tool_requires_vector_store_ids() -> None:
"""Test that empty vector_store_ids raises ValueError."""
with pytest.raises(ValueError, match="vector_store_ids"):
RawFoundryChatClient.get_file_search_tool(vector_store_ids=[])
def test_get_web_search_tool() -> None:
"""Test web search tool creation."""
tool_obj = RawFoundryChatClient.get_web_search_tool()
assert tool_obj is not None
def test_get_web_search_tool_with_location() -> None:
"""Test web search tool with user location."""
tool_obj = RawFoundryChatClient.get_web_search_tool(
user_location={"city": "Seattle", "country": "US"},
search_context_size="high",
)
assert tool_obj is not None
def test_get_image_generation_tool() -> None:
"""Test image generation tool creation."""
tool_obj = RawFoundryChatClient.get_image_generation_tool()
assert tool_obj is not None
def test_get_mcp_tool() -> None:
"""Test MCP tool creation."""
tool_obj = RawFoundryChatClient.get_mcp_tool(
name="my_mcp",
url="https://mcp.example.com",
)
assert tool_obj is not None
def test_get_mcp_tool_with_connection_id() -> None:
"""Test MCP tool with project connection ID."""
tool_obj = RawFoundryChatClient.get_mcp_tool(
name="github_mcp",
project_connection_id="conn_abc123",
description="GitHub MCP via Foundry",
)
assert tool_obj is not None
@@ -0,0 +1,501 @@
# Copyright (c) Microsoft. All rights reserved.
# pyright: reportPrivateUsage=false
from __future__ import annotations
import os
from unittest.mock import AsyncMock, Mock, patch
import pytest
from agent_framework import AGENT_FRAMEWORK_USER_AGENT, AgentResponse, Message
from agent_framework._sessions import AgentSession, SessionContext
from agent_framework_foundry._memory_provider import FoundryMemoryProvider
@pytest.fixture
def mock_project_client() -> AsyncMock:
"""Create a mock AIProjectClient."""
mock_client = AsyncMock()
mock_client.beta = AsyncMock()
mock_client.beta.memory_stores = AsyncMock()
mock_client.beta.memory_stores.search_memories = AsyncMock()
mock_client.beta.memory_stores.begin_update_memories = AsyncMock()
mock_client.__aenter__ = AsyncMock(return_value=mock_client)
mock_client.__aexit__ = AsyncMock()
return mock_client
@pytest.fixture
def mock_credential() -> Mock:
"""Create a mock Azure credential."""
return Mock()
# -- Initialization tests ------------------------------------------------------
def test_init_with_all_params(mock_project_client: AsyncMock) -> None:
provider = FoundryMemoryProvider(
source_id="custom_source",
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
context_prompt="Custom prompt",
update_delay=60,
)
assert provider.source_id == "custom_source"
assert provider.project_client is mock_project_client
assert provider.memory_store_name == "test_store"
assert provider.scope == "user_123"
assert provider.context_prompt == "Custom prompt"
assert provider.update_delay == 60
def test_init_default_source_id(mock_project_client: AsyncMock) -> None:
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
assert provider.source_id == FoundryMemoryProvider.DEFAULT_SOURCE_ID
def test_init_default_context_prompt(mock_project_client: AsyncMock) -> None:
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
assert provider.context_prompt == FoundryMemoryProvider.DEFAULT_CONTEXT_PROMPT
def test_init_default_update_delay(mock_project_client: AsyncMock) -> None:
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
assert provider.update_delay == 300
def test_init_with_project_endpoint_and_credential(mock_project_client: AsyncMock, mock_credential: Mock) -> None:
with patch("agent_framework_foundry._memory_provider.AIProjectClient") as mock_ai_project_client:
mock_ai_project_client.return_value = mock_project_client
provider = FoundryMemoryProvider(
project_endpoint="https://test.project.endpoint",
credential=mock_credential, # type: ignore[arg-type]
allow_preview=True,
memory_store_name="test_store",
scope="user_123",
)
assert provider.project_client is mock_project_client
mock_ai_project_client.assert_called_once_with(
endpoint="https://test.project.endpoint",
credential=mock_credential,
allow_preview=True,
user_agent=AGENT_FRAMEWORK_USER_AGENT,
)
def test_init_requires_project_endpoint_without_project_client() -> None:
with (
patch("agent_framework_foundry._memory_provider.load_settings") as mock_load_settings,
patch.dict(os.environ, {}, clear=True),
pytest.raises(ValueError, match="project endpoint is required"),
):
mock_load_settings.return_value = {"project_endpoint": None}
FoundryMemoryProvider(
memory_store_name="test_store",
scope="user_123",
)
def test_init_requires_credential_without_project_client() -> None:
with pytest.raises(ValueError, match="Azure credential is required"):
FoundryMemoryProvider(
project_endpoint="https://test.project.endpoint",
memory_store_name="test_store",
scope="user_123",
)
def test_init_requires_memory_store_name(mock_project_client: AsyncMock) -> None:
with pytest.raises(ValueError, match="memory_store_name is required"):
FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="",
scope="user_123",
)
def test_init_requires_scope(mock_project_client: AsyncMock) -> None:
with pytest.raises(ValueError, match="scope is required"):
FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="",
)
# -- before_run tests ----------------------------------------------------------
async def test_retrieves_static_memories_on_first_run(mock_project_client: AsyncMock) -> None:
mem1 = Mock()
mem1.memory_item.content = "User prefers Python"
mem2 = Mock()
mem2.memory_item.content = "User is based in Seattle"
mock_search_result = Mock()
mock_search_result.memories = [mem1, mem2]
mock_project_client.beta.memory_stores.search_memories.return_value = mock_search_result
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
session = AgentSession(session_id="test-session")
ctx = SessionContext(input_messages=[Message(role="user", text="Hello")], session_id="s1")
await provider.before_run( # type: ignore[arg-type]
agent=None, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
)
# Should call search_memories twice: once for static, once for contextual
assert mock_project_client.beta.memory_stores.search_memories.call_count == 2
# Static memories should be cached
assert len(session.state[provider.source_id]["static_memories"]) == 2
assert session.state[provider.source_id]["initialized"] is True
async def test_contextual_memories_added_to_context(mock_project_client: AsyncMock) -> None:
# Mock static search (first call)
static_mem = Mock()
static_mem.memory_item.content = "User prefers Python"
static_result = Mock()
static_result.memories = [static_mem]
# Mock contextual search (second call)
contextual_mem = Mock()
contextual_mem.memory_item.content = "Last discussed async patterns"
contextual_result = Mock()
contextual_result.memories = [contextual_mem]
contextual_result.search_id = "search-123"
mock_project_client.beta.memory_stores.search_memories.side_effect = [static_result, contextual_result]
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
session = AgentSession(session_id="test-session")
ctx = SessionContext(input_messages=[Message(role="user", text="Hello")], session_id="s1")
await provider.before_run( # type: ignore[arg-type]
agent=None, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
)
# Check that memories were added to context
assert provider.source_id in ctx.context_messages
added = ctx.context_messages[provider.source_id]
assert len(added) == 1
assert "User prefers Python" in added[0].text # type: ignore[operator]
assert "Last discussed async patterns" in added[0].text # type: ignore[operator]
assert provider.context_prompt in added[0].text # type: ignore[operator]
assert session.state[provider.source_id]["previous_search_id"] == "search-123"
async def test_empty_input_skips_contextual_search(mock_project_client: AsyncMock) -> None:
static_result = Mock()
static_result.memories = []
mock_project_client.beta.memory_stores.search_memories.return_value = static_result
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
session = AgentSession(session_id="test-session")
ctx = SessionContext(input_messages=[Message(role="user", text="")], session_id="s1")
await provider.before_run( # type: ignore[arg-type]
agent=None, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
)
# Should only call search_memories once for static memories
assert mock_project_client.beta.memory_stores.search_memories.call_count == 1
assert provider.source_id not in ctx.context_messages
async def test_empty_search_results_no_messages(mock_project_client: AsyncMock) -> None:
mock_search_result = Mock()
mock_search_result.memories = []
mock_project_client.beta.memory_stores.search_memories.return_value = mock_search_result
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
session = AgentSession(session_id="test-session")
ctx = SessionContext(input_messages=[Message(role="user", text="test")], session_id="s1")
await provider.before_run( # type: ignore[arg-type]
agent=None, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
)
assert provider.source_id not in ctx.context_messages
async def test_static_memories_only_retrieved_once(mock_project_client: AsyncMock) -> None:
static_mem = Mock()
static_mem.memory_item.content = "Static memory"
static_result = Mock()
static_result.memories = [static_mem]
contextual_result = Mock()
contextual_result.memories = []
mock_project_client.beta.memory_stores.search_memories.side_effect = [static_result, contextual_result]
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
session = AgentSession(session_id="test-session")
ctx = SessionContext(input_messages=[Message(role="user", text="Hello")], session_id="s1")
# First call
await provider.before_run( # type: ignore[arg-type]
agent=None, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
)
assert mock_project_client.beta.memory_stores.search_memories.call_count == 2
# Reset mock for second call
mock_project_client.beta.memory_stores.search_memories.reset_mock()
contextual_result2 = Mock()
contextual_result2.memories = []
mock_project_client.beta.memory_stores.search_memories.return_value = contextual_result2
# Second call - should only search contextual, not static
ctx2 = SessionContext(input_messages=[Message(role="user", text="World")], session_id="s1")
await provider.before_run( # type: ignore[arg-type]
agent=None, session=session, context=ctx2, state=session.state.setdefault(provider.source_id, {})
)
assert mock_project_client.beta.memory_stores.search_memories.call_count == 1
async def test_handles_search_exception_gracefully(mock_project_client: AsyncMock) -> None:
mock_project_client.beta.memory_stores.search_memories.side_effect = Exception("API error")
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
session = AgentSession(session_id="test-session")
ctx = SessionContext(input_messages=[Message(role="user", text="Hello")], session_id="s1")
# Should not raise exception
await provider.before_run( # type: ignore[arg-type]
agent=None, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
)
# No memories added
assert provider.source_id not in ctx.context_messages
# -- after_run tests -----------------------------------------------------------
async def test_stores_input_and_response(mock_project_client: AsyncMock) -> None:
mock_poller = Mock()
mock_poller.update_id = "update-456"
mock_project_client.beta.memory_stores.begin_update_memories.return_value = mock_poller
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
session = AgentSession(session_id="test-session")
ctx = SessionContext(input_messages=[Message(role="user", text="question")], session_id="s1")
ctx._response = AgentResponse(messages=[Message(role="assistant", text="answer")])
await provider.after_run( # type: ignore[arg-type]
agent=None, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
)
mock_project_client.beta.memory_stores.begin_update_memories.assert_awaited_once()
call_kwargs = mock_project_client.beta.memory_stores.begin_update_memories.call_args.kwargs
assert call_kwargs["name"] == "test_store"
assert call_kwargs["scope"] == "user_123"
assert len(call_kwargs["items"]) == 2
assert call_kwargs["items"][0]["content"] == "question"
assert call_kwargs["items"][1]["content"] == "answer"
assert session.state[provider.source_id]["previous_update_id"] == "update-456"
async def test_only_stores_user_assistant_system(mock_project_client: AsyncMock) -> None:
mock_poller = Mock()
mock_project_client.beta.memory_stores.begin_update_memories.return_value = mock_poller
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
session = AgentSession(session_id="test-session")
ctx = SessionContext(
input_messages=[
Message(role="user", text="hello"),
Message(role="tool", text="tool output"),
],
session_id="s1",
)
ctx._response = AgentResponse(messages=[Message(role="assistant", text="reply")])
await provider.after_run( # type: ignore[arg-type]
agent=None, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
)
call_kwargs = mock_project_client.beta.memory_stores.begin_update_memories.call_args.kwargs
items = call_kwargs["items"]
assert len(items) == 2
assert items[0]["content"] == "hello"
assert items[1]["content"] == "reply"
async def test_skips_empty_messages(mock_project_client: AsyncMock) -> None:
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
session = AgentSession(session_id="test-session")
ctx = SessionContext(
input_messages=[
Message(role="user", text=""),
Message(role="user", text=" "),
],
session_id="s1",
)
ctx._response = AgentResponse(messages=[])
await provider.after_run( # type: ignore[arg-type]
agent=None, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
)
mock_project_client.beta.memory_stores.begin_update_memories.assert_not_awaited()
async def test_uses_configured_update_delay(mock_project_client: AsyncMock) -> None:
mock_poller = Mock()
mock_project_client.beta.memory_stores.begin_update_memories.return_value = mock_poller
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
update_delay=60,
)
session = AgentSession(session_id="test-session")
ctx = SessionContext(input_messages=[Message(role="user", text="hi")], session_id="s1")
ctx._response = AgentResponse(messages=[Message(role="assistant", text="hey")])
await provider.after_run( # type: ignore[arg-type]
agent=None, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
)
call_kwargs = mock_project_client.beta.memory_stores.begin_update_memories.call_args.kwargs
assert call_kwargs["update_delay"] == 60
async def test_uses_previous_update_id_for_incremental_updates(mock_project_client: AsyncMock) -> None:
mock_poller1 = Mock()
mock_poller1.update_id = "update-1"
mock_poller2 = Mock()
mock_poller2.update_id = "update-2"
mock_project_client.beta.memory_stores.begin_update_memories.side_effect = [mock_poller1, mock_poller2]
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
session = AgentSession(session_id="test-session")
ctx1 = SessionContext(input_messages=[Message(role="user", text="first")], session_id="s1")
ctx1._response = AgentResponse(messages=[Message(role="assistant", text="response1")])
# First update
await provider.after_run( # type: ignore[arg-type]
agent=None, session=session, context=ctx1, state=session.state.setdefault(provider.source_id, {})
)
assert session.state[provider.source_id]["previous_update_id"] == "update-1"
# Second update should use previous_update_id
ctx2 = SessionContext(input_messages=[Message(role="user", text="second")], session_id="s1")
ctx2._response = AgentResponse(messages=[Message(role="assistant", text="response2")])
await provider.after_run( # type: ignore[arg-type]
agent=None, session=session, context=ctx2, state=session.state.setdefault(provider.source_id, {})
)
call_kwargs = mock_project_client.beta.memory_stores.begin_update_memories.call_args.kwargs
assert call_kwargs["previous_update_id"] == "update-1"
assert session.state[provider.source_id]["previous_update_id"] == "update-2"
async def test_handles_update_exception_gracefully(mock_project_client: AsyncMock) -> None:
mock_project_client.beta.memory_stores.begin_update_memories.side_effect = Exception("API error")
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
session = AgentSession(session_id="test-session")
ctx = SessionContext(input_messages=[Message(role="user", text="hi")], session_id="s1")
ctx._response = AgentResponse(messages=[Message(role="assistant", text="hey")])
# Should not raise exception
await provider.after_run( # type: ignore[arg-type]
agent=None, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
)
# -- Context manager tests -----------------------------------------------------
async def test_aenter_delegates_to_client(mock_project_client: AsyncMock) -> None:
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
result = await provider.__aenter__()
assert result is provider
mock_project_client.__aenter__.assert_awaited_once()
async def test_aexit_delegates_to_client(mock_project_client: AsyncMock) -> None:
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
await provider.__aexit__(None, None, None)
mock_project_client.__aexit__.assert_awaited_once()
async def test_async_with_syntax(mock_project_client: AsyncMock) -> None:
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
async with provider as p:
assert p is provider
@@ -1,374 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
"""Tests for FoundryAgentClient and FoundryAgent classes."""
from typing import Any
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from agent_framework._tools import tool
class TestRawFoundryAgentChatClient:
"""Tests for RawFoundryAgentChatClient."""
def test_init_requires_agent_name(self) -> None:
"""Test that agent_name is required."""
from agent_framework_foundry._foundry_agent_client import RawFoundryAgentChatClient
with pytest.raises(ValueError, match="Agent name is required"):
RawFoundryAgentChatClient(
project_client=MagicMock(),
)
def test_init_with_agent_name(self) -> None:
"""Test construction with agent_name and project_client."""
from agent_framework_foundry._foundry_agent_client import RawFoundryAgentChatClient
mock_project = MagicMock()
mock_project.get_openai_client.return_value = MagicMock()
client = RawFoundryAgentChatClient(
project_client=mock_project,
agent_name="test-agent",
agent_version="1.0",
)
assert client.agent_name == "test-agent"
assert client.agent_version == "1.0"
def test_get_agent_reference_with_version(self) -> None:
"""Test agent reference includes version when provided."""
from agent_framework_foundry._foundry_agent_client import RawFoundryAgentChatClient
mock_project = MagicMock()
mock_project.get_openai_client.return_value = MagicMock()
client = RawFoundryAgentChatClient(
project_client=mock_project,
agent_name="my-agent",
agent_version="2.0",
)
ref = client._get_agent_reference()
assert ref == {"name": "my-agent", "version": "2.0", "type": "agent_reference"}
def test_get_agent_reference_without_version(self) -> None:
"""Test agent reference omits version for HostedAgents."""
from agent_framework_foundry._foundry_agent_client import RawFoundryAgentChatClient
mock_project = MagicMock()
mock_project.get_openai_client.return_value = MagicMock()
client = RawFoundryAgentChatClient(
project_client=mock_project,
agent_name="hosted-agent",
)
ref = client._get_agent_reference()
assert ref == {"name": "hosted-agent", "type": "agent_reference"}
assert "version" not in ref
def test_as_agent_returns_foundry_agent_and_preserves_client_type(self) -> None:
"""Test that as_agent() wraps the client in FoundryAgent using the same client class."""
from agent_framework_foundry._foundry_agent import FoundryAgent
from agent_framework_foundry._foundry_agent_client import RawFoundryAgentChatClient
class CustomClient(RawFoundryAgentChatClient):
pass
mock_project = MagicMock()
mock_project.get_openai_client.return_value = MagicMock()
client = CustomClient(
project_client=mock_project,
agent_name="test-agent",
agent_version="1.0",
)
agent = client.as_agent(instructions="You are helpful.")
assert isinstance(agent, FoundryAgent)
assert agent.name == "test-agent"
assert isinstance(agent.client, CustomClient)
assert agent.client.project_client is mock_project
assert agent.client.agent_name == "test-agent"
assert agent.client.agent_version == "1.0"
named_agent = client.as_agent(name="display-name", instructions="You are helpful.")
assert named_agent.name == "display-name"
assert named_agent.client.agent_name == "test-agent"
async def test_prepare_options_validates_tools(self) -> None:
"""Test that _prepare_options rejects non-FunctionTool objects."""
from agent_framework import Message
from agent_framework_foundry._foundry_agent_client import RawFoundryAgentChatClient
mock_project = MagicMock()
mock_project.get_openai_client.return_value = MagicMock()
client = RawFoundryAgentChatClient(
project_client=mock_project,
agent_name="test-agent",
)
# A dict tool should be rejected
with pytest.raises(TypeError, match="Only FunctionTool objects are accepted"):
await client._prepare_options(
messages=[Message(role="user", contents="hi")],
options={"tools": [{"type": "function", "function": {"name": "bad"}}]},
)
async def test_prepare_options_accepts_function_tools(self) -> None:
"""Test that _prepare_options accepts FunctionTool objects."""
from agent_framework import Message
from agent_framework_foundry._foundry_agent_client import RawFoundryAgentChatClient
mock_project = MagicMock()
mock_openai = MagicMock()
mock_project.get_openai_client.return_value = mock_openai
client = RawFoundryAgentChatClient(
project_client=mock_project,
agent_name="test-agent",
)
@tool(approval_mode="never_require")
def my_func() -> str:
"""A test function."""
return "ok"
# Should not raise — patch the parent's _prepare_options
with patch(
"agent_framework_openai._chat_client.RawOpenAIChatClient._prepare_options",
new_callable=AsyncMock,
return_value={},
):
result = await client._prepare_options(
messages=[Message(role="user", contents="hi")],
options={"tools": [my_func]},
)
assert "extra_body" in result
assert result["extra_body"]["agent_reference"]["name"] == "test-agent"
def test_check_model_presence_is_noop(self) -> None:
"""Test that _check_model_presence does nothing (model is on service)."""
from agent_framework_foundry._foundry_agent_client import RawFoundryAgentChatClient
mock_project = MagicMock()
mock_project.get_openai_client.return_value = MagicMock()
client = RawFoundryAgentChatClient(
project_client=mock_project,
agent_name="test-agent",
)
options: dict[str, Any] = {}
client._check_model_presence(options)
assert "model" not in options
class TestFoundryAgentChatClient:
"""Tests for _FoundryAgentChatClient (full middleware)."""
def test_init(self) -> None:
"""Test construction of the full-middleware client."""
from agent_framework_foundry._foundry_agent_client import _FoundryAgentChatClient
mock_project = MagicMock()
mock_project.get_openai_client.return_value = MagicMock()
client = _FoundryAgentChatClient(
project_client=mock_project,
agent_name="test-agent",
agent_version="1.0",
)
assert client.agent_name == "test-agent"
class TestRawFoundryAgent:
"""Tests for RawFoundryAgent."""
def test_init_creates_client(self) -> None:
"""Test that RawFoundryAgent creates a client internally."""
from agent_framework_foundry._foundry_agent import RawFoundryAgent
mock_project = MagicMock()
mock_project.get_openai_client.return_value = MagicMock()
agent = RawFoundryAgent(
project_client=mock_project,
agent_name="test-agent",
agent_version="1.0",
)
assert agent.client is not None
assert agent.client.agent_name == "test-agent"
def test_init_with_custom_client_type(self) -> None:
"""Test that client_type parameter is respected."""
from agent_framework_foundry._foundry_agent import RawFoundryAgent
from agent_framework_foundry._foundry_agent_client import RawFoundryAgentChatClient
mock_project = MagicMock()
mock_project.get_openai_client.return_value = MagicMock()
agent = RawFoundryAgent(
project_client=mock_project,
agent_name="test-agent",
client_type=RawFoundryAgentChatClient,
)
assert isinstance(agent.client, RawFoundryAgentChatClient)
def test_init_rejects_invalid_client_type(self) -> None:
"""Test that invalid client_type raises TypeError."""
from agent_framework_foundry._foundry_agent import RawFoundryAgent
with pytest.raises(TypeError, match="must be a subclass of RawFoundryAgentChatClient"):
RawFoundryAgent(
project_client=MagicMock(),
agent_name="test-agent",
client_type=object, # type: ignore[arg-type]
)
def test_init_with_function_tools(self) -> None:
"""Test that FunctionTool and callables are accepted."""
from agent_framework_foundry._foundry_agent import RawFoundryAgent
mock_project = MagicMock()
mock_project.get_openai_client.return_value = MagicMock()
@tool(approval_mode="never_require")
def my_func() -> str:
"""A test function."""
return "ok"
agent = RawFoundryAgent(
project_client=mock_project,
agent_name="test-agent",
tools=[my_func],
)
assert agent.default_options.get("tools") is not None
class TestFoundryAgent:
"""Tests for FoundryAgent (full middleware)."""
def test_init(self) -> None:
"""Test construction of the full-middleware agent."""
from agent_framework_foundry._foundry_agent import FoundryAgent
mock_project = MagicMock()
mock_project.get_openai_client.return_value = MagicMock()
agent = FoundryAgent(
project_client=mock_project,
agent_name="test-agent",
agent_version="1.0",
)
assert agent.client is not None
assert agent.client.agent_name == "test-agent"
def test_init_with_middleware(self) -> None:
"""Test that agent-level middleware is accepted."""
from agent_framework import ChatContext, ChatMiddleware
from agent_framework_foundry._foundry_agent import FoundryAgent
mock_project = MagicMock()
mock_project.get_openai_client.return_value = MagicMock()
class MyMiddleware(ChatMiddleware):
async def process(self, context: ChatContext) -> None:
pass
agent = FoundryAgent(
project_client=mock_project,
agent_name="test-agent",
middleware=[MyMiddleware()],
)
assert agent.client is not None
class TestFoundryChatClientToolMethods:
"""Tests for RawFoundryChatClient tool factory methods."""
def test_get_code_interpreter_tool(self) -> None:
"""Test code interpreter tool creation."""
from agent_framework_foundry._foundry_chat_client import RawFoundryChatClient
tool_obj = RawFoundryChatClient.get_code_interpreter_tool()
assert tool_obj is not None
def test_get_code_interpreter_tool_with_file_ids(self) -> None:
"""Test code interpreter tool with file IDs."""
from agent_framework_foundry._foundry_chat_client import RawFoundryChatClient
tool_obj = RawFoundryChatClient.get_code_interpreter_tool(file_ids=["file-abc123"])
assert tool_obj is not None
def test_get_file_search_tool(self) -> None:
"""Test file search tool creation."""
from agent_framework_foundry._foundry_chat_client import RawFoundryChatClient
tool_obj = RawFoundryChatClient.get_file_search_tool(vector_store_ids=["vs_abc123"])
assert tool_obj is not None
def test_get_file_search_tool_requires_vector_store_ids(self) -> None:
"""Test that empty vector_store_ids raises ValueError."""
from agent_framework_foundry._foundry_chat_client import RawFoundryChatClient
with pytest.raises(ValueError, match="vector_store_ids"):
RawFoundryChatClient.get_file_search_tool(vector_store_ids=[])
def test_get_web_search_tool(self) -> None:
"""Test web search tool creation."""
from agent_framework_foundry._foundry_chat_client import RawFoundryChatClient
tool_obj = RawFoundryChatClient.get_web_search_tool()
assert tool_obj is not None
def test_get_web_search_tool_with_location(self) -> None:
"""Test web search tool with user location."""
from agent_framework_foundry._foundry_chat_client import RawFoundryChatClient
tool_obj = RawFoundryChatClient.get_web_search_tool(
user_location={"city": "Seattle", "country": "US"},
search_context_size="high",
)
assert tool_obj is not None
def test_get_image_generation_tool(self) -> None:
"""Test image generation tool creation."""
from agent_framework_foundry._foundry_chat_client import RawFoundryChatClient
tool_obj = RawFoundryChatClient.get_image_generation_tool()
assert tool_obj is not None
def test_get_mcp_tool(self) -> None:
"""Test MCP tool creation."""
from agent_framework_foundry._foundry_chat_client import RawFoundryChatClient
tool_obj = RawFoundryChatClient.get_mcp_tool(
name="my_mcp",
url="https://mcp.example.com",
)
assert tool_obj is not None
def test_get_mcp_tool_with_connection_id(self) -> None:
"""Test MCP tool with project connection ID."""
from agent_framework_foundry._foundry_chat_client import RawFoundryChatClient
tool_obj = RawFoundryChatClient.get_mcp_tool(
name="github_mcp",
project_connection_id="conn_abc123",
description="GitHub MCP via Foundry",
)
assert tool_obj is not None
@@ -1,507 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
# pyright: reportPrivateUsage=false
from __future__ import annotations
import os
from unittest.mock import AsyncMock, Mock, patch
import pytest
from agent_framework import AGENT_FRAMEWORK_USER_AGENT, AgentResponse, Message
from agent_framework._sessions import AgentSession, SessionContext
from agent_framework_foundry._foundry_memory_provider import FoundryMemoryProvider
@pytest.fixture
def mock_project_client() -> AsyncMock:
"""Create a mock AIProjectClient."""
mock_client = AsyncMock()
mock_client.beta = AsyncMock()
mock_client.beta.memory_stores = AsyncMock()
mock_client.beta.memory_stores.search_memories = AsyncMock()
mock_client.beta.memory_stores.begin_update_memories = AsyncMock()
mock_client.__aenter__ = AsyncMock(return_value=mock_client)
mock_client.__aexit__ = AsyncMock()
return mock_client
@pytest.fixture
def mock_credential() -> Mock:
"""Create a mock Azure credential."""
return Mock()
# -- Initialization tests ------------------------------------------------------
class TestInit:
"""Test FoundryMemoryProvider initialization."""
def test_init_with_all_params(self, mock_project_client: AsyncMock) -> None:
provider = FoundryMemoryProvider(
source_id="custom_source",
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
context_prompt="Custom prompt",
update_delay=60,
)
assert provider.source_id == "custom_source"
assert provider.project_client is mock_project_client
assert provider.memory_store_name == "test_store"
assert provider.scope == "user_123"
assert provider.context_prompt == "Custom prompt"
assert provider.update_delay == 60
def test_init_default_source_id(self, mock_project_client: AsyncMock) -> None:
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
assert provider.source_id == FoundryMemoryProvider.DEFAULT_SOURCE_ID
def test_init_default_context_prompt(self, mock_project_client: AsyncMock) -> None:
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
assert provider.context_prompt == FoundryMemoryProvider.DEFAULT_CONTEXT_PROMPT
def test_init_default_update_delay(self, mock_project_client: AsyncMock) -> None:
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
assert provider.update_delay == 300
def test_init_with_project_endpoint_and_credential(
self, mock_project_client: AsyncMock, mock_credential: Mock
) -> None:
with patch("agent_framework_foundry._foundry_memory_provider.AIProjectClient") as mock_ai_project_client:
mock_ai_project_client.return_value = mock_project_client
provider = FoundryMemoryProvider(
project_endpoint="https://test.project.endpoint",
credential=mock_credential, # type: ignore[arg-type]
allow_preview=True,
memory_store_name="test_store",
scope="user_123",
)
assert provider.project_client is mock_project_client
mock_ai_project_client.assert_called_once_with(
endpoint="https://test.project.endpoint",
credential=mock_credential,
allow_preview=True,
user_agent=AGENT_FRAMEWORK_USER_AGENT,
)
def test_init_requires_project_endpoint_without_project_client(self) -> None:
with (
patch("agent_framework_foundry._foundry_memory_provider.load_settings") as mock_load_settings,
patch.dict(os.environ, {}, clear=True),
pytest.raises(ValueError, match="project endpoint is required"),
):
mock_load_settings.return_value = {"project_endpoint": None}
FoundryMemoryProvider(
memory_store_name="test_store",
scope="user_123",
)
def test_init_requires_credential_without_project_client(self) -> None:
with pytest.raises(ValueError, match="Azure credential is required"):
FoundryMemoryProvider(
project_endpoint="https://test.project.endpoint",
memory_store_name="test_store",
scope="user_123",
)
def test_init_requires_memory_store_name(self, mock_project_client: AsyncMock) -> None:
with pytest.raises(ValueError, match="memory_store_name is required"):
FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="",
scope="user_123",
)
def test_init_requires_scope(self, mock_project_client: AsyncMock) -> None:
with pytest.raises(ValueError, match="scope is required"):
FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="",
)
# -- before_run tests ----------------------------------------------------------
class TestBeforeRun:
"""Test before_run hook."""
async def test_retrieves_static_memories_on_first_run(self, mock_project_client: AsyncMock) -> None:
"""First call retrieves static (user profile) memories."""
mem1 = Mock()
mem1.memory_item.content = "User prefers Python"
mem2 = Mock()
mem2.memory_item.content = "User is based in Seattle"
mock_search_result = Mock()
mock_search_result.memories = [mem1, mem2]
mock_project_client.beta.memory_stores.search_memories.return_value = mock_search_result
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
session = AgentSession(session_id="test-session")
ctx = SessionContext(input_messages=[Message(role="user", text="Hello")], session_id="s1")
await provider.before_run( # type: ignore[arg-type]
agent=None, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
)
# Should call search_memories twice: once for static, once for contextual
assert mock_project_client.beta.memory_stores.search_memories.call_count == 2
# Static memories should be cached
assert len(session.state[provider.source_id]["static_memories"]) == 2
assert session.state[provider.source_id]["initialized"] is True
async def test_contextual_memories_added_to_context(self, mock_project_client: AsyncMock) -> None:
"""Contextual search returns memories → messages added to context with prompt."""
# Mock static search (first call)
static_mem = Mock()
static_mem.memory_item.content = "User prefers Python"
static_result = Mock()
static_result.memories = [static_mem]
# Mock contextual search (second call)
contextual_mem = Mock()
contextual_mem.memory_item.content = "Last discussed async patterns"
contextual_result = Mock()
contextual_result.memories = [contextual_mem]
contextual_result.search_id = "search-123"
mock_project_client.beta.memory_stores.search_memories.side_effect = [static_result, contextual_result]
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
session = AgentSession(session_id="test-session")
ctx = SessionContext(input_messages=[Message(role="user", text="Hello")], session_id="s1")
await provider.before_run( # type: ignore[arg-type]
agent=None, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
)
# Check that memories were added to context
assert provider.source_id in ctx.context_messages
added = ctx.context_messages[provider.source_id]
assert len(added) == 1
assert "User prefers Python" in added[0].text # type: ignore[operator]
assert "Last discussed async patterns" in added[0].text # type: ignore[operator]
assert provider.context_prompt in added[0].text # type: ignore[operator]
assert session.state[provider.source_id]["previous_search_id"] == "search-123"
async def test_empty_input_skips_contextual_search(self, mock_project_client: AsyncMock) -> None:
"""Empty input messages → only static search performed, no contextual search."""
static_result = Mock()
static_result.memories = []
mock_project_client.beta.memory_stores.search_memories.return_value = static_result
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
session = AgentSession(session_id="test-session")
ctx = SessionContext(input_messages=[Message(role="user", text="")], session_id="s1")
await provider.before_run( # type: ignore[arg-type]
agent=None, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
)
# Should only call search_memories once for static memories
assert mock_project_client.beta.memory_stores.search_memories.call_count == 1
assert provider.source_id not in ctx.context_messages
async def test_empty_search_results_no_messages(self, mock_project_client: AsyncMock) -> None:
"""Empty search results → no messages added."""
mock_search_result = Mock()
mock_search_result.memories = []
mock_project_client.beta.memory_stores.search_memories.return_value = mock_search_result
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
session = AgentSession(session_id="test-session")
ctx = SessionContext(input_messages=[Message(role="user", text="test")], session_id="s1")
await provider.before_run( # type: ignore[arg-type]
agent=None, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
)
assert provider.source_id not in ctx.context_messages
async def test_static_memories_only_retrieved_once(self, mock_project_client: AsyncMock) -> None:
"""Static memories are only retrieved on the first call."""
static_mem = Mock()
static_mem.memory_item.content = "Static memory"
static_result = Mock()
static_result.memories = [static_mem]
contextual_result = Mock()
contextual_result.memories = []
mock_project_client.beta.memory_stores.search_memories.side_effect = [static_result, contextual_result]
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
session = AgentSession(session_id="test-session")
ctx = SessionContext(input_messages=[Message(role="user", text="Hello")], session_id="s1")
# First call
await provider.before_run( # type: ignore[arg-type]
agent=None, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
)
assert mock_project_client.beta.memory_stores.search_memories.call_count == 2
# Reset mock for second call
mock_project_client.beta.memory_stores.search_memories.reset_mock()
contextual_result2 = Mock()
contextual_result2.memories = []
mock_project_client.beta.memory_stores.search_memories.return_value = contextual_result2
# Second call - should only search contextual, not static
ctx2 = SessionContext(input_messages=[Message(role="user", text="World")], session_id="s1")
await provider.before_run( # type: ignore[arg-type]
agent=None, session=session, context=ctx2, state=session.state.setdefault(provider.source_id, {})
)
assert mock_project_client.beta.memory_stores.search_memories.call_count == 1
async def test_handles_search_exception_gracefully(self, mock_project_client: AsyncMock) -> None:
"""Search exception is logged but doesn't fail the operation."""
mock_project_client.beta.memory_stores.search_memories.side_effect = Exception("API error")
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
session = AgentSession(session_id="test-session")
ctx = SessionContext(input_messages=[Message(role="user", text="Hello")], session_id="s1")
# Should not raise exception
await provider.before_run( # type: ignore[arg-type]
agent=None, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
)
# No memories added
assert provider.source_id not in ctx.context_messages
# -- after_run tests -----------------------------------------------------------
class TestAfterRun:
"""Test after_run hook."""
async def test_stores_input_and_response(self, mock_project_client: AsyncMock) -> None:
"""Stores input+response messages via begin_update_memories."""
mock_poller = Mock()
mock_poller.update_id = "update-456"
mock_project_client.beta.memory_stores.begin_update_memories.return_value = mock_poller
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
session = AgentSession(session_id="test-session")
ctx = SessionContext(input_messages=[Message(role="user", text="question")], session_id="s1")
ctx._response = AgentResponse(messages=[Message(role="assistant", text="answer")])
await provider.after_run( # type: ignore[arg-type]
agent=None, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
)
mock_project_client.beta.memory_stores.begin_update_memories.assert_awaited_once()
call_kwargs = mock_project_client.beta.memory_stores.begin_update_memories.call_args.kwargs
assert call_kwargs["name"] == "test_store"
assert call_kwargs["scope"] == "user_123"
assert len(call_kwargs["items"]) == 2
assert call_kwargs["items"][0]["content"] == "question"
assert call_kwargs["items"][1]["content"] == "answer"
assert session.state[provider.source_id]["previous_update_id"] == "update-456"
async def test_only_stores_user_assistant_system(self, mock_project_client: AsyncMock) -> None:
"""Only stores user/assistant/system messages with text."""
mock_poller = Mock()
mock_project_client.beta.memory_stores.begin_update_memories.return_value = mock_poller
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
session = AgentSession(session_id="test-session")
ctx = SessionContext(
input_messages=[
Message(role="user", text="hello"),
Message(role="tool", text="tool output"),
],
session_id="s1",
)
ctx._response = AgentResponse(messages=[Message(role="assistant", text="reply")])
await provider.after_run( # type: ignore[arg-type]
agent=None, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
)
call_kwargs = mock_project_client.beta.memory_stores.begin_update_memories.call_args.kwargs
items = call_kwargs["items"]
assert len(items) == 2
assert items[0]["content"] == "hello"
assert items[1]["content"] == "reply"
async def test_skips_empty_messages(self, mock_project_client: AsyncMock) -> None:
"""Skips messages with empty text."""
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
session = AgentSession(session_id="test-session")
ctx = SessionContext(
input_messages=[
Message(role="user", text=""),
Message(role="user", text=" "),
],
session_id="s1",
)
ctx._response = AgentResponse(messages=[])
await provider.after_run( # type: ignore[arg-type]
agent=None, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
)
mock_project_client.beta.memory_stores.begin_update_memories.assert_not_awaited()
async def test_uses_configured_update_delay(self, mock_project_client: AsyncMock) -> None:
"""Uses the configured update_delay parameter."""
mock_poller = Mock()
mock_project_client.beta.memory_stores.begin_update_memories.return_value = mock_poller
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
update_delay=60,
)
session = AgentSession(session_id="test-session")
ctx = SessionContext(input_messages=[Message(role="user", text="hi")], session_id="s1")
ctx._response = AgentResponse(messages=[Message(role="assistant", text="hey")])
await provider.after_run( # type: ignore[arg-type]
agent=None, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
)
call_kwargs = mock_project_client.beta.memory_stores.begin_update_memories.call_args.kwargs
assert call_kwargs["update_delay"] == 60
async def test_uses_previous_update_id_for_incremental_updates(self, mock_project_client: AsyncMock) -> None:
"""Uses previous_update_id for incremental updates."""
mock_poller1 = Mock()
mock_poller1.update_id = "update-1"
mock_poller2 = Mock()
mock_poller2.update_id = "update-2"
mock_project_client.beta.memory_stores.begin_update_memories.side_effect = [mock_poller1, mock_poller2]
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
session = AgentSession(session_id="test-session")
ctx1 = SessionContext(input_messages=[Message(role="user", text="first")], session_id="s1")
ctx1._response = AgentResponse(messages=[Message(role="assistant", text="response1")])
# First update
await provider.after_run( # type: ignore[arg-type]
agent=None, session=session, context=ctx1, state=session.state.setdefault(provider.source_id, {})
)
assert session.state[provider.source_id]["previous_update_id"] == "update-1"
# Second update should use previous_update_id
ctx2 = SessionContext(input_messages=[Message(role="user", text="second")], session_id="s1")
ctx2._response = AgentResponse(messages=[Message(role="assistant", text="response2")])
await provider.after_run( # type: ignore[arg-type]
agent=None, session=session, context=ctx2, state=session.state.setdefault(provider.source_id, {})
)
call_kwargs = mock_project_client.beta.memory_stores.begin_update_memories.call_args.kwargs
assert call_kwargs["previous_update_id"] == "update-1"
assert session.state[provider.source_id]["previous_update_id"] == "update-2"
async def test_handles_update_exception_gracefully(self, mock_project_client: AsyncMock) -> None:
"""Update exception is logged but doesn't fail the operation."""
mock_project_client.beta.memory_stores.begin_update_memories.side_effect = Exception("API error")
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
session = AgentSession(session_id="test-session")
ctx = SessionContext(input_messages=[Message(role="user", text="hi")], session_id="s1")
ctx._response = AgentResponse(messages=[Message(role="assistant", text="hey")])
# Should not raise exception
await provider.after_run( # type: ignore[arg-type]
agent=None, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
)
# -- Context manager tests -----------------------------------------------------
class TestContextManager:
"""Test __aenter__/__aexit__ delegation."""
async def test_aenter_delegates_to_client(self, mock_project_client: AsyncMock) -> None:
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
result = await provider.__aenter__()
assert result is provider
mock_project_client.__aenter__.assert_awaited_once()
async def test_aexit_delegates_to_client(self, mock_project_client: AsyncMock) -> None:
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
await provider.__aexit__(None, None, None)
mock_project_client.__aexit__.assert_awaited_once()
async def test_async_with_syntax(self, mock_project_client: AsyncMock) -> None:
provider = FoundryMemoryProvider(
project_client=mock_project_client,
memory_store_name="test_store",
scope="user_123",
)
async with provider as p:
assert p is provider