mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
Python: [BREAKING] Python: Provider-leading client design & OpenAI package extraction (#4818)
* Python: Provider-leading client design & OpenAI package extraction Major refactoring of the Python Agent Framework client architecture: - Extract OpenAI clients into new `agent-framework-openai` package - Core package no longer depends on openai, azure-identity, azure-ai-projects - Rename clients for discoverability: OpenAIResponsesClient → OpenAIChatClient, OpenAIChatClient → OpenAIChatCompletionClient - Unify `model_id`/`deployment_name`/`model_deployment_name` → `model` param - New FoundryChatClient for Azure AI Foundry Responses API - New FoundryAgent/FoundryAgentClient for connecting to pre-configured Foundry agents - Remove OpenAIBase/OpenAIConfigMixin from non-deprecated client MRO - Deprecate AzureOpenAI* clients, AzureAIClient, OpenAIAssistantsClient - Reorganize samples: azure_openai+azure_ai+azure_ai_agent → azure/ - ADR-0020: Provider-Leading Client Design Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * fix: missing Agent imports in samples, .model_id → .model in foundry_local sample Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * fix: CI failures — mypy errors, coverage targets, sample imports - azure-ai mypy: add type ignores for TypedDict total=, model arg, forward ref - Coverage: replace core.azure/openai targets with openai package target - project_provider: add type annotation for opts dict Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * fix: populate openai .pyi stub, fix broken README links, coverage targets Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * fixes * updated observabilitty * reset azure init.pyi * fix errors * updated adr number * fix foundry local * fixed not renamed docstrings and comments, and added deprecated markers to old classes * fix tests and pyprojects * fix test vars * updated function tests * update durable * updated test setup for functions * Fix Foundry auth in workflow samples Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Stabilize Python integration workflows Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Update hosting samples for Foundry Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Trigger full CI rerun Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Trigger CI rerun again Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * trigger rerun * trigger rerun * fix for litellm * undo durabletask changes * Move Foundry APIs into foundry namespace Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Fix Foundry pyproject formatting Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Split provider samples by Foundry surface Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Restore hosting sample requirements Also fix the Foundry Local sample link after the provider sample move. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * updated tests * udpated foundry integration tests * removed dist from azurefunctions tests * Use separate Foundry clients for concurrent agents Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * fix client setup in azfunc and durable * disabled two tests * updated setup for some function and durable tests * improved azure openai setup with new clients * ignore deprecated * fixes * skip 11 * remove openai assistants int tests --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This commit is contained in:
Eduard van Valkenburg
committed by
GitHub
Unverified
GitHub
parent
4b533608b6
commit
5e056b672e
@@ -0,0 +1,21 @@
|
||||
MIT License
|
||||
|
||||
Copyright (c) Microsoft Corporation.
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the "Software"), to deal
|
||||
in the Software without restriction, including without limitation the rights
|
||||
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
copies of the Software, and to permit persons to whom the Software is
|
||||
furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in all
|
||||
copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
SOFTWARE
|
||||
@@ -0,0 +1,3 @@
|
||||
# Agent Framework Foundry
|
||||
|
||||
This package contains the cloud Azure AI Foundry integrations for Microsoft Agent Framework, including Foundry chat clients, preconfigured Foundry agents, and Foundry memory providers.
|
||||
@@ -0,0 +1,24 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
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
|
||||
|
||||
try:
|
||||
__version__ = importlib.metadata.version(__name__)
|
||||
except importlib.metadata.PackageNotFoundError:
|
||||
__version__ = "0.0.0"
|
||||
|
||||
__all__ = [
|
||||
"FoundryAgent",
|
||||
"FoundryChatClient",
|
||||
"FoundryChatOptions",
|
||||
"FoundryMemoryProvider",
|
||||
"RawFoundryAgent",
|
||||
"RawFoundryAgentChatClient",
|
||||
"RawFoundryChatClient",
|
||||
"__version__",
|
||||
]
|
||||
@@ -0,0 +1,67 @@
|
||||
# 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]
|
||||
@@ -0,0 +1,287 @@
|
||||
# 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]
|
||||
AgentTelemetryLayer,
|
||||
AgentMiddlewareLayer,
|
||||
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,
|
||||
)
|
||||
@@ -0,0 +1,395 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Microsoft Foundry Agent client for connecting to pre-configured agents in Foundry.
|
||||
|
||||
This module provides ``RawFoundryAgentClient`` and ``FoundryAgentClient`` for
|
||||
communicating with PromptAgents and HostedAgents via the Responses API.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import sys
|
||||
from collections.abc import 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_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__)
|
||||
|
||||
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, 12):
|
||||
from typing import override # type: ignore # pragma: no cover
|
||||
else:
|
||||
from typing_extensions import override # 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 import Agent, BaseContextProvider
|
||||
from agent_framework._middleware import (
|
||||
ChatMiddleware,
|
||||
ChatMiddlewareCallable,
|
||||
FunctionMiddleware,
|
||||
FunctionMiddlewareCallable,
|
||||
MiddlewareTypes,
|
||||
)
|
||||
from agent_framework._tools import ToolTypes
|
||||
|
||||
|
||||
class FoundryAgentSettings(TypedDict, total=False):
|
||||
"""Settings for Microsoft FoundryAgentClient resolved from args and environment.
|
||||
|
||||
Keyword Args:
|
||||
project_endpoint: The Foundry project endpoint URL.
|
||||
Can be set via environment variable FOUNDRY_PROJECT_ENDPOINT.
|
||||
agent_name: The name of the Foundry agent to connect to.
|
||||
Can be set via environment variable FOUNDRY_AGENT_NAME.
|
||||
agent_version: The version of the Foundry agent (for PromptAgents).
|
||||
Can be set via environment variable FOUNDRY_AGENT_VERSION.
|
||||
"""
|
||||
|
||||
project_endpoint: str | None
|
||||
agent_name: str | None
|
||||
agent_version: str | None
|
||||
|
||||
|
||||
FoundryAgentOptionsT = TypeVar(
|
||||
"FoundryAgentOptionsT",
|
||||
bound=TypedDict, # type: ignore[valid-type]
|
||||
default="OpenAIChatOptions",
|
||||
covariant=True,
|
||||
)
|
||||
|
||||
|
||||
class RawFoundryAgentChatClient( # type: ignore[misc]
|
||||
RawOpenAIChatClient[FoundryAgentOptionsT],
|
||||
Generic[FoundryAgentOptionsT],
|
||||
):
|
||||
"""Raw Microsoft Foundry Agent chat client for connecting to pre-configured agents in Foundry.
|
||||
|
||||
Connects to existing PromptAgents or HostedAgents via the Responses API.
|
||||
Does not create or delete agents — the agent must already exist in Foundry.
|
||||
|
||||
This is a raw client without function invocation, chat middleware, or telemetry layers.
|
||||
Tools passed in options are validated (only ``FunctionTool`` allowed) but **not invoked** —
|
||||
the function invocation loop is handled by ``_FoundryAgentChatClient`` or a custom subclass
|
||||
that includes ``FunctionInvocationLayer``.
|
||||
|
||||
Use this class as an extension point when building a custom client with specific middleware
|
||||
layers via subclassing::
|
||||
|
||||
from agent_framework._tools import FunctionInvocationLayer
|
||||
from agent_framework.foundry import RawFoundryAgentChatClient
|
||||
|
||||
|
||||
class MyClient(FunctionInvocationLayer, RawFoundryAgentChatClient):
|
||||
pass
|
||||
|
||||
|
||||
agent = FoundryAgent(..., client_type=MyClient)
|
||||
"""
|
||||
|
||||
OTEL_PROVIDER_NAME: ClassVar[str] = "azure.ai.foundry"
|
||||
|
||||
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,
|
||||
env_file_path: str | None = None,
|
||||
env_file_encoding: str | None = None,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
"""Initialize a raw Foundry Agent client.
|
||||
|
||||
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.
|
||||
env_file_path: Path to .env file for settings.
|
||||
env_file_encoding: Encoding for .env file.
|
||||
kwargs: Additional keyword arguments.
|
||||
"""
|
||||
settings = load_settings(
|
||||
FoundryAgentSettings,
|
||||
env_prefix="FOUNDRY_",
|
||||
project_endpoint=project_endpoint,
|
||||
agent_name=agent_name,
|
||||
agent_version=agent_version,
|
||||
env_file_path=env_file_path,
|
||||
env_file_encoding=env_file_encoding,
|
||||
)
|
||||
|
||||
resolved_endpoint = settings.get("project_endpoint")
|
||||
self.agent_name = settings.get("agent_name")
|
||||
self.agent_version = settings.get("agent_version")
|
||||
|
||||
if not self.agent_name:
|
||||
raise ValueError(
|
||||
"Agent name is required. Set via 'agent_name' parameter or 'FOUNDRY_AGENT_NAME' environment variable."
|
||||
)
|
||||
|
||||
# Create or use provided project client
|
||||
self._should_close_client = False
|
||||
if project_client is not None:
|
||||
self.project_client = project_client
|
||||
else:
|
||||
if not resolved_endpoint:
|
||||
raise ValueError(
|
||||
"Either 'project_endpoint' or 'project_client' is required. "
|
||||
"Set project_endpoint via parameter or 'FOUNDRY_PROJECT_ENDPOINT' environment variable."
|
||||
)
|
||||
if not credential:
|
||||
raise ValueError("Azure credential is required when using project_endpoint without a project_client.")
|
||||
project_client_kwargs: dict[str, Any] = {
|
||||
"endpoint": resolved_endpoint,
|
||||
"credential": credential,
|
||||
"user_agent": AGENT_FRAMEWORK_USER_AGENT,
|
||||
}
|
||||
if allow_preview is not None:
|
||||
project_client_kwargs["allow_preview"] = allow_preview
|
||||
self.project_client = AIProjectClient(**project_client_kwargs)
|
||||
self._should_close_client = True
|
||||
|
||||
# Get OpenAI client from project
|
||||
async_client = self.project_client.get_openai_client()
|
||||
|
||||
super().__init__(async_client=async_client, **kwargs)
|
||||
|
||||
def _get_agent_reference(self) -> dict[str, str]:
|
||||
"""Build the agent reference dict for the Responses API."""
|
||||
ref: dict[str, str] = {"name": self.agent_name, "type": "agent_reference"} # type: ignore[dict-item]
|
||||
if self.agent_version:
|
||||
ref["version"] = self.agent_version
|
||||
return ref
|
||||
|
||||
@override
|
||||
def as_agent(
|
||||
self,
|
||||
*,
|
||||
id: str | None = None,
|
||||
name: str | None = None,
|
||||
description: str | None = None,
|
||||
instructions: str | None = None,
|
||||
tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None,
|
||||
default_options: FoundryAgentOptionsT | Mapping[str, Any] | None = None,
|
||||
context_providers: Sequence[BaseContextProvider] | None = None,
|
||||
middleware: Sequence[MiddlewareTypes] | None = None,
|
||||
**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,
|
||||
)
|
||||
|
||||
return cast(
|
||||
"Agent[FoundryAgentOptionsT]",
|
||||
FoundryAgent(
|
||||
project_client=self.project_client,
|
||||
agent_name=self.agent_name,
|
||||
agent_version=self.agent_version,
|
||||
tools=function_tools,
|
||||
context_providers=context_providers,
|
||||
middleware=middleware,
|
||||
client_type=cast(type[RawFoundryAgentChatClient], self.__class__),
|
||||
id=id,
|
||||
name=self.agent_name if name is None else name,
|
||||
description=description,
|
||||
instructions=instructions,
|
||||
default_options=default_options,
|
||||
**kwargs,
|
||||
),
|
||||
)
|
||||
|
||||
@override
|
||||
async def _prepare_options(
|
||||
self,
|
||||
messages: Sequence[Message],
|
||||
options: Mapping[str, Any],
|
||||
**kwargs: Any,
|
||||
) -> dict[str, Any]:
|
||||
"""Prepare options for the Responses API, injecting agent reference and validating tools."""
|
||||
# Validate tools — only FunctionTool allowed
|
||||
tools = options.get("tools", [])
|
||||
if tools:
|
||||
for tool_item in tools:
|
||||
if not isinstance(tool_item, FunctionTool):
|
||||
raise TypeError(
|
||||
f"Only FunctionTool objects are accepted for Foundry agents, "
|
||||
f"got {type(tool_item).__name__}. Other tool types (MCPTool, dict schemas, "
|
||||
f"hosted tools) must be defined on the Foundry agent definition in the service."
|
||||
)
|
||||
|
||||
# Prepare messages: extract system/developer messages as instructions
|
||||
prepared_messages, _instructions = self._prepare_messages_for_azure_ai(messages)
|
||||
|
||||
# Call parent prepare_options (OpenAI Responses API format)
|
||||
run_options = await super()._prepare_options(prepared_messages, options, **kwargs)
|
||||
|
||||
# Apply Azure AI schema transforms
|
||||
if "input" in run_options and isinstance(run_options["input"], list):
|
||||
run_options["input"] = self._transform_input_for_azure_ai(cast(list[dict[str, Any]], run_options["input"]))
|
||||
|
||||
# Inject agent reference
|
||||
run_options["extra_body"] = {"agent_reference": self._get_agent_reference()}
|
||||
|
||||
return run_options
|
||||
|
||||
@override
|
||||
def _check_model_presence(self, options: dict[str, Any]) -> None:
|
||||
"""Skip model check — model is configured on the Foundry agent."""
|
||||
pass
|
||||
|
||||
def _prepare_messages_for_azure_ai(self, messages: Sequence[Message]) -> tuple[list[Message], str | None]:
|
||||
"""Extract system/developer messages as instructions for Azure AI.
|
||||
|
||||
Foundry agents may not support system/developer messages directly.
|
||||
Instead, extract them as instructions to prepend.
|
||||
"""
|
||||
prepared: list[Message] = []
|
||||
instructions_parts: list[str] = []
|
||||
for msg in messages:
|
||||
if msg.role in ("system", "developer"):
|
||||
if msg.text:
|
||||
instructions_parts.append(msg.text)
|
||||
else:
|
||||
prepared.append(msg)
|
||||
instructions = "\n".join(instructions_parts) if instructions_parts else None
|
||||
return prepared, instructions
|
||||
|
||||
def _transform_input_for_azure_ai(self, input_items: list[dict[str, Any]]) -> list[dict[str, Any]]:
|
||||
"""Transform input items to match Azure AI Projects expected schema.
|
||||
|
||||
Azure AI Projects 'create responses' API expects 'type' at item level
|
||||
and 'annotations' for output_text content items.
|
||||
"""
|
||||
transformed: list[dict[str, Any]] = []
|
||||
for item in input_items:
|
||||
new_item: dict[str, Any] = dict(item)
|
||||
|
||||
if "role" in new_item and "type" not in new_item:
|
||||
new_item["type"] = "message"
|
||||
|
||||
if (content := new_item.get("content")) and isinstance(content, list):
|
||||
new_content: list[Any] = []
|
||||
for content_item in content: # type: ignore[union-attr]
|
||||
if isinstance(content_item, MutableMapping):
|
||||
if content_item.get("type") == "output_text" and "annotations" not in content_item: # type: ignore[operator]
|
||||
content_item["annotations"] = []
|
||||
new_content.append(content_item)
|
||||
else:
|
||||
new_content.append(content_item)
|
||||
new_item["content"] = new_content
|
||||
|
||||
transformed.append(new_item)
|
||||
|
||||
return transformed
|
||||
|
||||
async def close(self) -> None:
|
||||
"""Close the project client if we created it."""
|
||||
if self._should_close_client:
|
||||
await self.project_client.close()
|
||||
|
||||
|
||||
class _FoundryAgentChatClient( # type: ignore[misc]
|
||||
FunctionInvocationLayer[FoundryAgentOptionsT],
|
||||
ChatMiddlewareLayer[FoundryAgentOptionsT],
|
||||
ChatTelemetryLayer[FoundryAgentOptionsT],
|
||||
RawFoundryAgentChatClient[FoundryAgentOptionsT],
|
||||
Generic[FoundryAgentOptionsT],
|
||||
):
|
||||
"""Microsoft Foundry Agent client with middleware, telemetry, and function invocation support.
|
||||
|
||||
Connects to existing PromptAgents or HostedAgents in Foundry.
|
||||
|
||||
Examples:
|
||||
.. code-block:: python
|
||||
|
||||
from agent_framework import Agent
|
||||
from agent_framework.foundry import FoundryAgentClient
|
||||
from azure.identity import AzureCliCredential
|
||||
|
||||
client = FoundryAgentClient(
|
||||
project_endpoint="https://your-project.services.ai.azure.com",
|
||||
agent_name="my-prompt-agent",
|
||||
agent_version="1.0",
|
||||
credential=AzureCliCredential(),
|
||||
)
|
||||
|
||||
agent = Agent(client=client, tools=[my_function_tool])
|
||||
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,
|
||||
env_file_path: str | None = None,
|
||||
env_file_encoding: str | None = None,
|
||||
middleware: (
|
||||
Sequence[ChatMiddleware | ChatMiddlewareCallable | FunctionMiddleware | FunctionMiddlewareCallable] | None
|
||||
) = None,
|
||||
function_invocation_configuration: FunctionInvocationConfiguration | None = None,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
"""Initialize a Foundry Agent client with full middleware support.
|
||||
|
||||
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.
|
||||
env_file_path: Path to .env file for settings.
|
||||
env_file_encoding: Encoding for .env file.
|
||||
middleware: Optional sequence of middleware.
|
||||
function_invocation_configuration: Optional function invocation configuration.
|
||||
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,
|
||||
env_file_path=env_file_path,
|
||||
env_file_encoding=env_file_encoding,
|
||||
middleware=middleware,
|
||||
function_invocation_configuration=function_invocation_configuration,
|
||||
**kwargs,
|
||||
)
|
||||
@@ -0,0 +1,530 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import sys
|
||||
from collections.abc import 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.observability import ChatTelemetryLayer
|
||||
from agent_framework_openai._chat_client import OpenAIChatOptions, RawOpenAIChatClient
|
||||
from azure.ai.projects.aio import AIProjectClient
|
||||
from azure.ai.projects.models import (
|
||||
AutoCodeInterpreterToolParam,
|
||||
CodeInterpreterTool,
|
||||
ImageGenTool,
|
||||
WebSearchApproximateLocation,
|
||||
WebSearchTool,
|
||||
WebSearchToolFilters,
|
||||
)
|
||||
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
|
||||
|
||||
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, 12):
|
||||
from typing import override # type: ignore # pragma: no cover
|
||||
else:
|
||||
from typing_extensions import override # 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 (
|
||||
ChatMiddleware,
|
||||
ChatMiddlewareCallable,
|
||||
FunctionMiddleware,
|
||||
FunctionMiddlewareCallable,
|
||||
)
|
||||
|
||||
logger: logging.Logger = logging.getLogger("agent_framework.foundry")
|
||||
|
||||
|
||||
class FoundrySettings(TypedDict, total=False):
|
||||
"""Settings for Microsoft FoundryChatClient resolved from args and environment.
|
||||
|
||||
Keyword Args:
|
||||
model: The model deployment name.
|
||||
Can be set via environment variable FOUNDRY_MODEL.
|
||||
project_endpoint: The Microsoft Foundry project endpoint URL.
|
||||
Can be set via environment variable FOUNDRY_PROJECT_ENDPOINT.
|
||||
"""
|
||||
|
||||
model: str | None
|
||||
project_endpoint: str | None
|
||||
|
||||
|
||||
FoundryChatOptionsT = TypeVar(
|
||||
"FoundryChatOptionsT",
|
||||
bound=TypedDict, # type: ignore[valid-type]
|
||||
default="OpenAIChatOptions",
|
||||
covariant=True,
|
||||
)
|
||||
|
||||
FoundryChatOptions = OpenAIChatOptions
|
||||
|
||||
|
||||
class RawFoundryChatClient( # type: ignore[misc]
|
||||
RawOpenAIChatClient[FoundryChatOptionsT],
|
||||
Generic[FoundryChatOptionsT],
|
||||
):
|
||||
"""Raw Microsoft Foundry chat client using the OpenAI Responses API via a Foundry project.
|
||||
|
||||
This client creates an OpenAI-compatible client from a Foundry project
|
||||
and delegates to ``RawOpenAIChatClient`` for request handling.
|
||||
|
||||
Environment variables:
|
||||
- ``FOUNDRY_PROJECT_ENDPOINT`` to provide the Foundry project endpoint.
|
||||
- ``FOUNDRY_MODEL`` to provide the Foundry model deployment name.
|
||||
|
||||
Warning:
|
||||
**This class should not normally be used directly.** Use ``FoundryChatClient``
|
||||
for a fully-featured client with middleware, telemetry, and function invocation.
|
||||
"""
|
||||
|
||||
OTEL_PROVIDER_NAME: ClassVar[str] = "azure.ai.foundry" # type: ignore[reportIncompatibleVariableOverride, misc]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
project_endpoint: str | None = None,
|
||||
project_client: AIProjectClient | None = None,
|
||||
model: str | None = None,
|
||||
credential: AzureCredentialTypes | AzureTokenProvider | None = None,
|
||||
allow_preview: bool | None = None,
|
||||
env_file_path: str | None = None,
|
||||
env_file_encoding: str | None = None,
|
||||
instruction_role: str | None = None,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
"""Initialize a raw Microsoft Foundry chat client.
|
||||
|
||||
Keyword Args:
|
||||
project_endpoint: The Foundry project endpoint URL.
|
||||
Can also be set via environment variable FOUNDRY_PROJECT_ENDPOINT.
|
||||
project_client: An existing AIProjectClient to use. If provided,
|
||||
the OpenAI client will be obtained via ``project_client.get_openai_client()``.
|
||||
model: The model deployment name.
|
||||
Can also be set via environment variable FOUNDRY_MODEL.
|
||||
credential: Azure credential or token provider for authentication.
|
||||
Required when using ``project_endpoint`` without a ``project_client``.
|
||||
allow_preview: Enables preview opt-in on internally-created AIProjectClient.
|
||||
env_file_path: Path to .env file for settings.
|
||||
env_file_encoding: Encoding for .env file.
|
||||
instruction_role: The role to use for 'instruction' messages.
|
||||
kwargs: Additional keyword arguments.
|
||||
"""
|
||||
foundry_settings = load_settings(
|
||||
FoundrySettings,
|
||||
env_prefix="FOUNDRY_",
|
||||
model=model,
|
||||
project_endpoint=project_endpoint,
|
||||
env_file_path=env_file_path,
|
||||
env_file_encoding=env_file_encoding,
|
||||
)
|
||||
|
||||
resolved_model = foundry_settings.get("model")
|
||||
if not resolved_model:
|
||||
raise ValueError("Model is required. Set via 'model' parameter or 'FOUNDRY_MODEL' environment variable.")
|
||||
|
||||
project_endpoint = foundry_settings.get("project_endpoint")
|
||||
|
||||
if project_endpoint is None and project_client is None:
|
||||
raise ValueError(
|
||||
"Either 'project_endpoint' or 'project_client' is required. "
|
||||
"Set project_endpoint via parameter or 'FOUNDRY_PROJECT_ENDPOINT' environment variable."
|
||||
)
|
||||
if not project_client:
|
||||
if not project_endpoint:
|
||||
raise ValueError(
|
||||
"Azure AI project endpoint is required. Set via 'project_endpoint' parameter "
|
||||
"or 'FOUNDRY_PROJECT_ENDPOINT' environment variable,"
|
||||
"or pass in a AIProjectClient."
|
||||
)
|
||||
if not credential:
|
||||
raise ValueError("Azure credential is required when using project_endpoint without a project_client.")
|
||||
project_client_kwargs: dict[str, Any] = {
|
||||
"endpoint": project_endpoint,
|
||||
"credential": credential, # type: ignore[arg-type]
|
||||
"user_agent": AGENT_FRAMEWORK_USER_AGENT,
|
||||
}
|
||||
if allow_preview is not None:
|
||||
project_client_kwargs["allow_preview"] = allow_preview
|
||||
project_client = AIProjectClient(**project_client_kwargs)
|
||||
|
||||
super().__init__(
|
||||
model=resolved_model,
|
||||
async_client=project_client.get_openai_client(),
|
||||
instruction_role=instruction_role,
|
||||
**kwargs,
|
||||
)
|
||||
self.project_client = project_client
|
||||
|
||||
@override
|
||||
def _check_model_presence(self, options: dict[str, Any]) -> None:
|
||||
if not options.get("model"):
|
||||
if not self.model:
|
||||
raise ValueError("model must be a non-empty string")
|
||||
options["model"] = self.model
|
||||
|
||||
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.
|
||||
|
||||
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().
|
||||
Common options include:
|
||||
- enable_live_metrics (bool): Enable Azure Monitor Live Metrics
|
||||
- credential (TokenCredential): Azure credential for Entra ID auth
|
||||
- resource (Resource): Custom OpenTelemetry resource
|
||||
|
||||
Raises:
|
||||
ImportError: If azure-monitor-opentelemetry-exporter is not installed.
|
||||
"""
|
||||
from azure.core.exceptions import ResourceNotFoundError
|
||||
|
||||
try:
|
||||
conn_string = await self.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)
|
||||
|
||||
# region Tool factory methods (override OpenAI defaults with Foundry versions)
|
||||
|
||||
@staticmethod
|
||||
def get_code_interpreter_tool( # type: ignore[override]
|
||||
*,
|
||||
file_ids: list[str | Content] | None = None,
|
||||
container: Literal["auto"] | dict[str, Any] = "auto",
|
||||
**kwargs: Any,
|
||||
) -> CodeInterpreterTool:
|
||||
"""Create a code interpreter tool configuration for Foundry.
|
||||
|
||||
Keyword Args:
|
||||
file_ids: Optional list of file IDs or Content objects to make available.
|
||||
container: Container configuration. Use "auto" for automatic management.
|
||||
**kwargs: Additional arguments passed to the SDK CodeInterpreterTool constructor.
|
||||
|
||||
Returns:
|
||||
A CodeInterpreterTool ready to pass to an Agent.
|
||||
"""
|
||||
if file_ids is None and isinstance(container, dict):
|
||||
file_ids = container.get("file_ids")
|
||||
resolved = resolve_file_ids(file_ids)
|
||||
tool_container = AutoCodeInterpreterToolParam(file_ids=resolved)
|
||||
return CodeInterpreterTool(container=tool_container, **kwargs)
|
||||
|
||||
@staticmethod
|
||||
def get_file_search_tool(
|
||||
*,
|
||||
vector_store_ids: list[str],
|
||||
max_num_results: int | None = None,
|
||||
ranking_options: dict[str, Any] | None = None,
|
||||
filters: dict[str, Any] | None = None,
|
||||
**kwargs: Any,
|
||||
) -> ProjectsFileSearchTool:
|
||||
"""Create a file search tool configuration for Foundry.
|
||||
|
||||
Keyword Args:
|
||||
vector_store_ids: List of vector store IDs to search.
|
||||
max_num_results: Maximum number of results to return (1-50).
|
||||
ranking_options: Ranking options for search results.
|
||||
filters: A filter to apply (ComparisonFilter or CompoundFilter).
|
||||
**kwargs: Additional arguments passed to the SDK FileSearchTool constructor.
|
||||
|
||||
Returns:
|
||||
A FileSearchTool ready to pass to an Agent.
|
||||
"""
|
||||
if not vector_store_ids:
|
||||
raise ValueError("File search tool requires 'vector_store_ids' to be specified.")
|
||||
return ProjectsFileSearchTool(
|
||||
vector_store_ids=vector_store_ids,
|
||||
max_num_results=max_num_results,
|
||||
ranking_options=ranking_options, # type: ignore[arg-type]
|
||||
filters=filters, # type: ignore[arg-type]
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def get_web_search_tool( # type: ignore[override]
|
||||
*,
|
||||
user_location: dict[str, str] | None = None,
|
||||
search_context_size: Literal["low", "medium", "high"] | None = None,
|
||||
allowed_domains: list[str] | None = None,
|
||||
custom_search_configuration: dict[str, Any] | None = None,
|
||||
**kwargs: Any,
|
||||
) -> WebSearchTool:
|
||||
"""Create a web search tool configuration for Microsoft Foundry.
|
||||
|
||||
Keyword Args:
|
||||
user_location: Location context with keys like "city", "country", "region", "timezone".
|
||||
search_context_size: Amount of context from search results ("low", "medium", "high").
|
||||
allowed_domains: List of domains to restrict search results to.
|
||||
custom_search_configuration: Custom Bing search configuration.
|
||||
**kwargs: Additional arguments passed to the SDK WebSearchTool constructor.
|
||||
|
||||
Returns:
|
||||
A WebSearchTool ready to pass to an Agent.
|
||||
"""
|
||||
ws_kwargs: dict[str, Any] = {**kwargs}
|
||||
if search_context_size:
|
||||
ws_kwargs["search_context_size"] = search_context_size
|
||||
if allowed_domains:
|
||||
ws_kwargs["filters"] = WebSearchToolFilters(allowed_domains=allowed_domains)
|
||||
if custom_search_configuration:
|
||||
ws_kwargs["custom_search_configuration"] = custom_search_configuration
|
||||
ws_tool = WebSearchTool(**ws_kwargs)
|
||||
if user_location:
|
||||
ws_tool.user_location = WebSearchApproximateLocation(
|
||||
city=user_location.get("city"),
|
||||
country=user_location.get("country"),
|
||||
region=user_location.get("region"),
|
||||
timezone=user_location.get("timezone"),
|
||||
)
|
||||
return ws_tool
|
||||
|
||||
@staticmethod
|
||||
def get_image_generation_tool( # type: ignore[override]
|
||||
*,
|
||||
model: Literal["gpt-image-1"] | str | None = None,
|
||||
size: Literal["1024x1024", "1024x1536", "1536x1024", "auto"] | None = None,
|
||||
output_format: Literal["png", "webp", "jpeg"] | None = None,
|
||||
quality: Literal["low", "medium", "high", "auto"] | None = None,
|
||||
background: Literal["transparent", "opaque", "auto"] | None = None,
|
||||
partial_images: int | None = None,
|
||||
moderation: Literal["auto", "low"] | None = None,
|
||||
output_compression: int | None = None,
|
||||
**kwargs: Any,
|
||||
) -> ImageGenTool:
|
||||
"""Create an image generation tool configuration for Foundry.
|
||||
|
||||
Keyword Args:
|
||||
model: The model to use for image generation.
|
||||
size: Output image size.
|
||||
output_format: Output image format.
|
||||
quality: Output image quality.
|
||||
background: Background transparency setting.
|
||||
partial_images: Number of partial images to return during generation.
|
||||
moderation: Moderation level.
|
||||
output_compression: Compression level.
|
||||
**kwargs: Additional arguments passed to the SDK ImageGenTool constructor.
|
||||
|
||||
Returns:
|
||||
An ImageGenTool ready to pass to an Agent.
|
||||
"""
|
||||
return ImageGenTool( # type: ignore[misc]
|
||||
model=model, # type: ignore[arg-type]
|
||||
size=size,
|
||||
output_format=output_format,
|
||||
quality=quality,
|
||||
background=background,
|
||||
partial_images=partial_images,
|
||||
moderation=moderation,
|
||||
output_compression=output_compression,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def get_mcp_tool(
|
||||
*,
|
||||
name: str,
|
||||
url: str | None = None,
|
||||
description: str | None = None,
|
||||
approval_mode: Literal["always_require", "never_require"] | dict[str, list[str]] | None = None,
|
||||
allowed_tools: list[str] | None = None,
|
||||
headers: dict[str, str] | None = None,
|
||||
project_connection_id: str | None = None,
|
||||
**kwargs: Any,
|
||||
) -> FoundryMCPTool:
|
||||
"""Create a hosted MCP tool configuration for Foundry.
|
||||
|
||||
This configures an MCP server that runs remotely on Azure AI, not locally.
|
||||
|
||||
Keyword Args:
|
||||
name: A label/name for the MCP server.
|
||||
url: The URL of the MCP server. Required if project_connection_id is not provided.
|
||||
description: A description of what the MCP server provides.
|
||||
approval_mode: Tool approval mode ("always_require", "never_require", or dict).
|
||||
allowed_tools: List of allowed tool names from this MCP server.
|
||||
headers: HTTP headers to include in requests to the MCP server.
|
||||
project_connection_id: Foundry connection ID for managed MCP connections.
|
||||
**kwargs: Additional arguments passed to the SDK MCPTool constructor.
|
||||
|
||||
Returns:
|
||||
An MCPTool configuration ready to pass to an Agent.
|
||||
"""
|
||||
mcp = FoundryMCPTool(server_label=name.replace(" ", "_"), server_url=url or "", **kwargs)
|
||||
|
||||
if description:
|
||||
mcp["server_description"] = description
|
||||
if project_connection_id:
|
||||
mcp["project_connection_id"] = project_connection_id
|
||||
elif headers:
|
||||
mcp["headers"] = headers
|
||||
if allowed_tools:
|
||||
mcp["allowed_tools"] = allowed_tools
|
||||
if approval_mode:
|
||||
if isinstance(approval_mode, str):
|
||||
mcp["require_approval"] = "always" if approval_mode == "always_require" else "never"
|
||||
else:
|
||||
if always_require := approval_mode.get("always_require_approval"):
|
||||
mcp["require_approval"] = {"always": {"tool_names": always_require}}
|
||||
if never_require := approval_mode.get("never_require_approval"):
|
||||
mcp["require_approval"] = {"never": {"tool_names": never_require}}
|
||||
|
||||
return mcp
|
||||
|
||||
# endregion
|
||||
|
||||
|
||||
class FoundryChatClient( # type: ignore[misc]
|
||||
FunctionInvocationLayer[FoundryChatOptionsT],
|
||||
ChatMiddlewareLayer[FoundryChatOptionsT],
|
||||
ChatTelemetryLayer[FoundryChatOptionsT],
|
||||
RawFoundryChatClient[FoundryChatOptionsT],
|
||||
Generic[FoundryChatOptionsT],
|
||||
):
|
||||
"""Microsoft Foundry chat client using the OpenAI Responses API.
|
||||
|
||||
Creates an OpenAI-compatible client from a Foundry project
|
||||
with middleware, telemetry, and function invocation support.
|
||||
|
||||
Environment variables:
|
||||
- ``FOUNDRY_PROJECT_ENDPOINT`` to provide the Foundry project endpoint.
|
||||
- ``FOUNDRY_MODEL`` to provide the Foundry model deployment name.
|
||||
|
||||
Keyword Args:
|
||||
project_endpoint: The Foundry project endpoint URL.
|
||||
Can also be set via environment variable ``FOUNDRY_PROJECT_ENDPOINT``.
|
||||
project_client: An existing AIProjectClient to use.
|
||||
model: The model deployment name.
|
||||
Can also be set via environment variable ``FOUNDRY_MODEL``.
|
||||
model_id: Deprecated alias for ``model``.
|
||||
credential: Azure credential or token provider for authentication.
|
||||
allow_preview: Enables preview opt-in on internally-created AIProjectClient.
|
||||
env_file_path: Path to .env file for settings.
|
||||
env_file_encoding: Encoding for .env file.
|
||||
instruction_role: The role to use for 'instruction' messages.
|
||||
middleware: Optional sequence of middleware.
|
||||
function_invocation_configuration: Optional function invocation configuration.
|
||||
|
||||
Examples:
|
||||
.. code-block:: python
|
||||
|
||||
from azure.identity import AzureCliCredential
|
||||
from agent_framework_foundry import FoundryChatClient
|
||||
|
||||
client = FoundryChatClient(
|
||||
project_endpoint="https://your-project.services.ai.azure.com",
|
||||
model="gpt-4o",
|
||||
credential=AzureCliCredential(),
|
||||
)
|
||||
|
||||
# Or using an existing AIProjectClient
|
||||
from azure.ai.projects.aio import AIProjectClient
|
||||
|
||||
project_client = AIProjectClient(
|
||||
endpoint="https://your-project.services.ai.azure.com",
|
||||
credential=AzureCliCredential(),
|
||||
)
|
||||
client = FoundryChatClient(
|
||||
project_client=project_client,
|
||||
model="gpt-4o",
|
||||
)
|
||||
"""
|
||||
|
||||
OTEL_PROVIDER_NAME: ClassVar[str] = "azure.ai.foundry" # type: ignore[reportIncompatibleVariableOverride, misc]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
project_endpoint: str | None = None,
|
||||
project_client: AIProjectClient | None = None,
|
||||
model: str | None = None,
|
||||
credential: AzureCredentialTypes | AzureTokenProvider | None = None,
|
||||
allow_preview: bool | None = None,
|
||||
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,
|
||||
function_invocation_configuration: FunctionInvocationConfiguration | None = None,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
"""Initialize a Foundry chat client.
|
||||
|
||||
Keyword Args:
|
||||
project_endpoint: The Foundry project endpoint URL.
|
||||
Can also be set via environment variable ``FOUNDRY_PROJECT_ENDPOINT``.
|
||||
project_client: An existing AIProjectClient to use.
|
||||
model: The model deployment name.
|
||||
Can also be set via environment variable ``FOUNDRY_MODEL``.
|
||||
credential: Azure credential or token provider for authentication.
|
||||
allow_preview: Enables preview opt-in on internally-created AIProjectClient.
|
||||
env_file_path: Path to .env file for settings.
|
||||
env_file_encoding: Encoding for .env file.
|
||||
instruction_role: The role to use for 'instruction' messages.
|
||||
middleware: Optional sequence of middleware.
|
||||
function_invocation_configuration: Optional function invocation configuration.
|
||||
kwargs: Additional keyword arguments.
|
||||
"""
|
||||
super().__init__(
|
||||
project_endpoint=project_endpoint,
|
||||
project_client=project_client,
|
||||
model=model,
|
||||
credential=credential,
|
||||
allow_preview=allow_preview,
|
||||
env_file_path=env_file_path,
|
||||
env_file_encoding=env_file_encoding,
|
||||
instruction_role=instruction_role,
|
||||
middleware=middleware,
|
||||
function_invocation_configuration=function_invocation_configuration,
|
||||
**kwargs,
|
||||
)
|
||||
@@ -0,0 +1,261 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Foundry Memory Context Provider using BaseContextProvider.
|
||||
|
||||
This module provides ``FoundryMemoryProvider``, built on
|
||||
:class:`BaseContextProvider`.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
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 azure.ai.projects.aio import AIProjectClient
|
||||
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
|
||||
else:
|
||||
from typing_extensions import Self # pragma: no cover
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from agent_framework._agents import SupportsAgentRun
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class FoundryMemoryProvider(BaseContextProvider):
|
||||
"""Foundry Memory context provider using the new BaseContextProvider hooks pattern.
|
||||
|
||||
Integrates Azure AI Foundry Memory Store for persistent semantic memory,
|
||||
searching and storing memories via the Azure AI Projects SDK.
|
||||
|
||||
Args:
|
||||
source_id: Unique identifier for this provider instance.
|
||||
project_client: Azure AI Project client for memory operations.
|
||||
memory_store_name: The name of the memory store to use.
|
||||
scope: The namespace that logically groups and isolates memories (e.g., user ID).
|
||||
context_prompt: The prompt to prepend to retrieved memories.
|
||||
update_delay: Timeout period before processing memory update in seconds.
|
||||
Defaults to 300 (5 minutes). Set to 0 to immediately trigger updates.
|
||||
"""
|
||||
|
||||
DEFAULT_SOURCE_ID: ClassVar[str] = "foundry_memory"
|
||||
DEFAULT_CONTEXT_PROMPT = "## Memories\nConsider the following memories when answering user questions:"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
source_id: str = DEFAULT_SOURCE_ID,
|
||||
*,
|
||||
project_client: AIProjectClient | None = None,
|
||||
project_endpoint: str | None = None,
|
||||
credential: AzureCredentialTypes | None = None,
|
||||
allow_preview: bool | None = None,
|
||||
memory_store_name: str,
|
||||
scope: str | None = None,
|
||||
context_prompt: str | None = None,
|
||||
update_delay: int = 300,
|
||||
env_file_path: str | None = None,
|
||||
env_file_encoding: str | None = None,
|
||||
) -> None:
|
||||
"""Initialize the Foundry Memory context provider.
|
||||
|
||||
Args:
|
||||
source_id: Unique identifier for this provider instance.
|
||||
project_client: Azure AI Project client for memory operations.
|
||||
project_endpoint: Foundry project endpoint URL. Used when project_client is not provided.
|
||||
credential: Azure credential for authentication. Accepts a TokenCredential,
|
||||
AsyncTokenCredential, or a callable token provider.
|
||||
Required when project_client is not provided.
|
||||
allow_preview: Enables preview opt-in on internally-created ``AIProjectClient``.
|
||||
memory_store_name: The name of the memory store to use.
|
||||
scope: The namespace that logically groups and isolates memories (e.g., user ID).
|
||||
If None, `session_id` will be used.
|
||||
context_prompt: The prompt to prepend to retrieved memories.
|
||||
update_delay: Timeout period before processing memory update in seconds.
|
||||
env_file_path: Path to environment file for loading settings.
|
||||
env_file_encoding: Encoding of the environment file.
|
||||
"""
|
||||
super().__init__(source_id)
|
||||
foundry_settings = load_settings(
|
||||
FoundryProjectSettings,
|
||||
env_prefix="FOUNDRY_",
|
||||
project_endpoint=project_endpoint,
|
||||
env_file_path=env_file_path,
|
||||
env_file_encoding=env_file_encoding,
|
||||
)
|
||||
|
||||
if project_client is None:
|
||||
resolved_endpoint = foundry_settings.get("project_endpoint")
|
||||
if not resolved_endpoint:
|
||||
raise ValueError(
|
||||
"Foundry project endpoint is required. Set via 'project_endpoint' parameter "
|
||||
"or 'FOUNDRY_PROJECT_ENDPOINT' environment variable."
|
||||
)
|
||||
if not credential:
|
||||
raise ValueError("Azure credential is required when project_client is not provided.")
|
||||
project_client_kwargs: dict[str, Any] = {
|
||||
"endpoint": resolved_endpoint,
|
||||
"credential": credential, # type: ignore[arg-type]
|
||||
"user_agent": AGENT_FRAMEWORK_USER_AGENT,
|
||||
}
|
||||
if allow_preview is not None:
|
||||
project_client_kwargs["allow_preview"] = allow_preview
|
||||
project_client = AIProjectClient(**project_client_kwargs)
|
||||
|
||||
if not memory_store_name:
|
||||
raise ValueError("memory_store_name is required")
|
||||
if not scope:
|
||||
raise ValueError("scope is required")
|
||||
|
||||
self.project_client = project_client
|
||||
self.memory_store_name = memory_store_name
|
||||
self.scope = scope
|
||||
self.context_prompt = context_prompt or self.DEFAULT_CONTEXT_PROMPT
|
||||
self.update_delay = update_delay
|
||||
|
||||
async def __aenter__(self) -> Self:
|
||||
"""Async context manager entry."""
|
||||
if self.project_client and isinstance(self.project_client, AbstractAsyncContextManager):
|
||||
await self.project_client.__aenter__()
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: Any) -> None:
|
||||
"""Async context manager exit."""
|
||||
if self.project_client and isinstance(self.project_client, AbstractAsyncContextManager):
|
||||
await self.project_client.__aexit__(exc_type, exc_val, exc_tb)
|
||||
|
||||
# -- Hooks pattern ---------------------------------------------------------
|
||||
|
||||
async def before_run(
|
||||
self,
|
||||
*,
|
||||
agent: SupportsAgentRun,
|
||||
session: AgentSession,
|
||||
context: SessionContext,
|
||||
state: dict[str, Any],
|
||||
) -> None:
|
||||
"""Search Foundry Memory for relevant memories and add to the session context.
|
||||
|
||||
This method:
|
||||
1. Retrieves static memories (user profile) on first call per session
|
||||
2. Searches for contextual memories based on input messages
|
||||
3. Combines and injects memories into the context
|
||||
"""
|
||||
# On first run, retrieve static memories (user profile memories)
|
||||
if not state.get("initialized"):
|
||||
try:
|
||||
static_search_result = await self.project_client.beta.memory_stores.search_memories(
|
||||
name=self.memory_store_name,
|
||||
scope=self.scope or context.session_id, # type: ignore[arg-type]
|
||||
)
|
||||
static_memories = [{"content": memory.memory_item.content} for memory in static_search_result.memories]
|
||||
state["static_memories"] = static_memories
|
||||
except Exception as e:
|
||||
# Log but don't fail - memory retrieval is non-critical
|
||||
logger.warning(f"Failed to retrieve static memories: {e}")
|
||||
state["static_memories"] = []
|
||||
finally:
|
||||
# Mark as initialized regardless of success to avoid repeated attempts
|
||||
state["initialized"] = True
|
||||
|
||||
# Search for contextual memories based on input messages
|
||||
# Check if there are any non-empty input messages
|
||||
has_input = any(msg and msg.text and msg.text.strip() for msg in context.input_messages)
|
||||
if not has_input:
|
||||
return
|
||||
|
||||
# Convert input messages to memory search item format
|
||||
items: list[ResponseInputItemParam] = [
|
||||
{"type": "message", "role": "user", "content": msg.text}
|
||||
for msg in context.input_messages
|
||||
if msg and msg.text and msg.text.strip()
|
||||
]
|
||||
|
||||
try:
|
||||
search_result = await self.project_client.beta.memory_stores.search_memories(
|
||||
name=self.memory_store_name,
|
||||
scope=self.scope or context.session_id, # type: ignore[arg-type]
|
||||
items=items,
|
||||
previous_search_id=state.get("previous_search_id"),
|
||||
)
|
||||
|
||||
# Extract search_id for next incremental search
|
||||
if search_result.memories:
|
||||
state["previous_search_id"] = search_result.search_id
|
||||
|
||||
# Combine static and contextual memories
|
||||
contextual_memories = [{"content": memory.memory_item.content} for memory in search_result.memories]
|
||||
|
||||
all_memories = state.get("static_memories", []) + contextual_memories
|
||||
|
||||
# Inject memories into context
|
||||
if all_memories:
|
||||
line_separated_memories = "\n".join(
|
||||
str(memory.get("content", "")) for memory in all_memories if memory.get("content")
|
||||
)
|
||||
if line_separated_memories:
|
||||
context.extend_messages(
|
||||
self.source_id,
|
||||
[Message(role="user", text=f"{self.context_prompt}\n{line_separated_memories}")],
|
||||
)
|
||||
except Exception as e:
|
||||
# Log but don't fail - memory retrieval is non-critical
|
||||
logger.warning(f"Failed to search contextual memories: {e}")
|
||||
|
||||
async def after_run(
|
||||
self,
|
||||
*,
|
||||
agent: SupportsAgentRun,
|
||||
session: AgentSession,
|
||||
context: SessionContext,
|
||||
state: dict[str, Any],
|
||||
) -> None:
|
||||
"""Store request/response messages to Foundry Memory for future retrieval.
|
||||
|
||||
This method updates the memory store with conversation messages.
|
||||
The update is debounced by the configured update_delay.
|
||||
"""
|
||||
messages_to_store: list[Message] = list(context.input_messages)
|
||||
if context.response and context.response.messages:
|
||||
messages_to_store.extend(context.response.messages)
|
||||
|
||||
# Filter and convert messages to memory update item format
|
||||
items: list[ResponseInputItemParam] = []
|
||||
for message in messages_to_store:
|
||||
if message.role in {"user", "assistant", "system"} and message.text and message.text.strip():
|
||||
if message.role == "user":
|
||||
items.append({"role": "user", "type": "message", "content": message.text})
|
||||
elif message.role == "assistant":
|
||||
items.append({"role": "assistant", "type": "message", "content": message.text})
|
||||
|
||||
if not items:
|
||||
return
|
||||
|
||||
try:
|
||||
# Fire and forget - don't wait for the update to complete
|
||||
update_poller = await self.project_client.beta.memory_stores.begin_update_memories(
|
||||
name=self.memory_store_name,
|
||||
scope=self.scope or context.session_id, # type: ignore[arg-type]
|
||||
items=items,
|
||||
previous_update_id=state.get("previous_update_id"),
|
||||
update_delay=self.update_delay,
|
||||
)
|
||||
# Store the update_id for next incremental update
|
||||
state["previous_update_id"] = update_poller.update_id
|
||||
|
||||
except Exception as e:
|
||||
# Log but don't fail - memory storage is non-critical
|
||||
logger.warning(f"Failed to update memories: {e}")
|
||||
|
||||
|
||||
__all__ = ["FoundryMemoryProvider"]
|
||||
@@ -0,0 +1,49 @@
|
||||
# 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
|
||||
@@ -0,0 +1,105 @@
|
||||
[project]
|
||||
name = "agent-framework-foundry"
|
||||
description = "Cloud Azure AI Foundry integration for Microsoft Agent Framework."
|
||||
authors = [{ name = "Microsoft", email = "af-support@microsoft.com"}]
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.10"
|
||||
version = "1.0.0rc5"
|
||||
license-files = ["LICENSE"]
|
||||
urls.homepage = "https://aka.ms/agent-framework"
|
||||
urls.source = "https://github.com/microsoft/agent-framework/tree/main/python"
|
||||
urls.release_notes = "https://github.com/microsoft/agent-framework/releases?q=tag%3Apython-1&expanded=true"
|
||||
urls.issues = "https://github.com/microsoft/agent-framework/issues"
|
||||
classifiers = [
|
||||
"License :: OSI Approved :: MIT License",
|
||||
"Development Status :: 4 - Beta",
|
||||
"Intended Audience :: Developers",
|
||||
"Programming Language :: Python :: 3",
|
||||
"Programming Language :: Python :: 3.10",
|
||||
"Programming Language :: Python :: 3.11",
|
||||
"Programming Language :: Python :: 3.12",
|
||||
"Programming Language :: Python :: 3.13",
|
||||
"Programming Language :: Python :: 3.14",
|
||||
"Typing :: Typed",
|
||||
]
|
||||
dependencies = [
|
||||
"agent-framework-core>=1.0.0rc5",
|
||||
"agent-framework-openai>=1.0.0rc5",
|
||||
"azure-ai-projects>=2.0.0,<3.0",
|
||||
]
|
||||
|
||||
[tool.uv]
|
||||
prerelease = "if-necessary-or-explicit"
|
||||
environments = [
|
||||
"sys_platform == 'darwin'",
|
||||
"sys_platform == 'linux'",
|
||||
"sys_platform == 'win32'"
|
||||
]
|
||||
|
||||
[tool.uv-dynamic-versioning]
|
||||
fallback-version = "0.0.0"
|
||||
|
||||
[tool.pytest.ini_options]
|
||||
testpaths = 'tests'
|
||||
addopts = "-ra -q -r fEX"
|
||||
asyncio_mode = "auto"
|
||||
asyncio_default_fixture_loop_scope = "function"
|
||||
filterwarnings = []
|
||||
timeout = 120
|
||||
markers = [
|
||||
"integration: marks tests as integration tests that require external services",
|
||||
]
|
||||
|
||||
[tool.ruff]
|
||||
extend = "../../pyproject.toml"
|
||||
|
||||
[tool.coverage.run]
|
||||
omit = [
|
||||
"**/__init__.py"
|
||||
]
|
||||
|
||||
[tool.pyright]
|
||||
extends = "../../pyproject.toml"
|
||||
include = ["agent_framework_foundry"]
|
||||
|
||||
[tool.mypy]
|
||||
plugins = ['pydantic.mypy']
|
||||
strict = true
|
||||
python_version = "3.10"
|
||||
ignore_missing_imports = true
|
||||
disallow_untyped_defs = true
|
||||
no_implicit_optional = true
|
||||
check_untyped_defs = true
|
||||
warn_return_any = true
|
||||
show_error_codes = true
|
||||
warn_unused_ignores = false
|
||||
disallow_incomplete_defs = true
|
||||
disallow_untyped_decorators = true
|
||||
|
||||
[tool.bandit]
|
||||
targets = ["agent_framework_foundry"]
|
||||
exclude_dirs = ["tests"]
|
||||
|
||||
[tool.poe]
|
||||
executor.type = "uv"
|
||||
include = "../../shared_tasks.toml"
|
||||
|
||||
[tool.poe.tasks.mypy]
|
||||
help = "Run MyPy for this package."
|
||||
cmd = "mypy --config-file $POE_ROOT/pyproject.toml agent_framework_foundry"
|
||||
|
||||
[tool.poe.tasks.test]
|
||||
help = "Run the default unit test suite for this package."
|
||||
cmd = 'pytest -m "not integration" --cov=agent_framework_foundry --cov-report=term-missing:skip-covered tests'
|
||||
|
||||
[tool.poe.tasks.integration-tests]
|
||||
help = "Run the package integration test suite."
|
||||
cmd = """
|
||||
pytest --import-mode=importlib
|
||||
-n logical --dist worksteal
|
||||
tests
|
||||
"""
|
||||
|
||||
[build-system]
|
||||
requires = ["flit-core >= 3.11,<4.0"]
|
||||
build-backend = "flit_core.buildapi"
|
||||
@@ -0,0 +1,78 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
from typing import Any
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
from pytest import fixture
|
||||
|
||||
|
||||
@fixture
|
||||
def exclude_list(request: Any) -> list[str]:
|
||||
"""Fixture that returns a list of environment variables to exclude."""
|
||||
return request.param if hasattr(request, "param") else []
|
||||
|
||||
|
||||
@fixture
|
||||
def override_env_param_dict(request: Any) -> dict[str, str]:
|
||||
"""Fixture that returns a dict of environment variables to override."""
|
||||
return request.param if hasattr(request, "param") else {}
|
||||
|
||||
|
||||
@fixture()
|
||||
def foundry_unit_test_env(monkeypatch, exclude_list, override_env_param_dict): # type: ignore
|
||||
"""Fixture to set environment variables for Foundry settings."""
|
||||
|
||||
if exclude_list is None:
|
||||
exclude_list = []
|
||||
|
||||
if override_env_param_dict is None:
|
||||
override_env_param_dict = {}
|
||||
|
||||
env_vars = {
|
||||
"FOUNDRY_PROJECT_ENDPOINT": "https://test-project.services.ai.azure.com/",
|
||||
"FOUNDRY_MODEL": "test-gpt-4o",
|
||||
}
|
||||
|
||||
env_vars.update(override_env_param_dict) # type: ignore
|
||||
|
||||
for key, value in env_vars.items():
|
||||
if key in exclude_list:
|
||||
monkeypatch.delenv(key, raising=False) # type: ignore
|
||||
continue
|
||||
monkeypatch.setenv(key, value) # type: ignore
|
||||
|
||||
return env_vars
|
||||
|
||||
|
||||
@fixture
|
||||
def mock_agents_client() -> MagicMock:
|
||||
"""Fixture that provides a mock AgentsClient."""
|
||||
mock_client = MagicMock()
|
||||
|
||||
# Mock agents property
|
||||
mock_client.create_agent = AsyncMock()
|
||||
mock_client.delete_agent = AsyncMock()
|
||||
|
||||
# Mock agent creation response
|
||||
mock_agent = MagicMock()
|
||||
mock_agent.id = "test-agent-id"
|
||||
mock_client.create_agent.return_value = mock_agent
|
||||
|
||||
# Mock threads property
|
||||
mock_client.threads = MagicMock()
|
||||
mock_client.threads.create = AsyncMock()
|
||||
mock_client.messages.create = AsyncMock()
|
||||
|
||||
# Mock runs property
|
||||
mock_client.runs = MagicMock()
|
||||
mock_client.runs.list = AsyncMock()
|
||||
mock_client.runs.cancel = AsyncMock()
|
||||
mock_client.runs.stream = AsyncMock()
|
||||
mock_client.runs.submit_tool_outputs_stream = AsyncMock()
|
||||
|
||||
return mock_client
|
||||
|
||||
|
||||
@fixture
|
||||
def mock_azure_credential() -> MagicMock:
|
||||
"""Fixture that provides a mock AsyncTokenCredential."""
|
||||
return MagicMock()
|
||||
@@ -0,0 +1,374 @@
|
||||
# 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
|
||||
@@ -0,0 +1,507 @@
|
||||
# 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