mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
[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:
committed by
GitHub
Unverified
parent
3611be82cf
commit
cc0cfaaac8
@@ -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__)
|
||||
|
||||
+266
-27
@@ -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,
|
||||
)
|
||||
+43
-18
@@ -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,
|
||||
)
|
||||
+22
-9
@@ -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
|
||||
Reference in New Issue
Block a user