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