From e9b3a5bbc78db3ac31198d55941b46521a84dcbb Mon Sep 17 00:00:00 2001 From: Eduard van Valkenburg Date: Fri, 13 Feb 2026 14:34:15 +0100 Subject: [PATCH 01/22] Python: fix: prevent repeating instructions in continued Responses API conversations (#3909) * fix: prevent repeating instructions in continued Responses API conversations - Instructions are now only prepended to messages on the first turn - When conversation_id/response_id exists (continuation), instructions are skipped - Covers OpenAI and Azure Responses API paths - Adds regression tests for all continuation scenarios Fixes #3498 * Apply lint fixes to continuation tests * Consolidate responses continuation tests --- .../openai/_responses_client.py | 6 +- .../openai/test_openai_responses_client.py | 84 +++++++++++++++++++ 2 files changed, 89 insertions(+), 1 deletion(-) diff --git a/python/packages/core/agent_framework/openai/_responses_client.py b/python/packages/core/agent_framework/openai/_responses_client.py index 5ab414dc85..55fdaeeda3 100644 --- a/python/packages/core/agent_framework/openai/_responses_client.py +++ b/python/packages/core/agent_framework/openai/_responses_client.py @@ -782,8 +782,12 @@ class RawOpenAIResponsesClient( # type: ignore[misc] # messages # Handle instructions by prepending to messages as system message - if instructions := options.get("instructions"): + # Only prepend instructions for the first turn (when no conversation/response ID exists) + conversation_id = self._get_current_conversation_id(options, **kwargs) + if (instructions := options.get("instructions")) and not conversation_id: + # First turn: prepend instructions as system message messages = prepend_instructions_to_messages(list(messages), instructions, role="system") + # Continuation turn: instructions already exist in conversation context, skip prepending request_input = self._prepare_messages_for_openai(messages) if not request_input: raise ServiceInvalidRequestError("Messages are required for chat completions") diff --git a/python/packages/core/tests/openai/test_openai_responses_client.py b/python/packages/core/tests/openai/test_openai_responses_client.py index a83c4a398b..749939783e 100644 --- a/python/packages/core/tests/openai/test_openai_responses_client.py +++ b/python/packages/core/tests/openai/test_openai_responses_client.py @@ -2168,6 +2168,90 @@ async def test_conversation_id_precedence_kwargs_over_options() -> None: assert "conversation" not in run_opts +def _create_mock_responses_text_response(*, response_id: str) -> MagicMock: + mock_response = MagicMock() + mock_response.id = response_id + mock_response.model = "test-model" + mock_response.created_at = 1000000000 + mock_response.output_parsed = None + mock_response.metadata = {} + mock_response.usage = None + mock_response.finish_reason = None + + mock_message_content = MagicMock() + mock_message_content.type = "output_text" + mock_message_content.text = "Hello! How can I help?" + mock_message_content.annotations = [] + + mock_message_item = MagicMock() + mock_message_item.type = "message" + mock_message_item.content = [mock_message_content] + + mock_response.output = [mock_message_item] + return mock_response + + +async def test_instructions_sent_first_turn_then_skipped_for_continuation() -> None: + client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") + mock_response = _create_mock_responses_text_response(response_id="resp_123") + + with patch.object(client.client.responses, "create", return_value=mock_response) as mock_create: + await client.get_response( + messages=[Message(role="user", text="Hello")], + options={"instructions": "Reply in uppercase."}, + ) + + first_input_messages = mock_create.call_args.kwargs["input"] + assert len(first_input_messages) == 2 + assert first_input_messages[0]["role"] == "system" + assert any("Reply in uppercase" in str(c) for c in first_input_messages[0]["content"]) + assert first_input_messages[1]["role"] == "user" + + await client.get_response( + messages=[Message(role="user", text="Tell me a joke")], + options={"instructions": "Reply in uppercase.", "conversation_id": "resp_123"}, + ) + + second_input_messages = mock_create.call_args.kwargs["input"] + assert len(second_input_messages) == 1 + assert second_input_messages[0]["role"] == "user" + assert not any(message["role"] == "system" for message in second_input_messages) + + +@pytest.mark.parametrize("conversation_id", ["resp_456", "conv_abc123"]) +async def test_instructions_not_repeated_for_continuation_ids(conversation_id: str) -> None: + client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") + mock_response = _create_mock_responses_text_response(response_id="resp_456") + + with patch.object(client.client.responses, "create", return_value=mock_response) as mock_create: + await client.get_response( + messages=[Message(role="user", text="Continue conversation")], + options={"instructions": "Be helpful.", "conversation_id": conversation_id}, + ) + + input_messages = mock_create.call_args.kwargs["input"] + assert len(input_messages) == 1 + assert input_messages[0]["role"] == "user" + assert not any(message["role"] == "system" for message in input_messages) + + +async def test_instructions_included_without_conversation_id() -> None: + client = OpenAIResponsesClient(model_id="test-model", api_key="test-key") + mock_response = _create_mock_responses_text_response(response_id="resp_new") + + with patch.object(client.client.responses, "create", return_value=mock_response) as mock_create: + await client.get_response( + messages=[Message(role="user", text="Hello")], + options={"instructions": "You are a helpful assistant."}, + ) + + input_messages = mock_create.call_args.kwargs["input"] + assert len(input_messages) == 2 + assert input_messages[0]["role"] == "system" + assert any("helpful assistant" in str(c) for c in input_messages[0]["content"]) + assert input_messages[1]["role"] == "user" + + def test_with_callable_api_key() -> None: """Test OpenAIResponsesClient initialization with callable API key.""" From f3ea8721563293befe0fc14174020b5f417950f9 Mon Sep 17 00:00:00 2001 From: Eduard van Valkenburg Date: Fri, 13 Feb 2026 14:55:39 +0100 Subject: [PATCH 02/22] Add default in-memory history provider for workflow agents (#3918) --- .../core/agent_framework/_workflows/_agent.py | 20 +++++++++++-- .../tests/workflow/test_workflow_agent.py | 30 +++++++++++++++++++ 2 files changed, 48 insertions(+), 2 deletions(-) diff --git a/python/packages/core/agent_framework/_workflows/_agent.py b/python/packages/core/agent_framework/_workflows/_agent.py index ef2b127d45..e4152b77d1 100644 --- a/python/packages/core/agent_framework/_workflows/_agent.py +++ b/python/packages/core/agent_framework/_workflows/_agent.py @@ -12,7 +12,13 @@ from datetime import datetime, timezone from typing import TYPE_CHECKING, Any, ClassVar, Literal, cast, overload from .._agents import BaseAgent -from .._sessions import AgentSession, BaseContextProvider, BaseHistoryProvider, SessionContext +from .._sessions import ( + AgentSession, + BaseContextProvider, + BaseHistoryProvider, + InMemoryHistoryProvider, + SessionContext, +) from .._types import ( AgentResponse, AgentResponseUpdate, @@ -112,7 +118,17 @@ class WorkflowAgent(BaseAgent): if not any(is_type_compatible(list[Message], input_type) for input_type in start_executor.input_types): raise ValueError("Workflow's start executor cannot handle list[Message]") - super().__init__(id=id, name=name, description=description, context_providers=context_providers, **kwargs) + resolved_context_providers = list(context_providers) if context_providers is not None else [] + if not resolved_context_providers: + resolved_context_providers.append(InMemoryHistoryProvider("memory")) + + super().__init__( + id=id, + name=name, + description=description, + context_providers=resolved_context_providers, + **kwargs, + ) self._workflow: Workflow = workflow self._pending_requests: dict[str, WorkflowEvent[Any]] = {} diff --git a/python/packages/core/tests/workflow/test_workflow_agent.py b/python/packages/core/tests/workflow/test_workflow_agent.py index 5adf82dd57..b2fbded39b 100644 --- a/python/packages/core/tests/workflow/test_workflow_agent.py +++ b/python/packages/core/tests/workflow/test_workflow_agent.py @@ -14,6 +14,7 @@ from agent_framework import ( AgentSession, Content, Executor, + InMemoryHistoryProvider, Message, ResponseStream, SupportsAgentRun, @@ -562,6 +563,35 @@ class TestWorkflowAgent: assert len(capturing_executor.received_messages) == 1 assert capturing_executor.received_messages[0].text == "Just a new message" + async def test_workflow_as_agent_adds_default_history_provider(self) -> None: + """Test that workflow.as_agent() defaults to in-memory history when no providers are configured.""" + capturing_executor = ConversationHistoryCapturingExecutor(id="default_history_provider_test") + workflow = WorkflowBuilder(start_executor=capturing_executor).build() + agent = workflow.as_agent(name="Default History Provider Agent") + session = AgentSession() + + await agent.run("first message", session=session) + await agent.run("second message", session=session) + + assert any(isinstance(provider, InMemoryHistoryProvider) for provider in agent.context_providers) + texts = [message.text for message in capturing_executor.received_messages] + assert "first message" in texts + assert "second message" in texts + + async def test_workflow_agent_keeps_explicit_context_providers(self) -> None: + """Test that WorkflowAgent does not append defaults when context providers are explicitly provided.""" + workflow = WorkflowBuilder( + start_executor=ConversationHistoryCapturingExecutor(id="explicit_provider_test") + ).build() + explicit_provider = InMemoryHistoryProvider("custom-memory") + agent = WorkflowAgent( + workflow=workflow, + name="Explicit Provider Agent", + context_providers=[explicit_provider], + ) + + assert agent.context_providers == [explicit_provider] + async def test_checkpoint_storage_passed_to_workflow(self) -> None: """Test that checkpoint_storage parameter is passed through to the workflow.""" from agent_framework import InMemoryCheckpointStorage From a39fd69f76965c927f19496f8d6ae0d230a1f2a3 Mon Sep 17 00:00:00 2001 From: Eduard van Valkenburg Date: Fri, 13 Feb 2026 14:56:19 +0100 Subject: [PATCH 03/22] Add memory run snippet tag for docs extraction (#3921) --- python/samples/01-get-started/04_memory.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/python/samples/01-get-started/04_memory.py b/python/samples/01-get-started/04_memory.py index d35ef21581..03ef75257c 100644 --- a/python/samples/01-get-started/04_memory.py +++ b/python/samples/01-get-started/04_memory.py @@ -78,6 +78,7 @@ async def main() -> None: ) # + # session = agent.create_session() # The provider doesn't know the user yet — it will ask for a name @@ -93,6 +94,7 @@ async def main() -> None: print(f"Agent: {result}\n") print(f"[Memory] Stored user name: {memory.user_name}") + # if __name__ == "__main__": From 4452997e8dcca60543113bfe69db41acfa23a7c8 Mon Sep 17 00:00:00 2001 From: Eduard van Valkenburg Date: Fri, 13 Feb 2026 15:02:36 +0100 Subject: [PATCH 04/22] Python: Replace wildcard imports with explicit imports (#3908) * Python: Replace wildcard imports with explicit imports - Replace all 'from ... import *' with explicit symbol imports - Add __all__ declarations to namespace packages for re-exports - Update CODING_STANDARD.md to prohibit wildcard imports - Maintain exported API and preserve all functionality fixes #3605 * Refine wildcard guidance example text * Simplify explicit exports without self-aliases --- python/CODING_STANDARD.md | 27 +- .../packages/core/agent_framework/__init__.py | 291 +++++++++++++++++- .../core/agent_framework/openai/__init__.py | 37 ++- .../agent_framework/lab/gaia/__init__.py | 26 +- .../agent_framework/lab/lightning/__init__.py | 4 +- .../agent_framework/lab/tau2/__init__.py | 18 +- 6 files changed, 378 insertions(+), 25 deletions(-) diff --git a/python/CODING_STANDARD.md b/python/CODING_STANDARD.md index e052b60669..df143ce62a 100644 --- a/python/CODING_STANDARD.md +++ b/python/CODING_STANDARD.md @@ -403,22 +403,37 @@ If in doubt, use the link above to read much more considerations of what to do a ### Explicit Exports -> **Note:** This convention is being adopted. See [#3605](https://github.com/microsoft/agent-framework/issues/3605) for progress. +**All wildcard imports (`from ... import *`) are prohibited** in production code, including both `.py` and `.pyi` files. Always use explicit import lists to maintain clarity and avoid namespace pollution. -Define `__all__` in each module to explicitly declare the public API. Avoid using `from module import *` in `__init__.py` files as it can impact performance and makes the public API unclear: +Define `__all__` in each module to explicitly declare the public API, then import specific symbols by name: ```python -# ✅ Preferred - explicit __all__ and imports +# ✅ Preferred - explicit __all__ and named imports __all__ = ["Agent", "Message", "ChatResponse"] from ._agents import Agent from ._types import Message, ChatResponse -# ❌ Avoid - star imports -from ._agents import * -from ._types import * +# ✅ For many exports, use parenthesized multi-line imports +from ._types import ( + AgentResponse, + ChatResponse, + Message, + ResponseStream, +) + +# ❌ Prohibited pattern: wildcard/star imports (do not use) +# from ._agents import +# from ._types import ``` +**Rationale:** +- **Clarity**: Explicit imports make it clear exactly what is being exported and used +- **IDE Support**: Enables better autocomplete, go-to-definition, and refactoring +- **Type Checking**: Improves static analysis and type checker accuracy +- **Maintenance**: Makes it easier to track symbol usage and detect breaking changes +- **Performance**: Avoids unnecessary symbol resolution during module import + ## Performance considerations ### Cache Expensive Computations diff --git a/python/packages/core/agent_framework/__init__.py b/python/packages/core/agent_framework/__init__.py index 48095326de..f433df94b8 100644 --- a/python/packages/core/agent_framework/__init__.py +++ b/python/packages/core/agent_framework/__init__.py @@ -9,13 +9,284 @@ except importlib.metadata.PackageNotFoundError: _version = "0.0.0" # Fallback for development mode __version__: Final[str] = _version -from ._agents import * # noqa: F403 -from ._clients import * # noqa: F403 -from ._logging import * # noqa: F403 -from ._mcp import * # noqa: F403 -from ._middleware import * # noqa: F403 -from ._sessions import * # noqa: F403 -from ._telemetry import * # noqa: F403 -from ._tools import * # noqa: F403 -from ._types import * # noqa: F403 -from ._workflows import * # noqa: F403 +from ._agents import Agent, BaseAgent, RawAgent, SupportsAgentRun +from ._clients import ( + BaseChatClient, + SupportsChatGetResponse, + SupportsCodeInterpreterTool, + SupportsFileSearchTool, + SupportsImageGenerationTool, + SupportsMCPTool, + SupportsWebSearchTool, +) +from ._logging import get_logger, setup_logging +from ._mcp import MCPStdioTool, MCPStreamableHTTPTool, MCPWebsocketTool +from ._middleware import ( + AgentContext, + AgentMiddleware, + AgentMiddlewareLayer, + AgentMiddlewareTypes, + ChatAndFunctionMiddlewareTypes, + ChatContext, + ChatMiddleware, + ChatMiddlewareLayer, + ChatMiddlewareTypes, + FunctionInvocationContext, + FunctionMiddleware, + FunctionMiddlewareTypes, + MiddlewareException, + MiddlewareTermination, + MiddlewareType, + MiddlewareTypes, + agent_middleware, + chat_middleware, + function_middleware, +) +from ._sessions import ( + AgentSession, + BaseContextProvider, + BaseHistoryProvider, + InMemoryHistoryProvider, + SessionContext, + register_state_type, +) +from ._telemetry import ( + AGENT_FRAMEWORK_USER_AGENT, + APP_INFO, + USER_AGENT_KEY, + USER_AGENT_TELEMETRY_DISABLED_ENV_VAR, + prepend_agent_framework_to_user_agent, +) +from ._tools import ( + FunctionInvocationConfiguration, + FunctionInvocationLayer, + FunctionTool, + normalize_function_invocation_configuration, + tool, +) +from ._types import ( + AgentResponse, + AgentResponseUpdate, + Annotation, + ChatOptions, + ChatResponse, + ChatResponseUpdate, + Content, + ContinuationToken, + FinalT, + FinishReason, + FinishReasonLiteral, + Message, + OuterFinalT, + OuterUpdateT, + ResponseStream, + Role, + RoleLiteral, + TextSpanRegion, + ToolMode, + UpdateT, + UsageDetails, + add_usage_details, + detect_media_type_from_base64, + map_chat_to_agent_update, + merge_chat_options, + normalize_messages, + normalize_tools, + prepend_instructions_to_messages, + validate_chat_options, + validate_tool_mode, + validate_tools, +) +from ._workflows import ( + DEFAULT_MAX_ITERATIONS, + AgentExecutor, + AgentExecutorRequest, + AgentExecutorResponse, + Case, + CheckpointStorage, + Default, + Edge, + EdgeCondition, + EdgeDuplicationError, + Executor, + FanInEdgeGroup, + FanOutEdgeGroup, + FileCheckpointStorage, + FunctionExecutor, + GraphConnectivityError, + InMemoryCheckpointStorage, + InProcRunnerContext, + Runner, + RunnerContext, + SingleEdgeGroup, + SubWorkflowRequestMessage, + SubWorkflowResponseMessage, + SwitchCaseEdgeGroup, + SwitchCaseEdgeGroupCase, + SwitchCaseEdgeGroupDefault, + TypeCompatibilityError, + ValidationTypeEnum, + Workflow, + WorkflowAgent, + WorkflowBuilder, + WorkflowCheckpoint, + WorkflowCheckpointException, + WorkflowContext, + WorkflowConvergenceException, + WorkflowErrorDetails, + WorkflowEvent, + WorkflowEventSource, + WorkflowEventType, + WorkflowException, + WorkflowExecutor, + WorkflowMessage, + WorkflowRunnerException, + WorkflowRunResult, + WorkflowRunState, + WorkflowValidationError, + WorkflowViz, + create_edge_runner, + executor, + handler, + resolve_agent_id, + response_handler, + validate_workflow_graph, +) + +__all__ = [ + "AGENT_FRAMEWORK_USER_AGENT", + "APP_INFO", + "DEFAULT_MAX_ITERATIONS", + "USER_AGENT_KEY", + "USER_AGENT_TELEMETRY_DISABLED_ENV_VAR", + "Agent", + "AgentContext", + "AgentExecutor", + "AgentExecutorRequest", + "AgentExecutorResponse", + "AgentMiddleware", + "AgentMiddlewareLayer", + "AgentMiddlewareTypes", + "AgentResponse", + "AgentResponseUpdate", + "AgentSession", + "Annotation", + "BaseAgent", + "BaseChatClient", + "BaseContextProvider", + "BaseHistoryProvider", + "Case", + "ChatAndFunctionMiddlewareTypes", + "ChatContext", + "ChatMiddleware", + "ChatMiddlewareLayer", + "ChatMiddlewareTypes", + "ChatOptions", + "ChatResponse", + "ChatResponseUpdate", + "CheckpointStorage", + "Content", + "ContinuationToken", + "Default", + "Edge", + "EdgeCondition", + "EdgeDuplicationError", + "Executor", + "FanInEdgeGroup", + "FanOutEdgeGroup", + "FileCheckpointStorage", + "FinalT", + "FinishReason", + "FinishReasonLiteral", + "FunctionExecutor", + "FunctionInvocationConfiguration", + "FunctionInvocationContext", + "FunctionInvocationLayer", + "FunctionMiddleware", + "FunctionMiddlewareTypes", + "FunctionTool", + "GraphConnectivityError", + "InMemoryCheckpointStorage", + "InMemoryHistoryProvider", + "InProcRunnerContext", + "MCPStdioTool", + "MCPStreamableHTTPTool", + "MCPWebsocketTool", + "Message", + "MiddlewareException", + "MiddlewareTermination", + "MiddlewareType", + "MiddlewareTypes", + "OuterFinalT", + "OuterUpdateT", + "RawAgent", + "ResponseStream", + "Role", + "RoleLiteral", + "Runner", + "RunnerContext", + "SessionContext", + "SingleEdgeGroup", + "SubWorkflowRequestMessage", + "SubWorkflowResponseMessage", + "SupportsAgentRun", + "SupportsChatGetResponse", + "SupportsCodeInterpreterTool", + "SupportsFileSearchTool", + "SupportsImageGenerationTool", + "SupportsMCPTool", + "SupportsWebSearchTool", + "SwitchCaseEdgeGroup", + "SwitchCaseEdgeGroupCase", + "SwitchCaseEdgeGroupDefault", + "TextSpanRegion", + "ToolMode", + "TypeCompatibilityError", + "UpdateT", + "UsageDetails", + "ValidationTypeEnum", + "Workflow", + "WorkflowAgent", + "WorkflowBuilder", + "WorkflowCheckpoint", + "WorkflowCheckpointException", + "WorkflowContext", + "WorkflowConvergenceException", + "WorkflowErrorDetails", + "WorkflowEvent", + "WorkflowEventSource", + "WorkflowEventType", + "WorkflowException", + "WorkflowExecutor", + "WorkflowMessage", + "WorkflowRunResult", + "WorkflowRunState", + "WorkflowRunnerException", + "WorkflowValidationError", + "WorkflowViz", + "add_usage_details", + "agent_middleware", + "chat_middleware", + "create_edge_runner", + "detect_media_type_from_base64", + "executor", + "function_middleware", + "get_logger", + "handler", + "map_chat_to_agent_update", + "merge_chat_options", + "normalize_function_invocation_configuration", + "normalize_messages", + "normalize_tools", + "prepend_agent_framework_to_user_agent", + "prepend_instructions_to_messages", + "register_state_type", + "resolve_agent_id", + "response_handler", + "setup_logging", + "tool", + "validate_chat_options", + "validate_tool_mode", + "validate_tools", + "validate_workflow_graph", +] diff --git a/python/packages/core/agent_framework/openai/__init__.py b/python/packages/core/agent_framework/openai/__init__.py index 008e2cb54c..c5e196b022 100644 --- a/python/packages/core/agent_framework/openai/__init__.py +++ b/python/packages/core/agent_framework/openai/__init__.py @@ -1,8 +1,33 @@ # Copyright (c) Microsoft. All rights reserved. -from ._assistant_provider import * # noqa: F403 -from ._assistants_client import * # noqa: F403 -from ._chat_client import * # noqa: F403 -from ._exceptions import * # noqa: F403 -from ._responses_client import * # noqa: F403 -from ._shared import * # noqa: F403 +from ._assistant_provider import OpenAIAssistantProvider +from ._assistants_client import ( + AssistantToolResources, + OpenAIAssistantsClient, + OpenAIAssistantsOptions, +) +from ._chat_client import OpenAIChatClient, OpenAIChatOptions +from ._exceptions import ContentFilterResultSeverity, OpenAIContentFilterException +from ._responses_client import ( + OpenAIContinuationToken, + OpenAIResponsesClient, + OpenAIResponsesOptions, + RawOpenAIResponsesClient, +) +from ._shared import OpenAISettings + +__all__ = [ + "AssistantToolResources", + "ContentFilterResultSeverity", + "OpenAIAssistantProvider", + "OpenAIAssistantsClient", + "OpenAIAssistantsOptions", + "OpenAIChatClient", + "OpenAIChatOptions", + "OpenAIContentFilterException", + "OpenAIContinuationToken", + "OpenAIResponsesClient", + "OpenAIResponsesOptions", + "OpenAISettings", + "RawOpenAIResponsesClient", +] diff --git a/python/packages/lab/namespace/agent_framework/lab/gaia/__init__.py b/python/packages/lab/namespace/agent_framework/lab/gaia/__init__.py index 391a57b4a3..a6e7da4cdc 100644 --- a/python/packages/lab/namespace/agent_framework/lab/gaia/__init__.py +++ b/python/packages/lab/namespace/agent_framework/lab/gaia/__init__.py @@ -1,4 +1,28 @@ # Copyright (c) Microsoft. All rights reserved. # Import and re-export from the actual implementation -from agent_framework_lab_gaia import * # noqa: F403 +from agent_framework_lab_gaia import ( + GAIA, + Evaluation, + Evaluator, + GAIATelemetryConfig, + Prediction, + Task, + TaskResult, + TaskRunner, + gaia_scorer, + viewer_main, +) + +__all__ = [ + "GAIA", + "Evaluation", + "Evaluator", + "GAIATelemetryConfig", + "Prediction", + "Task", + "TaskResult", + "TaskRunner", + "gaia_scorer", + "viewer_main", +] diff --git a/python/packages/lab/namespace/agent_framework/lab/lightning/__init__.py b/python/packages/lab/namespace/agent_framework/lab/lightning/__init__.py index 3a857f38da..21ad032b70 100644 --- a/python/packages/lab/namespace/agent_framework/lab/lightning/__init__.py +++ b/python/packages/lab/namespace/agent_framework/lab/lightning/__init__.py @@ -1,4 +1,6 @@ # Copyright (c) Microsoft. All rights reserved. # Import and re-export from the actual implementation -from agent_framework_lab_lightning import * # noqa: F403 +from agent_framework_lab_lightning import AgentFrameworkTracer + +__all__ = ["AgentFrameworkTracer"] diff --git a/python/packages/lab/namespace/agent_framework/lab/tau2/__init__.py b/python/packages/lab/namespace/agent_framework/lab/tau2/__init__.py index 56fda49217..7229589746 100644 --- a/python/packages/lab/namespace/agent_framework/lab/tau2/__init__.py +++ b/python/packages/lab/namespace/agent_framework/lab/tau2/__init__.py @@ -1,4 +1,20 @@ # Copyright (c) Microsoft. All rights reserved. # Import and re-export from the actual implementation -from agent_framework_lab_tau2 import * # noqa: F403 +from agent_framework_lab_tau2 import ( + ASSISTANT_AGENT_ID, + ORCHESTRATOR_ID, + USER_SIMULATOR_ID, + TaskRunner, + patch_env_set_state, + unpatch_env_set_state, +) + +__all__ = [ + "ASSISTANT_AGENT_ID", + "ORCHESTRATOR_ID", + "USER_SIMULATOR_ID", + "TaskRunner", + "patch_env_set_state", + "unpatch_env_set_state", +] From 3168eb4870a6dd452e9fede81508a46ee4a48148 Mon Sep 17 00:00:00 2001 From: westey <164392973+westey-m@users.noreply.github.com> Date: Fri, 13 Feb 2026 14:08:07 +0000 Subject: [PATCH 05/22] .NET: [BREAKING] Add session StateBag for state storage and support multiple providers on the Agent (#3806) * .NET: [BREAKING] Add session statebag to use for state storage instead of inside providers (#3737) * Add a StateBag to AgentSession and pass Agent and AgentSession to AIContextProvider and ChatHistoryProviders * Convert all AIContextProviders to use the statebag * Update InMemoryChatHistoryProvider to use StateBag * Update Comsos and Workflow ChatHistoryProviders * Update 3rd party chat history storage sample. * Remove serialize method from providers * Replacing provider factories with properties * Remove Providers from Session and flatten state bag serialization * Update samples to use getservice on agent * Updated additional session types to serialize statebag * Fix regression * Address PR comments * Address PR comments. * Fix formatting * Fix unit tests * Remove InMemoryAgentSession since it is not required anymore. * Address PR comments * Convert sessions for A2AAgent, ChatClientAgent, CopilotStudioAgent and GithubCopilotAgent to use regular json serialization. * Fix durable agent session jso usgae * Add jso to InMemory and Workflow ChatHistoryProviders * Update InMemoryChatHistoryProvider to use an options class for it's many optional settings. * Apply suggestions from code review Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * Address PR feedback * Fix verification bug. * Improve state bag thread safety * Address PR comments and fix unit tests * Address PR comments * Fix unit test --------- Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * Add a public StateKey property to providers (#3810) * .NET: [BREAKING] Update providers in such a way that they can participate in a pipeline (#3846) * Make providers pipeline capable * Fix unit tests * Move source stamping to providers from base class * Also update samples. * Address PR comments * Rename AsAgentRequestMessageSourcedMessage to WithAgentRequestMessageSource * .NET: [BREAKING] Add consistent message filtering to all providers. (#3851) * Add consistent message filtering to all providers. * Remove old chat history filtering classes * Fix merge issues * Fix unit test * Enforce non-nullable property * Fix merging bug and make troubleshooting source info easier by adding tostring implementation * .NET: [BREAKING] Add support for multiple AIContextProviders on a ChatClientAgent (#3863) * Add support for multiple AIContextProviders on a ChatClientAgent * Address PR comments and fix tests * Address PR comments. * .NET: [BREAKING]Delay AIContext Materialization until the end of the pipeline is reached. (#3883) * Delay AIContext Materialization until the end of the pipeline is reached. * Address PR comments. * Address PR comments * Modify InvokedContext to be immutable (#3888) * .NET: Address Feedback on StateBag feature branch PR (#3910) * Address Feedback on statebag feature branch PR * Update dotnet/src/Microsoft.Agents.AI.DurableTask/CHANGELOG.md Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * Address PR comments --------- Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --------- Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- .../Program.cs | 40 +- .../Program.cs | 19 +- .../Program.cs | 19 +- .../Program.cs | 68 +- .../Program.cs | 12 +- .../Program.cs | 2 +- .../Program.cs | 2 +- .../Program.cs | 93 +- .../Agent_Step16_ChatReduction/Program.cs | 7 +- .../Program.cs | 172 ++-- .../src/Microsoft.Agents.AI.A2A/A2AAgent.cs | 2 +- .../A2AAgentSession.cs | 51 +- .../A2AJsonUtilities.cs | 2 +- .../AIContext.cs | 33 +- .../AIContextProvider.cs | 146 +-- .../AgentAbstractionsJsonUtilities.cs | 5 +- .../AgentRequestMessageSourceAttribution.cs | 11 + .../AgentRequestMessageSourceType.cs | 6 + .../AgentSession.cs | 20 + .../AgentSessionStateBag.cs | 145 +++ .../AgentSessionStateBagJsonConverter.cs | 28 + .../AgentSessionStateBagValue.cs | 182 ++++ .../AgentSessionStateBagValueJsonConverter.cs | 27 + .../ChatHistoryProvider.cs | 119 +-- .../ChatHistoryProviderExtensions.cs | 52 -- .../ChatHistoryProviderMessageFilter.cs | 74 -- .../ChatMessageExtensions.cs | 2 +- .../InMemoryAgentSession.cs | 124 --- .../InMemoryChatHistoryProvider.cs | 266 ++---- .../InMemoryChatHistoryProviderOptions.cs | 93 ++ .../ServiceIdAgentSession.cs | 104 --- .../PersistentAgentsClientExtensions.cs | 4 +- .../AzureAIProjectChatClientExtensions.cs | 4 +- .../CopilotStudioAgent.cs | 2 +- .../CopilotStudioAgentSession.cs | 40 +- .../CopilotStudioJsonUtilities.cs | 48 + .../CosmosChatHistoryProvider.cs | 399 ++++----- .../CosmosDBChatExtensions.cs | 25 +- .../CHANGELOG.md | 1 + .../DurableAgentSession.cs | 26 +- .../GitHubCopilotAgent.cs | 2 +- .../GitHubCopilotAgentSession.cs | 44 +- .../GitHubCopilotJsonUtilities.cs | 2 +- .../Mem0JsonUtilities.cs | 2 +- .../Microsoft.Agents.AI.Mem0/Mem0Provider.cs | 229 ++--- .../Mem0ProviderOptions.cs | 30 + .../OpenAIAssistantClientExtensions.cs | 4 +- .../WorkflowChatHistoryProvider.cs | 86 +- .../WorkflowHostAgent.cs | 2 +- .../WorkflowSession.cs | 21 +- .../Microsoft.Agents.AI/AgentJsonUtilities.cs | 4 +- .../ChatClient/ChatClientAgent.cs | 309 ++++--- .../ChatClient/ChatClientAgentOptions.cs | 54 +- .../ChatClient/ChatClientAgentSession.cs | 162 +--- .../Memory/ChatHistoryMemoryProvider.cs | 283 +++--- .../ChatHistoryMemoryProviderOptions.cs | 33 + .../Microsoft.Agents.AI/TextSearchProvider.cs | 131 ++- .../TextSearchProviderOptions.cs | 29 + .../AnthropicChatCompletionFixture.cs | 6 +- .../AIProjectClientFixture.cs | 6 +- .../A2AAgentSessionTests.cs | 20 +- .../AIContextProviderTests.cs | 267 +----- .../AIContextTests.cs | 15 +- ...entRequestMessageSourceAttributionTests.cs | 43 + .../AgentRequestMessageSourceTypeTests.cs | 40 + .../AgentSessionStateBagTests.cs | 840 ++++++++++++++++++ .../AgentSessionTests.cs | 15 + .../ChatHistoryProviderExtensionsTests.cs | 141 --- .../ChatHistoryProviderMessageFilterTests.cs | 221 ----- .../ChatHistoryProviderTests.cs | 205 +---- .../ChatMessageExtensionsTests.cs | 20 +- .../InMemoryAgentSessionTests.cs | 155 ---- .../InMemoryChatHistoryProviderTests.cs | 621 +++++-------- .../ServiceIdAgentSessionTests.cs | 119 --- .../TestJsonSerializerContext.cs | 3 - ...AzureAIProjectChatClientExtensionsTests.cs | 28 +- .../CosmosChatHistoryProviderTests.cs | 379 +++++--- .../DurableAgentSessionTests.cs | 40 +- .../BasicStreamingTests.cs | 33 +- .../ForwardedPropertiesTests.cs | 21 +- .../SharedStateTests.cs | 21 +- ...AGUIEndpointRouteBuilderExtensionsTests.cs | 29 +- .../Mem0ProviderTests.cs | 93 +- .../Mem0ProviderTests.cs | 286 ++++-- .../ChatClient/ChatClientAgentOptionsTests.cs | 36 +- .../ChatClient/ChatClientAgentSessionTests.cs | 192 ++-- .../ChatClient/ChatClientAgentTests.cs | 570 ++++++++++-- ...hatClientAgent_BackgroundResponsesTests.cs | 52 +- ...tClientAgent_ChatHistoryManagementTests.cs | 74 +- .../ChatClientAgent_CreateSessionTests.cs | 71 -- ...ChatClientAgent_DeserializeSessionTests.cs | 80 -- .../Data/TextSearchProviderTests.cs | 419 ++++++--- .../Memory/ChatHistoryMemoryProviderTests.cs | 313 +++++-- .../TestJsonSerializerContext.cs | 1 + .../AgentWorkflowBuilderTests.cs | 2 +- .../InProcessExecutionTests.cs | 2 +- .../RoleCheckAgent.cs | 2 +- .../Sample/06_GroupChat_Workflow.cs | 2 +- .../TestEchoAgent.cs | 27 +- .../TestReplayAgent.cs | 2 +- .../TestRequestAgent.cs | 18 +- .../WorkflowHostSmokeTests.cs | 8 +- .../OpenAIChatCompletionFixture.cs | 6 +- .../OpenAIResponseFixture.cs | 6 +- 104 files changed, 5120 insertions(+), 4332 deletions(-) create mode 100644 dotnet/src/Microsoft.Agents.AI.Abstractions/AgentSessionStateBag.cs create mode 100644 dotnet/src/Microsoft.Agents.AI.Abstractions/AgentSessionStateBagJsonConverter.cs create mode 100644 dotnet/src/Microsoft.Agents.AI.Abstractions/AgentSessionStateBagValue.cs create mode 100644 dotnet/src/Microsoft.Agents.AI.Abstractions/AgentSessionStateBagValueJsonConverter.cs delete mode 100644 dotnet/src/Microsoft.Agents.AI.Abstractions/ChatHistoryProviderExtensions.cs delete mode 100644 dotnet/src/Microsoft.Agents.AI.Abstractions/ChatHistoryProviderMessageFilter.cs delete mode 100644 dotnet/src/Microsoft.Agents.AI.Abstractions/InMemoryAgentSession.cs create mode 100644 dotnet/src/Microsoft.Agents.AI.Abstractions/InMemoryChatHistoryProviderOptions.cs delete mode 100644 dotnet/src/Microsoft.Agents.AI.Abstractions/ServiceIdAgentSession.cs create mode 100644 dotnet/src/Microsoft.Agents.AI.CopilotStudio/CopilotStudioJsonUtilities.cs create mode 100644 dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/AgentSessionStateBagTests.cs delete mode 100644 dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/ChatHistoryProviderExtensionsTests.cs delete mode 100644 dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/ChatHistoryProviderMessageFilterTests.cs delete mode 100644 dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/InMemoryAgentSessionTests.cs delete mode 100644 dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/ServiceIdAgentSessionTests.cs delete mode 100644 dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_DeserializeSessionTests.cs diff --git a/dotnet/samples/GettingStarted/AgentProviders/Agent_With_CustomImplementation/Program.cs b/dotnet/samples/GettingStarted/AgentProviders/Agent_With_CustomImplementation/Program.cs index 1c9c9a3964..786fbea36b 100644 --- a/dotnet/samples/GettingStarted/AgentProviders/Agent_With_CustomImplementation/Program.cs +++ b/dotnet/samples/GettingStarted/AgentProviders/Agent_With_CustomImplementation/Program.cs @@ -6,6 +6,7 @@ using System.Runtime.CompilerServices; using System.Text.Json; +using System.Text.Json.Serialization; using Microsoft.Agents.AI; using Microsoft.Extensions.AI; using SampleApp; @@ -28,6 +29,8 @@ namespace SampleApp { public override string? Name => "UpperCaseParrotAgent"; + public readonly ChatHistoryProvider ChatHistoryProvider = new InMemoryChatHistoryProvider(); + protected override ValueTask CreateSessionCoreAsync(CancellationToken cancellationToken = default) => new(new CustomAgentSession()); @@ -38,11 +41,11 @@ namespace SampleApp throw new ArgumentException($"The provided session is not of type {nameof(CustomAgentSession)}.", nameof(session)); } - return new(typedSession.Serialize(jsonSerializerOptions)); + return new(JsonSerializer.SerializeToElement(typedSession, jsonSerializerOptions)); } protected override ValueTask DeserializeSessionCoreAsync(JsonElement serializedState, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default) - => new(new CustomAgentSession(serializedState, jsonSerializerOptions)); + => new(serializedState.Deserialize(jsonSerializerOptions)!); protected override async Task RunCoreAsync(IEnumerable messages, AgentSession? session = null, AgentRunOptions? options = null, CancellationToken cancellationToken = default) { @@ -56,17 +59,14 @@ namespace SampleApp // Get existing messages from the store var invokingContext = new ChatHistoryProvider.InvokingContext(this, session, messages); - var storeMessages = await typedSession.ChatHistoryProvider.InvokingAsync(invokingContext, cancellationToken); + var userAndChatHistoryMessages = await this.ChatHistoryProvider.InvokingAsync(invokingContext, cancellationToken); // Clone the input messages and turn them into response messages with upper case text. List responseMessages = CloneAndToUpperCase(messages, this.Name).ToList(); // Notify the session of the input and output messages. - var invokedContext = new ChatHistoryProvider.InvokedContext(this, session, messages) - { - ResponseMessages = responseMessages - }; - await typedSession.ChatHistoryProvider.InvokedAsync(invokedContext, cancellationToken); + var invokedContext = new ChatHistoryProvider.InvokedContext(this, session, userAndChatHistoryMessages, responseMessages); + await this.ChatHistoryProvider.InvokedAsync(invokedContext, cancellationToken); return new AgentResponse { @@ -88,17 +88,14 @@ namespace SampleApp // Get existing messages from the store var invokingContext = new ChatHistoryProvider.InvokingContext(this, session, messages); - var storeMessages = await typedSession.ChatHistoryProvider.InvokingAsync(invokingContext, cancellationToken); + var userAndChatHistoryMessages = await this.ChatHistoryProvider.InvokingAsync(invokingContext, cancellationToken); // Clone the input messages and turn them into response messages with upper case text. List responseMessages = CloneAndToUpperCase(messages, this.Name).ToList(); // Notify the session of the input and output messages. - var invokedContext = new ChatHistoryProvider.InvokedContext(this, session, messages) - { - ResponseMessages = responseMessages - }; - await typedSession.ChatHistoryProvider.InvokedAsync(invokedContext, cancellationToken); + var invokedContext = new ChatHistoryProvider.InvokedContext(this, session, userAndChatHistoryMessages, responseMessages); + await this.ChatHistoryProvider.InvokedAsync(invokedContext, cancellationToken); foreach (var message in responseMessages) { @@ -140,15 +137,16 @@ namespace SampleApp /// /// A session type for our custom agent that only supports in memory storage of messages. /// - internal sealed class CustomAgentSession : InMemoryAgentSession + internal sealed class CustomAgentSession : AgentSession { - internal CustomAgentSession() { } + internal CustomAgentSession() + { + } - internal CustomAgentSession(JsonElement serializedSessionState, JsonSerializerOptions? jsonSerializerOptions = null) - : base(serializedSessionState, jsonSerializerOptions) { } - - internal new JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null) - => base.Serialize(jsonSerializerOptions); + [JsonConstructor] + internal CustomAgentSession(AgentSessionStateBag stateBag) : base(stateBag) + { + } } } } diff --git a/dotnet/samples/GettingStarted/AgentWithMemory/AgentWithMemory_Step01_ChatHistoryMemory/Program.cs b/dotnet/samples/GettingStarted/AgentWithMemory/AgentWithMemory_Step01_ChatHistoryMemory/Program.cs index 4e2065e0eb..ff4628ef7a 100644 --- a/dotnet/samples/GettingStarted/AgentWithMemory/AgentWithMemory_Step01_ChatHistoryMemory/Program.cs +++ b/dotnet/samples/GettingStarted/AgentWithMemory/AgentWithMemory_Step01_ChatHistoryMemory/Program.cs @@ -37,16 +37,21 @@ AIAgent agent = new AzureOpenAIClient( { ChatOptions = new() { Instructions = "You are good at telling jokes." }, Name = "Joker", - AIContextProviderFactory = (ctx, ct) => new ValueTask(new ChatHistoryMemoryProvider( + AIContextProviders = [new ChatHistoryMemoryProvider( vectorStore, collectionName: "chathistory", vectorDimensions: 3072, - // Configure the scope values under which chat messages will be stored. - // In this case, we are using a fixed user ID and a unique session ID for each new session. - storageScope: new() { UserId = "UID1", SessionId = Guid.NewGuid().ToString() }, - // Configure the scope which would be used to search for relevant prior messages. - // In this case, we are searching for any messages for the user across all sessions. - searchScope: new() { UserId = "UID1" })) + // Callback to configure the initial state of the ChatHistoryMemoryProvider. + // The ChatHistoryMemoryProvider stores its state in the AgentSession and this callback + // will be called whenever the ChatHistoryMemoryProvider cannot find existing state in the session, + // typically the first time it is used with a new session. + session => new ChatHistoryMemoryProvider.State( + // Configure the scope values under which chat messages will be stored. + // In this case, we are using a fixed user ID and a unique session ID for each new session. + storageScope: new() { UserId = "UID1", SessionId = Guid.NewGuid().ToString() }, + // Configure the scope which would be used to search for relevant prior messages. + // In this case, we are searching for any messages for the user across all sessions. + searchScope: new() { UserId = "UID1" }))] }); // Start a new session for the agent conversation. diff --git a/dotnet/samples/GettingStarted/AgentWithMemory/AgentWithMemory_Step02_MemoryUsingMem0/Program.cs b/dotnet/samples/GettingStarted/AgentWithMemory/AgentWithMemory_Step02_MemoryUsingMem0/Program.cs index a81c496b5e..d1d6a60f9d 100644 --- a/dotnet/samples/GettingStarted/AgentWithMemory/AgentWithMemory_Step02_MemoryUsingMem0/Program.cs +++ b/dotnet/samples/GettingStarted/AgentWithMemory/AgentWithMemory_Step02_MemoryUsingMem0/Program.cs @@ -34,20 +34,21 @@ AIAgent agent = new AzureOpenAIClient( .AsAIAgent(new ChatClientAgentOptions() { ChatOptions = new() { Instructions = "You are a friendly travel assistant. Use known memories about the user when responding, and do not invent details." }, - AIContextProviderFactory = (ctx, ct) => new ValueTask(ctx.SerializedState.ValueKind is not JsonValueKind.Null and not JsonValueKind.Undefined - // If each session should have its own Mem0 scope, you can create a new id per session here: - // ? new Mem0Provider(mem0HttpClient, new Mem0ProviderScope() { ThreadId = Guid.NewGuid().ToString() }) - // In this case we are storing memories scoped by application and user instead so that memories are retained across threads. - ? new Mem0Provider(mem0HttpClient, new Mem0ProviderScope() { ApplicationId = "getting-started-agents", UserId = "sample-user" }) - // For cases where we are restoring from serialized state: - : new Mem0Provider(mem0HttpClient, ctx.SerializedState, ctx.JsonSerializerOptions)) + // The stateInitializer can be used to customize the Mem0 scope per session and it will be called each time a session + // is encountered by the Mem0Provider that does not already have Mem0Provider state stored on the session. + // If each session should have its own Mem0 scope, you can create a new id per session via the stateInitializer, e.g.: + // new Mem0Provider(mem0HttpClient, stateInitializer: _ => new(new Mem0ProviderScope() { ThreadId = Guid.NewGuid().ToString() })) + // In our case we are storing memories scoped by application and user instead so that memories are retained across threads. + AIContextProviders = [new Mem0Provider(mem0HttpClient, stateInitializer: _ => new(new Mem0ProviderScope() { ApplicationId = "getting-started-agents", UserId = "sample-user" }))] }); AgentSession session = await agent.CreateSessionAsync(); // Clear any existing memories for this scope to demonstrate fresh behavior. -Mem0Provider mem0Provider = session.GetService()!; -await mem0Provider.ClearStoredMemoriesAsync(); +// Note that the ClearStoredMemoriesAsync method will clear memories +// using the scope stored in the session, or provided via the stateInitializer. +Mem0Provider mem0Provider = agent.GetService()!; +await mem0Provider.ClearStoredMemoriesAsync(session); Console.WriteLine(await agent.RunAsync("Hi there! My name is Taylor and I'm planning a hiking trip to Patagonia in November.", session)); Console.WriteLine(await agent.RunAsync("I'm travelling with my sister and we love finding scenic viewpoints.", session)); diff --git a/dotnet/samples/GettingStarted/AgentWithMemory/AgentWithMemory_Step03_CustomMemory/Program.cs b/dotnet/samples/GettingStarted/AgentWithMemory/AgentWithMemory_Step03_CustomMemory/Program.cs index 4a736674fc..bac72d9b31 100644 --- a/dotnet/samples/GettingStarted/AgentWithMemory/AgentWithMemory_Step03_CustomMemory/Program.cs +++ b/dotnet/samples/GettingStarted/AgentWithMemory/AgentWithMemory_Step03_CustomMemory/Program.cs @@ -36,7 +36,7 @@ ChatClient chatClient = new AzureOpenAIClient( AIAgent agent = chatClient.AsAIAgent(new ChatClientAgentOptions() { ChatOptions = new() { Instructions = "You are a friendly assistant. Always address the user by their name." }, - AIContextProviderFactory = (ctx, ct) => new ValueTask(new UserInfoMemory(chatClient.AsIChatClient(), ctx.SerializedState, ctx.JsonSerializerOptions)) + AIContextProviders = [new UserInfoMemory(chatClient.AsIChatClient())] }); // Create a new session for the conversation. @@ -58,10 +58,10 @@ Console.WriteLine("\n>> Use deserialized session with previously created memorie var deserializedSession = await agent.DeserializeSessionAsync(sesionElement); Console.WriteLine(await agent.RunAsync("What is my name and age?", deserializedSession)); -Console.WriteLine("\n>> Read memories from memory component\n"); +Console.WriteLine("\n>> Read memories using memory component\n"); -// It's possible to access the memory component via the session's GetService method. -var userInfo = deserializedSession.GetService()?.UserInfo; +// It's possible to access the memory component via the agent's GetService method. +var userInfo = agent.GetService()?.GetUserInfo(deserializedSession); // Output the user info that was captured by the memory component. Console.WriteLine($"MEMORY - User Name: {userInfo?.UserName}"); @@ -69,12 +69,12 @@ Console.WriteLine($"MEMORY - User Age: {userInfo?.UserAge}"); Console.WriteLine("\n>> Use new session with previously created memories\n"); -// It is also possible to set the memories in a memory component on an individual session. +// It is also possible to set the memories using a memory component on an individual session. // This is useful if we want to start a new session, but have it share the same memories as a previous session. var newSession = await agent.CreateSessionAsync(); -if (userInfo is not null && newSession.GetService() is UserInfoMemory newSessionMemory) +if (userInfo is not null && agent.GetService() is UserInfoMemory newSessionMemory) { - newSessionMemory.UserInfo = userInfo; + newSessionMemory.SetUserInfo(newSession, userInfo); } // Invoke the agent and output the text result. @@ -89,28 +89,27 @@ namespace SampleApp internal sealed class UserInfoMemory : AIContextProvider { private readonly IChatClient _chatClient; + private readonly Func _stateInitializer; - public UserInfoMemory(IChatClient chatClient, UserInfo? userInfo = null) + public UserInfoMemory(IChatClient chatClient, Func? stateInitializer = null) { this._chatClient = chatClient; - this.UserInfo = userInfo ?? new UserInfo(); + this._stateInitializer = stateInitializer ?? (_ => new UserInfo()); } - public UserInfoMemory(IChatClient chatClient, JsonElement serializedState, JsonSerializerOptions? jsonSerializerOptions = null) - { - this._chatClient = chatClient; + public UserInfo GetUserInfo(AgentSession session) + => session.StateBag.GetValue(nameof(UserInfoMemory)) ?? new UserInfo(); - this.UserInfo = serializedState.ValueKind == JsonValueKind.Object ? - serializedState.Deserialize(jsonSerializerOptions)! : - new UserInfo(); - } - - public UserInfo UserInfo { get; set; } + public void SetUserInfo(AgentSession session, UserInfo userInfo) + => session.StateBag.SetValue(nameof(UserInfoMemory), userInfo); protected override async ValueTask InvokedCoreAsync(InvokedContext context, CancellationToken cancellationToken = default) { + var userInfo = context.Session?.StateBag.GetValue(nameof(UserInfoMemory)) + ?? this._stateInitializer.Invoke(context.Session); + // Try and extract the user name and age from the message if we don't have it already and it's a user message. - if ((this.UserInfo.UserName is null || this.UserInfo.UserAge is null) && context.RequestMessages.Any(x => x.Role == ChatRole.User)) + if ((userInfo.UserName is null || userInfo.UserAge is null) && context.RequestMessages.Any(x => x.Role == ChatRole.User)) { var result = await this._chatClient.GetResponseAsync( context.RequestMessages, @@ -120,36 +119,43 @@ namespace SampleApp }, cancellationToken: cancellationToken); - this.UserInfo.UserName ??= result.Result.UserName; - this.UserInfo.UserAge ??= result.Result.UserAge; + userInfo.UserName ??= result.Result.UserName; + userInfo.UserAge ??= result.Result.UserAge; } + + context.Session?.StateBag.SetValue(nameof(UserInfoMemory), userInfo); } protected override ValueTask InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default) { + var inputContext = context.AIContext; + var userInfo = context.Session?.StateBag.GetValue(nameof(UserInfoMemory)) + ?? this._stateInitializer.Invoke(context.Session); + StringBuilder instructions = new(); + if (!string.IsNullOrEmpty(inputContext.Instructions)) + { + instructions.AppendLine(inputContext.Instructions); + } // If we don't already know the user's name and age, add instructions to ask for them, otherwise just provide what we have to the context. instructions .AppendLine( - this.UserInfo.UserName is null ? + userInfo.UserName is null ? "Ask the user for their name and politely decline to answer any questions until they provide it." : - $"The user's name is {this.UserInfo.UserName}.") + $"The user's name is {userInfo.UserName}.") .AppendLine( - this.UserInfo.UserAge is null ? + userInfo.UserAge is null ? "Ask the user for their age and politely decline to answer any questions until they provide it." : - $"The user's age is {this.UserInfo.UserAge}."); + $"The user's age is {userInfo.UserAge}."); return new ValueTask(new AIContext { - Instructions = instructions.ToString() + Instructions = instructions.ToString(), + Messages = inputContext.Messages, + Tools = inputContext.Tools }); } - - public override JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null) - { - return JsonSerializer.SerializeToElement(this.UserInfo, jsonSerializerOptions); - } } internal sealed class UserInfo diff --git a/dotnet/samples/GettingStarted/AgentWithRAG/AgentWithRAG_Step01_BasicTextRAG/Program.cs b/dotnet/samples/GettingStarted/AgentWithRAG/AgentWithRAG_Step01_BasicTextRAG/Program.cs index 516585f7dc..c04601d940 100644 --- a/dotnet/samples/GettingStarted/AgentWithRAG/AgentWithRAG_Step01_BasicTextRAG/Program.cs +++ b/dotnet/samples/GettingStarted/AgentWithRAG/AgentWithRAG_Step01_BasicTextRAG/Program.cs @@ -65,12 +65,16 @@ AIAgent agent = azureOpenAIClient .AsAIAgent(new ChatClientAgentOptions { ChatOptions = new() { Instructions = "You are a helpful support specialist for Contoso Outdoors. Answer questions using the provided context and cite the source document when available." }, - AIContextProviderFactory = (ctx, ct) => new ValueTask(new TextSearchProvider(SearchAdapter, ctx.SerializedState, ctx.JsonSerializerOptions, textSearchOptions)), - // Since we are using ChatCompletion which stores chat history locally, we can also add a message removal policy + AIContextProviders = [new TextSearchProvider(SearchAdapter, textSearchOptions)], + // Since we are using ChatCompletion which stores chat history locally, we can also add a message filter // that removes messages produced by the TextSearchProvider before they are added to the chat history, so that // we don't bloat chat history with all the search result messages. - ChatHistoryProviderFactory = (ctx, ct) => new ValueTask(new InMemoryChatHistoryProvider(ctx.SerializedState, ctx.JsonSerializerOptions) - .WithAIContextProviderMessageRemoval()), + // By default the chat history provider will store all messages, except for those that came from chat history in the first place. + // We also want to maintain that exclusion here. + ChatHistoryProvider = new InMemoryChatHistoryProvider(new InMemoryChatHistoryProviderOptions + { + StorageInputMessageFilter = messages => messages.Where(m => m.GetAgentRequestMessageSourceType() != AgentRequestMessageSourceType.AIContextProvider && m.GetAgentRequestMessageSourceType() != AgentRequestMessageSourceType.ChatHistory) + }), }); AgentSession session = await agent.CreateSessionAsync(); diff --git a/dotnet/samples/GettingStarted/AgentWithRAG/AgentWithRAG_Step02_CustomVectorStoreRAG/Program.cs b/dotnet/samples/GettingStarted/AgentWithRAG/AgentWithRAG_Step02_CustomVectorStoreRAG/Program.cs index 4120f2d604..f06fda6a5b 100644 --- a/dotnet/samples/GettingStarted/AgentWithRAG/AgentWithRAG_Step02_CustomVectorStoreRAG/Program.cs +++ b/dotnet/samples/GettingStarted/AgentWithRAG/AgentWithRAG_Step02_CustomVectorStoreRAG/Program.cs @@ -74,7 +74,7 @@ AIAgent agent = azureOpenAIClient .AsAIAgent(new ChatClientAgentOptions { ChatOptions = new() { Instructions = "You are a helpful support specialist for the Microsoft Agent Framework. Answer questions using the provided context and cite the source document when available. Keep responses brief." }, - AIContextProviderFactory = (ctx, ct) => new ValueTask(new TextSearchProvider(SearchAdapter, ctx.SerializedState, ctx.JsonSerializerOptions, textSearchOptions)) + AIContextProviders = [new TextSearchProvider(SearchAdapter, textSearchOptions)] }); AgentSession session = await agent.CreateSessionAsync(); diff --git a/dotnet/samples/GettingStarted/AgentWithRAG/AgentWithRAG_Step03_CustomRAGDataSource/Program.cs b/dotnet/samples/GettingStarted/AgentWithRAG/AgentWithRAG_Step03_CustomRAGDataSource/Program.cs index 06da840df4..d4e3a40756 100644 --- a/dotnet/samples/GettingStarted/AgentWithRAG/AgentWithRAG_Step03_CustomRAGDataSource/Program.cs +++ b/dotnet/samples/GettingStarted/AgentWithRAG/AgentWithRAG_Step03_CustomRAGDataSource/Program.cs @@ -32,7 +32,7 @@ AIAgent agent = new AzureOpenAIClient( .AsAIAgent(new ChatClientAgentOptions { ChatOptions = new() { Instructions = "You are a helpful support specialist for Contoso Outdoors. Answer questions using the provided context and cite the source document when available." }, - AIContextProviderFactory = (ctx, ct) => new ValueTask(new TextSearchProvider(MockSearchAsync, ctx.SerializedState, ctx.JsonSerializerOptions, textSearchOptions)) + AIContextProviders = [new TextSearchProvider(MockSearchAsync, textSearchOptions)] }); AgentSession session = await agent.CreateSessionAsync(); diff --git a/dotnet/samples/GettingStarted/Agents/Agent_Step07_3rdPartyChatHistoryStorage/Program.cs b/dotnet/samples/GettingStarted/Agents/Agent_Step07_3rdPartyChatHistoryStorage/Program.cs index 0eaf3d8bc5..cc4ca1a6ed 100644 --- a/dotnet/samples/GettingStarted/Agents/Agent_Step07_3rdPartyChatHistoryStorage/Program.cs +++ b/dotnet/samples/GettingStarted/Agents/Agent_Step07_3rdPartyChatHistoryStorage/Program.cs @@ -3,7 +3,7 @@ #pragma warning disable CA1869 // Cache and reuse 'JsonSerializerOptions' instances // This sample shows how to create and use a simple AI agent with custom ChatHistoryProvider that stores chat history in a custom storage location. -// The state of the custom ChatHistoryProvider (SessionDbKey) is stored with the agent session, so that when the session is resumed later, +// The state of the custom ChatHistoryProvider (SessionDbKey) is stored in the AgentSession's StateBag, so that when the session is resumed later, // the chat history can be retrieved from the custom storage location. using System.Text.Json; @@ -36,11 +36,8 @@ AIAgent agent = new AzureOpenAIClient( { ChatOptions = new() { Instructions = "You are good at telling jokes." }, Name = "Joker", - ChatHistoryProviderFactory = (ctx, ct) => new ValueTask( - // Create a new ChatHistoryProvider for this agent that stores chat history in a vector store. - // Each session must get its own copy of the VectorChatHistoryProvider, since the provider - // also contains the id that the chat history is stored under. - new VectorChatHistoryProvider(vectorStore, ctx.SerializedState, ctx.JsonSerializerOptions)) + // Create a new ChatHistoryProvider for this agent that stores chat history in a vector store. + ChatHistoryProvider = new VectorChatHistoryProvider(vectorStore) }); // Start a new session for the agent conversation. @@ -66,48 +63,75 @@ AgentSession resumedSession = await agent.DeserializeSessionAsync(serializedSess // Run the agent with the session that stores chat history in the vector store a second time. Console.WriteLine(await agent.RunAsync("Now tell the same joke in the voice of a pirate, and add some emojis to the joke.", resumedSession)); -// We can access the VectorChatHistoryProvider via the session's GetService method if we need to read the key under which chat history is stored. -var chatHistoryProvider = resumedSession.GetService()!; -Console.WriteLine($"\nSession is stored in vector store under key: {chatHistoryProvider.SessionDbKey}"); +// We can access the VectorChatHistoryProvider via the agent's GetService method +// if we need to read the key under which chat history is stored. The key is stored +// in the session state, and therefore we need to provide the session when reading it. +var chatHistoryProvider = agent.GetService()!; +Console.WriteLine($"\nSession is stored in vector store under key: {chatHistoryProvider.GetSessionDbKey(resumedSession)}"); namespace SampleApp { /// /// A sample implementation of that stores chat history in a vector store. + /// State (the session DB key) is stored in the so it roundtrips + /// automatically with session serialization. /// internal sealed class VectorChatHistoryProvider : ChatHistoryProvider { private readonly VectorStore _vectorStore; + private readonly Func _stateInitializer; + private readonly string _stateKey; - public VectorChatHistoryProvider(VectorStore vectorStore, JsonElement serializedState, JsonSerializerOptions? jsonSerializerOptions = null) + /// + public override string StateKey => this._stateKey; + + public VectorChatHistoryProvider( + VectorStore vectorStore, + Func? stateInitializer = null, + string? stateKey = null) { this._vectorStore = vectorStore ?? throw new ArgumentNullException(nameof(vectorStore)); - - if (serializedState.ValueKind is JsonValueKind.String) - { - // Here we can deserialize the session id so that we can access the same messages as before the suspension. - this.SessionDbKey = serializedState.Deserialize(); - } + this._stateInitializer = stateInitializer ?? (_ => new State(Guid.NewGuid().ToString("N"))); + this._stateKey = stateKey ?? base.StateKey; } - public string? SessionDbKey { get; private set; } + public string GetSessionDbKey(AgentSession session) + => this.GetOrInitializeState(session).SessionDbKey; + + private State GetOrInitializeState(AgentSession? session) + { + if (session?.StateBag.TryGetValue(this._stateKey, out var state) is true && state is not null) + { + return state; + } + + state = this._stateInitializer(session); + if (session is not null) + { + session.StateBag.SetValue(this._stateKey, state); + } + + return state; + } protected override async ValueTask> InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default) { + var state = this.GetOrInitializeState(context.Session); var collection = this._vectorStore.GetCollection("ChatHistory"); await collection.EnsureCollectionExistsAsync(cancellationToken); var records = await collection .GetAsync( - x => x.SessionId == this.SessionDbKey, 10, + x => x.SessionId == state.SessionDbKey, 10, new() { OrderBy = x => x.Descending(y => y.Timestamp) }, cancellationToken) .ToListAsync(cancellationToken); - var messages = records.ConvertAll(x => JsonSerializer.Deserialize(x.SerializedMessage!)!) -; + var messages = records.ConvertAll(x => JsonSerializer.Deserialize(x.SerializedMessage!)!); messages.Reverse(); - return messages; + return messages + .Select(message => message.WithAgentRequestMessageSource(AgentRequestMessageSourceType.ChatHistory, this.GetType().FullName!)) + .Concat(context.RequestMessages); } protected override async ValueTask InvokedCoreAsync(InvokedContext context, CancellationToken cancellationToken = default) @@ -118,28 +142,39 @@ namespace SampleApp return; } - this.SessionDbKey ??= Guid.NewGuid().ToString("N"); + var state = this.GetOrInitializeState(context.Session); var collection = this._vectorStore.GetCollection("ChatHistory"); await collection.EnsureCollectionExistsAsync(cancellationToken); - // Add both request and response messages to the store + // Add both request and response messages to the store, excluding messages that came from chat history. // Optionally messages produced by the AIContextProvider can also be persisted (not shown). - var allNewMessages = context.RequestMessages.Concat(context.ResponseMessages ?? []); + var allNewMessages = context.RequestMessages + .Where(m => m.GetAgentRequestMessageSourceType() != AgentRequestMessageSourceType.ChatHistory) + .Concat(context.ResponseMessages ?? []); await collection.UpsertAsync(allNewMessages.Select(x => new ChatHistoryItem() { - Key = this.SessionDbKey + x.MessageId, + Key = state.SessionDbKey + x.MessageId, Timestamp = DateTimeOffset.UtcNow, - SessionId = this.SessionDbKey, + SessionId = state.SessionDbKey, SerializedMessage = JsonSerializer.Serialize(x), MessageText = x.Text }), cancellationToken); } - public override JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null) => - // We have to serialize the session id, so that on deserialization we can retrieve the messages using the same session id. - JsonSerializer.SerializeToElement(this.SessionDbKey); + /// + /// Represents the per-session state stored in the . + /// + public sealed class State + { + public State(string sessionDbKey) + { + this.SessionDbKey = sessionDbKey ?? throw new ArgumentNullException(nameof(sessionDbKey)); + } + + public string SessionDbKey { get; } + } /// /// The data structure used to store chat history items in the vector store. diff --git a/dotnet/samples/GettingStarted/Agents/Agent_Step16_ChatReduction/Program.cs b/dotnet/samples/GettingStarted/Agents/Agent_Step16_ChatReduction/Program.cs index 77abb7898a..875f3375a4 100644 --- a/dotnet/samples/GettingStarted/Agents/Agent_Step16_ChatReduction/Program.cs +++ b/dotnet/samples/GettingStarted/Agents/Agent_Step16_ChatReduction/Program.cs @@ -27,7 +27,7 @@ AIAgent agent = new AzureOpenAIClient( { ChatOptions = new() { Instructions = "You are good at telling jokes." }, Name = "Joker", - ChatHistoryProviderFactory = (ctx, ct) => new ValueTask(new InMemoryChatHistoryProvider(new MessageCountingChatReducer(2), ctx.SerializedState, ctx.JsonSerializerOptions)) + ChatHistoryProvider = new InMemoryChatHistoryProvider(new() { ChatReducer = new MessageCountingChatReducer(2) }) }); AgentSession session = await agent.CreateSessionAsync(); @@ -36,7 +36,10 @@ AgentSession session = await agent.CreateSessionAsync(); Console.WriteLine(await agent.RunAsync("Tell me a joke about a pirate.", session)); // Get the chat history to see how many messages are stored. -IList? chatHistory = session.GetService>(); +// We can use the ChatHistoryProvider, that is also used by the agent, to read the +// chat history from the session state, and see how the reducer is affecting the stored messages. +var provider = agent.GetService(); +List? chatHistory = provider?.GetMessages(session); Console.WriteLine($"\nChat history has {chatHistory?.Count} messages.\n"); // Invoke the agent a few more times. diff --git a/dotnet/samples/GettingStarted/Agents/Agent_Step20_AdditionalAIContext/Program.cs b/dotnet/samples/GettingStarted/Agents/Agent_Step20_AdditionalAIContext/Program.cs index b04cf836bb..ba56d94e93 100644 --- a/dotnet/samples/GettingStarted/Agents/Agent_Step20_AdditionalAIContext/Program.cs +++ b/dotnet/samples/GettingStarted/Agents/Agent_Step20_AdditionalAIContext/Program.cs @@ -1,13 +1,12 @@ // Copyright (c) Microsoft. All rights reserved. -// This sample shows how to inject additional AI context into a ChatClientAgent using a custom AIContextProvider component that is attached to the agent. -// The sample also shows how to combine the results from multiple providers into a single class, in order to attach multiple of these to an agent. +// This sample shows how to inject additional AI context into a ChatClientAgent using custom AIContextProvider components that are attached to the agent. +// Multiple providers can be attached to an agent, and they will be called in sequence, each receiving the accumulated context from the previous one. // This mechanism can be used for various purposes, such as injecting RAG search results or memories into the agent's context. // Also note that Agent Framework already provides built-in AIContextProviders for many of these scenarios. #pragma warning disable CA1869 // Cache and reuse 'JsonSerializerOptions' instances -using System.ComponentModel; using System.Text; using System.Text.Json; using Azure.AI.OpenAI; @@ -48,16 +47,20 @@ AIAgent agent = new AzureOpenAIClient( You manage a TODO list for the user. When the user has completed one of the tasks it can be removed from the TODO list. Only provide the list of TODO items if asked. You remind users of upcoming calendar events when the user interacts with you. """ }, - ChatHistoryProviderFactory = (ctx, ct) => new ValueTask(new InMemoryChatHistoryProvider() - // Use WithAIContextProviderMessageRemoval, so that we don't store the messages from the AI context provider in the chat history. + ChatHistoryProvider = new InMemoryChatHistoryProvider(new InMemoryChatHistoryProviderOptions + { + // Use StorageInputMessageFilter to provide a custom filter for messages stored in chat history. + // By default the chat history provider will store all messages, except for those that came from chat history in the first place. + // In this case, we want to also exclude messages that came from AI context providers. // You may want to store these messages, depending on their content and your requirements. - .WithAIContextProviderMessageRemoval()), - // Add an AI context provider that maintains a todo list for the agent and one that provides upcoming calendar entries. - // Wrap these in an AI context provider that aggregates the other two. - AIContextProviderFactory = (ctx, ct) => new ValueTask(new AggregatingAIContextProvider([ - AggregatingAIContextProvider.CreateFactory((jsonElement, jsonSerializerOptions) => new TodoListAIContextProvider(jsonElement, jsonSerializerOptions)), - AggregatingAIContextProvider.CreateFactory((_, _) => new CalendarSearchAIContextProvider(loadNextThreeCalendarEvents)) - ], ctx.SerializedState, ctx.JsonSerializerOptions)), + StorageInputMessageFilter = messages => messages.Where(m => m.GetAgentRequestMessageSourceType() != AgentRequestMessageSourceType.AIContextProvider && m.GetAgentRequestMessageSourceType() != AgentRequestMessageSourceType.ChatHistory) + }), + // Add multiple AI context providers: one that maintains a todo list and one that provides upcoming calendar entries. + // The agent will call each provider in sequence, accumulating context from each. + AIContextProviders = [ + new TodoListAIContextProvider(), + new CalendarSearchAIContextProvider(loadNextThreeCalendarEvents) + ], }); // Invoke the agent and output the text result. @@ -83,51 +86,67 @@ namespace SampleApp /// internal sealed class TodoListAIContextProvider : AIContextProvider { - private readonly List _todoItems = new(); + private static List GetTodoItems(AgentSession? session) + => session?.StateBag.GetValue>(nameof(TodoListAIContextProvider)) ?? new List(); - public TodoListAIContextProvider(JsonElement jsonElement, JsonSerializerOptions? jsonSerializerOptions = null) - { - // Only try and restore the state if we got an array, since any other json would be invalid or undefined/null meaning - // it's the first time we are running. - if (jsonElement.ValueKind == JsonValueKind.Array) - { - this._todoItems = JsonSerializer.Deserialize>(jsonElement.GetRawText(), jsonSerializerOptions) ?? new List(); - } - } + private static void SetTodoItems(AgentSession? session, List items) + => session?.StateBag.SetValue(nameof(TodoListAIContextProvider), items); protected override ValueTask InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default) { + var inputContext = context.AIContext; + var todoItems = GetTodoItems(context.Session); + StringBuilder outputMessageBuilder = new(); outputMessageBuilder.AppendLine("Your todo list contains the following items:"); - if (this._todoItems.Count == 0) + if (todoItems.Count == 0) { outputMessageBuilder.AppendLine(" (no items)"); } else { - for (int i = 0; i < this._todoItems.Count; i++) + for (int i = 0; i < todoItems.Count; i++) { - outputMessageBuilder.AppendLine($"{i}. {this._todoItems[i]}"); + outputMessageBuilder.AppendLine($"{i}. {todoItems[i]}"); } } return new ValueTask(new AIContext { - Tools = [AIFunctionFactory.Create(this.AddTodoItem), AIFunctionFactory.Create(this.RemoveTodoItem)], - Messages = [new MEAI.ChatMessage(ChatRole.User, outputMessageBuilder.ToString())] + Instructions = inputContext.Instructions, + Tools = (inputContext.Tools ?? []).Concat(new AITool[] + { + AIFunctionFactory.Create((string item) => AddTodoItem(context.Session, item), "AddTodoItem", "Adds an item to the todo list."), + AIFunctionFactory.Create((int index) => RemoveTodoItem(context.Session, index), "RemoveTodoItem", "Removes an item from the todo list. Index is zero based.") + }), + Messages = + (inputContext.Messages ?? []) + .Concat( + [ + new MEAI.ChatMessage(ChatRole.User, outputMessageBuilder.ToString()).WithAgentRequestMessageSource(AgentRequestMessageSourceType.AIContextProvider, this.GetType().FullName!) + ]) }); } - [Description("Adds an item to the todo list. Index is zero based.")] - private void RemoveTodoItem(int index) => - this._todoItems.RemoveAt(index); + private static void RemoveTodoItem(AgentSession? session, int index) + { + var items = GetTodoItems(session); + items.RemoveAt(index); + SetTodoItems(session, items); + } - private void AddTodoItem(string item) => - this._todoItems.Add(string.IsNullOrWhiteSpace(item) ? throw new ArgumentException("Item must have a value") : item); + private static void AddTodoItem(AgentSession? session, string item) + { + if (string.IsNullOrWhiteSpace(item)) + { + throw new ArgumentException("Item must have a value"); + } - public override JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null) => - JsonSerializer.SerializeToElement(this._todoItems, jsonSerializerOptions); + var items = GetTodoItems(session); + items.Add(item); + SetTodoItems(session, items); + } } /// @@ -137,6 +156,7 @@ namespace SampleApp { protected override async ValueTask InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default) { + var inputContext = context.AIContext; var events = await loadNextThreeCalendarEvents(); StringBuilder outputMessageBuilder = new(); @@ -148,84 +168,16 @@ namespace SampleApp return new() { + Instructions = inputContext.Instructions, Messages = - [ - new MEAI.ChatMessage(ChatRole.User, outputMessageBuilder.ToString()), - ] + (inputContext.Messages ?? []) + .Concat( + [ + new MEAI.ChatMessage(ChatRole.User, outputMessageBuilder.ToString()).WithAgentRequestMessageSource(AgentRequestMessageSourceType.AIContextProvider, this.GetType().FullName!) + ]) + .ToList(), + Tools = inputContext.Tools }; } } - - /// - /// An which aggregates multiple AI context providers into one. - /// Serialized state for the different providers are stored under their type name. - /// Tools and messages from all providers are combined, and instructions are concatenated. - /// - internal sealed class AggregatingAIContextProvider : AIContextProvider - { - private readonly List _providers = new(); - - public AggregatingAIContextProvider(ProviderFactory[] providerFactories, JsonElement jsonElement, JsonSerializerOptions? jsonSerializerOptions) - { - // We received a json object, so let's check if it has some previously serialized state that we can use. - if (jsonElement.ValueKind == JsonValueKind.Object) - { - this._providers = providerFactories - .Select(factory => factory.FactoryMethod(jsonElement.TryGetProperty(factory.ProviderType.Name, out var prop) ? prop : default, jsonSerializerOptions)) - .ToList(); - return; - } - - // We didn't receive any valid json, so we can just construct fresh providers. - this._providers = providerFactories - .Select(factory => factory.FactoryMethod(default, jsonSerializerOptions)) - .ToList(); - } - - protected override async ValueTask InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default) - { - // Invoke all the sub providers. - var tasks = this._providers.Select(provider => provider.InvokingAsync(context, cancellationToken).AsTask()); - var results = await Task.WhenAll(tasks); - - // Combine the results from each sub provider. - return new AIContext - { - Tools = results.SelectMany(r => r.Tools ?? []).ToList(), - Messages = results.SelectMany(r => r.Messages ?? []).ToList(), - Instructions = string.Join("\n", results.Select(r => r.Instructions).Where(s => !string.IsNullOrEmpty(s))) - }; - } - - public override JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null) - { - Dictionary elements = new(); - foreach (var provider in this._providers) - { - JsonElement element = provider.Serialize(jsonSerializerOptions); - - // Don't try to store state for any providers that aren't producing any. - if (element.ValueKind != JsonValueKind.Undefined && element.ValueKind != JsonValueKind.Null) - { - elements[provider.GetType().Name] = element; - } - } - - return JsonSerializer.SerializeToElement(elements, jsonSerializerOptions); - } - - public static ProviderFactory CreateFactory(Func factoryMethod) - where TProviderType : AIContextProvider => new() - { - FactoryMethod = (jsonElement, jsonSerializerOptions) => factoryMethod(jsonElement, jsonSerializerOptions), - ProviderType = typeof(TProviderType) - }; - - public readonly struct ProviderFactory - { - public Func FactoryMethod { get; init; } - - public Type ProviderType { get; init; } - } - } } diff --git a/dotnet/src/Microsoft.Agents.AI.A2A/A2AAgent.cs b/dotnet/src/Microsoft.Agents.AI.A2A/A2AAgent.cs index aea99b4e3d..5601653e8f 100644 --- a/dotnet/src/Microsoft.Agents.AI.A2A/A2AAgent.cs +++ b/dotnet/src/Microsoft.Agents.AI.A2A/A2AAgent.cs @@ -80,7 +80,7 @@ public sealed class A2AAgent : AIAgent /// protected override ValueTask DeserializeSessionCoreAsync(JsonElement serializedState, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default) - => new(new A2AAgentSession(serializedState, jsonSerializerOptions)); + => new(A2AAgentSession.Deserialize(serializedState, jsonSerializerOptions)); /// protected override async Task RunCoreAsync(IEnumerable messages, AgentSession? session = null, AgentRunOptions? options = null, CancellationToken cancellationToken = default) diff --git a/dotnet/src/Microsoft.Agents.AI.A2A/A2AAgentSession.cs b/dotnet/src/Microsoft.Agents.AI.A2A/A2AAgentSession.cs index cac9b43a30..045abc736a 100644 --- a/dotnet/src/Microsoft.Agents.AI.A2A/A2AAgentSession.cs +++ b/dotnet/src/Microsoft.Agents.AI.A2A/A2AAgentSession.cs @@ -1,66 +1,61 @@ // Copyright (c) Microsoft. All rights reserved. using System; +using System.Diagnostics; using System.Text.Json; +using System.Text.Json.Serialization; namespace Microsoft.Agents.AI.A2A; /// /// Session for A2A based agents. /// +[DebuggerDisplay("{DebuggerDisplay,nq}")] public sealed class A2AAgentSession : AgentSession { internal A2AAgentSession() { } - internal A2AAgentSession(JsonElement serializedSessionState, JsonSerializerOptions? jsonSerializerOptions = null) + [JsonConstructor] + internal A2AAgentSession(string? contextId, string? taskId, AgentSessionStateBag? stateBag) : base(stateBag ?? new()) { - if (serializedSessionState.ValueKind != JsonValueKind.Object) - { - throw new ArgumentException("The serialized session state must be a JSON object.", nameof(serializedSessionState)); - } - - var state = serializedSessionState.Deserialize( - A2AJsonUtilities.DefaultOptions.GetTypeInfo(typeof(A2AAgentSessionState))) as A2AAgentSessionState; - - if (state?.ContextId is string contextId) - { - this.ContextId = contextId; - } - - if (state?.TaskId is string taskId) - { - this.TaskId = taskId; - } + this.ContextId = contextId; + this.TaskId = taskId; } /// /// Gets the ID for the current conversation with the A2A agent. /// + [JsonPropertyName("contextId")] public string? ContextId { get; internal set; } /// /// Gets the ID for the task the agent is currently working on. /// + [JsonPropertyName("taskId")] public string? TaskId { get; internal set; } /// internal JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null) { - var state = new A2AAgentSessionState - { - ContextId = this.ContextId, - TaskId = this.TaskId - }; - - return JsonSerializer.SerializeToElement(state, A2AJsonUtilities.DefaultOptions.GetTypeInfo(typeof(A2AAgentSessionState))); + var jso = jsonSerializerOptions ?? A2AJsonUtilities.DefaultOptions; + return JsonSerializer.SerializeToElement(this, jso.GetTypeInfo(typeof(A2AAgentSession))); } - internal sealed class A2AAgentSessionState + internal static A2AAgentSession Deserialize(JsonElement serializedState, JsonSerializerOptions? jsonSerializerOptions = null) { - public string? ContextId { get; set; } + if (serializedState.ValueKind != JsonValueKind.Object) + { + throw new ArgumentException("The serialized session state must be a JSON object.", nameof(serializedState)); + } - public string? TaskId { get; set; } + var jso = jsonSerializerOptions ?? A2AJsonUtilities.DefaultOptions; + return serializedState.Deserialize(jso.GetTypeInfo(typeof(A2AAgentSession))) as A2AAgentSession + ?? new A2AAgentSession(); } + + [DebuggerBrowsable(DebuggerBrowsableState.Never)] + private string DebuggerDisplay => + $"ContextId = {this.ContextId}, TaskId = {this.TaskId}, StateBag Count = {this.StateBag.Count}"; } diff --git a/dotnet/src/Microsoft.Agents.AI.A2A/A2AJsonUtilities.cs b/dotnet/src/Microsoft.Agents.AI.A2A/A2AJsonUtilities.cs index 5f079d9573..3c25e350ae 100644 --- a/dotnet/src/Microsoft.Agents.AI.A2A/A2AJsonUtilities.cs +++ b/dotnet/src/Microsoft.Agents.AI.A2A/A2AJsonUtilities.cs @@ -74,7 +74,7 @@ public static partial class A2AJsonUtilities NumberHandling = JsonNumberHandling.AllowReadingFromString)] // A2A agent types - [JsonSerializable(typeof(A2AAgentSession.A2AAgentSessionState))] + [JsonSerializable(typeof(A2AAgentSession))] [ExcludeFromCodeCoverage] private sealed partial class JsonContext : JsonSerializerContext; } diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions/AIContext.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions/AIContext.cs index b05992d93e..9ccfc3e905 100644 --- a/dotnet/src/Microsoft.Agents.AI.Abstractions/AIContext.cs +++ b/dotnet/src/Microsoft.Agents.AI.Abstractions/AIContext.cs @@ -56,41 +56,44 @@ public sealed class AIContext public string? Instructions { get; set; } /// - /// Gets or sets a collection of messages to add to the conversation history. + /// Gets or sets the sequence of messages to use for the current invocation. /// /// - /// A list of instances to be permanently added to the conversation history, - /// or if no messages should be added. + /// A sequence of instances to be used for the current invocation, + /// or if no messages should be used. /// /// /// - /// Unlike and , messages added through this property become - /// permanent additions to the conversation history. They will persist beyond the current invocation and - /// will be available in future interactions within the same conversation thread. + /// Unlike and , messages added through this property may become + /// permanent additions to the conversation history. + /// If chat history is managed by the underlying AI service, these messages will become part of chat history. + /// If chat history is managed using a , these messages will be passed to the + /// method, + /// and the provider can choose which of these messages to permanently add to the conversation history. /// /// /// This property is useful for: /// - /// Injecting relevant historical context or background information + /// Injecting relevant historical context e.g. memories + /// Injecting relevant background information e.g. via Retrieval Augmented Generation /// Adding system messages that provide ongoing context - /// Including retrieved information that should be part of the conversation record - /// Inserting contextual exchanges that inform the current conversation /// /// /// - public IList? Messages { get; set; } + public IEnumerable? Messages { get; set; } /// - /// Gets or sets a collection of tools or functions to make available to the AI model for the current invocation. + /// Gets or sets a sequence of tools or functions to make available to the AI model for the current invocation. /// /// - /// A list of instances that will be available to the AI model during the current invocation, + /// A sequence of instances that will be available to the AI model during the current invocation, /// or if no additional tools should be provided. /// /// /// - /// These tools are transient and apply only to the current AI model invocation. They are combined with any - /// tools already configured for the agent to provide an expanded set of capabilities for the specific interaction. + /// These tools are transient and apply only to the current AI model invocation. Any existing tools + /// are provided as input to the instances, so context providers can choose to modify or replace the existing tools + /// as needed based on the current context. The resulting set of tools is then passed to the underlying AI model, which may choose to utilize them when generating responses. /// /// /// Context-specific tools enable: @@ -102,5 +105,5 @@ public sealed class AIContext /// /// /// - public IList? Tools { get; set; } + public IEnumerable? Tools { get; set; } } diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions/AIContextProvider.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions/AIContextProvider.cs index 76ff8752b9..b201e9f8e1 100644 --- a/dotnet/src/Microsoft.Agents.AI.Abstractions/AIContextProvider.cs +++ b/dotnet/src/Microsoft.Agents.AI.Abstractions/AIContextProvider.cs @@ -2,8 +2,6 @@ using System; using System.Collections.Generic; -using System.Linq; -using System.Text.Json; using System.Threading; using System.Threading.Tasks; using Microsoft.Extensions.AI; @@ -32,24 +30,15 @@ namespace Microsoft.Agents.AI; /// public abstract class AIContextProvider { - private readonly string _sourceId; - /// - /// Initializes a new instance of the class. + /// Gets the key used to store the provider state in the . /// - protected AIContextProvider() - { - this._sourceId = this.GetType().FullName!; - } - - /// - /// Initializes a new instance of the class with the specified source id. - /// - /// The source id to stamp on for each messages produced by the . - protected AIContextProvider(string sourceId) - { - this._sourceId = sourceId; - } + /// + /// The default value is the name of the concrete type (e.g. "TextSearchProvider"). + /// Implementations may override this to provide a custom key, for example when multiple + /// instances of the same provider type are used in the same session. + /// + public virtual string StateKey => this.GetType().Name; /// /// Called at the start of agent invocation to provide additional context. @@ -68,20 +57,8 @@ public abstract class AIContextProvider /// /// /// - public async ValueTask InvokingAsync(InvokingContext context, CancellationToken cancellationToken = default) - { - var aiContext = await this.InvokingCoreAsync(context, cancellationToken).ConfigureAwait(false); - if (aiContext.Messages is null) - { - return aiContext; - } - - aiContext.Messages = aiContext.Messages - .Select(message => message.AsAgentRequestMessageSourcedMessage(AgentRequestMessageSourceType.AIContextProvider, this._sourceId)) - .ToList(); - - return aiContext; - } + public ValueTask InvokingAsync(InvokingContext context, CancellationToken cancellationToken = default) + => this.InvokingCoreAsync(context, cancellationToken); /// /// Called at the start of agent invocation to provide additional context. @@ -112,13 +89,18 @@ public abstract class AIContextProvider /// /// Implementers can use the request and response messages in the provided to: /// - /// Update internal state based on conversation outcomes + /// Update state based on conversation outcomes /// Extract and store memories or preferences from user messages /// Log or audit conversation details /// Perform cleanup or finalization tasks /// /// /// + /// The is passed a reference to the via and + /// allowing it to store state in the . Since an is used with many different sessions, it should + /// not store any session-specific information within its own instance fields. Instead, any session-specific state should be stored in the associated . + /// + /// /// This method is called regardless of whether the invocation succeeded or failed. /// To check if the invocation was successful, inspect the property. /// @@ -150,18 +132,6 @@ public abstract class AIContextProvider protected virtual ValueTask InvokedCoreAsync(InvokedContext context, CancellationToken cancellationToken = default) => default; - /// - /// Serializes the current object's state to a using the specified serialization options. - /// - /// The JSON serialization options to use for the serialization process. - /// A representation of the object's state, or a default if the provider has no serializable state. - /// - /// The default implementation returns a default . Override this method if the provider - /// maintains state that should be preserved across sessions or distributed scenarios. - /// - public virtual JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null) - => default; - /// Asks the for an object of the specified type . /// The type of object being requested. /// An optional key that can be used to help identify the target service. @@ -203,20 +173,20 @@ public abstract class AIContextProvider public sealed class InvokingContext { /// - /// Initializes a new instance of the class with the specified request messages. + /// Initializes a new instance of the class. /// /// The agent being invoked. /// The session associated with the agent invocation. - /// The messages to be used by the agent for this invocation. - /// is . + /// The AI context to be used by the agent for this invocation. + /// or is . public InvokingContext( AIAgent agent, AgentSession? session, - IEnumerable requestMessages) + AIContext aiContext) { this.Agent = Throw.IfNull(agent); this.Session = session; - this.RequestMessages = Throw.IfNull(requestMessages); + this.AIContext = Throw.IfNull(aiContext); } /// @@ -230,39 +200,75 @@ public abstract class AIContextProvider public AgentSession? Session { get; } /// - /// Gets the caller provided messages that will be used by the agent for this invocation. + /// Gets the being built for the current invocation. Context providers can modify + /// and return or return a new instance to provide additional context for the invocation. /// - /// - /// A collection of instances representing new messages that were provided by the caller. - /// - public IEnumerable RequestMessages { get; set { field = Throw.IfNull(value); } } + /// + /// + /// If multiple instances are used in the same invocation, each + /// will receive the context returned by the previous allowing them to build on top of each other's context. + /// + /// + /// The first in the invocation pipeline will receive an instance + /// that already contains the caller provided messages that will be used by the agent for this invocation. + /// + /// + /// It may also contain messages from chat history, if a is being used. + /// + /// + public AIContext AIContext { get; } } /// /// Contains the context information provided to . /// /// - /// This class provides context about a completed agent invocation, including both the - /// request messages that were used and the response messages that were generated. It also indicates - /// whether the invocation succeeded or failed. + /// This class provides context about a completed agent invocation, including the accumulated + /// request messages (user input, chat history and any others provided by AI context providers) that were used + /// and the response messages that were generated. It also indicates whether the invocation succeeded or failed. /// public sealed class InvokedContext { /// - /// Initializes a new instance of the class with the specified request messages. + /// Initializes a new instance of the class for a successful invocation. /// - /// The agent being invoked. + /// The agent that was invoked. /// The session associated with the agent invocation. - /// The caller provided messages that were used by the agent for this invocation. - /// is . + /// The accumulated request messages (user input, chat history and any others provided by AI context providers) + /// that were used by the agent for this invocation. + /// The response messages generated during this invocation. + /// , , or is . public InvokedContext( AIAgent agent, AgentSession? session, - IEnumerable requestMessages) + IEnumerable requestMessages, + IEnumerable responseMessages) { this.Agent = Throw.IfNull(agent); this.Session = session; this.RequestMessages = Throw.IfNull(requestMessages); + this.ResponseMessages = Throw.IfNull(responseMessages); + } + + /// + /// Initializes a new instance of the class for a failed invocation. + /// + /// The agent that was invoked. + /// The session associated with the agent invocation. + /// The accumulated request messages (user input, chat history and any others provided by AI context providers) + /// that were used by the agent for this invocation. + /// The exception that caused the invocation to fail. + /// , , or is . + public InvokedContext( + AIAgent agent, + AgentSession? session, + IEnumerable requestMessages, + Exception invokeException) + { + this.Agent = Throw.IfNull(agent); + this.Session = session; + this.RequestMessages = Throw.IfNull(requestMessages); + this.InvokeException = Throw.IfNull(invokeException); } /// @@ -276,22 +282,22 @@ public abstract class AIContextProvider public AgentSession? Session { get; } /// - /// Gets the caller provided messages that were used by the agent for this invocation. + /// Gets the accumulated request messages (user input, chat history and any others provided by AI context providers) + /// that were used by the agent for this invocation. /// /// - /// A collection of instances representing new messages that were provided by the caller. - /// This does not include any supplied messages. + /// A collection of instances representing all messages that were used by the agent for this invocation. /// - public IEnumerable RequestMessages { get; set { field = Throw.IfNull(value); } } + public IEnumerable RequestMessages { get; } /// /// Gets the collection of response messages generated during this invocation if the invocation succeeded. /// /// /// A collection of instances representing the response, - /// or if the invocation failed or did not produce response messages. + /// or if the invocation failed. /// - public IEnumerable? ResponseMessages { get; set; } + public IEnumerable? ResponseMessages { get; } /// /// Gets the that was thrown during the invocation, if the invocation failed. @@ -299,6 +305,6 @@ public abstract class AIContextProvider /// /// The exception that caused the invocation to fail, or if the invocation succeeded. /// - public Exception? InvokeException { get; set; } + public Exception? InvokeException { get; } } } diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions/AgentAbstractionsJsonUtilities.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions/AgentAbstractionsJsonUtilities.cs index 17fbb9e4c6..f8c8aa9b98 100644 --- a/dotnet/src/Microsoft.Agents.AI.Abstractions/AgentAbstractionsJsonUtilities.cs +++ b/dotnet/src/Microsoft.Agents.AI.Abstractions/AgentAbstractionsJsonUtilities.cs @@ -1,5 +1,6 @@ // Copyright (c) Microsoft. All rights reserved. +using System.Collections.Concurrent; using System.Diagnostics.CodeAnalysis; using System.Text.Encodings.Web; using System.Text.Json; @@ -80,9 +81,9 @@ public static partial class AgentAbstractionsJsonUtilities [JsonSerializable(typeof(AgentResponse[]))] [JsonSerializable(typeof(AgentResponseUpdate))] [JsonSerializable(typeof(AgentResponseUpdate[]))] - [JsonSerializable(typeof(ServiceIdAgentSession.ServiceIdAgentSessionState))] - [JsonSerializable(typeof(InMemoryAgentSession.InMemoryAgentSessionState))] [JsonSerializable(typeof(InMemoryChatHistoryProvider.State))] + [JsonSerializable(typeof(AgentSessionStateBag))] + [JsonSerializable(typeof(ConcurrentDictionary))] [ExcludeFromCodeCoverage] private sealed partial class JsonContext : JsonSerializerContext; diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions/AgentRequestMessageSourceAttribution.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions/AgentRequestMessageSourceAttribution.cs index 2c606814ce..1515adec9a 100644 --- a/dotnet/src/Microsoft.Agents.AI.Abstractions/AgentRequestMessageSourceAttribution.cs +++ b/dotnet/src/Microsoft.Agents.AI.Abstractions/AgentRequestMessageSourceAttribution.cs @@ -63,6 +63,17 @@ public readonly struct AgentRequestMessageSourceAttribution : IEquatable + /// Returns a string representation of the current instance. + /// + /// A string containing the source type and source identifier. + public override string ToString() + { + return this.SourceId is null + ? $"{this.SourceType}" + : $"{this.SourceType}:{this.SourceId}"; + } + /// /// Returns a hash code for the current instance. /// diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions/AgentRequestMessageSourceType.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions/AgentRequestMessageSourceType.cs index 14bbbe3388..744f87bed6 100644 --- a/dotnet/src/Microsoft.Agents.AI.Abstractions/AgentRequestMessageSourceType.cs +++ b/dotnet/src/Microsoft.Agents.AI.Abstractions/AgentRequestMessageSourceType.cs @@ -58,6 +58,12 @@ public readonly struct AgentRequestMessageSourceType : IEquatable if is a and its value is the same as this instance; otherwise, . public override bool Equals(object? obj) => obj is AgentRequestMessageSourceType other && this.Equals(other); + /// + /// Returns the string representation of this instance. + /// + /// The string value representing the source of the agent request message. + public override string ToString() => this.Value; + /// /// Returns the hash code for this instance. /// diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions/AgentSession.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions/AgentSession.cs index 4c62eccc99..a154b0a9f5 100644 --- a/dotnet/src/Microsoft.Agents.AI.Abstractions/AgentSession.cs +++ b/dotnet/src/Microsoft.Agents.AI.Abstractions/AgentSession.cs @@ -1,7 +1,9 @@ // Copyright (c) Microsoft. All rights reserved. using System; +using System.Diagnostics; using System.Text.Json; +using System.Text.Json.Serialization; using Microsoft.Shared.Diagnostics; namespace Microsoft.Agents.AI; @@ -44,6 +46,7 @@ namespace Microsoft.Agents.AI; /// /// /// +[DebuggerDisplay("{DebuggerDisplay,nq}")] public abstract class AgentSession { /// @@ -53,6 +56,20 @@ public abstract class AgentSession { } + /// + /// Initializes a new instance of the class. + /// + protected AgentSession(AgentSessionStateBag stateBag) + { + this.StateBag = Throw.IfNull(stateBag); + } + + /// + /// Gets any arbitrary state associated with this session. + /// + [JsonPropertyName("stateBag")] + public AgentSessionStateBag StateBag { get; protected set; } = new(); + /// Asks the for an object of the specified type . /// The type of object being requested. /// An optional key that can be used to help identify the target service. @@ -82,4 +99,7 @@ public abstract class AgentSession /// public TService? GetService(object? serviceKey = null) => this.GetService(typeof(TService), serviceKey) is TService service ? service : default; + + [DebuggerBrowsable(DebuggerBrowsableState.Never)] + private string DebuggerDisplay => $"StateBag Count = {this.StateBag.Count}"; } diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions/AgentSessionStateBag.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions/AgentSessionStateBag.cs new file mode 100644 index 0000000000..d78a866b2c --- /dev/null +++ b/dotnet/src/Microsoft.Agents.AI.Abstractions/AgentSessionStateBag.cs @@ -0,0 +1,145 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Collections.Concurrent; +using System.Text.Json; +using System.Text.Json.Serialization; +using Microsoft.Shared.Diagnostics; + +namespace Microsoft.Agents.AI; + +/// +/// Provides a thread-safe key-value store for managing session-scoped state with support for type-safe access and JSON +/// serialization options. +/// +/// +/// SessionState enables storing and retrieving objects associated with a session using string keys. +/// Values can be accessed in a type-safe manner and are serialized or deserialized using configurable JSON serializer +/// options. This class is designed for concurrent access and is safe to use across multiple threads. +/// +[JsonConverter(typeof(AgentSessionStateBagJsonConverter))] +public class AgentSessionStateBag +{ + private readonly ConcurrentDictionary _state; + + /// + /// Initializes a new instance of the class. + /// + public AgentSessionStateBag() + { + this._state = new ConcurrentDictionary(); + } + + /// + /// Initializes a new instance of the class. + /// + /// The initial state dictionary. + internal AgentSessionStateBag(ConcurrentDictionary? state) + { + this._state = state ?? new ConcurrentDictionary(); + } + + /// + /// Gets the number of key-value pairs contained in the session state. + /// + public int Count => this._state.Count; + + /// + /// Tries to get a value from the session state. + /// + /// The type of the value to retrieve. + /// The key from which to retrieve the value. + /// The value if found and convertible to the required type; otherwise, null. + /// The JSON serializer options to use for serializing/deserializing the value. + /// if the value was successfully retrieved, otherwise. + public bool TryGetValue(string key, out T? value, JsonSerializerOptions? jsonSerializerOptions = null) + where T : class + { + _ = Throw.IfNullOrWhitespace(key); + var jso = jsonSerializerOptions ?? AgentAbstractionsJsonUtilities.DefaultOptions; + + if (this._state.TryGetValue(key, out var stateValue)) + { + return stateValue.TryReadDeserializedValue(out value, jso); + } + + value = null; + return false; + } + + /// + /// Gets a value from the session state. + /// + /// The type of value to get. + /// The key from which to retrieve the value. + /// The JSON serializer options to use for serializing/deserialing the value. + /// The retrieved value or null if not found. + /// The value could not be deserialized into the required type. + public T? GetValue(string key, JsonSerializerOptions? jsonSerializerOptions = null) + where T : class + { + _ = Throw.IfNullOrWhitespace(key); + var jso = jsonSerializerOptions ?? AgentAbstractionsJsonUtilities.DefaultOptions; + + if (this._state.TryGetValue(key, out var stateValue)) + { + return stateValue.ReadDeserializedValue(jso); + } + + return null; + } + + /// + /// Sets a value in the session state. + /// + /// The type of the value to set. + /// The key to store the value under. + /// The value to set. + /// The JSON serializer options to use for serializing the value. + public void SetValue(string key, T? value, JsonSerializerOptions? jsonSerializerOptions = null) + where T : class + { + _ = Throw.IfNullOrWhitespace(key); + var jso = jsonSerializerOptions ?? AgentAbstractionsJsonUtilities.DefaultOptions; + + var stateValue = this._state.GetOrAdd(key, _ => + new AgentSessionStateBagValue(value, typeof(T), jso)); + + stateValue.SetDeserialized(value, typeof(T), jso); + } + + /// + /// Tries to remove a value from the session state. + /// + /// The key of the value to remove. + /// if the value was successfully removed; otherwise, . + public bool TryRemoveValue(string key) + => this._state.TryRemove(Throw.IfNullOrWhitespace(key), out _); + + /// + /// Serializes all session state values to a JSON object. + /// + /// A representing the serialized session state. + /// Thrown when a session state value is not properly initialized. + public JsonElement Serialize() + { + return JsonSerializer.SerializeToElement(this._state, AgentAbstractionsJsonUtilities.DefaultOptions.GetTypeInfo(typeof(ConcurrentDictionary))); + } + + /// + /// Deserializes a JSON object into an instance. + /// + /// The element to deserialize. + /// The deserialized . + public static AgentSessionStateBag Deserialize(JsonElement jsonElement) + { + if (jsonElement.ValueKind is JsonValueKind.Undefined or JsonValueKind.Null) + { + return new AgentSessionStateBag(); + } + + return new AgentSessionStateBag( + jsonElement.Deserialize(AgentAbstractionsJsonUtilities.DefaultOptions.GetTypeInfo(typeof(ConcurrentDictionary))) as ConcurrentDictionary + ?? new ConcurrentDictionary()); + } +} diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions/AgentSessionStateBagJsonConverter.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions/AgentSessionStateBagJsonConverter.cs new file mode 100644 index 0000000000..bfb6904320 --- /dev/null +++ b/dotnet/src/Microsoft.Agents.AI.Abstractions/AgentSessionStateBagJsonConverter.cs @@ -0,0 +1,28 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Text.Json; +using System.Text.Json.Serialization; + +namespace Microsoft.Agents.AI; + +/// +/// Custom JSON converter for that serializes and deserializes +/// the internal dictionary contents rather than the container object's public properties. +/// +public sealed class AgentSessionStateBagJsonConverter : JsonConverter +{ + /// + public override AgentSessionStateBag Read(ref Utf8JsonReader reader, Type typeToConvert, JsonSerializerOptions options) + { + var element = JsonElement.ParseValue(ref reader); + return AgentSessionStateBag.Deserialize(element); + } + + /// + public override void Write(Utf8JsonWriter writer, AgentSessionStateBag value, JsonSerializerOptions options) + { + var element = value.Serialize(); + element.WriteTo(writer); + } +} diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions/AgentSessionStateBagValue.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions/AgentSessionStateBagValue.cs new file mode 100644 index 0000000000..0b4849aa1b --- /dev/null +++ b/dotnet/src/Microsoft.Agents.AI.Abstractions/AgentSessionStateBagValue.cs @@ -0,0 +1,182 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Text.Json; +using System.Text.Json.Serialization; + +namespace Microsoft.Agents.AI; + +/// +/// Used to store a value in session state. +/// +[JsonConverter(typeof(AgentSessionStateBagValueJsonConverter))] +internal class AgentSessionStateBagValue +{ + private readonly object _lock = new(); + private DeserializedCache? _cache; + private JsonElement _jsonValue; + + /// + /// Initializes a new instance of the SessionStateValue class with the specified value. + /// + /// The serialized value to associate with the session state. + public AgentSessionStateBagValue(JsonElement jsonValue) + { + this.JsonValue = jsonValue; + } + + /// + /// Initializes a new instance of the SessionStateValue class with the specified value. + /// + /// The value to associate with the session state. Can be any object, including null. + /// The type of the value. + /// The JSON serializer options to use for serializing the value. + public AgentSessionStateBagValue(object? deserializedValue, Type valueType, JsonSerializerOptions jsonSerializerOptions) + { + this._cache = new DeserializedCache(deserializedValue, valueType, jsonSerializerOptions); + } + + /// + /// Gets or sets the value associated with this instance. + /// + public JsonElement JsonValue + { + get + { + lock (this._lock) + { + // We are assuming here that JsonValue will only be read when the object is being serialized, + // which means that we will only call SerializeToElement when serializing and therefore it's + // OK to serialize on each read if the cache is set. + if (this._cache is { } cache) + { + this._jsonValue = JsonSerializer.SerializeToElement(cache.Value, cache.Options.GetTypeInfo(cache.ValueType)); + } + + return this._jsonValue; + } + } + set + { + lock (this._lock) + { + this._jsonValue = value; + this._cache = null; + } + } + } + + /// + /// Tries to read the deserialized value of this session state value. + /// Returns false if the value could not be deserialized into the required type, or if the value is undefined. + /// Returns true and sets the out parameter to null if the value is null. + /// + public bool TryReadDeserializedValue(out T? value, JsonSerializerOptions? jsonSerializerOptions = null) + where T : class + { + var jso = jsonSerializerOptions ?? AgentAbstractionsJsonUtilities.DefaultOptions; + + lock (this._lock) + { + switch (this._cache) + { + case DeserializedCache { Value: null, ValueType: Type cacheValueType } when cacheValueType == typeof(T): + value = null; + return true; + case DeserializedCache { Value: T cacheValue, ValueType: Type cacheValueType } when cacheValueType == typeof(T): + value = cacheValue; + return true; + case DeserializedCache { ValueType: Type cacheValueType } when cacheValueType != typeof(T): + value = null; + return false; + } + + switch (this._jsonValue) + { + case JsonElement jsonElement when jsonElement.ValueKind == JsonValueKind.Undefined: + value = null; + return false; + case JsonElement jsonElement when jsonElement.ValueKind == JsonValueKind.Null: + value = null; + return true; + default: + T? result = this._jsonValue.Deserialize(jso.GetTypeInfo(typeof(T))) as T; + if (result is null) + { + value = null; + return false; + } + + this._cache = new DeserializedCache(result, typeof(T), jso); + + value = result; + return true; + } + } + } + + /// + /// Reads the deserialized value of this session state value, throwing an exception if the value could not be deserialized into the required type or is undefined. + /// + public T? ReadDeserializedValue(JsonSerializerOptions? jsonSerializerOptions = null) + where T : class + { + var jso = jsonSerializerOptions ?? AgentAbstractionsJsonUtilities.DefaultOptions; + + lock (this._lock) + { + switch (this._cache) + { + case DeserializedCache { Value: null, ValueType: Type cacheValueType } when cacheValueType == typeof(T): + return null; + case DeserializedCache { Value: T cacheValue, ValueType: Type cacheValueType } when cacheValueType == typeof(T): + return cacheValue; + case DeserializedCache { ValueType: Type cacheValueType } when cacheValueType != typeof(T): + throw new InvalidOperationException($"The type of the cached value is {cacheValueType.FullName}, but the requested type is {typeof(T).FullName}."); + } + + switch (this._jsonValue) + { + case JsonElement jsonElement when jsonElement.ValueKind == JsonValueKind.Null || jsonElement.ValueKind == JsonValueKind.Undefined: + return null; + default: + T? result = this._jsonValue.Deserialize(jso.GetTypeInfo(typeof(T))) as T; + if (result is null) + { + throw new InvalidOperationException($"Failed to deserialize session state value to type {typeof(T).FullName}."); + } + + this._cache = new DeserializedCache(result, typeof(T), jso); + return result; + } + } + } + + /// + /// Sets the deserialized value of this session state value, updating the cache accordingly. + /// This does not update the JsonValue directly; the JsonValue will be updated on the next read or when the object is serialized. + /// + public void SetDeserialized(T? deserializedValue, Type valueType, JsonSerializerOptions jsonSerializerOptions) + { + lock (this._lock) + { + this._cache = new DeserializedCache(deserializedValue, valueType, jsonSerializerOptions); + } + } + + private readonly struct DeserializedCache + { + public DeserializedCache(object? value, Type valueType, JsonSerializerOptions options) + { + this.Value = value; + this.ValueType = valueType; + this.Options = options; + } + + public object? Value { get; } + + public Type ValueType { get; } + + public JsonSerializerOptions Options { get; } + } +} diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions/AgentSessionStateBagValueJsonConverter.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions/AgentSessionStateBagValueJsonConverter.cs new file mode 100644 index 0000000000..27c9dc08a8 --- /dev/null +++ b/dotnet/src/Microsoft.Agents.AI.Abstractions/AgentSessionStateBagValueJsonConverter.cs @@ -0,0 +1,27 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Text.Json; +using System.Text.Json.Serialization; + +namespace Microsoft.Agents.AI; + +/// +/// Custom JSON converter for that serializes and deserializes +/// the directly rather than wrapping it in a container object. +/// +internal sealed class AgentSessionStateBagValueJsonConverter : JsonConverter +{ + /// + public override AgentSessionStateBagValue Read(ref Utf8JsonReader reader, Type typeToConvert, JsonSerializerOptions options) + { + var element = JsonElement.ParseValue(ref reader); + return new AgentSessionStateBagValue(element); + } + + /// + public override void Write(Utf8JsonWriter writer, AgentSessionStateBagValue value, JsonSerializerOptions options) + { + value.JsonValue.WriteTo(writer); + } +} diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions/ChatHistoryProvider.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions/ChatHistoryProvider.cs index 086ad02061..d16ca69528 100644 --- a/dotnet/src/Microsoft.Agents.AI.Abstractions/ChatHistoryProvider.cs +++ b/dotnet/src/Microsoft.Agents.AI.Abstractions/ChatHistoryProvider.cs @@ -2,8 +2,6 @@ using System; using System.Collections.Generic; -using System.Linq; -using System.Text.Json; using System.Threading; using System.Threading.Tasks; using Microsoft.Extensions.AI; @@ -27,34 +25,29 @@ namespace Microsoft.Agents.AI; /// Storing chat messages with proper ordering and metadata preservation /// Retrieving messages in chronological order for agent context /// Managing storage limits through truncation, summarization, or other strategies -/// Supporting serialization for thread persistence and migration /// /// /// +/// The is passed a reference to the via and +/// allowing it to store state in the . Since a is used with many different sessions, it should +/// not store any session-specific information within its own instance fields. Instead, any session-specific state should be stored in the associated . +/// +/// /// A is only relevant for scenarios where the underlying AI service that the agent is using /// does not use in-service chat history storage. /// /// public abstract class ChatHistoryProvider { - private readonly string _sourceId; - /// - /// Initializes a new instance of the class. + /// Gets the key used to store the provider state in the . /// - protected ChatHistoryProvider() - { - this._sourceId = this.GetType().FullName!; - } - - /// - /// Initializes a new instance of the class with the specified source id. - /// - /// The source id to stamp on for each messages produced by the . - protected ChatHistoryProvider(string sourceId) - { - this._sourceId = sourceId; - } + /// + /// The default value is the name of the concrete type (e.g. "InMemoryChatHistoryProvider"). + /// Implementations may override this to provide a custom key, for example when multiple + /// instances of the same provider type are used in the same session. + /// + public virtual string StateKey => this.GetType().Name; /// /// Called at the start of agent invocation to provide messages from the chat history as context for the next agent invocation. @@ -80,17 +73,9 @@ public abstract class ChatHistoryProvider /// Archiving old messages while keeping active conversation context /// /// - /// - /// Each instance should be associated with a single to ensure proper message isolation - /// and context management. - /// /// - public async ValueTask> InvokingAsync(InvokingContext context, CancellationToken cancellationToken = default) - { - var messages = await this.InvokingCoreAsync(context, cancellationToken).ConfigureAwait(false); - - return messages.Select(message => message.AsAgentRequestMessageSourcedMessage(AgentRequestMessageSourceType.ChatHistory, this._sourceId)); - } + public ValueTask> InvokingAsync(InvokingContext context, CancellationToken cancellationToken = default) + => this.InvokingCoreAsync(context, cancellationToken); /// /// Called at the start of agent invocation to provide messages from the chat history as context for the next agent invocation. @@ -178,13 +163,6 @@ public abstract class ChatHistoryProvider /// protected abstract ValueTask InvokedCoreAsync(InvokedContext context, CancellationToken cancellationToken = default); - /// - /// Serializes the current object's state to a using the specified serialization options. - /// - /// The JSON serialization options to use. - /// A representation of the object's state. - public abstract JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null); - /// Asks the for an object of the specified type . /// The type of object being requested. /// An optional key that can be used to help identify the target service. @@ -229,7 +207,7 @@ public abstract class ChatHistoryProvider /// /// The agent being invoked. /// The session associated with the agent invocation. - /// The new messages to be used by the agent for this invocation. + /// The messages to be used by the agent for this invocation. /// is . public InvokingContext( AIAgent agent, @@ -252,11 +230,22 @@ public abstract class ChatHistoryProvider public AgentSession? Session { get; } /// - /// Gets the caller provided messages that will be used by the agent for this invocation. + /// Gets the messages that will be used by the agent for this invocation. instances can modify + /// and return or return a new message list to add additional messages for the invocation. /// /// - /// A collection of instances representing new messages that were provided by the caller. + /// A collection of instances representing the messages that will be used by the agent for this invocation. /// + /// + /// + /// If multiple instances are used in the same invocation, each + /// will receive the messages returned by the previous allowing them to build on top of each other's context. + /// + /// + /// The first in the invocation pipeline will receive the + /// caller provided messages. + /// + /// public IEnumerable RequestMessages { get; set { field = Throw.IfNull(value); } } } @@ -264,27 +253,52 @@ public abstract class ChatHistoryProvider /// Contains the context information provided to . /// /// - /// This class provides context about a completed agent invocation, including both the - /// request messages that were used and the response messages that were generated. It also indicates - /// whether the invocation succeeded or failed. + /// This class provides context about a completed agent invocation, including the accumulated + /// request messages (user input, chat history and any others provided by AI context providers) that were used + /// and the response messages that were generated. It also indicates whether the invocation succeeded or failed. /// public sealed class InvokedContext { /// - /// Initializes a new instance of the class with the specified request messages. + /// Initializes a new instance of the class for a successful invocation. /// - /// The agent being invoked. + /// The agent that was invoked. /// The session associated with the agent invocation. - /// The caller provided messages that were used by the agent for this invocation. - /// is . + /// The accumulated request messages (user input, chat history and any others provided by AI context providers) + /// that were used by the agent for this invocation. + /// The response messages generated during this invocation. + /// , , or is . public InvokedContext( AIAgent agent, AgentSession? session, - IEnumerable requestMessages) + IEnumerable requestMessages, + IEnumerable responseMessages) { this.Agent = Throw.IfNull(agent); this.Session = session; this.RequestMessages = Throw.IfNull(requestMessages); + this.ResponseMessages = Throw.IfNull(responseMessages); + } + + /// + /// Initializes a new instance of the class for a failed invocation. + /// + /// The agent that was invoked. + /// The session associated with the agent invocation. + /// The accumulated request messages (user input, chat history and any others provided by AI context providers) + /// that were used by the agent for this invocation. + /// The exception that caused the invocation to fail. + /// , , or is . + public InvokedContext( + AIAgent agent, + AgentSession? session, + IEnumerable requestMessages, + Exception invokeException) + { + this.Agent = Throw.IfNull(agent); + this.Session = session; + this.RequestMessages = Throw.IfNull(requestMessages); + this.InvokeException = Throw.IfNull(invokeException); } /// @@ -298,22 +312,23 @@ public abstract class ChatHistoryProvider public AgentSession? Session { get; } /// - /// Gets the caller provided messages that were used by the agent for this invocation. + /// Gets the accumulated request messages (user input, chat history and any others provided by AI context providers) + /// that were used by the agent for this invocation. /// /// /// A collection of instances representing new messages that were provided by the caller. /// This does not include any supplied messages. /// - public IEnumerable RequestMessages { get; set { field = Throw.IfNull(value); } } + public IEnumerable RequestMessages { get; } /// /// Gets the collection of response messages generated during this invocation if the invocation succeeded. /// /// /// A collection of instances representing the response, - /// or if the invocation failed or did not produce response messages. + /// or if the invocation failed. /// - public IEnumerable? ResponseMessages { get; set; } + public IEnumerable? ResponseMessages { get; } /// /// Gets the that was thrown during the invocation, if the invocation failed. @@ -321,6 +336,6 @@ public abstract class ChatHistoryProvider /// /// The exception that caused the invocation to fail, or if the invocation succeeded. /// - public Exception? InvokeException { get; set; } + public Exception? InvokeException { get; } } } diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions/ChatHistoryProviderExtensions.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions/ChatHistoryProviderExtensions.cs deleted file mode 100644 index 4c8ef2489f..0000000000 --- a/dotnet/src/Microsoft.Agents.AI.Abstractions/ChatHistoryProviderExtensions.cs +++ /dev/null @@ -1,52 +0,0 @@ -// Copyright (c) Microsoft. All rights reserved. - -using System; -using System.Collections.Generic; -using System.Linq; -using Microsoft.Extensions.AI; - -namespace Microsoft.Agents.AI; - -/// -/// Contains extension methods for the class. -/// -public static class ChatHistoryProviderExtensions -{ - /// - /// Adds message filtering to an existing , so that messages passed to the and messages - /// provided by the can be filtered, updated or replaced. - /// - /// The to add the message filter to. - /// An optional filter function to apply to messages produced by the . If null, no filter is applied at this - /// stage. - /// An optional filter function to apply to the invoked context messages before they are passed to the . If null, no - /// filter is applied at this stage. - /// The with filtering applied. - public static ChatHistoryProvider WithMessageFilters( - this ChatHistoryProvider provider, - Func, IEnumerable>? invokingMessagesFilter = null, - Func? invokedMessagesFilter = null) - { - return new ChatHistoryProviderMessageFilter( - innerProvider: provider, - invokingMessagesFilter: invokingMessagesFilter, - invokedMessagesFilter: invokedMessagesFilter); - } - - /// - /// Decorates the provided so that it does not add - /// messages with to chat history. - /// - /// The to add the message filter to. - /// A new instance that filters out messages so they do not get added. - public static ChatHistoryProvider WithAIContextProviderMessageRemoval(this ChatHistoryProvider provider) - { - return new ChatHistoryProviderMessageFilter( - innerProvider: provider, - invokedMessagesFilter: (ctx) => - { - ctx.RequestMessages = ctx.RequestMessages.Where(x => x.GetAgentRequestMessageSourceType() != AgentRequestMessageSourceType.AIContextProvider); - return ctx; - }); - } -} diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions/ChatHistoryProviderMessageFilter.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions/ChatHistoryProviderMessageFilter.cs deleted file mode 100644 index 6cee80986b..0000000000 --- a/dotnet/src/Microsoft.Agents.AI.Abstractions/ChatHistoryProviderMessageFilter.cs +++ /dev/null @@ -1,74 +0,0 @@ -// Copyright (c) Microsoft. All rights reserved. - -using System; -using System.Collections.Generic; -using System.Text.Json; -using System.Threading; -using System.Threading.Tasks; -using Microsoft.Extensions.AI; -using Microsoft.Shared.Diagnostics; - -namespace Microsoft.Agents.AI; - -/// -/// A decorator that allows filtering the messages -/// passed into and out of an inner . -/// -public sealed class ChatHistoryProviderMessageFilter : ChatHistoryProvider -{ - private readonly ChatHistoryProvider _innerProvider; - private readonly Func, IEnumerable>? _invokingMessagesFilter; - private readonly Func? _invokedMessagesFilter; - - /// - /// Initializes a new instance of the class. - /// - /// Use this constructor to customize how messages are filtered before and after invocation by - /// providing appropriate filter functions. If no filters are provided, the operates without - /// additional filtering. - /// The underlying to be wrapped. Cannot be null. - /// An optional filter function to apply to messages provided by the - /// before they are used by the agent. If null, no filter is applied at this stage. - /// An optional filter function to apply to the invocation context after messages have been produced. If null, no - /// filter is applied at this stage. - /// Thrown if is null. - public ChatHistoryProviderMessageFilter( - ChatHistoryProvider innerProvider, - Func, IEnumerable>? invokingMessagesFilter = null, - Func? invokedMessagesFilter = null) - { - this._innerProvider = Throw.IfNull(innerProvider); - - if (invokingMessagesFilter == null && invokedMessagesFilter == null) - { - throw new ArgumentException("At least one filter function, invokingMessagesFilter or invokedMessagesFilter, must be provided."); - } - - this._invokingMessagesFilter = invokingMessagesFilter; - this._invokedMessagesFilter = invokedMessagesFilter; - } - - /// - protected override async ValueTask> InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default) - { - var messages = await this._innerProvider.InvokingAsync(context, cancellationToken).ConfigureAwait(false); - return this._invokingMessagesFilter != null ? this._invokingMessagesFilter(messages) : messages; - } - - /// - protected override ValueTask InvokedCoreAsync(InvokedContext context, CancellationToken cancellationToken = default) - { - if (this._invokedMessagesFilter != null) - { - context = this._invokedMessagesFilter(context); - } - - return this._innerProvider.InvokedAsync(context, cancellationToken); - } - - /// - public override JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null) - { - return this._innerProvider.Serialize(jsonSerializerOptions); - } -} diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions/ChatMessageExtensions.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions/ChatMessageExtensions.cs index 052e48ce56..0ff4874732 100644 --- a/dotnet/src/Microsoft.Agents.AI.Abstractions/ChatMessageExtensions.cs +++ b/dotnet/src/Microsoft.Agents.AI.Abstractions/ChatMessageExtensions.cs @@ -54,7 +54,7 @@ public static class ChatMessageExtensions /// If the message is already tagged with the provided source type and source id, it is returned as is. /// Otherwise, a cloned message is returned with the appropriate tagging in the AdditionalProperties. /// - public static ChatMessage AsAgentRequestMessageSourcedMessage(this ChatMessage message, AgentRequestMessageSourceType sourceType, string? sourceId = null) + public static ChatMessage WithAgentRequestMessageSource(this ChatMessage message, AgentRequestMessageSourceType sourceType, string? sourceId = null) { if (message.AdditionalProperties != null // Check if the message was already tagged with the required source type and source id diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions/InMemoryAgentSession.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions/InMemoryAgentSession.cs deleted file mode 100644 index 05ffafaeb9..0000000000 --- a/dotnet/src/Microsoft.Agents.AI.Abstractions/InMemoryAgentSession.cs +++ /dev/null @@ -1,124 +0,0 @@ -// Copyright (c) Microsoft. All rights reserved. - -using System; -using System.Collections.Generic; -using System.Diagnostics; -using System.Text.Json; -using Microsoft.Extensions.AI; - -namespace Microsoft.Agents.AI; - -/// -/// Provides an abstract base class for an that maintain all chat history in local memory. -/// -/// -/// -/// is designed for scenarios where chat history should be stored locally -/// rather than in external services or databases. This approach provides high performance and simplicity while -/// maintaining full control over the conversation data. -/// -/// -/// In-memory threads do not persist conversation data across application restarts -/// unless explicitly serialized and restored. -/// -/// -[DebuggerDisplay("{DebuggerDisplay,nq}")] -public abstract class InMemoryAgentSession : AgentSession -{ - /// - /// Initializes a new instance of the class. - /// - /// - /// An optional instance to use for storing chat messages. - /// If , a new empty will be created. - /// - /// - /// This constructor allows sharing of between sessions or providing pre-configured - /// with specific reduction or processing logic. - /// - protected InMemoryAgentSession(InMemoryChatHistoryProvider? chatHistoryProvider = null) - { - this.ChatHistoryProvider = chatHistoryProvider ?? []; - } - - /// - /// Initializes a new instance of the class. - /// - /// The initial messages to populate the conversation history. - /// is . - /// - /// This constructor is useful for initializing sessions with existing conversation history or - /// for migrating conversations from other storage systems. - /// - protected InMemoryAgentSession(IEnumerable messages) - { - this.ChatHistoryProvider = [.. messages]; - } - - /// - /// Initializes a new instance of the class from previously serialized state. - /// - /// A representing the serialized state of the session. - /// Optional settings for customizing the JSON deserialization process. - /// - /// Optional factory function to create the from its serialized state. - /// If not provided, a default factory will be used that creates a basic . - /// - /// The is not a JSON object. - /// The is invalid or cannot be deserialized to the expected type. - /// - /// This constructor enables restoration of in-memory threads from previously saved state, allowing - /// conversations to be resumed across application restarts or migrated between different instances. - /// - protected InMemoryAgentSession( - JsonElement serializedState, - JsonSerializerOptions? jsonSerializerOptions = null, - Func? chatHistoryProviderFactory = null) - { - if (serializedState.ValueKind != JsonValueKind.Object) - { - throw new ArgumentException("The serialized session state must be a JSON object.", nameof(serializedState)); - } - - var state = serializedState.Deserialize( - AgentAbstractionsJsonUtilities.DefaultOptions.GetTypeInfo(typeof(InMemoryAgentSessionState))) as InMemoryAgentSessionState; - - this.ChatHistoryProvider = - chatHistoryProviderFactory?.Invoke(state?.ChatHistoryProviderState ?? default, jsonSerializerOptions) ?? - new InMemoryChatHistoryProvider(state?.ChatHistoryProviderState ?? default, jsonSerializerOptions); - } - - /// - /// Gets or sets the used by this thread. - /// - public InMemoryChatHistoryProvider ChatHistoryProvider { get; } - - /// - /// Serializes the current object's state to a using the specified serialization options. - /// - /// The JSON serialization options to use. - /// A representation of the object's state. - protected internal virtual JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null) - { - var chatHistoryProviderState = this.ChatHistoryProvider.Serialize(jsonSerializerOptions); - - var state = new InMemoryAgentSessionState - { - ChatHistoryProviderState = chatHistoryProviderState, - }; - - return JsonSerializer.SerializeToElement(state, AgentAbstractionsJsonUtilities.DefaultOptions.GetTypeInfo(typeof(InMemoryAgentSessionState))); - } - - /// - public override object? GetService(Type serviceType, object? serviceKey = null) => - base.GetService(serviceType, serviceKey) ?? this.ChatHistoryProvider?.GetService(serviceType, serviceKey); - - [DebuggerBrowsable(DebuggerBrowsableState.Never)] - private string DebuggerDisplay => $"Count = {this.ChatHistoryProvider.Count}"; - - internal sealed class InMemoryAgentSessionState - { - public JsonElement? ChatHistoryProviderState { get; set; } - } -} diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions/InMemoryChatHistoryProvider.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions/InMemoryChatHistoryProvider.cs index 001b3a3bcc..9c535923f4 100644 --- a/dotnet/src/Microsoft.Agents.AI.Abstractions/InMemoryChatHistoryProvider.cs +++ b/dotnet/src/Microsoft.Agents.AI.Abstractions/InMemoryChatHistoryProvider.cs @@ -1,11 +1,10 @@ // Copyright (c) Microsoft. All rights reserved. using System; -using System.Collections; using System.Collections.Generic; -using System.Diagnostics; using System.Linq; using System.Text.Json; +using System.Text.Json.Serialization; using System.Threading; using System.Threading.Tasks; using Microsoft.Extensions.AI; @@ -14,100 +13,49 @@ using Microsoft.Shared.Diagnostics; namespace Microsoft.Agents.AI; /// -/// Provides an in-memory implementation of with support for message reduction and collection semantics. +/// Provides an in-memory implementation of with support for message reduction. /// /// /// -/// stores chat messages entirely in local memory, providing fast access and manipulation -/// capabilities. It implements both for agent integration and -/// for direct collection manipulation. +/// stores chat messages in the , +/// providing fast access and manipulation capabilities integrated with session state management. /// /// /// This maintains all messages in memory. For long-running conversations or high-volume scenarios, consider using /// message reduction strategies or alternative storage implementations. /// /// -[DebuggerDisplay("Count = {Count}")] -[DebuggerTypeProxy(typeof(DebugView))] -public sealed class InMemoryChatHistoryProvider : ChatHistoryProvider, IList, IReadOnlyList +public sealed class InMemoryChatHistoryProvider : ChatHistoryProvider { - private List _messages; + private static IEnumerable DefaultExcludeChatHistoryFilter(IEnumerable messages) + => messages.Where(m => m.GetAgentRequestMessageSourceType() != AgentRequestMessageSourceType.ChatHistory); + + private readonly string _stateKey; + private readonly Func _stateInitializer; + private readonly JsonSerializerOptions _jsonSerializerOptions; + private readonly Func, IEnumerable> _storageInputMessageFilter; + private readonly Func, IEnumerable>? _retrievalOutputMessageFilter; /// /// Initializes a new instance of the class. /// - /// - /// This constructor creates a basic in-memory without message reduction capabilities. - /// Messages will be stored exactly as added without any automatic processing or reduction. - /// - public InMemoryChatHistoryProvider() - { - this._messages = []; - } - - /// - /// Initializes a new instance of the class from previously serialized state. - /// - /// A representing the serialized state of the provider. - /// Optional settings for customizing the JSON deserialization process. - /// The is not a valid JSON object or cannot be deserialized. - /// - /// This constructor enables restoration of messages from previously saved state, allowing - /// conversation history to be preserved across application restarts or migrated between instances. - /// - public InMemoryChatHistoryProvider(JsonElement serializedState, JsonSerializerOptions? jsonSerializerOptions = null) - : this(null, serializedState, jsonSerializerOptions, ChatReducerTriggerEvent.BeforeMessagesRetrieval) - { - } - - /// - /// Initializes a new instance of the class. - /// - /// - /// A instance used to process, reduce, or optimize chat messages. - /// This can be used to implement strategies like message summarization, truncation, or cleanup. + /// + /// Optional configuration options that control the provider's behavior, including state initialization, + /// message reduction, and serialization settings. If , default settings will be used. /// - /// - /// Specifies when the message reducer should be invoked. The default is , - /// which applies reduction logic when messages are retrieved for agent consumption. - /// - /// is . - /// - /// Message reducers enable automatic management of message storage by implementing strategies to - /// keep memory usage under control while preserving important conversation context. - /// - public InMemoryChatHistoryProvider(IChatReducer chatReducer, ChatReducerTriggerEvent reducerTriggerEvent = ChatReducerTriggerEvent.BeforeMessagesRetrieval) - : this(chatReducer, default, null, reducerTriggerEvent) + public InMemoryChatHistoryProvider(InMemoryChatHistoryProviderOptions? options = null) { - Throw.IfNull(chatReducer); + this._stateInitializer = options?.StateInitializer ?? (_ => new State()); + this.ChatReducer = options?.ChatReducer; + this.ReducerTriggerEvent = options?.ReducerTriggerEvent ?? InMemoryChatHistoryProviderOptions.ChatReducerTriggerEvent.BeforeMessagesRetrieval; + this._stateKey = options?.StateKey ?? base.StateKey; + this._jsonSerializerOptions = options?.JsonSerializerOptions ?? AgentAbstractionsJsonUtilities.DefaultOptions; + this._storageInputMessageFilter = options?.StorageInputMessageFilter ?? DefaultExcludeChatHistoryFilter; + this._retrievalOutputMessageFilter = options?.RetrievalOutputMessageFilter; } - /// - /// Initializes a new instance of the class, with an existing state from a serialized JSON element. - /// - /// An optional instance used to process or reduce chat messages. If null, no reduction logic will be applied. - /// A representing the serialized state of the provider. - /// Optional settings for customizing the JSON deserialization process. - /// The event that should trigger the reducer invocation. - public InMemoryChatHistoryProvider(IChatReducer? chatReducer, JsonElement serializedState, JsonSerializerOptions? jsonSerializerOptions = null, ChatReducerTriggerEvent reducerTriggerEvent = ChatReducerTriggerEvent.BeforeMessagesRetrieval) - { - this.ChatReducer = chatReducer; - this.ReducerTriggerEvent = reducerTriggerEvent; - - if (serializedState.ValueKind is JsonValueKind.Object) - { - var jso = jsonSerializerOptions ?? AgentAbstractionsJsonUtilities.DefaultOptions; - var state = serializedState.Deserialize( - jso.GetTypeInfo(typeof(State))) as State; - if (state?.Messages is { } messages) - { - this._messages = messages; - return; - } - } - - this._messages = []; - } + /// + public override string StateKey => this._stateKey; /// /// Gets the chat reducer used to process or reduce chat messages. If null, no reduction logic will be applied. @@ -117,19 +65,49 @@ public sealed class InMemoryChatHistoryProvider : ChatHistoryProvider, IList /// Gets the event that triggers the reducer invocation in this provider. /// - public ChatReducerTriggerEvent ReducerTriggerEvent { get; } + public InMemoryChatHistoryProviderOptions.ChatReducerTriggerEvent ReducerTriggerEvent { get; } - /// - public int Count => this._messages.Count; + /// + /// Gets the chat messages stored for the specified session. + /// + /// The agent session containing the state. + /// A list of chat messages, or an empty list if no state is found. + public List GetMessages(AgentSession? session) + => this.GetOrInitializeState(session).Messages; - /// - public bool IsReadOnly => ((IList)this._messages).IsReadOnly; - - /// - public ChatMessage this[int index] + /// + /// Sets the chat messages for the specified session. + /// + /// The agent session containing the state. + /// The messages to store. + /// is . + public void SetMessages(AgentSession? session, List messages) { - get => this._messages[index]; - set => this._messages[index] = value; + _ = Throw.IfNull(messages); + + var state = this.GetOrInitializeState(session); + state.Messages = messages; + } + + /// + /// Gets the state from the session's StateBag, or initializes it using the state initializer if not present. + /// + /// The agent session containing the StateBag. + /// The provider state, or null if no session is available. + private State GetOrInitializeState(AgentSession? session) + { + if (session?.StateBag.TryGetValue(this._stateKey, out var state, this._jsonSerializerOptions) is true && state is not null) + { + return state; + } + + state = this._stateInitializer(session); + if (session is not null) + { + session.StateBag.SetValue(this._stateKey, state, this._jsonSerializerOptions); + } + + return state; } /// @@ -137,12 +115,21 @@ public sealed class InMemoryChatHistoryProvider : ChatHistoryProvider, IList output = state.Messages; + if (this._retrievalOutputMessageFilter is not null) + { + output = this._retrievalOutputMessageFilter(output); + } + return output + .Select(message => message.WithAgentRequestMessageSource(AgentRequestMessageSourceType.ChatHistory, this.GetType().FullName!)) + .Concat(context.RequestMessages); } /// @@ -155,94 +142,27 @@ public sealed class InMemoryChatHistoryProvider : ChatHistoryProvider, IList - public override JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null) - { - State state = new() - { - Messages = this._messages, - }; - - var jso = jsonSerializerOptions ?? AgentAbstractionsJsonUtilities.DefaultOptions; - return JsonSerializer.SerializeToElement(state, jso.GetTypeInfo(typeof(State))); - } - - /// - public int IndexOf(ChatMessage item) - => this._messages.IndexOf(item); - - /// - public void Insert(int index, ChatMessage item) - => this._messages.Insert(index, item); - - /// - public void RemoveAt(int index) - => this._messages.RemoveAt(index); - - /// - public void Add(ChatMessage item) - => this._messages.Add(item); - - /// - public void Clear() - => this._messages.Clear(); - - /// - public bool Contains(ChatMessage item) - => this._messages.Contains(item); - - /// - public void CopyTo(ChatMessage[] array, int arrayIndex) - => this._messages.CopyTo(array, arrayIndex); - - /// - public bool Remove(ChatMessage item) - => this._messages.Remove(item); - - /// - public IEnumerator GetEnumerator() - => this._messages.GetEnumerator(); - - /// - IEnumerator IEnumerable.GetEnumerator() - => this.GetEnumerator(); - - internal sealed class State + /// + /// Represents the state of a stored in the . + /// + public sealed class State { + /// + /// Gets or sets the list of chat messages. + /// + [JsonPropertyName("messages")] public List Messages { get; set; } = []; } - - /// - /// Defines the events that can trigger a reducer in the . - /// - public enum ChatReducerTriggerEvent - { - /// - /// Trigger the reducer when a new message is added. - /// will only complete when reducer processing is done. - /// - AfterMessageAdded, - - /// - /// Trigger the reducer before messages are retrieved from the provider. - /// The reducer will process the messages before they are returned to the caller. - /// - BeforeMessagesRetrieval - } - - private sealed class DebugView(InMemoryChatHistoryProvider provider) - { - [DebuggerBrowsable(DebuggerBrowsableState.RootHidden)] - public ChatMessage[] Items => provider._messages.ToArray(); - } } diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions/InMemoryChatHistoryProviderOptions.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions/InMemoryChatHistoryProviderOptions.cs new file mode 100644 index 0000000000..41ab46321f --- /dev/null +++ b/dotnet/src/Microsoft.Agents.AI.Abstractions/InMemoryChatHistoryProviderOptions.cs @@ -0,0 +1,93 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Collections.Generic; +using System.Text.Json; +using Microsoft.Extensions.AI; + +namespace Microsoft.Agents.AI; + +/// +/// Represents configuration options for . +/// +public sealed class InMemoryChatHistoryProviderOptions +{ + /// + /// Gets or sets an optional delegate that initializes the provider state on the first invocation. + /// If , a default initializer that creates an empty state will be used. + /// + public Func? StateInitializer { get; set; } + + /// + /// Gets or sets an optional instance used to process, reduce, or optimize chat messages. + /// This can be used to implement strategies like message summarization, truncation, or cleanup. + /// + public IChatReducer? ChatReducer { get; set; } + + /// + /// Gets or sets when the message reducer should be invoked. + /// The default is , + /// which applies reduction logic when messages are retrieved for agent consumption. + /// + /// + /// Message reducers enable automatic management of message storage by implementing strategies to + /// keep memory usage under control while preserving important conversation context. + /// + public ChatReducerTriggerEvent ReducerTriggerEvent { get; set; } = ChatReducerTriggerEvent.BeforeMessagesRetrieval; + + /// + /// Gets or sets an optional key to use for storing the state in the . + /// If , a default key will be used. + /// + public string? StateKey { get; set; } + + /// + /// Gets or sets optional JSON serializer options for serializing the state of this provider. + /// This is valuable for cases like when the chat history contains custom types + /// and source generated serializers are required, or Native AOT / Trimming is required. + /// + public JsonSerializerOptions? JsonSerializerOptions { get; set; } + + /// + /// Gets or sets an optional filter function applied to request messages before they are added to storage + /// during . + /// + /// + /// When , the provider defaults to excluding messages with + /// source type to avoid + /// storing messages that came from chat history in the first place. + /// Depending on your requirements, you could provide a different filter, that also excludes + /// messages from e.g. AI context providers. + /// + public Func, IEnumerable>? StorageInputMessageFilter { get; set; } + + /// + /// Gets or sets an optional filter function applied to messages produced by this provider + /// during . + /// + /// + /// This filter is only applied to the messages that the provider itself produces (from its internal storage). + /// + /// + /// When , no filtering is applied to the output messages. + /// + public Func, IEnumerable>? RetrievalOutputMessageFilter { get; set; } + + /// + /// Defines the events that can trigger a reducer in the . + /// + public enum ChatReducerTriggerEvent + { + /// + /// Trigger the reducer when a new message is added. + /// will only complete when reducer processing is done. + /// + AfterMessageAdded, + + /// + /// Trigger the reducer before messages are retrieved from the provider. + /// The reducer will process the messages before they are returned to the caller. + /// + BeforeMessagesRetrieval + } +} diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions/ServiceIdAgentSession.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions/ServiceIdAgentSession.cs deleted file mode 100644 index cf00635984..0000000000 --- a/dotnet/src/Microsoft.Agents.AI.Abstractions/ServiceIdAgentSession.cs +++ /dev/null @@ -1,104 +0,0 @@ -// Copyright (c) Microsoft. All rights reserved. - -using System; -using System.Diagnostics; -using System.Text.Json; -using Microsoft.Shared.Diagnostics; - -namespace Microsoft.Agents.AI; - -/// -/// Provides a base class for agent sessions that store conversation state remotely in a service and maintain only an identifier reference locally. -/// -/// -/// This class is designed for scenarios where conversation state is managed by an external service (such as a cloud-based AI service) -/// rather than being stored locally. The session maintains only the service identifier needed to reference the remote conversation state. -/// -[DebuggerDisplay("ServiceSessionId = {ServiceSessionId}")] -public abstract class ServiceIdAgentSession : AgentSession -{ - /// - /// Initializes a new instance of the class without a service session identifier. - /// - /// - /// When using this constructor, the will be initially - /// and should be set by derived classes when the remote conversation is created. - /// - protected ServiceIdAgentSession() - { - } - - /// - /// Initializes a new instance of the class with the specified service session identifier. - /// - /// The unique identifier that references the conversation state stored in the remote service. - /// is . - /// is empty or contains only whitespace. - protected ServiceIdAgentSession(string serviceSessionId) - { - this.ServiceSessionId = Throw.IfNullOrEmpty(serviceSessionId); - } - - /// - /// Initializes a new instance of the class from previously serialized state. - /// - /// A representing the serialized state of the session. - /// Optional settings for customizing the JSON deserialization process. - /// The is not a JSON object. - /// The is invalid or cannot be deserialized to the expected type. - /// - /// This constructor enables restoration of a service-backed session from serialized state, typically used - /// when deserializing session information that was previously saved or transmitted across application boundaries. - /// - protected ServiceIdAgentSession( - JsonElement serializedState, - JsonSerializerOptions? jsonSerializerOptions = null) - { - if (serializedState.ValueKind != JsonValueKind.Object) - { - throw new ArgumentException("The serialized session state must be a JSON object.", nameof(serializedState)); - } - - var state = serializedState.Deserialize( - AgentAbstractionsJsonUtilities.DefaultOptions.GetTypeInfo(typeof(ServiceIdAgentSessionState))) as ServiceIdAgentSessionState; - - if (state?.ServiceSessionId is string serviceSessionId) - { - this.ServiceSessionId = serviceSessionId; - } - } - - /// - /// Gets or sets the unique identifier that references the conversation state stored in the remote service. - /// - /// - /// A string identifier that uniquely identifies the conversation within the remote service, - /// or if no remote conversation has been established yet. - /// - /// - /// This identifier is used by derived classes to reference the remote conversation state when making - /// API calls to the backing service. The exact format and meaning of this identifier depends on the - /// specific service implementation. - /// - protected string? ServiceSessionId { get; set; } - - /// - /// Serializes the current object's state to a using the specified serialization options. - /// - /// The JSON serialization options to use. - /// A representation of the object's state. - protected internal virtual JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null) - { - var state = new ServiceIdAgentSessionState - { - ServiceSessionId = this.ServiceSessionId, - }; - - return JsonSerializer.SerializeToElement(state, AgentAbstractionsJsonUtilities.DefaultOptions.GetTypeInfo(typeof(ServiceIdAgentSessionState))); - } - - internal sealed class ServiceIdAgentSessionState - { - public string? ServiceSessionId { get; set; } - } -} diff --git a/dotnet/src/Microsoft.Agents.AI.AzureAI.Persistent/PersistentAgentsClientExtensions.cs b/dotnet/src/Microsoft.Agents.AI.AzureAI.Persistent/PersistentAgentsClientExtensions.cs index 2058e6760b..660e874711 100644 --- a/dotnet/src/Microsoft.Agents.AI.AzureAI.Persistent/PersistentAgentsClientExtensions.cs +++ b/dotnet/src/Microsoft.Agents.AI.AzureAI.Persistent/PersistentAgentsClientExtensions.cs @@ -191,8 +191,8 @@ public static class PersistentAgentsClientExtensions Name = options.Name ?? persistentAgentMetadata.Name, Description = options.Description ?? persistentAgentMetadata.Description, ChatOptions = options.ChatOptions, - AIContextProviderFactory = options.AIContextProviderFactory, - ChatHistoryProviderFactory = options.ChatHistoryProviderFactory, + AIContextProviders = options.AIContextProviders, + ChatHistoryProvider = options.ChatHistoryProvider, UseProvidedChatClientAsIs = options.UseProvidedChatClientAsIs }; diff --git a/dotnet/src/Microsoft.Agents.AI.AzureAI/AzureAIProjectChatClientExtensions.cs b/dotnet/src/Microsoft.Agents.AI.AzureAI/AzureAIProjectChatClientExtensions.cs index 8433ff3b6f..c35f49b088 100644 --- a/dotnet/src/Microsoft.Agents.AI.AzureAI/AzureAIProjectChatClientExtensions.cs +++ b/dotnet/src/Microsoft.Agents.AI.AzureAI/AzureAIProjectChatClientExtensions.cs @@ -594,8 +594,8 @@ public static partial class AzureAIProjectChatClientExtensions var agentOptions = CreateChatClientAgentOptions(agentVersion, options?.ChatOptions, requireInvocableTools); if (options is not null) { - agentOptions.AIContextProviderFactory = options.AIContextProviderFactory; - agentOptions.ChatHistoryProviderFactory = options.ChatHistoryProviderFactory; + agentOptions.AIContextProviders = options.AIContextProviders; + agentOptions.ChatHistoryProvider = options.ChatHistoryProvider; agentOptions.UseProvidedChatClientAsIs = options.UseProvidedChatClientAsIs; } diff --git a/dotnet/src/Microsoft.Agents.AI.CopilotStudio/CopilotStudioAgent.cs b/dotnet/src/Microsoft.Agents.AI.CopilotStudio/CopilotStudioAgent.cs index c631f90f17..bb53f94b50 100644 --- a/dotnet/src/Microsoft.Agents.AI.CopilotStudio/CopilotStudioAgent.cs +++ b/dotnet/src/Microsoft.Agents.AI.CopilotStudio/CopilotStudioAgent.cs @@ -68,7 +68,7 @@ public class CopilotStudioAgent : AIAgent /// protected override ValueTask DeserializeSessionCoreAsync(JsonElement serializedState, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default) - => new(new CopilotStudioAgentSession(serializedState, jsonSerializerOptions)); + => new(CopilotStudioAgentSession.Deserialize(serializedState, jsonSerializerOptions)); /// protected override async Task RunCoreAsync( diff --git a/dotnet/src/Microsoft.Agents.AI.CopilotStudio/CopilotStudioAgentSession.cs b/dotnet/src/Microsoft.Agents.AI.CopilotStudio/CopilotStudioAgentSession.cs index ec3e21ca91..ba101df082 100644 --- a/dotnet/src/Microsoft.Agents.AI.CopilotStudio/CopilotStudioAgentSession.cs +++ b/dotnet/src/Microsoft.Agents.AI.CopilotStudio/CopilotStudioAgentSession.cs @@ -1,36 +1,58 @@ // Copyright (c) Microsoft. All rights reserved. +using System; +using System.Diagnostics; using System.Text.Json; +using System.Text.Json.Serialization; namespace Microsoft.Agents.AI.CopilotStudio; /// /// Session for CopilotStudio based agents. /// -public sealed class CopilotStudioAgentSession : ServiceIdAgentSession +[DebuggerDisplay("{DebuggerDisplay,nq}")] +public sealed class CopilotStudioAgentSession : AgentSession { internal CopilotStudioAgentSession() { } - internal CopilotStudioAgentSession(JsonElement serializedSessionState, JsonSerializerOptions? jsonSerializerOptions = null) : base(serializedSessionState, jsonSerializerOptions) + [JsonConstructor] + internal CopilotStudioAgentSession(string? conversationId, AgentSessionStateBag? stateBag) : base(stateBag ?? new()) { + this.ConversationId = conversationId; } /// /// Gets the ID for the current conversation with the Copilot Studio agent. /// - public string? ConversationId - { - get { return this.ServiceSessionId; } - internal set { this.ServiceSessionId = value; } - } + [JsonPropertyName("serviceSessionId")] + public string? ConversationId { get; internal set; } /// /// Serializes the current object's state to a using the specified serialization options. /// /// The JSON serialization options to use. /// A representation of the object's state. - internal new JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null) - => base.Serialize(jsonSerializerOptions); + internal JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null) + { + var jso = jsonSerializerOptions ?? CopilotStudioJsonUtilities.DefaultOptions; + return JsonSerializer.SerializeToElement(this, jso.GetTypeInfo(typeof(CopilotStudioAgentSession))); + } + + internal static CopilotStudioAgentSession Deserialize(JsonElement serializedState, JsonSerializerOptions? jsonSerializerOptions = null) + { + if (serializedState.ValueKind != JsonValueKind.Object) + { + throw new ArgumentException("The serialized session state must be a JSON object.", nameof(serializedState)); + } + + var jso = jsonSerializerOptions ?? CopilotStudioJsonUtilities.DefaultOptions; + return serializedState.Deserialize(jso.GetTypeInfo(typeof(CopilotStudioAgentSession))) as CopilotStudioAgentSession + ?? new CopilotStudioAgentSession(); + } + + [DebuggerBrowsable(DebuggerBrowsableState.Never)] + private string DebuggerDisplay => + $"ConversationId = {this.ConversationId}, StateBag Count = {this.StateBag.Count}"; } diff --git a/dotnet/src/Microsoft.Agents.AI.CopilotStudio/CopilotStudioJsonUtilities.cs b/dotnet/src/Microsoft.Agents.AI.CopilotStudio/CopilotStudioJsonUtilities.cs new file mode 100644 index 0000000000..44177b0708 --- /dev/null +++ b/dotnet/src/Microsoft.Agents.AI.CopilotStudio/CopilotStudioJsonUtilities.cs @@ -0,0 +1,48 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System.Diagnostics.CodeAnalysis; +using System.Text.Encodings.Web; +using System.Text.Json; +using System.Text.Json.Serialization; + +namespace Microsoft.Agents.AI.CopilotStudio; + +/// +/// Provides utility methods and configurations for JSON serialization operations within the Copilot Studio agent implementation. +/// +internal static partial class CopilotStudioJsonUtilities +{ + /// + /// Gets the default instance used for JSON serialization operations. + /// + public static JsonSerializerOptions DefaultOptions { get; } = CreateDefaultOptions(); + + /// + /// Creates and configures the default JSON serialization options. + /// + /// The configured options. + private static JsonSerializerOptions CreateDefaultOptions() + { + // Copy the configuration from the source generated context. + JsonSerializerOptions options = new(JsonContext.Default.Options) + { + Encoder = JavaScriptEncoder.UnsafeRelaxedJsonEscaping, + }; + + // Chain in the resolvers from both AgentAbstractionsJsonUtilities and our source generated context. + options.TypeInfoResolverChain.Clear(); + options.TypeInfoResolverChain.Add(AgentAbstractionsJsonUtilities.DefaultOptions.TypeInfoResolver!); + options.TypeInfoResolverChain.Add(JsonContext.Default.Options.TypeInfoResolver!); + + options.MakeReadOnly(); + return options; + } + + [JsonSourceGenerationOptions(JsonSerializerDefaults.Web, + UseStringEnumConverter = true, + DefaultIgnoreCondition = JsonIgnoreCondition.WhenWritingNull, + NumberHandling = JsonNumberHandling.AllowReadingFromString)] + [JsonSerializable(typeof(CopilotStudioAgentSession))] + [ExcludeFromCodeCoverage] + private sealed partial class JsonContext : JsonSerializerContext; +} diff --git a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatHistoryProvider.cs b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatHistoryProvider.cs index 85c5865f07..265f3a3675 100644 --- a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatHistoryProvider.cs +++ b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatHistoryProvider.cs @@ -21,17 +21,16 @@ namespace Microsoft.Agents.AI; [RequiresDynamicCode("The CosmosChatHistoryProvider uses JSON serialization which is incompatible with NativeAOT.")] public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable { + private static IEnumerable DefaultExcludeChatHistoryFilter(IEnumerable messages) + => messages.Where(m => m.GetAgentRequestMessageSourceType() != AgentRequestMessageSourceType.ChatHistory); + private readonly CosmosClient _cosmosClient; private readonly Container _container; private readonly bool _ownsClient; + private readonly string _stateKey; + private readonly Func _stateInitializer; private bool _disposed; - // Hierarchical partition key support - private readonly string? _tenantId; - private readonly string? _userId; - private readonly PartitionKey _partitionKey; - private readonly bool _useHierarchicalPartitioning; - /// /// Cached JSON serializer options for .NET 9.0 compatibility. /// @@ -47,6 +46,9 @@ public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable return options; } + /// + public override string StateKey => this._stateKey; + /// /// Gets or sets the maximum number of messages to return in a single query batch. /// Default is 100 for optimal performance. @@ -72,11 +74,6 @@ public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable /// public int? MessageTtlSeconds { get; set; } = 86400; - /// - /// Gets the conversation ID associated with this provider. - /// - public string ConversationId { get; init; } - /// /// Gets the database ID associated with this provider. /// @@ -88,36 +85,50 @@ public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable public string ContainerId { get; init; } /// - /// Internal primary constructor used by all public constructors. + /// A filter function applied to request messages before they are stored + /// during . The default filter excludes messages with the + /// source type. + /// + public Func, IEnumerable> StorageInputMessageFilter { get; set { field = Throw.IfNull(value); } } = DefaultExcludeChatHistoryFilter; + + /// + /// Gets or sets an optional filter function applied to messages produced by this provider + /// during . + /// + /// + /// This filter is only applied to the messages that the provider itself produces (from its internal storage). + /// + /// + /// When , no filtering is applied to the output messages. + /// + public Func, IEnumerable>? RetrievalOutputMessageFilter { get; set; } + + /// + /// Initializes a new instance of the class. /// /// The instance to use for Cosmos DB operations. /// The identifier of the Cosmos DB database. /// The identifier of the Cosmos DB container. - /// The unique identifier for this conversation thread. + /// A delegate that initializes the provider state on the first invocation, providing the conversation routing info (conversationId, tenantId, userId). /// Whether this instance owns the CosmosClient and should dispose it. - /// Optional tenant identifier for hierarchical partitioning. - /// Optional user identifier for hierarchical partitioning. - internal CosmosChatHistoryProvider(CosmosClient cosmosClient, string databaseId, string containerId, string conversationId, bool ownsClient, string? tenantId = null, string? userId = null) + /// An optional key to use for storing the state in the . + /// Thrown when or is . + /// Thrown when any string parameter is null or whitespace. + public CosmosChatHistoryProvider( + CosmosClient cosmosClient, + string databaseId, + string containerId, + Func stateInitializer, + bool ownsClient = false, + string? stateKey = null) { this._cosmosClient = Throw.IfNull(cosmosClient); - this._container = this._cosmosClient.GetContainer(Throw.IfNullOrWhitespace(databaseId), Throw.IfNullOrWhitespace(containerId)); - this.ConversationId = Throw.IfNullOrWhitespace(conversationId); - this.DatabaseId = databaseId; - this.ContainerId = containerId; + this.DatabaseId = Throw.IfNullOrWhitespace(databaseId); + this.ContainerId = Throw.IfNullOrWhitespace(containerId); + this._container = this._cosmosClient.GetContainer(databaseId, containerId); + this._stateInitializer = Throw.IfNull(stateInitializer); this._ownsClient = ownsClient; - - // Initialize partitioning mode - this._tenantId = tenantId; - this._userId = userId; - this._useHierarchicalPartitioning = tenantId != null && userId != null; - - this._partitionKey = this._useHierarchicalPartitioning - ? new PartitionKeyBuilder() - .Add(tenantId!) - .Add(userId!) - .Add(conversationId) - .Build() - : new PartitionKey(conversationId); + this._stateKey = stateKey ?? base.StateKey; } /// @@ -126,24 +137,17 @@ public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable /// The Cosmos DB connection string. /// The identifier of the Cosmos DB database. /// The identifier of the Cosmos DB container. + /// A delegate that initializes the provider state on the first invocation. + /// An optional key to use for storing the state in the . /// Thrown when any required parameter is null. /// Thrown when any string parameter is null or whitespace. - public CosmosChatHistoryProvider(string connectionString, string databaseId, string containerId) - : this(connectionString, databaseId, containerId, Guid.NewGuid().ToString("N")) - { - } - - /// - /// Initializes a new instance of the class using a connection string. - /// - /// The Cosmos DB connection string. - /// The identifier of the Cosmos DB database. - /// The identifier of the Cosmos DB container. - /// The unique identifier for this conversation thread. - /// Thrown when any required parameter is null. - /// Thrown when any string parameter is null or whitespace. - public CosmosChatHistoryProvider(string connectionString, string databaseId, string containerId, string conversationId) - : this(new CosmosClient(Throw.IfNullOrWhitespace(connectionString)), databaseId, containerId, conversationId, ownsClient: true) + public CosmosChatHistoryProvider( + string connectionString, + string databaseId, + string containerId, + Func stateInitializer, + string? stateKey = null) + : this(new CosmosClient(Throw.IfNullOrWhitespace(connectionString)), databaseId, containerId, stateInitializer, ownsClient: true, stateKey) { } @@ -154,136 +158,63 @@ public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable /// The TokenCredential to use for authentication (e.g., DefaultAzureCredential, ManagedIdentityCredential). /// The identifier of the Cosmos DB database. /// The identifier of the Cosmos DB container. + /// A delegate that initializes the provider state on the first invocation. + /// An optional key to use for storing the state in the . /// Thrown when any required parameter is null. /// Thrown when any string parameter is null or whitespace. - public CosmosChatHistoryProvider(string accountEndpoint, TokenCredential tokenCredential, string databaseId, string containerId) - : this(accountEndpoint, tokenCredential, databaseId, containerId, Guid.NewGuid().ToString("N")) + public CosmosChatHistoryProvider( + string accountEndpoint, + TokenCredential tokenCredential, + string databaseId, + string containerId, + Func stateInitializer, + string? stateKey = null) + : this(new CosmosClient(Throw.IfNullOrWhitespace(accountEndpoint), Throw.IfNull(tokenCredential)), databaseId, containerId, stateInitializer, ownsClient: true, stateKey) { } /// - /// Initializes a new instance of the class using a TokenCredential for authentication. + /// Gets the state from the session's StateBag, or initializes it using the state initializer if not present. /// - /// The Cosmos DB account endpoint URI. - /// The TokenCredential to use for authentication (e.g., DefaultAzureCredential, ManagedIdentityCredential). - /// The identifier of the Cosmos DB database. - /// The identifier of the Cosmos DB container. - /// The unique identifier for this conversation thread. - /// Thrown when any required parameter is null. - /// Thrown when any string parameter is null or whitespace. - public CosmosChatHistoryProvider(string accountEndpoint, TokenCredential tokenCredential, string databaseId, string containerId, string conversationId) - : this(new CosmosClient(Throw.IfNullOrWhitespace(accountEndpoint), Throw.IfNull(tokenCredential)), databaseId, containerId, conversationId, ownsClient: true) + /// The agent session containing the StateBag. + /// The provider state, or null if no session is available. + private State GetOrInitializeState(AgentSession? session) { - } - - /// - /// Initializes a new instance of the class using an existing . - /// - /// The instance to use for Cosmos DB operations. - /// The identifier of the Cosmos DB database. - /// The identifier of the Cosmos DB container. - /// Thrown when is null. - /// Thrown when any string parameter is null or whitespace. - public CosmosChatHistoryProvider(CosmosClient cosmosClient, string databaseId, string containerId) - : this(cosmosClient, databaseId, containerId, Guid.NewGuid().ToString("N")) - { - } - - /// - /// Initializes a new instance of the class using an existing . - /// - /// The instance to use for Cosmos DB operations. - /// The identifier of the Cosmos DB database. - /// The identifier of the Cosmos DB container. - /// The unique identifier for this conversation thread. - /// Thrown when is null. - /// Thrown when any string parameter is null or whitespace. - public CosmosChatHistoryProvider(CosmosClient cosmosClient, string databaseId, string containerId, string conversationId) - : this(cosmosClient, databaseId, containerId, conversationId, ownsClient: false) - { - } - - /// - /// Initializes a new instance of the class using a connection string with hierarchical partition keys. - /// - /// The Cosmos DB connection string. - /// The identifier of the Cosmos DB database. - /// The identifier of the Cosmos DB container. - /// The tenant identifier for hierarchical partitioning. - /// The user identifier for hierarchical partitioning. - /// The session identifier for hierarchical partitioning. - /// Thrown when any required parameter is null. - /// Thrown when any string parameter is null or whitespace. - public CosmosChatHistoryProvider(string connectionString, string databaseId, string containerId, string tenantId, string userId, string sessionId) - : this(new CosmosClient(Throw.IfNullOrWhitespace(connectionString)), databaseId, containerId, Throw.IfNullOrWhitespace(sessionId), ownsClient: true, Throw.IfNullOrWhitespace(tenantId), Throw.IfNullOrWhitespace(userId)) - { - } - - /// - /// Initializes a new instance of the class using a TokenCredential for authentication with hierarchical partition keys. - /// - /// The Cosmos DB account endpoint URI. - /// The TokenCredential to use for authentication (e.g., DefaultAzureCredential, ManagedIdentityCredential). - /// The identifier of the Cosmos DB database. - /// The identifier of the Cosmos DB container. - /// The tenant identifier for hierarchical partitioning. - /// The user identifier for hierarchical partitioning. - /// The session identifier for hierarchical partitioning. - /// Thrown when any required parameter is null. - /// Thrown when any string parameter is null or whitespace. - public CosmosChatHistoryProvider(string accountEndpoint, TokenCredential tokenCredential, string databaseId, string containerId, string tenantId, string userId, string sessionId) - : this(new CosmosClient(Throw.IfNullOrWhitespace(accountEndpoint), Throw.IfNull(tokenCredential)), databaseId, containerId, Throw.IfNullOrWhitespace(sessionId), ownsClient: true, Throw.IfNullOrWhitespace(tenantId), Throw.IfNullOrWhitespace(userId)) - { - } - - /// - /// Initializes a new instance of the class using an existing with hierarchical partition keys. - /// - /// The instance to use for Cosmos DB operations. - /// The identifier of the Cosmos DB database. - /// The identifier of the Cosmos DB container. - /// The tenant identifier for hierarchical partitioning. - /// The user identifier for hierarchical partitioning. - /// The session identifier for hierarchical partitioning. - /// Thrown when is null. - /// Thrown when any string parameter is null or whitespace. - public CosmosChatHistoryProvider(CosmosClient cosmosClient, string databaseId, string containerId, string tenantId, string userId, string sessionId) - : this(cosmosClient, databaseId, containerId, Throw.IfNullOrWhitespace(sessionId), ownsClient: false, Throw.IfNullOrWhitespace(tenantId), Throw.IfNullOrWhitespace(userId)) - { - } - - /// - /// Creates a new instance of the class from previously serialized state. - /// - /// The instance to use for Cosmos DB operations. - /// A representing the serialized state of the provider. - /// The identifier of the Cosmos DB database. - /// The identifier of the Cosmos DB container. - /// Optional settings for customizing the JSON deserialization process. - /// A new instance of initialized from the serialized state. - /// Thrown when is null. - /// Thrown when the serialized state cannot be deserialized. - public static CosmosChatHistoryProvider CreateFromSerializedState(CosmosClient cosmosClient, JsonElement serializedState, string databaseId, string containerId, JsonSerializerOptions? jsonSerializerOptions = null) - { - Throw.IfNull(cosmosClient); - Throw.IfNullOrWhitespace(databaseId); - Throw.IfNullOrWhitespace(containerId); - - if (serializedState.ValueKind is not JsonValueKind.Object) + if (session?.StateBag.TryGetValue(this._stateKey, out var state, AgentAbstractionsJsonUtilities.DefaultOptions) is true && state is not null) { - throw new ArgumentException("Invalid serialized state", nameof(serializedState)); + return state; } - var state = serializedState.Deserialize(jsonSerializerOptions); - if (state?.ConversationIdentifier is not { } conversationId) + state = this._stateInitializer(session); + if (session is not null) { - throw new ArgumentException("Invalid serialized state", nameof(serializedState)); + session.StateBag.SetValue(this._stateKey, state, AgentAbstractionsJsonUtilities.DefaultOptions); } - // Use the internal constructor with all parameters to ensure partition key logic is centralized - return state.UseHierarchicalPartitioning && state.TenantId != null && state.UserId != null - ? new CosmosChatHistoryProvider(cosmosClient, databaseId, containerId, conversationId, ownsClient: false, state.TenantId, state.UserId) - : new CosmosChatHistoryProvider(cosmosClient, databaseId, containerId, conversationId, ownsClient: false); + return state; + } + + /// + /// Determines whether hierarchical partitioning should be used based on the state. + /// + private static bool UseHierarchicalPartitioning(State state) => + state.TenantId is not null && state.UserId is not null; + + /// + /// Builds the partition key from the state. + /// + private static PartitionKey BuildPartitionKey(State state) + { + if (UseHierarchicalPartitioning(state)) + { + return new PartitionKeyBuilder() + .Add(state.TenantId) + .Add(state.UserId) + .Add(state.ConversationId) + .Build(); + } + + return new PartitionKey(state.ConversationId); } /// @@ -296,15 +227,20 @@ public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable } #pragma warning restore CA1513 + _ = Throw.IfNull(context); + + var state = this.GetOrInitializeState(context.Session); + var partitionKey = BuildPartitionKey(state); + // Fetch most recent messages in descending order when limit is set, then reverse to ascending var orderDirection = this.MaxMessagesToRetrieve.HasValue ? "DESC" : "ASC"; var query = new QueryDefinition($"SELECT * FROM c WHERE c.conversationId = @conversationId AND c.type = @type ORDER BY c.timestamp {orderDirection}") - .WithParameter("@conversationId", this.ConversationId) + .WithParameter("@conversationId", state.ConversationId) .WithParameter("@type", "ChatMessage"); var iterator = this._container.GetItemQueryIterator(query, requestOptions: new QueryRequestOptions { - PartitionKey = this._partitionKey, + PartitionKey = partitionKey, MaxItemCount = this.MaxItemCount // Configurable query performance }); @@ -343,7 +279,9 @@ public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable messages.Reverse(); } - return messages; + return (this.RetrievalOutputMessageFilter is not null ? this.RetrievalOutputMessageFilter(messages) : messages) + .Select(message => message.WithAgentRequestMessageSource(AgentRequestMessageSourceType.ChatHistory, this.GetType().FullName!)) + .Concat(context.RequestMessages); } /// @@ -364,27 +302,30 @@ public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable } #pragma warning restore CA1513 - var messageList = context.RequestMessages.Concat(context.ResponseMessages ?? []).ToList(); + var state = this.GetOrInitializeState(context.Session); + var messageList = this.StorageInputMessageFilter(context.RequestMessages).Concat(context.ResponseMessages ?? []).ToList(); if (messageList.Count == 0) { return; } + var partitionKey = BuildPartitionKey(state); + // Use transactional batch for atomic operations if (messageList.Count > 1) { - await this.AddMessagesInBatchAsync(messageList, cancellationToken).ConfigureAwait(false); + await this.AddMessagesInBatchAsync(partitionKey, state, messageList, cancellationToken).ConfigureAwait(false); } else { - await this.AddSingleMessageAsync(messageList.First(), cancellationToken).ConfigureAwait(false); + await this.AddSingleMessageAsync(partitionKey, state, messageList.First(), cancellationToken).ConfigureAwait(false); } } /// /// Adds multiple messages using transactional batch operations for atomicity. /// - private async Task AddMessagesInBatchAsync(List messages, CancellationToken cancellationToken) + private async Task AddMessagesInBatchAsync(PartitionKey partitionKey, State state, List messages, CancellationToken cancellationToken) { var currentTimestamp = DateTimeOffset.UtcNow.ToUnixTimeSeconds(); @@ -392,7 +333,7 @@ public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable for (int i = 0; i < messages.Count; i += this.MaxBatchSize) { var batchMessages = messages.Skip(i).Take(this.MaxBatchSize).ToList(); - await this.ExecuteBatchOperationAsync(batchMessages, currentTimestamp, cancellationToken).ConfigureAwait(false); + await this.ExecuteBatchOperationAsync(partitionKey, state, batchMessages, currentTimestamp, cancellationToken).ConfigureAwait(false); } } @@ -400,13 +341,13 @@ public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable /// Executes a single batch operation with enhanced error handling. /// Cosmos SDK handles throttling (429) retries automatically. /// - private async Task ExecuteBatchOperationAsync(List messages, long timestamp, CancellationToken cancellationToken) + private async Task ExecuteBatchOperationAsync(PartitionKey partitionKey, State state, List messages, long timestamp, CancellationToken cancellationToken) { // Create all documents upfront for validation and batch operation var documents = new List(messages.Count); foreach (var message in messages) { - documents.Add(this.CreateMessageDocument(message, timestamp)); + documents.Add(this.CreateMessageDocument(state, message, timestamp)); } // Defensive check: Verify all messages share the same partition key values @@ -414,7 +355,7 @@ public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable // In simple partitioning, this means same conversationId if (documents.Count > 0) { - if (this._useHierarchicalPartitioning) + if (UseHierarchicalPartitioning(state)) { // Verify all documents have matching hierarchical partition key components var firstDoc = documents[0]; @@ -436,7 +377,7 @@ public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable // All messages in this store share the same partition key by design // Transactional batches require all items to share the same partition key - var batch = this._container.CreateTransactionalBatch(this._partitionKey); + var batch = this._container.CreateTransactionalBatch(partitionKey); foreach (var document in documents) { @@ -457,7 +398,7 @@ public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable if (messages.Count == 1) { // Can't split further, use single operation - await this.AddSingleMessageAsync(messages[0], cancellationToken).ConfigureAwait(false); + await this.AddSingleMessageAsync(partitionKey, state, messages[0], cancellationToken).ConfigureAwait(false); return; } @@ -466,21 +407,21 @@ public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable var firstHalf = messages.Take(midpoint).ToList(); var secondHalf = messages.Skip(midpoint).ToList(); - await this.ExecuteBatchOperationAsync(firstHalf, timestamp, cancellationToken).ConfigureAwait(false); - await this.ExecuteBatchOperationAsync(secondHalf, timestamp, cancellationToken).ConfigureAwait(false); + await this.ExecuteBatchOperationAsync(partitionKey, state, firstHalf, timestamp, cancellationToken).ConfigureAwait(false); + await this.ExecuteBatchOperationAsync(partitionKey, state, secondHalf, timestamp, cancellationToken).ConfigureAwait(false); } } /// /// Adds a single message to the store. /// - private async Task AddSingleMessageAsync(ChatMessage message, CancellationToken cancellationToken) + private async Task AddSingleMessageAsync(PartitionKey partitionKey, State state, ChatMessage message, CancellationToken cancellationToken) { - var document = this.CreateMessageDocument(message, DateTimeOffset.UtcNow.ToUnixTimeSeconds()); + var document = this.CreateMessageDocument(state, message, DateTimeOffset.UtcNow.ToUnixTimeSeconds()); try { - await this._container.CreateItemAsync(document, this._partitionKey, cancellationToken: cancellationToken).ConfigureAwait(false); + await this._container.CreateItemAsync(document, partitionKey, cancellationToken: cancellationToken).ConfigureAwait(false); } catch (CosmosException ex) when (ex.StatusCode == System.Net.HttpStatusCode.RequestEntityTooLarge) { @@ -495,12 +436,14 @@ public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable /// /// Creates a message document with enhanced metadata. /// - private CosmosMessageDocument CreateMessageDocument(ChatMessage message, long timestamp) + private CosmosMessageDocument CreateMessageDocument(State state, ChatMessage message, long timestamp) { + var useHierarchical = UseHierarchicalPartitioning(state); + return new CosmosMessageDocument { Id = Guid.NewGuid().ToString(), - ConversationId = this.ConversationId, + ConversationId = state.ConversationId, Timestamp = timestamp, MessageId = message.MessageId, Role = message.Role.Value, @@ -508,41 +451,20 @@ public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable Type = "ChatMessage", // Type discriminator Ttl = this.MessageTtlSeconds, // Configurable TTL // Include hierarchical metadata when using hierarchical partitioning - TenantId = this._useHierarchicalPartitioning ? this._tenantId : null, - UserId = this._useHierarchicalPartitioning ? this._userId : null, - SessionId = this._useHierarchicalPartitioning ? this.ConversationId : null + TenantId = useHierarchical ? state.TenantId : null, + UserId = useHierarchical ? state.UserId : null, + SessionId = useHierarchical ? state.ConversationId : null }; } - /// - public override JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null) - { -#pragma warning disable CA1513 // Use ObjectDisposedException.ThrowIf - not available on all target frameworks - if (this._disposed) - { - throw new ObjectDisposedException(this.GetType().FullName); - } -#pragma warning restore CA1513 - - var state = new State - { - ConversationIdentifier = this.ConversationId, - TenantId = this._tenantId, - UserId = this._userId, - UseHierarchicalPartitioning = this._useHierarchicalPartitioning - }; - - var options = jsonSerializerOptions ?? s_defaultJsonOptions; - return JsonSerializer.SerializeToElement(state, options); - } - /// /// Gets the count of messages in this conversation. /// This is an additional utility method beyond the base contract. /// + /// The agent session to get state from. /// The cancellation token. /// The number of messages in the conversation. - public async Task GetMessageCountAsync(CancellationToken cancellationToken = default) + public async Task GetMessageCountAsync(AgentSession? session, CancellationToken cancellationToken = default) { #pragma warning disable CA1513 // Use ObjectDisposedException.ThrowIf - not available on all target frameworks if (this._disposed) @@ -551,14 +473,17 @@ public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable } #pragma warning restore CA1513 + var state = this.GetOrInitializeState(session); + var partitionKey = BuildPartitionKey(state); + // Efficient count query var query = new QueryDefinition("SELECT VALUE COUNT(1) FROM c WHERE c.conversationId = @conversationId AND c.Type = @type") - .WithParameter("@conversationId", this.ConversationId) + .WithParameter("@conversationId", state.ConversationId) .WithParameter("@type", "ChatMessage"); var iterator = this._container.GetItemQueryIterator(query, requestOptions: new QueryRequestOptions { - PartitionKey = this._partitionKey + PartitionKey = partitionKey }); // COUNT queries always return a result @@ -570,9 +495,10 @@ public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable /// Deletes all messages in this conversation. /// This is an additional utility method beyond the base contract. /// + /// The agent session to get state from. /// The cancellation token. /// The number of messages deleted. - public async Task ClearMessagesAsync(CancellationToken cancellationToken = default) + public async Task ClearMessagesAsync(AgentSession? session, CancellationToken cancellationToken = default) { #pragma warning disable CA1513 // Use ObjectDisposedException.ThrowIf - not available on all target frameworks if (this._disposed) @@ -581,14 +507,17 @@ public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable } #pragma warning restore CA1513 + var state = this.GetOrInitializeState(session); + var partitionKey = BuildPartitionKey(state); + // Batch delete for efficiency var query = new QueryDefinition("SELECT VALUE c.id FROM c WHERE c.conversationId = @conversationId AND c.Type = @type") - .WithParameter("@conversationId", this.ConversationId) + .WithParameter("@conversationId", state.ConversationId) .WithParameter("@type", "ChatMessage"); var iterator = this._container.GetItemQueryIterator(query, requestOptions: new QueryRequestOptions { - PartitionKey = this._partitionKey, + PartitionKey = partitionKey, MaxItemCount = this.MaxItemCount }); @@ -597,7 +526,7 @@ public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable while (iterator.HasMoreResults) { var response = await iterator.ReadNextAsync(cancellationToken).ConfigureAwait(false); - var batch = this._container.CreateTransactionalBatch(this._partitionKey); + var batch = this._container.CreateTransactionalBatch(partitionKey); var batchItemCount = 0; foreach (var itemId in response) @@ -632,12 +561,38 @@ public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable } } - private sealed class State + /// + /// Represents the per-session state of a stored in the . + /// + public sealed class State { - public string ConversationIdentifier { get; set; } = string.Empty; - public string? TenantId { get; set; } - public string? UserId { get; set; } - public bool UseHierarchicalPartitioning { get; set; } + /// + /// Initializes a new instance of the class. + /// + /// The unique identifier for this conversation thread. + /// Optional tenant identifier for hierarchical partitioning. + /// Optional user identifier for hierarchical partitioning. + public State(string conversationId, string? tenantId = null, string? userId = null) + { + this.ConversationId = Throw.IfNullOrWhitespace(conversationId); + this.TenantId = tenantId; + this.UserId = userId; + } + + /// + /// Gets the conversation ID associated with this state. + /// + public string ConversationId { get; } + + /// + /// Gets the tenant identifier for hierarchical partitioning, if any. + /// + public string? TenantId { get; } + + /// + /// Gets the user identifier for hierarchical partitioning, if any. + /// + public string? UserId { get; } } /// diff --git a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosDBChatExtensions.cs b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosDBChatExtensions.cs index 3d93e9dd6a..76b865e4c8 100644 --- a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosDBChatExtensions.cs +++ b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosDBChatExtensions.cs @@ -2,7 +2,6 @@ using System; using System.Diagnostics.CodeAnalysis; -using System.Threading.Tasks; using Azure.Core; using Microsoft.Azure.Cosmos; @@ -13,6 +12,9 @@ namespace Microsoft.Agents.AI; /// public static class CosmosDBChatExtensions { + private static readonly Func s_defaultStateInitializer = + _ => new CosmosChatHistoryProvider.State(Guid.NewGuid().ToString("N")); + /// /// Configures the agent to use Cosmos DB for message storage with connection string authentication. /// @@ -20,6 +22,7 @@ public static class CosmosDBChatExtensions /// The Cosmos DB connection string. /// The identifier of the Cosmos DB database. /// The identifier of the Cosmos DB container. + /// An optional delegate that initializes the provider state on the first invocation, providing the conversation routing info (conversationId, tenantId, userId). When not provided, a new conversation ID is generated automatically. /// The configured . /// Thrown when is null. /// Thrown when any string parameter is null or whitespace. @@ -29,14 +32,16 @@ public static class CosmosDBChatExtensions this ChatClientAgentOptions options, string connectionString, string databaseId, - string containerId) + string containerId, + Func? stateInitializer = null) { if (options is null) { throw new ArgumentNullException(nameof(options)); } - options.ChatHistoryProviderFactory = (context, ct) => new ValueTask(new CosmosChatHistoryProvider(connectionString, databaseId, containerId)); + options.ChatHistoryProvider = + new CosmosChatHistoryProvider(connectionString, databaseId, containerId, stateInitializer ?? s_defaultStateInitializer); return options; } @@ -48,6 +53,7 @@ public static class CosmosDBChatExtensions /// The identifier of the Cosmos DB database. /// The identifier of the Cosmos DB container. /// The TokenCredential to use for authentication (e.g., DefaultAzureCredential, ManagedIdentityCredential). + /// An optional delegate that initializes the provider state on the first invocation, providing the conversation routing info (conversationId, tenantId, userId). When not provided, a new conversation ID is generated automatically. /// The configured . /// Thrown when or is null. /// Thrown when any string parameter is null or whitespace. @@ -58,7 +64,8 @@ public static class CosmosDBChatExtensions string accountEndpoint, string databaseId, string containerId, - TokenCredential tokenCredential) + TokenCredential tokenCredential, + Func? stateInitializer = null) { if (options is null) { @@ -70,7 +77,8 @@ public static class CosmosDBChatExtensions throw new ArgumentNullException(nameof(tokenCredential)); } - options.ChatHistoryProviderFactory = (context, ct) => new ValueTask(new CosmosChatHistoryProvider(accountEndpoint, tokenCredential, databaseId, containerId)); + options.ChatHistoryProvider = + new CosmosChatHistoryProvider(accountEndpoint, tokenCredential, databaseId, containerId, stateInitializer ?? s_defaultStateInitializer); return options; } @@ -81,6 +89,7 @@ public static class CosmosDBChatExtensions /// The instance to use for Cosmos DB operations. /// The identifier of the Cosmos DB database. /// The identifier of the Cosmos DB container. + /// An optional delegate that initializes the provider state on the first invocation, providing the conversation routing info (conversationId, tenantId, userId). When not provided, a new conversation ID is generated automatically. /// The configured . /// Thrown when any required parameter is null. /// Thrown when any string parameter is null or whitespace. @@ -90,14 +99,16 @@ public static class CosmosDBChatExtensions this ChatClientAgentOptions options, CosmosClient cosmosClient, string databaseId, - string containerId) + string containerId, + Func? stateInitializer = null) { if (options is null) { throw new ArgumentNullException(nameof(options)); } - options.ChatHistoryProviderFactory = (context, ct) => new ValueTask(new CosmosChatHistoryProvider(cosmosClient, databaseId, containerId)); + options.ChatHistoryProvider = + new CosmosChatHistoryProvider(cosmosClient, databaseId, containerId, stateInitializer ?? s_defaultStateInitializer); return options; } } diff --git a/dotnet/src/Microsoft.Agents.AI.DurableTask/CHANGELOG.md b/dotnet/src/Microsoft.Agents.AI.DurableTask/CHANGELOG.md index d8260fcb84..ff886e2ebe 100644 --- a/dotnet/src/Microsoft.Agents.AI.DurableTask/CHANGELOG.md +++ b/dotnet/src/Microsoft.Agents.AI.DurableTask/CHANGELOG.md @@ -12,6 +12,7 @@ - Renamed serializedSession parameter to serializedState on DeserializeSessionAsync for consistency ([#3681](https://github.com/microsoft/agent-framework/pull/3681)) - Introduce Core method pattern for Session management methods on AIAgent ([#3699](https://github.com/microsoft/agent-framework/pull/3699)) - Changed AIAgent.SerializeSession to AIAgent.SerializeSessionAsync ([#3879](https://github.com/microsoft/agent-framework/pull/3879)) +- Changed ChatHistory and AIContext Providers to have pipeline semantics ([#3806](https://github.com/microsoft/agent-framework/pull/3806)) ## v1.0.0-preview.251204.1 diff --git a/dotnet/src/Microsoft.Agents.AI.DurableTask/DurableAgentSession.cs b/dotnet/src/Microsoft.Agents.AI.DurableTask/DurableAgentSession.cs index b9d9807728..ba33c15d32 100644 --- a/dotnet/src/Microsoft.Agents.AI.DurableTask/DurableAgentSession.cs +++ b/dotnet/src/Microsoft.Agents.AI.DurableTask/DurableAgentSession.cs @@ -7,17 +7,22 @@ using System.Text.Json.Serialization; namespace Microsoft.Agents.AI.DurableTask; /// -/// An agent thread implementation for durable agents. +/// An implementation for durable agents. /// -[DebuggerDisplay("{SessionId}")] +[DebuggerDisplay("{DebuggerDisplay,nq}")] public sealed class DurableAgentSession : AgentSession { - [JsonConstructor] internal DurableAgentSession(AgentSessionId sessionId) { this.SessionId = sessionId; } + [JsonConstructor] + internal DurableAgentSession(AgentSessionId sessionId, AgentSessionStateBag stateBag) : base(stateBag) + { + this.SessionId = sessionId; + } + /// /// Gets the agent session ID. /// @@ -28,9 +33,8 @@ public sealed class DurableAgentSession : AgentSession /// internal JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null) { - return JsonSerializer.SerializeToElement( - this, - DurableAgentJsonUtilities.DefaultOptions.GetTypeInfo(typeof(DurableAgentSession))); + var jso = jsonSerializerOptions ?? DurableAgentJsonUtilities.DefaultOptions; + return JsonSerializer.SerializeToElement(this, jso.GetTypeInfo(typeof(DurableAgentSession))); } /// @@ -49,7 +53,11 @@ public sealed class DurableAgentSession : AgentSession string sessionIdString = sessionIdElement.GetString() ?? throw new JsonException("sessionId property is null."); AgentSessionId sessionId = AgentSessionId.Parse(sessionIdString); - return new DurableAgentSession(sessionId); + AgentSessionStateBag stateBag = serializedSession.TryGetProperty("stateBag", out JsonElement stateBagElement) + ? AgentSessionStateBag.Deserialize(stateBagElement) + : new AgentSessionStateBag(); + + return new DurableAgentSession(sessionId, stateBag); } /// @@ -68,4 +76,8 @@ public sealed class DurableAgentSession : AgentSession { return this.SessionId.ToString(); } + + [DebuggerBrowsable(DebuggerBrowsableState.Never)] + private string DebuggerDisplay => + $"SessionId = {this.SessionId}, StateBag Count = {this.StateBag.Count}"; } diff --git a/dotnet/src/Microsoft.Agents.AI.GitHub.Copilot/GitHubCopilotAgent.cs b/dotnet/src/Microsoft.Agents.AI.GitHub.Copilot/GitHubCopilotAgent.cs index 8319832613..f0533c8461 100644 --- a/dotnet/src/Microsoft.Agents.AI.GitHub.Copilot/GitHubCopilotAgent.cs +++ b/dotnet/src/Microsoft.Agents.AI.GitHub.Copilot/GitHubCopilotAgent.cs @@ -115,7 +115,7 @@ public sealed class GitHubCopilotAgent : AIAgent, IAsyncDisposable JsonElement serializedState, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default) - => new(new GitHubCopilotAgentSession(serializedState, jsonSerializerOptions)); + => new(GitHubCopilotAgentSession.Deserialize(serializedState, jsonSerializerOptions)); /// protected override Task RunCoreAsync( diff --git a/dotnet/src/Microsoft.Agents.AI.GitHub.Copilot/GitHubCopilotAgentSession.cs b/dotnet/src/Microsoft.Agents.AI.GitHub.Copilot/GitHubCopilotAgentSession.cs index f514eeb71b..70fe43425e 100644 --- a/dotnet/src/Microsoft.Agents.AI.GitHub.Copilot/GitHubCopilotAgentSession.cs +++ b/dotnet/src/Microsoft.Agents.AI.GitHub.Copilot/GitHubCopilotAgentSession.cs @@ -1,17 +1,22 @@ // Copyright (c) Microsoft. All rights reserved. +using System; +using System.Diagnostics; using System.Text.Json; +using System.Text.Json.Serialization; namespace Microsoft.Agents.AI.GitHub.Copilot; /// /// Represents a session for a GitHub Copilot agent conversation. /// +[DebuggerDisplay("{DebuggerDisplay,nq}")] public sealed class GitHubCopilotAgentSession : AgentSession { /// /// Gets or sets the session ID for the GitHub Copilot conversation. /// + [JsonPropertyName("sessionId")] public string? SessionId { get; internal set; } /// @@ -21,35 +26,32 @@ public sealed class GitHubCopilotAgentSession : AgentSession { } - /// - /// Initializes a new instance of the class from serialized data. - /// - /// The serialized thread data. - /// Optional JSON serialization options. - internal GitHubCopilotAgentSession(JsonElement serializedThread, JsonSerializerOptions? jsonSerializerOptions = null) + [JsonConstructor] + internal GitHubCopilotAgentSession(string? sessionId, AgentSessionStateBag? stateBag) : base(stateBag ?? new()) { - // The JSON serialization uses camelCase - if (serializedThread.TryGetProperty("sessionId", out JsonElement sessionIdElement)) - { - this.SessionId = sessionIdElement.GetString(); - } + this.SessionId = sessionId; } /// internal JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null) { - State state = new() - { - SessionId = this.SessionId - }; - - return JsonSerializer.SerializeToElement( - state, - GitHubCopilotJsonUtilities.DefaultOptions.GetTypeInfo(typeof(State))); + var jso = jsonSerializerOptions ?? GitHubCopilotJsonUtilities.DefaultOptions; + return JsonSerializer.SerializeToElement(this, jso.GetTypeInfo(typeof(GitHubCopilotAgentSession))); } - internal sealed class State + internal static GitHubCopilotAgentSession Deserialize(JsonElement serializedState, JsonSerializerOptions? jsonSerializerOptions = null) { - public string? SessionId { get; set; } + if (serializedState.ValueKind != JsonValueKind.Object) + { + throw new ArgumentException("The serialized session state must be a JSON object.", nameof(serializedState)); + } + + var jso = jsonSerializerOptions ?? GitHubCopilotJsonUtilities.DefaultOptions; + return serializedState.Deserialize(jso.GetTypeInfo(typeof(GitHubCopilotAgentSession))) as GitHubCopilotAgentSession + ?? new GitHubCopilotAgentSession(); } + + [DebuggerBrowsable(DebuggerBrowsableState.Never)] + private string DebuggerDisplay => + $"SessionId = {this.SessionId}, StateBag Count = {this.StateBag.Count}"; } diff --git a/dotnet/src/Microsoft.Agents.AI.GitHub.Copilot/GitHubCopilotJsonUtilities.cs b/dotnet/src/Microsoft.Agents.AI.GitHub.Copilot/GitHubCopilotJsonUtilities.cs index c5254efd6b..9e97c0585b 100644 --- a/dotnet/src/Microsoft.Agents.AI.GitHub.Copilot/GitHubCopilotJsonUtilities.cs +++ b/dotnet/src/Microsoft.Agents.AI.GitHub.Copilot/GitHubCopilotJsonUtilities.cs @@ -42,7 +42,7 @@ internal static partial class GitHubCopilotJsonUtilities UseStringEnumConverter = true, DefaultIgnoreCondition = JsonIgnoreCondition.WhenWritingNull, NumberHandling = JsonNumberHandling.AllowReadingFromString)] - [JsonSerializable(typeof(GitHubCopilotAgentSession.State))] + [JsonSerializable(typeof(GitHubCopilotAgentSession))] [ExcludeFromCodeCoverage] private sealed partial class JsonContext : JsonSerializerContext; } diff --git a/dotnet/src/Microsoft.Agents.AI.Mem0/Mem0JsonUtilities.cs b/dotnet/src/Microsoft.Agents.AI.Mem0/Mem0JsonUtilities.cs index d139cb0f76..33f92f3ac2 100644 --- a/dotnet/src/Microsoft.Agents.AI.Mem0/Mem0JsonUtilities.cs +++ b/dotnet/src/Microsoft.Agents.AI.Mem0/Mem0JsonUtilities.cs @@ -65,7 +65,7 @@ public static partial class Mem0JsonUtilities NumberHandling = JsonNumberHandling.AllowReadingFromString)] // Agent abstraction types - [JsonSerializable(typeof(Mem0Provider.Mem0State))] + [JsonSerializable(typeof(Mem0Provider.State))] [ExcludeFromCodeCoverage] internal sealed partial class JsonContext : JsonSerializerContext; diff --git a/dotnet/src/Microsoft.Agents.AI.Mem0/Mem0Provider.cs b/dotnet/src/Microsoft.Agents.AI.Mem0/Mem0Provider.cs index a230eb0e4e..aaf2333553 100644 --- a/dotnet/src/Microsoft.Agents.AI.Mem0/Mem0Provider.cs +++ b/dotnet/src/Microsoft.Agents.AI.Mem0/Mem0Provider.cs @@ -4,7 +4,6 @@ using System; using System.Collections.Generic; using System.Linq; using System.Net.Http; -using System.Text.Json; using System.Text.Json.Serialization; using System.Threading; using System.Threading.Tasks; @@ -27,23 +26,27 @@ public sealed class Mem0Provider : AIContextProvider { private const string DefaultContextPrompt = "## Memories\nConsider the following memories when answering user questions:"; + private static IEnumerable DefaultExternalOnlyFilter(IEnumerable messages) + => messages.Where(m => m.GetAgentRequestMessageSourceType() == AgentRequestMessageSourceType.External); + private readonly string _contextPrompt; private readonly bool _enableSensitiveTelemetryData; + private readonly string _stateKey; + private readonly Func _stateInitializer; + private readonly Func, IEnumerable> _searchInputMessageFilter; + private readonly Func, IEnumerable> _storageInputMessageFilter; private readonly Mem0Client _client; private readonly ILogger? _logger; - private readonly Mem0ProviderScope _storageScope; - private readonly Mem0ProviderScope _searchScope; - /// /// Initializes a new instance of the class. /// /// Configured (base address + auth). - /// Optional values to scope the memory storage with. - /// Optional values to scope the memory search with. Defaults to if not provided. + /// A delegate that initializes the provider state on the first invocation, providing the storage and search scopes. /// Provider options. /// Optional logger factory. + /// Thrown when or is . /// /// The base address of the required mem0 service, and any authentication headers, should be set on the /// already, when passed as a parameter here. E.g.: @@ -51,83 +54,60 @@ public sealed class Mem0Provider : AIContextProvider /// using var httpClient = new HttpClient(); /// httpClient.BaseAddress = new Uri("https://api.mem0.ai"); /// httpClient.DefaultRequestHeaders.Authorization = new AuthenticationHeaderValue("Token", "<Your APIKey>"); - /// new Mem0AIContextProvider(httpClient); + /// new Mem0Provider(httpClient); /// /// - public Mem0Provider(HttpClient httpClient, Mem0ProviderScope storageScope, Mem0ProviderScope? searchScope = null, Mem0ProviderOptions? options = null, ILoggerFactory? loggerFactory = null) + public Mem0Provider(HttpClient httpClient, Func stateInitializer, Mem0ProviderOptions? options = null, ILoggerFactory? loggerFactory = null) { + Throw.IfNull(httpClient); if (string.IsNullOrWhiteSpace(httpClient.BaseAddress?.AbsoluteUri)) { throw new ArgumentException("The HttpClient BaseAddress must be set for Mem0 operations.", nameof(httpClient)); } + this._stateInitializer = Throw.IfNull(stateInitializer); this._logger = loggerFactory?.CreateLogger(); this._client = new Mem0Client(httpClient); this._contextPrompt = options?.ContextPrompt ?? DefaultContextPrompt; this._enableSensitiveTelemetryData = options?.EnableSensitiveTelemetryData ?? false; - this._storageScope = new Mem0ProviderScope(Throw.IfNull(storageScope)); - this._searchScope = searchScope ?? storageScope; - - if (string.IsNullOrWhiteSpace(this._storageScope.ApplicationId) - && string.IsNullOrWhiteSpace(this._storageScope.AgentId) - && string.IsNullOrWhiteSpace(this._storageScope.ThreadId) - && string.IsNullOrWhiteSpace(this._storageScope.UserId)) - { - throw new ArgumentException("At least one of ApplicationId, AgentId, ThreadId, or UserId must be provided for the storage scope."); - } - - if (string.IsNullOrWhiteSpace(this._searchScope.ApplicationId) - && string.IsNullOrWhiteSpace(this._searchScope.AgentId) - && string.IsNullOrWhiteSpace(this._searchScope.ThreadId) - && string.IsNullOrWhiteSpace(this._searchScope.UserId)) - { - throw new ArgumentException("At least one of ApplicationId, AgentId, ThreadId, or UserId must be provided for the search scope."); - } + this._stateKey = options?.StateKey ?? base.StateKey; + this._searchInputMessageFilter = options?.SearchInputMessageFilter ?? DefaultExternalOnlyFilter; + this._storageInputMessageFilter = options?.StorageInputMessageFilter ?? DefaultExternalOnlyFilter; } + /// + public override string StateKey => this._stateKey; + /// - /// Initializes a new instance of the class, with existing state from a serialized JSON element. + /// Gets the state from the session's StateBag, or initializes it using the StateInitializer if not present. /// - /// Configured (base address + auth). - /// A representing the serialized state of the store. - /// Optional settings for customizing the JSON deserialization process. - /// Provider options. - /// Optional logger factory. - /// - /// - /// The base address of the required mem0 service, and any authentication headers, should be set on the - /// already, when passed as a parameter here. E.g.: - /// - /// using var httpClient = new HttpClient(); - /// httpClient.BaseAddress = new Uri("https://api.mem0.ai"); - /// httpClient.DefaultRequestHeaders.Authorization = new AuthenticationHeaderValue("Token", "<Your APIKey>"); - /// new Mem0AIContextProvider(httpClient, state); - /// - /// - public Mem0Provider(HttpClient httpClient, JsonElement serializedState, JsonSerializerOptions? jsonSerializerOptions = null, Mem0ProviderOptions? options = null, ILoggerFactory? loggerFactory = null) + /// The agent session containing the StateBag. + /// The provider state, or null if no session is available. + private State? GetOrInitializeState(AgentSession? session) { - if (string.IsNullOrWhiteSpace(httpClient.BaseAddress?.AbsoluteUri)) + if (session?.StateBag.TryGetValue(this._stateKey, out var state, Mem0JsonUtilities.DefaultOptions) is true && state is not null) { - throw new ArgumentException("The HttpClient BaseAddress must be set for Mem0 operations.", nameof(httpClient)); + return state; } - this._logger = loggerFactory?.CreateLogger(); - this._client = new Mem0Client(httpClient); + state = this._stateInitializer(session); - this._contextPrompt = options?.ContextPrompt ?? DefaultContextPrompt; - this._enableSensitiveTelemetryData = options?.EnableSensitiveTelemetryData ?? false; - - var jso = jsonSerializerOptions ?? Mem0JsonUtilities.DefaultOptions; - var state = serializedState.Deserialize(jso.GetTypeInfo(typeof(Mem0State))) as Mem0State; - - if (state == null || state.StorageScope == null || state.SearchScope == null) + if (state is null + || state.StorageScope is null + || (state.StorageScope.AgentId is null && state.StorageScope.ThreadId is null && state.StorageScope.UserId is null && state.StorageScope.ApplicationId is null) + || state.SearchScope is null + || (state.SearchScope.AgentId is null && state.SearchScope.ThreadId is null && state.SearchScope.UserId is null && state.SearchScope.ApplicationId is null)) { - throw new InvalidOperationException("The Mem0Provider state did not contain the required scope properties."); + throw new InvalidOperationException("State initializer must return a non-null state with valid storage and search scopes, where at lest one scoping parameter is set for each."); } - this._storageScope = state.StorageScope; - this._searchScope = state.SearchScope; + if (session is not null) + { + session.StateBag.SetValue(this._stateKey, state, Mem0JsonUtilities.DefaultOptions); + } + + return state; } /// @@ -135,36 +115,42 @@ public sealed class Mem0Provider : AIContextProvider { Throw.IfNull(context); + var inputContext = context.AIContext; + var state = this.GetOrInitializeState(context.Session); + var searchScope = state?.SearchScope ?? new Mem0ProviderScope(); + string queryText = string.Join( Environment.NewLine, - context.RequestMessages - .Where(m => m.GetAgentRequestMessageSourceType() == AgentRequestMessageSourceType.External) + this._searchInputMessageFilter(inputContext.Messages ?? []) .Where(m => !string.IsNullOrWhiteSpace(m.Text)) .Select(m => m.Text)); try { var memories = (await this._client.SearchAsync( - this._searchScope.ApplicationId, - this._searchScope.AgentId, - this._searchScope.ThreadId, - this._searchScope.UserId, + searchScope.ApplicationId, + searchScope.AgentId, + searchScope.ThreadId, + searchScope.UserId, queryText, cancellationToken).ConfigureAwait(false)).ToList(); var outputMessageText = memories.Count == 0 ? null : $"{this._contextPrompt}\n{string.Join(Environment.NewLine, memories)}"; + var outputMessage = memories.Count == 0 + ? null + : new ChatMessage(ChatRole.User, outputMessageText!).WithAgentRequestMessageSource(AgentRequestMessageSourceType.AIContextProvider, this.GetType().FullName!); if (this._logger?.IsEnabled(LogLevel.Information) is true) { this._logger.LogInformation( "Mem0AIContextProvider: Retrieved {Count} memories. ApplicationId: '{ApplicationId}', AgentId: '{AgentId}', ThreadId: '{ThreadId}', UserId: '{UserId}'.", memories.Count, - this._searchScope.ApplicationId, - this._searchScope.AgentId, - this._searchScope.ThreadId, - this.SanitizeLogData(this._searchScope.UserId)); + searchScope.ApplicationId, + searchScope.AgentId, + searchScope.ThreadId, + this.SanitizeLogData(searchScope.UserId)); if (outputMessageText is not null && this._logger.IsEnabled(LogLevel.Trace)) { @@ -172,16 +158,20 @@ public sealed class Mem0Provider : AIContextProvider "Mem0AIContextProvider: Search Results\nInput:{Input}\nOutput:{MessageText}\nApplicationId: '{ApplicationId}', AgentId: '{AgentId}', ThreadId: '{ThreadId}', UserId: '{UserId}'.", this.SanitizeLogData(queryText), this.SanitizeLogData(outputMessageText), - this._searchScope.ApplicationId, - this._searchScope.AgentId, - this._searchScope.ThreadId, - this.SanitizeLogData(this._searchScope.UserId)); + searchScope.ApplicationId, + searchScope.AgentId, + searchScope.ThreadId, + this.SanitizeLogData(searchScope.UserId)); } } return new AIContext { - Messages = [new ChatMessage(ChatRole.User, outputMessageText)] + Instructions = inputContext.Instructions, + Messages = + (inputContext.Messages ?? []) + .Concat(outputMessage is not null ? [outputMessage] : []), + Tools = inputContext.Tools }; } catch (ArgumentException) @@ -195,12 +185,12 @@ public sealed class Mem0Provider : AIContextProvider this._logger.LogError( ex, "Mem0AIContextProvider: Failed to search Mem0 for memories due to error. ApplicationId: '{ApplicationId}', AgentId: '{AgentId}', ThreadId: '{ThreadId}', UserId: '{UserId}'.", - this._searchScope.ApplicationId, - this._searchScope.AgentId, - this._searchScope.ThreadId, - this.SanitizeLogData(this._searchScope.UserId)); + searchScope.ApplicationId, + searchScope.AgentId, + searchScope.ThreadId, + this.SanitizeLogData(searchScope.UserId)); } - return new AIContext(); + return inputContext; } } @@ -212,12 +202,15 @@ public sealed class Mem0Provider : AIContextProvider return; // Do not update memory on failed invocations. } + var state = this.GetOrInitializeState(context.Session); + var storageScope = state?.StorageScope ?? new Mem0ProviderScope(); + try { // Persist request and response messages after invocation. await this.PersistMessagesAsync( - context.RequestMessages - .Where(m => m.GetAgentRequestMessageSourceType() == AgentRequestMessageSourceType.External) + storageScope, + this._storageInputMessageFilter(context.RequestMessages) .Concat(context.ResponseMessages ?? []), cancellationToken).ConfigureAwait(false); } @@ -228,36 +221,39 @@ public sealed class Mem0Provider : AIContextProvider this._logger.LogError( ex, "Mem0AIContextProvider: Failed to send messages to Mem0 due to error. ApplicationId: '{ApplicationId}', AgentId: '{AgentId}', ThreadId: '{ThreadId}', UserId: '{UserId}'.", - this._storageScope.ApplicationId, - this._storageScope.AgentId, - this._storageScope.ThreadId, - this.SanitizeLogData(this._storageScope.UserId)); + storageScope.ApplicationId, + storageScope.AgentId, + storageScope.ThreadId, + this.SanitizeLogData(storageScope.UserId)); } } } /// - /// Clears stored memories for the configured scopes. + /// Clears stored memories for the specified scope. /// + /// The session containing the scope state to clear memories for. /// Cancellation token. - public Task ClearStoredMemoriesAsync(CancellationToken cancellationToken = default) => - this._client.ClearMemoryAsync( - this._storageScope.ApplicationId, - this._storageScope.AgentId, - this._storageScope.ThreadId, - this._storageScope.UserId, - cancellationToken); - - /// - public override JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null) + public Task ClearStoredMemoriesAsync(AgentSession session, CancellationToken cancellationToken = default) { - var state = new Mem0State(this._storageScope, this._searchScope); + Throw.IfNull(session); + var state = this.GetOrInitializeState(session); + var storageScope = state?.StorageScope; - var jso = jsonSerializerOptions ?? Mem0JsonUtilities.DefaultOptions; - return JsonSerializer.SerializeToElement(state, jso.GetTypeInfo(typeof(Mem0State))); + if (storageScope is null) + { + return Task.CompletedTask; // Nothing to clear if there is no state. + } + + return this._client.ClearMemoryAsync( + storageScope.ApplicationId, + storageScope.AgentId, + storageScope.ThreadId, + storageScope.UserId, + cancellationToken); } - private async Task PersistMessagesAsync(IEnumerable messages, CancellationToken cancellationToken) + private async Task PersistMessagesAsync(Mem0ProviderScope storageScope, IEnumerable messages, CancellationToken cancellationToken) { foreach (var message in messages) { @@ -277,27 +273,42 @@ public sealed class Mem0Provider : AIContextProvider } await this._client.CreateMemoryAsync( - this._storageScope.ApplicationId, - this._storageScope.AgentId, - this._storageScope.ThreadId, - this._storageScope.UserId, + storageScope.ApplicationId, + storageScope.AgentId, + storageScope.ThreadId, + storageScope.UserId, message.Text, message.Role.Value, cancellationToken).ConfigureAwait(false); } } - internal sealed class Mem0State + /// + /// Represents the state of a stored in the . + /// + public sealed class State { + /// + /// Initializes a new instance of the class with the specified storage and search scopes. + /// + /// The scope to use when storing memories. + /// The scope to use when searching for memories. If null, the storage scope will be used for searching as well. [JsonConstructor] - public Mem0State(Mem0ProviderScope storageScope, Mem0ProviderScope searchScope) + public State(Mem0ProviderScope storageScope, Mem0ProviderScope? searchScope = null) { - this.StorageScope = storageScope; - this.SearchScope = searchScope; + this.StorageScope = Throw.IfNull(storageScope); + this.SearchScope = searchScope ?? storageScope; } - public Mem0ProviderScope StorageScope { get; set; } - public Mem0ProviderScope SearchScope { get; set; } + /// + /// Gets the scope used when storing memories. + /// + public Mem0ProviderScope StorageScope { get; } + + /// + /// Gets the scope used when searching memories. + /// + public Mem0ProviderScope SearchScope { get; } } private string? SanitizeLogData(string? data) => this._enableSensitiveTelemetryData ? data : ""; diff --git a/dotnet/src/Microsoft.Agents.AI.Mem0/Mem0ProviderOptions.cs b/dotnet/src/Microsoft.Agents.AI.Mem0/Mem0ProviderOptions.cs index f2d3d89e16..f7d14028d9 100644 --- a/dotnet/src/Microsoft.Agents.AI.Mem0/Mem0ProviderOptions.cs +++ b/dotnet/src/Microsoft.Agents.AI.Mem0/Mem0ProviderOptions.cs @@ -1,5 +1,9 @@ // Copyright (c) Microsoft. All rights reserved. +using System; +using System.Collections.Generic; +using Microsoft.Extensions.AI; + namespace Microsoft.Agents.AI.Mem0; /// @@ -18,4 +22,30 @@ public sealed class Mem0ProviderOptions /// /// Defaults to . public bool EnableSensitiveTelemetryData { get; set; } + + /// + /// Gets or sets the key used to store the provider state in the session's . + /// + /// Defaults to the provider's type name. + public string? StateKey { get; set; } + + /// + /// Gets or sets an optional filter function applied to request messages when building the search text to use when + /// searching for relevant memories during . + /// + /// + /// When , the provider defaults to including only + /// messages. + /// + public Func, IEnumerable>? SearchInputMessageFilter { get; set; } + + /// + /// Gets or sets an optional filter function applied to request messages when determining which messages to + /// extract memories from during . + /// + /// + /// When , the provider defaults to including only + /// messages. + /// + public Func, IEnumerable>? StorageInputMessageFilter { get; set; } } diff --git a/dotnet/src/Microsoft.Agents.AI.OpenAI/Extensions/OpenAIAssistantClientExtensions.cs b/dotnet/src/Microsoft.Agents.AI.OpenAI/Extensions/OpenAIAssistantClientExtensions.cs index d167d1f0b4..c56a63c76e 100644 --- a/dotnet/src/Microsoft.Agents.AI.OpenAI/Extensions/OpenAIAssistantClientExtensions.cs +++ b/dotnet/src/Microsoft.Agents.AI.OpenAI/Extensions/OpenAIAssistantClientExtensions.cs @@ -204,8 +204,8 @@ public static class OpenAIAssistantClientExtensions Name = options.Name ?? assistantMetadata.Name, Description = options.Description ?? assistantMetadata.Description, ChatOptions = options.ChatOptions, - AIContextProviderFactory = options.AIContextProviderFactory, - ChatHistoryProviderFactory = options.ChatHistoryProviderFactory, + AIContextProviders = options.AIContextProviders, + ChatHistoryProvider = options.ChatHistoryProvider, UseProvidedChatClientAsIs = options.UseProvidedChatClientAsIs }; diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowChatHistoryProvider.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowChatHistoryProvider.cs index f631de8e8a..6672b9e2a3 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowChatHistoryProvider.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowChatHistoryProvider.cs @@ -6,48 +6,56 @@ using System.Text.Json; using System.Threading; using System.Threading.Tasks; using Microsoft.Extensions.AI; -using Microsoft.Shared.Diagnostics; namespace Microsoft.Agents.AI.Workflows; internal sealed class WorkflowChatHistoryProvider : ChatHistoryProvider { - private int _bookmark; - private readonly List _chatMessages = []; + private readonly JsonSerializerOptions _jsonSerializerOptions; - public WorkflowChatHistoryProvider() + /// + /// Initializes a new instance of the class. + /// + /// + /// Optional JSON serializer options for serializing the state of this provider. + /// This is valuable for cases like when the chat history contains custom types + /// and source generated serializers are required, or Native AOT / Trimming is required. + /// + public WorkflowChatHistoryProvider(JsonSerializerOptions? jsonSerializerOptions = null) { - } - - public WorkflowChatHistoryProvider(StoreState state) - { - this.ImportStoreState(Throw.IfNull(state)); - } - - private void ImportStoreState(StoreState state, bool clearMessages = false) - { - if (clearMessages) - { - this._chatMessages.Clear(); - } - - if (state?.Messages is not null) - { - this._chatMessages.AddRange(state.Messages); - } - this._bookmark = state?.Bookmark ?? 0; + this._jsonSerializerOptions = jsonSerializerOptions ?? AgentAbstractionsJsonUtilities.DefaultOptions; } internal sealed class StoreState { public int Bookmark { get; set; } - public IList Messages { get; set; } = []; + public List Messages { get; set; } = []; } - internal void AddMessages(params IEnumerable messages) => this._chatMessages.AddRange(messages); + private StoreState GetOrInitializeState(AgentSession? session) + { + if (session?.StateBag.TryGetValue(this.StateKey, out var state, this._jsonSerializerOptions) is true && state is not null) + { + return state; + } + + state = new(); + if (session is not null) + { + session.StateBag.SetValue(this.StateKey, state, this._jsonSerializerOptions); + } + + return state; + } + + internal void AddMessages(AgentSession session, params IEnumerable messages) + => this.GetOrInitializeState(session).Messages.AddRange(messages); protected override ValueTask> InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default) - => new(this._chatMessages.AsReadOnly()); + => new(this.GetOrInitializeState(context.Session) + .Messages + .Select(message => message.WithAgentRequestMessageSource(AgentRequestMessageSourceType.ChatHistory, this.GetType().FullName!)) + .Concat(context.RequestMessages)); protected override ValueTask InvokedCoreAsync(InvokedContext context, CancellationToken cancellationToken = default) { @@ -56,29 +64,27 @@ internal sealed class WorkflowChatHistoryProvider : ChatHistoryProvider return default; } - var allNewMessages = context.RequestMessages.Concat(context.ResponseMessages ?? []); - this._chatMessages.AddRange(allNewMessages); + var allNewMessages = context.RequestMessages + .Where(m => m.GetAgentRequestMessageSourceType() != AgentRequestMessageSourceType.ChatHistory) + .Concat(context.ResponseMessages ?? []); + this.GetOrInitializeState(context.Session).Messages.AddRange(allNewMessages); return default; } - public IEnumerable GetFromBookmark() + public IEnumerable GetFromBookmark(AgentSession session) { - for (int i = this._bookmark; i < this._chatMessages.Count; i++) + var state = this.GetOrInitializeState(session); + + for (int i = state.Bookmark; i < state.Messages.Count; i++) { - yield return this._chatMessages[i]; + yield return state.Messages[i]; } } - public void UpdateBookmark() => this._bookmark = this._chatMessages.Count; - - public override JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null) + public void UpdateBookmark(AgentSession session) { - StoreState state = this.ExportStoreState(); - - return JsonSerializer.SerializeToElement(state, - WorkflowsJsonUtilities.DefaultOptions.GetTypeInfo(typeof(StoreState))); + var state = this.GetOrInitializeState(session); + state.Bookmark = state.Messages.Count; } - - internal StoreState ExportStoreState() => new() { Bookmark = this._bookmark, Messages = this._chatMessages }; } diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowHostAgent.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowHostAgent.cs index c08ba5c3f4..7438ce1a34 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowHostAgent.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowHostAgent.cs @@ -94,7 +94,7 @@ internal sealed class WorkflowHostAgent : AIAgent // For workflow threads, messages are added directly via the internal AddMessages method // The MessageStore methods are used for agent invocation scenarios - workflowSession.ChatHistoryProvider.AddMessages(messages); + workflowSession.ChatHistoryProvider.AddMessages(session, messages); return workflowSession; } diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowSession.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowSession.cs index 7730067e60..4ea88cec0e 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowSession.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowSession.cs @@ -70,7 +70,8 @@ internal sealed class WorkflowSession : AgentSession this.RunId = sessionState.RunId; this.LastCheckpoint = sessionState.LastCheckpoint; - this.ChatHistoryProvider = new WorkflowChatHistoryProvider(sessionState.ChatHistoryProviderState); + this.ChatHistoryProvider = new WorkflowChatHistoryProvider(); + this.StateBag = sessionState.StateBag; } public CheckpointInfo? LastCheckpoint { get; set; } @@ -81,8 +82,8 @@ internal sealed class WorkflowSession : AgentSession SessionState info = new( this.RunId, this.LastCheckpoint, - this.ChatHistoryProvider.ExportStoreState(), - this._inMemoryCheckpointManager); + this._inMemoryCheckpointManager, + this.StateBag); return marshaller.Marshal(info); } @@ -100,7 +101,7 @@ internal sealed class WorkflowSession : AgentSession RawRepresentation = raw }; - this.ChatHistoryProvider.AddMessages(update.ToChatMessage()); + this.ChatHistoryProvider.AddMessages(this, update.ToChatMessage()); return update; } @@ -117,7 +118,7 @@ internal sealed class WorkflowSession : AgentSession RawRepresentation = raw }; - this.ChatHistoryProvider.AddMessages(update.ToChatMessage()); + this.ChatHistoryProvider.AddMessages(this, update.ToChatMessage()); return update; } @@ -156,7 +157,7 @@ internal sealed class WorkflowSession : AgentSession try { this.LastResponseId = Guid.NewGuid().ToString("N"); - List messages = this.ChatHistoryProvider.GetFromBookmark().ToList(); + List messages = this.ChatHistoryProvider.GetFromBookmark(this).ToList(); #pragma warning disable CA2007 // Analyzer misfiring and not seeing .ConfigureAwait(false) below. await using Checkpointed checkpointed = @@ -240,7 +241,7 @@ internal sealed class WorkflowSession : AgentSession finally { // Do we want to try to undo the step, and not update the bookmark? - this.ChatHistoryProvider.UpdateBookmark(); + this.ChatHistoryProvider.UpdateBookmark(this); } } @@ -254,12 +255,12 @@ internal sealed class WorkflowSession : AgentSession internal sealed class SessionState( string runId, CheckpointInfo? lastCheckpoint, - WorkflowChatHistoryProvider.StoreState chatHistoryProviderState, - InMemoryCheckpointManager? checkpointManager = null) + InMemoryCheckpointManager? checkpointManager = null, + AgentSessionStateBag? stateBag = null) { public string RunId { get; } = runId; public CheckpointInfo? LastCheckpoint { get; } = lastCheckpoint; - public WorkflowChatHistoryProvider.StoreState ChatHistoryProviderState { get; } = chatHistoryProviderState; public InMemoryCheckpointManager? CheckpointManager { get; } = checkpointManager; + public AgentSessionStateBag StateBag { get; } = stateBag ?? new(); } } diff --git a/dotnet/src/Microsoft.Agents.AI/AgentJsonUtilities.cs b/dotnet/src/Microsoft.Agents.AI/AgentJsonUtilities.cs index 5c915b6b01..96ec6dbecb 100644 --- a/dotnet/src/Microsoft.Agents.AI/AgentJsonUtilities.cs +++ b/dotnet/src/Microsoft.Agents.AI/AgentJsonUtilities.cs @@ -65,9 +65,9 @@ internal static partial class AgentJsonUtilities NumberHandling = JsonNumberHandling.AllowReadingFromString)] // Agent abstraction types - [JsonSerializable(typeof(ChatClientAgentSession.SessionState))] + [JsonSerializable(typeof(ChatClientAgentSession))] [JsonSerializable(typeof(TextSearchProvider.TextSearchProviderState))] - [JsonSerializable(typeof(ChatHistoryMemoryProvider.ChatHistoryMemoryProviderState))] + [JsonSerializable(typeof(ChatHistoryMemoryProvider.State))] [ExcludeFromCodeCoverage] internal sealed partial class JsonContext : JsonSerializerContext; diff --git a/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgent.cs b/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgent.cs index dc462cf501..6dc6d175a7 100644 --- a/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgent.cs +++ b/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgent.cs @@ -20,6 +20,7 @@ namespace Microsoft.Agents.AI; public sealed partial class ChatClientAgent : AIAgent { private readonly ChatClientAgentOptions? _agentOptions; + private readonly HashSet _aiContextProviderStateKeys; private readonly AIAgentMetadata _agentMetadata; private readonly ILogger _logger; private readonly Type _chatClientType; @@ -105,6 +106,15 @@ public sealed partial class ChatClientAgent : AIAgent // If the user has not opted out of using our default decorators, we wrap the chat client. this.ChatClient = options?.UseProvidedChatClientAsIs is true ? chatClient : chatClient.WithDefaultAgentMiddleware(options, services); + // Use the ChatHistoryProvider from options if provided. + // If one was not provided, and we later find out that the underlying service does not manage chat history server-side, + // we will use the default InMemoryChatHistoryProvider at that time. + this.ChatHistoryProvider = options?.ChatHistoryProvider; + this.AIContextProviders = this._agentOptions?.AIContextProviders as IReadOnlyList ?? this._agentOptions?.AIContextProviders?.ToList(); + + // Validate that no two providers share the same StateKey, since they would overwrite each other's state in the session. + this._aiContextProviderStateKeys = ValidateAndCollectStateKeys(this._agentOptions?.AIContextProviders, this.ChatHistoryProvider); + this._logger = (loggerFactory ?? chatClient.GetService() ?? NullLoggerFactory.Instance).CreateLogger(); } @@ -120,6 +130,22 @@ public sealed partial class ChatClientAgent : AIAgent /// public IChatClient ChatClient { get; } + /// + /// Gets the used by this agent, to support cases where the chat history is not stored by the agent service. + /// + /// + /// This property may be null in case the agent stores messages in the underlying agent service. + /// + public ChatHistoryProvider? ChatHistoryProvider { get; private set; } + + /// + /// Gets the list of instances used by this agent, to support cases where additional context is needed for each agent run. + /// + /// + /// This property may be null in case no additional context providers were configured. + /// + public IReadOnlyList? AIContextProviders { get; } + /// protected override string? IdCore => this._agentOptions?.Id; @@ -206,7 +232,6 @@ public sealed partial class ChatClientAgent : AIAgent (ChatClientAgentSession safeSession, ChatOptions? chatOptions, - List inputMessagesForProviders, List inputMessagesForChatClient, ChatClientAgentContinuationToken? continuationToken) = await this.PrepareSessionAndMessagesAsync(session, inputMessages, options, cancellationToken).ConfigureAwait(false); @@ -230,8 +255,8 @@ public sealed partial class ChatClientAgent : AIAgent } catch (Exception ex) { - await this.NotifyChatHistoryProviderOfFailureAsync(safeSession, ex, GetInputMessages(inputMessagesForProviders, continuationToken), chatOptions, cancellationToken).ConfigureAwait(false); - await this.NotifyAIContextProviderOfFailureAsync(safeSession, ex, GetInputMessages(inputMessagesForProviders, continuationToken), cancellationToken).ConfigureAwait(false); + await this.NotifyChatHistoryProviderOfFailureAsync(safeSession, ex, GetInputMessages(inputMessagesForChatClient, continuationToken), chatOptions, cancellationToken).ConfigureAwait(false); + await this.NotifyAIContextProviderOfFailureAsync(safeSession, ex, GetInputMessages(inputMessagesForChatClient, continuationToken), cancellationToken).ConfigureAwait(false); throw; } @@ -245,8 +270,8 @@ public sealed partial class ChatClientAgent : AIAgent } catch (Exception ex) { - await this.NotifyChatHistoryProviderOfFailureAsync(safeSession, ex, GetInputMessages(inputMessagesForProviders, continuationToken), chatOptions, cancellationToken).ConfigureAwait(false); - await this.NotifyAIContextProviderOfFailureAsync(safeSession, ex, GetInputMessages(inputMessagesForProviders, continuationToken), cancellationToken).ConfigureAwait(false); + await this.NotifyChatHistoryProviderOfFailureAsync(safeSession, ex, GetInputMessages(inputMessagesForChatClient, continuationToken), chatOptions, cancellationToken).ConfigureAwait(false); + await this.NotifyAIContextProviderOfFailureAsync(safeSession, ex, GetInputMessages(inputMessagesForChatClient, continuationToken), cancellationToken).ConfigureAwait(false); throw; } @@ -272,8 +297,8 @@ public sealed partial class ChatClientAgent : AIAgent } catch (Exception ex) { - await this.NotifyChatHistoryProviderOfFailureAsync(safeSession, ex, GetInputMessages(inputMessagesForProviders, continuationToken), chatOptions, cancellationToken).ConfigureAwait(false); - await this.NotifyAIContextProviderOfFailureAsync(safeSession, ex, GetInputMessages(inputMessagesForProviders, continuationToken), cancellationToken).ConfigureAwait(false); + await this.NotifyChatHistoryProviderOfFailureAsync(safeSession, ex, GetInputMessages(inputMessagesForChatClient, continuationToken), chatOptions, cancellationToken).ConfigureAwait(false); + await this.NotifyAIContextProviderOfFailureAsync(safeSession, ex, GetInputMessages(inputMessagesForChatClient, continuationToken), cancellationToken).ConfigureAwait(false); throw; } } @@ -282,13 +307,13 @@ public sealed partial class ChatClientAgent : AIAgent // We can derive the type of supported session from whether we have a conversation id, // so let's update it and set the conversation id for the service session case. - await this.UpdateSessionWithTypeAndConversationIdAsync(safeSession, chatResponse.ConversationId, cancellationToken).ConfigureAwait(false); + this.UpdateSessionConversationId(safeSession, chatResponse.ConversationId, cancellationToken); // To avoid inconsistent state we only notify the session of the input messages if no error occurs after the initial request. - await this.NotifyChatHistoryProviderOfNewMessagesAsync(safeSession, GetInputMessages(inputMessagesForProviders, continuationToken), chatResponse.Messages, chatOptions, cancellationToken).ConfigureAwait(false); + await this.NotifyChatHistoryProviderOfNewMessagesAsync(safeSession, GetInputMessages(inputMessagesForChatClient, continuationToken), chatResponse.Messages, chatOptions, cancellationToken).ConfigureAwait(false); // Notify the AIContextProvider of all new messages. - await this.NotifyAIContextProviderOfSuccessAsync(safeSession, GetInputMessages(inputMessagesForProviders, continuationToken), chatResponse.Messages, cancellationToken).ConfigureAwait(false); + await this.NotifyAIContextProviderOfSuccessAsync(safeSession, GetInputMessages(inputMessagesForChatClient, continuationToken), chatResponse.Messages, cancellationToken).ConfigureAwait(false); } /// @@ -298,24 +323,14 @@ public sealed partial class ChatClientAgent : AIAgent : serviceType == typeof(IChatClient) ? this.ChatClient : serviceType == typeof(ChatOptions) ? this._agentOptions?.ChatOptions : serviceType == typeof(ChatClientAgentOptions) ? this._agentOptions - : this.ChatClient.GetService(serviceType, serviceKey)); + : this.AIContextProviders?.Select(provider => provider.GetService(serviceType, serviceKey)).FirstOrDefault(s => s is not null) + ?? this.ChatHistoryProvider?.GetService(serviceType, serviceKey) + ?? this.ChatClient.GetService(serviceType, serviceKey)); /// - protected override async ValueTask CreateSessionCoreAsync(CancellationToken cancellationToken = default) + protected override ValueTask CreateSessionCoreAsync(CancellationToken cancellationToken = default) { - ChatHistoryProvider? chatHistoryProvider = this._agentOptions?.ChatHistoryProviderFactory is not null - ? await this._agentOptions.ChatHistoryProviderFactory.Invoke(new() { SerializedState = default, JsonSerializerOptions = null }, cancellationToken).ConfigureAwait(false) - : null; - - AIContextProvider? contextProvider = this._agentOptions?.AIContextProviderFactory is not null - ? await this._agentOptions.AIContextProviderFactory.Invoke(new() { SerializedState = default, JsonSerializerOptions = null }, cancellationToken).ConfigureAwait(false) - : null; - - return new ChatClientAgentSession - { - ChatHistoryProvider = chatHistoryProvider, - AIContextProvider = contextProvider - }; + return new(new ChatClientAgentSession()); } /// @@ -336,52 +351,12 @@ public sealed partial class ChatClientAgent : AIAgent /// instances that support server-side conversation storage through their underlying . /// /// - public async ValueTask CreateSessionAsync(string conversationId, CancellationToken cancellationToken = default) + public ValueTask CreateSessionAsync(string conversationId, CancellationToken cancellationToken = default) { - AIContextProvider? contextProvider = this._agentOptions?.AIContextProviderFactory is not null - ? await this._agentOptions.AIContextProviderFactory.Invoke(new() { SerializedState = default, JsonSerializerOptions = null }, cancellationToken).ConfigureAwait(false) - : null; - - return new ChatClientAgentSession() + return new(new ChatClientAgentSession() { ConversationId = conversationId, - AIContextProvider = contextProvider - }; - } - - /// - /// Creates a new agent session instance using an existing to continue a conversation. - /// - /// The instance to use for managing the conversation's message history. - /// The to monitor for cancellation requests. - /// - /// A value task representing the asynchronous operation. The task result contains a new instance configured to work with the provided . - /// - /// - /// - /// This method creates threads that do not support server-side conversation storage. - /// Some AI services require server-side conversation storage to function properly, and creating a session - /// with a may not be compatible with these services. - /// - /// - /// Where a service requires server-side conversation storage, use . - /// - /// - /// If the agent detects, during the first run, that the underlying AI service requires server-side conversation storage, - /// the session will throw an exception to indicate that it cannot continue using the provided . - /// - /// - public async ValueTask CreateSessionAsync(ChatHistoryProvider chatHistoryProvider, CancellationToken cancellationToken = default) - { - AIContextProvider? contextProvider = this._agentOptions?.AIContextProviderFactory is not null - ? await this._agentOptions.AIContextProviderFactory.Invoke(new() { SerializedState = default, JsonSerializerOptions = null }, cancellationToken).ConfigureAwait(false) - : null; - - return new ChatClientAgentSession() - { - ChatHistoryProvider = Throw.IfNull(chatHistoryProvider), - AIContextProvider = contextProvider - }; + }); } /// @@ -398,22 +373,9 @@ public sealed partial class ChatClientAgent : AIAgent } /// - protected override async ValueTask DeserializeSessionCoreAsync(JsonElement serializedState, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default) + protected override ValueTask DeserializeSessionCoreAsync(JsonElement serializedState, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default) { - Func>? chatHistoryProviderFactory = this._agentOptions?.ChatHistoryProviderFactory is null ? - null : - (jse, jso, ct) => this._agentOptions.ChatHistoryProviderFactory.Invoke(new() { SerializedState = jse, JsonSerializerOptions = jso }, ct); - - Func>? aiContextProviderFactory = this._agentOptions?.AIContextProviderFactory is null ? - null : - (jse, jso, ct) => this._agentOptions.AIContextProviderFactory.Invoke(new() { SerializedState = jse, JsonSerializerOptions = jso }, ct); - - return await ChatClientAgentSession.DeserializeAsync( - serializedState, - jsonSerializerOptions, - chatHistoryProviderFactory, - aiContextProviderFactory, - cancellationToken).ConfigureAwait(false); + return new(ChatClientAgentSession.Deserialize(serializedState, jsonSerializerOptions)); } #region Private @@ -432,7 +394,6 @@ public sealed partial class ChatClientAgent : AIAgent (ChatClientAgentSession safeSession, ChatOptions? chatOptions, - List inputMessagesForProviders, List inputMessagesForChatClient, ChatClientAgentContinuationToken? _) = await this.PrepareSessionAndMessagesAsync(session, inputMessages, options, cancellationToken).ConfigureAwait(false); @@ -453,8 +414,8 @@ public sealed partial class ChatClientAgent : AIAgent } catch (Exception ex) { - await this.NotifyChatHistoryProviderOfFailureAsync(safeSession, ex, inputMessagesForProviders, chatOptions, cancellationToken).ConfigureAwait(false); - await this.NotifyAIContextProviderOfFailureAsync(safeSession, ex, inputMessagesForProviders, cancellationToken).ConfigureAwait(false); + await this.NotifyChatHistoryProviderOfFailureAsync(safeSession, ex, inputMessagesForChatClient, chatOptions, cancellationToken).ConfigureAwait(false); + await this.NotifyAIContextProviderOfFailureAsync(safeSession, ex, inputMessagesForChatClient, cancellationToken).ConfigureAwait(false); throw; } @@ -462,7 +423,7 @@ public sealed partial class ChatClientAgent : AIAgent // We can derive the type of supported session from whether we have a conversation id, // so let's update it and set the conversation id for the service session case. - await this.UpdateSessionWithTypeAndConversationIdAsync(safeSession, chatResponse.ConversationId, cancellationToken).ConfigureAwait(false); + this.UpdateSessionConversationId(safeSession, chatResponse.ConversationId, cancellationToken); // Ensure that the author name is set for each message in the response. foreach (ChatMessage chatResponseMessage in chatResponse.Messages) @@ -471,10 +432,10 @@ public sealed partial class ChatClientAgent : AIAgent } // Only notify the session of new messages if the chatResponse was successful to avoid inconsistent message state in the session. - await this.NotifyChatHistoryProviderOfNewMessagesAsync(safeSession, inputMessagesForProviders, chatResponse.Messages, chatOptions, cancellationToken).ConfigureAwait(false); + await this.NotifyChatHistoryProviderOfNewMessagesAsync(safeSession, inputMessagesForChatClient, chatResponse.Messages, chatOptions, cancellationToken).ConfigureAwait(false); // Notify the AIContextProvider of all new messages. - await this.NotifyAIContextProviderOfSuccessAsync(safeSession, inputMessagesForProviders, chatResponse.Messages, cancellationToken).ConfigureAwait(false); + await this.NotifyAIContextProviderOfSuccessAsync(safeSession, inputMessagesForChatClient, chatResponse.Messages, cancellationToken).ConfigureAwait(false); var agentResponse = agentResponseFactoryFunc(chatResponse); @@ -492,10 +453,14 @@ public sealed partial class ChatClientAgent : AIAgent IEnumerable responseMessages, CancellationToken cancellationToken) { - if (session.AIContextProvider is not null) + if (this.AIContextProviders is { Count: > 0 } contextProviders) { - await session.AIContextProvider.InvokedAsync(new(this, session, inputMessages) { ResponseMessages = responseMessages }, - cancellationToken).ConfigureAwait(false); + AIContextProvider.InvokedContext invokedContext = new(this, session, inputMessages, responseMessages); + + foreach (var contextProvider in contextProviders) + { + await contextProvider.InvokedAsync(invokedContext, cancellationToken).ConfigureAwait(false); + } } } @@ -508,10 +473,14 @@ public sealed partial class ChatClientAgent : AIAgent IEnumerable inputMessages, CancellationToken cancellationToken) { - if (session.AIContextProvider is not null) + if (this.AIContextProviders is { Count: > 0 } contextProviders) { - await session.AIContextProvider.InvokedAsync(new(this, session, inputMessages) { InvokeException = ex }, - cancellationToken).ConfigureAwait(false); + AIContextProvider.InvokedContext invokedContext = new(this, session, inputMessages, ex); + + foreach (var contextProvider in contextProviders) + { + await contextProvider.InvokedAsync(invokedContext, cancellationToken).ConfigureAwait(false); + } } } @@ -679,7 +648,6 @@ public sealed partial class ChatClientAgent : AIAgent <( ChatClientAgentSession AgentSession, ChatOptions? ChatOptions, - List inputMessagesForProviders, List InputMessagesForChatClient, ChatClientAgentContinuationToken? ContinuationToken )> PrepareSessionAndMessagesAsync( @@ -709,52 +677,53 @@ public sealed partial class ChatClientAgent : AIAgent throw new InvalidOperationException("Input messages are not allowed when continuing a background response using a continuation token."); } - List inputMessagesForProviders = []; - List inputMessagesForChatClient = []; + IEnumerable inputMessagesForChatClient = inputMessages; // Populate the session messages only if we are not continuing an existing response as it's not allowed if (chatOptions?.ContinuationToken is null) { - ChatHistoryProvider? chatHistoryProvider = ResolveChatHistoryProvider(typedSession, chatOptions); + ChatHistoryProvider? chatHistoryProvider = this.ResolveChatHistoryProvider(chatOptions, typedSession); // Add any existing messages from the session to the messages to be sent to the chat client. + // The ChatHistoryProvider returns the merged result (history + input messages). if (chatHistoryProvider is not null) { - var invokingContext = new ChatHistoryProvider.InvokingContext(this, typedSession, inputMessages); - var providerMessages = await chatHistoryProvider.InvokingAsync(invokingContext, cancellationToken).ConfigureAwait(false); - inputMessagesForChatClient.AddRange(providerMessages); + var invokingContext = new ChatHistoryProvider.InvokingContext(this, typedSession, inputMessagesForChatClient); + inputMessagesForChatClient = await chatHistoryProvider.InvokingAsync(invokingContext, cancellationToken).ConfigureAwait(false); } - // Add the input messages before getting context from AIContextProvider. - inputMessagesForProviders.AddRange(inputMessages); - inputMessagesForChatClient.AddRange(inputMessages); - // If we have an AIContextProvider, we should get context from it, and update our // messages and options with the additional context. - if (typedSession.AIContextProvider is not null) + // The AIContextProvider returns the accumulated AIContext (original + new contributions). + if (this.AIContextProviders is { Count: > 0 } aiContextProviders) { - var invokingContext = new AIContextProvider.InvokingContext(this, typedSession, inputMessages); - var aiContext = await typedSession.AIContextProvider.InvokingAsync(invokingContext, cancellationToken).ConfigureAwait(false); - if (aiContext.Messages is { Count: > 0 }) + var aiContext = new AIContext { - inputMessagesForProviders.AddRange(aiContext.Messages); - inputMessagesForChatClient.AddRange(aiContext.Messages); + Instructions = chatOptions?.Instructions, + Messages = inputMessagesForChatClient, + Tools = chatOptions?.Tools + }; + + foreach (var aiContextProvider in aiContextProviders) + { + var invokingContext = new AIContextProvider.InvokingContext(this, typedSession, aiContext); + aiContext = await aiContextProvider.InvokingAsync(invokingContext, cancellationToken).ConfigureAwait(false); } - if (aiContext.Tools is { Count: > 0 }) + // Materialize the accumulated messages and tools once at the end of the provider pipeline. + inputMessagesForChatClient = aiContext.Messages ?? []; + + var tools = aiContext.Tools as IList ?? aiContext.Tools?.ToList(); + if (chatOptions?.Tools is { Count: > 0 } || tools is { Count: > 0 }) { chatOptions ??= new(); - chatOptions.Tools ??= []; - foreach (AITool tool in aiContext.Tools) - { - chatOptions.Tools.Add(tool); - } + chatOptions.Tools = tools; } - if (aiContext.Instructions is not null) + if (chatOptions?.Instructions is not null || aiContext.Instructions is not null) { chatOptions ??= new(); - chatOptions.Instructions = string.IsNullOrWhiteSpace(chatOptions.Instructions) ? aiContext.Instructions : $"{chatOptions.Instructions}\n{aiContext.Instructions}"; + chatOptions.Instructions = aiContext.Instructions; } } } @@ -777,10 +746,13 @@ public sealed partial class ChatClientAgent : AIAgent chatOptions.ConversationId = typedSession.ConversationId; } - return (typedSession, chatOptions, inputMessagesForProviders, inputMessagesForChatClient, continuationToken); + // Materialize the accumulated messages once at the end of the provider pipeline, reusing the existing list if possible. + List messagesList = inputMessagesForChatClient as List ?? inputMessagesForChatClient.ToList(); + + return (typedSession, chatOptions, messagesList, continuationToken); } - private async Task UpdateSessionWithTypeAndConversationIdAsync(ChatClientAgentSession session, string? responseConversationId, CancellationToken cancellationToken) + private void UpdateSessionConversationId(ChatClientAgentSession session, string? responseConversationId, CancellationToken cancellationToken) { if (string.IsNullOrWhiteSpace(responseConversationId) && !string.IsNullOrWhiteSpace(session.ConversationId)) { @@ -791,6 +763,14 @@ public sealed partial class ChatClientAgent : AIAgent if (!string.IsNullOrWhiteSpace(responseConversationId)) { + if (this.ChatHistoryProvider is not null) + { + // The agent has a ChatHistoryProvider configured, but the service returned a conversation id, + // meaning the service manages chat history server-side. Both cannot be used simultaneously. + throw new InvalidOperationException( + $"Only {nameof(ChatClientAgentSession.ConversationId)} or {nameof(this.ChatHistoryProvider)} may be used, but not both. The service returned a conversation id indicating server-side chat history management, but the agent has a {nameof(this.ChatHistoryProvider)} configured."); + } + // If we got a conversation id back from the chat client, it means that the service supports server side session storage // so we should update the session with the new id. session.ConversationId = responseConversationId; @@ -798,11 +778,9 @@ public sealed partial class ChatClientAgent : AIAgent else { // If the service doesn't use service side chat history storage (i.e. we got no id back from invocation), and - // the session has no ChatHistoryProvider yet, we should update the session with the custom ChatHistoryProvider or - // default InMemoryChatHistoryProvider so that it has somewhere to store the chat history. - session.ChatHistoryProvider ??= this._agentOptions?.ChatHistoryProviderFactory is not null - ? await this._agentOptions.ChatHistoryProviderFactory.Invoke(new() { SerializedState = default, JsonSerializerOptions = null }, cancellationToken).ConfigureAwait(false) - : new InMemoryChatHistoryProvider(); + // the agent has no ChatHistoryProvider yet, we should use the default InMemoryChatHistoryProvider so that + // we have somewhere to store the chat history. + this.ChatHistoryProvider ??= new InMemoryChatHistoryProvider(); } } @@ -813,16 +791,13 @@ public sealed partial class ChatClientAgent : AIAgent ChatOptions? chatOptions, CancellationToken cancellationToken) { - ChatHistoryProvider? provider = ResolveChatHistoryProvider(session, chatOptions); + ChatHistoryProvider? provider = this.ResolveChatHistoryProvider(chatOptions, session); // Only notify the provider if we have one. // If we don't have one, it means that the chat history is service managed and the underlying service is responsible for storing messages. if (provider is not null) { - var invokedContext = new ChatHistoryProvider.InvokedContext(this, session, requestMessages) - { - InvokeException = ex - }; + var invokedContext = new ChatHistoryProvider.InvokedContext(this, session, requestMessages, ex); return provider.InvokedAsync(invokedContext, cancellationToken).AsTask(); } @@ -837,29 +812,45 @@ public sealed partial class ChatClientAgent : AIAgent ChatOptions? chatOptions, CancellationToken cancellationToken) { - ChatHistoryProvider? provider = ResolveChatHistoryProvider(session, chatOptions); + ChatHistoryProvider? provider = this.ResolveChatHistoryProvider(chatOptions, session); // Only notify the provider if we have one. // If we don't have one, it means that the chat history is service managed and the underlying service is responsible for storing messages. if (provider is not null) { - var invokedContext = new ChatHistoryProvider.InvokedContext(this, session, requestMessages) - { - ResponseMessages = responseMessages - }; + var invokedContext = new ChatHistoryProvider.InvokedContext(this, session, requestMessages, responseMessages); return provider.InvokedAsync(invokedContext, cancellationToken).AsTask(); } return Task.CompletedTask; } - private static ChatHistoryProvider? ResolveChatHistoryProvider(ChatClientAgentSession session, ChatOptions? chatOptions) + private ChatHistoryProvider? ResolveChatHistoryProvider(ChatOptions? chatOptions, ChatClientAgentSession session) { - ChatHistoryProvider? provider = session.ChatHistoryProvider; + ChatHistoryProvider? provider = this.ChatHistoryProvider; - // If someone provided an override ChatHistoryProvider via AdditionalProperties, we should use that instead of the one on the session. + if (session.ConversationId is not null && provider is not null) + { + throw new InvalidOperationException( + $"Only {nameof(ChatClientAgentSession.ConversationId)} or {nameof(this.ChatHistoryProvider)} may be used, but not both. The current {nameof(ChatClientAgentSession)} has a {nameof(ChatClientAgentSession.ConversationId)} indicating server-side chat history management, but the agent has a {nameof(this.ChatHistoryProvider)} configured."); + } + + // If someone provided an override ChatHistoryProvider via AdditionalProperties, we should use that instead. if (chatOptions?.AdditionalProperties?.TryGetValue(out ChatHistoryProvider? overrideProvider) is true) { + if (session.ConversationId is not null && overrideProvider is not null) + { + throw new InvalidOperationException( + $"Only {nameof(ChatClientAgentSession.ConversationId)} or {nameof(this.ChatHistoryProvider)} may be used, but not both. The current {nameof(ChatClientAgentSession)} has a {nameof(ChatClientAgentSession.ConversationId)} indicating server-side chat history management, but an override {nameof(this.ChatHistoryProvider)} was provided via {nameof(AgentRunOptions.AdditionalProperties)}."); + } + + // Validate that the override provider's StateKey does not clash with any AIContextProvider's StateKey. + if (overrideProvider is not null && this._aiContextProviderStateKeys.Contains(overrideProvider.StateKey)) + { + throw new InvalidOperationException( + $"The ChatHistoryProvider '{overrideProvider.GetType().Name}' uses the state key '{overrideProvider.StateKey}' which is already used by one of the configured AIContextProviders. Each provider must use a unique state key to avoid overwriting each other's state."); + } + provider = overrideProvider; } @@ -906,5 +897,43 @@ public sealed partial class ChatClientAgent : AIAgent } private string GetLoggingAgentName() => this.Name ?? "UnnamedAgent"; + + /// + /// Validates that all configured providers have unique values + /// and returns a of the AIContextProvider state keys. + /// + private static HashSet ValidateAndCollectStateKeys(IEnumerable? aiContextProviders, ChatHistoryProvider? chatHistoryProvider) + { + HashSet stateKeys = new(StringComparer.Ordinal); + + if (aiContextProviders is not null) + { + foreach (var provider in aiContextProviders) + { + if (!stateKeys.Add(provider.StateKey)) + { + throw new InvalidOperationException( + $"Multiple providers use the same state key '{provider.StateKey}'. Each provider must use a unique state key to avoid overwriting each other's state."); + } + } + } + + if (chatHistoryProvider is null + && stateKeys.Contains(nameof(InMemoryChatHistoryProvider))) + { + throw new InvalidOperationException( + $"The default {nameof(InMemoryChatHistoryProvider)} uses the state key '{nameof(InMemoryChatHistoryProvider)}', which is already used by one of the configured AIContextProviders. Each provider must use a unique state key to avoid overwriting each other's state. To resolve this, either configure a different state key for the AIContextProvider that is using '{nameof(InMemoryChatHistoryProvider)}' as its state key, or provide a custom ChatHistoryProvider with a unique state key."); + } + + if (chatHistoryProvider is not null + && stateKeys.Contains(chatHistoryProvider.StateKey)) + { + throw new InvalidOperationException( + $"The ChatHistoryProvider '{chatHistoryProvider.GetType().Name}' uses the state key '{chatHistoryProvider.StateKey}' which is already used by one of the configured AIContextProviders. Each provider must use a unique state key to avoid overwriting each other's state. To resolve this, either configure a different state key for the AIContextProvider that is using '{chatHistoryProvider.StateKey}' as its state key, or reconfigure the custom ChatHistoryProvider with a unique state key."); + } + + return stateKeys; + } + #endregion } diff --git a/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgentOptions.cs b/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgentOptions.cs index 6f8451e2b8..ddca9197ab 100644 --- a/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgentOptions.cs +++ b/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgentOptions.cs @@ -1,9 +1,6 @@ // Copyright (c) Microsoft. All rights reserved. -using System; -using System.Text.Json; -using System.Threading; -using System.Threading.Tasks; +using System.Collections.Generic; using Microsoft.Extensions.AI; namespace Microsoft.Agents.AI; @@ -39,17 +36,14 @@ public sealed class ChatClientAgentOptions public ChatOptions? ChatOptions { get; set; } /// - /// Gets or sets a factory function to create an instance of - /// which will be used to provide chat history for this agent. + /// Gets or sets the instance to use for providing chat history for this agent. /// - public Func>? ChatHistoryProviderFactory { get; set; } + public ChatHistoryProvider? ChatHistoryProvider { get; set; } /// - /// Gets or sets a factory function to create an instance of - /// which will be used to create a context provider for each new thread, and can then - /// provide additional context for each agent run. + /// Gets or sets the list of instances to use for providing additional context for each agent run. /// - public Func>? AIContextProviderFactory { get; set; } + public IEnumerable? AIContextProviders { get; set; } /// /// Gets or sets a value indicating whether to use the provided instance as is, @@ -75,41 +69,7 @@ public sealed class ChatClientAgentOptions Name = this.Name, Description = this.Description, ChatOptions = this.ChatOptions?.Clone(), - ChatHistoryProviderFactory = this.ChatHistoryProviderFactory, - AIContextProviderFactory = this.AIContextProviderFactory, + ChatHistoryProvider = this.ChatHistoryProvider, + AIContextProviders = this.AIContextProviders is null ? null : new List(this.AIContextProviders), }; - - /// - /// Context object passed to the to create a new instance of . - /// - public sealed class AIContextProviderFactoryContext - { - /// - /// Gets or sets the serialized state of the , if any. - /// - /// if there is no state, e.g. when the is first created. - public JsonElement SerializedState { get; set; } - - /// - /// Gets or sets the JSON serialization options to use when deserializing the . - /// - public JsonSerializerOptions? JsonSerializerOptions { get; set; } - } - - /// - /// Context object passed to the to create a new instance of . - /// - public sealed class ChatHistoryProviderFactoryContext - { - /// - /// Gets or sets the serialized state of the , if any. - /// - /// if there is no state, e.g. when the is first created. - public JsonElement SerializedState { get; set; } - - /// - /// Gets or sets the JSON serialization options to use when deserializing the . - /// - public JsonSerializerOptions? JsonSerializerOptions { get; set; } - } } diff --git a/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgentSession.cs b/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgentSession.cs index 1a79ae64d1..400bfbcaf6 100644 --- a/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgentSession.cs +++ b/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgentSession.cs @@ -3,8 +3,7 @@ using System; using System.Diagnostics; using System.Text.Json; -using System.Threading; -using System.Threading.Tasks; +using System.Text.Json.Serialization; using Microsoft.Shared.Diagnostics; namespace Microsoft.Agents.AI; @@ -15,8 +14,6 @@ namespace Microsoft.Agents.AI; [DebuggerDisplay("{DebuggerDisplay,nq}")] public sealed class ChatClientAgentSession : AgentSession { - private ChatHistoryProvider? _chatHistoryProvider; - /// /// Initializes a new instance of the class. /// @@ -24,29 +21,30 @@ public sealed class ChatClientAgentSession : AgentSession { } + [JsonConstructor] + internal ChatClientAgentSession(string? conversationId, AgentSessionStateBag? stateBag) : base(stateBag ?? new()) + { + this.ConversationId = conversationId; + } + /// - /// Gets or sets the ID of the underlying service thread to support cases where the chat history is stored by the agent service. + /// Gets or sets the ID of the underlying service chat history to support cases where the chat history is stored by the agent service. /// /// /// - /// Note that either or may be set, but not both. - /// If is not null, setting will throw an - /// exception. - /// - /// /// This property may be null in the following cases: /// - /// The thread stores messages via the and not in the agent service. - /// This thread object is new and a server managed thread has not yet been created in the agent service. + /// The agent stores messages via a and not in the agent service. + /// This session object is new and server managed chat history has not yet been created in the agent service. /// /// /// - /// The id may also change over time where the id is pointing at a - /// agent service managed thread, and the default behavior of a service is - /// to fork the thread with each iteration. + /// The id may also change over time where the id is pointing at + /// agent service managed chat history, and the default behavior of a service is + /// to fork the chat history with each iteration. /// /// - /// Attempted to set a conversation ID but a is already set. + [JsonPropertyName("conversationId")] public string? ConversationId { get; @@ -57,149 +55,37 @@ public sealed class ChatClientAgentSession : AgentSession return; } - if (this._chatHistoryProvider is not null) - { - // If we have a ChatHistoryProvider already, we shouldn't switch the session to use a conversation id - // since it means that the session contents will essentially be deleted, and the session will not work - // with the original agent anymore. - throw new InvalidOperationException("Only the ConversationId or ChatHistoryProvider may be set, but not both and switching from one to another is not supported."); - } - field = Throw.IfNullOrWhitespace(value); } } - /// - /// Gets or sets the used by this thread, for cases where messages should be stored in a custom location. - /// - /// - /// - /// Note that either or may be set, but not both. - /// If is not null, and is set, - /// will be reverted to null, and vice versa. - /// - /// - /// This property may be null in the following cases: - /// - /// The thread stores messages in the agent service and just has an id to the remove thread, instead of in an . - /// This thread object is new it is not yet clear whether it will be backed by a server managed thread or an . - /// - /// - /// - public ChatHistoryProvider? ChatHistoryProvider - { - get => this._chatHistoryProvider; - internal set - { - if (this._chatHistoryProvider is null && value is null) - { - return; - } - - if (!string.IsNullOrWhiteSpace(this.ConversationId)) - { - // If we have a conversation id already, we shouldn't switch the session to use a ChatHistoryProvider - // since it means that the session will not work with the original agent anymore. - throw new InvalidOperationException("Only the ConversationId or ChatHistoryProvider may be set, but not both and switching from one to another is not supported."); - } - - this._chatHistoryProvider = Throw.IfNull(value); - } - } - - /// - /// Gets or sets the used by this thread to provide additional context to the AI model before each invocation. - /// - public AIContextProvider? AIContextProvider { get; internal set; } - /// /// Creates a new instance of the class from previously serialized state. /// /// A representing the serialized state of the session. - /// Optional settings for customizing the JSON deserialization process. - /// - /// An optional factory function to create a custom from its serialized state. - /// If not provided, the default will be used. - /// - /// - /// An optional factory function to create a custom from its serialized state. - /// If not provided, no context provider will be configured. - /// - /// The to monitor for cancellation requests. - /// A task representing the asynchronous operation. The task result contains the deserialized . - internal static async Task DeserializeAsync( - JsonElement serializedState, - JsonSerializerOptions? jsonSerializerOptions = null, - Func>? chatHistoryProviderFactory = null, - Func>? aiContextProviderFactory = null, - CancellationToken cancellationToken = default) + /// Optional JSON serialization options to use instead of the default options. + /// The deserialized . + internal static ChatClientAgentSession Deserialize(JsonElement serializedState, JsonSerializerOptions? jsonSerializerOptions = null) { if (serializedState.ValueKind != JsonValueKind.Object) { throw new ArgumentException("The serialized session state must be a JSON object.", nameof(serializedState)); } - var state = serializedState.Deserialize( - AgentJsonUtilities.DefaultOptions.GetTypeInfo(typeof(SessionState))) as SessionState; - - var session = new ChatClientAgentSession(); - - session.AIContextProvider = aiContextProviderFactory is not null - ? await aiContextProviderFactory.Invoke(state?.AIContextProviderState ?? default, jsonSerializerOptions, cancellationToken).ConfigureAwait(false) - : null; - - if (state?.ConversationId is string sessionId) - { - session.ConversationId = sessionId; - - // Since we have an ID, we should not have a ChatHistoryProvider and we can return here. - return session; - } - - session._chatHistoryProvider = - chatHistoryProviderFactory is not null - ? await chatHistoryProviderFactory.Invoke(state?.ChatHistoryProviderState ?? default, jsonSerializerOptions, cancellationToken).ConfigureAwait(false) - : new InMemoryChatHistoryProvider(state?.ChatHistoryProviderState ?? default, jsonSerializerOptions); // default to an in-memory ChatHistoryProvider - - return session; + var jso = jsonSerializerOptions ?? AgentJsonUtilities.DefaultOptions; + return serializedState.Deserialize(jso.GetTypeInfo(typeof(ChatClientAgentSession))) as ChatClientAgentSession + ?? new ChatClientAgentSession(); } /// internal JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null) { - JsonElement? chatHistoryProviderState = this._chatHistoryProvider?.Serialize(jsonSerializerOptions); - - JsonElement? aiContextProviderState = this.AIContextProvider?.Serialize(jsonSerializerOptions); - - var state = new SessionState - { - ConversationId = this.ConversationId, - ChatHistoryProviderState = chatHistoryProviderState is { ValueKind: not JsonValueKind.Undefined } ? chatHistoryProviderState : null, - AIContextProviderState = aiContextProviderState is { ValueKind: not JsonValueKind.Undefined } ? aiContextProviderState : null, - }; - - return JsonSerializer.SerializeToElement(state, AgentJsonUtilities.DefaultOptions.GetTypeInfo(typeof(SessionState))); + var jso = jsonSerializerOptions ?? AgentJsonUtilities.DefaultOptions; + return JsonSerializer.SerializeToElement(this, jso.GetTypeInfo(typeof(ChatClientAgentSession))); } - /// - public override object? GetService(Type serviceType, object? serviceKey = null) => - base.GetService(serviceType, serviceKey) - ?? this.AIContextProvider?.GetService(serviceType, serviceKey) - ?? this.ChatHistoryProvider?.GetService(serviceType, serviceKey); - [DebuggerBrowsable(DebuggerBrowsableState.Never)] private string DebuggerDisplay => - this.ConversationId is { } conversationId ? $"ConversationId = {conversationId}" : - this._chatHistoryProvider is InMemoryChatHistoryProvider inMemoryChatHistoryProvider ? $"Count = {inMemoryChatHistoryProvider.Count}" : - this._chatHistoryProvider is { } chatHistoryProvider ? $"ChatHistoryProvider = {chatHistoryProvider.GetType().Name}" : - "Count = 0"; - - internal sealed class SessionState - { - public string? ConversationId { get; set; } - - public JsonElement? ChatHistoryProviderState { get; set; } - - public JsonElement? AIContextProviderState { get; set; } - } + this.ConversationId is { } conversationId ? $"ConversationId = {conversationId}, StateBag Count = {this.StateBag.Count}" : + $"StateBag Count = {this.StateBag.Count}"; } diff --git a/dotnet/src/Microsoft.Agents.AI/Memory/ChatHistoryMemoryProvider.cs b/dotnet/src/Microsoft.Agents.AI/Memory/ChatHistoryMemoryProvider.cs index b8e495152c..9d163f79cf 100644 --- a/dotnet/src/Microsoft.Agents.AI/Memory/ChatHistoryMemoryProvider.cs +++ b/dotnet/src/Microsoft.Agents.AI/Memory/ChatHistoryMemoryProvider.cs @@ -4,7 +4,6 @@ using System; using System.Collections.Generic; using System.Linq; using System.Linq.Expressions; -using System.Text.Json; using System.Threading; using System.Threading.Tasks; using Microsoft.Extensions.AI; @@ -42,17 +41,24 @@ public sealed class ChatHistoryMemoryProvider : AIContextProvider, IDisposable private const string DefaultFunctionToolName = "Search"; private const string DefaultFunctionToolDescription = "Allows searching for related previous chat history to help answer the user question."; + private static IEnumerable DefaultExternalOnlyFilter(IEnumerable messages) + => messages.Where(m => m.GetAgentRequestMessageSourceType() == AgentRequestMessageSourceType.External); + +#pragma warning disable CA2213 // VectorStore is not owned by this class - caller is responsible for disposal private readonly VectorStore _vectorStore; +#pragma warning restore CA2213 private readonly VectorStoreCollection> _collection; private readonly int _maxResults; private readonly string _contextPrompt; private readonly bool _enableSensitiveTelemetryData; private readonly ChatHistoryMemoryProviderOptions.SearchBehavior _searchTime; - private readonly AITool[] _tools; + private readonly string _toolName; + private readonly string _toolDescription; private readonly ILogger? _logger; - - private readonly ChatHistoryMemoryProviderScope _storageScope; - private readonly ChatHistoryMemoryProviderScope _searchScope; + private readonly string _stateKey; + private readonly Func _stateInitializer; + private readonly Func, IEnumerable> _searchInputMessageFilter; + private readonly Func, IEnumerable> _storageInputMessageFilter; private bool _collectionInitialized; private readonly SemaphoreSlim _initializationLock = new(1, 1); @@ -64,93 +70,32 @@ public sealed class ChatHistoryMemoryProvider : AIContextProvider, IDisposable /// The vector store to use for storing and retrieving chat history. /// The name of the collection for storing chat history in the vector store. /// The number of dimensions to use for the chat history vector store embeddings. - /// Optional values to scope the chat history storage with. - /// Optional values to scope the chat history search with. Where values are null, no filtering is done using those values. Defaults to if not provided. + /// A delegate that initializes the provider state on the first invocation, providing the storage and search scopes. /// Optional configuration options. /// Optional logger factory. - /// Thrown when is . + /// Thrown when or is . public ChatHistoryMemoryProvider( VectorStore vectorStore, string collectionName, int vectorDimensions, - ChatHistoryMemoryProviderScope storageScope, - ChatHistoryMemoryProviderScope? searchScope = null, + Func stateInitializer, ChatHistoryMemoryProviderOptions? options = null, ILoggerFactory? loggerFactory = null) - : this( - vectorStore, - collectionName, - vectorDimensions, - new ChatHistoryMemoryProviderState - { - StorageScope = new(Throw.IfNull(storageScope)), - SearchScope = searchScope ?? new(storageScope), - }, - options, - loggerFactory) { - } + this._vectorStore = Throw.IfNull(vectorStore); + this._stateInitializer = Throw.IfNull(stateInitializer); - /// - /// Initializes a new instance of the class from previously serialized state. - /// - /// The vector store to use for storing and retrieving chat history. - /// The name of the collection for storing chat history in the vector store. - /// The number of dimensions to use for the chat history vector store embeddings. - /// A representing the serialized state of the provider. - /// Optional settings for customizing the JSON deserialization process. - /// Optional configuration options. - /// Optional logger factory. - public ChatHistoryMemoryProvider( - VectorStore vectorStore, - string collectionName, - int vectorDimensions, - JsonElement serializedState, - JsonSerializerOptions? jsonSerializerOptions = null, - ChatHistoryMemoryProviderOptions? options = null, - ILoggerFactory? loggerFactory = null) - : this( - vectorStore, - collectionName, - vectorDimensions, - DeserializeState(serializedState, jsonSerializerOptions), - options, - loggerFactory) - { - } - - private ChatHistoryMemoryProvider( - VectorStore vectorStore, - string collectionName, - int vectorDimensions, - ChatHistoryMemoryProviderState? state = null, - ChatHistoryMemoryProviderOptions? options = null, - ILoggerFactory? loggerFactory = null) - { - this._vectorStore = vectorStore ?? throw new ArgumentNullException(nameof(vectorStore)); options ??= new ChatHistoryMemoryProviderOptions(); this._maxResults = options.MaxResults.HasValue ? Throw.IfLessThanOrEqual(options.MaxResults.Value, 0) : DefaultMaxResults; this._contextPrompt = options.ContextPrompt ?? DefaultContextPrompt; this._enableSensitiveTelemetryData = options.EnableSensitiveTelemetryData; this._searchTime = options.SearchTime; + this._stateKey = options.StateKey ?? base.StateKey; this._logger = loggerFactory?.CreateLogger(); - - if (state == null || state.StorageScope == null || state.SearchScope == null) - { - throw new InvalidOperationException($"The {nameof(ChatHistoryMemoryProvider)} state did not contain the required scope properties."); - } - - this._storageScope = state.StorageScope; - this._searchScope = state.SearchScope; - - // Create on-demand search tool (only used when behavior is OnDemandFunctionCalling) - this._tools = - [ - AIFunctionFactory.Create( - (Func>)this.SearchTextAsync, - name: options.FunctionToolName ?? DefaultFunctionToolName, - description: options.FunctionToolDescription ?? DefaultFunctionToolDescription) - ]; + this._toolName = options.FunctionToolName ?? DefaultFunctionToolName; + this._toolDescription = options.FunctionToolDescription ?? DefaultFunctionToolDescription; + this._searchInputMessageFilter = options.SearchInputMessageFilter ?? DefaultExternalOnlyFilter; + this._storageInputMessageFilter = options.StorageInputMessageFilter ?? DefaultExternalOnlyFilter; // Create a definition so that we can use the dimensions provided at runtime. var definition = new VectorStoreCollectionDefinition @@ -174,41 +119,93 @@ public sealed class ChatHistoryMemoryProvider : AIContextProvider, IDisposable this._collection = this._vectorStore.GetDynamicCollection(Throw.IfNullOrWhitespace(collectionName), definition); } + /// + public override string StateKey => this._stateKey; + + /// + /// Gets the state from the session's StateBag, or initializes it using the StateInitializer if not present. + /// + /// The agent session containing the StateBag. + /// The provider state, or null if no session is available. + private State? GetOrInitializeState(AgentSession? session) + { + if (session?.StateBag.TryGetValue(this._stateKey, out var state, AgentJsonUtilities.DefaultOptions) is true && state is not null) + { + return state; + } + + state = this._stateInitializer(session); + if (state is not null && session is not null) + { + session.StateBag.SetValue(this._stateKey, state, AgentJsonUtilities.DefaultOptions); + } + + return state; + } + /// protected override async ValueTask InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default) { _ = Throw.IfNull(context); + var inputContext = context.AIContext; + var state = this.GetOrInitializeState(context.Session); + var searchScope = state?.SearchScope ?? new ChatHistoryMemoryProviderScope(); + if (this._searchTime == ChatHistoryMemoryProviderOptions.SearchBehavior.OnDemandFunctionCalling) { - // Expose search tool for on-demand invocation by the model - return new AIContext { Tools = this._tools }; + Task InlineSearchAsync(string userQuestion, CancellationToken ct) + => this.SearchTextAsync(userQuestion, searchScope, ct); + + // Create on-demand search tool (only used when behavior is OnDemandFunctionCalling) + AITool[] tools = + [ + AIFunctionFactory.Create( + InlineSearchAsync, + name: this._toolName, + description: this._toolDescription) + ]; + + // Expose search tool for on-demand invocation by the model, accumulated with the input context + return new AIContext + { + Instructions = inputContext.Instructions, + Messages = inputContext.Messages, + Tools = (inputContext.Tools ?? []).Concat(tools) + }; } try { // Get the text from the current request messages - var requestText = string.Join("\n", context.RequestMessages - .Where(m => m.GetAgentRequestMessageSourceType() == AgentRequestMessageSourceType.External) + var requestText = string.Join("\n", + this._searchInputMessageFilter(inputContext.Messages ?? []) .Where(m => m != null && !string.IsNullOrWhiteSpace(m.Text)) .Select(m => m.Text)); if (string.IsNullOrWhiteSpace(requestText)) { - return new AIContext(); + return inputContext; } // Search for relevant chat history - var contextText = await this.SearchTextAsync(requestText, cancellationToken).ConfigureAwait(false); + var contextText = await this.SearchTextAsync(requestText, searchScope, cancellationToken).ConfigureAwait(false); if (string.IsNullOrWhiteSpace(contextText)) { - return new AIContext(); + return inputContext; } return new AIContext { - Messages = [new ChatMessage(ChatRole.User, contextText)] + Instructions = inputContext.Instructions, + Messages = + (inputContext.Messages ?? []) + .Concat( + [ + new ChatMessage(ChatRole.User, contextText).WithAgentRequestMessageSource(AgentRequestMessageSourceType.AIContextProvider, this.GetType().FullName!) + ]), + Tools = inputContext.Tools }; } catch (Exception ex) @@ -218,13 +215,13 @@ public sealed class ChatHistoryMemoryProvider : AIContextProvider, IDisposable this._logger.LogError( ex, "ChatHistoryMemoryProvider: Failed to search for chat history due to error. ApplicationId: '{ApplicationId}', AgentId: '{AgentId}', SessionId: '{SessionId}', UserId: '{UserId}'.", - this._searchScope.ApplicationId, - this._searchScope.AgentId, - this._searchScope.SessionId, - this.SanitizeLogData(this._searchScope.UserId)); + searchScope.ApplicationId, + searchScope.AgentId, + searchScope.SessionId, + this.SanitizeLogData(searchScope.UserId)); } - return new AIContext(); + return inputContext; } } @@ -239,13 +236,15 @@ public sealed class ChatHistoryMemoryProvider : AIContextProvider, IDisposable return; } + var state = this.GetOrInitializeState(context.Session); + var storageScope = state?.StorageScope ?? new ChatHistoryMemoryProviderScope(); + try { // Ensure the collection is initialized var collection = await this.EnsureCollectionExistsAsync(cancellationToken).ConfigureAwait(false); - List> itemsToStore = context.RequestMessages - .Where(m => m.GetAgentRequestMessageSourceType() == AgentRequestMessageSourceType.External) + List> itemsToStore = this._storageInputMessageFilter(context.RequestMessages) .Concat(context.ResponseMessages ?? []) .Select(message => new Dictionary { @@ -253,10 +252,10 @@ public sealed class ChatHistoryMemoryProvider : AIContextProvider, IDisposable ["Role"] = message.Role.ToString(), ["MessageId"] = message.MessageId, ["AuthorName"] = message.AuthorName, - ["ApplicationId"] = this._storageScope?.ApplicationId, - ["AgentId"] = this._storageScope?.AgentId, - ["UserId"] = this._storageScope?.UserId, - ["SessionId"] = this._storageScope?.SessionId, + ["ApplicationId"] = storageScope.ApplicationId, + ["AgentId"] = storageScope.AgentId, + ["UserId"] = storageScope.UserId, + ["SessionId"] = storageScope.SessionId, ["Content"] = message.Text, ["CreatedAt"] = message.CreatedAt?.ToString("O") ?? DateTimeOffset.UtcNow.ToString("O"), ["ContentEmbedding"] = message.Text, @@ -275,10 +274,10 @@ public sealed class ChatHistoryMemoryProvider : AIContextProvider, IDisposable this._logger.LogError( ex, "ChatHistoryMemoryProvider: Failed to add messages to chat history vector store due to error. ApplicationId: '{ApplicationId}', AgentId: '{AgentId}', SessionId: '{SessionId}', UserId: '{UserId}'.", - this._searchScope.ApplicationId, - this._searchScope.AgentId, - this._searchScope.SessionId, - this.SanitizeLogData(this._searchScope.UserId)); + storageScope.ApplicationId, + storageScope.AgentId, + storageScope.SessionId, + this.SanitizeLogData(storageScope.UserId)); } } } @@ -287,16 +286,17 @@ public sealed class ChatHistoryMemoryProvider : AIContextProvider, IDisposable /// Function callable by the AI model (when enabled) to perform an ad-hoc chat history search. /// /// The query text. + /// The scope to filter search results with. /// Cancellation token. /// Formatted search results (may be empty). - internal async Task SearchTextAsync(string userQuestion, CancellationToken cancellationToken = default) + private async Task SearchTextAsync(string userQuestion, ChatHistoryMemoryProviderScope searchScope, CancellationToken cancellationToken = default) { if (string.IsNullOrWhiteSpace(userQuestion)) { return string.Empty; } - var results = await this.SearchChatHistoryAsync(userQuestion, this._maxResults, cancellationToken).ConfigureAwait(false); + var results = await this.SearchChatHistoryAsync(userQuestion, searchScope, this._maxResults, cancellationToken).ConfigureAwait(false); if (!results.Any()) { return string.Empty; @@ -317,10 +317,10 @@ public sealed class ChatHistoryMemoryProvider : AIContextProvider, IDisposable "ChatHistoryMemoryProvider: Search Results\nInput:{Input}\nOutput:{MessageText}\n ApplicationId: '{ApplicationId}', AgentId: '{AgentId}', SessionId: '{SessionId}', UserId: '{UserId}'.", this.SanitizeLogData(userQuestion), this.SanitizeLogData(formatted), - this._searchScope.ApplicationId, - this._searchScope.AgentId, - this._searchScope.SessionId, - this.SanitizeLogData(this._searchScope.UserId)); + searchScope.ApplicationId, + searchScope.AgentId, + searchScope.SessionId, + this.SanitizeLogData(searchScope.UserId)); } return formatted; @@ -330,11 +330,13 @@ public sealed class ChatHistoryMemoryProvider : AIContextProvider, IDisposable /// Searches for relevant chat history items based on the provided query text. /// /// The text to search for. + /// The scope to filter search results with. /// The maximum number of results to return. /// The cancellation token. /// A list of relevant chat history items. private async Task>> SearchChatHistoryAsync( string queryText, + ChatHistoryMemoryProviderScope searchScope, int top, CancellationToken cancellationToken = default) { @@ -345,10 +347,10 @@ public sealed class ChatHistoryMemoryProvider : AIContextProvider, IDisposable var collection = await this.EnsureCollectionExistsAsync(cancellationToken).ConfigureAwait(false); - string? applicationId = this._searchScope.ApplicationId; - string? agentId = this._searchScope.AgentId; - string? userId = this._searchScope.UserId; - string? sessionId = this._searchScope.SessionId; + string? applicationId = searchScope.ApplicationId; + string? agentId = searchScope.AgentId; + string? userId = searchScope.UserId; + string? sessionId = searchScope.SessionId; Expression, bool>>? filter = null; if (applicationId != null) @@ -401,10 +403,10 @@ public sealed class ChatHistoryMemoryProvider : AIContextProvider, IDisposable this._logger.LogInformation( "ChatHistoryMemoryProvider: Retrieved {Count} search results. ApplicationId: '{ApplicationId}', AgentId: '{AgentId}', SessionId: '{SessionId}', UserId: '{UserId}'.", results.Count, - this._searchScope.ApplicationId, - this._searchScope.AgentId, - this._searchScope.SessionId, - this.SanitizeLogData(this._searchScope.UserId)); + searchScope.ApplicationId, + searchScope.AgentId, + searchScope.SessionId, + this.SanitizeLogData(searchScope.UserId)); } return results; @@ -465,39 +467,32 @@ public sealed class ChatHistoryMemoryProvider : AIContextProvider, IDisposable GC.SuppressFinalize(this); } - /// - /// Serializes the current provider state to a including storage and search scopes. - /// - /// Optional serializer options. - /// Serialized provider state. - public override JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null) - { - var state = new ChatHistoryMemoryProviderState - { - StorageScope = this._storageScope, - SearchScope = this._searchScope, - }; - - var jso = jsonSerializerOptions ?? AgentJsonUtilities.DefaultOptions; - return JsonSerializer.SerializeToElement(state, jso.GetTypeInfo(typeof(ChatHistoryMemoryProviderState))); - } - - private static ChatHistoryMemoryProviderState? DeserializeState(JsonElement serializedState, JsonSerializerOptions? jsonSerializerOptions) - { - if (serializedState.ValueKind != JsonValueKind.Object) - { - return null; - } - - var jso = jsonSerializerOptions ?? AgentJsonUtilities.DefaultOptions; - return serializedState.Deserialize(jso.GetTypeInfo(typeof(ChatHistoryMemoryProviderState))) as ChatHistoryMemoryProviderState; - } - private string? SanitizeLogData(string? data) => this._enableSensitiveTelemetryData ? data : ""; - internal sealed class ChatHistoryMemoryProviderState + /// + /// Represents the state of a stored in the . + /// + public sealed class State { - public ChatHistoryMemoryProviderScope? StorageScope { get; set; } - public ChatHistoryMemoryProviderScope? SearchScope { get; set; } + /// + /// Initializes a new instance of the class with the specified storage and search scopes. + /// + /// The scope to use when storing chat history messages. + /// The scope to use when searching for relevant chat history messages. If null, the storage scope will be used for searching as well. + public State(ChatHistoryMemoryProviderScope storageScope, ChatHistoryMemoryProviderScope? searchScope = null) + { + this.StorageScope = Throw.IfNull(storageScope); + this.SearchScope = searchScope ?? storageScope; + } + + /// + /// Gets or sets the scope used when storing chat history messages. + /// + public ChatHistoryMemoryProviderScope StorageScope { get; } + + /// + /// Gets or sets the scope used when searching chat history messages. + /// + public ChatHistoryMemoryProviderScope SearchScope { get; } } } diff --git a/dotnet/src/Microsoft.Agents.AI/Memory/ChatHistoryMemoryProviderOptions.cs b/dotnet/src/Microsoft.Agents.AI/Memory/ChatHistoryMemoryProviderOptions.cs index e09de68a59..6c92a426f3 100644 --- a/dotnet/src/Microsoft.Agents.AI/Memory/ChatHistoryMemoryProviderOptions.cs +++ b/dotnet/src/Microsoft.Agents.AI/Memory/ChatHistoryMemoryProviderOptions.cs @@ -1,5 +1,9 @@ // Copyright (c) Microsoft. All rights reserved. +using System; +using System.Collections.Generic; +using Microsoft.Extensions.AI; + namespace Microsoft.Agents.AI; /// @@ -44,6 +48,35 @@ public sealed class ChatHistoryMemoryProviderOptions /// Defaults to . public bool EnableSensitiveTelemetryData { get; set; } + /// + /// Gets or sets the key used to store provider state in the . + /// + /// + /// Defaults to the provider's type name. Override this if you need multiple + /// instances with separate state in the same session. + /// + public string? StateKey { get; set; } + + /// + /// Gets or sets an optional filter function applied to request messages when constructing the search text to use + /// to search for relevant chat history during . + /// + /// + /// When , the provider defaults to including only + /// messages. + /// + public Func, IEnumerable>? SearchInputMessageFilter { get; set; } + + /// + /// Gets or sets an optional filter function applied to request messages when storing recent chat history + /// during . + /// + /// + /// When , the provider defaults to including only + /// messages. + /// + public Func, IEnumerable>? StorageInputMessageFilter { get; set; } + /// /// Behavior choices for the provider. /// diff --git a/dotnet/src/Microsoft.Agents.AI/TextSearchProvider.cs b/dotnet/src/Microsoft.Agents.AI/TextSearchProvider.cs index f29aadf808..f038fa3c38 100644 --- a/dotnet/src/Microsoft.Agents.AI/TextSearchProvider.cs +++ b/dotnet/src/Microsoft.Agents.AI/TextSearchProvider.cs @@ -4,7 +4,6 @@ using System; using System.Collections.Generic; using System.Linq; using System.Text; -using System.Text.Json; using System.Threading; using System.Threading.Tasks; using Microsoft.Extensions.AI; @@ -40,30 +39,31 @@ public sealed class TextSearchProvider : AIContextProvider private const string DefaultContextPrompt = "## Additional Context\nConsider the following information from source documents when responding to the user:"; private const string DefaultCitationsPrompt = "Include citations to the source document with document name and link if document name and link is available."; + private static IEnumerable DefaultExternalOnlyFilter(IEnumerable messages) + => messages.Where(m => m.GetAgentRequestMessageSourceType() == AgentRequestMessageSourceType.External); + private readonly Func>> _searchAsync; private readonly ILogger? _logger; private readonly AITool[] _tools; - private readonly Queue _recentMessagesText; private readonly List _recentMessageRolesIncluded; private readonly int _recentMessageMemoryLimit; private readonly TextSearchProviderOptions.TextSearchBehavior _searchTime; private readonly string _contextPrompt; private readonly string _citationsPrompt; + private readonly string _stateKey; private readonly Func, string>? _contextFormatter; + private readonly Func, IEnumerable> _searchInputMessageFilter; + private readonly Func, IEnumerable> _storageInputMessageFilter; /// /// Initializes a new instance of the class. /// /// Delegate that executes the search logic. Must not be . - /// A representing the serialized provider state. - /// Optional serializer options (unused - source generated context is used). /// Optional configuration options. /// Optional logger factory. /// Thrown when is . public TextSearchProvider( Func>> searchAsync, - JsonElement serializedState, - JsonSerializerOptions? jsonSerializerOptions = null, TextSearchProviderOptions? options = null, ILoggerFactory? loggerFactory = null) { @@ -75,26 +75,10 @@ public sealed class TextSearchProvider : AIContextProvider this._searchTime = options?.SearchTime ?? TextSearchProviderOptions.TextSearchBehavior.BeforeAIInvoke; this._contextPrompt = options?.ContextPrompt ?? DefaultContextPrompt; this._citationsPrompt = options?.CitationsPrompt ?? DefaultCitationsPrompt; + this._stateKey = options?.StateKey ?? base.StateKey; this._contextFormatter = options?.ContextFormatter; - - // Restore recent messages from serialized state if provided - List? restoredMessages = null; - if (serializedState.ValueKind is JsonValueKind.Null or JsonValueKind.Undefined) - { - this._recentMessagesText = new(); - } - else - { - var jso = jsonSerializerOptions ?? AgentJsonUtilities.DefaultOptions; - var state = serializedState.Deserialize(jso.GetTypeInfo(typeof(TextSearchProviderState))) as TextSearchProviderState; - if (state?.RecentMessagesText is { Count: > 0 }) - { - restoredMessages = state.RecentMessagesText; - } - - // Restore recent messages respecting the limit (may truncate if limit changed afterwards). - this._recentMessagesText = restoredMessages is null ? new() : new(restoredMessages.Take(this._recentMessageMemoryLimit)); - } + this._searchInputMessageFilter = options?.SearchInputMessageFilter ?? DefaultExternalOnlyFilter; + this._storageInputMessageFilter = options?.StorageInputMessageFilter ?? DefaultExternalOnlyFilter; // Create the on-demand search tool (only used if behavior is OnDemandFunctionCalling) this._tools = @@ -106,21 +90,35 @@ public sealed class TextSearchProvider : AIContextProvider ]; } + /// + public override string StateKey => this._stateKey; + /// protected override async ValueTask InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default) { + var inputContext = context.AIContext; + if (this._searchTime != TextSearchProviderOptions.TextSearchBehavior.BeforeAIInvoke) { - // Expose the search tool for on-demand invocation. - return new AIContext { Tools = this._tools }; // No automatic message injection. + // Expose the search tool for on-demand invocation, accumulated with the input context. + return new AIContext + { + Instructions = inputContext.Instructions, + Messages = inputContext.Messages, + Tools = (inputContext.Tools ?? []).Concat(this._tools) + }; } + // Retrieve recent messages from the session state bag. + var recentMessagesText = context.Session?.StateBag.GetValue(this._stateKey, AgentJsonUtilities.DefaultOptions)?.RecentMessagesText + ?? []; + // Aggregate text from memory + current request messages. var sbInput = new StringBuilder(); - var requestMessagesText = context.RequestMessages - .Where(m => m.GetAgentRequestMessageSourceType() == AgentRequestMessageSourceType.External) + var requestMessagesText = + this._searchInputMessageFilter(inputContext.Messages ?? []) .Where(x => !string.IsNullOrWhiteSpace(x?.Text)).Select(x => x.Text); - foreach (var messageText in this._recentMessagesText.Concat(requestMessagesText)) + foreach (var messageText in recentMessagesText.Concat(requestMessagesText)) { if (sbInput.Length > 0) { @@ -144,7 +142,7 @@ public sealed class TextSearchProvider : AIContextProvider if (materialized.Count == 0) { - return new AIContext(); + return inputContext; } // Format search results @@ -157,13 +155,20 @@ public sealed class TextSearchProvider : AIContextProvider return new AIContext { - Messages = [new ChatMessage(ChatRole.User, formatted) { AdditionalProperties = new AdditionalPropertiesDictionary() { ["IsTextSearchProviderOutput"] = true } }] + Instructions = inputContext.Instructions, + Messages = + (inputContext.Messages ?? []) + .Concat( + [ + new ChatMessage(ChatRole.User, formatted).WithAgentRequestMessageSource(AgentRequestMessageSourceType.AIContextProvider, this.GetType().FullName!) + ]), + Tools = inputContext.Tools }; } catch (Exception ex) { this._logger?.LogError(ex, "TextSearchProvider: Failed to search for data due to error"); - return new AIContext(); + return inputContext; } } @@ -176,58 +181,42 @@ public sealed class TextSearchProvider : AIContextProvider return default; // Memory disabled. } + if (context.Session is null) + { + return default; // No session to store state in. + } + if (context.InvokeException is not null) { return default; // Do not update memory on failed invocations. } - var messagesText = context.RequestMessages - .Where(m => m.GetAgentRequestMessageSourceType() == AgentRequestMessageSourceType.External) + // Retrieve existing recent messages from the session state bag. + var recentMessagesText = context.Session.StateBag.GetValue(this._stateKey, AgentJsonUtilities.DefaultOptions)?.RecentMessagesText + ?? []; + + var newMessagesText = this._storageInputMessageFilter(context.RequestMessages) .Concat(context.ResponseMessages ?? []) .Where(m => this._recentMessageRolesIncluded.Contains(m.Role) && - !string.IsNullOrWhiteSpace(m.Text) && - // Filter out any messages that were added by this class in InvokingAsync, since we don't want - // a feedback loop where previous search results are used to find new search results. - (m.AdditionalProperties == null || m.AdditionalProperties.TryGetValue("IsTextSearchProviderOutput", out bool isTextSearchProviderOutput) == false || !isTextSearchProviderOutput)) - .Select(m => m.Text) - .ToList(); - if (messagesText.Count > limit) - { - // If the current request/response exceeds the limit, only keep the most recent messages from it. - messagesText = messagesText.Skip(messagesText.Count - limit).ToList(); - } + !string.IsNullOrWhiteSpace(m.Text)) + .Select(m => m.Text); - foreach (var message in messagesText) - { - this._recentMessagesText.Enqueue(message); - } + // Combine existing messages with new messages, then take the most recent up to the limit. + var allMessages = recentMessagesText.Concat(newMessagesText).ToList(); + var updatedMessages = allMessages.Count > limit + ? allMessages.Skip(allMessages.Count - limit).ToList() + : allMessages; - while (this._recentMessagesText.Count > limit) - { - this._recentMessagesText.Dequeue(); - } + // Store updated state back to the session state bag. + context.Session.StateBag.SetValue( + this._stateKey, + new TextSearchProviderState { RecentMessagesText = updatedMessages }, + AgentJsonUtilities.DefaultOptions); return default; } - /// - /// Serializes the current provider state to a containing any overridden prompts or descriptions. - /// - /// Optional serializer options (ignored, source generated context is used). - /// A with overridden values, or default if nothing was overridden. - public override JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null) - { - // Only persist values that differ from defaults plus recent memory configuration & messages. - TextSearchProviderState state = new(); - if (this._recentMessageMemoryLimit > 0 && this._recentMessagesText.Count > 0) - { - state.RecentMessagesText = this._recentMessagesText.Take(this._recentMessageMemoryLimit).ToList(); - } - - return JsonSerializer.SerializeToElement(state, AgentJsonUtilities.DefaultOptions.GetTypeInfo(typeof(TextSearchProviderState))); - } - /// /// Function callable by the AI model (when enabled) to perform an ad-hoc search. /// diff --git a/dotnet/src/Microsoft.Agents.AI/TextSearchProviderOptions.cs b/dotnet/src/Microsoft.Agents.AI/TextSearchProviderOptions.cs index e90a6efa63..837470b776 100644 --- a/dotnet/src/Microsoft.Agents.AI/TextSearchProviderOptions.cs +++ b/dotnet/src/Microsoft.Agents.AI/TextSearchProviderOptions.cs @@ -59,6 +59,35 @@ public sealed class TextSearchProviderOptions /// public int RecentMessageMemoryLimit { get; set; } + /// + /// Gets or sets the key used to store provider state in the . + /// + /// + /// Defaults to the provider's type name. Override this if you need multiple + /// instances with separate state in the same session. + /// + public string? StateKey { get; set; } + + /// + /// Gets or sets an optional filter function applied to request messages when constructing the search input + /// text during . + /// + /// + /// When , the provider defaults to including only + /// messages. + /// + public Func, IEnumerable>? SearchInputMessageFilter { get; set; } + + /// + /// Gets or sets an optional filter function applied to request messages when updating the recent message + /// memory during . + /// + /// + /// When , the provider defaults to including only + /// messages. + /// + public Func, IEnumerable>? StorageInputMessageFilter { get; set; } + /// /// Gets or sets the list of types to filter recent messages to /// when deciding which recent messages to include when constructing the search input. diff --git a/dotnet/tests/AnthropicChatCompletion.IntegrationTests/AnthropicChatCompletionFixture.cs b/dotnet/tests/AnthropicChatCompletion.IntegrationTests/AnthropicChatCompletionFixture.cs index a0e4b64763..f36cb119d9 100644 --- a/dotnet/tests/AnthropicChatCompletion.IntegrationTests/AnthropicChatCompletionFixture.cs +++ b/dotnet/tests/AnthropicChatCompletion.IntegrationTests/AnthropicChatCompletionFixture.cs @@ -37,14 +37,14 @@ public class AnthropicChatCompletionFixture : IChatClientAgentFixture public async Task> GetChatHistoryAsync(AIAgent agent, AgentSession session) { - var typedSession = (ChatClientAgentSession)session; + var chatHistoryProvider = agent.GetService(); - if (typedSession.ChatHistoryProvider is null) + if (chatHistoryProvider is null) { return []; } - return (await typedSession.ChatHistoryProvider.InvokingAsync(new(agent, session, []))).ToList(); + return (await chatHistoryProvider.InvokingAsync(new(agent, session, []))).ToList(); } public Task CreateChatClientAgentAsync( diff --git a/dotnet/tests/AzureAI.IntegrationTests/AIProjectClientFixture.cs b/dotnet/tests/AzureAI.IntegrationTests/AIProjectClientFixture.cs index c655fd7a58..dd926174c0 100644 --- a/dotnet/tests/AzureAI.IntegrationTests/AIProjectClientFixture.cs +++ b/dotnet/tests/AzureAI.IntegrationTests/AIProjectClientFixture.cs @@ -48,12 +48,14 @@ public class AIProjectClientFixture : IChatClientAgentFixture return await this.GetChatHistoryFromResponsesChainAsync(chatClientSession.ConversationId); } - if (chatClientSession.ChatHistoryProvider is null) + var chatHistoryProvider = agent.GetService(); + + if (chatHistoryProvider is null) { return []; } - return (await chatClientSession.ChatHistoryProvider.InvokingAsync(new(agent, session, []))).ToList(); + return (await chatHistoryProvider.InvokingAsync(new(agent, session, []))).ToList(); } private async Task> GetChatHistoryFromResponsesChainAsync(string conversationId) diff --git a/dotnet/tests/Microsoft.Agents.AI.A2A.UnitTests/A2AAgentSessionTests.cs b/dotnet/tests/Microsoft.Agents.AI.A2A.UnitTests/A2AAgentSessionTests.cs index 66018e0131..8c3e89adf6 100644 --- a/dotnet/tests/Microsoft.Agents.AI.A2A.UnitTests/A2AAgentSessionTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.A2A.UnitTests/A2AAgentSessionTests.cs @@ -21,10 +21,28 @@ public sealed class A2AAgentSessionTests // Act JsonElement serialized = originalSession.Serialize(); - A2AAgentSession deserializedSession = new(serialized); + A2AAgentSession deserializedSession = A2AAgentSession.Deserialize(serialized); // Assert Assert.Equal(originalSession.ContextId, deserializedSession.ContextId); Assert.Equal(originalSession.TaskId, deserializedSession.TaskId); } + + [Fact] + public void Constructor_RoundTrip_SerializationPreservesStateBag() + { + // Arrange + A2AAgentSession originalSession = new() { ContextId = "ctx-1", TaskId = "task-1" }; + originalSession.StateBag.SetValue("testKey", "testValue"); + + // Act + JsonElement serialized = originalSession.Serialize(); + A2AAgentSession deserializedSession = A2AAgentSession.Deserialize(serialized); + + // Assert + Assert.Equal("ctx-1", deserializedSession.ContextId); + Assert.Equal("task-1", deserializedSession.TaskId); + Assert.True(deserializedSession.StateBag.TryGetValue("testKey", out var value)); + Assert.Equal("testValue", value); + } } diff --git a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/AIContextProviderTests.cs b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/AIContextProviderTests.cs index aa41a03efc..14a0f81e08 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/AIContextProviderTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/AIContextProviderTests.cs @@ -3,7 +3,6 @@ using System; using System.Collections.Generic; using System.Collections.ObjectModel; -using System.Linq; using System.Threading; using System.Threading.Tasks; using Microsoft.Extensions.AI; @@ -16,110 +15,6 @@ public class AIContextProviderTests private static readonly AIAgent s_mockAgent = new Mock().Object; private static readonly AgentSession s_mockSession = new Mock().Object; - #region InvokingAsync Message Stamping Tests - - [Fact] - public async Task InvokingAsync_StampsMessagesWithSourceTypeAndSourceIdAsync() - { - // Arrange - var provider = new TestAIContextProviderWithMessages(); - var context = new AIContextProvider.InvokingContext(s_mockAgent, s_mockSession, [new ChatMessage(ChatRole.User, "Request")]); - - // Act - AIContext aiContext = await provider.InvokingAsync(context); - - // Assert - Assert.NotNull(aiContext.Messages); - ChatMessage message = aiContext.Messages.Single(); - Assert.NotNull(message.AdditionalProperties); - Assert.True(message.AdditionalProperties.TryGetValue(AgentRequestMessageSourceAttribution.AdditionalPropertiesKey, out object? attribution)); - var typedAttribution = Assert.IsType(attribution); - Assert.Equal(AgentRequestMessageSourceType.AIContextProvider, typedAttribution.SourceType); - Assert.Equal(typeof(TestAIContextProviderWithMessages).FullName, typedAttribution.SourceId); - } - - [Fact] - public async Task InvokingAsync_WithCustomSourceId_StampsMessagesWithCustomSourceIdAsync() - { - // Arrange - const string CustomSourceId = "CustomContextSource"; - var provider = new TestAIContextProviderWithCustomSource(CustomSourceId); - var context = new AIContextProvider.InvokingContext(s_mockAgent, s_mockSession, [new ChatMessage(ChatRole.User, "Request")]); - - // Act - AIContext aiContext = await provider.InvokingAsync(context); - - // Assert - Assert.NotNull(aiContext.Messages); - ChatMessage message = aiContext.Messages.Single(); - Assert.NotNull(message.AdditionalProperties); - Assert.True(message.AdditionalProperties.TryGetValue(AgentRequestMessageSourceAttribution.AdditionalPropertiesKey, out object? attribution)); - var typedAttribution = Assert.IsType(attribution); - Assert.Equal(AgentRequestMessageSourceType.AIContextProvider, typedAttribution.SourceType); - Assert.Equal(CustomSourceId, typedAttribution.SourceId); - } - - [Fact] - public async Task InvokingAsync_DoesNotReStampAlreadyStampedMessagesAsync() - { - // Arrange - var provider = new TestAIContextProviderWithPreStampedMessages(); - var context = new AIContextProvider.InvokingContext(s_mockAgent, s_mockSession, [new ChatMessage(ChatRole.User, "Request")]); - - // Act - AIContext aiContext = await provider.InvokingAsync(context); - - // Assert - Assert.NotNull(aiContext.Messages); - ChatMessage message = aiContext.Messages.Single(); - Assert.NotNull(message.AdditionalProperties); - Assert.True(message.AdditionalProperties.TryGetValue(AgentRequestMessageSourceAttribution.AdditionalPropertiesKey, out object? attribution)); - var typedAttribution = Assert.IsType(attribution); - Assert.Equal(AgentRequestMessageSourceType.AIContextProvider, typedAttribution.SourceType); - Assert.Equal(typeof(TestAIContextProviderWithPreStampedMessages).FullName, typedAttribution.SourceId); - } - - [Fact] - public async Task InvokingAsync_StampsMultipleMessagesAsync() - { - // Arrange - var provider = new TestAIContextProviderWithMultipleMessages(); - var context = new AIContextProvider.InvokingContext(s_mockAgent, s_mockSession, [new ChatMessage(ChatRole.User, "Request")]); - - // Act - AIContext aiContext = await provider.InvokingAsync(context); - - // Assert - Assert.NotNull(aiContext.Messages); - List messageList = aiContext.Messages.ToList(); - Assert.Equal(3, messageList.Count); - - foreach (ChatMessage message in messageList) - { - Assert.NotNull(message.AdditionalProperties); - Assert.True(message.AdditionalProperties.TryGetValue(AgentRequestMessageSourceAttribution.AdditionalPropertiesKey, out object? attribution)); - var typedAttribution = Assert.IsType(attribution); - Assert.Equal(AgentRequestMessageSourceType.AIContextProvider, typedAttribution.SourceType); - Assert.Equal(typeof(TestAIContextProviderWithMultipleMessages).FullName, typedAttribution.SourceId); - } - } - - [Fact] - public async Task InvokingAsync_WithNullMessages_ReturnsContextWithoutStampingAsync() - { - // Arrange - var provider = new TestAIContextProvider(); - var context = new AIContextProvider.InvokingContext(s_mockAgent, s_mockSession, [new ChatMessage(ChatRole.User, "Request")]); - - // Act - AIContext aiContext = await provider.InvokingAsync(context); - - // Assert - Assert.Null(aiContext.Messages); - } - - #endregion - #region Basic Tests [Fact] @@ -130,25 +25,12 @@ public class AIContextProviderTests var messages = new ReadOnlyCollection([]); // Act - ValueTask task = provider.InvokedAsync(new(s_mockAgent, s_mockSession, messages)); + ValueTask task = provider.InvokedAsync(new(s_mockAgent, s_mockSession, messages, [])); // Assert Assert.Equal(default, task); } - [Fact] - public void Serialize_ReturnsEmptyElement() - { - // Arrange - var provider = new TestAIContextProvider(); - - // Act - var actual = provider.Serialize(); - - // Assert - Assert.Equal(default, actual); - } - [Fact] public void InvokingContext_Constructor_ThrowsForNullMessages() { @@ -160,7 +42,7 @@ public class AIContextProviderTests public void InvokedContext_Constructor_ThrowsForNullMessages() { // Act & Assert - Assert.Throws(() => new AIContextProvider.InvokedContext(s_mockAgent, s_mockSession, null!)); + Assert.Throws(() => new AIContextProvider.InvokedContext(s_mockAgent, s_mockSession, null!, [])); } #endregion @@ -284,39 +166,33 @@ public class AIContextProviderTests #region InvokingContext Tests [Fact] - public void InvokingContext_RequestMessages_SetterThrowsForNull() + public void InvokingContext_Constructor_ThrowsForNullAIContext() { - // Arrange - var messages = new ReadOnlyCollection([new(ChatRole.User, "Hello")]); - var context = new AIContextProvider.InvokingContext(s_mockAgent, s_mockSession, messages); - // Act & Assert - Assert.Throws(() => context.RequestMessages = null!); + Assert.Throws(() => new AIContextProvider.InvokingContext(s_mockAgent, s_mockSession, null!)); } [Fact] - public void InvokingContext_RequestMessages_SetterRoundtrips() + public void InvokingContext_AIContext_ConstructorValueRoundtrips() { // Arrange - var initialMessages = new ReadOnlyCollection([new(ChatRole.User, "Hello")]); - var newMessages = new List { new(ChatRole.User, "New message") }; - var context = new AIContextProvider.InvokingContext(s_mockAgent, s_mockSession, initialMessages); + var aiContext = new AIContext { Messages = [new ChatMessage(ChatRole.User, "Hello")] }; // Act - context.RequestMessages = newMessages; + var context = new AIContextProvider.InvokingContext(s_mockAgent, s_mockSession, aiContext); // Assert - Assert.Same(newMessages, context.RequestMessages); + Assert.Same(aiContext, context.AIContext); } [Fact] public void InvokingContext_Agent_ReturnsConstructorValue() { // Arrange - var messages = new ReadOnlyCollection([new(ChatRole.User, "Hello")]); + var aiContext = new AIContext { Messages = [new ChatMessage(ChatRole.User, "Hello")] }; // Act - var context = new AIContextProvider.InvokingContext(s_mockAgent, s_mockSession, messages); + var context = new AIContextProvider.InvokingContext(s_mockAgent, s_mockSession, aiContext); // Assert Assert.Same(s_mockAgent, context.Agent); @@ -326,10 +202,10 @@ public class AIContextProviderTests public void InvokingContext_Session_ReturnsConstructorValue() { // Arrange - var messages = new ReadOnlyCollection([new(ChatRole.User, "Hello")]); + var aiContext = new AIContext { Messages = [new ChatMessage(ChatRole.User, "Hello")] }; // Act - var context = new AIContextProvider.InvokingContext(s_mockAgent, s_mockSession, messages); + var context = new AIContextProvider.InvokingContext(s_mockAgent, s_mockSession, aiContext); // Assert Assert.Same(s_mockSession, context.Session); @@ -339,10 +215,10 @@ public class AIContextProviderTests public void InvokingContext_Session_CanBeNull() { // Arrange - var messages = new ReadOnlyCollection([new(ChatRole.User, "Hello")]); + var aiContext = new AIContext { Messages = [new ChatMessage(ChatRole.User, "Hello")] }; // Act - var context = new AIContextProvider.InvokingContext(s_mockAgent, null, messages); + var context = new AIContextProvider.InvokingContext(s_mockAgent, null, aiContext); // Assert Assert.Null(context.Session); @@ -352,52 +228,25 @@ public class AIContextProviderTests public void InvokingContext_Constructor_ThrowsForNullAgent() { // Arrange - var messages = new ReadOnlyCollection([new(ChatRole.User, "Hello")]); + var aiContext = new AIContext { Messages = [new ChatMessage(ChatRole.User, "Hello")] }; // Act & Assert - Assert.Throws(() => new AIContextProvider.InvokingContext(null!, s_mockSession, messages)); + Assert.Throws(() => new AIContextProvider.InvokingContext(null!, s_mockSession, aiContext)); } #endregion #region InvokedContext Tests - [Fact] - public void InvokedContext_RequestMessages_SetterThrowsForNull() - { - // Arrange - var messages = new ReadOnlyCollection([new(ChatRole.User, "Hello")]); - var context = new AIContextProvider.InvokedContext(s_mockAgent, s_mockSession, messages); - - // Act & Assert - Assert.Throws(() => context.RequestMessages = null!); - } - - [Fact] - public void InvokedContext_RequestMessages_SetterRoundtrips() - { - // Arrange - var initialMessages = new ReadOnlyCollection([new(ChatRole.User, "Hello")]); - var newMessages = new List { new(ChatRole.User, "New message") }; - var context = new AIContextProvider.InvokedContext(s_mockAgent, s_mockSession, initialMessages); - - // Act - context.RequestMessages = newMessages; - - // Assert - Assert.Same(newMessages, context.RequestMessages); - } - [Fact] public void InvokedContext_ResponseMessages_Roundtrips() { // Arrange var requestMessages = new ReadOnlyCollection([new(ChatRole.User, "Hello")]); var responseMessages = new List { new(ChatRole.Assistant, "Response message") }; - var context = new AIContextProvider.InvokedContext(s_mockAgent, s_mockSession, requestMessages); // Act - context.ResponseMessages = responseMessages; + var context = new AIContextProvider.InvokedContext(s_mockAgent, s_mockSession, requestMessages, responseMessages); // Assert Assert.Same(responseMessages, context.ResponseMessages); @@ -409,10 +258,9 @@ public class AIContextProviderTests // Arrange var requestMessages = new ReadOnlyCollection([new(ChatRole.User, "Hello")]); var exception = new InvalidOperationException("Test exception"); - var context = new AIContextProvider.InvokedContext(s_mockAgent, s_mockSession, requestMessages); // Act - context.InvokeException = exception; + var context = new AIContextProvider.InvokedContext(s_mockAgent, s_mockSession, requestMessages, exception); // Assert Assert.Same(exception, context.InvokeException); @@ -425,7 +273,7 @@ public class AIContextProviderTests var requestMessages = new ReadOnlyCollection([new(ChatRole.User, "Hello")]); // Act - var context = new AIContextProvider.InvokedContext(s_mockAgent, s_mockSession, requestMessages); + var context = new AIContextProvider.InvokedContext(s_mockAgent, s_mockSession, requestMessages, []); // Assert Assert.Same(s_mockAgent, context.Agent); @@ -438,7 +286,7 @@ public class AIContextProviderTests var requestMessages = new ReadOnlyCollection([new(ChatRole.User, "Hello")]); // Act - var context = new AIContextProvider.InvokedContext(s_mockAgent, s_mockSession, requestMessages); + var context = new AIContextProvider.InvokedContext(s_mockAgent, s_mockSession, requestMessages, []); // Assert Assert.Same(s_mockSession, context.Session); @@ -451,7 +299,7 @@ public class AIContextProviderTests var requestMessages = new ReadOnlyCollection([new(ChatRole.User, "Hello")]); // Act - var context = new AIContextProvider.InvokedContext(s_mockAgent, null, requestMessages); + var context = new AIContextProvider.InvokedContext(s_mockAgent, null, requestMessages, []); // Assert Assert.Null(context.Session); @@ -464,7 +312,27 @@ public class AIContextProviderTests var requestMessages = new ReadOnlyCollection([new(ChatRole.User, "Hello")]); // Act & Assert - Assert.Throws(() => new AIContextProvider.InvokedContext(null!, s_mockSession, requestMessages)); + Assert.Throws(() => new AIContextProvider.InvokedContext(null!, s_mockSession, requestMessages, [])); + } + + [Fact] + public void InvokedContext_SuccessConstructor_ThrowsForNullResponseMessages() + { + // Arrange + var requestMessages = new ReadOnlyCollection([new(ChatRole.User, "Hello")]); + + // Act & Assert + Assert.Throws(() => new AIContextProvider.InvokedContext(s_mockAgent, s_mockSession, requestMessages, (IEnumerable)null!)); + } + + [Fact] + public void InvokedContext_FailureConstructor_ThrowsForNullException() + { + // Arrange + var requestMessages = new ReadOnlyCollection([new(ChatRole.User, "Hello")]); + + // Act & Assert + Assert.Throws(() => new AIContextProvider.InvokedContext(s_mockAgent, s_mockSession, requestMessages, (Exception)null!)); } #endregion @@ -474,55 +342,4 @@ public class AIContextProviderTests protected override ValueTask InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default) => new(new AIContext()); } - - private sealed class TestAIContextProviderWithMessages : AIContextProvider - { - protected override ValueTask InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default) - => new(new AIContext - { - Messages = [new ChatMessage(ChatRole.System, "Context Message")] - }); - } - - private sealed class TestAIContextProviderWithCustomSource : AIContextProvider - { - public TestAIContextProviderWithCustomSource(string sourceId) : base(sourceId) - { - } - - protected override ValueTask InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default) - => new(new AIContext - { - Messages = [new ChatMessage(ChatRole.System, "Context Message")] - }); - } - - private sealed class TestAIContextProviderWithPreStampedMessages : AIContextProvider - { - protected override ValueTask InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default) - { - var message = new ChatMessage(ChatRole.System, "Pre-stamped Message"); - message.AdditionalProperties = new AdditionalPropertiesDictionary - { - [AgentRequestMessageSourceAttribution.AdditionalPropertiesKey] = new AgentRequestMessageSourceAttribution(AgentRequestMessageSourceType.AIContextProvider, this.GetType().FullName!) - }; - return new(new AIContext - { - Messages = [message] - }); - } - } - - private sealed class TestAIContextProviderWithMultipleMessages : AIContextProvider - { - protected override ValueTask InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default) - => new(new AIContext - { - Messages = [ - new ChatMessage(ChatRole.System, "Message 1"), - new ChatMessage(ChatRole.User, "Message 2"), - new ChatMessage(ChatRole.Assistant, "Message 3") - ] - }); - } } diff --git a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/AIContextTests.cs b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/AIContextTests.cs index b1ba6060ea..c925f098b3 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/AIContextTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/AIContextTests.cs @@ -1,5 +1,6 @@ // Copyright (c) Microsoft. All rights reserved. +using System.Linq; using Microsoft.Extensions.AI; namespace Microsoft.Agents.AI.Abstractions.UnitTests; @@ -33,9 +34,10 @@ public class AIContextTests }; Assert.NotNull(context.Messages); - Assert.Equal(2, context.Messages.Count); - Assert.Equal("Hello", context.Messages[0].Text); - Assert.Equal("Hi there!", context.Messages[1].Text); + var messages = context.Messages.ToList(); + Assert.Equal(2, messages.Count); + Assert.Equal("Hello", messages[0].Text); + Assert.Equal("Hi there!", messages[1].Text); } [Fact] @@ -51,8 +53,9 @@ public class AIContextTests }; Assert.NotNull(context.Tools); - Assert.Equal(2, context.Tools.Count); - Assert.Equal("Function1", context.Tools[0].Name); - Assert.Equal("Function2", context.Tools[1].Name); + var tools = context.Tools.ToList(); + Assert.Equal(2, tools.Count); + Assert.Equal("Function1", tools[0].Name); + Assert.Equal("Function2", tools[1].Name); } } diff --git a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/AgentRequestMessageSourceAttributionTests.cs b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/AgentRequestMessageSourceAttributionTests.cs index 5aee121097..70c7a2e06f 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/AgentRequestMessageSourceAttributionTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/AgentRequestMessageSourceAttributionTests.cs @@ -390,6 +390,49 @@ public sealed class AgentRequestMessageSourceAttributionTests #endregion + #region ToString Tests + + [Fact] + public void ToString_WithSourceId_ReturnsTypeColonId() + { + // Arrange + AgentRequestMessageSourceAttribution attribution = new(AgentRequestMessageSourceType.AIContextProvider, "MyProvider"); + + // Act + string result = attribution.ToString(); + + // Assert + Assert.Equal("AIContextProvider:MyProvider", result); + } + + [Fact] + public void ToString_WithNullSourceId_ReturnsTypeOnly() + { + // Arrange + AgentRequestMessageSourceAttribution attribution = new(AgentRequestMessageSourceType.ChatHistory, null); + + // Act + string result = attribution.ToString(); + + // Assert + Assert.Equal("ChatHistory", result); + } + + [Fact] + public void ToString_Default_ReturnsExternalOnly() + { + // Arrange + AgentRequestMessageSourceAttribution attribution = default; + + // Act + string result = attribution.ToString(); + + // Assert + Assert.Equal("External", result); + } + + #endregion + #region Inequality Operator Tests [Fact] diff --git a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/AgentRequestMessageSourceTypeTests.cs b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/AgentRequestMessageSourceTypeTests.cs index 000505fe32..973228828b 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/AgentRequestMessageSourceTypeTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/AgentRequestMessageSourceTypeTests.cs @@ -414,6 +414,46 @@ public sealed class AgentRequestMessageSourceTypeTests #endregion + #region ToString Tests + + [Fact] + public void ToString_ReturnsValue() + { + // Arrange + AgentRequestMessageSourceType source = new("CustomSource"); + + // Act + string result = source.ToString(); + + // Assert + Assert.Equal("CustomSource", result); + } + + [Fact] + public void ToString_StaticExternal_ReturnsExternal() + { + // Arrange & Act + string result = AgentRequestMessageSourceType.External.ToString(); + + // Assert + Assert.Equal("External", result); + } + + [Fact] + public void ToString_Default_ReturnsExternal() + { + // Arrange + AgentRequestMessageSourceType source = default; + + // Act + string result = source.ToString(); + + // Assert + Assert.Equal("External", result); + } + + #endregion + #region IEquatable Tests [Fact] diff --git a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/AgentSessionStateBagTests.cs b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/AgentSessionStateBagTests.cs new file mode 100644 index 0000000000..a51f6dcb86 --- /dev/null +++ b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/AgentSessionStateBagTests.cs @@ -0,0 +1,840 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Text.Json; +using Microsoft.Agents.AI.Abstractions.UnitTests.Models; + +namespace Microsoft.Agents.AI.Abstractions.UnitTests; + +/// +/// Contains tests for the class. +/// +public sealed class AgentSessionStateBagTests +{ + #region Constructor Tests + + [Fact] + public void Constructor_Default_CreatesEmptyStateBag() + { + // Act + var stateBag = new AgentSessionStateBag(); + + // Assert + Assert.False(stateBag.TryGetValue("nonexistent", out _)); + } + + #endregion + + #region SetValue Tests + + [Fact] + public void SetValue_WithValidKeyAndValue_StoresValue() + { + // Arrange + var stateBag = new AgentSessionStateBag(); + + // Act + stateBag.SetValue("key1", "value1"); + + // Assert + Assert.True(stateBag.TryGetValue("key1", out var result)); + Assert.Equal("value1", result); + } + + [Fact] + public void SetValue_WithNullKey_ThrowsArgumentException() + { + // Arrange + var stateBag = new AgentSessionStateBag(); + + // Act & Assert + Assert.Throws(() => stateBag.SetValue(null!, "value")); + } + + [Fact] + public void SetValue_WithEmptyKey_ThrowsArgumentException() + { + // Arrange + var stateBag = new AgentSessionStateBag(); + + // Act & Assert + Assert.Throws(() => stateBag.SetValue("", "value")); + } + + [Fact] + public void SetValue_WithWhitespaceKey_ThrowsArgumentException() + { + // Arrange + var stateBag = new AgentSessionStateBag(); + + // Act & Assert + Assert.Throws(() => stateBag.SetValue(" ", "value")); + } + + [Fact] + public void SetValue_OverwritesExistingValue() + { + // Arrange + var stateBag = new AgentSessionStateBag(); + stateBag.SetValue("key1", "originalValue"); + + // Act + stateBag.SetValue("key1", "newValue"); + + // Assert + Assert.Equal("newValue", stateBag.GetValue("key1")); + } + + #endregion + + #region GetValue Tests + + [Fact] + public void GetValue_WithExistingKey_ReturnsValue() + { + // Arrange + var stateBag = new AgentSessionStateBag(); + stateBag.SetValue("key1", "value1"); + + // Act + var result = stateBag.GetValue("key1"); + + // Assert + Assert.Equal("value1", result); + } + + [Fact] + public void GetValue_WithNonexistentKey_ReturnsNull() + { + // Arrange + var stateBag = new AgentSessionStateBag(); + + // Act + var result = stateBag.GetValue("nonexistent"); + + // Assert + Assert.Null(result); + } + + [Fact] + public void GetValue_WithNullKey_ThrowsArgumentException() + { + // Arrange + var stateBag = new AgentSessionStateBag(); + + // Act & Assert + Assert.Throws(() => stateBag.GetValue(null!)); + } + + [Fact] + public void GetValue_WithEmptyKey_ThrowsArgumentException() + { + // Arrange + var stateBag = new AgentSessionStateBag(); + + // Act & Assert + Assert.Throws(() => stateBag.GetValue("")); + } + + [Fact] + public void GetValue_CachesDeserializedValue() + { + // Arrange + var stateBag = new AgentSessionStateBag(); + stateBag.SetValue("key1", "value1"); + + // Act + var result1 = stateBag.GetValue("key1"); + var result2 = stateBag.GetValue("key1"); + + // Assert + Assert.Same(result1, result2); + } + + #endregion + + #region TryGetValue Tests + + [Fact] + public void TryGetValue_WithExistingKey_ReturnsTrueAndValue() + { + // Arrange + var stateBag = new AgentSessionStateBag(); + stateBag.SetValue("key1", "value1"); + + // Act + var found = stateBag.TryGetValue("key1", out var result); + + // Assert + Assert.True(found); + Assert.Equal("value1", result); + } + + [Fact] + public void TryGetValue_WithNonexistentKey_ReturnsFalseAndNull() + { + // Arrange + var stateBag = new AgentSessionStateBag(); + + // Act + var found = stateBag.TryGetValue("nonexistent", out var result); + + // Assert + Assert.False(found); + Assert.Null(result); + } + + [Fact] + public void TryGetValue_WithNullKey_ThrowsArgumentException() + { + // Arrange + var stateBag = new AgentSessionStateBag(); + + // Act & Assert + Assert.Throws(() => stateBag.TryGetValue(null!, out _)); + } + + [Fact] + public void TryGetValue_WithEmptyKey_ThrowsArgumentException() + { + // Arrange + var stateBag = new AgentSessionStateBag(); + + // Act & Assert + Assert.Throws(() => stateBag.TryGetValue("", out _)); + } + + #endregion + + #region Null Value Tests + + [Fact] + public void SetValue_WithNullValue_StoresNull() + { + // Arrange + var stateBag = new AgentSessionStateBag(); + + // Act + stateBag.SetValue("key1", null); + + // Assert + Assert.Equal(1, stateBag.Count); + } + + [Fact] + public void TryGetValue_WithNullValue_ReturnsTrueAndNull() + { + // Arrange + var stateBag = new AgentSessionStateBag(); + stateBag.SetValue("key1", null); + + // Act + var found = stateBag.TryGetValue("key1", out var result); + + // Assert + Assert.True(found); + Assert.Null(result); + } + + [Fact] + public void GetValue_WithNullValue_ReturnsNull() + { + // Arrange + var stateBag = new AgentSessionStateBag(); + stateBag.SetValue("key1", null); + + // Act + var result = stateBag.GetValue("key1"); + + // Assert + Assert.Null(result); + } + + [Fact] + public void SetValue_OverwriteWithNull_ReturnsNull() + { + // Arrange + var stateBag = new AgentSessionStateBag(); + stateBag.SetValue("key1", "value1"); + + // Act + stateBag.SetValue("key1", null); + + // Assert + Assert.True(stateBag.TryGetValue("key1", out var result)); + Assert.Null(result); + } + + [Fact] + public void SetValue_OverwriteNullWithValue_ReturnsValue() + { + // Arrange + var stateBag = new AgentSessionStateBag(); + stateBag.SetValue("key1", null); + + // Act + stateBag.SetValue("key1", "newValue"); + + // Assert + Assert.True(stateBag.TryGetValue("key1", out var result)); + Assert.Equal("newValue", result); + } + + [Fact] + public void SerializeDeserialize_WithNullValue_SerializesAsNull() + { + // Arrange + var stateBag = new AgentSessionStateBag(); + stateBag.SetValue("nullKey", null); + + // Act + var json = stateBag.Serialize(); + + // Assert - null values are serialized as JSON null + Assert.Equal(JsonValueKind.Object, json.ValueKind); + Assert.True(json.TryGetProperty("nullKey", out var nullElement)); + Assert.Equal(JsonValueKind.Null, nullElement.ValueKind); + } + + #endregion + + #region TryRemoveValue Tests + + [Fact] + public void TryRemoveValue_ExistingKey_ReturnsTrueAndRemoves() + { + // Arrange + var stateBag = new AgentSessionStateBag(); + stateBag.SetValue("key1", "value1"); + + // Act + var removed = stateBag.TryRemoveValue("key1"); + + // Assert + Assert.True(removed); + Assert.Equal(0, stateBag.Count); + Assert.False(stateBag.TryGetValue("key1", out _)); + } + + [Fact] + public void TryRemoveValue_NonexistentKey_ReturnsFalse() + { + // Arrange + var stateBag = new AgentSessionStateBag(); + + // Act + var removed = stateBag.TryRemoveValue("nonexistent"); + + // Assert + Assert.False(removed); + } + + [Fact] + public void TryRemoveValue_WithNullKey_ThrowsArgumentException() + { + // Arrange + var stateBag = new AgentSessionStateBag(); + + // Act & Assert + Assert.Throws(() => stateBag.TryRemoveValue(null!)); + } + + [Fact] + public void TryRemoveValue_WithEmptyKey_ThrowsArgumentException() + { + // Arrange + var stateBag = new AgentSessionStateBag(); + + // Act & Assert + Assert.Throws(() => stateBag.TryRemoveValue("")); + } + + [Fact] + public void TryRemoveValue_WithWhitespaceKey_ThrowsArgumentException() + { + // Arrange + var stateBag = new AgentSessionStateBag(); + + // Act & Assert + Assert.Throws(() => stateBag.TryRemoveValue(" ")); + } + + [Fact] + public void TryRemoveValue_DoesNotAffectOtherKeys() + { + // Arrange + var stateBag = new AgentSessionStateBag(); + stateBag.SetValue("key1", "value1"); + stateBag.SetValue("key2", "value2"); + + // Act + stateBag.TryRemoveValue("key1"); + + // Assert + Assert.Equal(1, stateBag.Count); + Assert.False(stateBag.TryGetValue("key1", out _)); + Assert.True(stateBag.TryGetValue("key2", out var value)); + Assert.Equal("value2", value); + } + + [Fact] + public void TryRemoveValue_ThenSetValue_Works() + { + // Arrange + var stateBag = new AgentSessionStateBag(); + stateBag.SetValue("key1", "original"); + + // Act + stateBag.TryRemoveValue("key1"); + stateBag.SetValue("key1", "replacement"); + + // Assert + Assert.True(stateBag.TryGetValue("key1", out var result)); + Assert.Equal("replacement", result); + } + + #endregion + + #region Serialize/Deserialize Tests + + [Fact] + public void Serialize_EmptyStateBag_ReturnsEmptyObject() + { + // Arrange + var stateBag = new AgentSessionStateBag(); + + // Act + var json = stateBag.Serialize(); + + // Assert + Assert.Equal(JsonValueKind.Object, json.ValueKind); + } + + [Fact] + public void Serialize_WithStringValue_ReturnsJsonWithValue() + { + // Arrange + var stateBag = new AgentSessionStateBag(); + stateBag.SetValue("stringKey", "stringValue"); + + // Act + var json = stateBag.Serialize(); + + // Assert + Assert.Equal(JsonValueKind.Object, json.ValueKind); + Assert.True(json.TryGetProperty("stringKey", out _)); + } + + [Fact] + public void Deserialize_FromJsonDocument_ReturnsEmptyStateBag() + { + // Arrange + var emptyJson = JsonDocument.Parse("{}").RootElement; + + // Act + var stateBag = AgentSessionStateBag.Deserialize(emptyJson); + + // Assert + Assert.False(stateBag.TryGetValue("nonexistent", out _)); + } + + [Fact] + public void Deserialize_NullElement_ReturnsEmptyStateBag() + { + // Arrange + var nullJson = default(JsonElement); + + // Act + var stateBag = AgentSessionStateBag.Deserialize(nullJson); + + // Assert + Assert.False(stateBag.TryGetValue("nonexistent", out _)); + } + + [Fact] + public void SerializeDeserialize_WithStringValue_Roundtrips() + { + // Arrange + var originalStateBag = new AgentSessionStateBag(); + originalStateBag.SetValue("stringKey", "stringValue"); + + // Act + var json = originalStateBag.Serialize(); + var restoredStateBag = AgentSessionStateBag.Deserialize(json); + + // Assert + Assert.Equal("stringValue", restoredStateBag.GetValue("stringKey")); + } + + #endregion + + #region Thread Safety Tests + + [Fact] + public async System.Threading.Tasks.Task SetValue_MultipleConcurrentWrites_DoesNotThrowAsync() + { + // Arrange + var stateBag = new AgentSessionStateBag(); + var tasks = new System.Threading.Tasks.Task[100]; + + // Act + for (int i = 0; i < 100; i++) + { + int index = i; + tasks[i] = System.Threading.Tasks.Task.Run(() => stateBag.SetValue($"key{index}", $"value{index}")); + } + + await System.Threading.Tasks.Task.WhenAll(tasks); + + // Assert + for (int i = 0; i < 100; i++) + { + Assert.True(stateBag.TryGetValue($"key{i}", out var value)); + Assert.Equal($"value{i}", value); + } + } + + [Fact] + public async System.Threading.Tasks.Task ConcurrentWritesAndSerialize_DoesNotThrowAsync() + { + // Arrange + var stateBag = new AgentSessionStateBag(); + stateBag.SetValue("shared", "initial"); + var tasks = new System.Threading.Tasks.Task[100]; + + // Act - concurrently write and serialize the same key + for (int i = 0; i < 100; i++) + { + int index = i; + tasks[i] = System.Threading.Tasks.Task.Run(() => + { + stateBag.SetValue("shared", $"value{index}"); + _ = stateBag.Serialize(); + }); + } + + await System.Threading.Tasks.Task.WhenAll(tasks); + + // Assert - should have some value and serialize without error + Assert.True(stateBag.TryGetValue("shared", out var result)); + Assert.NotNull(result); + var json = stateBag.Serialize(); + Assert.Equal(JsonValueKind.Object, json.ValueKind); + } + + [Fact] + public async System.Threading.Tasks.Task ConcurrentReadsAndWrites_DoesNotThrowAsync() + { + // Arrange + var stateBag = new AgentSessionStateBag(); + stateBag.SetValue("key", "initial"); + var tasks = new System.Threading.Tasks.Task[200]; + + // Act - half readers, half writers on the same key + for (int i = 0; i < 200; i++) + { + int index = i; + tasks[i] = (index % 2 == 0) + ? System.Threading.Tasks.Task.Run(() => stateBag.GetValue("key")) + : System.Threading.Tasks.Task.Run(() => stateBag.SetValue("key", $"value{index}")); + } + + await System.Threading.Tasks.Task.WhenAll(tasks); + + // Assert - should have a consistent value + Assert.True(stateBag.TryGetValue("key", out var result)); + Assert.NotNull(result); + } + + #endregion + + #region Complex Object Tests + + [Fact] + public void SetValue_WithComplexObject_StoresValue() + { + // Arrange + var stateBag = new AgentSessionStateBag(); + var animal = new Animal { Id = 1, FullName = "Buddy", Species = Species.Bear }; + + // Act + stateBag.SetValue("animal", animal, TestJsonSerializerContext.Default.Options); + + // Assert + Animal? result = stateBag.GetValue("animal", TestJsonSerializerContext.Default.Options); + Assert.NotNull(result); + Assert.Equal(1, result.Id); + Assert.Equal("Buddy", result.FullName); + Assert.Equal(Species.Bear, result.Species); + } + + [Fact] + public void GetValue_WithComplexObject_CachesDeserializedValue() + { + // Arrange + var stateBag = new AgentSessionStateBag(); + var animal = new Animal { Id = 2, FullName = "Whiskers", Species = Species.Tiger }; + stateBag.SetValue("animal", animal, TestJsonSerializerContext.Default.Options); + + // Act + Animal? result1 = stateBag.GetValue("animal", TestJsonSerializerContext.Default.Options); + Animal? result2 = stateBag.GetValue("animal", TestJsonSerializerContext.Default.Options); + + // Assert + Assert.Same(result1, result2); + } + + [Fact] + public void TryGetValue_WithComplexObject_ReturnsTrueAndValue() + { + // Arrange + var stateBag = new AgentSessionStateBag(); + var animal = new Animal { Id = 3, FullName = "Goldie", Species = Species.Walrus }; + stateBag.SetValue("animal", animal, TestJsonSerializerContext.Default.Options); + + // Act + bool found = stateBag.TryGetValue("animal", out Animal? result, TestJsonSerializerContext.Default.Options); + + // Assert + Assert.True(found); + Assert.NotNull(result); + Assert.Equal(3, result.Id); + Assert.Equal("Goldie", result.FullName); + Assert.Equal(Species.Walrus, result.Species); + } + + [Fact] + public void SerializeDeserialize_WithComplexObject_Roundtrips() + { + // Arrange + var originalStateBag = new AgentSessionStateBag(); + var animal = new Animal { Id = 4, FullName = "Polly", Species = Species.Bear }; + originalStateBag.SetValue("animal", animal, TestJsonSerializerContext.Default.Options); + + // Act + JsonElement json = originalStateBag.Serialize(); + AgentSessionStateBag restoredStateBag = AgentSessionStateBag.Deserialize(json); + + // Assert + Animal? restoredAnimal = restoredStateBag.GetValue("animal", TestJsonSerializerContext.Default.Options); + Assert.NotNull(restoredAnimal); + Assert.Equal(4, restoredAnimal.Id); + Assert.Equal("Polly", restoredAnimal.FullName); + Assert.Equal(Species.Bear, restoredAnimal.Species); + } + + [Fact] + public void Serialize_WithComplexObject_ReturnsJsonWithProperties() + { + // Arrange + var stateBag = new AgentSessionStateBag(); + var animal = new Animal { Id = 7, FullName = "Spot", Species = Species.Walrus }; + stateBag.SetValue("animal", animal, TestJsonSerializerContext.Default.Options); + + // Act + JsonElement json = stateBag.Serialize(); + + // Assert + Assert.Equal(JsonValueKind.Object, json.ValueKind); + Assert.True(json.TryGetProperty("animal", out JsonElement animalElement)); + Assert.Equal(JsonValueKind.Object, animalElement.ValueKind); + Assert.Equal(7, animalElement.GetProperty("id").GetInt32()); + Assert.Equal("Spot", animalElement.GetProperty("fullName").GetString()); + Assert.Equal("Walrus", animalElement.GetProperty("species").GetString()); + } + + #endregion + + #region Type Mismatch Tests + + [Fact] + public void TryGetValue_WithDifferentTypeAfterSet_ReturnsFalse() + { + // Arrange + var stateBag = new AgentSessionStateBag(); + stateBag.SetValue("key1", "hello"); + + // Act + var found = stateBag.TryGetValue("key1", out var result, TestJsonSerializerContext.Default.Options); + + // Assert + Assert.False(found); + Assert.Null(result); + } + + [Fact] + public void GetValue_WithDifferentTypeAfterSet_ThrowsInvalidOperationException() + { + // Arrange + var stateBag = new AgentSessionStateBag(); + stateBag.SetValue("key1", "hello"); + + // Act & Assert + Assert.Throws(() => stateBag.GetValue("key1", TestJsonSerializerContext.Default.Options)); + } + + [Fact] + public void TryGetValue_WithDifferentTypeAfterDeserializedRead_ReturnsFalse() + { + // Arrange + var stateBag = new AgentSessionStateBag(); + stateBag.SetValue("key1", "hello"); + + // First read caches the value as string + var cachedValue = stateBag.GetValue("key1"); + Assert.Equal("hello", cachedValue); + + // Act - request as a different type + var found = stateBag.TryGetValue("key1", out var result, TestJsonSerializerContext.Default.Options); + + // Assert + Assert.False(found); + Assert.Null(result); + } + + [Fact] + public void GetValue_WithDifferentTypeAfterDeserializedRoundtrip_ThrowsInvalidOperationException() + { + // Arrange + var originalStateBag = new AgentSessionStateBag(); + originalStateBag.SetValue("key1", "hello"); + + // Round-trip through serialization + var json = originalStateBag.Serialize(); + var restoredStateBag = AgentSessionStateBag.Deserialize(json); + + // First read caches the value as string + var cachedValue = restoredStateBag.GetValue("key1"); + Assert.Equal("hello", cachedValue); + + // Act & Assert - request as a different type + Assert.Throws(() => restoredStateBag.GetValue("key1", TestJsonSerializerContext.Default.Options)); + } + + [Fact] + public void TryGetValue_ComplexTypeAfterSetString_ReturnsFalse() + { + // Arrange + var stateBag = new AgentSessionStateBag(); + stateBag.SetValue("animal", "not an animal"); + + // Act + var found = stateBag.TryGetValue("animal", out var result, TestJsonSerializerContext.Default.Options); + + // Assert + Assert.False(found); + Assert.Null(result); + } + + [Fact] + public void GetValue_TypeMismatch_ExceptionMessageContainsBothTypeNames() + { + // Arrange + var stateBag = new AgentSessionStateBag(); + stateBag.SetValue("key1", "hello"); + + // Act + var exception = Assert.Throws(() => stateBag.GetValue("key1", TestJsonSerializerContext.Default.Options)); + + // Assert + Assert.Contains(typeof(string).FullName!, exception.Message); + Assert.Contains(typeof(Animal).FullName!, exception.Message); + } + + #endregion + + #region JsonSerializer Integration Tests + + [Fact] + public void JsonSerializerSerialize_EmptyStateBag_ReturnsEmptyObject() + { + // Arrange + var stateBag = new AgentSessionStateBag(); + + // Act + var json = JsonSerializer.Serialize(stateBag, AgentAbstractionsJsonUtilities.DefaultOptions); + + // Assert + Assert.Equal("{}", json); + } + + [Fact] + public void JsonSerializerSerialize_WithStringValue_ProducesSameOutputAsSerializeMethod() + { + // Arrange + var stateBag = new AgentSessionStateBag(); + stateBag.SetValue("stringKey", "stringValue"); + + // Act + var jsonFromSerializer = JsonSerializer.Serialize(stateBag, AgentAbstractionsJsonUtilities.DefaultOptions); + var jsonFromMethod = stateBag.Serialize().GetRawText(); + + // Assert + Assert.Equal(jsonFromMethod, jsonFromSerializer); + } + + [Fact] + public void JsonSerializerRoundtrip_WithStringValue_PreservesData() + { + // Arrange + var stateBag = new AgentSessionStateBag(); + stateBag.SetValue("greeting", "hello world"); + + // Act + var json = JsonSerializer.Serialize(stateBag, AgentAbstractionsJsonUtilities.DefaultOptions); + var restored = JsonSerializer.Deserialize(json, AgentAbstractionsJsonUtilities.DefaultOptions); + + // Assert + Assert.NotNull(restored); + Assert.Equal("hello world", restored!.GetValue("greeting")); + } + + [Fact] + public void JsonSerializerRoundtrip_WithComplexObject_PreservesData() + { + // Arrange + var stateBag = new AgentSessionStateBag(); + var animal = new Animal { Id = 10, FullName = "Rex", Species = Species.Tiger }; + stateBag.SetValue("animal", animal, TestJsonSerializerContext.Default.Options); + + // Act + var json = JsonSerializer.Serialize(stateBag, AgentAbstractionsJsonUtilities.DefaultOptions); + var restored = JsonSerializer.Deserialize(json, AgentAbstractionsJsonUtilities.DefaultOptions); + + // Assert + Assert.NotNull(restored); + var restoredAnimal = restored!.GetValue("animal", TestJsonSerializerContext.Default.Options); + Assert.NotNull(restoredAnimal); + Assert.Equal(10, restoredAnimal!.Id); + Assert.Equal("Rex", restoredAnimal.FullName); + Assert.Equal(Species.Tiger, restoredAnimal.Species); + } + + [Fact] + public void JsonSerializerDeserialize_NullJson_ReturnsNull() + { + // Arrange + const string Json = "null"; + + // Act + var stateBag = JsonSerializer.Deserialize(Json, AgentAbstractionsJsonUtilities.DefaultOptions); + + // Assert + Assert.Null(stateBag); + } + +#if NET10_0_OR_GREATER + [Fact] + public void JsonSerializerSerialize_WithUnknownType_Throws() + { + // Arrange + var stateBag = new AgentSessionStateBag(); + stateBag.SetValue("key", new { Name = "Test" }); // Anonymous type which cannot be deserialized + + // Act & Assert + Assert.Throws(() => JsonSerializer.Serialize(stateBag, AgentAbstractionsJsonUtilities.DefaultOptions)); + } +#endif + + #endregion +} diff --git a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/AgentSessionTests.cs b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/AgentSessionTests.cs index 5a776c9fb0..b80f0a4fd2 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/AgentSessionTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/AgentSessionTests.cs @@ -11,6 +11,21 @@ namespace Microsoft.Agents.AI.Abstractions.UnitTests; /// public class AgentSessionTests { + #region StateBag Tests + + [Fact] + public void StateBag_Values_Roundtrips() + { + // Arrange + var session = new TestAgentSession(); + + // Act & Assert + session.StateBag.SetValue("key1", "value1"); + Assert.Equal("value1", session.StateBag.GetValue("key1")); + } + + #endregion + #region GetService Method Tests /// diff --git a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/ChatHistoryProviderExtensionsTests.cs b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/ChatHistoryProviderExtensionsTests.cs deleted file mode 100644 index 1fcbe37e25..0000000000 --- a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/ChatHistoryProviderExtensionsTests.cs +++ /dev/null @@ -1,141 +0,0 @@ -// Copyright (c) Microsoft. All rights reserved. - -using System.Collections.Generic; -using System.Linq; -using System.Threading; -using System.Threading.Tasks; -using Microsoft.Extensions.AI; -using Moq; -using Moq.Protected; - -namespace Microsoft.Agents.AI.Abstractions.UnitTests; - -/// -/// Contains tests for the class. -/// -public sealed class ChatHistoryProviderExtensionsTests -{ - private static readonly AIAgent s_mockAgent = new Mock().Object; - private static readonly AgentSession s_mockSession = new Mock().Object; - - [Fact] - public void WithMessageFilters_ReturnsChatHistoryProviderMessageFilter() - { - // Arrange - Mock providerMock = new(); - - // Act - ChatHistoryProvider result = providerMock.Object.WithMessageFilters( - invokingMessagesFilter: msgs => msgs, - invokedMessagesFilter: ctx => ctx); - - // Assert - Assert.IsType(result); - } - - [Fact] - public async Task WithMessageFilters_InvokingFilter_IsAppliedAsync() - { - // Arrange - Mock providerMock = new(); - List innerMessages = [new(ChatRole.User, "Hello"), new(ChatRole.Assistant, "Hi")]; - ChatHistoryProvider.InvokingContext context = new(s_mockAgent, s_mockSession, [new ChatMessage(ChatRole.User, "Test")]); - - providerMock - .Protected() - .Setup>>("InvokingCoreAsync", ItExpr.IsAny(), ItExpr.IsAny()) - .ReturnsAsync(innerMessages); - - ChatHistoryProvider filtered = providerMock.Object.WithMessageFilters( - invokingMessagesFilter: msgs => msgs.Where(m => m.Role == ChatRole.User)); - - // Act - List result = (await filtered.InvokingAsync(context, CancellationToken.None)).ToList(); - - // Assert - Assert.Single(result); - Assert.Equal(ChatRole.User, result[0].Role); - } - - [Fact] - public async Task WithMessageFilters_InvokedFilter_IsAppliedAsync() - { - // Arrange - Mock providerMock = new(); - List requestMessages = - [ - new(ChatRole.System, "System") { AdditionalProperties = new() { { AgentRequestMessageSourceAttribution.AdditionalPropertiesKey, new AgentRequestMessageSourceAttribution(AgentRequestMessageSourceType.ChatHistory, "TestSource") } } }, - new(ChatRole.User, "Hello") - ]; - ChatHistoryProvider.InvokedContext context = new(s_mockAgent, s_mockSession, requestMessages) - { - ResponseMessages = [new ChatMessage(ChatRole.Assistant, "Response")] - }; - - ChatHistoryProvider.InvokedContext? capturedContext = null; - providerMock - .Protected() - .Setup("InvokedCoreAsync", ItExpr.IsAny(), ItExpr.IsAny()) - .Callback((ctx, _) => capturedContext = ctx) - .Returns(default(ValueTask)); - - ChatHistoryProvider filtered = providerMock.Object.WithMessageFilters( - invokedMessagesFilter: ctx => - { - ctx.ResponseMessages = null; - return ctx; - }); - - // Act - await filtered.InvokedAsync(context, CancellationToken.None); - - // Assert - Assert.NotNull(capturedContext); - Assert.Null(capturedContext.ResponseMessages); - } - - [Fact] - public void WithAIContextProviderMessageRemoval_ReturnsChatHistoryProviderMessageFilter() - { - // Arrange - Mock providerMock = new(); - - // Act - ChatHistoryProvider result = providerMock.Object.WithAIContextProviderMessageRemoval(); - - // Assert - Assert.IsType(result); - } - - [Fact] - public async Task WithAIContextProviderMessageRemoval_RemovesAIContextProviderMessagesAsync() - { - // Arrange - Mock providerMock = new(); - List requestMessages = - [ - new(ChatRole.System, "System") { AdditionalProperties = new() { { AgentRequestMessageSourceAttribution.AdditionalPropertiesKey, new AgentRequestMessageSourceAttribution(AgentRequestMessageSourceType.ChatHistory, "TestSource") } } }, - new(ChatRole.User, "Hello"), - new(ChatRole.System, "Context") { AdditionalProperties = new() { { AgentRequestMessageSourceAttribution.AdditionalPropertiesKey, new AgentRequestMessageSourceAttribution(AgentRequestMessageSourceType.AIContextProvider, "TestContextSource") } } } - ]; - ChatHistoryProvider.InvokedContext context = new(s_mockAgent, s_mockSession, requestMessages); - - ChatHistoryProvider.InvokedContext? capturedContext = null; - providerMock - .Protected() - .Setup("InvokedCoreAsync", ItExpr.IsAny(), ItExpr.IsAny()) - .Callback((ctx, _) => capturedContext = ctx) - .Returns(default(ValueTask)); - - ChatHistoryProvider filtered = providerMock.Object.WithAIContextProviderMessageRemoval(); - - // Act - await filtered.InvokedAsync(context, CancellationToken.None); - - // Assert - Assert.NotNull(capturedContext); - Assert.Equal(2, capturedContext.RequestMessages.Count()); - Assert.Contains("System", capturedContext.RequestMessages.Select(x => x.Text)); - Assert.Contains("Hello", capturedContext.RequestMessages.Select(x => x.Text)); - } -} diff --git a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/ChatHistoryProviderMessageFilterTests.cs b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/ChatHistoryProviderMessageFilterTests.cs deleted file mode 100644 index 75d8f554c5..0000000000 --- a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/ChatHistoryProviderMessageFilterTests.cs +++ /dev/null @@ -1,221 +0,0 @@ -// Copyright (c) Microsoft. All rights reserved. - -using System; -using System.Collections.Generic; -using System.Linq; -using System.Text.Json; -using System.Threading; -using System.Threading.Tasks; -using Microsoft.Extensions.AI; -using Moq; -using Moq.Protected; - -namespace Microsoft.Agents.AI.Abstractions.UnitTests; - -/// -/// Contains tests for the class. -/// -public sealed class ChatHistoryProviderMessageFilterTests -{ - private static readonly AIAgent s_mockAgent = new Mock().Object; - private static readonly AgentSession s_mockSession = new Mock().Object; - - [Fact] - public void Constructor_WithNullInnerProvider_ThrowsArgumentNullException() - { - // Arrange, Act & Assert - Assert.Throws(() => new ChatHistoryProviderMessageFilter(null!)); - } - - [Fact] - public void Constructor_WithOnlyInnerProvider_Throws() - { - // Arrange - var innerProviderMock = new Mock(); - - // Act & Assert - Assert.Throws(() => new ChatHistoryProviderMessageFilter(innerProviderMock.Object)); - } - - [Fact] - public void Constructor_WithAllParameters_CreatesInstance() - { - // Arrange - var innerProviderMock = new Mock(); - - IEnumerable InvokingFilter(IEnumerable msgs) => msgs; - ChatHistoryProvider.InvokedContext InvokedFilter(ChatHistoryProvider.InvokedContext ctx) => ctx; - - // Act - var filter = new ChatHistoryProviderMessageFilter(innerProviderMock.Object, InvokingFilter, InvokedFilter); - - // Assert - Assert.NotNull(filter); - } - - [Fact] - public async Task InvokingAsync_WithNoOpFilters_ReturnsInnerProviderMessagesAsync() - { - // Arrange - var innerProviderMock = new Mock(); - var expectedMessages = new List - { - new(ChatRole.User, "Hello"), - new(ChatRole.Assistant, "Hi there!") - }; - var context = new ChatHistoryProvider.InvokingContext(s_mockAgent, s_mockSession, [new ChatMessage(ChatRole.User, "Test")]); - - innerProviderMock - .Protected() - .Setup>>("InvokingCoreAsync", ItExpr.IsAny(), ItExpr.IsAny()) - .ReturnsAsync(expectedMessages); - - var filter = new ChatHistoryProviderMessageFilter(innerProviderMock.Object, x => x, x => x); - - // Act - var result = (await filter.InvokingAsync(context, CancellationToken.None)).ToList(); - - // Assert - Assert.Equal(2, result.Count); - Assert.Equal("Hello", result[0].Text); - Assert.Equal("Hi there!", result[1].Text); - innerProviderMock - .Protected() - .Verify>>("InvokingCoreAsync", Times.Once(), ItExpr.IsAny(), ItExpr.IsAny()); - } - - [Fact] - public async Task InvokingAsync_WithInvokingFilter_AppliesFilterAsync() - { - // Arrange - var innerProviderMock = new Mock(); - var innerMessages = new List - { - new(ChatRole.User, "Hello"), - new(ChatRole.Assistant, "Hi there!"), - new(ChatRole.User, "How are you?") - }; - var context = new ChatHistoryProvider.InvokingContext(s_mockAgent, s_mockSession, [new ChatMessage(ChatRole.User, "Test")]); - - innerProviderMock - .Protected() - .Setup>>("InvokingCoreAsync", ItExpr.IsAny(), ItExpr.IsAny()) - .ReturnsAsync(innerMessages); - - // Filter to only user messages - IEnumerable InvokingFilter(IEnumerable msgs) => msgs.Where(m => m.Role == ChatRole.User); - - var filter = new ChatHistoryProviderMessageFilter(innerProviderMock.Object, InvokingFilter); - - // Act - var result = (await filter.InvokingAsync(context, CancellationToken.None)).ToList(); - - // Assert - Assert.Equal(2, result.Count); - Assert.All(result, msg => Assert.Equal(ChatRole.User, msg.Role)); - innerProviderMock - .Protected() - .Verify>>("InvokingCoreAsync", Times.Once(), ItExpr.IsAny(), ItExpr.IsAny()); - } - - [Fact] - public async Task InvokingAsync_WithInvokingFilter_CanModifyMessagesAsync() - { - // Arrange - var innerProviderMock = new Mock(); - var innerMessages = new List - { - new(ChatRole.User, "Hello"), - new(ChatRole.Assistant, "Hi there!") - }; - var context = new ChatHistoryProvider.InvokingContext(s_mockAgent, s_mockSession, [new ChatMessage(ChatRole.User, "Test")]); - - innerProviderMock - .Protected() - .Setup>>("InvokingCoreAsync", ItExpr.IsAny(), ItExpr.IsAny()) - .ReturnsAsync(innerMessages); - - // Filter that transforms messages - IEnumerable InvokingFilter(IEnumerable msgs) => - msgs.Select(m => new ChatMessage(m.Role, $"[FILTERED] {m.Text}")); - - var filter = new ChatHistoryProviderMessageFilter(innerProviderMock.Object, InvokingFilter); - - // Act - var result = (await filter.InvokingAsync(context, CancellationToken.None)).ToList(); - - // Assert - Assert.Equal(2, result.Count); - Assert.Equal("[FILTERED] Hello", result[0].Text); - Assert.Equal("[FILTERED] Hi there!", result[1].Text); - } - - [Fact] - public async Task InvokedAsync_WithInvokedFilter_AppliesFilterAsync() - { - // Arrange - var innerProviderMock = new Mock(); - List requestMessages = - [ - new(ChatRole.System, "System") { AdditionalProperties = new() { { AgentRequestMessageSourceAttribution.AdditionalPropertiesKey, new AgentRequestMessageSourceAttribution(AgentRequestMessageSourceType.ChatHistory, "TestSource") } } }, - new(ChatRole.User, "Hello"), - ]; - var responseMessages = new List { new(ChatRole.Assistant, "Response") }; - var context = new ChatHistoryProvider.InvokedContext(s_mockAgent, s_mockSession, requestMessages) - { - ResponseMessages = responseMessages - }; - - ChatHistoryProvider.InvokedContext? capturedContext = null; - innerProviderMock - .Protected() - .Setup("InvokedCoreAsync", ItExpr.IsAny(), ItExpr.IsAny()) - .Callback((ctx, ct) => capturedContext = ctx) - .Returns(default(ValueTask)); - - // Filter that modifies the context - ChatHistoryProvider.InvokedContext InvokedFilter(ChatHistoryProvider.InvokedContext ctx) - { - var modifiedRequestMessages = ctx.RequestMessages.Where(x => x.GetAgentRequestMessageSourceType() == AgentRequestMessageSourceType.External).Select(m => new ChatMessage(m.Role, $"[FILTERED] {m.Text}")).ToList(); - return new ChatHistoryProvider.InvokedContext(s_mockAgent, s_mockSession, modifiedRequestMessages) - { - ResponseMessages = ctx.ResponseMessages, - InvokeException = ctx.InvokeException - }; - } - - var filter = new ChatHistoryProviderMessageFilter(innerProviderMock.Object, invokedMessagesFilter: InvokedFilter); - - // Act - await filter.InvokedAsync(context, CancellationToken.None); - - // Assert - Assert.NotNull(capturedContext); - Assert.Single(capturedContext.RequestMessages); - Assert.Equal("[FILTERED] Hello", capturedContext.RequestMessages.First().Text); - innerProviderMock - .Protected() - .Verify("InvokedCoreAsync", Times.Once(), ItExpr.IsAny(), ItExpr.IsAny()); - } - - [Fact] - public void Serialize_DelegatesToInnerProvider() - { - // Arrange - var innerProviderMock = new Mock(); - var expectedJson = JsonSerializer.SerializeToElement("data", TestJsonSerializerContext.Default.String); - - innerProviderMock - .Setup(s => s.Serialize(It.IsAny())) - .Returns(expectedJson); - - var filter = new ChatHistoryProviderMessageFilter(innerProviderMock.Object, x => x, x => x); - - // Act - var result = filter.Serialize(); - - // Assert - Assert.Equal(expectedJson.GetRawText(), result.GetRawText()); - innerProviderMock.Verify(s => s.Serialize(null), Times.Once); - } -} diff --git a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/ChatHistoryProviderTests.cs b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/ChatHistoryProviderTests.cs index 8b07366b03..7fccca8d1b 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/ChatHistoryProviderTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/ChatHistoryProviderTests.cs @@ -3,7 +3,6 @@ using System; using System.Collections.Generic; using System.Linq; -using System.Text.Json; using System.Threading; using System.Threading.Tasks; using Microsoft.Extensions.AI; @@ -19,92 +18,6 @@ public class ChatHistoryProviderTests private static readonly AIAgent s_mockAgent = new Mock().Object; private static readonly AgentSession s_mockSession = new Mock().Object; - #region InvokingAsync Message Stamping Tests - - [Fact] - public async Task InvokingAsync_StampsMessagesWithSourceTypeAndSourceIdAsync() - { - // Arrange - var provider = new TestChatHistoryProvider(); - var context = new ChatHistoryProvider.InvokingContext(s_mockAgent, s_mockSession, [new ChatMessage(ChatRole.User, "Request")]); - - // Act - IEnumerable messages = await provider.InvokingAsync(context); - - // Assert - ChatMessage message = messages.Single(); - Assert.NotNull(message.AdditionalProperties); - Assert.True(message.AdditionalProperties.TryGetValue(AgentRequestMessageSourceAttribution.AdditionalPropertiesKey, out object? attribution)); - var typedAttribution = Assert.IsType(attribution); - Assert.Equal(AgentRequestMessageSourceType.ChatHistory, typedAttribution.SourceType); - Assert.Equal(typeof(TestChatHistoryProvider).FullName, typedAttribution.SourceId); - } - - [Fact] - public async Task InvokingAsync_WithCustomSourceId_StampsMessagesWithCustomSourceIdAsync() - { - // Arrange - const string CustomSourceId = "CustomHistorySource"; - var provider = new TestChatHistoryProviderWithCustomSource(CustomSourceId); - var context = new ChatHistoryProvider.InvokingContext(s_mockAgent, s_mockSession, [new ChatMessage(ChatRole.User, "Request")]); - - // Act - IEnumerable messages = await provider.InvokingAsync(context); - - // Assert - ChatMessage message = messages.Single(); - Assert.NotNull(message.AdditionalProperties); - Assert.True(message.AdditionalProperties.TryGetValue(AgentRequestMessageSourceAttribution.AdditionalPropertiesKey, out object? attribution)); - var typedAttribution = Assert.IsType(attribution); - Assert.Equal(AgentRequestMessageSourceType.ChatHistory, typedAttribution.SourceType); - Assert.Equal(CustomSourceId, typedAttribution.SourceId); - } - - [Fact] - public async Task InvokingAsync_DoesNotReStampAlreadyStampedMessagesAsync() - { - // Arrange - var provider = new TestChatHistoryProviderWithPreStampedMessages(); - var context = new ChatHistoryProvider.InvokingContext(s_mockAgent, s_mockSession, [new ChatMessage(ChatRole.User, "Request")]); - - // Act - IEnumerable messages = await provider.InvokingAsync(context); - - // Assert - ChatMessage message = messages.Single(); - Assert.NotNull(message.AdditionalProperties); - Assert.True(message.AdditionalProperties.TryGetValue(AgentRequestMessageSourceAttribution.AdditionalPropertiesKey, out object? attribution)); - var typedAttribution = Assert.IsType(attribution); - Assert.Equal(AgentRequestMessageSourceType.ChatHistory, typedAttribution.SourceType); - Assert.Equal(typeof(TestChatHistoryProviderWithPreStampedMessages).FullName, typedAttribution.SourceId); - } - - [Fact] - public async Task InvokingAsync_StampsMultipleMessagesAsync() - { - // Arrange - var provider = new TestChatHistoryProviderWithMultipleMessages(); - var context = new ChatHistoryProvider.InvokingContext(s_mockAgent, s_mockSession, [new ChatMessage(ChatRole.User, "Request")]); - - // Act - IEnumerable messages = await provider.InvokingAsync(context); - - // Assert - List messageList = messages.ToList(); - Assert.Equal(3, messageList.Count); - - foreach (ChatMessage message in messageList) - { - Assert.NotNull(message.AdditionalProperties); - Assert.True(message.AdditionalProperties.TryGetValue(AgentRequestMessageSourceAttribution.AdditionalPropertiesKey, out object? attribution)); - var typedAttribution = Assert.IsType(attribution); - Assert.Equal(AgentRequestMessageSourceType.ChatHistory, typedAttribution.SourceType); - Assert.Equal(typeof(TestChatHistoryProviderWithMultipleMessages).FullName, typedAttribution.SourceId); - } - } - - #endregion - #region GetService Method Tests [Fact] @@ -259,33 +172,7 @@ public class ChatHistoryProviderTests public void InvokedContext_Constructor_ThrowsForNullRequestMessages() { // Arrange & Act & Assert - Assert.Throws(() => new ChatHistoryProvider.InvokedContext(s_mockAgent, s_mockSession, null!)); - } - - [Fact] - public void InvokedContext_RequestMessages_SetterThrowsForNull() - { - // Arrange - var requestMessages = new List { new(ChatRole.User, "Hello") }; - var context = new ChatHistoryProvider.InvokedContext(s_mockAgent, s_mockSession, requestMessages); - - // Act & Assert - Assert.Throws(() => context.RequestMessages = null!); - } - - [Fact] - public void InvokedContext_RequestMessages_SetterRoundtrips() - { - // Arrange - var initialMessages = new List { new(ChatRole.User, "Hello") }; - var newMessages = new List { new(ChatRole.User, "New message") }; - var context = new ChatHistoryProvider.InvokedContext(s_mockAgent, s_mockSession, initialMessages); - - // Act - context.RequestMessages = newMessages; - - // Assert - Assert.Same(newMessages, context.RequestMessages); + Assert.Throws(() => new ChatHistoryProvider.InvokedContext(s_mockAgent, s_mockSession, null!, [])); } [Fact] @@ -294,10 +181,9 @@ public class ChatHistoryProviderTests // Arrange var requestMessages = new List { new(ChatRole.User, "Hello") }; var responseMessages = new List { new(ChatRole.Assistant, "Response message") }; - var context = new ChatHistoryProvider.InvokedContext(s_mockAgent, s_mockSession, requestMessages); // Act - context.ResponseMessages = responseMessages; + var context = new ChatHistoryProvider.InvokedContext(s_mockAgent, s_mockSession, requestMessages, responseMessages); // Assert Assert.Same(responseMessages, context.ResponseMessages); @@ -309,10 +195,9 @@ public class ChatHistoryProviderTests // Arrange var requestMessages = new List { new(ChatRole.User, "Hello") }; var exception = new InvalidOperationException("Test exception"); - var context = new ChatHistoryProvider.InvokedContext(s_mockAgent, s_mockSession, requestMessages); // Act - context.InvokeException = exception; + var context = new ChatHistoryProvider.InvokedContext(s_mockAgent, s_mockSession, requestMessages, exception); // Assert Assert.Same(exception, context.InvokeException); @@ -325,7 +210,7 @@ public class ChatHistoryProviderTests var requestMessages = new List { new(ChatRole.User, "Hello") }; // Act - var context = new ChatHistoryProvider.InvokedContext(s_mockAgent, s_mockSession, requestMessages); + var context = new ChatHistoryProvider.InvokedContext(s_mockAgent, s_mockSession, requestMessages, []); // Assert Assert.Same(s_mockAgent, context.Agent); @@ -338,7 +223,7 @@ public class ChatHistoryProviderTests var requestMessages = new List { new(ChatRole.User, "Hello") }; // Act - var context = new ChatHistoryProvider.InvokedContext(s_mockAgent, s_mockSession, requestMessages); + var context = new ChatHistoryProvider.InvokedContext(s_mockAgent, s_mockSession, requestMessages, []); // Assert Assert.Same(s_mockSession, context.Session); @@ -351,7 +236,7 @@ public class ChatHistoryProviderTests var requestMessages = new List { new(ChatRole.User, "Hello") }; // Act - var context = new ChatHistoryProvider.InvokedContext(s_mockAgent, null, requestMessages); + var context = new ChatHistoryProvider.InvokedContext(s_mockAgent, null, requestMessages, []); // Assert Assert.Null(context.Session); @@ -364,7 +249,27 @@ public class ChatHistoryProviderTests var requestMessages = new List { new(ChatRole.User, "Hello") }; // Act & Assert - Assert.Throws(() => new ChatHistoryProvider.InvokedContext(null!, s_mockSession, requestMessages)); + Assert.Throws(() => new ChatHistoryProvider.InvokedContext(null!, s_mockSession, requestMessages, [])); + } + + [Fact] + public void InvokedContext_SuccessConstructor_ThrowsForNullResponseMessages() + { + // Arrange + var requestMessages = new List { new(ChatRole.User, "Hello") }; + + // Act & Assert + Assert.Throws(() => new ChatHistoryProvider.InvokedContext(s_mockAgent, s_mockSession, requestMessages, (IEnumerable)null!)); + } + + [Fact] + public void InvokedContext_FailureConstructor_ThrowsForNullException() + { + // Arrange + var requestMessages = new List { new(ChatRole.User, "Hello") }; + + // Act & Assert + Assert.Throws(() => new ChatHistoryProvider.InvokedContext(s_mockAgent, s_mockSession, requestMessages, (Exception)null!)); } #endregion @@ -372,63 +277,9 @@ public class ChatHistoryProviderTests private sealed class TestChatHistoryProvider : ChatHistoryProvider { protected override ValueTask> InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default) - => new([new ChatMessage(ChatRole.User, "Test Message")]); + => new(new ChatMessage[] { new(ChatRole.User, "Test Message") }.Concat(context.RequestMessages)); protected override ValueTask InvokedCoreAsync(InvokedContext context, CancellationToken cancellationToken = default) => default; - - public override JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null) - => default; - } - - private sealed class TestChatHistoryProviderWithCustomSource : ChatHistoryProvider - { - public TestChatHistoryProviderWithCustomSource(string sourceId) : base(sourceId) - { - } - - protected override ValueTask> InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default) - => new([new ChatMessage(ChatRole.User, "Test Message")]); - - protected override ValueTask InvokedCoreAsync(InvokedContext context, CancellationToken cancellationToken = default) - => default; - - public override JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null) - => default; - } - - private sealed class TestChatHistoryProviderWithPreStampedMessages : ChatHistoryProvider - { - protected override ValueTask> InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default) - { - var message = new ChatMessage(ChatRole.User, "Pre-stamped Message"); - message.AdditionalProperties = new AdditionalPropertiesDictionary - { - [AgentRequestMessageSourceAttribution.AdditionalPropertiesKey] = new AgentRequestMessageSourceAttribution(AgentRequestMessageSourceType.ChatHistory, this.GetType().FullName!) - }; - return new([message]); - } - - protected override ValueTask InvokedCoreAsync(InvokedContext context, CancellationToken cancellationToken = default) - => default; - - public override JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null) - => default; - } - - private sealed class TestChatHistoryProviderWithMultipleMessages : ChatHistoryProvider - { - protected override ValueTask> InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default) - => new([ - new ChatMessage(ChatRole.User, "Message 1"), - new ChatMessage(ChatRole.Assistant, "Message 2"), - new ChatMessage(ChatRole.User, "Message 3") - ]); - - protected override ValueTask InvokedCoreAsync(InvokedContext context, CancellationToken cancellationToken = default) - => default; - - public override JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null) - => default; } } diff --git a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/ChatMessageExtensionsTests.cs b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/ChatMessageExtensionsTests.cs index 97050c3071..05fd576798 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/ChatMessageExtensionsTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/ChatMessageExtensionsTests.cs @@ -350,7 +350,7 @@ public sealed class ChatMessageExtensionsTests ChatMessage message = new(ChatRole.User, "Hello"); // Act - ChatMessage result = message.AsAgentRequestMessageSourcedMessage(AgentRequestMessageSourceType.External, "TestSourceId"); + ChatMessage result = message.WithAgentRequestMessageSource(AgentRequestMessageSourceType.External, "TestSourceId"); // Assert Assert.NotSame(message, result); @@ -368,7 +368,7 @@ public sealed class ChatMessageExtensionsTests }; // Act - ChatMessage result = message.AsAgentRequestMessageSourcedMessage(AgentRequestMessageSourceType.AIContextProvider, "ProviderSourceId"); + ChatMessage result = message.WithAgentRequestMessageSource(AgentRequestMessageSourceType.AIContextProvider, "ProviderSourceId"); // Assert Assert.NotSame(message, result); @@ -389,7 +389,7 @@ public sealed class ChatMessageExtensionsTests }; // Act - ChatMessage result = message.AsAgentRequestMessageSourcedMessage(AgentRequestMessageSourceType.ChatHistory, "HistoryId"); + ChatMessage result = message.WithAgentRequestMessageSource(AgentRequestMessageSourceType.ChatHistory, "HistoryId"); // Assert Assert.Same(message, result); @@ -408,7 +408,7 @@ public sealed class ChatMessageExtensionsTests }; // Act - ChatMessage result = message.AsAgentRequestMessageSourcedMessage(AgentRequestMessageSourceType.AIContextProvider, "SourceId"); + ChatMessage result = message.WithAgentRequestMessageSource(AgentRequestMessageSourceType.AIContextProvider, "SourceId"); // Assert Assert.NotSame(message, result); @@ -429,7 +429,7 @@ public sealed class ChatMessageExtensionsTests }; // Act - ChatMessage result = message.AsAgentRequestMessageSourcedMessage(AgentRequestMessageSourceType.External, "NewId"); + ChatMessage result = message.WithAgentRequestMessageSource(AgentRequestMessageSourceType.External, "NewId"); // Assert Assert.NotSame(message, result); @@ -444,7 +444,7 @@ public sealed class ChatMessageExtensionsTests ChatMessage message = new(ChatRole.User, "Hello"); // Act - ChatMessage result = message.AsAgentRequestMessageSourcedMessage(AgentRequestMessageSourceType.ChatHistory); + ChatMessage result = message.WithAgentRequestMessageSource(AgentRequestMessageSourceType.ChatHistory); // Assert Assert.NotSame(message, result); @@ -465,7 +465,7 @@ public sealed class ChatMessageExtensionsTests }; // Act - ChatMessage result = message.AsAgentRequestMessageSourcedMessage(AgentRequestMessageSourceType.External); + ChatMessage result = message.WithAgentRequestMessageSource(AgentRequestMessageSourceType.External); // Assert Assert.Same(message, result); @@ -478,7 +478,7 @@ public sealed class ChatMessageExtensionsTests ChatMessage message = new(ChatRole.User, "Hello"); // Act - ChatMessage result = message.AsAgentRequestMessageSourcedMessage(AgentRequestMessageSourceType.AIContextProvider, "ProviderId"); + ChatMessage result = message.WithAgentRequestMessageSource(AgentRequestMessageSourceType.AIContextProvider, "ProviderId"); // Assert Assert.Null(message.AdditionalProperties); @@ -499,7 +499,7 @@ public sealed class ChatMessageExtensionsTests }; // Act - ChatMessage result = message.AsAgentRequestMessageSourcedMessage(AgentRequestMessageSourceType.External, "SourceId"); + ChatMessage result = message.WithAgentRequestMessageSource(AgentRequestMessageSourceType.External, "SourceId"); // Assert Assert.NotSame(message, result); @@ -514,7 +514,7 @@ public sealed class ChatMessageExtensionsTests ChatMessage message = new(ChatRole.Assistant, "Test content"); // Act - ChatMessage result = message.AsAgentRequestMessageSourcedMessage(AgentRequestMessageSourceType.ChatHistory, "HistoryId"); + ChatMessage result = message.WithAgentRequestMessageSource(AgentRequestMessageSourceType.ChatHistory, "HistoryId"); // Assert Assert.Equal(ChatRole.Assistant, result.Role); diff --git a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/InMemoryAgentSessionTests.cs b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/InMemoryAgentSessionTests.cs deleted file mode 100644 index a3a4bf7e0e..0000000000 --- a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/InMemoryAgentSessionTests.cs +++ /dev/null @@ -1,155 +0,0 @@ -// Copyright (c) Microsoft. All rights reserved. - -using System; -using System.Collections.Generic; -using System.Linq; -using System.Text.Json; -using Microsoft.Extensions.AI; - -namespace Microsoft.Agents.AI.Abstractions.UnitTests; - -/// -/// Contains tests for . -/// -public class InMemoryAgentSessionTests -{ - #region Constructor and Property Tests - - [Fact] - public void Constructor_SetsDefaultChatHistoryProvider() - { - // Arrange & Act - var session = new TestInMemoryAgentSession(); - - // Assert - Assert.NotNull(session.GetChatHistoryProvider()); - Assert.Empty(session.GetChatHistoryProvider()); - } - - [Fact] - public void Constructor_WithChatHistoryProvider_SetsProperty() - { - // Arrange - InMemoryChatHistoryProvider provider = [new(ChatRole.User, "Hello")]; - - // Act - var session = new TestInMemoryAgentSession(provider); - - // Assert - Assert.Same(provider, session.GetChatHistoryProvider()); - Assert.Single(session.GetChatHistoryProvider()); - Assert.Equal("Hello", session.GetChatHistoryProvider()[0].Text); - } - - [Fact] - public void Constructor_WithMessages_SetsProperty() - { - // Arrange - var messages = new List { new(ChatRole.User, "Hi") }; - - // Act - var session = new TestInMemoryAgentSession(messages); - - // Assert - Assert.NotNull(session.GetChatHistoryProvider()); - Assert.Single(session.GetChatHistoryProvider()); - Assert.Equal("Hi", session.GetChatHistoryProvider()[0].Text); - } - - [Fact] - public void Constructor_WithSerializedState_SetsProperty() - { - // Arrange - InMemoryChatHistoryProvider provider = [new(ChatRole.User, "TestMsg")]; - var providerState = provider.Serialize(); - var sessionStateWrapper = new InMemoryAgentSession.InMemoryAgentSessionState { ChatHistoryProviderState = providerState }; - var json = JsonSerializer.SerializeToElement(sessionStateWrapper, TestJsonSerializerContext.Default.InMemoryAgentSessionState); - - // Act - var session = new TestInMemoryAgentSession(json); - - // Assert - Assert.NotNull(session.GetChatHistoryProvider()); - Assert.Single(session.GetChatHistoryProvider()); - Assert.Equal("TestMsg", session.GetChatHistoryProvider()[0].Text); - } - - [Fact] - public void Constructor_WithInvalidJson_ThrowsArgumentException() - { - // Arrange - var invalidJson = JsonSerializer.SerializeToElement(42, TestJsonSerializerContext.Default.Int32); - - // Act & Assert - Assert.Throws(() => new TestInMemoryAgentSession(invalidJson)); - } - - #endregion - - #region SerializeAsync Tests - - [Fact] - public void Serialize_ReturnsCorrectJson_WhenMessagesExist() - { - // Arrange - var session = new TestInMemoryAgentSession([new(ChatRole.User, "TestContent")]); - - // Act - var json = session.Serialize(); - - // Assert - Assert.Equal(JsonValueKind.Object, json.ValueKind); - Assert.True(json.TryGetProperty("chatHistoryProviderState", out var providerStateProperty)); - Assert.Equal(JsonValueKind.Object, providerStateProperty.ValueKind); - Assert.True(providerStateProperty.TryGetProperty("messages", out var messagesProperty)); - Assert.Equal(JsonValueKind.Array, messagesProperty.ValueKind); - var messagesList = messagesProperty.EnumerateArray().ToList(); - Assert.Single(messagesList); - } - - [Fact] - public void Serialize_ReturnsEmptyMessages_WhenNoMessages() - { - // Arrange - var session = new TestInMemoryAgentSession(); - - // Act - var json = session.Serialize(); - - // Assert - Assert.Equal(JsonValueKind.Object, json.ValueKind); - Assert.True(json.TryGetProperty("chatHistoryProviderState", out var providerStateProperty)); - Assert.Equal(JsonValueKind.Object, providerStateProperty.ValueKind); - Assert.True(providerStateProperty.TryGetProperty("messages", out var messagesProperty)); - Assert.Equal(JsonValueKind.Array, messagesProperty.ValueKind); - Assert.Empty(messagesProperty.EnumerateArray()); - } - - #endregion - - #region GetService Tests - - [Fact] - public void GetService_RequestingChatHistoryProvider_ReturnsChatHistoryProvider() - { - // Arrange - var session = new TestInMemoryAgentSession(); - - // Act & Assert - Assert.NotNull(session.GetService(typeof(ChatHistoryProvider))); - Assert.Same(session.GetChatHistoryProvider(), session.GetService(typeof(ChatHistoryProvider))); - Assert.Same(session.GetChatHistoryProvider(), session.GetService(typeof(InMemoryChatHistoryProvider))); - } - - #endregion - - // Sealed test subclass to expose protected members for testing - private sealed class TestInMemoryAgentSession : InMemoryAgentSession - { - public TestInMemoryAgentSession() { } - public TestInMemoryAgentSession(InMemoryChatHistoryProvider? provider) : base(provider) { } - public TestInMemoryAgentSession(IEnumerable messages) : base(messages) { } - public TestInMemoryAgentSession(JsonElement serializedSessionState) : base(serializedSessionState) { } - public InMemoryChatHistoryProvider GetChatHistoryProvider() => this.ChatHistoryProvider; - } -} diff --git a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/InMemoryChatHistoryProviderTests.cs b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/InMemoryChatHistoryProviderTests.cs index cecce5bea1..b907529241 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/InMemoryChatHistoryProviderTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/InMemoryChatHistoryProviderTests.cs @@ -3,9 +3,6 @@ using System; using System.Collections.Generic; using System.Linq; -using System.Text.Encodings.Web; -using System.Text.Json; -using System.Text.Json.Serialization.Metadata; using System.Threading; using System.Threading.Tasks; using Microsoft.Extensions.AI; @@ -19,22 +16,18 @@ namespace Microsoft.Agents.AI.Abstractions.UnitTests; public class InMemoryChatHistoryProviderTests { private static readonly AIAgent s_mockAgent = new Mock().Object; - private static readonly AgentSession s_mockSession = new Mock().Object; - [Fact] - public void Constructor_Throws_ForNullReducer() => - // Arrange & Act & Assert - Assert.Throws(() => new InMemoryChatHistoryProvider(null!)); + private static AgentSession CreateMockSession() => new Mock().Object; [Fact] public void Constructor_DefaultsToBeforeMessageRetrieval_ForNotProvidedTriggerEvent() { // Arrange & Act var reducerMock = new Mock(); - var provider = new InMemoryChatHistoryProvider(reducerMock.Object); + var provider = new InMemoryChatHistoryProvider(new() { ChatReducer = reducerMock.Object }); // Assert - Assert.Equal(InMemoryChatHistoryProvider.ChatReducerTriggerEvent.BeforeMessagesRetrieval, provider.ReducerTriggerEvent); + Assert.Equal(InMemoryChatHistoryProviderOptions.ChatReducerTriggerEvent.BeforeMessagesRetrieval, provider.ReducerTriggerEvent); } [Fact] @@ -42,20 +35,43 @@ public class InMemoryChatHistoryProviderTests { // Arrange & Act var reducerMock = new Mock(); - var provider = new InMemoryChatHistoryProvider(reducerMock.Object, InMemoryChatHistoryProvider.ChatReducerTriggerEvent.AfterMessageAdded); + var provider = new InMemoryChatHistoryProvider(new() { ChatReducer = reducerMock.Object, ReducerTriggerEvent = InMemoryChatHistoryProviderOptions.ChatReducerTriggerEvent.AfterMessageAdded }); // Assert Assert.Same(reducerMock.Object, provider.ChatReducer); - Assert.Equal(InMemoryChatHistoryProvider.ChatReducerTriggerEvent.AfterMessageAdded, provider.ReducerTriggerEvent); + Assert.Equal(InMemoryChatHistoryProviderOptions.ChatReducerTriggerEvent.AfterMessageAdded, provider.ReducerTriggerEvent); + } + + [Fact] + public void StateKey_ReturnsDefaultKey_WhenNoOptionsProvided() + { + // Arrange & Act + var provider = new InMemoryChatHistoryProvider(); + + // Assert + Assert.Equal("InMemoryChatHistoryProvider", provider.StateKey); + } + + [Fact] + public void StateKey_ReturnsCustomKey_WhenSetViaOptions() + { + // Arrange & Act + var provider = new InMemoryChatHistoryProvider(new() { StateKey = "custom-key" }); + + // Assert + Assert.Equal("custom-key", provider.StateKey); } [Fact] public async Task InvokedAsyncAddsMessagesAsync() { + var session = CreateMockSession(); + + // Arrange var requestMessages = new List { new(ChatRole.User, "Hello"), - new(ChatRole.System, "additional context") { AdditionalProperties = new() { { AgentRequestMessageSourceAttribution.AdditionalPropertiesKey, new AgentRequestMessageSourceAttribution(AgentRequestMessageSourceType.ChatHistory, "TestSource") } } }, + new(ChatRole.System, "additional context") { AdditionalProperties = new() { { AgentRequestMessageSourceAttribution.AdditionalPropertiesKey, new AgentRequestMessageSourceAttribution(AgentRequestMessageSourceType.AIContextProvider, "TestSource") } } }, }; var responseMessages = new List { @@ -67,125 +83,146 @@ public class InMemoryChatHistoryProviderTests }; var provider = new InMemoryChatHistoryProvider(); - provider.Add(providerMessages[0]); - var context = new ChatHistoryProvider.InvokedContext(s_mockAgent, s_mockSession, requestMessages) - { - ResponseMessages = responseMessages - }; + provider.SetMessages(session, providerMessages); + var context = new ChatHistoryProvider.InvokedContext(s_mockAgent, session, requestMessages, responseMessages); await provider.InvokedAsync(context, CancellationToken.None); - Assert.Equal(4, provider.Count); - Assert.Equal("original instructions", provider[0].Text); - Assert.Equal("Hello", provider[1].Text); - Assert.Equal("additional context", provider[2].Text); - Assert.Equal("Hi there!", provider[3].Text); + // Assert + var messages = provider.GetMessages(session); + Assert.Equal(4, messages.Count); + Assert.Equal("original instructions", messages[0].Text); + Assert.Equal("Hello", messages[1].Text); + Assert.Equal("additional context", messages[2].Text); + Assert.Equal("Hi there!", messages[3].Text); } [Fact] public async Task InvokedAsyncWithEmptyDoesNotFailAsync() { + var session = CreateMockSession(); + + // Arrange var provider = new InMemoryChatHistoryProvider(); - var context = new ChatHistoryProvider.InvokedContext(s_mockAgent, s_mockSession, []); + var context = new ChatHistoryProvider.InvokedContext(s_mockAgent, session, [], []); await provider.InvokedAsync(context, CancellationToken.None); - - Assert.Empty(provider); + // Assert + Assert.Empty(provider.GetMessages(session)); } [Fact] public async Task InvokingAsyncReturnsAllMessagesAsync() { - var provider = new InMemoryChatHistoryProvider + var session = CreateMockSession(); + + // Arrange + var requestMessages = new List { + new(ChatRole.User, "Hello"), + }; + + var provider = new InMemoryChatHistoryProvider(); + provider.SetMessages(session, + [ new ChatMessage(ChatRole.User, "Test1"), new ChatMessage(ChatRole.Assistant, "Test2") - }; + ]); - var context = new ChatHistoryProvider.InvokingContext(s_mockAgent, s_mockSession, []); + var context = new ChatHistoryProvider.InvokingContext(s_mockAgent, session, requestMessages); var result = (await provider.InvokingAsync(context, CancellationToken.None)).ToList(); - Assert.Equal(2, result.Count); + // Assert + Assert.Equal(3, result.Count); Assert.Contains(result, m => m.Text == "Test1"); Assert.Contains(result, m => m.Text == "Test2"); + Assert.Contains(result, m => m.Text == "Hello"); + + Assert.Equal(AgentRequestMessageSourceType.ChatHistory, result[0].GetAgentRequestMessageSourceType()); + Assert.Equal(AgentRequestMessageSourceType.ChatHistory, result[1].GetAgentRequestMessageSourceType()); + Assert.Equal(AgentRequestMessageSourceType.External, result[2].GetAgentRequestMessageSourceType()); } [Fact] - public async Task DeserializeConstructorWithEmptyElementAsync() + public void StateInitializer_IsInvoked_WhenSessionHasNoState() { - var emptyObject = JsonSerializer.Deserialize("{}", TestJsonSerializerContext.Default.JsonElement); + // Arrange + var initialMessages = new List + { + new(ChatRole.User, "Initial message") + }; + var provider = new InMemoryChatHistoryProvider(new() + { + StateInitializer = _ => new InMemoryChatHistoryProvider.State { Messages = initialMessages } + }); - var newProvider = new InMemoryChatHistoryProvider(emptyObject); + // Act + var messages = provider.GetMessages(CreateMockSession()); - Assert.Empty(newProvider); + // Assert + Assert.Single(messages); + Assert.Equal("Initial message", messages[0].Text); } [Fact] - public async Task SerializeAndDeserializeConstructorRoundtripsAsync() + public void GetMessages_ReturnsEmptyList_WhenNullSession() { - var provider = new InMemoryChatHistoryProvider - { - new ChatMessage(ChatRole.User, "A"), - new ChatMessage(ChatRole.Assistant, "B") - }; + // Arrange + var provider = new InMemoryChatHistoryProvider(); - var jsonElement = provider.Serialize(); - var newProvider = new InMemoryChatHistoryProvider(jsonElement); + // Act + var messages = provider.GetMessages(null); - Assert.Equal(2, newProvider.Count); - Assert.Equal("A", newProvider[0].Text); - Assert.Equal("B", newProvider[1].Text); + // Assert + Assert.Empty(messages); } [Fact] - public async Task SerializeAndDeserializeConstructorRoundtripsWithCustomAIContentAsync() + public void SetMessages_ThrowsForNullMessages() { - JsonSerializerOptions options = new(TestJsonSerializerContext.Default.Options) - { - TypeInfoResolver = JsonTypeInfoResolver.Combine(AgentAbstractionsJsonUtilities.DefaultOptions.TypeInfoResolver, TestJsonSerializerContext.Default), - Encoder = JavaScriptEncoder.UnsafeRelaxedJsonEscaping, - }; - options.AddAIContentType(typeDiscriminatorId: "testContent"); + // Arrange + var provider = new InMemoryChatHistoryProvider(); - var provider = new InMemoryChatHistoryProvider - { - new ChatMessage(ChatRole.User, [new TestAIContent("foo data")]), - }; - - var jsonElement = provider.Serialize(options); - var newProvider = new InMemoryChatHistoryProvider(jsonElement, options); - - Assert.Single(newProvider); - var actualTestAIContent = Assert.IsType(newProvider[0].Contents[0]); - Assert.Equal("foo data", actualTestAIContent.TestData); + // Act & Assert + Assert.Throws(() => provider.SetMessages(CreateMockSession(), null!)); } [Fact] - public async Task SerializeAndDeserializeWorksWithExperimentalContentTypesAsync() + public void SetMessages_UpdatesState() { - var provider = new InMemoryChatHistoryProvider + var session = CreateMockSession(); + + // Arrange + var provider = new InMemoryChatHistoryProvider(); + var messages = new List { - new ChatMessage(ChatRole.User, [new FunctionApprovalRequestContent("call123", new FunctionCallContent("call123", "some_func"))]), - new ChatMessage(ChatRole.Assistant, [new FunctionApprovalResponseContent("call123", true, new FunctionCallContent("call123", "some_func"))]) + new(ChatRole.User, "Hello"), + new(ChatRole.Assistant, "World") }; - var jsonElement = provider.Serialize(); - var newProvider = new InMemoryChatHistoryProvider(jsonElement); + // Act + provider.SetMessages(session, messages); + var retrieved = provider.GetMessages(session); - Assert.Equal(2, newProvider.Count); - Assert.IsType(newProvider[0].Contents[0]); - Assert.IsType(newProvider[1].Contents[0]); + // Assert + Assert.Equal(2, retrieved.Count); + Assert.Equal("Hello", retrieved[0].Text); + Assert.Equal("World", retrieved[1].Text); } [Fact] public async Task InvokedAsyncWithEmptyMessagesDoesNotChangeProviderAsync() { + var session = CreateMockSession(); + + // Arrange var provider = new InMemoryChatHistoryProvider(); var messages = new List(); - var context = new ChatHistoryProvider.InvokedContext(s_mockAgent, s_mockSession, messages); + var context = new ChatHistoryProvider.InvokedContext(s_mockAgent, session, messages, []); await provider.InvokedAsync(context, CancellationToken.None); - Assert.Empty(provider); + // Assert + Assert.Empty(provider.GetMessages(session)); } [Fact] @@ -198,308 +235,11 @@ public class InMemoryChatHistoryProviderTests await Assert.ThrowsAsync(() => provider.InvokedAsync(null!, CancellationToken.None).AsTask()); } - [Fact] - public void DeserializeContructor_WithNullSerializedState_CreatesEmptyProvider() - { - // Act - var provider = new InMemoryChatHistoryProvider(new JsonElement()); - - // Assert - Assert.Empty(provider); - } - - [Fact] - public async Task DeserializeContructor_WithEmptyMessages_DoesNotAddMessagesAsync() - { - // Arrange - var stateWithEmptyMessages = JsonSerializer.SerializeToElement( - new Dictionary { ["messages"] = new List() }, - TestJsonSerializerContext.Default.IDictionaryStringObject); - - // Act - var provider = new InMemoryChatHistoryProvider(stateWithEmptyMessages); - - // Assert - Assert.Empty(provider); - } - - [Fact] - public async Task DeserializeConstructor_WithNullMessages_DoesNotAddMessagesAsync() - { - // Arrange - var stateWithNullMessages = JsonSerializer.SerializeToElement( - new Dictionary { ["messages"] = null! }, - TestJsonSerializerContext.Default.DictionaryStringObject); - - // Act - var provider = new InMemoryChatHistoryProvider(stateWithNullMessages); - - // Assert - Assert.Empty(provider); - } - - [Fact] - public async Task DeserializeConstructor_WithValidMessages_AddsMessagesAsync() - { - // Arrange - var messages = new List - { - new(ChatRole.User, "User message"), - new(ChatRole.Assistant, "Assistant message") - }; - var state = new Dictionary { ["messages"] = messages }; - var serializedState = JsonSerializer.SerializeToElement( - state, - TestJsonSerializerContext.Default.DictionaryStringObject); - - // Act - var provider = new InMemoryChatHistoryProvider(serializedState); - - // Assert - Assert.Equal(2, provider.Count); - Assert.Equal("User message", provider[0].Text); - Assert.Equal("Assistant message", provider[1].Text); - } - - [Fact] - public void IndexerGet_ReturnsCorrectMessage() - { - // Arrange - var provider = new InMemoryChatHistoryProvider(); - var message1 = new ChatMessage(ChatRole.User, "First"); - var message2 = new ChatMessage(ChatRole.Assistant, "Second"); - provider.Add(message1); - provider.Add(message2); - - // Act & Assert - Assert.Same(message1, provider[0]); - Assert.Same(message2, provider[1]); - } - - [Fact] - public void IndexerSet_UpdatesMessage() - { - // Arrange - var provider = new InMemoryChatHistoryProvider(); - var originalMessage = new ChatMessage(ChatRole.User, "Original"); - var newMessage = new ChatMessage(ChatRole.User, "Updated"); - provider.Add(originalMessage); - - // Act - provider[0] = newMessage; - - // Assert - Assert.Same(newMessage, provider[0]); - Assert.Equal("Updated", provider[0].Text); - } - - [Fact] - public void IsReadOnly_ReturnsFalse() - { - // Arrange - var provider = new InMemoryChatHistoryProvider(); - - // Act & Assert - Assert.False(provider.IsReadOnly); - } - - [Fact] - public void IndexOf_ReturnsCorrectIndex() - { - // Arrange - var provider = new InMemoryChatHistoryProvider(); - var message1 = new ChatMessage(ChatRole.User, "First"); - var message2 = new ChatMessage(ChatRole.Assistant, "Second"); - var message3 = new ChatMessage(ChatRole.User, "Third"); - provider.Add(message1); - provider.Add(message2); - - // Act & Assert - Assert.Equal(0, provider.IndexOf(message1)); - Assert.Equal(1, provider.IndexOf(message2)); - Assert.Equal(-1, provider.IndexOf(message3)); // Not in provider - } - - [Fact] - public void Insert_InsertsMessageAtCorrectIndex() - { - // Arrange - var provider = new InMemoryChatHistoryProvider(); - var message1 = new ChatMessage(ChatRole.User, "First"); - var message2 = new ChatMessage(ChatRole.Assistant, "Second"); - var insertMessage = new ChatMessage(ChatRole.User, "Inserted"); - provider.Add(message1); - provider.Add(message2); - - // Act - provider.Insert(1, insertMessage); - - // Assert - Assert.Equal(3, provider.Count); - Assert.Same(message1, provider[0]); - Assert.Same(insertMessage, provider[1]); - Assert.Same(message2, provider[2]); - } - - [Fact] - public void RemoveAt_RemovesMessageAtIndex() - { - // Arrange - var provider = new InMemoryChatHistoryProvider(); - var message1 = new ChatMessage(ChatRole.User, "First"); - var message2 = new ChatMessage(ChatRole.Assistant, "Second"); - var message3 = new ChatMessage(ChatRole.User, "Third"); - provider.Add(message1); - provider.Add(message2); - provider.Add(message3); - - // Act - provider.RemoveAt(1); - - // Assert - Assert.Equal(2, provider.Count); - Assert.Same(message1, provider[0]); - Assert.Same(message3, provider[1]); - } - - [Fact] - public void Clear_RemovesAllMessages() - { - // Arrange - var provider = new InMemoryChatHistoryProvider - { - new ChatMessage(ChatRole.User, "First"), - new ChatMessage(ChatRole.Assistant, "Second") - }; - - // Act - provider.Clear(); - - // Assert - Assert.Empty(provider); - } - - [Fact] - public void Contains_ReturnsTrueForExistingMessage() - { - // Arrange - var provider = new InMemoryChatHistoryProvider(); - var message1 = new ChatMessage(ChatRole.User, "First"); - var message2 = new ChatMessage(ChatRole.Assistant, "Second"); - provider.Add(message1); - - // Act & Assert - Assert.Contains(message1, provider); - Assert.DoesNotContain(message2, provider); - } - - [Fact] - public void CopyTo_CopiesMessagesToArray() - { - // Arrange - var provider = new InMemoryChatHistoryProvider(); - var message1 = new ChatMessage(ChatRole.User, "First"); - var message2 = new ChatMessage(ChatRole.Assistant, "Second"); - provider.Add(message1); - provider.Add(message2); - var array = new ChatMessage[4]; - - // Act - provider.CopyTo(array, 1); - - // Assert - Assert.Null(array[0]); - Assert.Same(message1, array[1]); - Assert.Same(message2, array[2]); - Assert.Null(array[3]); - } - - [Fact] - public void Remove_RemovesSpecificMessage() - { - // Arrange - var provider = new InMemoryChatHistoryProvider(); - var message1 = new ChatMessage(ChatRole.User, "First"); - var message2 = new ChatMessage(ChatRole.Assistant, "Second"); - var message3 = new ChatMessage(ChatRole.User, "Third"); - provider.Add(message1); - provider.Add(message2); - provider.Add(message3); - - // Act - var removed = provider.Remove(message2); - - // Assert - Assert.True(removed); - Assert.Equal(2, provider.Count); - Assert.Same(message1, provider[0]); - Assert.Same(message3, provider[1]); - } - - [Fact] - public void Remove_ReturnsFalseForNonExistentMessage() - { - // Arrange - var provider = new InMemoryChatHistoryProvider(); - var message1 = new ChatMessage(ChatRole.User, "First"); - var message2 = new ChatMessage(ChatRole.Assistant, "Second"); - provider.Add(message1); - - // Act - var removed = provider.Remove(message2); - - // Assert - Assert.False(removed); - Assert.Single(provider); - } - - [Fact] - public void GetEnumerator_Generic_ReturnsAllMessages() - { - // Arrange - var provider = new InMemoryChatHistoryProvider(); - var message1 = new ChatMessage(ChatRole.User, "First"); - var message2 = new ChatMessage(ChatRole.Assistant, "Second"); - provider.Add(message1); - provider.Add(message2); - - // Act - var messages = new List(); - messages.AddRange(provider); - - // Assert - Assert.Equal(2, messages.Count); - Assert.Same(message1, messages[0]); - Assert.Same(message2, messages[1]); - } - - [Fact] - public void GetEnumerator_NonGeneric_ReturnsAllMessages() - { - // Arrange - var provider = new InMemoryChatHistoryProvider(); - var message1 = new ChatMessage(ChatRole.User, "First"); - var message2 = new ChatMessage(ChatRole.Assistant, "Second"); - provider.Add(message1); - provider.Add(message2); - - // Act - var messages = new List(); - var enumerator = ((System.Collections.IEnumerable)provider).GetEnumerator(); - while (enumerator.MoveNext()) - { - messages.Add((ChatMessage)enumerator.Current); - } - - // Assert - Assert.Equal(2, messages.Count); - Assert.Same(message1, messages[0]); - Assert.Same(message2, messages[1]); - } - [Fact] public async Task AddMessagesAsync_WithReducer_AfterMessageAdded_InvokesReducerAsync() { + var session = CreateMockSession(); + // Arrange var originalMessages = new List { @@ -516,21 +256,24 @@ public class InMemoryChatHistoryProviderTests .Setup(r => r.ReduceAsync(It.Is>(x => x.SequenceEqual(originalMessages)), It.IsAny())) .ReturnsAsync(reducedMessages); - var provider = new InMemoryChatHistoryProvider(reducerMock.Object, InMemoryChatHistoryProvider.ChatReducerTriggerEvent.AfterMessageAdded); + var provider = new InMemoryChatHistoryProvider(new() { ChatReducer = reducerMock.Object, ReducerTriggerEvent = InMemoryChatHistoryProviderOptions.ChatReducerTriggerEvent.AfterMessageAdded }); // Act - var context = new ChatHistoryProvider.InvokedContext(s_mockAgent, s_mockSession, originalMessages); + var context = new ChatHistoryProvider.InvokedContext(s_mockAgent, session, originalMessages, []); await provider.InvokedAsync(context, CancellationToken.None); // Assert - Assert.Single(provider); - Assert.Equal("Reduced", provider[0].Text); + var messages = provider.GetMessages(session); + Assert.Single(messages); + Assert.Equal("Reduced", messages[0].Text); reducerMock.Verify(r => r.ReduceAsync(It.Is>(x => x.SequenceEqual(originalMessages)), It.IsAny()), Times.Once); } [Fact] public async Task GetMessagesAsync_WithReducer_BeforeMessagesRetrieval_InvokesReducerAsync() { + var session = CreateMockSession(); + // Arrange var originalMessages = new List { @@ -547,15 +290,11 @@ public class InMemoryChatHistoryProviderTests .Setup(r => r.ReduceAsync(It.Is>(x => x.SequenceEqual(originalMessages)), It.IsAny())) .ReturnsAsync(reducedMessages); - var provider = new InMemoryChatHistoryProvider(reducerMock.Object, InMemoryChatHistoryProvider.ChatReducerTriggerEvent.BeforeMessagesRetrieval); - // Add messages directly to the provider for this test - foreach (var msg in originalMessages) - { - provider.Add(msg); - } + var provider = new InMemoryChatHistoryProvider(new() { ChatReducer = reducerMock.Object, ReducerTriggerEvent = InMemoryChatHistoryProviderOptions.ChatReducerTriggerEvent.BeforeMessagesRetrieval }); + provider.SetMessages(session, new List(originalMessages)); // Act - var invokingContext = new ChatHistoryProvider.InvokingContext(s_mockAgent, s_mockSession, Array.Empty()); + var invokingContext = new ChatHistoryProvider.InvokingContext(s_mockAgent, session, Array.Empty()); var result = (await provider.InvokingAsync(invokingContext, CancellationToken.None)).ToList(); // Assert @@ -567,6 +306,8 @@ public class InMemoryChatHistoryProviderTests [Fact] public async Task AddMessagesAsync_WithReducer_ButWrongTrigger_DoesNotInvokeReducerAsync() { + var session = CreateMockSession(); + // Arrange var originalMessages = new List { @@ -575,21 +316,24 @@ public class InMemoryChatHistoryProviderTests var reducerMock = new Mock(); - var provider = new InMemoryChatHistoryProvider(reducerMock.Object, InMemoryChatHistoryProvider.ChatReducerTriggerEvent.BeforeMessagesRetrieval); + var provider = new InMemoryChatHistoryProvider(new() { ChatReducer = reducerMock.Object, ReducerTriggerEvent = InMemoryChatHistoryProviderOptions.ChatReducerTriggerEvent.BeforeMessagesRetrieval }); // Act - var context = new ChatHistoryProvider.InvokedContext(s_mockAgent, s_mockSession, originalMessages); + var context = new ChatHistoryProvider.InvokedContext(s_mockAgent, session, originalMessages, []); await provider.InvokedAsync(context, CancellationToken.None); // Assert - Assert.Single(provider); - Assert.Equal("Hello", provider[0].Text); + var messages = provider.GetMessages(session); + Assert.Single(messages); + Assert.Equal("Hello", messages[0].Text); reducerMock.Verify(r => r.ReduceAsync(It.IsAny>(), It.IsAny()), Times.Never); } [Fact] public async Task GetMessagesAsync_WithReducer_ButWrongTrigger_DoesNotInvokeReducerAsync() { + var session = CreateMockSession(); + // Arrange var originalMessages = new List { @@ -598,13 +342,11 @@ public class InMemoryChatHistoryProviderTests var reducerMock = new Mock(); - var provider = new InMemoryChatHistoryProvider(reducerMock.Object, InMemoryChatHistoryProvider.ChatReducerTriggerEvent.AfterMessageAdded) - { - originalMessages[0] - }; + var provider = new InMemoryChatHistoryProvider(new() { ChatReducer = reducerMock.Object, ReducerTriggerEvent = InMemoryChatHistoryProviderOptions.ChatReducerTriggerEvent.AfterMessageAdded }); + provider.SetMessages(session, new List(originalMessages)); // Act - var invokingContext = new ChatHistoryProvider.InvokingContext(s_mockAgent, s_mockSession, Array.Empty()); + var invokingContext = new ChatHistoryProvider.InvokingContext(s_mockAgent, session, Array.Empty()); var result = (await provider.InvokingAsync(invokingContext, CancellationToken.None)).ToList(); // Assert @@ -616,27 +358,21 @@ public class InMemoryChatHistoryProviderTests [Fact] public async Task InvokedAsync_WithException_DoesNotAddMessagesAsync() { + var session = CreateMockSession(); + // Arrange var provider = new InMemoryChatHistoryProvider(); var requestMessages = new List { new(ChatRole.User, "Hello") }; - var responseMessages = new List - { - new(ChatRole.Assistant, "Hi there!") - }; - var context = new ChatHistoryProvider.InvokedContext(s_mockAgent, s_mockSession, requestMessages) - { - ResponseMessages = responseMessages, - InvokeException = new InvalidOperationException("Test exception") - }; + var context = new ChatHistoryProvider.InvokedContext(s_mockAgent, session, requestMessages, new InvalidOperationException("Test exception")); // Act await provider.InvokedAsync(context, CancellationToken.None); // Assert - Assert.Empty(provider); + Assert.Empty(provider.GetMessages(session)); } [Fact] @@ -649,6 +385,85 @@ public class InMemoryChatHistoryProviderTests await Assert.ThrowsAsync(() => provider.InvokingAsync(null!, CancellationToken.None).AsTask()); } + [Fact] + public async Task InvokedAsync_DefaultFilter_ExcludesChatHistoryMessagesAsync() + { + // Arrange + var session = CreateMockSession(); + var provider = new InMemoryChatHistoryProvider(); + var requestMessages = new List + { + new(ChatRole.User, "External message"), + new(ChatRole.System, "From history") { AdditionalProperties = new() { { AgentRequestMessageSourceAttribution.AdditionalPropertiesKey, new AgentRequestMessageSourceAttribution(AgentRequestMessageSourceType.ChatHistory, "HistorySource") } } }, + new(ChatRole.System, "From context provider") { AdditionalProperties = new() { { AgentRequestMessageSourceAttribution.AdditionalPropertiesKey, new AgentRequestMessageSourceAttribution(AgentRequestMessageSourceType.AIContextProvider, "ContextSource") } } }, + }; + + var context = new ChatHistoryProvider.InvokedContext(s_mockAgent, session, requestMessages, [new ChatMessage(ChatRole.Assistant, "Response")]); + + // Act + await provider.InvokedAsync(context, CancellationToken.None); + + // Assert - ChatHistory message excluded, AIContextProvider message included + var messages = provider.GetMessages(session); + Assert.Equal(3, messages.Count); + Assert.Equal("External message", messages[0].Text); + Assert.Equal("From context provider", messages[1].Text); + Assert.Equal("Response", messages[2].Text); + } + + [Fact] + public async Task InvokedAsync_CustomFilter_OverridesDefaultAsync() + { + // Arrange + var session = CreateMockSession(); + var provider = new InMemoryChatHistoryProvider(new InMemoryChatHistoryProviderOptions + { + StorageInputMessageFilter = messages => messages.Where(m => m.GetAgentRequestMessageSourceType() == AgentRequestMessageSourceType.External) + }); + var requestMessages = new List + { + new(ChatRole.User, "External message"), + new(ChatRole.System, "From history") { AdditionalProperties = new() { { AgentRequestMessageSourceAttribution.AdditionalPropertiesKey, new AgentRequestMessageSourceAttribution(AgentRequestMessageSourceType.ChatHistory, "HistorySource") } } }, + new(ChatRole.System, "From context provider") { AdditionalProperties = new() { { AgentRequestMessageSourceAttribution.AdditionalPropertiesKey, new AgentRequestMessageSourceAttribution(AgentRequestMessageSourceType.AIContextProvider, "ContextSource") } } }, + }; + + var context = new ChatHistoryProvider.InvokedContext(s_mockAgent, session, requestMessages, [new ChatMessage(ChatRole.Assistant, "Response")]); + + // Act + await provider.InvokedAsync(context, CancellationToken.None); + + // Assert - Custom filter keeps only External messages (both ChatHistory and AIContextProvider excluded) + var messages = provider.GetMessages(session); + Assert.Equal(2, messages.Count); + Assert.Equal("External message", messages[0].Text); + Assert.Equal("Response", messages[1].Text); + } + + [Fact] + public async Task InvokingAsync_OutputFilter_FiltersOutputMessagesAsync() + { + // Arrange + var session = CreateMockSession(); + var provider = new InMemoryChatHistoryProvider(new InMemoryChatHistoryProviderOptions + { + RetrievalOutputMessageFilter = messages => messages.Where(m => m.Role == ChatRole.User) + }); + provider.SetMessages(session, + [ + new ChatMessage(ChatRole.User, "User message"), + new ChatMessage(ChatRole.Assistant, "Assistant message"), + new ChatMessage(ChatRole.System, "System message") + ]); + + // Act + var context = new ChatHistoryProvider.InvokingContext(s_mockAgent, session, []); + var result = (await provider.InvokingAsync(context, CancellationToken.None)).ToList(); + + // Assert - Only user messages pass through the output filter + Assert.Single(result); + Assert.Equal("User message", result[0].Text); + } + public class TestAIContent(string testData) : AIContent { public string TestData => testData; diff --git a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/ServiceIdAgentSessionTests.cs b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/ServiceIdAgentSessionTests.cs deleted file mode 100644 index e4a6626f72..0000000000 --- a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/ServiceIdAgentSessionTests.cs +++ /dev/null @@ -1,119 +0,0 @@ -// Copyright (c) Microsoft. All rights reserved. - -using System; -using System.Text.Json; - -namespace Microsoft.Agents.AI.Abstractions.UnitTests; - -/// -/// Tests for . -/// -public class ServiceIdAgentSessionTests -{ - #region Constructor and Property Tests - - [Fact] - public void Constructor_SetsDefaults() - { - // Arrange & Act - var session = new TestServiceIdAgentSession(); - - // Assert - Assert.Null(session.GetServiceSessionId()); - } - - [Fact] - public void Constructor_WithServiceSessionId_SetsProperty() - { - // Arrange & Act - var session = new TestServiceIdAgentSession("service-id-123"); - - // Assert - Assert.Equal("service-id-123", session.GetServiceSessionId()); - } - - [Fact] - public void Constructor_WithSerializedId_SetsProperty() - { - // Arrange - var serviceSessionWrapper = new ServiceIdAgentSession.ServiceIdAgentSessionState { ServiceSessionId = "service-id-456" }; - var json = JsonSerializer.SerializeToElement(serviceSessionWrapper, TestJsonSerializerContext.Default.ServiceIdAgentSessionState); - - // Act - var session = new TestServiceIdAgentSession(json); - - // Assert - Assert.Equal("service-id-456", session.GetServiceSessionId()); - } - - [Fact] - public void Constructor_WithSerializedUndefinedId_SetsProperty() - { - // Arrange - var emptyObject = new EmptyObject(); - var json = JsonSerializer.SerializeToElement(emptyObject, TestJsonSerializerContext.Default.EmptyObject); - - // Act - var session = new TestServiceIdAgentSession(json); - - // Assert - Assert.Null(session.GetServiceSessionId()); - } - - [Fact] - public void Constructor_WithInvalidJson_ThrowsArgumentException() - { - // Arrange - var invalidJson = JsonSerializer.SerializeToElement(42, TestJsonSerializerContext.Default.Int32); - - // Act & Assert - Assert.Throws(() => new TestServiceIdAgentSession(invalidJson)); - } - - #endregion - - #region SerializeAsync Tests - - [Fact] - public void Serialize_ReturnsCorrectJson_WhenServiceSessionIdIsSet() - { - // Arrange - var session = new TestServiceIdAgentSession("service-id-789"); - - // Act - var json = session.Serialize(); - - // Assert - Assert.Equal(JsonValueKind.Object, json.ValueKind); - Assert.True(json.TryGetProperty("serviceSessionId", out var idProperty)); - Assert.Equal("service-id-789", idProperty.GetString()); - } - - [Fact] - public void Serialize_ReturnsUndefinedServiceSessionId_WhenNotSet() - { - // Arrange - var session = new TestServiceIdAgentSession(); - - // Act - var json = session.Serialize(); - - // Assert - Assert.Equal(JsonValueKind.Object, json.ValueKind); - Assert.False(json.TryGetProperty("serviceSessionId", out _)); - } - - #endregion - - // Sealed test subclass to expose protected members for testing - private sealed class TestServiceIdAgentSession : ServiceIdAgentSession - { - public TestServiceIdAgentSession() { } - public TestServiceIdAgentSession(string serviceSessionId) : base(serviceSessionId) { } - public TestServiceIdAgentSession(JsonElement serializedSessionState) : base(serializedSessionState) { } - public string? GetServiceSessionId() => this.ServiceSessionId; - } - - // Helper class to represent empty objects - internal sealed class EmptyObject; -} diff --git a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/TestJsonSerializerContext.cs b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/TestJsonSerializerContext.cs index 05d13c2e95..c4f3b7511a 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/TestJsonSerializerContext.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/TestJsonSerializerContext.cs @@ -19,8 +19,5 @@ namespace Microsoft.Agents.AI.Abstractions.UnitTests; [JsonSerializable(typeof(Dictionary))] [JsonSerializable(typeof(string[]))] [JsonSerializable(typeof(int))] -[JsonSerializable(typeof(InMemoryAgentSession.InMemoryAgentSessionState))] -[JsonSerializable(typeof(ServiceIdAgentSession.ServiceIdAgentSessionState))] -[JsonSerializable(typeof(ServiceIdAgentSessionTests.EmptyObject))] [JsonSerializable(typeof(InMemoryChatHistoryProviderTests.TestAIContent))] internal sealed partial class TestJsonSerializerContext : JsonSerializerContext; diff --git a/dotnet/tests/Microsoft.Agents.AI.AzureAI.UnitTests/AzureAIProjectChatClientExtensionsTests.cs b/dotnet/tests/Microsoft.Agents.AI.AzureAI.UnitTests/AzureAIProjectChatClientExtensionsTests.cs index e2cbb16b1b..2f2e276ae9 100644 --- a/dotnet/tests/Microsoft.Agents.AI.AzureAI.UnitTests/AzureAIProjectChatClientExtensionsTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.AzureAI.UnitTests/AzureAIProjectChatClientExtensionsTests.cs @@ -2310,23 +2310,18 @@ public sealed class AzureAIProjectChatClientExtensionsTests #region CreateChatClientAgentOptions - Options Preservation Tests /// - /// Verify that CreateChatClientAgentOptions preserves AIContextProviderFactory. + /// Verify that CreateChatClientAgentOptions preserves AIContextProviders. /// [Fact] - public async Task GetAIAgentAsync_WithAIContextProviderFactory_PreservesFactoryAsync() + public async Task GetAIAgentAsync_WithAIContextProviders_PreservesProviderAsync() { // Arrange AIProjectClient client = this.CreateTestAgentClient(); - bool factoryInvoked = false; var options = new ChatClientAgentOptions { Name = "test-agent", ChatOptions = new ChatOptions { Instructions = "Test" }, - AIContextProviderFactory = (_, _) => - { - factoryInvoked = true; - return new ValueTask(new TestAIContextProvider()); - } + AIContextProviders = [new TestAIContextProvider()] }; // Act @@ -2334,15 +2329,13 @@ public sealed class AzureAIProjectChatClientExtensionsTests // Assert Assert.NotNull(agent); - // Verify the factory was captured (though not necessarily invoked yet) - Assert.False(factoryInvoked); // Factory is not invoked during creation } /// - /// Verify that CreateChatClientAgentOptions preserves ChatHistoryProviderFactory. + /// Verify that CreateChatClientAgentOptions preserves ChatHistoryProvider. /// [Fact] - public async Task GetAIAgentAsync_WithChatHistoryProviderFactory_PreservesFactoryAsync() + public async Task GetAIAgentAsync_WithChatHistoryProvider_PreservesProviderAsync() { // Arrange AIProjectClient client = this.CreateTestAgentClient(); @@ -2350,7 +2343,7 @@ public sealed class AzureAIProjectChatClientExtensionsTests { Name = "test-agent", ChatOptions = new ChatOptions { Instructions = "Test" }, - ChatHistoryProviderFactory = (_, _) => new ValueTask(new TestChatHistoryProvider()) + ChatHistoryProvider = new TestChatHistoryProvider() }; // Act @@ -3142,7 +3135,7 @@ public sealed class AzureAIProjectChatClientExtensionsTests { protected override ValueTask InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default) { - return new ValueTask(new AIContext()); + return new ValueTask(context.AIContext); } } @@ -3153,18 +3146,13 @@ public sealed class AzureAIProjectChatClientExtensionsTests { protected override ValueTask> InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default) { - return new ValueTask>(Array.Empty()); + return new ValueTask>(context.RequestMessages); } protected override ValueTask InvokedCoreAsync(InvokedContext context, CancellationToken cancellationToken = default) { return default; } - - public override JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null) - { - return default; - } } } diff --git a/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatHistoryProviderTests.cs b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatHistoryProviderTests.cs index e7276f70b1..12bd467e71 100644 --- a/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatHistoryProviderTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatHistoryProviderTests.cs @@ -3,8 +3,6 @@ using System; using System.Collections.Generic; using System.Linq; -using System.Text.Json; -using System.Text.Json.Serialization.Metadata; using System.Threading.Tasks; using Azure.Core; using Azure.Identity; @@ -42,7 +40,8 @@ namespace Microsoft.Agents.AI.CosmosNoSql.UnitTests; public sealed class CosmosChatHistoryProviderTests : IAsyncLifetime, IDisposable { private static readonly AIAgent s_mockAgent = new Moq.Mock().Object; - private static readonly AgentSession s_mockSession = new Moq.Mock().Object; + + private static AgentSession CreateMockSession() => new Moq.Mock().Object; // Cosmos DB Emulator connection settings private const string EmulatorEndpoint = "https://localhost:8081"; @@ -149,6 +148,35 @@ public sealed class CosmosChatHistoryProviderTests : IAsyncLifetime, IDisposable #region Constructor Tests + [SkippableFact] + [Trait("Category", "CosmosDB")] + public void StateKey_ReturnsDefaultKey_WhenNoStateKeyProvided() + { + // Arrange & Act + this.SkipIfEmulatorNotAvailable(); + + using var provider = new CosmosChatHistoryProvider(this._connectionString, s_testDatabaseId, TestContainerId, + _ => new CosmosChatHistoryProvider.State("test-conversation")); + + // Assert + Assert.Equal("CosmosChatHistoryProvider", provider.StateKey); + } + + [SkippableFact] + [Trait("Category", "CosmosDB")] + public void StateKey_ReturnsCustomKey_WhenSetViaConstructor() + { + // Arrange & Act + this.SkipIfEmulatorNotAvailable(); + + using var provider = new CosmosChatHistoryProvider(this._connectionString, s_testDatabaseId, TestContainerId, + _ => new CosmosChatHistoryProvider.State("test-conversation"), + stateKey: "custom-key"); + + // Assert + Assert.Equal("custom-key", provider.StateKey); + } + [SkippableFact] [Trait("Category", "CosmosDB")] public void Constructor_WithConnectionString_ShouldCreateInstance() @@ -157,28 +185,11 @@ public sealed class CosmosChatHistoryProviderTests : IAsyncLifetime, IDisposable this.SkipIfEmulatorNotAvailable(); // Act - using var provider = new CosmosChatHistoryProvider(this._connectionString, s_testDatabaseId, TestContainerId, "test-conversation"); + using var provider = new CosmosChatHistoryProvider(this._connectionString, s_testDatabaseId, TestContainerId, + _ => new CosmosChatHistoryProvider.State("test-conversation")); // Assert Assert.NotNull(provider); - Assert.Equal("test-conversation", provider.ConversationId); - Assert.Equal(s_testDatabaseId, provider.DatabaseId); - Assert.Equal(TestContainerId, provider.ContainerId); - } - - [SkippableFact] - [Trait("Category", "CosmosDB")] - public void Constructor_WithConnectionStringNoConversationId_ShouldCreateInstance() - { - // Arrange - this.SkipIfEmulatorNotAvailable(); - - // Act - using var provider = new CosmosChatHistoryProvider(this._connectionString, s_testDatabaseId, TestContainerId); - - // Assert - Assert.NotNull(provider); - Assert.NotNull(provider.ConversationId); Assert.Equal(s_testDatabaseId, provider.DatabaseId); Assert.Equal(TestContainerId, provider.ContainerId); } @@ -189,18 +200,19 @@ public sealed class CosmosChatHistoryProviderTests : IAsyncLifetime, IDisposable { // Arrange & Act & Assert Assert.Throws(() => - new CosmosChatHistoryProvider((string)null!, s_testDatabaseId, TestContainerId, "test-conversation")); + new CosmosChatHistoryProvider((string)null!, s_testDatabaseId, TestContainerId, + _ => new CosmosChatHistoryProvider.State("test-conversation"))); } [SkippableFact] [Trait("Category", "CosmosDB")] - public void Constructor_WithEmptyConversationId_ShouldThrowArgumentException() + public void Constructor_WithNullStateInitializer_ShouldThrowArgumentNullException() { // Arrange & Act & Assert this.SkipIfEmulatorNotAvailable(); - Assert.Throws(() => - new CosmosChatHistoryProvider(this._connectionString, s_testDatabaseId, TestContainerId, "")); + Assert.Throws(() => + new CosmosChatHistoryProvider(this._connectionString, s_testDatabaseId, TestContainerId, null!)); } #endregion @@ -213,14 +225,13 @@ public sealed class CosmosChatHistoryProviderTests : IAsyncLifetime, IDisposable { // Arrange this.SkipIfEmulatorNotAvailable(); + var session = CreateMockSession(); var conversationId = Guid.NewGuid().ToString(); - using var provider = new CosmosChatHistoryProvider(this._connectionString, s_testDatabaseId, TestContainerId, conversationId); + using var provider = new CosmosChatHistoryProvider(this._connectionString, s_testDatabaseId, TestContainerId, + _ => new CosmosChatHistoryProvider.State(conversationId)); var message = new ChatMessage(ChatRole.User, "Hello, world!"); - var context = new ChatHistoryProvider.InvokedContext(s_mockAgent, s_mockSession, [message]) - { - ResponseMessages = [] - }; + var context = new ChatHistoryProvider.InvokedContext(s_mockAgent, session, [message], []); // Act await provider.InvokedAsync(context); @@ -229,7 +240,7 @@ public sealed class CosmosChatHistoryProviderTests : IAsyncLifetime, IDisposable await Task.Delay(100); // Assert - var invokingContext = new ChatHistoryProvider.InvokingContext(s_mockAgent, s_mockSession, []); + var invokingContext = new ChatHistoryProvider.InvokingContext(s_mockAgent, session, []); var messages = await provider.InvokingAsync(invokingContext); var messageList = messages.ToList(); @@ -279,8 +290,10 @@ public sealed class CosmosChatHistoryProviderTests : IAsyncLifetime, IDisposable { // Arrange this.SkipIfEmulatorNotAvailable(); + var session = CreateMockSession(); var conversationId = Guid.NewGuid().ToString(); - using var provider = new CosmosChatHistoryProvider(this._connectionString, s_testDatabaseId, TestContainerId, conversationId); + using var provider = new CosmosChatHistoryProvider(this._connectionString, s_testDatabaseId, TestContainerId, + _ => new CosmosChatHistoryProvider.State(conversationId)); var requestMessages = new[] { new ChatMessage(ChatRole.User, "First message"), @@ -293,16 +306,13 @@ public sealed class CosmosChatHistoryProviderTests : IAsyncLifetime, IDisposable new ChatMessage(ChatRole.Assistant, "Response message") }; - var context = new ChatHistoryProvider.InvokedContext(s_mockAgent, s_mockSession, requestMessages) - { - ResponseMessages = responseMessages - }; + var context = new ChatHistoryProvider.InvokedContext(s_mockAgent, session, requestMessages, responseMessages); // Act await provider.InvokedAsync(context); // Assert - var invokingContext = new ChatHistoryProvider.InvokingContext(s_mockAgent, s_mockSession, []); + var invokingContext = new ChatHistoryProvider.InvokingContext(s_mockAgent, session, []); var retrievedMessages = await provider.InvokingAsync(invokingContext); var messageList = retrievedMessages.ToList(); Assert.Equal(5, messageList.Count); @@ -323,10 +333,12 @@ public sealed class CosmosChatHistoryProviderTests : IAsyncLifetime, IDisposable { // Arrange this.SkipIfEmulatorNotAvailable(); - using var provider = new CosmosChatHistoryProvider(this._connectionString, s_testDatabaseId, TestContainerId, Guid.NewGuid().ToString()); + var session = CreateMockSession(); + using var provider = new CosmosChatHistoryProvider(this._connectionString, s_testDatabaseId, TestContainerId, + _ => new CosmosChatHistoryProvider.State(Guid.NewGuid().ToString())); // Act - var invokingContext = new ChatHistoryProvider.InvokingContext(s_mockAgent, s_mockSession, []); + var invokingContext = new ChatHistoryProvider.InvokingContext(s_mockAgent, session, []); var messages = await provider.InvokingAsync(invokingContext); // Assert @@ -339,21 +351,25 @@ public sealed class CosmosChatHistoryProviderTests : IAsyncLifetime, IDisposable { // Arrange this.SkipIfEmulatorNotAvailable(); + var session = CreateMockSession(); var conversation1 = Guid.NewGuid().ToString(); var conversation2 = Guid.NewGuid().ToString(); - using var store1 = new CosmosChatHistoryProvider(this._connectionString, s_testDatabaseId, TestContainerId, conversation1); - using var store2 = new CosmosChatHistoryProvider(this._connectionString, s_testDatabaseId, TestContainerId, conversation2); + // Use different stateKey values so the providers don't overwrite each other's state in the shared session + using var store1 = new CosmosChatHistoryProvider(this._connectionString, s_testDatabaseId, TestContainerId, + _ => new CosmosChatHistoryProvider.State(conversation1), stateKey: "conv1"); + using var store2 = new CosmosChatHistoryProvider(this._connectionString, s_testDatabaseId, TestContainerId, + _ => new CosmosChatHistoryProvider.State(conversation2), stateKey: "conv2"); - var context1 = new ChatHistoryProvider.InvokedContext(s_mockAgent, s_mockSession, [new ChatMessage(ChatRole.User, "Message for conversation 1")]); - var context2 = new ChatHistoryProvider.InvokedContext(s_mockAgent, s_mockSession, [new ChatMessage(ChatRole.User, "Message for conversation 2")]); + var context1 = new ChatHistoryProvider.InvokedContext(s_mockAgent, session, [new ChatMessage(ChatRole.User, "Message for conversation 1")], []); + var context2 = new ChatHistoryProvider.InvokedContext(s_mockAgent, session, [new ChatMessage(ChatRole.User, "Message for conversation 2")], []); await store1.InvokedAsync(context1); await store2.InvokedAsync(context2); // Act - var invokingContext1 = new ChatHistoryProvider.InvokingContext(s_mockAgent, s_mockSession, []); - var invokingContext2 = new ChatHistoryProvider.InvokingContext(s_mockAgent, s_mockSession, []); + var invokingContext1 = new ChatHistoryProvider.InvokingContext(s_mockAgent, session, []); + var invokingContext2 = new ChatHistoryProvider.InvokingContext(s_mockAgent, session, []); var messages1 = await store1.InvokingAsync(invokingContext1); var messages2 = await store2.InvokingAsync(invokingContext2); @@ -365,6 +381,8 @@ public sealed class CosmosChatHistoryProviderTests : IAsyncLifetime, IDisposable Assert.Single(messageList2); Assert.Equal("Message for conversation 1", messageList1[0].Text); Assert.Equal("Message for conversation 2", messageList2[0].Text); + Assert.Equal(AgentRequestMessageSourceType.ChatHistory, messageList1[0].GetAgentRequestMessageSourceType()); + Assert.Equal(AgentRequestMessageSourceType.ChatHistory, messageList2[0].GetAgentRequestMessageSourceType()); } #endregion @@ -377,8 +395,10 @@ public sealed class CosmosChatHistoryProviderTests : IAsyncLifetime, IDisposable { // Arrange this.SkipIfEmulatorNotAvailable(); + var session = CreateMockSession(); var conversationId = $"test-conversation-{Guid.NewGuid():N}"; // Use unique conversation ID - using var originalStore = new CosmosChatHistoryProvider(this._connectionString, s_testDatabaseId, TestContainerId, conversationId); + using var originalStore = new CosmosChatHistoryProvider(this._connectionString, s_testDatabaseId, TestContainerId, + _ => new CosmosChatHistoryProvider.State(conversationId)); var messages = new[] { @@ -390,18 +410,21 @@ public sealed class CosmosChatHistoryProviderTests : IAsyncLifetime, IDisposable }; // Act 1: Add messages - var invokedContext = new ChatHistoryProvider.InvokedContext(s_mockAgent, s_mockSession, messages); + var invokedContext = new ChatHistoryProvider.InvokedContext(s_mockAgent, session, messages, []); await originalStore.InvokedAsync(invokedContext); // Act 2: Verify messages were added - var invokingContext = new ChatHistoryProvider.InvokingContext(s_mockAgent, s_mockSession, []); + var invokingContext = new ChatHistoryProvider.InvokingContext(s_mockAgent, session, []); var retrievedMessages = await originalStore.InvokingAsync(invokingContext); var retrievedList = retrievedMessages.ToList(); Assert.Equal(5, retrievedList.Count); // Act 3: Create new provider instance for same conversation (test persistence) - using var newProvider = new CosmosChatHistoryProvider(this._connectionString, s_testDatabaseId, TestContainerId, conversationId); - var persistedMessages = await newProvider.InvokingAsync(invokingContext); + using var newProvider = new CosmosChatHistoryProvider(this._connectionString, s_testDatabaseId, TestContainerId, + _ => new CosmosChatHistoryProvider.State(conversationId)); + var newSession = CreateMockSession(); + var newInvokingContext = new ChatHistoryProvider.InvokingContext(s_mockAgent, newSession, []); + var persistedMessages = await newProvider.InvokingAsync(newInvokingContext); var persistedList = persistedMessages.ToList(); // Assert final state @@ -423,7 +446,8 @@ public sealed class CosmosChatHistoryProviderTests : IAsyncLifetime, IDisposable { // Arrange this.SkipIfEmulatorNotAvailable(); - var provider = new CosmosChatHistoryProvider(this._connectionString, s_testDatabaseId, TestContainerId, Guid.NewGuid().ToString()); + var provider = new CosmosChatHistoryProvider(this._connectionString, s_testDatabaseId, TestContainerId, + _ => new CosmosChatHistoryProvider.State(Guid.NewGuid().ToString())); // Act & Assert provider.Dispose(); // Should not throw @@ -435,7 +459,8 @@ public sealed class CosmosChatHistoryProviderTests : IAsyncLifetime, IDisposable { // Arrange this.SkipIfEmulatorNotAvailable(); - var provider = new CosmosChatHistoryProvider(this._connectionString, s_testDatabaseId, TestContainerId, Guid.NewGuid().ToString()); + var provider = new CosmosChatHistoryProvider(this._connectionString, s_testDatabaseId, TestContainerId, + _ => new CosmosChatHistoryProvider.State(Guid.NewGuid().ToString())); // Act & Assert provider.Dispose(); // First call @@ -454,11 +479,11 @@ public sealed class CosmosChatHistoryProviderTests : IAsyncLifetime, IDisposable this.SkipIfEmulatorNotAvailable(); // Act - using var provider = new CosmosChatHistoryProvider(this._connectionString, s_testDatabaseId, HierarchicalTestContainerId, "tenant-123", "user-456", "session-789"); + using var provider = new CosmosChatHistoryProvider(this._connectionString, s_testDatabaseId, HierarchicalTestContainerId, + _ => new CosmosChatHistoryProvider.State("session-789", "tenant-123", "user-456")); // Assert Assert.NotNull(provider); - Assert.Equal("session-789", provider.ConversationId); Assert.Equal(s_testDatabaseId, provider.DatabaseId); Assert.Equal(HierarchicalTestContainerId, provider.ContainerId); } @@ -472,11 +497,11 @@ public sealed class CosmosChatHistoryProviderTests : IAsyncLifetime, IDisposable // Act TokenCredential credential = new DefaultAzureCredential(); - using var provider = new CosmosChatHistoryProvider(EmulatorEndpoint, credential, s_testDatabaseId, HierarchicalTestContainerId, "tenant-123", "user-456", "session-789"); + using var provider = new CosmosChatHistoryProvider(EmulatorEndpoint, credential, s_testDatabaseId, HierarchicalTestContainerId, + _ => new CosmosChatHistoryProvider.State("session-789", "tenant-123", "user-456")); // Assert Assert.NotNull(provider); - Assert.Equal("session-789", provider.ConversationId); Assert.Equal(s_testDatabaseId, provider.DatabaseId); Assert.Equal(HierarchicalTestContainerId, provider.ContainerId); } @@ -489,46 +514,31 @@ public sealed class CosmosChatHistoryProviderTests : IAsyncLifetime, IDisposable this.SkipIfEmulatorNotAvailable(); using var cosmosClient = new CosmosClient(EmulatorEndpoint, EmulatorKey); - using var provider = new CosmosChatHistoryProvider(cosmosClient, s_testDatabaseId, HierarchicalTestContainerId, "tenant-123", "user-456", "session-789"); + using var provider = new CosmosChatHistoryProvider(cosmosClient, s_testDatabaseId, HierarchicalTestContainerId, + _ => new CosmosChatHistoryProvider.State("session-789", "tenant-123", "user-456")); // Assert Assert.NotNull(provider); - Assert.Equal("session-789", provider.ConversationId); Assert.Equal(s_testDatabaseId, provider.DatabaseId); Assert.Equal(HierarchicalTestContainerId, provider.ContainerId); } [SkippableFact] [Trait("Category", "CosmosDB")] - public void Constructor_WithHierarchicalNullTenantId_ShouldThrowArgumentException() + public void State_WithEmptyConversationId_ShouldThrowArgumentException() { // Arrange & Act & Assert - this.SkipIfEmulatorNotAvailable(); - - Assert.Throws(() => - new CosmosChatHistoryProvider(this._connectionString, s_testDatabaseId, TestContainerId, null!, "user-456", "session-789")); + Assert.Throws(() => + new CosmosChatHistoryProvider.State("")); } [SkippableFact] [Trait("Category", "CosmosDB")] - public void Constructor_WithHierarchicalEmptyUserId_ShouldThrowArgumentException() + public void State_WithWhitespaceConversationId_ShouldThrowArgumentException() { // Arrange & Act & Assert - this.SkipIfEmulatorNotAvailable(); - Assert.Throws(() => - new CosmosChatHistoryProvider(this._connectionString, s_testDatabaseId, HierarchicalTestContainerId, "tenant-123", "", "session-789")); - } - - [SkippableFact] - [Trait("Category", "CosmosDB")] - public void Constructor_WithHierarchicalWhitespaceSessionId_ShouldThrowArgumentException() - { - // Arrange & Act & Assert - this.SkipIfEmulatorNotAvailable(); - - Assert.Throws(() => - new CosmosChatHistoryProvider(this._connectionString, s_testDatabaseId, HierarchicalTestContainerId, "tenant-123", "user-456", " ")); + new CosmosChatHistoryProvider.State(" ")); } [SkippableFact] @@ -537,14 +547,16 @@ public sealed class CosmosChatHistoryProviderTests : IAsyncLifetime, IDisposable { // Arrange this.SkipIfEmulatorNotAvailable(); + var session = CreateMockSession(); const string TenantId = "tenant-123"; const string UserId = "user-456"; const string SessionId = "session-789"; // Test hierarchical partitioning constructor with connection string - using var provider = new CosmosChatHistoryProvider(this._connectionString, s_testDatabaseId, HierarchicalTestContainerId, TenantId, UserId, SessionId); + using var provider = new CosmosChatHistoryProvider(this._connectionString, s_testDatabaseId, HierarchicalTestContainerId, + _ => new CosmosChatHistoryProvider.State(SessionId, TenantId, UserId)); var message = new ChatMessage(ChatRole.User, "Hello from hierarchical partitioning!"); - var context = new ChatHistoryProvider.InvokedContext(s_mockAgent, s_mockSession, [message]); + var context = new ChatHistoryProvider.InvokedContext(s_mockAgent, session, [message], []); // Act await provider.InvokedAsync(context); @@ -553,7 +565,7 @@ public sealed class CosmosChatHistoryProviderTests : IAsyncLifetime, IDisposable await Task.Delay(100); // Assert - var invokingContext = new ChatHistoryProvider.InvokingContext(s_mockAgent, s_mockSession, []); + var invokingContext = new ChatHistoryProvider.InvokingContext(s_mockAgent, session, []); var messages = await provider.InvokingAsync(invokingContext); var messageList = messages.ToList(); @@ -589,11 +601,13 @@ public sealed class CosmosChatHistoryProviderTests : IAsyncLifetime, IDisposable { // Arrange this.SkipIfEmulatorNotAvailable(); + var session = CreateMockSession(); const string TenantId = "tenant-batch"; const string UserId = "user-batch"; const string SessionId = "session-batch"; // Test hierarchical partitioning constructor with connection string - using var provider = new CosmosChatHistoryProvider(this._connectionString, s_testDatabaseId, HierarchicalTestContainerId, TenantId, UserId, SessionId); + using var provider = new CosmosChatHistoryProvider(this._connectionString, s_testDatabaseId, HierarchicalTestContainerId, + _ => new CosmosChatHistoryProvider.State(SessionId, TenantId, UserId)); var messages = new[] { new ChatMessage(ChatRole.User, "First hierarchical message"), @@ -601,7 +615,7 @@ public sealed class CosmosChatHistoryProviderTests : IAsyncLifetime, IDisposable new ChatMessage(ChatRole.User, "Third hierarchical message") }; - var context = new ChatHistoryProvider.InvokedContext(s_mockAgent, s_mockSession, messages); + var context = new ChatHistoryProvider.InvokedContext(s_mockAgent, session, messages, []); // Act await provider.InvokedAsync(context); @@ -610,7 +624,7 @@ public sealed class CosmosChatHistoryProviderTests : IAsyncLifetime, IDisposable await Task.Delay(100); // Assert - var invokingContext = new ChatHistoryProvider.InvokingContext(s_mockAgent, s_mockSession, []); + var invokingContext = new ChatHistoryProvider.InvokingContext(s_mockAgent, session, []); var retrievedMessages = await provider.InvokingAsync(invokingContext); var messageList = retrievedMessages.ToList(); @@ -626,18 +640,22 @@ public sealed class CosmosChatHistoryProviderTests : IAsyncLifetime, IDisposable { // Arrange this.SkipIfEmulatorNotAvailable(); + var session = CreateMockSession(); const string TenantId = "tenant-isolation"; const string UserId1 = "user-1"; const string UserId2 = "user-2"; const string SessionId = "session-isolation"; // Different userIds create different hierarchical partitions, providing proper isolation - using var store1 = new CosmosChatHistoryProvider(this._connectionString, s_testDatabaseId, HierarchicalTestContainerId, TenantId, UserId1, SessionId); - using var store2 = new CosmosChatHistoryProvider(this._connectionString, s_testDatabaseId, HierarchicalTestContainerId, TenantId, UserId2, SessionId); + // Use different stateKey values so the providers don't overwrite each other's state in the shared session + using var store1 = new CosmosChatHistoryProvider(this._connectionString, s_testDatabaseId, HierarchicalTestContainerId, + _ => new CosmosChatHistoryProvider.State(SessionId, TenantId, UserId1), stateKey: "user1"); + using var store2 = new CosmosChatHistoryProvider(this._connectionString, s_testDatabaseId, HierarchicalTestContainerId, + _ => new CosmosChatHistoryProvider.State(SessionId, TenantId, UserId2), stateKey: "user2"); // Add messages to both stores - var context1 = new ChatHistoryProvider.InvokedContext(s_mockAgent, s_mockSession, [new ChatMessage(ChatRole.User, "Message from user 1")]); - var context2 = new ChatHistoryProvider.InvokedContext(s_mockAgent, s_mockSession, [new ChatMessage(ChatRole.User, "Message from user 2")]); + var context1 = new ChatHistoryProvider.InvokedContext(s_mockAgent, session, [new ChatMessage(ChatRole.User, "Message from user 1")], []); + var context2 = new ChatHistoryProvider.InvokedContext(s_mockAgent, session, [new ChatMessage(ChatRole.User, "Message from user 2")], []); await store1.InvokedAsync(context1); await store2.InvokedAsync(context2); @@ -646,8 +664,8 @@ public sealed class CosmosChatHistoryProviderTests : IAsyncLifetime, IDisposable await Task.Delay(100); // Act & Assert - var invokingContext1 = new ChatHistoryProvider.InvokingContext(s_mockAgent, s_mockSession, []); - var invokingContext2 = new ChatHistoryProvider.InvokingContext(s_mockAgent, s_mockSession, []); + var invokingContext1 = new ChatHistoryProvider.InvokingContext(s_mockAgent, session, []); + var invokingContext2 = new ChatHistoryProvider.InvokingContext(s_mockAgent, session, []); var messages1 = await store1.InvokingAsync(invokingContext1); var messageList1 = messages1.ToList(); @@ -664,43 +682,37 @@ public sealed class CosmosChatHistoryProviderTests : IAsyncLifetime, IDisposable [SkippableFact] [Trait("Category", "CosmosDB")] - public async Task SerializeDeserialize_WithHierarchicalPartitioning_ShouldPreserveStateAsync() + public async Task StateBag_WithHierarchicalPartitioning_ShouldPreserveStateAcrossProviderInstancesAsync() { // Arrange this.SkipIfEmulatorNotAvailable(); + var session = CreateMockSession(); const string TenantId = "tenant-serialize"; const string UserId = "user-serialize"; const string SessionId = "session-serialize"; - using var originalStore = new CosmosChatHistoryProvider(this._connectionString, s_testDatabaseId, HierarchicalTestContainerId, TenantId, UserId, SessionId); + using var originalStore = new CosmosChatHistoryProvider(this._connectionString, s_testDatabaseId, HierarchicalTestContainerId, + _ => new CosmosChatHistoryProvider.State(SessionId, TenantId, UserId)); - var context = new ChatHistoryProvider.InvokedContext(s_mockAgent, s_mockSession, [new ChatMessage(ChatRole.User, "Test serialization message")]); + var context = new ChatHistoryProvider.InvokedContext(s_mockAgent, session, [new ChatMessage(ChatRole.User, "Test serialization message")], []); await originalStore.InvokedAsync(context); - // Act - Serialize the provider state - var serializedState = originalStore.Serialize(); - - // Create a new provider from the serialized state - using var cosmosClient = new CosmosClient(EmulatorEndpoint, EmulatorKey); - var serializerOptions = new JsonSerializerOptions - { - TypeInfoResolver = new DefaultJsonTypeInfoResolver() - }; - using var deserializedStore = CosmosChatHistoryProvider.CreateFromSerializedState(cosmosClient, serializedState, s_testDatabaseId, HierarchicalTestContainerId, serializerOptions); - // Wait a moment for eventual consistency await Task.Delay(100); - // Assert - The deserialized provider should have the same functionality - var invokingContext = new ChatHistoryProvider.InvokingContext(s_mockAgent, s_mockSession, []); - var messages = await deserializedStore.InvokingAsync(invokingContext); + // Act - Create a new provider that uses a different intializer, but we will use the same session. + using var newStore = new CosmosChatHistoryProvider(this._connectionString, s_testDatabaseId, HierarchicalTestContainerId, + _ => new CosmosChatHistoryProvider.State(Guid.NewGuid().ToString())); + + // Assert - The new provider should read the same messages from Cosmos DB + var invokingContext = new ChatHistoryProvider.InvokingContext(s_mockAgent, session, []); + var messages = await newStore.InvokingAsync(invokingContext); var messageList = messages.ToList(); Assert.Single(messageList); Assert.Equal("Test serialization message", messageList[0].Text); - Assert.Equal(SessionId, deserializedStore.ConversationId); - Assert.Equal(s_testDatabaseId, deserializedStore.DatabaseId); - Assert.Equal(HierarchicalTestContainerId, deserializedStore.ContainerId); + Assert.Equal(s_testDatabaseId, newStore.DatabaseId); + Assert.Equal(HierarchicalTestContainerId, newStore.ContainerId); } [SkippableFact] @@ -711,13 +723,17 @@ public sealed class CosmosChatHistoryProviderTests : IAsyncLifetime, IDisposable this.SkipIfEmulatorNotAvailable(); const string SessionId = "coexist-session"; + var session = CreateMockSession(); // Create simple provider using simple partitioning container and hierarchical provider using hierarchical container - using var simpleProvider = new CosmosChatHistoryProvider(this._connectionString, s_testDatabaseId, TestContainerId, SessionId); - using var hierarchicalProvider = new CosmosChatHistoryProvider(this._connectionString, s_testDatabaseId, HierarchicalTestContainerId, "tenant-coexist", "user-coexist", SessionId); + // Use different stateKey values so the providers don't overwrite each other's state in the shared session + using var simpleProvider = new CosmosChatHistoryProvider(this._connectionString, s_testDatabaseId, TestContainerId, + _ => new CosmosChatHistoryProvider.State(SessionId), stateKey: "simple"); + using var hierarchicalProvider = new CosmosChatHistoryProvider(this._connectionString, s_testDatabaseId, HierarchicalTestContainerId, + _ => new CosmosChatHistoryProvider.State(SessionId, "tenant-coexist", "user-coexist"), stateKey: "hierarchical"); // Add messages to both - var simpleContext = new ChatHistoryProvider.InvokedContext(s_mockAgent, s_mockSession, [new ChatMessage(ChatRole.User, "Simple partitioning message")]); - var hierarchicalContext = new ChatHistoryProvider.InvokedContext(s_mockAgent, s_mockSession, [new ChatMessage(ChatRole.User, "Hierarchical partitioning message")]); + var simpleContext = new ChatHistoryProvider.InvokedContext(s_mockAgent, session, [new ChatMessage(ChatRole.User, "Simple partitioning message")], []); + var hierarchicalContext = new ChatHistoryProvider.InvokedContext(s_mockAgent, session, [new ChatMessage(ChatRole.User, "Hierarchical partitioning message")], []); await simpleProvider.InvokedAsync(simpleContext); await hierarchicalProvider.InvokedAsync(hierarchicalContext); @@ -726,7 +742,7 @@ public sealed class CosmosChatHistoryProviderTests : IAsyncLifetime, IDisposable await Task.Delay(100); // Act & Assert - var invokingContext = new ChatHistoryProvider.InvokingContext(s_mockAgent, s_mockSession, []); + var invokingContext = new ChatHistoryProvider.InvokingContext(s_mockAgent, session, []); var simpleMessages = await simpleProvider.InvokingAsync(invokingContext); var simpleMessageList = simpleMessages.ToList(); @@ -747,9 +763,11 @@ public sealed class CosmosChatHistoryProviderTests : IAsyncLifetime, IDisposable { // Arrange this.SkipIfEmulatorNotAvailable(); + var session = CreateMockSession(); const string ConversationId = "max-messages-test"; - using var provider = new CosmosChatHistoryProvider(this._connectionString, s_testDatabaseId, TestContainerId, ConversationId); + using var provider = new CosmosChatHistoryProvider(this._connectionString, s_testDatabaseId, TestContainerId, + _ => new CosmosChatHistoryProvider.State(ConversationId)); // Add 10 messages var messages = new List(); @@ -759,7 +777,7 @@ public sealed class CosmosChatHistoryProviderTests : IAsyncLifetime, IDisposable await Task.Delay(10); // Small delay to ensure different timestamps } - var context = new ChatHistoryProvider.InvokedContext(s_mockAgent, s_mockSession, messages); + var context = new ChatHistoryProvider.InvokedContext(s_mockAgent, session, messages, []); await provider.InvokedAsync(context); // Wait for eventual consistency @@ -767,7 +785,7 @@ public sealed class CosmosChatHistoryProviderTests : IAsyncLifetime, IDisposable // Act - Set max to 5 and retrieve provider.MaxMessagesToRetrieve = 5; - var invokingContext = new ChatHistoryProvider.InvokingContext(s_mockAgent, s_mockSession, []); + var invokingContext = new ChatHistoryProvider.InvokingContext(s_mockAgent, session, []); var retrievedMessages = await provider.InvokingAsync(invokingContext); var messageList = retrievedMessages.ToList(); @@ -786,9 +804,11 @@ public sealed class CosmosChatHistoryProviderTests : IAsyncLifetime, IDisposable { // Arrange this.SkipIfEmulatorNotAvailable(); + var session = CreateMockSession(); const string ConversationId = "max-messages-null-test"; - using var provider = new CosmosChatHistoryProvider(this._connectionString, s_testDatabaseId, TestContainerId, ConversationId); + using var provider = new CosmosChatHistoryProvider(this._connectionString, s_testDatabaseId, TestContainerId, + _ => new CosmosChatHistoryProvider.State(ConversationId)); // Add 10 messages var messages = new List(); @@ -797,14 +817,14 @@ public sealed class CosmosChatHistoryProviderTests : IAsyncLifetime, IDisposable messages.Add(new ChatMessage(ChatRole.User, $"Message {i}")); } - var context = new ChatHistoryProvider.InvokedContext(s_mockAgent, s_mockSession, messages); + var context = new ChatHistoryProvider.InvokedContext(s_mockAgent, session, messages, []); await provider.InvokedAsync(context); // Wait for eventual consistency await Task.Delay(100); // Act - No limit set (default null) - var invokingContext = new ChatHistoryProvider.InvokingContext(s_mockAgent, s_mockSession, []); + var invokingContext = new ChatHistoryProvider.InvokingContext(s_mockAgent, session, []); var retrievedMessages = await provider.InvokingAsync(invokingContext); var messageList = retrievedMessages.ToList(); @@ -815,4 +835,119 @@ public sealed class CosmosChatHistoryProviderTests : IAsyncLifetime, IDisposable } #endregion + + #region Message Filter Tests + + [SkippableFact] + [Trait("Category", "CosmosDB")] + public async Task InvokedAsync_DefaultFilter_ExcludesChatHistoryMessagesFromStorageAsync() + { + // Arrange + this.SkipIfEmulatorNotAvailable(); + var session = CreateMockSession(); + var conversationId = Guid.NewGuid().ToString(); + using var provider = new CosmosChatHistoryProvider(this._connectionString, s_testDatabaseId, TestContainerId, + _ => new CosmosChatHistoryProvider.State(conversationId)); + + var requestMessages = new[] + { + new ChatMessage(ChatRole.User, "External message"), + new ChatMessage(ChatRole.System, "From history") { AdditionalProperties = new() { { AgentRequestMessageSourceAttribution.AdditionalPropertiesKey, new AgentRequestMessageSourceAttribution(AgentRequestMessageSourceType.ChatHistory, "HistorySource") } } }, + new ChatMessage(ChatRole.System, "From context provider") { AdditionalProperties = new() { { AgentRequestMessageSourceAttribution.AdditionalPropertiesKey, new AgentRequestMessageSourceAttribution(AgentRequestMessageSourceType.AIContextProvider, "ContextSource") } } }, + }; + + var context = new ChatHistoryProvider.InvokedContext(s_mockAgent, session, requestMessages, [new ChatMessage(ChatRole.Assistant, "Response")]); + + // Act + await provider.InvokedAsync(context); + + // Wait for eventual consistency + await Task.Delay(100); + + // Assert - ChatHistory message excluded, External + AIContextProvider + Response stored + var invokingContext = new ChatHistoryProvider.InvokingContext(s_mockAgent, session, []); + var messages = (await provider.InvokingAsync(invokingContext)).ToList(); + Assert.Equal(3, messages.Count); + Assert.Equal("External message", messages[0].Text); + Assert.Equal("From context provider", messages[1].Text); + Assert.Equal("Response", messages[2].Text); + } + + [SkippableFact] + [Trait("Category", "CosmosDB")] + public async Task InvokedAsync_CustomStorageInputFilter_OverridesDefaultAsync() + { + // Arrange + this.SkipIfEmulatorNotAvailable(); + var session = CreateMockSession(); + var conversationId = Guid.NewGuid().ToString(); + using var provider = new CosmosChatHistoryProvider(this._connectionString, s_testDatabaseId, TestContainerId, + _ => new CosmosChatHistoryProvider.State(conversationId)) + { + // Custom filter: only store External messages (also exclude AIContextProvider) + StorageInputMessageFilter = messages => messages.Where(m => m.GetAgentRequestMessageSourceType() == AgentRequestMessageSourceType.External) + }; + + var requestMessages = new[] + { + new ChatMessage(ChatRole.User, "External message"), + new ChatMessage(ChatRole.System, "From history") { AdditionalProperties = new() { { AgentRequestMessageSourceAttribution.AdditionalPropertiesKey, new AgentRequestMessageSourceAttribution(AgentRequestMessageSourceType.ChatHistory, "HistorySource") } } }, + new ChatMessage(ChatRole.System, "From context provider") { AdditionalProperties = new() { { AgentRequestMessageSourceAttribution.AdditionalPropertiesKey, new AgentRequestMessageSourceAttribution(AgentRequestMessageSourceType.AIContextProvider, "ContextSource") } } }, + }; + + var context = new ChatHistoryProvider.InvokedContext(s_mockAgent, session, requestMessages, [new ChatMessage(ChatRole.Assistant, "Response")]); + + // Act + await provider.InvokedAsync(context); + + // Wait for eventual consistency + await Task.Delay(100); + + // Assert - Custom filter: only External + Response stored (both ChatHistory and AIContextProvider excluded) + var invokingContext = new ChatHistoryProvider.InvokingContext(s_mockAgent, session, []); + var messages = (await provider.InvokingAsync(invokingContext)).ToList(); + Assert.Equal(2, messages.Count); + Assert.Equal("External message", messages[0].Text); + Assert.Equal("Response", messages[1].Text); + } + + [SkippableFact] + [Trait("Category", "CosmosDB")] + public async Task InvokingAsync_RetrievalOutputFilter_FiltersRetrievedMessagesAsync() + { + // Arrange + this.SkipIfEmulatorNotAvailable(); + var session = CreateMockSession(); + var conversationId = Guid.NewGuid().ToString(); + using var provider = new CosmosChatHistoryProvider(this._connectionString, s_testDatabaseId, TestContainerId, + _ => new CosmosChatHistoryProvider.State(conversationId)) + { + // Only return User messages when retrieving + RetrievalOutputMessageFilter = messages => messages.Where(m => m.Role == ChatRole.User) + }; + + var requestMessages = new[] + { + new ChatMessage(ChatRole.User, "User message"), + new ChatMessage(ChatRole.System, "System message"), + }; + + var context = new ChatHistoryProvider.InvokedContext(s_mockAgent, session, requestMessages, [new ChatMessage(ChatRole.Assistant, "Assistant response")]); + + await provider.InvokedAsync(context); + + // Wait for eventual consistency + await Task.Delay(100); + + // Act + var invokingContext = new ChatHistoryProvider.InvokingContext(s_mockAgent, session, []); + var messages = (await provider.InvokingAsync(invokingContext)).ToList(); + + // Assert - Only User messages returned (System and Assistant filtered by RetrievalOutputMessageFilter) + Assert.Single(messages); + Assert.Equal("User message", messages[0].Text); + Assert.Equal(ChatRole.User, messages[0].Role); + } + + #endregion } diff --git a/dotnet/tests/Microsoft.Agents.AI.DurableTask.UnitTests/DurableAgentSessionTests.cs b/dotnet/tests/Microsoft.Agents.AI.DurableTask.UnitTests/DurableAgentSessionTests.cs index 4bf8ebc718..bc06c35ab8 100644 --- a/dotnet/tests/Microsoft.Agents.AI.DurableTask.UnitTests/DurableAgentSessionTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.DurableTask.UnitTests/DurableAgentSessionTests.cs @@ -15,7 +15,7 @@ public sealed class DurableAgentSessionTests JsonElement serializedSession = session.Serialize(); // Expected format: "{\"sessionId\":\"@dafx-test-agent@\"}" - string expectedSerializedSession = $"{{\"sessionId\":\"@dafx-{sessionId.Name}@{sessionId.Key}\"}}"; + string expectedSerializedSession = $"{{\"sessionId\":\"@dafx-{sessionId.Name}@{sessionId.Key}\",\"stateBag\":{{}}}}"; Assert.Equal(expectedSerializedSession, serializedSession.ToString()); DurableAgentSession deserializedSession = DurableAgentSession.Deserialize(serializedSession); @@ -33,11 +33,47 @@ public sealed class DurableAgentSessionTests string serializedSession = JsonSerializer.Serialize(session, typeof(DurableAgentSession)); // Expected format: "{\"sessionId\":\"@dafx-test-agent@\"}" - string expectedSerializedSession = $"{{\"sessionId\":\"@dafx-{sessionId.Name}@{sessionId.Key}\"}}"; + string expectedSerializedSession = $"{{\"sessionId\":\"@dafx-{sessionId.Name}@{sessionId.Key}\",\"stateBag\":{{}}}}"; Assert.Equal(expectedSerializedSession, serializedSession); DurableAgentSession? deserializedSession = JsonSerializer.Deserialize(serializedSession); Assert.NotNull(deserializedSession); Assert.Equal(sessionId, deserializedSession.SessionId); } + + [Fact] + public void BuiltInSerialization_RoundTrip_PreservesStateBag() + { + // Arrange + AgentSessionId sessionId = AgentSessionId.WithRandomKey("test-agent"); + DurableAgentSession session = new(sessionId); + session.StateBag.SetValue("durableKey", "durableValue"); + + // Act + JsonElement serializedSession = session.Serialize(); + DurableAgentSession deserializedSession = DurableAgentSession.Deserialize(serializedSession); + + // Assert + Assert.Equal(sessionId, deserializedSession.SessionId); + Assert.True(deserializedSession.StateBag.TryGetValue("durableKey", out var value)); + Assert.Equal("durableValue", value); + } + + [Fact] + public void STJSerialization_RoundTrip_PreservesStateBag() + { + // Arrange + AgentSessionId sessionId = AgentSessionId.WithRandomKey("test-agent"); + DurableAgentSession session = new(sessionId); + session.StateBag.SetValue("stjKey", "stjValue"); + + // Act + string serializedSession = JsonSerializer.Serialize(session, typeof(DurableAgentSession)); + DurableAgentSession? deserializedSession = JsonSerializer.Deserialize(serializedSession); + + // Assert + Assert.NotNull(deserializedSession); + Assert.True(deserializedSession.StateBag.TryGetValue("stjKey", out var value)); + Assert.Equal("stjValue", value); + } } diff --git a/dotnet/tests/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore.IntegrationTests/BasicStreamingTests.cs b/dotnet/tests/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore.IntegrationTests/BasicStreamingTests.cs index 03dfe63d99..1be1eddef1 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore.IntegrationTests/BasicStreamingTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore.IntegrationTests/BasicStreamingTests.cs @@ -7,6 +7,7 @@ using System.Linq; using System.Net.Http; using System.Runtime.CompilerServices; using System.Text.Json; +using System.Text.Json.Serialization; using System.Threading; using System.Threading.Tasks; using FluentAssertions; @@ -281,10 +282,10 @@ internal sealed class FakeChatClientAgent : AIAgent public override string? Description => "A fake agent for testing"; protected override ValueTask CreateSessionCoreAsync(CancellationToken cancellationToken = default) => - new(new FakeInMemoryAgentSession()); + new(new FakeAgentSession()); protected override ValueTask DeserializeSessionCoreAsync(JsonElement serializedState, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default) => - new(new FakeInMemoryAgentSession(serializedState, jsonSerializerOptions)); + new(serializedState.Deserialize(jsonSerializerOptions)!); protected override ValueTask SerializeSessionCoreAsync(AgentSession session, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default) => throw new NotImplementedException(); @@ -326,15 +327,14 @@ internal sealed class FakeChatClientAgent : AIAgent } } - private sealed class FakeInMemoryAgentSession : InMemoryAgentSession + private sealed class FakeAgentSession : AgentSession { - public FakeInMemoryAgentSession() - : base() + public FakeAgentSession() { } - public FakeInMemoryAgentSession(JsonElement serializedSession, JsonSerializerOptions? jsonSerializerOptions = null) - : base(serializedSession, jsonSerializerOptions) + [JsonConstructor] + public FakeAgentSession(AgentSessionStateBag stateBag) : base(stateBag) { } } @@ -348,19 +348,19 @@ internal sealed class FakeMultiMessageAgent : AIAgent public override string? Description => "A fake agent that sends multiple messages for testing"; protected override ValueTask CreateSessionCoreAsync(CancellationToken cancellationToken = default) => - new(new FakeInMemoryAgentSession()); + new(new FakeAgentSession()); protected override ValueTask DeserializeSessionCoreAsync(JsonElement serializedState, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default) => - new(new FakeInMemoryAgentSession(serializedState, jsonSerializerOptions)); + new(serializedState.Deserialize(jsonSerializerOptions)!); protected override ValueTask SerializeSessionCoreAsync(AgentSession session, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default) { - if (session is not FakeInMemoryAgentSession fakeSession) + if (session is not FakeAgentSession fakeSession) { throw new InvalidOperationException("The provided session is not compatible with the agent. Only sessions created by the agent can be serialized."); } - return new(fakeSession.Serialize(jsonSerializerOptions)); + return new(JsonSerializer.SerializeToElement(fakeSession, jsonSerializerOptions)); } protected override async Task RunCoreAsync( @@ -427,19 +427,16 @@ internal sealed class FakeMultiMessageAgent : AIAgent } } - private sealed class FakeInMemoryAgentSession : InMemoryAgentSession + private sealed class FakeAgentSession : AgentSession { - public FakeInMemoryAgentSession() - : base() + public FakeAgentSession() { } - public FakeInMemoryAgentSession(JsonElement serializedSession, JsonSerializerOptions? jsonSerializerOptions = null) - : base(serializedSession, jsonSerializerOptions) + [JsonConstructor] + public FakeAgentSession(AgentSessionStateBag stateBag) : base(stateBag) { } - internal new JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null) - => base.Serialize(jsonSerializerOptions); } public override object? GetService(Type serviceType, object? serviceKey = null) => null; diff --git a/dotnet/tests/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore.IntegrationTests/ForwardedPropertiesTests.cs b/dotnet/tests/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore.IntegrationTests/ForwardedPropertiesTests.cs index 67108676ac..844900372b 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore.IntegrationTests/ForwardedPropertiesTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore.IntegrationTests/ForwardedPropertiesTests.cs @@ -9,6 +9,7 @@ using System.Net.ServerSentEvents; using System.Runtime.CompilerServices; using System.Text; using System.Text.Json; +using System.Text.Json.Serialization; using System.Threading; using System.Threading.Tasks; using FluentAssertions; @@ -335,35 +336,31 @@ internal sealed class FakeForwardedPropsAgent : AIAgent } protected override ValueTask CreateSessionCoreAsync(CancellationToken cancellationToken = default) => - new(new FakeInMemoryAgentSession()); + new(new FakeAgentSession()); protected override ValueTask DeserializeSessionCoreAsync(JsonElement serializedState, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default) => - new(new FakeInMemoryAgentSession(serializedState, jsonSerializerOptions)); + new(serializedState.Deserialize(jsonSerializerOptions)!); protected override ValueTask SerializeSessionCoreAsync(AgentSession session, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default) { - if (session is not FakeInMemoryAgentSession fakeSession) + if (session is not FakeAgentSession fakeSession) { throw new InvalidOperationException("The provided session is not compatible with the agent. Only sessions created by the agent can be serialized."); } - return new(fakeSession.Serialize(jsonSerializerOptions)); + return new(JsonSerializer.SerializeToElement(fakeSession, jsonSerializerOptions)); } - private sealed class FakeInMemoryAgentSession : InMemoryAgentSession + private sealed class FakeAgentSession : AgentSession { - public FakeInMemoryAgentSession() - : base() + public FakeAgentSession() { } - public FakeInMemoryAgentSession(JsonElement serializedSession, JsonSerializerOptions? jsonSerializerOptions = null) - : base(serializedSession, jsonSerializerOptions) + [JsonConstructor] + public FakeAgentSession(AgentSessionStateBag stateBag) : base(stateBag) { } - - internal new JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null) - => base.Serialize(jsonSerializerOptions); } public override object? GetService(Type serviceType, object? serviceKey = null) => null; diff --git a/dotnet/tests/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore.IntegrationTests/SharedStateTests.cs b/dotnet/tests/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore.IntegrationTests/SharedStateTests.cs index 9ff3dde1a4..78441254d9 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore.IntegrationTests/SharedStateTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore.IntegrationTests/SharedStateTests.cs @@ -7,6 +7,7 @@ using System.Linq; using System.Net.Http; using System.Runtime.CompilerServices; using System.Text.Json; +using System.Text.Json.Serialization; using System.Threading; using System.Threading.Tasks; using FluentAssertions; @@ -418,35 +419,31 @@ internal sealed class FakeStateAgent : AIAgent } protected override ValueTask CreateSessionCoreAsync(CancellationToken cancellationToken = default) => - new(new FakeInMemoryAgentSession()); + new(new FakeAgentSession()); protected override ValueTask DeserializeSessionCoreAsync(JsonElement serializedState, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default) => - new(new FakeInMemoryAgentSession(serializedState, jsonSerializerOptions)); + new(serializedState.Deserialize(jsonSerializerOptions)!); protected override ValueTask SerializeSessionCoreAsync(AgentSession session, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default) { - if (session is not FakeInMemoryAgentSession fakeSession) + if (session is not FakeAgentSession fakeSession) { throw new InvalidOperationException("The provided session is not compatible with the agent. Only sessions created by the agent can be serialized."); } - return new(fakeSession.Serialize(jsonSerializerOptions)); + return new(JsonSerializer.SerializeToElement(fakeSession, jsonSerializerOptions)); } - private sealed class FakeInMemoryAgentSession : InMemoryAgentSession + private sealed class FakeAgentSession : AgentSession { - public FakeInMemoryAgentSession() - : base() + public FakeAgentSession() { } - public FakeInMemoryAgentSession(JsonElement serializedSession, JsonSerializerOptions? jsonSerializerOptions = null) - : base(serializedSession, jsonSerializerOptions) + [JsonConstructor] + public FakeAgentSession(AgentSessionStateBag stateBag) : base(stateBag) { } - - internal new JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null) - => base.Serialize(jsonSerializerOptions); } public override object? GetService(Type serviceType, object? serviceKey = null) => null; diff --git a/dotnet/tests/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore.UnitTests/AGUIEndpointRouteBuilderExtensionsTests.cs b/dotnet/tests/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore.UnitTests/AGUIEndpointRouteBuilderExtensionsTests.cs index 16efbd0e6e..383a7d1062 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore.UnitTests/AGUIEndpointRouteBuilderExtensionsTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore.UnitTests/AGUIEndpointRouteBuilderExtensionsTests.cs @@ -6,6 +6,7 @@ using System.IO; using System.Linq; using System.Text; using System.Text.Json; +using System.Text.Json.Serialization; using System.Threading; using System.Threading.Tasks; using Microsoft.Agents.AI.Hosting.AGUI.AspNetCore.Shared; @@ -426,19 +427,19 @@ public sealed class AGUIEndpointRouteBuilderExtensionsTests public override string? Description => "Agent that produces multiple text chunks"; protected override ValueTask CreateSessionCoreAsync(CancellationToken cancellationToken = default) => - new(new TestInMemoryAgentSession()); + new(new TestAgentSession()); protected override ValueTask DeserializeSessionCoreAsync(JsonElement serializedState, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default) => - new(new TestInMemoryAgentSession(serializedState, jsonSerializerOptions)); + new(serializedState.Deserialize(jsonSerializerOptions)!); protected override ValueTask SerializeSessionCoreAsync(AgentSession session, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default) { - if (session is not TestInMemoryAgentSession testSession) + if (session is not TestAgentSession testSession) { throw new InvalidOperationException("The provided session is not compatible with the agent. Only sessions created by the agent can be serialized."); } - return new(testSession.Serialize(jsonSerializerOptions)); + return new(JsonSerializer.SerializeToElement(testSession, jsonSerializerOptions)); } protected override Task RunCoreAsync(IEnumerable messages, AgentSession? session = null, AgentRunOptions? options = null, CancellationToken cancellationToken = default) @@ -506,20 +507,16 @@ public sealed class AGUIEndpointRouteBuilderExtensionsTests }; } - private sealed class TestInMemoryAgentSession : InMemoryAgentSession + private sealed class TestAgentSession : AgentSession { - public TestInMemoryAgentSession() - : base() + public TestAgentSession() { } - public TestInMemoryAgentSession(JsonElement serializedSessionState, JsonSerializerOptions? jsonSerializerOptions = null) - : base(serializedSessionState, jsonSerializerOptions, null) + [JsonConstructor] + public TestAgentSession(AgentSessionStateBag stateBag) : base(stateBag) { } - - internal new JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null) - => base.Serialize(jsonSerializerOptions); } private sealed class TestAgent : AIAgent @@ -529,19 +526,19 @@ public sealed class AGUIEndpointRouteBuilderExtensionsTests public override string? Description => "Test agent"; protected override ValueTask CreateSessionCoreAsync(CancellationToken cancellationToken = default) => - new(new TestInMemoryAgentSession()); + new(new TestAgentSession()); protected override ValueTask DeserializeSessionCoreAsync(JsonElement serializedState, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default) => - new(new TestInMemoryAgentSession(serializedState, jsonSerializerOptions)); + new(serializedState.Deserialize(jsonSerializerOptions)!); protected override ValueTask SerializeSessionCoreAsync(AgentSession session, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default) { - if (session is not TestInMemoryAgentSession testSession) + if (session is not TestAgentSession testSession) { throw new InvalidOperationException("The provided session is not compatible with the agent. Only sessions created by the agent can be serialized."); } - return new(testSession.Serialize(jsonSerializerOptions)); + return new(JsonSerializer.SerializeToElement(testSession, jsonSerializerOptions)); } protected override Task RunCoreAsync(IEnumerable messages, AgentSession? session = null, AgentRunOptions? options = null, CancellationToken cancellationToken = default) diff --git a/dotnet/tests/Microsoft.Agents.AI.Mem0.IntegrationTests/Mem0ProviderTests.cs b/dotnet/tests/Microsoft.Agents.AI.Mem0.IntegrationTests/Mem0ProviderTests.cs index a10f1246aa..72e9f6bdff 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Mem0.IntegrationTests/Mem0ProviderTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Mem0.IntegrationTests/Mem0ProviderTests.cs @@ -1,6 +1,8 @@ // Copyright (c) Microsoft. All rights reserved. using System; +using System.Collections.Generic; +using System.Linq; using System.Net.Http; using System.Net.Http.Headers; using System.Threading; @@ -19,7 +21,6 @@ public sealed class Mem0ProviderTests : IDisposable private const string SkipReason = "Requires a Mem0 service configured"; // Set to null to enable. private static readonly AIAgent s_mockAgent = new Moq.Mock().Object; - private static readonly AgentSession s_mockSession = new Moq.Mock().Object; private readonly HttpClient _httpClient; @@ -49,21 +50,22 @@ public sealed class Mem0ProviderTests : IDisposable var question = new ChatMessage(ChatRole.User, "What is my name?"); var input = new ChatMessage(ChatRole.User, "Hello, my name is Caoimhe."); var storageScope = new Mem0ProviderScope { ThreadId = "it-thread-1", UserId = "it-user-1" }; - var sut = new Mem0Provider(this._httpClient, storageScope); + var mockSession = new TestAgentSession(); + var sut = new Mem0Provider(this._httpClient, _ => new Mem0Provider.State(storageScope)); - await sut.ClearStoredMemoriesAsync(); - var ctxBefore = await sut.InvokingAsync(new AIContextProvider.InvokingContext(s_mockAgent, s_mockSession, [question])); - Assert.DoesNotContain("Caoimhe", ctxBefore.Messages?[0].Text ?? string.Empty); + await sut.ClearStoredMemoriesAsync(mockSession); + var ctxBefore = await sut.InvokingAsync(new AIContextProvider.InvokingContext(s_mockAgent, mockSession, new AIContext { Messages = new List { question } })); + Assert.DoesNotContain("Caoimhe", ctxBefore.Messages?.LastOrDefault()?.Text ?? string.Empty); // Act - await sut.InvokedAsync(new AIContextProvider.InvokedContext(s_mockAgent, s_mockSession, [input])); - var ctxAfterAdding = await GetContextWithRetryAsync(sut, question); - await sut.ClearStoredMemoriesAsync(); - var ctxAfterClearing = await sut.InvokingAsync(new AIContextProvider.InvokingContext(s_mockAgent, s_mockSession, [question])); + await sut.InvokedAsync(new AIContextProvider.InvokedContext(s_mockAgent, mockSession, [input], [])); + var ctxAfterAdding = await GetContextWithRetryAsync(sut, mockSession, question); + await sut.ClearStoredMemoriesAsync(mockSession); + var ctxAfterClearing = await sut.InvokingAsync(new AIContextProvider.InvokingContext(s_mockAgent, mockSession, new AIContext { Messages = new List { question } })); // Assert - Assert.Contains("Caoimhe", ctxAfterAdding.Messages?[0].Text ?? string.Empty); - Assert.DoesNotContain("Caoimhe", ctxAfterClearing.Messages?[0].Text ?? string.Empty); + Assert.Contains("Caoimhe", ctxAfterAdding.Messages?.LastOrDefault()?.Text ?? string.Empty); + Assert.DoesNotContain("Caoimhe", ctxAfterClearing.Messages?.LastOrDefault()?.Text ?? string.Empty); } [Fact(Skip = SkipReason)] @@ -73,21 +75,22 @@ public sealed class Mem0ProviderTests : IDisposable var question = new ChatMessage(ChatRole.User, "What is your name?"); var assistantIntro = new ChatMessage(ChatRole.Assistant, "Hello, I'm a friendly assistant and my name is Caoimhe."); var storageScope = new Mem0ProviderScope { AgentId = "it-agent-1" }; - var sut = new Mem0Provider(this._httpClient, storageScope); + var mockSession = new TestAgentSession(); + var sut = new Mem0Provider(this._httpClient, _ => new Mem0Provider.State(storageScope)); - await sut.ClearStoredMemoriesAsync(); - var ctxBefore = await sut.InvokingAsync(new AIContextProvider.InvokingContext(s_mockAgent, s_mockSession, [question])); - Assert.DoesNotContain("Caoimhe", ctxBefore.Messages?[0].Text ?? string.Empty); + await sut.ClearStoredMemoriesAsync(mockSession); + var ctxBefore = await sut.InvokingAsync(new AIContextProvider.InvokingContext(s_mockAgent, mockSession, new AIContext { Messages = new List { question } })); + Assert.DoesNotContain("Caoimhe", ctxBefore.Messages?.LastOrDefault()?.Text ?? string.Empty); // Act - await sut.InvokedAsync(new AIContextProvider.InvokedContext(s_mockAgent, s_mockSession, [assistantIntro])); - var ctxAfterAdding = await GetContextWithRetryAsync(sut, question); - await sut.ClearStoredMemoriesAsync(); - var ctxAfterClearing = await sut.InvokingAsync(new AIContextProvider.InvokingContext(s_mockAgent, s_mockSession, [question])); + await sut.InvokedAsync(new AIContextProvider.InvokedContext(s_mockAgent, mockSession, [assistantIntro], [])); + var ctxAfterAdding = await GetContextWithRetryAsync(sut, mockSession, question); + await sut.ClearStoredMemoriesAsync(mockSession); + var ctxAfterClearing = await sut.InvokingAsync(new AIContextProvider.InvokingContext(s_mockAgent, mockSession, new AIContext { Messages = new List { question } })); // Assert - Assert.Contains("Caoimhe", ctxAfterAdding.Messages?[0].Text ?? string.Empty); - Assert.DoesNotContain("Caoimhe", ctxAfterClearing.Messages?[0].Text ?? string.Empty); + Assert.Contains("Caoimhe", ctxAfterAdding.Messages?.LastOrDefault()?.Text ?? string.Empty); + Assert.DoesNotContain("Caoimhe", ctxAfterClearing.Messages?.LastOrDefault()?.Text ?? string.Empty); } [Fact(Skip = SkipReason)] @@ -96,38 +99,42 @@ public sealed class Mem0ProviderTests : IDisposable // Arrange var question = new ChatMessage(ChatRole.User, "What is your name?"); var assistantIntro = new ChatMessage(ChatRole.Assistant, "I'm an AI tutor and my name is Caoimhe."); - var sut1 = new Mem0Provider(this._httpClient, new Mem0ProviderScope { AgentId = "it-agent-a" }); - var sut2 = new Mem0Provider(this._httpClient, new Mem0ProviderScope { AgentId = "it-agent-b" }); + var storageScope1 = new Mem0ProviderScope { AgentId = "it-agent-a" }; + var storageScope2 = new Mem0ProviderScope { AgentId = "it-agent-b" }; + var mockSession1 = new TestAgentSession(); + var mockSession2 = new TestAgentSession(); + var sut1 = new Mem0Provider(this._httpClient, _ => new Mem0Provider.State(storageScope1)); + var sut2 = new Mem0Provider(this._httpClient, _ => new Mem0Provider.State(storageScope2)); - await sut1.ClearStoredMemoriesAsync(); - await sut2.ClearStoredMemoriesAsync(); + await sut1.ClearStoredMemoriesAsync(mockSession1); + await sut2.ClearStoredMemoriesAsync(mockSession2); - var ctxBefore1 = await sut1.InvokingAsync(new AIContextProvider.InvokingContext(s_mockAgent, s_mockSession, [question])); - var ctxBefore2 = await sut2.InvokingAsync(new AIContextProvider.InvokingContext(s_mockAgent, s_mockSession, [question])); - Assert.DoesNotContain("Caoimhe", ctxBefore1.Messages?[0].Text ?? string.Empty); - Assert.DoesNotContain("Caoimhe", ctxBefore2.Messages?[0].Text ?? string.Empty); + var ctxBefore1 = await sut1.InvokingAsync(new AIContextProvider.InvokingContext(s_mockAgent, mockSession1, new AIContext { Messages = new List { question } })); + var ctxBefore2 = await sut2.InvokingAsync(new AIContextProvider.InvokingContext(s_mockAgent, mockSession2, new AIContext { Messages = new List { question } })); + Assert.DoesNotContain("Caoimhe", ctxBefore1.Messages?.LastOrDefault()?.Text ?? string.Empty); + Assert.DoesNotContain("Caoimhe", ctxBefore2.Messages?.LastOrDefault()?.Text ?? string.Empty); // Act - await sut1.InvokedAsync(new AIContextProvider.InvokedContext(s_mockAgent, s_mockSession, [assistantIntro])); - var ctxAfterAdding1 = await GetContextWithRetryAsync(sut1, question); - var ctxAfterAdding2 = await GetContextWithRetryAsync(sut2, question); + await sut1.InvokedAsync(new AIContextProvider.InvokedContext(s_mockAgent, mockSession1, [assistantIntro], [])); + var ctxAfterAdding1 = await GetContextWithRetryAsync(sut1, mockSession1, question); + var ctxAfterAdding2 = await GetContextWithRetryAsync(sut2, mockSession2, question); // Assert - Assert.Contains("Caoimhe", ctxAfterAdding1.Messages?[0].Text ?? string.Empty); - Assert.DoesNotContain("Caoimhe", ctxAfterAdding2.Messages?[0].Text ?? string.Empty); + Assert.Contains("Caoimhe", ctxAfterAdding1.Messages?.LastOrDefault()?.Text ?? string.Empty); + Assert.DoesNotContain("Caoimhe", ctxAfterAdding2.Messages?.LastOrDefault()?.Text ?? string.Empty); // Cleanup - await sut1.ClearStoredMemoriesAsync(); - await sut2.ClearStoredMemoriesAsync(); + await sut1.ClearStoredMemoriesAsync(mockSession1); + await sut2.ClearStoredMemoriesAsync(mockSession2); } - private static async Task GetContextWithRetryAsync(Mem0Provider provider, ChatMessage question, int attempts = 5, int delayMs = 1000) + private static async Task GetContextWithRetryAsync(Mem0Provider provider, AgentSession session, ChatMessage question, int attempts = 5, int delayMs = 1000) { AIContext? ctx = null; for (int i = 0; i < attempts; i++) { - ctx = await provider.InvokingAsync(new AIContextProvider.InvokingContext(s_mockAgent, s_mockSession, [question]), CancellationToken.None); - var text = ctx.Messages?[0].Text; + ctx = await provider.InvokingAsync(new AIContextProvider.InvokingContext(s_mockAgent, session, new AIContext { Messages = new List { question } }), CancellationToken.None); + var text = ctx.Messages?.LastOrDefault()?.Text; if (!string.IsNullOrEmpty(text) && text.IndexOf("Caoimhe", StringComparison.OrdinalIgnoreCase) >= 0) { break; @@ -141,4 +148,12 @@ public sealed class Mem0ProviderTests : IDisposable { this._httpClient.Dispose(); } + + private sealed class TestAgentSession : AgentSession + { + public TestAgentSession() + { + this.StateBag = new AgentSessionStateBag(); + } + } } diff --git a/dotnet/tests/Microsoft.Agents.AI.Mem0.UnitTests/Mem0ProviderTests.cs b/dotnet/tests/Microsoft.Agents.AI.Mem0.UnitTests/Mem0ProviderTests.cs index 53c87b09ba..7ad77b7df5 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Mem0.UnitTests/Mem0ProviderTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Mem0.UnitTests/Mem0ProviderTests.cs @@ -19,7 +19,6 @@ namespace Microsoft.Agents.AI.Mem0.UnitTests; public sealed class Mem0ProviderTests : IDisposable { private static readonly AIAgent s_mockAgent = new Mock().Object; - private static readonly AgentSession s_mockSession = new Mock().Object; private readonly Mock> _loggerMock; private readonly Mock _loggerFactoryMock; @@ -55,35 +54,39 @@ public sealed class Mem0ProviderTests : IDisposable using HttpClient client = new(); // Act & Assert - var ex = Assert.Throws(() => new Mem0Provider(client, new Mem0ProviderScope() { ThreadId = "tid" })); + var ex = Assert.Throws(() => new Mem0Provider(client, _ => new Mem0Provider.State(new Mem0ProviderScope { ThreadId = "tid" }))); Assert.StartsWith("The HttpClient BaseAddress must be set for Mem0 operations.", ex.Message); } [Fact] - public void Constructor_Throws_WhenNoStorageScopeValueIsSet() + public void Constructor_Throws_WhenStateInitializerIsNull() { // Act & Assert - var ex = Assert.Throws(() => new Mem0Provider(this._httpClient, new Mem0ProviderScope())); - Assert.StartsWith("At least one of ApplicationId, AgentId, ThreadId, or UserId must be provided for the storage scope.", ex.Message); + var ex = Assert.Throws(() => new Mem0Provider(this._httpClient, null!)); + Assert.Contains("stateInitializer", ex.Message); } [Fact] - public void Constructor_Throws_WhenNoSearchScopeValueIsSet() + public void StateKey_ReturnsDefaultKey_WhenNoOptionsProvided() { - // Act & Assert - var ex = Assert.Throws(() => new Mem0Provider(this._httpClient, new Mem0ProviderScope() { ThreadId = "tid" }, new Mem0ProviderScope())); - Assert.StartsWith("At least one of ApplicationId, AgentId, ThreadId, or UserId must be provided for the search scope.", ex.Message); + // Arrange & Act + var provider = new Mem0Provider(this._httpClient, _ => new Mem0Provider.State(new Mem0ProviderScope { ThreadId = "tid" })); + + // Assert + Assert.Equal("Mem0Provider", provider.StateKey); } [Fact] - public void DeserializingConstructor_Throws_WithEmptyJsonElement() + public void StateKey_ReturnsCustomKey_WhenSetViaOptions() { - // Arrange - var jsonElement = JsonSerializer.SerializeToElement(new object(), Mem0JsonUtilities.DefaultOptions); + // Arrange & Act + var provider = new Mem0Provider( + this._httpClient, + _ => new Mem0Provider.State(new Mem0ProviderScope { ThreadId = "tid" }), + new Mem0ProviderOptions { StateKey = "custom-key" }); - // Act & Assert - var ex = Assert.Throws(() => new Mem0Provider(this._httpClient, jsonElement)); - Assert.StartsWith("The Mem0Provider state did not contain the required scope properties.", ex.Message); + // Assert + Assert.Equal("custom-key", provider.StateKey); } [Fact] @@ -98,8 +101,9 @@ public sealed class Mem0ProviderTests : IDisposable ThreadId = "session", UserId = "user" }; - var sut = new Mem0Provider(this._httpClient, storageScope, options: new() { EnableSensitiveTelemetryData = true }, loggerFactory: this._loggerFactoryMock.Object); - var invokingContext = new AIContextProvider.InvokingContext(s_mockAgent, s_mockSession, [new ChatMessage(ChatRole.User, "What is my name?")]); + var mockSession = new TestAgentSession(); + var sut = new Mem0Provider(this._httpClient, _ => new Mem0Provider.State(storageScope), options: new() { EnableSensitiveTelemetryData = true }, loggerFactory: this._loggerFactoryMock.Object); + var invokingContext = new AIContextProvider.InvokingContext(s_mockAgent, mockSession, new AIContext { Messages = new List { new(ChatRole.User, "What is my name?") } }); // Act var aiContext = await sut.InvokingAsync(invokingContext); @@ -114,9 +118,13 @@ public sealed class Mem0ProviderTests : IDisposable Assert.Equal("What is my name?", doc.RootElement.GetProperty("query").GetString()); Assert.NotNull(aiContext.Messages); - var contextMessage = Assert.Single(aiContext.Messages); + var messages = aiContext.Messages.ToList(); + Assert.Equal(2, messages.Count); + Assert.Equal(AgentRequestMessageSourceType.External, messages[0].GetAgentRequestMessageSourceType()); + var contextMessage = messages[1]; Assert.Equal(ChatRole.User, contextMessage.Role); Assert.Contains("Name is Caoimhe", contextMessage.Text); + Assert.Equal(AgentRequestMessageSourceType.AIContextProvider, contextMessage.GetAgentRequestMessageSourceType()); this._loggerMock.Verify( l => l.Log( @@ -162,9 +170,10 @@ public sealed class Mem0ProviderTests : IDisposable UserId = "user" }; var options = new Mem0ProviderOptions { EnableSensitiveTelemetryData = enableSensitiveTelemetryData }; + var mockSession = new TestAgentSession(); - var sut = new Mem0Provider(this._httpClient, storageScope, options: options, loggerFactory: this._loggerFactoryMock.Object); - var invokingContext = new AIContextProvider.InvokingContext(s_mockAgent, s_mockSession, [new ChatMessage(ChatRole.User, "Who am I?")]); + var sut = new Mem0Provider(this._httpClient, _ => new Mem0Provider.State(storageScope), options: options, loggerFactory: this._loggerFactoryMock.Object); + var invokingContext = new AIContextProvider.InvokingContext(s_mockAgent, mockSession, new AIContext { Messages = new List { new(ChatRole.User, "Who am I?") } }); // Act await sut.InvokingAsync(invokingContext, CancellationToken.None); @@ -204,7 +213,8 @@ public sealed class Mem0ProviderTests : IDisposable this._handler.EnqueueEmptyOk(); // For second CreateMemory this._handler.EnqueueEmptyOk(); // For third CreateMemory var storageScope = new Mem0ProviderScope { ApplicationId = "a", AgentId = "b", ThreadId = "c", UserId = "d" }; - var sut = new Mem0Provider(this._httpClient, storageScope); + var mockSession = new TestAgentSession(); + var sut = new Mem0Provider(this._httpClient, _ => new Mem0Provider.State(storageScope)); var requestMessages = new List { @@ -218,7 +228,7 @@ public sealed class Mem0ProviderTests : IDisposable }; // Act - await sut.InvokedAsync(new AIContextProvider.InvokedContext(s_mockAgent, s_mockSession, requestMessages) { ResponseMessages = responseMessages }); + await sut.InvokedAsync(new AIContextProvider.InvokedContext(s_mockAgent, mockSession, requestMessages, responseMessages)); // Assert var memoryPosts = this._handler.Requests.Where(r => r.RequestMessage.RequestUri!.AbsolutePath == "/v1/memories/" && r.RequestMessage.Method == HttpMethod.Post).ToList(); @@ -235,7 +245,8 @@ public sealed class Mem0ProviderTests : IDisposable { // Arrange var storageScope = new Mem0ProviderScope { ApplicationId = "a", AgentId = "b", ThreadId = "c", UserId = "d" }; - var sut = new Mem0Provider(this._httpClient, storageScope); + var mockSession = new TestAgentSession(); + var sut = new Mem0Provider(this._httpClient, _ => new Mem0Provider.State(storageScope)); var requestMessages = new List { @@ -245,7 +256,7 @@ public sealed class Mem0ProviderTests : IDisposable }; // Act - await sut.InvokedAsync(new AIContextProvider.InvokedContext(s_mockAgent, s_mockSession, requestMessages) { ResponseMessages = null, InvokeException = new InvalidOperationException("Request Failed") }); + await sut.InvokedAsync(new AIContextProvider.InvokedContext(s_mockAgent, mockSession, requestMessages, new InvalidOperationException("Request Failed"))); // Assert Assert.Empty(this._handler.Requests); @@ -256,7 +267,8 @@ public sealed class Mem0ProviderTests : IDisposable { // Arrange var storageScope = new Mem0ProviderScope { ApplicationId = "a", AgentId = "b", ThreadId = "c", UserId = "d" }; - var sut = new Mem0Provider(this._httpClient, storageScope, loggerFactory: this._loggerFactoryMock.Object); + var mockSession = new TestAgentSession(); + var sut = new Mem0Provider(this._httpClient, _ => new Mem0Provider.State(storageScope), loggerFactory: this._loggerFactoryMock.Object); this._handler.EnqueueEmptyInternalServerError(); var requestMessages = new List @@ -271,7 +283,7 @@ public sealed class Mem0ProviderTests : IDisposable }; // Act - await sut.InvokedAsync(new AIContextProvider.InvokedContext(s_mockAgent, s_mockSession, requestMessages) { ResponseMessages = responseMessages }); + await sut.InvokedAsync(new AIContextProvider.InvokedContext(s_mockAgent, mockSession, requestMessages, responseMessages)); // Assert this._loggerMock.Verify( @@ -310,7 +322,8 @@ public sealed class Mem0ProviderTests : IDisposable }; var options = new Mem0ProviderOptions { EnableSensitiveTelemetryData = enableSensitiveTelemetryData }; - var sut = new Mem0Provider(this._httpClient, storageScope, options: options, loggerFactory: this._loggerFactoryMock.Object); + var mockSession = new TestAgentSession(); + var sut = new Mem0Provider(this._httpClient, _ => new Mem0Provider.State(storageScope), options: options, loggerFactory: this._loggerFactoryMock.Object); var requestMessages = new List { new(ChatRole.User, "User text") @@ -321,7 +334,7 @@ public sealed class Mem0ProviderTests : IDisposable }; // Act - await sut.InvokedAsync(new AIContextProvider.InvokedContext(s_mockAgent, s_mockSession, requestMessages) { ResponseMessages = responseMessages }); + await sut.InvokedAsync(new AIContextProvider.InvokedContext(s_mockAgent, mockSession, requestMessages, responseMessages)); // Assert Assert.Equal(expectedLogCount, this._loggerMock.Invocations.Count); @@ -343,73 +356,33 @@ public sealed class Mem0ProviderTests : IDisposable { // Arrange var storageScope = new Mem0ProviderScope { ApplicationId = "app", AgentId = "agent", ThreadId = "session", UserId = "user" }; - var sut = new Mem0Provider(this._httpClient, storageScope); + var sut = new Mem0Provider(this._httpClient, _ => new Mem0Provider.State(storageScope)); this._handler.EnqueueEmptyOk(); // for DELETE + var mockSession = new TestAgentSession(); // Act - await sut.ClearStoredMemoriesAsync(); + await sut.ClearStoredMemoriesAsync(mockSession); // Assert var delete = Assert.Single(this._handler.Requests, r => r.RequestMessage.Method == HttpMethod.Delete); Assert.Equal("https://localhost/v1/memories/?app_id=app&agent_id=agent&run_id=session&user_id=user", delete.RequestMessage.RequestUri!.AbsoluteUri); } - [Fact] - public void Serialize_RoundTripsScopes() - { - // Arrange - var storageScope = new Mem0ProviderScope { ApplicationId = "app", AgentId = "agent", ThreadId = "session", UserId = "user" }; - var sut = new Mem0Provider(this._httpClient, storageScope, options: new() { ContextPrompt = "Custom:" }, loggerFactory: this._loggerFactoryMock.Object); - - // Act - var stateElement = sut.Serialize(); - using JsonDocument doc = JsonDocument.Parse(stateElement.GetRawText()); - var storageScopeElement = doc.RootElement.GetProperty("storageScope"); - Assert.Equal("app", storageScopeElement.GetProperty("applicationId").GetString()); - Assert.Equal("agent", storageScopeElement.GetProperty("agentId").GetString()); - Assert.Equal("session", storageScopeElement.GetProperty("threadId").GetString()); - Assert.Equal("user", storageScopeElement.GetProperty("userId").GetString()); - - var sut2 = new Mem0Provider(this._httpClient, stateElement); - var stateElement2 = sut2.Serialize(); - - // Assert - using JsonDocument doc2 = JsonDocument.Parse(stateElement2.GetRawText()); - var storageScopeElement2 = doc2.RootElement.GetProperty("storageScope"); - Assert.Equal("app", storageScopeElement2.GetProperty("applicationId").GetString()); - Assert.Equal("agent", storageScopeElement2.GetProperty("agentId").GetString()); - Assert.Equal("session", storageScopeElement2.GetProperty("threadId").GetString()); - Assert.Equal("user", storageScopeElement2.GetProperty("userId").GetString()); - } - - [Fact] - public void Serialize_DoesNotIncludeDefaultContextPrompt() - { - // Arrange - var storageScope = new Mem0ProviderScope { ApplicationId = "app" }; - var sut = new Mem0Provider(this._httpClient, storageScope); - - // Act - var stateElement = sut.Serialize(); - - // Assert - using JsonDocument doc = JsonDocument.Parse(stateElement.GetRawText()); - Assert.False(doc.RootElement.TryGetProperty("contextPrompt", out _)); - } - [Fact] public async Task InvokingAsync_ShouldNotThrow_WhenSearchFailsAsync() { // Arrange var storageScope = new Mem0ProviderScope { ApplicationId = "app" }; - var provider = new Mem0Provider(this._httpClient, storageScope, loggerFactory: this._loggerFactoryMock.Object); - var invokingContext = new AIContextProvider.InvokingContext(s_mockAgent, s_mockSession, [new ChatMessage(ChatRole.User, "Q?")]); + var mockSession = new TestAgentSession(); + var provider = new Mem0Provider(this._httpClient, _ => new Mem0Provider.State(storageScope), loggerFactory: this._loggerFactoryMock.Object); + var invokingContext = new AIContextProvider.InvokingContext(s_mockAgent, mockSession, new AIContext { Messages = new List { new(ChatRole.User, "Q?") } }); // Act var aiContext = await provider.InvokingAsync(invokingContext, CancellationToken.None); // Assert - Assert.Null(aiContext.Messages); + Assert.NotNull(aiContext.Messages); + Assert.Single(aiContext.Messages); Assert.Null(aiContext.Tools); this._loggerMock.Verify( l => l.Log( @@ -421,6 +394,159 @@ public sealed class Mem0ProviderTests : IDisposable Times.Once); } + [Fact] + public async Task StateInitializer_IsCalledOnceAndStoredInStateBagAsync() + { + // Arrange + this._handler.EnqueueJsonResponse("[]"); + this._handler.EnqueueJsonResponse("[]"); + var storageScope = new Mem0ProviderScope { ApplicationId = "app" }; + var mockSession = new TestAgentSession(); + int initializerCallCount = 0; + var sut = new Mem0Provider(this._httpClient, _ => + { + initializerCallCount++; + return new Mem0Provider.State(storageScope); + }); + var invokingContext = new AIContextProvider.InvokingContext(s_mockAgent, mockSession, new AIContext { Messages = new List { new(ChatRole.User, "Q?") } }); + + // Act + await sut.InvokingAsync(invokingContext, CancellationToken.None); + await sut.InvokingAsync(invokingContext, CancellationToken.None); + + // Assert + Assert.Equal(1, initializerCallCount); + } + + [Fact] + public async Task StateKey_CanBeConfiguredViaOptionsAsync() + { + // Arrange + this._handler.EnqueueJsonResponse("[]"); + var storageScope = new Mem0ProviderScope { ApplicationId = "app" }; + var mockSession = new TestAgentSession(); + const string CustomKey = "MyCustomKey"; + var sut = new Mem0Provider(this._httpClient, _ => new Mem0Provider.State(storageScope), options: new() { StateKey = CustomKey }); + var invokingContext = new AIContextProvider.InvokingContext(s_mockAgent, mockSession, new AIContext { Messages = new List { new(ChatRole.User, "Q?") } }); + + // Act + await sut.InvokingAsync(invokingContext, CancellationToken.None); + + // Assert + Assert.True(mockSession.StateBag.TryGetValue(CustomKey, out var state, Mem0JsonUtilities.DefaultOptions)); + Assert.NotNull(state); + } + + [Fact] + public async Task InvokingAsync_DefaultFilter_ExcludesNonExternalMessagesFromSearchAsync() + { + // Arrange + this._handler.EnqueueJsonResponse("[]"); // Empty search results + var storageScope = new Mem0ProviderScope { ApplicationId = "app", AgentId = "agent", ThreadId = "session", UserId = "user" }; + var mockSession = new TestAgentSession(); + var sut = new Mem0Provider(this._httpClient, _ => new Mem0Provider.State(storageScope)); + + var requestMessages = new List + { + new(ChatRole.User, "External message"), + new(ChatRole.System, "From history") { AdditionalProperties = new() { { AgentRequestMessageSourceAttribution.AdditionalPropertiesKey, new AgentRequestMessageSourceAttribution(AgentRequestMessageSourceType.ChatHistory, "HistorySource") } } }, + new(ChatRole.System, "From context provider") { AdditionalProperties = new() { { AgentRequestMessageSourceAttribution.AdditionalPropertiesKey, new AgentRequestMessageSourceAttribution(AgentRequestMessageSourceType.AIContextProvider, "ContextSource") } } }, + }; + + var invokingContext = new AIContextProvider.InvokingContext(s_mockAgent, mockSession, new AIContext { Messages = requestMessages }); + + // Act + await sut.InvokingAsync(invokingContext, CancellationToken.None); + + // Assert - Search query should only contain the External message + var searchRequest = Assert.Single(this._handler.Requests, r => r.RequestMessage.Method == HttpMethod.Post); + using JsonDocument doc = JsonDocument.Parse(searchRequest.RequestBody); + Assert.Equal("External message", doc.RootElement.GetProperty("query").GetString()); + } + + [Fact] + public async Task InvokingAsync_CustomSearchInputFilter_OverridesDefaultAsync() + { + // Arrange + this._handler.EnqueueJsonResponse("[]"); // Empty search results + var storageScope = new Mem0ProviderScope { ApplicationId = "app", AgentId = "agent", ThreadId = "session", UserId = "user" }; + var mockSession = new TestAgentSession(); + var sut = new Mem0Provider(this._httpClient, _ => new Mem0Provider.State(storageScope), options: new Mem0ProviderOptions + { + SearchInputMessageFilter = messages => messages // No filtering + }); + + var requestMessages = new List + { + new(ChatRole.User, "External message"), + new(ChatRole.System, "From history") { AdditionalProperties = new() { { AgentRequestMessageSourceAttribution.AdditionalPropertiesKey, new AgentRequestMessageSourceAttribution(AgentRequestMessageSourceType.ChatHistory, "HistorySource") } } }, + }; + + var invokingContext = new AIContextProvider.InvokingContext(s_mockAgent, mockSession, new AIContext { Messages = requestMessages }); + + // Act + await sut.InvokingAsync(invokingContext, CancellationToken.None); + + // Assert - Search query should contain all messages (custom identity filter) + var searchRequest = Assert.Single(this._handler.Requests, r => r.RequestMessage.Method == HttpMethod.Post); + using JsonDocument doc = JsonDocument.Parse(searchRequest.RequestBody); + var queryText = doc.RootElement.GetProperty("query").GetString(); + Assert.Contains("External message", queryText); + Assert.Contains("From history", queryText); + } + + [Fact] + public async Task InvokedAsync_DefaultFilter_ExcludesNonExternalMessagesFromStorageAsync() + { + // Arrange + this._handler.EnqueueEmptyOk(); // For the one message that should be stored + var storageScope = new Mem0ProviderScope { ApplicationId = "a", AgentId = "b", ThreadId = "c", UserId = "d" }; + var mockSession = new TestAgentSession(); + var sut = new Mem0Provider(this._httpClient, _ => new Mem0Provider.State(storageScope)); + + var requestMessages = new List + { + new(ChatRole.User, "External message"), + new(ChatRole.System, "From history") { AdditionalProperties = new() { { AgentRequestMessageSourceAttribution.AdditionalPropertiesKey, new AgentRequestMessageSourceAttribution(AgentRequestMessageSourceType.ChatHistory, "HistorySource") } } }, + }; + + // Act + await sut.InvokedAsync(new AIContextProvider.InvokedContext(s_mockAgent, mockSession, requestMessages, [])); + + // Assert - Only the External message should be persisted + var memoryPosts = this._handler.Requests.Where(r => r.RequestMessage.RequestUri!.AbsolutePath == "/v1/memories/" && r.RequestMessage.Method == HttpMethod.Post).ToList(); + Assert.Single(memoryPosts); + Assert.Contains("External message", memoryPosts[0].RequestBody); + Assert.DoesNotContain(memoryPosts, r => ContainsOrdinal(r.RequestBody, "From history")); + } + + [Fact] + public async Task InvokedAsync_CustomStorageInputFilter_OverridesDefaultAsync() + { + // Arrange + this._handler.EnqueueEmptyOk(); // For first CreateMemory + this._handler.EnqueueEmptyOk(); // For second CreateMemory + var storageScope = new Mem0ProviderScope { ApplicationId = "a", AgentId = "b", ThreadId = "c", UserId = "d" }; + var mockSession = new TestAgentSession(); + var sut = new Mem0Provider(this._httpClient, _ => new Mem0Provider.State(storageScope), options: new Mem0ProviderOptions + { + StorageInputMessageFilter = messages => messages // No filtering - store everything + }); + + var requestMessages = new List + { + new(ChatRole.User, "External message"), + new(ChatRole.System, "From history") { AdditionalProperties = new() { { AgentRequestMessageSourceAttribution.AdditionalPropertiesKey, new AgentRequestMessageSourceAttribution(AgentRequestMessageSourceType.ChatHistory, "HistorySource") } } }, + }; + + // Act + await sut.InvokedAsync(new AIContextProvider.InvokedContext(s_mockAgent, mockSession, requestMessages, [])); + + // Assert - Both messages should be persisted (identity filter overrides default) + var memoryPosts = this._handler.Requests.Where(r => r.RequestMessage.RequestUri!.AbsolutePath == "/v1/memories/" && r.RequestMessage.Method == HttpMethod.Post).ToList(); + Assert.Equal(2, memoryPosts.Count); + } + private static bool ContainsOrdinal(string source, string value) => source.IndexOf(value, StringComparison.Ordinal) >= 0; public void Dispose() @@ -465,4 +591,12 @@ public sealed class Mem0ProviderTests : IDisposable public void EnqueueEmptyInternalServerError() => this._responses.Enqueue(new HttpResponseMessage(System.Net.HttpStatusCode.InternalServerError)); } + + private sealed class TestAgentSession : AgentSession + { + public TestAgentSession() + { + this.StateBag = new AgentSessionStateBag(); + } + } } diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentOptionsTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentOptionsTests.cs index 8502550d2c..f69fb3d636 100644 --- a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentOptionsTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentOptionsTests.cs @@ -1,8 +1,6 @@ // Copyright (c) Microsoft. All rights reserved. using System.Collections.Generic; -using System.Threading; -using System.Threading.Tasks; using Microsoft.Extensions.AI; using Moq; @@ -23,8 +21,8 @@ public class ChatClientAgentOptionsTests Assert.Null(options.Name); Assert.Null(options.Description); Assert.Null(options.ChatOptions); - Assert.Null(options.ChatHistoryProviderFactory); - Assert.Null(options.AIContextProviderFactory); + Assert.Null(options.ChatHistoryProvider); + Assert.Null(options.AIContextProviders); } [Fact] @@ -36,8 +34,8 @@ public class ChatClientAgentOptionsTests // Assert Assert.Null(options.Name); Assert.Null(options.Description); - Assert.Null(options.AIContextProviderFactory); - Assert.Null(options.ChatHistoryProviderFactory); + Assert.Null(options.AIContextProviders); + Assert.Null(options.ChatHistoryProvider); Assert.NotNull(options.ChatOptions); Assert.Null(options.ChatOptions.Instructions); Assert.Null(options.ChatOptions.Tools); @@ -117,11 +115,8 @@ public class ChatClientAgentOptionsTests const string Description = "Test description"; var tools = new List { AIFunctionFactory.Create(() => "test") }; - static ValueTask ChatHistoryProviderFactoryAsync( - ChatClientAgentOptions.ChatHistoryProviderFactoryContext ctx, CancellationToken ct) => new(new Mock().Object); - - static ValueTask AIContextProviderFactoryAsync( - ChatClientAgentOptions.AIContextProviderFactoryContext ctx, CancellationToken ct) => new(new Mock().Object); + var mockChatHistoryProvider = new Mock().Object; + var mockAIContextProvider = new Mock().Object; var original = new ChatClientAgentOptions() { @@ -129,8 +124,8 @@ public class ChatClientAgentOptionsTests Description = Description, ChatOptions = new() { Tools = tools }, Id = "test-id", - ChatHistoryProviderFactory = ChatHistoryProviderFactoryAsync, - AIContextProviderFactory = AIContextProviderFactoryAsync + ChatHistoryProvider = mockChatHistoryProvider, + AIContextProviders = [mockAIContextProvider] }; // Act @@ -141,8 +136,8 @@ public class ChatClientAgentOptionsTests Assert.Equal(original.Id, clone.Id); Assert.Equal(original.Name, clone.Name); Assert.Equal(original.Description, clone.Description); - Assert.Same(original.ChatHistoryProviderFactory, clone.ChatHistoryProviderFactory); - Assert.Same(original.AIContextProviderFactory, clone.AIContextProviderFactory); + Assert.Same(original.ChatHistoryProvider, clone.ChatHistoryProvider); + Assert.Equal(original.AIContextProviders, clone.AIContextProviders); // ChatOptions should be cloned, not the same reference Assert.NotSame(original.ChatOptions, clone.ChatOptions); @@ -154,11 +149,16 @@ public class ChatClientAgentOptionsTests public void Clone_WithoutProvidingChatOptions_ClonesCorrectly() { // Arrange + var mockChatHistoryProvider = new Mock().Object; + var mockAIContextProvider = new Mock().Object; + var original = new ChatClientAgentOptions { Id = "test-id", Name = "Test name", - Description = "Test description" + Description = "Test description", + ChatHistoryProvider = mockChatHistoryProvider, + AIContextProviders = [mockAIContextProvider] }; // Act @@ -170,8 +170,8 @@ public class ChatClientAgentOptionsTests Assert.Equal(original.Name, clone.Name); Assert.Equal(original.Description, clone.Description); Assert.Null(original.ChatOptions); - Assert.Null(clone.ChatHistoryProviderFactory); - Assert.Null(clone.AIContextProviderFactory); + Assert.Same(original.ChatHistoryProvider, clone.ChatHistoryProvider); + Assert.Equal(original.AIContextProviders, clone.AIContextProviders); } private static void AssertSameTools(IList? expected, IList? actual) diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentSessionTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentSessionTests.cs index fd311f9225..1f5e5aa9cd 100644 --- a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentSessionTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentSessionTests.cs @@ -1,12 +1,9 @@ // Copyright (c) Microsoft. All rights reserved. using System; -using System.Collections.Generic; using System.Linq; using System.Text.Json; -using System.Threading.Tasks; using Microsoft.Extensions.AI; -using Moq; #pragma warning disable CA1861 // Avoid constant arrays as arguments @@ -24,7 +21,6 @@ public class ChatClientAgentSessionTests // Assert Assert.Null(session.ConversationId); - Assert.Null(session.ChatHistoryProvider); } [Fact] @@ -39,53 +35,6 @@ public class ChatClientAgentSessionTests // Assert Assert.Equal(ConversationId, session.ConversationId); - Assert.Null(session.ChatHistoryProvider); - } - - [Fact] - public void SetChatHistoryProviderRoundtrips() - { - // Arrange - var session = new ChatClientAgentSession(); - var chatHistoryProvider = new InMemoryChatHistoryProvider(); - - // Act - session.ChatHistoryProvider = chatHistoryProvider; - - // Assert - Assert.Same(chatHistoryProvider, session.ChatHistoryProvider); - Assert.Null(session.ConversationId); - } - - [Fact] - public void SetConversationIdThrowsWhenChatHistoryProviderIsSet() - { - // Arrange - var session = new ChatClientAgentSession - { - ChatHistoryProvider = new InMemoryChatHistoryProvider() - }; - - // Act & Assert - var exception = Assert.Throws(() => session.ConversationId = "new-session-id"); - Assert.Equal("Only the ConversationId or ChatHistoryProvider may be set, but not both and switching from one to another is not supported.", exception.Message); - Assert.NotNull(session.ChatHistoryProvider); - } - - [Fact] - public void SetChatHistoryProviderThrowsWhenConversationIdIsSet() - { - // Arrange - var session = new ChatClientAgentSession - { - ConversationId = "existing-session-id" - }; - var provider = new InMemoryChatHistoryProvider(); - - // Act & Assert - var exception = Assert.Throws(() => session.ChatHistoryProvider = provider); - Assert.Equal("Only the ConversationId or ChatHistoryProvider may be set, but not both and switching from one to another is not supported.", exception.Message); - Assert.NotNull(session.ConversationId); } #endregion Constructor and Property Tests @@ -93,29 +42,33 @@ public class ChatClientAgentSessionTests #region Deserialize Tests [Fact] - public async Task VerifyDeserializeWithMessagesAsync() + public void VerifyDeserializeWithMessages() { // Arrange var json = JsonSerializer.Deserialize(""" { - "chatHistoryProviderState": { "messages": [{"authorName": "testAuthor"}] } + "stateBag": { + "InMemoryChatHistoryProvider": { + "messages": [{"authorName": "testAuthor"}] + } + } } """, TestJsonSerializerContext.Default.JsonElement); // Act. - var session = await ChatClientAgentSession.DeserializeAsync(json); + var session = ChatClientAgentSession.Deserialize(json, TestJsonSerializerContext.Default.Options); // Assert Assert.Null(session.ConversationId); - var chatHistoryProvider = session.ChatHistoryProvider as InMemoryChatHistoryProvider; - Assert.NotNull(chatHistoryProvider); - Assert.Single(chatHistoryProvider); - Assert.Equal("testAuthor", chatHistoryProvider[0].AuthorName); + var chatHistoryProvider = new InMemoryChatHistoryProvider(); + var messages = chatHistoryProvider.GetMessages(session); + Assert.Single(messages); + Assert.Equal("testAuthor", messages[0].AuthorName); } [Fact] - public async Task VerifyDeserializeWithIdAsync() + public void VerifyDeserializeWithId() { // Arrange var json = JsonSerializer.Deserialize(""" @@ -125,42 +78,43 @@ public class ChatClientAgentSessionTests """, TestJsonSerializerContext.Default.JsonElement); // Act - var session = await ChatClientAgentSession.DeserializeAsync(json); + var session = ChatClientAgentSession.Deserialize(json); // Assert Assert.Equal("TestConvId", session.ConversationId); - Assert.Null(session.ChatHistoryProvider); } [Fact] - public async Task VerifyDeserializeWithAIContextProviderAsync() + public void VerifyDeserializeWithStateBag() { // Arrange var json = JsonSerializer.Deserialize(""" { "conversationId": "TestConvId", - "aiContextProviderState": ["CP1"] + "stateBag": { + "dog": { + "name": "Fido" + } + } } """, TestJsonSerializerContext.Default.JsonElement); - Mock mockProvider = new(); - // Act - var session = await ChatClientAgentSession.DeserializeAsync(json, aiContextProviderFactory: (_, _, _) => new(mockProvider.Object)); + var session = ChatClientAgentSession.Deserialize(json); // Assert - Assert.Null(session.ChatHistoryProvider); - Assert.Same(session.AIContextProvider, mockProvider.Object); + var dog = session.StateBag.GetValue("dog", TestJsonSerializerContext.Default.Options); + Assert.NotNull(dog); + Assert.Equal("Fido", dog.Name); } [Fact] - public async Task DeserializeWithInvalidJsonThrowsAsync() + public void DeserializeWithInvalidJsonThrows() { // Arrange var invalidJson = JsonSerializer.Deserialize("[42]", TestJsonSerializerContext.Default.JsonElement); - var session = new ChatClientAgentSession(); // Act & Assert - await Assert.ThrowsAsync(() => ChatClientAgentSession.DeserializeAsync(invalidJson)); + Assert.Throws(() => ChatClientAgentSession.Deserialize(invalidJson)); } #endregion Deserialize Tests @@ -195,8 +149,9 @@ public class ChatClientAgentSessionTests public void VerifySessionSerializationWithMessages() { // Arrange - InMemoryChatHistoryProvider provider = [new(ChatRole.User, "TestContent") { AuthorName = "TestAuthor" }]; - var session = new ChatClientAgentSession { ChatHistoryProvider = provider }; + var provider = new InMemoryChatHistoryProvider(); + var session = new ChatClientAgentSession(); + provider.SetMessages(session, [new(ChatRole.User, "TestContent") { AuthorName = "TestAuthor" }]); // Act var json = session.Serialize(); @@ -206,10 +161,12 @@ public class ChatClientAgentSessionTests Assert.False(json.TryGetProperty("conversationId", out _)); - Assert.True(json.TryGetProperty("chatHistoryProviderState", out var chatHistoryProviderStateProperty)); - Assert.Equal(JsonValueKind.Object, chatHistoryProviderStateProperty.ValueKind); - - Assert.True(chatHistoryProviderStateProperty.TryGetProperty("messages", out var messagesProperty)); + // Messages should be stored in the stateBag + Assert.True(json.TryGetProperty("stateBag", out var stateBagProperty)); + Assert.Equal(JsonValueKind.Object, stateBagProperty.ValueKind); + Assert.True(stateBagProperty.TryGetProperty("InMemoryChatHistoryProvider", out var providerStateProperty)); + Assert.Equal(JsonValueKind.Object, providerStateProperty.ValueKind); + Assert.True(providerStateProperty.TryGetProperty("messages", out var messagesProperty)); Assert.Equal(JsonValueKind.Array, messagesProperty.ValueKind); Assert.Single(messagesProperty.EnumerateArray()); @@ -224,29 +181,23 @@ public class ChatClientAgentSessionTests } [Fact] - public void VerifySessionSerializationWithWithAIContextProvider() + public void VerifySessionSerializationWithWithStateBag() { // Arrange - Mock mockProvider = new(); - mockProvider - .Setup(m => m.Serialize(It.IsAny())) - .Returns(JsonSerializer.SerializeToElement(["CP1"], TestJsonSerializerContext.Default.StringArray)); - - var session = new ChatClientAgentSession - { - AIContextProvider = mockProvider.Object - }; + var session = new ChatClientAgentSession(); + session.StateBag.SetValue("dog", new Animal { Name = "Fido" }, TestJsonSerializerContext.Default.Options); // Act var json = session.Serialize(); // Assert Assert.Equal(JsonValueKind.Object, json.ValueKind); - Assert.True(json.TryGetProperty("aiContextProviderState", out var providerStateProperty)); - Assert.Equal(JsonValueKind.Array, providerStateProperty.ValueKind); - Assert.Single(providerStateProperty.EnumerateArray()); - Assert.Equal("CP1", providerStateProperty.EnumerateArray().First().GetString()); - mockProvider.Verify(m => m.Serialize(It.IsAny()), Times.Once); + Assert.True(json.TryGetProperty("stateBag", out var stateBagProperty)); + Assert.Equal(JsonValueKind.Object, stateBagProperty.ValueKind); + Assert.True(stateBagProperty.TryGetProperty("dog", out var dogProperty)); + Assert.Equal(JsonValueKind.Object, dogProperty.ValueKind); + Assert.True(dogProperty.TryGetProperty("name", out var nameProperty)); + Assert.Equal("Fido", nameProperty.GetString()); } /// @@ -258,17 +209,7 @@ public class ChatClientAgentSessionTests // Arrange var session = new ChatClientAgentSession(); JsonSerializerOptions options = new() { PropertyNamingPolicy = JsonNamingPolicy.SnakeCaseLower }; - options.TypeInfoResolverChain.Add(AgentAbstractionsJsonUtilities.DefaultOptions.TypeInfoResolver!); - - var chatHistoryProviderStateElement = JsonSerializer.SerializeToElement( - new Dictionary { ["Key"] = "TestValue" }, - TestJsonSerializerContext.Default.DictionaryStringObject); - - var chatHistoryProviderMock = new Mock(); - chatHistoryProviderMock - .Setup(m => m.Serialize(options)) - .Returns(chatHistoryProviderStateElement); - session.ChatHistoryProvider = chatHistoryProviderMock.Object; + options.TypeInfoResolverChain.Add(AgentJsonUtilities.DefaultOptions.TypeInfoResolver!); // Act var json = session.Serialize(options); @@ -276,54 +217,29 @@ public class ChatClientAgentSessionTests // Assert Assert.Equal(JsonValueKind.Object, json.ValueKind); - Assert.False(json.TryGetProperty("conversationId", out var idProperty)); - - Assert.True(json.TryGetProperty("chatHistoryProviderState", out var chatHistoryProviderStateProperty)); - Assert.Equal(JsonValueKind.Object, chatHistoryProviderStateProperty.ValueKind); - - Assert.True(chatHistoryProviderStateProperty.TryGetProperty("Key", out var keyProperty)); - Assert.Equal("TestValue", keyProperty.GetString()); - - chatHistoryProviderMock.Verify(m => m.Serialize(options), Times.Once); + // [JsonPropertyName] takes precedence over naming policy + Assert.True(json.TryGetProperty("conversationId", out var _)); } #endregion Serialize Tests - #region GetService Tests + #region StateBag Roundtrip Tests [Fact] - public void GetService_RequestingAIContextProvider_ReturnsAIContextProvider() + public void VerifyStateBagRoundtrips() { // Arrange var session = new ChatClientAgentSession(); - var mockProvider = new Mock(); - mockProvider - .Setup(m => m.GetService(It.Is(x => x == typeof(AIContextProvider)), null)) - .Returns(mockProvider.Object); - session.AIContextProvider = mockProvider.Object; + session.StateBag.SetValue("dog", new Animal { Name = "Fido" }, TestJsonSerializerContext.Default.Options); // Act - var result = session.GetService(typeof(AIContextProvider)); + var serializedSession = session.Serialize(); + var deserializedSession = ChatClientAgentSession.Deserialize(serializedSession); // Assert - Assert.NotNull(result); - Assert.Same(mockProvider.Object, result); - } - - [Fact] - public void GetService_RequestingChatHistoryProvider_ReturnsChatHistoryProvider() - { - // Arrange - var session = new ChatClientAgentSession(); - var chatHistoryProvider = new InMemoryChatHistoryProvider(); - session.ChatHistoryProvider = chatHistoryProvider; - - // Act - var result = session.GetService(typeof(ChatHistoryProvider)); - - // Assert - Assert.NotNull(result); - Assert.Same(chatHistoryProvider, result); + var dog = deserializedSession.StateBag.GetValue("dog", TestJsonSerializerContext.Default.Options); + Assert.NotNull(dog); + Assert.Equal("Fido", dog.Name); } #endregion diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentTests.cs index c23e8cffaf..7b33cbbd5f 100644 --- a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentTests.cs @@ -45,6 +45,154 @@ public partial class ChatClientAgentTests Assert.Equal("FunctionInvokingChatClient", agent.ChatClient.GetType().Name); } + /// + /// Verify that the constructor throws when two AIContextProviders use the same StateKey. + /// + [Fact] + public void Constructor_ThrowsWhenDuplicateAIContextProviderStateKeys() + { + // Arrange + var chatClient = new Mock().Object; + var provider1 = new TestAIContextProvider("SharedKey"); + var provider2 = new TestAIContextProvider("SharedKey"); + + // Act & Assert + var ex = Assert.Throws(() => + new ChatClientAgent(chatClient, options: new() + { + AIContextProviders = [provider1, provider2] + })); + + Assert.Contains("SharedKey", ex.Message); + } + + /// + /// Verify that the constructor throws when an AIContextProvider uses the same StateKey as the default InMemoryChatHistoryProvider + /// and no explicit ChatHistoryProvider is configured. + /// + [Fact] + public void Constructor_ThrowsWhenAIContextProviderStateKeyClashesWithDefaultInMemoryChatHistoryProvider() + { + // Arrange + var chatClient = new Mock().Object; + var contextProvider = new TestAIContextProvider(nameof(InMemoryChatHistoryProvider)); + + // Act & Assert + var ex = Assert.Throws(() => + new ChatClientAgent(chatClient, options: new() + { + AIContextProviders = [contextProvider] + })); + + Assert.Contains(nameof(InMemoryChatHistoryProvider), ex.Message); + } + + /// + /// Verify that the constructor throws when a ChatHistoryProvider uses the same StateKey as an AIContextProvider. + /// + [Fact] + public void Constructor_ThrowsWhenChatHistoryProviderStateKeyClashesWithAIContextProvider() + { + // Arrange + var chatClient = new Mock().Object; + var contextProvider = new TestAIContextProvider("SharedKey"); + var historyProvider = new TestChatHistoryProvider("SharedKey"); + + // Act & Assert + var ex = Assert.Throws(() => + new ChatClientAgent(chatClient, options: new() + { + AIContextProviders = [contextProvider], + ChatHistoryProvider = historyProvider + })); + + Assert.Contains("SharedKey", ex.Message); + Assert.Contains(nameof(ChatHistoryProvider), ex.Message); + } + + /// + /// Verify that the constructor succeeds when all providers use unique StateKeys. + /// + [Fact] + public void Constructor_SucceedsWithUniqueProviderStateKeys() + { + // Arrange + var chatClient = new Mock().Object; + var contextProvider1 = new TestAIContextProvider("Key1"); + var contextProvider2 = new TestAIContextProvider("Key2"); + var historyProvider = new TestChatHistoryProvider("Key3"); + + // Act & Assert - should not throw + _ = new ChatClientAgent(chatClient, options: new() + { + AIContextProviders = [contextProvider1, contextProvider2], + ChatHistoryProvider = historyProvider + }); + } + + /// + /// Verify that RunAsync throws when an override ChatHistoryProvider's StateKey clashes with an AIContextProvider. + /// + [Fact] + public async Task RunAsync_ThrowsWhenOverrideChatHistoryProviderStateKeyClashesWithAIContextProviderAsync() + { + // Arrange + Mock mockService = new(); + mockService.Setup( + s => s.GetResponseAsync( + It.IsAny>(), + It.IsAny(), + It.IsAny())).ReturnsAsync(new ChatResponse([new(ChatRole.Assistant, "response")])); + + var contextProvider = new TestAIContextProvider("SharedKey"); + var overrideHistoryProvider = new TestChatHistoryProvider("SharedKey"); + + ChatClientAgent agent = new(mockService.Object, options: new() + { + AIContextProviders = [contextProvider] + }); + + // Act & Assert + ChatClientAgentSession? session = await agent.CreateSessionAsync() as ChatClientAgentSession; + AdditionalPropertiesDictionary additionalProperties = new(); + additionalProperties.Add(overrideHistoryProvider); + + var ex = await Assert.ThrowsAsync(() => + agent.RunAsync([new(ChatRole.User, "test")], session, options: new AgentRunOptions { AdditionalProperties = additionalProperties })); + + Assert.Contains("SharedKey", ex.Message); + } + + /// + /// Verify that RunAsync succeeds when an override ChatHistoryProvider uses the same StateKey as the default ChatHistoryProvider. + /// + [Fact] + public async Task RunAsync_SucceedsWhenOverrideChatHistoryProviderSharesKeyWithDefaultAsync() + { + // Arrange + Mock mockService = new(); + mockService.Setup( + s => s.GetResponseAsync( + It.IsAny>(), + It.IsAny(), + It.IsAny())).ReturnsAsync(new ChatResponse([new(ChatRole.Assistant, "response")])); + + var defaultHistoryProvider = new TestChatHistoryProvider("SameKey"); + var overrideHistoryProvider = new TestChatHistoryProvider("SameKey"); + + ChatClientAgent agent = new(mockService.Object, options: new() + { + ChatHistoryProvider = defaultHistoryProvider + }); + + // Act & Assert - should not throw + ChatClientAgentSession? session = await agent.CreateSessionAsync() as ChatClientAgentSession; + AdditionalPropertiesDictionary additionalProperties = new(); + additionalProperties.Add(overrideHistoryProvider); + + await agent.RunAsync([new(ChatRole.User, "test")], session, options: new AgentRunOptions { AdditionalProperties = additionalProperties }); + } + #endregion #region RunAsync Tests @@ -345,18 +493,19 @@ public partial class ChatClientAgentTests mockProvider .Protected() .Setup>("InvokingCoreAsync", ItExpr.IsAny(), ItExpr.IsAny()) - .ReturnsAsync(new AIContext - { - Messages = aiContextProviderMessages, - Instructions = "context provider instructions", - Tools = [AIFunctionFactory.Create(() => { }, "context provider function")] - }); + .Returns((AIContextProvider.InvokingContext ctx, CancellationToken _) => + new ValueTask(new AIContext + { + Messages = (ctx.AIContext.Messages ?? []).Concat(aiContextProviderMessages), + Instructions = ctx.AIContext.Instructions + "\ncontext provider instructions", + Tools = (ctx.AIContext.Tools ?? []).Concat(new[] { AIFunctionFactory.Create(() => { }, "context provider function") }) + })); mockProvider .Protected() .Setup("InvokedCoreAsync", ItExpr.IsAny(), ItExpr.IsAny()) .Returns(new ValueTask()); - ChatClientAgent agent = new(mockService.Object, options: new() { AIContextProviderFactory = (_, _) => new(mockProvider.Object), ChatOptions = new() { Instructions = "base instructions", Tools = [AIFunctionFactory.Create(() => { }, "base function")] } }); + ChatClientAgent agent = new(mockService.Object, options: new() { AIContextProviders = [mockProvider.Object], ChatOptions = new() { Instructions = "base instructions", Tools = [AIFunctionFactory.Create(() => { }, "base function")] } }); // Act var session = await agent.CreateSessionAsync() as ChatClientAgentSession; @@ -375,11 +524,13 @@ public partial class ChatClientAgentTests Assert.Contains(capturedTools, t => t.Name == "context provider function"); // Verify that the session was updated with the ai context provider, input and response messages - var chatHistoryProvider = Assert.IsType(session!.ChatHistoryProvider); - Assert.Equal(3, chatHistoryProvider.Count); - Assert.Equal("user message", chatHistoryProvider[0].Text); - Assert.Equal("context provider message", chatHistoryProvider[1].Text); - Assert.Equal("response", chatHistoryProvider[2].Text); + var chatHistoryProvider = agent.ChatHistoryProvider as InMemoryChatHistoryProvider; + Assert.NotNull(chatHistoryProvider); + var messages = chatHistoryProvider.GetMessages(session); + Assert.Equal(3, messages.Count); + Assert.Equal("user message", messages[0].Text); + Assert.Equal("context provider message", messages[1].Text); + Assert.Equal("response", messages[2].Text); mockProvider .Protected() @@ -413,16 +564,17 @@ public partial class ChatClientAgentTests mockProvider .Protected() .Setup>("InvokingCoreAsync", ItExpr.IsAny(), ItExpr.IsAny()) - .ReturnsAsync(new AIContext - { - Messages = aiContextProviderMessages, - }); + .Returns((AIContextProvider.InvokingContext ctx, CancellationToken _) => + new ValueTask(new AIContext + { + Messages = (ctx.AIContext.Messages ?? []).Concat(aiContextProviderMessages), + })); mockProvider .Protected() .Setup("InvokedCoreAsync", ItExpr.IsAny(), ItExpr.IsAny()) .Returns(new ValueTask()); - ChatClientAgent agent = new(mockService.Object, options: new() { AIContextProviderFactory = (_, _) => new(mockProvider.Object), ChatOptions = new() { Instructions = "base instructions", Tools = [AIFunctionFactory.Create(() => { }, "base function")] } }); + ChatClientAgent agent = new(mockService.Object, options: new() { AIContextProviders = [mockProvider.Object], ChatOptions = new() { Instructions = "base instructions", Tools = [AIFunctionFactory.Create(() => { }, "base function")] } }); // Act await Assert.ThrowsAsync(() => agent.RunAsync(requestMessages)); @@ -470,9 +622,15 @@ public partial class ChatClientAgentTests mockProvider .Protected() .Setup>("InvokingCoreAsync", ItExpr.IsAny(), ItExpr.IsAny()) - .ReturnsAsync(new AIContext()); + .Returns((AIContextProvider.InvokingContext ctx, CancellationToken _) => + new ValueTask(new AIContext + { + Instructions = ctx.AIContext.Instructions, + Messages = ctx.AIContext.Messages, + Tools = ctx.AIContext.Tools + })); - ChatClientAgent agent = new(mockService.Object, options: new() { AIContextProviderFactory = (_, _) => new(mockProvider.Object), ChatOptions = new() { Instructions = "base instructions", Tools = [AIFunctionFactory.Create(() => { }, "base function")] } }); + ChatClientAgent agent = new(mockService.Object, options: new() { AIContextProviders = [mockProvider.Object], ChatOptions = new() { Instructions = "base instructions", Tools = [AIFunctionFactory.Create(() => { }, "base function")] } }); // Act await agent.RunAsync([new(ChatRole.User, "user message")]); @@ -490,6 +648,299 @@ public partial class ChatClientAgentTests .Verify>("InvokingCoreAsync", Times.Once(), ItExpr.IsAny(), ItExpr.IsAny()); } + /// + /// Verify that RunAsync invokes multiple AIContextProviders in sequence, each receiving the accumulated context. + /// + [Fact] + public async Task RunAsyncInvokesMultipleAIContextProvidersInOrderAsync() + { + // Arrange + ChatMessage[] requestMessages = [new(ChatRole.User, "user message")]; + ChatMessage[] responseMessages = [new(ChatRole.Assistant, "response")]; + Mock mockService = new(); + List capturedMessages = []; + string capturedInstructions = string.Empty; + List capturedTools = []; + mockService + .Setup(s => s.GetResponseAsync( + It.IsAny>(), + It.IsAny(), + It.IsAny())) + .Callback, ChatOptions, CancellationToken>((msgs, opts, ct) => + { + capturedMessages.AddRange(msgs); + capturedInstructions = opts.Instructions ?? string.Empty; + if (opts.Tools is not null) + { + capturedTools.AddRange(opts.Tools); + } + }) + .ReturnsAsync(new ChatResponse(responseMessages)); + + // Provider 1: adds a system message and a tool + var mockProvider1 = new Mock(); + mockProvider1.SetupGet(p => p.StateKey).Returns("Provider1"); + mockProvider1 + .Protected() + .Setup>("InvokingCoreAsync", ItExpr.IsAny(), ItExpr.IsAny()) + .Returns((AIContextProvider.InvokingContext ctx, CancellationToken _) => + new ValueTask(new AIContext + { + Messages = (ctx.AIContext.Messages ?? []).Concat([new ChatMessage(ChatRole.System, "provider1 context")]).ToList(), + Instructions = ctx.AIContext.Instructions + "\nprovider1 instructions", + Tools = (ctx.AIContext.Tools ?? []).Concat([AIFunctionFactory.Create(() => { }, "provider1 function")]).ToList() + })); + mockProvider1 + .Protected() + .Setup("InvokedCoreAsync", ItExpr.IsAny(), ItExpr.IsAny()) + .Returns(new ValueTask()); + + // Provider 2: adds another system message and verifies it receives accumulated context from provider 1 + AIContext? provider2ReceivedContext = null; + var mockProvider2 = new Mock(); + mockProvider2.SetupGet(p => p.StateKey).Returns("Provider2"); + mockProvider2 + .Protected() + .Setup>("InvokingCoreAsync", ItExpr.IsAny(), ItExpr.IsAny()) + .Returns((AIContextProvider.InvokingContext ctx, CancellationToken _) => + { + provider2ReceivedContext = ctx.AIContext; + return new ValueTask(new AIContext + { + Messages = (ctx.AIContext.Messages ?? []).Concat([new ChatMessage(ChatRole.System, "provider2 context")]).ToList(), + Instructions = ctx.AIContext.Instructions + "\nprovider2 instructions", + Tools = (ctx.AIContext.Tools ?? []).Concat([AIFunctionFactory.Create(() => { }, "provider2 function")]).ToList() + }); + }); + mockProvider2 + .Protected() + .Setup("InvokedCoreAsync", ItExpr.IsAny(), ItExpr.IsAny()) + .Returns(new ValueTask()); + + ChatClientAgent agent = new(mockService.Object, options: new() + { + AIContextProviders = [mockProvider1.Object, mockProvider2.Object], + ChatOptions = new() { Instructions = "base instructions", Tools = [AIFunctionFactory.Create(() => { }, "base function")] } + }); + + // Act + var session = await agent.CreateSessionAsync() as ChatClientAgentSession; + await agent.RunAsync(requestMessages, session); + + // Assert + // Provider 2 should have received accumulated context from provider 1 + Assert.NotNull(provider2ReceivedContext); + Assert.Contains(provider2ReceivedContext.Messages!, m => m.Text == "provider1 context"); + Assert.Contains("provider1 instructions", provider2ReceivedContext.Instructions); + + // Final captured messages should contain user message + both provider contexts + Assert.Equal(3, capturedMessages.Count); + Assert.Equal("user message", capturedMessages[0].Text); + Assert.Equal("provider1 context", capturedMessages[1].Text); + Assert.Equal("provider2 context", capturedMessages[2].Text); + + // Instructions should be accumulated + Assert.Equal("base instructions\nprovider1 instructions\nprovider2 instructions", capturedInstructions); + + // Tools should contain base + both provider tools + Assert.Equal(3, capturedTools.Count); + Assert.Contains(capturedTools, t => t.Name == "base function"); + Assert.Contains(capturedTools, t => t.Name == "provider1 function"); + Assert.Contains(capturedTools, t => t.Name == "provider2 function"); + + // Both providers should have been invoked + mockProvider1 + .Protected() + .Verify>("InvokingCoreAsync", Times.Once(), ItExpr.IsAny(), ItExpr.IsAny()); + mockProvider2 + .Protected() + .Verify>("InvokingCoreAsync", Times.Once(), ItExpr.IsAny(), ItExpr.IsAny()); + + // Both providers should have been notified of success + mockProvider1 + .Protected() + .Verify("InvokedCoreAsync", Times.Once(), ItExpr.Is(x => + x.ResponseMessages == responseMessages && + x.InvokeException == null), ItExpr.IsAny()); + mockProvider2 + .Protected() + .Verify("InvokedCoreAsync", Times.Once(), ItExpr.Is(x => + x.ResponseMessages == responseMessages && + x.InvokeException == null), ItExpr.IsAny()); + } + + /// + /// Verify that RunAsync invokes InvokedCoreAsync on all AIContextProviders when the downstream GetResponse call fails. + /// + [Fact] + public async Task RunAsyncInvokesMultipleAIContextProvidersOnFailureAsync() + { + // Arrange + ChatMessage[] requestMessages = [new(ChatRole.User, "user message")]; + Mock mockService = new(); + mockService + .Setup(s => s.GetResponseAsync( + It.IsAny>(), + It.IsAny(), + It.IsAny())) + .ThrowsAsync(new InvalidOperationException("downstream failure")); + + var mockProvider1 = new Mock(); + mockProvider1.SetupGet(p => p.StateKey).Returns("Provider1"); + mockProvider1 + .Protected() + .Setup>("InvokingCoreAsync", ItExpr.IsAny(), ItExpr.IsAny()) + .Returns((AIContextProvider.InvokingContext ctx, CancellationToken _) => + new ValueTask(new AIContext + { + Messages = ctx.AIContext.Messages?.ToList(), + Instructions = ctx.AIContext.Instructions, + Tools = ctx.AIContext.Tools + })); + mockProvider1 + .Protected() + .Setup("InvokedCoreAsync", ItExpr.IsAny(), ItExpr.IsAny()) + .Returns(new ValueTask()); + + var mockProvider2 = new Mock(); + mockProvider2.SetupGet(p => p.StateKey).Returns("Provider2"); + mockProvider2 + .Protected() + .Setup>("InvokingCoreAsync", ItExpr.IsAny(), ItExpr.IsAny()) + .Returns((AIContextProvider.InvokingContext ctx, CancellationToken _) => + new ValueTask(new AIContext + { + Messages = ctx.AIContext.Messages?.ToList(), + Instructions = ctx.AIContext.Instructions, + Tools = ctx.AIContext.Tools + })); + mockProvider2 + .Protected() + .Setup("InvokedCoreAsync", ItExpr.IsAny(), ItExpr.IsAny()) + .Returns(new ValueTask()); + + ChatClientAgent agent = new(mockService.Object, options: new() + { + AIContextProviders = [mockProvider1.Object, mockProvider2.Object], + ChatOptions = new() { Instructions = "base instructions" } + }); + + // Act + await Assert.ThrowsAsync(() => agent.RunAsync(requestMessages)); + + // Assert - both providers should have been notified of the failure + mockProvider1 + .Protected() + .Verify>("InvokingCoreAsync", Times.Once(), ItExpr.IsAny(), ItExpr.IsAny()); + mockProvider2 + .Protected() + .Verify>("InvokingCoreAsync", Times.Once(), ItExpr.IsAny(), ItExpr.IsAny()); + + mockProvider1 + .Protected() + .Verify("InvokedCoreAsync", Times.Once(), ItExpr.Is(x => + x.InvokeException is InvalidOperationException), ItExpr.IsAny()); + mockProvider2 + .Protected() + .Verify("InvokedCoreAsync", Times.Once(), ItExpr.Is(x => + x.InvokeException is InvalidOperationException), ItExpr.IsAny()); + } + + /// + /// Verify that RunStreamingAsync invokes multiple AIContextProviders in sequence. + /// + [Fact] + public async Task RunStreamingAsyncInvokesMultipleAIContextProvidersAsync() + { + // Arrange + ChatMessage[] requestMessages = [new(ChatRole.User, "user message")]; + ChatResponseUpdate[] responseUpdates = [new(ChatRole.Assistant, "response")]; + Mock mockService = new(); + List capturedMessages = []; + string capturedInstructions = string.Empty; + mockService + .Setup(s => s.GetStreamingResponseAsync( + It.IsAny>(), + It.IsAny(), + It.IsAny())) + .Callback, ChatOptions, CancellationToken>((msgs, opts, ct) => + { + capturedMessages.AddRange(msgs); + capturedInstructions = opts.Instructions ?? string.Empty; + }) + .Returns(ToAsyncEnumerableAsync(responseUpdates)); + + var mockProvider1 = new Mock(); + mockProvider1.SetupGet(p => p.StateKey).Returns("Provider1"); + mockProvider1 + .Protected() + .Setup>("InvokingCoreAsync", ItExpr.IsAny(), ItExpr.IsAny()) + .Returns((AIContextProvider.InvokingContext ctx, CancellationToken _) => + new ValueTask(new AIContext + { + Messages = (ctx.AIContext.Messages ?? []).Concat([new ChatMessage(ChatRole.System, "provider1 context")]).ToList(), + Instructions = ctx.AIContext.Instructions + "\nprovider1 instructions", + Tools = ctx.AIContext.Tools + })); + mockProvider1 + .Protected() + .Setup("InvokedCoreAsync", ItExpr.IsAny(), ItExpr.IsAny()) + .Returns(new ValueTask()); + + var mockProvider2 = new Mock(); + mockProvider2.SetupGet(p => p.StateKey).Returns("Provider2"); + mockProvider2 + .Protected() + .Setup>("InvokingCoreAsync", ItExpr.IsAny(), ItExpr.IsAny()) + .Returns((AIContextProvider.InvokingContext ctx, CancellationToken _) => + new ValueTask(new AIContext + { + Messages = (ctx.AIContext.Messages ?? []).Concat([new ChatMessage(ChatRole.System, "provider2 context")]).ToList(), + Instructions = ctx.AIContext.Instructions + "\nprovider2 instructions", + Tools = ctx.AIContext.Tools + })); + mockProvider2 + .Protected() + .Setup("InvokedCoreAsync", ItExpr.IsAny(), ItExpr.IsAny()) + .Returns(new ValueTask()); + + ChatClientAgent agent = new( + mockService.Object, + options: new() + { + ChatOptions = new() { Instructions = "base instructions" }, + AIContextProviders = [mockProvider1.Object, mockProvider2.Object] + }); + + // Act + var session = await agent.CreateSessionAsync() as ChatClientAgentSession; + var updates = agent.RunStreamingAsync(requestMessages, session); + _ = await updates.ToAgentResponseAsync(); + + // Assert + Assert.Equal(3, capturedMessages.Count); + Assert.Equal("user message", capturedMessages[0].Text); + Assert.Equal("provider1 context", capturedMessages[1].Text); + Assert.Equal("provider2 context", capturedMessages[2].Text); + Assert.Equal("base instructions\nprovider1 instructions\nprovider2 instructions", capturedInstructions); + + // Both providers should have been invoked and notified + mockProvider1 + .Protected() + .Verify>("InvokingCoreAsync", Times.Once(), ItExpr.IsAny(), ItExpr.IsAny()); + mockProvider2 + .Protected() + .Verify>("InvokingCoreAsync", Times.Once(), ItExpr.IsAny(), ItExpr.IsAny()); + mockProvider1 + .Protected() + .Verify("InvokedCoreAsync", Times.Once(), ItExpr.Is(x => + x.InvokeException == null), ItExpr.IsAny()); + mockProvider2 + .Protected() + .Verify("InvokedCoreAsync", Times.Once(), ItExpr.Is(x => + x.InvokeException == null), ItExpr.IsAny()); + } + #endregion #region RunAsync Structured Output Tests @@ -1300,12 +1751,10 @@ public partial class ChatClientAgentTests It.IsAny>(), It.IsAny(), It.IsAny())).Returns(ToAsyncEnumerableAsync(returnUpdates)); - Mock>> mockFactory = new(); - mockFactory.Setup(f => f(It.IsAny(), It.IsAny())).ReturnsAsync(new InMemoryChatHistoryProvider()); ChatClientAgent agent = new(mockService.Object, options: new() { ChatOptions = new() { Instructions = "test instructions" }, - ChatHistoryProviderFactory = mockFactory.Object + ChatHistoryProvider = new InMemoryChatHistoryProvider() }); // Act @@ -1313,11 +1762,11 @@ public partial class ChatClientAgentTests await agent.RunStreamingAsync([new(ChatRole.User, "test")], session).ToListAsync(); // Assert - var chatHistoryProvider = Assert.IsType(session!.ChatHistoryProvider); - Assert.Equal(2, chatHistoryProvider.Count); - Assert.Equal("test", chatHistoryProvider[0].Text); - Assert.Equal("what?", chatHistoryProvider[1].Text); - mockFactory.Verify(f => f(It.IsAny(), It.IsAny()), Times.Once); + var chatHistoryProvider = Assert.IsType(agent.GetService(typeof(ChatHistoryProvider))); + var historyMessages = chatHistoryProvider.GetMessages(session); + Assert.Equal(2, historyMessages.Count); + Assert.Equal("test", historyMessages[0].Text); + Assert.Equal("what?", historyMessages[1].Text); } /// @@ -1360,10 +1809,10 @@ public partial class ChatClientAgentTests } /// - /// Verify that RunStreamingAsync throws when a factory is provided and the chat client returns a conversation id. + /// Verify that RunStreamingAsync throws when a is provided and the chat client returns a conversation id. /// [Fact] - public async Task RunStreamingAsyncThrowsWhenChatHistoryProviderFactoryProvidedAndConversationIdReturnedByChatClientAsync() + public async Task RunStreamingAsyncThrowsWhenChatHistoryProviderProvidedAndConversationIdReturnedByChatClientAsync() { // Arrange Mock mockService = new(); @@ -1377,18 +1826,16 @@ public partial class ChatClientAgentTests It.IsAny>(), It.IsAny(), It.IsAny())).Returns(ToAsyncEnumerableAsync(returnUpdates)); - Mock>> mockFactory = new(); - mockFactory.Setup(f => f(It.IsAny(), It.IsAny())).ReturnsAsync(new InMemoryChatHistoryProvider()); ChatClientAgent agent = new(mockService.Object, options: new() { ChatOptions = new() { Instructions = "test instructions" }, - ChatHistoryProviderFactory = mockFactory.Object + ChatHistoryProvider = new InMemoryChatHistoryProvider() }); // Act & Assert ChatClientAgentSession? session = await agent.CreateSessionAsync() as ChatClientAgentSession; var exception = await Assert.ThrowsAsync(async () => await agent.RunStreamingAsync([new(ChatRole.User, "test")], session).ToListAsync()); - Assert.Equal("Only the ConversationId or ChatHistoryProvider may be set, but not both and switching from one to another is not supported.", exception.Message); + Assert.Equal("Only ConversationId or ChatHistoryProvider may be used, but not both. The service returned a conversation id indicating server-side chat history management, but the agent has a ChatHistoryProvider configured.", exception.Message); } /// @@ -1425,12 +1872,13 @@ public partial class ChatClientAgentTests mockProvider .Protected() .Setup>("InvokingCoreAsync", ItExpr.IsAny(), ItExpr.IsAny()) - .ReturnsAsync(new AIContext - { - Messages = aiContextProviderMessages, - Instructions = "context provider instructions", - Tools = [AIFunctionFactory.Create(() => { }, "context provider function")] - }); + .Returns((AIContextProvider.InvokingContext ctx, CancellationToken _) => + new ValueTask(new AIContext + { + Messages = (ctx.AIContext.Messages ?? []).Concat(aiContextProviderMessages), + Instructions = ctx.AIContext.Instructions + "\ncontext provider instructions", + Tools = (ctx.AIContext.Tools ?? []).Concat(new[] { AIFunctionFactory.Create(() => { }, "context provider function") }) + })); mockProvider .Protected() .Setup("InvokedCoreAsync", ItExpr.IsAny(), ItExpr.IsAny()) @@ -1441,7 +1889,7 @@ public partial class ChatClientAgentTests options: new() { ChatOptions = new() { Instructions = "base instructions", Tools = [AIFunctionFactory.Create(() => { }, "base function")] }, - AIContextProviderFactory = (_, _) => new(mockProvider.Object) + AIContextProviders = [mockProvider.Object] }); // Act @@ -1462,11 +1910,13 @@ public partial class ChatClientAgentTests Assert.Contains(capturedTools, t => t.Name == "context provider function"); // Verify that the session was updated with the input, ai context provider, and response messages - var chatHistoryProvider = Assert.IsType(session!.ChatHistoryProvider); - Assert.Equal(3, chatHistoryProvider.Count); - Assert.Equal("user message", chatHistoryProvider[0].Text); - Assert.Equal("context provider message", chatHistoryProvider[1].Text); - Assert.Equal("response", chatHistoryProvider[2].Text); + var chatHistoryProvider = agent.ChatHistoryProvider as InMemoryChatHistoryProvider; + Assert.NotNull(chatHistoryProvider); + var historyMessages2 = chatHistoryProvider.GetMessages(session); + Assert.Equal(3, historyMessages2.Count); + Assert.Equal("user message", historyMessages2[0].Text); + Assert.Equal("context provider message", historyMessages2[1].Text); + Assert.Equal("response", historyMessages2[2].Text); mockProvider .Protected() @@ -1501,10 +1951,11 @@ public partial class ChatClientAgentTests mockProvider .Protected() .Setup>("InvokingCoreAsync", ItExpr.IsAny(), ItExpr.IsAny()) - .ReturnsAsync(new AIContext - { - Messages = aiContextProviderMessages, - }); + .Returns((AIContextProvider.InvokingContext ctx, CancellationToken _) => + new ValueTask(new AIContext + { + Messages = (ctx.AIContext.Messages ?? []).Concat(aiContextProviderMessages), + })); mockProvider .Protected() .Setup("InvokedCoreAsync", ItExpr.IsAny(), ItExpr.IsAny()) @@ -1515,7 +1966,7 @@ public partial class ChatClientAgentTests options: new() { ChatOptions = new() { Instructions = "base instructions", Tools = [AIFunctionFactory.Create(() => { }, "base function")] }, - AIContextProviderFactory = (_, _) => new(mockProvider.Object) + AIContextProviders = [mockProvider.Object] }); // Act @@ -1565,4 +2016,23 @@ public partial class ChatClientAgentTests [JsonSourceGenerationOptions(UseStringEnumConverter = true, PropertyNamingPolicy = JsonKnownNamingPolicy.CamelCase)] [JsonSerializable(typeof(Animal))] private sealed partial class JsonContext2 : JsonSerializerContext; + + private sealed class TestAIContextProvider(string stateKey) : AIContextProvider + { + public override string StateKey => stateKey; + + protected override ValueTask InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default) + => new(context.AIContext); + } + + private sealed class TestChatHistoryProvider(string stateKey) : ChatHistoryProvider + { + public override string StateKey => stateKey; + + protected override ValueTask> InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default) + => new(context.RequestMessages); + + protected override ValueTask InvokedCoreAsync(InvokedContext context, CancellationToken cancellationToken = default) + => default; + } } diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_BackgroundResponsesTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_BackgroundResponsesTests.cs index 2eed890292..84285ff9c4 100644 --- a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_BackgroundResponsesTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_BackgroundResponsesTests.cs @@ -339,6 +339,7 @@ public class ChatClientAgent_BackgroundResponsesTests // Create a mock chat history provider that would normally provide messages var mockChatHistoryProvider = new Mock(); + mockChatHistoryProvider.SetupGet(p => p.StateKey).Returns("ChatHistoryProvider"); mockChatHistoryProvider .Protected() .Setup>>("InvokingCoreAsync", ItExpr.IsAny(), ItExpr.IsAny()) @@ -346,6 +347,7 @@ public class ChatClientAgent_BackgroundResponsesTests // Create a mock AI context provider that would normally provide context var mockContextProvider = new Mock(); + mockContextProvider.SetupGet(p => p.StateKey).Returns("Provider1"); mockContextProvider .Protected() .Setup>("InvokingCoreAsync", ItExpr.IsAny(), ItExpr.IsAny()) @@ -365,14 +367,14 @@ public class ChatClientAgent_BackgroundResponsesTests capturedMessages.AddRange(msgs)) .ReturnsAsync(new ChatResponse([new(ChatRole.Assistant, "continued response")])); - ChatClientAgent agent = new(mockChatClient.Object); - - // Create a session with both chat history provider and AI context provider - ChatClientAgentSession? session = new() + ChatClientAgent agent = new(mockChatClient.Object, options: new() { ChatHistoryProvider = mockChatHistoryProvider.Object, - AIContextProvider = mockContextProvider.Object - }; + AIContextProviders = [mockContextProvider.Object] + }); + + // Create a session + ChatClientAgentSession? session = new(); AgentRunOptions runOptions = new() { @@ -406,6 +408,7 @@ public class ChatClientAgent_BackgroundResponsesTests // Create a mock chat history provider that would normally provide messages var mockChatHistoryProvider = new Mock(); + mockChatHistoryProvider.SetupGet(p => p.StateKey).Returns("ChatHistoryProvider"); mockChatHistoryProvider .Protected() .Setup>>("InvokingCoreAsync", ItExpr.IsAny(), ItExpr.IsAny()) @@ -413,6 +416,7 @@ public class ChatClientAgent_BackgroundResponsesTests // Create a mock AI context provider that would normally provide context var mockContextProvider = new Mock(); + mockContextProvider.SetupGet(p => p.StateKey).Returns("Provider1"); mockContextProvider .Protected() .Setup>("InvokingCoreAsync", ItExpr.IsAny(), ItExpr.IsAny()) @@ -432,14 +436,14 @@ public class ChatClientAgent_BackgroundResponsesTests capturedMessages.AddRange(msgs)) .Returns(ToAsyncEnumerableAsync([new ChatResponseUpdate(role: ChatRole.Assistant, content: "continued response")])); - ChatClientAgent agent = new(mockChatClient.Object); - - // Create a session with both chat history provider and AI context provider - ChatClientAgentSession? session = new() + ChatClientAgent agent = new(mockChatClient.Object, options: new() { ChatHistoryProvider = mockChatHistoryProvider.Object, - AIContextProvider = mockContextProvider.Object - }; + AIContextProviders = [mockContextProvider.Object] + }); + + // Create a session + ChatClientAgentSession? session = new(); AgentRunOptions runOptions = new() { @@ -633,10 +637,9 @@ public class ChatClientAgent_BackgroundResponsesTests It.IsAny())) .Returns(ToAsyncEnumerableAsync(returnUpdates)); - ChatClientAgent agent = new(mockChatClient.Object); - List capturedMessagesAddedToProvider = []; var mockChatHistoryProvider = new Mock(); + mockChatHistoryProvider.SetupGet(p => p.StateKey).Returns("ChatHistoryProvider"); mockChatHistoryProvider .Protected() .Setup("InvokedCoreAsync", ItExpr.IsAny(), ItExpr.IsAny()) @@ -645,17 +648,20 @@ public class ChatClientAgent_BackgroundResponsesTests AIContextProvider.InvokedContext? capturedInvokedContext = null; var mockContextProvider = new Mock(); + mockContextProvider.SetupGet(p => p.StateKey).Returns("Provider1"); mockContextProvider .Protected() .Setup("InvokedCoreAsync", ItExpr.IsAny(), ItExpr.IsAny()) .Callback((context, ct) => capturedInvokedContext = context) .Returns(new ValueTask()); - ChatClientAgentSession? session = new() + ChatClientAgent agent = new(mockChatClient.Object, options: new() { ChatHistoryProvider = mockChatHistoryProvider.Object, - AIContextProvider = mockContextProvider.Object - }; + AIContextProviders = [mockContextProvider.Object] + }); + + ChatClientAgentSession? session = new(); AgentRunOptions runOptions = new() { @@ -695,10 +701,9 @@ public class ChatClientAgent_BackgroundResponsesTests It.IsAny())) .Returns(ToAsyncEnumerableAsync(Array.Empty())); - ChatClientAgent agent = new(mockChatClient.Object); - List capturedMessagesAddedToProvider = []; var mockChatHistoryProvider = new Mock(); + mockChatHistoryProvider.SetupGet(p => p.StateKey).Returns("ChatHistoryProvider"); mockChatHistoryProvider .Protected() .Setup("InvokedCoreAsync", ItExpr.IsAny(), ItExpr.IsAny()) @@ -707,17 +712,20 @@ public class ChatClientAgent_BackgroundResponsesTests AIContextProvider.InvokedContext? capturedInvokedContext = null; var mockContextProvider = new Mock(); + mockContextProvider.SetupGet(p => p.StateKey).Returns("Provider1"); mockContextProvider .Protected() .Setup("InvokedCoreAsync", ItExpr.IsAny(), ItExpr.IsAny()) .Callback((context, ct) => capturedInvokedContext = context) .Returns(new ValueTask()); - ChatClientAgentSession? session = new() + ChatClientAgent agent = new(mockChatClient.Object, options: new() { ChatHistoryProvider = mockChatHistoryProvider.Object, - AIContextProvider = mockContextProvider.Object - }; + AIContextProviders = [mockContextProvider.Object] + }); + + ChatClientAgentSession? session = new(); AgentRunOptions runOptions = new() { diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_ChatHistoryManagementTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_ChatHistoryManagementTests.cs index 4de8f01f8e..4731805af7 100644 --- a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_ChatHistoryManagementTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_ChatHistoryManagementTests.cs @@ -163,17 +163,19 @@ public class ChatClientAgent_ChatHistoryManagementTests await agent.RunAsync([new(ChatRole.User, "test")], session); // Assert - InMemoryChatHistoryProvider chatHistoryProvider = Assert.IsType(session!.ChatHistoryProvider); - Assert.Equal(2, chatHistoryProvider.Count); - Assert.Equal("test", chatHistoryProvider[0].Text); - Assert.Equal("response", chatHistoryProvider[1].Text); + var inMemoryProvider = agent.ChatHistoryProvider as InMemoryChatHistoryProvider; + Assert.NotNull(inMemoryProvider); + var messages = inMemoryProvider.GetMessages(session!); + Assert.Equal(2, messages.Count); + Assert.Equal("test", messages[0].Text); + Assert.Equal("response", messages[1].Text); } /// - /// Verify that RunAsync uses the ChatHistoryProvider factory when the chat client returns no conversation id. + /// Verify that RunAsync uses the ChatHistoryProvider when the chat client returns no conversation id. /// [Fact] - public async Task RunAsync_UsesChatHistoryProviderFactory_WhenProvidedAndNoConversationIdReturnedByChatClientAsync() + public async Task RunAsync_UsesChatHistoryProvider_WhenProvidedAndNoConversationIdReturnedByChatClientAsync() { // Arrange Mock mockService = new(); @@ -187,19 +189,17 @@ public class ChatClientAgent_ChatHistoryManagementTests mockChatHistoryProvider .Protected() .Setup>>("InvokingCoreAsync", ItExpr.IsAny(), ItExpr.IsAny()) - .ReturnsAsync([new ChatMessage(ChatRole.User, "Existing Chat History")]); + .Returns((ChatHistoryProvider.InvokingContext ctx, CancellationToken _) => + new ValueTask>(new List { new(ChatRole.User, "Existing Chat History") }.Concat(ctx.RequestMessages).ToList())); mockChatHistoryProvider .Protected() .Setup("InvokedCoreAsync", ItExpr.IsAny(), ItExpr.IsAny()) .Returns(new ValueTask()); - Mock>> mockFactory = new(); - mockFactory.Setup(f => f(It.IsAny(), It.IsAny())).ReturnsAsync(mockChatHistoryProvider.Object); - ChatClientAgent agent = new(mockService.Object, options: new() { ChatOptions = new() { Instructions = "test instructions" }, - ChatHistoryProviderFactory = mockFactory.Object + ChatHistoryProvider = mockChatHistoryProvider.Object }); // Act @@ -207,7 +207,7 @@ public class ChatClientAgent_ChatHistoryManagementTests await agent.RunAsync([new(ChatRole.User, "test")], session); // Assert - Assert.IsType(session!.ChatHistoryProvider, exactMatch: false); + Assert.Same(mockChatHistoryProvider.Object, agent.ChatHistoryProvider); mockService.Verify( x => x.GetResponseAsync( It.Is>(msgs => msgs.Count() == 2 && msgs.Any(m => m.Text == "Existing Chat History") && msgs.Any(m => m.Text == "test")), @@ -222,9 +222,8 @@ public class ChatClientAgent_ChatHistoryManagementTests mockChatHistoryProvider .Protected() .Verify("InvokedCoreAsync", Times.Once(), - ItExpr.Is(x => x.RequestMessages.Count() == 1 && x.ResponseMessages!.Count() == 1), + ItExpr.Is(x => x.RequestMessages.Count() == 2 && x.ResponseMessages!.Count() == 1), ItExpr.IsAny()); - mockFactory.Verify(f => f(It.IsAny(), It.IsAny()), Times.Once); } /// @@ -242,14 +241,16 @@ public class ChatClientAgent_ChatHistoryManagementTests It.IsAny())).Throws(new InvalidOperationException("Test Error")); Mock mockChatHistoryProvider = new(); - - Mock>> mockFactory = new(); - mockFactory.Setup(f => f(It.IsAny(), It.IsAny())).ReturnsAsync(mockChatHistoryProvider.Object); + mockChatHistoryProvider + .Protected() + .Setup>>("InvokingCoreAsync", ItExpr.IsAny(), ItExpr.IsAny()) + .Returns((ChatHistoryProvider.InvokingContext ctx, CancellationToken _) => + new ValueTask>(ctx.RequestMessages.ToList())); ChatClientAgent agent = new(mockService.Object, options: new() { ChatOptions = new() { Instructions = "test instructions" }, - ChatHistoryProviderFactory = mockFactory.Object + ChatHistoryProvider = mockChatHistoryProvider.Object }); // Act @@ -257,20 +258,19 @@ public class ChatClientAgent_ChatHistoryManagementTests await Assert.ThrowsAsync(() => agent.RunAsync([new(ChatRole.User, "test")], session)); // Assert - Assert.IsType(session!.ChatHistoryProvider, exactMatch: false); + Assert.Same(mockChatHistoryProvider.Object, agent.ChatHistoryProvider); mockChatHistoryProvider .Protected() .Verify("InvokedCoreAsync", Times.Once(), ItExpr.Is(x => x.RequestMessages.Count() == 1 && x.ResponseMessages == null && x.InvokeException!.Message == "Test Error"), ItExpr.IsAny()); - mockFactory.Verify(f => f(It.IsAny(), It.IsAny()), Times.Once); } /// - /// Verify that RunAsync throws when a ChatHistoryProvider Factory is provided and the chat client returns a conversation id. + /// Verify that RunAsync throws when a ChatHistoryProvider is provided and the chat client returns a conversation id. /// [Fact] - public async Task RunAsync_Throws_WhenChatHistoryProviderFactoryProvidedAndConversationIdReturnedByChatClientAsync() + public async Task RunAsync_Throws_WhenChatHistoryProviderProvidedAndConversationIdReturnedByChatClientAsync() { // Arrange Mock mockService = new(); @@ -279,18 +279,16 @@ public class ChatClientAgent_ChatHistoryManagementTests It.IsAny>(), It.IsAny(), It.IsAny())).ReturnsAsync(new ChatResponse([new(ChatRole.Assistant, "response")]) { ConversationId = "ConvId" }); - Mock>> mockFactory = new(); - mockFactory.Setup(f => f(It.IsAny(), It.IsAny())).ReturnsAsync(new InMemoryChatHistoryProvider()); ChatClientAgent agent = new(mockService.Object, options: new() { ChatOptions = new() { Instructions = "test instructions" }, - ChatHistoryProviderFactory = mockFactory.Object + ChatHistoryProvider = new InMemoryChatHistoryProvider() }); // Act & Assert ChatClientAgentSession? session = await agent.CreateSessionAsync() as ChatClientAgentSession; InvalidOperationException exception = await Assert.ThrowsAsync(() => agent.RunAsync([new(ChatRole.User, "test")], session)); - Assert.Equal("Only the ConversationId or ChatHistoryProvider may be set, but not both and switching from one to another is not supported.", exception.Message); + Assert.Equal("Only ConversationId or ChatHistoryProvider may be used, but not both. The service returned a conversation id indicating server-side chat history management, but the agent has a ChatHistoryProvider configured.", exception.Message); } #endregion @@ -317,31 +315,29 @@ public class ChatClientAgent_ChatHistoryManagementTests mockOverrideChatHistoryProvider .Protected() .Setup>>("InvokingCoreAsync", ItExpr.IsAny(), ItExpr.IsAny()) - .ReturnsAsync([new ChatMessage(ChatRole.User, "Existing Chat History")]); + .Returns((ChatHistoryProvider.InvokingContext ctx, CancellationToken _) => + new ValueTask>(new List { new(ChatRole.User, "Existing Chat History") }.Concat(ctx.RequestMessages).ToList())); mockOverrideChatHistoryProvider .Protected() .Setup("InvokedCoreAsync", ItExpr.IsAny(), ItExpr.IsAny()) .Returns(new ValueTask()); - // Arrange a chat history provider to provide to the agent via a factory at construction time. + // Arrange a chat history provider to provide to the agent at construction time. // This one shouldn't be used since it is being overridden. - Mock mockFactoryChatHistoryProvider = new(); - mockFactoryChatHistoryProvider + Mock mockAgentOptionsChatHistoryProvider = new(); + mockAgentOptionsChatHistoryProvider .Protected() .Setup>>("InvokingCoreAsync", ItExpr.IsAny(), ItExpr.IsAny()) .ThrowsAsync(FailException.ForFailure("Base ChatHistoryProvider shouldn't be used.")); - mockFactoryChatHistoryProvider + mockAgentOptionsChatHistoryProvider .Protected() .Setup("InvokedCoreAsync", ItExpr.IsAny(), ItExpr.IsAny()) .Throws(FailException.ForFailure("Base ChatHistoryProvider shouldn't be used.")); - Mock>> mockFactory = new(); - mockFactory.Setup(f => f(It.IsAny(), It.IsAny())).ReturnsAsync(mockFactoryChatHistoryProvider.Object); - ChatClientAgent agent = new(mockService.Object, options: new() { ChatOptions = new() { Instructions = "test instructions" }, - ChatHistoryProviderFactory = mockFactory.Object + ChatHistoryProvider = mockAgentOptionsChatHistoryProvider.Object }); // Act @@ -351,7 +347,7 @@ public class ChatClientAgent_ChatHistoryManagementTests await agent.RunAsync([new(ChatRole.User, "test")], session, options: new AgentRunOptions { AdditionalProperties = additionalProperties }); // Assert - Assert.Same(mockFactoryChatHistoryProvider.Object, session!.ChatHistoryProvider); + Assert.Same(mockAgentOptionsChatHistoryProvider.Object, agent.ChatHistoryProvider); mockService.Verify( x => x.GetResponseAsync( It.Is>(msgs => msgs.Count() == 2 && msgs.Any(m => m.Text == "Existing Chat History") && msgs.Any(m => m.Text == "test")), @@ -366,15 +362,15 @@ public class ChatClientAgent_ChatHistoryManagementTests mockOverrideChatHistoryProvider .Protected() .Verify("InvokedCoreAsync", Times.Once(), - ItExpr.Is(x => x.RequestMessages.Count() == 1 && x.ResponseMessages!.Count() == 1), + ItExpr.Is(x => x.RequestMessages.Count() == 2 && x.ResponseMessages!.Count() == 1), ItExpr.IsAny()); - mockFactoryChatHistoryProvider + mockAgentOptionsChatHistoryProvider .Protected() .Verify>>("InvokingCoreAsync", Times.Never(), ItExpr.IsAny(), ItExpr.IsAny()); - mockFactoryChatHistoryProvider + mockAgentOptionsChatHistoryProvider .Protected() .Verify("InvokedCoreAsync", Times.Never(), ItExpr.IsAny(), diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_CreateSessionTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_CreateSessionTests.cs index 86220a6462..68fd008e9e 100644 --- a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_CreateSessionTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_CreateSessionTests.cs @@ -11,77 +11,6 @@ namespace Microsoft.Agents.AI.UnitTests; /// public class ChatClientAgent_CreateSessionTests { - [Fact] - public async Task CreateSession_UsesAIContextProviderFactory_IfProvidedAsync() - { - // Arrange - var mockChatClient = new Mock(); - var mockContextProvider = new Mock(); - var factoryCalled = false; - var agent = new ChatClientAgent(mockChatClient.Object, new ChatClientAgentOptions - { - ChatOptions = new() { Instructions = "Test instructions" }, - AIContextProviderFactory = (_, _) => - { - factoryCalled = true; - return new ValueTask(mockContextProvider.Object); - } - }); - - // Act - var session = await agent.CreateSessionAsync(); - - // Assert - Assert.True(factoryCalled, "AIContextProviderFactory was not called."); - Assert.IsType(session); - var typedSession = (ChatClientAgentSession)session; - Assert.Same(mockContextProvider.Object, typedSession.AIContextProvider); - } - - [Fact] - public async Task CreateSession_UsesChatHistoryProviderFactory_IfProvidedAsync() - { - // Arrange - var mockChatClient = new Mock(); - var mockChatHistoryProvider = new Mock(); - var factoryCalled = false; - var agent = new ChatClientAgent(mockChatClient.Object, new ChatClientAgentOptions - { - ChatOptions = new() { Instructions = "Test instructions" }, - ChatHistoryProviderFactory = (_, _) => - { - factoryCalled = true; - return new ValueTask(mockChatHistoryProvider.Object); - } - }); - - // Act - var session = await agent.CreateSessionAsync(); - - // Assert - Assert.True(factoryCalled, "ChatHistoryProviderFactory was not called."); - Assert.IsType(session); - var typedSession = (ChatClientAgentSession)session; - Assert.Same(mockChatHistoryProvider.Object, typedSession.ChatHistoryProvider); - } - - [Fact] - public async Task CreateSession_UsesChatHistoryProvider_FromTypedOverloadAsync() - { - // Arrange - var mockChatClient = new Mock(); - var mockChatHistoryProvider = new Mock(); - var agent = new ChatClientAgent(mockChatClient.Object); - - // Act - var session = await agent.CreateSessionAsync(mockChatHistoryProvider.Object); - - // Assert - Assert.IsType(session); - var typedSession = (ChatClientAgentSession)session; - Assert.Same(mockChatHistoryProvider.Object, typedSession.ChatHistoryProvider); - } - [Fact] public async Task CreateSession_UsesConversationId_FromTypedOverloadAsync() { diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_DeserializeSessionTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_DeserializeSessionTests.cs deleted file mode 100644 index 014cb1483b..0000000000 --- a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_DeserializeSessionTests.cs +++ /dev/null @@ -1,80 +0,0 @@ -// Copyright (c) Microsoft. All rights reserved. - -using System.Text.Json; -using System.Threading.Tasks; -using Microsoft.Extensions.AI; -using Moq; - -namespace Microsoft.Agents.AI.UnitTests; - -/// -/// Contains unit tests for the ChatClientAgent.DeserializeSession methods. -/// -public class ChatClientAgent_DeserializeSessionTests -{ - [Fact] - public async Task DeserializeSession_UsesAIContextProviderFactory_IfProvidedAsync() - { - // Arrange - var mockChatClient = new Mock(); - var mockContextProvider = new Mock(); - var factoryCalled = false; - var agent = new ChatClientAgent(mockChatClient.Object, new ChatClientAgentOptions - { - ChatOptions = new() { Instructions = "Test instructions" }, - AIContextProviderFactory = (_, _) => - { - factoryCalled = true; - return new ValueTask(mockContextProvider.Object); - } - }); - - var json = JsonSerializer.Deserialize(""" - { - "aiContextProviderState": ["CP1"] - } - """, TestJsonSerializerContext.Default.JsonElement); - - // Act - var session = await agent.DeserializeSessionAsync(json); - - // Assert - Assert.True(factoryCalled, "AIContextProviderFactory was not called."); - Assert.IsType(session); - var typedSession = (ChatClientAgentSession)session; - Assert.Same(mockContextProvider.Object, typedSession.AIContextProvider); - } - - [Fact] - public async Task DeserializeSession_UsesChatHistoryProviderFactory_IfProvidedAsync() - { - // Arrange - var mockChatClient = new Mock(); - var mockChatHistoryProvider = new Mock(); - var factoryCalled = false; - var agent = new ChatClientAgent(mockChatClient.Object, new ChatClientAgentOptions - { - ChatOptions = new() { Instructions = "Test instructions" }, - ChatHistoryProviderFactory = (_, _) => - { - factoryCalled = true; - return new ValueTask(mockChatHistoryProvider.Object); - } - }); - - var json = JsonSerializer.Deserialize(""" - { - "chatHistoryProviderState": { } - } - """, TestJsonSerializerContext.Default.JsonElement); - - // Act - var session = await agent.DeserializeSessionAsync(json); - - // Assert - Assert.True(factoryCalled, "ChatHistoryProviderFactory was not called."); - Assert.IsType(session); - var typedSession = (ChatClientAgentSession)session; - Assert.Same(mockChatHistoryProvider.Object, typedSession.ChatHistoryProvider); - } -} diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/Data/TextSearchProviderTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/Data/TextSearchProviderTests.cs index ec8dda3c45..602fe40e08 100644 --- a/dotnet/tests/Microsoft.Agents.AI.UnitTests/Data/TextSearchProviderTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/Data/TextSearchProviderTests.cs @@ -18,7 +18,6 @@ namespace Microsoft.Agents.AI.UnitTests.Data; public sealed class TextSearchProviderTests { private static readonly AIAgent s_mockAgent = new Mock().Object; - private static readonly AgentSession s_mockSession = new Mock().Object; private readonly Mock> _loggerMock; private readonly Mock _loggerFactoryMock; @@ -39,6 +38,28 @@ public sealed class TextSearchProviderTests .Returns(true); } + [Fact] + public void StateKey_ReturnsDefaultKey_WhenNoOptionsProvided() + { + // Arrange & Act + var provider = new TextSearchProvider((_, _) => Task.FromResult>([])); + + // Assert + Assert.Equal("TextSearchProvider", provider.StateKey); + } + + [Fact] + public void StateKey_ReturnsCustomKey_WhenSetViaOptions() + { + // Arrange & Act + var provider = new TextSearchProvider( + (_, _) => Task.FromResult>([]), + new TextSearchProviderOptions { StateKey = "custom-key" }); + + // Assert + Assert.Equal("custom-key", provider.StateKey); + } + [Theory] [InlineData(null, null, true)] [InlineData("Custom context prompt", "Custom citations prompt", false)] @@ -64,15 +85,19 @@ public sealed class TextSearchProviderTests ContextPrompt = overrideContextPrompt, CitationsPrompt = overrideCitationsPrompt }; - var provider = new TextSearchProvider(SearchDelegateAsync, default, null, options, withLogging ? this._loggerFactoryMock.Object : null); + var provider = new TextSearchProvider(SearchDelegateAsync, options, withLogging ? this._loggerFactoryMock.Object : null); var invokingContext = new AIContextProvider.InvokingContext( s_mockAgent, - s_mockSession, - [ - new ChatMessage(ChatRole.User, "Sample user question?"), - new ChatMessage(ChatRole.User, "Additional part") - ]); + new TestAgentSession(), + new AIContext + { + Messages = new List + { + new(ChatRole.User, "Sample user question?"), + new(ChatRole.User, "Additional part") + } + }); // Act var aiContext = await provider.InvokingAsync(invokingContext, CancellationToken.None); @@ -81,9 +106,15 @@ public sealed class TextSearchProviderTests Assert.Equal("Sample user question?\nAdditional part", capturedInput); Assert.Null(aiContext.Instructions); // TextSearchProvider uses a user message for context injection. Assert.NotNull(aiContext.Messages); - Assert.Single(aiContext.Messages!); - var message = aiContext.Messages!.Single(); + var messages = aiContext.Messages!.ToList(); + Assert.Equal(3, messages.Count); // 2 input messages + 1 search result message + Assert.Equal("Sample user question?", messages[0].Text); + Assert.Equal("Additional part", messages[1].Text); + Assert.Equal(AgentRequestMessageSourceType.External, messages[0].GetAgentRequestMessageSourceType()); + Assert.Equal(AgentRequestMessageSourceType.External, messages[1].GetAgentRequestMessageSourceType()); + var message = messages.Last(); Assert.Equal(ChatRole.User, message.Role); + Assert.Equal(AgentRequestMessageSourceType.AIContextProvider, message.GetAgentRequestMessageSourceType()); string text = message.Text!; if (overrideContextPrompt is null) @@ -143,17 +174,21 @@ public sealed class TextSearchProviderTests FunctionToolName = overrideName, FunctionToolDescription = overrideDescription }; - var provider = new TextSearchProvider(this.NoResultSearchAsync, default, null, options); - var invokingContext = new AIContextProvider.InvokingContext(s_mockAgent, s_mockSession, [new ChatMessage(ChatRole.User, "Q?")]); + var provider = new TextSearchProvider(this.NoResultSearchAsync, options); + var invokingContext = new AIContextProvider.InvokingContext(s_mockAgent, new TestAgentSession(), new AIContext { Messages = new List { new(ChatRole.User, "Q?") } }); // Act var aiContext = await provider.InvokingAsync(invokingContext, CancellationToken.None); // Assert - Assert.Null(aiContext.Messages); // No automatic injection. + Assert.NotNull(aiContext.Messages); // Input messages are preserved. + var messages = aiContext.Messages!.ToList(); + Assert.Single(messages); + Assert.Equal("Q?", messages[0].Text); Assert.NotNull(aiContext.Tools); - Assert.Single(aiContext.Tools); - var tool = aiContext.Tools.Single(); + var tools = aiContext.Tools!.ToList(); + Assert.Single(tools); + var tool = tools[0]; Assert.Equal(expectedName, tool.Name); Assert.Equal(expectedDescription, tool.Description); } @@ -162,14 +197,17 @@ public sealed class TextSearchProviderTests public async Task InvokingAsync_ShouldNotThrow_WhenSearchFailsAsync() { // Arrange - var provider = new TextSearchProvider(this.FailingSearchAsync, default, null, loggerFactory: this._loggerFactoryMock.Object); - var invokingContext = new AIContextProvider.InvokingContext(s_mockAgent, s_mockSession, [new ChatMessage(ChatRole.User, "Q?")]); + var provider = new TextSearchProvider(this.FailingSearchAsync, loggerFactory: this._loggerFactoryMock.Object); + var invokingContext = new AIContextProvider.InvokingContext(s_mockAgent, new TestAgentSession(), new AIContext { Messages = new List { new(ChatRole.User, "Q?") } }); // Act var aiContext = await provider.InvokingAsync(invokingContext, CancellationToken.None); // Assert - Assert.Null(aiContext.Messages); + Assert.NotNull(aiContext.Messages); // Input messages are preserved on error. + var messages = aiContext.Messages!.ToList(); + Assert.Single(messages); + Assert.Equal("Q?", messages[0].Text); Assert.Null(aiContext.Tools); this._loggerMock.Verify( l => l.Log( @@ -203,7 +241,7 @@ public sealed class TextSearchProviderTests ContextPrompt = overrideContextPrompt, CitationsPrompt = overrideCitationsPrompt }; - var provider = new TextSearchProvider(SearchDelegateAsync, default, null, options); + var provider = new TextSearchProvider(SearchDelegateAsync, options); // Act var formatted = await provider.SearchAsync("Sample user question?", CancellationToken.None); @@ -255,16 +293,18 @@ public sealed class TextSearchProviderTests SearchTime = TextSearchProviderOptions.TextSearchBehavior.BeforeAIInvoke, ContextFormatter = r => $"Custom formatted context with {r.Count} results." }; - var provider = new TextSearchProvider(SearchDelegateAsync, default, null, options); - var invokingContext = new AIContextProvider.InvokingContext(s_mockAgent, s_mockSession, [new ChatMessage(ChatRole.User, "Q?")]); + var provider = new TextSearchProvider(SearchDelegateAsync, options); + var invokingContext = new AIContextProvider.InvokingContext(s_mockAgent, new TestAgentSession(), new AIContext { Messages = new List { new(ChatRole.User, "Q?") } }); // Act var aiContext = await provider.InvokingAsync(invokingContext, CancellationToken.None); // Assert Assert.NotNull(aiContext.Messages); - Assert.Single(aiContext.Messages!); - Assert.Equal("Custom formatted context with 2 results.", aiContext.Messages![0].Text); + var messages = aiContext.Messages!.ToList(); + Assert.Equal(2, messages.Count); // 1 input message + 1 formatted result message + Assert.Equal("Q?", messages[0].Text); + Assert.Equal("Custom formatted context with 2 results.", messages[1].Text); } [Fact] @@ -289,16 +329,18 @@ public sealed class TextSearchProviderTests SearchTime = TextSearchProviderOptions.TextSearchBehavior.BeforeAIInvoke, ContextFormatter = r => string.Join(",", r.Select(x => ((RawPayload)x.RawRepresentation!).Id)) }; - var provider = new TextSearchProvider(SearchDelegateAsync, default, null, options); - var invokingContext = new AIContextProvider.InvokingContext(s_mockAgent, s_mockSession, [new ChatMessage(ChatRole.User, "Q?")]); + var provider = new TextSearchProvider(SearchDelegateAsync, options); + var invokingContext = new AIContextProvider.InvokingContext(s_mockAgent, new TestAgentSession(), new AIContext { Messages = new List { new(ChatRole.User, "Q?") } }); // Act var aiContext = await provider.InvokingAsync(invokingContext, CancellationToken.None); // Assert Assert.NotNull(aiContext.Messages); - Assert.Single(aiContext.Messages!); - Assert.Equal("R1,R2", aiContext.Messages![0].Text); + var messages = aiContext.Messages!.ToList(); + Assert.Equal(2, messages.Count); // 1 input message + 1 formatted result message + Assert.Equal("Q?", messages[0].Text); + Assert.Equal("R1,R2", messages[1].Text); } [Fact] @@ -306,18 +348,155 @@ public sealed class TextSearchProviderTests { // Arrange var options = new TextSearchProviderOptions { SearchTime = TextSearchProviderOptions.TextSearchBehavior.BeforeAIInvoke }; - var provider = new TextSearchProvider(this.NoResultSearchAsync, default, null, options); - var invokingContext = new AIContextProvider.InvokingContext(s_mockAgent, s_mockSession, [new ChatMessage(ChatRole.User, "Q?")]); + var provider = new TextSearchProvider(this.NoResultSearchAsync, options); + var invokingContext = new AIContextProvider.InvokingContext(s_mockAgent, new TestAgentSession(), new AIContext { Messages = new List { new(ChatRole.User, "Q?") } }); // Act var aiContext = await provider.InvokingAsync(invokingContext, CancellationToken.None); // Assert - Assert.Null(aiContext.Messages); + Assert.NotNull(aiContext.Messages); // Input messages are preserved when no results found. + var messages = aiContext.Messages!.ToList(); + Assert.Single(messages); + Assert.Equal("Q?", messages[0].Text); Assert.Null(aiContext.Instructions); Assert.Null(aiContext.Tools); } + #region Message Filter Tests + + [Fact] + public async Task InvokingAsync_DefaultFilter_ExcludesNonExternalMessagesFromSearchInputAsync() + { + // Arrange + string? capturedInput = null; + Task> SearchDelegateAsync(string input, CancellationToken ct) + { + capturedInput = input; + return Task.FromResult>([]); + } + + var provider = new TextSearchProvider(SearchDelegateAsync); + var requestMessages = new List + { + new(ChatRole.User, "External message"), + new(ChatRole.System, "From history") { AdditionalProperties = new() { { AgentRequestMessageSourceAttribution.AdditionalPropertiesKey, new AgentRequestMessageSourceAttribution(AgentRequestMessageSourceType.ChatHistory, "HistorySource") } } }, + new(ChatRole.System, "From context provider") { AdditionalProperties = new() { { AgentRequestMessageSourceAttribution.AdditionalPropertiesKey, new AgentRequestMessageSourceAttribution(AgentRequestMessageSourceType.AIContextProvider, "ContextSource") } } }, + }; + + var invokingContext = new AIContextProvider.InvokingContext(s_mockAgent, new TestAgentSession(), new AIContext { Messages = requestMessages }); + + // Act + await provider.InvokingAsync(invokingContext, CancellationToken.None); + + // Assert - Only external messages should be used for search input + Assert.Equal("External message", capturedInput); + } + + [Fact] + public async Task InvokingAsync_CustomSearchInputFilter_OverridesDefaultAsync() + { + // Arrange + string? capturedInput = null; + Task> SearchDelegateAsync(string input, CancellationToken ct) + { + capturedInput = input; + return Task.FromResult>([]); + } + + var provider = new TextSearchProvider(SearchDelegateAsync, new TextSearchProviderOptions + { + SearchInputMessageFilter = messages => messages.Where(m => m.Role == ChatRole.System) + }); + var requestMessages = new List + { + new(ChatRole.User, "User message"), + new(ChatRole.System, "System message"), + }; + + var invokingContext = new AIContextProvider.InvokingContext(s_mockAgent, new TestAgentSession(), new AIContext { Messages = requestMessages }); + + // Act + await provider.InvokingAsync(invokingContext, CancellationToken.None); + + // Assert - Custom filter keeps only System messages + Assert.Equal("System message", capturedInput); + } + + [Fact] + public async Task InvokedAsync_DefaultFilter_ExcludesNonExternalMessagesFromStorageAsync() + { + // Arrange + var options = new TextSearchProviderOptions + { + RecentMessageMemoryLimit = 10, + RecentMessageRolesIncluded = [ChatRole.User, ChatRole.System] + }; + string? capturedInput = null; + Task> SearchDelegateAsync(string input, CancellationToken ct) + { + capturedInput = input; + return Task.FromResult>([]); + } + var provider = new TextSearchProvider(SearchDelegateAsync, options); + var session = new TestAgentSession(); + + var requestMessages = new List + { + new(ChatRole.User, "External message"), + new(ChatRole.System, "From history") { AdditionalProperties = new() { { AgentRequestMessageSourceAttribution.AdditionalPropertiesKey, new AgentRequestMessageSourceAttribution(AgentRequestMessageSourceType.ChatHistory, "HistorySource") } } }, + new(ChatRole.System, "From context provider") { AdditionalProperties = new() { { AgentRequestMessageSourceAttribution.AdditionalPropertiesKey, new AgentRequestMessageSourceAttribution(AgentRequestMessageSourceType.AIContextProvider, "ContextSource") } } }, + }; + + // Store messages via InvokedAsync + await provider.InvokedAsync(new(s_mockAgent, session, requestMessages, [])); + + // Now invoke to read stored memory + var invokingContext = new AIContextProvider.InvokingContext(s_mockAgent, session, new AIContext { Messages = [new ChatMessage(ChatRole.User, "Next")] }); + await provider.InvokingAsync(invokingContext, CancellationToken.None); + + // Assert - Only "External message" was stored in memory, so search input = "External message" + "Next" + Assert.Equal("External message\nNext", capturedInput); + } + + [Fact] + public async Task InvokedAsync_CustomStorageInputFilter_OverridesDefaultAsync() + { + // Arrange + var options = new TextSearchProviderOptions + { + RecentMessageMemoryLimit = 10, + RecentMessageRolesIncluded = [ChatRole.User, ChatRole.System], + StorageInputMessageFilter = messages => messages // No filtering - store everything + }; + string? capturedInput = null; + Task> SearchDelegateAsync(string input, CancellationToken ct) + { + capturedInput = input; + return Task.FromResult>([]); + } + var provider = new TextSearchProvider(SearchDelegateAsync, options); + var session = new TestAgentSession(); + + var requestMessages = new List + { + new(ChatRole.User, "External message"), + new(ChatRole.System, "From history") { AdditionalProperties = new() { { AgentRequestMessageSourceAttribution.AdditionalPropertiesKey, new AgentRequestMessageSourceAttribution(AgentRequestMessageSourceType.ChatHistory, "HistorySource") } } }, + }; + + // Store messages via InvokedAsync + await provider.InvokedAsync(new(s_mockAgent, session, requestMessages, [])); + + // Now invoke to read stored memory + var invokingContext = new AIContextProvider.InvokingContext(s_mockAgent, session, new AIContext { Messages = [new ChatMessage(ChatRole.User, "Next")] }); + await provider.InvokingAsync(invokingContext, CancellationToken.None); + + // Assert - Both messages stored (identity filter), so search input includes all + current + Assert.Equal("External message\nFrom history\nNext", capturedInput); + } + + #endregion + #region Recent Message Memory Tests [Fact] @@ -335,7 +514,7 @@ public sealed class TextSearchProviderTests capturedInput = input; return Task.FromResult>([]); // No results needed. } - var provider = new TextSearchProvider(SearchDelegateAsync, default, null, options); + var provider = new TextSearchProvider(SearchDelegateAsync, options); // Populate memory with more messages than the limit (A,B,C,D) -> should retain B,C,D var initialMessages = new[] @@ -345,14 +524,14 @@ public sealed class TextSearchProviderTests new ChatMessage(ChatRole.User, "C"), new ChatMessage(ChatRole.Assistant, "D"), }; - await provider.InvokedAsync(new(s_mockAgent, s_mockSession, initialMessages) { InvokeException = new InvalidOperationException("Request Failed") }); + + var session = new TestAgentSession(); + await provider.InvokedAsync(new(s_mockAgent, session, initialMessages, new InvalidOperationException("Request Failed"))); var invokingContext = new AIContextProvider.InvokingContext( s_mockAgent, - s_mockSession, - [ - new ChatMessage(ChatRole.User, "E") - ]); + session, + new AIContext { Messages = new List { new(ChatRole.User, "E") } }); // Act await provider.InvokingAsync(invokingContext, CancellationToken.None); @@ -377,7 +556,8 @@ public sealed class TextSearchProviderTests capturedInput = input; return Task.FromResult>([]); // No results needed. } - var provider = new TextSearchProvider(SearchDelegateAsync, default, null, options); + var provider = new TextSearchProvider(SearchDelegateAsync, options); + var session = new TestAgentSession(); // Populate memory with more messages than the limit (A,B,C,D) -> should retain B,C,D var initialMessages = new[] @@ -387,14 +567,12 @@ public sealed class TextSearchProviderTests new ChatMessage(ChatRole.User, "C"), new ChatMessage(ChatRole.Assistant, "D"), }; - await provider.InvokedAsync(new(s_mockAgent, s_mockSession, initialMessages)); + await provider.InvokedAsync(new(s_mockAgent, session, initialMessages, [])); var invokingContext = new AIContextProvider.InvokingContext( s_mockAgent, - s_mockSession, - [ - new ChatMessage(ChatRole.User, "E") - ]); + session, + new AIContext { Messages = new List { new(ChatRole.User, "E") } }); // Act await provider.InvokingAsync(invokingContext, CancellationToken.None); @@ -419,28 +597,31 @@ public sealed class TextSearchProviderTests capturedInput = input; return Task.FromResult>([]); } - var provider = new TextSearchProvider(SearchDelegateAsync, default, null, options); + var provider = new TextSearchProvider(SearchDelegateAsync, options); + var session = new TestAgentSession(); // First memory update (A,B) await provider.InvokedAsync(new( - s_mockAgent, - s_mockSession, - [ - new ChatMessage(ChatRole.User, "A"), - new ChatMessage(ChatRole.Assistant, "B"), - ])); + s_mockAgent, + session, + [ + new ChatMessage(ChatRole.User, "A"), + new ChatMessage(ChatRole.Assistant, "B"), + ], + [])); // Second memory update (C,D,E) await provider.InvokedAsync(new( - s_mockAgent, - s_mockSession, - [ - new ChatMessage(ChatRole.User, "C"), - new ChatMessage(ChatRole.Assistant, "D"), - new ChatMessage(ChatRole.User, "E"), - ])); + s_mockAgent, + session, + [ + new ChatMessage(ChatRole.User, "C"), + new ChatMessage(ChatRole.Assistant, "D"), + new ChatMessage(ChatRole.User, "E"), + ], + [])); - var invokingContext = new AIContextProvider.InvokingContext(s_mockAgent, s_mockSession, [new ChatMessage(ChatRole.User, "F")]); + var invokingContext = new AIContextProvider.InvokingContext(s_mockAgent, session, new AIContext { Messages = new List { new(ChatRole.User, "F") } }); // Act await provider.InvokingAsync(invokingContext, CancellationToken.None); @@ -465,7 +646,8 @@ public sealed class TextSearchProviderTests capturedInput = input; return Task.FromResult>([]); // No results needed for this test. } - var provider = new TextSearchProvider(SearchDelegateAsync, default, null, options); + var provider = new TextSearchProvider(SearchDelegateAsync, options); + var session = new TestAgentSession(); // Populate memory with mixed roles; only Assistant messages (A1,A2) should be retained. var initialMessages = new[] @@ -475,14 +657,12 @@ public sealed class TextSearchProviderTests new ChatMessage(ChatRole.User, "U2"), new ChatMessage(ChatRole.Assistant, "A2"), }; - await provider.InvokedAsync(new(s_mockAgent, s_mockSession, initialMessages)); + await provider.InvokedAsync(new(s_mockAgent, session, initialMessages, [])); var invokingContext = new AIContextProvider.InvokingContext( s_mockAgent, - s_mockSession, - [ - new ChatMessage(ChatRole.User, "Question?") // Current request message always appended. - ]); + session, + new AIContext { Messages = new List { new(ChatRole.User, "Question?") } }); // Current request message always appended. // Act await provider.InvokingAsync(invokingContext, CancellationToken.None); @@ -496,26 +676,7 @@ public sealed class TextSearchProviderTests #region Serialization Tests [Fact] - public void Serialize_WithNoRecentMessages_ShouldReturnEmptyState() - { - // Arrange - var options = new TextSearchProviderOptions - { - SearchTime = TextSearchProviderOptions.TextSearchBehavior.BeforeAIInvoke, - RecentMessageMemoryLimit = 3 - }; - var provider = new TextSearchProvider(this.NoResultSearchAsync, default, null, options); - - // Act - var state = provider.Serialize(); - - // Assert - Assert.Equal(JsonValueKind.Object, state.ValueKind); - Assert.False(state.TryGetProperty("recentMessagesText", out _)); - } - - [Fact] - public async Task Serialize_WithRecentMessages_ShouldPersistMessagesUpToLimitAsync() + public async Task InvokedAsync_ShouldPersistMessagesToSessionStateBagAsync() { // Arrange var options = new TextSearchProviderOptions @@ -524,7 +685,8 @@ public sealed class TextSearchProviderTests RecentMessageMemoryLimit = 3, RecentMessageRolesIncluded = [ChatRole.User, ChatRole.Assistant] }; - var provider = new TextSearchProvider(this.NoResultSearchAsync, default, null, options); + var provider = new TextSearchProvider(this.NoResultSearchAsync, options); + var session = new TestAgentSession(); var messages = new[] { new ChatMessage(ChatRole.User, "M1"), @@ -533,11 +695,12 @@ public sealed class TextSearchProviderTests }; // Act - await provider.InvokedAsync(new(s_mockAgent, s_mockSession, messages)); // Populate recent memory. - var state = provider.Serialize(); + await provider.InvokedAsync(new(s_mockAgent, session, messages, [])); // Populate recent memory. - // Assert - Assert.True(state.TryGetProperty("recentMessagesText", out var recentProperty)); + // Assert - State should be in the session's StateBag + var stateBagSerialized = session.StateBag.Serialize(); + Assert.True(stateBagSerialized.TryGetProperty("TextSearchProvider", out var stateProperty)); + Assert.True(stateProperty.TryGetProperty("recentMessagesText", out var recentProperty)); Assert.Equal(JsonValueKind.Array, recentProperty.ValueKind); var list = recentProperty.EnumerateArray().Select(e => e.GetString()).ToList(); Assert.Equal(3, list.Count); @@ -545,7 +708,7 @@ public sealed class TextSearchProviderTests } [Fact] - public async Task SerializeAndDeserialize_RoundtripRestoresMessagesAsync() + public async Task StateBag_RoundtripRestoresMessagesAsync() { // Arrange var options = new TextSearchProviderOptions @@ -554,7 +717,8 @@ public sealed class TextSearchProviderTests RecentMessageMemoryLimit = 4, RecentMessageRolesIncluded = [ChatRole.User, ChatRole.Assistant] }; - var provider = new TextSearchProvider(this.NoResultSearchAsync, default, null, options); + var provider = new TextSearchProvider(this.NoResultSearchAsync, options); + var session = new TestAgentSession(); var messages = new[] { new ChatMessage(ChatRole.User, "A"), @@ -562,23 +726,24 @@ public sealed class TextSearchProviderTests new ChatMessage(ChatRole.User, "C"), new ChatMessage(ChatRole.Assistant, "D"), }; - await provider.InvokedAsync(new(s_mockAgent, s_mockSession, messages)); + await provider.InvokedAsync(new(s_mockAgent, session, messages, [])); + + // Act - Serialize and deserialize the StateBag + var serializedStateBag = session.StateBag.Serialize(); + var restoredSession = new TestAgentSession(AgentSessionStateBag.Deserialize(serializedStateBag)); - // Act - var state = provider.Serialize(); string? capturedInput = null; Task> SearchDelegate2Async(string input, CancellationToken ct) { capturedInput = input; return Task.FromResult>([]); } - var roundTrippedProvider = new TextSearchProvider(SearchDelegate2Async, state, options: new TextSearchProviderOptions + var newProvider = new TextSearchProvider(SearchDelegate2Async, new TextSearchProviderOptions { SearchTime = TextSearchProviderOptions.TextSearchBehavior.BeforeAIInvoke, RecentMessageMemoryLimit = 4 }); - var emptyMessages = Array.Empty(); - await roundTrippedProvider.InvokingAsync(new(s_mockAgent, s_mockSession, emptyMessages), CancellationToken.None); // Trigger search to read memory. + await newProvider.InvokingAsync(new(s_mockAgent, restoredSession, new AIContext()), CancellationToken.None); // Trigger search to read memory. // Assert Assert.NotNull(capturedInput); @@ -586,25 +751,10 @@ public sealed class TextSearchProviderTests } [Fact] - public async Task Deserialize_WithChangedLowerLimit_ShouldTruncateToNewLimitAsync() + public async Task InvokingAsync_WithEmptyStateBag_ShouldHaveNoMessagesAsync() { // Arrange - var initialProvider = new TextSearchProvider(this.NoResultSearchAsync, default, null, new TextSearchProviderOptions - { - SearchTime = TextSearchProviderOptions.TextSearchBehavior.BeforeAIInvoke, - RecentMessageMemoryLimit = 5, - RecentMessageRolesIncluded = [ChatRole.User, ChatRole.Assistant] - }); - var messages = new[] - { - new ChatMessage(ChatRole.User, "L1"), - new ChatMessage(ChatRole.Assistant, "L2"), - new ChatMessage(ChatRole.User, "L3"), - new ChatMessage(ChatRole.Assistant, "L4"), - new ChatMessage(ChatRole.User, "L5"), - }; - await initialProvider.InvokedAsync(new(s_mockAgent, s_mockSession, messages)); - var state = initialProvider.Serialize(); + var session = new TestAgentSession(); // Fresh session with empty StateBag string? capturedInput = null; Task> SearchDelegate2Async(string input, CancellationToken ct) @@ -614,43 +764,16 @@ public sealed class TextSearchProviderTests } // Act - var restoredProvider = new TextSearchProvider(SearchDelegate2Async, state, options: new TextSearchProviderOptions - { - SearchTime = TextSearchProviderOptions.TextSearchBehavior.BeforeAIInvoke, - RecentMessageMemoryLimit = 3 // Lower limit - }); - await restoredProvider.InvokingAsync(new(s_mockAgent, s_mockSession, Array.Empty()), CancellationToken.None); - - // Assert - Assert.NotNull(capturedInput); - Assert.Equal("L1\nL2\nL3", capturedInput); - } - - [Fact] - public async Task Deserialize_WithEmptyState_ShouldHaveNoMessagesAsync() - { - // Arrange - var emptyState = JsonSerializer.Deserialize("{}", TestJsonSerializerContext.Default.JsonElement); - - string? capturedInput = null; - Task> SearchDelegate2Async(string input, CancellationToken ct) - { - capturedInput = input; - return Task.FromResult>([]); - } - - // Act - var provider = new TextSearchProvider(SearchDelegate2Async, emptyState, options: new TextSearchProviderOptions + var provider = new TextSearchProvider(SearchDelegate2Async, new TextSearchProviderOptions { SearchTime = TextSearchProviderOptions.TextSearchBehavior.BeforeAIInvoke, RecentMessageMemoryLimit = 3 }); - var emptyMessages = Array.Empty(); - await provider.InvokingAsync(new(s_mockAgent, s_mockSession, emptyMessages), CancellationToken.None); + await provider.InvokingAsync(new(s_mockAgent, session, new AIContext()), CancellationToken.None); // Assert Assert.NotNull(capturedInput); - Assert.Equal(string.Empty, capturedInput); // No recent messages serialized => empty input. + Assert.Equal(string.Empty, capturedInput); // No recent messages in StateBag => empty input. } #endregion @@ -669,4 +792,16 @@ public sealed class TextSearchProviderTests { public string Id { get; set; } = string.Empty; } + + private sealed class TestAgentSession : AgentSession + { + public TestAgentSession() + { + } + + public TestAgentSession(AgentSessionStateBag stateBag) + { + this.StateBag = stateBag; + } + } } diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/Memory/ChatHistoryMemoryProviderTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/Memory/ChatHistoryMemoryProviderTests.cs index 0a11e74528..be73260b15 100644 --- a/dotnet/tests/Microsoft.Agents.AI.UnitTests/Memory/ChatHistoryMemoryProviderTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/Memory/ChatHistoryMemoryProviderTests.cs @@ -3,7 +3,6 @@ using System; using System.Collections.Generic; using System.Linq; -using System.Text.Json; using System.Threading; using System.Threading.Tasks; using Microsoft.Extensions.AI; @@ -19,7 +18,6 @@ namespace Microsoft.Agents.AI.Memory.UnitTests; public class ChatHistoryMemoryProviderTests { private static readonly AIAgent s_mockAgent = new Mock().Object; - private static readonly AgentSession s_mockSession = new Mock().Object; private readonly Mock> _loggerMock; private readonly Mock _loggerFactoryMock; @@ -57,33 +55,82 @@ public class ChatHistoryMemoryProviderTests .Returns(this._vectorStoreCollectionMock.Object); } + [Fact] + public void StateKey_ReturnsDefaultKey_WhenNoOptionsProvided() + { + // Arrange & Act + var provider = new ChatHistoryMemoryProvider( + this._vectorStoreMock.Object, + TestCollectionName, + 1, + _ => new ChatHistoryMemoryProvider.State(new ChatHistoryMemoryProviderScope { UserId = "UID" })); + + // Assert + Assert.Equal("ChatHistoryMemoryProvider", provider.StateKey); + } + + [Fact] + public void StateKey_ReturnsCustomKey_WhenSetViaOptions() + { + // Arrange & Act + var provider = new ChatHistoryMemoryProvider( + this._vectorStoreMock.Object, + TestCollectionName, + 1, + _ => new ChatHistoryMemoryProvider.State(new ChatHistoryMemoryProviderScope { UserId = "UID" }), + new ChatHistoryMemoryProviderOptions { StateKey = "custom-key" }); + + // Assert + Assert.Equal("custom-key", provider.StateKey); + } + [Fact] public void Constructor_Throws_ForNullVectorStore() { // Act & Assert - Assert.Throws(() => new ChatHistoryMemoryProvider(null!, "testcollection", 1, new ChatHistoryMemoryProviderScope() { UserId = "UID" })); + Assert.Throws(() => new ChatHistoryMemoryProvider( + null!, + "testcollection", + 1, + _ => new ChatHistoryMemoryProvider.State(new ChatHistoryMemoryProviderScope { UserId = "UID" }))); } [Fact] public void Constructor_Throws_ForNullCollectionName() { // Act & Assert - Assert.Throws(() => new ChatHistoryMemoryProvider(this._vectorStoreMock.Object, null!, 1, new ChatHistoryMemoryProviderScope() { UserId = "UID" })); + Assert.Throws(() => new ChatHistoryMemoryProvider( + this._vectorStoreMock.Object, + null!, + 1, + _ => new ChatHistoryMemoryProvider.State(new ChatHistoryMemoryProviderScope { UserId = "UID" }))); } [Fact] - public void Constructor_Throws_ForNullStorageScope() + public void Constructor_Throws_ForNullStateInitializer() { // Act & Assert - Assert.Throws(() => new ChatHistoryMemoryProvider(this._vectorStoreMock.Object, "testcollection", 1, null!)); + Assert.Throws(() => new ChatHistoryMemoryProvider( + this._vectorStoreMock.Object, + "testcollection", + 1, + null!)); } [Fact] public void Constructor_Throws_ForInvalidVectorDimensions() { // Act & Assert - Assert.Throws(() => new ChatHistoryMemoryProvider(this._vectorStoreMock.Object, "testcollection", 0, new ChatHistoryMemoryProviderScope() { UserId = "UID" })); - Assert.Throws(() => new ChatHistoryMemoryProvider(this._vectorStoreMock.Object, "testcollection", -5, new ChatHistoryMemoryProviderScope() { UserId = "UID" })); + Assert.Throws(() => new ChatHistoryMemoryProvider( + this._vectorStoreMock.Object, + "testcollection", + 0, + _ => new ChatHistoryMemoryProvider.State(new ChatHistoryMemoryProviderScope { UserId = "UID" }))); + Assert.Throws(() => new ChatHistoryMemoryProvider( + this._vectorStoreMock.Object, + "testcollection", + -5, + _ => new ChatHistoryMemoryProvider.State(new ChatHistoryMemoryProviderScope { UserId = "UID" }))); } #region InvokedAsync Tests @@ -113,16 +160,17 @@ public class ChatHistoryMemoryProviderTests UserId = "user1" }; - var provider = new ChatHistoryMemoryProvider(this._vectorStoreMock.Object, TestCollectionName, 1, storeScope); + var provider = new ChatHistoryMemoryProvider( + this._vectorStoreMock.Object, + TestCollectionName, + 1, + _ => new ChatHistoryMemoryProvider.State(storeScope)); var requestMsgWithValues = new ChatMessage(ChatRole.User, "request text") { MessageId = "req-1", AuthorName = "user1", CreatedAt = new DateTimeOffset(new DateTime(2000, 1, 1), TimeSpan.Zero) }; var requestMsgWithNulls = new ChatMessage(ChatRole.User, "request text nulls"); var responseMsg = new ChatMessage(ChatRole.Assistant, "response text") { MessageId = "resp-1", AuthorName = "assistant" }; - var invokedContext = new AIContextProvider.InvokedContext(s_mockAgent, s_mockSession, [requestMsgWithValues, requestMsgWithNulls]) - { - ResponseMessages = [responseMsg] - }; + var invokedContext = new AIContextProvider.InvokedContext(s_mockAgent, new TestAgentSession(), [requestMsgWithValues, requestMsgWithNulls], [responseMsg]); // Act await provider.InvokedAsync(invokedContext, CancellationToken.None); @@ -175,12 +223,9 @@ public class ChatHistoryMemoryProviderTests this._vectorStoreMock.Object, TestCollectionName, 1, - new ChatHistoryMemoryProviderScope() { UserId = "UID" }); + _ => new ChatHistoryMemoryProvider.State(new ChatHistoryMemoryProviderScope { UserId = "UID" })); var requestMsg = new ChatMessage(ChatRole.User, "request text") { MessageId = "req-1" }; - var invokedContext = new AIContextProvider.InvokedContext(s_mockAgent, s_mockSession, [requestMsg]) - { - InvokeException = new InvalidOperationException("Invoke failed") - }; + var invokedContext = new AIContextProvider.InvokedContext(s_mockAgent, new TestAgentSession(), [requestMsg], new InvalidOperationException("Invoke failed")); // Act await provider.InvokedAsync(invokedContext, CancellationToken.None); @@ -203,10 +248,10 @@ public class ChatHistoryMemoryProviderTests this._vectorStoreMock.Object, TestCollectionName, 1, - new ChatHistoryMemoryProviderScope() { UserId = "UID" }, + _ => new ChatHistoryMemoryProvider.State(new ChatHistoryMemoryProviderScope { UserId = "UID" }), loggerFactory: this._loggerFactoryMock.Object); var requestMsg = new ChatMessage(ChatRole.User, "request text") { MessageId = "req-1" }; - var invokedContext = new AIContextProvider.InvokedContext(s_mockAgent, s_mockSession, [requestMsg]); + var invokedContext = new AIContextProvider.InvokedContext(s_mockAgent, new TestAgentSession(), [requestMsg], []); // Act await provider.InvokedAsync(invokedContext, CancellationToken.None); @@ -252,12 +297,12 @@ public class ChatHistoryMemoryProviderTests this._vectorStoreMock.Object, TestCollectionName, 1, - new ChatHistoryMemoryProviderScope { UserId = "user1" }, + _ => new ChatHistoryMemoryProvider.State(new ChatHistoryMemoryProviderScope { UserId = "user1" }), options: options, loggerFactory: this._loggerFactoryMock.Object); var requestMsg = new ChatMessage(ChatRole.User, "request text"); - var invokedContext = new AIContextProvider.InvokedContext(s_mockAgent, s_mockSession, [requestMsg]); + var invokedContext = new AIContextProvider.InvokedContext(s_mockAgent, new TestAgentSession(), [requestMsg], []); // Act await provider.InvokedAsync(invokedContext, CancellationToken.None); @@ -326,14 +371,14 @@ public class ChatHistoryMemoryProviderTests this._vectorStoreMock.Object, TestCollectionName, 1, - new ChatHistoryMemoryProviderScope() { UserId = "UID" }, + _ => new ChatHistoryMemoryProvider.State(new ChatHistoryMemoryProviderScope { UserId = "UID" }), options: providerOptions); var requestMsg = new ChatMessage(ChatRole.User, "requesting relevant history"); - var invokingContext = new AIContextProvider.InvokingContext(s_mockAgent, s_mockSession, [requestMsg]); + var invokingContext = new AIContextProvider.InvokingContext(s_mockAgent, new TestAgentSession(), new AIContext { Messages = new List { requestMsg } }); // Act - await provider.InvokingAsync(invokingContext, CancellationToken.None); + var aiContext = await provider.InvokingAsync(invokingContext, CancellationToken.None); // Assert this._vectorStoreCollectionMock.Verify( @@ -343,6 +388,12 @@ public class ChatHistoryMemoryProviderTests It.IsAny>>(), It.IsAny()), Times.Once); + + Assert.NotNull(aiContext.Messages); + var messages = aiContext.Messages.ToList(); + Assert.Equal(2, messages.Count); + Assert.Equal(AgentRequestMessageSourceType.External, messages[0].GetAgentRequestMessageSourceType()); + Assert.Equal(AgentRequestMessageSourceType.AIContextProvider, messages[1].GetAgentRequestMessageSourceType()); } [Fact] @@ -378,10 +429,15 @@ public class ChatHistoryMemoryProviderTests }) .Returns(ToAsyncEnumerableAsync(new List>>())); - var provider = new ChatHistoryMemoryProvider(this._vectorStoreMock.Object, TestCollectionName, 1, options: providerOptions, storageScope: searchScope, searchScope: searchScope); + var provider = new ChatHistoryMemoryProvider( + this._vectorStoreMock.Object, + TestCollectionName, + 1, + _ => new ChatHistoryMemoryProvider.State(searchScope, searchScope), + options: providerOptions); var requestMsg = new ChatMessage(ChatRole.User, "requesting relevant history"); - var invokingContext = new AIContextProvider.InvokingContext(s_mockAgent, s_mockSession, [requestMsg]); + var invokingContext = new AIContextProvider.InvokingContext(s_mockAgent, new TestAgentSession(), new AIContext { Messages = new List { requestMsg } }); // Act await provider.InvokingAsync(invokingContext, CancellationToken.None); @@ -440,12 +496,11 @@ public class ChatHistoryMemoryProviderTests this._vectorStoreMock.Object, TestCollectionName, 1, - storageScope: scope, - searchScope: scope, + _ => new ChatHistoryMemoryProvider.State(scope, scope), options: options, loggerFactory: this._loggerFactoryMock.Object); - var invokingContext = new AIContextProvider.InvokingContext(s_mockAgent, s_mockSession, [new ChatMessage(ChatRole.User, "requesting relevant history")]); + var invokingContext = new AIContextProvider.InvokingContext(s_mockAgent, new TestAgentSession(), new AIContext { Messages = new List { new(ChatRole.User, "requesting relevant history") } }); // Act await provider.InvokingAsync(invokingContext, CancellationToken.None); @@ -479,52 +534,178 @@ public class ChatHistoryMemoryProviderTests #endregion - #region Serialization Tests + #region Message Filter Tests [Fact] - public void Serialize_Deserialize_RoundtripsScopes() + public async Task InvokingAsync_DefaultFilter_ExcludesNonExternalMessagesFromSearchAsync() { // Arrange - var storageScope = new ChatHistoryMemoryProviderScope + var providerOptions = new ChatHistoryMemoryProviderOptions { - ApplicationId = "app", - AgentId = "agent", - SessionId = "session", - UserId = "user" + SearchTime = ChatHistoryMemoryProviderOptions.SearchBehavior.BeforeAIInvoke, }; - var searchScope = new ChatHistoryMemoryProviderScope + string? capturedQuery = null; + this._vectorStoreCollectionMock + .Setup(c => c.SearchAsync( + It.IsAny(), + It.IsAny(), + It.IsAny>>(), + It.IsAny())) + .Callback>, CancellationToken>((query, _, _, _) => capturedQuery = query) + .Returns(ToAsyncEnumerableAsync(new List>>())); + + var provider = new ChatHistoryMemoryProvider( + this._vectorStoreMock.Object, + TestCollectionName, + 1, + _ => new ChatHistoryMemoryProvider.State(new ChatHistoryMemoryProviderScope { UserId = "UID" }), + options: providerOptions); + + var requestMessages = new List { - ApplicationId = "app2", - AgentId = "agent2", - SessionId = "session2", - UserId = "user2" + new(ChatRole.User, "External message"), + new(ChatRole.System, "From history") { AdditionalProperties = new() { { AgentRequestMessageSourceAttribution.AdditionalPropertiesKey, new AgentRequestMessageSourceAttribution(AgentRequestMessageSourceType.ChatHistory, "HistorySource") } } }, + new(ChatRole.System, "From context provider") { AdditionalProperties = new() { { AgentRequestMessageSourceAttribution.AdditionalPropertiesKey, new AgentRequestMessageSourceAttribution(AgentRequestMessageSourceType.AIContextProvider, "ContextSource") } } }, }; - var provider = new ChatHistoryMemoryProvider(this._vectorStoreMock.Object, TestCollectionName, 1, storageScope: storageScope, searchScope: searchScope); + var invokingContext = new AIContextProvider.InvokingContext(s_mockAgent, new TestAgentSession(), new AIContext { Messages = requestMessages }); // Act - var stateElement = provider.Serialize(); + await provider.InvokingAsync(invokingContext, CancellationToken.None); - using JsonDocument doc = JsonDocument.Parse(stateElement.GetRawText()); - var storage = doc.RootElement.GetProperty("storageScope"); - Assert.Equal("app", storage.GetProperty("applicationId").GetString()); - Assert.Equal("agent", storage.GetProperty("agentId").GetString()); - Assert.Equal("session", storage.GetProperty("sessionId").GetString()); - Assert.Equal("user", storage.GetProperty("userId").GetString()); + // Assert - Only External message used for search query + Assert.Equal("External message", capturedQuery); + } - var search = doc.RootElement.GetProperty("searchScope"); - Assert.Equal("app2", search.GetProperty("applicationId").GetString()); - Assert.Equal("agent2", search.GetProperty("agentId").GetString()); - Assert.Equal("session2", search.GetProperty("sessionId").GetString()); - Assert.Equal("user2", search.GetProperty("userId").GetString()); + [Fact] + public async Task InvokingAsync_CustomSearchInputFilter_OverridesDefaultAsync() + { + // Arrange + var providerOptions = new ChatHistoryMemoryProviderOptions + { + SearchTime = ChatHistoryMemoryProviderOptions.SearchBehavior.BeforeAIInvoke, + SearchInputMessageFilter = messages => messages // No filtering + }; - // Act - deserialize and serialize again - var provider2 = new ChatHistoryMemoryProvider(this._vectorStoreMock.Object, TestCollectionName, 1, serializedState: stateElement); - var stateElement2 = provider2.Serialize(); + string? capturedQuery = null; + this._vectorStoreCollectionMock + .Setup(c => c.SearchAsync( + It.IsAny(), + It.IsAny(), + It.IsAny>>(), + It.IsAny())) + .Callback>, CancellationToken>((query, _, _, _) => capturedQuery = query) + .Returns(ToAsyncEnumerableAsync(new List>>())); - // Assert - roundtrip the state - Assert.Equal(stateElement.GetRawText(), stateElement2.GetRawText()); + var provider = new ChatHistoryMemoryProvider( + this._vectorStoreMock.Object, + TestCollectionName, + 1, + _ => new ChatHistoryMemoryProvider.State(new ChatHistoryMemoryProviderScope { UserId = "UID" }), + options: providerOptions); + + var requestMessages = new List + { + new(ChatRole.User, "External message"), + new(ChatRole.System, "From history") { AdditionalProperties = new() { { AgentRequestMessageSourceAttribution.AdditionalPropertiesKey, new AgentRequestMessageSourceAttribution(AgentRequestMessageSourceType.ChatHistory, "HistorySource") } } }, + }; + + var invokingContext = new AIContextProvider.InvokingContext(s_mockAgent, new TestAgentSession(), new AIContext { Messages = requestMessages }); + + // Act + await provider.InvokingAsync(invokingContext, CancellationToken.None); + + // Assert - Both messages should be included in search query (identity filter) + Assert.NotNull(capturedQuery); + Assert.Contains("External message", capturedQuery); + Assert.Contains("From history", capturedQuery); + } + + [Fact] + public async Task InvokedAsync_DefaultFilter_ExcludesNonExternalMessagesFromStorageAsync() + { + // Arrange + var stored = new List>(); + + this._vectorStoreCollectionMock + .Setup(c => c.UpsertAsync(It.IsAny>>(), It.IsAny())) + .Callback>, CancellationToken>((items, ct) => + { + if (items != null) + { + stored.AddRange(items); + } + }) + .Returns(Task.CompletedTask); + + var provider = new ChatHistoryMemoryProvider( + this._vectorStoreMock.Object, + TestCollectionName, + 1, + _ => new ChatHistoryMemoryProvider.State(new ChatHistoryMemoryProviderScope { UserId = "UID" })); + + var requestMessages = new List + { + new(ChatRole.User, "External message"), + new(ChatRole.System, "From history") { AdditionalProperties = new() { { AgentRequestMessageSourceAttribution.AdditionalPropertiesKey, new AgentRequestMessageSourceAttribution(AgentRequestMessageSourceType.ChatHistory, "HistorySource") } } }, + new(ChatRole.System, "From context provider") { AdditionalProperties = new() { { AgentRequestMessageSourceAttribution.AdditionalPropertiesKey, new AgentRequestMessageSourceAttribution(AgentRequestMessageSourceType.AIContextProvider, "ContextSource") } } }, + }; + + var invokedContext = new AIContextProvider.InvokedContext(s_mockAgent, new TestAgentSession(), requestMessages, [new ChatMessage(ChatRole.Assistant, "Response")]); + + // Act + await provider.InvokedAsync(invokedContext, CancellationToken.None); + + // Assert - Only External message + response stored (ChatHistory and AIContextProvider excluded by default) + Assert.Equal(2, stored.Count); + Assert.Equal("External message", stored[0]["Content"]); + Assert.Equal("Response", stored[1]["Content"]); + } + + [Fact] + public async Task InvokedAsync_CustomStorageInputFilter_OverridesDefaultAsync() + { + // Arrange + var stored = new List>(); + + this._vectorStoreCollectionMock + .Setup(c => c.UpsertAsync(It.IsAny>>(), It.IsAny())) + .Callback>, CancellationToken>((items, ct) => + { + if (items != null) + { + stored.AddRange(items); + } + }) + .Returns(Task.CompletedTask); + + var provider = new ChatHistoryMemoryProvider( + this._vectorStoreMock.Object, + TestCollectionName, + 1, + _ => new ChatHistoryMemoryProvider.State(new ChatHistoryMemoryProviderScope { UserId = "UID" }), + options: new ChatHistoryMemoryProviderOptions + { + StorageInputMessageFilter = messages => messages // No filtering - store everything + }); + + var requestMessages = new List + { + new(ChatRole.User, "External message"), + new(ChatRole.System, "From history") { AdditionalProperties = new() { { AgentRequestMessageSourceAttribution.AdditionalPropertiesKey, new AgentRequestMessageSourceAttribution(AgentRequestMessageSourceType.ChatHistory, "HistorySource") } } }, + }; + + var invokedContext = new AIContextProvider.InvokedContext(s_mockAgent, new TestAgentSession(), requestMessages, [new ChatMessage(ChatRole.Assistant, "Response")]); + + // Act + await provider.InvokedAsync(invokedContext, CancellationToken.None); + + // Assert - All messages stored (identity filter overrides default) + Assert.Equal(3, stored.Count); + Assert.Equal("External message", stored[0]["Content"]); + Assert.Equal("From history", stored[1]["Content"]); + Assert.Equal("Response", stored[2]["Content"]); } #endregion @@ -537,4 +718,16 @@ public class ChatHistoryMemoryProviderTests yield return update; } } + + private sealed class TestAgentSession : AgentSession + { + public TestAgentSession() + { + } + + public TestAgentSession(AgentSessionStateBag stateBag) + { + this.StateBag = stateBag; + } + } } diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/TestJsonSerializerContext.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/TestJsonSerializerContext.cs index 0ac3ab9fbf..c07dd6eb8d 100644 --- a/dotnet/tests/Microsoft.Agents.AI.UnitTests/TestJsonSerializerContext.cs +++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/TestJsonSerializerContext.cs @@ -15,4 +15,5 @@ namespace Microsoft.Agents.AI.UnitTests; [JsonSerializable(typeof(string[]))] [JsonSerializable(typeof(Dictionary))] [JsonSerializable(typeof(ChatClientAgentSessionTests.Animal))] +[JsonSerializable(typeof(ChatClientAgentSession))] internal sealed partial class TestJsonSerializerContext : JsonSerializerContext; diff --git a/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/AgentWorkflowBuilderTests.cs b/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/AgentWorkflowBuilderTests.cs index 3dfe605f2e..eb017f07a3 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/AgentWorkflowBuilderTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/AgentWorkflowBuilderTests.cs @@ -161,7 +161,7 @@ public class AgentWorkflowBuilderTests } } - private sealed class DoubleEchoAgentSession() : InMemoryAgentSession(); + private sealed class DoubleEchoAgentSession() : AgentSession(); [Fact] public async Task BuildConcurrent_AgentsRunInParallelAsync() diff --git a/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/InProcessExecutionTests.cs b/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/InProcessExecutionTests.cs index dc51338aa3..e45cc39d2f 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/InProcessExecutionTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/InProcessExecutionTests.cs @@ -195,5 +195,5 @@ public class InProcessExecutionTests /// /// Simple session implementation for SimpleTestAgent. /// - private sealed class SimpleTestAgentSession : InMemoryAgentSession; + private sealed class SimpleTestAgentSession : AgentSession; } diff --git a/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/RoleCheckAgent.cs b/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/RoleCheckAgent.cs index dde1d1feed..385c427b37 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/RoleCheckAgent.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/RoleCheckAgent.cs @@ -46,5 +46,5 @@ internal sealed class RoleCheckAgent(bool allowOtherAssistantRoles, string? id = }; } - private sealed class RoleCheckAgentSession : InMemoryAgentSession; + private sealed class RoleCheckAgentSession : AgentSession; } diff --git a/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/Sample/06_GroupChat_Workflow.cs b/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/Sample/06_GroupChat_Workflow.cs index afade362a1..1e00e2e649 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/Sample/06_GroupChat_Workflow.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/Sample/06_GroupChat_Workflow.cs @@ -90,4 +90,4 @@ internal sealed class HelloAgent(string id = nameof(HelloAgent)) : AIAgent } } -internal sealed class HelloAgentSession() : InMemoryAgentSession(); +internal sealed class HelloAgentSession() : AgentSession(); diff --git a/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/TestEchoAgent.cs b/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/TestEchoAgent.cs index 9d5eca42bf..b1c7c8cbe3 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/TestEchoAgent.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/TestEchoAgent.cs @@ -5,6 +5,7 @@ using System.Collections.Generic; using System.Linq; using System.Runtime.CompilerServices; using System.Text.Json; +using System.Text.Json.Serialization; using System.Threading; using System.Threading.Tasks; using Microsoft.Extensions.AI; @@ -16,6 +17,8 @@ internal class TestEchoAgent(string? id = null, string? name = null, string? pre protected override string? IdCore => id; public override string? Name => name ?? base.Name; + public InMemoryChatHistoryProvider ChatHistoryProvider { get; } = new(); + protected override async ValueTask DeserializeSessionCoreAsync(JsonElement serializedState, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default) { return serializedState.Deserialize(jsonSerializerOptions) ?? await this.CreateSessionAsync(cancellationToken); @@ -28,15 +31,15 @@ internal class TestEchoAgent(string? id = null, string? name = null, string? pre throw new InvalidOperationException("The provided session is not compatible with the agent. Only sessions created by the agent can be serialized."); } - return new(typedSession.Serialize(jsonSerializerOptions)); + return new(JsonSerializer.SerializeToElement(typedSession, jsonSerializerOptions)); } protected override ValueTask CreateSessionCoreAsync(CancellationToken cancellationToken = default) => new(new EchoAgentSession()); - private static ChatMessage UpdateSession(ChatMessage message, InMemoryAgentSession? session = null) + private ChatMessage UpdateSession(ChatMessage message, AgentSession? session = null) { - session?.ChatHistoryProvider.Add(message); + this.ChatHistoryProvider.GetMessages(session).Add(message); return message; } @@ -45,7 +48,7 @@ internal class TestEchoAgent(string? id = null, string? name = null, string? pre { foreach (ChatMessage message in messages) { - UpdateSession(message, session as InMemoryAgentSession); + this.UpdateSession(message, session); } IEnumerable echoMessages @@ -53,14 +56,14 @@ internal class TestEchoAgent(string? id = null, string? name = null, string? pre where message.Role == ChatRole.User && !string.IsNullOrEmpty(message.Text) select - UpdateSession(new ChatMessage(ChatRole.Assistant, $"{prefix}{message.Text}") + this.UpdateSession(new ChatMessage(ChatRole.Assistant, $"{prefix}{message.Text}") { AuthorName = this.Name ?? this.Id, CreatedAt = DateTimeOffset.Now, MessageId = Guid.NewGuid().ToString("N") - }, session as InMemoryAgentSession); + }, session); - return echoMessages.Concat(this.GetEpilogueMessages(options).Select(m => UpdateSession(m, session as InMemoryAgentSession))); + return echoMessages.Concat(this.GetEpilogueMessages(options).Select(m => this.UpdateSession(m, session))); } protected virtual IEnumerable GetEpilogueMessages(AgentRunOptions? options = null) @@ -99,11 +102,11 @@ internal class TestEchoAgent(string? id = null, string? name = null, string? pre } } - private sealed class EchoAgentSession : InMemoryAgentSession + private sealed class EchoAgentSession : AgentSession { - internal new JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null) - { - return base.Serialize(jsonSerializerOptions); - } + internal EchoAgentSession() { } + + [JsonConstructor] + internal EchoAgentSession(AgentSessionStateBag stateBag) : base(stateBag) { } } } diff --git a/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/TestReplayAgent.cs b/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/TestReplayAgent.cs index 2dd33a67e4..3a117b9492 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/TestReplayAgent.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/TestReplayAgent.cs @@ -104,5 +104,5 @@ public class TestReplayAgent(List? messages = null, string? id = nu return candidateMessages; } - private sealed class ReplayAgentSession() : InMemoryAgentSession(); + private sealed class ReplayAgentSession() : AgentSession(); } diff --git a/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/TestRequestAgent.cs b/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/TestRequestAgent.cs index 65a49add96..0b49f4de05 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/TestRequestAgent.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/TestRequestAgent.cs @@ -330,7 +330,7 @@ internal sealed class TestRequestAgent(TestAgentRequestType requestType, int unp } } - private sealed class TestRequestAgentSession : InMemoryAgentSession + private sealed class TestRequestAgentSession : AgentSession where TRequest : AIContent where TResponse : AIContent { @@ -343,19 +343,13 @@ internal sealed class TestRequestAgent(TestAgentRequestType requestType, int unp public HashSet ServicedRequests { get; } = new(); public HashSet PairedRequests { get; } = new(); - private static JsonElement DeserializeAndExtractState(JsonElement serializedState, - out TestRequestAgentSessionState state, - JsonSerializerOptions? jsonSerializerOptions = null) + public TestRequestAgentSession(JsonElement element, JsonSerializerOptions? jsonSerializerOptions = null) { - state = JsonSerializer.Deserialize(serializedState, jsonSerializerOptions) + var state = JsonSerializer.Deserialize(element, jsonSerializerOptions) ?? throw new ArgumentException("Unable to deserialize session state."); - return state.SessionState; - } + this.StateBag = AgentSessionStateBag.Deserialize(state.SessionState); - public TestRequestAgentSession(JsonElement element, JsonSerializerOptions? jsonSerializerOptions = null) - : base(DeserializeAndExtractState(element, out TestRequestAgentSessionState state, jsonSerializerOptions)) - { this.UnservicedRequests = state.UnservicedRequests.ToDictionary( keySelector: item => item.Key, elementSelector: item => item.Value.As()!); @@ -364,9 +358,9 @@ internal sealed class TestRequestAgent(TestAgentRequestType requestType, int unp this.PairedRequests = state.PairedRequests; } - protected override JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null) + internal JsonElement Serialize(JsonSerializerOptions? jsonSerializerOptions = null) { - JsonElement sessionState = base.Serialize(jsonSerializerOptions); + JsonElement sessionState = this.StateBag.Serialize(); Dictionary portableUnservicedRequests = this.UnservicedRequests.ToDictionary( diff --git a/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/WorkflowHostSmokeTests.cs b/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/WorkflowHostSmokeTests.cs index 5a041699d1..a68a5bac75 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/WorkflowHostSmokeTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/WorkflowHostSmokeTests.cs @@ -32,18 +32,16 @@ public class WorkflowHostSmokeTests { private sealed class AlwaysFailsAIAgent(bool failByThrowing) : AIAgent { - private sealed class Session : InMemoryAgentSession + private sealed class Session : AgentSession { public Session() { } - public Session(JsonElement serializedSession, JsonSerializerOptions? jsonSerializerOptions = null) - : base(serializedSession, jsonSerializerOptions) - { } + public Session(AgentSessionStateBag stateBag) : base(stateBag) { } } protected override ValueTask DeserializeSessionCoreAsync(JsonElement serializedState, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default) { - return new(new Session(serializedState, jsonSerializerOptions)); + return new(serializedState.Deserialize(jsonSerializerOptions)!); } protected override ValueTask CreateSessionCoreAsync(CancellationToken cancellationToken = default) diff --git a/dotnet/tests/OpenAIChatCompletion.IntegrationTests/OpenAIChatCompletionFixture.cs b/dotnet/tests/OpenAIChatCompletion.IntegrationTests/OpenAIChatCompletionFixture.cs index 304df28fba..96a9a17ae8 100644 --- a/dotnet/tests/OpenAIChatCompletion.IntegrationTests/OpenAIChatCompletionFixture.cs +++ b/dotnet/tests/OpenAIChatCompletion.IntegrationTests/OpenAIChatCompletionFixture.cs @@ -30,14 +30,14 @@ public class OpenAIChatCompletionFixture : IChatClientAgentFixture public async Task> GetChatHistoryAsync(AIAgent agent, AgentSession session) { - var typedSession = (ChatClientAgentSession)session; + var chatHistoryProvider = agent.GetService(); - if (typedSession.ChatHistoryProvider is null) + if (chatHistoryProvider is null) { return []; } - return (await typedSession.ChatHistoryProvider.InvokingAsync(new(agent, session, []))).ToList(); + return (await chatHistoryProvider.InvokingAsync(new(agent, session, []))).ToList(); } public Task CreateChatClientAgentAsync( diff --git a/dotnet/tests/OpenAIResponse.IntegrationTests/OpenAIResponseFixture.cs b/dotnet/tests/OpenAIResponse.IntegrationTests/OpenAIResponseFixture.cs index 719db6a0b0..e36f8990f6 100644 --- a/dotnet/tests/OpenAIResponse.IntegrationTests/OpenAIResponseFixture.cs +++ b/dotnet/tests/OpenAIResponse.IntegrationTests/OpenAIResponseFixture.cs @@ -50,12 +50,14 @@ public class OpenAIResponseFixture(bool store) : IChatClientAgentFixture return [.. previousMessages, responseMessage]; } - if (typedSession.ChatHistoryProvider is null) + var chatHistoryProvider = agent.GetService(); + + if (chatHistoryProvider is null) { return []; } - return (await typedSession.ChatHistoryProvider.InvokingAsync(new(agent, session, []))).ToList(); + return (await chatHistoryProvider.InvokingAsync(new(agent, session, []))).ToList(); } private static ChatMessage ConvertToChatMessage(ResponseItem item) From 9506fb28f6994f22b671a55dc26b6fc23811f73b Mon Sep 17 00:00:00 2001 From: SergeyMenshykh <68852919+SergeyMenshykh@users.noreply.github.com> Date: Fri, 13 Feb 2026 17:03:51 +0000 Subject: [PATCH 06/22] .NET: [Breaking] Structured Output improvements (#3761) * .NET: Delete AgentResponse.{Try}Deserialize methods (#3518) * delete deserialize method of agent response * order usings * Update dotnet/samples/GettingStarted/FoundryAgents/FoundryAgents_Step05_StructuredOutput/Program.cs Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * Update dotnet/samples/GettingStarted/Workflows/_Foundational/08_WriterCriticWorkflow/Program.cs Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * Update dotnet/samples/GettingStarted/AGUI/Step05_StateManagement/Server/SharedStateAgent.cs Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * Update dotnet/samples/AGUIClientServer/AGUIDojoServer/SharedState/SharedStateAgent.cs Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * Update dotnet/samples/M365Agent/Agents/WeatherForecastAgent.cs Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --------- Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * .NET:[Breaking] Add support for structured output (#3658) * add support for so * restore lost xml comment part * fix using ordering * Update dotnet/src/Microsoft.Agents.AI.Abstractions/AIAgentStructuredOutput.cs Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * Update dotnet/src/Microsoft.Agents.AI.Abstractions/AIAgentStructuredOutput.cs Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * Update dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_SO_WithFormatResponseTests.cs Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * addressw pr review comments * address pr review feedback * address pr review comments * fix compilation issues after the latest merge with main * remove unnecessry options * remove RunAsync methods * address code review feedback * address pr review feedback * make copy constructor protected * address pr review feedback --------- Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * .NET: Add decorator for structured output support (#3694) * add decorator that adds structured output support to agents that don't natively support it. * Update dotnet/src/Microsoft.Agents.AI/StructuredOutput/StructuredOutputAgentResponse.cs Co-authored-by: westey <164392973+westey-m@users.noreply.github.com> * Update dotnet/samples/GettingStarted/Agents/Agent_Step05_StructuredOutput/Program.cs Co-authored-by: westey <164392973+westey-m@users.noreply.github.com> * address pr review feedback --------- Co-authored-by: westey <164392973+westey-m@users.noreply.github.com> * .NET: Support primitives and arrays for SO (#3696) * wrap primitives and arrays * fix file encoding * address review comments * add adr * add missed change * fix compilation issue * address review comments * rename adr file name * reflect decision to have SO decorator as a reference implementation in samples * .NET: Move SO agent to samples (#3820) * move SO agent to samples * change file encoding * fix files encoding * .NET: Preserve caller context (#3803) * fix stuck orchestration * add previously removed RunAsync method to DurableAIAgent * suppress IDE0005 warning * update changelog and remove unused constructor of AgentResponse * updatge the changelog * address PR review feedback * .NET: Disable irrelevant integration test (#3913) * disable irrelevant integration test * Update dotnet/tests/AzureAI.IntegrationTests/AIProjectClientAgentStructuredOutputRunTests.cs Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --------- Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * forgotten change * address pr review feedback * disable intermittently failing integration test. --------- Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> Co-authored-by: westey <164392973+westey-m@users.noreply.github.com> --- docs/decisions/0016-structured-output.md | 658 ++++++++++++++++++ dotnet/agent-framework-dotnet.slnx | 3 + dotnet/eng/MSBuild/Shared.props | 3 + .../SharedState/SharedStateAgent.cs | 23 +- .../Server/SharedStateAgent.cs | 23 +- .../AIAgentBuilderExtensions.cs | 49 ++ .../Agent_Step05_StructuredOutput/Program.cs | 177 ++++- .../Agent_Step05_StructuredOutput/README.md | 52 ++ .../StructuredOutputAgent.cs | 88 +++ .../StructuredOutputAgentOptions.cs | 31 + .../StructuredOutputAgentResponse.cs | 28 + .../Program.cs | 3 +- .../08_WriterCriticWorkflow/Program.cs | 3 +- .../M365Agent/Agents/WeatherForecastAgent.cs | 23 +- .../AIAgent.cs | 2 +- .../AIAgentStructuredOutput.cs} | 91 +-- .../AgentResponse.cs | 147 +--- .../AgentResponse{T}.cs | 88 ++- .../AgentRunOptions.cs | 34 +- .../Microsoft.Agents.AI.Abstractions.csproj | 1 + .../CHANGELOG.md | 1 + .../DurableAIAgent.cs | 189 ++--- .../DurableAIAgentProxy.cs | 7 +- .../DurableAgentRunOptions.cs | 29 +- .../Microsoft.Agents.AI.DurableTask.csproj | 5 + .../ChatClient/ChatClientAgent.cs | 6 + .../ChatClientAgentCustomOptions.cs | 44 +- .../ChatClient/ChatClientAgentRunOptions.cs | 14 + .../ChatClientAgentRunResponse{T}.cs | 45 -- .../StructuredOutputSchemaUtilities.cs | 104 +++ .../StructuredOutputRunTests.cs | 110 +++ .../Support/Constants.cs | 2 +- .../Support/SessionCleanup.cs | 2 +- ...jectClientAgentStructuredOutputRunTests.cs | 99 +++ .../AIProjectClientFixture.cs | 15 +- ...gentsPersistentStructuredOutputRunTests.cs | 13 + .../AIAgentStructuredOutputTests.cs | 391 +++++++++++ .../AgentResponseTests.cs | 130 +--- .../AgentRunOptionsTests.cs | 61 +- .../TestJsonSerializerContext.cs | 1 + .../DurableAgentRunOptionsTests.cs | 94 +++ .../ChatClientAgentRunOptionsTests.cs | 87 +++ .../ChatClient/ChatClientAgentTests.cs | 54 -- ...tructuredOutput_WithFormatResponseTests.cs | 212 ++++++ ...gent_StructuredOutput_WithRunAsyncTests.cs | 56 ++ .../Models/Animal.cs | 10 + .../Models/Species.cs | 10 + ...OpenAIAssistantStructuredOutputRunTests.cs | 9 + ...IChatCompletionStructuredOutputRunTests.cs | 9 + .../OpenAIResponseStructuredOutputRunTests.cs | 9 + 50 files changed, 2751 insertions(+), 594 deletions(-) create mode 100644 docs/decisions/0016-structured-output.md create mode 100644 dotnet/samples/GettingStarted/Agents/Agent_Step05_StructuredOutput/AIAgentBuilderExtensions.cs create mode 100644 dotnet/samples/GettingStarted/Agents/Agent_Step05_StructuredOutput/README.md create mode 100644 dotnet/samples/GettingStarted/Agents/Agent_Step05_StructuredOutput/StructuredOutputAgent.cs create mode 100644 dotnet/samples/GettingStarted/Agents/Agent_Step05_StructuredOutput/StructuredOutputAgentOptions.cs create mode 100644 dotnet/samples/GettingStarted/Agents/Agent_Step05_StructuredOutput/StructuredOutputAgentResponse.cs rename dotnet/src/{Microsoft.Agents.AI/ChatClient/ChatClientAgentStructuredOutput.cs => Microsoft.Agents.AI.Abstractions/AIAgentStructuredOutput.cs} (58%) delete mode 100644 dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgentRunResponse{T}.cs create mode 100644 dotnet/src/Shared/StructuredOutput/StructuredOutputSchemaUtilities.cs create mode 100644 dotnet/tests/AgentConformance.IntegrationTests/StructuredOutputRunTests.cs create mode 100644 dotnet/tests/AzureAI.IntegrationTests/AIProjectClientAgentStructuredOutputRunTests.cs create mode 100644 dotnet/tests/AzureAIAgentsPersistent.IntegrationTests/AzureAIAgentsPersistentStructuredOutputRunTests.cs create mode 100644 dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/AIAgentStructuredOutputTests.cs create mode 100644 dotnet/tests/Microsoft.Agents.AI.DurableTask.UnitTests/DurableAgentRunOptionsTests.cs create mode 100644 dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_StructuredOutput_WithFormatResponseTests.cs create mode 100644 dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_StructuredOutput_WithRunAsyncTests.cs create mode 100644 dotnet/tests/Microsoft.Agents.AI.UnitTests/Models/Animal.cs create mode 100644 dotnet/tests/Microsoft.Agents.AI.UnitTests/Models/Species.cs create mode 100644 dotnet/tests/OpenAIAssistant.IntegrationTests/OpenAIAssistantStructuredOutputRunTests.cs create mode 100644 dotnet/tests/OpenAIChatCompletion.IntegrationTests/OpenAIChatCompletionStructuredOutputRunTests.cs create mode 100644 dotnet/tests/OpenAIResponse.IntegrationTests/OpenAIResponseStructuredOutputRunTests.cs diff --git a/docs/decisions/0016-structured-output.md b/docs/decisions/0016-structured-output.md new file mode 100644 index 0000000000..4fdae3c77e --- /dev/null +++ b/docs/decisions/0016-structured-output.md @@ -0,0 +1,658 @@ +--- +status: proposed +contact: sergeymenshykh +date: 2026-01-22 +deciders: rbarreto, westey-m, stephentoub +informed: {} +--- + +# Structured Output + +Structured output is a valuable aspect of any agent system, since it forces an agent to produce output in a required format that may include required fields. +This allows easily turning unstructured data into structured data using a general-purpose language model. + +## Context and Problem Statement + +Structured output is currently supported only by `ChatClientAgent` and can be configured in two ways: + +**Approach 1: ResponseFormat + Deserialize** + +Specify the SO type schema via the `ChatClientAgent{Run}Options.ChatOptions.ResponseFormat` property at agent creation or invocation time, then use `JsonSerializer.Deserialize` to extract the structured data from the response text. + + ```csharp + // SO type can be provided at agent creation time + ChatClientAgent agent = chatClient.AsAIAgent(new ChatClientAgentOptions() + { + Name = "...", + ChatOptions = new() { ResponseFormat = ChatResponseFormat.ForJsonSchema() } + }); + + AgentResponse response = await agent.RunAsync("..."); + + PersonInfo personInfo = response.Deserialize(JsonSerializerOptions.Web); + + Console.WriteLine($"Name: {personInfo.Name}"); + Console.WriteLine($"Age: {personInfo.Age}"); + Console.WriteLine($"Occupation: {personInfo.Occupation}"); + + // Alternatively, SO type can be provided at agent invocation time + response = await agent.RunAsync("...", new ChatClientAgentRunOptions() + { + ChatOptions = new() { ResponseFormat = ChatResponseFormat.ForJsonSchema() } + }); + + personInfo = response.Deserialize(JsonSerializerOptions.Web); + + Console.WriteLine($"Name: {personInfo.Name}"); + Console.WriteLine($"Age: {personInfo.Age}"); + Console.WriteLine($"Occupation: {personInfo.Occupation}"); + ``` + +**Approach 2: Generic RunAsync** + +Supply the SO type as a generic parameter to `RunAsync` and access the parsed result directly via the `Result` property. + + ```csharp + ChatClientAgent agent = ...; + + AgentResponse response = await agent.RunAsync("..."); + + Console.WriteLine($"Name: {response.Result.Name}"); + Console.WriteLine($"Age: {response.Result.Age}"); + Console.WriteLine($"Occupation: {response.Result.Occupation}"); + ``` + Note: `RunAsync` is an instance method of `ChatClientAgent` and not part of the `AIAgent` base class since not all agents support structured output. + +Approach 1 is perceived as cumbersome by the community, as it requires additional effort when using primitive or collection types - the SO schema may need to be wrapped in an artificial JSON object. Otherwise, the caller will encounter an error like _Invalid schema for response_format 'Movie': schema must be a JSON Schema of 'type: "object"', got 'type: "array"'_. +This occurs because OpenAI and compatible APIs require a JSON object as the root schema. + +Approach 1 is also necessary in scenarios where (a) agents can only be configured with SO at creation time (such as with `AIProjectClient`), (b) the SO type is not known at compile time, or (c) the JSON schema is represented as text (for declarative agents) or as a `JsonElement`. + +Approach 2 is more convenient and works seamlessly with primitives and collections. However, it requires the SO type to be known at compile time, making it less flexible. + +Additionally, since the `RunAsync` methods are instance methods of `ChatClientAgent` and are not part of the `AIAgent` base class, applying decorators like `OpenTelemetryAgent` on top of `ChatClientAgent` prevents users from accessing `RunAsync`, meaning structured output is not available with decorated agents. + +Given the different scenarios above in which structured output can be used, there is no one-size-fits-all solution. Each approach has its own advantages and limitations, +and the two can complement each other to provide a comprehensive structured output experience across various use cases. + +## Approaches Overview + +1. SO usage via `ResponseFormat` property +2. SO usage via `RunAsync` generic method + +## 1. SO usage via `ResponseFormat` property + +This approach should be used in the following scenarios: + - 1.1 SO result as text is sufficient as is, and deserialization is not required + - 1.2 SO for inter-agent collaboration + - 1.3 SO can only be configured at agent creation time (such as with `AIProjectClient`) + - 1.4 SO type is not known at compile time and represented by System.Type + - 1.5 SO is represented by JSON schema and there's no corresponding .NET type either at compile time or at runtime + - 1.6 SO in streaming scenarios, where the SO response is produced in parts + +**Note: Primitives and arrays are not supported by this approach.** + +When a caller provides a schema via `ResponseFormat`, they are explicitly telling the framework what schema to use. The framework passes that schema through as-is and +is not responsible for transforming it. Because the framework does not own the schema, it cannot wrap primitives or arrays into a JSON object to satisfy API requirements, +nor can it unwrap the response afterward - the caller controls the schema and is responsible for ensuring it is compatible with the underlying API. + +This is in contrast to the `RunAsync` approach (section 2), where the caller provides a type `T` and says "make it work." In that case, the caller does not +dictate the schema - the framework infers the schema from `T`, owns the end-to-end pipeline (schema generation, API invocation, and deserialization), and can +therefore wrap and unwrap primitives and arrays transparently. + +Additionally, in streaming scenarios (1.6), the framework cannot reliably unwrap a response it did not wrap, since it has no way of knowing whether the caller wrapped the schema.Wrapping and unwrapping can only be done safely when the framework owns the entire lifecycle - from schema creation through deserialization — which is only the case with `RunAsync`. + +If a caller needs to work with primitives or arrays via the `ResponseFormat` approach, they can easily create a wrapper type around them: + +```csharp +public class MovieListWrapper +{ + public List Movies { get; set; } +} +``` + +### 1.1 SO result as text is sufficient as is, and deserialization is not required + +In this scenario, the caller only needs the raw JSON text returned by the model and does not need to deserialize it into a .NET type. +The SO schema is specified via `ResponseFormat` at agent creation or invocation time, and the response text is consumed directly from the `AgentResponse`. + +```csharp +AIAgent agent = chatClient.AsAIAgent(); + +AgentRunOptions runOptions = new() +{ + ResponseFormat = ChatResponseFormat.ForJsonSchema() +}; + +AgentResponse response = await agent.RunAsync("...", options: runOptions); + +Console.WriteLine(response.Text); +``` + +### 1.2 SO for inter-agent collaboration + +This scenario assumes a multi-agent setup where agents collaborate by passing messages to each other. +One agent produces structured output as text that is then passed directly as input to the next agent, without intermediate deserialization. + +```csharp +// First agent extracts structured data from unstructured input +AIAgent extractionAgent = chatClient.AsAIAgent(new ChatClientAgentOptions() +{ + Name = "ExtractionAgent", + ChatOptions = new() + { + Instructions = "Extract person information from the provided text.", + ResponseFormat = ChatResponseFormat.ForJsonSchema() + } +}); + +AgentResponse extractionResponse = await extractionAgent.RunAsync("John Smith is a 35-year-old software engineer."); + +// Pass the message with structured output text directly to the next agent +ChatMessage soMessage = extractionResponse.Messages.Last(); + +AIAgent summaryAgent = chatClient.AsAIAgent(new ChatClientAgentOptions() +{ + Name = "SummaryAgent", + ChatOptions = new() { Instructions = "Given the following structured person data, write a short professional bio." } +}); + +AgentResponse summaryResponse = await summaryAgent.RunAsync(soMessage); + +Console.WriteLine(summaryResponse); +``` + +### 1.3 SO configured at agent creation time + +In this scenario, the SO schema can only be configured at agent creation time (such as with `AIProjectClient`) and cannot be changed on a per-run basis. +The caller specifies the `ResponseFormat` when creating the agent, and all subsequent invocations use the same schema. + +```csharp +AIProjectClient client = ...; + +AIAgent agent = await client.CreateAIAgentAsync(model: "", new ChatClientAgentOptions() +{ + Name = "...", + ChatOptions = new() { ResponseFormat = ChatResponseFormat.ForJsonSchema() } +}); + +AgentResponse response = await agent.RunAsync("Please provide information about John Smith."); + +PersonInfo personInfo = JsonSerializer.Deserialize(response.Text, JsonSerializerOptions.Web)!; + +Console.WriteLine($"Name: {personInfo.Name}"); +Console.WriteLine($"Age: {personInfo.Age}"); +Console.WriteLine($"Occupation: {personInfo.Occupation}"); +``` + +### 1.4 SO type not known at compile time and represented by System.Type + +In this scenario, the SO type is not known at compile time and is provided as a `System.Type` at runtime. This is useful for dynamic scenarios where the schema is determined programmatically, +such as when building tooling or frameworks that work with user-defined types. + +```csharp +Type soType = GetStructuredOutputTypeFromConfiguration(); // e.g., typeof(PersonInfo) + +ChatResponseFormat responseFormat = ChatResponseFormat.ForJsonSchema(soType); + +AgentResponse response = await agent.RunAsync("...", new ChatClientAgentRunOptions() +{ + ChatOptions = new() { ResponseFormat = responseFormat } +}); + +PersonInfo personInfo = (PersonInfo)JsonSerializer.Deserialize(response.Text, soType, JsonSerializerOptions.Web)!; +``` + +### 1.5 SO represented by JSON schema with no corresponding .NET type + +In this scenario, the SO schema is represented as raw JSON schema text or a `JsonElement`, and there is no corresponding .NET type available at compile time or runtime. +This is typical for declarative agents or scenarios where schemas are loaded from external configuration. + +```csharp +// JSON schema provided as a string, e.g., loaded from a configuration file +string jsonSchema = """ +{ + "type": "object", + "properties": { + "name": { "type": "string" }, + "age": { "type": "integer" }, + "occupation": { "type": "string" } + }, + "required": ["name", "age", "occupation"] +} +"""; + +ChatResponseFormat responseFormat = ChatResponseFormat.ForJsonSchema( + jsonSchemaName: "PersonInfo", + jsonSchema: BinaryData.FromString(jsonSchema)); + +AgentResponse response = await agent.RunAsync("...", new ChatClientAgentRunOptions() +{ + ChatOptions = new() { ResponseFormat = responseFormat } +}); + +// Consume the SO result as text since there's no .NET type to deserialize into +Console.WriteLine(response.Text); +``` + +### 1.6 SO in streaming scenarios + +In this scenario, the SO response is produced incrementally in parts via streaming. The caller specifies the `ResponseFormat` and consumes the response chunks as they arrive. +Deserialization is performed after all chunks have been received. + +```csharp +AIAgent agent = chatClient.AsAIAgent(new ChatClientAgentOptions() +{ + Name = "HelpfulAssistant", + ChatOptions = new() + { + Instructions = "You are a helpful assistant.", + ResponseFormat = ChatResponseFormat.ForJsonSchema() + } +}); + +IAsyncEnumerable updates = agent.RunStreamingAsync("Please provide information about John Smith, who is a 35-year-old software engineer."); + +AgentResponse response = await updates.ToAgentResponseAsync(); + +// Deserialize the complete SO result after streaming is finished +PersonInfo personInfo = JsonSerializer.Deserialize(response.Text)!; +``` + +## 2. SO usage via `RunAsync` generic method + +This approach provides a convenient way to work with structured output on a per-run basis when the target type is known at compile time and a typed instance of the result +is required. + +### Decision Drivers + +1. Support arrays and primitives as SO types +2. Support complex types as SO types +3. Work with `AIAgent` decorators (e.g., `OpenTelemetryAgent`) +4. Enable SO for all AI agents, regardless of whether they natively support it + +### Considered Options + +1. `RunAsync` as an instance method of `AIAgent` class delegating to virtual `RunCoreAsync` +2. `RunAsync` as an extension method using feature collection +3. `RunAsync` as a method of the new `ITypedAIAgent` interface +4. `RunAsync` as an instance method of `AIAgent` class working via the new `AgentRunOptions.ResponseFormat` property + +### 1. `RunAsync` as an instance method of `AIAgent` class delegating to virtual `RunCoreAsync` + +This option adds the `RunAsync` method directly to the `AIAgent` base class. + +```csharp +public abstract class AIAgent +{ + public Task> RunAsync( + IEnumerable messages, + AgentSession? session = null, + JsonSerializerOptions? serializerOptions = null, + AgentRunOptions? options = null, + CancellationToken cancellationToken = default) + => this.RunCoreAsync(messages, session, serializerOptions, options, cancellationToken); + + protected virtual Task> RunCoreAsync( + IEnumerable messages, + AgentSession? session = null, + JsonSerializerOptions? serializerOptions = null, + AgentRunOptions? options = null, + CancellationToken cancellationToken = default) + { + throw new NotSupportedException($"The agent of type '{this.GetType().FullName}' does not support typed responses."); + } +} +``` + +Agents with native SO support override the `RunCoreAsync` method to provide their implementation. If not overridden, the method throws a `NotSupportedException`. + +Users will call the generic `RunAsync` method directly on the agent: + +```csharp +AIAgent agent = chatClient.AsAIAgent(name: "HelpfulAssistant", instructions: "You are a helpful assistant."); + +AgentResponse response = await agent.RunAsync("Please provide information about John Smith, who is a 35-year-old software engineer."); +``` + +Decision drivers satisfied: +1. Support arrays and primitives as SO types +2. Support complex types as SO types +3. Work with `AIAgent` decorators (e.g., `OpenTelemetryAgent`) +4. Enable SO for all AI agents, regardless of whether they natively support it + +Pros: +- The `AIAgent.RunAsync` method is easily discoverable. +- Both the SO decorator and `ChatClientAgent` have compile-time access to the type `T`, allowing them to use the native `IChatClient.GetResponseAsync` API, which handles primitives and collections seamlessly. + +Cons: +- Agents without native SO support will still expose `RunAsync`, which may be misleading. +- `ChatClientAgent` exposing `RunAsync` may be misleading when the underlying chat client does not support SO. +- All `AIAgent` decorators must override `RunCoreAsync` to properly handle `RunAsync` calls. + +### 2. `RunAsync` as an extension method using feature collection + +This option uses the Agent Framework feature collection (implemented via `AgentRunOptions.AdditionalProperties`) to pass a `StructuredOutputFeature` to agents, signaling that SO is requested. + +Agents with native SO support check for this feature. If present, they read the target type, build the schema, invoke the underlying API, and store the response back in the feature. +```csharp +public class StructuredOutputFeature +{ + public StructuredOutputFeature(Type outputType) + { + this.OutputType = outputType; + } + + [JsonIgnore] + public Type OutputType { get; set; } + + public JsonSerializerOptions? SerializerOptions { get; set; } + + public AgentResponse? Response { get; set; } +} +``` + +The `RunAsync` extension method for `AIAgent` adds this feature to the collection. +```csharp +public static async Task> RunAsync( + this AIAgent agent, + IEnumerable messages, + AgentSession? session = null, + JsonSerializerOptions? serializerOptions = null, + AgentRunOptions? options = null, + CancellationToken cancellationToken = default) +{ + // Create the structured output feature. + StructuredOutputFeature structuredOutputFeature = new(typeof(T)) + { + SerializerOptions = serializerOptions, + }; + + // Register it in the feature collection. + ((options ??= new AgentRunOptions()).AdditionalProperties ??= []).Add(typeof(StructuredOutputFeature).FullName!, structuredOutputFeature); + + var response = await agent.RunAsync(messages, session, options, cancellationToken).ConfigureAwait(false); + + if (structuredOutputFeature.Response is not null) + { + return new StructuredOutputResponse(structuredOutputFeature.Response, response, serializerOptions); + } + + throw new InvalidOperationException("No structured output response was generated by the agent."); +} +``` + +Users will call the `RunAsync` extension method directly on the agent: + +```csharp +AIAgent agent = chatClient.AsAIAgent(name: "HelpfulAssistant", instructions: "You are a helpful assistant."); + +AgentResponse response = await agent.RunAsync("Please provide information about John Smith, who is a 35-year-old software engineer."); +``` + +Decision drivers satisfied: +1. Support arrays and primitives as SO types +2. Support complex types as SO types +3. Work with `AIAgent` decorators (e.g., `OpenTelemetryAgent`) +4. Enable SO for all AI agents, regardless of whether they natively support it + +Pros: +- The `RunAsync` extension method is easily discoverable. +- The `AIAgent` public API surface remains unchanged. +- No changes required to `AIAgent` decorators. + +Cons: +- Agents without native SO support will still expose `RunAsync`, which may be misleading. +- `ChatClientAgent` exposing `RunAsync` may be misleading when the underlying chat client does not support SO. + +### 3. `RunAsync` as a method of the new `ITypedAIAgent` interface + +This option defines a new `ITypedAIAgent` interface that agents with SO support implement. Agents without SO support do not implement it, allowing users to check for SO capability via interface detection. + +The interface: +```csharp +public interface ITypedAIAgent +{ + Task> RunAsync( + IEnumerable messages, + AgentSession? session = null, + JsonSerializerOptions? serializerOptions = null, + AgentRunOptions? options = null, + CancellationToken cancellationToken = default); + + ... +} +``` + +Agents with SO support implement this interface: +```csharp +public sealed partial class ChatClientAgent : AIAgent, ITypedAIAgent +{ + public async Task> RunAsync( + IEnumerable messages, + AgentSession? session = null, + JsonSerializerOptions? serializerOptions = null, + AgentRunOptions? options = null, + CancellationToken cancellationToken = default) + { + ... + } +} +``` + +However, `ChatClientAgent` presents a challenge: it can work with chat clients that either support or do not support SO. Implementing the interface does not guarantee +the underlying chat client supports SO, which undermines the core idea of using interface detection to determine SO capability. + +Additionally, to allow users to access interface methods on decorated agents, all decorators must implement `ITypedAIAgent`. This makes it difficult for users to +determine whether the underlying agent actually supports SO, further weakening the purpose of this approach. + +Furthermore, users would have to probe the agent type to check if it implements the `ITypedAIAgent` interface and cast it accordingly to access the `RunAsync` methods. +This adds friction to the user experience. A `RunAsync` extension method for `AIAgent` could be provided to alleviate that. + +Given these drawbacks, this option is more complex to implement than the others without providing clear benefits. + +Decision drivers satisfied: +1. Support arrays and primitives as SO types +2. Support complex types as SO types +3. Work with `AIAgent` decorators (e.g., `OpenTelemetryAgent`) +4. Enable SO for all AI agents, regardless of whether they natively support it + +Pros: +- Both the SO decorator and `ChatClientAgent` have compile-time access to the type `T`, allowing them to use the native `IChatClient.GetResponseAsync` API, which handles primitives and collections seamlessly. + +Cons: +- `ChatClientAgent` implementing `ITypedAIAgent` may be misleading when the underlying chat client does not support SO. +- All `AIAgent` decorators must implement `ITypedAIAgent` to handle `RunAsync` calls. +- Decorators implementing the interface may mislead users into thinking the underlying agent natively supports SO. +- Agents must implement all members of `ITypedAIAgent`, not just a core method. +- Users must check the agent type and cast to `ITypedAIAgent` to access `RunAsync`. + +### 4. `RunAsync` as an instance method of `AIAgent` class working via the new `AgentRunOptions.ResponseFormat` property + +This option adds a `ResponseFormat` property of type `ChatResponseFormat` to `AgentRunOptions`. Agents that support SO check for the presence of +this property in the options passed to `RunAsync` to determine whether structured output is requested. If present, they use the schema from `ResponseFormat` +to invoke the underlying API and obtain the SO response. + +```csharp +public class AgentRunOptions +{ + public ChatResponseFormat? ResponseFormat { get; set; } +} +``` + +Additionally, a generic `RunAsync` method is added to `AIAgent` that initializes the `ResponseFormat` based on the type `T` and delegates to the non-generic `RunAsync`. + +```csharp +public abstract class AIAgent +{ + public async Task> RunAsync( + IEnumerable messages, + AgentSession? session = null, + JsonSerializerOptions? serializerOptions = null, + AgentRunOptions? options = null, + CancellationToken cancellationToken = default) + { + serializerOptions ??= AgentAbstractionsJsonUtilities.DefaultOptions; + + var responseFormat = ChatResponseFormat.ForJsonSchema(serializerOptions); + + options = options?.Clone() ?? new AgentRunOptions(); + options.ResponseFormat = responseFormat; + + AgentResponse response = await this.RunAsync(messages, session, options, cancellationToken).ConfigureAwait(false); + + return new AgentResponse(response, serializerOptions); + } +} +``` + +Users call the generic `RunAsync` method directly on the agent: + +```csharp +AIAgent agent = chatClient.AsAIAgent(name: "HelpfulAssistant", instructions: "You are a helpful assistant."); + +AgentResponse response = await agent.RunAsync("Please provide information about John Smith, who is a 35-year-old software engineer."); +``` + +Decision drivers satisfied: +1. Support arrays and primitives as SO types +2. Support complex types as SO types +3. Work with `AIAgent` decorators (e.g., `OpenTelemetryAgent`) +4. Enable SO for all AI agents, regardless of whether they natively support it + +Pros: +- The `AIAgent.RunAsync` method is easily discoverable. +- No changes required to `AIAgent` decorators + +Cons: +- Agents without native SO support will still expose `RunAsync`, which may be misleading. +- `ChatClientAgent` exposing `RunAsync` may be misleading when the underlying chat client does not support SO. + +### Decision Table + +| | Option 1: Instance method + RunCoreAsync | Option 2: Extension method + feature collection | Option 3: ITypedAIAgent Interface | Option 4: Instance method + AgentRunOptions.ResponseFormat | +|---|---|---|---|---| +| Discoverability | ✅ `RunAsync` easily discoverable | ✅ `RunAsync` easily discoverable | ❌ Requires type check and cast | ✅ `RunAsync` easily discoverable | +| Decorator changes | ❌ All decorators must override `RunCoreAsync` | ✅ No changes required | ❌ All decorators must implement `ITypedAIAgent` | ✅ No changes required to decorators | +| Primitives/collections handling | ✅ Native support via `IChatClient.GetResponseAsync` | ❌ Must wrap/unwrap internally | ✅ Native support via `IChatClient.GetResponseAsync` | ❌ Must wrap/unwrap internally | +| Misleading API exposure | ❌ Agents without SO still expose `RunAsync` | ❌ Agents without SO still expose `RunAsync` | ❌ Interface on `ChatClientAgent` may be misleading | ❌ Agents without SO still expose `RunAsync` | +| Implementation burden | ❌ Decorators must override method | ❌ Must handle schema wrapping | ❌ Agents must implement all interface members | ✅ Delegates to existing `RunAsync` via `ResponseFormat` | + +## Cross-Cutting Aspects + +1. **The `useJsonSchemaResponseFormat` parameter**: The `ChatClientAgent.RunAsync` method has this parameter to enable structured output on LLMs that do not natively support it. + It works by adding a user message like "Respond with a JSON value conforming to the following schema:" along with the JSON schema. However, this approach has not been reliable historically. The recommendation is not to carry this parameter forward, regardless of which option is chosen. + +2. **Primitives and array types handling**: There are a few options for how primitive and array types can be handled in the Agent Framework: + + 1. **Never wrap**, regardless of whether the schema is provided via `ResponseFormat` or `RunAsync`. + - Pro: No changes needed; user has full control. + - Pro: No issues with unwrapping in streaming scenarios. + - Con: User must wrap manually. + + 2. **Always wrap**, regardless of whether the schema is provided via `ResponseFormat` or `RunAsync`. + - Pro: Consistent wrapping behavior; no manual wrapping needed. + - Con: Inconsistent unwrapping behavior; it may be unexpected to have SO result wrapped when schema is provided via `ResponseFormat`. + - Con: Impossible to know if SO result is wrapped to unwrap it in streaming scenarios. + + 3. **Wrap only for `RunAsync`** and do not wrap the schema provided via `ResponseFormat`. + - Pro: No unexpectedly wrapped result when schema is provided via `ResponseFormat`. + - Pro: Solves the problem with unwrapping in streaming scenarios. + + 4. **User decides** whether to wrap schema provided via `ResponseFormat` using a new `wrapPrimitivesAndArrays` property of `ChatResponseFormatJson`. For SO provided via `RunAsync`, AF always wraps. + - Pro: No manual wrapping needed; just flip a switch. + - Pro: Solves the problem with unwrapping in streaming scenarios. + - Con: Extends the public API surface. + +3. **Structured output for agents without native SO support**: Some AI agents in AF do not support structured output natively. This is either because it is not part of the protocol (e.g., A2A agent) or because the agents use LLMs without structured output capabilities. + To address this gap, AF can provide the `StructuredOutputAgent` decorator. This decorator wraps any `AIAgent` and adds structured output support by obtaining the text response from the decorated agent and delegating it to a configured chat client for JSON transformation. + + ```csharp + public class StructuredOutputAgent : DelegatingAIAgent + { + private readonly IChatClient _chatClient; + + public StructuredOutputAgent(AIAgent innerAgent, IChatClient chatClient) + : base(innerAgent) + { + this._chatClient = Throw.IfNull(chatClient); + } + + protected override async Task> RunCoreAsync( + IEnumerable messages, + AgentSession? session = null, + JsonSerializerOptions? serializerOptions = null, + AgentRunOptions? options = null, + CancellationToken cancellationToken = default) + { + // Run the inner agent first, to get back the text response we want to convert. + var textResponse = await this.InnerAgent.RunAsync(messages, session, options, cancellationToken).ConfigureAwait(false); + + // Invoke the chat client to transform the text output into structured data. + ChatResponse soResponse = await this._chatClient.GetResponseAsync( + messages: + [ + new ChatMessage(ChatRole.System, "You are a json expert and when provided with any text, will convert it to the requested json format."), + new ChatMessage(ChatRole.User, textResponse.Text) + ], + serializerOptions: serializerOptions ?? AgentJsonUtilities.DefaultOptions, + cancellationToken: cancellationToken).ConfigureAwait(false); + + return new StructuredOutputAgentResponse(soResponse, textResponse); + } + } + ``` + + The decorator preserves the original response from the decorated agent and surfaces it via the `OriginalResponse` property on the returned `StructuredOutputAgentResponse`. + This allows users to access both the original unstructured response and the new structured response when using this decorator. + ```csharp + public class StructuredOutputAgentResponse : AgentResponse + { + internal StructuredOutputAgentResponse(ChatResponse chatResponse, AgentResponse agentResponse) : base(chatResponse) + { + this.OriginalResponse = agentResponse; + } + + public AgentResponse OriginalResponse { get; } + } + ``` + + The decorator can be registered during the agent configuration step using the `UseStructuredOutput` extension method on `AIAgentBuilder`. + + ```csharp + IChatClient meaiChatClient = chatClient.AsIChatClient(); + + AIAgent baseAgent = meaiChatClient.AsAIAgent(name: "HelpfulAssistant", instructions: "You are a helpful assistant."); + + // Register the StructuredOutputAgent decorator during agent building + AIAgent agent = baseAgent + .AsBuilder() + .UseStructuredOutput(meaiChatClient) + .Build(); + + AgentResponse response = await agent.RunAsync("Please provide information about John Smith, who is a 35-year-old software engineer."); + + Console.WriteLine($"Name: {response.Result.Name}"); + Console.WriteLine($"Age: {response.Result.Age}"); + Console.WriteLine($"Occupation: {response.Result.Occupation}"); + + var originalResponse = ((StructuredOutputAgentResponse)response.RawRepresentation!).OriginalResponse; + Console.WriteLine($"Original unstructured response: {originalResponse.Text}"); + + ``` + +## Decision Outcome + +It was decided to keep both approaches for structured output - via `ResponseFormat` and via `RunAsync` since they serve different scenarios and use cases. + +For the `RunAsync` approach, option 4 was selected, which adds a generic `RunAsync` method to `AIAgent` that works via the new `AgentRunOptions.ResponseFormat` property. +This was chosen for its simplicity and because no changes are required to existing `AIAgent` decorators. + +For cross-cutting aspects, the `useJsonSchemaResponseFormat` parameter will not be carried forward due to reliability issues. + +For handling primitives and array types, option 3 was selected: wrap only for `RunAsync` and do not wrap the schema provided via `ResponseFormat`. +This avoids the issues described in the Approach 1 section note. + +Finally, it was decided not to include the `StructuredOutputAgent` decorator in the framework, since the reliability of producing structured output via an additional +LLM call may not be sufficient for all scenarios. Instead, this pattern is provided as a sample to demonstrate how structured output can be achieved for agents without native support, +giving users a reference implementation they can adapt to their own requirements. \ No newline at end of file diff --git a/dotnet/agent-framework-dotnet.slnx b/dotnet/agent-framework-dotnet.slnx index a52e026a8f..e592e80803 100644 --- a/dotnet/agent-framework-dotnet.slnx +++ b/dotnet/agent-framework-dotnet.slnx @@ -389,6 +389,9 @@ + + + diff --git a/dotnet/eng/MSBuild/Shared.props b/dotnet/eng/MSBuild/Shared.props index da8806a1f3..95eb5e13da 100644 --- a/dotnet/eng/MSBuild/Shared.props +++ b/dotnet/eng/MSBuild/Shared.props @@ -20,4 +20,7 @@ + + + diff --git a/dotnet/samples/AGUIClientServer/AGUIDojoServer/SharedState/SharedStateAgent.cs b/dotnet/samples/AGUIClientServer/AGUIDojoServer/SharedState/SharedStateAgent.cs index 51f4791272..af1a54b103 100644 --- a/dotnet/samples/AGUIClientServer/AGUIDojoServer/SharedState/SharedStateAgent.cs +++ b/dotnet/samples/AGUIClientServer/AGUIDojoServer/SharedState/SharedStateAgent.cs @@ -78,7 +78,7 @@ internal sealed class SharedStateAgent : DelegatingAIAgent var response = allUpdates.ToAgentResponse(); - if (response.TryDeserialize(this._jsonSerializerOptions, out JsonElement stateSnapshot)) + if (TryDeserialize(response.Text, this._jsonSerializerOptions, out JsonElement stateSnapshot)) { byte[] stateBytes = JsonSerializer.SerializeToUtf8Bytes( stateSnapshot, @@ -103,4 +103,25 @@ internal sealed class SharedStateAgent : DelegatingAIAgent yield return update; } } + + private static bool TryDeserialize(string json, JsonSerializerOptions jsonSerializerOptions, out T structuredOutput) + { + try + { + T? result = JsonSerializer.Deserialize(json, jsonSerializerOptions); + if (result is null) + { + structuredOutput = default!; + return false; + } + + structuredOutput = result; + return true; + } + catch + { + structuredOutput = default!; + return false; + } + } } diff --git a/dotnet/samples/GettingStarted/AGUI/Step05_StateManagement/Server/SharedStateAgent.cs b/dotnet/samples/GettingStarted/AGUI/Step05_StateManagement/Server/SharedStateAgent.cs index 66e8da6eed..17bde7d215 100644 --- a/dotnet/samples/GettingStarted/AGUI/Step05_StateManagement/Server/SharedStateAgent.cs +++ b/dotnet/samples/GettingStarted/AGUI/Step05_StateManagement/Server/SharedStateAgent.cs @@ -107,7 +107,7 @@ internal sealed class SharedStateAgent : DelegatingAIAgent var response = allUpdates.ToAgentResponse(); // Try to deserialize the structured state response - if (response.TryDeserialize(this._jsonSerializerOptions, out JsonElement stateSnapshot)) + if (TryDeserialize(response.Text, this._jsonSerializerOptions, out JsonElement stateSnapshot)) { // Serialize and emit as STATE_SNAPSHOT via DataContent byte[] stateBytes = JsonSerializer.SerializeToUtf8Bytes( @@ -134,4 +134,25 @@ internal sealed class SharedStateAgent : DelegatingAIAgent yield return update; } } + + private static bool TryDeserialize(string json, JsonSerializerOptions jsonSerializerOptions, out T structuredOutput) + { + try + { + T? deserialized = JsonSerializer.Deserialize(json, jsonSerializerOptions); + if (deserialized is null) + { + structuredOutput = default!; + return false; + } + + structuredOutput = deserialized; + return true; + } + catch + { + structuredOutput = default!; + return false; + } + } } diff --git a/dotnet/samples/GettingStarted/Agents/Agent_Step05_StructuredOutput/AIAgentBuilderExtensions.cs b/dotnet/samples/GettingStarted/Agents/Agent_Step05_StructuredOutput/AIAgentBuilderExtensions.cs new file mode 100644 index 0000000000..987869e175 --- /dev/null +++ b/dotnet/samples/GettingStarted/Agents/Agent_Step05_StructuredOutput/AIAgentBuilderExtensions.cs @@ -0,0 +1,49 @@ +// Copyright (c) Microsoft. All rights reserved. + +using Microsoft.Agents.AI; +using Microsoft.Extensions.AI; +using Microsoft.Extensions.DependencyInjection; + +namespace SampleApp; + +/// +/// Provides extension methods for adding structured output capabilities to instances. +/// +internal static class AIAgentBuilderExtensions +{ + /// + /// Adds structured output capabilities to the agent pipeline, enabling conversion of text responses to structured JSON format. + /// + /// The to which structured output support will be added. + /// + /// The chat client used to transform text responses into structured JSON format. + /// If , the chat client will be resolved from the service provider. + /// + /// + /// An optional factory function that returns the instance to use. + /// This allows for fine-tuning the structured output behavior such as setting the response format or system message. + /// + /// The with structured output capabilities added, enabling method chaining. + /// + /// + /// A must be specified either through the + /// at runtime or the + /// provided during configuration. + /// + /// + public static AIAgentBuilder UseStructuredOutput( + this AIAgentBuilder builder, + IChatClient? chatClient = null, + Func? optionsFactory = null) + { + ArgumentNullException.ThrowIfNull(builder); + + return builder.Use((innerAgent, services) => + { + chatClient ??= services?.GetService() + ?? throw new InvalidOperationException($"No {nameof(IChatClient)} was provided and none could be resolved from the service provider. Either provide an {nameof(IChatClient)} explicitly or register one in the dependency injection container."); + + return new StructuredOutputAgent(innerAgent, chatClient, optionsFactory?.Invoke()); + }); + } +} diff --git a/dotnet/samples/GettingStarted/Agents/Agent_Step05_StructuredOutput/Program.cs b/dotnet/samples/GettingStarted/Agents/Agent_Step05_StructuredOutput/Program.cs index 851b3340d5..7e74315e7d 100644 --- a/dotnet/samples/GettingStarted/Agents/Agent_Step05_StructuredOutput/Program.cs +++ b/dotnet/samples/GettingStarted/Agents/Agent_Step05_StructuredOutput/Program.cs @@ -8,11 +8,13 @@ using System.Text.Json.Serialization; using Azure.AI.OpenAI; using Azure.Identity; using Microsoft.Agents.AI; +using Microsoft.Extensions.AI; using OpenAI.Chat; using SampleApp; +using ChatMessage = Microsoft.Extensions.AI.ChatMessage; -var endpoint = Environment.GetEnvironmentVariable("AZURE_OPENAI_ENDPOINT") ?? throw new InvalidOperationException("AZURE_OPENAI_ENDPOINT is not set."); -var deploymentName = Environment.GetEnvironmentVariable("AZURE_OPENAI_DEPLOYMENT_NAME") ?? "gpt-4o-mini"; +string endpoint = Environment.GetEnvironmentVariable("AZURE_OPENAI_ENDPOINT") ?? throw new InvalidOperationException("AZURE_OPENAI_ENDPOINT is not set."); +string deploymentName = Environment.GetEnvironmentVariable("AZURE_OPENAI_DEPLOYMENT_NAME") ?? "gpt-4o-mini"; // Create chat client to be used by chat client agents. // WARNING: DefaultAzureCredential is convenient for development but requires careful consideration in production. @@ -23,52 +25,159 @@ ChatClient chatClient = new AzureOpenAIClient( new DefaultAzureCredential()) .GetChatClient(deploymentName); -// Create the ChatClientAgent with the specified name and instructions. -ChatClientAgent agent = chatClient.AsAIAgent(name: "HelpfulAssistant", instructions: "You are a helpful assistant."); +// Demonstrates how to work with structured output via ResponseFormat with the non-generic RunAsync method. +// This approach is useful when: +// a. Structured output is used for inter-agent communication, where one agent produces structured output +// and passes it as text to another agent as input, without the need for the caller to directly work with the structured output. +// b. The type of the structured output is not known at compile time, so the generic RunAsync method cannot be used. +// c. The type of the structured output is represented by JSON schema only, without a corresponding class or type in the code. +await UseStructuredOutputWithResponseFormatAsync(chatClient); -// Set PersonInfo as the type parameter of RunAsync method to specify the expected structured output from the agent and invoke the agent with some unstructured input. -AgentResponse response = await agent.RunAsync("Please provide information about John Smith, who is a 35-year-old software engineer."); +// Demonstrates how to work with structured output via the generic RunAsync method. +// This approach is useful when the caller needs to directly work with the structured output in the code +// via an instance of the corresponding class or type and the type is known at compile time. +await UseStructuredOutputWithRunAsync(chatClient); -// Access the structured output via the Result property of the agent response. -Console.WriteLine("Assistant Output:"); -Console.WriteLine($"Name: {response.Result.Name}"); -Console.WriteLine($"Age: {response.Result.Age}"); -Console.WriteLine($"Occupation: {response.Result.Occupation}"); +// Demonstrates how to work with structured output when streaming using the RunStreamingAsync method. +await UseStructuredOutputWithRunStreamingAsync(chatClient); -// Create the ChatClientAgent with the specified name, instructions, and expected structured output the agent should produce. -ChatClientAgent agentWithPersonInfo = chatClient.AsAIAgent(new ChatClientAgentOptions() +// Demonstrates how to add structured output support to agents that don't natively support it using the structured output middleware. +// This approach is useful when working with agents that don't support structured output natively, or agents using models +// that don't have the capability to produce structured output, allowing you to still leverage structured output features by transforming +// the text output from the agent into structured data using a chat client. +await UseStructuredOutputWithMiddlewareAsync(chatClient); + +static async Task UseStructuredOutputWithResponseFormatAsync(ChatClient chatClient) { - Name = "HelpfulAssistant", - ChatOptions = new() { Instructions = "You are a helpful assistant.", ResponseFormat = Microsoft.Extensions.AI.ChatResponseFormat.ForJsonSchema() } -}); + Console.WriteLine("=== Structured Output with ResponseFormat ==="); -// Invoke the agent with some unstructured input while streaming, to extract the structured information from. -var updates = agentWithPersonInfo.RunStreamingAsync("Please provide information about John Smith, who is a 35-year-old software engineer."); + // Create the agent + AIAgent agent = chatClient.AsAIAgent(new ChatClientAgentOptions() + { + Name = "HelpfulAssistant", + ChatOptions = new() + { + Instructions = "You are a helpful assistant.", + // Specify CityInfo as the type parameter of ForJsonSchema to indicate the expected structured output from the agent. + ResponseFormat = Microsoft.Extensions.AI.ChatResponseFormat.ForJsonSchema() + } + }); -// Assemble all the parts of the streamed output, since we can only deserialize once we have the full json, -// then deserialize the response into the PersonInfo class. -PersonInfo personInfo = (await updates.ToAgentResponseAsync()).Deserialize(JsonSerializerOptions.Web); + // Invoke the agent with some unstructured input to extract the structured information from. + AgentResponse response = await agent.RunAsync("Provide information about the capital of France."); -Console.WriteLine("Assistant Output:"); -Console.WriteLine($"Name: {personInfo.Name}"); -Console.WriteLine($"Age: {personInfo.Age}"); -Console.WriteLine($"Occupation: {personInfo.Occupation}"); + // Access the structured output via the Text property of the agent response as JSON in scenarios when JSON as text is required + // and no object instance is needed (e.g., for logging, forwarding to another service, or storing in a database). + Console.WriteLine("Assistant Output (JSON):"); + Console.WriteLine(response.Text); + Console.WriteLine(); + + // Deserialize the JSON text to work with the structured object in scenarios when you need to access properties, + // perform operations, or pass the data to methods that require the typed object instance. + CityInfo cityInfo = JsonSerializer.Deserialize(response.Text)!; + + Console.WriteLine("Assistant Output (Deserialized):"); + Console.WriteLine($"Name: {cityInfo.Name}"); + Console.WriteLine(); +} + +static async Task UseStructuredOutputWithRunAsync(ChatClient chatClient) +{ + Console.WriteLine("=== Structured Output with RunAsync ==="); + + // Create the agent + AIAgent agent = chatClient.AsAIAgent(name: "HelpfulAssistant", instructions: "You are a helpful assistant."); + + // Set CityInfo as the type parameter of RunAsync method to specify the expected structured output from the agent and invoke it with some unstructured input. + AgentResponse response = await agent.RunAsync("Provide information about the capital of France."); + + // Access the structured output via the Result property of the agent response. + CityInfo cityInfo = response.Result; + + Console.WriteLine("Assistant Output:"); + Console.WriteLine($"Name: {cityInfo.Name}"); + Console.WriteLine(); +} + +static async Task UseStructuredOutputWithRunStreamingAsync(ChatClient chatClient) +{ + Console.WriteLine("=== Structured Output with RunStreamingAsync ==="); + + // Create the agent + AIAgent agent = chatClient.AsAIAgent(new ChatClientAgentOptions() + { + Name = "HelpfulAssistant", + ChatOptions = new() + { + Instructions = "You are a helpful assistant.", + // Specify CityInfo as the type parameter of ForJsonSchema to indicate the expected structured output from the agent. + ResponseFormat = Microsoft.Extensions.AI.ChatResponseFormat.ForJsonSchema() + } + }); + + // Invoke the agent with some unstructured input while streaming, to extract the structured information from. + IAsyncEnumerable updates = agent.RunStreamingAsync("Provide information about the capital of France."); + + // Assemble all the parts of the streamed output. + AgentResponse nonGenericResponse = await updates.ToAgentResponseAsync(); + + // Access the structured output by deserializing JSON in the Text property. + CityInfo cityInfo = JsonSerializer.Deserialize(nonGenericResponse.Text)!; + + Console.WriteLine("Assistant Output:"); + Console.WriteLine($"Name: {cityInfo.Name}"); + Console.WriteLine(); +} + +static async Task UseStructuredOutputWithMiddlewareAsync(ChatClient chatClient) +{ + Console.WriteLine("=== Structured Output with UseStructuredOutput Middleware ==="); + + // Create chat client that will transform the agent text response into structured output. + IChatClient meaiChatClient = chatClient.AsIChatClient(); + + // Create the agent + AIAgent agent = meaiChatClient.AsAIAgent(name: "HelpfulAssistant", instructions: "You are a helpful assistant."); + + // Add structured output middleware via UseStructuredOutput method to add structured output support to the agent. + // This middleware transforms the agent's text response into structured data using a chat client. + // Since our agent does support structured output natively, we will add a middleware that removes ResponseFormat + // from the AgentRunOptions to emulate an agent that doesn't support structured output natively + agent = agent + .AsBuilder() + .UseStructuredOutput(meaiChatClient) + .Use(ResponseFormatRemovalMiddleware, null) + .Build(); + + // Set CityInfo as the type parameter of RunAsync method to specify the expected structured output from the agent and invoke it with some unstructured input. + AgentResponse response = await agent.RunAsync("Provide information about the capital of France."); + + // Access the structured output via the Result property of the agent response. + CityInfo cityInfo = response.Result; + + Console.WriteLine("Assistant Output:"); + Console.WriteLine($"Name: {cityInfo.Name}"); + Console.WriteLine(); +} + +static Task ResponseFormatRemovalMiddleware(IEnumerable messages, AgentSession? session, AgentRunOptions? options, AIAgent innerAgent, CancellationToken cancellationToken) +{ + // Remove any ResponseFormat from the options to emulate an agent that doesn't support structured output natively. + options = options?.Clone(); + options?.ResponseFormat = null; + + return innerAgent.RunAsync(messages, session, options, cancellationToken); +} namespace SampleApp { /// - /// Represents information about a person, including their name, age, and occupation, matched to the JSON schema used in the agent. + /// Represents information about a city, including its name. /// - [Description("Information about a person including their name, age, and occupation")] - public class PersonInfo + [Description("Information about a city")] + public sealed class CityInfo { [JsonPropertyName("name")] public string? Name { get; set; } - - [JsonPropertyName("age")] - public int? Age { get; set; } - - [JsonPropertyName("occupation")] - public string? Occupation { get; set; } } } diff --git a/dotnet/samples/GettingStarted/Agents/Agent_Step05_StructuredOutput/README.md b/dotnet/samples/GettingStarted/Agents/Agent_Step05_StructuredOutput/README.md new file mode 100644 index 0000000000..2471e92194 --- /dev/null +++ b/dotnet/samples/GettingStarted/Agents/Agent_Step05_StructuredOutput/README.md @@ -0,0 +1,52 @@ +# Structured Output with ChatClientAgent + +This sample demonstrates how to configure ChatClientAgent to produce structured output in JSON format using various approaches. + +## What this sample demonstrates + +- **ResponseFormat approach**: Configuring agents with JSON schema response format via `ChatResponseFormat.ForJsonSchema()` for inter-agent communication or when the type is not known at compile time +- **Generic RunAsync method**: Using the generic `RunAsync` method for structured output when the caller needs to work directly with typed objects +- **Structured output with Streaming**: Using `RunStreamingAsync` to stream responses while still obtaining structured output by assembling and deserializing the streamed content +- **StructuredOutput middleware**: Adding structured output support to agents that don't natively support it (like A2A agents or models without structured output capability) by transforming text output into structured data using a chat client + +## Prerequisites + +Before you begin, ensure you have the following prerequisites: + +- .NET 10 SDK or later +- Azure OpenAI service endpoint and deployment configured +- Azure CLI installed and authenticated (for Azure credential authentication) +- User has the `Cognitive Services OpenAI Contributor` role for the Azure OpenAI resource + +**Note**: This sample uses Azure OpenAI models. For more information, see [how to deploy Azure OpenAI models with Azure AI Foundry](https://learn.microsoft.com/en-us/azure/ai-foundry/how-to/deploy-models-openai). + +**Note**: This demo uses Azure CLI credentials for authentication. Make sure you're logged in with `az login` and have access to the Azure OpenAI resource and have the `Cognitive Services OpenAI Contributor` role. For more information, see the [Azure CLI documentation](https://learn.microsoft.com/cli/azure/authenticate-azure-cli-interactively). + +## Environment Variables + +Set the following environment variables: + +```powershell +$env:AZURE_OPENAI_ENDPOINT="https://your-resource.openai.azure.com/" # Replace with your Azure OpenAI resource endpoint +$env:AZURE_OPENAI_DEPLOYMENT_NAME="gpt-4o-mini" # Optional, defaults to gpt-4o-mini +``` + +## Run the sample + +Navigate to the sample directory and run: + +```powershell +cd dotnet/samples/GettingStarted/Agents/Agent_Step05_StructuredOutput +dotnet run +``` + +## Expected behavior + +The sample will demonstrate four different approaches to structured output: + +1. **Structured Output with ResponseFormat**: Creates an agent with `ResponseFormat` set to `ForJsonSchema()`, invokes it with unstructured input, and accesses the structured output via the `Text` property +2. **Structured Output with RunAsync**: Creates an agent and uses the generic `RunAsync()` method to get a typed `AgentResponse` with the result accessible via the `Result` property +3. **Structured Output with RunStreamingAsync**: Creates an agent with JSON schema response format, streams the response using `RunStreamingAsync`, assembles the updates using `ToAgentResponseAsync()`, and deserializes the JSON text into a typed object +4. **Structured Output with StructuredOutput Middleware**: Uses the `UseStructuredOutput` method on `AIAgentBuilder` to add structured output support to agents that don't natively support it + +Each approach will output information about the capital of France (Paris) in a structured format. diff --git a/dotnet/samples/GettingStarted/Agents/Agent_Step05_StructuredOutput/StructuredOutputAgent.cs b/dotnet/samples/GettingStarted/Agents/Agent_Step05_StructuredOutput/StructuredOutputAgent.cs new file mode 100644 index 0000000000..641e0adfc4 --- /dev/null +++ b/dotnet/samples/GettingStarted/Agents/Agent_Step05_StructuredOutput/StructuredOutputAgent.cs @@ -0,0 +1,88 @@ +// Copyright (c) Microsoft. All rights reserved. + +using Microsoft.Agents.AI; +using Microsoft.Extensions.AI; + +namespace SampleApp; + +/// +/// A delegating AI agent that converts text responses from an inner AI agent into structured output using a chat client. +/// +/// +/// +/// The wraps an inner agent and uses a chat client to transform +/// the inner agent's text response into a structured JSON format based on the specified response format. +/// +/// +/// This agent requires a to be specified either through the +/// or the +/// provided during construction. +/// +/// +internal sealed class StructuredOutputAgent : DelegatingAIAgent +{ + private readonly IChatClient _chatClient; + private readonly StructuredOutputAgentOptions? _agentOptions; + + /// + /// Initializes a new instance of the class. + /// + /// The underlying agent that generates text responses to be converted to structured output. + /// The chat client used to transform text responses into structured JSON format. + /// Optional configuration options for the structured output agent. + public StructuredOutputAgent(AIAgent innerAgent, IChatClient chatClient, StructuredOutputAgentOptions? options = null) + : base(innerAgent) + { + this._chatClient = chatClient ?? throw new ArgumentNullException(nameof(chatClient)); + this._agentOptions = options; + } + + /// + protected override async Task RunCoreAsync( + IEnumerable messages, + AgentSession? session = null, + AgentRunOptions? options = null, + CancellationToken cancellationToken = default) + { + // Run the inner agent first, to get back the text response we want to convert. + var textResponse = await this.InnerAgent.RunAsync(messages, session, options, cancellationToken).ConfigureAwait(false); + + // Invoke the chat client to transform the text output into structured data. + ChatResponse soResponse = await this._chatClient.GetResponseAsync( + messages: this.GetChatMessages(textResponse.Text), + options: this.GetChatOptions(options), + cancellationToken: cancellationToken).ConfigureAwait(false); + + return new StructuredOutputAgentResponse(soResponse, textResponse); + } + + private List GetChatMessages(string? textResponseText) + { + List chatMessages = []; + + if (this._agentOptions?.ChatClientSystemMessage is not null) + { + chatMessages.Add(new ChatMessage(ChatRole.System, this._agentOptions.ChatClientSystemMessage)); + } + + chatMessages.Add(new ChatMessage(ChatRole.User, textResponseText)); + + return chatMessages; + } + + private ChatOptions GetChatOptions(AgentRunOptions? options) + { + ChatResponseFormat responseFormat = options?.ResponseFormat + ?? this._agentOptions?.ChatOptions?.ResponseFormat + ?? throw new InvalidOperationException($"A response format of type '{nameof(ChatResponseFormatJson)}' must be specified, but none was specified."); + + if (responseFormat is not ChatResponseFormatJson jsonResponseFormat) + { + throw new NotSupportedException($"A response format of type '{nameof(ChatResponseFormatJson)}' must be specified, but was '{responseFormat.GetType().Name}'."); + } + + var chatOptions = this._agentOptions?.ChatOptions?.Clone() ?? new ChatOptions(); + chatOptions.ResponseFormat = jsonResponseFormat; + return chatOptions; + } +} diff --git a/dotnet/samples/GettingStarted/Agents/Agent_Step05_StructuredOutput/StructuredOutputAgentOptions.cs b/dotnet/samples/GettingStarted/Agents/Agent_Step05_StructuredOutput/StructuredOutputAgentOptions.cs new file mode 100644 index 0000000000..c5613d2015 --- /dev/null +++ b/dotnet/samples/GettingStarted/Agents/Agent_Step05_StructuredOutput/StructuredOutputAgentOptions.cs @@ -0,0 +1,31 @@ +// Copyright (c) Microsoft. All rights reserved. + +using Microsoft.Agents.AI; +using Microsoft.Extensions.AI; + +namespace SampleApp; + +/// +/// Represents configuration options for a . +/// +#pragma warning disable CA1812 // Instantiated via AIAgentBuilderExtensions.UseStructuredOutput optionsFactory parameter +internal sealed class StructuredOutputAgentOptions +#pragma warning restore CA1812 +{ + /// + /// Gets or sets the system message to use when invoking the chat client for structured output conversion. + /// + public string? ChatClientSystemMessage { get; set; } + + /// + /// Gets or sets the chat options to use for the structured output conversion by the chat client + /// used by the agent. + /// + /// + /// This property is optional. The should be set to a + /// instance to specify the expected JSON schema for the structured output. + /// Note that if is provided when running the agent, + /// it will take precedence and override the specified here. + /// + public ChatOptions? ChatOptions { get; set; } +} diff --git a/dotnet/samples/GettingStarted/Agents/Agent_Step05_StructuredOutput/StructuredOutputAgentResponse.cs b/dotnet/samples/GettingStarted/Agents/Agent_Step05_StructuredOutput/StructuredOutputAgentResponse.cs new file mode 100644 index 0000000000..c903b9f3ca --- /dev/null +++ b/dotnet/samples/GettingStarted/Agents/Agent_Step05_StructuredOutput/StructuredOutputAgentResponse.cs @@ -0,0 +1,28 @@ +// Copyright (c) Microsoft. All rights reserved. + +using Microsoft.Agents.AI; +using Microsoft.Extensions.AI; + +namespace SampleApp; + +/// +/// Represents an agent response that contains structured output and +/// the original agent response from which the structured output was generated. +/// +internal sealed class StructuredOutputAgentResponse : AgentResponse +{ + /// + /// Initializes a new instance of the class. + /// + /// The containing the structured output. + /// The original from the inner agent. + public StructuredOutputAgentResponse(ChatResponse chatResponse, AgentResponse agentResponse) : base(chatResponse) + { + this.OriginalResponse = agentResponse; + } + + /// + /// Gets the original non-structured response from the inner agent used by chat client to produce the structured output. + /// + public AgentResponse OriginalResponse { get; } +} diff --git a/dotnet/samples/GettingStarted/FoundryAgents/FoundryAgents_Step05_StructuredOutput/Program.cs b/dotnet/samples/GettingStarted/FoundryAgents/FoundryAgents_Step05_StructuredOutput/Program.cs index fc26d09ea7..9d031bc64b 100644 --- a/dotnet/samples/GettingStarted/FoundryAgents/FoundryAgents_Step05_StructuredOutput/Program.cs +++ b/dotnet/samples/GettingStarted/FoundryAgents/FoundryAgents_Step05_StructuredOutput/Program.cs @@ -64,7 +64,8 @@ IAsyncEnumerable updates = agentWithPersonInfo.RunStreaming // Assemble all the parts of the streamed output, since we can only deserialize once we have the full json, // then deserialize the response into the PersonInfo class. -PersonInfo personInfo = (await updates.ToAgentResponseAsync()).Deserialize(JsonSerializerOptions.Web); +PersonInfo personInfo = JsonSerializer.Deserialize((await updates.ToAgentResponseAsync()).Text, JsonSerializerOptions.Web) + ?? throw new InvalidOperationException("Failed to deserialize the streamed response into PersonInfo."); Console.WriteLine("Assistant Output:"); Console.WriteLine($"Name: {personInfo.Name}"); diff --git a/dotnet/samples/GettingStarted/Workflows/_Foundational/08_WriterCriticWorkflow/Program.cs b/dotnet/samples/GettingStarted/Workflows/_Foundational/08_WriterCriticWorkflow/Program.cs index 38bb80dddc..de20cf31ad 100644 --- a/dotnet/samples/GettingStarted/Workflows/_Foundational/08_WriterCriticWorkflow/Program.cs +++ b/dotnet/samples/GettingStarted/Workflows/_Foundational/08_WriterCriticWorkflow/Program.cs @@ -330,7 +330,8 @@ internal sealed class CriticExecutor : Executor // Convert the stream to a response and deserialize the structured output AgentResponse response = await updates.ToAgentResponseAsync(cancellationToken); - CriticDecision decision = response.Deserialize(JsonSerializerOptions.Web); + CriticDecision decision = JsonSerializer.Deserialize(response.Text, JsonSerializerOptions.Web) + ?? throw new JsonException("Failed to deserialize CriticDecision from response text."); Console.WriteLine($"Decision: {(decision.Approved ? "✅ APPROVED" : "❌ NEEDS REVISION")}"); if (!string.IsNullOrEmpty(decision.Feedback)) diff --git a/dotnet/samples/M365Agent/Agents/WeatherForecastAgent.cs b/dotnet/samples/M365Agent/Agents/WeatherForecastAgent.cs index 5968e4eb28..11caea6939 100644 --- a/dotnet/samples/M365Agent/Agents/WeatherForecastAgent.cs +++ b/dotnet/samples/M365Agent/Agents/WeatherForecastAgent.cs @@ -54,7 +54,7 @@ public class WeatherForecastAgent : DelegatingAIAgent // If the agent returned a valid structured output response // we might be able to enhance the response with an adaptive card. - if (response.TryDeserialize(JsonSerializerOptions.Web, out var structuredOutput)) + if (TryDeserialize(response.Text, JsonSerializerOptions.Web, out var structuredOutput)) { var textContentMessage = response.Messages.FirstOrDefault(x => x.Contents.OfType().Any()); if (textContentMessage is not null) @@ -112,4 +112,25 @@ public class WeatherForecastAgent : DelegatingAIAgent }); return card; } + + private static bool TryDeserialize(string json, JsonSerializerOptions jsonSerializerOptions, out T structuredOutput) + { + try + { + T? result = JsonSerializer.Deserialize(json, jsonSerializerOptions); + if (result is null) + { + structuredOutput = default!; + return false; + } + + structuredOutput = result; + return true; + } + catch + { + structuredOutput = default!; + return false; + } + } } diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions/AIAgent.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions/AIAgent.cs index 6258937cd2..6ebdfa7978 100644 --- a/dotnet/src/Microsoft.Agents.AI.Abstractions/AIAgent.cs +++ b/dotnet/src/Microsoft.Agents.AI.Abstractions/AIAgent.cs @@ -22,7 +22,7 @@ namespace Microsoft.Agents.AI; /// may involve multiple agents working together. /// [DebuggerDisplay("{DebuggerDisplay,nq}")] -public abstract class AIAgent +public abstract partial class AIAgent { private static readonly AsyncLocal s_currentContext = new(); diff --git a/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgentStructuredOutput.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions/AIAgentStructuredOutput.cs similarity index 58% rename from dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgentStructuredOutput.cs rename to dotnet/src/Microsoft.Agents.AI.Abstractions/AIAgentStructuredOutput.cs index d933939884..f93b43157c 100644 --- a/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgentStructuredOutput.cs +++ b/dotnet/src/Microsoft.Agents.AI.Abstractions/AIAgentStructuredOutput.cs @@ -11,155 +11,130 @@ using Microsoft.Shared.Diagnostics; namespace Microsoft.Agents.AI; /// -/// Provides an that delegates to an implementation. +/// Provides structured output methods for that enable requesting responses in a specific type format. /// -public sealed partial class ChatClientAgent +public abstract partial class AIAgent { /// /// Run the agent with no message assuming that all required instructions are already provided to the agent or on the session, and requesting a response of the specified type . /// + /// The type of structured output to request. /// /// The conversation session to use for this invocation. If , a new session will be created. /// The session will be updated with any response messages generated during invocation. /// - /// The JSON serialization options to use. + /// Optional JSON serializer options to use for deserializing the response. /// Optional configuration parameters for controlling the agent's invocation behavior. - /// - /// to set a JSON schema on the ; otherwise, . The default is . - /// Using a JSON schema improves reliability if the underlying model supports native structured output with a schema, but might cause an error if the model does not support it. - /// /// The to monitor for cancellation requests. The default is . - /// A task that represents the asynchronous operation. The task result contains an with the agent's output. + /// A task that represents the asynchronous operation. The task result contains an with the agent's output. /// /// This overload is useful when the agent has sufficient context from previous messages in the session /// or from its initial configuration to generate a meaningful response without additional input. /// - public Task> RunAsync( + public Task> RunAsync( AgentSession? session = null, JsonSerializerOptions? serializerOptions = null, AgentRunOptions? options = null, - bool? useJsonSchemaResponseFormat = null, CancellationToken cancellationToken = default) => - this.RunAsync([], session, serializerOptions, options, useJsonSchemaResponseFormat, cancellationToken); + this.RunAsync([], session, serializerOptions, options, cancellationToken); /// /// Runs the agent with a text message from the user, requesting a response of the specified type . /// + /// The type of structured output to request. /// The user message to send to the agent. /// /// The conversation session to use for this invocation. If , a new session will be created. /// The session will be updated with the input message and any response messages generated during invocation. /// - /// The JSON serialization options to use. + /// Optional JSON serializer options to use for deserializing the response. /// Optional configuration parameters for controlling the agent's invocation behavior. - /// - /// to set a JSON schema on the ; otherwise, . The default is . - /// Using a JSON schema improves reliability if the underlying model supports native structured output with a schema, but might cause an error if the model does not support it. - /// /// The to monitor for cancellation requests. The default is . - /// A task that represents the asynchronous operation. The task result contains an with the agent's output. + /// A task that represents the asynchronous operation. The task result contains an with the agent's output. /// is , empty, or contains only whitespace. /// /// The provided text will be wrapped in a with the role /// before being sent to the agent. This is a convenience method for simple text-based interactions. /// - public Task> RunAsync( + public Task> RunAsync( string message, AgentSession? session = null, JsonSerializerOptions? serializerOptions = null, AgentRunOptions? options = null, - bool? useJsonSchemaResponseFormat = null, CancellationToken cancellationToken = default) { _ = Throw.IfNullOrWhitespace(message); - return this.RunAsync(new ChatMessage(ChatRole.User, message), session, serializerOptions, options, useJsonSchemaResponseFormat, cancellationToken); + return this.RunAsync(new ChatMessage(ChatRole.User, message), session, serializerOptions, options, cancellationToken); } /// /// Runs the agent with a single chat message, requesting a response of the specified type . /// + /// The type of structured output to request. /// The chat message to send to the agent. /// /// The conversation session to use for this invocation. If , a new session will be created. /// The session will be updated with the input message and any response messages generated during invocation. /// - /// The JSON serialization options to use. + /// Optional JSON serializer options to use for deserializing the response. /// Optional configuration parameters for controlling the agent's invocation behavior. - /// - /// to set a JSON schema on the ; otherwise, . The default is . - /// Using a JSON schema improves reliability if the underlying model supports native structured output with a schema, but might cause an error if the model does not support it. - /// /// The to monitor for cancellation requests. The default is . - /// A task that represents the asynchronous operation. The task result contains an with the agent's output. + /// A task that represents the asynchronous operation. The task result contains an with the agent's output. /// is . - public Task> RunAsync( + public Task> RunAsync( ChatMessage message, AgentSession? session = null, JsonSerializerOptions? serializerOptions = null, AgentRunOptions? options = null, - bool? useJsonSchemaResponseFormat = null, CancellationToken cancellationToken = default) { _ = Throw.IfNull(message); - return this.RunAsync([message], session, serializerOptions, options, useJsonSchemaResponseFormat, cancellationToken); + return this.RunAsync([message], session, serializerOptions, options, cancellationToken); } /// /// Runs the agent with a collection of chat messages, requesting a response of the specified type . /// + /// The type of structured output to request. /// The collection of messages to send to the agent for processing. /// /// The conversation session to use for this invocation. If , a new session will be created. /// The session will be updated with the input messages and any response messages generated during invocation. /// - /// The JSON serialization options to use. + /// Optional JSON serializer options to use for deserializing the response. /// Optional configuration parameters for controlling the agent's invocation behavior. - /// - /// to set a JSON schema on the ; otherwise, . The default is . - /// Using a JSON schema improves reliability if the underlying model supports native structured output with a schema, but might cause an error if the model does not support it. - /// /// The to monitor for cancellation requests. The default is . - /// A task that represents the asynchronous operation. The task result contains an with the agent's output. - /// The type of structured output to request. + /// A task that represents the asynchronous operation. The task result contains an with the agent's output. /// /// - /// This is the primary invocation method that implementations must override. It handles collections of messages, - /// allowing for complex conversational scenarios including multi-turn interactions, function calls, and - /// context-rich conversations. + /// This method handles collections of messages, allowing for complex conversational scenarios including + /// multi-turn interactions, function calls, and context-rich conversations. /// /// /// The messages are processed in the order provided and become part of the conversation history. /// The agent's response will also be added to if one is provided. /// /// - public Task> RunAsync( + public async Task> RunAsync( IEnumerable messages, AgentSession? session = null, JsonSerializerOptions? serializerOptions = null, AgentRunOptions? options = null, - bool? useJsonSchemaResponseFormat = null, CancellationToken cancellationToken = default) { - async Task> GetResponseAsync(IChatClient chatClient, List threadMessages, ChatOptions? chatOptions, CancellationToken ct) - { - return await chatClient.GetResponseAsync( - threadMessages, - serializerOptions ?? AgentJsonUtilities.DefaultOptions, - chatOptions, - useJsonSchemaResponseFormat, - ct).ConfigureAwait(false); - } + serializerOptions ??= AgentAbstractionsJsonUtilities.DefaultOptions; - static ChatClientAgentResponse CreateResponse(ChatResponse chatResponse) - { - return new ChatClientAgentResponse(chatResponse) - { - ContinuationToken = WrapContinuationToken(chatResponse.ContinuationToken) - }; - } + var responseFormat = ChatResponseFormat.ForJsonSchema(serializerOptions); - return this.RunCoreAsync(GetResponseAsync, CreateResponse, messages, session, options, cancellationToken); + (responseFormat, bool isWrappedInObject) = StructuredOutputSchemaUtilities.WrapNonObjectSchema(responseFormat); + + options = options?.Clone() ?? new AgentRunOptions(); + options.ResponseFormat = responseFormat; + + AgentResponse response = await this.RunAsync(messages, session, options, cancellationToken).ConfigureAwait(false); + + return new AgentResponse(response, serializerOptions) { IsWrappedInObject = isWrappedInObject }; } } diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions/AgentResponse.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions/AgentResponse.cs index c93c8e9184..168f607491 100644 --- a/dotnet/src/Microsoft.Agents.AI.Abstractions/AgentResponse.cs +++ b/dotnet/src/Microsoft.Agents.AI.Abstractions/AgentResponse.cs @@ -1,20 +1,11 @@ // Copyright (c) Microsoft. All rights reserved. using System; -#if NET -using System.Buffers; -#endif using System.Collections.Generic; using System.Diagnostics.CodeAnalysis; - -#if NET -using System.Text; -#endif -using System.Text.Json; using System.Text.Json.Serialization; -using System.Text.Json.Serialization.Metadata; -using Microsoft.Shared.Diagnostics; using Microsoft.Extensions.AI; +using Microsoft.Shared.Diagnostics; namespace Microsoft.Agents.AI; @@ -76,6 +67,29 @@ public class AgentResponse this.ContinuationToken = response.ContinuationToken; } + /// + /// Initializes a new instance of the class from an existing . + /// + /// The from which to copy properties. + /// is . + /// + /// This constructor creates a copy of an existing agent response, preserving all + /// metadata and storing the original response in for access to + /// the underlying implementation details. + /// + protected AgentResponse(AgentResponse response) + { + _ = Throw.IfNull(response); + + this.AdditionalProperties = response.AdditionalProperties; + this.CreatedAt = response.CreatedAt; + this.Messages = response.Messages; + this.RawRepresentation = response; + this.ResponseId = response.ResponseId; + this.Usage = response.Usage; + this.ContinuationToken = response.ContinuationToken; + } + /// /// Initializes a new instance of the class with the specified collection of messages. /// @@ -274,117 +288,4 @@ public class AgentResponse return updates; } - - /// - /// Deserializes the response text into the given type. - /// - /// The output type to deserialize into. - /// The result as the requested type. - /// The result is not parsable into the requested type. - public T Deserialize() => - this.Deserialize(AgentAbstractionsJsonUtilities.DefaultOptions); - - /// - /// Deserializes the response text into the given type using the specified serializer options. - /// - /// The output type to deserialize into. - /// The JSON serialization options to use. - /// The result as the requested type. - /// The result is not parsable into the requested type. - public T Deserialize(JsonSerializerOptions serializerOptions) - { - _ = Throw.IfNull(serializerOptions); - - var structuredOutput = this.GetResultCore(serializerOptions, out var failureReason); - return failureReason switch - { - FailureReason.ResultDidNotContainJson => throw new InvalidOperationException("The response did not contain JSON to be deserialized."), - FailureReason.DeserializationProducedNull => throw new InvalidOperationException("The deserialized response is null."), - _ => structuredOutput!, - }; - } - - /// - /// Tries to deserialize response text into the given type. - /// - /// The output type to deserialize into. - /// The parsed structured output. - /// if parsing was successful; otherwise, . - public bool TryDeserialize([NotNullWhen(true)] out T? structuredOutput) => - this.TryDeserialize(AgentAbstractionsJsonUtilities.DefaultOptions, out structuredOutput); - - /// - /// Tries to deserialize response text into the given type using the specified serializer options. - /// - /// The output type to deserialize into. - /// The JSON serialization options to use. - /// The parsed structured output. - /// if parsing was successful; otherwise, . - public bool TryDeserialize(JsonSerializerOptions serializerOptions, [NotNullWhen(true)] out T? structuredOutput) - { - _ = Throw.IfNull(serializerOptions); - - try - { - structuredOutput = this.GetResultCore(serializerOptions, out var failureReason); - return failureReason is null; - } - catch - { - structuredOutput = default; - return false; - } - } - - private static T? DeserializeFirstTopLevelObject(string json, JsonTypeInfo typeInfo) - { -#if NET - // We need to deserialize only the first top-level object as a workaround for a common LLM backend - // issue. GPT 3.5 Turbo commonly returns multiple top-level objects after doing a function call. - // See https://community.openai.com/t/2-json-objects-returned-when-using-function-calling-and-json-mode/574348 - var utf8ByteLength = Encoding.UTF8.GetByteCount(json); - var buffer = ArrayPool.Shared.Rent(utf8ByteLength); - try - { - var utf8SpanLength = Encoding.UTF8.GetBytes(json, 0, json.Length, buffer, 0); - var reader = new Utf8JsonReader(new ReadOnlySpan(buffer, 0, utf8SpanLength), new() { AllowMultipleValues = true }); - return JsonSerializer.Deserialize(ref reader, typeInfo); - } - finally - { - ArrayPool.Shared.Return(buffer); - } -#else - return JsonSerializer.Deserialize(json, typeInfo); -#endif - } - - private T? GetResultCore(JsonSerializerOptions serializerOptions, out FailureReason? failureReason) - { - var json = this.Text; - if (string.IsNullOrEmpty(json)) - { - failureReason = FailureReason.ResultDidNotContainJson; - return default; - } - - // If there's an exception here, we want it to propagate, since the Result property is meant to throw directly - - T? deserialized = DeserializeFirstTopLevelObject(json!, (JsonTypeInfo)serializerOptions.GetTypeInfo(typeof(T))); - - if (deserialized is null) - { - failureReason = FailureReason.DeserializationProducedNull; - return default; - } - - failureReason = default; - return deserialized; - } - - private enum FailureReason - { - ResultDidNotContainJson, - DeserializationProducedNull - } } diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions/AgentResponse{T}.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions/AgentResponse{T}.cs index 2a18aadb37..7f12aaed5f 100644 --- a/dotnet/src/Microsoft.Agents.AI.Abstractions/AgentResponse{T}.cs +++ b/dotnet/src/Microsoft.Agents.AI.Abstractions/AgentResponse{T}.cs @@ -1,6 +1,17 @@ // Copyright (c) Microsoft. All rights reserved. -using Microsoft.Extensions.AI; +using System; +#if NET +using System.Buffers; +#endif + +#if NET +using System.Text; +#endif +using System.Text.Json; +using System.Text.Json.Serialization; +using System.Text.Json.Serialization.Metadata; +using Microsoft.Shared.Diagnostics; namespace Microsoft.Agents.AI; @@ -8,23 +19,80 @@ namespace Microsoft.Agents.AI; /// Represents the response of the specified type to an run request. /// /// The type of value expected from the agent. -public abstract class AgentResponse : AgentResponse +public class AgentResponse : AgentResponse { - /// Initializes a new instance of the class. - protected AgentResponse() + private readonly JsonSerializerOptions _serializerOptions; + + /// + /// Initializes a new instance of the class. + /// + /// The from which to populate this . + /// The to use when deserializing the result. + /// is . + public AgentResponse(AgentResponse response, JsonSerializerOptions serializerOptions) : base(response) { + _ = Throw.IfNull(serializerOptions); + + this._serializerOptions = serializerOptions; } /// - /// Initializes a new instance of the class from an existing . + /// Gets or sets a value indicating whether the JSON schema has an extra object wrapper. /// - /// The from which to populate this . - protected AgentResponse(ChatResponse response) : base(response) - { - } + /// + /// The wrapper is required for any non-JSON-object-typed values such as numbers, enum values, and arrays. + /// + public bool IsWrappedInObject { get; init; } /// /// Gets the result value of the agent response as an instance of . /// - public abstract T Result { get; } + [JsonIgnore] + public virtual T Result + { + get + { + var json = this.Text; + if (string.IsNullOrEmpty(json)) + { + throw new InvalidOperationException("The response did not contain JSON to be deserialized."); + } + + if (this.IsWrappedInObject) + { + json = StructuredOutputSchemaUtilities.UnwrapResponseData(json!); + } + + T? deserialized = DeserializeFirstTopLevelObject(json!, (JsonTypeInfo)this._serializerOptions.GetTypeInfo(typeof(T))); + if (deserialized is null) + { + throw new InvalidOperationException("The deserialized response is null."); + } + + return deserialized; + } + } + + private static T? DeserializeFirstTopLevelObject(string json, JsonTypeInfo typeInfo) + { +#if NET + // We need to deserialize only the first top-level object as a workaround for a common LLM backend + // issue. GPT 3.5 Turbo commonly returns multiple top-level objects after doing a function call. + // See https://community.openai.com/t/2-json-objects-returned-when-using-function-calling-and-json-mode/574348 + var utf8ByteLength = Encoding.UTF8.GetByteCount(json); + var buffer = ArrayPool.Shared.Rent(utf8ByteLength); + try + { + var utf8SpanLength = Encoding.UTF8.GetBytes(json, 0, json.Length, buffer, 0); + var reader = new Utf8JsonReader(new ReadOnlySpan(buffer, 0, utf8SpanLength), new() { AllowMultipleValues = true }); + return JsonSerializer.Deserialize(ref reader, typeInfo); + } + finally + { + ArrayPool.Shared.Return(buffer); + } +#else + return JsonSerializer.Deserialize(json, typeInfo); +#endif + } } diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions/AgentRunOptions.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions/AgentRunOptions.cs index 611d15036c..08811df288 100644 --- a/dotnet/src/Microsoft.Agents.AI.Abstractions/AgentRunOptions.cs +++ b/dotnet/src/Microsoft.Agents.AI.Abstractions/AgentRunOptions.cs @@ -28,12 +28,13 @@ public class AgentRunOptions /// /// The options instance from which to copy values. /// is . - public AgentRunOptions(AgentRunOptions options) + protected AgentRunOptions(AgentRunOptions options) { _ = Throw.IfNull(options); this.ContinuationToken = options.ContinuationToken; this.AllowBackgroundResponses = options.AllowBackgroundResponses; this.AdditionalProperties = options.AdditionalProperties?.Clone(); + this.ResponseFormat = options.ResponseFormat; } /// @@ -90,4 +91,35 @@ public class AgentRunOptions /// preserving implementation-specific details or extending the options with custom data. /// public AdditionalPropertiesDictionary? AdditionalProperties { get; set; } + + /// + /// Gets or sets the response format. + /// + /// + /// If , no response format is specified and the agent will use its default. + /// This property can be set to to specify that the response should be unstructured text, + /// to to specify that the response should be structured JSON data, or + /// an instance of constructed with a specific JSON schema to request that the + /// response be structured JSON data according to that schema. It is up to the agent implementation if or how + /// to honor the request. If the agent implementation doesn't recognize the specific kind of , + /// it can be ignored. + /// + public ChatResponseFormat? ResponseFormat { get; set; } + + /// + /// Produces a clone of the current instance. + /// + /// + /// A clone of the current instance. + /// + /// + /// + /// The clone will have the same values for all properties as the original instance. Any collections, like , + /// are shallow-cloned, meaning a new collection instance is created, but any references contained by the collections are shared with the original. + /// + /// + /// Derived types should override to return an instance of the derived type. + /// + /// + public virtual AgentRunOptions Clone() => new(this); } diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions/Microsoft.Agents.AI.Abstractions.csproj b/dotnet/src/Microsoft.Agents.AI.Abstractions/Microsoft.Agents.AI.Abstractions.csproj index 6b6f9d44f2..0f7ae06530 100644 --- a/dotnet/src/Microsoft.Agents.AI.Abstractions/Microsoft.Agents.AI.Abstractions.csproj +++ b/dotnet/src/Microsoft.Agents.AI.Abstractions/Microsoft.Agents.AI.Abstractions.csproj @@ -8,6 +8,7 @@ true + true true true true diff --git a/dotnet/src/Microsoft.Agents.AI.DurableTask/CHANGELOG.md b/dotnet/src/Microsoft.Agents.AI.DurableTask/CHANGELOG.md index ff886e2ebe..b0124d3de7 100644 --- a/dotnet/src/Microsoft.Agents.AI.DurableTask/CHANGELOG.md +++ b/dotnet/src/Microsoft.Agents.AI.DurableTask/CHANGELOG.md @@ -13,6 +13,7 @@ - Introduce Core method pattern for Session management methods on AIAgent ([#3699](https://github.com/microsoft/agent-framework/pull/3699)) - Changed AIAgent.SerializeSession to AIAgent.SerializeSessionAsync ([#3879](https://github.com/microsoft/agent-framework/pull/3879)) - Changed ChatHistory and AIContext Providers to have pipeline semantics ([#3806](https://github.com/microsoft/agent-framework/pull/3806)) +- Marked all `RunAsync` overloads as `new`, added missing ones, and added support for primitives and arrays ([#3803](https://github.com/microsoft/agent-framework/pull/3803)) ## v1.0.0-preview.251204.1 diff --git a/dotnet/src/Microsoft.Agents.AI.DurableTask/DurableAIAgent.cs b/dotnet/src/Microsoft.Agents.AI.DurableTask/DurableAIAgent.cs index c790222e50..ee0c457deb 100644 --- a/dotnet/src/Microsoft.Agents.AI.DurableTask/DurableAIAgent.cs +++ b/dotnet/src/Microsoft.Agents.AI.DurableTask/DurableAIAgent.cs @@ -1,12 +1,11 @@ // Copyright (c) Microsoft. All rights reserved. -using System.Diagnostics.CodeAnalysis; using System.Runtime.CompilerServices; using System.Text.Json; -using System.Text.Json.Serialization.Metadata; using Microsoft.DurableTask; using Microsoft.DurableTask.Entities; using Microsoft.Extensions.AI; +using Microsoft.Shared.Diagnostics; namespace Microsoft.Agents.AI.DurableTask; @@ -114,7 +113,6 @@ public sealed class DurableAIAgent : AIAgent { enableToolCalls = durableOptions.EnableToolCalls; enableToolNames = durableOptions.EnableToolNames; - responseFormat = durableOptions.ResponseFormat; } else if (options is ChatClientAgentRunOptions chatClientOptions && chatClientOptions.ChatOptions?.Tools != null) { @@ -122,6 +120,12 @@ public sealed class DurableAIAgent : AIAgent responseFormat = chatClientOptions.ChatOptions?.ResponseFormat; } + // Override the response format if specified in the agent run options + if (options?.ResponseFormat is { } format) + { + responseFormat = format; + } + RunRequest request = new([.. messages], responseFormat, enableToolCalls, enableToolNames) { OrchestrationId = this._context.InstanceId @@ -168,108 +172,125 @@ public sealed class DurableAIAgent : AIAgent } /// - /// Runs the agent with a message and returns the deserialized output as an instance of . + /// Run the agent with no message assuming that all required instructions are already provided to the agent or on the session, and requesting a response of the specified type . /// - /// The message to send to the agent. - /// The agent session to use. - /// Optional JSON serializer options. - /// Optional run options. - /// The cancellation token. - /// The type of the output. - /// - /// Thrown when the provided already contains a response schema. - /// Thrown when the provided is not a . - /// - /// - /// Thrown when the agent response is empty or cannot be deserialized. - /// - /// The output from the agent. - public async Task> RunAsync( + /// The type of structured output to request. + /// + /// The conversation session to use for this invocation. If , a new session will be created. + /// The session will be updated with any response messages generated during invocation. + /// + /// Optional JSON serializer options to use for deserializing the response. + /// Optional configuration parameters for controlling the agent's invocation behavior. + /// The to monitor for cancellation requests. The default is . + /// A task that represents the asynchronous operation. The task result contains an with the agent's output. + /// + /// This method is specific to durable agents because the Durable Task Framework uses a custom + /// synchronization context for orchestration execution, and all continuations must run on the + /// orchestration thread to avoid breaking the durable orchestration and potential deadlocks. + /// + public new Task> RunAsync( + AgentSession? session = null, + JsonSerializerOptions? serializerOptions = null, + AgentRunOptions? options = null, + CancellationToken cancellationToken = default) => + this.RunAsync([], session, serializerOptions, options, cancellationToken); + + /// + /// Runs the agent with a text message from the user, requesting a response of the specified type . + /// + /// The type of structured output to request. + /// The user message to send to the agent. + /// + /// The conversation session to use for this invocation. If , a new session will be created. + /// The session will be updated with the input message and any response messages generated during invocation. + /// + /// Optional JSON serializer options to use for deserializing the response. + /// Optional configuration parameters for controlling the agent's invocation behavior. + /// The to monitor for cancellation requests. The default is . + /// A task that represents the asynchronous operation. The task result contains an with the agent's output. + /// is , empty, or contains only whitespace. + /// + /// + /// + public new Task> RunAsync( string message, AgentSession? session = null, JsonSerializerOptions? serializerOptions = null, AgentRunOptions? options = null, CancellationToken cancellationToken = default) { - return await this.RunAsync( - messages: [new ChatMessage(ChatRole.User, message) { CreatedAt = DateTimeOffset.UtcNow }], - session, - serializerOptions, - options, - cancellationToken); + _ = Throw.IfNull(message); + + return this.RunAsync(new ChatMessage(ChatRole.User, message), session, serializerOptions, options, cancellationToken); } /// - /// Runs the agent with messages and returns the deserialized output as an instance of . + /// Runs the agent with a single chat message, requesting a response of the specified type . /// - /// The messages to send to the agent. - /// The agent session to use. - /// Optional JSON serializer options. - /// Optional run options. - /// The cancellation token. - /// The type of the output. - /// - /// Thrown when the provided already contains a response schema. - /// Thrown when the provided is not a . - /// - /// - /// Thrown when the agent response is empty or cannot be deserialized. - /// - /// The output from the agent. - [UnconditionalSuppressMessage("Trimming", "IL2026", Justification = "Fallback to reflection-based deserialization is intentional for library flexibility with user-defined types.")] - [UnconditionalSuppressMessage("ReflectionAnalysis", "IL3050", Justification = "Fallback to reflection-based deserialization is intentional for library flexibility with user-defined types.")] - public async Task> RunAsync( + /// The type of structured output to request. + /// The chat message to send to the agent. + /// + /// The conversation session to use for this invocation. If , a new session will be created. + /// The session will be updated with the input message and any response messages generated during invocation. + /// + /// Optional JSON serializer options to use for deserializing the response. + /// Optional configuration parameters for controlling the agent's invocation behavior. + /// The to monitor for cancellation requests. The default is . + /// A task that represents the asynchronous operation. The task result contains an with the agent's output. + /// is . + /// + /// + /// + public new Task> RunAsync( + ChatMessage message, + AgentSession? session = null, + JsonSerializerOptions? serializerOptions = null, + AgentRunOptions? options = null, + CancellationToken cancellationToken = default) + { + _ = Throw.IfNull(message); + + return this.RunAsync([message], session, serializerOptions, options, cancellationToken); + } + + /// + /// Runs the agent with a collection of chat messages, requesting a response of the specified type . + /// + /// The type of structured output to request. + /// The collection of messages to send to the agent for processing. + /// + /// The conversation session to use for this invocation. If , a new session will be created. + /// The session will be updated with the input messages and any response messages generated during invocation. + /// + /// Optional JSON serializer options to use for deserializing the response. + /// Optional configuration parameters for controlling the agent's invocation behavior. + /// The to monitor for cancellation requests. The default is . + /// A task that represents the asynchronous operation. The task result contains an with the agent's output. + /// + /// + /// + public new async Task> RunAsync( IEnumerable messages, AgentSession? session = null, JsonSerializerOptions? serializerOptions = null, AgentRunOptions? options = null, CancellationToken cancellationToken = default) { - options ??= new DurableAgentRunOptions(); - if (options is not DurableAgentRunOptions durableOptions) - { - throw new ArgumentException( - "Response schema is only supported with DurableAgentRunOptions when using durable agents. " + - "Cannot specify a response schema when calling RunAsync.", - paramName: nameof(options)); - } + serializerOptions ??= AgentAbstractionsJsonUtilities.DefaultOptions; - if (durableOptions.ResponseFormat is not null) - { - throw new ArgumentException( - "A response schema is already defined in the provided DurableAgentRunOptions. " + - "Cannot specify a response schema when calling RunAsync.", - paramName: nameof(options)); - } + var responseFormat = ChatResponseFormat.ForJsonSchema(serializerOptions); - // Create the JSON schema for the response type - durableOptions.ResponseFormat = ChatResponseFormat.ForJsonSchema(); + (responseFormat, bool isWrappedInObject) = StructuredOutputSchemaUtilities.WrapNonObjectSchema(responseFormat); - AgentResponse response = await this.RunAsync(messages, session, durableOptions, cancellationToken); + options = options?.Clone() ?? new DurableAgentRunOptions(); + options.ResponseFormat = responseFormat; - // Deserialize the response text to the requested type - if (string.IsNullOrEmpty(response.Text)) - { - throw new InvalidOperationException("Agent response is empty and cannot be deserialized."); - } + // ConfigureAwait(false) cannot be used here because the Durable Task Framework uses + // a custom synchronization context that requires all continuations to execute on the + // orchestration thread. Scheduling the continuation on an arbitrary thread would break + // the orchestration. + AgentResponse response = await this.RunAsync(messages, session, options, cancellationToken); - serializerOptions ??= DurableAgentJsonUtilities.DefaultOptions; - - // Prefer source-generated metadata when available to support AOT/trimming scenarios. - // Fallback to reflection-based deserialization for types without source-generated metadata. - // This is necessary since T is a user-provided type that may not have [JsonSerializable] coverage. - JsonTypeInfo? typeInfo = serializerOptions.GetTypeInfo(typeof(T)); - T? result = (typeInfo is JsonTypeInfo typedInfo - ? (T?)JsonSerializer.Deserialize(response.Text, typedInfo) - : JsonSerializer.Deserialize(response.Text, serializerOptions)) - ?? throw new InvalidOperationException($"Failed to deserialize agent response to type {typeof(T).Name}."); - - return new DurableAIAgentResponse(response, result); - } - - private sealed class DurableAIAgentResponse(AgentResponse response, T result) - : AgentResponse(response.AsChatResponse()) - { - public override T Result { get; } = result; + return new AgentResponse(response, serializerOptions) { IsWrappedInObject = isWrappedInObject }; } } diff --git a/dotnet/src/Microsoft.Agents.AI.DurableTask/DurableAIAgentProxy.cs b/dotnet/src/Microsoft.Agents.AI.DurableTask/DurableAIAgentProxy.cs index f0f7a4ffd4..8987343c60 100644 --- a/dotnet/src/Microsoft.Agents.AI.DurableTask/DurableAIAgentProxy.cs +++ b/dotnet/src/Microsoft.Agents.AI.DurableTask/DurableAIAgentProxy.cs @@ -62,7 +62,6 @@ internal class DurableAIAgentProxy(string name, IDurableAgentClient agentClient) { enableToolCalls = durableOptions.EnableToolCalls; enableToolNames = durableOptions.EnableToolNames; - responseFormat = durableOptions.ResponseFormat; isFireAndForget = durableOptions.IsFireAndForget; } else if (options is ChatClientAgentRunOptions chatClientOptions) @@ -71,6 +70,12 @@ internal class DurableAIAgentProxy(string name, IDurableAgentClient agentClient) responseFormat = chatClientOptions.ChatOptions?.ResponseFormat; } + // Override the response format if specified in the agent run options + if (options?.ResponseFormat is { } format) + { + responseFormat = format; + } + RunRequest request = new([.. messages], responseFormat, enableToolCalls, enableToolNames); AgentSessionId sessionId = durableSession.SessionId; diff --git a/dotnet/src/Microsoft.Agents.AI.DurableTask/DurableAgentRunOptions.cs b/dotnet/src/Microsoft.Agents.AI.DurableTask/DurableAgentRunOptions.cs index 0f1984ad62..f698eab3d8 100644 --- a/dotnet/src/Microsoft.Agents.AI.DurableTask/DurableAgentRunOptions.cs +++ b/dotnet/src/Microsoft.Agents.AI.DurableTask/DurableAgentRunOptions.cs @@ -1,7 +1,5 @@ // Copyright (c) Microsoft. All rights reserved. -using Microsoft.Extensions.AI; - namespace Microsoft.Agents.AI.DurableTask; /// @@ -9,6 +7,25 @@ namespace Microsoft.Agents.AI.DurableTask; /// public sealed class DurableAgentRunOptions : AgentRunOptions { + /// + /// Initializes a new instance of the class. + /// + public DurableAgentRunOptions() + { + } + + /// + /// Initializes a new instance of the class by copying values from the specified options. + /// + /// The options instance from which to copy values. + private DurableAgentRunOptions(DurableAgentRunOptions options) + : base(options) + { + this.EnableToolCalls = options.EnableToolCalls; + this.EnableToolNames = options.EnableToolNames is not null ? new List(options.EnableToolNames) : null; + this.IsFireAndForget = options.IsFireAndForget; + } + /// /// Gets or sets whether to enable tool calls for this request. /// @@ -19,11 +36,6 @@ public sealed class DurableAgentRunOptions : AgentRunOptions /// public IList? EnableToolNames { get; set; } - /// - /// Gets or sets the response format for the agent's response. - /// - public ChatResponseFormat? ResponseFormat { get; set; } - /// /// Gets or sets whether to fire and forget the agent run request. /// @@ -33,4 +45,7 @@ public sealed class DurableAgentRunOptions : AgentRunOptions /// long-running tasks where the caller does not need to wait for the agent to complete the run. /// public bool IsFireAndForget { get; set; } + + /// + public override AgentRunOptions Clone() => new DurableAgentRunOptions(this); } diff --git a/dotnet/src/Microsoft.Agents.AI.DurableTask/Microsoft.Agents.AI.DurableTask.csproj b/dotnet/src/Microsoft.Agents.AI.DurableTask/Microsoft.Agents.AI.DurableTask.csproj index 43ebe9c61f..8233afe964 100644 --- a/dotnet/src/Microsoft.Agents.AI.DurableTask/Microsoft.Agents.AI.DurableTask.csproj +++ b/dotnet/src/Microsoft.Agents.AI.DurableTask/Microsoft.Agents.AI.DurableTask.csproj @@ -17,6 +17,11 @@ README.md + + true + true + + diff --git a/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgent.cs b/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgent.cs index 6dc6d175a7..acf90fc16d 100644 --- a/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgent.cs +++ b/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgent.cs @@ -612,6 +612,12 @@ public sealed partial class ChatClientAgent : AIAgent chatOptions.AllowBackgroundResponses = agentRunOptions.AllowBackgroundResponses; } + if (agentRunOptions?.ResponseFormat is not null) + { + chatOptions ??= new ChatOptions(); + chatOptions.ResponseFormat = agentRunOptions.ResponseFormat; + } + ChatClientAgentContinuationToken? agentContinuationToken = null; if ((agentRunOptions?.ContinuationToken ?? chatOptions?.ContinuationToken) is { } continuationToken) diff --git a/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgentCustomOptions.cs b/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgentCustomOptions.cs index 25ae5a903e..e5d6296700 100644 --- a/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgentCustomOptions.cs +++ b/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgentCustomOptions.cs @@ -162,19 +162,14 @@ public partial class ChatClientAgent /// /// The JSON serialization options to use. /// Configuration parameters for controlling the agent's invocation behavior. - /// - /// to set a JSON schema on the ; otherwise, . The default is . - /// Using a JSON schema improves reliability if the underlying model supports native structured output with a schema, but might cause an error if the model does not support it. - /// /// The to monitor for cancellation requests. The default is . - /// A task that represents the asynchronous operation. The task result contains an with the agent's output. - public Task> RunAsync( + /// A task that represents the asynchronous operation. The task result contains an with the agent's output. + public Task> RunAsync( AgentSession? session, JsonSerializerOptions? serializerOptions, ChatClientAgentRunOptions? options, - bool? useJsonSchemaResponseFormat = null, CancellationToken cancellationToken = default) => - this.RunAsync(session, serializerOptions, (AgentRunOptions?)options, useJsonSchemaResponseFormat, cancellationToken); + this.RunAsync(session, serializerOptions, (AgentRunOptions?)options, cancellationToken); /// /// Runs the agent with a text message from the user, requesting a response of the specified type . @@ -186,20 +181,15 @@ public partial class ChatClientAgent /// /// The JSON serialization options to use. /// Configuration parameters for controlling the agent's invocation behavior. - /// - /// to set a JSON schema on the ; otherwise, . The default is . - /// Using a JSON schema improves reliability if the underlying model supports native structured output with a schema, but might cause an error if the model does not support it. - /// /// The to monitor for cancellation requests. The default is . - /// A task that represents the asynchronous operation. The task result contains an with the agent's output. - public Task> RunAsync( + /// A task that represents the asynchronous operation. The task result contains an with the agent's output. + public Task> RunAsync( string message, AgentSession? session, JsonSerializerOptions? serializerOptions, ChatClientAgentRunOptions? options, - bool? useJsonSchemaResponseFormat = null, CancellationToken cancellationToken = default) => - this.RunAsync(message, session, serializerOptions, (AgentRunOptions?)options, useJsonSchemaResponseFormat, cancellationToken); + this.RunAsync(message, session, serializerOptions, (AgentRunOptions?)options, cancellationToken); /// /// Runs the agent with a single chat message, requesting a response of the specified type . @@ -211,20 +201,15 @@ public partial class ChatClientAgent /// /// The JSON serialization options to use. /// Configuration parameters for controlling the agent's invocation behavior. - /// - /// to set a JSON schema on the ; otherwise, . The default is . - /// Using a JSON schema improves reliability if the underlying model supports native structured output with a schema, but might cause an error if the model does not support it. - /// /// The to monitor for cancellation requests. The default is . - /// A task that represents the asynchronous operation. The task result contains an with the agent's output. - public Task> RunAsync( + /// A task that represents the asynchronous operation. The task result contains an with the agent's output. + public Task> RunAsync( ChatMessage message, AgentSession? session, JsonSerializerOptions? serializerOptions, ChatClientAgentRunOptions? options, - bool? useJsonSchemaResponseFormat = null, CancellationToken cancellationToken = default) => - this.RunAsync(message, session, serializerOptions, (AgentRunOptions?)options, useJsonSchemaResponseFormat, cancellationToken); + this.RunAsync(message, session, serializerOptions, (AgentRunOptions?)options, cancellationToken); /// /// Runs the agent with a collection of chat messages, requesting a response of the specified type . @@ -236,18 +221,13 @@ public partial class ChatClientAgent /// /// The JSON serialization options to use. /// Configuration parameters for controlling the agent's invocation behavior. - /// - /// to set a JSON schema on the ; otherwise, . The default is . - /// Using a JSON schema improves reliability if the underlying model supports native structured output with a schema, but might cause an error if the model does not support it. - /// /// The to monitor for cancellation requests. The default is . - /// A task that represents the asynchronous operation. The task result contains an with the agent's output. - public Task> RunAsync( + /// A task that represents the asynchronous operation. The task result contains an with the agent's output. + public Task> RunAsync( IEnumerable messages, AgentSession? session, JsonSerializerOptions? serializerOptions, ChatClientAgentRunOptions? options, - bool? useJsonSchemaResponseFormat = null, CancellationToken cancellationToken = default) => - this.RunAsync(messages, session, serializerOptions, (AgentRunOptions?)options, useJsonSchemaResponseFormat, cancellationToken); + this.RunAsync(messages, session, serializerOptions, (AgentRunOptions?)options, cancellationToken); } diff --git a/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgentRunOptions.cs b/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgentRunOptions.cs index 0f2c9da485..cf35aa80a1 100644 --- a/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgentRunOptions.cs +++ b/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgentRunOptions.cs @@ -26,6 +26,17 @@ public sealed class ChatClientAgentRunOptions : AgentRunOptions this.ChatOptions = chatOptions; } + /// + /// Initializes a new instance of the class by copying values from the specified options. + /// + /// The options instance from which to copy values. + private ChatClientAgentRunOptions(ChatClientAgentRunOptions options) + : base(options) + { + this.ChatOptions = options.ChatOptions?.Clone(); + this.ChatClientFactory = options.ChatClientFactory; + } + /// /// Gets or sets the chat options to apply to the agent invocation. /// @@ -50,4 +61,7 @@ public sealed class ChatClientAgentRunOptions : AgentRunOptions /// chat client will be used without modification. /// public Func? ChatClientFactory { get; set; } + + /// + public override AgentRunOptions Clone() => new ChatClientAgentRunOptions(this); } diff --git a/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgentRunResponse{T}.cs b/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgentRunResponse{T}.cs deleted file mode 100644 index a4fadff0c7..0000000000 --- a/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgentRunResponse{T}.cs +++ /dev/null @@ -1,45 +0,0 @@ -// Copyright (c) Microsoft. All rights reserved. - -using System; -using Microsoft.Extensions.AI; -using Microsoft.Shared.Diagnostics; - -namespace Microsoft.Agents.AI; - -/// -/// Represents the response of the specified type to an run request. -/// -/// The type of value expected from the chat response. -/// -/// Language models are not guaranteed to honor the requested schema. If the model's output is not -/// parsable as the expected type, you can access the underlying JSON response on the property. -/// -public sealed class ChatClientAgentResponse : AgentResponse -{ - private readonly ChatResponse _response; - - /// - /// Initializes a new instance of the class from an existing . - /// - /// The from which to populate this . - /// is . - /// - /// This constructor creates an agent response that wraps an existing , preserving all - /// metadata and storing the original response in for access to - /// the underlying implementation details. - /// - public ChatClientAgentResponse(ChatResponse response) : base(response) - { - _ = Throw.IfNull(response); - - this._response = response; - } - - /// - /// Gets the result value of the agent response as an instance of . - /// - /// - /// If the response did not contain JSON, or if deserialization fails, this property will throw. - /// - public override T Result => this._response.Result; -} diff --git a/dotnet/src/Shared/StructuredOutput/StructuredOutputSchemaUtilities.cs b/dotnet/src/Shared/StructuredOutput/StructuredOutputSchemaUtilities.cs new file mode 100644 index 0000000000..95836b95c4 --- /dev/null +++ b/dotnet/src/Shared/StructuredOutput/StructuredOutputSchemaUtilities.cs @@ -0,0 +1,104 @@ +// Copyright (c) Microsoft. All rights reserved. + +#pragma warning disable IDE0005 // Using directive is unnecessary. + +using System; +using System.Text.Json; +using System.Text.Json.Nodes; +using Microsoft.Extensions.AI; + +namespace Microsoft.Agents.AI; + +/// +/// Internal utilities for working with structured output JSON schemas. +/// +internal static class StructuredOutputSchemaUtilities +{ + private const string DataPropertyName = "data"; + + /// + /// Ensures the given response format has an object schema at the root, wrapping non-object schemas if necessary. + /// + /// The response format to check. + /// A tuple containing the (possibly wrapped) response format and whether wrapping occurred. + /// The response format does not have a valid JSON schema. + internal static (ChatResponseFormatJson ResponseFormat, bool IsWrappedInObject) WrapNonObjectSchema(ChatResponseFormatJson responseFormat) + { + if (responseFormat.Schema is null) + { + throw new InvalidOperationException("The response format must have a valid JSON schema."); + } + + var schema = responseFormat.Schema.Value; + bool isWrappedInObject = false; + + if (!SchemaRepresentsObject(responseFormat.Schema)) + { + // For non-object-representing schemas, we wrap them in an object schema, because all + // the real LLM providers today require an object schema as the root. This is currently + // true even for providers that support native structured output. + isWrappedInObject = true; + schema = JsonSerializer.SerializeToElement(new JsonObject + { + { "$schema", "https://json-schema.org/draft/2020-12/schema" }, + { "type", "object" }, + { "properties", new JsonObject { { DataPropertyName, JsonElementToJsonNode(schema) } } }, + { "additionalProperties", false }, + { "required", new JsonArray(DataPropertyName) }, + }, AIJsonUtilities.DefaultOptions.GetTypeInfo(typeof(JsonObject))); + + responseFormat = ChatResponseFormat.ForJsonSchema(schema, responseFormat.SchemaName, responseFormat.SchemaDescription); + } + + return (responseFormat, isWrappedInObject); + } + + /// + /// Unwraps the "data" property from a JSON object that was previously wrapped by . + /// + /// The JSON string to unwrap. + /// The raw JSON text of the "data" property, or the original JSON if no wrapping is detected. + internal static string UnwrapResponseData(string json) + { + using var document = JsonDocument.Parse(json); + if (document.RootElement.ValueKind == JsonValueKind.Object && + document.RootElement.TryGetProperty(DataPropertyName, out JsonElement dataElement)) + { + return dataElement.GetRawText(); + } + + // If root is not an object or "data" property is not found, return the original JSON as a fallback + return json; + } + + private static bool SchemaRepresentsObject(JsonElement? schema) + { + if (schema is not { } schemaElement) + { + return false; + } + + if (schemaElement.ValueKind is JsonValueKind.Object) + { + foreach (var property in schemaElement.EnumerateObject()) + { + if (property.NameEquals("type"u8)) + { + return property.Value.ValueKind == JsonValueKind.String + && property.Value.ValueEquals("object"u8); + } + } + } + + return false; + } + + private static JsonNode? JsonElementToJsonNode(JsonElement element) => + element.ValueKind switch + { + JsonValueKind.Null => null, + JsonValueKind.Array => JsonArray.Create(element), + JsonValueKind.Object => JsonObject.Create(element), + _ => JsonValue.Create(element) + }; +} diff --git a/dotnet/tests/AgentConformance.IntegrationTests/StructuredOutputRunTests.cs b/dotnet/tests/AgentConformance.IntegrationTests/StructuredOutputRunTests.cs new file mode 100644 index 0000000000..6b7556b456 --- /dev/null +++ b/dotnet/tests/AgentConformance.IntegrationTests/StructuredOutputRunTests.cs @@ -0,0 +1,110 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Text.Json; +using System.Threading.Tasks; +using AgentConformance.IntegrationTests.Support; +using Microsoft.Agents.AI; +using Microsoft.Extensions.AI; + +namespace AgentConformance.IntegrationTests; + +/// +/// Conformance tests for structured output handling for run methods on agents. +/// +/// The type of test fixture used by the concrete test implementation. +/// Function to create the test fixture with. +public abstract class StructuredOutputRunTests(Func createAgentFixture) : AgentTests(createAgentFixture) + where TAgentFixture : IAgentFixture +{ + [RetryFact(Constants.RetryCount, Constants.RetryDelay)] + public virtual async Task RunWithResponseFormatReturnsExpectedResultAsync() + { + // Arrange + var agent = this.Fixture.Agent; + var session = await agent.CreateSessionAsync(); + await using var cleanup = new SessionCleanup(session, this.Fixture); + + var options = new AgentRunOptions + { + ResponseFormat = ChatResponseFormat.ForJsonSchema(AgentAbstractionsJsonUtilities.DefaultOptions) + }; + + // Act + var response = await agent.RunAsync(new ChatMessage(ChatRole.User, "Provide information about the capital of France."), session, options); + + // Assert + Assert.NotNull(response); + Assert.Single(response.Messages); + Assert.Contains("Paris", response.Text); + Assert.True(TryDeserialize(response.Text, AgentAbstractionsJsonUtilities.DefaultOptions, out CityInfo cityInfo)); + Assert.Equal("Paris", cityInfo.Name); + } + + [RetryFact(Constants.RetryCount, Constants.RetryDelay)] + public virtual async Task RunWithGenericTypeReturnsExpectedResultAsync() + { + // Arrange + var agent = this.Fixture.Agent; + var session = await agent.CreateSessionAsync(); + await using var cleanup = new SessionCleanup(session, this.Fixture); + + // Act + AgentResponse response = await agent.RunAsync( + new ChatMessage(ChatRole.User, "Provide information about the capital of France."), + session); + + // Assert + Assert.NotNull(response); + Assert.Single(response.Messages); + Assert.Contains("Paris", response.Text); + + Assert.NotNull(response.Result); + Assert.Equal("Paris", response.Result.Name); + } + + [RetryFact(Constants.RetryCount, Constants.RetryDelay)] + public virtual async Task RunWithPrimitiveTypeReturnsExpectedResultAsync() + { + // Arrange + var agent = this.Fixture.Agent; + var session = await agent.CreateSessionAsync(); + await using var cleanup = new SessionCleanup(session, this.Fixture); + + // Act - Request a primitive type, which requires wrapping in an object schema + AgentResponse response = await agent.RunAsync( + new ChatMessage(ChatRole.User, "What is the sum of 15 and 27? Respond with just the number."), + session); + + // Assert + Assert.NotNull(response); + Assert.Single(response.Messages); + Assert.Equal(42, response.Result); + } + + protected static bool TryDeserialize(string json, JsonSerializerOptions jsonSerializerOptions, out T structuredOutput) + { + try + { + T? deserialized = JsonSerializer.Deserialize(json, jsonSerializerOptions); + if (deserialized is null) + { + structuredOutput = default!; + return false; + } + + structuredOutput = deserialized; + return true; + } + catch + { + structuredOutput = default!; + return false; + } + } +} + +public sealed class CityInfo +{ + public string? Name { get; set; } +} diff --git a/dotnet/tests/AgentConformance.IntegrationTests/Support/Constants.cs b/dotnet/tests/AgentConformance.IntegrationTests/Support/Constants.cs index 178b1951ba..232b5fdb10 100644 --- a/dotnet/tests/AgentConformance.IntegrationTests/Support/Constants.cs +++ b/dotnet/tests/AgentConformance.IntegrationTests/Support/Constants.cs @@ -2,7 +2,7 @@ namespace AgentConformance.IntegrationTests.Support; -internal static class Constants +public static class Constants { public const int RetryCount = 3; public const int RetryDelay = 5000; diff --git a/dotnet/tests/AgentConformance.IntegrationTests/Support/SessionCleanup.cs b/dotnet/tests/AgentConformance.IntegrationTests/Support/SessionCleanup.cs index c59b999fd2..91e858e53f 100644 --- a/dotnet/tests/AgentConformance.IntegrationTests/Support/SessionCleanup.cs +++ b/dotnet/tests/AgentConformance.IntegrationTests/Support/SessionCleanup.cs @@ -11,7 +11,7 @@ namespace AgentConformance.IntegrationTests.Support; /// /// The session to delete. /// The fixture that provides agent specific capabilities. -internal sealed class SessionCleanup(AgentSession session, IAgentFixture fixture) : IAsyncDisposable +public sealed class SessionCleanup(AgentSession session, IAgentFixture fixture) : IAsyncDisposable { public async ValueTask DisposeAsync() => await fixture.DeleteSessionAsync(session); diff --git a/dotnet/tests/AzureAI.IntegrationTests/AIProjectClientAgentStructuredOutputRunTests.cs b/dotnet/tests/AzureAI.IntegrationTests/AIProjectClientAgentStructuredOutputRunTests.cs new file mode 100644 index 0000000000..94ce01e221 --- /dev/null +++ b/dotnet/tests/AzureAI.IntegrationTests/AIProjectClientAgentStructuredOutputRunTests.cs @@ -0,0 +1,99 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Threading.Tasks; +using AgentConformance.IntegrationTests; +using AgentConformance.IntegrationTests.Support; +using Microsoft.Agents.AI; +using Microsoft.Extensions.AI; + +namespace AzureAI.IntegrationTests; + +public class AIProjectClientAgentStructuredOutputRunTests() : StructuredOutputRunTests>(() => new AIProjectClientStructuredOutputFixture()) +{ + private const string NotSupported = "AIProjectClient does not support specifying structured output type at invocation time."; + + /// + /// Verifies that response format provided at agent initialization is used when invoking RunAsync. + /// + /// + [RetryFact(Constants.RetryCount, Constants.RetryDelay)] + public async Task RunWithResponseFormatAtAgentInitializationReturnsExpectedResultAsync() + { + // Arrange + var agent = this.Fixture.Agent; + var session = await agent.CreateSessionAsync(); + await using var cleanup = new SessionCleanup(session, this.Fixture); + + // Act + var response = await agent.RunAsync(new ChatMessage(ChatRole.User, "Provide information about the capital of France."), session); + + // Assert + Assert.NotNull(response); + Assert.Single(response.Messages); + Assert.Contains("Paris", response.Text); + Assert.True(TryDeserialize(response.Text, AgentAbstractionsJsonUtilities.DefaultOptions, out CityInfo cityInfo)); + Assert.Equal("Paris", cityInfo.Name); + } + + /// + /// Verifies that generic RunAsync works with AIProjectClient when structured output is configured at agent initialization. + /// + /// + /// AIProjectClient does not support specifying the structured output type at invocation time yet. + /// The type T provided to RunAsync<T> is ignored by AzureAIProjectChatClient and is only used + /// for deserializing the agent response by AgentResponse<T>.Result. + /// + [RetryFact(Constants.RetryCount, Constants.RetryDelay)] + public async Task RunGenericWithResponseFormatAtAgentInitializationReturnsExpectedResultAsync() + { + // Arrange + var agent = this.Fixture.Agent; + var session = await agent.CreateSessionAsync(); + await using var cleanup = new SessionCleanup(session, this.Fixture); + + // Act + AgentResponse response = await agent.RunAsync( + new ChatMessage(ChatRole.User, "Provide information about the capital of France."), + session); + + // Assert + Assert.NotNull(response); + Assert.Single(response.Messages); + Assert.Contains("Paris", response.Text); + + Assert.NotNull(response.Result); + Assert.Equal("Paris", response.Result.Name); + } + + [Fact(Skip = NotSupported)] + public override Task RunWithGenericTypeReturnsExpectedResultAsync() => + base.RunWithGenericTypeReturnsExpectedResultAsync(); + + [Fact(Skip = NotSupported)] + public override Task RunWithResponseFormatReturnsExpectedResultAsync() => + base.RunWithResponseFormatReturnsExpectedResultAsync(); + + [Fact(Skip = NotSupported)] + public override Task RunWithPrimitiveTypeReturnsExpectedResultAsync() => + base.RunWithPrimitiveTypeReturnsExpectedResultAsync(); +} + +/// +/// Represents a fixture for testing AIProjectClient with structured output of type provided at agent initialization. +/// +public class AIProjectClientStructuredOutputFixture : AIProjectClientFixture +{ + public override Task InitializeAsync() + { + var agentOptions = new ChatClientAgentOptions + { + ChatOptions = new ChatOptions() + { + ResponseFormat = ChatResponseFormat.ForJsonSchema(AgentAbstractionsJsonUtilities.DefaultOptions) + }, + }; + + return this.InitializeAsync(agentOptions); + } +} diff --git a/dotnet/tests/AzureAI.IntegrationTests/AIProjectClientFixture.cs b/dotnet/tests/AzureAI.IntegrationTests/AIProjectClientFixture.cs index dd926174c0..a13af9c940 100644 --- a/dotnet/tests/AzureAI.IntegrationTests/AIProjectClientFixture.cs +++ b/dotnet/tests/AzureAI.IntegrationTests/AIProjectClientFixture.cs @@ -121,6 +121,13 @@ public class AIProjectClientFixture : IChatClientAgentFixture return await this._client.CreateAIAgentAsync(GenerateUniqueAgentName(name), model: s_config.DeploymentName, instructions: instructions, tools: aiTools); } + public async Task CreateChatClientAgentAsync(ChatClientAgentOptions options) + { + options.Name ??= GenerateUniqueAgentName("HelpfulAssistant"); + + return await this._client.CreateAIAgentAsync(model: s_config.DeploymentName, options); + } + public static string GenerateUniqueAgentName(string baseName) => $"{baseName}-{Guid.NewGuid().ToString("N").Substring(0, 8)}"; @@ -161,9 +168,15 @@ public class AIProjectClientFixture : IChatClientAgentFixture return Task.CompletedTask; } - public async Task InitializeAsync() + public virtual async Task InitializeAsync() { this._client = new(new Uri(s_config.Endpoint), new AzureCliCredential()); this._agent = await this.CreateChatClientAgentAsync(); } + + public async Task InitializeAsync(ChatClientAgentOptions options) + { + this._client = new(new Uri(s_config.Endpoint), new AzureCliCredential()); + this._agent = await this.CreateChatClientAgentAsync(options); + } } diff --git a/dotnet/tests/AzureAIAgentsPersistent.IntegrationTests/AzureAIAgentsPersistentStructuredOutputRunTests.cs b/dotnet/tests/AzureAIAgentsPersistent.IntegrationTests/AzureAIAgentsPersistentStructuredOutputRunTests.cs new file mode 100644 index 0000000000..3e28e025d6 --- /dev/null +++ b/dotnet/tests/AzureAIAgentsPersistent.IntegrationTests/AzureAIAgentsPersistentStructuredOutputRunTests.cs @@ -0,0 +1,13 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System.Threading.Tasks; +using AgentConformance.IntegrationTests; + +namespace AzureAIAgentsPersistent.IntegrationTests; + +public class AzureAIAgentsPersistentStructuredOutputRunTests() : StructuredOutputRunTests(() => new()) +{ + [Fact(Skip = "Fails intermittently, at build agent")] + public override Task RunWithResponseFormatReturnsExpectedResultAsync() => + base.RunWithResponseFormatReturnsExpectedResultAsync(); +} diff --git a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/AIAgentStructuredOutputTests.cs b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/AIAgentStructuredOutputTests.cs new file mode 100644 index 0000000000..a8881ca761 --- /dev/null +++ b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/AIAgentStructuredOutputTests.cs @@ -0,0 +1,391 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System.Collections.Generic; +using System.Text.Json; +using System.Threading; +using System.Threading.Tasks; +using Microsoft.Agents.AI.Abstractions.UnitTests.Models; +using Microsoft.Extensions.AI; +using Moq; +using Moq.Protected; + +namespace Microsoft.Agents.AI.Abstractions.UnitTests; + +/// +/// Unit tests for the structured output functionality in . +/// +public class AIAgentStructuredOutputTests +{ + private readonly Mock _agentMock; + + public AIAgentStructuredOutputTests() + { + this._agentMock = new Mock { CallBase = true }; + } + + #region Schema Wrapping Tests + + /// + /// Verifies that when requesting an object type, the schema is NOT wrapped. + /// + [Fact] + public async Task RunAsyncGeneric_WithObjectType_DoesNotWrapSchemaAsync() + { + // Arrange + Animal expectedAnimal = new() { Id = 1, FullName = "Test", Species = Species.Tiger }; + string responseJson = JsonSerializer.Serialize(expectedAnimal, TestJsonSerializerContext.Default.Animal); + AgentResponse response = new(new ChatMessage(ChatRole.Assistant, responseJson)); + + this._agentMock + .Protected() + .Setup>("RunCoreAsync", + ItExpr.IsAny>(), + ItExpr.IsAny(), + ItExpr.IsAny(), + ItExpr.IsAny()) + .ReturnsAsync(response); + + // Act + AgentResponse result = await this._agentMock.Object.RunAsync( + "Get me an animal", + serializerOptions: TestJsonSerializerContext.Default.Options); + + // Assert - Verify the result is NOT marked as wrapped + Assert.False(result.IsWrappedInObject); + } + + /// + /// Verifies that when requesting a primitive type (int), the schema IS wrapped. + /// + [Fact] + public async Task RunAsyncGeneric_WithPrimitiveType_WrapsSchemaAsync() + { + // Arrange + const string ResponseJson = "{\"data\":42}"; + AgentResponse response = new(new ChatMessage(ChatRole.Assistant, ResponseJson)); + + this._agentMock + .Protected() + .Setup>("RunCoreAsync", + ItExpr.IsAny>(), + ItExpr.IsAny(), + ItExpr.IsAny(), + ItExpr.IsAny()) + .ReturnsAsync(response); + + // Act + AgentResponse result = await this._agentMock.Object.RunAsync( + "Give me a number", + serializerOptions: TestJsonSerializerContext.Default.Options); + + // Assert - Verify the result is marked as wrapped + Assert.True(result.IsWrappedInObject); + } + + /// + /// Verifies that when requesting an array type, the schema IS wrapped. + /// + [Fact] + public async Task RunAsyncGeneric_WithArrayType_WrapsSchemaAsync() + { + // Arrange + const string ResponseJson = "{\"data\":[\"a\",\"b\",\"c\"]}"; + AgentResponse response = new(new ChatMessage(ChatRole.Assistant, ResponseJson)); + + this._agentMock + .Protected() + .Setup>("RunCoreAsync", + ItExpr.IsAny>(), + ItExpr.IsAny(), + ItExpr.IsAny(), + ItExpr.IsAny()) + .ReturnsAsync(response); + + // Act + AgentResponse result = await this._agentMock.Object.RunAsync( + "Give me an array of strings", + serializerOptions: TestJsonSerializerContext.Default.Options); + + // Assert - Verify the result is marked as wrapped + Assert.True(result.IsWrappedInObject); + } + + /// + /// Verifies that when requesting an enum type, the schema IS wrapped. + /// + [Fact] + public async Task RunAsyncGeneric_WithEnumType_WrapsSchemaAsync() + { + // Arrange + const string ResponseJson = "{\"data\":\"Tiger\"}"; + AgentResponse response = new(new ChatMessage(ChatRole.Assistant, ResponseJson)); + + this._agentMock + .Protected() + .Setup>("RunCoreAsync", + ItExpr.IsAny>(), + ItExpr.IsAny(), + ItExpr.IsAny(), + ItExpr.IsAny()) + .ReturnsAsync(response); + + // Act + AgentResponse result = await this._agentMock.Object.RunAsync( + "Give me a species", + serializerOptions: TestJsonSerializerContext.Default.Options); + + // Assert - Verify the result is marked as wrapped + Assert.True(result.IsWrappedInObject); + } + + #endregion + + #region AgentResponse.Result Unwrapping Tests + + /// + /// Verifies that AgentResponse{T}.Result correctly deserializes an object without unwrapping. + /// + [Fact] + public void AgentResponseGeneric_Result_DeserializesObjectWithoutUnwrapping() + { + // Arrange + Animal expectedAnimal = new() { Id = 1, FullName = "Tigger", Species = Species.Tiger }; + string responseJson = JsonSerializer.Serialize(expectedAnimal, TestJsonSerializerContext.Default.Animal); + AgentResponse response = new(new ChatMessage(ChatRole.Assistant, responseJson)); + AgentResponse typedResponse = new(response, TestJsonSerializerContext.Default.Options); + + // Act + Animal result = typedResponse.Result; + + // Assert + Assert.Equal(expectedAnimal.Id, result.Id); + Assert.Equal(expectedAnimal.FullName, result.FullName); + Assert.Equal(expectedAnimal.Species, result.Species); + } + + /// + /// Verifies that AgentResponse{T}.Result correctly unwraps and deserializes a primitive value. + /// + [Fact] + public void AgentResponseGeneric_Result_UnwrapsPrimitiveFromDataProperty() + { + // Arrange + const string ResponseJson = "{\"data\":42}"; + AgentResponse response = new(new ChatMessage(ChatRole.Assistant, ResponseJson)); + AgentResponse typedResponse = new(response, TestJsonSerializerContext.Default.Options) { IsWrappedInObject = true }; + + // Act + int result = typedResponse.Result; + + // Assert + Assert.Equal(42, result); + } + + /// + /// Verifies that AgentResponse{T}.Result correctly unwraps and deserializes an array. + /// + [Fact] + public void AgentResponseGeneric_Result_UnwrapsArrayFromDataProperty() + { + // Arrange + const string ResponseJson = "{\"data\":[\"apple\",\"banana\",\"cherry\"]}"; + AgentResponse response = new(new ChatMessage(ChatRole.Assistant, ResponseJson)); + AgentResponse typedResponse = new(response, TestJsonSerializerContext.Default.Options) { IsWrappedInObject = true }; + + // Act + string[] result = typedResponse.Result; + + // Assert + Assert.Equal(["apple", "banana", "cherry"], result); + } + + /// + /// Verifies that AgentResponse{T}.Result correctly unwraps and deserializes an enum. + /// + [Fact] + public void AgentResponseGeneric_Result_UnwrapsEnumFromDataProperty() + { + // Arrange + const string ResponseJson = "{\"data\":\"Walrus\"}"; + AgentResponse response = new(new ChatMessage(ChatRole.Assistant, ResponseJson)); + AgentResponse typedResponse = new(response, TestJsonSerializerContext.Default.Options) { IsWrappedInObject = true }; + + // Act + Species result = typedResponse.Result; + + // Assert + Assert.Equal(Species.Walrus, result); + } + + /// + /// Verifies that AgentResponse{T}.Result falls back to original JSON when data property is missing. + /// + [Fact] + public void AgentResponseGeneric_Result_FallsBackWhenDataPropertyMissing() + { + // Arrange - simulate a case where wrapping was expected but response does not have data + const string ResponseJson = "42"; + AgentResponse response = new(new ChatMessage(ChatRole.Assistant, ResponseJson)); + AgentResponse typedResponse = new(response, TestJsonSerializerContext.Default.Options) { IsWrappedInObject = true }; + + // Act + int result = typedResponse.Result; + + // Assert - should still work by falling back to original JSON + Assert.Equal(42, result); + } + + /// + /// Verifies that AgentResponse{T}.Result throws when response text is empty. + /// + [Fact] + public void AgentResponseGeneric_Result_ThrowsWhenTextIsEmpty() + { + // Arrange + AgentResponse response = new(new ChatMessage(ChatRole.Assistant, string.Empty)); + AgentResponse typedResponse = new(response, TestJsonSerializerContext.Default.Options); + + // Act and Assert + Assert.Throws(() => typedResponse.Result); + } + + /// + /// Verifies that AgentResponse{T}.Result throws when deserialized value is null. + /// + [Fact] + public void AgentResponseGeneric_Result_ThrowsWhenDeserializedValueIsNull() + { + // Arrange + const string ResponseJson = "null"; + AgentResponse response = new(new ChatMessage(ChatRole.Assistant, ResponseJson)); + AgentResponse typedResponse = new(response, TestJsonSerializerContext.Default.Options); + + // Act and Assert + Assert.Throws(() => typedResponse.Result); + } + + #endregion + + #region End-to-End Tests + + /// + /// End-to-end test: Request a primitive type, verify wrapping, and verify correct deserialization. + /// + [Fact] + public async Task RunAsyncGeneric_PrimitiveEndToEnd_WrapsAndDeserializesCorrectlyAsync() + { + // Arrange + const string ResponseJson = "{\"data\":123}"; + AgentResponse response = new(new ChatMessage(ChatRole.Assistant, ResponseJson)); + + this._agentMock + .Protected() + .Setup>("RunCoreAsync", + ItExpr.IsAny>(), + ItExpr.IsAny(), + ItExpr.IsAny(), + ItExpr.IsAny()) + .ReturnsAsync(response); + + // Act + AgentResponse result = await this._agentMock.Object.RunAsync( + "Give me a number", + serializerOptions: TestJsonSerializerContext.Default.Options); + + // Assert + Assert.True(result.IsWrappedInObject); + Assert.Equal(123, result.Result); + } + + /// + /// End-to-end test: Request an array type, verify wrapping, and verify correct deserialization. + /// + [Fact] + public async Task RunAsyncGeneric_ArrayEndToEnd_WrapsAndDeserializesCorrectlyAsync() + { + // Arrange + const string ResponseJson = "{\"data\":[\"one\",\"two\",\"three\"]}"; + AgentResponse response = new(new ChatMessage(ChatRole.Assistant, ResponseJson)); + + this._agentMock + .Protected() + .Setup>("RunCoreAsync", + ItExpr.IsAny>(), + ItExpr.IsAny(), + ItExpr.IsAny(), + ItExpr.IsAny()) + .ReturnsAsync(response); + + // Act + AgentResponse result = await this._agentMock.Object.RunAsync( + "Give me an array of strings", + serializerOptions: TestJsonSerializerContext.Default.Options); + + // Assert + Assert.True(result.IsWrappedInObject); + Assert.Equal(["one", "two", "three"], result.Result); + } + + /// + /// End-to-end test: Request an object type, verify no wrapping, and verify correct deserialization. + /// + [Fact] + public async Task RunAsyncGeneric_ObjectEndToEnd_NoWrappingAndDeserializesCorrectlyAsync() + { + // Arrange + Animal expectedAnimal = new() { Id = 99, FullName = "Leo", Species = Species.Bear }; + string responseJson = JsonSerializer.Serialize(expectedAnimal, TestJsonSerializerContext.Default.Animal); + AgentResponse response = new(new ChatMessage(ChatRole.Assistant, responseJson)); + + this._agentMock + .Protected() + .Setup>("RunCoreAsync", + ItExpr.IsAny>(), + ItExpr.IsAny(), + ItExpr.IsAny(), + ItExpr.IsAny()) + .ReturnsAsync(response); + + // Act + AgentResponse result = await this._agentMock.Object.RunAsync( + "Give me an animal", + serializerOptions: TestJsonSerializerContext.Default.Options); + + // Assert + Assert.False(result.IsWrappedInObject); + Assert.Equal(expectedAnimal.Id, result.Result.Id); + Assert.Equal(expectedAnimal.FullName, result.Result.FullName); + Assert.Equal(expectedAnimal.Species, result.Result.Species); + } + + /// + /// End-to-end test: Request an enum type, verify wrapping, and verify correct deserialization. + /// + [Fact] + public async Task RunAsyncGeneric_EnumEndToEnd_WrapsAndDeserializesCorrectlyAsync() + { + // Arrange + const string ResponseJson = "{\"data\":\"Bear\"}"; + AgentResponse response = new(new ChatMessage(ChatRole.Assistant, ResponseJson)); + + this._agentMock + .Protected() + .Setup>("RunCoreAsync", + ItExpr.IsAny>(), + ItExpr.IsAny(), + ItExpr.IsAny(), + ItExpr.IsAny()) + .ReturnsAsync(response); + + // Act + AgentResponse result = await this._agentMock.Object.RunAsync( + "Give me a species", + serializerOptions: TestJsonSerializerContext.Default.Options); + + // Assert + Assert.True(result.IsWrappedInObject); + Assert.Equal(Species.Bear, result.Result); + } + + #endregion +} diff --git a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/AgentResponseTests.cs b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/AgentResponseTests.cs index 8300701c5b..e1425b3144 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/AgentResponseTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/AgentResponseTests.cs @@ -214,30 +214,6 @@ public class AgentResponseTests Assert.Equal(100, usageContent.Details.TotalTokenCount); } -#if NETFRAMEWORK - /// - /// Since Json Serialization using reflection is disabled in .net core builds, and we are using a custom type here that wouldn't - /// be registered with the default source generated serializer, this test will only pass in .net framework builds where reflection-based - /// serialization is available. - /// - [Fact] - public void ParseAsStructuredOutputSuccess() - { - // Arrange. - var expectedResult = new Animal { Id = 1, FullName = "Tigger", Species = Species.Tiger }; - var response = new AgentResponse(new ChatMessage(ChatRole.Assistant, JsonSerializer.Serialize(expectedResult, TestJsonSerializerContext.Default.Animal))); - - // Act. - var animal = response.Deserialize(); - - // Assert. - Assert.NotNull(animal); - Assert.Equal(expectedResult.Id, animal.Id); - Assert.Equal(expectedResult.FullName, animal.FullName); - Assert.Equal(expectedResult.Species, animal.Species); - } -#endif - [Fact] public void ParseAsStructuredOutputWithJSOSuccess() { @@ -246,7 +222,7 @@ public class AgentResponseTests var response = new AgentResponse(new ChatMessage(ChatRole.Assistant, JsonSerializer.Serialize(expectedResult, TestJsonSerializerContext.Default.Animal))); // Act. - var animal = response.Deserialize(TestJsonSerializerContext.Default.Options); + var animal = JsonSerializer.Deserialize(response.Text, TestJsonSerializerContext.Default.Options); // Assert. Assert.NotNull(animal); @@ -255,98 +231,6 @@ public class AgentResponseTests Assert.Equal(expectedResult.Species, animal.Species); } - [Fact] - public void ParseAsStructuredOutputFailsWithEmptyString() - { - // Arrange. - var response = new AgentResponse(new ChatMessage(ChatRole.Assistant, string.Empty)); - - // Act & Assert. - var exception = Assert.Throws(() => response.Deserialize(TestJsonSerializerContext.Default.Options)); - Assert.Equal("The response did not contain JSON to be deserialized.", exception.Message); - } - - [Fact] - public void ParseAsStructuredOutputFailsWithInvalidJson() - { - // Arrange. - var response = new AgentResponse(new ChatMessage(ChatRole.Assistant, "invalid json")); - - // Act & Assert. - Assert.Throws(() => response.Deserialize(TestJsonSerializerContext.Default.Options)); - } - - [Fact] - public void ParseAsStructuredOutputFailsWithIncorrectTypedJson() - { - // Arrange. - var response = new AgentResponse(new ChatMessage(ChatRole.Assistant, "[]")); - - // Act & Assert. - Assert.Throws(() => response.Deserialize(TestJsonSerializerContext.Default.Options)); - } - -#if NETFRAMEWORK - /// - /// Since Json Serialization using reflection is disabled in .net core builds, and we are using a custom type here that wouldn't - /// be registered with the default source generated serializer, this test will only pass in .net framework builds where reflection-based - /// serialization is available. - /// - [Fact] - public void TryParseAsStructuredOutputSuccess() - { - // Arrange. - var expectedResult = new Animal { Id = 1, FullName = "Tigger", Species = Species.Tiger }; - var response = new AgentResponse(new ChatMessage(ChatRole.Assistant, JsonSerializer.Serialize(expectedResult, TestJsonSerializerContext.Default.Animal))); - - // Act. - response.TryDeserialize(out Animal? animal); - - // Assert. - Assert.NotNull(animal); - Assert.Equal(expectedResult.Id, animal.Id); - Assert.Equal(expectedResult.FullName, animal.FullName); - Assert.Equal(expectedResult.Species, animal.Species); - } -#endif - - [Fact] - public void TryParseAsStructuredOutputWithJSOSuccess() - { - // Arrange. - var expectedResult = new Animal { Id = 1, FullName = "Tigger", Species = Species.Tiger }; - var response = new AgentResponse(new ChatMessage(ChatRole.Assistant, JsonSerializer.Serialize(expectedResult, TestJsonSerializerContext.Default.Animal))); - - // Act. - response.TryDeserialize(TestJsonSerializerContext.Default.Options, out Animal? animal); - - // Assert. - Assert.NotNull(animal); - Assert.Equal(expectedResult.Id, animal.Id); - Assert.Equal(expectedResult.FullName, animal.FullName); - Assert.Equal(expectedResult.Species, animal.Species); - } - - [Fact] - public void TryParseAsStructuredOutputFailsWithEmptyText() - { - // Arrange. - var response = new AgentResponse(new ChatMessage(ChatRole.Assistant, string.Empty)); - - // Act & Assert. - Assert.False(response.TryDeserialize(TestJsonSerializerContext.Default.Options, out _)); - } - - [Fact] - public void TryParseAsStructuredOutputFailsWithIncorrectTypedJson() - { - // Arrange. - var response = new AgentResponse(new ChatMessage(ChatRole.Assistant, "[]")); - - // Act & Assert. - Assert.False(response.TryDeserialize(TestJsonSerializerContext.Default.Options, out _)); - } - [Fact] public void ToAgentResponseUpdatesWithNoMessagesProducesEmptyArray() { @@ -395,16 +279,4 @@ public class AgentResponseTests Assert.NotNull(update.AdditionalProperties); Assert.Equal("value", update.AdditionalProperties!["key"]); } - - [Fact] - public void Deserialize_ThrowsWhenDeserializationReturnsNull() - { - // Arrange - AgentResponse response = new(new ChatMessage(ChatRole.Assistant, "null")); - - // Act & Assert - InvalidOperationException exception = Assert.Throws( - () => response.Deserialize(TestJsonSerializerContext.Default.Options)); - Assert.Equal("The deserialized response is null.", exception.Message); - } } diff --git a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/AgentRunOptionsTests.cs b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/AgentRunOptionsTests.cs index 7460ea4623..028828c520 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/AgentRunOptionsTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/AgentRunOptionsTests.cs @@ -1,6 +1,5 @@ // Copyright (c) Microsoft. All rights reserved. -using System; using System.Text.Json; using Microsoft.Extensions.AI; @@ -27,7 +26,7 @@ public class AgentRunOptionsTests }; // Act - var clone = new AgentRunOptions(options); + var clone = options.Clone(); // Assert Assert.NotNull(clone); @@ -39,11 +38,6 @@ public class AgentRunOptionsTests Assert.Equal(42, clone.AdditionalProperties["key2"]); } - [Fact] - public void CloningConstructorThrowsIfNull() => - // Act & Assert - Assert.Throws(() => new AgentRunOptions(null!)); - [Fact] public void JsonSerializationRoundtrips() { @@ -77,4 +71,57 @@ public class AgentRunOptionsTests Assert.IsType(value2); Assert.Equal(42, ((JsonElement)value2!).GetInt32()); } + + [Fact] + public void CloneReturnsNewInstanceWithSameValues() + { + // Arrange + var options = new AgentRunOptions + { + ContinuationToken = ResponseContinuationToken.FromBytes(new byte[] { 1, 2, 3 }), + AllowBackgroundResponses = true, + AdditionalProperties = new AdditionalPropertiesDictionary + { + ["key1"] = "value1", + ["key2"] = 42 + }, + ResponseFormat = ChatResponseFormat.Json + }; + + // Act + AgentRunOptions clone = options.Clone(); + + // Assert + Assert.NotNull(clone); + Assert.IsType(clone); + Assert.NotSame(options, clone); + Assert.Same(options.ContinuationToken, clone.ContinuationToken); + Assert.Equal(options.AllowBackgroundResponses, clone.AllowBackgroundResponses); + Assert.NotNull(clone.AdditionalProperties); + Assert.NotSame(options.AdditionalProperties, clone.AdditionalProperties); + Assert.Equal("value1", clone.AdditionalProperties["key1"]); + Assert.Equal(42, clone.AdditionalProperties["key2"]); + Assert.Same(options.ResponseFormat, clone.ResponseFormat); + } + + [Fact] + public void CloneCreatesIndependentAdditionalPropertiesDictionary() + { + // Arrange + var options = new AgentRunOptions + { + AdditionalProperties = new AdditionalPropertiesDictionary + { + ["key1"] = "value1" + } + }; + + // Act + AgentRunOptions clone = options.Clone(); + clone.AdditionalProperties!["key2"] = "value2"; + + // Assert + Assert.True(clone.AdditionalProperties.ContainsKey("key2")); + Assert.False(options.AdditionalProperties.ContainsKey("key2")); + } } diff --git a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/TestJsonSerializerContext.cs b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/TestJsonSerializerContext.cs index c4f3b7511a..3de33e31d9 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/TestJsonSerializerContext.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/TestJsonSerializerContext.cs @@ -15,6 +15,7 @@ namespace Microsoft.Agents.AI.Abstractions.UnitTests; [JsonSerializable(typeof(AgentResponseUpdate))] [JsonSerializable(typeof(AgentRunOptions))] [JsonSerializable(typeof(Animal))] +[JsonSerializable(typeof(Species))] [JsonSerializable(typeof(JsonElement))] [JsonSerializable(typeof(Dictionary))] [JsonSerializable(typeof(string[]))] diff --git a/dotnet/tests/Microsoft.Agents.AI.DurableTask.UnitTests/DurableAgentRunOptionsTests.cs b/dotnet/tests/Microsoft.Agents.AI.DurableTask.UnitTests/DurableAgentRunOptionsTests.cs new file mode 100644 index 0000000000..77012f4957 --- /dev/null +++ b/dotnet/tests/Microsoft.Agents.AI.DurableTask.UnitTests/DurableAgentRunOptionsTests.cs @@ -0,0 +1,94 @@ +// Copyright (c) Microsoft. All rights reserved. + +using Microsoft.Extensions.AI; + +namespace Microsoft.Agents.AI.DurableTask.UnitTests; + +/// +/// Unit tests for the class. +/// +public sealed class DurableAgentRunOptionsTests +{ + [Fact] + public void CloneReturnsNewInstanceWithSameValues() + { + // Arrange + DurableAgentRunOptions options = new() + { + EnableToolCalls = false, + EnableToolNames = new List { "tool1", "tool2" }, + IsFireAndForget = true, + AllowBackgroundResponses = true, + ContinuationToken = ResponseContinuationToken.FromBytes(new byte[] { 1, 2, 3 }), + AdditionalProperties = new AdditionalPropertiesDictionary + { + ["key1"] = "value1", + ["key2"] = 42 + }, + ResponseFormat = ChatResponseFormat.Json + }; + + // Act + AgentRunOptions cloneAsBase = options.Clone(); + + // Assert + Assert.NotNull(cloneAsBase); + Assert.IsType(cloneAsBase); + DurableAgentRunOptions clone = (DurableAgentRunOptions)cloneAsBase; + Assert.NotSame(options, clone); + Assert.Equal(options.EnableToolCalls, clone.EnableToolCalls); + Assert.NotNull(clone.EnableToolNames); + Assert.NotSame(options.EnableToolNames, clone.EnableToolNames); + Assert.Equal(2, clone.EnableToolNames.Count); + Assert.Contains("tool1", clone.EnableToolNames); + Assert.Contains("tool2", clone.EnableToolNames); + Assert.Equal(options.IsFireAndForget, clone.IsFireAndForget); + Assert.Equal(options.AllowBackgroundResponses, clone.AllowBackgroundResponses); + Assert.Same(options.ContinuationToken, clone.ContinuationToken); + Assert.NotNull(clone.AdditionalProperties); + Assert.NotSame(options.AdditionalProperties, clone.AdditionalProperties); + Assert.Equal("value1", clone.AdditionalProperties["key1"]); + Assert.Equal(42, clone.AdditionalProperties["key2"]); + Assert.Same(options.ResponseFormat, clone.ResponseFormat); + } + + [Fact] + public void CloneCreatesIndependentEnableToolNamesList() + { + // Arrange + DurableAgentRunOptions options = new() + { + EnableToolNames = new List { "tool1" } + }; + + // Act + DurableAgentRunOptions clone = (DurableAgentRunOptions)options.Clone(); + clone.EnableToolNames!.Add("tool2"); + + // Assert + Assert.Equal(2, clone.EnableToolNames.Count); + Assert.Single(options.EnableToolNames); + Assert.DoesNotContain("tool2", options.EnableToolNames); + } + + [Fact] + public void CloneCreatesIndependentAdditionalPropertiesDictionary() + { + // Arrange + DurableAgentRunOptions options = new() + { + AdditionalProperties = new AdditionalPropertiesDictionary + { + ["key1"] = "value1" + } + }; + + // Act + DurableAgentRunOptions clone = (DurableAgentRunOptions)options.Clone(); + clone.AdditionalProperties!["key2"] = "value2"; + + // Assert + Assert.True(clone.AdditionalProperties.ContainsKey("key2")); + Assert.False(options.AdditionalProperties.ContainsKey("key2")); + } +} diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentRunOptionsTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentRunOptionsTests.cs index 1aa49dc328..7a00a9f796 100644 --- a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentRunOptionsTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentRunOptionsTests.cs @@ -332,4 +332,91 @@ public class ChatClientAgentRunOptionsTests } #endregion + + #region Clone Tests + + /// + /// Verify that Clone returns a new instance with the same property values. + /// + [Fact] + public void CloneReturnsNewInstanceWithSameValues() + { + // Arrange + var chatOptions = new ChatOptions { MaxOutputTokens = 100, Temperature = 0.7f }; + Func factory = c => c; + var runOptions = new ChatClientAgentRunOptions(chatOptions) + { + ChatClientFactory = factory, + AllowBackgroundResponses = true, + ContinuationToken = ResponseContinuationToken.FromBytes(new byte[] { 1, 2, 3 }), + AdditionalProperties = new AdditionalPropertiesDictionary + { + ["key1"] = "value1" + } + }; + + // Act + AgentRunOptions cloneAsBase = runOptions.Clone(); + + // Assert + Assert.NotNull(cloneAsBase); + Assert.IsType(cloneAsBase); + ChatClientAgentRunOptions clone = (ChatClientAgentRunOptions)cloneAsBase; + Assert.NotSame(runOptions, clone); + Assert.NotNull(clone.ChatOptions); + Assert.NotSame(runOptions.ChatOptions, clone.ChatOptions); + Assert.Equal(100, clone.ChatOptions!.MaxOutputTokens); + Assert.Equal(0.7f, clone.ChatOptions.Temperature); + Assert.Same(factory, clone.ChatClientFactory); + Assert.Equal(runOptions.AllowBackgroundResponses, clone.AllowBackgroundResponses); + Assert.Same(runOptions.ContinuationToken, clone.ContinuationToken); + Assert.NotNull(clone.AdditionalProperties); + Assert.NotSame(runOptions.AdditionalProperties, clone.AdditionalProperties); + Assert.Equal("value1", clone.AdditionalProperties["key1"]); + } + + /// + /// Verify that modifying the cloned ChatOptions does not affect the original. + /// + [Fact] + public void CloneCreatesIndependentChatOptions() + { + // Arrange + var chatOptions = new ChatOptions { MaxOutputTokens = 100 }; + var runOptions = new ChatClientAgentRunOptions(chatOptions); + + // Act + ChatClientAgentRunOptions clone = (ChatClientAgentRunOptions)runOptions.Clone(); + clone.ChatOptions!.MaxOutputTokens = 200; + + // Assert + Assert.Equal(100, runOptions.ChatOptions!.MaxOutputTokens); + Assert.Equal(200, clone.ChatOptions.MaxOutputTokens); + } + + /// + /// Verify that modifying the cloned AdditionalProperties does not affect the original. + /// + [Fact] + public void CloneCreatesIndependentAdditionalPropertiesDictionary() + { + // Arrange + var runOptions = new ChatClientAgentRunOptions + { + AdditionalProperties = new AdditionalPropertiesDictionary + { + ["key1"] = "value1" + } + }; + + // Act + ChatClientAgentRunOptions clone = (ChatClientAgentRunOptions)runOptions.Clone(); + clone.AdditionalProperties!["key2"] = "value2"; + + // Assert + Assert.True(clone.AdditionalProperties.ContainsKey("key2")); + Assert.False(runOptions.AdditionalProperties.ContainsKey("key2")); + } + + #endregion } diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentTests.cs index 7b33cbbd5f..d90fe5153a 100644 --- a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentTests.cs @@ -3,7 +3,6 @@ using System; using System.Collections.Generic; using System.Linq; -using System.Text.Json; using System.Text.Json.Serialization; using System.Threading; using System.Threading.Tasks; @@ -943,45 +942,6 @@ public partial class ChatClientAgentTests #endregion - #region RunAsync Structured Output Tests - - /// - /// Verify the invocation of with specified type parameter is - /// propagated to the underlying call and the expected structured output is returned. - /// - [Fact] - public async Task RunAsyncWithTypeParameterInvokesChatClientMethodForStructuredOutputAsync() - { - // Arrange - Animal expectedSO = new() { Id = 1, FullName = "Tigger", Species = Species.Tiger }; - - Mock mockService = new(); - mockService.Setup(s => s - .GetResponseAsync( - It.IsAny>(), - It.IsAny(), - It.IsAny())) - .ReturnsAsync(new ChatResponse(new ChatMessage(ChatRole.Assistant, JsonSerializer.Serialize(expectedSO, JsonContext2.Default.Animal))) - { - ResponseId = "test", - }); - - ChatClientAgent agent = new(mockService.Object, options: new()); - - // Act - AgentResponse agentResponse = await agent.RunAsync(messages: [new(ChatRole.User, "Hello")], serializerOptions: JsonContext2.Default.Options); - - // Assert - Assert.Single(agentResponse.Messages); - - Assert.NotNull(agentResponse.Result); - Assert.Equal(expectedSO.Id, agentResponse.Result.Id); - Assert.Equal(expectedSO.FullName, agentResponse.Result.FullName); - Assert.Equal(expectedSO.Species, agentResponse.Result.Species); - } - - #endregion - #region Property Override Tests /// @@ -1999,20 +1959,6 @@ public partial class ChatClientAgentTests } } - private sealed class Animal - { - public int Id { get; set; } - public string? FullName { get; set; } - public Species Species { get; set; } - } - - private enum Species - { - Bear, - Tiger, - Walrus, - } - [JsonSourceGenerationOptions(UseStringEnumConverter = true, PropertyNamingPolicy = JsonKnownNamingPolicy.CamelCase)] [JsonSerializable(typeof(Animal))] private sealed partial class JsonContext2 : JsonSerializerContext; diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_StructuredOutput_WithFormatResponseTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_StructuredOutput_WithFormatResponseTests.cs new file mode 100644 index 0000000000..7ec4a4f59f --- /dev/null +++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_StructuredOutput_WithFormatResponseTests.cs @@ -0,0 +1,212 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System.Collections.Generic; +using System.Text.Json; +using System.Text.Json.Serialization; +using System.Threading; +using System.Threading.Tasks; +using Microsoft.Extensions.AI; +using Moq; + +namespace Microsoft.Agents.AI.UnitTests; + +public partial class ChatClientAgent_StructuredOutput_WithFormatResponseTests +{ + [Fact] + public async Task RunAsync_ResponseFormatProvidedAtAgentInitialization_IsPropagatedToChatClientAsync() + { + // Arrange + ChatResponseFormat? capturedResponseFormat = null; + + Mock mockService = new(); + mockService.Setup(s => s + .GetResponseAsync( + It.IsAny>(), + It.IsAny(), + It.IsAny())) + .Callback, ChatOptions, CancellationToken>((msgs, opts, ct) => capturedResponseFormat = opts?.ResponseFormat) + .ReturnsAsync(new ChatResponse(new ChatMessage(ChatRole.Assistant, "test")) + { + ResponseId = "test", + }); + + ChatResponseFormatJson responseFormat = ChatResponseFormat.ForJsonSchema(JsonContext4.Default.Options); + + ChatClientAgent agent = new(mockService.Object, options: new ChatClientAgentOptions + { + ChatOptions = new ChatOptions() + { + ResponseFormat = responseFormat + } + }); + + // Act + await agent.RunAsync(messages: [new(ChatRole.User, "Hello")]); + + // Assert + Assert.NotNull(capturedResponseFormat); + Assert.Same(responseFormat, capturedResponseFormat); + } + + [Fact] + public async Task RunAsync_ResponseFormatProvidedAtAgentInvocation_IsPropagatedToChatClientAsync() + { + // Arrange + ChatResponseFormat? capturedResponseFormat = null; + + Mock mockService = new(); + mockService.Setup(s => s + .GetResponseAsync( + It.IsAny>(), + It.IsAny(), + It.IsAny())) + .Callback, ChatOptions, CancellationToken>((msgs, opts, ct) => capturedResponseFormat = opts?.ResponseFormat) + .ReturnsAsync(new ChatResponse(new ChatMessage(ChatRole.Assistant, "test")) + { + ResponseId = "test", + }); + + ChatResponseFormatJson responseFormat = ChatResponseFormat.ForJsonSchema(JsonContext4.Default.Options); + + ChatClientAgent agent = new(mockService.Object); + + ChatClientAgentRunOptions runOptions = new() + { + ResponseFormat = responseFormat + }; + + // Act + await agent.RunAsync(messages: [new(ChatRole.User, "Hello")], options: runOptions); + + // Assert + Assert.NotNull(capturedResponseFormat); + Assert.Same(responseFormat, capturedResponseFormat); + } + + [Fact] + public async Task RunAsync_ResponseFormatProvidedAtAgentInvocation_OverridesOneProvidedAtAgentInitializationAsync() + { + // Arrange + ChatResponseFormat? capturedResponseFormat = null; + + Mock mockService = new(); + mockService.Setup(s => s + .GetResponseAsync( + It.IsAny>(), + It.IsAny(), + It.IsAny())) + .Callback, ChatOptions, CancellationToken>((msgs, opts, ct) => capturedResponseFormat = opts?.ResponseFormat) + .ReturnsAsync(new ChatResponse(new ChatMessage(ChatRole.Assistant, "test")) + { + ResponseId = "test", + }); + + ChatResponseFormatJson initializationResponseFormat = ChatResponseFormat.ForJsonSchema(JsonContext4.Default.Options); + ChatResponseFormatJson invocationResponseFormat = ChatResponseFormat.ForJsonSchema(JsonContext4.Default.Options); + + ChatClientAgent agent = new(mockService.Object, options: new ChatClientAgentOptions + { + ChatOptions = new ChatOptions() + { + ResponseFormat = initializationResponseFormat + }, + }); + + ChatClientAgentRunOptions runOptions = new() + { + ResponseFormat = invocationResponseFormat + }; + + // Act + await agent.RunAsync(messages: [new(ChatRole.User, "Hello")], options: runOptions); + + // Assert + Assert.NotNull(capturedResponseFormat); + Assert.Same(invocationResponseFormat, capturedResponseFormat); + Assert.NotSame(initializationResponseFormat, capturedResponseFormat); + } + + [Fact] + public async Task RunAsync_ResponseFormatProvidedAtAgentRunOptions_OverridesOneProvidedViaChatOptionsAsync() + { + // Arrange + ChatResponseFormat? capturedResponseFormat = null; + + Mock mockService = new(); + mockService.Setup(s => s + .GetResponseAsync( + It.IsAny>(), + It.IsAny(), + It.IsAny())) + .Callback, ChatOptions, CancellationToken>((msgs, opts, ct) => capturedResponseFormat = opts?.ResponseFormat) + .ReturnsAsync(new ChatResponse(new ChatMessage(ChatRole.Assistant, "test")) + { + ResponseId = "test", + }); + + ChatResponseFormatJson chatOptionsResponseFormat = ChatResponseFormat.ForJsonSchema(JsonContext4.Default.Options); + ChatResponseFormatJson runOptionsResponseFormat = ChatResponseFormat.ForJsonSchema(JsonContext4.Default.Options); + + ChatClientAgent agent = new(mockService.Object); + + ChatClientAgentRunOptions runOptions = new() + { + ChatOptions = new ChatOptions + { + ResponseFormat = chatOptionsResponseFormat + }, + ResponseFormat = runOptionsResponseFormat + }; + + // Act + await agent.RunAsync(messages: [new(ChatRole.User, "Hello")], options: runOptions); + + // Assert + Assert.NotNull(capturedResponseFormat); + Assert.Same(runOptionsResponseFormat, capturedResponseFormat); + Assert.NotSame(chatOptionsResponseFormat, capturedResponseFormat); + } + + [Fact] + public async Task RunAsync_StructuredOutputResponse_IsAvailableAsTextOnAgentResponseAsync() + { + // Arrange + Animal expectedAnimal = new() { FullName = "Wally the Walrus", Id = 1, Species = Species.Walrus }; + + Mock mockService = new(); + mockService.Setup(s => s + .GetResponseAsync( + It.IsAny>(), + It.IsAny(), + It.IsAny())) + .ReturnsAsync(new ChatResponse(new ChatMessage(ChatRole.Assistant, JsonSerializer.Serialize(expectedAnimal, JsonContext4.Default.Animal))) + { + ResponseId = "test", + }); + + ChatResponseFormatJson responseFormat = ChatResponseFormat.ForJsonSchema(JsonContext4.Default.Options); + + ChatClientAgent agent = new(mockService.Object, options: new ChatClientAgentOptions + { + ChatOptions = new ChatOptions() + { + ResponseFormat = responseFormat + }, + }); + + // Act + AgentResponse agentResponse = await agent.RunAsync(messages: [new(ChatRole.User, "Hello")]); + + // Assert + Assert.NotNull(agentResponse?.Text); + + Animal? deserialised = JsonSerializer.Deserialize(agentResponse.Text, JsonContext4.Default.Animal); + Assert.NotNull(deserialised); + Assert.Equal(expectedAnimal.Id, deserialised.Id); + Assert.Equal(expectedAnimal.FullName, deserialised.FullName); + Assert.Equal(expectedAnimal.Species, deserialised.Species); + } + + [JsonSerializable(typeof(Animal))] + private sealed partial class JsonContext4 : JsonSerializerContext; +} diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_StructuredOutput_WithRunAsyncTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_StructuredOutput_WithRunAsyncTests.cs new file mode 100644 index 0000000000..57e1bbf371 --- /dev/null +++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_StructuredOutput_WithRunAsyncTests.cs @@ -0,0 +1,56 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System.Collections.Generic; +using System.Text.Json; +using System.Text.Json.Serialization; +using System.Threading; +using System.Threading.Tasks; +using Microsoft.Extensions.AI; +using Moq; + +namespace Microsoft.Agents.AI.UnitTests; + +public partial class ChatClientAgent_StructuredOutput_WithRunAsyncTests +{ + [Fact] + public async Task RunAsync_WithGenericType_SetsJsonSchemaResponseFormatAndDeserializesResultAsync() + { + // Arrange + ChatResponseFormat? capturedResponseFormat = null; + ChatResponseFormatJson expectedResponseFormat = ChatResponseFormat.ForJsonSchema(JsonContext3.Default.Options); + Animal expectedSO = new() { Id = 1, FullName = "Tigger", Species = Species.Tiger }; + + Mock mockService = new(); + mockService.Setup(s => s + .GetResponseAsync( + It.IsAny>(), + It.IsAny(), + It.IsAny())) + .Callback, ChatOptions, CancellationToken>((msgs, opts, ct) => capturedResponseFormat = opts?.ResponseFormat) + .ReturnsAsync(new ChatResponse(new ChatMessage(ChatRole.Assistant, JsonSerializer.Serialize(expectedSO, JsonContext3.Default.Animal))) + { + ResponseId = "test", + }); + + ChatClientAgent agent = new(mockService.Object); + + // Act + AgentResponse agentResponse = await agent.RunAsync( + messages: [new(ChatRole.User, "Hello")], + serializerOptions: JsonContext3.Default.Options); + + // Assert + Assert.NotNull(capturedResponseFormat); + Assert.Equal(expectedResponseFormat.Schema?.GetRawText(), ((ChatResponseFormatJson)capturedResponseFormat).Schema?.GetRawText()); + + Animal animal = agentResponse.Result; + Assert.NotNull(animal); + Assert.Equal(expectedSO.Id, animal.Id); + Assert.Equal(expectedSO.FullName, animal.FullName); + Assert.Equal(expectedSO.Species, animal.Species); + } + + [JsonSourceGenerationOptions(UseStringEnumConverter = true, PropertyNamingPolicy = JsonKnownNamingPolicy.CamelCase)] + [JsonSerializable(typeof(Animal))] + private sealed partial class JsonContext3 : JsonSerializerContext; +} diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/Models/Animal.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/Models/Animal.cs new file mode 100644 index 0000000000..331f336b8e --- /dev/null +++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/Models/Animal.cs @@ -0,0 +1,10 @@ +// Copyright (c) Microsoft. All rights reserved. + +namespace Microsoft.Agents.AI.UnitTests; + +internal sealed class Animal +{ + public int Id { get; set; } + public string? FullName { get; set; } + public Species Species { get; set; } +} diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/Models/Species.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/Models/Species.cs new file mode 100644 index 0000000000..14c493be72 --- /dev/null +++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/Models/Species.cs @@ -0,0 +1,10 @@ +// Copyright (c) Microsoft. All rights reserved. + +namespace Microsoft.Agents.AI.UnitTests; + +internal enum Species +{ + Bear, + Tiger, + Walrus, +} diff --git a/dotnet/tests/OpenAIAssistant.IntegrationTests/OpenAIAssistantStructuredOutputRunTests.cs b/dotnet/tests/OpenAIAssistant.IntegrationTests/OpenAIAssistantStructuredOutputRunTests.cs new file mode 100644 index 0000000000..caa42ecc8d --- /dev/null +++ b/dotnet/tests/OpenAIAssistant.IntegrationTests/OpenAIAssistantStructuredOutputRunTests.cs @@ -0,0 +1,9 @@ +// Copyright (c) Microsoft. All rights reserved. + +using AgentConformance.IntegrationTests; + +namespace OpenAIAssistant.IntegrationTests; + +public class OpenAIAssistantStructuredOutputRunTests() : StructuredOutputRunTests(() => new()) +{ +} diff --git a/dotnet/tests/OpenAIChatCompletion.IntegrationTests/OpenAIChatCompletionStructuredOutputRunTests.cs b/dotnet/tests/OpenAIChatCompletion.IntegrationTests/OpenAIChatCompletionStructuredOutputRunTests.cs new file mode 100644 index 0000000000..b7c66f1f26 --- /dev/null +++ b/dotnet/tests/OpenAIChatCompletion.IntegrationTests/OpenAIChatCompletionStructuredOutputRunTests.cs @@ -0,0 +1,9 @@ +// Copyright (c) Microsoft. All rights reserved. + +using AgentConformance.IntegrationTests; + +namespace OpenAIChatCompletion.IntegrationTests; + +public class OpenAIChatCompletionStructuredOutputRunTests() : StructuredOutputRunTests(() => new(useReasoningChatModel: false)) +{ +} diff --git a/dotnet/tests/OpenAIResponse.IntegrationTests/OpenAIResponseStructuredOutputRunTests.cs b/dotnet/tests/OpenAIResponse.IntegrationTests/OpenAIResponseStructuredOutputRunTests.cs new file mode 100644 index 0000000000..497c16eb5a --- /dev/null +++ b/dotnet/tests/OpenAIResponse.IntegrationTests/OpenAIResponseStructuredOutputRunTests.cs @@ -0,0 +1,9 @@ +// Copyright (c) Microsoft. All rights reserved. + +using AgentConformance.IntegrationTests; + +namespace ResponseResult.IntegrationTests; + +public class OpenAIResponseStructuredOutputRunTests() : StructuredOutputRunTests(() => new(store: false)) +{ +} From e563849be330cd242b0fb5d3cb6ddfddb86cb025 Mon Sep 17 00:00:00 2001 From: Eduard van Valkenburg Date: Fri, 13 Feb 2026 20:02:34 +0100 Subject: [PATCH 07/22] Align Python hosting get-started sample with Azure Functions (#3922) --- .../01-get-started/06_host_your_agent.py | 64 +++++++------------ python/samples/01-get-started/README.md | 2 +- 2 files changed, 24 insertions(+), 42 deletions(-) diff --git a/python/samples/01-get-started/06_host_your_agent.py b/python/samples/01-get-started/06_host_your_agent.py index 916944cdb6..7f80607dde 100644 --- a/python/samples/01-get-started/06_host_your_agent.py +++ b/python/samples/01-get-started/06_host_your_agent.py @@ -1,60 +1,42 @@ # Copyright (c) Microsoft. All rights reserved. -import asyncio -import os +"""Host your agent with Azure Functions. -from agent_framework.azure import AzureOpenAIResponsesClient -from azure.identity import AzureCliCredential - -""" -Host Your Agent — Minimal A2A hosting stub - -This sample shows the pattern for exposing an agent via the Agent-to-Agent -(A2A) protocol. It creates the agent and demonstrates how to wrap it with -the A2A hosting layer. +This sample shows the Python hosting pattern used in docs: +- Create an agent with `AzureOpenAIChatClient` +- Register it with `AgentFunctionApp` +- Run with Azure Functions Core Tools (`func start`) Prerequisites: - pip install agent-framework[a2a] --pre + pip install agent-framework-azurefunctions --pre Environment variables: - AZURE_AI_PROJECT_ENDPOINT — Your Azure AI Foundry project endpoint - AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME — Model deployment name (e.g. gpt-4o) - -To run a full A2A server, see samples/04-hosting/a2a/ for a complete example. + AZURE_OPENAI_ENDPOINT + AZURE_OPENAI_CHAT_DEPLOYMENT_NAME """ +from typing import Any -async def main() -> None: - # - credential = AzureCliCredential() - client = AzureOpenAIResponsesClient( - project_endpoint=os.environ["AZURE_AI_PROJECT_ENDPOINT"], - deployment_name=os.environ["AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME"], - credential=credential, - ) +from agent_framework.azure import AgentFunctionApp, AzureOpenAIChatClient +from azure.identity import AzureCliCredential - agent = client.as_agent( + +# +def _create_agent() -> Any: + """Create a hosted agent backed by Azure OpenAI.""" + return AzureOpenAIChatClient(credential=AzureCliCredential()).as_agent( name="HostedAgent", - instructions="You are a helpful assistant exposed via A2A.", + instructions="You are a helpful assistant hosted in Azure Functions.", ) - # - # - # The A2A hosting integration wraps your agent behind an HTTP endpoint. - # Import is gated so this sample can run without the a2a extra installed. - try: - from agent_framework.a2a import A2AAgent # noqa: F401 - print("A2A support is available.") - print("See samples/04-hosting/a2a/ for a runnable A2A server example.") - except ImportError: - print("Install a2a extras: pip install agent-framework[a2a] --pre") +# - # Quick smoke-test: run the agent locally to verify it works - result = await agent.run("Hello! What can you do?") - print(f"Agent: {result}") - # +# +app = AgentFunctionApp(agents=[_create_agent()], enable_health_check=True, max_poll_retries=50) +# if __name__ == "__main__": - asyncio.run(main()) + print("Start the Functions host with: func start") + print("Then call: POST /api/agents/HostedAgent/run") diff --git a/python/samples/01-get-started/README.md b/python/samples/01-get-started/README.md index f13eeb26c5..774368e15f 100644 --- a/python/samples/01-get-started/README.md +++ b/python/samples/01-get-started/README.md @@ -25,7 +25,7 @@ export OPENAI_RESPONSES_MODEL_ID="gpt-4o" # optional, defaults to gpt-4o | 3 | [03_multi_turn.py](03_multi_turn.py) | Keep conversation history across turns with `AgentThread`. | | 4 | [04_memory.py](04_memory.py) | Add dynamic context with a custom `ContextProvider`. | | 5 | [05_first_workflow.py](05_first_workflow.py) | Chain executors into a workflow with edges. | -| 6 | [06_host_your_agent.py](06_host_your_agent.py) | Prepare your agent for A2A hosting. | +| 6 | [06_host_your_agent.py](06_host_your_agent.py) | Host a single agent with Azure Functions. | Run any sample with: From cd1e3110aa9ffec1dadbb53578fab3d914cbc984 Mon Sep 17 00:00:00 2001 From: Copilot <198982749+Copilot@users.noreply.github.com> Date: Fri, 13 Feb 2026 12:13:30 -0800 Subject: [PATCH 08/22] Python: Achieve 85%+ unit test coverage for azurefunctions package (#3866) * Initial plan * Initial analysis: azurefunctions package at 80% coverage, need 85% Co-authored-by: larohra <41490930+larohra@users.noreply.github.com> * Add comprehensive unit tests to achieve 86% coverage for azurefunctions package Co-authored-by: larohra <41490930+larohra@users.noreply.github.com> * Add comprehensive coverage report documentation for azurefunctions package Co-authored-by: larohra <41490930+larohra@users.noreply.github.com> * Fix linting errors: combine nested with statements in test_entities.py Co-authored-by: larohra <41490930+larohra@users.noreply.github.com> * Remove COVERAGE_REPORT.md and coverage.json files as requested Co-authored-by: larohra <41490930+larohra@users.noreply.github.com> * Address PR review feedback: fix unused variables, remove line numbers from docstrings, improve test clarity Co-authored-by: larohra <41490930+larohra@users.noreply.github.com> --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: larohra <41490930+larohra@users.noreply.github.com> Co-authored-by: Laveesh Rohra Co-authored-by: Tao Chen --- .../packages/azurefunctions/tests/test_app.py | 153 ++++++++++++++++++ .../azurefunctions/tests/test_entities.py | 108 +++++++++++++ .../azurefunctions/tests/test_errors.py | 38 +++++ .../tests/test_orchestration.py | 40 +++++ 4 files changed, 339 insertions(+) create mode 100644 python/packages/azurefunctions/tests/test_errors.py diff --git a/python/packages/azurefunctions/tests/test_app.py b/python/packages/azurefunctions/tests/test_app.py index cc68176c93..2729d73e7e 100644 --- a/python/packages/azurefunctions/tests/test_app.py +++ b/python/packages/azurefunctions/tests/test_app.py @@ -1164,5 +1164,158 @@ class TestMCPToolEndpoint: assert body["agents"][0]["mcp_tool_enabled"] is True +class TestAgentFunctionAppErrorPaths: + """Test suite for error handling paths.""" + + def test_init_with_invalid_max_poll_retries(self) -> None: + """Test initialization handles invalid max_poll_retries by falling back to default.""" + mock_agent = Mock() + mock_agent.name = "TestAgent" + + # Test with invalid type + app = AgentFunctionApp(agents=[mock_agent], max_poll_retries="invalid") + assert app.max_poll_retries >= 1 # Should use default + + # Test with None + app2 = AgentFunctionApp(agents=[mock_agent], max_poll_retries=None) + assert app2.max_poll_retries >= 1 # Should use default + + def test_init_with_invalid_poll_interval_seconds(self) -> None: + """Test initialization handles invalid poll_interval_seconds by falling back to default.""" + mock_agent = Mock() + mock_agent.name = "TestAgent" + + # Test with invalid type + app = AgentFunctionApp(agents=[mock_agent], poll_interval_seconds="invalid") + assert app.poll_interval_seconds > 0 # Should use default + + # Test with None + app2 = AgentFunctionApp(agents=[mock_agent], poll_interval_seconds=None) + assert app2.poll_interval_seconds > 0 # Should use default + + def test_get_agent_raises_for_unregistered_agent(self) -> None: + """Test get_agent raises ValueError for unregistered agent.""" + mock_agent = Mock() + mock_agent.name = "RegisteredAgent" + + app = AgentFunctionApp(agents=[mock_agent], enable_http_endpoints=False) + + # Create mock orchestration context + mock_context = Mock() + + # Should raise ValueError for unregistered agent + with pytest.raises(ValueError, match="Agent 'UnknownAgent' is not registered"): + app.get_agent(mock_context, "UnknownAgent") + + def test_convert_payload_to_text_with_response_key(self) -> None: + """Test _convert_payload_to_text returns response key value.""" + app = AgentFunctionApp(enable_http_endpoints=False, enable_health_check=False) + + # Test with response key + payload = {"response": "Test response"} + result = app._convert_payload_to_text(payload) + assert result == "Test response" + + # Test with error key + payload = {"error": "Error message"} + result = app._convert_payload_to_text(payload) + assert result == "Error message" + + # Test with message key + payload = {"message": "Message text"} + result = app._convert_payload_to_text(payload) + assert result == "Message text" + + # Test with no matching keys - should return JSON string + payload = {"other": "value"} + result = app._convert_payload_to_text(payload) + assert "other" in result + assert "value" in result + + def test_create_session_id_with_thread_id(self) -> None: + """Test _create_session_id with provided thread_id.""" + app = AgentFunctionApp(enable_http_endpoints=False, enable_health_check=False) + + # With thread_id provided + session_id = app._create_session_id("TestAgent", "my-thread-123") + assert session_id.key == "my-thread-123" + + # Without thread_id (None) - should generate random + session_id = app._create_session_id("TestAgent", None) + assert session_id.key is not None + assert len(session_id.key) > 0 + + def test_resolve_thread_id_from_body(self) -> None: + """Test _resolve_thread_id extracts from body.""" + app = AgentFunctionApp(enable_http_endpoints=False, enable_health_check=False) + + mock_req = Mock() + mock_req.params = {} + + # Thread ID in body - field name is "thread_id" + req_body = {"thread_id": "body-thread-123"} + result = app._resolve_thread_id(mock_req, req_body) + assert result == "body-thread-123" + + def test_select_body_parser_json_content_type(self) -> None: + """Test _select_body_parser for JSON content type.""" + app = AgentFunctionApp(enable_http_endpoints=False, enable_health_check=False) + + # Test with application/json + parser, format_str = app._select_body_parser("application/json") + assert parser == app._parse_json_body + assert format_str == "json" + + # Test with +json suffix + parser, format_str = app._select_body_parser("application/vnd.api+json") + assert parser == app._parse_json_body + assert format_str == "json" + + def test_accepts_json_response_with_accept_header(self) -> None: + """Test _accepts_json_response checks accept header.""" + app = AgentFunctionApp(enable_http_endpoints=False, enable_health_check=False) + + # With application/json in accept header + headers = {"accept": "application/json"} + result = app._accepts_json_response(headers) + assert result is True + + # Without accept header + headers = {} + result = app._accepts_json_response(headers) + assert result is False + + def test_parse_json_body_invalid_type(self) -> None: + """Test _parse_json_body raises error for invalid JSON.""" + from agent_framework_azurefunctions._errors import IncomingRequestError + + app = AgentFunctionApp(enable_http_endpoints=False, enable_health_check=False) + + # Mock request with non-dict JSON + mock_req = Mock() + mock_req.get_json.return_value = ["not", "a", "dict"] + + with pytest.raises(IncomingRequestError, match="Invalid JSON payload"): + app._parse_json_body(mock_req) + + def test_coerce_to_bool_with_none(self) -> None: + """Test _coerce_to_bool handles None and various value types.""" + app = AgentFunctionApp(enable_http_endpoints=False, enable_health_check=False) + + # None returns False + assert app._coerce_to_bool(None) is False + + # Integer + assert app._coerce_to_bool(1) is True + assert app._coerce_to_bool(0) is False + + # String + assert app._coerce_to_bool("true") is True + assert app._coerce_to_bool("false") is False + + # Other type returns False + assert app._coerce_to_bool([]) is False + + if __name__ == "__main__": pytest.main([__file__, "-v", "--tb=short"]) diff --git a/python/packages/azurefunctions/tests/test_entities.py b/python/packages/azurefunctions/tests/test_entities.py index 2fdbc3463e..a9db60c9dc 100644 --- a/python/packages/azurefunctions/tests/test_entities.py +++ b/python/packages/azurefunctions/tests/test_entities.py @@ -198,6 +198,114 @@ class TestCreateAgentEntity: persisted_state = mock_context.set_state.call_args[0][0] assert persisted_state["data"]["conversationHistory"] == [] + def test_entity_function_handles_string_input(self) -> None: + """Test that the entity function handles non-dict input by converting to string.""" + mock_agent = Mock() + mock_agent.run = AsyncMock(return_value=_agent_response("String response")) + + entity_function = create_agent_entity(mock_agent) + + # Mock context with non-dict input (like a number) + mock_context = Mock() + mock_context.operation_name = "run" + mock_context.entity_key = "conv-456" + # Use a number to test the str() conversion path + mock_context.get_input.return_value = 12345 + mock_context.get_state.return_value = None + + # Execute - entity will convert non-dict input to string + entity_function(mock_context) + + # Verify the result was set + assert mock_context.set_result.called + + def test_entity_function_handles_none_input(self) -> None: + """Test that the entity function handles None input by converting to empty string.""" + mock_agent = Mock() + mock_agent.run = AsyncMock(return_value=_agent_response("Empty response")) + + entity_function = create_agent_entity(mock_agent) + + # Mock context with None input + mock_context = Mock() + mock_context.operation_name = "run" + mock_context.entity_key = "conv-789" + mock_context.get_input.return_value = None + mock_context.get_state.return_value = None + + # Execute - should hit error path since entity expects dict or valid JSON string + entity_function(mock_context) + + # Verify the result was set (likely error result) + assert mock_context.set_result.called + + def test_entity_function_handles_event_loop_runtime_error(self) -> None: + """Test that the entity function handles RuntimeError from get_event_loop by creating a new loop.""" + from unittest.mock import patch + + mock_agent = Mock() + mock_agent.run = AsyncMock(return_value=_agent_response("Response")) + + entity_function = create_agent_entity(mock_agent) + + mock_context = Mock() + mock_context.operation_name = "run" + mock_context.entity_key = "conv-loop-test" + mock_context.get_input.return_value = {"message": "Test"} + mock_context.get_state.return_value = None + + # Simulate RuntimeError when getting event loop + with ( + patch("asyncio.get_event_loop", side_effect=RuntimeError("No event loop")), + patch("asyncio.new_event_loop") as mock_new_loop, + patch("asyncio.set_event_loop") as mock_set_loop, + ): + mock_loop = Mock() + mock_loop.is_running.return_value = False + mock_loop.run_until_complete = Mock() + mock_new_loop.return_value = mock_loop + + # Execute + entity_function(mock_context) + + # Verify new event loop was created + mock_new_loop.assert_called_once() + mock_set_loop.assert_called_once_with(mock_loop) + + def test_entity_function_handles_running_event_loop(self) -> None: + """Test that the entity function handles a running event loop by creating a temporary loop.""" + from unittest.mock import patch + + mock_agent = Mock() + mock_agent.run = AsyncMock(return_value=_agent_response("Response")) + + entity_function = create_agent_entity(mock_agent) + + mock_context = Mock() + mock_context.operation_name = "run" + mock_context.entity_key = "conv-running-loop" + mock_context.get_input.return_value = {"message": "Test"} + mock_context.get_state.return_value = None + + # Simulate a running event loop + mock_existing_loop = Mock() + mock_existing_loop.is_running.return_value = True + + mock_temp_loop = Mock() + mock_temp_loop.run_until_complete = Mock() + mock_temp_loop.close = Mock() + + with ( + patch("asyncio.get_event_loop", return_value=mock_existing_loop), + patch("asyncio.new_event_loop", return_value=mock_temp_loop), + ): + # Execute + entity_function(mock_context) + + # Verify temporary loop was created and closed + mock_temp_loop.run_until_complete.assert_called_once() + mock_temp_loop.close.assert_called_once() + if __name__ == "__main__": pytest.main([__file__, "-v", "--tb=short"]) diff --git a/python/packages/azurefunctions/tests/test_errors.py b/python/packages/azurefunctions/tests/test_errors.py new file mode 100644 index 0000000000..09bf8797c6 --- /dev/null +++ b/python/packages/azurefunctions/tests/test_errors.py @@ -0,0 +1,38 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Unit tests for custom exception types.""" + +import pytest + +from agent_framework_azurefunctions._errors import IncomingRequestError + + +class TestIncomingRequestError: + """Test suite for IncomingRequestError exception.""" + + def test_incoming_request_error_default_status_code(self) -> None: + """Test that IncomingRequestError has a default status code of 400.""" + error = IncomingRequestError("Invalid request") + + assert str(error) == "Invalid request" + assert error.status_code == 400 + + def test_incoming_request_error_custom_status_code(self) -> None: + """Test that IncomingRequestError can have a custom status code.""" + error = IncomingRequestError("Unauthorized", status_code=401) + + assert str(error) == "Unauthorized" + assert error.status_code == 401 + + def test_incoming_request_error_is_value_error(self) -> None: + """Test that IncomingRequestError inherits from ValueError.""" + error = IncomingRequestError("Test error") + + assert isinstance(error, ValueError) + + def test_incoming_request_error_can_be_raised_and_caught(self) -> None: + """Test that IncomingRequestError can be raised and caught.""" + with pytest.raises(IncomingRequestError) as exc_info: + raise IncomingRequestError("Bad request", status_code=400) + + assert exc_info.value.status_code == 400 diff --git a/python/packages/azurefunctions/tests/test_orchestration.py b/python/packages/azurefunctions/tests/test_orchestration.py index d1be7d9a77..45e57f0dcf 100644 --- a/python/packages/azurefunctions/tests/test_orchestration.py +++ b/python/packages/azurefunctions/tests/test_orchestration.py @@ -129,6 +129,25 @@ def executor_with_context(mock_context_with_uuid: tuple[Mock, str]) -> tuple[Any class TestAgentResponseHelpers: """Tests for response handling through public AgentTask API.""" + def test_try_set_value_exception_handling(self) -> None: + """Test try_set_value handles exceptions raised when converting a successful task result to AgentResponse.""" + entity_task = _create_entity_task() + task = AgentTask(entity_task, None, "correlation-id") + + # Simulate successful entity task with invalid result that causes exception + entity_task.state = TaskState.SUCCEEDED + entity_task.result = {"invalid": "format"} # Missing required fields for AgentResponse + + # Clear pending_tasks to simulate that parent has processed the child + task.pending_tasks.clear() + + # Call try_set_value - should catch exception and set error + task.try_set_value(entity_task) + + # Verify task failed due to conversion exception + assert task.state == TaskState.FAILED + assert isinstance(task.result, Exception) + def test_try_set_value_success(self) -> None: """Test try_set_value correctly processes successful task completion.""" entity_task = _create_entity_task() @@ -279,6 +298,27 @@ class TestAzureFunctionsFireAndForget: assert isinstance(result, AgentTask) +class TestAzureFunctionsAgentExecutor: + """Tests for AzureFunctionsAgentExecutor.""" + + def test_generate_unique_id(self, mock_context_with_uuid: tuple[Mock, str]) -> None: + """Test generate_unique_id method returns UUID from orchestration context.""" + from agent_framework_azurefunctions._orchestration import AzureFunctionsAgentExecutor + + context, _ = mock_context_with_uuid + executor = AzureFunctionsAgentExecutor(context) + + # Call generate_unique_id + unique_id = executor.generate_unique_id() + + # Verify it returns the UUID from context (as string with dashes) + # The UUID is returned in standard format with dashes + context.new_uuid.assert_called_once() + # Just verify it's a string representation of UUID + assert isinstance(unique_id, str) + assert len(unique_id) > 0 + + class TestOrchestrationIntegration: """Integration tests for orchestration scenarios.""" From fc9c81b0b11170bdab8fa2d42bb96981e65fd270 Mon Sep 17 00:00:00 2001 From: Eduard van Valkenburg Date: Sat, 14 Feb 2026 11:12:21 +0100 Subject: [PATCH 09/22] Python: [BREAKING] Remove FunctionTool[Any] compatibility shim for schema passthrough (#3600) (#3907) * Fix #3600: Pass JSON schemas through without Pydantic conversion This change optimizes FunctionTool and MCP flows by passing JSON schemas directly to providers without converting them to Pydantic models first. Key changes: - Store JSON schema as-is when supplied to FunctionTool - Skip Pydantic model_validate for schema-supplied tools in invoke() - Return MCP tool schemas directly without conversion - Add comprehensive tests for schema passthrough behavior Performance benefits: - Eliminates expensive Pydantic model creation for supplied schemas - Preserves exact schema structure (additionalProperties, custom fields, etc.) - Reduces memory overhead and initialization time Maintains backward compatibility: - Function signature inference still uses Pydantic models - Explicit Pydantic models passed as input_model work as before - All existing tests pass * Fix schema passthrough validation and remove helper * Simplify FunctionTool without generic model dependency * Fix FunctionTool typing fallout in 3600 * Remove FunctionTool[Any] compatibility shim * Use serializable kwargs in OTEL tool args --- .../ag-ui/agent_framework_ag_ui/_client.py | 2 +- .../ag-ui/agent_framework_ag_ui/_utils.py | 6 +- .../agents/ui_generator_agent.py | 10 +- .../claude/agent_framework_claude/_agent.py | 2 +- .../packages/core/agent_framework/_agents.py | 17 +- python/packages/core/agent_framework/_mcp.py | 33 +- .../core/agent_framework/_middleware.py | 4 +- .../packages/core/agent_framework/_tools.py | 258 +++++-- .../core/agent_framework/observability.py | 2 +- python/packages/core/tests/core/test_mcp.py | 658 +++++++++--------- .../core/tests/core/test_middleware.py | 28 +- .../core/test_middleware_context_result.py | 13 +- python/packages/core/tests/core/test_tools.py | 86 ++- .../agent_framework_github_copilot/_agent.py | 2 +- .../agent_framework_lab_tau2/_tau2_utils.py | 2 +- .../_handoff.py | 4 +- 16 files changed, 645 insertions(+), 482 deletions(-) diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_client.py b/python/packages/ag-ui/agent_framework_ag_ui/_client.py index 373e6321bb..a63b6ac50c 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_client.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_client.py @@ -267,7 +267,7 @@ class AGUIChatClient( if any(getattr(tool, "name", None) == tool_name for tool in additional_tools): return - placeholder: FunctionTool[Any] = FunctionTool( + placeholder: FunctionTool = FunctionTool( name=tool_name, description="Server-managed tool placeholder (AG-UI)", func=None, diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_utils.py b/python/packages/ag-ui/agent_framework_ag_ui/_utils.py index abbfb88562..bfda3948ec 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_utils.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_utils.py @@ -162,7 +162,7 @@ def make_json_safe(obj: Any) -> Any: # noqa: ANN401 def convert_agui_tools_to_agent_framework( agui_tools: list[dict[str, Any]] | None, -) -> list[FunctionTool[Any]] | None: +) -> list[FunctionTool] | None: """Convert AG-UI tool definitions to Agent Framework FunctionTool declarations. Creates declaration-only FunctionTool instances (no executable implementation). @@ -181,13 +181,13 @@ def convert_agui_tools_to_agent_framework( if not agui_tools: return None - result: list[FunctionTool[Any]] = [] + result: list[FunctionTool] = [] for tool_def in agui_tools: # Create declaration-only FunctionTool (func=None means no implementation) # When func=None, the declaration_only property returns True, # which tells the function invocation mixin to return the function call # without executing it (so it can be sent back to the client) - func: FunctionTool[Any] = FunctionTool( + func: FunctionTool = FunctionTool( name=tool_def.get("name", ""), description=tool_def.get("description", ""), func=None, # CRITICAL: Makes declaration_only=True diff --git a/python/packages/ag-ui/agent_framework_ag_ui_examples/agents/ui_generator_agent.py b/python/packages/ag-ui/agent_framework_ag_ui_examples/agents/ui_generator_agent.py index 7f5a4b0f2c..9ea92dd24e 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui_examples/agents/ui_generator_agent.py +++ b/python/packages/ag-ui/agent_framework_ag_ui_examples/agents/ui_generator_agent.py @@ -5,7 +5,7 @@ from __future__ import annotations import sys -from typing import TYPE_CHECKING, Any, TypedDict +from typing import TYPE_CHECKING, TypedDict from agent_framework import Agent, FunctionTool, SupportsChatGetResponse from agent_framework.ag_ui import AgentFrameworkAgent @@ -23,7 +23,7 @@ if TYPE_CHECKING: from agent_framework import ChatOptions # Declaration-only tools (func=None) - actual rendering happens on the client side -generate_haiku = FunctionTool[Any]( +generate_haiku = FunctionTool( name="generate_haiku", description="""Generate a haiku with image and gradient background (FRONTEND_RENDER). @@ -71,7 +71,7 @@ generate_haiku = FunctionTool[Any]( }, ) -create_chart = FunctionTool[Any]( +create_chart = FunctionTool( name="create_chart", description="""Create an interactive chart (FRONTEND_RENDER). @@ -99,7 +99,7 @@ create_chart = FunctionTool[Any]( }, ) -display_timeline = FunctionTool[Any]( +display_timeline = FunctionTool( name="display_timeline", description="""Display an interactive timeline (FRONTEND_RENDER). @@ -127,7 +127,7 @@ display_timeline = FunctionTool[Any]( }, ) -show_comparison_table = FunctionTool[Any]( +show_comparison_table = FunctionTool( name="show_comparison_table", description="""Show a comparison table (FRONTEND_RENDER). diff --git a/python/packages/claude/agent_framework_claude/_agent.py b/python/packages/claude/agent_framework_claude/_agent.py index 10ef5cbf45..ad6f1b3e03 100644 --- a/python/packages/claude/agent_framework_claude/_agent.py +++ b/python/packages/claude/agent_framework_claude/_agent.py @@ -484,7 +484,7 @@ class ClaudeAgent(BaseAgent, Generic[OptionsT]): return create_sdk_mcp_server(name=TOOLS_MCP_SERVER_NAME, tools=sdk_tools), tool_names - def _function_tool_to_sdk_mcp_tool(self, func_tool: FunctionTool[Any]) -> SdkMcpTool[Any]: + def _function_tool_to_sdk_mcp_tool(self, func_tool: FunctionTool) -> SdkMcpTool[Any]: """Convert a FunctionTool to an SDK MCP tool. Args: diff --git a/python/packages/core/agent_framework/_agents.py b/python/packages/core/agent_framework/_agents.py index 49bdd64387..d11f0e2c7d 100644 --- a/python/packages/core/agent_framework/_agents.py +++ b/python/packages/core/agent_framework/_agents.py @@ -439,7 +439,7 @@ class BaseAgent(SerializationMixin): stream_callback: Callable[[AgentResponseUpdate], None] | Callable[[AgentResponseUpdate], Awaitable[None]] | None = None, - ) -> FunctionTool[BaseModel]: + ) -> FunctionTool: """Create a FunctionTool that wraps this agent. Keyword Args: @@ -513,7 +513,7 @@ class BaseAgent(SerializationMixin): # Create final text from accumulated updates return AgentResponse.from_updates(response_updates).text - agent_tool: FunctionTool[BaseModel] = FunctionTool( + agent_tool: FunctionTool = FunctionTool( name=tool_name, description=tool_description, func=agent_wrapper, @@ -1258,17 +1258,12 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc] @server.list_tools() # type: ignore async def _list_tools() -> list[types.Tool]: # type: ignore """List all tools in the agent.""" - # Get the JSON schema from the Pydantic model - schema = agent_tool.input_model.model_json_schema() + schema = agent_tool.parameters() tool = types.Tool( name=agent_tool.name, description=agent_tool.description, - inputSchema={ - "type": "object", - "properties": schema.get("properties", {}), - "required": schema.get("required", []), - }, + inputSchema=schema, ) await _log(level="debug", data=f"Agent tool: {agent_tool}") @@ -1291,7 +1286,9 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc] # Create an instance of the input model with the arguments try: - args_instance = agent_tool.input_model(**arguments) + args_instance: BaseModel | dict[str, Any] = ( + agent_tool.input_model(**arguments) if agent_tool.input_model is not None else arguments + ) result = await agent_tool.invoke(arguments=args_instance) except Exception as e: raise McpError( diff --git a/python/packages/core/agent_framework/_mcp.py b/python/packages/core/agent_framework/_mcp.py index d9a2b5579d..f9d2cd9971 100644 --- a/python/packages/core/agent_framework/_mcp.py +++ b/python/packages/core/agent_framework/_mcp.py @@ -24,11 +24,9 @@ from mcp.client.websocket import websocket_client from mcp.shared.context import RequestContext from mcp.shared.exceptions import McpError from mcp.shared.session import RequestResponder -from pydantic import BaseModel, create_model from ._tools import ( FunctionTool, - _build_pydantic_model_from_json_schema, ) from ._types import ( Content, @@ -355,11 +353,14 @@ def _prepare_message_for_mcp( return messages -def _get_input_model_from_mcp_prompt(prompt: types.Prompt) -> type[BaseModel]: - """Creates a Pydantic model from a prompt's parameters.""" +def _get_input_model_from_mcp_prompt(prompt: types.Prompt) -> dict[str, Any]: + """Get the input model from an MCP prompt. + + Returns a JSON schema dictionary for prompt arguments. + """ # Check if 'arguments' is missing or empty if not prompt.arguments: - return create_model(f"{prompt.name}_input") + return {"type": "object", "properties": {}} # Convert prompt arguments to JSON schema format properties: dict[str, Any] = {} @@ -374,13 +375,10 @@ def _get_input_model_from_mcp_prompt(prompt: types.Prompt) -> type[BaseModel]: if prompt_argument.required: required.append(prompt_argument.name) - schema = {"properties": properties, "required": required} - return _build_pydantic_model_from_json_schema(prompt.name, schema) - - -def _get_input_model_from_mcp_tool(tool: types.Tool) -> type[BaseModel]: - """Creates a Pydantic model from a tools parameters.""" - return _build_pydantic_model_from_json_schema(tool.name, tool.inputSchema) + schema: dict[str, Any] = {"type": "object", "properties": properties} + if required: + schema["required"] = required + return schema def _normalize_mcp_name(name: str) -> str: @@ -467,7 +465,7 @@ class MCPTool: self.session = session self.request_timeout = request_timeout self.client = client - self._functions: list[FunctionTool[Any]] = [] + self._functions: list[FunctionTool] = [] self.is_connected: bool = False self._tools_loaded: bool = False self._prompts_loaded: bool = False @@ -476,7 +474,7 @@ class MCPTool: return f"MCPTool(name={self.name}, description={self.description})" @property - def functions(self) -> list[FunctionTool[Any]]: + def functions(self) -> list[FunctionTool]: """Get the list of functions that are allowed.""" if not self.allowed_tools: return self._functions @@ -744,7 +742,7 @@ class MCPTool: input_model = _get_input_model_from_mcp_prompt(prompt) approval_mode = self._determine_approval_mode(local_name) - func: FunctionTool[BaseModel] = FunctionTool( + func: FunctionTool = FunctionTool( func=partial(self.get_prompt, prompt.name), name=local_name, description=prompt.description or "", @@ -785,15 +783,14 @@ class MCPTool: if local_name in existing_names: continue - input_model = _get_input_model_from_mcp_tool(tool) approval_mode = self._determine_approval_mode(local_name) # Create FunctionTools out of each tool - func: FunctionTool[BaseModel] = FunctionTool( + func: FunctionTool = FunctionTool( func=partial(self.call_tool, tool.name), name=local_name, description=tool.description or "", approval_mode=approval_mode, - input_model=input_model, + input_model=tool.inputSchema, ) self._functions.append(func) existing_names.add(local_name) diff --git a/python/packages/core/agent_framework/_middleware.py b/python/packages/core/agent_framework/_middleware.py index 969ef1efc9..299da150ab 100644 --- a/python/packages/core/agent_framework/_middleware.py +++ b/python/packages/core/agent_framework/_middleware.py @@ -234,8 +234,8 @@ class FunctionInvocationContext: def __init__( self, - function: FunctionTool[Any], - arguments: BaseModel, + function: FunctionTool, + arguments: BaseModel | Mapping[str, Any], metadata: Mapping[str, Any] | None = None, result: Any = None, kwargs: Mapping[str, Any] | None = None, diff --git a/python/packages/core/agent_framework/_tools.py b/python/packages/core/agent_framework/_tools.py index 8374d18e60..4a4e30d324 100644 --- a/python/packages/core/agent_framework/_tools.py +++ b/python/packages/core/agent_framework/_tools.py @@ -26,7 +26,6 @@ from typing import ( Literal, TypedDict, Union, - cast, get_args, get_origin, overload, @@ -89,8 +88,6 @@ DEFAULT_MAX_CONSECUTIVE_ERRORS_PER_REQUEST: Final[int] = 3 ChatClientT = TypeVar("ChatClientT", bound="SupportsChatGetResponse[Any]") # region Helpers -ArgsT = TypeVar("ArgsT", bound=BaseModel, default=BaseModel) - def _parse_inputs( inputs: Content | dict[str, Any] | str | list[Content | dict[str, Any] | str] | None, @@ -183,11 +180,7 @@ def _default_histogram() -> Histogram: ClassT = TypeVar("ClassT", bound="SerializationMixin") -class EmptyInputModel(BaseModel): - """An empty input model for functions with no parameters.""" - - -class FunctionTool(SerializationMixin, Generic[ArgsT]): +class FunctionTool(SerializationMixin): """A tool that wraps a Python function to make it callable by AI models. This class wraps a Python function to make it callable by AI models with automatic @@ -240,6 +233,8 @@ class FunctionTool(SerializationMixin, Generic[ArgsT]): "input_model", "_invocation_duration_histogram", "_cached_parameters", + "_input_schema", + "_schema_supplied", } def __init__( @@ -252,7 +247,7 @@ class FunctionTool(SerializationMixin, Generic[ArgsT]): max_invocation_exceptions: int | None = None, additional_properties: dict[str, Any] | None = None, func: Callable[..., Any] | None = None, - input_model: type[ArgsT] | Mapping[str, Any] | None = None, + input_model: type[BaseModel] | Mapping[str, Any] | None = None, result_parser: Callable[[Any], str] | None = None, **kwargs: Any, ) -> None: @@ -299,7 +294,16 @@ class FunctionTool(SerializationMixin, Generic[ArgsT]): # FunctionTool-specific attributes self.func = func self._instance = None # Store the instance for bound methods - self.input_model = self._resolve_input_model(input_model) + + # Track if schema was supplied as JSON dict (for optimization) + if isinstance(input_model, Mapping): + self._schema_supplied = True + self._input_schema: dict[str, Any] = dict(input_model) + self.input_model: type[BaseModel] | None = None + else: + self._schema_supplied = False + self.input_model = self._resolve_input_model(input_model) + self._input_schema = self.input_model.model_json_schema() self._cached_parameters: dict[str, Any] | None = None self.approval_mode = approval_mode or "never_require" if max_invocations is not None and max_invocations < 1: @@ -335,7 +339,7 @@ class FunctionTool(SerializationMixin, Generic[ArgsT]): return True return self.func is None - def __get__(self, obj: Any, objtype: type | None = None) -> FunctionTool[ArgsT]: + def __get__(self, obj: Any, objtype: type | None = None) -> FunctionTool: """Implement the descriptor protocol to support bound methods. When a FunctionTool is accessed as an attribute of a class instance, @@ -366,17 +370,30 @@ class FunctionTool(SerializationMixin, Generic[ArgsT]): return self - def _resolve_input_model(self, input_model: type[ArgsT] | Mapping[str, Any] | None) -> type[ArgsT]: + def _resolve_input_model(self, input_model: type[BaseModel] | None) -> type[BaseModel]: """Resolve the input model for the function.""" - if input_model is None: - if self.func is None: - return cast(type[ArgsT], EmptyInputModel) - return cast(type[ArgsT], _create_input_model_from_func(func=self.func, name=self.name)) - if inspect.isclass(input_model) and issubclass(input_model, BaseModel): - return input_model - if isinstance(input_model, Mapping): - return cast(type[ArgsT], _create_model_from_json_schema(self.name, input_model)) - raise TypeError("input_model must be a Pydantic BaseModel subclass or a JSON schema dict.") + if input_model is not None: + if inspect.isclass(input_model) and issubclass(input_model, BaseModel): + return input_model + raise TypeError("input_model must be a Pydantic BaseModel subclass or a JSON schema dict.") + + if self.func is None: + return create_model(f"{self.name}_input") + + func = self.func.func if isinstance(self.func, FunctionTool) else self.func + if func is None: + return create_model(f"{self.name}_input") + sig = inspect.signature(func) + fields: dict[str, Any] = { + pname: ( + _parse_annotation(param.annotation) if param.annotation is not inspect.Parameter.empty else str, + param.default if param.default is not inspect.Parameter.empty else ..., + ) + for pname, param in sig.parameters.items() + if pname not in {"self", "cls"} + and param.kind not in {inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD} + } + return create_model(f"{self.name}_input", **fields) def __call__(self, *args: Any, **kwargs: Any) -> Any: """Call the wrapped function with the provided arguments.""" @@ -407,7 +424,7 @@ class FunctionTool(SerializationMixin, Generic[ArgsT]): async def invoke( self, *, - arguments: ArgsT | None = None, + arguments: BaseModel | Mapping[str, Any] | None = None, **kwargs: Any, ) -> str: """Run the AI function with the provided arguments as a Pydantic model. @@ -417,14 +434,14 @@ class FunctionTool(SerializationMixin, Generic[ArgsT]): ``result_parser`` if one was provided. Keyword Args: - arguments: A Pydantic model instance containing the arguments for the function. + arguments: A mapping or model instance containing the arguments for the function. kwargs: Keyword arguments to pass to the function, will not be used if ``arguments`` is provided. Returns: The parsed result as a string — either plain text or serialized JSON. Raises: - TypeError: If arguments is not an instance of the expected input model. + TypeError: If arguments is not mapping-like or fails schema checks. """ if self.declaration_only: raise ToolException(f"Function '{self.name}' is declaration only and cannot be invoked.") @@ -436,9 +453,32 @@ class FunctionTool(SerializationMixin, Generic[ArgsT]): original_kwargs = dict(kwargs) tool_call_id = original_kwargs.pop("tool_call_id", None) if arguments is not None: - if not isinstance(arguments, self.input_model): - raise TypeError(f"Expected {self.input_model.__name__}, got {type(arguments).__name__}") - kwargs = arguments.model_dump(exclude_none=True) + try: + if isinstance(arguments, Mapping): + parsed_arguments = dict(arguments) + if self.input_model is not None and not self._schema_supplied: + parsed_arguments = self.input_model.model_validate(parsed_arguments).model_dump( + exclude_none=True + ) + elif isinstance(arguments, BaseModel): + if ( + self.input_model is not None + and not self._schema_supplied + and not isinstance(arguments, self.input_model) + ): + raise TypeError(f"Expected {self.input_model.__name__}, got {type(arguments).__name__}") + parsed_arguments = arguments.model_dump(exclude_none=True) + else: + raise TypeError( + f"Expected mapping-like arguments for tool '{self.name}', got {type(arguments).__name__}" + ) + except ValidationError as exc: + raise TypeError(f"Invalid arguments for '{self.name}': {exc}") from exc + kwargs = _validate_arguments_against_schema( + arguments=parsed_arguments, + schema=self.parameters(), + tool_name=self.name, + ) if getattr(self, "_forward_runtime_kwargs", False) and original_kwargs: kwargs.update(original_kwargs) else: @@ -458,34 +498,34 @@ class FunctionTool(SerializationMixin, Generic[ArgsT]): return parsed attributes = get_function_span_attributes(self, tool_call_id=tool_call_id) - if OBSERVABILITY_SETTINGS.SENSITIVE_DATA_ENABLED: # type: ignore[name-defined] - # Filter out framework kwargs that are not JSON serializable - serializable_kwargs = { - k: v - for k, v in kwargs.items() - if k - not in { - "chat_options", - "tools", - "tool_choice", - "session", - "conversation_id", - "options", - "response_format", - } + # Filter out framework kwargs that are not JSON serializable. + serializable_kwargs = { + k: v + for k, v in kwargs.items() + if k + not in { + "chat_options", + "tools", + "tool_choice", + "session", + "conversation_id", + "options", + "response_format", } + } + if OBSERVABILITY_SETTINGS.SENSITIVE_DATA_ENABLED: # type: ignore[name-defined] attributes.update({ - OtelAttr.TOOL_ARGUMENTS: arguments.model_dump_json(ensure_ascii=False) - if arguments - else json.dumps(serializable_kwargs, default=str, ensure_ascii=False) - if serializable_kwargs - else "None" + OtelAttr.TOOL_ARGUMENTS: ( + json.dumps(serializable_kwargs, default=str, ensure_ascii=False) + if serializable_kwargs + else "None" + ) }) with get_function_span(attributes=attributes) as span: attributes[OtelAttr.MEASUREMENT_FUNCTION_TAG_NAME] = self.name logger.info(f"Function name: {self.name}") if OBSERVABILITY_SETTINGS.SENSITIVE_DATA_ENABLED: # type: ignore[name-defined] - logger.debug(f"Function arguments: {kwargs}") + logger.debug(f"Function arguments: {serializable_kwargs}") start_time_stamp = perf_counter() end_time_stamp: float | None = None try: @@ -523,7 +563,7 @@ class FunctionTool(SerializationMixin, Generic[ArgsT]): The result is cached after the first call for performance. """ if self._cached_parameters is None: - self._cached_parameters = self.input_model.model_json_schema() + self._cached_parameters = self._input_schema return self._cached_parameters @staticmethod @@ -677,23 +717,79 @@ def _parse_annotation(annotation: Any) -> Any: return annotation -def _create_input_model_from_func(func: Callable[..., Any], name: str) -> type[BaseModel]: - """Create a Pydantic model from a function's signature.""" - # Unwrap FunctionTool objects to get the underlying function - if isinstance(func, FunctionTool): - func = func.func # type: ignore[assignment] +def _matches_json_schema_type(value: Any, schema_type: str) -> bool: + """Check a value against a simple JSON schema primitive type.""" + match schema_type: + case "string": + return isinstance(value, str) + case "integer": + return isinstance(value, int) and not isinstance(value, bool) + case "number": + return (isinstance(value, int | float)) and not isinstance(value, bool) + case "boolean": + return isinstance(value, bool) + case "array": + return isinstance(value, list) + case "object": + return isinstance(value, dict) + case "null": + return value is None + case _: + return True - sig = inspect.signature(func) - fields = { - pname: ( - _parse_annotation(param.annotation) if param.annotation is not inspect.Parameter.empty else str, - param.default if param.default is not inspect.Parameter.empty else ..., - ) - for pname, param in sig.parameters.items() - if pname not in {"self", "cls"} - and param.kind not in {inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD} - } - return create_model(f"{name}_input", **fields) # type: ignore[call-overload, no-any-return] + +def _validate_arguments_against_schema( + *, + arguments: Mapping[str, Any], + schema: Mapping[str, Any], + tool_name: str, +) -> dict[str, Any]: + """Run lightweight argument checks for schema-supplied tools.""" + parsed_arguments = dict(arguments) + + required_raw = schema.get("required", []) + required_fields = [field for field in required_raw if isinstance(field, str)] + missing_fields = [field for field in required_fields if field not in parsed_arguments] + if missing_fields: + raise TypeError(f"Missing required argument(s) for '{tool_name}': {', '.join(sorted(missing_fields))}") + + properties_raw = schema.get("properties") + properties = properties_raw if isinstance(properties_raw, Mapping) else {} + + if schema.get("additionalProperties") is False: + unexpected_fields = sorted(field for field in parsed_arguments if field not in properties) + if unexpected_fields: + raise TypeError(f"Unexpected argument(s) for '{tool_name}': {', '.join(unexpected_fields)}") + + for field_name, field_value in parsed_arguments.items(): + field_schema = properties.get(field_name) + if not isinstance(field_schema, Mapping): + continue + + enum_values = field_schema.get("enum") + if isinstance(enum_values, list) and enum_values and field_value not in enum_values: + raise TypeError( + f"Invalid value for '{field_name}' in '{tool_name}': {field_value!r} is not in {enum_values!r}" + ) + + schema_type = field_schema.get("type") + if isinstance(schema_type, str): + if not _matches_json_schema_type(field_value, schema_type): + raise TypeError( + f"Invalid type for '{field_name}' in '{tool_name}': " + f"expected {schema_type}, got {type(field_value).__name__}" + ) + continue + + if isinstance(schema_type, list): + allowed_types = [item for item in schema_type if isinstance(item, str)] + if allowed_types and not any(_matches_json_schema_type(field_value, item) for item in allowed_types): + raise TypeError( + f"Invalid type for '{field_name}' in '{tool_name}': expected one of " + f"{allowed_types}, got {type(field_value).__name__}" + ) + + return parsed_arguments # Map JSON Schema types to Pydantic types @@ -942,7 +1038,7 @@ def tool( max_invocation_exceptions: int | None = None, additional_properties: dict[str, Any] | None = None, result_parser: Callable[[Any], str] | None = None, -) -> FunctionTool[Any]: ... +) -> FunctionTool: ... @overload @@ -957,7 +1053,7 @@ def tool( max_invocation_exceptions: int | None = None, additional_properties: dict[str, Any] | None = None, result_parser: Callable[[Any], str] | None = None, -) -> Callable[[Callable[..., Any]], FunctionTool[Any]]: ... +) -> Callable[[Callable[..., Any]], FunctionTool]: ... def tool( @@ -971,7 +1067,7 @@ def tool( max_invocation_exceptions: int | None = None, additional_properties: dict[str, Any] | None = None, result_parser: Callable[[Any], str] | None = None, -) -> FunctionTool[Any] | Callable[[Callable[..., Any]], FunctionTool[Any]]: +) -> FunctionTool | Callable[[Callable[..., Any]], FunctionTool]: """Decorate a function to turn it into a FunctionTool that can be passed to models and executed automatically. This decorator creates a Pydantic model from the function's signature, @@ -1095,12 +1191,12 @@ def tool( """ - def decorator(func: Callable[..., Any]) -> FunctionTool[Any]: + def decorator(func: Callable[..., Any]) -> FunctionTool: @wraps(func) - def wrapper(f: Callable[..., Any]) -> FunctionTool[Any]: + def wrapper(f: Callable[..., Any]) -> FunctionTool: tool_name: str = name or getattr(f, "__name__", "unknown_function") # type: ignore[assignment] tool_desc: str = description or (f.__doc__ or "") - return FunctionTool[Any]( + return FunctionTool( name=tool_name, description=tool_desc, approval_mode=approval_mode, @@ -1193,7 +1289,7 @@ async def _auto_invoke_function( custom_args: dict[str, Any] | None = None, *, config: FunctionInvocationConfiguration, - tool_map: dict[str, FunctionTool[BaseModel]], + tool_map: dict[str, FunctionTool], sequence_index: int | None = None, request_index: int | None = None, middleware_pipeline: FunctionMiddlewarePipeline | None = None, # Optional MiddlewarePipeline @@ -1225,7 +1321,7 @@ async def _auto_invoke_function( # this function is called. This function only handles the actual execution of approved, # non-declaration-only functions. - tool: FunctionTool[BaseModel] | None = None + tool: FunctionTool | None = None if function_call_content.type == "function_call": tool = tool_map.get(function_call_content.name) # type: ignore[arg-type] # Tool should exist because _try_execute_function_calls validates this @@ -1258,8 +1354,16 @@ async def _auto_invoke_function( if key not in {"_function_middleware_pipeline", "middleware", "conversation_id"} } try: - args = tool.input_model.model_validate(parsed_args) - except ValidationError as exc: + if not tool._schema_supplied and tool.input_model is not None: + args = tool.input_model.model_validate(parsed_args).model_dump(exclude_none=True) + else: + args = dict(parsed_args) + args = _validate_arguments_against_schema( + arguments=args, + schema=tool.parameters(), + tool_name=tool.name, + ) + except (TypeError, ValidationError) as exc: message = "Error: Argument parsing failed." if config["include_detailed_errors"]: message = f"{message} Exception: {exc}" @@ -1340,8 +1444,8 @@ def _get_tool_map( | Callable[..., Any] | MutableMapping[str, Any] | Sequence[FunctionTool | Callable[..., Any] | MutableMapping[str, Any]], -) -> dict[str, FunctionTool[Any]]: - tool_list: dict[str, FunctionTool[Any]] = {} +) -> dict[str, FunctionTool]: + tool_list: dict[str, FunctionTool] = {} for tool_item in tools if isinstance(tools, list) else [tools]: if isinstance(tool_item, FunctionTool): tool_list[tool_item.name] = tool_item diff --git a/python/packages/core/agent_framework/observability.py b/python/packages/core/agent_framework/observability.py index 5417c8988d..d19683d9e2 100644 --- a/python/packages/core/agent_framework/observability.py +++ b/python/packages/core/agent_framework/observability.py @@ -1448,7 +1448,7 @@ class AgentTelemetryLayer: # region Otel Helpers -def get_function_span_attributes(function: FunctionTool[Any], tool_call_id: str | None = None) -> dict[str, str]: +def get_function_span_attributes(function: FunctionTool, tool_call_id: str | None = None) -> dict[str, str]: """Get the span attributes for the given function. Args: diff --git a/python/packages/core/tests/core/test_mcp.py b/python/packages/core/tests/core/test_mcp.py index 38cb243412..8b213476aa 100644 --- a/python/packages/core/tests/core/test_mcp.py +++ b/python/packages/core/tests/core/test_mcp.py @@ -10,7 +10,7 @@ import pytest from mcp import types from mcp.client.session import ClientSession from mcp.shared.exceptions import McpError -from pydantic import AnyUrl, BaseModel, ValidationError +from pydantic import AnyUrl, BaseModel from agent_framework import ( Content, @@ -22,7 +22,6 @@ from agent_framework import ( from agent_framework._mcp import ( MCPTool, _get_input_model_from_mcp_prompt, - _get_input_model_from_mcp_tool, _normalize_mcp_name, _parse_content_from_mcp, _parse_message_from_mcp, @@ -276,363 +275,338 @@ def test_prepare_message_for_mcp(): @pytest.mark.parametrize( - "test_id,input_schema,valid_data,expected_values,invalid_data,validation_check", + "test_id,input_schema", [ - # Basic types with required/optional fields - ( - "basic_types", - { - "type": "object", - "properties": {"param1": {"type": "string"}, "param2": {"type": "number"}}, - "required": ["param1"], - }, - {"param1": "test", "param2": 42}, - {"param1": "test", "param2": 42}, - {"param2": 42}, # Missing required param1 - None, - ), - # Nested object - ( - "nested_object", - { - "type": "object", - "properties": { - "params": { - "type": "object", - "properties": {"customer_id": {"type": "integer"}}, - "required": ["customer_id"], - } + (test_id, input_schema) + for test_id, input_schema, _, _, _, _ in [ + # Basic types with required/optional fields + ( + "basic_types", + { + "type": "object", + "properties": {"param1": {"type": "string"}, "param2": {"type": "number"}}, + "required": ["param1"], }, - "required": ["params"], - }, - {"params": {"customer_id": 251}}, - {"params.customer_id": 251}, - {"params": {}}, # Missing required customer_id - lambda instance: isinstance(instance.params, BaseModel), - ), - # $ref resolution - ( - "ref_schema", - { - "type": "object", - "properties": {"params": {"$ref": "#/$defs/CustomerIdParam"}}, - "required": ["params"], - "$defs": { - "CustomerIdParam": { - "type": "object", - "properties": {"customer_id": {"type": "integer"}}, - "required": ["customer_id"], - } - }, - }, - {"params": {"customer_id": 251}}, - {"params.customer_id": 251}, - {"params": {}}, # Missing required customer_id - lambda instance: isinstance(instance.params, BaseModel), - ), - # Array of strings (typed) - ( - "array_of_strings", - { - "type": "object", - "properties": { - "tags": { - "type": "array", - "description": "List of tags", - "items": {"type": "string"}, - } - }, - "required": ["tags"], - }, - {"tags": ["tag1", "tag2", "tag3"]}, - {"tags": ["tag1", "tag2", "tag3"]}, - None, # No validation error test for this case - None, - ), - # Array of integers (typed) - ( - "array_of_integers", - { - "type": "object", - "properties": { - "numbers": { - "type": "array", - "description": "List of integers", - "items": {"type": "integer"}, - } - }, - "required": ["numbers"], - }, - {"numbers": [1, 2, 3]}, - {"numbers": [1, 2, 3]}, - None, - None, - ), - # Array of objects (complex nested) - ( - "array_of_objects", - { - "type": "object", - "properties": { - "users": { - "type": "array", - "description": "List of users", - "items": { + {"param1": "test", "param2": 42}, + {"param1": "test", "param2": 42}, + {"param2": 42}, # Missing required param1 + None, + ), + # Nested object + ( + "nested_object", + { + "type": "object", + "properties": { + "params": { "type": "object", - "properties": { - "id": {"type": "integer", "description": "User ID"}, - "name": {"type": "string", "description": "User name"}, - }, - "required": ["id", "name"], - }, - } + "properties": {"customer_id": {"type": "integer"}}, + "required": ["customer_id"], + } + }, + "required": ["params"], }, - "required": ["users"], - }, - {"users": [{"id": 1, "name": "Alice"}, {"id": 2, "name": "Bob"}]}, - {"users[0].id": 1, "users[0].name": "Alice", "users[1].id": 2, "users[1].name": "Bob"}, - {"users": [{"id": 1}]}, # Missing required 'name' - lambda instance: all(isinstance(user, BaseModel) for user in instance.users), - ), - # Deeply nested objects (3+ levels) - ( - "deeply_nested", - { - "type": "object", - "properties": { - "query": { - "type": "object", - "properties": { - "filters": { + {"params": {"customer_id": 251}}, + {"params.customer_id": 251}, + {"params": {}}, # Missing required customer_id + lambda instance: isinstance(instance.params, BaseModel), + ), + # $ref resolution + ( + "ref_schema", + { + "type": "object", + "properties": {"params": {"$ref": "#/$defs/CustomerIdParam"}}, + "required": ["params"], + "$defs": { + "CustomerIdParam": { + "type": "object", + "properties": {"customer_id": {"type": "integer"}}, + "required": ["customer_id"], + } + }, + }, + {"params": {"customer_id": 251}}, + {"params.customer_id": 251}, + {"params": {}}, # Missing required customer_id + lambda instance: isinstance(instance.params, BaseModel), + ), + # Array of strings (typed) + ( + "array_of_strings", + { + "type": "object", + "properties": { + "tags": { + "type": "array", + "description": "List of tags", + "items": {"type": "string"}, + } + }, + "required": ["tags"], + }, + {"tags": ["tag1", "tag2", "tag3"]}, + {"tags": ["tag1", "tag2", "tag3"]}, + None, # No validation error test for this case + None, + ), + # Array of integers (typed) + ( + "array_of_integers", + { + "type": "object", + "properties": { + "numbers": { + "type": "array", + "description": "List of integers", + "items": {"type": "integer"}, + } + }, + "required": ["numbers"], + }, + {"numbers": [1, 2, 3]}, + {"numbers": [1, 2, 3]}, + None, + None, + ), + # Array of objects (complex nested) + ( + "array_of_objects", + { + "type": "object", + "properties": { + "users": { + "type": "array", + "description": "List of users", + "items": { "type": "object", "properties": { - "date_range": { - "type": "object", - "properties": { - "start": {"type": "string"}, - "end": {"type": "string"}, - }, - "required": ["start", "end"], - }, - "categories": {"type": "array", "items": {"type": "string"}}, + "id": {"type": "integer", "description": "User ID"}, + "name": {"type": "string", "description": "User name"}, }, - "required": ["date_range"], - } - }, - "required": ["filters"], + "required": ["id", "name"], + }, + } + }, + "required": ["users"], + }, + {"users": [{"id": 1, "name": "Alice"}, {"id": 2, "name": "Bob"}]}, + {"users[0].id": 1, "users[0].name": "Alice", "users[1].id": 2, "users[1].name": "Bob"}, + {"users": [{"id": 1}]}, # Missing required 'name' + lambda instance: all(isinstance(user, BaseModel) for user in instance.users), + ), + # Deeply nested objects (3+ levels) + ( + "deeply_nested", + { + "type": "object", + "properties": { + "query": { + "type": "object", + "properties": { + "filters": { + "type": "object", + "properties": { + "date_range": { + "type": "object", + "properties": { + "start": {"type": "string"}, + "end": {"type": "string"}, + }, + "required": ["start", "end"], + }, + "categories": {"type": "array", "items": {"type": "string"}}, + }, + "required": ["date_range"], + } + }, + "required": ["filters"], + } + }, + "required": ["query"], + }, + { + "query": { + "filters": { + "date_range": {"start": "2024-01-01", "end": "2024-12-31"}, + "categories": ["tech", "science"], + } } }, - "required": ["query"], - }, - { - "query": { - "filters": { - "date_range": {"start": "2024-01-01", "end": "2024-12-31"}, - "categories": ["tech", "science"], + { + "query.filters.date_range.start": "2024-01-01", + "query.filters.date_range.end": "2024-12-31", + "query.filters.categories": ["tech", "science"], + }, + {"query": {"filters": {"date_range": {}}}}, # Missing required start and end + None, + ), + # Complex $ref with nested structure + ( + "ref_nested_structure", + { + "type": "object", + "properties": {"order": {"$ref": "#/$defs/OrderParams"}}, + "required": ["order"], + "$defs": { + "OrderParams": { + "type": "object", + "properties": { + "customer": {"$ref": "#/$defs/Customer"}, + "items": {"type": "array", "items": {"$ref": "#/$defs/OrderItem"}}, + }, + "required": ["customer", "items"], + }, + "Customer": { + "type": "object", + "properties": {"id": {"type": "integer"}, "email": {"type": "string"}}, + "required": ["id", "email"], + }, + "OrderItem": { + "type": "object", + "properties": {"product_id": {"type": "string"}, "quantity": {"type": "integer"}}, + "required": ["product_id", "quantity"], + }, + }, + }, + { + "order": { + "customer": {"id": 123, "email": "test@example.com"}, + "items": [{"product_id": "prod1", "quantity": 2}], } - } - }, - { - "query.filters.date_range.start": "2024-01-01", - "query.filters.date_range.end": "2024-12-31", - "query.filters.categories": ["tech", "science"], - }, - {"query": {"filters": {"date_range": {}}}}, # Missing required start and end - None, - ), - # Complex $ref with nested structure - ( - "ref_nested_structure", - { - "type": "object", - "properties": {"order": {"$ref": "#/$defs/OrderParams"}}, - "required": ["order"], - "$defs": { - "OrderParams": { - "type": "object", - "properties": { - "customer": {"$ref": "#/$defs/Customer"}, - "items": {"type": "array", "items": {"$ref": "#/$defs/OrderItem"}}, + }, + { + "order.customer.id": 123, + "order.customer.email": "test@example.com", + "order.items[0].product_id": "prod1", + "order.items[0].quantity": 2, + }, + {"order": {"customer": {"id": 123}, "items": []}}, # Missing email + lambda instance: isinstance(instance.order.customer, BaseModel), + ), + # Mixed types (primitives, arrays, nested objects) + ( + "mixed_types", + { + "type": "object", + "properties": { + "simple_string": {"type": "string"}, + "simple_number": {"type": "integer"}, + "string_array": {"type": "array", "items": {"type": "string"}}, + "nested_config": { + "type": "object", + "properties": { + "enabled": {"type": "boolean"}, + "options": {"type": "array", "items": {"type": "string"}}, + }, + "required": ["enabled"], }, - "required": ["customer", "items"], }, - "Customer": { - "type": "object", - "properties": {"id": {"type": "integer"}, "email": {"type": "string"}}, - "required": ["id", "email"], - }, - "OrderItem": { - "type": "object", - "properties": {"product_id": {"type": "string"}, "quantity": {"type": "integer"}}, - "required": ["product_id", "quantity"], + "required": ["simple_string", "nested_config"], + }, + { + "simple_string": "test", + "simple_number": 42, + "string_array": ["a", "b"], + "nested_config": {"enabled": True, "options": ["opt1", "opt2"]}, + }, + { + "simple_string": "test", + "simple_number": 42, + "string_array": ["a", "b"], + "nested_config.enabled": True, + "nested_config.options": ["opt1", "opt2"], + }, + None, + None, + ), + # Empty schema (no properties) + ( + "empty_schema", + {"type": "object", "properties": {}}, + {}, + {}, + None, + None, + ), + # All primitive types + ( + "all_primitives", + { + "type": "object", + "properties": { + "string_field": {"type": "string"}, + "integer_field": {"type": "integer"}, + "number_field": {"type": "number"}, + "boolean_field": {"type": "boolean"}, }, }, - }, - { - "order": { - "customer": {"id": 123, "email": "test@example.com"}, - "items": [{"product_id": "prod1", "quantity": 2}], - } - }, - { - "order.customer.id": 123, - "order.customer.email": "test@example.com", - "order.items[0].product_id": "prod1", - "order.items[0].quantity": 2, - }, - {"order": {"customer": {"id": 123}, "items": []}}, # Missing email - lambda instance: isinstance(instance.order.customer, BaseModel), - ), - # Mixed types (primitives, arrays, nested objects) - ( - "mixed_types", - { - "type": "object", - "properties": { - "simple_string": {"type": "string"}, - "simple_number": {"type": "integer"}, - "string_array": {"type": "array", "items": {"type": "string"}}, - "nested_config": { - "type": "object", - "properties": { - "enabled": {"type": "boolean"}, - "options": {"type": "array", "items": {"type": "string"}}, - }, - "required": ["enabled"], - }, + {"string_field": "test", "integer_field": 42, "number_field": 3.14, "boolean_field": True}, + {"string_field": "test", "integer_field": 42, "number_field": 3.14, "boolean_field": True}, + None, + None, + ), + # Edge case: unresolvable $ref (fallback to dict) + ( + "unresolvable_ref", + { + "type": "object", + "properties": {"data": {"$ref": "#/$defs/NonExistent"}}, + "$defs": {}, }, - "required": ["simple_string", "nested_config"], - }, - { - "simple_string": "test", - "simple_number": 42, - "string_array": ["a", "b"], - "nested_config": {"enabled": True, "options": ["opt1", "opt2"]}, - }, - { - "simple_string": "test", - "simple_number": 42, - "string_array": ["a", "b"], - "nested_config.enabled": True, - "nested_config.options": ["opt1", "opt2"], - }, - None, - None, - ), - # Empty schema (no properties) - ( - "empty_schema", - {"type": "object", "properties": {}}, - {}, - {}, - None, - None, - ), - # All primitive types - ( - "all_primitives", - { - "type": "object", - "properties": { - "string_field": {"type": "string"}, - "integer_field": {"type": "integer"}, - "number_field": {"type": "number"}, - "boolean_field": {"type": "boolean"}, + {"data": {"key": "value"}}, + {"data": {"key": "value"}}, + None, + None, + ), + # Edge case: array without items schema (fallback to bare list) + ( + "array_no_items", + { + "type": "object", + "properties": {"items": {"type": "array"}}, }, - }, - {"string_field": "test", "integer_field": 42, "number_field": 3.14, "boolean_field": True}, - {"string_field": "test", "integer_field": 42, "number_field": 3.14, "boolean_field": True}, - None, - None, - ), - # Edge case: unresolvable $ref (fallback to dict) - ( - "unresolvable_ref", - { - "type": "object", - "properties": {"data": {"$ref": "#/$defs/NonExistent"}}, - "$defs": {}, - }, - {"data": {"key": "value"}}, - {"data": {"key": "value"}}, - None, - None, - ), - # Edge case: array without items schema (fallback to bare list) - ( - "array_no_items", - { - "type": "object", - "properties": {"items": {"type": "array"}}, - }, - {"items": [1, "two", 3.0]}, - {"items": [1, "two", 3.0]}, - None, - None, - ), - # Edge case: object without properties (fallback to dict) - ( - "object_no_properties", - { - "type": "object", - "properties": {"config": {"type": "object"}}, - }, - {"config": {"arbitrary": "data", "nested": {"key": "value"}}}, - {"config": {"arbitrary": "data", "nested": {"key": "value"}}}, - None, - None, - ), + {"items": [1, "two", 3.0]}, + {"items": [1, "two", 3.0]}, + None, + None, + ), + # Edge case: object without properties (fallback to dict) + ( + "object_no_properties", + { + "type": "object", + "properties": {"config": {"type": "object"}}, + }, + {"config": {"arbitrary": "data", "nested": {"key": "value"}}}, + {"config": {"arbitrary": "data", "nested": {"key": "value"}}}, + None, + None, + ), + ] ], ) -def test_get_input_model_from_mcp_tool_parametrized( - test_id, input_schema, valid_data, expected_values, invalid_data, validation_check -): - """Parametrized test for JSON schema to Pydantic model conversion. +def test_get_input_model_from_mcp_tool_parametrized(test_id: str, input_schema: dict[str, Any]) -> None: + """Parametrized test for MCP tool input schema passthrough. - This test covers various edge cases including: - - Basic types with required/optional fields - - Nested objects - - $ref resolution - - Typed arrays (strings, integers, objects) - - Deeply nested structures - - Complex $ref with nested structures - - Mixed types + This test verifies that MCP tool schemas are passed through as-is + without Pydantic conversion, which improves performance and preserves + the original schema structure. To add a new test case, add a tuple to the parametrize decorator with: - test_id: A descriptive name for the test case - input_schema: The JSON schema (inputSchema dict) - - valid_data: Valid data to instantiate the model - - expected_values: Dict of expected values (supports dot notation for nested access) - - invalid_data: Invalid data to test validation errors (None to skip) - - validation_check: Optional callable to perform additional validation checks """ tool = types.Tool(name="test_tool", description="A test tool", inputSchema=input_schema) - model = _get_input_model_from_mcp_tool(tool) + schema = tool.inputSchema - # Test valid data - instance = model(**valid_data) - - # Check expected values - for field_path, expected_value in expected_values.items(): - # Support dot notation and array indexing for nested access - current = instance - parts = field_path.replace("]", "").replace("[", ".").split(".") - for part in parts: - current = current[int(part)] if part.isdigit() else getattr(current, part) - assert current == expected_value, f"Field {field_path} = {current}, expected {expected_value}" - - # Run additional validation checks if provided - if validation_check: - assert validation_check(instance), f"Validation check failed for {test_id}" - - # Test invalid data if provided - if invalid_data is not None: - with pytest.raises(ValidationError): - model(**invalid_data) + # Verify schema is returned as-is (dict) + assert isinstance(schema, dict), f"Expected dict, got {type(schema)}" + assert schema == input_schema, "Schema should be passed through unchanged" def test_get_input_model_from_mcp_prompt(): - """Test creation of input model from MCP prompt.""" + """Test creation of input schema from MCP prompt.""" prompt = types.Prompt( name="test_prompt", description="A test prompt", @@ -641,16 +615,24 @@ def test_get_input_model_from_mcp_prompt(): types.PromptArgument(name="arg2", description="Second argument", required=False), ], ) - model = _get_input_model_from_mcp_prompt(prompt) + result = _get_input_model_from_mcp_prompt(prompt) - # Create an instance to verify the model works - instance = model(arg1="test", arg2="optional") - assert instance.arg1 == "test" - assert instance.arg2 == "optional" + # Should return a dict (schema) + assert isinstance(result, dict), f"Expected dict, got {type(result)}" + assert result["type"] == "object" + assert "arg1" in result["properties"] + assert "arg2" in result["properties"] + assert "arg1" in result["required"] + assert "arg2" not in result["required"] - # Test validation - with pytest.raises(ValidationError): # Missing required arg1 - model(arg2="optional") + +def test_get_input_model_from_mcp_prompt_without_arguments(): + """Test prompt schema generation when no prompt arguments are defined.""" + prompt = types.Prompt(name="empty_prompt", description="No args prompt", arguments=[]) + result = _get_input_model_from_mcp_prompt(prompt) + + assert isinstance(result, dict) + assert result == {"type": "object", "properties": {}} # MCPTool tests diff --git a/python/packages/core/tests/core/test_middleware.py b/python/packages/core/tests/core/test_middleware.py index 4ac4f22f1c..6c559c40d4 100644 --- a/python/packages/core/tests/core/test_middleware.py +++ b/python/packages/core/tests/core/test_middleware.py @@ -74,7 +74,7 @@ class TestAgentContext: class TestFunctionInvocationContext: """Test cases for FunctionInvocationContext.""" - def test_init_with_defaults(self, mock_function: FunctionTool[Any]) -> None: + def test_init_with_defaults(self, mock_function: FunctionTool) -> None: """Test FunctionInvocationContext initialization with default values.""" arguments = FunctionTestArgs(name="test") context = FunctionInvocationContext(function=mock_function, arguments=arguments) @@ -83,7 +83,7 @@ class TestFunctionInvocationContext: assert context.arguments == arguments assert context.metadata == {} - def test_init_with_custom_metadata(self, mock_function: FunctionTool[Any]) -> None: + def test_init_with_custom_metadata(self, mock_function: FunctionTool) -> None: """Test FunctionInvocationContext initialization with custom metadata.""" arguments = FunctionTestArgs(name="test") metadata = {"key": "value"} @@ -420,7 +420,7 @@ class TestFunctionMiddlewarePipeline: await call_next() raise MiddlewareTermination - async def test_execute_with_pre_next_termination(self, mock_function: FunctionTool[Any]) -> None: + async def test_execute_with_pre_next_termination(self, mock_function: FunctionTool) -> None: """Test pipeline execution with termination before next() raises MiddlewareTermination.""" middleware = self.PreNextTerminateFunctionMiddleware() pipeline = FunctionMiddlewarePipeline(middleware) @@ -439,7 +439,7 @@ class TestFunctionMiddlewarePipeline: # Handler should not be called when terminated before next() assert execution_order == [] - async def test_execute_with_post_next_termination(self, mock_function: FunctionTool[Any]) -> None: + async def test_execute_with_post_next_termination(self, mock_function: FunctionTool) -> None: """Test pipeline execution with termination after next() raises MiddlewareTermination.""" middleware = self.PostNextTerminateFunctionMiddleware() pipeline = FunctionMiddlewarePipeline(middleware) @@ -480,7 +480,7 @@ class TestFunctionMiddlewarePipeline: pipeline = FunctionMiddlewarePipeline(test_middleware) assert pipeline.has_middlewares - async def test_execute_no_middleware(self, mock_function: FunctionTool[Any]) -> None: + async def test_execute_no_middleware(self, mock_function: FunctionTool) -> None: """Test pipeline execution with no middleware.""" pipeline = FunctionMiddlewarePipeline() arguments = FunctionTestArgs(name="test") @@ -494,7 +494,7 @@ class TestFunctionMiddlewarePipeline: result = await pipeline.execute(context, final_handler) assert result == expected_result - async def test_execute_with_middleware(self, mock_function: FunctionTool[Any]) -> None: + async def test_execute_with_middleware(self, mock_function: FunctionTool) -> None: """Test pipeline execution with middleware.""" execution_order: list[str] = [] @@ -787,7 +787,7 @@ class TestClassBasedMiddleware: assert context.metadata["after"] is True assert metadata_updates == ["before", "handler", "after"] - async def test_function_middleware_execution(self, mock_function: FunctionTool[Any]) -> None: + async def test_function_middleware_execution(self, mock_function: FunctionTool) -> None: """Test class-based function middleware execution.""" metadata_updates: list[str] = [] @@ -847,7 +847,7 @@ class TestFunctionBasedMiddleware: assert context.metadata["function_middleware"] is True assert execution_order == ["function_before", "handler", "function_after"] - async def test_function_function_middleware(self, mock_function: FunctionTool[Any]) -> None: + async def test_function_function_middleware(self, mock_function: FunctionTool) -> None: """Test function-based function middleware.""" execution_order: list[str] = [] @@ -905,7 +905,7 @@ class TestMixedMiddleware: assert result is not None assert execution_order == ["class_before", "function_before", "handler", "function_after", "class_after"] - async def test_mixed_function_middleware(self, mock_function: FunctionTool[Any]) -> None: + async def test_mixed_function_middleware(self, mock_function: FunctionTool) -> None: """Test mixed class and function-based function middleware.""" execution_order: list[str] = [] @@ -1017,7 +1017,7 @@ class TestMultipleMiddlewareOrdering: ] assert execution_order == expected_order - async def test_function_middleware_execution_order(self, mock_function: FunctionTool[Any]) -> None: + async def test_function_middleware_execution_order(self, mock_function: FunctionTool) -> None: """Test that multiple function middleware execute in registration order.""" execution_order: list[str] = [] @@ -1143,7 +1143,7 @@ class TestContextContentValidation: result = await pipeline.execute(context, final_handler) assert result is not None - async def test_function_context_validation(self, mock_function: FunctionTool[Any]) -> None: + async def test_function_context_validation(self, mock_function: FunctionTool) -> None: """Test that function context contains expected data.""" class ContextValidationMiddleware(FunctionMiddleware): @@ -1489,7 +1489,7 @@ class TestMiddlewareExecutionControl: assert not handler_called assert context.result is None - async def test_function_middleware_no_next_no_execution(self, mock_function: FunctionTool[Any]) -> None: + async def test_function_middleware_no_next_no_execution(self, mock_function: FunctionTool) -> None: """Test that when function middleware doesn't call next(), no execution happens.""" class FunctionTestArgs(BaseModel): @@ -1666,9 +1666,9 @@ def mock_agent() -> SupportsAgentRun: @pytest.fixture -def mock_function() -> FunctionTool[Any]: +def mock_function() -> FunctionTool: """Mock function for testing.""" - function = MagicMock(spec=FunctionTool[Any]) + function = MagicMock(spec=FunctionTool) function.name = "test_function" return function diff --git a/python/packages/core/tests/core/test_middleware_context_result.py b/python/packages/core/tests/core/test_middleware_context_result.py index ba6bfb9c4a..6d9eec351b 100644 --- a/python/packages/core/tests/core/test_middleware_context_result.py +++ b/python/packages/core/tests/core/test_middleware_context_result.py @@ -1,7 +1,6 @@ # Copyright (c) Microsoft. All rights reserved. from collections.abc import AsyncIterable, Awaitable, Callable -from typing import Any from unittest.mock import MagicMock import pytest @@ -103,7 +102,7 @@ class TestResultOverrideMiddleware: assert updates[0].text == "overridden" assert updates[1].text == " stream" - async def test_function_middleware_result_override(self, mock_function: FunctionTool[Any]) -> None: + async def test_function_middleware_result_override(self, mock_function: FunctionTool) -> None: """Test that function middleware can override result.""" override_result = "overridden function result" @@ -252,7 +251,7 @@ class TestResultOverrideMiddleware: assert execute_result.messages[0].text == "executed response" assert handler_called - async def test_function_middleware_conditional_no_next(self, mock_function: FunctionTool[Any]) -> None: + async def test_function_middleware_conditional_no_next(self, mock_function: FunctionTool) -> None: """Test that when function middleware conditionally doesn't call next(), no execution happens.""" class ConditionalNoNextFunctionMiddleware(FunctionMiddleware): @@ -335,7 +334,7 @@ class TestResultObservability: assert observed_responses[0].messages[0].text == "executed response" assert result == observed_responses[0] - async def test_function_middleware_result_observability(self, mock_function: FunctionTool[Any]) -> None: + async def test_function_middleware_result_observability(self, mock_function: FunctionTool) -> None: """Test that middleware can observe function result after execution.""" observed_results: list[str] = [] @@ -402,7 +401,7 @@ class TestResultObservability: assert result is not None assert result.messages[0].text == "modified after execution" - async def test_function_middleware_post_execution_override(self, mock_function: FunctionTool[Any]) -> None: + async def test_function_middleware_post_execution_override(self, mock_function: FunctionTool) -> None: """Test that middleware can override function result after observing execution.""" class PostExecutionOverrideMiddleware(FunctionMiddleware): @@ -444,8 +443,8 @@ def mock_agent() -> SupportsAgentRun: @pytest.fixture -def mock_function() -> FunctionTool[Any]: +def mock_function() -> FunctionTool: """Mock function for testing.""" - function = MagicMock(spec=FunctionTool[Any]) + function = MagicMock(spec=FunctionTool) function.name = "test_function" return function diff --git a/python/packages/core/tests/core/test_tools.py b/python/packages/core/tests/core/test_tools.py index dbcf7aac6f..8d74dc181d 100644 --- a/python/packages/core/tests/core/test_tools.py +++ b/python/packages/core/tests/core/test_tools.py @@ -108,6 +108,90 @@ def test_tool_decorator_with_json_schema_dict(): assert search("hello") == "Searching for: hello (max 10)" +async def test_tool_decorator_with_json_schema_invoke_uses_mapping(): + """Test that schema-based tools can be invoked directly with mapping arguments.""" + + json_schema = { + "type": "object", + "properties": { + "query": {"type": "string"}, + "max_results": {"type": "integer"}, + }, + "required": ["query"], + } + + @tool(name="search", description="Search tool", schema=json_schema) + def search(query: str, max_results: int = 10) -> str: + return f"{query}:{max_results}" + + result = await search.invoke(arguments={"query": "hello", "max_results": 3}) + assert result == "hello:3" + + +async def test_tool_decorator_with_json_schema_invoke_missing_required(): + """Test schema-required fields are checked for mapping arguments.""" + + json_schema = { + "type": "object", + "properties": { + "query": {"type": "string"}, + }, + "required": ["query"], + } + + @tool(name="search", description="Search tool", schema=json_schema) + def search(query: str) -> str: + return query + + with pytest.raises(TypeError, match="Missing required argument"): + await search.invoke(arguments={}) + + +async def test_tool_decorator_with_json_schema_invoke_invalid_type(): + """Test schema type checks run for mapping arguments.""" + + json_schema = { + "type": "object", + "properties": { + "query": {"type": "string"}, + "max_results": {"type": "integer"}, + }, + "required": ["query"], + } + + @tool(name="search", description="Search tool", schema=json_schema) + def search(query: str, max_results: int = 10) -> str: + return f"{query}:{max_results}" + + with pytest.raises(TypeError, match="Invalid type for 'max_results'"): + await search.invoke(arguments={"query": "hello", "max_results": "three"}) + + +def test_tool_decorator_with_json_schema_preserves_custom_properties(): + """Test schema passthrough keeps custom JSON schema properties.""" + + json_schema = { + "type": "object", + "properties": { + "priority": { + "type": "string", + "enum": ["low", "medium", "high"], + "x-custom-field": "custom-value", + }, + }, + "required": ["priority"], + "additionalProperties": False, + } + + @tool(name="process", description="Process tool", schema=json_schema) + def process(priority: str) -> str: + return priority + + params = process.parameters() + assert not params.get("additionalProperties") + assert params["properties"]["priority"]["x-custom-field"] == "custom-value" + + def test_tool_decorator_schema_none_default(): """Test that schema=None (default) still infers from function signature.""" @@ -555,7 +639,7 @@ async def test_tool_invoke_telemetry_with_pydantic_args(span_exporter: InMemoryS assert span.attributes[OtelAttr.TOOL_CALL_ID] == "pydantic_call" assert span.attributes[OtelAttr.TOOL_TYPE] == "function" assert span.attributes[OtelAttr.TOOL_DESCRIPTION] == "A test tool with Pydantic args" - assert span.attributes[OtelAttr.TOOL_ARGUMENTS] == '{"x":5,"y":10}' + assert span.attributes[OtelAttr.TOOL_ARGUMENTS] == '{"x": 5, "y": 10}' async def test_tool_invoke_telemetry_with_exception(span_exporter: InMemorySpanExporter): diff --git a/python/packages/github_copilot/agent_framework_github_copilot/_agent.py b/python/packages/github_copilot/agent_framework_github_copilot/_agent.py index 42de197014..1cafc1ba85 100644 --- a/python/packages/github_copilot/agent_framework_github_copilot/_agent.py +++ b/python/packages/github_copilot/agent_framework_github_copilot/_agent.py @@ -499,7 +499,7 @@ class GitHubCopilotAgent(BaseAgent, Generic[OptionsT]): return copilot_tools - def _tool_to_copilot_tool(self, ai_func: FunctionTool[Any]) -> CopilotTool: + def _tool_to_copilot_tool(self, ai_func: FunctionTool) -> CopilotTool: """Convert an FunctionTool to a Copilot SDK tool.""" async def handler(invocation: ToolInvocation) -> ToolResult: diff --git a/python/packages/lab/tau2/agent_framework_lab_tau2/_tau2_utils.py b/python/packages/lab/tau2/agent_framework_lab_tau2/_tau2_utils.py index 42d03393f8..75c0676cb6 100644 --- a/python/packages/lab/tau2/agent_framework_lab_tau2/_tau2_utils.py +++ b/python/packages/lab/tau2/agent_framework_lab_tau2/_tau2_utils.py @@ -27,7 +27,7 @@ from tau2.environment.tool import Tool # type: ignore[import-untyped] _original_set_state = Environment.set_state -def convert_tau2_tool_to_function_tool(tau2_tool: Tool) -> FunctionTool[Any]: +def convert_tau2_tool_to_function_tool(tau2_tool: Tool) -> FunctionTool: """Convert a tau2 Tool to a FunctionTool for agent framework compatibility. Creates a wrapper that preserves the tool's interface while ensuring diff --git a/python/packages/orchestrations/agent_framework_orchestrations/_handoff.py b/python/packages/orchestrations/agent_framework_orchestrations/_handoff.py index ea5f8bd201..e0dc3e8ea9 100644 --- a/python/packages/orchestrations/agent_framework_orchestrations/_handoff.py +++ b/python/packages/orchestrations/agent_framework_orchestrations/_handoff.py @@ -324,7 +324,7 @@ class HandoffAgentExecutor(AgentExecutor): existing_tools = list(default_options.get("tools") or []) existing_names = {getattr(tool, "name", "") for tool in existing_tools if hasattr(tool, "name")} - new_tools: list[FunctionTool[Any]] = [] + new_tools: list[FunctionTool] = [] for target in targets: handoff_tool = self._create_handoff_tool(target.target_id, target.description) if handoff_tool.name in existing_names: @@ -340,7 +340,7 @@ class HandoffAgentExecutor(AgentExecutor): else: default_options["tools"] = existing_tools - def _create_handoff_tool(self, target_id: str, description: str | None = None) -> FunctionTool[Any]: + def _create_handoff_tool(self, target_id: str, description: str | None = None) -> FunctionTool: """Construct the synthetic handoff tool that signals routing to `target_id`.""" tool_name = get_handoff_tool_name(target_id) doc = description or f"Handoff to the {target_id} agent." From b68d0f93e3fc7504d1cd29089b956beefa531cbc Mon Sep 17 00:00:00 2001 From: Eduard van Valkenburg Date: Mon, 16 Feb 2026 13:46:51 +0100 Subject: [PATCH 10/22] Python: Warn on unsupported AzureAIClient runtime tool/structured_output overrides (#3919) * Guard AzureAIClient runtime tool and structured output overrides * Simplify AzureAI runtime option pruning logic * small fix * slight update * fix error message in test * fix test var * Move Azure AI runtime override checks Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../agent_framework_azure_ai/_client.py | 121 +++++++++++--- .../azure-ai/tests/test_azure_ai_client.py | 149 ++++++++++++++++++ 2 files changed, 249 insertions(+), 21 deletions(-) diff --git a/python/packages/azure-ai/agent_framework_azure_ai/_client.py b/python/packages/azure-ai/agent_framework_azure_ai/_client.py index 79a30b0d81..82338f3d8d 100644 --- a/python/packages/azure-ai/agent_framework_azure_ai/_client.py +++ b/python/packages/azure-ai/agent_framework_azure_ai/_client.py @@ -2,8 +2,10 @@ from __future__ import annotations +import json import sys from collections.abc import Callable, Mapping, MutableMapping, Sequence +from contextlib import suppress from typing import Any, ClassVar, Generic, Literal, TypedDict, TypeVar, cast from agent_framework import ( @@ -218,6 +220,10 @@ class RawAzureAIClient(RawOpenAIResponsesClient[AzureAIClientOptionsT], Generic[ self._is_application_endpoint = "/applications/" in project_client._config.endpoint # type: ignore # Track whether we should close client connection self._should_close_client = should_close_client + # Track creation-time agent configuration for runtime mismatch warnings. + self.warn_runtime_tools_and_structure_changed = False + self._created_agent_tool_names: set[str] = set() + self._created_agent_structured_output_signature: str | None = None async def configure_azure_monitor( self, @@ -341,18 +347,18 @@ class RawAzureAIClient(RawOpenAIResponsesClient[AzureAIClientOptionsT], Generic[ "Agent name is required. Provide 'agent_name' when initializing AzureAIClient " "or 'name' when initializing Agent." ) + # If the agent exists and we do not want to track agent configuration, return early + if self.agent_version is not None and not self.warn_runtime_tools_and_structure_changed: + return {"name": self.agent_name, "version": self.agent_version, "type": "agent_reference"} # If no agent_version is provided, either use latest version or create a new agent: if self.agent_version is None: # Try to use latest version if requested and agent exists if self.use_latest_version: - try: + with suppress(ResourceNotFoundError): existing_agent = await self.project_client.agents.get(self.agent_name) self.agent_version = existing_agent.versions.latest.version return {"name": self.agent_name, "version": self.agent_version, "type": "agent_reference"} - except ResourceNotFoundError: - # Agent doesn't exist, fall through to creation logic - pass if "model" not in run_options or not run_options["model"]: raise ServiceInitializationError( @@ -395,7 +401,9 @@ class RawAzureAIClient(RawOpenAIResponsesClient[AzureAIClientOptionsT], Generic[ ) self.agent_version = created_agent.version - + self.warn_runtime_tools_and_structure_changed = True + self._created_agent_tool_names = self._extract_tool_names(run_options.get("tools")) + self._created_agent_structured_output_signature = self._get_structured_output_signature(chat_options) return {"name": self.agent_name, "version": self.agent_version, "type": "agent_reference"} async def _close_client_if_needed(self) -> None: @@ -403,6 +411,91 @@ class RawAzureAIClient(RawOpenAIResponsesClient[AzureAIClientOptionsT], Generic[ if self._should_close_client: await self.project_client.close() + def _extract_tool_names(self, tools: Any) -> set[str]: + """Extract comparable tool names from runtime tool payloads.""" + if not isinstance(tools, Sequence) or isinstance(tools, str | bytes): + return set() + return {self._get_tool_name(tool) for tool in tools} + + def _get_tool_name(self, tool: Any) -> str: + """Get a stable name for a tool for runtime comparison.""" + if isinstance(tool, FunctionTool): + return tool.name + if isinstance(tool, Mapping): + tool_type = tool.get("type") + if tool_type == "function": + if isinstance(function_data := tool.get("function"), Mapping) and function_data.get("name"): + return str(function_data["name"]) + if tool.get("name"): + return str(tool["name"]) + if tool.get("name"): + return str(tool["name"]) + if tool.get("server_label"): + return f"mcp:{tool['server_label']}" + if tool_type: + return str(tool_type) + if getattr(tool, "name", None): + return str(tool.name) + if getattr(tool, "server_label", None): + return f"mcp:{tool.server_label}" + if getattr(tool, "type", None): + return str(tool.type) + return type(tool).__name__ + + def _get_structured_output_signature(self, chat_options: Mapping[str, Any] | None) -> str | None: + """Build a stable signature for structured_output/response_format values.""" + if not chat_options: + return None + response_format = chat_options.get("response_format") + if response_format is None: + return None + if isinstance(response_format, type): + return f"{response_format.__module__}.{response_format.__qualname__}" + if isinstance(response_format, Mapping): + return json.dumps(response_format, sort_keys=True, default=str) + return str(response_format) + + def _remove_agent_level_run_options( + self, + run_options: dict[str, Any], + chat_options: Mapping[str, Any] | None = None, + ) -> None: + """Remove request-level options that Azure AI only supports at agent creation time.""" + runtime_tools = run_options.get("tools") + runtime_structured_output = self._get_structured_output_signature(chat_options) + + if runtime_tools is not None or runtime_structured_output is not None: + tools_changed = runtime_tools is not None + structured_output_changed = runtime_structured_output is not None + + if self.warn_runtime_tools_and_structure_changed: + if runtime_tools is not None: + tools_changed = self._extract_tool_names(runtime_tools) != self._created_agent_tool_names + if runtime_structured_output is not None: + structured_output_changed = ( + runtime_structured_output != self._created_agent_structured_output_signature + ) + + if tools_changed or structured_output_changed: + logger.warning( + "AzureAIClient does not support runtime tools or structured_output overrides after agent creation. " + "Use AzureOpenAIResponsesClient instead." + ) + + agent_level_option_to_run_keys = { + "model_id": ("model",), + "tools": ("tools",), + "response_format": ("response_format", "text", "text_format"), + "rai_config": ("rai_config",), + "temperature": ("temperature",), + "top_p": ("top_p",), + "reasoning": ("reasoning",), + } + + for run_keys in agent_level_option_to_run_keys.values(): + for run_key in run_keys: + run_options.pop(run_key, None) + @override async def _prepare_options( self, @@ -427,22 +520,8 @@ class RawAzureAIClient(RawOpenAIResponsesClient[AzureAIClientOptionsT], Generic[ agent_reference = await self._get_agent_reference_or_create(run_options, instructions, options) run_options["extra_body"] = {"agent": agent_reference} - # Remove properties that are not supported on request level - # but were configured on agent level - exclude = [ - "model", - "tools", - "response_format", - "rai_config", - "temperature", - "top_p", - "text", - "text_format", - "reasoning", - ] - - for property in exclude: - run_options.pop(property, None) + # Remove only keys that map to this client's declared options TypedDict. + self._remove_agent_level_run_options(run_options, options) return run_options diff --git a/python/packages/azure-ai/tests/test_azure_ai_client.py b/python/packages/azure-ai/tests/test_azure_ai_client.py index 1114747d1b..73b4d3394e 100644 --- a/python/packages/azure-ai/tests/test_azure_ai_client.py +++ b/python/packages/azure-ai/tests/test_azure_ai_client.py @@ -130,6 +130,9 @@ def create_test_azure_ai_client( client.conversation_id = conversation_id client._is_application_endpoint = False # type: ignore client._should_close_client = should_close_client # type: ignore + client.warn_runtime_tools_and_structure_changed = False # type: ignore + client._created_agent_tool_names = set() # type: ignore + client._created_agent_structured_output_signature = None # type: ignore client.additional_properties = {} client.middleware = None @@ -773,6 +776,82 @@ async def test_agent_creation_with_tools( assert call_args[1]["definition"].tools == test_tools +async def test_runtime_tools_override_logs_warning( + mock_project_client: MagicMock, +) -> None: + """Test warning is logged when runtime tools differ from creation-time tools.""" + client = create_test_azure_ai_client(mock_project_client, agent_name="test-agent") + + mock_agent = MagicMock() + mock_agent.name = "test-agent" + mock_agent.version = "1.0" + mock_project_client.agents.create_version = AsyncMock(return_value=mock_agent) + messages = [Message(role="user", contents=[Content.from_text(text="Hello")])] + + with patch( + "agent_framework.openai._responses_client.RawOpenAIResponsesClient._prepare_options", + return_value={"model": "test-model", "tools": [{"type": "function", "name": "tool_one"}]}, + ): + await client._prepare_options(messages, {}) + + with ( + patch( + "agent_framework.openai._responses_client.RawOpenAIResponsesClient._prepare_options", + return_value={"model": "test-model", "tools": [{"type": "function", "name": "tool_two"}]}, + ), + patch("agent_framework_azure_ai._client.logger.warning") as mock_warning, + ): + await client._prepare_options(messages, {}) + mock_warning.assert_called_once() + assert "Use AzureOpenAIResponsesClient instead." in mock_warning.call_args[0][0] + + +async def test_prepare_options_logs_warning_for_tools_with_existing_agent_version( + mock_project_client: MagicMock, +) -> None: + """Test warning is logged when tools are supplied against an existing agent version.""" + client = create_test_azure_ai_client(mock_project_client, agent_name="test-agent", agent_version="1.0") + messages = [Message(role="user", contents=[Content.from_text(text="Hello")])] + + with ( + patch( + "agent_framework.openai._responses_client.RawOpenAIResponsesClient._prepare_options", + return_value={"model": "test-model", "tools": [{"type": "function", "name": "tool_one"}]}, + ), + patch("agent_framework_azure_ai._client.logger.warning") as mock_warning, + ): + run_options = await client._prepare_options(messages, {}) + + mock_warning.assert_called_once() + assert "Use AzureOpenAIResponsesClient instead." in mock_warning.call_args[0][0] + assert "tools" not in run_options + + +async def test_prepare_options_logs_warning_for_tools_on_application_endpoint( + mock_project_client: MagicMock, +) -> None: + """Test warning is logged when runtime tools are removed for application endpoints.""" + client = create_test_azure_ai_client(mock_project_client) + client._is_application_endpoint = True # type: ignore + messages = [Message(role="user", contents=[Content.from_text(text="Hello")])] + + with ( + patch( + "agent_framework.openai._responses_client.RawOpenAIResponsesClient._prepare_options", + return_value={"model": "test-model", "tools": [{"type": "function", "name": "tool_one"}]}, + ), + patch.object(client, "_get_agent_reference_or_create", new_callable=AsyncMock) as mock_get_agent_reference, + patch("agent_framework_azure_ai._client.logger.warning") as mock_warning, + ): + run_options = await client._prepare_options(messages, {}) + + mock_get_agent_reference.assert_not_called() + mock_warning.assert_called_once() + assert "Use AzureOpenAIResponsesClient instead." in mock_warning.call_args[0][0] + assert "tools" not in run_options + assert "extra_body" not in run_options + + async def test_use_latest_version_existing_agent( mock_project_client: MagicMock, ) -> None: @@ -872,6 +951,13 @@ class ResponseFormatModel(BaseModel): model_config = ConfigDict(extra="forbid") +class AlternateResponseFormatModel(BaseModel): + """Alternate model for structured output warning checks.""" + + summary: str + confidence: float + + async def test_agent_creation_with_response_format( mock_project_client: MagicMock, ) -> None: @@ -964,6 +1050,36 @@ async def test_agent_creation_with_mapping_response_format( assert format_config.strict is True +async def test_runtime_structured_output_override_logs_warning( + mock_project_client: MagicMock, +) -> None: + """Test warning is logged when runtime structured_output differs from creation-time configuration.""" + client = create_test_azure_ai_client(mock_project_client, agent_name="test-agent") + + mock_agent = MagicMock() + mock_agent.name = "test-agent" + mock_agent.version = "1.0" + mock_project_client.agents.create_version = AsyncMock(return_value=mock_agent) + messages = [Message(role="user", contents=[Content.from_text(text="Hello")])] + + with patch( + "agent_framework.openai._responses_client.RawOpenAIResponsesClient._prepare_options", + return_value={"model": "test-model"}, + ): + await client._prepare_options(messages, {"response_format": ResponseFormatModel}) + + with ( + patch( + "agent_framework.openai._responses_client.RawOpenAIResponsesClient._prepare_options", + return_value={"model": "test-model"}, + ), + patch("agent_framework_azure_ai._client.logger.warning") as mock_warning, + ): + await client._prepare_options(messages, {"response_format": AlternateResponseFormatModel}) + mock_warning.assert_called_once() + assert "Use AzureOpenAIResponsesClient instead." in mock_warning.call_args[0][0] + + async def test_prepare_options_excludes_response_format( mock_project_client: MagicMock, ) -> None: @@ -1001,6 +1117,39 @@ async def test_prepare_options_excludes_response_format( assert run_options["extra_body"]["agent"]["name"] == "test-agent" +async def test_prepare_options_keeps_values_for_unsupported_option_keys( + mock_project_client: MagicMock, +) -> None: + """Test that run_options removal only applies to known AzureAI agent-level option mappings.""" + client = create_test_azure_ai_client(mock_project_client, agent_name="test-agent", agent_version="1.0") + messages = [Message(role="user", contents=[Content.from_text(text="Hello")])] + + with ( + patch( + "agent_framework.openai._responses_client.RawOpenAIResponsesClient._prepare_options", + return_value={ + "model": "test-model", + "tools": [{"type": "function", "name": "weather"}], + "text": {"format": {"type": "json_schema", "name": "schema"}}, + "text_format": ResponseFormatModel, + "custom_option": "keep-me", + }, + ), + patch.object( + client, + "_get_agent_reference_or_create", + return_value={"name": "test-agent", "version": "1.0", "type": "agent_reference"}, + ), + ): + run_options = await client._prepare_options(messages, {}) + + assert "model" not in run_options + assert "tools" not in run_options + assert "text" not in run_options + assert "text_format" not in run_options + assert run_options["custom_option"] == "keep-me" + + def test_get_conversation_id_with_store_true_and_conversation_id() -> None: """Test _get_conversation_id returns conversation ID when store is True and conversation exists.""" client = create_test_azure_ai_client(MagicMock()) From 7ae4b7b5370a811f1fda8be6339c563289334ff9 Mon Sep 17 00:00:00 2001 From: Copilot <198982749+Copilot@users.noreply.github.com> Date: Mon, 16 Feb 2026 14:13:40 +0000 Subject: [PATCH 11/22] .NET: Add CreateSessionAsync overload with taskId for A2AAgent session resumption (#3924) * Initial plan * Add CreateSessionAsync overload with contextId and taskId parameters Co-authored-by: westey-m <164392973+westey-m@users.noreply.github.com> * Update dotnet/tests/Microsoft.Agents.AI.A2A.UnitTests/A2AAgentTests.cs Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * Add parameter validation to CreateSessionAsync methods Co-authored-by: westey-m <164392973+westey-m@users.noreply.github.com> * Inline parameter validation in CreateSessionAsync methods Co-authored-by: westey-m <164392973+westey-m@users.noreply.github.com> --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: westey-m <164392973+westey-m@users.noreply.github.com> Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- .../src/Microsoft.Agents.AI.A2A/A2AAgent.cs | 11 ++- .../A2AAgentTests.cs | 94 +++++++++++++++++++ 2 files changed, 104 insertions(+), 1 deletion(-) diff --git a/dotnet/src/Microsoft.Agents.AI.A2A/A2AAgent.cs b/dotnet/src/Microsoft.Agents.AI.A2A/A2AAgent.cs index 5601653e8f..18d8bfacdd 100644 --- a/dotnet/src/Microsoft.Agents.AI.A2A/A2AAgent.cs +++ b/dotnet/src/Microsoft.Agents.AI.A2A/A2AAgent.cs @@ -63,7 +63,16 @@ public sealed class A2AAgent : AIAgent /// The context id to continue. /// A value task representing the asynchronous operation. The task result contains a new instance. public ValueTask CreateSessionAsync(string contextId) - => new(new A2AAgentSession() { ContextId = contextId }); + => new(new A2AAgentSession() { ContextId = Throw.IfNullOrWhitespace(contextId) }); + + /// + /// Get a new instance using an existing context id and task id, to resume that conversation from a specific task. + /// + /// The context id to continue. + /// The task id to resume from. + /// A value task representing the asynchronous operation. The task result contains a new instance. + public ValueTask CreateSessionAsync(string contextId, string taskId) + => new(new A2AAgentSession() { ContextId = Throw.IfNullOrWhitespace(contextId), TaskId = Throw.IfNullOrWhitespace(taskId) }); /// protected override ValueTask SerializeSessionCoreAsync(AgentSession session, JsonSerializerOptions? jsonSerializerOptions = null, CancellationToken cancellationToken = default) diff --git a/dotnet/tests/Microsoft.Agents.AI.A2A.UnitTests/A2AAgentTests.cs b/dotnet/tests/Microsoft.Agents.AI.A2A.UnitTests/A2AAgentTests.cs index d6e736a528..50d83c140d 100644 --- a/dotnet/tests/Microsoft.Agents.AI.A2A.UnitTests/A2AAgentTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.A2A.UnitTests/A2AAgentTests.cs @@ -1146,6 +1146,100 @@ public sealed class A2AAgentTests : IDisposable Assert.Equal("a2a", metadata.ProviderName); } + /// + /// Verify that CreateSessionAsync with contextId creates a session with the correct context ID. + /// + [Fact] + public async Task CreateSessionAsync_WithContextId_CreatesSessionWithContextIdAsync() + { + // Arrange + const string ContextId = "test-context-123"; + + // Act + var session = await this._agent.CreateSessionAsync(ContextId); + + // Assert + Assert.NotNull(session); + Assert.IsType(session); + var typedSession = (A2AAgentSession)session; + Assert.Equal(ContextId, typedSession.ContextId); + Assert.Null(typedSession.TaskId); + } + + /// + /// Verify that CreateSessionAsync with contextId and taskId creates a session with both IDs set correctly. + /// + [Fact] + public async Task CreateSessionAsync_WithContextIdAndTaskId_CreatesSessionWithBothIdsAsync() + { + // Arrange + const string ContextId = "test-context-456"; + const string TaskId = "test-task-789"; + + // Act + var session = await this._agent.CreateSessionAsync(ContextId, TaskId); + + // Assert + Assert.NotNull(session); + Assert.IsType(session); + var typedSession = (A2AAgentSession)session; + Assert.Equal(ContextId, typedSession.ContextId); + Assert.Equal(TaskId, typedSession.TaskId); + } + + /// + /// Verify that CreateSessionAsync throws when contextId is null, empty, or whitespace. + /// + [Theory] + [InlineData(null)] + [InlineData("")] + [InlineData(" ")] + [InlineData("\t")] + [InlineData("\r\n")] + public async Task CreateSessionAsync_WithInvalidContextId_ThrowsArgumentExceptionAsync(string? contextId) + { + // Act & Assert + await Assert.ThrowsAnyAsync(async () => + await this._agent.CreateSessionAsync(contextId!)); + } + + /// + /// Verify that CreateSessionAsync with both parameters throws when contextId is null, empty, or whitespace. + /// + [Theory] + [InlineData(null)] + [InlineData("")] + [InlineData(" ")] + [InlineData("\t")] + [InlineData("\r\n")] + public async Task CreateSessionAsync_WithInvalidContextIdAndValidTaskId_ThrowsArgumentExceptionAsync(string? contextId) + { + // Arrange + const string TaskId = "valid-task-id"; + + // Act & Assert + await Assert.ThrowsAnyAsync(async () => + await this._agent.CreateSessionAsync(contextId!, TaskId)); + } + + /// + /// Verify that CreateSessionAsync with both parameters throws when taskId is null, empty, or whitespace. + /// + [Theory] + [InlineData(null)] + [InlineData("")] + [InlineData(" ")] + [InlineData("\t")] + [InlineData("\r\n")] + public async Task CreateSessionAsync_WithValidContextIdAndInvalidTaskId_ThrowsArgumentExceptionAsync(string? taskId) + { + // Arrange + const string ContextId = "valid-context-id"; + + // Act & Assert + await Assert.ThrowsAnyAsync(async () => + await this._agent.CreateSessionAsync(ContextId, taskId!)); + } #endregion public void Dispose() From 503eb10fdd7e455ee4b05e2cf9988fc65c2bfed0 Mon Sep 17 00:00:00 2001 From: westey <164392973+westey-m@users.noreply.github.com> Date: Mon, 16 Feb 2026 14:53:04 +0000 Subject: [PATCH 12/22] .NET: Add skill to verify samples (#3931) * Add skill to verify samples * Apply suggestions from code review Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * Apply suggestions from code review Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * tweak formatting --------- Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- .../skills/verify-dotnet-samples/SKILL.md | 82 +++++++++++++++++++ 1 file changed, 82 insertions(+) create mode 100644 dotnet/.github/skills/verify-dotnet-samples/SKILL.md diff --git a/dotnet/.github/skills/verify-dotnet-samples/SKILL.md b/dotnet/.github/skills/verify-dotnet-samples/SKILL.md new file mode 100644 index 0000000000..36e8a1d3bb --- /dev/null +++ b/dotnet/.github/skills/verify-dotnet-samples/SKILL.md @@ -0,0 +1,82 @@ +--- +name: verify-dotnet-samples +description: > How to build, run and verify the .NET sample projects in the Agent Framework repository. Use this when a user wants to verify that the samples still function as expected. +--- + +# Verifying .NET Sample Projects + +## Sample Pre-requisites + +We should only support verifying samples that: +1. Use environment variables for configuration. +2. Have no complex setup requirements, e.g., where multiple applications need to be run together, or where we need to launch a browser, etc. + +Always report to the user which samples were run and which were not, and why. + +## Verifying a sample + +Samples should be verified to ensure that they actually work as intended and that their output matches what is expected. +For each sample that is run, output should be produced that shows the result and explains the reasoning about what output +was expected, what was produced, and why it didn't match what the sample was expected to produce. + +Steps to verify a sample: +1. Read the code for the sample +1. Check what environment variables are required for the sample +1. Check if each environment variable has been set +1. If there are any missing, give the user a list of missing environment variables to set and terminate +1. Summarize what the expected output of the sample should be +1. Run the sample +1. Show the user any output from the sample run as it gets produced, so that they can see the run progress +1. Check the output of the run against expectations +1. After running all requested samples, produce output for each sample that was verified: + 1. If expectations were matched, output the following: + ```text + [Sample Name] Succeeded + ``` + 1. If expectations were not matched, output the following: + ```text + [Sample Name] Failed + Actual Output: + [What the sample produced] + Expected Output: + [Explanation of what was expected and why the actual output didn't match expectations] + ``` + +## Environment Variables + +Most samples use environment variables to configure settings. + +```csharp +var endpoint = Environment.GetEnvironmentVariable("AZURE_OPENAI_ENDPOINT") ?? throw new InvalidOperationException("AZURE_OPENAI_ENDPOINT is not set."); +var deploymentName = Environment.GetEnvironmentVariable("AZURE_OPENAI_DEPLOYMENT_NAME") ?? "gpt-4o-mini"; +``` + +To run a sample, the environment variables should be set first. +Before running a sample, check whether each environment variable in the sample has a value and +then give the user a list of environment variables to set. + +You can provide the user some examples of how to set the variables like this: + +```bash +export AZURE_OPENAI_ENDPOINT="https://my-openai-instance.openai.azure.com/" +export AZURE_OPENAI_DEPLOYMENT_NAME="gpt-4o-mini" +``` + +To check if a variable has a value use e.g.: + +```bash +echo $AZURE_OPENAI_ENDPOINT +``` + +## How to Run a Sample (General Pattern) + +```bash +cd dotnet/samples// +dotnet run +``` + +For multi-targeted projects (e.g., Durable console apps), specify the framework: + +```bash +dotnet run --framework net10.0 +``` From dc9439a75a2a51b1c908ae7cfb9f0dc17b646528 Mon Sep 17 00:00:00 2001 From: Eduard van Valkenburg Date: Mon, 16 Feb 2026 16:27:25 +0100 Subject: [PATCH 13/22] Python: [BREAKING] Fix #3613 chat/agent message typing alignment (#3920) * Fix #3613 message typing across chat and agents * Address #3613 review feedback and sample input style * refactor: use shared AgentRunMessages aliases (#3613) * refactor: rename agent run input aliases for #3613 * samples: inline image content in run calls * core: export AgentRunInputs from package init * core: use explicit init re-exports without __all__ * updated logging and inits * Fix core mypy export and samples XML note * Remove AgentRunInputsOrNone and dedupe loggers * Remove prepare_messages helper * fix integration tests --- .../skills/python-development/SKILL.md | 7 +- python/CODING_STANDARD.md | 47 ++++----- .../a2a/agent_framework_a2a/_agent.py | 7 +- .../ag-ui/agent_framework_ag_ui/_client.py | 8 +- .../_message_adapters.py | 3 - .../packages/ag-ui/getting_started/client.py | 4 +- .../ag-ui/getting_started/client_advanced.py | 18 ++-- python/packages/ag-ui/tests/ag_ui/conftest.py | 14 +-- .../agent_framework_anthropic/_chat_client.py | 4 +- .../_context_provider.py | 4 +- .../agent_framework_azure_ai/_chat_client.py | 4 +- .../agent_framework_azure_ai/_client.py | 4 +- .../_project_provider.py | 4 +- .../agent_framework_azure_ai/_shared.py | 4 +- .../azure-ai/tests/test_azure_ai_client.py | 17 +++- .../agent_framework_azurefunctions/_app.py | 5 +- .../_entities.py | 5 +- .../_orchestration.py | 5 +- .../agent_framework_bedrock/_chat_client.py | 4 +- .../claude/agent_framework_claude/_agent.py | 14 +-- .../agent_framework_copilotstudio/_agent.py | 11 ++- .../packages/core/agent_framework/__init__.py | 8 +- .../packages/core/agent_framework/_agents.py | 24 +++-- .../packages/core/agent_framework/_clients.py | 38 +++---- .../packages/core/agent_framework/_logging.py | 29 ------ python/packages/core/agent_framework/_mcp.py | 6 -- .../core/agent_framework/_middleware.py | 44 +++------ .../core/agent_framework/_serialization.py | 5 +- .../core/agent_framework/_sessions.py | 10 -- .../core/agent_framework/_settings.py | 1 - .../core/agent_framework/_telemetry.py | 11 +-- .../packages/core/agent_framework/_tools.py | 27 ++--- .../packages/core/agent_framework/_types.py | 78 +-------------- .../core/agent_framework/_workflows/_agent.py | 11 ++- .../_workflows/_checkpoint_encoding.py | 5 +- .../_workflows/_message_utils.py | 19 ++-- .../azure/_assistants_client.py | 2 - .../agent_framework/azure/_chat_client.py | 1 - .../azure/_responses_client.py | 2 - .../core/agent_framework/observability.py | 25 +++-- .../openai/_assistant_provider.py | 1 - .../openai/_assistants_client.py | 6 -- .../agent_framework/openai/_chat_client.py | 6 +- .../agent_framework/openai/_exceptions.py | 2 - .../openai/_responses_client.py | 7 +- .../core/agent_framework/openai/_shared.py | 6 +- .../azure/test_azure_responses_client.py | 19 ++-- .../packages/core/tests/core/test_agents.py | 6 ++ .../packages/core/tests/core/test_clients.py | 12 ++- .../core/test_function_invocation_logic.py | 98 ++++++++++++++----- .../packages/core/tests/core/test_logging.py | 39 -------- .../tests/openai/test_openai_chat_client.py | 9 +- .../openai/test_openai_responses_client.py | 9 +- .../test_agent_executor_tool_calls.py | 2 +- .../tests/workflow/test_full_conversation.py | 4 +- .../core/tests/workflow/test_workflow.py | 2 +- .../tests/workflow/test_workflow_kwargs.py | 4 +- .../agent_framework_declarative/_models.py | 4 +- .../_workflows/_actions_agents.py | 4 +- .../_workflows/_actions_basic.py | 5 +- .../_workflows/_actions_control_flow.py | 5 +- .../_workflows/_actions_error.py | 5 +- .../_workflows/_factory.py | 4 +- .../_workflows/_handlers.py | 5 +- .../_workflows/_human_input.py | 5 +- .../_workflows/_state.py | 5 +- .../agent_framework_durabletask/_client.py | 6 +- .../_durable_agent_state.py | 4 +- .../agent_framework_durabletask/_entities.py | 4 +- .../agent_framework_durabletask/_executors.py | 5 +- .../_orchestration_context.py | 5 +- .../_response_utils.py | 5 +- .../agent_framework_durabletask/_shim.py | 34 +++---- .../agent_framework_durabletask/_worker.py | 5 +- .../agent_framework_github_copilot/_agent.py | 12 +-- .../agent_framework_ollama/_chat_client.py | 4 +- .../orchestrations/tests/test_group_chat.py | 8 +- .../orchestrations/tests/test_magentic.py | 4 +- .../agent_framework_purview/_client.py | 4 +- .../agent_framework_purview/_middleware.py | 4 +- .../agent_framework_purview/_models.py | 4 +- .../agent_framework_purview/_processor.py | 4 +- .../samples/02-agents/observability/README.md | 9 +- ...onfigure_otel_providers_with_parameters.py | 12 ++- .../azure_responses_client_image_analysis.py | 25 ++--- .../openai_responses_client_image_analysis.py | 25 ++--- python/samples/README.md | 4 + 87 files changed, 422 insertions(+), 578 deletions(-) delete mode 100644 python/packages/core/agent_framework/_logging.py delete mode 100644 python/packages/core/tests/core/test_logging.py diff --git a/python/.github/skills/python-development/SKILL.md b/python/.github/skills/python-development/SKILL.md index c19f273588..ad34c2561c 100644 --- a/python/.github/skills/python-development/SKILL.md +++ b/python/.github/skills/python-development/SKILL.md @@ -81,7 +81,12 @@ from agent_framework.azure import AzureOpenAIChatClient ## Public API and Exports -Define `__all__` in each module. Avoid `from module import *` in `__init__.py` files: +In `__init__.py` files that define package-level public APIs, use direct re-export imports plus an explicit +`__all__`. Avoid identity aliases like `from ._agents import ChatAgent as ChatAgent`, and avoid +`from module import *`. + +Do not define `__all__` in internal non-`__init__.py` modules. Exception: modules intentionally exposed as a +public import surface (for example, `agent_framework.observability`) should define `__all__`. ```python __all__ = ["ChatAgent", "Message", "ChatResponse"] diff --git a/python/CODING_STANDARD.md b/python/CODING_STANDARD.md index df143ce62a..028b914ba9 100644 --- a/python/CODING_STANDARD.md +++ b/python/CODING_STANDARD.md @@ -130,27 +130,6 @@ user_msg = UserMessage(content="Hello, world!") asst_msg = AssistantMessage(content="Hello, world!") ``` -### Logging - -Use the centralized logging system: - -```python -from agent_framework import get_logger - -# For main package -logger = get_logger() - -# For subpackages -logger = get_logger('agent_framework.azure') -``` - -**Do not use** direct logging module imports: -```python -# ❌ Avoid this -import logging -logger = logging.getLogger(__name__) -``` - ### Import Structure The package follows a flat import structure: @@ -189,8 +168,6 @@ python/ │ │ ├── _clients.py # Chat client protocols and base classes │ │ ├── _tools.py # Tool definitions │ │ ├── _types.py # Type definitions -│ │ ├── _logging.py # Logging utilities -│ │ │ │ │ │ # Provider folders - lazy load from connector packages │ │ ├── openai/ # OpenAI clients (built into core) │ │ ├── azure/ # Lazy loads from azure-ai, azure-ai-search, azurefunctions @@ -405,12 +382,15 @@ If in doubt, use the link above to read much more considerations of what to do a **All wildcard imports (`from ... import *`) are prohibited** in production code, including both `.py` and `.pyi` files. Always use explicit import lists to maintain clarity and avoid namespace pollution. -Define `__all__` in each module to explicitly declare the public API, then import specific symbols by name: +Do not use ``__all__`` in internal modules. Define it in the ``__init__`` file of the level you want to expose. +If a non-``__init__`` module is intentionally part of the public API surface (for example, ``observability.py``), +it should define ``__all__`` as well. + +Also avoid identity alias imports in ``__init__`` files. Use ``from ._module import Symbol`` instead of +``from ._module import Symbol as Symbol``. ```python # ✅ Preferred - explicit __all__ and named imports -__all__ = ["Agent", "Message", "ChatResponse"] - from ._agents import Agent from ._types import Message, ChatResponse @@ -422,9 +402,20 @@ from ._types import ( ResponseStream, ) +__all__ = [ + "Agent", + "AgentResponse", + "ChatResponse", + "Message", + "ResponseStream", +] + # ❌ Prohibited pattern: wildcard/star imports (do not use) -# from ._agents import -# from ._types import +# from ._agents import * +# from ._types import * + +# ❌ Prohibited pattern: identity alias imports (do not use) +# from ._agents import Agent as Agent ``` **Rationale:** diff --git a/python/packages/a2a/agent_framework_a2a/_agent.py b/python/packages/a2a/agent_framework_a2a/_agent.py index fa53d8d675..2eec8a41db 100644 --- a/python/packages/a2a/agent_framework_a2a/_agent.py +++ b/python/packages/a2a/agent_framework_a2a/_agent.py @@ -40,6 +40,7 @@ from agent_framework import ( normalize_messages, prepend_agent_framework_to_user_agent, ) +from agent_framework._types import AgentRunInputs from agent_framework.observability import AgentTelemetryLayer __all__ = ["A2AAgent", "A2AContinuationToken"] @@ -208,7 +209,7 @@ class A2AAgent(AgentTelemetryLayer, BaseAgent): @overload def run( self, - messages: str | Message | Sequence[str | Message] | None = None, + messages: AgentRunInputs | None = None, *, stream: Literal[False] = ..., session: AgentSession | None = None, @@ -220,7 +221,7 @@ class A2AAgent(AgentTelemetryLayer, BaseAgent): @overload def run( self, - messages: str | Message | Sequence[str | Message] | None = None, + messages: AgentRunInputs | None = None, *, stream: Literal[True], session: AgentSession | None = None, @@ -231,7 +232,7 @@ class A2AAgent(AgentTelemetryLayer, BaseAgent): def run( self, - messages: str | Message | Sequence[str | Message] | None = None, + messages: AgentRunInputs | None = None, *, stream: bool = False, session: AgentSession | None = None, diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_client.py b/python/packages/ag-ui/agent_framework_ag_ui/_client.py index a63b6ac50c..21c9e484f3 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_client.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_client.py @@ -49,7 +49,7 @@ if TYPE_CHECKING: from ._types import AGUIChatOptions -logger: logging.Logger = logging.getLogger(__name__) +logger: logging.Logger = logging.getLogger("agent_framework.ag_ui") def _unwrap_server_function_call_contents(contents: MutableSequence[Content | dict[str, Any]]) -> None: @@ -277,9 +277,6 @@ class AGUIChatClient( registered: set[str] = getattr(self, "_registered_server_tools", set()) registered.add(tool_name) self._registered_server_tools = registered # type: ignore[attr-defined] - from agent_framework._logging import get_logger - - logger = get_logger() logger.debug(f"[AGUIChatClient] Registered server placeholder: {tool_name}") def _extract_state_from_messages(self, messages: Sequence[Message]) -> tuple[list[Message], dict[str, Any] | None]: @@ -310,9 +307,6 @@ class AGUIChatClient( messages_without_state = list(messages[:-1]) if len(messages) > 1 else [] return messages_without_state, state except (json.JSONDecodeError, ValueError, KeyError) as e: - from agent_framework._logging import get_logger - - logger = get_logger() logger.warning(f"Failed to extract state from message: {e}") return list(messages), None diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_message_adapters.py b/python/packages/ag-ui/agent_framework_ag_ui/_message_adapters.py index efdb3a0f53..a0a86ab5de 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_message_adapters.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_message_adapters.py @@ -408,9 +408,6 @@ def agui_messages_to_agent_framework(messages: list[dict[str, Any]]) -> list[Mes approval_payload_text = result_content if isinstance(result_content, str) else json.dumps(parsed) # Log the full approval payload to debug modified arguments - import logging - - logger = logging.getLogger(__name__) logger.info(f"Approval payload received: {parsed}") approval_call_id = tool_call_id diff --git a/python/packages/ag-ui/getting_started/client.py b/python/packages/ag-ui/getting_started/client.py index d75aedc3df..fc0dfc3884 100644 --- a/python/packages/ag-ui/getting_started/client.py +++ b/python/packages/ag-ui/getting_started/client.py @@ -11,7 +11,7 @@ import asyncio import os from typing import cast -from agent_framework import ChatResponse, ChatResponseUpdate, ResponseStream +from agent_framework import ChatResponse, ChatResponseUpdate, Message, ResponseStream from agent_framework.ag_ui import AGUIChatClient @@ -44,7 +44,7 @@ async def main(): metadata = {"thread_id": thread_id} if thread_id else None stream = client.get_response( - message, + [Message(role="user", text=message)], stream=True, options={"metadata": metadata} if metadata else None, ) diff --git a/python/packages/ag-ui/getting_started/client_advanced.py b/python/packages/ag-ui/getting_started/client_advanced.py index dcb4e5ca3c..25f1edc7a4 100644 --- a/python/packages/ag-ui/getting_started/client_advanced.py +++ b/python/packages/ag-ui/getting_started/client_advanced.py @@ -15,7 +15,7 @@ import asyncio import os from typing import cast -from agent_framework import ChatResponse, ChatResponseUpdate, ResponseStream, tool +from agent_framework import ChatResponse, ChatResponseUpdate, Message, ResponseStream, tool from agent_framework.ag_ui import AGUIChatClient @@ -73,7 +73,7 @@ async def streaming_example(client: AGUIChatClient, thread_id: str | None = None print("Assistant: ", end="", flush=True) stream = client.get_response( - "Tell me a short joke", + [Message(role="user", text="Tell me a short joke")], stream=True, options={"metadata": metadata} if metadata else None, ) @@ -100,7 +100,7 @@ async def non_streaming_example(client: AGUIChatClient, thread_id: str | None = print("\nUser: What is 2 + 2?\n") - response = await client.get_response("What is 2 + 2?", metadata=metadata) + response = await client.get_response([Message(role="user", text="What is 2 + 2?")], metadata=metadata) print(f"Assistant: {response.text}") @@ -139,7 +139,7 @@ async def tool_example(client: AGUIChatClient, thread_id: str | None = None): print("(Server must be configured with matching tools to execute them)\n") response = await client.get_response( - "What's the weather in Seattle?", tools=[get_weather, calculate], metadata=metadata + [Message(role="user", text="What's the weather in Seattle?")], tools=[get_weather, calculate], metadata=metadata ) print(f"Assistant: {response.text}") @@ -174,14 +174,16 @@ async def conversation_example(client: AGUIChatClient): # First turn print("User: My name is Alice\n") - response1 = await client.get_response("My name is Alice") + response1 = await client.get_response([Message(role="user", text="My name is Alice")]) print(f"Assistant: {response1.text}") thread_id = response1.additional_properties.get("thread_id") print(f"\n[Thread: {thread_id}]") # Second turn - using same thread print("\nUser: What's my name?\n") - response2 = await client.get_response("What's my name?", options={"metadata": {"thread_id": thread_id}}) + response2 = await client.get_response( + [Message(role="user", text="What's my name?")], options={"metadata": {"thread_id": thread_id}} + ) print(f"Assistant: {response2.text}") # Check if context was maintained @@ -191,7 +193,9 @@ async def conversation_example(client: AGUIChatClient): # Third turn print("\nUser: Can you also tell me what 10 * 5 is?\n") response3 = await client.get_response( - "Can you also tell me what 10 * 5 is?", options={"metadata": {"thread_id": thread_id}}, tools=[calculate] + [Message(role="user", text="Can you also tell me what 10 * 5 is?")], + options={"metadata": {"thread_id": thread_id}}, + tools=[calculate], ) print(f"Assistant: {response3.text}") diff --git a/python/packages/ag-ui/tests/ag_ui/conftest.py b/python/packages/ag-ui/tests/ag_ui/conftest.py index 09a4ff57f1..d86ebb1720 100644 --- a/python/packages/ag-ui/tests/ag_ui/conftest.py +++ b/python/packages/ag-ui/tests/ag_ui/conftest.py @@ -55,7 +55,7 @@ class StreamingChatClientStub( @overload def get_response( self, - messages: str | Message | Sequence[str | Message], + messages: Sequence[Message], *, stream: Literal[False] = ..., options: ChatOptions[Any], @@ -65,7 +65,7 @@ class StreamingChatClientStub( @overload def get_response( self, - messages: str | Message | Sequence[str | Message], + messages: Sequence[Message], *, stream: Literal[False] = ..., options: OptionsCoT | ChatOptions[None] | None = ..., @@ -75,7 +75,7 @@ class StreamingChatClientStub( @overload def get_response( self, - messages: str | Message | Sequence[str | Message], + messages: Sequence[Message], *, stream: Literal[True], options: OptionsCoT | ChatOptions[Any] | None = ..., @@ -84,7 +84,7 @@ class StreamingChatClientStub( def get_response( self, - messages: str | Message | Sequence[str | Message], + messages: Sequence[Message], *, stream: bool = False, options: OptionsCoT | ChatOptions[Any] | None = None, @@ -175,7 +175,7 @@ class StubAgent(SupportsAgentRun): @overload def run( self, - messages: str | Message | Sequence[str | Message] | None = None, + messages: str | Content | Message | Sequence[str | Content | Message] | None = None, *, stream: Literal[False] = ..., session: AgentSession | None = None, @@ -185,7 +185,7 @@ class StubAgent(SupportsAgentRun): @overload def run( self, - messages: str | Message | Sequence[str | Message] | None = None, + messages: str | Content | Message | Sequence[str | Content | Message] | None = None, *, stream: Literal[True], session: AgentSession | None = None, @@ -194,7 +194,7 @@ class StubAgent(SupportsAgentRun): def run( self, - messages: str | Message | Sequence[str | Message] | None = None, + messages: str | Content | Message | Sequence[str | Content | Message] | None = None, *, stream: bool = False, session: AgentSession | None = None, diff --git a/python/packages/anthropic/agent_framework_anthropic/_chat_client.py b/python/packages/anthropic/agent_framework_anthropic/_chat_client.py index d3ea19dfa0..a6b0c8e7d3 100644 --- a/python/packages/anthropic/agent_framework_anthropic/_chat_client.py +++ b/python/packages/anthropic/agent_framework_anthropic/_chat_client.py @@ -2,6 +2,7 @@ from __future__ import annotations +import logging import sys from collections.abc import AsyncIterable, Awaitable, Mapping, MutableMapping, Sequence from typing import Any, ClassVar, Final, Generic, Literal, TypedDict @@ -24,7 +25,6 @@ from agent_framework import ( ResponseStream, TextSpanRegion, UsageDetails, - get_logger, ) from agent_framework._settings import SecretString, load_settings from agent_framework._types import _get_data_bytes_as_str # type: ignore @@ -68,7 +68,7 @@ __all__ = [ "ThinkingConfig", ] -logger = get_logger("agent_framework.anthropic") +logger = logging.getLogger("agent_framework.anthropic") ANTHROPIC_DEFAULT_MAX_TOKENS: Final[int] = 1024 BETA_FLAGS: Final[list[str]] = ["mcp-client-2025-04-04", "code-execution-2025-08-25"] diff --git a/python/packages/azure-ai-search/agent_framework_azure_ai_search/_context_provider.py b/python/packages/azure-ai-search/agent_framework_azure_ai_search/_context_provider.py index 9edf73fdc3..e6779bd7a2 100644 --- a/python/packages/azure-ai-search/agent_framework_azure_ai_search/_context_provider.py +++ b/python/packages/azure-ai-search/agent_framework_azure_ai_search/_context_provider.py @@ -8,12 +8,12 @@ This module provides ``AzureAISearchContextProvider``, built on the new from __future__ import annotations +import logging import sys from collections.abc import Awaitable, Callable from typing import TYPE_CHECKING, Any, ClassVar, Literal, TypedDict from agent_framework import AGENT_FRAMEWORK_USER_AGENT, Message -from agent_framework._logging import get_logger from agent_framework._sessions import AgentSession, BaseContextProvider, SessionContext from agent_framework._settings import SecretString, load_settings from agent_framework.exceptions import ServiceInitializationError @@ -103,7 +103,7 @@ try: except ImportError: _agentic_retrieval_available = False -logger = get_logger(__name__) +logger = logging.getLogger("agent_framework.azure_ai_search") _DEFAULT_AGENTIC_MESSAGE_HISTORY_COUNT = 10 diff --git a/python/packages/azure-ai/agent_framework_azure_ai/_chat_client.py b/python/packages/azure-ai/agent_framework_azure_ai/_chat_client.py index 22d77c76b8..b548885dae 100644 --- a/python/packages/azure-ai/agent_framework_azure_ai/_chat_client.py +++ b/python/packages/azure-ai/agent_framework_azure_ai/_chat_client.py @@ -4,6 +4,7 @@ from __future__ import annotations import ast import json +import logging import os import re import sys @@ -31,7 +32,6 @@ from agent_framework import ( Role, TextSpanRegion, UsageDetails, - get_logger, ) from agent_framework._settings import load_settings from agent_framework.exceptions import ServiceInitializationError, ServiceInvalidRequestError, ServiceResponseException @@ -102,7 +102,7 @@ else: from typing_extensions import Self, TypedDict # type: ignore # pragma: no cover -logger = get_logger("agent_framework.azure") +logger = logging.getLogger("agent_framework.azure") __all__ = ["AzureAIAgentClient", "AzureAIAgentOptions"] diff --git a/python/packages/azure-ai/agent_framework_azure_ai/_client.py b/python/packages/azure-ai/agent_framework_azure_ai/_client.py index 82338f3d8d..3b2f3c0e10 100644 --- a/python/packages/azure-ai/agent_framework_azure_ai/_client.py +++ b/python/packages/azure-ai/agent_framework_azure_ai/_client.py @@ -3,6 +3,7 @@ from __future__ import annotations import json +import logging import sys from collections.abc import Callable, Mapping, MutableMapping, Sequence from contextlib import suppress @@ -19,7 +20,6 @@ from agent_framework import ( FunctionTool, Message, MiddlewareTypes, - get_logger, ) from agent_framework._settings import load_settings from agent_framework.exceptions import ServiceInitializationError @@ -59,7 +59,7 @@ else: from typing_extensions import Self, TypedDict # type: ignore # pragma: no cover -logger = get_logger("agent_framework.azure") +logger = logging.getLogger("agent_framework.azure") class AzureAIProjectAgentOptions(OpenAIResponsesOptions, total=False): diff --git a/python/packages/azure-ai/agent_framework_azure_ai/_project_provider.py b/python/packages/azure-ai/agent_framework_azure_ai/_project_provider.py index f2faffdb99..1a71df1f79 100644 --- a/python/packages/azure-ai/agent_framework_azure_ai/_project_provider.py +++ b/python/packages/azure-ai/agent_framework_azure_ai/_project_provider.py @@ -2,6 +2,7 @@ from __future__ import annotations +import logging import sys from collections.abc import Callable, MutableMapping, Sequence from typing import Any, Generic @@ -12,7 +13,6 @@ from agent_framework import ( BaseContextProvider, FunctionTool, MiddlewareTypes, - get_logger, normalize_tools, ) from agent_framework._mcp import MCPTool @@ -43,7 +43,7 @@ else: from typing_extensions import Self, TypedDict # type: ignore # pragma: no cover -logger = get_logger("agent_framework.azure") +logger = logging.getLogger("agent_framework.azure") # Type variable for options - allows typed Agent[OptionsT] returns diff --git a/python/packages/azure-ai/agent_framework_azure_ai/_shared.py b/python/packages/azure-ai/agent_framework_azure_ai/_shared.py index 81c113a1e4..a709b3f66c 100644 --- a/python/packages/azure-ai/agent_framework_azure_ai/_shared.py +++ b/python/packages/azure-ai/agent_framework_azure_ai/_shared.py @@ -2,13 +2,13 @@ from __future__ import annotations +import logging import sys from collections.abc import Mapping, MutableMapping, Sequence from typing import Any, cast from agent_framework import ( FunctionTool, - get_logger, ) from agent_framework.exceptions import ServiceInvalidRequestError from azure.ai.agents.models import ( @@ -37,7 +37,7 @@ if sys.version_info >= (3, 11): else: from typing_extensions import TypedDict # type: ignore # pragma: no cover -logger = get_logger("agent_framework.azure") +logger = logging.getLogger("agent_framework.azure") class AzureAISettings(TypedDict, total=False): diff --git a/python/packages/azure-ai/tests/test_azure_ai_client.py b/python/packages/azure-ai/tests/test_azure_ai_client.py index 73b4d3394e..38839b32d0 100644 --- a/python/packages/azure-ai/tests/test_azure_ai_client.py +++ b/python/packages/azure-ai/tests/test_azure_ai_client.py @@ -134,7 +134,7 @@ def create_test_azure_ai_client( client._created_agent_tool_names = set() # type: ignore client._created_agent_structured_output_signature = None # type: ignore client.additional_properties = {} - client.middleware = None + client.chat_middleware = [] # Mock the OpenAI client attribute mock_openai_client = MagicMock() @@ -1546,7 +1546,12 @@ async def test_integration_web_search() -> None: async with temporary_chat_client(agent_name="af-int-test-web-search") as client: for streaming in [False, True]: content = { - "messages": "Who are the main characters of Kpop Demon Hunters? Do a web search to find the answer.", + "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": [client.get_web_search_tool()], @@ -1565,7 +1570,9 @@ async def test_integration_web_search() -> None: # Test that the client will use the web search tool with location content = { - "messages": "What is the current weather? Do not ask for my current location.", + "messages": [ + Message(role="user", text="What is the current weather? Do not ask for my current location.") + ], "options": { "tool_choice": "auto", "tools": [client.get_web_search_tool(user_location={"country": "US", "city": "Seattle"})], @@ -1584,7 +1591,7 @@ async def test_integration_agent_hosted_mcp_tool() -> None: """Integration test for MCP tool with Azure Response Agent using Microsoft Learn MCP.""" async with temporary_chat_client(agent_name="af-int-test-mcp") as client: response = await client.get_response( - "How to create an Azure storage account using az cli?", + messages=[Message(role="user", text="How to create an Azure storage account using az cli?")], options={ # this needs to be high enough to handle the full MCP tool response. "max_tokens": 5000, @@ -1608,7 +1615,7 @@ async def test_integration_agent_hosted_code_interpreter_tool(): """Test Azure Responses Client agent with code interpreter tool through AzureAIClient.""" async with temporary_chat_client(agent_name="af-int-test-code-interpreter") as client: response = await client.get_response( - "Calculate the sum of numbers from 1 to 10 using Python code.", + messages=[Message(role="user", text="Calculate the sum of numbers from 1 to 10 using Python code.")], options={ "tools": [client.get_code_interpreter_tool()], }, diff --git a/python/packages/azurefunctions/agent_framework_azurefunctions/_app.py b/python/packages/azurefunctions/agent_framework_azurefunctions/_app.py index 886e0d8588..1c77d22e4d 100644 --- a/python/packages/azurefunctions/agent_framework_azurefunctions/_app.py +++ b/python/packages/azurefunctions/agent_framework_azurefunctions/_app.py @@ -9,6 +9,7 @@ with Azure Durable Entities, enabling stateful and durable AI agent execution. from __future__ import annotations import json +import logging import re import uuid from collections.abc import Callable, Mapping @@ -18,7 +19,7 @@ from typing import TYPE_CHECKING, Any, TypeVar, cast import azure.durable_functions as df import azure.functions as func -from agent_framework import SupportsAgentRun, get_logger +from agent_framework import SupportsAgentRun from agent_framework_durabletask import ( DEFAULT_MAX_POLL_RETRIES, DEFAULT_POLL_INTERVAL_SECONDS, @@ -42,7 +43,7 @@ from ._entities import create_agent_entity from ._errors import IncomingRequestError from ._orchestration import AgentOrchestrationContextType, AgentTask, AzureFunctionsAgentExecutor -logger = get_logger("agent_framework.azurefunctions") +logger = logging.getLogger("agent_framework.azurefunctions") EntityHandler = Callable[[df.DurableEntityContext], None] HandlerT = TypeVar("HandlerT", bound=Callable[..., Any]) diff --git a/python/packages/azurefunctions/agent_framework_azurefunctions/_entities.py b/python/packages/azurefunctions/agent_framework_azurefunctions/_entities.py index 23ea1e0f5c..d734950979 100644 --- a/python/packages/azurefunctions/agent_framework_azurefunctions/_entities.py +++ b/python/packages/azurefunctions/agent_framework_azurefunctions/_entities.py @@ -10,18 +10,19 @@ allows for long-running agent conversations. from __future__ import annotations import asyncio +import logging from collections.abc import Callable from typing import Any, cast import azure.durable_functions as df -from agent_framework import SupportsAgentRun, get_logger +from agent_framework import SupportsAgentRun from agent_framework_durabletask import ( AgentEntity, AgentEntityStateProviderMixin, AgentResponseCallbackProtocol, ) -logger = get_logger("agent_framework.azurefunctions.entities") +logger = logging.getLogger("agent_framework.azurefunctions") class AzureFunctionEntityStateProvider(AgentEntityStateProviderMixin): diff --git a/python/packages/azurefunctions/agent_framework_azurefunctions/_orchestration.py b/python/packages/azurefunctions/agent_framework_azurefunctions/_orchestration.py index 5875f9119b..be6c2d015e 100644 --- a/python/packages/azurefunctions/agent_framework_azurefunctions/_orchestration.py +++ b/python/packages/azurefunctions/agent_framework_azurefunctions/_orchestration.py @@ -5,11 +5,12 @@ This module provides support for using agents inside Durable Function orchestrations. """ +import logging from collections.abc import Callable from typing import TYPE_CHECKING, Any, TypeAlias import azure.durable_functions as df -from agent_framework import AgentSession, get_logger +from agent_framework import AgentSession from agent_framework_durabletask import ( DurableAgentExecutor, RunRequest, @@ -21,7 +22,7 @@ from azure.durable_functions.models.actions.NoOpAction import NoOpAction from azure.durable_functions.models.Task import CompoundTask, TaskState from pydantic import BaseModel -logger = get_logger("agent_framework.azurefunctions.orchestration") +logger = logging.getLogger("agent_framework.azurefunctions") CompoundActionConstructor: TypeAlias = Callable[[list[Any]], Any] | None diff --git a/python/packages/bedrock/agent_framework_bedrock/_chat_client.py b/python/packages/bedrock/agent_framework_bedrock/_chat_client.py index 3520d7b1a1..73bf9cc428 100644 --- a/python/packages/bedrock/agent_framework_bedrock/_chat_client.py +++ b/python/packages/bedrock/agent_framework_bedrock/_chat_client.py @@ -4,6 +4,7 @@ from __future__ import annotations import asyncio import json +import logging import sys from collections import deque from collections.abc import AsyncIterable, Awaitable, Mapping, MutableMapping, Sequence @@ -26,7 +27,6 @@ from agent_framework import ( Message, ResponseStream, UsageDetails, - get_logger, validate_tool_mode, ) from agent_framework._settings import SecretString, load_settings @@ -50,7 +50,7 @@ if sys.version_info >= (3, 11): else: from typing_extensions import TypedDict # type: ignore # pragma: no cover -logger = get_logger("agent_framework.bedrock") +logger = logging.getLogger("agent_framework.bedrock") __all__ = [ diff --git a/python/packages/claude/agent_framework_claude/_agent.py b/python/packages/claude/agent_framework_claude/_agent.py index ad6f1b3e03..6ad805c5d5 100644 --- a/python/packages/claude/agent_framework_claude/_agent.py +++ b/python/packages/claude/agent_framework_claude/_agent.py @@ -3,6 +3,7 @@ from __future__ import annotations import contextlib +import logging import sys from collections.abc import AsyncIterable, Awaitable, Callable, MutableMapping, Sequence from pathlib import Path @@ -19,11 +20,10 @@ from agent_framework import ( FunctionTool, Message, ResponseStream, - get_logger, normalize_messages, ) from agent_framework._settings import load_settings -from agent_framework._types import normalize_tools +from agent_framework._types import AgentRunInputs, normalize_tools from agent_framework.exceptions import ServiceException from claude_agent_sdk import ( AssistantMessage, @@ -61,7 +61,7 @@ if TYPE_CHECKING: __all__ = ["ClaudeAgent", "ClaudeAgentOptions"] -logger = get_logger("agent_framework.claude") +logger = logging.getLogger("agent_framework.claude") # Name of the in-process MCP server that hosts Agent Framework tools. # FunctionTool instances are converted to SDK MCP tools and served @@ -557,7 +557,7 @@ class ClaudeAgent(BaseAgent, Generic[OptionsT]): @overload def run( self, - messages: str | Message | Sequence[str | Message] | None = None, + messages: AgentRunInputs | None = None, *, stream: Literal[True], session: AgentSession | None = None, @@ -568,7 +568,7 @@ class ClaudeAgent(BaseAgent, Generic[OptionsT]): @overload async def run( self, - messages: str | Message | Sequence[str | Message] | None = None, + messages: AgentRunInputs | None = None, *, stream: Literal[False] = ..., session: AgentSession | None = None, @@ -578,7 +578,7 @@ class ClaudeAgent(BaseAgent, Generic[OptionsT]): def run( self, - messages: str | Message | Sequence[str | Message] | None = None, + messages: AgentRunInputs | None = None, *, stream: bool = False, session: AgentSession | None = None, @@ -612,7 +612,7 @@ class ClaudeAgent(BaseAgent, Generic[OptionsT]): async def _get_stream( self, - messages: str | Message | Sequence[str | Message] | None = None, + messages: AgentRunInputs | None = None, *, session: AgentSession | None = None, options: OptionsT | MutableMapping[str, Any] | None = None, diff --git a/python/packages/copilotstudio/agent_framework_copilotstudio/_agent.py b/python/packages/copilotstudio/agent_framework_copilotstudio/_agent.py index a3729d325d..fb832e4468 100644 --- a/python/packages/copilotstudio/agent_framework_copilotstudio/_agent.py +++ b/python/packages/copilotstudio/agent_framework_copilotstudio/_agent.py @@ -18,6 +18,7 @@ from agent_framework import ( normalize_messages, ) from agent_framework._settings import load_settings +from agent_framework._types import AgentRunInputs from agent_framework.exceptions import ServiceException, ServiceInitializationError from microsoft_agents.copilotstudio.client import AgentType, ConnectionSettings, CopilotClient, PowerPlatformCloud @@ -187,7 +188,7 @@ class CopilotStudioAgent(BaseAgent): @overload def run( self, - messages: str | Message | list[str] | list[Message] | None = None, + messages: AgentRunInputs | None = None, *, stream: Literal[False] = False, session: AgentSession | None = None, @@ -197,7 +198,7 @@ class CopilotStudioAgent(BaseAgent): @overload def run( self, - messages: str | Message | list[str] | list[Message] | None = None, + messages: AgentRunInputs | None = None, *, stream: Literal[True], session: AgentSession | None = None, @@ -206,7 +207,7 @@ class CopilotStudioAgent(BaseAgent): def run( self, - messages: str | Message | list[str] | list[Message] | None = None, + messages: AgentRunInputs | None = None, *, stream: bool = False, session: AgentSession | None = None, @@ -236,7 +237,7 @@ class CopilotStudioAgent(BaseAgent): async def _run_impl( self, - messages: str | Message | list[str] | list[Message] | None = None, + messages: AgentRunInputs | None = None, *, session: AgentSession | None = None, **kwargs: Any, @@ -261,7 +262,7 @@ class CopilotStudioAgent(BaseAgent): def _run_stream_impl( self, - messages: str | Message | list[str] | list[Message] | None = None, + messages: AgentRunInputs | None = None, *, session: AgentSession | None = None, **kwargs: Any, diff --git a/python/packages/core/agent_framework/__init__.py b/python/packages/core/agent_framework/__init__.py index f433df94b8..bc53b10d16 100644 --- a/python/packages/core/agent_framework/__init__.py +++ b/python/packages/core/agent_framework/__init__.py @@ -19,7 +19,6 @@ from ._clients import ( SupportsMCPTool, SupportsWebSearchTool, ) -from ._logging import get_logger, setup_logging from ._mcp import MCPStdioTool, MCPStreamableHTTPTool, MCPWebsocketTool from ._middleware import ( AgentContext, @@ -34,7 +33,6 @@ from ._middleware import ( FunctionInvocationContext, FunctionMiddleware, FunctionMiddlewareTypes, - MiddlewareException, MiddlewareTermination, MiddlewareType, MiddlewareTypes, @@ -67,6 +65,7 @@ from ._tools import ( from ._types import ( AgentResponse, AgentResponseUpdate, + AgentRunInputs, Annotation, ChatOptions, ChatResponse, @@ -152,6 +151,7 @@ from ._workflows import ( response_handler, validate_workflow_graph, ) +from .exceptions import MiddlewareException __all__ = [ "AGENT_FRAMEWORK_USER_AGENT", @@ -169,6 +169,7 @@ __all__ = [ "AgentMiddlewareTypes", "AgentResponse", "AgentResponseUpdate", + "AgentRunInputs", "AgentSession", "Annotation", "BaseAgent", @@ -264,6 +265,7 @@ __all__ = [ "WorkflowRunnerException", "WorkflowValidationError", "WorkflowViz", + "__version__", "add_usage_details", "agent_middleware", "chat_middleware", @@ -271,7 +273,6 @@ __all__ = [ "detect_media_type_from_base64", "executor", "function_middleware", - "get_logger", "handler", "map_chat_to_agent_update", "merge_chat_options", @@ -283,7 +284,6 @@ __all__ = [ "register_state_type", "resolve_agent_id", "response_handler", - "setup_logging", "tool", "validate_chat_options", "validate_tool_mode", diff --git a/python/packages/core/agent_framework/_agents.py b/python/packages/core/agent_framework/_agents.py index d11f0e2c7d..0b1b20bed8 100644 --- a/python/packages/core/agent_framework/_agents.py +++ b/python/packages/core/agent_framework/_agents.py @@ -3,6 +3,7 @@ from __future__ import annotations import inspect +import logging import re import sys from collections.abc import Awaitable, Callable, Mapping, MutableMapping, Sequence @@ -29,7 +30,6 @@ from mcp.shared.exceptions import McpError from pydantic import BaseModel, Field, create_model from ._clients import BaseChatClient, SupportsChatGetResponse -from ._logging import get_logger from ._mcp import LOG_LEVEL_MAPPING, MCPTool from ._middleware import AgentMiddlewareLayer, MiddlewareTypes from ._serialization import SerializationMixin @@ -41,6 +41,7 @@ from ._tools import ( from ._types import ( AgentResponse, AgentResponseUpdate, + AgentRunInputs, ChatResponse, ChatResponseUpdate, Message, @@ -67,7 +68,7 @@ else: if TYPE_CHECKING: from ._types import ChatOptions -logger = get_logger("agent_framework") +logger = logging.getLogger("agent_framework") ResponseModelBoundT = TypeVar("ResponseModelBoundT", bound=BaseModel) OptionsCoT = TypeVar( @@ -159,9 +160,6 @@ class _RunContext(TypedDict): finalize_kwargs: dict[str, Any] -__all__ = ["Agent", "BaseAgent", "RawAgent", "SupportsAgentRun"] - - # region Agent Protocol @@ -230,7 +228,7 @@ class SupportsAgentRun(Protocol): @overload def run( self, - messages: str | Message | Sequence[str | Message] | None = None, + messages: AgentRunInputs | None = None, *, stream: Literal[False] = ..., session: AgentSession | None = None, @@ -242,7 +240,7 @@ class SupportsAgentRun(Protocol): @overload def run( self, - messages: str | Message | Sequence[str | Message] | None = None, + messages: AgentRunInputs | None = None, *, stream: Literal[True], session: AgentSession | None = None, @@ -253,7 +251,7 @@ class SupportsAgentRun(Protocol): def run( self, - messages: str | Message | Sequence[str | Message] | None = None, + messages: AgentRunInputs | None = None, *, stream: bool = False, session: AgentSession | None = None, @@ -763,7 +761,7 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc] @overload def run( self, - messages: str | Message | Sequence[str | Message] | None = None, + messages: AgentRunInputs | None = None, *, stream: Literal[False] = ..., session: AgentSession | None = None, @@ -780,7 +778,7 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc] @overload def run( self, - messages: str | Message | Sequence[str | Message] | None = None, + messages: AgentRunInputs | None = None, *, stream: Literal[False] = ..., session: AgentSession | None = None, @@ -797,7 +795,7 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc] @overload def run( self, - messages: str | Message | Sequence[str | Message] | None = None, + messages: AgentRunInputs | None = None, *, stream: Literal[True], session: AgentSession | None = None, @@ -813,7 +811,7 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc] def run( self, - messages: str | Message | Sequence[str | Message] | None = None, + messages: AgentRunInputs | None = None, *, stream: bool = False, session: AgentSession | None = None, @@ -1000,7 +998,7 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc] async def _prepare_run_context( self, *, - messages: str | Message | Sequence[str | Message] | None, + messages: AgentRunInputs | None, session: AgentSession | None, tools: FunctionTool | Callable[..., Any] diff --git a/python/packages/core/agent_framework/_clients.py b/python/packages/core/agent_framework/_clients.py index 57daed6286..c8956a67e3 100644 --- a/python/packages/core/agent_framework/_clients.py +++ b/python/packages/core/agent_framework/_clients.py @@ -2,6 +2,7 @@ from __future__ import annotations +import logging import sys from abc import ABC, abstractmethod from collections.abc import ( @@ -27,7 +28,6 @@ from typing import ( from pydantic import BaseModel -from ._logging import get_logger from ._serialization import SerializationMixin from ._tools import ( FunctionInvocationConfiguration, @@ -38,7 +38,6 @@ from ._types import ( ChatResponseUpdate, Message, ResponseStream, - prepare_messages, validate_chat_options, ) @@ -61,17 +60,7 @@ InputT = TypeVar("InputT", contravariant=True) EmbeddingT = TypeVar("EmbeddingT") BaseChatClientT = TypeVar("BaseChatClientT", bound="BaseChatClient") -logger = get_logger() - -__all__ = [ - "BaseChatClient", - "SupportsChatGetResponse", - "SupportsCodeInterpreterTool", - "SupportsFileSearchTool", - "SupportsImageGenerationTool", - "SupportsMCPTool", - "SupportsWebSearchTool", -] +logger = logging.getLogger("agent_framework") # region SupportsChatGetResponse Protocol @@ -139,7 +128,7 @@ class SupportsChatGetResponse(Protocol[OptionsContraT]): @overload def get_response( self, - messages: str | Message | Sequence[str | Message], + messages: Sequence[Message], *, stream: Literal[False] = ..., options: ChatOptions[ResponseModelBoundT], @@ -149,7 +138,7 @@ class SupportsChatGetResponse(Protocol[OptionsContraT]): @overload def get_response( self, - messages: str | Message | Sequence[str | Message], + messages: Sequence[Message], *, stream: Literal[False] = ..., options: OptionsContraT | ChatOptions[None] | None = None, @@ -159,7 +148,7 @@ class SupportsChatGetResponse(Protocol[OptionsContraT]): @overload def get_response( self, - messages: str | Message | Sequence[str | Message], + messages: Sequence[Message], *, stream: Literal[True], options: OptionsContraT | ChatOptions[Any] | None = None, @@ -168,7 +157,7 @@ class SupportsChatGetResponse(Protocol[OptionsContraT]): def get_response( self, - messages: str | Message | Sequence[str | Message], + messages: Sequence[Message], *, stream: bool = False, options: OptionsContraT | ChatOptions[Any] | None = None, @@ -254,9 +243,9 @@ class BaseChatClient(SerializationMixin, ABC, Generic[OptionsCoT]): client = CustomChatClient() # Use the client to get responses - response = await client.get_response("Hello, how are you?") + response = await client.get_response([Message(role="user", text="Hello, how are you?")]) # Or stream responses - async for update in client.get_response("Hello!", stream=True): + async for update in client.get_response([Message(role="user", text="Hello!")], stream=True): print(update) """ @@ -376,7 +365,7 @@ class BaseChatClient(SerializationMixin, ABC, Generic[OptionsCoT]): @overload def get_response( self, - messages: str | Message | Sequence[str | Message], + messages: Sequence[Message], *, stream: Literal[False] = ..., options: ChatOptions[ResponseModelBoundT], @@ -386,7 +375,7 @@ class BaseChatClient(SerializationMixin, ABC, Generic[OptionsCoT]): @overload def get_response( self, - messages: str | Message | Sequence[str | Message], + messages: Sequence[Message], *, stream: Literal[False] = ..., options: OptionsCoT | ChatOptions[None] | None = None, @@ -396,7 +385,7 @@ class BaseChatClient(SerializationMixin, ABC, Generic[OptionsCoT]): @overload def get_response( self, - messages: str | Message | Sequence[str | Message], + messages: Sequence[Message], *, stream: Literal[True], options: OptionsCoT | ChatOptions[Any] | None = None, @@ -405,7 +394,7 @@ class BaseChatClient(SerializationMixin, ABC, Generic[OptionsCoT]): def get_response( self, - messages: str | Message | Sequence[str | Message], + messages: Sequence[Message], *, stream: bool = False, options: OptionsCoT | ChatOptions[Any] | None = None, @@ -422,9 +411,8 @@ class BaseChatClient(SerializationMixin, ABC, Generic[OptionsCoT]): Returns: When streaming a response stream of ChatResponseUpdates, otherwise an Awaitable ChatResponse. """ - prepared_messages = prepare_messages(messages) return self._inner_get_response( - messages=prepared_messages, + messages=messages, stream=stream, options=options or {}, # type: ignore[arg-type] **kwargs, diff --git a/python/packages/core/agent_framework/_logging.py b/python/packages/core/agent_framework/_logging.py deleted file mode 100644 index 012de28bf1..0000000000 --- a/python/packages/core/agent_framework/_logging.py +++ /dev/null @@ -1,29 +0,0 @@ -# Copyright (c) Microsoft. All rights reserved. - -import logging - -from .exceptions import AgentFrameworkException - -__all__ = ["get_logger", "setup_logging"] - - -def setup_logging() -> None: - """Setup the logging configuration for the agent framework.""" - logging.basicConfig( - format="[%(asctime)s - %(pathname)s:%(lineno)d - %(levelname)s] %(message)s", - datefmt="%Y-%m-%d %H:%M:%S", - ) - - -def get_logger(name: str = "agent_framework") -> logging.Logger: - """Get a logger with the specified name, defaulting to 'agent_framework'. - - Args: - name (str): The name of the logger. Defaults to 'agent_framework'. - - Returns: - logging.Logger: The configured logger instance. - """ - if not name.startswith("agent_framework"): - raise AgentFrameworkException("Logger name must start with 'agent_framework'.") - return logging.getLogger(name) diff --git a/python/packages/core/agent_framework/_mcp.py b/python/packages/core/agent_framework/_mcp.py index f9d2cd9971..2be5e452ce 100644 --- a/python/packages/core/agent_framework/_mcp.py +++ b/python/packages/core/agent_framework/_mcp.py @@ -72,12 +72,6 @@ LOG_LEVEL_MAPPING: dict[types.LoggingLevel, int] = { "emergency": logging.CRITICAL, } -__all__ = [ - "MCPStdioTool", - "MCPStreamableHTTPTool", - "MCPWebsocketTool", -] - def _parse_prompt_result_from_mcp( mcp_type: types.GetPromptResult, diff --git a/python/packages/core/agent_framework/_middleware.py b/python/packages/core/agent_framework/_middleware.py index 299da150ab..1f0f9e3338 100644 --- a/python/packages/core/agent_framework/_middleware.py +++ b/python/packages/core/agent_framework/_middleware.py @@ -14,11 +14,12 @@ from ._clients import SupportsChatGetResponse from ._types import ( AgentResponse, AgentResponseUpdate, + AgentRunInputs, ChatResponse, ChatResponseUpdate, Message, ResponseStream, - prepare_messages, + normalize_messages, ) from .exceptions import MiddlewareException @@ -42,27 +43,6 @@ if TYPE_CHECKING: ResponseModelBoundT = TypeVar("ResponseModelBoundT", bound=BaseModel) -__all__ = [ - "AgentContext", - "AgentMiddleware", - "AgentMiddlewareLayer", - "AgentMiddlewareTypes", - "ChatAndFunctionMiddlewareTypes", - "ChatContext", - "ChatMiddleware", - "ChatMiddlewareLayer", - "ChatMiddlewareTypes", - "FunctionInvocationContext", - "FunctionMiddleware", - "FunctionMiddlewareTypes", - "MiddlewareException", - "MiddlewareTermination", - "MiddlewareType", - "MiddlewareTypes", - "agent_middleware", - "chat_middleware", - "function_middleware", -] AgentT = TypeVar("AgentT", bound="SupportsAgentRun") ContextT = TypeVar("ContextT") @@ -978,7 +958,7 @@ class ChatMiddlewareLayer(Generic[OptionsCoT]): @overload def get_response( self, - messages: str | Message | Sequence[str | Message], + messages: Sequence[Message], *, stream: Literal[False] = ..., options: ChatOptions[ResponseModelBoundT], @@ -988,7 +968,7 @@ class ChatMiddlewareLayer(Generic[OptionsCoT]): @overload def get_response( self, - messages: str | Message | Sequence[str | Message], + messages: Sequence[Message], *, stream: Literal[False] = ..., options: OptionsCoT | ChatOptions[None] | None = None, @@ -998,7 +978,7 @@ class ChatMiddlewareLayer(Generic[OptionsCoT]): @overload def get_response( self, - messages: str | Message | Sequence[str | Message], + messages: Sequence[Message], *, stream: Literal[True], options: OptionsCoT | ChatOptions[Any] | None = None, @@ -1007,7 +987,7 @@ class ChatMiddlewareLayer(Generic[OptionsCoT]): def get_response( self, - messages: str | Message | Sequence[str | Message], + messages: Sequence[Message], *, stream: bool = False, options: OptionsCoT | ChatOptions[Any] | None = None, @@ -1034,7 +1014,7 @@ class ChatMiddlewareLayer(Generic[OptionsCoT]): context = ChatContext( client=self, # type: ignore[arg-type] - messages=prepare_messages(messages), + messages=list(messages), options=options, stream=stream, kwargs=kwargs, @@ -1095,7 +1075,7 @@ class AgentMiddlewareLayer: @overload def run( self, - messages: str | Message | Sequence[str | Message] | None = None, + messages: AgentRunInputs | None = None, *, stream: Literal[False] = ..., session: AgentSession | None = None, @@ -1107,7 +1087,7 @@ class AgentMiddlewareLayer: @overload def run( self, - messages: str | Message | Sequence[str | Message] | None = None, + messages: AgentRunInputs | None = None, *, stream: Literal[False] = ..., session: AgentSession | None = None, @@ -1119,7 +1099,7 @@ class AgentMiddlewareLayer: @overload def run( self, - messages: str | Message | Sequence[str | Message] | None = None, + messages: AgentRunInputs | None = None, *, stream: Literal[True], session: AgentSession | None = None, @@ -1130,7 +1110,7 @@ class AgentMiddlewareLayer: def run( self, - messages: str | Message | Sequence[str | Message] | None = None, + messages: AgentRunInputs | None = None, *, stream: bool = False, session: AgentSession | None = None, @@ -1161,7 +1141,7 @@ class AgentMiddlewareLayer: context = AgentContext( agent=self, # type: ignore[arg-type] - messages=prepare_messages(messages), # type: ignore[arg-type] + messages=normalize_messages(messages), session=session, options=options, stream=stream, diff --git a/python/packages/core/agent_framework/_serialization.py b/python/packages/core/agent_framework/_serialization.py index 1259ca2a1a..7934477298 100644 --- a/python/packages/core/agent_framework/_serialization.py +++ b/python/packages/core/agent_framework/_serialization.py @@ -3,13 +3,12 @@ from __future__ import annotations import json +import logging import re from collections.abc import Mapping, MutableMapping from typing import Any, ClassVar, Protocol, TypeVar, runtime_checkable -from ._logging import get_logger - -logger = get_logger() +logger = logging.getLogger("agent_framework") ClassT = TypeVar("ClassT", bound="SerializationMixin") ProtocolT = TypeVar("ProtocolT", bound="SerializationProtocol") diff --git a/python/packages/core/agent_framework/_sessions.py b/python/packages/core/agent_framework/_sessions.py index 6240c632c2..2498bfdd69 100644 --- a/python/packages/core/agent_framework/_sessions.py +++ b/python/packages/core/agent_framework/_sessions.py @@ -24,16 +24,6 @@ if TYPE_CHECKING: from ._agents import SupportsAgentRun -__all__ = [ - "AgentSession", - "BaseContextProvider", - "BaseHistoryProvider", - "InMemoryHistoryProvider", - "SessionContext", - "register_state_type", -] - - # Registry of known types for state deserialization _STATE_TYPE_REGISTRY: dict[str, type] = {} diff --git a/python/packages/core/agent_framework/_settings.py b/python/packages/core/agent_framework/_settings.py index 30fd8c8508..ada0c146f0 100644 --- a/python/packages/core/agent_framework/_settings.py +++ b/python/packages/core/agent_framework/_settings.py @@ -45,7 +45,6 @@ if sys.version_info >= (3, 13): else: from typing_extensions import TypeVar # type: ignore # pragma: no cover -__all__ = ["SecretString", "load_settings"] SettingsT = TypeVar("SettingsT", default=dict[str, Any]) diff --git a/python/packages/core/agent_framework/_telemetry.py b/python/packages/core/agent_framework/_telemetry.py index 5e3ee57222..a044fc9d02 100644 --- a/python/packages/core/agent_framework/_telemetry.py +++ b/python/packages/core/agent_framework/_telemetry.py @@ -2,21 +2,14 @@ from __future__ import annotations +import logging import os from typing import Any, Final from . import __version__ as version_info -from ._logging import get_logger -logger = get_logger() +logger = logging.getLogger("agent_framework") -__all__ = [ - "AGENT_FRAMEWORK_USER_AGENT", - "APP_INFO", - "USER_AGENT_KEY", - "USER_AGENT_TELEMETRY_DISABLED_ENV_VAR", - "prepend_agent_framework_to_user_agent", -] # Note that if this environment variable does not exist, user agent telemetry is enabled. USER_AGENT_TELEMETRY_DISABLED_ENV_VAR = "AGENT_FRAMEWORK_USER_AGENT_DISABLED" diff --git a/python/packages/core/agent_framework/_tools.py b/python/packages/core/agent_framework/_tools.py index 4a4e30d324..8f941b41e0 100644 --- a/python/packages/core/agent_framework/_tools.py +++ b/python/packages/core/agent_framework/_tools.py @@ -5,6 +5,7 @@ from __future__ import annotations import asyncio import inspect import json +import logging import sys from collections.abc import ( AsyncIterable, @@ -34,7 +35,6 @@ from typing import ( from opentelemetry.metrics import Histogram, NoOpHistogram from pydantic import BaseModel, Field, ValidationError, create_model -from ._logging import get_logger from ._serialization import SerializationMixin from .exceptions import ToolException from .observability import ( @@ -71,18 +71,8 @@ if TYPE_CHECKING: ResponseModelBoundT = TypeVar("ResponseModelBoundT", bound=BaseModel) -logger = get_logger() +logger = logging.getLogger("agent_framework") -__all__ = [ - "FunctionInvocationConfiguration", - "FunctionInvocationLayer", - "FunctionTool", - "normalize_function_invocation_configuration", - "tool", -] - - -logger = get_logger() DEFAULT_MAX_ITERATIONS: Final[int] = 40 DEFAULT_MAX_CONSECUTIVE_ERRORS_PER_REQUEST: Final[int] = 3 ChatClientT = TypeVar("ChatClientT", bound="SupportsChatGetResponse[Any]") @@ -1941,7 +1931,7 @@ class FunctionInvocationLayer(Generic[OptionsCoT]): @overload def get_response( self, - messages: str | Message | Sequence[str | Message], + messages: Sequence[Message], *, stream: Literal[False] = ..., options: ChatOptions[ResponseModelBoundT], @@ -1951,7 +1941,7 @@ class FunctionInvocationLayer(Generic[OptionsCoT]): @overload def get_response( self, - messages: str | Message | Sequence[str | Message], + messages: Sequence[Message], *, stream: Literal[False] = ..., options: OptionsCoT | ChatOptions[None] | None = None, @@ -1961,7 +1951,7 @@ class FunctionInvocationLayer(Generic[OptionsCoT]): @overload def get_response( self, - messages: str | Message | Sequence[str | Message], + messages: Sequence[Message], *, stream: Literal[True], options: OptionsCoT | ChatOptions[Any] | None = None, @@ -1970,7 +1960,7 @@ class FunctionInvocationLayer(Generic[OptionsCoT]): def get_response( self, - messages: str | Message | Sequence[str | Message], + messages: Sequence[Message], *, stream: bool = False, options: OptionsCoT | ChatOptions[Any] | None = None, @@ -1982,7 +1972,6 @@ class FunctionInvocationLayer(Generic[OptionsCoT]): ChatResponse, ChatResponseUpdate, ResponseStream, - prepare_messages, ) super_get_response = super().get_response # type: ignore[misc] @@ -2014,7 +2003,7 @@ class FunctionInvocationLayer(Generic[OptionsCoT]): nonlocal mutable_options nonlocal filtered_kwargs errors_in_a_row: int = 0 - prepped_messages = prepare_messages(messages) + prepped_messages = list(messages) fcc_messages: list[Message] = [] response: ChatResponse | None = None @@ -2108,7 +2097,7 @@ class FunctionInvocationLayer(Generic[OptionsCoT]): nonlocal mutable_options nonlocal stream_result_hooks errors_in_a_row: int = 0 - prepped_messages = prepare_messages(messages) + prepped_messages = list(messages) fcc_messages: list[Message] = [] response: ChatResponse | None = None diff --git a/python/packages/core/agent_framework/_types.py b/python/packages/core/agent_framework/_types.py index 71a635f1cc..a3786e8838 100644 --- a/python/packages/core/agent_framework/_types.py +++ b/python/packages/core/agent_framework/_types.py @@ -4,6 +4,7 @@ from __future__ import annotations import base64 import json +import logging import re import sys from asyncio import iscoroutine @@ -13,7 +14,6 @@ from typing import TYPE_CHECKING, Any, ClassVar, Final, Generic, Literal, NewTyp from pydantic import BaseModel -from ._logging import get_logger from ._serialization import SerializationMixin from ._tools import FunctionTool, tool from .exceptions import AdditionItemMismatch, ContentError @@ -27,41 +27,7 @@ if sys.version_info >= (3, 11): else: from typing_extensions import TypedDict # type: ignore # pragma: no cover -__all__ = [ - "AgentResponse", - "AgentResponseUpdate", - "Annotation", - "ChatOptions", - "ChatResponse", - "ChatResponseUpdate", - "Content", - "ContinuationToken", - "FinalT", - "FinishReason", - "FinishReasonLiteral", - "Message", - "OuterFinalT", - "OuterUpdateT", - "ResponseStream", - "Role", - "RoleLiteral", - "TextSpanRegion", - "ToolMode", - "UpdateT", - "UsageDetails", - "add_usage_details", - "detect_media_type_from_base64", - "map_chat_to_agent_update", - "merge_chat_options", - "normalize_messages", - "normalize_tools", - "prepend_instructions_to_messages", - "validate_chat_options", - "validate_tool_mode", - "validate_tools", -] - -logger = get_logger("agent_framework") +logger = logging.getLogger("agent_framework") # region Content Parsing Utilities @@ -1536,47 +1502,11 @@ class Message(SerializationMixin): return " ".join(content.text for content in self.contents if content.type == "text") # type: ignore[misc] -def prepare_messages( - messages: str | Content | Message | Sequence[str | Content | Message], - system_instructions: str | Sequence[str] | None = None, -) -> list[Message]: - """Convert various message input formats into a list of Message objects. - - Args: - messages: The input messages in various supported formats. Can be: - - A string (converted to a user message) - - A Content object (wrapped in a user Message) - - A Message object - - A sequence containing any mix of the above - system_instructions: The system instructions. They will be inserted to the start of the messages list. - - Returns: - A list of Message objects. - """ - if system_instructions is not None: - if isinstance(system_instructions, str): - system_instructions = [system_instructions] - system_instruction_messages = [Message("system", [instr]) for instr in system_instructions] - else: - system_instruction_messages = [] - - if isinstance(messages, str): - return [*system_instruction_messages, Message("user", [messages])] - if isinstance(messages, Content): - return [*system_instruction_messages, Message("user", [messages])] - if isinstance(messages, Message): - return [*system_instruction_messages, messages] - - return_messages: list[Message] = system_instruction_messages - for msg in messages: - if isinstance(msg, (str, Content)): - msg = Message("user", [msg]) - return_messages.append(msg) - return return_messages +AgentRunInputs = str | Content | Message | Sequence[str | Content | Message] def normalize_messages( - messages: str | Content | Message | Sequence[str | Content | Message] | None = None, + messages: AgentRunInputs | None = None, ) -> list[Message]: """Normalize message inputs to a list of Message objects. diff --git a/python/packages/core/agent_framework/_workflows/_agent.py b/python/packages/core/agent_framework/_workflows/_agent.py index e4152b77d1..765da7f15c 100644 --- a/python/packages/core/agent_framework/_workflows/_agent.py +++ b/python/packages/core/agent_framework/_workflows/_agent.py @@ -22,6 +22,7 @@ from .._sessions import ( from .._types import ( AgentResponse, AgentResponseUpdate, + AgentRunInputs, Content, Message, ResponseStream, @@ -145,7 +146,7 @@ class WorkflowAgent(BaseAgent): @overload def run( self, - messages: str | Message | Sequence[str | Message] | None = None, + messages: AgentRunInputs | None = None, *, stream: Literal[True], session: AgentSession | None = None, @@ -157,7 +158,7 @@ class WorkflowAgent(BaseAgent): @overload async def run( self, - messages: str | Message | Sequence[str | Message] | None = None, + messages: AgentRunInputs | None = None, *, stream: Literal[False] = ..., session: AgentSession | None = None, @@ -168,7 +169,7 @@ class WorkflowAgent(BaseAgent): def run( self, - messages: str | Message | Sequence[str | Message] | None = None, + messages: AgentRunInputs | None = None, *, stream: bool = False, session: AgentSession | None = None, @@ -214,7 +215,7 @@ class WorkflowAgent(BaseAgent): async def _run_impl( self, - messages: str | Message | Sequence[str | Message], + messages: AgentRunInputs, response_id: str, session: AgentSession | None, checkpoint_id: str | None = None, @@ -270,7 +271,7 @@ class WorkflowAgent(BaseAgent): async def _run_stream_impl( self, - messages: str | Message | Sequence[str | Message], + messages: AgentRunInputs, response_id: str, session: AgentSession | None, checkpoint_id: str | None = None, diff --git a/python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py b/python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py index 524f291c5e..310451229e 100644 --- a/python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py +++ b/python/packages/core/agent_framework/_workflows/_checkpoint_encoding.py @@ -3,11 +3,10 @@ from __future__ import annotations import base64 +import logging import pickle # nosec # noqa: S403 from typing import Any -from agent_framework import get_logger - """Checkpoint encoding using JSON structure with pickle+base64 for arbitrary data. This hybrid approach provides: @@ -20,7 +19,7 @@ from trusted sources. Loading a malicious checkpoint file can execute arbitrary """ -logger = get_logger(__name__) +logger = logging.getLogger("agent_framework") # Marker to identify pickled values in serialized JSON _PICKLE_MARKER = "__pickled__" diff --git a/python/packages/core/agent_framework/_workflows/_message_utils.py b/python/packages/core/agent_framework/_workflows/_message_utils.py index 6d27a905ee..ba69fa340d 100644 --- a/python/packages/core/agent_framework/_workflows/_message_utils.py +++ b/python/packages/core/agent_framework/_workflows/_message_utils.py @@ -2,18 +2,17 @@ """Shared helpers for normalizing workflow message inputs.""" -from collections.abc import Sequence - -from agent_framework import Message +from agent_framework import Content, Message +from agent_framework._types import AgentRunInputs def normalize_messages_input( - messages: str | Message | Sequence[str | Message] | None = None, + messages: AgentRunInputs | None = None, ) -> list[Message]: """Normalize heterogeneous message inputs to a list of Message objects. Args: - messages: String, Message, or sequence of either. None yields empty list. + messages: String, Content, Message, or sequence of those values. None yields empty list. Returns: List of Message instances suitable for workflow consumption. @@ -24,6 +23,9 @@ def normalize_messages_input( if isinstance(messages, str): return [Message(role="user", text=messages)] + if isinstance(messages, Content): + return [Message(role="user", contents=[messages])] + if isinstance(messages, Message): return [messages] @@ -31,13 +33,12 @@ def normalize_messages_input( for item in messages: if isinstance(item, str): normalized.append(Message(role="user", text=item)) + elif isinstance(item, Content): + normalized.append(Message(role="user", contents=[item])) elif isinstance(item, Message): normalized.append(item) else: raise TypeError( - f"Messages sequence must contain only str or Message instances; found {type(item).__name__}." + f"Messages sequence must contain only str, Content, or Message instances; found {type(item).__name__}." ) return normalized - - -__all__ = ["normalize_messages_input"] diff --git a/python/packages/core/agent_framework/azure/_assistants_client.py b/python/packages/core/agent_framework/azure/_assistants_client.py index 89399ed833..ee3e05c4ed 100644 --- a/python/packages/core/agent_framework/azure/_assistants_client.py +++ b/python/packages/core/agent_framework/azure/_assistants_client.py @@ -27,8 +27,6 @@ if sys.version_info >= (3, 11): else: from typing_extensions import TypedDict # type: ignore # pragma: no cover -__all__ = ["AzureOpenAIAssistantsClient"] - # region Azure OpenAI Assistants Options TypedDict diff --git a/python/packages/core/agent_framework/azure/_chat_client.py b/python/packages/core/agent_framework/azure/_chat_client.py index 485df4a6ed..b195969563 100644 --- a/python/packages/core/agent_framework/azure/_chat_client.py +++ b/python/packages/core/agent_framework/azure/_chat_client.py @@ -53,7 +53,6 @@ if TYPE_CHECKING: logger: logging.Logger = logging.getLogger(__name__) -__all__ = ["AzureOpenAIChatClient", "AzureOpenAIChatOptions", "AzureUserSecurityContext"] ResponseModelT = TypeVar("ResponseModelT", bound=BaseModel | None, default=None) diff --git a/python/packages/core/agent_framework/azure/_responses_client.py b/python/packages/core/agent_framework/azure/_responses_client.py index 0a6c0cd8c8..cf6be12af5 100644 --- a/python/packages/core/agent_framework/azure/_responses_client.py +++ b/python/packages/core/agent_framework/azure/_responses_client.py @@ -42,8 +42,6 @@ if TYPE_CHECKING: from .._middleware import MiddlewareTypes from ..openai._responses_client import OpenAIResponsesOptions -__all__ = ["AzureOpenAIResponsesClient"] - AzureOpenAIResponsesOptionsT = TypeVar( "AzureOpenAIResponsesOptionsT", diff --git a/python/packages/core/agent_framework/observability.py b/python/packages/core/agent_framework/observability.py index d19683d9e2..7973a781a4 100644 --- a/python/packages/core/agent_framework/observability.py +++ b/python/packages/core/agent_framework/observability.py @@ -20,7 +20,6 @@ from opentelemetry.semconv.attributes import service_attributes from opentelemetry.semconv_ai import Meters, SpanAttributes from . import __version__ as version_info -from ._logging import get_logger from ._settings import load_settings if sys.version_info >= (3, 13): @@ -44,6 +43,7 @@ if TYPE_CHECKING: # pragma: no cover from ._types import ( AgentResponse, AgentResponseUpdate, + AgentRunInputs, ChatOptions, ChatResponse, ChatResponseUpdate, @@ -73,7 +73,7 @@ AgentT = TypeVar("AgentT", bound="SupportsAgentRun") ChatClientT = TypeVar("ChatClientT", bound="SupportsChatGetResponse[Any]") -logger = get_logger() +logger = logging.getLogger("agent_framework") OTEL_METRICS: Final[str] = "__otel_metrics__" @@ -747,7 +747,6 @@ class ObservabilitySettings: for log_exporter in log_exporters: logger_provider.add_log_record_processor(BatchLogRecordProcessor(log_exporter)) # Attach a handler with the provider to the root logger - logger = logging.getLogger() handler = LoggingHandler(logger_provider=logger_provider) logger.addHandler(handler) set_logger_provider(logger_provider) @@ -1084,7 +1083,7 @@ class ChatTelemetryLayer(Generic[OptionsCoT]): @overload def get_response( self, - messages: str | Message | Sequence[str | Message], + messages: Sequence[Message], *, stream: Literal[False] = ..., options: ChatOptions[ResponseModelBoundT], @@ -1094,7 +1093,7 @@ class ChatTelemetryLayer(Generic[OptionsCoT]): @overload def get_response( self, - messages: str | Message | Sequence[str | Message], + messages: Sequence[Message], *, stream: Literal[False] = ..., options: OptionsCoT | ChatOptions[None] | None = None, @@ -1104,7 +1103,7 @@ class ChatTelemetryLayer(Generic[OptionsCoT]): @overload def get_response( self, - messages: str | Message | Sequence[str | Message], + messages: Sequence[Message], *, stream: Literal[True], options: OptionsCoT | ChatOptions[Any] | None = None, @@ -1113,7 +1112,7 @@ class ChatTelemetryLayer(Generic[OptionsCoT]): def get_response( self, - messages: str | Message | Sequence[str | Message], + messages: Sequence[Message], *, stream: bool = False, options: OptionsCoT | ChatOptions[Any] | None = None, @@ -1277,7 +1276,7 @@ class AgentTelemetryLayer: @overload def run( self, - messages: str | Message | Sequence[str | Message] | None = None, + messages: AgentRunInputs | None = None, *, stream: Literal[False] = ..., session: AgentSession | None = None, @@ -1287,7 +1286,7 @@ class AgentTelemetryLayer: @overload def run( self, - messages: str | Message | Sequence[str | Message] | None = None, + messages: AgentRunInputs | None = None, *, stream: Literal[True], session: AgentSession | None = None, @@ -1296,7 +1295,7 @@ class AgentTelemetryLayer: def run( self, - messages: str | Message | Sequence[str | Message] | None = None, + messages: AgentRunInputs | None = None, *, stream: bool = False, session: AgentSession | None = None, @@ -1614,15 +1613,15 @@ def capture_exception(span: trace.Span, exception: Exception, timestamp: int | N def _capture_messages( span: trace.Span, provider_name: str, - messages: str | Message | Sequence[str | Message], + messages: AgentRunInputs, system_instructions: str | list[str] | None = None, output: bool = False, finish_reason: FinishReason | None = None, ) -> None: """Log messages with extra information.""" - from ._types import prepare_messages + from ._types import normalize_messages, prepend_instructions_to_messages - prepped = prepare_messages(messages, system_instructions=system_instructions) + prepped = prepend_instructions_to_messages(normalize_messages(messages), system_instructions) otel_messages: list[dict[str, Any]] = [] for index, message in enumerate(prepped): # Reuse the otel message representation for logging instead of calling to_dict() diff --git a/python/packages/core/agent_framework/openai/_assistant_provider.py b/python/packages/core/agent_framework/openai/_assistant_provider.py index 8082a4ad9b..6b6c2ef687 100644 --- a/python/packages/core/agent_framework/openai/_assistant_provider.py +++ b/python/packages/core/agent_framework/openai/_assistant_provider.py @@ -33,7 +33,6 @@ if sys.version_info >= (3, 11): else: from typing_extensions import Self, TypedDict # type:ignore # pragma: no cover -__all__ = ["OpenAIAssistantProvider"] # Type variable for options - allows typed OpenAIAssistantProvider[OptionsCoT] returns # Default matches OpenAIAssistantsClient's default options type diff --git a/python/packages/core/agent_framework/openai/_assistants_client.py b/python/packages/core/agent_framework/openai/_assistants_client.py index 03899d4891..6135f9cac2 100644 --- a/python/packages/core/agent_framework/openai/_assistants_client.py +++ b/python/packages/core/agent_framework/openai/_assistants_client.py @@ -68,12 +68,6 @@ else: if TYPE_CHECKING: from .._middleware import MiddlewareTypes -__all__ = [ - "AssistantToolResources", - "OpenAIAssistantsClient", - "OpenAIAssistantsOptions", -] - # region OpenAI Assistants Options TypedDict diff --git a/python/packages/core/agent_framework/openai/_chat_client.py b/python/packages/core/agent_framework/openai/_chat_client.py index fa232c20a1..ee8f60999c 100644 --- a/python/packages/core/agent_framework/openai/_chat_client.py +++ b/python/packages/core/agent_framework/openai/_chat_client.py @@ -3,6 +3,7 @@ from __future__ import annotations import json +import logging import sys from collections.abc import AsyncIterable, Awaitable, Callable, Mapping, MutableMapping, Sequence from datetime import datetime, timezone @@ -20,7 +21,6 @@ from openai.types.chat.completion_create_params import WebSearchOptions from pydantic import BaseModel from .._clients import BaseChatClient -from .._logging import get_logger from .._middleware import ChatAndFunctionMiddlewareTypes, ChatMiddlewareLayer from .._settings import load_settings from .._tools import ( @@ -60,9 +60,7 @@ if sys.version_info >= (3, 11): else: from typing_extensions import TypedDict # type: ignore # pragma: no cover -__all__ = ["OpenAIChatClient", "OpenAIChatOptions"] - -logger = get_logger("agent_framework.openai") +logger = logging.getLogger("agent_framework.openai") ResponseModelT = TypeVar("ResponseModelT", bound=BaseModel | None, default=None) diff --git a/python/packages/core/agent_framework/openai/_exceptions.py b/python/packages/core/agent_framework/openai/_exceptions.py index 4b4944fd7a..dc134de257 100644 --- a/python/packages/core/agent_framework/openai/_exceptions.py +++ b/python/packages/core/agent_framework/openai/_exceptions.py @@ -10,8 +10,6 @@ from openai import BadRequestError from ..exceptions import ServiceContentFilterException -__all__ = ["ContentFilterResultSeverity", "OpenAIContentFilterException"] - class ContentFilterResultSeverity(Enum): """The severity of the content filter result.""" diff --git a/python/packages/core/agent_framework/openai/_responses_client.py b/python/packages/core/agent_framework/openai/_responses_client.py index 55fdaeeda3..0d970e38d2 100644 --- a/python/packages/core/agent_framework/openai/_responses_client.py +++ b/python/packages/core/agent_framework/openai/_responses_client.py @@ -2,6 +2,7 @@ from __future__ import annotations +import logging import sys from collections.abc import ( AsyncIterable, @@ -36,7 +37,6 @@ from openai.types.responses.web_search_tool_param import WebSearchToolParam from pydantic import BaseModel from .._clients import BaseChatClient -from .._logging import get_logger from .._middleware import ChatMiddlewareLayer from .._settings import load_settings from .._tools import ( @@ -90,10 +90,7 @@ if TYPE_CHECKING: FunctionMiddlewareCallable, ) -logger = get_logger("agent_framework.openai") - - -__all__ = ["OpenAIContinuationToken", "OpenAIResponsesClient", "OpenAIResponsesOptions", "RawOpenAIResponsesClient"] +logger = logging.getLogger("agent_framework.openai") class OpenAIContinuationToken(ContinuationToken): diff --git a/python/packages/core/agent_framework/openai/_shared.py b/python/packages/core/agent_framework/openai/_shared.py index c41f4b6247..d7b0be3723 100644 --- a/python/packages/core/agent_framework/openai/_shared.py +++ b/python/packages/core/agent_framework/openai/_shared.py @@ -22,14 +22,13 @@ from openai.types.responses.response import Response from openai.types.responses.response_stream_event import ResponseStreamEvent from packaging.version import parse -from .._logging import get_logger from .._serialization import SerializationMixin from .._settings import SecretString from .._telemetry import APP_INFO, USER_AGENT_KEY, prepend_agent_framework_to_user_agent from .._tools import FunctionTool from ..exceptions import ServiceInitializationError -logger: logging.Logger = get_logger("agent_framework.openai") +logger: logging.Logger = logging.getLogger("agent_framework.openai") RESPONSE_TYPE = Union[ @@ -53,9 +52,6 @@ else: from typing_extensions import TypedDict # type: ignore # pragma: no cover -__all__ = ["OpenAISettings"] - - def _check_openai_version_for_callable_api_key() -> None: """Check if OpenAI version supports callable API keys. diff --git a/python/packages/core/tests/azure/test_azure_responses_client.py b/python/packages/core/tests/azure/test_azure_responses_client.py index ef8a7df479..5bd7744652 100644 --- a/python/packages/core/tests/azure/test_azure_responses_client.py +++ b/python/packages/core/tests/azure/test_azure_responses_client.py @@ -334,8 +334,10 @@ async def test_integration_options( messages = [Message(role="user", text="What is the weather in Seattle?")] elif option_name == "response_format": # Use prompt that works well with structured output - messages = [Message(role="user", text="The weather in Seattle is sunny")] - messages.append(Message(role="user", text="What is the weather in Seattle?")) + messages = [ + Message(role="user", text="The weather in Seattle is sunny"), + Message(role="user", text="What is the weather in Seattle?"), + ] else: # Generic prompt for simple options messages = [Message(role="user", text="Say 'Hello World' briefly.")] @@ -396,7 +398,12 @@ async def test_integration_web_search() -> None: for streaming in [False, True]: content = { - "messages": "Who are the main characters of Kpop Demon Hunters? Do a web search to find the answer.", + "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": [AzureOpenAIResponsesClient.get_web_search_tool()], @@ -416,7 +423,7 @@ async def test_integration_web_search() -> None: # Test that the client will use the web search tool with location content = { - "messages": "What is the current weather? Do not ask for my current location.", + "messages": [Message(role="user", text="What is the current weather? Do not ask for my current location.")], "options": { "tool_choice": "auto", "tools": [ @@ -498,7 +505,7 @@ async def test_integration_client_agent_hosted_mcp_tool() -> None: """Integration test for MCP tool with Azure Response Agent using Microsoft Learn MCP.""" client = AzureOpenAIResponsesClient(credential=AzureCliCredential()) response = await client.get_response( - "How to create an Azure storage account using az cli?", + messages=[Message(role="user", text="How to create an Azure storage account using az cli?")], options={ # this needs to be high enough to handle the full MCP tool response. "max_tokens": 5000, @@ -523,7 +530,7 @@ async def test_integration_client_agent_hosted_code_interpreter_tool(): client = AzureOpenAIResponsesClient(credential=AzureCliCredential()) response = await client.get_response( - "Calculate the sum of numbers from 1 to 10 using Python code.", + messages=[Message(role="user", text="Calculate the sum of numbers from 1 to 10 using Python code.")], options={ "tools": [AzureOpenAIResponsesClient.get_code_interpreter_tool()], }, diff --git a/python/packages/core/tests/core/test_agents.py b/python/packages/core/tests/core/test_agents.py index 6d776a8b43..dbe386a22b 100644 --- a/python/packages/core/tests/core/test_agents.py +++ b/python/packages/core/tests/core/test_agents.py @@ -43,6 +43,12 @@ async def test_agent_run(agent: SupportsAgentRun) -> None: assert response.messages[0].text == "Response" +async def test_agent_run_with_content(agent: SupportsAgentRun) -> None: + response = await agent.run(Content.from_text("test")) + assert response.messages[0].role == "assistant" + assert response.messages[0].text == "Response" + + async def test_agent_run_streaming(agent: SupportsAgentRun) -> None: async def collect_updates(updates: AsyncIterable[AgentResponseUpdate]) -> list[AgentResponseUpdate]: return [u async for u in updates] diff --git a/python/packages/core/tests/core/test_clients.py b/python/packages/core/tests/core/test_clients.py index 0f87828baa..a23b1d2a5f 100644 --- a/python/packages/core/tests/core/test_clients.py +++ b/python/packages/core/tests/core/test_clients.py @@ -21,13 +21,13 @@ def test_chat_client_type(client: SupportsChatGetResponse): async def test_chat_client_get_response(client: SupportsChatGetResponse): - response = await client.get_response(Message(role="user", text="Hello")) + response = await client.get_response([Message(role="user", text="Hello")]) assert response.text == "test response" assert response.messages[0].role == "assistant" async def test_chat_client_get_response_streaming(client: SupportsChatGetResponse): - async for update in client.get_response(Message(role="user", text="Hello"), stream=True): + async for update in client.get_response([Message(role="user", text="Hello")], stream=True): assert update.text == "test streaming response " or update.text == "another update" assert update.role == "assistant" @@ -38,13 +38,13 @@ def test_base_client(chat_client_base: SupportsChatGetResponse): async def test_base_client_get_response(chat_client_base: SupportsChatGetResponse): - response = await chat_client_base.get_response(Message(role="user", text="Hello")) + response = await chat_client_base.get_response([Message(role="user", text="Hello")]) assert response.messages[0].role == "assistant" assert response.messages[0].text == "test response - Hello" async def test_base_client_get_response_streaming(chat_client_base: SupportsChatGetResponse): - async for update in chat_client_base.get_response(Message(role="user", text="Hello"), stream=True): + async for update in chat_client_base.get_response([Message(role="user", text="Hello")], stream=True): assert update.text == "update - Hello" or update.text == "another update" @@ -59,7 +59,9 @@ async def test_chat_client_instructions_handling(chat_client_base: SupportsChatG "_inner_get_response", side_effect=fake_inner_get_response, ) as mock_inner_get_response: - await chat_client_base.get_response("hello", options={"instructions": instructions}) + await chat_client_base.get_response( + [Message(role="user", text="hello")], options={"instructions": instructions} + ) mock_inner_get_response.assert_called_once() _, kwargs = mock_inner_get_response.call_args messages = kwargs.get("messages", []) diff --git a/python/packages/core/tests/core/test_function_invocation_logic.py b/python/packages/core/tests/core/test_function_invocation_logic.py index 9e498cae76..593aa6f737 100644 --- a/python/packages/core/tests/core/test_function_invocation_logic.py +++ b/python/packages/core/tests/core/test_function_invocation_logic.py @@ -38,7 +38,9 @@ async def test_base_client_with_function_calling(chat_client_base: SupportsChatG ), ChatResponse(messages=Message(role="assistant", text="done")), ] - response = await chat_client_base.get_response("hello", options={"tool_choice": "auto", "tools": [ai_func]}) + response = await chat_client_base.get_response( + [Message(role="user", text="hello")], options={"tool_choice": "auto", "tools": [ai_func]} + ) assert exec_counter == 1 assert len(response.messages) == 3 assert response.messages[0].role == "assistant" @@ -83,7 +85,9 @@ async def test_base_client_with_function_calling_resets(chat_client_base: Suppor ), ChatResponse(messages=Message(role="assistant", text="done")), ] - response = await chat_client_base.get_response("hello", options={"tool_choice": "auto", "tools": [ai_func]}) + response = await chat_client_base.get_response( + [Message(role="user", text="hello")], options={"tool_choice": "auto", "tools": [ai_func]} + ) assert exec_counter == 2 assert len(response.messages) == 5 assert response.messages[0].role == "assistant" @@ -388,11 +392,13 @@ async def test_function_invocation_scenarios( options["conversation_id"] = conversation_id if not streaming: - response = await chat_client_base.get_response("hello", options=options) + response = await chat_client_base.get_response([Message(role="user", text="hello")], options=options) messages = response.messages else: updates = [] - async for update in chat_client_base.get_response("hello", options=options, stream=True): + async for update in chat_client_base.get_response( + [Message(role="user", text="hello")], options=options, stream=True + ): updates.append(update) messages = updates @@ -776,7 +782,9 @@ async def test_max_iterations_limit(chat_client_base: SupportsChatGetResponse): # Set max_iterations to 1 in additional_properties chat_client_base.function_invocation_configuration["max_iterations"] = 1 - response = await chat_client_base.get_response("hello", options={"tool_choice": "auto", "tools": [ai_func]}) + response = await chat_client_base.get_response( + [Message(role="user", text="hello")], options={"tool_choice": "auto", "tools": [ai_func]} + ) # With max_iterations=1, we should: # 1. Execute first function call (exec_counter=1) @@ -803,7 +811,9 @@ async def test_function_invocation_config_enabled_false(chat_client_base: Suppor # Disable function invocation chat_client_base.function_invocation_configuration["enabled"] = False - response = await chat_client_base.get_response("hello", options={"tool_choice": "auto", "tools": [ai_func]}) + response = await chat_client_base.get_response( + [Message(role="user", text="hello")], options={"tool_choice": "auto", "tools": [ai_func]} + ) # Function should not be executed - when enabled=False, the loop doesn't run assert exec_counter == 0 @@ -859,7 +869,9 @@ async def test_function_invocation_config_max_consecutive_errors(chat_client_bas # Set max_consecutive_errors to 2 chat_client_base.function_invocation_configuration["max_consecutive_errors_per_request"] = 2 - response = await chat_client_base.get_response("hello", options={"tool_choice": "auto", "tools": [error_func]}) + response = await chat_client_base.get_response( + [Message(role="user", text="hello")], options={"tool_choice": "auto", "tools": [error_func]} + ) # Should stop after 2 consecutive errors and force a non-tool response error_results = [ @@ -904,7 +916,9 @@ async def test_function_invocation_config_terminate_on_unknown_calls_false(chat_ # Set terminate_on_unknown_calls to False (default) chat_client_base.function_invocation_configuration["terminate_on_unknown_calls"] = False - response = await chat_client_base.get_response("hello", options={"tool_choice": "auto", "tools": [known_func]}) + response = await chat_client_base.get_response( + [Message(role="user", text="hello")], options={"tool_choice": "auto", "tools": [known_func]} + ) # Should have a result message indicating the tool wasn't found assert len(response.messages) == 3 @@ -940,7 +954,9 @@ async def test_function_invocation_config_terminate_on_unknown_calls_true(chat_c # Should raise an exception when encountering an unknown function with pytest.raises(KeyError, match='Error: Requested function "unknown_function" not found'): - await chat_client_base.get_response("hello", options={"tool_choice": "auto", "tools": [known_func]}) + await chat_client_base.get_response( + [Message(role="user", text="hello")], options={"tool_choice": "auto", "tools": [known_func]} + ) assert exec_counter == 0 @@ -978,7 +994,9 @@ async def test_function_invocation_config_additional_tools(chat_client_base: Sup chat_client_base.function_invocation_configuration["additional_tools"] = [hidden_func] # Only pass visible_func in the tools parameter - response = await chat_client_base.get_response("hello", options={"tool_choice": "auto", "tools": [visible_func]}) + response = await chat_client_base.get_response( + [Message(role="user", text="hello")], options={"tool_choice": "auto", "tools": [visible_func]} + ) # Additional tools are treated as declaration_only, so not executed # The function call should be in the messages but not executed @@ -1016,7 +1034,9 @@ async def test_function_invocation_config_include_detailed_errors_false(chat_cli # Set include_detailed_errors to False (default) chat_client_base.function_invocation_configuration["include_detailed_errors"] = False - response = await chat_client_base.get_response("hello", options={"tool_choice": "auto", "tools": [error_func]}) + response = await chat_client_base.get_response( + [Message(role="user", text="hello")], options={"tool_choice": "auto", "tools": [error_func]} + ) # Should have a generic error message error_result = next( @@ -1050,7 +1070,9 @@ async def test_function_invocation_config_include_detailed_errors_true(chat_clie # Set include_detailed_errors to True chat_client_base.function_invocation_configuration["include_detailed_errors"] = True - response = await chat_client_base.get_response("hello", options={"tool_choice": "auto", "tools": [error_func]}) + response = await chat_client_base.get_response( + [Message(role="user", text="hello")], options={"tool_choice": "auto", "tools": [error_func]} + ) # Should have detailed error message error_result = next( @@ -1120,7 +1142,9 @@ async def test_argument_validation_error_with_detailed_errors(chat_client_base: # Set include_detailed_errors to True chat_client_base.function_invocation_configuration["include_detailed_errors"] = True - response = await chat_client_base.get_response("hello", options={"tool_choice": "auto", "tools": [typed_func]}) + response = await chat_client_base.get_response( + [Message(role="user", text="hello")], options={"tool_choice": "auto", "tools": [typed_func]} + ) # Should have detailed validation error error_result = next( @@ -1154,7 +1178,9 @@ async def test_argument_validation_error_without_detailed_errors(chat_client_bas # Set include_detailed_errors to False (default) chat_client_base.function_invocation_configuration["include_detailed_errors"] = False - response = await chat_client_base.get_response("hello", options={"tool_choice": "auto", "tools": [typed_func]}) + response = await chat_client_base.get_response( + [Message(role="user", text="hello")], options={"tool_choice": "auto", "tools": [typed_func]} + ) # Should have generic validation error error_result = next( @@ -1219,7 +1245,9 @@ async def test_unapproved_tool_execution_raises_exception(chat_client_base: Supp ] # Get approval request - response1 = await chat_client_base.get_response("hello", options={"tool_choice": "auto", "tools": [test_func]}) + response1 = await chat_client_base.get_response( + [Message(role="user", text="hello")], options={"tool_choice": "auto", "tools": [test_func]} + ) approval_req = [c for c in response1.messages[0].contents if c.type == "function_approval_request"][0] @@ -1277,7 +1305,9 @@ async def test_approved_function_call_with_error_without_detailed_errors(chat_cl chat_client_base.function_invocation_configuration["include_detailed_errors"] = False # Get approval request - response1 = await chat_client_base.get_response("hello", options={"tool_choice": "auto", "tools": [error_func]}) + response1 = await chat_client_base.get_response( + [Message(role="user", text="hello")], options={"tool_choice": "auto", "tools": [error_func]} + ) approval_req = [c for c in response1.messages[0].contents if c.type == "function_approval_request"][0] @@ -1340,7 +1370,9 @@ async def test_approved_function_call_with_error_with_detailed_errors(chat_clien chat_client_base.function_invocation_configuration["include_detailed_errors"] = True # Get approval request - response1 = await chat_client_base.get_response("hello", options={"tool_choice": "auto", "tools": [error_func]}) + response1 = await chat_client_base.get_response( + [Message(role="user", text="hello")], options={"tool_choice": "auto", "tools": [error_func]} + ) approval_req = [c for c in response1.messages[0].contents if c.type == "function_approval_request"][0] @@ -1403,7 +1435,9 @@ async def test_approved_function_call_with_validation_error(chat_client_base: Su chat_client_base.function_invocation_configuration["include_detailed_errors"] = True # Get approval request - response1 = await chat_client_base.get_response("hello", options={"tool_choice": "auto", "tools": [typed_func]}) + response1 = await chat_client_base.get_response( + [Message(role="user", text="hello")], options={"tool_choice": "auto", "tools": [typed_func]} + ) approval_req = [c for c in response1.messages[0].contents if c.type == "function_approval_request"][0] @@ -1459,7 +1493,9 @@ async def test_approved_function_call_successful_execution(chat_client_base: Sup ] # Get approval request - response1 = await chat_client_base.get_response("hello", options={"tool_choice": "auto", "tools": [success_func]}) + response1 = await chat_client_base.get_response( + [Message(role="user", text="hello")], options={"tool_choice": "auto", "tools": [success_func]} + ) approval_req = [c for c in response1.messages[0].contents if c.type == "function_approval_request"][0] @@ -1575,7 +1611,9 @@ async def test_multiple_function_calls_parallel_execution(chat_client_base: Supp ChatResponse(messages=Message(role="assistant", text="done")), ] - response = await chat_client_base.get_response("hello", options={"tool_choice": "auto", "tools": [func1, func2]}) + response = await chat_client_base.get_response( + [Message(role="user", text="hello")], options={"tool_choice": "auto", "tools": [func1, func2]} + ) # Both functions should have been executed assert "func1_start" in exec_order @@ -1612,7 +1650,9 @@ async def test_callable_function_converted_to_tool(chat_client_base: SupportsCha ] # Pass plain function (will be auto-converted) - response = await chat_client_base.get_response("hello", options={"tool_choice": "auto", "tools": [plain_function]}) + response = await chat_client_base.get_response( + [Message(role="user", text="hello")], options={"tool_choice": "auto", "tools": [plain_function]} + ) # Function should be executed assert exec_counter == 1 @@ -1644,7 +1684,9 @@ async def test_conversation_id_handling(chat_client_base: SupportsChatGetRespons ), ] - response = await chat_client_base.get_response("hello", options={"tool_choice": "auto", "tools": [test_func]}) + response = await chat_client_base.get_response( + [Message(role="user", text="hello")], options={"tool_choice": "auto", "tools": [test_func]} + ) # Should have executed the function results = [content for msg in response.messages for content in msg.contents if content.type == "function_result"] @@ -1671,7 +1713,9 @@ async def test_function_result_appended_to_existing_assistant_message(chat_clien ChatResponse(messages=Message(role="assistant", text="done")), ] - response = await chat_client_base.get_response("hello", options={"tool_choice": "auto", "tools": [test_func]}) + response = await chat_client_base.get_response( + [Message(role="user", text="hello")], options={"tool_choice": "auto", "tools": [test_func]} + ) # Should have messages with both function call and function result assert len(response.messages) >= 2 @@ -1716,7 +1760,9 @@ async def test_error_recovery_resets_counter(chat_client_base: SupportsChatGetRe ChatResponse(messages=Message(role="assistant", text="done")), ] - response = await chat_client_base.get_response("hello", options={"tool_choice": "auto", "tools": [sometimes_fails]}) + response = await chat_client_base.get_response( + [Message(role="user", text="hello")], options={"tool_choice": "auto", "tools": [sometimes_fails]} + ) # Should have both an error and a success error_results = [ @@ -1990,7 +2036,9 @@ async def test_streaming_function_invocation_config_terminate_on_unknown_calls_t # Should raise an exception when encountering an unknown function with pytest.raises(KeyError, match='Error: Requested function "unknown_function" not found'): - async for _ in chat_client_base.get_response("hello", options={"tool_choice": "auto", "tools": [known_func]}): + async for _ in chat_client_base.get_response( + [Message(role="user", text="hello")], options={"tool_choice": "auto", "tools": [known_func]} + ): pass assert exec_counter == 0 diff --git a/python/packages/core/tests/core/test_logging.py b/python/packages/core/tests/core/test_logging.py deleted file mode 100644 index 6565834596..0000000000 --- a/python/packages/core/tests/core/test_logging.py +++ /dev/null @@ -1,39 +0,0 @@ -# Copyright (c) Microsoft. All rights reserved. - - -import pytest - -from agent_framework import get_logger -from agent_framework.exceptions import AgentFrameworkException - - -def test_get_logger(): - """Test that the logger is created with the correct name.""" - logger = get_logger() - assert logger.name == "agent_framework" - - -def test_get_logger_custom_name(): - """Test that the logger can be created with a custom name.""" - custom_name = "agent_framework.custom" - logger = get_logger(custom_name) - assert logger.name == custom_name - - -def test_get_logger_invalid_name(): - """Test that an exception is raised for an invalid logger name.""" - with pytest.raises(AgentFrameworkException): - get_logger("invalid_name") - - -def test_log(caplog): - """Test that the logger can log messages and adheres to the expected format.""" - logger = get_logger() - with caplog.at_level("DEBUG"): - logger.debug("This is a debug message") - assert len(caplog.records) == 1 - record = caplog.records[0] - assert record.levelname == "DEBUG" - assert record.message == "This is a debug message" - assert record.name == "agent_framework" - assert record.pathname.endswith("test_logging.py") diff --git a/python/packages/core/tests/openai/test_openai_chat_client.py b/python/packages/core/tests/openai/test_openai_chat_client.py index e6e5de8314..c5ee81bfce 100644 --- a/python/packages/core/tests/openai/test_openai_chat_client.py +++ b/python/packages/core/tests/openai/test_openai_chat_client.py @@ -1083,7 +1083,12 @@ async def test_integration_web_search() -> None: # Use static method for web search tool web_search_tool = OpenAIChatClient.get_web_search_tool() content = { - "messages": "Who are the main characters of Kpop Demon Hunters? Do a web search to find the answer.", + "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], @@ -1110,7 +1115,7 @@ async def test_integration_web_search() -> None: } ) content = { - "messages": "What is the current weather? Do not ask for my current location.", + "messages": [Message(role="user", text="What is the current weather? Do not ask for my current location.")], "options": { "tool_choice": "auto", "tools": [web_search_tool_with_location], diff --git a/python/packages/core/tests/openai/test_openai_responses_client.py b/python/packages/core/tests/openai/test_openai_responses_client.py index 749939783e..240f09c4a0 100644 --- a/python/packages/core/tests/openai/test_openai_responses_client.py +++ b/python/packages/core/tests/openai/test_openai_responses_client.py @@ -2416,7 +2416,12 @@ async def test_integration_web_search() -> None: # Use static method for web search tool web_search_tool = OpenAIResponsesClient.get_web_search_tool() content = { - "messages": "Who are the main characters of Kpop Demon Hunters? Do a web search to find the answer.", + "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], @@ -2438,7 +2443,7 @@ async def test_integration_web_search() -> None: user_location={"country": "US", "city": "Seattle"}, ) content = { - "messages": "What is the current weather? Do not ask for my current location.", + "messages": [Message(role="user", text="What is the current weather? Do not ask for my current location.")], "options": { "tool_choice": "auto", "tools": [web_search_tool_with_location], diff --git a/python/packages/core/tests/workflow/test_agent_executor_tool_calls.py b/python/packages/core/tests/workflow/test_agent_executor_tool_calls.py index 47356638a6..cae5ea4e3b 100644 --- a/python/packages/core/tests/workflow/test_agent_executor_tool_calls.py +++ b/python/packages/core/tests/workflow/test_agent_executor_tool_calls.py @@ -39,7 +39,7 @@ class _ToolCallingAgent(BaseAgent): def run( self, - messages: str | Message | Sequence[str | Message] | None = None, + messages: str | Content | Message | Sequence[str | Content | Message] | None = None, *, stream: bool = False, session: AgentSession | None = None, diff --git a/python/packages/core/tests/workflow/test_full_conversation.py b/python/packages/core/tests/workflow/test_full_conversation.py index 80a10347b6..74fe419e71 100644 --- a/python/packages/core/tests/workflow/test_full_conversation.py +++ b/python/packages/core/tests/workflow/test_full_conversation.py @@ -35,7 +35,7 @@ class _SimpleAgent(BaseAgent): def run( self, - messages: str | Message | Sequence[str | Message] | None = None, + messages: str | Content | Message | Sequence[str | Content | Message] | None = None, *, stream: bool = False, session: AgentSession | None = None, @@ -105,7 +105,7 @@ class _CaptureAgent(BaseAgent): def run( self, - messages: str | Message | Sequence[str | Message] | None = None, + messages: str | Content | Message | Sequence[str | Content | Message] | None = None, *, stream: bool = False, session: AgentSession | None = None, diff --git a/python/packages/core/tests/workflow/test_workflow.py b/python/packages/core/tests/workflow/test_workflow.py index fb90df6b39..8bbf11fa6a 100644 --- a/python/packages/core/tests/workflow/test_workflow.py +++ b/python/packages/core/tests/workflow/test_workflow.py @@ -835,7 +835,7 @@ class _StreamingTestAgent(BaseAgent): def run( self, - messages: str | Message | Sequence[str | Message] | None = None, + messages: str | Content | Message | Sequence[str | Content | Message] | None = None, *, stream: bool = False, session: AgentSession | None = None, diff --git a/python/packages/core/tests/workflow/test_workflow_kwargs.py b/python/packages/core/tests/workflow/test_workflow_kwargs.py index 36aece0bb3..bf1fd00974 100644 --- a/python/packages/core/tests/workflow/test_workflow_kwargs.py +++ b/python/packages/core/tests/workflow/test_workflow_kwargs.py @@ -52,7 +52,7 @@ class _KwargsCapturingAgent(BaseAgent): def run( self, - messages: str | Message | Sequence[str | Message] | None = None, + messages: str | Content | Message | Sequence[str | Content | Message] | None = None, *, stream: bool = False, session: AgentSession | None = None, @@ -85,7 +85,7 @@ class _OptionsAwareAgent(BaseAgent): def run( self, - messages: str | Message | Sequence[str | Message] | None = None, + messages: str | Content | Message | Sequence[str | Content | Message] | None = None, *, stream: bool = False, session: AgentSession | None = None, diff --git a/python/packages/declarative/agent_framework_declarative/_models.py b/python/packages/declarative/agent_framework_declarative/_models.py index 38bcbdd855..181b3307b7 100644 --- a/python/packages/declarative/agent_framework_declarative/_models.py +++ b/python/packages/declarative/agent_framework_declarative/_models.py @@ -1,12 +1,12 @@ # Copyright (c) Microsoft. All rights reserved. from __future__ import annotations +import logging import os from collections.abc import MutableMapping from contextvars import ContextVar from typing import Any, Literal, TypeVar, Union -from agent_framework import get_logger from agent_framework._serialization import SerializationMixin try: @@ -20,7 +20,7 @@ except (ImportError, RuntimeError): from typing import overload -logger = get_logger("agent_framework.declarative") +logger = logging.getLogger("agent_framework.declarative") # Context variable for safe_mode setting. # When True (default), environment variables are NOT accessible in PowerFx expressions. diff --git a/python/packages/declarative/agent_framework_declarative/_workflows/_actions_agents.py b/python/packages/declarative/agent_framework_declarative/_workflows/_actions_agents.py index e34f6e06f8..32a165b89f 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_actions_agents.py +++ b/python/packages/declarative/agent_framework_declarative/_workflows/_actions_agents.py @@ -10,10 +10,10 @@ This module implements handlers for: from __future__ import annotations import json +import logging from collections.abc import AsyncGenerator from typing import Any, cast -from agent_framework import get_logger from agent_framework._types import AgentResponse, Message from ._handlers import ( @@ -25,7 +25,7 @@ from ._handlers import ( ) from ._human_input import ExternalLoopEvent, QuestionRequest -logger = get_logger("agent_framework.declarative.workflows.actions") +logger = logging.getLogger("agent_framework.declarative") def _extract_json_from_response(text: str) -> Any: diff --git a/python/packages/declarative/agent_framework_declarative/_workflows/_actions_basic.py b/python/packages/declarative/agent_framework_declarative/_workflows/_actions_basic.py index d5132e125c..1813c62233 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_actions_basic.py +++ b/python/packages/declarative/agent_framework_declarative/_workflows/_actions_basic.py @@ -18,11 +18,10 @@ actually yielding any events. from __future__ import annotations +import logging from collections.abc import AsyncGenerator from typing import TYPE_CHECKING, Any, cast -from agent_framework import get_logger - from ._handlers import ( ActionContext, AttachmentOutputEvent, @@ -35,7 +34,7 @@ from ._handlers import ( if TYPE_CHECKING: from ._state import WorkflowState -logger = get_logger("agent_framework.declarative.workflows.actions") +logger = logging.getLogger("agent_framework.declarative") @action_handler("SetValue") diff --git a/python/packages/declarative/agent_framework_declarative/_workflows/_actions_control_flow.py b/python/packages/declarative/agent_framework_declarative/_workflows/_actions_control_flow.py index 0bd04369f7..7328afa970 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_actions_control_flow.py +++ b/python/packages/declarative/agent_framework_declarative/_workflows/_actions_control_flow.py @@ -11,10 +11,9 @@ This module implements handlers for: - ContinueLoop: Skip to the next iteration """ +import logging from collections.abc import AsyncGenerator -from agent_framework import get_logger - from ._handlers import ( ActionContext, LoopControlSignal, @@ -22,7 +21,7 @@ from ._handlers import ( action_handler, ) -logger = get_logger("agent_framework.declarative.workflows.actions") +logger = logging.getLogger("agent_framework.declarative") @action_handler("Foreach") diff --git a/python/packages/declarative/agent_framework_declarative/_workflows/_actions_error.py b/python/packages/declarative/agent_framework_declarative/_workflows/_actions_error.py index 9e32bd8647..ece5d02953 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_actions_error.py +++ b/python/packages/declarative/agent_framework_declarative/_workflows/_actions_error.py @@ -9,18 +9,17 @@ This module implements handlers for: from __future__ import annotations +import logging from collections.abc import AsyncGenerator from dataclasses import dataclass -from agent_framework import get_logger - from ._handlers import ( ActionContext, WorkflowEvent, action_handler, ) -logger = get_logger("agent_framework.declarative.workflows.actions") +logger = logging.getLogger("agent_framework.declarative") class WorkflowActionError(Exception): diff --git a/python/packages/declarative/agent_framework_declarative/_workflows/_factory.py b/python/packages/declarative/agent_framework_declarative/_workflows/_factory.py index 576ef73ac6..319e47dc3b 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_factory.py +++ b/python/packages/declarative/agent_framework_declarative/_workflows/_factory.py @@ -12,6 +12,7 @@ enabling checkpointing, visualization, and pause/resume capabilities. from __future__ import annotations +import logging from collections.abc import Mapping from pathlib import Path from typing import Any, cast @@ -22,13 +23,12 @@ from agent_framework import ( CheckpointStorage, SupportsAgentRun, Workflow, - get_logger, ) from .._loader import AgentFactory from ._declarative_builder import DeclarativeWorkflowBuilder -logger = get_logger("agent_framework.declarative.workflows") +logger = logging.getLogger("agent_framework.declarative") class DeclarativeWorkflowError(Exception): diff --git a/python/packages/declarative/agent_framework_declarative/_workflows/_handlers.py b/python/packages/declarative/agent_framework_declarative/_workflows/_handlers.py index c8a2039073..bd321f615e 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_handlers.py +++ b/python/packages/declarative/agent_framework_declarative/_workflows/_handlers.py @@ -9,16 +9,15 @@ has a corresponding handler registered via the @action_handler decorator. from __future__ import annotations +import logging from collections.abc import AsyncGenerator, Callable from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable -from agent_framework import get_logger - if TYPE_CHECKING: from ._state import WorkflowState -logger = get_logger("agent_framework.declarative.workflows") +logger = logging.getLogger("agent_framework.declarative") @dataclass diff --git a/python/packages/declarative/agent_framework_declarative/_workflows/_human_input.py b/python/packages/declarative/agent_framework_declarative/_workflows/_human_input.py index e7e1da97e4..f0baae8e5c 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_human_input.py +++ b/python/packages/declarative/agent_framework_declarative/_workflows/_human_input.py @@ -10,12 +10,11 @@ This module implements handlers for human input patterns: from __future__ import annotations +import logging from collections.abc import AsyncGenerator from dataclasses import dataclass from typing import TYPE_CHECKING, Any, cast -from agent_framework import get_logger - from ._handlers import ( ActionContext, WorkflowEvent, @@ -25,7 +24,7 @@ from ._handlers import ( if TYPE_CHECKING: from ._state import WorkflowState -logger = get_logger("agent_framework.declarative.workflows.human_input") +logger = logging.getLogger("agent_framework.declarative") @dataclass diff --git a/python/packages/declarative/agent_framework_declarative/_workflows/_state.py b/python/packages/declarative/agent_framework_declarative/_workflows/_state.py index 31ad4124da..3d85e27069 100644 --- a/python/packages/declarative/agent_framework_declarative/_workflows/_state.py +++ b/python/packages/declarative/agent_framework_declarative/_workflows/_state.py @@ -11,11 +11,10 @@ This module provides state management for declarative workflows, handling: from __future__ import annotations +import logging from collections.abc import Mapping from typing import Any, cast -from agent_framework import get_logger - try: from powerfx import Engine @@ -25,7 +24,7 @@ except (ImportError, RuntimeError): # RuntimeError: .NET runtime not available or misconfigured _powerfx_engine = None -logger = get_logger("agent_framework.declarative.workflows") +logger = logging.getLogger("agent_framework.declarative") class WorkflowState: diff --git a/python/packages/durabletask/agent_framework_durabletask/_client.py b/python/packages/durabletask/agent_framework_durabletask/_client.py index 258b710c94..b7eef35252 100644 --- a/python/packages/durabletask/agent_framework_durabletask/_client.py +++ b/python/packages/durabletask/agent_framework_durabletask/_client.py @@ -8,14 +8,16 @@ with durable agents via gRPC. from __future__ import annotations -from agent_framework import AgentResponse, get_logger +import logging + +from agent_framework import AgentResponse from durabletask.client import TaskHubGrpcClient from ._constants import DEFAULT_MAX_POLL_RETRIES, DEFAULT_POLL_INTERVAL_SECONDS from ._executors import ClientAgentExecutor from ._shim import DurableAgentProvider, DurableAIAgent -logger = get_logger("agent_framework.durabletask.client") +logger = logging.getLogger("agent_framework.durabletask") class DurableAIAgentClient(DurableAgentProvider[AgentResponse]): diff --git a/python/packages/durabletask/agent_framework_durabletask/_durable_agent_state.py b/python/packages/durabletask/agent_framework_durabletask/_durable_agent_state.py index 4fd59df051..a10fceefa4 100644 --- a/python/packages/durabletask/agent_framework_durabletask/_durable_agent_state.py +++ b/python/packages/durabletask/agent_framework_durabletask/_durable_agent_state.py @@ -30,6 +30,7 @@ All classes support bidirectional conversion between: from __future__ import annotations import json +import logging from collections.abc import MutableMapping from datetime import datetime, timezone from enum import Enum @@ -40,14 +41,13 @@ from agent_framework import ( Content, Message, UsageDetails, - get_logger, ) from dateutil import parser as date_parser from ._constants import ContentTypes, DurableStateFields from ._models import RunRequest, serialize_response_format -logger = get_logger("agent_framework.durabletask.durable_agent_state") +logger = logging.getLogger("agent_framework.durabletask") class DurableAgentStateEntryJsonType(str, Enum): diff --git a/python/packages/durabletask/agent_framework_durabletask/_entities.py b/python/packages/durabletask/agent_framework_durabletask/_entities.py index 186561e3f4..650e1b8013 100644 --- a/python/packages/durabletask/agent_framework_durabletask/_entities.py +++ b/python/packages/durabletask/agent_framework_durabletask/_entities.py @@ -5,6 +5,7 @@ from __future__ import annotations import inspect +import logging from datetime import datetime, timezone from typing import Any, cast @@ -15,7 +16,6 @@ from agent_framework import ( Message, ResponseStream, SupportsAgentRun, - get_logger, ) from durabletask.entities import DurableEntity @@ -28,7 +28,7 @@ from ._durable_agent_state import ( ) from ._models import RunRequest -logger = get_logger("agent_framework.durabletask.entities") +logger = logging.getLogger("agent_framework.durabletask") class AgentEntityStateProviderMixin: diff --git a/python/packages/durabletask/agent_framework_durabletask/_executors.py b/python/packages/durabletask/agent_framework_durabletask/_executors.py index 7adb79875a..0a7cf50b0b 100644 --- a/python/packages/durabletask/agent_framework_durabletask/_executors.py +++ b/python/packages/durabletask/agent_framework_durabletask/_executors.py @@ -10,13 +10,14 @@ and are injected into the shim. from __future__ import annotations +import logging import time import uuid from abc import ABC, abstractmethod from datetime import datetime, timezone from typing import Any, Generic, TypeVar -from agent_framework import AgentResponse, AgentSession, Content, Message, get_logger +from agent_framework import AgentResponse, AgentSession, Content, Message from durabletask.client import TaskHubGrpcClient from durabletask.entities import EntityInstanceId from durabletask.task import CompletableTask, CompositeTask, OrchestrationContext, Task @@ -27,7 +28,7 @@ from ._durable_agent_state import DurableAgentState from ._models import AgentSessionId, DurableAgentSession, RunRequest from ._response_utils import ensure_response_format, load_agent_response -logger = get_logger("agent_framework.durabletask.executors") +logger = logging.getLogger("agent_framework.durabletask") # TypeVar for the task type returned by executors TaskT = TypeVar("TaskT") diff --git a/python/packages/durabletask/agent_framework_durabletask/_orchestration_context.py b/python/packages/durabletask/agent_framework_durabletask/_orchestration_context.py index 2dd78efe1c..d5bd648482 100644 --- a/python/packages/durabletask/agent_framework_durabletask/_orchestration_context.py +++ b/python/packages/durabletask/agent_framework_durabletask/_orchestration_context.py @@ -8,13 +8,14 @@ orchestration functions to interact with durable agents. from __future__ import annotations -from agent_framework import get_logger +import logging + from durabletask.task import OrchestrationContext from ._executors import DurableAgentTask, OrchestrationAgentExecutor from ._shim import DurableAgentProvider, DurableAIAgent -logger = get_logger("agent_framework.durabletask.orchestration_context") +logger = logging.getLogger("agent_framework.durabletask") class DurableAIAgentOrchestrationContext(DurableAgentProvider[DurableAgentTask]): diff --git a/python/packages/durabletask/agent_framework_durabletask/_response_utils.py b/python/packages/durabletask/agent_framework_durabletask/_response_utils.py index 1085d4b51d..fe371b592f 100644 --- a/python/packages/durabletask/agent_framework_durabletask/_response_utils.py +++ b/python/packages/durabletask/agent_framework_durabletask/_response_utils.py @@ -2,12 +2,13 @@ """Shared utilities for handling AgentResponse parsing and validation.""" +import logging from typing import Any -from agent_framework import AgentResponse, get_logger +from agent_framework import AgentResponse from pydantic import BaseModel -logger = get_logger("agent_framework.durabletask.response_utils") +logger = logging.getLogger("agent_framework.durabletask") def load_agent_response(agent_response: AgentResponse | dict[str, Any] | None) -> AgentResponse: diff --git a/python/packages/durabletask/agent_framework_durabletask/_shim.py b/python/packages/durabletask/agent_framework_durabletask/_shim.py index 1f6165133c..5693876ad7 100644 --- a/python/packages/durabletask/agent_framework_durabletask/_shim.py +++ b/python/packages/durabletask/agent_framework_durabletask/_shim.py @@ -12,7 +12,8 @@ from __future__ import annotations from abc import ABC, abstractmethod from typing import Any, Generic, Literal, TypeVar -from agent_framework import AgentSession, Message, SupportsAgentRun +from agent_framework import AgentSession, SupportsAgentRun, normalize_messages +from agent_framework._types import AgentRunInputs from ._executors import DurableAgentExecutor from ._models import DurableAgentSession @@ -86,7 +87,7 @@ class DurableAIAgent(SupportsAgentRun, Generic[TaskT]): def run( # type: ignore[override] self, - messages: str | Message | list[str] | list[Message] | None = None, + messages: AgentRunInputs | None = None, *, stream: Literal[False] = False, session: AgentSession | None = None, @@ -143,7 +144,7 @@ class DurableAIAgent(SupportsAgentRun, Generic[TaskT]): """ return self._executor.get_new_session(self.name, **kwargs) - def _normalize_messages(self, messages: str | Message | list[str] | list[Message] | None) -> str: + def _normalize_messages(self, messages: AgentRunInputs | None) -> str: """Convert supported message inputs to a single string. Args: @@ -151,19 +152,18 @@ class DurableAIAgent(SupportsAgentRun, Generic[TaskT]): Returns: A single string representation of the messages + + Raises: + ValueError: If normalized messages contain non-text content only. """ - if messages is None: + normalized_messages = normalize_messages(messages) + if not normalized_messages: return "" - if isinstance(messages, str): - return messages - if isinstance(messages, Message): - return messages.text or "" - if isinstance(messages, list): - if not messages: - return "" - first_item = messages[0] - if isinstance(first_item, str): - return "\n".join(messages) # type: ignore[arg-type] - # List of Message - return "\n".join([msg.text or "" for msg in messages]) # type: ignore[union-attr] - return "" + + message_texts: list[str] = [] + for message in normalized_messages: + if not message.text: + raise ValueError("DurableAIAgent only supports text message inputs.") + message_texts.append(message.text) + + return "\n".join(message_texts) diff --git a/python/packages/durabletask/agent_framework_durabletask/_worker.py b/python/packages/durabletask/agent_framework_durabletask/_worker.py index 636dadff2a..96aa058d6f 100644 --- a/python/packages/durabletask/agent_framework_durabletask/_worker.py +++ b/python/packages/durabletask/agent_framework_durabletask/_worker.py @@ -9,15 +9,16 @@ and enables registration of agents as durable entities. from __future__ import annotations import asyncio +import logging from typing import Any -from agent_framework import SupportsAgentRun, get_logger +from agent_framework import SupportsAgentRun from durabletask.worker import TaskHubGrpcWorker from ._callbacks import AgentResponseCallbackProtocol from ._entities import AgentEntity, DurableTaskEntityStateProvider -logger = get_logger("agent_framework.durabletask.worker") +logger = logging.getLogger("agent_framework.durabletask") class DurableAIAgentWorker: diff --git a/python/packages/github_copilot/agent_framework_github_copilot/_agent.py b/python/packages/github_copilot/agent_framework_github_copilot/_agent.py index 1cafc1ba85..26323e915d 100644 --- a/python/packages/github_copilot/agent_framework_github_copilot/_agent.py +++ b/python/packages/github_copilot/agent_framework_github_copilot/_agent.py @@ -23,7 +23,7 @@ from agent_framework import ( ) from agent_framework._settings import load_settings from agent_framework._tools import FunctionTool -from agent_framework._types import normalize_tools +from agent_framework._types import AgentRunInputs, normalize_tools from agent_framework.exceptions import ServiceException from copilot import CopilotClient, CopilotSession from copilot.generated.session_events import SessionEvent, SessionEventType @@ -277,7 +277,7 @@ class GitHubCopilotAgent(BaseAgent, Generic[OptionsT]): @overload def run( self, - messages: str | Message | Sequence[str | Message] | None = None, + messages: AgentRunInputs | None = None, *, stream: Literal[False] = False, session: AgentSession | None = None, @@ -288,7 +288,7 @@ class GitHubCopilotAgent(BaseAgent, Generic[OptionsT]): @overload def run( self, - messages: str | Message | Sequence[str | Message] | None = None, + messages: AgentRunInputs | None = None, *, stream: Literal[True], session: AgentSession | None = None, @@ -298,7 +298,7 @@ class GitHubCopilotAgent(BaseAgent, Generic[OptionsT]): def run( self, - messages: str | Message | Sequence[str | Message] | None = None, + messages: AgentRunInputs | None = None, *, stream: bool = False, session: AgentSession | None = None, @@ -340,7 +340,7 @@ class GitHubCopilotAgent(BaseAgent, Generic[OptionsT]): async def _run_impl( self, - messages: str | Message | Sequence[str | Message] | None = None, + messages: AgentRunInputs | None = None, *, session: AgentSession | None = None, options: OptionsT | None = None, @@ -388,7 +388,7 @@ class GitHubCopilotAgent(BaseAgent, Generic[OptionsT]): async def _stream_updates( self, - messages: str | Message | Sequence[str | Message] | None = None, + messages: AgentRunInputs | None = None, *, session: AgentSession | None = None, options: OptionsT | None = None, diff --git a/python/packages/ollama/agent_framework_ollama/_chat_client.py b/python/packages/ollama/agent_framework_ollama/_chat_client.py index a074562a83..cf7f8931c4 100644 --- a/python/packages/ollama/agent_framework_ollama/_chat_client.py +++ b/python/packages/ollama/agent_framework_ollama/_chat_client.py @@ -3,6 +3,7 @@ from __future__ import annotations import json +import logging import sys from collections.abc import ( AsyncIterable, @@ -28,7 +29,6 @@ from agent_framework import ( Message, ResponseStream, UsageDetails, - get_logger, ) from agent_framework._settings import load_settings from agent_framework.exceptions import ( @@ -281,7 +281,7 @@ class OllamaSettings(TypedDict, total=False): model_id: str | None -logger = get_logger("agent_framework.ollama") +logger = logging.getLogger("agent_framework.ollama") class OllamaChatClient( diff --git a/python/packages/orchestrations/tests/test_group_chat.py b/python/packages/orchestrations/tests/test_group_chat.py index 7115e9c7d7..7550f820c7 100644 --- a/python/packages/orchestrations/tests/test_group_chat.py +++ b/python/packages/orchestrations/tests/test_group_chat.py @@ -38,7 +38,7 @@ class StubAgent(BaseAgent): def run( # type: ignore[override] self, - messages: str | Message | Sequence[str | Message] | None = None, + messages: str | Content | Message | Sequence[str | Content | Message] | None = None, *, stream: bool = False, session: AgentSession | None = None, @@ -76,7 +76,7 @@ class StubManagerAgent(Agent): async def run( self, - messages: str | Message | Sequence[str | Message] | None = None, + messages: str | Content | Message | Sequence[str | Content | Message] | None = None, *, session: AgentSession | None = None, **kwargs: Any, @@ -130,7 +130,7 @@ class ConcatenatedJsonManagerAgent(Agent): async def run( self, - messages: str | Message | Sequence[str | Message] | None = None, + messages: str | Content | Message | Sequence[str | Content | Message] | None = None, *, session: AgentSession | None = None, **kwargs: Any, @@ -896,7 +896,7 @@ async def test_group_chat_with_orchestrator_factory_returning_chat_agent(): async def run( self, - messages: str | Message | Sequence[str | Message] | None = None, + messages: str | Content | Message | Sequence[str | Content | Message] | None = None, *, session: AgentSession | None = None, **kwargs: Any, diff --git a/python/packages/orchestrations/tests/test_magentic.py b/python/packages/orchestrations/tests/test_magentic.py index 25c068b9fe..e1d0ef8c32 100644 --- a/python/packages/orchestrations/tests/test_magentic.py +++ b/python/packages/orchestrations/tests/test_magentic.py @@ -150,7 +150,7 @@ class StubAgent(BaseAgent): def run( # type: ignore[override] self, - messages: str | Message | Sequence[str | Message] | None = None, + messages: str | Content | Message | Sequence[str | Content | Message] | None = None, *, stream: bool = False, session: AgentSession | None = None, @@ -411,7 +411,7 @@ class StubManagerAgent(BaseAgent): def run( self, - messages: str | Message | Sequence[str | Message] | None = None, + messages: str | Content | Message | Sequence[str | Content | Message] | None = None, *, stream: bool = False, session: Any = None, diff --git a/python/packages/purview/agent_framework_purview/_client.py b/python/packages/purview/agent_framework_purview/_client.py index b2b07b2637..2de5340cb7 100644 --- a/python/packages/purview/agent_framework_purview/_client.py +++ b/python/packages/purview/agent_framework_purview/_client.py @@ -5,12 +5,12 @@ from __future__ import annotations import base64 import inspect import json +import logging from typing import Any, cast from uuid import uuid4 import httpx from agent_framework import AGENT_FRAMEWORK_USER_AGENT -from agent_framework._logging import get_logger from agent_framework.observability import get_tracer from azure.core.credentials import TokenCredential from azure.core.credentials_async import AsyncTokenCredential @@ -33,7 +33,7 @@ from ._models import ( ) from ._settings import PurviewSettings, get_purview_scopes -logger = get_logger("agent_framework.purview") +logger = logging.getLogger("agent_framework.purview") class PurviewClient: diff --git a/python/packages/purview/agent_framework_purview/_middleware.py b/python/packages/purview/agent_framework_purview/_middleware.py index f0165704d8..3a4a628294 100644 --- a/python/packages/purview/agent_framework_purview/_middleware.py +++ b/python/packages/purview/agent_framework_purview/_middleware.py @@ -1,9 +1,9 @@ # Copyright (c) Microsoft. All rights reserved. +import logging from collections.abc import Awaitable, Callable from agent_framework import AgentContext, AgentMiddleware, ChatContext, ChatMiddleware, MiddlewareTermination -from agent_framework._logging import get_logger from azure.core.credentials import TokenCredential from azure.core.credentials_async import AsyncTokenCredential @@ -14,7 +14,7 @@ from ._models import Activity from ._processor import ScopedContentProcessor from ._settings import PurviewSettings -logger = get_logger("agent_framework.purview") +logger = logging.getLogger("agent_framework.purview") class PurviewPolicyMiddleware(AgentMiddleware): diff --git a/python/packages/purview/agent_framework_purview/_models.py b/python/packages/purview/agent_framework_purview/_models.py index 0e4985689e..ad6cc5b331 100644 --- a/python/packages/purview/agent_framework_purview/_models.py +++ b/python/packages/purview/agent_framework_purview/_models.py @@ -2,16 +2,16 @@ from __future__ import annotations +import logging from collections.abc import Mapping, MutableMapping, Sequence from datetime import datetime from enum import Enum, Flag, auto from typing import Any, ClassVar, TypeVar, cast from uuid import uuid4 -from agent_framework._logging import get_logger from agent_framework._serialization import SerializationMixin -logger = get_logger("agent_framework.purview") +logger = logging.getLogger("agent_framework.purview") # -------------------------------------------------------------------------------------- # Enums & flag helpers diff --git a/python/packages/purview/agent_framework_purview/_processor.py b/python/packages/purview/agent_framework_purview/_processor.py index bc8cd045c0..a7fc030cbf 100644 --- a/python/packages/purview/agent_framework_purview/_processor.py +++ b/python/packages/purview/agent_framework_purview/_processor.py @@ -1,13 +1,13 @@ # Copyright (c) Microsoft. All rights reserved. import asyncio +import logging import time import uuid from collections.abc import Iterable, MutableMapping from typing import Any from agent_framework import Message -from agent_framework._logging import get_logger from ._cache import CacheProvider, InMemoryCacheProvider, create_protection_scopes_cache_key from ._client import PurviewClient @@ -37,7 +37,7 @@ from ._models import ( ) from ._settings import PurviewSettings -logger = get_logger("agent_framework.purview") +logger = logging.getLogger("agent_framework.purview") def _is_valid_guid(value: str | None) -> bool: diff --git a/python/samples/02-agents/observability/README.md b/python/samples/02-agents/observability/README.md index d0e4ea06d1..fa6ba0fec1 100644 --- a/python/samples/02-agents/observability/README.md +++ b/python/samples/02-agents/observability/README.md @@ -178,12 +178,15 @@ The `configure_otel_providers()` function automatically reads **standard OpenTel > **Note**: These are standard OpenTelemetry environment variables. See the [OpenTelemetry spec](https://opentelemetry.io/docs/specs/otel/configuration/sdk-environment-variables/) for more details. #### Logging -Agent Framework has a built-in logging configuration that works well with telemetry. It sets the format to a standard format that includes timestamp, pathname, line number, and log level. You can use that by calling the `setup_logging()` function from the `agent_framework` module. +Use standard Python logging configuration to align logs with telemetry output. ```python -from agent_framework import setup_logging +import logging -setup_logging() +logging.basicConfig( + format="[%(asctime)s - %(pathname)s:%(lineno)d - %(levelname)s] %(message)s", + datefmt="%Y-%m-%d %H:%M:%S", +) ``` You can control at what level logging happens and thus what logs get exported, you can do this, by adding this: diff --git a/python/samples/02-agents/observability/configure_otel_providers_with_parameters.py b/python/samples/02-agents/observability/configure_otel_providers_with_parameters.py index 5ec2698607..4ad525d354 100644 --- a/python/samples/02-agents/observability/configure_otel_providers_with_parameters.py +++ b/python/samples/02-agents/observability/configure_otel_providers_with_parameters.py @@ -2,11 +2,12 @@ import argparse import asyncio +import logging from contextlib import suppress from random import randint from typing import TYPE_CHECKING, Annotated, Literal -from agent_framework import setup_logging, tool +from agent_framework import tool from agent_framework.observability import configure_otel_providers, get_tracer from agent_framework.openai import OpenAIResponsesClient from opentelemetry import trace @@ -31,7 +32,9 @@ Use this approach when you need custom exporter configuration beyond what enviro SCENARIOS = ["client", "client_stream", "tool", "all"] -# NOTE: approval_mode="never_require" is for sample brevity. Use "always_require" in production; see samples/02-agents/tools/function_tool_with_approval.py and samples/02-agents/tools/function_tool_with_approval_and_sessions.py. +# NOTE: approval_mode="never_require" is for sample brevity. +# Use "always_require" in production; see samples/02-agents/tools/function_tool_with_approval.py +# and samples/02-agents/tools/function_tool_with_approval_and_sessions.py. @tool(approval_mode="never_require") async def get_weather( location: Annotated[str, Field(description="The location to get the weather for.")], @@ -101,7 +104,10 @@ async def main(scenario: Literal["client", "client_stream", "tool", "all"] = "al """Run the selected scenario(s).""" # Setup the logging with the more complete format - setup_logging() + logging.basicConfig( + format="[%(asctime)s - %(pathname)s:%(lineno)d - %(levelname)s] %(message)s", + datefmt="%Y-%m-%d %H:%M:%S", + ) # Create custom OTLP exporters with specific configuration # Note: You need to install opentelemetry-exporter-otlp-proto-grpc or -http separately diff --git a/python/samples/02-agents/providers/azure_openai/azure_responses_client_image_analysis.py b/python/samples/02-agents/providers/azure_openai/azure_responses_client_image_analysis.py index e9bedfd474..f23222b018 100644 --- a/python/samples/02-agents/providers/azure_openai/azure_responses_client_image_analysis.py +++ b/python/samples/02-agents/providers/azure_openai/azure_responses_client_image_analysis.py @@ -2,7 +2,7 @@ import asyncio -from agent_framework import Content, Message +from agent_framework import Content from agent_framework.azure import AzureOpenAIResponsesClient from azure.identity import AzureCliCredential @@ -20,24 +20,17 @@ async def main(): # 1. Create an Azure Responses agent with vision capabilities agent = AzureOpenAIResponsesClient(credential=AzureCliCredential()).as_agent( name="VisionAgent", - instructions="You are a helpful agent that can analyze images.", + instructions="You are a image analysist, you get a image and need to respond with what you see in the picture.", ) - # 2. Create a simple message with both text and image content - user_message = Message( - role="user", - contents=[ - Content.from_text("What do you see in this image?"), - Content.from_uri( - uri="https://images.unsplash.com/photo-1506905925346-21bda4d32df4?w=800", - media_type="image/jpeg", - ), - ], - ) - - # 3. Get the agent's response + # 2. Get the agent's response print("User: What do you see in this image? [Image provided]") - result = await agent.run(user_message) + result = await agent.run( + Content.from_uri( + uri="https://images.unsplash.com/photo-1506905925346-21bda4d32df4?w=800", + media_type="image/jpeg", + ) + ) print(f"Agent: {result.text}") print() diff --git a/python/samples/02-agents/providers/openai/openai_responses_client_image_analysis.py b/python/samples/02-agents/providers/openai/openai_responses_client_image_analysis.py index 93c517b97b..7e297524d1 100644 --- a/python/samples/02-agents/providers/openai/openai_responses_client_image_analysis.py +++ b/python/samples/02-agents/providers/openai/openai_responses_client_image_analysis.py @@ -2,7 +2,7 @@ import asyncio -from agent_framework import Content, Message +from agent_framework import Content from agent_framework.openai import OpenAIResponsesClient """ @@ -19,24 +19,17 @@ async def main(): # 1. Create an OpenAI Responses agent with vision capabilities agent = OpenAIResponsesClient().as_agent( name="VisionAgent", - instructions="You are a helpful agent that can analyze images.", + instructions="You are a image analysist, you get a image and need to respond with what you see in the picture.", ) - # 2. Create a simple message with both text and image content - user_message = Message( - role="user", - contents=[ - Content.from_text(text="What do you see in this image?"), - Content.from_uri( - uri="https://images.unsplash.com/photo-1506905925346-21bda4d32df4?w=800", - media_type="image/jpeg", - ), - ], - ) - - # 3. Get the agent's response + # 2. Get the agent's response print("User: What do you see in this image? [Image provided]") - result = await agent.run(user_message) + result = await agent.run( + Content.from_uri( + uri="https://images.unsplash.com/photo-1506905925346-21bda4d32df4?w=800", + media_type="image/jpeg", + ) + ) print(f"Agent: {result.text}") print() diff --git a/python/samples/README.md b/python/samples/README.md index 36df843e6a..eaeb47e7fe 100644 --- a/python/samples/README.md +++ b/python/samples/README.md @@ -38,6 +38,10 @@ export AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME="gpt-4o" For Azure authentication, run `az login` before running samples. +## Note on XML tags + +Some sample files include XML-style snippet tags (for example `` and ``). These are used by our documentation tooling and can be ignored or removed when you use the samples outside this repository. + ## Additional Resources - [Agent Framework Documentation](https://learn.microsoft.com/agent-framework/) From ed113f941c128acf00291a27c9ec094d31c8d747 Mon Sep 17 00:00:00 2001 From: westey <164392973+westey-m@users.noreply.github.com> Date: Mon, 16 Feb 2026 15:47:44 +0000 Subject: [PATCH 14/22] Fixes for bug bash. (#3927) --- .../Program.cs | 15 ++++++++------- .../Agents/Agent_Step10_AsMcpTool/README.md | 4 ++-- .../Agents/Agent_Step16_ChatReduction/Program.cs | 15 +++++++++++++-- 3 files changed, 23 insertions(+), 11 deletions(-) diff --git a/dotnet/samples/GettingStarted/Agents/Agent_Step06_PersistedConversations/Program.cs b/dotnet/samples/GettingStarted/Agents/Agent_Step06_PersistedConversations/Program.cs index e22c377929..d3331cb2b8 100644 --- a/dotnet/samples/GettingStarted/Agents/Agent_Step06_PersistedConversations/Program.cs +++ b/dotnet/samples/GettingStarted/Agents/Agent_Step06_PersistedConversations/Program.cs @@ -1,5 +1,7 @@ // Copyright (c) Microsoft. All rights reserved. +#pragma warning disable CA1869 // Cache and reuse 'JsonSerializerOptions' instances + // This sample shows how to create and use a simple AI agent with a conversation that can be persisted to disk. using System.Text.Json; @@ -30,15 +32,14 @@ Console.WriteLine(await agent.RunAsync("Tell me a joke about a pirate.", session // Serialize the session state to a JsonElement, so it can be stored for later use. JsonElement serializedSession = await agent.SerializeSessionAsync(session); -// Save the serialized session to a temporary file (for demonstration purposes). -string tempFilePath = Path.GetTempFileName(); -await File.WriteAllTextAsync(tempFilePath, JsonSerializer.Serialize(serializedSession)); - -// Load the serialized session from the temporary file (for demonstration purposes). -JsonElement reloadedSerializedSession = JsonElement.Parse(await File.ReadAllTextAsync(tempFilePath)); +// In a real application, you would typically write the serialized session to a file or +// database for persistence, and read it back when resuming the conversation. +// Here we'll just write the serialized session to console (for demonstration purposes). +Console.WriteLine("\n--- Serialized session ---\n"); +Console.WriteLine(JsonSerializer.Serialize(serializedSession, new JsonSerializerOptions { WriteIndented = true }) + "\n"); // Deserialize the session state after loading from storage. -AgentSession resumedSession = await agent.DeserializeSessionAsync(reloadedSerializedSession); +AgentSession resumedSession = await agent.DeserializeSessionAsync(serializedSession); // Run the agent again with the resumed session. Console.WriteLine(await agent.RunAsync("Now tell the same joke in the voice of a pirate, and add some emojis to the joke.", resumedSession)); diff --git a/dotnet/samples/GettingStarted/Agents/Agent_Step10_AsMcpTool/README.md b/dotnet/samples/GettingStarted/Agents/Agent_Step10_AsMcpTool/README.md index c56c9a7a68..56701066d1 100644 --- a/dotnet/samples/GettingStarted/Agents/Agent_Step10_AsMcpTool/README.md +++ b/dotnet/samples/GettingStarted/Agents/Agent_Step10_AsMcpTool/README.md @@ -11,9 +11,9 @@ Alternatively, use the QuickstartClient sample from this repository: https://git To use the [MCP Inspector](https://modelcontextprotocol.io/docs/tools/inspector), follow these steps: 1. Open a terminal in the Agent_Step10_AsMcpTool project directory. -1. Run the `npx @modelcontextprotocol/inspector dotnet run` command to start the MCP Inspector. Make sure you have [node.js](https://nodejs.org/en/download/) and npm installed. +1. Run the `npx @modelcontextprotocol/inspector dotnet run --framework net10.0` command to start the MCP Inspector. Make sure you have [node.js](https://nodejs.org/en/download/) and npm installed. ```bash - npx @modelcontextprotocol/inspector dotnet run + npx @modelcontextprotocol/inspector dotnet run --framework net10.0 ``` 1. When the inspector is running, it will display a URL in the terminal, like this: ``` diff --git a/dotnet/samples/GettingStarted/Agents/Agent_Step16_ChatReduction/Program.cs b/dotnet/samples/GettingStarted/Agents/Agent_Step16_ChatReduction/Program.cs index 875f3375a4..c95f46a6c0 100644 --- a/dotnet/samples/GettingStarted/Agents/Agent_Step16_ChatReduction/Program.cs +++ b/dotnet/samples/GettingStarted/Agents/Agent_Step16_ChatReduction/Program.cs @@ -38,18 +38,29 @@ Console.WriteLine(await agent.RunAsync("Tell me a joke about a pirate.", session // Get the chat history to see how many messages are stored. // We can use the ChatHistoryProvider, that is also used by the agent, to read the // chat history from the session state, and see how the reducer is affecting the stored messages. +// Here we expect to see 2 messages, the original user message and the agent response message. var provider = agent.GetService(); List? chatHistory = provider?.GetMessages(session); Console.WriteLine($"\nChat history has {chatHistory?.Count} messages.\n"); // Invoke the agent a few more times. Console.WriteLine(await agent.RunAsync("Tell me a joke about a robot.", session)); + +// Now we expect to see 4 messages in the chat history, 2 input and 2 output. +// While the target number of messages is 2, the default time for the InMemoryChatHistoryProvider +// to trigger the reducer is just before messages are contributed to a new agent run. +// So at this time, we have not yet triggered the reducer for the most recently added messages, +// and they are still in the chat history. +chatHistory = provider?.GetMessages(session); Console.WriteLine($"\nChat history has {chatHistory?.Count} messages.\n"); + Console.WriteLine(await agent.RunAsync("Tell me a joke about a lemur.", session)); +chatHistory = provider?.GetMessages(session); Console.WriteLine($"\nChat history has {chatHistory?.Count} messages.\n"); // At this point, the chat history has exceeded the limit and the original message will not exist anymore, -// so asking a follow up question about it will not work as expected. -Console.WriteLine(await agent.RunAsync("Tell me the joke about the pirate again, but add emojis and use the voice of a parrot.", session)); +// so asking a follow up question about it may not work as expected. +Console.WriteLine(await agent.RunAsync("What was the first joke I asked you to tell again?", session)); +chatHistory = provider?.GetMessages(session); Console.WriteLine($"\nChat history has {chatHistory?.Count} messages.\n"); From aab621f5eb1ea17d106b8232a9256cddc2af82ee Mon Sep 17 00:00:00 2001 From: Eduard van Valkenburg Date: Mon, 16 Feb 2026 17:30:38 +0100 Subject: [PATCH 15/22] Python: Fix tool normalization and provider sample consolidation (#3953) * Fix tool normalization and provider samples - restore callable/single-tool normalization paths and unset tool-choice behavior\n- consolidate and expand chat/provider samples (OpenAI/Azure/Anthropic/Ollama/Bedrock)\n- migrate Bedrock lazy import surface to agent_framework.amazon and move provider samples Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * small fix in sample * Finalize provider, samples, and core cleanup Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Fix CopilotTool passthrough in agent Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * fix link --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- python/CODING_STANDARD.md | 15 ++ .../agent_framework_anthropic/_chat_client.py | 2 + .../_agent_provider.py | 36 ++-- .../agent_framework_azure_ai/_chat_client.py | 7 +- .../agent_framework_azure_ai/_client.py | 9 +- .../_project_provider.py | 35 ++-- python/packages/bedrock/AGENTS.md | 4 +- python/packages/bedrock/README.md | 2 +- .../agent_framework_bedrock/_chat_client.py | 2 +- python/packages/bedrock/samples/__init__.py | 0 .../bedrock/samples/bedrock_sample.py | 45 ----- .../claude/agent_framework_claude/_agent.py | 25 +-- .../packages/core/agent_framework/__init__.py | 8 + .../packages/core/agent_framework/_agents.py | 69 ++------ .../packages/core/agent_framework/_clients.py | 9 +- .../packages/core/agent_framework/_tools.py | 108 ++++++------ .../packages/core/agent_framework/_types.py | 61 ++----- .../agent_framework/_workflows/__init__.py | 13 ++ .../core/agent_framework/a2a/__init__.py | 11 +- .../core/agent_framework/a2a/__init__.pyi | 2 - .../core/agent_framework/ag_ui/__init__.py | 14 +- .../core/agent_framework/ag_ui/__init__.pyi | 2 - .../core/agent_framework/amazon/__init__.py | 35 ++++ .../core/agent_framework/amazon/__init__.pyi | 15 ++ .../agent_framework/anthropic/__init__.py | 31 +++- .../agent_framework/anthropic/__init__.pyi | 5 +- .../core/agent_framework/azure/__init__.py | 14 ++ .../core/agent_framework/chatkit/__init__.py | 13 +- .../core/agent_framework/chatkit/__init__.pyi | 2 - .../agent_framework/declarative/__init__.py | 13 +- .../agent_framework/declarative/__init__.pyi | 2 - .../core/agent_framework/devui/__init__.py | 15 +- .../core/agent_framework/devui/__init__.pyi | 2 - .../core/agent_framework/exceptions.py | 2 + .../core/agent_framework/github/__init__.py | 12 +- .../core/agent_framework/github/__init__.pyi | 2 - .../core/agent_framework/lab/__init__.py | 6 + .../core/agent_framework/mem0/__init__.py | 11 +- .../core/agent_framework/mem0/__init__.pyi | 2 - .../agent_framework/microsoft/__init__.py | 30 +++- .../agent_framework/microsoft/__init__.pyi | 10 +- .../core/agent_framework/observability.py | 18 +- .../core/agent_framework/ollama/__init__.py | 12 +- .../core/agent_framework/ollama/__init__.pyi | 2 - .../core/agent_framework/openai/__init__.py | 12 ++ .../openai/_assistant_provider.py | 29 ++-- .../openai/_assistants_client.py | 33 ++-- .../agent_framework/openai/_chat_client.py | 30 ++-- .../openai/_responses_client.py | 19 ++- .../orchestrations/__init__.py | 14 +- .../orchestrations/__init__.pyi | 2 - .../core/agent_framework/redis/__init__.py | 12 +- .../core/agent_framework/redis/__init__.pyi | 2 - python/packages/core/pyproject.toml | 3 + .../core/test_function_invocation_logic.py | 30 ++++ python/packages/core/tests/core/test_types.py | 4 +- .../openai/test_openai_assistants_client.py | 47 ++++++ .../tests/openai/test_openai_chat_client.py | 15 ++ python/packages/foundry_local/README.md | 4 + .../agent_framework_github_copilot/_agent.py | 18 +- .../samples/01-get-started/01_hello_agent.py | 2 + python/samples/01-get-started/README.md | 6 +- .../samples/02-agents/chat_client/README.md | 67 ++++++-- .../chat_client/azure_ai_chat_client.py | 49 ------ .../chat_client/azure_assistants_client.py | 49 ------ .../chat_client/azure_chat_client.py | 49 ------ .../chat_client/azure_responses_client.py | 95 ----------- .../chat_client/built_in_chat_clients.py | 156 ++++++++++++++++++ .../chat_client/custom_chat_client.py | 55 +++--- .../chat_client/openai_assistants_client.py | 47 ------ .../chat_client/openai_chat_client.py | 47 ------ .../chat_client/openai_responses_client.py | 47 ------ .../azure_ai_with_search_context_agentic.py | 2 + .../azure_ai_with_search_context_semantic.py | 1 + .../context_providers/mem0/mem0_basic.py | 2 +- .../context_providers/mem0/mem0_oss.py | 6 +- .../context_providers/mem0/mem0_sessions.py | 42 +++-- .../context_providers/redis/README.md | 19 ++- .../redis/azure_redis_conversation.py | 23 +-- .../context_providers/redis/redis_basics.py | 27 ++- .../redis/redis_conversation.py | 13 +- .../context_providers/redis/redis_sessions.py | 39 ++--- .../simple_context_provider.py | 119 ++++++------- python/samples/02-agents/providers/README.md | 19 +++ .../02-agents/providers/amazon/README.md | 17 ++ .../providers/amazon/bedrock_chat_client.py | 61 +++++++ .../anthropic/anthropic_claude_basic.py | 2 +- .../anthropic/anthropic_claude_with_mcp.py | 2 +- ...hropic_claude_with_multiple_permissions.py | 2 +- .../anthropic_claude_with_session.py | 2 +- .../anthropic/anthropic_claude_with_shell.py | 2 +- .../anthropic/anthropic_claude_with_tools.py | 2 +- .../anthropic/anthropic_claude_with_url.py | 2 +- .../providers/foundry_local/README.md | 22 +++ .../foundry_local}/foundry_local_agent.py | 2 +- ...penai_responses_client_image_generation.py | 20 ++- python/samples/02-agents/response_stream.py | 8 +- python/samples/02-agents/typed_options.py | 11 +- python/uv.lock | 16 +- 99 files changed, 1190 insertions(+), 969 deletions(-) delete mode 100644 python/packages/bedrock/samples/__init__.py delete mode 100644 python/packages/bedrock/samples/bedrock_sample.py create mode 100644 python/packages/core/agent_framework/amazon/__init__.py create mode 100644 python/packages/core/agent_framework/amazon/__init__.pyi delete mode 100644 python/samples/02-agents/chat_client/azure_ai_chat_client.py delete mode 100644 python/samples/02-agents/chat_client/azure_assistants_client.py delete mode 100644 python/samples/02-agents/chat_client/azure_chat_client.py delete mode 100644 python/samples/02-agents/chat_client/azure_responses_client.py create mode 100644 python/samples/02-agents/chat_client/built_in_chat_clients.py delete mode 100644 python/samples/02-agents/chat_client/openai_assistants_client.py delete mode 100644 python/samples/02-agents/chat_client/openai_chat_client.py delete mode 100644 python/samples/02-agents/chat_client/openai_responses_client.py create mode 100644 python/samples/02-agents/providers/README.md create mode 100644 python/samples/02-agents/providers/amazon/README.md create mode 100644 python/samples/02-agents/providers/amazon/bedrock_chat_client.py create mode 100644 python/samples/02-agents/providers/foundry_local/README.md rename python/{packages/foundry_local/samples => samples/02-agents/providers/foundry_local}/foundry_local_agent.py (97%) diff --git a/python/CODING_STANDARD.md b/python/CODING_STANDARD.md index 028b914ba9..846d29aac1 100644 --- a/python/CODING_STANDARD.md +++ b/python/CODING_STANDARD.md @@ -10,6 +10,21 @@ We use [ruff](https://github.com/astral-sh/ruff) for both linting and formatting - **Target Python version**: 3.10+ - **Google-style docstrings**: All public functions, classes, and modules should have docstrings following Google conventions +### Module Docstrings + +Public modules must include a module-level docstring, including `__init__.py` files. + +- Namespace-style `__init__.py` modules (for example under `agent_framework//`) should use a structured + docstring that includes: + - A one-line summary of the namespace + - A short "This module lazily re-exports objects from:" section that lists only pip install package names + (for example `agent-framework-a2a`) + - A short "Supported classes:" (or "Supported classes and functions:") section +- The main `agent_framework/__init__.py` should include a concise background-oriented docstring rather than a long + per-symbol list. +- Core modules with broad surface area, including `agent_framework/exceptions.py` and + `agent_framework/observability.py`, should always have explicit module docstrings. + ## Type Annotations ### Future Annotations diff --git a/python/packages/anthropic/agent_framework_anthropic/_chat_client.py b/python/packages/anthropic/agent_framework_anthropic/_chat_client.py index a6b0c8e7d3..4e7ffafc2d 100644 --- a/python/packages/anthropic/agent_framework_anthropic/_chat_client.py +++ b/python/packages/anthropic/agent_framework_anthropic/_chat_client.py @@ -720,6 +720,8 @@ class AnthropicClient( if options.get("tool_choice") is None: return result or None tool_mode = validate_tool_mode(options.get("tool_choice")) + if tool_mode is None: + return result or None allow_multiple = options.get("allow_multiple_tool_calls") match tool_mode.get("mode"): case "auto": diff --git a/python/packages/azure-ai/agent_framework_azure_ai/_agent_provider.py b/python/packages/azure-ai/agent_framework_azure_ai/_agent_provider.py index f5c5201531..65bcb9d6c0 100644 --- a/python/packages/azure-ai/agent_framework_azure_ai/_agent_provider.py +++ b/python/packages/azure-ai/agent_framework_azure_ai/_agent_provider.py @@ -16,6 +16,7 @@ from agent_framework import ( ) from agent_framework._mcp import MCPTool from agent_framework._settings import load_settings +from agent_framework._tools import ToolTypes from agent_framework.exceptions import ServiceInitializationError from azure.ai.agents.aio import AgentsClient from azure.ai.agents.models import Agent as AzureAgent @@ -169,11 +170,7 @@ class AzureAIAgentsProvider(Generic[OptionsCoT]): model: str | None = None, instructions: str | None = None, description: str | None = None, - tools: FunctionTool - | Callable[..., Any] - | MutableMapping[str, Any] - | Sequence[FunctionTool | Callable[..., Any] | MutableMapping[str, Any]] - | None = None, + tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None, default_options: OptionsCoT | None = None, middleware: Sequence[MiddlewareTypes] | None = None, context_providers: Sequence[BaseContextProvider] | None = None, @@ -242,7 +239,12 @@ class AzureAIAgentsProvider(Generic[OptionsCoT]): normalized_tools = normalize_tools(tools) if normalized_tools: # Only convert non-MCP tools to Azure AI format - non_mcp_tools = [t for t in normalized_tools if not isinstance(t, MCPTool)] + non_mcp_tools: list[FunctionTool | MutableMapping[str, Any]] = [] + for normalized_tool in normalized_tools: + if isinstance(normalized_tool, MCPTool): + continue + if isinstance(normalized_tool, (FunctionTool, MutableMapping)): + non_mcp_tools.append(normalized_tool) if non_mcp_tools: # Pass run_options to capture tool_resources (e.g., for file search vector stores) run_options: dict[str, Any] = {} @@ -266,11 +268,7 @@ class AzureAIAgentsProvider(Generic[OptionsCoT]): self, id: str, *, - tools: FunctionTool - | Callable[..., Any] - | MutableMapping[str, Any] - | Sequence[FunctionTool | Callable[..., Any] | MutableMapping[str, Any]] - | None = None, + tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None, default_options: OptionsCoT | None = None, middleware: Sequence[MiddlewareTypes] | None = None, context_providers: Sequence[BaseContextProvider] | None = None, @@ -322,11 +320,7 @@ class AzureAIAgentsProvider(Generic[OptionsCoT]): def as_agent( self, agent: AzureAgent, - tools: FunctionTool - | Callable[..., Any] - | MutableMapping[str, Any] - | Sequence[FunctionTool | Callable[..., Any] | MutableMapping[str, Any]] - | None = None, + tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None, default_options: OptionsCoT | None = None, middleware: Sequence[MiddlewareTypes] | None = None, context_providers: Sequence[BaseContextProvider] | None = None, @@ -379,7 +373,7 @@ class AzureAIAgentsProvider(Generic[OptionsCoT]): def _to_chat_agent_from_agent( self, agent: AzureAgent, - provided_tools: Sequence[FunctionTool | MutableMapping[str, Any]] | None = None, + provided_tools: Sequence[ToolTypes] | None = None, default_options: OptionsCoT | None = None, middleware: Sequence[MiddlewareTypes] | None = None, context_providers: Sequence[BaseContextProvider] | None = None, @@ -422,8 +416,8 @@ class AzureAIAgentsProvider(Generic[OptionsCoT]): def _merge_tools( self, agent_tools: Sequence[Any] | None, - provided_tools: Sequence[FunctionTool | MutableMapping[str, Any]] | None, - ) -> list[FunctionTool | dict[str, Any]]: + provided_tools: Sequence[ToolTypes] | None, + ) -> list[ToolTypes]: """Merge hosted tools from agent with user-provided function tools. Args: @@ -433,7 +427,7 @@ class AzureAIAgentsProvider(Generic[OptionsCoT]): Returns: Combined list of tools for the Agent. """ - merged: list[FunctionTool | dict[str, Any]] = [] + merged: list[ToolTypes] = [] # Convert hosted tools from agent definition hosted_tools = from_azure_ai_agent_tools(agent_tools) @@ -459,7 +453,7 @@ class AzureAIAgentsProvider(Generic[OptionsCoT]): def _validate_function_tools( self, agent_tools: Sequence[Any] | None, - provided_tools: Sequence[FunctionTool | MutableMapping[str, Any]] | None, + provided_tools: Sequence[ToolTypes] | None, ) -> None: """Validate that required function tools are provided. diff --git a/python/packages/azure-ai/agent_framework_azure_ai/_chat_client.py b/python/packages/azure-ai/agent_framework_azure_ai/_chat_client.py index b548885dae..c6028442fc 100644 --- a/python/packages/azure-ai/agent_framework_azure_ai/_chat_client.py +++ b/python/packages/azure-ai/agent_framework_azure_ai/_chat_client.py @@ -34,6 +34,7 @@ from agent_framework import ( UsageDetails, ) from agent_framework._settings import load_settings +from agent_framework._tools import ToolTypes from agent_framework.exceptions import ServiceInitializationError, ServiceInvalidRequestError, ServiceResponseException from agent_framework.observability import ChatTelemetryLayer from azure.ai.agents.aio import AgentsClient @@ -1428,11 +1429,7 @@ class AzureAIAgentClient( name: str | None = None, description: str | None = None, instructions: str | None = None, - tools: FunctionTool - | Callable[..., Any] - | MutableMapping[str, Any] - | Sequence[FunctionTool | Callable[..., Any] | MutableMapping[str, Any]] - | None = None, + tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None, default_options: AzureAIAgentOptionsT | Mapping[str, Any] | None = None, context_providers: Sequence[BaseContextProvider] | None = None, middleware: Sequence[MiddlewareTypes] | None = None, diff --git a/python/packages/azure-ai/agent_framework_azure_ai/_client.py b/python/packages/azure-ai/agent_framework_azure_ai/_client.py index 3b2f3c0e10..afbbf6cea3 100644 --- a/python/packages/azure-ai/agent_framework_azure_ai/_client.py +++ b/python/packages/azure-ai/agent_framework_azure_ai/_client.py @@ -5,7 +5,7 @@ from __future__ import annotations import json import logging import sys -from collections.abc import Callable, Mapping, MutableMapping, Sequence +from collections.abc import Callable, Mapping, Sequence from contextlib import suppress from typing import Any, ClassVar, Generic, Literal, TypedDict, TypeVar, cast @@ -22,6 +22,7 @@ from agent_framework import ( MiddlewareTypes, ) from agent_framework._settings import load_settings +from agent_framework._tools import ToolTypes from agent_framework.exceptions import ServiceInitializationError from agent_framework.observability import ChatTelemetryLayer from agent_framework.openai import OpenAIResponsesOptions @@ -880,11 +881,7 @@ class RawAzureAIClient(RawOpenAIResponsesClient[AzureAIClientOptionsT], Generic[ name: str | None = None, description: str | None = None, instructions: str | None = None, - tools: FunctionTool - | Callable[..., Any] - | MutableMapping[str, Any] - | Sequence[FunctionTool | Callable[..., Any] | MutableMapping[str, Any]] - | None = None, + tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None, default_options: AzureAIClientOptionsT | Mapping[str, Any] | None = None, context_providers: Sequence[BaseContextProvider] | None = None, middleware: Sequence[MiddlewareTypes] | None = None, diff --git a/python/packages/azure-ai/agent_framework_azure_ai/_project_provider.py b/python/packages/azure-ai/agent_framework_azure_ai/_project_provider.py index 1a71df1f79..b3f7b35147 100644 --- a/python/packages/azure-ai/agent_framework_azure_ai/_project_provider.py +++ b/python/packages/azure-ai/agent_framework_azure_ai/_project_provider.py @@ -17,6 +17,7 @@ from agent_framework import ( ) from agent_framework._mcp import MCPTool from agent_framework._settings import load_settings +from agent_framework._tools import ToolTypes from agent_framework.exceptions import ServiceInitializationError from azure.ai.projects.aio import AIProjectClient from azure.ai.projects.models import ( @@ -161,11 +162,7 @@ class AzureAIProjectAgentProvider(Generic[OptionsCoT]): model: str | None = None, instructions: str | None = None, description: str | None = None, - tools: FunctionTool - | Callable[..., Any] - | MutableMapping[str, Any] - | Sequence[FunctionTool | Callable[..., Any] | MutableMapping[str, Any]] - | None = None, + tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None, default_options: OptionsCoT | None = None, middleware: Sequence[MiddlewareTypes] | None = None, context_providers: Sequence[BaseContextProvider] | None = None, @@ -226,7 +223,7 @@ class AzureAIProjectAgentProvider(Generic[OptionsCoT]): for tool in normalized_tools: if isinstance(tool, MCPTool): mcp_tools.append(tool) - else: + elif isinstance(tool, (FunctionTool, MutableMapping)): non_mcp_tools.append(tool) # Connect MCP tools and discover their functions BEFORE creating the agent @@ -263,11 +260,7 @@ class AzureAIProjectAgentProvider(Generic[OptionsCoT]): *, name: str | None = None, reference: AgentReference | None = None, - tools: FunctionTool - | Callable[..., Any] - | MutableMapping[str, Any] - | Sequence[FunctionTool | Callable[..., Any] | MutableMapping[str, Any]] - | None = None, + tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None, default_options: OptionsCoT | None = None, middleware: Sequence[MiddlewareTypes] | None = None, context_providers: Sequence[BaseContextProvider] | None = None, @@ -323,11 +316,7 @@ class AzureAIProjectAgentProvider(Generic[OptionsCoT]): def as_agent( self, details: AgentVersionDetails, - tools: FunctionTool - | Callable[..., Any] - | MutableMapping[str, Any] - | Sequence[FunctionTool | Callable[..., Any] | MutableMapping[str, Any]] - | None = None, + tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None, default_options: OptionsCoT | None = None, middleware: Sequence[MiddlewareTypes] | None = None, context_providers: Sequence[BaseContextProvider] | None = None, @@ -367,7 +356,7 @@ class AzureAIProjectAgentProvider(Generic[OptionsCoT]): def _to_chat_agent_from_details( self, details: AgentVersionDetails, - provided_tools: Sequence[FunctionTool | MutableMapping[str, Any]] | None = None, + provided_tools: Sequence[ToolTypes] | None = None, default_options: OptionsCoT | None = None, middleware: Sequence[MiddlewareTypes] | None = None, context_providers: Sequence[BaseContextProvider] | None = None, @@ -415,8 +404,8 @@ class AzureAIProjectAgentProvider(Generic[OptionsCoT]): def _merge_tools( self, definition_tools: Sequence[Any] | None, - provided_tools: Sequence[FunctionTool | MutableMapping[str, Any]] | None, - ) -> list[FunctionTool | dict[str, Any]]: + provided_tools: Sequence[ToolTypes] | None, + ) -> list[ToolTypes]: """Merge hosted tools from definition with user-provided function tools. Args: @@ -426,7 +415,7 @@ class AzureAIProjectAgentProvider(Generic[OptionsCoT]): Returns: Combined list of tools for the Agent. """ - merged: list[FunctionTool | dict[str, Any]] = [] + merged: list[ToolTypes] = [] # Convert hosted tools from definition (MCP, code interpreter, file search, web search) # Function tools from the definition are skipped - we use user-provided implementations instead @@ -450,11 +439,7 @@ class AzureAIProjectAgentProvider(Generic[OptionsCoT]): def _validate_function_tools( self, agent_tools: Sequence[Any] | None, - provided_tools: FunctionTool - | Callable[..., Any] - | MutableMapping[str, Any] - | Sequence[FunctionTool | Callable[..., Any] | MutableMapping[str, Any]] - | None, + provided_tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None, ) -> None: """Validate that required function tools are provided.""" # Normalize and validate function tools diff --git a/python/packages/bedrock/AGENTS.md b/python/packages/bedrock/AGENTS.md index 69c0a1a692..0ab072ed24 100644 --- a/python/packages/bedrock/AGENTS.md +++ b/python/packages/bedrock/AGENTS.md @@ -12,7 +12,7 @@ Integration with AWS Bedrock for LLM inference. ## Usage ```python -from agent_framework_bedrock import BedrockChatClient +from agent_framework.amazon import BedrockChatClient client = BedrockChatClient(model_id="anthropic.claude-3-sonnet-20240229-v1:0") response = await client.get_response("Hello") @@ -21,5 +21,5 @@ response = await client.get_response("Hello") ## Import Path ```python -from agent_framework_bedrock import BedrockChatClient +from agent_framework.amazon import BedrockChatClient ``` diff --git a/python/packages/bedrock/README.md b/python/packages/bedrock/README.md index 6bcd9ff53a..10a3bd9f25 100644 --- a/python/packages/bedrock/README.md +++ b/python/packages/bedrock/README.md @@ -12,7 +12,7 @@ The Bedrock integration enables Microsoft Agent Framework applications to call A ### Basic Usage Example -See the [Bedrock sample script](samples/bedrock_sample.py) for a runnable end-to-end script that: +See the [Bedrock sample](../../samples/02-agents/providers/amazon/bedrock_chat_client.py) for a runnable end-to-end script that: - Loads credentials from the `BEDROCK_*` environment variables - Instantiates `BedrockChatClient` diff --git a/python/packages/bedrock/agent_framework_bedrock/_chat_client.py b/python/packages/bedrock/agent_framework_bedrock/_chat_client.py index 73bf9cc428..334cd25248 100644 --- a/python/packages/bedrock/agent_framework_bedrock/_chat_client.py +++ b/python/packages/bedrock/agent_framework_bedrock/_chat_client.py @@ -260,7 +260,7 @@ class BedrockChatClient( Examples: .. code-block:: python - from agent_framework.bedrock import BedrockChatClient + from agent_framework.amazon import BedrockChatClient # Basic usage with default credentials client = BedrockChatClient(model_id="") diff --git a/python/packages/bedrock/samples/__init__.py b/python/packages/bedrock/samples/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/python/packages/bedrock/samples/bedrock_sample.py b/python/packages/bedrock/samples/bedrock_sample.py deleted file mode 100644 index 188e6bf1da..0000000000 --- a/python/packages/bedrock/samples/bedrock_sample.py +++ /dev/null @@ -1,45 +0,0 @@ -# Copyright (c) Microsoft. All rights reserved. - -import asyncio -import logging - -from agent_framework import Agent, tool - -from agent_framework_bedrock import BedrockChatClient - - -@tool(approval_mode="never_require") -def get_weather(city: str) -> dict[str, str]: - """Return a mock forecast for the requested city.""" - normalized = city.strip() or "New York" - return {"city": normalized, "forecast": "72F and sunny"} - - -async def main() -> None: - """Run the Bedrock sample agent, invoke the weather tool, and log the response.""" - agent = Agent( - client=BedrockChatClient(), - instructions="You are a concise travel assistant.", - name="BedrockWeatherAgent", - tool_choice="auto", - tools=[get_weather], - ) - - response = await agent.run("Use the weather tool to check the forecast for new york.") - logging.info("\nAssistant reply:", response.text or "") - logging.info("\nConversation transcript:") - for message in response.messages: - for idx, content in enumerate(message.contents, start=1): - match content.type: - case "text": - logging.info(f" {idx}. text -> {content.text}") - case "function_call": - logging.info(f" {idx}. function_call ({content.name}) -> {content.arguments}") - case "function_result": - logging.info(f" {idx}. function_result ({content.call_id}) -> {content.result}") - case _: - logging.info(f" {idx}. {content.type}") - - -if __name__ == "__main__": - asyncio.run(main()) diff --git a/python/packages/claude/agent_framework_claude/_agent.py b/python/packages/claude/agent_framework_claude/_agent.py index 6ad805c5d5..204b364962 100644 --- a/python/packages/claude/agent_framework_claude/_agent.py +++ b/python/packages/claude/agent_framework_claude/_agent.py @@ -23,6 +23,7 @@ from agent_framework import ( normalize_messages, ) from agent_framework._settings import load_settings +from agent_framework._tools import ToolTypes from agent_framework._types import AgentRunInputs, normalize_tools from agent_framework.exceptions import ServiceException from claude_agent_sdk import ( @@ -217,12 +218,7 @@ class ClaudeAgent(BaseAgent, Generic[OptionsT]): description: str | None = None, context_providers: Sequence[BaseContextProvider] | None = None, middleware: Sequence[AgentMiddlewareTypes] | None = None, - tools: FunctionTool - | Callable[..., Any] - | MutableMapping[str, Any] - | str - | Sequence[FunctionTool | Callable[..., Any] | MutableMapping[str, Any] | str] - | None = None, + tools: ToolTypes | Callable[..., Any] | str | Sequence[ToolTypes | Callable[..., Any] | str] | None = None, default_options: OptionsT | MutableMapping[str, Any] | None = None, env_file_path: str | None = None, env_file_encoding: str | None = None, @@ -289,7 +285,7 @@ class ClaudeAgent(BaseAgent, Generic[OptionsT]): # Separate built-in tools (strings) from custom tools (callables/FunctionTool) self._builtin_tools: list[str] = [] - self._custom_tools: list[FunctionTool | MutableMapping[str, Any]] = [] + self._custom_tools: list[ToolTypes] = [] self._normalize_tools(tools) self._default_options = opts @@ -298,12 +294,7 @@ class ClaudeAgent(BaseAgent, Generic[OptionsT]): def _normalize_tools( self, - tools: FunctionTool - | Callable[..., Any] - | MutableMapping[str, Any] - | str - | Sequence[FunctionTool | Callable[..., Any] | MutableMapping[str, Any] | str] - | None, + tools: ToolTypes | Callable[..., Any] | str | Sequence[ToolTypes | Callable[..., Any] | str] | None, ) -> None: """Separate built-in tools (strings) from custom tools. @@ -316,10 +307,10 @@ class ClaudeAgent(BaseAgent, Generic[OptionsT]): # Normalize to sequence if isinstance(tools, str): tools_list: Sequence[Any] = [tools] - elif isinstance(tools, (FunctionTool, MutableMapping)) or callable(tools): - tools_list = [tools] - else: + elif isinstance(tools, Sequence): tools_list = list(tools) + else: + tools_list = [tools] for tool in tools_list: if isinstance(tool, str): @@ -457,7 +448,7 @@ class ClaudeAgent(BaseAgent, Generic[OptionsT]): def _prepare_tools( self, - tools: list[FunctionTool | MutableMapping[str, Any]], + tools: Sequence[ToolTypes], ) -> tuple[Any, list[str]]: """Convert Agent Framework tools to SDK MCP server. diff --git a/python/packages/core/agent_framework/__init__.py b/python/packages/core/agent_framework/__init__.py index bc53b10d16..2cc7d0d713 100644 --- a/python/packages/core/agent_framework/__init__.py +++ b/python/packages/core/agent_framework/__init__.py @@ -1,5 +1,13 @@ # Copyright (c) Microsoft. All rights reserved. +"""Public API surface for Agent Framework core. + +This module exposes the primary abstractions for agents, chat clients, tools, sessions, +middleware, observability, and workflows. Connector namespaces such as +``agent_framework.azure`` and ``agent_framework.anthropic`` provide provider-specific +integrations, many of which are lazy-loaded from optional packages. +""" + import importlib.metadata from typing import Final diff --git a/python/packages/core/agent_framework/_agents.py b/python/packages/core/agent_framework/_agents.py index 0b1b20bed8..d64bb8ced4 100644 --- a/python/packages/core/agent_framework/_agents.py +++ b/python/packages/core/agent_framework/_agents.py @@ -37,6 +37,8 @@ from ._sessions import AgentSession, BaseContextProvider, BaseHistoryProvider, I from ._tools import ( FunctionInvocationLayer, FunctionTool, + ToolTypes, + normalize_tools, ) from ._types import ( AgentResponse, @@ -614,12 +616,7 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc] id: str | None = None, name: str | None = None, description: str | None = None, - tools: FunctionTool - | Callable[..., Any] - | MutableMapping[str, Any] - | Any - | Sequence[FunctionTool | Callable[..., Any] | MutableMapping[str, Any] | Any] - | None = None, + tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None, default_options: OptionsCoT | None = None, context_providers: Sequence[BaseContextProvider] | None = None, **kwargs: Any, @@ -665,24 +662,14 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc] # Get tools from options or named parameter (named param takes precedence) tools_ = tools if tools is not None else opts.pop("tools", None) - tools_ = cast( - FunctionTool - | Callable[..., Any] - | MutableMapping[str, Any] - | list[FunctionTool | Callable[..., Any] | MutableMapping[str, Any]] - | None, - tools_, - ) # Handle instructions - named parameter takes precedence over options instructions_ = instructions if instructions is not None else opts.pop("instructions", None) # We ignore the MCP Servers here and store them separately, # we add their functions to the tools list at runtime - normalized_tools: list[FunctionTool | Callable[..., Any] | MutableMapping[str, Any]] = ( # type:ignore[reportUnknownVariableType] - [] if tools_ is None else tools_ if isinstance(tools_, list) else [tools_] # type: ignore[list-item] - ) - self.mcp_tools: list[MCPTool] = [tool for tool in normalized_tools if isinstance(tool, MCPTool)] # type: ignore[misc] + normalized_tools = normalize_tools(tools_) + self.mcp_tools: list[MCPTool] = [tool for tool in normalized_tools if isinstance(tool, MCPTool)] agent_tools = [tool for tool in normalized_tools if not isinstance(tool, MCPTool)] # Build chat options dict @@ -765,12 +752,7 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc] *, stream: Literal[False] = ..., session: AgentSession | None = None, - tools: FunctionTool - | Callable[..., Any] - | MutableMapping[str, Any] - | Any - | list[FunctionTool | Callable[..., Any] | MutableMapping[str, Any] | Any] - | None = None, + tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None, options: ChatOptions[ResponseModelBoundT], **kwargs: Any, ) -> Awaitable[AgentResponse[ResponseModelBoundT]]: ... @@ -782,12 +764,7 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc] *, stream: Literal[False] = ..., session: AgentSession | None = None, - tools: FunctionTool - | Callable[..., Any] - | MutableMapping[str, Any] - | Any - | list[FunctionTool | Callable[..., Any] | MutableMapping[str, Any] | Any] - | None = None, + tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None, options: OptionsCoT | ChatOptions[None] | None = None, **kwargs: Any, ) -> Awaitable[AgentResponse[Any]]: ... @@ -799,12 +776,7 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc] *, stream: Literal[True], session: AgentSession | None = None, - tools: FunctionTool - | Callable[..., Any] - | MutableMapping[str, Any] - | Any - | list[FunctionTool | Callable[..., Any] | MutableMapping[str, Any] | Any] - | None = None, + tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None, options: OptionsCoT | ChatOptions[Any] | None = None, **kwargs: Any, ) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ... @@ -815,12 +787,7 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc] *, stream: bool = False, session: AgentSession | None = None, - tools: FunctionTool - | Callable[..., Any] - | MutableMapping[str, Any] - | Any - | list[FunctionTool | Callable[..., Any] | MutableMapping[str, Any] | Any] - | None = None, + tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None, options: OptionsCoT | ChatOptions[Any] | None = None, **kwargs: Any, ) -> Awaitable[AgentResponse[Any]] | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: @@ -1000,12 +967,7 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc] *, messages: AgentRunInputs | None, session: AgentSession | None, - tools: FunctionTool - | Callable[..., Any] - | MutableMapping[str, Any] - | Any - | list[FunctionTool | Callable[..., Any] | MutableMapping[str, Any] | Any] - | None, + tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None, options: Mapping[str, Any] | None, kwargs: dict[str, Any], ) -> _RunContext: @@ -1035,9 +997,7 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc] ) # Normalize tools - normalized_tools: list[FunctionTool | Callable[..., Any] | MutableMapping[str, Any] | Any] = ( - [] if tools_ is None else tools_ if isinstance(tools_, list) else [tools_] - ) + normalized_tools = normalize_tools(tools_) agent_name = self._get_agent_name() # Resolve final tool list (runtime provided tools + local MCP server tools) @@ -1343,12 +1303,7 @@ class Agent( id: str | None = None, name: str | None = None, description: str | None = None, - tools: FunctionTool - | Callable[..., Any] - | MutableMapping[str, Any] - | Any - | Sequence[FunctionTool | Callable[..., Any] | MutableMapping[str, Any] | Any] - | None = None, + tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None, default_options: OptionsCoT | None = None, context_providers: Sequence[BaseContextProvider] | None = None, middleware: Sequence[MiddlewareTypes] | None = None, diff --git a/python/packages/core/agent_framework/_clients.py b/python/packages/core/agent_framework/_clients.py index c8956a67e3..b407be11cf 100644 --- a/python/packages/core/agent_framework/_clients.py +++ b/python/packages/core/agent_framework/_clients.py @@ -10,7 +10,6 @@ from collections.abc import ( Awaitable, Callable, Mapping, - MutableMapping, Sequence, ) from typing import ( @@ -31,7 +30,7 @@ from pydantic import BaseModel from ._serialization import SerializationMixin from ._tools import ( FunctionInvocationConfiguration, - FunctionTool, + ToolTypes, ) from ._types import ( ChatResponse, @@ -436,11 +435,7 @@ class BaseChatClient(SerializationMixin, ABC, Generic[OptionsCoT]): name: str | None = None, description: str | None = None, instructions: str | None = None, - tools: FunctionTool - | Callable[..., Any] - | MutableMapping[str, Any] - | Sequence[FunctionTool | Callable[..., Any] | MutableMapping[str, Any]] - | None = None, + tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None, default_options: OptionsCoT | Mapping[str, Any] | None = None, context_providers: Sequence[Any] | None = None, middleware: Sequence[MiddlewareTypes] | None = None, diff --git a/python/packages/core/agent_framework/_tools.py b/python/packages/core/agent_framework/_tools.py index 8f941b41e0..b5d110aad1 100644 --- a/python/packages/core/agent_framework/_tools.py +++ b/python/packages/core/agent_framework/_tools.py @@ -12,7 +12,6 @@ from collections.abc import ( Awaitable, Callable, Mapping, - MutableMapping, Sequence, ) from functools import partial, wraps @@ -25,6 +24,7 @@ from typing import ( Final, Generic, Literal, + TypeAlias, TypedDict, Union, get_args, @@ -58,6 +58,7 @@ else: if TYPE_CHECKING: from ._clients import SupportsChatGetResponse + from ._mcp import MCPTool from ._middleware import FunctionMiddlewarePipeline, FunctionMiddlewareTypes from ._types import ( ChatOptions, @@ -69,6 +70,8 @@ if TYPE_CHECKING: ) ResponseModelBoundT = TypeVar("ResponseModelBoundT", bound=BaseModel) +else: + MCPTool = Any # type: ignore[assignment,misc] logger = logging.getLogger("agent_framework") @@ -506,9 +509,7 @@ class FunctionTool(SerializationMixin): if OBSERVABILITY_SETTINGS.SENSITIVE_DATA_ENABLED: # type: ignore[name-defined] attributes.update({ OtelAttr.TOOL_ARGUMENTS: ( - json.dumps(serializable_kwargs, default=str, ensure_ascii=False) - if serializable_kwargs - else "None" + json.dumps(serializable_kwargs, default=str, ensure_ascii=False) if serializable_kwargs else "None" ) }) with get_function_span(attributes=attributes) as span: @@ -623,14 +624,46 @@ class FunctionTool(SerializationMixin): return as_dict +ToolTypes: TypeAlias = FunctionTool | MCPTool | Mapping[str, Any] | Any + + +def normalize_tools( + tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None, +) -> list[ToolTypes]: + """Normalize tool inputs while preserving non-callable tool objects. + + Args: + tools: A single tool or sequence of tools. + + Returns: + A normalized list where callable inputs are converted to ``FunctionTool`` + using :func:`tool`, and existing tool objects are passed through unchanged. + """ + if not tools: + return [] + + tool_items = ( + list(tools) + if isinstance(tools, Sequence) and not isinstance(tools, (str, bytes, bytearray, Mapping)) + else [tools] + ) + from ._mcp import MCPTool + + normalized: list[ToolTypes] = [] + for tool_item in tool_items: + # check known types, these are also callable, so we need to do that first + if isinstance(tool_item, (FunctionTool, Mapping, MCPTool)): + normalized.append(tool_item) + continue + if callable(tool_item): + normalized.append(tool(tool_item)) + continue + normalized.append(tool_item) + return normalized + + def _tools_to_dict( - tools: ( - FunctionTool - | Callable[..., Any] - | MutableMapping[str, Any] - | Sequence[FunctionTool | Callable[..., Any] | MutableMapping[str, Any]] - | None - ), + tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None, ) -> list[str | dict[str, Any]] | None: """Parse the tools to a dict. @@ -640,32 +673,20 @@ def _tools_to_dict( Returns: A list of tool specifications as dictionaries, or None if no tools provided. """ - if not tools: - return None - if not isinstance(tools, list): - if isinstance(tools, FunctionTool): - return [tools.to_json_schema_spec()] - if isinstance(tools, SerializationMixin): - return [tools.to_dict()] - if isinstance(tools, dict): - return [tools] - if callable(tools): - return [tool(tools).to_json_schema_spec()] - logger.warning("Can't parse tool.") + normalized_tools = normalize_tools(tools) + if not normalized_tools: return None + results: list[str | dict[str, Any]] = [] - for tool_item in tools: + for tool_item in normalized_tools: if isinstance(tool_item, FunctionTool): results.append(tool_item.to_json_schema_spec()) continue if isinstance(tool_item, SerializationMixin): results.append(tool_item.to_dict()) continue - if isinstance(tool_item, dict): - results.append(tool_item) - continue - if callable(tool_item): - results.append(tool(tool_item).to_json_schema_spec()) + if isinstance(tool_item, Mapping): + results.append(dict(tool_item)) continue logger.warning("Can't parse tool.") return results @@ -1430,20 +1451,12 @@ async def _auto_invoke_function( def _get_tool_map( - tools: FunctionTool - | Callable[..., Any] - | MutableMapping[str, Any] - | Sequence[FunctionTool | Callable[..., Any] | MutableMapping[str, Any]], + tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]], ) -> dict[str, FunctionTool]: tool_list: dict[str, FunctionTool] = {} - for tool_item in tools if isinstance(tools, list) else [tools]: + for tool_item in normalize_tools(tools): if isinstance(tool_item, FunctionTool): tool_list[tool_item.name] = tool_item - continue - if callable(tool_item): - # Convert to AITool if it's a function or callable - ai_tool = tool(tool_item) - tool_list[ai_tool.name] = ai_tool return tool_list @@ -1451,10 +1464,7 @@ async def _try_execute_function_calls( custom_args: dict[str, Any], attempt_idx: int, function_calls: Sequence[Content], - tools: FunctionTool - | Callable[..., Any] - | MutableMapping[str, Any] - | Sequence[FunctionTool | Callable[..., Any] | MutableMapping[str, Any]], + tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]], config: FunctionInvocationConfiguration, middleware_pipeline: Any = None, # Optional MiddlewarePipeline to avoid circular imports ) -> tuple[Sequence[Content], bool]: @@ -1633,15 +1643,16 @@ async def _ensure_response_stream( return stream -def _extract_tools(options: dict[str, Any] | None) -> Any: +def _extract_tools( + options: dict[str, Any] | None, +) -> ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None: """Extract tools from options dict. Args: options: The options dict containing chat options. Returns: - FunctionTool | Callable[..., Any] | MutableMapping[str, Any] | - Sequence[FunctionTool | Callable[..., Any] | MutableMapping[str, Any]] | None + ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None """ if options and isinstance(options, dict): return options.get("tools") @@ -1996,6 +2007,11 @@ class FunctionInvocationLayer(Generic[OptionsCoT]): # Remove additional_function_arguments from options passed to underlying chat client # It's for tool invocation only and not recognized by chat service APIs mutable_options.pop("additional_function_arguments", None) + # Support tools passed via kwargs in direct client.get_response(...) calls. + if "tools" in filtered_kwargs: + if mutable_options.get("tools") is None: + mutable_options["tools"] = filtered_kwargs["tools"] + filtered_kwargs.pop("tools", None) if not stream: diff --git a/python/packages/core/agent_framework/_types.py b/python/packages/core/agent_framework/_types.py index a3786e8838..fb43b4e9e4 100644 --- a/python/packages/core/agent_framework/_types.py +++ b/python/packages/core/agent_framework/_types.py @@ -15,7 +15,8 @@ from typing import TYPE_CHECKING, Any, ClassVar, Final, Generic, Literal, NewTyp from pydantic import BaseModel from ._serialization import SerializationMixin -from ._tools import FunctionTool, tool +from ._tools import ToolTypes +from ._tools import normalize_tools as _normalize_tools from .exceptions import AdditionItemMismatch, ContentError if sys.version_info >= (3, 13): @@ -2871,10 +2872,9 @@ class _ChatOptionsBase(TypedDict, total=False): # Tool configuration (forward reference to avoid circular import) tools: ( - FunctionTool + ToolTypes | Callable[..., Any] - | MutableMapping[str, Any] - | Sequence[FunctionTool | Callable[..., Any] | MutableMapping[str, Any]] + | Sequence[ToolTypes | Callable[..., Any]] | None ) tool_choice: ToolMode | Literal["auto", "required", "none"] @@ -2963,18 +2963,11 @@ async def validate_chat_options(options: dict[str, Any]) -> dict[str, Any]: def normalize_tools( - tools: ( - FunctionTool - | Callable[..., Any] - | MutableMapping[str, Any] - | Sequence[FunctionTool | Callable[..., Any] | MutableMapping[str, Any]] - | None - ), -) -> list[FunctionTool | MutableMapping[str, Any]]: + tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None, +) -> list[ToolTypes]: """Normalize tools into a list. - Converts callables to FunctionTool objects and ensures all tools are either - FunctionTool instances or MutableMappings. + Converts callables to FunctionTool objects and preserves existing tool objects. Args: tools: Tools to normalize - can be a single tool, callable, or sequence. @@ -2999,37 +2992,16 @@ def normalize_tools( # List of tools tools = normalize_tools([my_tool, another_tool]) """ - final_tools: list[FunctionTool | MutableMapping[str, Any]] = [] - if not tools: - return final_tools - if not isinstance(tools, Sequence) or isinstance(tools, (str, MutableMapping)): - # Single tool (not a sequence, or is a mapping which shouldn't be treated as sequence) - if not isinstance(tools, (FunctionTool, MutableMapping)): - return [tool(tools)] - return [tools] - for tool_item in tools: - if isinstance(tool_item, (FunctionTool, MutableMapping)): - final_tools.append(tool_item) - else: - # Convert callable to FunctionTool - final_tools.append(tool(tool_item)) - return final_tools + return _normalize_tools(tools) async def validate_tools( - tools: ( - FunctionTool - | Callable[..., Any] - | MutableMapping[str, Any] - | Sequence[FunctionTool | Callable[..., Any] | MutableMapping[str, Any]] - | None - ), -) -> list[FunctionTool | MutableMapping[str, Any]]: + tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None, +) -> list[ToolTypes]: """Validate and normalize tools into a list. Converts callables to FunctionTool objects, expands MCP tools to their constituent - functions (connecting them if needed), and ensures all tools are either FunctionTool - instances or MutableMappings. + functions (connecting them if needed), while preserving non-callable tool objects. Args: tools: Tools to validate - can be a single tool, callable, or sequence. @@ -3058,7 +3030,7 @@ async def validate_tools( normalized = normalize_tools(tools) # Handle MCP tool expansion (async-only) - final_tools: list[FunctionTool | MutableMapping[str, Any]] = [] + final_tools: list[ToolTypes] = [] for tool_ in normalized: # Import MCPTool here to avoid circular imports from ._mcp import MCPTool @@ -3076,20 +3048,21 @@ async def validate_tools( def validate_tool_mode( tool_choice: ToolMode | Literal["auto", "required", "none"] | None, -) -> ToolMode: +) -> ToolMode | None: """Validate and normalize tool_choice to a ToolMode dict. Args: tool_choice: The tool choice value to validate. Returns: - A ToolMode dict (contains keys: "mode", and optionally "required_function_name"). + A ToolMode dict (contains keys: "mode", and optionally + "required_function_name"), or ``None`` when not provided. Raises: ContentError: If the tool_choice string is invalid. """ - if not tool_choice: - return {"mode": "none"} + if tool_choice is None: + return None if isinstance(tool_choice, str): if tool_choice not in ("auto", "required", "none"): raise ContentError(f"Invalid tool choice: {tool_choice}") diff --git a/python/packages/core/agent_framework/_workflows/__init__.py b/python/packages/core/agent_framework/_workflows/__init__.py index 3eb65335c9..2dfa029840 100644 --- a/python/packages/core/agent_framework/_workflows/__init__.py +++ b/python/packages/core/agent_framework/_workflows/__init__.py @@ -1,5 +1,18 @@ # Copyright (c) Microsoft. All rights reserved. +"""Workflow namespace for built-in Agent Framework orchestration primitives. + +This module re-exports objects from workflow implementation modules under +``agent_framework._workflows``. + +Supported classes include: +- Workflow +- WorkflowBuilder +- AgentExecutor +- Runner +- WorkflowExecutor +""" + from ._agent import WorkflowAgent from ._agent_executor import ( AgentExecutor, diff --git a/python/packages/core/agent_framework/a2a/__init__.py b/python/packages/core/agent_framework/a2a/__init__.py index f06ee08a0b..7c7de63456 100644 --- a/python/packages/core/agent_framework/a2a/__init__.py +++ b/python/packages/core/agent_framework/a2a/__init__.py @@ -1,11 +1,20 @@ # Copyright (c) Microsoft. All rights reserved. +"""A2A integration namespace for optional Agent Framework connectors. + +This module lazily re-exports objects from: +- ``agent-framework-a2a`` + +Supported classes: +- A2AAgent +""" + import importlib from typing import Any IMPORT_PATH = "agent_framework_a2a" PACKAGE_NAME = "agent-framework-a2a" -_IMPORTS = ["__version__", "A2AAgent"] +_IMPORTS = ["A2AAgent"] def __getattr__(name: str) -> Any: diff --git a/python/packages/core/agent_framework/a2a/__init__.pyi b/python/packages/core/agent_framework/a2a/__init__.pyi index 8121fc0eb7..5a54bb22a9 100644 --- a/python/packages/core/agent_framework/a2a/__init__.pyi +++ b/python/packages/core/agent_framework/a2a/__init__.pyi @@ -2,10 +2,8 @@ from agent_framework_a2a import ( A2AAgent, - __version__, ) __all__ = [ "A2AAgent", - "__version__", ] diff --git a/python/packages/core/agent_framework/ag_ui/__init__.py b/python/packages/core/agent_framework/ag_ui/__init__.py index b469bb8a60..0f6991c9b9 100644 --- a/python/packages/core/agent_framework/ag_ui/__init__.py +++ b/python/packages/core/agent_framework/ag_ui/__init__.py @@ -1,12 +1,24 @@ # Copyright (c) Microsoft. All rights reserved. +"""AG-UI integration namespace for optional Agent Framework connectors. + +This module lazily re-exports objects from: +- ``agent-framework-ag-ui`` + +Supported classes and functions: +- AgentFrameworkAgent +- AGUIChatClient +- AGUIEventConverter +- AGUIHttpService +- add_agent_framework_fastapi_endpoint +""" + import importlib from typing import Any IMPORT_PATH = "agent_framework_ag_ui" PACKAGE_NAME = "agent-framework-ag-ui" _IMPORTS = [ - "__version__", "AgentFrameworkAgent", "add_agent_framework_fastapi_endpoint", "AGUIChatClient", diff --git a/python/packages/core/agent_framework/ag_ui/__init__.pyi b/python/packages/core/agent_framework/ag_ui/__init__.pyi index d7b6acafec..6e1948ce86 100644 --- a/python/packages/core/agent_framework/ag_ui/__init__.pyi +++ b/python/packages/core/agent_framework/ag_ui/__init__.pyi @@ -5,7 +5,6 @@ from agent_framework_ag_ui import ( AGUIChatClient, AGUIEventConverter, AGUIHttpService, - __version__, add_agent_framework_fastapi_endpoint, ) @@ -14,6 +13,5 @@ __all__ = [ "AGUIEventConverter", "AGUIHttpService", "AgentFrameworkAgent", - "__version__", "add_agent_framework_fastapi_endpoint", ] diff --git a/python/packages/core/agent_framework/amazon/__init__.py b/python/packages/core/agent_framework/amazon/__init__.py new file mode 100644 index 0000000000..42324acf96 --- /dev/null +++ b/python/packages/core/agent_framework/amazon/__init__.py @@ -0,0 +1,35 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Amazon Bedrock integration namespace for optional Agent Framework connectors. + +This module lazily re-exports objects from: +- ``agent-framework-bedrock`` + +Supported classes: +- BedrockChatClient +- BedrockChatOptions +- BedrockGuardrailConfig +- BedrockSettings +""" + +import importlib +from typing import Any + +IMPORT_PATH = "agent_framework_bedrock" +PACKAGE_NAME = "agent-framework-bedrock" +_IMPORTS = ["BedrockChatClient", "BedrockChatOptions", "BedrockGuardrailConfig", "BedrockSettings"] + + +def __getattr__(name: str) -> Any: + if name in _IMPORTS: + try: + return getattr(importlib.import_module(IMPORT_PATH), name) + except ModuleNotFoundError as exc: + raise ModuleNotFoundError( + f"The '{PACKAGE_NAME}' package is not installed, please do `pip install {PACKAGE_NAME}`" + ) from exc + raise AttributeError(f"Module {IMPORT_PATH} has no attribute {name}.") + + +def __dir__() -> list[str]: + return _IMPORTS diff --git a/python/packages/core/agent_framework/amazon/__init__.pyi b/python/packages/core/agent_framework/amazon/__init__.pyi new file mode 100644 index 0000000000..c691334da5 --- /dev/null +++ b/python/packages/core/agent_framework/amazon/__init__.pyi @@ -0,0 +1,15 @@ +# Copyright (c) Microsoft. All rights reserved. + +from agent_framework_bedrock import ( + BedrockChatClient, + BedrockChatOptions, + BedrockGuardrailConfig, + BedrockSettings, +) + +__all__ = [ + "BedrockChatClient", + "BedrockChatOptions", + "BedrockGuardrailConfig", + "BedrockSettings", +] diff --git a/python/packages/core/agent_framework/anthropic/__init__.py b/python/packages/core/agent_framework/anthropic/__init__.py index ea03e6cdf0..242554cf16 100644 --- a/python/packages/core/agent_framework/anthropic/__init__.py +++ b/python/packages/core/agent_framework/anthropic/__init__.py @@ -1,23 +1,40 @@ # Copyright (c) Microsoft. All rights reserved. +"""Anthropic integration namespace for optional Agent Framework connectors. + +This module lazily re-exports objects from: +- ``agent-framework-anthropic`` +- ``agent-framework-claude`` + +Supported classes: +- AnthropicClient +- AnthropicChatOptions +- ClaudeAgent +- ClaudeAgentOptions +""" + import importlib from typing import Any -IMPORT_PATH = "agent_framework_anthropic" -PACKAGE_NAME = "agent-framework-anthropic" -_IMPORTS = ["__version__", "AnthropicClient", "AnthropicChatOptions"] +_IMPORTS: dict[str, tuple[str, str]] = { + "AnthropicClient": ("agent_framework_anthropic", "agent-framework-anthropic"), + "AnthropicChatOptions": ("agent_framework_anthropic", "agent-framework-anthropic"), + "ClaudeAgent": ("agent_framework_claude", "agent-framework-claude"), + "ClaudeAgentOptions": ("agent_framework_claude", "agent-framework-claude"), +} def __getattr__(name: str) -> Any: if name in _IMPORTS: + import_path, package_name = _IMPORTS[name] try: - return getattr(importlib.import_module(IMPORT_PATH), name) + return getattr(importlib.import_module(import_path), name) except ModuleNotFoundError as exc: raise ModuleNotFoundError( - f"The '{PACKAGE_NAME}' package is not installed, please do `pip install {PACKAGE_NAME}`" + f"The '{package_name}' package is not installed, please do `pip install {package_name}`" ) from exc - raise AttributeError(f"Module {IMPORT_PATH} has no attribute {name}.") + raise AttributeError(f"Module `anthropic` has no attribute {name}.") def __dir__() -> list[str]: - return _IMPORTS + return list(_IMPORTS.keys()) diff --git a/python/packages/core/agent_framework/anthropic/__init__.pyi b/python/packages/core/agent_framework/anthropic/__init__.pyi index 3d790ebb07..29d70f62b8 100644 --- a/python/packages/core/agent_framework/anthropic/__init__.pyi +++ b/python/packages/core/agent_framework/anthropic/__init__.pyi @@ -3,11 +3,12 @@ from agent_framework_anthropic import ( AnthropicChatOptions, AnthropicClient, - __version__, ) +from agent_framework_claude import ClaudeAgent, ClaudeAgentOptions __all__ = [ "AnthropicChatOptions", "AnthropicClient", - "__version__", + "ClaudeAgent", + "ClaudeAgentOptions", ] diff --git a/python/packages/core/agent_framework/azure/__init__.py b/python/packages/core/agent_framework/azure/__init__.py index 93d7dc1e0d..7ac132bf3a 100644 --- a/python/packages/core/agent_framework/azure/__init__.py +++ b/python/packages/core/agent_framework/azure/__init__.py @@ -1,5 +1,19 @@ # Copyright (c) Microsoft. All rights reserved. +"""Azure integration namespace for optional Agent Framework connectors. + +This module lazily re-exports objects from optional Azure connector packages and +built-in core Azure OpenAI modules. + +Supported classes include: +- AzureAIClient +- AzureAIAgentClient +- AzureOpenAIChatClient +- AzureOpenAIResponsesClient +- AzureAISearchContextProvider +- DurableAIAgent +""" + import importlib from typing import Any diff --git a/python/packages/core/agent_framework/chatkit/__init__.py b/python/packages/core/agent_framework/chatkit/__init__.py index 024454be5c..7363b82361 100644 --- a/python/packages/core/agent_framework/chatkit/__init__.py +++ b/python/packages/core/agent_framework/chatkit/__init__.py @@ -1,11 +1,22 @@ # Copyright (c) Microsoft. All rights reserved. +"""ChatKit integration namespace for optional Agent Framework connectors. + +This module lazily re-exports objects from: +- ``agent-framework-chatkit`` + +Supported classes and functions: +- ThreadItemConverter +- simple_to_agent_input +- stream_agent_response +""" + import importlib from typing import Any IMPORT_PATH = "agent_framework_chatkit" PACKAGE_NAME = "agent-framework-chatkit" -_IMPORTS = ["__version__", "ThreadItemConverter", "simple_to_agent_input", "stream_agent_response"] +_IMPORTS = ["ThreadItemConverter", "simple_to_agent_input", "stream_agent_response"] def __getattr__(name: str) -> Any: diff --git a/python/packages/core/agent_framework/chatkit/__init__.pyi b/python/packages/core/agent_framework/chatkit/__init__.pyi index 2fba862a1b..0fdeff2523 100644 --- a/python/packages/core/agent_framework/chatkit/__init__.pyi +++ b/python/packages/core/agent_framework/chatkit/__init__.pyi @@ -2,14 +2,12 @@ from agent_framework_chatkit import ( ThreadItemConverter, - __version__, simple_to_agent_input, stream_agent_response, ) __all__ = [ "ThreadItemConverter", - "__version__", "simple_to_agent_input", "stream_agent_response", ] diff --git a/python/packages/core/agent_framework/declarative/__init__.py b/python/packages/core/agent_framework/declarative/__init__.py index 4ad557c0f7..a4fa5ba9f8 100644 --- a/python/packages/core/agent_framework/declarative/__init__.py +++ b/python/packages/core/agent_framework/declarative/__init__.py @@ -1,12 +1,23 @@ # Copyright (c) Microsoft. All rights reserved. +"""Declarative integration namespace for optional Agent Framework connectors. + +This module lazily re-exports objects from: +- ``agent-framework-declarative`` + +Supported classes include: +- AgentFactory +- WorkflowFactory +- ExternalInputRequest +- ExternalInputResponse +""" + import importlib from typing import Any IMPORT_PATH = "agent_framework_declarative" PACKAGE_NAME = "agent-framework-declarative" _IMPORTS = [ - "__version__", "AgentFactory", "AgentExternalInputRequest", "AgentExternalInputResponse", diff --git a/python/packages/core/agent_framework/declarative/__init__.pyi b/python/packages/core/agent_framework/declarative/__init__.pyi index 8d2b717c99..214bb132ab 100644 --- a/python/packages/core/agent_framework/declarative/__init__.pyi +++ b/python/packages/core/agent_framework/declarative/__init__.pyi @@ -13,7 +13,6 @@ from agent_framework_declarative import ( ProviderTypeMapping, WorkflowFactory, WorkflowState, - __version__, ) __all__ = [ @@ -29,5 +28,4 @@ __all__ = [ "ProviderTypeMapping", "WorkflowFactory", "WorkflowState", - "__version__", ] diff --git a/python/packages/core/agent_framework/devui/__init__.py b/python/packages/core/agent_framework/devui/__init__.py index cd18b0c5da..ce8215f941 100644 --- a/python/packages/core/agent_framework/devui/__init__.py +++ b/python/packages/core/agent_framework/devui/__init__.py @@ -1,5 +1,19 @@ # Copyright (c) Microsoft. All rights reserved. +"""DevUI integration namespace for optional Agent Framework connectors. + +This module lazily re-exports objects from: +- ``agent-framework-devui`` + +Supported classes and functions include: +- DevServer +- AgentFrameworkRequest +- DiscoveryResponse +- ResponseStreamEvent +- serve +- main +""" + import importlib from typing import Any @@ -16,7 +30,6 @@ _IMPORTS = [ "main", "register_cleanup", "serve", - "__version__", ] diff --git a/python/packages/core/agent_framework/devui/__init__.pyi b/python/packages/core/agent_framework/devui/__init__.pyi index 9396af54bb..d6828ce96f 100644 --- a/python/packages/core/agent_framework/devui/__init__.pyi +++ b/python/packages/core/agent_framework/devui/__init__.pyi @@ -8,7 +8,6 @@ from agent_framework_devui import ( OpenAIError, OpenAIResponse, ResponseStreamEvent, - __version__, main, register_cleanup, serve, @@ -22,7 +21,6 @@ __all__ = [ "OpenAIError", "OpenAIResponse", "ResponseStreamEvent", - "__version__", "main", "register_cleanup", "serve", diff --git a/python/packages/core/agent_framework/exceptions.py b/python/packages/core/agent_framework/exceptions.py index 21e50e571e..5e9cba7572 100644 --- a/python/packages/core/agent_framework/exceptions.py +++ b/python/packages/core/agent_framework/exceptions.py @@ -1,5 +1,7 @@ # Copyright (c) Microsoft. All rights reserved. +"""Exception hierarchy used across Agent Framework core and connectors.""" + import logging from typing import Any, Literal diff --git a/python/packages/core/agent_framework/github/__init__.py b/python/packages/core/agent_framework/github/__init__.py index f763ebf6d4..831a838c0d 100644 --- a/python/packages/core/agent_framework/github/__init__.py +++ b/python/packages/core/agent_framework/github/__init__.py @@ -1,5 +1,16 @@ # Copyright (c) Microsoft. All rights reserved. +"""GitHub integration namespace for optional Agent Framework connectors. + +This module lazily re-exports objects from: +- ``agent-framework-github-copilot`` + +Supported classes: +- GitHubCopilotAgent +- GitHubCopilotOptions +- GitHubCopilotSettings +""" + import importlib from typing import Any @@ -7,7 +18,6 @@ _IMPORTS: dict[str, tuple[str, str]] = { "GitHubCopilotAgent": ("agent_framework_github_copilot", "agent-framework-github-copilot"), "GitHubCopilotOptions": ("agent_framework_github_copilot", "agent-framework-github-copilot"), "GitHubCopilotSettings": ("agent_framework_github_copilot", "agent-framework-github-copilot"), - "__version__": ("agent_framework_github_copilot", "agent-framework-github-copilot"), } diff --git a/python/packages/core/agent_framework/github/__init__.pyi b/python/packages/core/agent_framework/github/__init__.pyi index 60a32ed620..567ab9490d 100644 --- a/python/packages/core/agent_framework/github/__init__.pyi +++ b/python/packages/core/agent_framework/github/__init__.pyi @@ -4,12 +4,10 @@ from agent_framework_github_copilot import ( GitHubCopilotAgent, GitHubCopilotOptions, GitHubCopilotSettings, - __version__, ) __all__ = [ "GitHubCopilotAgent", "GitHubCopilotOptions", "GitHubCopilotSettings", - "__version__", ] diff --git a/python/packages/core/agent_framework/lab/__init__.py b/python/packages/core/agent_framework/lab/__init__.py index 0d8d8fbaa4..80ed9251a8 100644 --- a/python/packages/core/agent_framework/lab/__init__.py +++ b/python/packages/core/agent_framework/lab/__init__.py @@ -1,4 +1,10 @@ # Copyright (c) Microsoft. All rights reserved. +"""Lab namespace package for experimental Agent Framework integrations. + +This module extends the package path so experimental lab integrations can be +distributed in separate packages under the ``agent_framework.lab`` namespace. +""" + # This makes agent_framework.lab a namespace package __path__ = __import__("pkgutil").extend_path(__path__, __name__) diff --git a/python/packages/core/agent_framework/mem0/__init__.py b/python/packages/core/agent_framework/mem0/__init__.py index dddc742ef0..56714d2224 100644 --- a/python/packages/core/agent_framework/mem0/__init__.py +++ b/python/packages/core/agent_framework/mem0/__init__.py @@ -1,11 +1,20 @@ # Copyright (c) Microsoft. All rights reserved. +"""Mem0 integration namespace for optional Agent Framework connectors. + +This module lazily re-exports objects from: +- ``agent-framework-mem0`` + +Supported classes: +- Mem0ContextProvider +""" + import importlib from typing import Any IMPORT_PATH = "agent_framework_mem0" PACKAGE_NAME = "agent-framework-mem0" -_IMPORTS = ["__version__", "Mem0ContextProvider"] +_IMPORTS = ["Mem0ContextProvider"] def __getattr__(name: str) -> Any: diff --git a/python/packages/core/agent_framework/mem0/__init__.pyi b/python/packages/core/agent_framework/mem0/__init__.pyi index 18ee3bf2bd..20a68e1b06 100644 --- a/python/packages/core/agent_framework/mem0/__init__.pyi +++ b/python/packages/core/agent_framework/mem0/__init__.pyi @@ -2,10 +2,8 @@ from agent_framework_mem0 import ( Mem0ContextProvider, - __version__, ) __all__ = [ "Mem0ContextProvider", - "__version__", ] diff --git a/python/packages/core/agent_framework/microsoft/__init__.py b/python/packages/core/agent_framework/microsoft/__init__.py index 689faf6fd7..cf61e518d9 100644 --- a/python/packages/core/agent_framework/microsoft/__init__.py +++ b/python/packages/core/agent_framework/microsoft/__init__.py @@ -1,11 +1,36 @@ # Copyright (c) Microsoft. All rights reserved. +"""Microsoft integration namespace for optional Agent Framework connectors. + +This module lazily re-exports objects from: +- ``agent-framework-copilotstudio`` +- ``agent-framework-purview`` +- ``agent-framework-foundry-local`` + +Supported classes: +- CopilotStudioAgent +- PurviewPolicyMiddleware +- PurviewChatPolicyMiddleware +- PurviewSettings +- PurviewAppLocation +- PurviewLocationType +- PurviewAuthenticationError +- PurviewPaymentRequiredError +- PurviewRateLimitError +- PurviewRequestError +- PurviewServiceError +- CacheProvider +- FoundryLocalChatOptions +- FoundryLocalClient +- FoundryLocalSettings + +""" + import importlib from typing import Any _IMPORTS: dict[str, tuple[str, str]] = { "CopilotStudioAgent": ("agent_framework_copilotstudio", "agent-framework-copilotstudio"), - "__version__": ("agent_framework_copilotstudio", "agent-framework-copilotstudio"), "acquire_token": ("agent_framework_copilotstudio", "agent-framework-copilotstudio"), "PurviewPolicyMiddleware": ("agent_framework_purview", "agent-framework-purview"), "PurviewChatPolicyMiddleware": ("agent_framework_purview", "agent-framework-purview"), @@ -18,6 +43,9 @@ _IMPORTS: dict[str, tuple[str, str]] = { "PurviewRequestError": ("agent_framework_purview", "agent-framework-purview"), "PurviewServiceError": ("agent_framework_purview", "agent-framework-purview"), "CacheProvider": ("agent_framework_purview", "agent-framework-purview"), + "FoundryLocalChatOptions": ("agent_framework_foundry_local", "agent-framework-foundry-local"), + "FoundryLocalClient": ("agent_framework_foundry_local", "agent-framework-foundry-local"), + "FoundryLocalSettings": ("agent_framework_foundry_local", "agent-framework-foundry-local"), } diff --git a/python/packages/core/agent_framework/microsoft/__init__.pyi b/python/packages/core/agent_framework/microsoft/__init__.pyi index 2d2ec42d4d..be3abe523c 100644 --- a/python/packages/core/agent_framework/microsoft/__init__.pyi +++ b/python/packages/core/agent_framework/microsoft/__init__.pyi @@ -2,9 +2,13 @@ from agent_framework_copilotstudio import ( CopilotStudioAgent, - __version__, acquire_token, ) +from agent_framework_foundry_local import ( + FoundryLocalChatOptions, + FoundryLocalClient, + FoundryLocalSettings, +) from agent_framework_purview import ( CacheProvider, PurviewAppLocation, @@ -22,6 +26,9 @@ from agent_framework_purview import ( __all__ = [ "CacheProvider", "CopilotStudioAgent", + "FoundryLocalChatOptions", + "FoundryLocalClient", + "FoundryLocalSettings", "PurviewAppLocation", "PurviewAuthenticationError", "PurviewChatPolicyMiddleware", @@ -32,6 +39,5 @@ __all__ = [ "PurviewRequestError", "PurviewServiceError", "PurviewSettings", - "__version__", "acquire_token", ] diff --git a/python/packages/core/agent_framework/observability.py b/python/packages/core/agent_framework/observability.py index 7973a781a4..0127a101df 100644 --- a/python/packages/core/agent_framework/observability.py +++ b/python/packages/core/agent_framework/observability.py @@ -1,5 +1,16 @@ # Copyright (c) Microsoft. All rights reserved. +"""Observability and OpenTelemetry helpers for Agent Framework. + +Commonly used exports: +- enable_instrumentation +- configure_otel_providers +- AgentTelemetryLayer +- ChatTelemetryLayer +- get_tracer +- get_meter +""" + from __future__ import annotations import contextlib @@ -1128,11 +1139,8 @@ class ChatTelemetryLayer(Generic[OptionsCoT]): opts: dict[str, Any] = options or {} # type: ignore[assignment] provider_name = str(self.otel_provider_name) model_id = kwargs.get("model_id") or opts.get("model_id") or getattr(self, "model_id", None) or "unknown" - service_url = str( - service_url_func() - if (service_url_func := getattr(self, "service_url", None)) and callable(service_url_func) - else "unknown" - ) + service_url_func = getattr(self, "service_url", None) + service_url = str(service_url_func() if callable(service_url_func) else "unknown") attributes = _get_span_attributes( operation_name=OtelAttr.CHAT_COMPLETION_OPERATION, provider_name=provider_name, diff --git a/python/packages/core/agent_framework/ollama/__init__.py b/python/packages/core/agent_framework/ollama/__init__.py index eae73853c2..1c6ba6820c 100644 --- a/python/packages/core/agent_framework/ollama/__init__.py +++ b/python/packages/core/agent_framework/ollama/__init__.py @@ -1,11 +1,21 @@ # Copyright (c) Microsoft. All rights reserved. +"""Ollama integration namespace for optional Agent Framework connectors. + +This module lazily re-exports objects from: +- ``agent-framework-ollama`` + +Supported classes: +- OllamaChatClient +- OllamaSettings +""" + import importlib from typing import Any IMPORT_PATH = "agent_framework_ollama" PACKAGE_NAME = "agent-framework-ollama" -_IMPORTS = ["__version__", "OllamaChatClient", "OllamaSettings"] +_IMPORTS = ["OllamaChatClient", "OllamaSettings"] def __getattr__(name: str) -> Any: diff --git a/python/packages/core/agent_framework/ollama/__init__.pyi b/python/packages/core/agent_framework/ollama/__init__.pyi index 3a1e7824d6..ed439f3b36 100644 --- a/python/packages/core/agent_framework/ollama/__init__.pyi +++ b/python/packages/core/agent_framework/ollama/__init__.pyi @@ -3,11 +3,9 @@ from agent_framework_ollama import ( OllamaChatClient, OllamaSettings, - __version__, ) __all__ = [ "OllamaChatClient", "OllamaSettings", - "__version__", ] diff --git a/python/packages/core/agent_framework/openai/__init__.py b/python/packages/core/agent_framework/openai/__init__.py index c5e196b022..2d9cf09648 100644 --- a/python/packages/core/agent_framework/openai/__init__.py +++ b/python/packages/core/agent_framework/openai/__init__.py @@ -1,5 +1,17 @@ # Copyright (c) Microsoft. All rights reserved. +"""OpenAI namespace for built-in Agent Framework clients. + +This module re-exports objects from the core OpenAI implementation modules in +``agent_framework.openai``. + +Supported classes include: +- OpenAIChatClient +- OpenAIResponsesClient +- OpenAIAssistantsClient +- OpenAIAssistantProvider +""" + from ._assistant_provider import OpenAIAssistantProvider from ._assistants_client import ( AssistantToolResources, diff --git a/python/packages/core/agent_framework/openai/_assistant_provider.py b/python/packages/core/agent_framework/openai/_assistant_provider.py index 6b6c2ef687..66f477dcc1 100644 --- a/python/packages/core/agent_framework/openai/_assistant_provider.py +++ b/python/packages/core/agent_framework/openai/_assistant_provider.py @@ -15,8 +15,7 @@ from agent_framework._settings import SecretString, load_settings from .._agents import Agent from .._middleware import MiddlewareTypes from .._sessions import BaseContextProvider -from .._tools import FunctionTool -from .._types import normalize_tools +from .._tools import FunctionTool, ToolTypes, normalize_tools from ..exceptions import ServiceInitializationError from ._assistants_client import OpenAIAssistantsClient from ._shared import OpenAISettings, from_assistant_tools, to_assistant_tools @@ -43,13 +42,6 @@ OptionsCoT = TypeVar( covariant=True, ) -_ToolsType = ( - FunctionTool - | Callable[..., Any] - | MutableMapping[str, Any] - | Sequence[FunctionTool | Callable[..., Any] | MutableMapping[str, Any]] -) - class OpenAIAssistantProvider(Generic[OptionsCoT]): """Provider for creating Agent instances from OpenAI Assistants API. @@ -203,7 +195,7 @@ class OpenAIAssistantProvider(Generic[OptionsCoT]): model: str, instructions: str | None = None, description: str | None = None, - tools: _ToolsType | None = None, + tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None, metadata: dict[str, str] | None = None, default_options: OptionsCoT | None = None, middleware: Sequence[MiddlewareTypes] | None = None, @@ -259,7 +251,8 @@ class OpenAIAssistantProvider(Generic[OptionsCoT]): """ # Normalize tools normalized_tools = normalize_tools(tools) - api_tools = to_assistant_tools(normalized_tools) if normalized_tools else [] + assistant_tools = [tool for tool in normalized_tools if isinstance(tool, (FunctionTool, MutableMapping))] + api_tools = to_assistant_tools(assistant_tools) if assistant_tools else [] # Extract response_format from default_options if present opts = dict(default_options) if default_options else {} @@ -311,7 +304,7 @@ class OpenAIAssistantProvider(Generic[OptionsCoT]): self, assistant_id: str, *, - tools: _ToolsType | None = None, + tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None, instructions: str | None = None, default_options: OptionsCoT | None = None, middleware: Sequence[MiddlewareTypes] | None = None, @@ -377,7 +370,7 @@ class OpenAIAssistantProvider(Generic[OptionsCoT]): self, assistant: Assistant, *, - tools: _ToolsType | None = None, + tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None, instructions: str | None = None, default_options: OptionsCoT | None = None, middleware: Sequence[MiddlewareTypes] | None = None, @@ -442,7 +435,7 @@ class OpenAIAssistantProvider(Generic[OptionsCoT]): def _validate_function_tools( self, assistant_tools: list[Any], - provided_tools: _ToolsType | None, + provided_tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None, ) -> None: """Validate that required function tools are provided. @@ -493,8 +486,8 @@ class OpenAIAssistantProvider(Generic[OptionsCoT]): def _merge_tools( self, assistant_tools: list[Any], - user_tools: _ToolsType | None, - ) -> list[FunctionTool | MutableMapping[str, Any]]: + user_tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None, + ) -> list[FunctionTool | MutableMapping[str, Any] | Any]: """Merge hosted tools from assistant with user-provided function tools. Args: @@ -504,7 +497,7 @@ class OpenAIAssistantProvider(Generic[OptionsCoT]): Returns: A list of all tools (hosted tools + user function implementations). """ - merged: list[FunctionTool | MutableMapping[str, Any]] = [] + merged: list[FunctionTool | MutableMapping[str, Any] | Any] = [] # Add hosted tools from assistant using shared conversion hosted_tools = from_assistant_tools(assistant_tools) @@ -520,7 +513,7 @@ class OpenAIAssistantProvider(Generic[OptionsCoT]): def _create_chat_agent_from_assistant( self, assistant: Assistant, - tools: list[FunctionTool | MutableMapping[str, Any]] | None, + tools: list[FunctionTool | MutableMapping[str, Any] | Any] | None, instructions: str | None, middleware: Sequence[MiddlewareTypes] | None, context_providers: Sequence[BaseContextProvider] | None, diff --git a/python/packages/core/agent_framework/openai/_assistants_client.py b/python/packages/core/agent_framework/openai/_assistants_client.py index 6135f9cac2..49ba32166e 100644 --- a/python/packages/core/agent_framework/openai/_assistants_client.py +++ b/python/packages/core/agent_framework/openai/_assistants_client.py @@ -36,6 +36,7 @@ from .._tools import ( FunctionInvocationConfiguration, FunctionInvocationLayer, FunctionTool, + normalize_tools, ) from .._types import ( ChatOptions, @@ -686,26 +687,26 @@ class OpenAIAssistantsClient( # type: ignore[misc] tool_definitions: list[MutableMapping[str, Any]] = [] # Always include tools if provided, regardless of tool_choice # tool_choice="none" means the model won't call tools, but tools should still be available - if tools is not None: - for tool in tools: - if isinstance(tool, FunctionTool): - tool_definitions.append(tool.to_json_schema_spec()) # type: ignore[reportUnknownArgumentType] - elif isinstance(tool, MutableMapping): - # Pass through dict-based tools directly (from static factory methods) - tool_definitions.append(tool) + for tool in normalize_tools(tools): + if isinstance(tool, FunctionTool): + tool_definitions.append(tool.to_json_schema_spec()) # type: ignore[reportUnknownArgumentType] + elif isinstance(tool, MutableMapping): + # Pass through dict-based tools directly (from static factory methods) + tool_definitions.append(tool) if len(tool_definitions) > 0: run_options["tools"] = tool_definitions - if (mode := tool_mode["mode"]) == "required" and ( - func_name := tool_mode.get("required_function_name") - ) is not None: - run_options["tool_choice"] = { - "type": "function", - "function": {"name": func_name}, - } - else: - run_options["tool_choice"] = mode + if tool_mode is not None: + if (mode := tool_mode["mode"]) == "required" and ( + func_name := tool_mode.get("required_function_name") + ) is not None: + run_options["tool_choice"] = { + "type": "function", + "function": {"name": func_name}, + } + else: + run_options["tool_choice"] = mode if response_format is not None: if isinstance(response_format, dict): diff --git a/python/packages/core/agent_framework/openai/_chat_client.py b/python/packages/core/agent_framework/openai/_chat_client.py index ee8f60999c..750910d651 100644 --- a/python/packages/core/agent_framework/openai/_chat_client.py +++ b/python/packages/core/agent_framework/openai/_chat_client.py @@ -27,6 +27,8 @@ from .._tools import ( FunctionInvocationConfiguration, FunctionInvocationLayer, FunctionTool, + ToolTypes, + normalize_tools, ) from .._types import ( ChatOptions, @@ -271,21 +273,24 @@ class RawOpenAIChatClient( # type: ignore[misc] # region content creation - def _prepare_tools_for_openai(self, tools: Sequence[Any]) -> dict[str, Any]: + def _prepare_tools_for_openai( + self, + tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None, + ) -> dict[str, Any]: """Prepare tools for the OpenAI Chat Completions API. Converts FunctionTool to JSON schema format. Web search tools are routed to web_search_options parameter. All other tools pass through unchanged. Args: - tools: Sequence of tools to prepare. + tools: Tool(s) to prepare. Returns: Dict containing tools and optionally web_search_options. """ chat_tools: list[Any] = [] web_search_options: dict[str, Any] | None = None - for tool in tools: + for tool in normalize_tools(tools): if isinstance(tool, FunctionTool): chat_tools.append(tool.to_json_schema_spec()) elif isinstance(tool, MutableMapping) and tool.get("type") == "web_search": @@ -338,15 +343,16 @@ class RawOpenAIChatClient( # type: ignore[misc] run_options.pop("tool_choice", None) elif tool_choice := run_options.pop("tool_choice", None): tool_mode = validate_tool_mode(tool_choice) - if (mode := tool_mode.get("mode")) == "required" and ( - func_name := tool_mode.get("required_function_name") - ) is not None: - run_options["tool_choice"] = { - "type": "function", - "function": {"name": func_name}, - } - else: - run_options["tool_choice"] = mode + if tool_mode is not None: + if (mode := tool_mode.get("mode")) == "required" and ( + func_name := tool_mode.get("required_function_name") + ) is not None: + run_options["tool_choice"] = { + "type": "function", + "function": {"name": func_name}, + } + else: + run_options["tool_choice"] = mode # response format if response_format := options.get("response_format"): diff --git a/python/packages/core/agent_framework/openai/_responses_client.py b/python/packages/core/agent_framework/openai/_responses_client.py index 0d970e38d2..bdda41a173 100644 --- a/python/packages/core/agent_framework/openai/_responses_client.py +++ b/python/packages/core/agent_framework/openai/_responses_client.py @@ -822,15 +822,16 @@ class RawOpenAIResponsesClient( # type: ignore[misc] # tool_choice: convert ToolMode to appropriate format if tool_choice := options.get("tool_choice"): tool_mode = validate_tool_mode(tool_choice) - if (mode := tool_mode.get("mode")) == "required" and ( - func_name := tool_mode.get("required_function_name") - ) is not None: - run_options["tool_choice"] = { - "type": "function", - "name": func_name, - } - else: - run_options["tool_choice"] = mode + if tool_mode is not None: + if (mode := tool_mode.get("mode")) == "required" and ( + func_name := tool_mode.get("required_function_name") + ) is not None: + run_options["tool_choice"] = { + "type": "function", + "name": func_name, + } + else: + run_options["tool_choice"] = mode else: run_options.pop("parallel_tool_calls", None) run_options.pop("tool_choice", None) diff --git a/python/packages/core/agent_framework/orchestrations/__init__.py b/python/packages/core/agent_framework/orchestrations/__init__.py index 6e220bac93..fa3561f22f 100644 --- a/python/packages/core/agent_framework/orchestrations/__init__.py +++ b/python/packages/core/agent_framework/orchestrations/__init__.py @@ -1,12 +1,24 @@ # Copyright (c) Microsoft. All rights reserved. +"""Orchestrations integration namespace for optional Agent Framework connectors. + +This module lazily re-exports objects from: +- ``agent-framework-orchestrations`` + +Supported classes include: +- SequentialBuilder +- ConcurrentBuilder +- GroupChatBuilder +- MagenticBuilder +- HandoffBuilder +""" + import importlib from typing import Any IMPORT_PATH = "agent_framework_orchestrations" PACKAGE_NAME = "agent-framework-orchestrations" _IMPORTS = [ - "__version__", # Sequential "SequentialBuilder", # Concurrent diff --git a/python/packages/core/agent_framework/orchestrations/__init__.pyi b/python/packages/core/agent_framework/orchestrations/__init__.pyi index cf26847972..79e9378dec 100644 --- a/python/packages/core/agent_framework/orchestrations/__init__.pyi +++ b/python/packages/core/agent_framework/orchestrations/__init__.pyi @@ -35,7 +35,6 @@ from agent_framework_orchestrations import ( MagenticResetSignal, SequentialBuilder, StandardMagenticManager, - __version__, ) __all__ = [ @@ -73,5 +72,4 @@ __all__ = [ "MagenticResetSignal", "SequentialBuilder", "StandardMagenticManager", - "__version__", ] diff --git a/python/packages/core/agent_framework/redis/__init__.py b/python/packages/core/agent_framework/redis/__init__.py index 9f96b3455f..d5e5de238c 100644 --- a/python/packages/core/agent_framework/redis/__init__.py +++ b/python/packages/core/agent_framework/redis/__init__.py @@ -1,11 +1,21 @@ # Copyright (c) Microsoft. All rights reserved. +"""Redis integration namespace for optional Agent Framework connectors. + +This module lazily re-exports objects from: +- ``agent-framework-redis`` + +Supported classes: +- RedisContextProvider +- RedisHistoryProvider +""" + import importlib from typing import Any IMPORT_PATH = "agent_framework_redis" PACKAGE_NAME = "agent-framework-redis" -_IMPORTS = ["__version__", "RedisContextProvider", "RedisHistoryProvider"] +_IMPORTS = ["RedisContextProvider", "RedisHistoryProvider"] def __getattr__(name: str) -> Any: diff --git a/python/packages/core/agent_framework/redis/__init__.pyi b/python/packages/core/agent_framework/redis/__init__.pyi index fc62badb76..bdf36624ba 100644 --- a/python/packages/core/agent_framework/redis/__init__.pyi +++ b/python/packages/core/agent_framework/redis/__init__.pyi @@ -3,11 +3,9 @@ from agent_framework_redis import ( RedisContextProvider, RedisHistoryProvider, - __version__, ) __all__ = [ "RedisContextProvider", "RedisHistoryProvider", - "__version__", ] diff --git a/python/packages/core/pyproject.toml b/python/packages/core/pyproject.toml index bb3bb394b9..a883be8979 100644 --- a/python/packages/core/pyproject.toml +++ b/python/packages/core/pyproject.toml @@ -45,13 +45,16 @@ all = [ "agent-framework-ag-ui", "agent-framework-azure-ai-search", "agent-framework-anthropic", + "agent-framework-claude", "agent-framework-azure-ai", "agent-framework-azurefunctions", + "agent-framework-bedrock", "agent-framework-chatkit", "agent-framework-copilotstudio", "agent-framework-declarative", "agent-framework-devui", "agent-framework-durabletask", + "agent-framework-foundry-local", "agent-framework-github-copilot", "agent-framework-lab", "agent-framework-mem0", diff --git a/python/packages/core/tests/core/test_function_invocation_logic.py b/python/packages/core/tests/core/test_function_invocation_logic.py index 593aa6f737..2e7601f787 100644 --- a/python/packages/core/tests/core/test_function_invocation_logic.py +++ b/python/packages/core/tests/core/test_function_invocation_logic.py @@ -56,6 +56,36 @@ async def test_base_client_with_function_calling(chat_client_base: SupportsChatG assert response.messages[2].text == "done" +async def test_base_client_with_function_calling_tools_in_kwargs(chat_client_base: SupportsChatGetResponse): + exec_counter = 0 + + @tool(name="test_function", approval_mode="never_require") + def ai_func(arg1: str) -> str: + nonlocal exec_counter + exec_counter += 1 + return f"Processed {arg1}" + + chat_client_base.run_responses = [ + ChatResponse( + messages=Message( + role="assistant", + contents=[ + Content.from_function_call(call_id="1", name="test_function", arguments='{"arg1": "value1"}') + ], + ) + ), + ChatResponse(messages=Message(role="assistant", text="done")), + ] + + response = await chat_client_base.get_response("hello", tools=[ai_func]) + + assert exec_counter == 1 + assert len(response.messages) == 3 + assert response.messages[1].role == "tool" + assert response.messages[1].contents[0].type == "function_result" + assert response.messages[1].contents[0].result == "Processed value1" + + @pytest.mark.parametrize("max_iterations", [3]) async def test_base_client_with_function_calling_resets(chat_client_base: SupportsChatGetResponse): exec_counter = 0 diff --git a/python/packages/core/tests/core/test_types.py b/python/packages/core/tests/core/test_types.py index 7a5acdedf7..8a8885b919 100644 --- a/python/packages/core/tests/core/test_types.py +++ b/python/packages/core/tests/core/test_types.py @@ -921,8 +921,8 @@ def test_chat_options_tool_choice_validation(): } assert validate_tool_mode({"mode": "none"}) == {"mode": "none"} - # None should return mode==none - assert validate_tool_mode(None) == {"mode": "none"} + # None should remain unset + assert validate_tool_mode(None) is None with raises(ContentError): validate_tool_mode("invalid_mode") diff --git a/python/packages/core/tests/openai/test_openai_assistants_client.py b/python/packages/core/tests/openai/test_openai_assistants_client.py index 80e2020d02..4521a3c693 100644 --- a/python/packages/core/tests/openai/test_openai_assistants_client.py +++ b/python/packages/core/tests/openai/test_openai_assistants_client.py @@ -701,6 +701,7 @@ def test_prepare_options_basic(mock_async_openai: MagicMock) -> None: assert run_options["model"] == "gpt-4" assert run_options["temperature"] == 0.7 assert run_options["top_p"] == 0.9 + assert "tool_choice" not in run_options assert tool_results is None @@ -733,6 +734,52 @@ def test_prepare_options_with_tool_tool(mock_async_openai: MagicMock) -> None: assert run_options["tool_choice"] == "auto" +def test_prepare_options_with_tools_without_tool_choice(mock_async_openai: MagicMock) -> None: + """Test _prepare_options keeps tool_choice unset when not provided.""" + + client = create_test_openai_assistants_client(mock_async_openai) + + @tool(approval_mode="never_require") + def test_function(query: str) -> str: + """A test function.""" + return f"Result for {query}" + + options = { + "tools": [test_function], + } + + messages = [Message(role="user", text="Hello")] + run_options, _ = client._prepare_options(messages, options) # type: ignore + + assert "tools" in run_options + assert "tool_choice" not in run_options + + +def test_prepare_options_with_single_tool_tool(mock_async_openai: MagicMock) -> None: + """Test _prepare_options with a single FunctionTool (non-sequence).""" + client = create_test_openai_assistants_client(mock_async_openai) + + @tool(approval_mode="never_require") + def test_function(query: str) -> str: + """A test function.""" + return f"Result for {query}" + + options = { + "tools": test_function, + "tool_choice": "auto", + } + + messages = [Message(role="user", text="Hello")] + run_options, tool_results = client._prepare_options(messages, options) # type: ignore + + assert "tools" in run_options + assert len(run_options["tools"]) == 1 + assert run_options["tools"][0]["type"] == "function" + assert "function" in run_options["tools"][0] + assert run_options["tool_choice"] == "auto" + assert tool_results is None + + def test_prepare_options_with_code_interpreter(mock_async_openai: MagicMock) -> None: """Test _prepare_options with code interpreter tool.""" client = create_test_openai_assistants_client(mock_async_openai) diff --git a/python/packages/core/tests/openai/test_openai_chat_client.py b/python/packages/core/tests/openai/test_openai_chat_client.py index c5ee81bfce..a6482367e9 100644 --- a/python/packages/core/tests/openai/test_openai_chat_client.py +++ b/python/packages/core/tests/openai/test_openai_chat_client.py @@ -190,6 +190,21 @@ def test_unsupported_tool_handling(openai_unit_test_env: dict[str, str]) -> None assert result["tools"] == [dict_tool] +def test_prepare_tools_with_single_function_tool(openai_unit_test_env: dict[str, str]) -> None: + """Test that a single FunctionTool is accepted for tool preparation.""" + client = OpenAIChatClient() + + @tool(approval_mode="never_require") + def test_function(query: str) -> str: + """A test function.""" + return f"Result for {query}" + + result = client._prepare_tools_for_openai(test_function) + assert "tools" in result + assert len(result["tools"]) == 1 + assert result["tools"][0]["type"] == "function" + + @tool(approval_mode="never_require") def get_story_text() -> str: """Returns a story about Emily and David.""" diff --git a/python/packages/foundry_local/README.md b/python/packages/foundry_local/README.md index c65e5a0386..cf4f340bd3 100644 --- a/python/packages/foundry_local/README.md +++ b/python/packages/foundry_local/README.md @@ -7,3 +7,7 @@ pip install agent-framework-foundry-local --pre ``` and see the [README](https://github.com/microsoft/agent-framework/tree/main/python/README.md) for more information. + +## Foundry Local Sample + +See the [Foundry Local provider sample](../../samples/02-agents/providers/foundry_local/foundry_local_agent.py) for a runnable example. diff --git a/python/packages/github_copilot/agent_framework_github_copilot/_agent.py b/python/packages/github_copilot/agent_framework_github_copilot/_agent.py index 26323e915d..36a2ee80b7 100644 --- a/python/packages/github_copilot/agent_framework_github_copilot/_agent.py +++ b/python/packages/github_copilot/agent_framework_github_copilot/_agent.py @@ -22,7 +22,7 @@ from agent_framework import ( normalize_messages, ) from agent_framework._settings import load_settings -from agent_framework._tools import FunctionTool +from agent_framework._tools import FunctionTool, ToolTypes from agent_framework._types import AgentRunInputs, normalize_tools from agent_framework.exceptions import ServiceException from copilot import CopilotClient, CopilotSession @@ -151,11 +151,7 @@ class GitHubCopilotAgent(BaseAgent, Generic[OptionsT]): description: str | None = None, context_providers: Sequence[BaseContextProvider] | None = None, middleware: Sequence[AgentMiddlewareTypes] | None = None, - tools: FunctionTool - | Callable[..., Any] - | MutableMapping[str, Any] - | Sequence[FunctionTool | Callable[..., Any] | MutableMapping[str, Any]] - | None = None, + tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None, default_options: OptionsT | None = None, env_file_path: str | None = None, env_file_encoding: str | None = None, @@ -478,7 +474,7 @@ class GitHubCopilotAgent(BaseAgent, Generic[OptionsT]): def _prepare_tools( self, - tools: list[FunctionTool | MutableMapping[str, Any]], + tools: Sequence[ToolTypes | CopilotTool], ) -> list[CopilotTool]: """Convert Agent Framework tools to Copilot SDK tools. @@ -491,10 +487,12 @@ class GitHubCopilotAgent(BaseAgent, Generic[OptionsT]): copilot_tools: list[CopilotTool] = [] for tool in tools: - if isinstance(tool, FunctionTool): - copilot_tools.append(self._tool_to_copilot_tool(tool)) # type: ignore - elif isinstance(tool, CopilotTool): + if isinstance(tool, CopilotTool): copilot_tools.append(tool) + elif isinstance(tool, FunctionTool): + copilot_tools.append(self._tool_to_copilot_tool(tool)) # type: ignore + elif isinstance(tool, MutableMapping): + copilot_tools.append(tool) # type: ignore[arg-type] # Note: Other tool types (e.g., dict-based hosted tools) are skipped return copilot_tools diff --git a/python/samples/01-get-started/01_hello_agent.py b/python/samples/01-get-started/01_hello_agent.py index a3353a5e87..45b0cf410a 100644 --- a/python/samples/01-get-started/01_hello_agent.py +++ b/python/samples/01-get-started/01_hello_agent.py @@ -12,6 +12,8 @@ Hello Agent — Simplest possible agent This sample creates a minimal agent using AzureOpenAIResponsesClient via an Azure AI Foundry project endpoint, and runs it in both non-streaming and streaming modes. +There are XML tags in all of the get started samples, those are used to display the same code in the docs repo. + Environment variables: AZURE_AI_PROJECT_ENDPOINT — Your Azure AI Foundry project endpoint AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME — Model deployment name (e.g. gpt-4o) diff --git a/python/samples/01-get-started/README.md b/python/samples/01-get-started/README.md index 774368e15f..5ba119e016 100644 --- a/python/samples/01-get-started/README.md +++ b/python/samples/01-get-started/README.md @@ -12,8 +12,8 @@ pip install agent-framework --pre Set the required environment variables: ```bash -export OPENAI_API_KEY="sk-..." -export OPENAI_RESPONSES_MODEL_ID="gpt-4o" # optional, defaults to gpt-4o +export AZURE_AI_PROJECT_ENDPOINT="https://your-project-endpoint" +export AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME="gpt-4o" # optional, defaults to gpt-4o ``` ## Samples @@ -32,3 +32,5 @@ Run any sample with: ```bash python 01_hello_agent.py ``` + +These samples use Azure Foundry models with the Responses API. To switch providers, just replace the client, see [all providers](../02-agents/providers/README.md) diff --git a/python/samples/02-agents/chat_client/README.md b/python/samples/02-agents/chat_client/README.md index 5bf9b471ad..e03d532812 100644 --- a/python/samples/02-agents/chat_client/README.md +++ b/python/samples/02-agents/chat_client/README.md @@ -1,41 +1,74 @@ # Chat Client Examples -This folder contains simple examples demonstrating direct usage of various chat clients. +This folder contains examples for direct chat client usage patterns. ## Examples | File | Description | |------|-------------| -| [`azure_assistants_client.py`](azure_assistants_client.py) | Direct usage of Azure Assistants Client for basic chat interactions with Azure OpenAI assistants. | -| [`azure_chat_client.py`](azure_chat_client.py) | Direct usage of Azure Chat Client for chat interactions with Azure OpenAI models. | -| [`azure_responses_client.py`](azure_responses_client.py) | Direct usage of Azure Responses Client for structured response generation with Azure OpenAI models. | +| [`built_in_chat_clients.py`](built_in_chat_clients.py) | Consolidated sample for built-in chat clients. Uses `get_client()` to create the selected client and pass it to `main()`. | | [`chat_response_cancellation.py`](chat_response_cancellation.py) | Demonstrates how to cancel chat responses during streaming, showing proper cancellation handling and cleanup. | -| [`azure_ai_chat_client.py`](azure_ai_chat_client.py) | Direct usage of Azure AI Chat Client for chat interactions with Azure AI models. | -| [`openai_assistants_client.py`](openai_assistants_client.py) | Direct usage of OpenAI Assistants Client for basic chat interactions with OpenAI assistants. | -| [`openai_chat_client.py`](openai_chat_client.py) | Direct usage of OpenAI Chat Client for chat interactions with OpenAI models. | -| [`openai_responses_client.py`](openai_responses_client.py) | Direct usage of OpenAI Responses Client for structured response generation with OpenAI models. | | [`custom_chat_client.py`](custom_chat_client.py) | Demonstrates how to create custom chat clients by extending the `BaseChatClient` class. Shows a `EchoingChatClient` implementation and how to integrate it with `Agent` using the `as_agent()` method. | +## Selecting a built-in client + +`built_in_chat_clients.py` starts with: + +```python +asyncio.run(main("openai_chat")) +``` + +Change the argument to pick a client: + +- `openai_chat` +- `openai_responses` +- `openai_assistants` +- `anthropic` +- `ollama` +- `bedrock` +- `azure_openai_chat` +- `azure_openai_responses` +- `azure_openai_responses_foundry` +- `azure_openai_assistants` +- `azure_ai_agent` + +Example: + +```bash +uv run samples/02-agents/chat_client/built_in_chat_clients.py +``` + ## Environment Variables -Depending on which client you're using, set the appropriate environment variables: +Depending on the selected client, set the appropriate environment variables: **For Azure clients:** - `AZURE_OPENAI_ENDPOINT`: Your Azure OpenAI endpoint - `AZURE_OPENAI_CHAT_DEPLOYMENT_NAME`: The name of your Azure OpenAI chat deployment - `AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME`: The name of your Azure OpenAI responses deployment -**For Azure AI client:** +**For Azure OpenAI Foundry responses client (`azure_openai_responses_foundry`):** - `AZURE_AI_PROJECT_ENDPOINT`: Your Azure AI project endpoint -- `AZURE_AI_MODEL_DEPLOYMENT_NAME`: The name of your model deployment +- `AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME`: The name of your Azure OpenAI responses deployment + +**For Azure AI agent client (`azure_ai_agent`):** +- `AZURE_AI_PROJECT_ENDPOINT`: Your Azure AI project endpoint +- `AZURE_AI_MODEL_DEPLOYMENT_NAME`: The name of your model deployment (used by `azure_ai_agent`) **For OpenAI clients:** - `OPENAI_API_KEY`: Your OpenAI API key -- `OPENAI_CHAT_MODEL_ID`: The OpenAI model to use for chat clients (e.g., `gpt-4o`, `gpt-4o-mini`, `gpt-3.5-turbo`) -- `OPENAI_RESPONSES_MODEL_ID`: The OpenAI model to use for responses clients (e.g., `gpt-4o`, `gpt-4o-mini`, `gpt-3.5-turbo`) +- `OPENAI_CHAT_MODEL_ID`: The OpenAI model for `openai_chat` and `openai_assistants` +- `OPENAI_RESPONSES_MODEL_ID`: The OpenAI model for `openai_responses` -**For Ollama client:** -- `OLLAMA_HOST`: Your Ollama server URL (defaults to `http://localhost:11434` if not set) -- `OLLAMA_MODEL_ID`: The Ollama model to use for chat (e.g., `llama3.2`, `llama2`, `codellama`) +**For Anthropic client (`anthropic`):** +- `ANTHROPIC_API_KEY`: Your Anthropic API key +- `ANTHROPIC_CHAT_MODEL_ID`: The Anthropic model ID (for example, `claude-sonnet-4-5`) -> **Note**: For Ollama, ensure you have Ollama installed and running locally with at least one model downloaded. Visit [https://ollama.com/](https://ollama.com/) for installation instructions. +**For Ollama client (`ollama`):** +- `OLLAMA_HOST`: Ollama server URL (defaults to `http://localhost:11434` if unset) +- `OLLAMA_MODEL_ID`: Ollama model name (for example, `mistral`, `qwen2.5:8b`) + +**For Bedrock client (`bedrock`):** +- `BEDROCK_CHAT_MODEL_ID`: Bedrock model ID (for example, `anthropic.claude-3-5-sonnet-20240620-v1:0`) +- `BEDROCK_REGION`: AWS region (defaults to `us-east-1` if unset) +- AWS credentials via standard environment variables (for example, `AWS_ACCESS_KEY_ID`, `AWS_SECRET_ACCESS_KEY`) diff --git a/python/samples/02-agents/chat_client/azure_ai_chat_client.py b/python/samples/02-agents/chat_client/azure_ai_chat_client.py deleted file mode 100644 index 236d93f1a7..0000000000 --- a/python/samples/02-agents/chat_client/azure_ai_chat_client.py +++ /dev/null @@ -1,49 +0,0 @@ -# Copyright (c) Microsoft. All rights reserved. - -import asyncio -from random import randint -from typing import Annotated - -from agent_framework import tool -from agent_framework.azure import AzureAIAgentClient -from azure.identity.aio import AzureCliCredential -from pydantic import Field - -""" -Azure AI Chat Client Direct Usage Example - -Demonstrates direct AzureAIChatClient usage for chat interactions with Azure AI models. -Shows function calling capabilities with custom business logic. -""" - - -# NOTE: approval_mode="never_require" is for sample brevity. Use "always_require" in production; see samples/02-agents/tools/function_tool_with_approval.py and samples/02-agents/tools/function_tool_with_approval_and_sessions.py. -@tool(approval_mode="never_require") -def get_weather( - location: Annotated[str, Field(description="The location to get the weather for.")], -) -> str: - """Get the weather for a given location.""" - conditions = ["sunny", "cloudy", "rainy", "stormy"] - return f"The weather in {location} is {conditions[randint(0, 3)]} with a high of {randint(10, 30)}°C." - - -async def main() -> None: - # For authentication, run `az login` command in terminal or replace AzureCliCredential with preferred - # authentication option. - async with AzureAIAgentClient(credential=AzureCliCredential()) as client: - message = "What's the weather in Amsterdam and in Paris?" - stream = False - print(f"User: {message}") - if stream: - print("Assistant: ", end="") - async for chunk in client.get_response(message, tools=get_weather, stream=True): - if str(chunk): - print(str(chunk), end="") - print("") - else: - response = await client.get_response(message, tools=get_weather) - print(f"Assistant: {response}") - - -if __name__ == "__main__": - asyncio.run(main()) diff --git a/python/samples/02-agents/chat_client/azure_assistants_client.py b/python/samples/02-agents/chat_client/azure_assistants_client.py deleted file mode 100644 index 66034e8eee..0000000000 --- a/python/samples/02-agents/chat_client/azure_assistants_client.py +++ /dev/null @@ -1,49 +0,0 @@ -# Copyright (c) Microsoft. All rights reserved. - -import asyncio -from random import randint -from typing import Annotated - -from agent_framework import tool -from agent_framework.azure import AzureOpenAIAssistantsClient -from azure.identity import AzureCliCredential -from pydantic import Field - -""" -Azure Assistants Client Direct Usage Example - -Demonstrates direct AzureAssistantsClient usage for chat interactions with Azure OpenAI assistants. -Shows function calling capabilities and automatic assistant creation. -""" - - -# NOTE: approval_mode="never_require" is for sample brevity. Use "always_require" in production; see samples/02-agents/tools/function_tool_with_approval.py and samples/02-agents/tools/function_tool_with_approval_and_sessions.py. -@tool(approval_mode="never_require") -def get_weather( - location: Annotated[str, Field(description="The location to get the weather for.")], -) -> str: - """Get the weather for a given location.""" - conditions = ["sunny", "cloudy", "rainy", "stormy"] - return f"The weather in {location} is {conditions[randint(0, 3)]} with a high of {randint(10, 30)}°C." - - -async def main() -> None: - # For authentication, run `az login` command in terminal or replace AzureCliCredential with preferred - # authentication option. - async with AzureOpenAIAssistantsClient(credential=AzureCliCredential()) as client: - message = "What's the weather in Amsterdam and in Paris?" - stream = False - print(f"User: {message}") - if stream: - print("Assistant: ", end="") - async for chunk in client.get_response(message, tools=get_weather, stream=True): - if str(chunk): - print(str(chunk), end="") - print("") - else: - response = await client.get_response(message, tools=get_weather) - print(f"Assistant: {response}") - - -if __name__ == "__main__": - asyncio.run(main()) diff --git a/python/samples/02-agents/chat_client/azure_chat_client.py b/python/samples/02-agents/chat_client/azure_chat_client.py deleted file mode 100644 index 675df29774..0000000000 --- a/python/samples/02-agents/chat_client/azure_chat_client.py +++ /dev/null @@ -1,49 +0,0 @@ -# Copyright (c) Microsoft. All rights reserved. - -import asyncio -from random import randint -from typing import Annotated - -from agent_framework import tool -from agent_framework.azure import AzureOpenAIChatClient -from azure.identity import AzureCliCredential -from pydantic import Field - -""" -Azure Chat Client Direct Usage Example - -Demonstrates direct AzureChatClient usage for chat interactions with Azure OpenAI models. -Shows function calling capabilities with custom business logic. -""" - - -# NOTE: approval_mode="never_require" is for sample brevity. Use "always_require" in production; see samples/02-agents/tools/function_tool_with_approval.py and samples/02-agents/tools/function_tool_with_approval_and_sessions.py. -@tool(approval_mode="never_require") -def get_weather( - location: Annotated[str, Field(description="The location to get the weather for.")], -) -> str: - """Get the weather for a given location.""" - conditions = ["sunny", "cloudy", "rainy", "stormy"] - return f"The weather in {location} is {conditions[randint(0, 3)]} with a high of {randint(10, 30)}°C." - - -async def main() -> None: - # For authentication, run `az login` command in terminal or replace AzureCliCredential with preferred - # authentication option. - client = AzureOpenAIChatClient(credential=AzureCliCredential()) - message = "What's the weather in Amsterdam and in Paris?" - stream = False - print(f"User: {message}") - if stream: - print("Assistant: ", end="") - async for chunk in client.get_response(message, tools=get_weather, stream=True): - if str(chunk): - print(str(chunk), end="") - print("") - else: - response = await client.get_response(message, tools=get_weather) - print(f"Assistant: {response}") - - -if __name__ == "__main__": - asyncio.run(main()) diff --git a/python/samples/02-agents/chat_client/azure_responses_client.py b/python/samples/02-agents/chat_client/azure_responses_client.py deleted file mode 100644 index 7ab4212a3a..0000000000 --- a/python/samples/02-agents/chat_client/azure_responses_client.py +++ /dev/null @@ -1,95 +0,0 @@ -# Copyright (c) Microsoft. All rights reserved. - -import asyncio -from random import randint -from typing import Annotated - -from agent_framework import tool -from agent_framework.azure import AzureOpenAIResponsesClient -from azure.identity import AzureCliCredential -from pydantic import BaseModel - -""" -Azure Responses Client Direct Usage Example - -Demonstrates direct AzureResponsesClient usage for structured response generation with Azure OpenAI models. -Shows function calling capabilities with custom business logic. -""" - - -# NOTE: approval_mode="never_require" is for sample brevity. Use "always_require" in production; see samples/02-agents/tools/function_tool_with_approval.py and samples/02-agents/tools/function_tool_with_approval_and_sessions.py. -@tool(approval_mode="never_require") -def get_weather( - location: Annotated[str, "The location to get the weather for."], -) -> str: - """Get the weather for a given location.""" - conditions = ["sunny", "cloudy", "rainy", "stormy"] - return f"The weather in {location} is {conditions[randint(0, 3)]} with a high of {randint(10, 30)}°C." - - -@tool(approval_mode="never_require") -def get_time(): - """Get the current time.""" - from datetime import datetime - - now = datetime.now() - return f"The current date time is {now.strftime('%Y-%m-%d - %H:%M:%S')}." - - -class WeatherDetail(BaseModel): - """Structured output for weather information.""" - - location: str - weather: str - - -class Weather(BaseModel): - """Container for multiple outputs.""" - - date_time: str - weather_details: list[WeatherDetail] - - -async def main() -> None: - # For authentication, run `az login` command in terminal or replace AzureCliCredential with preferred - # authentication option. - client = AzureOpenAIResponsesClient(credential=AzureCliCredential(), api_version="preview") - message = "What's the weather in Amsterdam and in Paris?" - stream = True - print(f"User: {message}") - response = client.get_response( - message, - options={"response_format": Weather, "tools": [get_weather, get_time]}, - stream=stream, - ) - if stream: - response = await response.get_final_response() - else: - response = await response - if result := response.value: - print(f"Assistant: {result.model_dump_json(indent=2)}") - else: - print(f"Assistant: {response.text}") - - -# Expected output (time will be different): -""" -User: What's the weather in Amsterdam and in Paris? -Assistant: { - "date_time": "2026-02-06 - 13:30:40", - "weather_details": [ - { - "location": "Amsterdam", - "weather": "The weather in Amsterdam is cloudy with a high of 21°C." - }, - { - "location": "Paris", - "weather": "The weather in Paris is sunny with a high of 27°C." - } - ] -} -""" - - -if __name__ == "__main__": - asyncio.run(main()) diff --git a/python/samples/02-agents/chat_client/built_in_chat_clients.py b/python/samples/02-agents/chat_client/built_in_chat_clients.py new file mode 100644 index 0000000000..06be683781 --- /dev/null +++ b/python/samples/02-agents/chat_client/built_in_chat_clients.py @@ -0,0 +1,156 @@ +# Copyright (c) Microsoft. All rights reserved. + +import asyncio +import os +from random import randint +from typing import Annotated, Any, Literal + +from agent_framework import SupportsChatGetResponse, tool +from agent_framework.azure import ( + AzureAIAgentClient, + AzureOpenAIAssistantsClient, +) +from agent_framework.openai import OpenAIAssistantsClient +from azure.identity import AzureCliCredential +from azure.identity.aio import AzureCliCredential as AsyncAzureCliCredential +from pydantic import Field + +""" +Built-in Chat Clients Example + +This sample demonstrates how to run the same prompt flow against different built-in +chat clients using a single `get_client` factory. + +Select one of these client names: +- openai_chat +- openai_responses +- openai_assistants +- anthropic +- ollama +- bedrock +- azure_openai_chat +- azure_openai_responses +- azure_openai_responses_foundry +- azure_openai_assistants +- azure_ai_agent +""" + +ClientName = Literal[ + "openai_chat", + "openai_responses", + "openai_assistants", + "anthropic", + "ollama", + "bedrock", + "azure_openai_chat", + "azure_openai_responses", + "azure_openai_responses_foundry", + "azure_openai_assistants", + "azure_ai_agent", +] + + +# NOTE: approval_mode="never_require" is for sample brevity. +# Use "always_require" in production; see samples/02-agents/tools/function_tool_with_approval.py +# and samples/02-agents/tools/function_tool_with_approval_and_sessions.py. +@tool(approval_mode="never_require") +def get_weather( + location: Annotated[str, Field(description="The location to get the weather for.")], +) -> str: + """Get the weather for a given location.""" + conditions = ["sunny", "cloudy", "rainy", "stormy"] + return f"The weather in {location} is {conditions[randint(0, 3)]} with a high of {randint(10, 30)}°C." + + +def get_client(client_name: ClientName) -> SupportsChatGetResponse[Any]: + """Create a built-in chat client from a name.""" + from agent_framework.amazon import BedrockChatClient + from agent_framework.anthropic import AnthropicClient + from agent_framework.azure import ( + AzureOpenAIChatClient, + AzureOpenAIResponsesClient, + ) + from agent_framework.ollama import OllamaChatClient + from agent_framework.openai import OpenAIChatClient, OpenAIResponsesClient + + # 1. Create OpenAI clients. + if client_name == "openai_chat": + return OpenAIChatClient() + if client_name == "openai_responses": + return OpenAIResponsesClient() + if client_name == "openai_assistants": + return OpenAIAssistantsClient() + if client_name == "anthropic": + return AnthropicClient() + if client_name == "ollama": + return OllamaChatClient() + if client_name == "bedrock": + return BedrockChatClient() + + # 2. Create Azure OpenAI clients. + if client_name == "azure_openai_chat": + return AzureOpenAIChatClient(credential=AzureCliCredential()) + if client_name == "azure_openai_responses": + return AzureOpenAIResponsesClient(credential=AzureCliCredential(), api_version="preview") + if client_name == "azure_openai_responses_foundry": + return AzureOpenAIResponsesClient( + project_endpoint=os.environ["AZURE_AI_PROJECT_ENDPOINT"], + deployment_name=os.environ["AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME"], + credential=AzureCliCredential(), + ) + if client_name == "azure_openai_assistants": + return AzureOpenAIAssistantsClient(credential=AzureCliCredential()) + + # 3. Create Azure AI client. + if client_name == "azure_ai_agent": + return AzureAIAgentClient(credential=AsyncAzureCliCredential()) + + raise ValueError(f"Unsupported client name: {client_name}") + + +async def main(client_name: ClientName = "openai_chat") -> None: + """Run a basic prompt using a selected built-in client.""" + client = get_client(client_name) + + # 1. Configure prompt and streaming mode. + message = "What's the weather in Amsterdam and in Paris?" + stream = os.getenv("STREAM", "false").lower() == "true" + print(f"Client: {client_name}") + print(f"User: {message}") + + # 2. Run with context-managed clients. + if isinstance(client, OpenAIAssistantsClient | AzureOpenAIAssistantsClient | AzureAIAgentClient): + async with client: + if stream: + response_stream = client.get_response(message, stream=True, options={"tools": get_weather}) + print("Assistant: ", end="") + async for chunk in response_stream: + if chunk.text: + print(chunk.text, end="") + print("") + else: + print(f"Assistant: {await client.get_response(message, stream=False, options={'tools': get_weather})}") + return + + # 3. Run with non-context-managed clients. + if stream: + response_stream = client.get_response(message, stream=True, options={"tools": get_weather}) + print("Assistant: ", end="") + async for chunk in response_stream: + if chunk.text: + print(chunk.text, end="") + print("") + else: + print(f"Assistant: {await client.get_response(message, stream=False, options={'tools': get_weather})}") + + +if __name__ == "__main__": + asyncio.run(main("openai_chat")) + + +""" +Sample output: +User: What's the weather in Amsterdam and in Paris? +Assistant: The weather in Amsterdam is sunny with a high of 25°C. +...and in Paris it is cloudy with a high of 19°C. +""" diff --git a/python/samples/02-agents/chat_client/custom_chat_client.py b/python/samples/02-agents/chat_client/custom_chat_client.py index b6c69bd0ac..ae009c059a 100644 --- a/python/samples/02-agents/chat_client/custom_chat_client.py +++ b/python/samples/02-agents/chat_client/custom_chat_client.py @@ -4,7 +4,7 @@ import asyncio import random import sys from collections.abc import AsyncIterable, Awaitable, Mapping, Sequence -from typing import Any, ClassVar, Generic +from typing import Any, ClassVar, TypeAlias, TypedDict from agent_framework import ( BaseChatClient, @@ -15,15 +15,9 @@ from agent_framework import ( FunctionInvocationLayer, Message, ResponseStream, - Role, ) -from agent_framework._clients import OptionsCoT from agent_framework.observability import ChatTelemetryLayer -if sys.version_info >= (3, 13): - pass -else: - pass if sys.version_info >= (3, 12): from typing import override # type: ignore # pragma: no cover else: @@ -38,7 +32,18 @@ middleware, telemetry, and function invocation layers explicitly. """ -class EchoingChatClient(BaseChatClient[OptionsCoT], Generic[OptionsCoT]): +class EchoingChatClientOptions(TypedDict, total=False): + """Custom options for EchoingChatClient.""" + + uppercase: bool + suffix: str + stream_delay_seconds: float + + +OptionsT: TypeAlias = EchoingChatClientOptions + + +class EchoingChatClient(BaseChatClient[OptionsT]): """A custom chat client that echoes messages back with modifications. This demonstrates how to implement a custom chat client by extending BaseChatClient @@ -73,7 +78,7 @@ class EchoingChatClient(BaseChatClient[OptionsCoT], Generic[OptionsCoT]): # Echo the last user message last_user_message = None for message in reversed(messages): - if message.role == Role.USER: + if message.role == "user": last_user_message = message break @@ -82,7 +87,13 @@ class EchoingChatClient(BaseChatClient[OptionsCoT], Generic[OptionsCoT]): else: response_text = f"{self.prefix} [No text message found]" - response_message = Message(role=Role.ASSISTANT, contents=[Content.from_text(response_text)]) + if options.get("uppercase"): + response_text = response_text.upper() + if suffix := options.get("suffix"): + response_text = f"{response_text} {suffix}" + stream_delay_seconds = float(options.get("stream_delay_seconds", 0.05)) + + response_message = Message(role="assistant", contents=[Content.from_text(response_text)]) response = ChatResponse( messages=[response_message], @@ -102,21 +113,20 @@ class EchoingChatClient(BaseChatClient[OptionsCoT], Generic[OptionsCoT]): for char in response_text_local: yield ChatResponseUpdate( contents=[Content.from_text(char)], - role=Role.ASSISTANT, + role="assistant", response_id=f"echo-stream-resp-{random.randint(1000, 9999)}", model_id="echo-model-v1", ) - await asyncio.sleep(0.05) + await asyncio.sleep(stream_delay_seconds) return ResponseStream(_stream(), finalizer=lambda updates: response) -class EchoingChatClientWithLayers( # type: ignore[misc,type-var] - ChatMiddlewareLayer[OptionsCoT], - ChatTelemetryLayer[OptionsCoT], - FunctionInvocationLayer[OptionsCoT], - EchoingChatClient[OptionsCoT], - Generic[OptionsCoT], +class EchoingChatClientWithLayers( # type: ignore[misc] + ChatMiddlewareLayer[OptionsT], + ChatTelemetryLayer[OptionsT], + FunctionInvocationLayer[OptionsT], + EchoingChatClient, ): """Echoing chat client that explicitly composes middleware, telemetry, and function layers.""" @@ -134,7 +144,14 @@ async def main() -> None: # Use the chat client directly print("Using chat client directly:") - direct_response = await echo_client.get_response("Hello, custom chat client!") + direct_response = await echo_client.get_response( + "Hello, custom chat client!", + options={ + "uppercase": True, + "suffix": "(CUSTOM OPTIONS)", + "stream_delay_seconds": 0.02, + }, + ) print(f"Direct response: {direct_response.messages[0].text}") # Create an agent using the custom chat client diff --git a/python/samples/02-agents/chat_client/openai_assistants_client.py b/python/samples/02-agents/chat_client/openai_assistants_client.py deleted file mode 100644 index 7783743950..0000000000 --- a/python/samples/02-agents/chat_client/openai_assistants_client.py +++ /dev/null @@ -1,47 +0,0 @@ -# Copyright (c) Microsoft. All rights reserved. - -import asyncio -from random import randint -from typing import Annotated - -from agent_framework import tool -from agent_framework.openai import OpenAIAssistantsClient -from pydantic import Field - -""" -OpenAI Assistants Client Direct Usage Example - -Demonstrates direct OpenAIAssistantsClient usage for chat interactions with OpenAI assistants. -Shows function calling capabilities and automatic assistant creation. - -""" - - -# NOTE: approval_mode="never_require" is for sample brevity. Use "always_require" in production; see samples/02-agents/tools/function_tool_with_approval.py and samples/02-agents/tools/function_tool_with_approval_and_sessions.py. -@tool(approval_mode="never_require") -def get_weather( - location: Annotated[str, Field(description="The location to get the weather for.")], -) -> str: - """Get the weather for a given location.""" - conditions = ["sunny", "cloudy", "rainy", "stormy"] - return f"The weather in {location} is {conditions[randint(0, 3)]} with a high of {randint(10, 30)}°C." - - -async def main() -> None: - async with OpenAIAssistantsClient() as client: - message = "What's the weather in Amsterdam and in Paris?" - stream = False - print(f"User: {message}") - if stream: - print("Assistant: ", end="") - async for chunk in client.get_response(message, tools=get_weather, stream=True): - if str(chunk): - print(str(chunk), end="") - print("") - else: - response = await client.get_response(message, tools=get_weather) - print(f"Assistant: {response}") - - -if __name__ == "__main__": - asyncio.run(main()) diff --git a/python/samples/02-agents/chat_client/openai_chat_client.py b/python/samples/02-agents/chat_client/openai_chat_client.py deleted file mode 100644 index e784c17ae2..0000000000 --- a/python/samples/02-agents/chat_client/openai_chat_client.py +++ /dev/null @@ -1,47 +0,0 @@ -# Copyright (c) Microsoft. All rights reserved. - -import asyncio -from random import randint -from typing import Annotated - -from agent_framework import tool -from agent_framework.openai import OpenAIChatClient -from pydantic import Field - -""" -OpenAI Chat Client Direct Usage Example - -Demonstrates direct OpenAIChatClient usage for chat interactions with OpenAI models. -Shows function calling capabilities with custom business logic. - -""" - - -# NOTE: approval_mode="never_require" is for sample brevity. Use "always_require" in production; see samples/02-agents/tools/function_tool_with_approval.py and samples/02-agents/tools/function_tool_with_approval_and_sessions.py. -@tool(approval_mode="never_require") -def get_weather( - location: Annotated[str, Field(description="The location to get the weather for.")], -) -> str: - """Get the weather for a given location.""" - conditions = ["sunny", "cloudy", "rainy", "stormy"] - return f"The weather in {location} is {conditions[randint(0, 3)]} with a high of {randint(10, 30)}°C." - - -async def main() -> None: - client = OpenAIChatClient() - message = "What's the weather in Amsterdam and in Paris?" - stream = True - print(f"User: {message}") - if stream: - print("Assistant: ", end="") - async for chunk in client.get_response(message, tools=get_weather, stream=True): - if chunk.text: - print(chunk.text, end="") - print("") - else: - response = await client.get_response(message, tools=get_weather) - print(f"Assistant: {response}") - - -if __name__ == "__main__": - asyncio.run(main()) diff --git a/python/samples/02-agents/chat_client/openai_responses_client.py b/python/samples/02-agents/chat_client/openai_responses_client.py deleted file mode 100644 index ba589e1c2f..0000000000 --- a/python/samples/02-agents/chat_client/openai_responses_client.py +++ /dev/null @@ -1,47 +0,0 @@ -# Copyright (c) Microsoft. All rights reserved. - -import asyncio -from random import randint -from typing import Annotated - -from agent_framework import tool -from agent_framework.openai import OpenAIResponsesClient -from pydantic import Field - -""" -OpenAI Responses Client Direct Usage Example - -Demonstrates direct OpenAIResponsesClient usage for structured response generation with OpenAI models. -Shows function calling capabilities with custom business logic. - -""" - - -# NOTE: approval_mode="never_require" is for sample brevity. Use "always_require" in production; see samples/02-agents/tools/function_tool_with_approval.py and samples/02-agents/tools/function_tool_with_approval_and_sessions.py. -@tool(approval_mode="never_require") -def get_weather( - location: Annotated[str, Field(description="The location to get the weather for.")], -) -> str: - """Get the weather for a given location.""" - conditions = ["sunny", "cloudy", "rainy", "stormy"] - return f"The weather in {location} is {conditions[randint(0, 3)]} with a high of {randint(10, 30)}°C." - - -async def main() -> None: - client = OpenAIResponsesClient() - message = "What's the weather in Amsterdam and in Paris?" - stream = True - print(f"User: {message}") - print("Assistant: ", end="") - response = client.get_response(message, stream=stream, options={"tools": get_weather}) - if stream: - # TODO: review names of the methods, could be related to things like HTTP clients? - response.with_transform_hook(lambda chunk: print(chunk.text, end="")) - await response.get_final_response() - else: - response = await response - print(f"Assistant: {response}") - - -if __name__ == "__main__": - asyncio.run(main()) diff --git a/python/samples/02-agents/context_providers/azure_ai_search/azure_ai_with_search_context_agentic.py b/python/samples/02-agents/context_providers/azure_ai_search/azure_ai_with_search_context_agentic.py index 5a4503f920..02bac618df 100644 --- a/python/samples/02-agents/context_providers/azure_ai_search/azure_ai_with_search_context_agentic.py +++ b/python/samples/02-agents/context_providers/azure_ai_search/azure_ai_with_search_context_agentic.py @@ -75,6 +75,7 @@ async def main() -> None: if knowledge_base_name: # Use existing Knowledge Base - simplest approach search_provider = AzureAISearchContextProvider( + source_id="search_provider", endpoint=search_endpoint, api_key=search_key, credential=AzureCliCredential() if not search_key else None, @@ -91,6 +92,7 @@ async def main() -> None: if not azure_openai_resource_url: raise ValueError("AZURE_OPENAI_RESOURCE_URL required when using index_name") search_provider = AzureAISearchContextProvider( + source_id="search_provider", endpoint=search_endpoint, index_name=index_name, api_key=search_key, diff --git a/python/samples/02-agents/context_providers/azure_ai_search/azure_ai_with_search_context_semantic.py b/python/samples/02-agents/context_providers/azure_ai_search/azure_ai_with_search_context_semantic.py index 8309d5197c..e9763531fb 100644 --- a/python/samples/02-agents/context_providers/azure_ai_search/azure_ai_with_search_context_semantic.py +++ b/python/samples/02-agents/context_providers/azure_ai_search/azure_ai_with_search_context_semantic.py @@ -53,6 +53,7 @@ async def main() -> None: # Create Azure AI Search context provider with semantic mode (recommended, fast) print("Using SEMANTIC mode (hybrid search + semantic ranking, fast)\n") search_provider = AzureAISearchContextProvider( + source_id="search_provider", endpoint=search_endpoint, index_name=index_name, api_key=search_key, # Use api_key for API key auth, or credential for managed identity diff --git a/python/samples/02-agents/context_providers/mem0/mem0_basic.py b/python/samples/02-agents/context_providers/mem0/mem0_basic.py index b4e99e0a9f..773443f2be 100644 --- a/python/samples/02-agents/context_providers/mem0/mem0_basic.py +++ b/python/samples/02-agents/context_providers/mem0/mem0_basic.py @@ -39,7 +39,7 @@ async def main() -> None: name="FriendlyAssistant", instructions="You are a friendly assistant.", tools=retrieve_company_report, - context_providers=[Mem0ContextProvider(user_id=user_id)], + context_providers=[Mem0ContextProvider(source_id="mem0", user_id=user_id)], ) as agent, ): # First ask the agent to retrieve a company report with no previous context. diff --git a/python/samples/02-agents/context_providers/mem0/mem0_oss.py b/python/samples/02-agents/context_providers/mem0/mem0_oss.py index 1b03ac5fc1..8bcedc5214 100644 --- a/python/samples/02-agents/context_providers/mem0/mem0_oss.py +++ b/python/samples/02-agents/context_providers/mem0/mem0_oss.py @@ -10,7 +10,9 @@ from azure.identity.aio import AzureCliCredential from mem0 import AsyncMemory -# NOTE: approval_mode="never_require" is for sample brevity. Use "always_require" in production; see samples/02-agents/tools/function_tool_with_approval.py and samples/02-agents/tools/function_tool_with_approval_and_sessions.py. +# NOTE: approval_mode="never_require" is for sample brevity. +# Use "always_require" in production; see samples/02-agents/tools/function_tool_with_approval.py +# and samples/02-agents/tools/function_tool_with_approval_and_sessions.py. @tool(approval_mode="never_require") def retrieve_company_report(company_code: str, detailed: bool) -> str: if company_code != "CNTS": @@ -42,7 +44,7 @@ async def main() -> None: name="FriendlyAssistant", instructions="You are a friendly assistant.", tools=retrieve_company_report, - context_providers=[Mem0ContextProvider(user_id=user_id, mem0_client=local_mem0_client)], + context_providers=[Mem0ContextProvider(source_id="mem0", user_id=user_id, mem0_client=local_mem0_client)], ) as agent, ): # First ask the agent to retrieve a company report with no previous context. diff --git a/python/samples/02-agents/context_providers/mem0/mem0_sessions.py b/python/samples/02-agents/context_providers/mem0/mem0_sessions.py index cc5548e979..c2f1ddd69b 100644 --- a/python/samples/02-agents/context_providers/mem0/mem0_sessions.py +++ b/python/samples/02-agents/context_providers/mem0/mem0_sessions.py @@ -34,11 +34,14 @@ async def example_global_thread_scope() -> None: name="GlobalMemoryAssistant", instructions="You are an assistant that remembers user preferences across conversations.", tools=get_user_preferences, - context_providers=[Mem0ContextProvider( - user_id=user_id, - thread_id=global_thread_id, - scope_to_per_operation_thread_id=False, # Share memories across all sessions - )], + context_providers=[ + Mem0ContextProvider( + source_id="mem0", + user_id=user_id, + thread_id=global_thread_id, + scope_to_per_operation_thread_id=False, # Share memories across all sessions + ) + ], ) as global_agent, ): # Store some preferences in the global scope @@ -72,10 +75,13 @@ async def example_per_operation_thread_scope() -> None: name="ScopedMemoryAssistant", instructions="You are an assistant with thread-scoped memory.", tools=get_user_preferences, - context_providers=[Mem0ContextProvider( - user_id=user_id, - scope_to_per_operation_thread_id=True, # Isolate memories per session - )], + context_providers=[ + Mem0ContextProvider( + source_id="mem0", + user_id=user_id, + scope_to_per_operation_thread_id=True, # Isolate memories per session + ) + ], ) as scoped_agent, ): # Create a specific session for this scoped provider @@ -119,16 +125,22 @@ async def example_multiple_agents() -> None: AzureAIAgentClient(credential=credential).as_agent( name="PersonalAssistant", instructions="You are a personal assistant that helps with personal tasks.", - context_providers=[Mem0ContextProvider( - agent_id=agent_id_1, - )], + context_providers=[ + Mem0ContextProvider( + source_id="mem0", + agent_id=agent_id_1, + ) + ], ) as personal_agent, AzureAIAgentClient(credential=credential).as_agent( name="WorkAssistant", instructions="You are a work assistant that helps with professional tasks.", - context_providers=[Mem0ContextProvider( - agent_id=agent_id_2, - )], + context_providers=[ + Mem0ContextProvider( + source_id="mem0", + agent_id=agent_id_2, + ) + ], ) as work_agent, ): # Store personal information diff --git a/python/samples/02-agents/context_providers/redis/README.md b/python/samples/02-agents/context_providers/redis/README.md index b7b25c8d77..060061c908 100644 --- a/python/samples/02-agents/context_providers/redis/README.md +++ b/python/samples/02-agents/context_providers/redis/README.md @@ -20,7 +20,8 @@ This folder contains an example demonstrating how to use the Redis context provi 1. A running Redis with RediSearch (Redis Stack or a managed service) 2. Python environment with Agent Framework Redis extra installed -3. Optional: OpenAI API key if using vector embeddings +3. Azure AI Foundry project endpoint and Azure OpenAI Responses deployment +4. Optional: OpenAI API key if using vector embeddings ### Install the package @@ -50,6 +51,8 @@ See quickstart: `https://learn.microsoft.com/azure/redis/quickstart-create-manag ### Environment variables +- `AZURE_AI_PROJECT_ENDPOINT` (required): Azure AI Foundry project endpoint for `AzureOpenAIResponsesClient` +- `AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME` (required): Azure OpenAI Responses deployment name - `OPENAI_API_KEY` (optional): Required only if you set `vectorizer_choice="openai"` to enable hybrid search. ### Provider configuration highlights @@ -70,19 +73,26 @@ The provider supports both full‑text only and hybrid vector search: 2. Agent integration: teaches the agent a preference and verifies it is remembered across turns. 3. Agent + tool: calls a sample tool (flight search) and then asks the agent to recall details remembered from the tool output. -It uses OpenAI for both chat (via `OpenAIChatClient`) and, in some steps, optional embeddings for hybrid search. +It uses `AzureOpenAIResponsesClient` (Foundry project endpoint setup) for chat and, in some steps, optional OpenAI embeddings for hybrid search. ## How to run 1) Start Redis (see options above). For local default, ensure it's reachable at `redis://localhost:6379`. -2) Set your OpenAI key if using embeddings and for the chat client used in the sample: +2) Set Azure Foundry/OpenAI responses environment variables: + +```bash +export AZURE_AI_PROJECT_ENDPOINT="https://.services.ai.azure.com/api/projects/" +export AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME="" +``` + +3) (Optional) Set your OpenAI key if using embeddings: ```bash export OPENAI_API_KEY="" ``` -3) Run the example: +4) Run the example: ```bash python redis_basics.py @@ -109,5 +119,6 @@ You should see the agent responses and, when using embeddings, context retrieved ## Troubleshooting - Ensure at least one of `application_id`, `agent_id`, `user_id`, or `thread_id` is set; the provider requires a scope. +- Verify `AZURE_AI_PROJECT_ENDPOINT` and `AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME` are set for the chat client. - If using embeddings, verify `OPENAI_API_KEY` is set and reachable. - Make sure Redis exposes RediSearch (Redis Stack image or managed service with search enabled). diff --git a/python/samples/02-agents/context_providers/redis/azure_redis_conversation.py b/python/samples/02-agents/context_providers/redis/azure_redis_conversation.py index ce569be8cb..d5e3bd6f8e 100644 --- a/python/samples/02-agents/context_providers/redis/azure_redis_conversation.py +++ b/python/samples/02-agents/context_providers/redis/azure_redis_conversation.py @@ -13,24 +13,25 @@ Requirements: Environment Variables: - AZURE_REDIS_HOST: Your Azure Managed Redis host (e.g., myredis.redis.cache.windows.net) - - OPENAI_API_KEY: Your OpenAI API key - - OPENAI_CHAT_MODEL_ID: OpenAI model (e.g., gpt-4o-mini) + - AZURE_AI_PROJECT_ENDPOINT: Your Azure AI Foundry project endpoint + - AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME: Azure OpenAI Responses deployment name - AZURE_USER_OBJECT_ID: Your Azure AD User Object ID for authentication """ import asyncio import os -from agent_framework.openai import OpenAIChatClient +from agent_framework.azure import AzureOpenAIResponsesClient from agent_framework.redis import RedisHistoryProvider -from azure.identity.aio import AzureCliCredential +from azure.identity import AzureCliCredential +from azure.identity.aio import AzureCliCredential as AsyncAzureCliCredential from redis.credentials import CredentialProvider class AzureCredentialProvider(CredentialProvider): """Credential provider for Azure AD authentication with Redis Enterprise.""" - def __init__(self, azure_credential: AzureCliCredential, user_object_id: str): + def __init__(self, azure_credential: AsyncAzureCliCredential, user_object_id: str): self.azure_credential = azure_credential self.user_object_id = user_object_id @@ -57,24 +58,26 @@ async def main() -> None: return # Create Azure CLI credential provider (uses 'az login' credentials) - azure_credential = AzureCliCredential() + azure_credential = AsyncAzureCliCredential() credential_provider = AzureCredentialProvider(azure_credential, user_object_id) - session_id = "azure_test_session" - # Create Azure Redis history provider history_provider = RedisHistoryProvider( + source_id="redis_memory", credential_provider=credential_provider, host=redis_host, port=10000, ssl=True, - thread_id=session_id, key_prefix="chat_messages", max_messages=100, ) # Create chat client - client = OpenAIChatClient() + client = AzureOpenAIResponsesClient( + project_endpoint=os.environ["AZURE_AI_PROJECT_ENDPOINT"], + deployment_name=os.environ["AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME"], + credential=AzureCliCredential(), + ) # Create agent with Azure Redis history provider agent = client.as_agent( diff --git a/python/samples/02-agents/context_providers/redis/redis_basics.py b/python/samples/02-agents/context_providers/redis/redis_basics.py index 5f78d65320..079108bd15 100644 --- a/python/samples/02-agents/context_providers/redis/redis_basics.py +++ b/python/samples/02-agents/context_providers/redis/redis_basics.py @@ -31,13 +31,16 @@ import asyncio import os from agent_framework import Message, tool -from agent_framework.openai import OpenAIChatClient +from agent_framework.azure import AzureOpenAIResponsesClient from agent_framework.redis import RedisContextProvider +from azure.identity import AzureCliCredential from redisvl.extensions.cache.embeddings import EmbeddingsCache from redisvl.utils.vectorize import OpenAITextVectorizer -# NOTE: approval_mode="never_require" is for sample brevity. Use "always_require" in production; see samples/02-agents/tools/function_tool_with_approval.py and samples/02-agents/tools/function_tool_with_approval_and_sessions.py. +# NOTE: approval_mode="never_require" is for sample brevity. +# Use "always_require" in production; see samples/02-agents/tools/function_tool_with_approval.py +# and samples/02-agents/tools/function_tool_with_approval_and_sessions.py. @tool(approval_mode="never_require") def search_flights(origin_airport_code: str, destination_airport_code: str, detailed: bool = False) -> str: """Simulated flight-search tool to demonstrate tool memory. @@ -88,6 +91,15 @@ def search_flights(origin_airport_code: str, destination_airport_code: str, deta ) +def create_chat_client() -> AzureOpenAIResponsesClient: + """Create an Azure OpenAI Responses client using a Foundry project endpoint.""" + return AzureOpenAIResponsesClient( + project_endpoint=os.environ["AZURE_AI_PROJECT_ENDPOINT"], + deployment_name=os.environ["AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME"], + credential=AzureCliCredential(), + ) + + async def main() -> None: """Walk through provider-only, agent integration, and tool-memory scenarios. @@ -100,8 +112,8 @@ async def main() -> None: print("-" * 40) # Create a provider with partition scope and OpenAI embeddings - # Please set the OPENAI_API_KEY and OPENAI_CHAT_MODEL_ID environment variables to use the OpenAI vectorizer - # Recommend default for OPENAI_CHAT_MODEL_ID is gpt-4o-mini + # Please set OPENAI_API_KEY to use the OpenAI vectorizer. + # For chat responses, also set AZURE_AI_PROJECT_ENDPOINT and AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME. # We attach an embedding vectorizer so the provider can perform hybrid (text + vector) # retrieval. If you prefer text-only retrieval, instantiate RedisContextProvider without the @@ -115,6 +127,7 @@ async def main() -> None: # scope data for multi-tenant separation; thread_id (set later) narrows to a # specific conversation. provider = RedisContextProvider( + source_id="redis_context", redis_url="redis://localhost:6379", index_name="redis_basics", application_id="matrix_of_kermits", @@ -170,6 +183,7 @@ async def main() -> None: ) # Recreate a clean index so the next scenario starts fresh provider = RedisContextProvider( + source_id="redis_context", redis_url="redis://localhost:6379", index_name="redis_basics_2", prefix="context_2", @@ -183,7 +197,7 @@ async def main() -> None: ) # Create chat client for the agent - client = OpenAIChatClient(model_id=os.getenv("OPENAI_CHAT_MODEL_ID"), api_key=os.getenv("OPENAI_API_KEY")) + client = create_chat_client() # Create agent wired to the Redis context provider. The provider automatically # persists conversational details and surfaces relevant context on each turn. agent = client.as_agent( @@ -217,6 +231,7 @@ async def main() -> None: print("-" * 40) # Text-only provider (full-text search only). Omits vectorizer and related params. provider = RedisContextProvider( + source_id="redis_context", redis_url="redis://localhost:6379", index_name="redis_basics_3", prefix="context_3", @@ -227,7 +242,7 @@ async def main() -> None: # Create agent exposing the flight search tool. Tool outputs are captured by the # provider and become retrievable context for later turns. - client = OpenAIChatClient(model_id=os.getenv("OPENAI_CHAT_MODEL_ID"), api_key=os.getenv("OPENAI_API_KEY")) + client = create_chat_client() agent = client.as_agent( name="MemoryEnhancedAssistant", instructions=( diff --git a/python/samples/02-agents/context_providers/redis/redis_conversation.py b/python/samples/02-agents/context_providers/redis/redis_conversation.py index 2d345d9930..f25dbbbe52 100644 --- a/python/samples/02-agents/context_providers/redis/redis_conversation.py +++ b/python/samples/02-agents/context_providers/redis/redis_conversation.py @@ -17,8 +17,9 @@ Run: import asyncio import os -from agent_framework.openai import OpenAIChatClient +from agent_framework.azure import AzureOpenAIResponsesClient from agent_framework.redis import RedisContextProvider +from azure.identity import AzureCliCredential from redisvl.extensions.cache.embeddings import EmbeddingsCache from redisvl.utils.vectorize import OpenAITextVectorizer @@ -36,9 +37,8 @@ async def main() -> None: cache=EmbeddingsCache(name="openai_embeddings_cache", redis_url="redis://localhost:6379"), ) - session_id = "test_session" - provider = RedisContextProvider( + source_id="redis_context", redis_url="redis://localhost:6379", index_name="redis_conversation", prefix="redis_conversation", @@ -49,11 +49,14 @@ async def main() -> None: vector_field_name="vector", vector_algorithm="hnsw", vector_distance_metric="cosine", - thread_id=session_id, ) # Create chat client for the agent - client = OpenAIChatClient(model_id=os.getenv("OPENAI_CHAT_MODEL_ID"), api_key=os.getenv("OPENAI_API_KEY")) + client = AzureOpenAIResponsesClient( + project_endpoint=os.environ["AZURE_AI_PROJECT_ENDPOINT"], + deployment_name=os.environ["AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME"], + credential=AzureCliCredential(), + ) # Create agent wired to the Redis context provider. The provider automatically # persists conversational details and surfaces relevant context on each turn. agent = client.as_agent( diff --git a/python/samples/02-agents/context_providers/redis/redis_sessions.py b/python/samples/02-agents/context_providers/redis/redis_sessions.py index 34179048d9..aa1b7501f8 100644 --- a/python/samples/02-agents/context_providers/redis/redis_sessions.py +++ b/python/samples/02-agents/context_providers/redis/redis_sessions.py @@ -28,15 +28,24 @@ Run: import asyncio import os -import uuid -from agent_framework.openai import OpenAIChatClient +from agent_framework.azure import AzureOpenAIResponsesClient from agent_framework.redis import RedisContextProvider +from azure.identity import AzureCliCredential from redisvl.extensions.cache.embeddings import EmbeddingsCache from redisvl.utils.vectorize import OpenAITextVectorizer -# Please set the OPENAI_API_KEY and OPENAI_CHAT_MODEL_ID environment variables to use the OpenAI vectorizer -# Recommend default for OPENAI_CHAT_MODEL_ID is gpt-4o-mini +# Please set OPENAI_API_KEY to use the OpenAI vectorizer. +# For chat responses, also set AZURE_AI_PROJECT_ENDPOINT and AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME. + + +def create_chat_client() -> AzureOpenAIResponsesClient: + """Create an Azure OpenAI Responses client using a Foundry project endpoint.""" + return AzureOpenAIResponsesClient( + project_endpoint=os.environ["AZURE_AI_PROJECT_ENDPOINT"], + deployment_name=os.environ["AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME"], + credential=AzureCliCredential(), + ) async def example_global_thread_scope() -> None: @@ -44,20 +53,15 @@ async def example_global_thread_scope() -> None: print("1. Global Thread Scope Example:") print("-" * 40) - global_thread_id = str(uuid.uuid4()) - - client = OpenAIChatClient( - model_id=os.getenv("OPENAI_CHAT_MODEL_ID", "gpt-4o-mini"), - api_key=os.getenv("OPENAI_API_KEY"), - ) + client = create_chat_client() provider = RedisContextProvider( + source_id="redis_context", redis_url="redis://localhost:6379", index_name="redis_threads_global", application_id="threads_demo_app", agent_id="threads_demo_agent", user_id="threads_demo_user", - thread_id=global_thread_id, scope_to_per_operation_thread_id=False, # Share memories across all sessions ) @@ -97,10 +101,7 @@ async def example_per_operation_thread_scope() -> None: print("2. Per-Operation Thread Scope Example:") print("-" * 40) - client = OpenAIChatClient( - model_id=os.getenv("OPENAI_CHAT_MODEL_ID", "gpt-4o-mini"), - api_key=os.getenv("OPENAI_API_KEY"), - ) + client = create_chat_client() vectorizer = OpenAITextVectorizer( model="text-embedding-ada-002", @@ -109,6 +110,7 @@ async def example_per_operation_thread_scope() -> None: ) provider = RedisContextProvider( + source_id="redis_context", redis_url="redis://localhost:6379", index_name="redis_threads_dynamic", # overwrite_redis_index=True, @@ -165,10 +167,7 @@ async def example_multiple_agents() -> None: print("3. Multiple Agents with Different Thread Configurations:") print("-" * 40) - client = OpenAIChatClient( - model_id=os.getenv("OPENAI_CHAT_MODEL_ID", "gpt-4o-mini"), - api_key=os.getenv("OPENAI_API_KEY"), - ) + client = create_chat_client() vectorizer = OpenAITextVectorizer( model="text-embedding-ada-002", @@ -177,6 +176,7 @@ async def example_multiple_agents() -> None: ) personal_provider = RedisContextProvider( + source_id="redis_context", redis_url="redis://localhost:6379", index_name="redis_threads_agents", application_id="threads_demo_app", @@ -195,6 +195,7 @@ async def example_multiple_agents() -> None: ) work_provider = RedisContextProvider( + source_id="redis_context", redis_url="redis://localhost:6379", index_name="redis_threads_agents", application_id="threads_demo_app", diff --git a/python/samples/02-agents/context_providers/simple_context_provider.py b/python/samples/02-agents/context_providers/simple_context_provider.py index fd2a7ce747..db9e93bc88 100644 --- a/python/samples/02-agents/context_providers/simple_context_provider.py +++ b/python/samples/02-agents/context_providers/simple_context_provider.py @@ -1,11 +1,13 @@ # Copyright (c) Microsoft. All rights reserved. import asyncio +import os +from contextlib import suppress from typing import Any from agent_framework import Agent, AgentSession, BaseContextProvider, SessionContext, SupportsChatGetResponse -from agent_framework.azure import AzureAIClient -from azure.identity.aio import AzureCliCredential +from agent_framework.azure import AzureOpenAIResponsesClient +from azure.identity import AzureCliCredential from pydantic import BaseModel @@ -15,19 +17,13 @@ class UserInfo(BaseModel): class UserInfoMemory(BaseContextProvider): - def __init__(self, client: SupportsChatGetResponse, user_info: UserInfo | None = None, **kwargs: Any): + def __init__(self, source_id: str = "user-info-memory", *, client: SupportsChatGetResponse, **kwargs: Any): """Create the memory. If you pass in kwargs, they will be attempted to be used to create a UserInfo object. """ - super().__init__("user-info-memory") + super().__init__(source_id) self._chat_client = client - if user_info: - self.user_info = user_info - elif kwargs: - self.user_info = UserInfo.model_validate(kwargs) - else: - self.user_info = UserInfo() async def after_run( self, @@ -38,12 +34,15 @@ class UserInfoMemory(BaseContextProvider): state: dict[str, Any], ) -> None: """Extract user information from messages after each agent call.""" - request_messages = context.get_messages() + # ensure you get all the messages you want to parse from, including the input in this case. + request_messages = context.get_messages(include_input=True, include_response=True) # Check if we need to extract user info from user messages user_messages = [msg for msg in request_messages if hasattr(msg, "role") and msg.role == "user"] # type: ignore - if (self.user_info.name is None or self.user_info.age is None) and user_messages: - try: + if ( + state[self.source_id]["user_info"].name is None or state[self.source_id]["user_info"].age is None + ) and user_messages: + with suppress(Exception): # Use the chat client to extract structured information result = await self._chat_client.get_response( messages=request_messages, # type: ignore @@ -53,17 +52,12 @@ class UserInfoMemory(BaseContextProvider): ) # Update user info with extracted data - try: + with suppress(Exception): extracted = result.value - if self.user_info.name is None and extracted.name: - self.user_info.name = extracted.name - if self.user_info.age is None and extracted.age: - self.user_info.age = extracted.age - except Exception: - pass # Failed to extract, continue without updating - - except Exception: - pass # Failed to extract, continue without updating + if state[self.source_id]["user_info"].name is None and extracted.name: + state[self.source_id]["user_info"].name = extracted.name + if state[self.source_id]["user_info"].age is None and extracted.age: + state[self.source_id]["user_info"].age = extracted.age async def before_run( self, @@ -74,55 +68,52 @@ class UserInfoMemory(BaseContextProvider): state: dict[str, Any], ) -> None: """Provide user information context before each agent call.""" - instructions: list[str] = [] + if state.setdefault(self.source_id, None) is None: + state[self.source_id] = {"user_info": UserInfo()} - if self.user_info.name is None: - instructions.append( - "Ask the user for their name and politely decline to answer any questions until they provide it." - ) - else: - instructions.append(f"The user's name is {self.user_info.name}.") - - if self.user_info.age is None: - instructions.append( - "Ask the user for their age and politely decline to answer any questions until they provide it." - ) - else: - instructions.append(f"The user's age is {self.user_info.age}.") - - # Add context with additional instructions - context.extend_instructions(self.source_id, " ".join(instructions)) - - def serialize(self) -> str: - """Serialize the user info for session persistence.""" - return self.user_info.model_dump_json() + context.extend_instructions( + self.source_id, + "Ask the user for their name and politely decline to answer any questions until they provide it." + if state[self.source_id]["user_info"].name is None + else f"The user's name is {state[self.source_id]['user_info'].name}.", + ) + context.extend_instructions( + self.source_id, + "Ask the user for their age and politely decline to answer any questions until they provide it." + if state[self.source_id]["user_info"].age is None + else f"The user's age is {state[self.source_id]['user_info'].age}.", + ) async def main(): - async with AzureCliCredential() as credential: - client = AzureAIClient(credential=credential) + client = AzureOpenAIResponsesClient( + project_endpoint=os.environ["AZURE_AI_PROJECT_ENDPOINT"], + deployment_name=os.environ["AZURE_OPENAI_RESPONSES_DEPLOYMENT_NAME"], + credential=AzureCliCredential(), + ) - # Create the memory provider - memory_provider = UserInfoMemory(client) + context_name = "user-info-memory" - # Create the agent with memory - async with Agent( - client=client, - instructions="You are a friendly assistant. Always address the user by their name.", - context_providers=[memory_provider], - ) as agent: - # Create a new session for the conversation - session = agent.create_session() + # Create the memory provider + memory_provider = UserInfoMemory(context_name, client=client) - print(await agent.run("Hello, what is the square root of 9?", session=session)) - print(await agent.run("My name is Ruaidhrí", session=session)) - print(await agent.run("I am 20 years old", session=session)) + # Create the agent with memory + async with Agent( + client=client, + instructions="You are a friendly assistant. Always address the user by their name.", + context_providers=[memory_provider], + ) as agent: + # Create a new session for the conversation + session = agent.create_session() - # Access the memory component and inspect the memories - if memory_provider: - print() - print(f"MEMORY - User Name: {memory_provider.user_info.name}") - print(f"MEMORY - User Age: {memory_provider.user_info.age}") + for msg in ["Hello, what is the square root of 9?", "My name is Ruaidhrí", "I am 20 years old"]: + print(f"User: {msg}") + print(f"Assistant: {await agent.run(msg, session=session)}") + + # Access the memory component and inspect the memories + print() + print(f"MEMORY - User Name: {session.state[context_name]['user_info'].name}") + print(f"MEMORY - User Age: {session.state[context_name]['user_info'].age}") if __name__ == "__main__": diff --git a/python/samples/02-agents/providers/README.md b/python/samples/02-agents/providers/README.md new file mode 100644 index 0000000000..20a598b08e --- /dev/null +++ b/python/samples/02-agents/providers/README.md @@ -0,0 +1,19 @@ +# Provider Samples Overview + +This directory groups provider-specific samples for Agent Framework. + +| Folder | What you will find | +| --- | --- | +| [`anthropic/`](anthropic/) | Anthropic Claude samples using both `AnthropicClient` and `ClaudeAgent`, including tools, MCP, sessions, and Foundry Anthropic integration. | +| [`amazon/`](amazon/) | AWS Bedrock samples using `BedrockChatClient`, including tool-enabled agent usage. | +| [`azure_ai/`](azure_ai/) | Azure AI Foundry V2 (`azure-ai-projects`) samples with `AzureAIClient`, from basic setup to advanced patterns like search, memory, A2A, MCP, and provider methods. | +| [`azure_ai_agent/`](azure_ai_agent/) | Azure AI Foundry V1 (`azure-ai-agents`) samples with `AzureAIAgentsProvider`, including provider methods and common hosted tool integrations. | +| [`azure_openai/`](azure_openai/) | Azure OpenAI samples for Assistants, Chat, and Responses clients, with examples for sessions, tools, MCP, file search, and code interpreter. | +| [`copilotstudio/`](copilotstudio/) | Microsoft Copilot Studio agent samples, including required environment/app registration setup and explicit authentication patterns. | +| [`custom/`](custom/) | Framework extensibility samples for building custom `BaseAgent` and `BaseChatClient` implementations, including layer-composition guidance. | +| [`foundry_local/`](foundry_local/) | Foundry Local samples using `FoundryLocalClient` for local model inference with streaming, non-streaming, and tool-calling patterns. | +| [`github_copilot/`](github_copilot/) | `GitHubCopilotAgent` samples showing basic usage, session handling, permission-scoped shell/file/url access, and MCP integration. | +| [`ollama/`](ollama/) | Local Ollama samples using `OllamaChatClient` (recommended) plus OpenAI-compatible Ollama setup, including reasoning and multimodal examples. | +| [`openai/`](openai/) | OpenAI provider samples for Assistants, Chat, and Responses clients, including tools, structured output, sessions, MCP, web search, and multimodal tasks. | + +Each folder has its own README with setup requirements and file-by-file details. diff --git a/python/samples/02-agents/providers/amazon/README.md b/python/samples/02-agents/providers/amazon/README.md new file mode 100644 index 0000000000..ab10aeecfd --- /dev/null +++ b/python/samples/02-agents/providers/amazon/README.md @@ -0,0 +1,17 @@ +# Bedrock Examples + +This folder contains examples demonstrating how to use AWS Bedrock models with the Agent Framework. The sample +uses `BEDROCK_CHAT_MODEL_ID`, `BEDROCK_REGION`, and AWS credentials (`AWS_ACCESS_KEY_ID`, +`AWS_SECRET_ACCESS_KEY`, optional `AWS_SESSION_TOKEN`). + +## Examples + +| File | Description | +|------|-------------| +| [`bedrock_chat_client.py`](bedrock_chat_client.py) | Uses `BedrockChatClient` with a simple tool-enabled `Agent` to demonstrate direct Bedrock chat integration. | + +## Environment Variables + +- `BEDROCK_CHAT_MODEL_ID`: Bedrock model ID (for example, `anthropic.claude-3-5-sonnet-20240620-v1:0`) +- `BEDROCK_REGION`: AWS region (defaults to `us-east-1` if unset) +- AWS credentials via standard variables (`AWS_ACCESS_KEY_ID`, `AWS_SECRET_ACCESS_KEY`, optional `AWS_SESSION_TOKEN`) diff --git a/python/samples/02-agents/providers/amazon/bedrock_chat_client.py b/python/samples/02-agents/providers/amazon/bedrock_chat_client.py new file mode 100644 index 0000000000..1913d55dea --- /dev/null +++ b/python/samples/02-agents/providers/amazon/bedrock_chat_client.py @@ -0,0 +1,61 @@ +# Copyright (c) Microsoft. All rights reserved. + +import asyncio +from typing import Annotated + +from agent_framework import Agent, tool +from agent_framework.amazon import BedrockChatClient +from pydantic import Field + +""" +Bedrock Chat Client Example + +This sample demonstrates using `BedrockChatClient` with an agent and a simple tool. + +Environment variables used: +- `BEDROCK_CHAT_MODEL_ID` +- `BEDROCK_REGION` (defaults to `us-east-1` if unset) +- AWS credentials via standard variables (`AWS_ACCESS_KEY_ID`, `AWS_SECRET_ACCESS_KEY`, + optional `AWS_SESSION_TOKEN`) +""" + + +# NOTE: approval_mode="never_require" is for sample brevity. +# Use "always_require" in production; see samples/02-agents/tools/function_tool_with_approval.py +# and samples/02-agents/tools/function_tool_with_approval_and_sessions.py. +@tool(approval_mode="never_require") +def get_weather( + city: Annotated[str, Field(description="The city to get the weather for.")], +) -> dict[str, str]: + """Return a mock forecast for the requested city.""" + normalized_city = city.strip() or "New York" + return {"city": normalized_city, "forecast": "72F and sunny"} + + +async def main() -> None: + """Run a Bedrock-backed agent with one tool call.""" + # 1. Create an agent with Bedrock chat client and one tool. + agent = Agent( + client=BedrockChatClient(), + instructions="You are a concise travel assistant.", + name="BedrockWeatherAgent", + tool_choice="auto", + tools=[get_weather], + ) + + # 2. Run a query that uses the weather tool. + query = "Use the weather tool to check the forecast for New York." + print(f"User: {query}") + response = await agent.run(query) + print(f"Assistant: {response.text}") + + +if __name__ == "__main__": + asyncio.run(main()) + + +""" +Sample output: +User: Use the weather tool to check the forecast for New York. +Assistant: The forecast for New York is 72F and sunny. +""" diff --git a/python/samples/02-agents/providers/anthropic/anthropic_claude_basic.py b/python/samples/02-agents/providers/anthropic/anthropic_claude_basic.py index 8bea9263de..b040bbd299 100644 --- a/python/samples/02-agents/providers/anthropic/anthropic_claude_basic.py +++ b/python/samples/02-agents/providers/anthropic/anthropic_claude_basic.py @@ -19,7 +19,7 @@ import asyncio from typing import Annotated from agent_framework import tool -from agent_framework_claude import ClaudeAgent +from agent_framework.anthropic import ClaudeAgent @tool diff --git a/python/samples/02-agents/providers/anthropic/anthropic_claude_with_mcp.py b/python/samples/02-agents/providers/anthropic/anthropic_claude_with_mcp.py index f47dbd1648..455f4fd190 100644 --- a/python/samples/02-agents/providers/anthropic/anthropic_claude_with_mcp.py +++ b/python/samples/02-agents/providers/anthropic/anthropic_claude_with_mcp.py @@ -19,7 +19,7 @@ servers you trust. Use permission handlers to control what actions are allowed. import asyncio from typing import Any -from agent_framework_claude import ClaudeAgent +from agent_framework.anthropic import ClaudeAgent from claude_agent_sdk import PermissionResultAllow, PermissionResultDeny diff --git a/python/samples/02-agents/providers/anthropic/anthropic_claude_with_multiple_permissions.py b/python/samples/02-agents/providers/anthropic/anthropic_claude_with_multiple_permissions.py index e4e2d10605..e4165089c0 100644 --- a/python/samples/02-agents/providers/anthropic/anthropic_claude_with_multiple_permissions.py +++ b/python/samples/02-agents/providers/anthropic/anthropic_claude_with_multiple_permissions.py @@ -22,7 +22,7 @@ More permissions mean more potential for unintended actions. import asyncio from typing import Any -from agent_framework_claude import ClaudeAgent +from agent_framework.anthropic import ClaudeAgent from claude_agent_sdk import PermissionResultAllow, PermissionResultDeny diff --git a/python/samples/02-agents/providers/anthropic/anthropic_claude_with_session.py b/python/samples/02-agents/providers/anthropic/anthropic_claude_with_session.py index 623be4f299..6ff1fff25a 100644 --- a/python/samples/02-agents/providers/anthropic/anthropic_claude_with_session.py +++ b/python/samples/02-agents/providers/anthropic/anthropic_claude_with_session.py @@ -13,7 +13,7 @@ from random import randint from typing import Annotated from agent_framework import tool -from agent_framework_claude import ClaudeAgent +from agent_framework.anthropic import ClaudeAgent from pydantic import Field diff --git a/python/samples/02-agents/providers/anthropic/anthropic_claude_with_shell.py b/python/samples/02-agents/providers/anthropic/anthropic_claude_with_shell.py index 849a96c593..86f8be9794 100644 --- a/python/samples/02-agents/providers/anthropic/anthropic_claude_with_shell.py +++ b/python/samples/02-agents/providers/anthropic/anthropic_claude_with_shell.py @@ -14,7 +14,7 @@ Shell commands have full access to your system within the permissions of the run import asyncio from typing import Any -from agent_framework_claude import ClaudeAgent +from agent_framework.anthropic import ClaudeAgent from claude_agent_sdk import PermissionResultAllow, PermissionResultDeny diff --git a/python/samples/02-agents/providers/anthropic/anthropic_claude_with_tools.py b/python/samples/02-agents/providers/anthropic/anthropic_claude_with_tools.py index 15b9cbc5dc..8ce13ef54e 100644 --- a/python/samples/02-agents/providers/anthropic/anthropic_claude_with_tools.py +++ b/python/samples/02-agents/providers/anthropic/anthropic_claude_with_tools.py @@ -17,7 +17,7 @@ Available built-in tools: import asyncio -from agent_framework_claude import ClaudeAgent +from agent_framework.anthropic import ClaudeAgent async def main() -> None: diff --git a/python/samples/02-agents/providers/anthropic/anthropic_claude_with_url.py b/python/samples/02-agents/providers/anthropic/anthropic_claude_with_url.py index 102785ca94..bcb6b8b03e 100644 --- a/python/samples/02-agents/providers/anthropic/anthropic_claude_with_url.py +++ b/python/samples/02-agents/providers/anthropic/anthropic_claude_with_url.py @@ -16,7 +16,7 @@ URL fetching allows the agent to access any URL accessible from your network. import asyncio -from agent_framework_claude import ClaudeAgent +from agent_framework.anthropic import ClaudeAgent async def main() -> None: diff --git a/python/samples/02-agents/providers/foundry_local/README.md b/python/samples/02-agents/providers/foundry_local/README.md new file mode 100644 index 0000000000..72451a9a69 --- /dev/null +++ b/python/samples/02-agents/providers/foundry_local/README.md @@ -0,0 +1,22 @@ +# Foundry Local Examples + +This folder contains examples demonstrating how to run local models with `FoundryLocalClient` via `agent_framework.microsoft`. + +## Prerequisites + +1. Install Foundry Local and required local runtime components. +2. Install the connector package: + + ```bash + pip install agent-framework-foundry-local --pre + ``` + +## Examples + +| File | Description | +|------|-------------| +| [`foundry_local_agent.py`](foundry_local_agent.py) | Basic Foundry Local agent usage with streaming and non-streaming responses, plus function tool calling. | + +## Environment Variables + +- `FOUNDRY_LOCAL_MODEL_ID`: Optional model alias/ID to use by default when `model_id` is not passed to `FoundryLocalClient`. diff --git a/python/packages/foundry_local/samples/foundry_local_agent.py b/python/samples/02-agents/providers/foundry_local/foundry_local_agent.py similarity index 97% rename from python/packages/foundry_local/samples/foundry_local_agent.py rename to python/samples/02-agents/providers/foundry_local/foundry_local_agent.py index bca1d469d9..0ea0c15bc2 100644 --- a/python/packages/foundry_local/samples/foundry_local_agent.py +++ b/python/samples/02-agents/providers/foundry_local/foundry_local_agent.py @@ -7,7 +7,7 @@ import asyncio from random import randint from typing import TYPE_CHECKING, Annotated -from agent_framework_foundry_local import FoundryLocalClient +from agent_framework.microsoft import FoundryLocalClient if TYPE_CHECKING: from agent_framework import Agent diff --git a/python/samples/02-agents/providers/openai/openai_responses_client_image_generation.py b/python/samples/02-agents/providers/openai/openai_responses_client_image_generation.py index 1e015b3762..e6f54677e4 100644 --- a/python/samples/02-agents/providers/openai/openai_responses_client_image_generation.py +++ b/python/samples/02-agents/providers/openai/openai_responses_client_image_generation.py @@ -6,7 +6,6 @@ import tempfile import urllib.request as urllib_request from pathlib import Path -import aiofiles # pyright: ignore[reportMissingModuleSource] from agent_framework import Content from agent_framework.openai import OpenAIResponsesClient @@ -20,8 +19,11 @@ and automated visual asset generation. """ -async def save_image(output: Content) -> None: - """Save the generated image to a temporary directory.""" +def save_image(output: Content) -> None: + """Save the generated image to a temporary directory. + + This sample is simplified, usually a async aware storing method would be better. + """ filename = "generated_image.webp" file_path = Path(tempfile.gettempdir()) / filename @@ -37,15 +39,15 @@ async def save_image(output: Content) -> None: data_bytes = None else: try: - data_bytes = await asyncio.to_thread(lambda: urllib_request.urlopen(uri).read()) + data_bytes = urllib_request.urlopen(uri).read() except Exception: data_bytes = None if data_bytes is None: raise RuntimeError("Image output present but could not retrieve bytes.") - async with aiofiles.open(file_path, "wb") as f: - await f.write(data_bytes) + with open(file_path, "wb") as f: + f.write(data_bytes) print(f"Image downloaded and saved to: {file_path}") @@ -76,15 +78,15 @@ async def main() -> None: image_saved = False for message in result.messages: for content in message.contents: - if content.type == "image_generation_tool_result_tool_result" and content.outputs: + if content.type == "image_generation_tool_result" and content.outputs: output = content.outputs if isinstance(output, Content) and output.uri: - await save_image(output) + save_image(output) image_saved = True elif isinstance(output, list): for out in output: if isinstance(out, Content) and out.uri: - await save_image(out) + save_image(out) image_saved = True break if image_saved: diff --git a/python/samples/02-agents/response_stream.py b/python/samples/02-agents/response_stream.py index 1b26ac5e90..840dc18b26 100644 --- a/python/samples/02-agents/response_stream.py +++ b/python/samples/02-agents/response_stream.py @@ -3,7 +3,7 @@ import asyncio from collections.abc import AsyncIterable, Sequence -from agent_framework import ChatResponse, ChatResponseUpdate, Content, ResponseStream, Role +from agent_framework import ChatResponse, ChatResponseUpdate, Content, Message, ResponseStream """ResponseStream: A Deep Dive @@ -256,8 +256,7 @@ async def main() -> None: """Result hook that wraps the response text in quotes.""" if response.text: return ChatResponse( - messages=f'"{response.text}"', - role=Role.ASSISTANT, + messages=[Message(text=f'"{response.text}"', role="assistant")], additional_properties=response.additional_properties, ) return response @@ -294,8 +293,7 @@ async def main() -> None: # In real code, this would create an AgentResponse text = "".join(u.text or "" for u in updates) return ChatResponse( - text=f"[AGENT FINAL] {text}", - role=Role.ASSISTANT, + messages=[Message(text=f"[AGENT FINAL] {text}", role="assistant")], additional_properties={"layer": "agent"}, ) diff --git a/python/samples/02-agents/typed_options.py b/python/samples/02-agents/typed_options.py index e111222601..58bb116055 100644 --- a/python/samples/02-agents/typed_options.py +++ b/python/samples/02-agents/typed_options.py @@ -22,6 +22,11 @@ which provides: The sample shows usage with both OpenAI and Anthropic clients, demonstrating how provider-specific options work for ChatClient and Agent. But the same approach works for other providers too. + +The following environment variables are used: + - ANTHROPIC_API_KEY=... + - OPENAI_API_KEY=... + """ @@ -109,14 +114,13 @@ async def demo_openai_chat_client_reasoning_models() -> None: print("\n=== OpenAI ChatClient with TypedDict Options ===\n") # Create OpenAI client - client = OpenAIChatClient[OpenAIReasoningChatOptions]() + client = OpenAIChatClient[OpenAIReasoningChatOptions](model_id="o3") # With specific options, you get full IDE autocomplete! # Try typing `client.get_response("Hello", options={` and see the suggestions response = await client.get_response( "What is 2 + 2?", options={ - "model_id": "o3", "max_tokens": 100, "allow_multiple_tool_calls": True, # OpenAI-specific options work: @@ -140,12 +144,11 @@ async def demo_openai_agent() -> None: # or on the client when constructing the client instance: # client = OpenAIChatClient[OpenAIReasoningChatOptions]() agent = Agent[OpenAIReasoningChatOptions]( - client=OpenAIChatClient(), + client=OpenAIChatClient(model_id="o3"), name="weather-assistant", instructions="You are a helpful assistant. Answer concisely.", # Options can be set at construction time default_options={ - "model_id": "o3", "max_tokens": 100, "allow_multiple_tool_calls": True, # OpenAI-specific options work: diff --git a/python/uv.lock b/python/uv.lock index 40c021172b..a5bccab41e 100644 --- a/python/uv.lock +++ b/python/uv.lock @@ -343,11 +343,14 @@ all = [ { name = "agent-framework-azure-ai", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, { name = "agent-framework-azure-ai-search", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, { name = "agent-framework-azurefunctions", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, + { name = "agent-framework-bedrock", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, { name = "agent-framework-chatkit", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, + { name = "agent-framework-claude", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, { name = "agent-framework-copilotstudio", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, { name = "agent-framework-declarative", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, { name = "agent-framework-devui", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, { name = "agent-framework-durabletask", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, + { name = "agent-framework-foundry-local", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, { name = "agent-framework-github-copilot", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, { name = "agent-framework-lab", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, { name = "agent-framework-mem0", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, @@ -365,11 +368,14 @@ requires-dist = [ { name = "agent-framework-azure-ai", marker = "extra == 'all'", editable = "packages/azure-ai" }, { name = "agent-framework-azure-ai-search", marker = "extra == 'all'", editable = "packages/azure-ai-search" }, { name = "agent-framework-azurefunctions", marker = "extra == 'all'", editable = "packages/azurefunctions" }, + { name = "agent-framework-bedrock", marker = "extra == 'all'", editable = "packages/bedrock" }, { name = "agent-framework-chatkit", marker = "extra == 'all'", editable = "packages/chatkit" }, + { name = "agent-framework-claude", marker = "extra == 'all'", editable = "packages/claude" }, { name = "agent-framework-copilotstudio", marker = "extra == 'all'", editable = "packages/copilotstudio" }, { name = "agent-framework-declarative", marker = "extra == 'all'", editable = "packages/declarative" }, { name = "agent-framework-devui", marker = "extra == 'all'", editable = "packages/devui" }, { name = "agent-framework-durabletask", marker = "extra == 'all'", editable = "packages/durabletask" }, + { name = "agent-framework-foundry-local", marker = "extra == 'all'", editable = "packages/foundry_local" }, { name = "agent-framework-github-copilot", marker = "extra == 'all'", editable = "packages/github_copilot" }, { name = "agent-framework-lab", marker = "extra == 'all'", editable = "packages/lab" }, { name = "agent-framework-mem0", marker = "extra == 'all'", editable = "packages/mem0" }, @@ -1356,7 +1362,7 @@ name = "clr-loader" version = "0.2.10" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "cffi", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, + { name = "cffi", marker = "(python_full_version < '3.14' and sys_platform == 'darwin') or (python_full_version < '3.14' and sys_platform == 'linux') or (python_full_version < '3.14' and sys_platform == 'win32')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/18/24/c12faf3f61614b3131b5c98d3bf0d376b49c7feaa73edca559aeb2aee080/clr_loader-0.2.10.tar.gz", hash = "sha256:81f114afbc5005bafc5efe5af1341d400e22137e275b042a8979f3feb9fc9446", size = 83605, upload-time = "2026-01-03T23:13:06.984Z" } wheels = [ @@ -1835,7 +1841,7 @@ name = "exceptiongroup" version = "1.3.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "typing-extensions", marker = "(python_full_version < '3.13' and sys_platform == 'darwin') or (python_full_version < '3.13' and sys_platform == 'linux') or (python_full_version < '3.13' and sys_platform == 'win32')" }, + { name = "typing-extensions", marker = "(python_full_version < '3.11' and sys_platform == 'darwin') or (python_full_version < '3.11' and sys_platform == 'linux') or (python_full_version < '3.11' and sys_platform == 'win32')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/50/79/66800aadf48771f6b62f7eb014e352e5d06856655206165d775e675a02c9/exceptiongroup-1.3.1.tar.gz", hash = "sha256:8b412432c6055b0b7d14c310000ae93352ed6754f70fa8f7c34141f91c4e3219", size = 30371, upload-time = "2025-11-21T23:01:54.787Z" } wheels = [ @@ -4571,8 +4577,8 @@ name = "powerfx" version = "0.0.34" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "cffi", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, - { name = "pythonnet", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, + { name = "cffi", marker = "(python_full_version < '3.14' and sys_platform == 'darwin') or (python_full_version < '3.14' and sys_platform == 'linux') or (python_full_version < '3.14' and sys_platform == 'win32')" }, + { name = "pythonnet", marker = "(python_full_version < '3.14' and sys_platform == 'darwin') or (python_full_version < '3.14' and sys_platform == 'linux') or (python_full_version < '3.14' and sys_platform == 'win32')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/9f/fb/6c4bf87e0c74ca1c563921ce89ca1c5785b7576bca932f7255cdf81082a7/powerfx-0.0.34.tar.gz", hash = "sha256:956992e7afd272657ed16d80f4cad24ec95d9e4a79fb9dfa4a068a09e136af32", size = 3237555, upload-time = "2025-12-22T15:50:59.682Z" } wheels = [ @@ -5221,7 +5227,7 @@ name = "pythonnet" version = "3.0.5" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "clr-loader", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, + { name = "clr-loader", marker = "(python_full_version < '3.14' and sys_platform == 'darwin') or (python_full_version < '3.14' and sys_platform == 'linux') or (python_full_version < '3.14' and sys_platform == 'win32')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/9a/d6/1afd75edd932306ae9bd2c2d961d603dc2b52fcec51b04afea464f1f6646/pythonnet-3.0.5.tar.gz", hash = "sha256:48e43ca463941b3608b32b4e236db92d8d40db4c58a75ace902985f76dac21cf", size = 239212, upload-time = "2024-12-13T08:30:44.393Z" } wheels = [ From 54a67d96cd297c7d49ff7de9c8641534034a5468 Mon Sep 17 00:00:00 2001 From: westey <164392973+westey-m@users.noreply.github.com> Date: Mon, 16 Feb 2026 19:02:47 +0000 Subject: [PATCH 16/22] .NET: [BREAKING] Refactor providers to move common functionality to base (#3900) * Move common functionality to provider base classes Co-Authored-By: Copilot <175728472+Copilot@users.noreply.github.com> * Address PR comments. --------- Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- .../Program.cs | 30 +- .../Program.cs | 46 +-- .../AIContextProvider.cs | 159 ++++++++- .../AIContextProvider{TState}.cs | 87 +++++ .../ChatHistoryProvider.cs | 142 +++++++- .../ChatHistoryProvider{TState}.cs | 87 +++++ .../InMemoryChatHistoryProvider.cs | 71 +--- .../InMemoryChatHistoryProviderOptions.cs | 2 +- .../CosmosChatHistoryProvider.cs | 95 ++---- .../Microsoft.Agents.AI.Mem0/Mem0Provider.cs | 98 ++---- .../WorkflowChatHistoryProvider.cs | 41 +-- .../Memory/ChatHistoryMemoryProvider.cs | 82 +---- .../Microsoft.Agents.AI/TextSearchProvider.cs | 74 ++--- .../AIContextProviderTStateTests.cs | 199 +++++++++++ .../AIContextProviderTests.cs | 310 +++++++++++++++++- .../ChatHistoryProviderTStateTests.cs | 195 +++++++++++ .../ChatHistoryProviderTests.cs | 275 +++++++++++++++- .../InMemoryChatHistoryProviderTests.cs | 2 +- .../CosmosChatHistoryProviderTests.cs | 26 +- .../ChatClient/ChatClientAgentOptionsTests.cs | 8 +- .../ChatClient/ChatClientAgentTests.cs | 22 +- ...hatClientAgent_BackgroundResponsesTests.cs | 16 +- ...tClientAgent_ChatHistoryManagementTests.cs | 8 +- 23 files changed, 1593 insertions(+), 482 deletions(-) create mode 100644 dotnet/src/Microsoft.Agents.AI.Abstractions/AIContextProvider{TState}.cs create mode 100644 dotnet/src/Microsoft.Agents.AI.Abstractions/ChatHistoryProvider{TState}.cs create mode 100644 dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/AIContextProviderTStateTests.cs create mode 100644 dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/ChatHistoryProviderTStateTests.cs diff --git a/dotnet/samples/GettingStarted/AgentWithMemory/AgentWithMemory_Step03_CustomMemory/Program.cs b/dotnet/samples/GettingStarted/AgentWithMemory/AgentWithMemory_Step03_CustomMemory/Program.cs index bac72d9b31..9311e56c88 100644 --- a/dotnet/samples/GettingStarted/AgentWithMemory/AgentWithMemory_Step03_CustomMemory/Program.cs +++ b/dotnet/samples/GettingStarted/AgentWithMemory/AgentWithMemory_Step03_CustomMemory/Program.cs @@ -86,27 +86,25 @@ namespace SampleApp /// /// Sample memory component that can remember a user's name and age. /// - internal sealed class UserInfoMemory : AIContextProvider + internal sealed class UserInfoMemory : AIContextProvider { private readonly IChatClient _chatClient; - private readonly Func _stateInitializer; public UserInfoMemory(IChatClient chatClient, Func? stateInitializer = null) + : base(stateInitializer ?? (_ => new UserInfo()), null, null, null, null) { this._chatClient = chatClient; - this._stateInitializer = stateInitializer ?? (_ => new UserInfo()); } public UserInfo GetUserInfo(AgentSession session) - => session.StateBag.GetValue(nameof(UserInfoMemory)) ?? new UserInfo(); + => this.GetOrInitializeState(session); public void SetUserInfo(AgentSession session, UserInfo userInfo) - => session.StateBag.SetValue(nameof(UserInfoMemory), userInfo); + => this.SaveState(session, userInfo); - protected override async ValueTask InvokedCoreAsync(InvokedContext context, CancellationToken cancellationToken = default) + protected override async ValueTask StoreAIContextAsync(InvokedContext context, CancellationToken cancellationToken = default) { - var userInfo = context.Session?.StateBag.GetValue(nameof(UserInfoMemory)) - ?? this._stateInitializer.Invoke(context.Session); + var userInfo = this.GetOrInitializeState(context.Session); // Try and extract the user name and age from the message if we don't have it already and it's a user message. if ((userInfo.UserName is null || userInfo.UserAge is null) && context.RequestMessages.Any(x => x.Role == ChatRole.User)) @@ -123,20 +121,14 @@ namespace SampleApp userInfo.UserAge ??= result.Result.UserAge; } - context.Session?.StateBag.SetValue(nameof(UserInfoMemory), userInfo); + this.SaveState(context.Session, userInfo); } - protected override ValueTask InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default) + protected override ValueTask ProvideAIContextAsync(InvokingContext context, CancellationToken cancellationToken = default) { - var inputContext = context.AIContext; - var userInfo = context.Session?.StateBag.GetValue(nameof(UserInfoMemory)) - ?? this._stateInitializer.Invoke(context.Session); + var userInfo = this.GetOrInitializeState(context.Session); StringBuilder instructions = new(); - if (!string.IsNullOrEmpty(inputContext.Instructions)) - { - instructions.AppendLine(inputContext.Instructions); - } // If we don't already know the user's name and age, add instructions to ask for them, otherwise just provide what we have to the context. instructions @@ -151,9 +143,7 @@ namespace SampleApp return new ValueTask(new AIContext { - Instructions = instructions.ToString(), - Messages = inputContext.Messages, - Tools = inputContext.Tools + Instructions = instructions.ToString() }); } } diff --git a/dotnet/samples/GettingStarted/Agents/Agent_Step07_3rdPartyChatHistoryStorage/Program.cs b/dotnet/samples/GettingStarted/Agents/Agent_Step07_3rdPartyChatHistoryStorage/Program.cs index cc4ca1a6ed..1c25c36d22 100644 --- a/dotnet/samples/GettingStarted/Agents/Agent_Step07_3rdPartyChatHistoryStorage/Program.cs +++ b/dotnet/samples/GettingStarted/Agents/Agent_Step07_3rdPartyChatHistoryStorage/Program.cs @@ -76,45 +76,23 @@ namespace SampleApp /// State (the session DB key) is stored in the so it roundtrips /// automatically with session serialization. /// - internal sealed class VectorChatHistoryProvider : ChatHistoryProvider + internal sealed class VectorChatHistoryProvider : ChatHistoryProvider { private readonly VectorStore _vectorStore; - private readonly Func _stateInitializer; - private readonly string _stateKey; - - /// - public override string StateKey => this._stateKey; public VectorChatHistoryProvider( VectorStore vectorStore, Func? stateInitializer = null, string? stateKey = null) + : base(stateInitializer: stateInitializer ?? (_ => new State(Guid.NewGuid().ToString("N"))), stateKey: stateKey, jsonSerializerOptions: null, provideOutputMessageFilter: null, storeInputMessageFilter: null) { this._vectorStore = vectorStore ?? throw new ArgumentNullException(nameof(vectorStore)); - this._stateInitializer = stateInitializer ?? (_ => new State(Guid.NewGuid().ToString("N"))); - this._stateKey = stateKey ?? base.StateKey; } public string GetSessionDbKey(AgentSession session) => this.GetOrInitializeState(session).SessionDbKey; - private State GetOrInitializeState(AgentSession? session) - { - if (session?.StateBag.TryGetValue(this._stateKey, out var state) is true && state is not null) - { - return state; - } - - state = this._stateInitializer(session); - if (session is not null) - { - session.StateBag.SetValue(this._stateKey, state); - } - - return state; - } - - protected override async ValueTask> InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default) + protected override async ValueTask> ProvideChatHistoryAsync(InvokingContext context, CancellationToken cancellationToken = default) { var state = this.GetOrInitializeState(context.Session); var collection = this._vectorStore.GetCollection("ChatHistory"); @@ -129,29 +107,17 @@ namespace SampleApp var messages = records.ConvertAll(x => JsonSerializer.Deserialize(x.SerializedMessage!)!); messages.Reverse(); - return messages - .Select(message => message.WithAgentRequestMessageSource(AgentRequestMessageSourceType.ChatHistory, this.GetType().FullName!)) - .Concat(context.RequestMessages); + return messages; } - protected override async ValueTask InvokedCoreAsync(InvokedContext context, CancellationToken cancellationToken = default) + protected override async ValueTask StoreChatHistoryAsync(InvokedContext context, CancellationToken cancellationToken = default) { - // Don't store messages if the request failed. - if (context.InvokeException is not null) - { - return; - } - var state = this.GetOrInitializeState(context.Session); var collection = this._vectorStore.GetCollection("ChatHistory"); await collection.EnsureCollectionExistsAsync(cancellationToken); - // Add both request and response messages to the store, excluding messages that came from chat history. - // Optionally messages produced by the AIContextProvider can also be persisted (not shown). - var allNewMessages = context.RequestMessages - .Where(m => m.GetAgentRequestMessageSourceType() != AgentRequestMessageSourceType.ChatHistory) - .Concat(context.ResponseMessages ?? []); + var allNewMessages = context.RequestMessages.Concat(context.ResponseMessages ?? []); await collection.UpsertAsync(allNewMessages.Select(x => new ChatHistoryItem() { diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions/AIContextProvider.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions/AIContextProvider.cs index b201e9f8e1..023f6c8e5f 100644 --- a/dotnet/src/Microsoft.Agents.AI.Abstractions/AIContextProvider.cs +++ b/dotnet/src/Microsoft.Agents.AI.Abstractions/AIContextProvider.cs @@ -2,6 +2,7 @@ using System; using System.Collections.Generic; +using System.Linq; using System.Threading; using System.Threading.Tasks; using Microsoft.Extensions.AI; @@ -10,7 +11,7 @@ using Microsoft.Shared.Diagnostics; namespace Microsoft.Agents.AI; /// -/// Provides an abstract base class for components that enhance AI context management during agent invocations. +/// Provides an abstract base class for components that enhance AI context during agent invocations. /// /// /// @@ -30,6 +31,25 @@ namespace Microsoft.Agents.AI; /// public abstract class AIContextProvider { + private static IEnumerable DefaultExternalOnlyFilter(IEnumerable messages) + => messages.Where(m => m.GetAgentRequestMessageSourceType() == AgentRequestMessageSourceType.External); + + private readonly Func, IEnumerable> _provideInputMessageFilter; + private readonly Func, IEnumerable> _storeInputMessageFilter; + + /// + /// Initializes a new instance of the class. + /// + /// An optional filter function to apply to input messages before providing context via . If not set, defaults to including only messages. + /// An optional filter function to apply to request messages before storing context via . If not set, defaults to including only messages. + protected AIContextProvider( + Func, IEnumerable>? provideInputMessageFilter = null, + Func, IEnumerable>? storeInputMessageFilter = null) + { + this._provideInputMessageFilter = provideInputMessageFilter ?? DefaultExternalOnlyFilter; + this._storeInputMessageFilter = storeInputMessageFilter ?? DefaultExternalOnlyFilter; + } + /// /// Gets the key used to store the provider state in the . /// @@ -58,7 +78,7 @@ public abstract class AIContextProvider /// /// public ValueTask InvokingAsync(InvokingContext context, CancellationToken cancellationToken = default) - => this.InvokingCoreAsync(context, cancellationToken); + => this.InvokingCoreAsync(Throw.IfNull(context), cancellationToken); /// /// Called at the start of agent invocation to provide additional context. @@ -76,8 +96,96 @@ public abstract class AIContextProvider /// Injecting contextual messages from conversation history /// /// + /// + /// The default implementation of this method filters the input messages using the configured provide-input message filter + /// (which defaults to including only messages), + /// then calls to get additional context, + /// stamps any messages from the returned context with source attribution, + /// and merges the returned context with the original (unfiltered) input context (concatenating instructions, messages, and tools). + /// For most scenarios, overriding is sufficient to provide additional context, + /// while still benefiting from the default filtering, merging and source stamping behavior. + /// However, for scenarios that require more control over context filtering, merging or source stamping, overriding this method + /// allows you to directly control the full returned for the invocation. + /// /// - protected abstract ValueTask InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default); + protected virtual async ValueTask InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default) + { + var inputContext = context.AIContext; + + // Create a filtered context for ProvideAIContextAsync, filtering input messages + // to exclude non-external messages (e.g. chat history, other AI context provider messages). + var filteredContext = new InvokingContext( + context.Agent, + context.Session, + new AIContext + { + Instructions = inputContext.Instructions, + Messages = inputContext.Messages is not null ? this._provideInputMessageFilter(inputContext.Messages) : null, + Tools = inputContext.Tools + }); + + var provided = await this.ProvideAIContextAsync(filteredContext, cancellationToken).ConfigureAwait(false); + + var mergedInstructions = (inputContext.Instructions, provided.Instructions) switch + { + (null, null) => null, + (string a, null) => a, + (null, string b) => b, + (string a, string b) => a + "\n" + b + }; + + var providedMessages = provided.Messages is not null + ? provided.Messages.Select(m => m.WithAgentRequestMessageSource(AgentRequestMessageSourceType.AIContextProvider, this.GetType().FullName!)) + : null; + + var mergedMessages = (inputContext.Messages, providedMessages) switch + { + (null, null) => null, + (var a, null) => a, + (null, var b) => b, + (var a, var b) => a.Concat(b) + }; + + var mergedTools = (inputContext.Tools, provided.Tools) switch + { + (null, null) => null, + (var a, null) => a, + (null, var b) => b, + (var a, var b) => a.Concat(b) + }; + + return new AIContext + { + Instructions = mergedInstructions, + Messages = mergedMessages, + Tools = mergedTools + }; + } + + /// + /// When overridden in a derived class, provides additional AI context to be merged with the input context for the current invocation. + /// + /// + /// + /// This method is called from . + /// Note that can be overridden to directly control context merging and source stamping, in which case + /// it is up to the implementer to call this method as needed to retrieve the additional context. + /// + /// + /// In contrast with , this method only returns additional context to be merged with the input, + /// while is responsible for returning the full merged for the invocation. + /// + /// + /// Contains the request context including the caller provided messages that will be used by the agent for this invocation. + /// The to monitor for cancellation requests. The default is . + /// + /// A task that represents the asynchronous operation. The task result contains an + /// with additional context to be merged with the input context. + /// + protected virtual ValueTask ProvideAIContextAsync(InvokingContext context, CancellationToken cancellationToken = default) + { + return new ValueTask(new AIContext()); + } /// /// Called at the end of the agent invocation to process the invocation results. @@ -106,7 +214,7 @@ public abstract class AIContextProvider /// /// public ValueTask InvokedAsync(InvokedContext context, CancellationToken cancellationToken = default) - => this.InvokedCoreAsync(context, cancellationToken); + => this.InvokedCoreAsync(Throw.IfNull(context), cancellationToken); /// /// Called at the end of the agent invocation to process the invocation results. @@ -128,9 +236,50 @@ public abstract class AIContextProvider /// This method is called regardless of whether the invocation succeeded or failed. /// To check if the invocation was successful, inspect the property. /// + /// + /// The default implementation of this method skips execution for any invocation failures, + /// filters the request messages using the configured store-input message filter + /// (which defaults to including only messages), + /// and calls to process the invocation results. + /// For most scenarios, overriding is sufficient to process invocation results, + /// while still benefiting from the default error handling and filtering behavior. + /// However, for scenarios that require more control over error handling or message filtering, overriding this method + /// allows you to directly control the processing of invocation results. + /// /// protected virtual ValueTask InvokedCoreAsync(InvokedContext context, CancellationToken cancellationToken = default) - => default; + { + if (context.InvokeException is not null) + { + return default; + } + + var subContext = new InvokedContext(context.Agent, context.Session, this._storeInputMessageFilter(context.RequestMessages), context.ResponseMessages!); + return this.StoreAIContextAsync(subContext, cancellationToken); + } + + /// + /// When overridden in a derived class, processes invocation results at the end of the agent invocation. + /// + /// Contains the invocation context including request messages, response messages, and any exception that occurred. + /// The to monitor for cancellation requests. The default is . + /// A task that represents the asynchronous operation. + /// + /// + /// This method is called from . + /// Note that can be overridden to directly control error handling, in which case + /// it is up to the implementer to call this method as needed to process the invocation results. + /// + /// + /// In contrast with , this method only processes the invocation results, + /// while is also responsible for error handling. + /// + /// + /// The default implementation of only calls this method if the invocation succeeded. + /// + /// + protected virtual ValueTask StoreAIContextAsync(InvokedContext context, CancellationToken cancellationToken = default) => + default; /// Asks the for an object of the specified type . /// The type of object being requested. diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions/AIContextProvider{TState}.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions/AIContextProvider{TState}.cs new file mode 100644 index 0000000000..136724f44b --- /dev/null +++ b/dotnet/src/Microsoft.Agents.AI.Abstractions/AIContextProvider{TState}.cs @@ -0,0 +1,87 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Collections.Generic; +using System.Text.Json; +using Microsoft.Extensions.AI; + +namespace Microsoft.Agents.AI; + +/// +/// Provides an abstract base class for components that enhance AI context during agent invocations with support for maintaining provider state of type . +/// +/// The type of the state to be maintained by the context provider. Must be a reference type. +/// +/// This class extends by introducing a strongly-typed state management mechanism, allowing derived classes to maintain and persist custom state information across invocations. +/// The state is stored in the session's StateBag using a configurable key and JSON serialization options, enabling seamless integration with the agent session lifecycle. +/// +public abstract class AIContextProvider : AIContextProvider + where TState : class +{ + private readonly Func _stateInitializer; + private readonly string _stateKey; + private readonly JsonSerializerOptions _jsonSerializerOptions; + + /// + /// Initializes a new instance of the class. + /// + /// A function to initialize the state for the context provider. + /// The key used to store the state in the session's StateBag. + /// Options for JSON serialization and deserialization of the state. + /// An optional filter function to apply to input messages before providing context. If not set, defaults to including only messages. + /// An optional filter function to apply to request messages before storing context. If not set, defaults to including only messages. + protected AIContextProvider( + Func stateInitializer, + string? stateKey, + JsonSerializerOptions? jsonSerializerOptions, + Func, IEnumerable>? provideInputMessageFilter, + Func, IEnumerable>? storeInputMessageFilter) + : base(provideInputMessageFilter, storeInputMessageFilter) + { + this._stateInitializer = stateInitializer; + this._jsonSerializerOptions = jsonSerializerOptions ?? AgentAbstractionsJsonUtilities.DefaultOptions; + this._stateKey = stateKey ?? this.GetType().Name; + } + + /// + public override string StateKey => this._stateKey; + + /// + /// Gets the state from the session's StateBag, or initializes it using the state initializer if not present. + /// + /// The agent session containing the StateBag. + /// The provider state. + protected virtual TState GetOrInitializeState(AgentSession? session) + { + if (session?.StateBag.TryGetValue(this._stateKey, out var state, this._jsonSerializerOptions) is true && state is not null) + { + return state; + } + + state = this._stateInitializer(session); + if (session is not null) + { + session.StateBag.SetValue(this._stateKey, state, this._jsonSerializerOptions); + } + + return state; + } + + /// + /// Saves the specified state to the session's StateBag using the configured state key and JSON serializer options. + /// If the session is null, this method does nothing. + /// + /// + /// This method provides a convenient way for derived classes to persist state changes back to the session after processing. + /// It abstracts away the details of how state is stored in the session, allowing derived classes to focus on their specific logic. + /// + /// The agent session containing the StateBag. + /// The state to be saved. + protected virtual void SaveState(AgentSession? session, TState state) + { + if (session is not null) + { + session.StateBag.SetValue(this._stateKey, state, this._jsonSerializerOptions); + } + } +} diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions/ChatHistoryProvider.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions/ChatHistoryProvider.cs index d16ca69528..ad3f3aacfb 100644 --- a/dotnet/src/Microsoft.Agents.AI.Abstractions/ChatHistoryProvider.cs +++ b/dotnet/src/Microsoft.Agents.AI.Abstractions/ChatHistoryProvider.cs @@ -2,6 +2,7 @@ using System; using System.Collections.Generic; +using System.Linq; using System.Threading; using System.Threading.Tasks; using Microsoft.Extensions.AI; @@ -39,6 +40,25 @@ namespace Microsoft.Agents.AI; /// public abstract class ChatHistoryProvider { + private static IEnumerable DefaultExcludeChatHistoryFilter(IEnumerable messages) + => messages.Where(m => m.GetAgentRequestMessageSourceType() != AgentRequestMessageSourceType.ChatHistory); + + private readonly Func, IEnumerable>? _provideOutputMessageFilter; + private readonly Func, IEnumerable> _storeInputMessageFilter; + + /// + /// Initializes a new instance of the class. + /// + /// An optional filter function to apply to messages when retrieving them from the chat history. + /// An optional filter function to apply to messages before storing them in the chat history. If not set, defaults to excluding messages with source type . + protected ChatHistoryProvider( + Func, IEnumerable>? provideOutputMessageFilter = null, + Func, IEnumerable>? storeInputMessageFilter = null) + { + this._provideOutputMessageFilter = provideOutputMessageFilter; + this._storeInputMessageFilter = storeInputMessageFilter ?? DefaultExcludeChatHistoryFilter; + } + /// /// Gets the key used to store the provider state in the . /// @@ -50,20 +70,16 @@ public abstract class ChatHistoryProvider public virtual string StateKey => this.GetType().Name; /// - /// Called at the start of agent invocation to provide messages from the chat history as context for the next agent invocation. + /// Called at the start of agent invocation to provide messages for the next agent invocation. /// /// Contains the request context including the caller provided messages that will be used by the agent for this invocation. /// The to monitor for cancellation requests. The default is . /// /// A task that represents the asynchronous operation. The task result contains a collection of - /// instances in ascending chronological order (oldest first). + /// instances that will be used for the agent invocation. /// /// /// - /// Messages are returned in chronological order to maintain proper conversation flow and context for the agent. - /// The oldest messages appear first in the collection, followed by more recent messages. - /// - /// /// If the total message history becomes very large, implementations should apply appropriate strategies to manage /// storage constraints, such as: /// @@ -75,23 +91,19 @@ public abstract class ChatHistoryProvider /// /// public ValueTask> InvokingAsync(InvokingContext context, CancellationToken cancellationToken = default) - => this.InvokingCoreAsync(context, cancellationToken); + => this.InvokingCoreAsync(Throw.IfNull(context), cancellationToken); /// - /// Called at the start of agent invocation to provide messages from the chat history as context for the next agent invocation. + /// Called at the start of agent invocation to provide messages for the next agent invocation. /// /// Contains the request context including the caller provided messages that will be used by the agent for this invocation. /// The to monitor for cancellation requests. The default is . /// /// A task that represents the asynchronous operation. The task result contains a collection of - /// instances in ascending chronological order (oldest first). + /// instances that will be used for the agent invocation. /// /// /// - /// Messages are returned in chronological order to maintain proper conversation flow and context for the agent. - /// The oldest messages appear first in the collection, followed by more recent messages. - /// - /// /// If the total message history becomes very large, implementations should apply appropriate strategies to manage /// storage constraints, such as: /// @@ -102,11 +114,54 @@ public abstract class ChatHistoryProvider /// /// /// - /// Each instance should be associated with a single to ensure proper message isolation - /// and context management. + /// The default implementation of this method, calls to get the chat history messages, applies the optional retrieval output filter, + /// and merges the returned messages with the caller provided messages (with chat history messages appearing first) before returning the full message list to be used for the invocation. + /// For most scenarios, overriding is sufficient to return the desired chat history messages, while still benefiting from the default merging and filtering behavior. + /// However, for scenarios that require more control over message filtering, merging or source stamping, overriding this method allows you to directly control the full set of messages returned for the invocation. /// /// - protected abstract ValueTask> InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default); + protected virtual async ValueTask> InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default) + { + var output = await this.ProvideChatHistoryAsync(context, cancellationToken).ConfigureAwait(false); + + if (this._provideOutputMessageFilter is not null) + { + output = this._provideOutputMessageFilter(output); + } + + return output + .Select(message => message.WithAgentRequestMessageSource(AgentRequestMessageSourceType.ChatHistory, this.GetType().FullName!)) + .Concat(context.RequestMessages); + } + + /// + /// When overridden in a derived class, provides the chat history messages to be used for the current invocation. + /// + /// + /// + /// This method is called from . + /// Note that can be overridden to directly control message filtering, merging and source stamping, in which case + /// it is up to the implementer to call this method as needed to retrieve the unfiltered/unmerged chat history messages. + /// + /// + /// In contrast with , this method only returns additional messages to be added to the request, + /// while is responsible for returning the full set of messages to be used for the invocation (including caller provided messages). + /// + /// + /// Messages are returned in chronological order to maintain proper conversation flow and context for the agent. + /// The oldest messages appear first in the collection, followed by more recent messages. + /// + /// + /// Contains the request context including the caller provided messages that will be used by the agent for this invocation. + /// The to monitor for cancellation requests. The default is . + /// + /// A task that represents the asynchronous operation. The task result contains a collection of + /// instances in ascending chronological order (oldest first). + /// + protected virtual ValueTask> ProvideChatHistoryAsync(InvokingContext context, CancellationToken cancellationToken = default) + { + return new ValueTask>([]); + } /// /// Called at the end of the agent invocation to add new messages to the chat history. @@ -134,7 +189,7 @@ public abstract class ChatHistoryProvider /// /// public ValueTask InvokedAsync(InvokedContext context, CancellationToken cancellationToken = default) => - this.InvokedCoreAsync(context, cancellationToken); + this.InvokedCoreAsync(Throw.IfNull(context), cancellationToken); /// /// Called at the end of the agent invocation to add new messages to the chat history. @@ -160,8 +215,59 @@ public abstract class ChatHistoryProvider /// This method is called regardless of whether the invocation succeeded or failed. /// To check if the invocation was successful, inspect the property. /// + /// + /// The default implementation of this method, skips execution for any invocation failures, filters messages using the optional storage input message filter + /// and calls to store new chat history messages. + /// For most scenarios, overriding is sufficient to store chat history messages, while still benefiting from the default error handling and filtering behavior. + /// However, for scenarios that require more control over error handling or message filtering, overriding this method allows you to directly control the messages that are stored for the invocation. + /// /// - protected abstract ValueTask InvokedCoreAsync(InvokedContext context, CancellationToken cancellationToken = default); + protected virtual ValueTask InvokedCoreAsync(InvokedContext context, CancellationToken cancellationToken = default) + { + if (context.InvokeException is not null) + { + return default; + } + + var subContext = new InvokedContext(context.Agent, context.Session, this._storeInputMessageFilter(context.RequestMessages), context.ResponseMessages!); + return this.StoreChatHistoryAsync(subContext, cancellationToken); + } + + /// + /// When overridden in a derived class, adds new messages to the chat history at the end of the agent invocation. + /// + /// Contains the invocation context including request messages, response messages, and any exception that occurred. + /// The to monitor for cancellation requests. The default is . + /// A task that represents the asynchronous add operation. + /// + /// + /// Messages should be added in the order they were generated to maintain proper chronological sequence. + /// The is responsible for preserving message ordering and ensuring that subsequent calls to + /// return messages in the correct chronological order. + /// + /// + /// Implementations may perform additional processing during message addition, such as: + /// + /// Validating message content and metadata + /// Applying storage optimizations or compression + /// Triggering background maintenance operations + /// + /// + /// + /// This method is called from . + /// Note that can be overridden to directly control message filtering and error handling, in which case + /// it is up to the implementer to call this method as needed to store messages. + /// + /// + /// In contrast with , this method only stores messages, + /// while is also responsible for messages filtering and error handling. + /// + /// + /// The default implementation of only calls this method if the invocation succeeded. + /// + /// + protected virtual ValueTask StoreChatHistoryAsync(InvokedContext context, CancellationToken cancellationToken = default) => + default; /// Asks the for an object of the specified type . /// The type of object being requested. diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions/ChatHistoryProvider{TState}.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions/ChatHistoryProvider{TState}.cs new file mode 100644 index 0000000000..ea1280eb1e --- /dev/null +++ b/dotnet/src/Microsoft.Agents.AI.Abstractions/ChatHistoryProvider{TState}.cs @@ -0,0 +1,87 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Collections.Generic; +using System.Text.Json; +using Microsoft.Extensions.AI; + +namespace Microsoft.Agents.AI; + +/// +/// Provides an abstract base class for fetching chat messages from, and adding chat messages to, chat history for the purposes of agent execution with support for maintaining provider state of type . +/// +/// The type of the state to be maintained by the chat history provider. Must be a reference type. +/// +/// This class extends by introducing a strongly-typed state management mechanism, allowing derived classes to maintain and persist custom state information across invocations. +/// The state is stored in the session's StateBag using a configurable key and JSON serialization options, enabling seamless integration with the agent session lifecycle. +/// +public abstract class ChatHistoryProvider : ChatHistoryProvider + where TState : class +{ + private readonly Func _stateInitializer; + private readonly string _stateKey; + private readonly JsonSerializerOptions _jsonSerializerOptions; + + /// + /// Initializes a new instance of the class. + /// + /// A function to initialize the state for the chat history provider. + /// The key used to store the state in the session's StateBag. + /// Options for JSON serialization and deserialization of the state. + /// A filter function to apply to messages when retrieving them from the chat history. + /// A filter function to apply to messages before storing them in the chat history. + protected ChatHistoryProvider( + Func stateInitializer, + string? stateKey, + JsonSerializerOptions? jsonSerializerOptions, + Func, IEnumerable>? provideOutputMessageFilter, + Func, IEnumerable>? storeInputMessageFilter) + : base(provideOutputMessageFilter, storeInputMessageFilter) + { + this._stateInitializer = stateInitializer; + this._jsonSerializerOptions = jsonSerializerOptions ?? AgentAbstractionsJsonUtilities.DefaultOptions; + this._stateKey = stateKey ?? this.GetType().Name; + } + + /// + public override string StateKey => this._stateKey; + + /// + /// Gets the state from the session's StateBag, or initializes it using the state initializer if not present. + /// + /// The agent session containing the StateBag. + /// The provider state, or null if no session is available. + protected virtual TState GetOrInitializeState(AgentSession? session) + { + if (session?.StateBag.TryGetValue(this._stateKey, out var state, this._jsonSerializerOptions) is true && state is not null) + { + return state; + } + + state = this._stateInitializer(session); + if (session is not null) + { + session.StateBag.SetValue(this._stateKey, state, this._jsonSerializerOptions); + } + + return state; + } + + /// + /// Saves the specified state to the session's StateBag using the configured state key and JSON serializer options. + /// If the session is null, this method does nothing. + /// + /// + /// This method provides a convenient way for derived classes to persist state changes back to the session after processing. + /// It abstracts away the details of how state is stored in the session, allowing derived classes to focus on their specific logic. + /// + /// The agent session containing the StateBag. + /// The state to be saved. + protected virtual void SaveState(AgentSession? session, TState state) + { + if (session is not null) + { + session.StateBag.SetValue(this._stateKey, state, this._jsonSerializerOptions); + } + } +} diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions/InMemoryChatHistoryProvider.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions/InMemoryChatHistoryProvider.cs index 9c535923f4..2517586c17 100644 --- a/dotnet/src/Microsoft.Agents.AI.Abstractions/InMemoryChatHistoryProvider.cs +++ b/dotnet/src/Microsoft.Agents.AI.Abstractions/InMemoryChatHistoryProvider.cs @@ -3,7 +3,6 @@ using System; using System.Collections.Generic; using System.Linq; -using System.Text.Json; using System.Text.Json.Serialization; using System.Threading; using System.Threading.Tasks; @@ -25,17 +24,8 @@ namespace Microsoft.Agents.AI; /// message reduction strategies or alternative storage implementations. /// /// -public sealed class InMemoryChatHistoryProvider : ChatHistoryProvider +public sealed class InMemoryChatHistoryProvider : ChatHistoryProvider { - private static IEnumerable DefaultExcludeChatHistoryFilter(IEnumerable messages) - => messages.Where(m => m.GetAgentRequestMessageSourceType() != AgentRequestMessageSourceType.ChatHistory); - - private readonly string _stateKey; - private readonly Func _stateInitializer; - private readonly JsonSerializerOptions _jsonSerializerOptions; - private readonly Func, IEnumerable> _storageInputMessageFilter; - private readonly Func, IEnumerable>? _retrievalOutputMessageFilter; - /// /// Initializes a new instance of the class. /// @@ -44,19 +34,17 @@ public sealed class InMemoryChatHistoryProvider : ChatHistoryProvider /// message reduction, and serialization settings. If , default settings will be used. /// public InMemoryChatHistoryProvider(InMemoryChatHistoryProviderOptions? options = null) + : base( + options?.StateInitializer ?? (_ => new State()), + options?.StateKey, + options?.JsonSerializerOptions, + options?.ProvideOutputMessageFilter, + options?.StorageInputMessageFilter) { - this._stateInitializer = options?.StateInitializer ?? (_ => new State()); this.ChatReducer = options?.ChatReducer; this.ReducerTriggerEvent = options?.ReducerTriggerEvent ?? InMemoryChatHistoryProviderOptions.ChatReducerTriggerEvent.BeforeMessagesRetrieval; - this._stateKey = options?.StateKey ?? base.StateKey; - this._jsonSerializerOptions = options?.JsonSerializerOptions ?? AgentAbstractionsJsonUtilities.DefaultOptions; - this._storageInputMessageFilter = options?.StorageInputMessageFilter ?? DefaultExcludeChatHistoryFilter; - this._retrievalOutputMessageFilter = options?.RetrievalOutputMessageFilter; } - /// - public override string StateKey => this._stateKey; - /// /// Gets the chat reducer used to process or reduce chat messages. If null, no reduction logic will be applied. /// @@ -89,32 +77,9 @@ public sealed class InMemoryChatHistoryProvider : ChatHistoryProvider state.Messages = messages; } - /// - /// Gets the state from the session's StateBag, or initializes it using the state initializer if not present. - /// - /// The agent session containing the StateBag. - /// The provider state, or null if no session is available. - private State GetOrInitializeState(AgentSession? session) - { - if (session?.StateBag.TryGetValue(this._stateKey, out var state, this._jsonSerializerOptions) is true && state is not null) - { - return state; - } - - state = this._stateInitializer(session); - if (session is not null) - { - session.StateBag.SetValue(this._stateKey, state, this._jsonSerializerOptions); - } - - return state; - } - /// - protected override async ValueTask> InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default) + protected override async ValueTask> ProvideChatHistoryAsync(InvokingContext context, CancellationToken cancellationToken = default) { - _ = Throw.IfNull(context); - var state = this.GetOrInitializeState(context.Session); if (this.ReducerTriggerEvent is InMemoryChatHistoryProviderOptions.ChatReducerTriggerEvent.BeforeMessagesRetrieval && this.ChatReducer is not null) @@ -122,30 +87,16 @@ public sealed class InMemoryChatHistoryProvider : ChatHistoryProvider state.Messages = (await this.ChatReducer.ReduceAsync(state.Messages, cancellationToken).ConfigureAwait(false)).ToList(); } - IEnumerable output = state.Messages; - if (this._retrievalOutputMessageFilter is not null) - { - output = this._retrievalOutputMessageFilter(output); - } - return output - .Select(message => message.WithAgentRequestMessageSource(AgentRequestMessageSourceType.ChatHistory, this.GetType().FullName!)) - .Concat(context.RequestMessages); + return state.Messages; } /// - protected override async ValueTask InvokedCoreAsync(InvokedContext context, CancellationToken cancellationToken = default) + protected override async ValueTask StoreChatHistoryAsync(InvokedContext context, CancellationToken cancellationToken = default) { - _ = Throw.IfNull(context); - - if (context.InvokeException is not null) - { - return; - } - var state = this.GetOrInitializeState(context.Session); // Add request and response messages to the provider - var allNewMessages = this._storageInputMessageFilter(context.RequestMessages).Concat(context.ResponseMessages ?? []); + var allNewMessages = context.RequestMessages.Concat(context.ResponseMessages ?? []); state.Messages.AddRange(allNewMessages); if (this.ReducerTriggerEvent is InMemoryChatHistoryProviderOptions.ChatReducerTriggerEvent.AfterMessageAdded && this.ChatReducer is not null) diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions/InMemoryChatHistoryProviderOptions.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions/InMemoryChatHistoryProviderOptions.cs index 41ab46321f..ba24f55ded 100644 --- a/dotnet/src/Microsoft.Agents.AI.Abstractions/InMemoryChatHistoryProviderOptions.cs +++ b/dotnet/src/Microsoft.Agents.AI.Abstractions/InMemoryChatHistoryProviderOptions.cs @@ -71,7 +71,7 @@ public sealed class InMemoryChatHistoryProviderOptions /// /// When , no filtering is applied to the output messages. /// - public Func, IEnumerable>? RetrievalOutputMessageFilter { get; set; } + public Func, IEnumerable>? ProvideOutputMessageFilter { get; set; } /// /// Defines the events that can trigger a reducer in the . diff --git a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatHistoryProvider.cs b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatHistoryProvider.cs index 265f3a3675..b12a465905 100644 --- a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatHistoryProvider.cs +++ b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatHistoryProvider.cs @@ -19,16 +19,11 @@ namespace Microsoft.Agents.AI; /// [RequiresUnreferencedCode("The CosmosChatHistoryProvider uses JSON serialization which is incompatible with trimming.")] [RequiresDynamicCode("The CosmosChatHistoryProvider uses JSON serialization which is incompatible with NativeAOT.")] -public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable +public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable { - private static IEnumerable DefaultExcludeChatHistoryFilter(IEnumerable messages) - => messages.Where(m => m.GetAgentRequestMessageSourceType() != AgentRequestMessageSourceType.ChatHistory); - private readonly CosmosClient _cosmosClient; private readonly Container _container; private readonly bool _ownsClient; - private readonly string _stateKey; - private readonly Func _stateInitializer; private bool _disposed; /// @@ -46,9 +41,6 @@ public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable return options; } - /// - public override string StateKey => this._stateKey; - /// /// Gets or sets the maximum number of messages to return in a single query batch. /// Default is 100 for optimal performance. @@ -84,25 +76,6 @@ public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable /// public string ContainerId { get; init; } - /// - /// A filter function applied to request messages before they are stored - /// during . The default filter excludes messages with the - /// source type. - /// - public Func, IEnumerable> StorageInputMessageFilter { get; set { field = Throw.IfNull(value); } } = DefaultExcludeChatHistoryFilter; - - /// - /// Gets or sets an optional filter function applied to messages produced by this provider - /// during . - /// - /// - /// This filter is only applied to the messages that the provider itself produces (from its internal storage). - /// - /// - /// When , no filtering is applied to the output messages. - /// - public Func, IEnumerable>? RetrievalOutputMessageFilter { get; set; } - /// /// Initializes a new instance of the class. /// @@ -112,6 +85,8 @@ public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable /// A delegate that initializes the provider state on the first invocation, providing the conversation routing info (conversationId, tenantId, userId). /// Whether this instance owns the CosmosClient and should dispose it. /// An optional key to use for storing the state in the . + /// An optional filter function to apply to messages when retrieving them from the chat history. + /// An optional filter function to apply to messages before storing them in the chat history. If not set, defaults to excluding messages with source type . /// Thrown when or is . /// Thrown when any string parameter is null or whitespace. public CosmosChatHistoryProvider( @@ -120,15 +95,16 @@ public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable string containerId, Func stateInitializer, bool ownsClient = false, - string? stateKey = null) + string? stateKey = null, + Func, IEnumerable>? provideOutputMessageFilter = null, + Func, IEnumerable>? storeInputMessageFilter = null) + : base(Throw.IfNull(stateInitializer), stateKey, null, provideOutputMessageFilter, storeInputMessageFilter) { this._cosmosClient = Throw.IfNull(cosmosClient); this.DatabaseId = Throw.IfNullOrWhitespace(databaseId); this.ContainerId = Throw.IfNullOrWhitespace(containerId); this._container = this._cosmosClient.GetContainer(databaseId, containerId); - this._stateInitializer = Throw.IfNull(stateInitializer); this._ownsClient = ownsClient; - this._stateKey = stateKey ?? base.StateKey; } /// @@ -139,6 +115,8 @@ public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable /// The identifier of the Cosmos DB container. /// A delegate that initializes the provider state on the first invocation. /// An optional key to use for storing the state in the . + /// An optional filter function to apply to messages when retrieving them from the chat history. + /// An optional filter function to apply to messages before storing them in the chat history. If not set, defaults to excluding messages with source type . /// Thrown when any required parameter is null. /// Thrown when any string parameter is null or whitespace. public CosmosChatHistoryProvider( @@ -146,8 +124,10 @@ public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable string databaseId, string containerId, Func stateInitializer, - string? stateKey = null) - : this(new CosmosClient(Throw.IfNullOrWhitespace(connectionString)), databaseId, containerId, stateInitializer, ownsClient: true, stateKey) + string? stateKey = null, + Func, IEnumerable>? provideOutputMessageFilter = null, + Func, IEnumerable>? storeInputMessageFilter = null) + : this(new CosmosClient(Throw.IfNullOrWhitespace(connectionString)), databaseId, containerId, stateInitializer, ownsClient: true, stateKey, provideOutputMessageFilter, storeInputMessageFilter) { } @@ -160,6 +140,8 @@ public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable /// The identifier of the Cosmos DB container. /// A delegate that initializes the provider state on the first invocation. /// An optional key to use for storing the state in the . + /// An optional filter function to apply to messages when retrieving them from the chat history. + /// An optional filter function to apply to messages before storing them in the chat history. If not set, defaults to excluding messages with source type . /// Thrown when any required parameter is null. /// Thrown when any string parameter is null or whitespace. public CosmosChatHistoryProvider( @@ -168,32 +150,13 @@ public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable string databaseId, string containerId, Func stateInitializer, - string? stateKey = null) - : this(new CosmosClient(Throw.IfNullOrWhitespace(accountEndpoint), Throw.IfNull(tokenCredential)), databaseId, containerId, stateInitializer, ownsClient: true, stateKey) + string? stateKey = null, + Func, IEnumerable>? provideOutputMessageFilter = null, + Func, IEnumerable>? storeInputMessageFilter = null) + : this(new CosmosClient(Throw.IfNullOrWhitespace(accountEndpoint), Throw.IfNull(tokenCredential)), databaseId, containerId, stateInitializer, ownsClient: true, stateKey, provideOutputMessageFilter, storeInputMessageFilter) { } - /// - /// Gets the state from the session's StateBag, or initializes it using the state initializer if not present. - /// - /// The agent session containing the StateBag. - /// The provider state, or null if no session is available. - private State GetOrInitializeState(AgentSession? session) - { - if (session?.StateBag.TryGetValue(this._stateKey, out var state, AgentAbstractionsJsonUtilities.DefaultOptions) is true && state is not null) - { - return state; - } - - state = this._stateInitializer(session); - if (session is not null) - { - session.StateBag.SetValue(this._stateKey, state, AgentAbstractionsJsonUtilities.DefaultOptions); - } - - return state; - } - /// /// Determines whether hierarchical partitioning should be used based on the state. /// @@ -218,7 +181,7 @@ public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable } /// - protected override async ValueTask> InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default) + protected override async ValueTask> ProvideChatHistoryAsync(InvokingContext context, CancellationToken cancellationToken = default) { #pragma warning disable CA1513 // Use ObjectDisposedException.ThrowIf - not available on all target frameworks if (this._disposed) @@ -227,8 +190,6 @@ public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable } #pragma warning restore CA1513 - _ = Throw.IfNull(context); - var state = this.GetOrInitializeState(context.Session); var partitionKey = BuildPartitionKey(state); @@ -279,22 +240,12 @@ public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable messages.Reverse(); } - return (this.RetrievalOutputMessageFilter is not null ? this.RetrievalOutputMessageFilter(messages) : messages) - .Select(message => message.WithAgentRequestMessageSource(AgentRequestMessageSourceType.ChatHistory, this.GetType().FullName!)) - .Concat(context.RequestMessages); + return messages; } /// - protected override async ValueTask InvokedCoreAsync(InvokedContext context, CancellationToken cancellationToken = default) + protected override async ValueTask StoreChatHistoryAsync(InvokedContext context, CancellationToken cancellationToken = default) { - Throw.IfNull(context); - - if (context.InvokeException is not null) - { - // Do not store messages if there was an exception during invocation - return; - } - #pragma warning disable CA1513 // Use ObjectDisposedException.ThrowIf - not available on all target frameworks if (this._disposed) { @@ -303,7 +254,7 @@ public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable #pragma warning restore CA1513 var state = this.GetOrInitializeState(context.Session); - var messageList = this.StorageInputMessageFilter(context.RequestMessages).Concat(context.ResponseMessages ?? []).ToList(); + var messageList = context.RequestMessages.Concat(context.ResponseMessages ?? []).ToList(); if (messageList.Count == 0) { return; diff --git a/dotnet/src/Microsoft.Agents.AI.Mem0/Mem0Provider.cs b/dotnet/src/Microsoft.Agents.AI.Mem0/Mem0Provider.cs index aaf2333553..4ef096ba6e 100644 --- a/dotnet/src/Microsoft.Agents.AI.Mem0/Mem0Provider.cs +++ b/dotnet/src/Microsoft.Agents.AI.Mem0/Mem0Provider.cs @@ -22,19 +22,12 @@ namespace Microsoft.Agents.AI.Mem0; /// for new invocations using a semantic search endpoint. Retrieved memories are injected as user messages /// to the model, prefixed by a configurable context prompt. /// -public sealed class Mem0Provider : AIContextProvider +public sealed class Mem0Provider : AIContextProvider { private const string DefaultContextPrompt = "## Memories\nConsider the following memories when answering user questions:"; - private static IEnumerable DefaultExternalOnlyFilter(IEnumerable messages) - => messages.Where(m => m.GetAgentRequestMessageSourceType() == AgentRequestMessageSourceType.External); - private readonly string _contextPrompt; private readonly bool _enableSensitiveTelemetryData; - private readonly string _stateKey; - private readonly Func _stateInitializer; - private readonly Func, IEnumerable> _searchInputMessageFilter; - private readonly Func, IEnumerable> _storageInputMessageFilter; private readonly Mem0Client _client; private readonly ILogger? _logger; @@ -58,6 +51,7 @@ public sealed class Mem0Provider : AIContextProvider /// /// public Mem0Provider(HttpClient httpClient, Func stateInitializer, Mem0ProviderOptions? options = null, ILoggerFactory? loggerFactory = null) + : base(ValidateStateInitializer(Throw.IfNull(stateInitializer)), options?.StateKey, Mem0JsonUtilities.DefaultOptions, options?.SearchInputMessageFilter, options?.StorageInputMessageFilter) { Throw.IfNull(httpClient); if (string.IsNullOrWhiteSpace(httpClient.BaseAddress?.AbsoluteUri)) @@ -65,63 +59,41 @@ public sealed class Mem0Provider : AIContextProvider throw new ArgumentException("The HttpClient BaseAddress must be set for Mem0 operations.", nameof(httpClient)); } - this._stateInitializer = Throw.IfNull(stateInitializer); this._logger = loggerFactory?.CreateLogger(); this._client = new Mem0Client(httpClient); this._contextPrompt = options?.ContextPrompt ?? DefaultContextPrompt; this._enableSensitiveTelemetryData = options?.EnableSensitiveTelemetryData ?? false; - this._stateKey = options?.StateKey ?? base.StateKey; - this._searchInputMessageFilter = options?.SearchInputMessageFilter ?? DefaultExternalOnlyFilter; - this._storageInputMessageFilter = options?.StorageInputMessageFilter ?? DefaultExternalOnlyFilter; } - /// - public override string StateKey => this._stateKey; - - /// - /// Gets the state from the session's StateBag, or initializes it using the StateInitializer if not present. - /// - /// The agent session containing the StateBag. - /// The provider state, or null if no session is available. - private State? GetOrInitializeState(AgentSession? session) - { - if (session?.StateBag.TryGetValue(this._stateKey, out var state, Mem0JsonUtilities.DefaultOptions) is true && state is not null) + private static Func ValidateStateInitializer(Func stateInitializer) => + session => { + var state = stateInitializer(session); + + if (state is null + || state.StorageScope is null + || (state.StorageScope.AgentId is null && state.StorageScope.ThreadId is null && state.StorageScope.UserId is null && state.StorageScope.ApplicationId is null) + || state.SearchScope is null + || (state.SearchScope.AgentId is null && state.SearchScope.ThreadId is null && state.SearchScope.UserId is null && state.SearchScope.ApplicationId is null)) + { + throw new InvalidOperationException("State initializer must return a non-null state with valid storage and search scopes, where at least one scoping parameter is set for each."); + } + return state; - } - - state = this._stateInitializer(session); - - if (state is null - || state.StorageScope is null - || (state.StorageScope.AgentId is null && state.StorageScope.ThreadId is null && state.StorageScope.UserId is null && state.StorageScope.ApplicationId is null) - || state.SearchScope is null - || (state.SearchScope.AgentId is null && state.SearchScope.ThreadId is null && state.SearchScope.UserId is null && state.SearchScope.ApplicationId is null)) - { - throw new InvalidOperationException("State initializer must return a non-null state with valid storage and search scopes, where at lest one scoping parameter is set for each."); - } - - if (session is not null) - { - session.StateBag.SetValue(this._stateKey, state, Mem0JsonUtilities.DefaultOptions); - } - - return state; - } + }; /// - protected override async ValueTask InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default) + protected override async ValueTask ProvideAIContextAsync(InvokingContext context, CancellationToken cancellationToken = default) { Throw.IfNull(context); - var inputContext = context.AIContext; var state = this.GetOrInitializeState(context.Session); - var searchScope = state?.SearchScope ?? new Mem0ProviderScope(); + var searchScope = state.SearchScope; string queryText = string.Join( Environment.NewLine, - this._searchInputMessageFilter(inputContext.Messages ?? []) + (context.AIContext.Messages ?? []) .Where(m => !string.IsNullOrWhiteSpace(m.Text)) .Select(m => m.Text)); @@ -138,9 +110,6 @@ public sealed class Mem0Provider : AIContextProvider var outputMessageText = memories.Count == 0 ? null : $"{this._contextPrompt}\n{string.Join(Environment.NewLine, memories)}"; - var outputMessage = memories.Count == 0 - ? null - : new ChatMessage(ChatRole.User, outputMessageText!).WithAgentRequestMessageSource(AgentRequestMessageSourceType.AIContextProvider, this.GetType().FullName!); if (this._logger?.IsEnabled(LogLevel.Information) is true) { @@ -167,11 +136,9 @@ public sealed class Mem0Provider : AIContextProvider return new AIContext { - Instructions = inputContext.Instructions, - Messages = - (inputContext.Messages ?? []) - .Concat(outputMessage is not null ? [outputMessage] : []), - Tools = inputContext.Tools + Messages = outputMessageText is not null + ? [new ChatMessage(ChatRole.User, outputMessageText)] + : null }; } catch (ArgumentException) @@ -190,27 +157,23 @@ public sealed class Mem0Provider : AIContextProvider searchScope.ThreadId, this.SanitizeLogData(searchScope.UserId)); } - return inputContext; + + return new AIContext(); } } /// - protected override async ValueTask InvokedCoreAsync(InvokedContext context, CancellationToken cancellationToken = default) + protected override async ValueTask StoreAIContextAsync(InvokedContext context, CancellationToken cancellationToken = default) { - if (context.InvokeException is not null) - { - return; // Do not update memory on failed invocations. - } - var state = this.GetOrInitializeState(context.Session); - var storageScope = state?.StorageScope ?? new Mem0ProviderScope(); + var storageScope = state.StorageScope; try { // Persist request and response messages after invocation. await this.PersistMessagesAsync( storageScope, - this._storageInputMessageFilter(context.RequestMessages) + context.RequestMessages .Concat(context.ResponseMessages ?? []), cancellationToken).ConfigureAwait(false); } @@ -238,12 +201,7 @@ public sealed class Mem0Provider : AIContextProvider { Throw.IfNull(session); var state = this.GetOrInitializeState(session); - var storageScope = state?.StorageScope; - - if (storageScope is null) - { - return Task.CompletedTask; // Nothing to clear if there is no state. - } + var storageScope = state.StorageScope; return this._client.ClearMemoryAsync( storageScope.ApplicationId, diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowChatHistoryProvider.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowChatHistoryProvider.cs index 6672b9e2a3..74e2fae3fa 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowChatHistoryProvider.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowChatHistoryProvider.cs @@ -9,10 +9,8 @@ using Microsoft.Extensions.AI; namespace Microsoft.Agents.AI.Workflows; -internal sealed class WorkflowChatHistoryProvider : ChatHistoryProvider +internal sealed class WorkflowChatHistoryProvider : ChatHistoryProvider { - private readonly JsonSerializerOptions _jsonSerializerOptions; - /// /// Initializes a new instance of the class. /// @@ -22,8 +20,8 @@ internal sealed class WorkflowChatHistoryProvider : ChatHistoryProvider /// and source generated serializers are required, or Native AOT / Trimming is required. /// public WorkflowChatHistoryProvider(JsonSerializerOptions? jsonSerializerOptions = null) + : base(stateInitializer: _ => new StoreState(), stateKey: null, jsonSerializerOptions: jsonSerializerOptions, provideOutputMessageFilter: null, storeInputMessageFilter: null) { - this._jsonSerializerOptions = jsonSerializerOptions ?? AgentAbstractionsJsonUtilities.DefaultOptions; } internal sealed class StoreState @@ -32,43 +30,16 @@ internal sealed class WorkflowChatHistoryProvider : ChatHistoryProvider public List Messages { get; set; } = []; } - private StoreState GetOrInitializeState(AgentSession? session) - { - if (session?.StateBag.TryGetValue(this.StateKey, out var state, this._jsonSerializerOptions) is true && state is not null) - { - return state; - } - - state = new(); - if (session is not null) - { - session.StateBag.SetValue(this.StateKey, state, this._jsonSerializerOptions); - } - - return state; - } - internal void AddMessages(AgentSession session, params IEnumerable messages) => this.GetOrInitializeState(session).Messages.AddRange(messages); - protected override ValueTask> InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default) - => new(this.GetOrInitializeState(context.Session) - .Messages - .Select(message => message.WithAgentRequestMessageSource(AgentRequestMessageSourceType.ChatHistory, this.GetType().FullName!)) - .Concat(context.RequestMessages)); + protected override ValueTask> ProvideChatHistoryAsync(InvokingContext context, CancellationToken cancellationToken = default) + => new(this.GetOrInitializeState(context.Session).Messages); - protected override ValueTask InvokedCoreAsync(InvokedContext context, CancellationToken cancellationToken = default) + protected override ValueTask StoreChatHistoryAsync(InvokedContext context, CancellationToken cancellationToken = default) { - if (context.InvokeException is not null) - { - return default; - } - - var allNewMessages = context.RequestMessages - .Where(m => m.GetAgentRequestMessageSourceType() != AgentRequestMessageSourceType.ChatHistory) - .Concat(context.ResponseMessages ?? []); + var allNewMessages = context.RequestMessages.Concat(context.ResponseMessages ?? []); this.GetOrInitializeState(context.Session).Messages.AddRange(allNewMessages); - return default; } diff --git a/dotnet/src/Microsoft.Agents.AI/Memory/ChatHistoryMemoryProvider.cs b/dotnet/src/Microsoft.Agents.AI/Memory/ChatHistoryMemoryProvider.cs index 9d163f79cf..346365e8cc 100644 --- a/dotnet/src/Microsoft.Agents.AI/Memory/ChatHistoryMemoryProvider.cs +++ b/dotnet/src/Microsoft.Agents.AI/Memory/ChatHistoryMemoryProvider.cs @@ -24,8 +24,8 @@ namespace Microsoft.Agents.AI; /// abstractions to work with any compatible vector store implementation. /// /// -/// Messages are stored during the method and retrieved during the -/// method using semantic similarity search. +/// Messages are stored during the method and retrieved during the +/// method using semantic similarity search. /// /// /// Behavior is configurable through . When @@ -34,16 +34,13 @@ namespace Microsoft.Agents.AI; /// injecting them automatically on each invocation. /// /// -public sealed class ChatHistoryMemoryProvider : AIContextProvider, IDisposable +public sealed class ChatHistoryMemoryProvider : AIContextProvider, IDisposable { private const string DefaultContextPrompt = "## Memories\nConsider the following memories when answering user questions:"; private const int DefaultMaxResults = 3; private const string DefaultFunctionToolName = "Search"; private const string DefaultFunctionToolDescription = "Allows searching for related previous chat history to help answer the user question."; - private static IEnumerable DefaultExternalOnlyFilter(IEnumerable messages) - => messages.Where(m => m.GetAgentRequestMessageSourceType() == AgentRequestMessageSourceType.External); - #pragma warning disable CA2213 // VectorStore is not owned by this class - caller is responsible for disposal private readonly VectorStore _vectorStore; #pragma warning restore CA2213 @@ -55,10 +52,6 @@ public sealed class ChatHistoryMemoryProvider : AIContextProvider, IDisposable private readonly string _toolName; private readonly string _toolDescription; private readonly ILogger? _logger; - private readonly string _stateKey; - private readonly Func _stateInitializer; - private readonly Func, IEnumerable> _searchInputMessageFilter; - private readonly Func, IEnumerable> _storageInputMessageFilter; private bool _collectionInitialized; private readonly SemaphoreSlim _initializationLock = new(1, 1); @@ -81,21 +74,18 @@ public sealed class ChatHistoryMemoryProvider : AIContextProvider, IDisposable Func stateInitializer, ChatHistoryMemoryProviderOptions? options = null, ILoggerFactory? loggerFactory = null) + : base(Throw.IfNull(stateInitializer), options?.StateKey, AgentJsonUtilities.DefaultOptions, options?.SearchInputMessageFilter, options?.StorageInputMessageFilter) { this._vectorStore = Throw.IfNull(vectorStore); - this._stateInitializer = Throw.IfNull(stateInitializer); options ??= new ChatHistoryMemoryProviderOptions(); this._maxResults = options.MaxResults.HasValue ? Throw.IfLessThanOrEqual(options.MaxResults.Value, 0) : DefaultMaxResults; this._contextPrompt = options.ContextPrompt ?? DefaultContextPrompt; this._enableSensitiveTelemetryData = options.EnableSensitiveTelemetryData; this._searchTime = options.SearchTime; - this._stateKey = options.StateKey ?? base.StateKey; this._logger = loggerFactory?.CreateLogger(); this._toolName = options.FunctionToolName ?? DefaultFunctionToolName; this._toolDescription = options.FunctionToolDescription ?? DefaultFunctionToolDescription; - this._searchInputMessageFilter = options.SearchInputMessageFilter ?? DefaultExternalOnlyFilter; - this._storageInputMessageFilter = options.StorageInputMessageFilter ?? DefaultExternalOnlyFilter; // Create a definition so that we can use the dimensions provided at runtime. var definition = new VectorStoreCollectionDefinition @@ -120,37 +110,12 @@ public sealed class ChatHistoryMemoryProvider : AIContextProvider, IDisposable } /// - public override string StateKey => this._stateKey; - - /// - /// Gets the state from the session's StateBag, or initializes it using the StateInitializer if not present. - /// - /// The agent session containing the StateBag. - /// The provider state, or null if no session is available. - private State? GetOrInitializeState(AgentSession? session) - { - if (session?.StateBag.TryGetValue(this._stateKey, out var state, AgentJsonUtilities.DefaultOptions) is true && state is not null) - { - return state; - } - - state = this._stateInitializer(session); - if (state is not null && session is not null) - { - session.StateBag.SetValue(this._stateKey, state, AgentJsonUtilities.DefaultOptions); - } - - return state; - } - - /// - protected override async ValueTask InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default) + protected override async ValueTask ProvideAIContextAsync(InvokingContext context, CancellationToken cancellationToken = default) { _ = Throw.IfNull(context); - var inputContext = context.AIContext; var state = this.GetOrInitializeState(context.Session); - var searchScope = state?.SearchScope ?? new ChatHistoryMemoryProviderScope(); + var searchScope = state.SearchScope; if (this._searchTime == ChatHistoryMemoryProviderOptions.SearchBehavior.OnDemandFunctionCalling) { @@ -166,12 +131,10 @@ public sealed class ChatHistoryMemoryProvider : AIContextProvider, IDisposable description: this._toolDescription) ]; - // Expose search tool for on-demand invocation by the model, accumulated with the input context + // Expose search tool for on-demand invocation by the model return new AIContext { - Instructions = inputContext.Instructions, - Messages = inputContext.Messages, - Tools = (inputContext.Tools ?? []).Concat(tools) + Tools = tools }; } @@ -179,13 +142,13 @@ public sealed class ChatHistoryMemoryProvider : AIContextProvider, IDisposable { // Get the text from the current request messages var requestText = string.Join("\n", - this._searchInputMessageFilter(inputContext.Messages ?? []) + (context.AIContext.Messages ?? []) .Where(m => m != null && !string.IsNullOrWhiteSpace(m.Text)) .Select(m => m.Text)); if (string.IsNullOrWhiteSpace(requestText)) { - return inputContext; + return new AIContext(); } // Search for relevant chat history @@ -193,19 +156,12 @@ public sealed class ChatHistoryMemoryProvider : AIContextProvider, IDisposable if (string.IsNullOrWhiteSpace(contextText)) { - return inputContext; + return new AIContext(); } return new AIContext { - Instructions = inputContext.Instructions, - Messages = - (inputContext.Messages ?? []) - .Concat( - [ - new ChatMessage(ChatRole.User, contextText).WithAgentRequestMessageSource(AgentRequestMessageSourceType.AIContextProvider, this.GetType().FullName!) - ]), - Tools = inputContext.Tools + Messages = [new ChatMessage(ChatRole.User, contextText)] }; } catch (Exception ex) @@ -221,30 +177,24 @@ public sealed class ChatHistoryMemoryProvider : AIContextProvider, IDisposable this.SanitizeLogData(searchScope.UserId)); } - return inputContext; + return new AIContext(); } } /// - protected override async ValueTask InvokedCoreAsync(InvokedContext context, CancellationToken cancellationToken = default) + protected override async ValueTask StoreAIContextAsync(InvokedContext context, CancellationToken cancellationToken = default) { _ = Throw.IfNull(context); - // Only store if invocation was successful - if (context.InvokeException != null) - { - return; - } - var state = this.GetOrInitializeState(context.Session); - var storageScope = state?.StorageScope ?? new ChatHistoryMemoryProviderScope(); + var storageScope = state.StorageScope; try { // Ensure the collection is initialized var collection = await this.EnsureCollectionExistsAsync(cancellationToken).ConfigureAwait(false); - List> itemsToStore = this._storageInputMessageFilter(context.RequestMessages) + List> itemsToStore = context.RequestMessages .Concat(context.ResponseMessages ?? []) .Select(message => new Dictionary { diff --git a/dotnet/src/Microsoft.Agents.AI/TextSearchProvider.cs b/dotnet/src/Microsoft.Agents.AI/TextSearchProvider.cs index f038fa3c38..099059d65e 100644 --- a/dotnet/src/Microsoft.Agents.AI/TextSearchProvider.cs +++ b/dotnet/src/Microsoft.Agents.AI/TextSearchProvider.cs @@ -32,16 +32,13 @@ namespace Microsoft.Agents.AI; /// multi-turn context to the retrieval layer without permanently altering the conversation history. /// /// -public sealed class TextSearchProvider : AIContextProvider +public sealed class TextSearchProvider : AIContextProvider { private const string DefaultPluginSearchFunctionName = "Search"; private const string DefaultPluginSearchFunctionDescription = "Allows searching for additional information to help answer the user question."; private const string DefaultContextPrompt = "## Additional Context\nConsider the following information from source documents when responding to the user:"; private const string DefaultCitationsPrompt = "Include citations to the source document with document name and link if document name and link is available."; - private static IEnumerable DefaultExternalOnlyFilter(IEnumerable messages) - => messages.Where(m => m.GetAgentRequestMessageSourceType() == AgentRequestMessageSourceType.External); - private readonly Func>> _searchAsync; private readonly ILogger? _logger; private readonly AITool[] _tools; @@ -50,10 +47,7 @@ public sealed class TextSearchProvider : AIContextProvider private readonly TextSearchProviderOptions.TextSearchBehavior _searchTime; private readonly string _contextPrompt; private readonly string _citationsPrompt; - private readonly string _stateKey; private readonly Func, string>? _contextFormatter; - private readonly Func, IEnumerable> _searchInputMessageFilter; - private readonly Func, IEnumerable> _storageInputMessageFilter; /// /// Initializes a new instance of the class. @@ -66,6 +60,7 @@ public sealed class TextSearchProvider : AIContextProvider Func>> searchAsync, TextSearchProviderOptions? options = null, ILoggerFactory? loggerFactory = null) + : base(_ => new TextSearchProviderState(), options?.StateKey, AgentJsonUtilities.DefaultOptions, options?.SearchInputMessageFilter, options?.StorageInputMessageFilter) { // Validate and assign parameters this._searchAsync = Throw.IfNull(searchAsync); @@ -75,10 +70,7 @@ public sealed class TextSearchProvider : AIContextProvider this._searchTime = options?.SearchTime ?? TextSearchProviderOptions.TextSearchBehavior.BeforeAIInvoke; this._contextPrompt = options?.ContextPrompt ?? DefaultContextPrompt; this._citationsPrompt = options?.CitationsPrompt ?? DefaultCitationsPrompt; - this._stateKey = options?.StateKey ?? base.StateKey; this._contextFormatter = options?.ContextFormatter; - this._searchInputMessageFilter = options?.SearchInputMessageFilter ?? DefaultExternalOnlyFilter; - this._storageInputMessageFilter = options?.StorageInputMessageFilter ?? DefaultExternalOnlyFilter; // Create the on-demand search tool (only used if behavior is OnDemandFunctionCalling) this._tools = @@ -91,32 +83,25 @@ public sealed class TextSearchProvider : AIContextProvider } /// - public override string StateKey => this._stateKey; - - /// - protected override async ValueTask InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default) + protected override async ValueTask ProvideAIContextAsync(InvokingContext context, CancellationToken cancellationToken = default) { - var inputContext = context.AIContext; - if (this._searchTime != TextSearchProviderOptions.TextSearchBehavior.BeforeAIInvoke) { - // Expose the search tool for on-demand invocation, accumulated with the input context. + // Expose the search tool for on-demand invocation. return new AIContext { - Instructions = inputContext.Instructions, - Messages = inputContext.Messages, - Tools = (inputContext.Tools ?? []).Concat(this._tools) + Tools = this._tools }; } - // Retrieve recent messages from the session state bag. - var recentMessagesText = context.Session?.StateBag.GetValue(this._stateKey, AgentJsonUtilities.DefaultOptions)?.RecentMessagesText + // Retrieve recent messages from the session state. + var recentMessagesText = this.GetOrInitializeState(context.Session).RecentMessagesText ?? []; // Aggregate text from memory + current request messages. var sbInput = new StringBuilder(); var requestMessagesText = - this._searchInputMessageFilter(inputContext.Messages ?? []) + (context.AIContext.Messages ?? []) .Where(x => !string.IsNullOrWhiteSpace(x?.Text)).Select(x => x.Text); foreach (var messageText in recentMessagesText.Concat(requestMessagesText)) { @@ -142,7 +127,7 @@ public sealed class TextSearchProvider : AIContextProvider if (materialized.Count == 0) { - return inputContext; + return new AIContext(); } // Format search results @@ -155,25 +140,18 @@ public sealed class TextSearchProvider : AIContextProvider return new AIContext { - Instructions = inputContext.Instructions, - Messages = - (inputContext.Messages ?? []) - .Concat( - [ - new ChatMessage(ChatRole.User, formatted).WithAgentRequestMessageSource(AgentRequestMessageSourceType.AIContextProvider, this.GetType().FullName!) - ]), - Tools = inputContext.Tools + Messages = [new ChatMessage(ChatRole.User, formatted)] }; } catch (Exception ex) { this._logger?.LogError(ex, "TextSearchProvider: Failed to search for data due to error"); - return inputContext; + return new AIContext(); } } /// - protected override ValueTask InvokedCoreAsync(InvokedContext context, CancellationToken cancellationToken = default) + protected override ValueTask StoreAIContextAsync(InvokedContext context, CancellationToken cancellationToken = default) { int limit = this._recentMessageMemoryLimit; if (limit <= 0) @@ -186,16 +164,11 @@ public sealed class TextSearchProvider : AIContextProvider return default; // No session to store state in. } - if (context.InvokeException is not null) - { - return default; // Do not update memory on failed invocations. - } - - // Retrieve existing recent messages from the session state bag. - var recentMessagesText = context.Session.StateBag.GetValue(this._stateKey, AgentJsonUtilities.DefaultOptions)?.RecentMessagesText + // Retrieve existing recent messages from the session state. + var recentMessagesText = this.GetOrInitializeState(context.Session).RecentMessagesText ?? []; - var newMessagesText = this._storageInputMessageFilter(context.RequestMessages) + var newMessagesText = context.RequestMessages .Concat(context.ResponseMessages ?? []) .Where(m => this._recentMessageRolesIncluded.Contains(m.Role) && @@ -208,11 +181,10 @@ public sealed class TextSearchProvider : AIContextProvider ? allMessages.Skip(allMessages.Count - limit).ToList() : allMessages; - // Store updated state back to the session state bag. - context.Session.StateBag.SetValue( - this._stateKey, - new TextSearchProviderState { RecentMessagesText = updatedMessages }, - AgentJsonUtilities.DefaultOptions); + // Store updated state back to the session. + this.SaveState( + context.Session, + new TextSearchProviderState { RecentMessagesText = updatedMessages }); return default; } @@ -311,8 +283,14 @@ public sealed class TextSearchProvider : AIContextProvider public object? RawRepresentation { get; set; } } - internal sealed class TextSearchProviderState + /// + /// Represents the per-session state of a stored in the . + /// + public sealed class TextSearchProviderState { + /// + /// Gets or sets the list of recent message texts retained for multi-turn search context. + /// public List? RecentMessagesText { get; set; } } } diff --git a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/AIContextProviderTStateTests.cs b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/AIContextProviderTStateTests.cs new file mode 100644 index 0000000000..b457bf1945 --- /dev/null +++ b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/AIContextProviderTStateTests.cs @@ -0,0 +1,199 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Linq; +using System.Threading; +using System.Threading.Tasks; +using Microsoft.Extensions.AI; +using Moq; + +namespace Microsoft.Agents.AI.Abstractions.UnitTests; + +/// +/// Contains tests for the class. +/// +public class AIContextProviderTStateTests +{ + private static readonly AIAgent s_mockAgent = new Mock().Object; + + #region GetOrInitializeState Tests + + [Fact] + public void GetOrInitializeState_InitializesFromStateInitializerOnFirstCall() + { + // Arrange + var expectedState = new TestState { Value = "initialized" }; + var provider = new TestAIContextProvider(_ => expectedState); + var session = new TestAgentSession(); + + // Act + var state = provider.GetState(session); + + // Assert + Assert.Same(expectedState, state); + } + + [Fact] + public void GetOrInitializeState_ReturnsCachedStateFromStateBagOnSecondCall() + { + // Arrange + var callCount = 0; + var provider = new TestAIContextProvider(_ => + { + callCount++; + return new TestState { Value = $"init-{callCount}" }; + }); + var session = new TestAgentSession(); + + // Act + var state1 = provider.GetState(session); + var state2 = provider.GetState(session); + + // Assert - initializer called only once; second call reads from StateBag + Assert.Equal(1, callCount); + Assert.Equal("init-1", state1.Value); + Assert.Equal("init-1", state2.Value); + } + + [Fact] + public void GetOrInitializeState_WorksWhenSessionIsNull() + { + // Arrange + var provider = new TestAIContextProvider(_ => new TestState { Value = "no-session" }); + + // Act + var state = provider.GetState(null); + + // Assert + Assert.Equal("no-session", state.Value); + } + + [Fact] + public void GetOrInitializeState_ReInitializesWhenSessionIsNull() + { + // Arrange - without a session, state can't be cached in StateBag + var callCount = 0; + var provider = new TestAIContextProvider(_ => + { + callCount++; + return new TestState { Value = $"init-{callCount}" }; + }); + + // Act + provider.GetState(null); + provider.GetState(null); + + // Assert - initializer called each time since there's no session to cache in + Assert.Equal(2, callCount); + } + + #endregion + + #region SaveState Tests + + [Fact] + public void SaveState_SavesToStateBag() + { + // Arrange + var provider = new TestAIContextProvider(_ => new TestState()); + var session = new TestAgentSession(); + var state = new TestState { Value = "saved" }; + + // Act + provider.DoSaveState(session, state); + var retrieved = provider.GetState(session); + + // Assert + Assert.Equal("saved", retrieved.Value); + } + + [Fact] + public void SaveState_NoOpWhenSessionIsNull() + { + // Arrange + var provider = new TestAIContextProvider(_ => new TestState { Value = "default" }); + + // Act - should not throw + provider.DoSaveState(null, new TestState { Value = "saved" }); + + // Assert - no exception; can't verify further without a session + } + + #endregion + + #region StateKey Tests + + [Fact] + public void StateKey_DefaultsToTypeName() + { + // Arrange + var provider = new TestAIContextProvider(_ => new TestState()); + + // Act & Assert + Assert.Equal(nameof(TestAIContextProvider), provider.StateKey); + } + + [Fact] + public void StateKey_UsesCustomKeyWhenProvided() + { + // Arrange + var provider = new TestAIContextProvider(_ => new TestState(), stateKey: "custom-key"); + + // Act & Assert + Assert.Equal("custom-key", provider.StateKey); + } + + #endregion + + #region Integration Tests + + [Fact] + public async Task InvokingCoreAsync_CanUseStateInProvideAIContextAsync() + { + // Arrange + var provider = new TestAIContextProvider(_ => new TestState { Value = "state-value" }); + var session = new TestAgentSession(); + var inputContext = new AIContext { Messages = [new ChatMessage(ChatRole.User, "Hi")] }; + var context = new AIContextProvider.InvokingContext(s_mockAgent, session, inputContext); + + // Act + var result = await provider.InvokingAsync(context); + + // Assert - the provider uses state to produce context messages + var messages = result.Messages!.ToList(); + Assert.Equal(2, messages.Count); + Assert.Contains("state-value", messages[1].Text); + } + + #endregion + + public sealed class TestState + { + public string Value { get; set; } = string.Empty; + } + + private sealed class TestAIContextProvider : AIContextProvider + { + public TestAIContextProvider( + Func stateInitializer, + string? stateKey = null) + : base(stateInitializer, stateKey, null, null, null) + { + } + + public TestState GetState(AgentSession? session) => this.GetOrInitializeState(session); + + public void DoSaveState(AgentSession? session, TestState state) => this.SaveState(session, state); + + protected override ValueTask ProvideAIContextAsync(InvokingContext context, CancellationToken cancellationToken = default) + { + var state = this.GetOrInitializeState(context.Session); + return new(new AIContext + { + Messages = [new ChatMessage(ChatRole.System, $"Context from state: {state.Value}")] + }); + } + } + + private sealed class TestAgentSession : AgentSession; +} diff --git a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/AIContextProviderTests.cs b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/AIContextProviderTests.cs index 14a0f81e08..811f9a3216 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/AIContextProviderTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/AIContextProviderTests.cs @@ -3,6 +3,7 @@ using System; using System.Collections.Generic; using System.Collections.ObjectModel; +using System.Linq; using System.Threading; using System.Threading.Tasks; using Microsoft.Extensions.AI; @@ -337,9 +338,314 @@ public class AIContextProviderTests #endregion + #region InvokingAsync / InvokedAsync Null Check Tests + + [Fact] + public async Task InvokingAsync_NullContext_ThrowsArgumentNullExceptionAsync() + { + // Arrange + var provider = new TestAIContextProvider(); + + // Act & Assert + await Assert.ThrowsAsync(() => provider.InvokingAsync(null!).AsTask()); + } + + [Fact] + public async Task InvokedAsync_NullContext_ThrowsArgumentNullExceptionAsync() + { + // Arrange + var provider = new TestAIContextProvider(); + + // Act & Assert + await Assert.ThrowsAsync(() => provider.InvokedAsync(null!).AsTask()); + } + + #endregion + + #region InvokingCoreAsync Tests + + [Fact] + public async Task InvokingCoreAsync_CallsProvideAIContextAndReturnsMergedContextAsync() + { + // Arrange + var providedMessages = new[] { new ChatMessage(ChatRole.System, "Context message") }; + var provider = new TestAIContextProvider(provideContext: new AIContext { Messages = providedMessages }); + var inputContext = new AIContext { Messages = [new ChatMessage(ChatRole.User, "User input")] }; + var context = new AIContextProvider.InvokingContext(s_mockAgent, s_mockSession, inputContext); + + // Act + var result = await provider.InvokingAsync(context); + + // Assert - input messages + provided messages merged + var messages = result.Messages!.ToList(); + Assert.Equal(2, messages.Count); + Assert.Equal("User input", messages[0].Text); + Assert.Equal("Context message", messages[1].Text); + } + + [Fact] + public async Task InvokingCoreAsync_FiltersInputToExternalOnlyByDefaultAsync() + { + // Arrange + var provider = new TestAIContextProvider(captureFilteredContext: true); + var externalMsg = new ChatMessage(ChatRole.User, "External"); + var chatHistoryMsg = new ChatMessage(ChatRole.User, "History") + .WithAgentRequestMessageSource(AgentRequestMessageSourceType.ChatHistory, "src"); + var contextProviderMsg = new ChatMessage(ChatRole.User, "ContextProvider") + .WithAgentRequestMessageSource(AgentRequestMessageSourceType.AIContextProvider, "src"); + var inputContext = new AIContext { Messages = [externalMsg, chatHistoryMsg, contextProviderMsg] }; + var context = new AIContextProvider.InvokingContext(s_mockAgent, s_mockSession, inputContext); + + // Act + await provider.InvokingAsync(context); + + // Assert - ProvideAIContextAsync received only External messages + Assert.NotNull(provider.LastProvidedContext); + var filteredMessages = provider.LastProvidedContext!.AIContext.Messages!.ToList(); + Assert.Single(filteredMessages); + Assert.Equal("External", filteredMessages[0].Text); + } + + [Fact] + public async Task InvokingCoreAsync_StampsProvidedMessagesWithAIContextProviderSourceAsync() + { + // Arrange + var providedMessages = new[] { new ChatMessage(ChatRole.System, "Provided") }; + var provider = new TestAIContextProvider(provideContext: new AIContext { Messages = providedMessages }); + var inputContext = new AIContext { Messages = [] }; + var context = new AIContextProvider.InvokingContext(s_mockAgent, s_mockSession, inputContext); + + // Act + var result = await provider.InvokingAsync(context); + + // Assert + var messages = result.Messages!.ToList(); + Assert.Single(messages); + Assert.Equal(AgentRequestMessageSourceType.AIContextProvider, messages[0].GetAgentRequestMessageSourceType()); + } + + [Fact] + public async Task InvokingCoreAsync_MergesInstructionsAsync() + { + // Arrange + var provider = new TestAIContextProvider(provideContext: new AIContext { Instructions = "Provided instructions" }); + var inputContext = new AIContext { Instructions = "Input instructions" }; + var context = new AIContextProvider.InvokingContext(s_mockAgent, s_mockSession, inputContext); + + // Act + var result = await provider.InvokingAsync(context); + + // Assert - instructions are joined with newline + Assert.Equal("Input instructions\nProvided instructions", result.Instructions); + } + + [Fact] + public async Task InvokingCoreAsync_MergesToolsAsync() + { + // Arrange + var inputTool = AIFunctionFactory.Create(() => "a", "inputTool"); + var providedTool = AIFunctionFactory.Create(() => "b", "providedTool"); + var provider = new TestAIContextProvider(provideContext: new AIContext { Tools = [providedTool] }); + var inputContext = new AIContext { Tools = [inputTool] }; + var context = new AIContextProvider.InvokingContext(s_mockAgent, s_mockSession, inputContext); + + // Act + var result = await provider.InvokingAsync(context); + + // Assert - both tools present + var tools = result.Tools!.ToList(); + Assert.Equal(2, tools.Count); + } + + [Fact] + public async Task InvokingCoreAsync_UsesCustomProvideInputFilterAsync() + { + // Arrange - filter that keeps all messages (not just External) + var provider = new TestAIContextProvider( + captureFilteredContext: true, + provideInputMessageFilter: msgs => msgs); + var externalMsg = new ChatMessage(ChatRole.User, "External"); + var chatHistoryMsg = new ChatMessage(ChatRole.User, "History") + .WithAgentRequestMessageSource(AgentRequestMessageSourceType.ChatHistory, "src"); + var inputContext = new AIContext { Messages = [externalMsg, chatHistoryMsg] }; + var context = new AIContextProvider.InvokingContext(s_mockAgent, s_mockSession, inputContext); + + // Act + await provider.InvokingAsync(context); + + // Assert - ProvideAIContextAsync received ALL messages (custom filter keeps everything) + Assert.NotNull(provider.LastProvidedContext); + var filteredMessages = provider.LastProvidedContext!.AIContext.Messages!.ToList(); + Assert.Equal(2, filteredMessages.Count); + } + + [Fact] + public async Task InvokingCoreAsync_ReturnsEmptyContextByDefaultAsync() + { + // Arrange - provider that doesn't override ProvideAIContextAsync + var provider = new DefaultAIContextProvider(); + var inputContext = new AIContext { Messages = [new ChatMessage(ChatRole.User, "Hello")] }; + var context = new AIContextProvider.InvokingContext(s_mockAgent, s_mockSession, inputContext); + + // Act + var result = await provider.InvokingAsync(context); + + // Assert - only the input messages (no additional provided) + var messages = result.Messages!.ToList(); + Assert.Single(messages); + Assert.Equal("Hello", messages[0].Text); + } + + [Fact] + public async Task InvokingCoreAsync_MergesWithOriginalUnfilteredMessagesAsync() + { + // Arrange - default filter is External-only, but the MERGED result should include + // the original unfiltered input messages plus the provided messages + var providedMessages = new[] { new ChatMessage(ChatRole.System, "Provided") }; + var provider = new TestAIContextProvider(provideContext: new AIContext { Messages = providedMessages }); + var externalMsg = new ChatMessage(ChatRole.User, "External"); + var chatHistoryMsg = new ChatMessage(ChatRole.User, "History") + .WithAgentRequestMessageSource(AgentRequestMessageSourceType.ChatHistory, "src"); + var inputContext = new AIContext { Messages = [externalMsg, chatHistoryMsg] }; + var context = new AIContextProvider.InvokingContext(s_mockAgent, s_mockSession, inputContext); + + // Act + var result = await provider.InvokingAsync(context); + + // Assert - original 2 input messages + 1 provided message + var messages = result.Messages!.ToList(); + Assert.Equal(3, messages.Count); + Assert.Equal("External", messages[0].Text); + Assert.Equal("History", messages[1].Text); + Assert.Equal("Provided", messages[2].Text); + } + + #endregion + + #region InvokedCoreAsync Tests + + [Fact] + public async Task InvokedCoreAsync_CallsStoreAIContextWithFilteredMessagesAsync() + { + // Arrange + var provider = new TestAIContextProvider(); + var externalMessage = new ChatMessage(ChatRole.User, "External"); + var chatHistoryMessage = new ChatMessage(ChatRole.User, "History") + .WithAgentRequestMessageSource(AgentRequestMessageSourceType.ChatHistory, "src"); + var responseMessages = new[] { new ChatMessage(ChatRole.Assistant, "Response") }; + var context = new AIContextProvider.InvokedContext(s_mockAgent, s_mockSession, new[] { externalMessage, chatHistoryMessage }, responseMessages); + + // Act + await provider.InvokedAsync(context); + + // Assert - default filter keeps only External messages + Assert.NotNull(provider.LastStoredContext); + var storedRequest = provider.LastStoredContext!.RequestMessages.ToList(); + Assert.Single(storedRequest); + Assert.Equal("External", storedRequest[0].Text); + Assert.Same(responseMessages, provider.LastStoredContext.ResponseMessages); + } + + [Fact] + public async Task InvokedCoreAsync_SkipsStorageWhenInvokeExceptionIsNotNullAsync() + { + // Arrange + var provider = new TestAIContextProvider(); + var context = new AIContextProvider.InvokedContext(s_mockAgent, s_mockSession, [new ChatMessage(ChatRole.User, "msg")], new InvalidOperationException("Failed")); + + // Act + await provider.InvokedAsync(context); + + // Assert - StoreAIContextAsync was NOT called + Assert.Null(provider.LastStoredContext); + } + + [Fact] + public async Task InvokedCoreAsync_UsesCustomStoreInputFilterAsync() + { + // Arrange - filter that only keeps System messages + var provider = new TestAIContextProvider( + storeInputMessageFilter: msgs => msgs.Where(m => m.Role == ChatRole.System)); + var messages = new[] + { + new ChatMessage(ChatRole.User, "User msg"), + new ChatMessage(ChatRole.System, "System msg") + }; + var context = new AIContextProvider.InvokedContext(s_mockAgent, s_mockSession, messages, [new ChatMessage(ChatRole.Assistant, "Response")]); + + // Act + await provider.InvokedAsync(context); + + // Assert - only System messages were passed to store + Assert.NotNull(provider.LastStoredContext); + var storedRequest = provider.LastStoredContext!.RequestMessages.ToList(); + Assert.Single(storedRequest); + Assert.Equal("System msg", storedRequest[0].Text); + } + + [Fact] + public async Task InvokedCoreAsync_DefaultFilterExcludesNonExternalMessagesAsync() + { + // Arrange + var provider = new TestAIContextProvider(); + var external = new ChatMessage(ChatRole.User, "External"); + var fromHistory = new ChatMessage(ChatRole.User, "History") + .WithAgentRequestMessageSource(AgentRequestMessageSourceType.ChatHistory, "src"); + var fromContext = new ChatMessage(ChatRole.User, "Context") + .WithAgentRequestMessageSource(AgentRequestMessageSourceType.AIContextProvider, "src"); + var context = new AIContextProvider.InvokedContext(s_mockAgent, s_mockSession, [external, fromHistory, fromContext], []); + + // Act + await provider.InvokedAsync(context); + + // Assert - only External messages kept + Assert.NotNull(provider.LastStoredContext); + var storedRequest = provider.LastStoredContext!.RequestMessages.ToList(); + Assert.Single(storedRequest); + Assert.Equal("External", storedRequest[0].Text); + } + + #endregion + private sealed class TestAIContextProvider : AIContextProvider { - protected override ValueTask InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default) - => new(new AIContext()); + private readonly AIContext? _provideContext; + private readonly bool _captureFilteredContext; + + public InvokedContext? LastStoredContext { get; private set; } + + public InvokingContext? LastProvidedContext { get; private set; } + + public TestAIContextProvider( + AIContext? provideContext = null, + bool captureFilteredContext = false, + Func, IEnumerable>? provideInputMessageFilter = null, + Func, IEnumerable>? storeInputMessageFilter = null) + : base(provideInputMessageFilter, storeInputMessageFilter) + { + this._provideContext = provideContext; + this._captureFilteredContext = captureFilteredContext; + } + + protected override ValueTask ProvideAIContextAsync(InvokingContext context, CancellationToken cancellationToken = default) + { + if (this._captureFilteredContext) + { + this.LastProvidedContext = context; + } + + return new(this._provideContext ?? new AIContext()); + } + + protected override ValueTask StoreAIContextAsync(InvokedContext context, CancellationToken cancellationToken = default) + { + this.LastStoredContext = context; + return default; + } } + + /// + /// A provider that uses only base class defaults (no overrides of ProvideAIContextAsync/StoreAIContextAsync). + /// + private sealed class DefaultAIContextProvider : AIContextProvider; } diff --git a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/ChatHistoryProviderTStateTests.cs b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/ChatHistoryProviderTStateTests.cs new file mode 100644 index 0000000000..e3c0507e8d --- /dev/null +++ b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/ChatHistoryProviderTStateTests.cs @@ -0,0 +1,195 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Collections.Generic; +using System.Linq; +using System.Threading; +using System.Threading.Tasks; +using Microsoft.Extensions.AI; +using Moq; + +namespace Microsoft.Agents.AI.Abstractions.UnitTests; + +/// +/// Contains tests for the class. +/// +public class ChatHistoryProviderTStateTests +{ + private static readonly AIAgent s_mockAgent = new Mock().Object; + + #region GetOrInitializeState Tests + + [Fact] + public void GetOrInitializeState_InitializesFromStateInitializerOnFirstCall() + { + // Arrange + var expectedState = new TestState { Value = "initialized" }; + var provider = new TestChatHistoryProvider(_ => expectedState); + var session = new TestAgentSession(); + + // Act + var state = provider.GetState(session); + + // Assert + Assert.Same(expectedState, state); + } + + [Fact] + public void GetOrInitializeState_ReturnsCachedStateFromStateBagOnSecondCall() + { + // Arrange + var callCount = 0; + var provider = new TestChatHistoryProvider(_ => + { + callCount++; + return new TestState { Value = $"init-{callCount}" }; + }); + var session = new TestAgentSession(); + + // Act + var state1 = provider.GetState(session); + var state2 = provider.GetState(session); + + // Assert - initializer called only once; second call reads from StateBag + Assert.Equal(1, callCount); + Assert.Equal("init-1", state1.Value); + Assert.Equal("init-1", state2.Value); + } + + [Fact] + public void GetOrInitializeState_WorksWhenSessionIsNull() + { + // Arrange + var provider = new TestChatHistoryProvider(_ => new TestState { Value = "no-session" }); + + // Act + var state = provider.GetState(null); + + // Assert + Assert.Equal("no-session", state.Value); + } + + [Fact] + public void GetOrInitializeState_ReInitializesWhenSessionIsNull() + { + // Arrange - without a session, state can't be cached in StateBag + var callCount = 0; + var provider = new TestChatHistoryProvider(_ => + { + callCount++; + return new TestState { Value = $"init-{callCount}" }; + }); + + // Act + _ = provider.GetState(null); + provider.GetState(null); + + // Assert - initializer called each time since there's no session to cache in + Assert.Equal(2, callCount); + } + + #endregion + + #region SaveState Tests + + [Fact] + public void SaveState_SavesToStateBag() + { + // Arrange + var provider = new TestChatHistoryProvider(_ => new TestState()); + var session = new TestAgentSession(); + var state = new TestState { Value = "saved" }; + + // Act + provider.DoSaveState(session, state); + var retrieved = provider.GetState(session); + + // Assert + Assert.Equal("saved", retrieved.Value); + } + + [Fact] + public void SaveState_NoOpWhenSessionIsNull() + { + // Arrange + var provider = new TestChatHistoryProvider(_ => new TestState { Value = "default" }); + + // Act - should not throw + provider.DoSaveState(null, new TestState { Value = "saved" }); + + // Assert - no exception; can't verify further without a session + } + + #endregion + + #region StateKey Tests + + [Fact] + public void StateKey_DefaultsToTypeName() + { + // Arrange + var provider = new TestChatHistoryProvider(_ => new TestState()); + + // Act & Assert + Assert.Equal(nameof(TestChatHistoryProvider), provider.StateKey); + } + + [Fact] + public void StateKey_UsesCustomKeyWhenProvided() + { + // Arrange + var provider = new TestChatHistoryProvider(_ => new TestState(), stateKey: "custom-key"); + + // Act & Assert + Assert.Equal("custom-key", provider.StateKey); + } + + #endregion + + #region Integration Tests + + [Fact] + public async Task InvokingCoreAsync_CanUseStateInProvideChatHistoryAsync() + { + // Arrange + var provider = new TestChatHistoryProvider(_ => new TestState { Value = "state-value" }); + var session = new TestAgentSession(); + var context = new ChatHistoryProvider.InvokingContext(s_mockAgent, session, [new ChatMessage(ChatRole.User, "Hi")]); + + // Act + var result = (await provider.InvokingAsync(context)).ToList(); + + // Assert - the provider uses state to produce history messages + Assert.Equal(2, result.Count); + Assert.Contains("state-value", result[0].Text); + } + + #endregion + + public sealed class TestState + { + public string Value { get; set; } = string.Empty; + } + + private sealed class TestChatHistoryProvider : ChatHistoryProvider + { + public TestChatHistoryProvider( + Func stateInitializer, + string? stateKey = null) + : base(stateInitializer, stateKey, null, null, null) + { + } + + public TestState GetState(AgentSession? session) => this.GetOrInitializeState(session); + + public void DoSaveState(AgentSession? session, TestState state) => this.SaveState(session, state); + + protected override ValueTask> ProvideChatHistoryAsync(InvokingContext context, CancellationToken cancellationToken = default) + { + var state = this.GetOrInitializeState(context.Session); + return new(new[] { new ChatMessage(ChatRole.System, $"History from state: {state.Value}") }); + } + } + + private sealed class TestAgentSession : AgentSession; +} diff --git a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/ChatHistoryProviderTests.cs b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/ChatHistoryProviderTests.cs index 7fccca8d1b..5df661f009 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/ChatHistoryProviderTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/ChatHistoryProviderTests.cs @@ -274,12 +274,279 @@ public class ChatHistoryProviderTests #endregion + #region InvokingAsync / InvokedAsync Null Check Tests + + [Fact] + public async Task InvokingAsync_NullContext_ThrowsArgumentNullExceptionAsync() + { + // Arrange + var provider = new TestChatHistoryProvider(); + + // Act & Assert + await Assert.ThrowsAsync(() => provider.InvokingAsync(null!).AsTask()); + } + + [Fact] + public async Task InvokedAsync_NullContext_ThrowsArgumentNullExceptionAsync() + { + // Arrange + var provider = new TestChatHistoryProvider(); + + // Act & Assert + await Assert.ThrowsAsync(() => provider.InvokedAsync(null!).AsTask()); + } + + #endregion + + #region InvokingCoreAsync Tests + + [Fact] + public async Task InvokingCoreAsync_CallsProvideChatHistoryAndReturnsMessagesAsync() + { + // Arrange + var historyMessages = new[] { new ChatMessage(ChatRole.User, "History message") }; + var provider = new TestChatHistoryProvider(provideMessages: historyMessages); + var requestMessages = new[] { new ChatMessage(ChatRole.User, "Request message") }; + var context = new ChatHistoryProvider.InvokingContext(s_mockAgent, s_mockSession, requestMessages); + + // Act + var result = (await provider.InvokingAsync(context)).ToList(); + + // Assert + Assert.Equal(2, result.Count); + Assert.Equal("History message", result[0].Text); + Assert.Equal("Request message", result[1].Text); + } + + [Fact] + public async Task InvokingCoreAsync_HistoryAppearsBeforeRequestMessagesAsync() + { + // Arrange + var historyMessages = new[] + { + new ChatMessage(ChatRole.User, "Hist1"), + new ChatMessage(ChatRole.Assistant, "Hist2") + }; + var provider = new TestChatHistoryProvider(provideMessages: historyMessages); + var requestMessages = new[] { new ChatMessage(ChatRole.User, "Req1") }; + var context = new ChatHistoryProvider.InvokingContext(s_mockAgent, s_mockSession, requestMessages); + + // Act + var result = (await provider.InvokingAsync(context)).ToList(); + + // Assert + Assert.Equal(3, result.Count); + Assert.Equal("Hist1", result[0].Text); + Assert.Equal("Hist2", result[1].Text); + Assert.Equal("Req1", result[2].Text); + } + + [Fact] + public async Task InvokingCoreAsync_StampsHistoryMessagesWithChatHistorySourceAsync() + { + // Arrange + var historyMessages = new[] { new ChatMessage(ChatRole.User, "History") }; + var provider = new TestChatHistoryProvider(provideMessages: historyMessages); + var context = new ChatHistoryProvider.InvokingContext(s_mockAgent, s_mockSession, []); + + // Act + var result = (await provider.InvokingAsync(context)).ToList(); + + // Assert + Assert.Single(result); + Assert.Equal(AgentRequestMessageSourceType.ChatHistory, result[0].GetAgentRequestMessageSourceType()); + } + + [Fact] + public async Task InvokingCoreAsync_NoFilterAppliedWhenProvideOutputFilterIsNullAsync() + { + // Arrange + var historyMessages = new[] + { + new ChatMessage(ChatRole.User, "User msg"), + new ChatMessage(ChatRole.System, "System msg"), + new ChatMessage(ChatRole.Assistant, "Assistant msg") + }; + var provider = new TestChatHistoryProvider(provideMessages: historyMessages); + var context = new ChatHistoryProvider.InvokingContext(s_mockAgent, s_mockSession, []); + + // Act + var result = (await provider.InvokingAsync(context)).ToList(); + + // Assert - all 3 history messages returned (no filter) + Assert.Equal(3, result.Count); + } + + [Fact] + public async Task InvokingCoreAsync_AppliesProvideOutputFilterWhenProvidedAsync() + { + // Arrange + var historyMessages = new[] + { + new ChatMessage(ChatRole.User, "User msg"), + new ChatMessage(ChatRole.System, "System msg"), + new ChatMessage(ChatRole.Assistant, "Assistant msg") + }; + var provider = new TestChatHistoryProvider( + provideMessages: historyMessages, + provideOutputMessageFilter: msgs => msgs.Where(m => m.Role == ChatRole.User)); + var context = new ChatHistoryProvider.InvokingContext(s_mockAgent, s_mockSession, []); + + // Act + var result = (await provider.InvokingAsync(context)).ToList(); + + // Assert - only User messages remain after filter + Assert.Single(result); + Assert.Equal("User msg", result[0].Text); + } + + [Fact] + public async Task InvokingCoreAsync_ReturnsEmptyHistoryByDefaultAsync() + { + // Arrange - provider that doesn't override ProvideChatHistoryAsync (uses base default) + var provider = new DefaultChatHistoryProvider(); + var requestMessages = new[] { new ChatMessage(ChatRole.User, "Hello") }; + var context = new ChatHistoryProvider.InvokingContext(s_mockAgent, s_mockSession, requestMessages); + + // Act + var result = (await provider.InvokingAsync(context)).ToList(); + + // Assert - only the request message (no history) + Assert.Single(result); + Assert.Equal("Hello", result[0].Text); + } + + #endregion + + #region InvokedCoreAsync Tests + + [Fact] + public async Task InvokedCoreAsync_CallsStoreChatHistoryWithFilteredMessagesAsync() + { + // Arrange + var provider = new TestChatHistoryProvider(); + var externalMessage = new ChatMessage(ChatRole.User, "External"); + var chatHistoryMessage = new ChatMessage(ChatRole.User, "From history") + .WithAgentRequestMessageSource(AgentRequestMessageSourceType.ChatHistory, "source"); + var responseMessages = new[] { new ChatMessage(ChatRole.Assistant, "Response") }; + var context = new ChatHistoryProvider.InvokedContext(s_mockAgent, s_mockSession, new[] { externalMessage, chatHistoryMessage }, responseMessages); + + // Act + await provider.InvokedAsync(context); + + // Assert - default filter excludes ChatHistory-sourced messages + Assert.NotNull(provider.LastStoredContext); + var storedRequest = provider.LastStoredContext!.RequestMessages.ToList(); + Assert.Single(storedRequest); + Assert.Equal("External", storedRequest[0].Text); + Assert.Same(responseMessages, provider.LastStoredContext.ResponseMessages); + } + + [Fact] + public async Task InvokedCoreAsync_SkipsStorageWhenInvokeExceptionIsNotNullAsync() + { + // Arrange + var provider = new TestChatHistoryProvider(); + var context = new ChatHistoryProvider.InvokedContext(s_mockAgent, s_mockSession, [new ChatMessage(ChatRole.User, "msg")], new InvalidOperationException("Failed")); + + // Act + await provider.InvokedAsync(context); + + // Assert - StoreChatHistoryAsync was NOT called + Assert.Null(provider.LastStoredContext); + } + + [Fact] + public async Task InvokedCoreAsync_UsesCustomStoreInputFilterAsync() + { + // Arrange - filter that only keeps System messages + var provider = new TestChatHistoryProvider( + storeInputMessageFilter: msgs => msgs.Where(m => m.Role == ChatRole.System)); + var messages = new[] + { + new ChatMessage(ChatRole.User, "User msg"), + new ChatMessage(ChatRole.System, "System msg") + }; + var context = new ChatHistoryProvider.InvokedContext(s_mockAgent, s_mockSession, messages, [new ChatMessage(ChatRole.Assistant, "Response")]); + + // Act + await provider.InvokedAsync(context); + + // Assert - only System messages were passed to store + Assert.NotNull(provider.LastStoredContext); + var storedRequest = provider.LastStoredContext!.RequestMessages.ToList(); + Assert.Single(storedRequest); + Assert.Equal("System msg", storedRequest[0].Text); + } + + [Fact] + public async Task InvokedCoreAsync_DefaultFilterExcludesChatHistorySourcedMessagesAsync() + { + // Arrange + var provider = new TestChatHistoryProvider(); + var external = new ChatMessage(ChatRole.User, "External"); + var fromHistory = new ChatMessage(ChatRole.User, "History") + .WithAgentRequestMessageSource(AgentRequestMessageSourceType.ChatHistory, "src"); + var fromContext = new ChatMessage(ChatRole.User, "Context") + .WithAgentRequestMessageSource(AgentRequestMessageSourceType.AIContextProvider, "src"); + var context = new ChatHistoryProvider.InvokedContext(s_mockAgent, s_mockSession, [external, fromHistory, fromContext], []); + + // Act + await provider.InvokedAsync(context); + + // Assert - External and AIContextProvider messages kept, ChatHistory excluded + Assert.NotNull(provider.LastStoredContext); + var storedRequest = provider.LastStoredContext!.RequestMessages.ToList(); + Assert.Equal(2, storedRequest.Count); + Assert.Equal("External", storedRequest[0].Text); + Assert.Equal("Context", storedRequest[1].Text); + } + + [Fact] + public async Task InvokedCoreAsync_PassesResponseMessagesToStoreAsync() + { + // Arrange + var provider = new TestChatHistoryProvider(); + var responseMessages = new[] { new ChatMessage(ChatRole.Assistant, "Resp1"), new ChatMessage(ChatRole.Assistant, "Resp2") }; + var context = new ChatHistoryProvider.InvokedContext(s_mockAgent, s_mockSession, [new ChatMessage(ChatRole.User, "msg")], responseMessages); + + // Act + await provider.InvokedAsync(context); + + // Assert + Assert.NotNull(provider.LastStoredContext); + Assert.Same(responseMessages, provider.LastStoredContext!.ResponseMessages); + } + + #endregion + private sealed class TestChatHistoryProvider : ChatHistoryProvider { - protected override ValueTask> InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default) - => new(new ChatMessage[] { new(ChatRole.User, "Test Message") }.Concat(context.RequestMessages)); + private readonly IEnumerable? _provideMessages; - protected override ValueTask InvokedCoreAsync(InvokedContext context, CancellationToken cancellationToken = default) - => default; + public InvokedContext? LastStoredContext { get; private set; } + + public TestChatHistoryProvider( + IEnumerable? provideMessages = null, + Func, IEnumerable>? provideOutputMessageFilter = null, + Func, IEnumerable>? storeInputMessageFilter = null) + : base(provideOutputMessageFilter, storeInputMessageFilter) + { + this._provideMessages = provideMessages; + } + + protected override ValueTask> ProvideChatHistoryAsync(InvokingContext context, CancellationToken cancellationToken = default) + => new(this._provideMessages ?? []); + + protected override ValueTask StoreChatHistoryAsync(InvokedContext context, CancellationToken cancellationToken = default) + { + this.LastStoredContext = context; + return default; + } } + + /// + /// A provider that uses only base class defaults (no overrides of ProvideChatHistoryAsync/StoreChatHistoryAsync). + /// + private sealed class DefaultChatHistoryProvider : ChatHistoryProvider; } diff --git a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/InMemoryChatHistoryProviderTests.cs b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/InMemoryChatHistoryProviderTests.cs index b907529241..ebe1131ab7 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/InMemoryChatHistoryProviderTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/InMemoryChatHistoryProviderTests.cs @@ -446,7 +446,7 @@ public class InMemoryChatHistoryProviderTests var session = CreateMockSession(); var provider = new InMemoryChatHistoryProvider(new InMemoryChatHistoryProviderOptions { - RetrievalOutputMessageFilter = messages => messages.Where(m => m.Role == ChatRole.User) + ProvideOutputMessageFilter = messages => messages.Where(m => m.Role == ChatRole.User) }); provider.SetMessages(session, [ diff --git a/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatHistoryProviderTests.cs b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatHistoryProviderTests.cs index 12bd467e71..4c142d3c1e 100644 --- a/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatHistoryProviderTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatHistoryProviderTests.cs @@ -881,12 +881,12 @@ public sealed class CosmosChatHistoryProviderTests : IAsyncLifetime, IDisposable this.SkipIfEmulatorNotAvailable(); var session = CreateMockSession(); var conversationId = Guid.NewGuid().ToString(); - using var provider = new CosmosChatHistoryProvider(this._connectionString, s_testDatabaseId, TestContainerId, - _ => new CosmosChatHistoryProvider.State(conversationId)) - { - // Custom filter: only store External messages (also exclude AIContextProvider) - StorageInputMessageFilter = messages => messages.Where(m => m.GetAgentRequestMessageSourceType() == AgentRequestMessageSourceType.External) - }; + using var provider = new CosmosChatHistoryProvider( + this._connectionString, + s_testDatabaseId, + TestContainerId, + _ => new CosmosChatHistoryProvider.State(conversationId), + storeInputMessageFilter: messages => messages.Where(m => m.GetAgentRequestMessageSourceType() == AgentRequestMessageSourceType.External)); var requestMessages = new[] { @@ -919,12 +919,12 @@ public sealed class CosmosChatHistoryProviderTests : IAsyncLifetime, IDisposable this.SkipIfEmulatorNotAvailable(); var session = CreateMockSession(); var conversationId = Guid.NewGuid().ToString(); - using var provider = new CosmosChatHistoryProvider(this._connectionString, s_testDatabaseId, TestContainerId, - _ => new CosmosChatHistoryProvider.State(conversationId)) - { - // Only return User messages when retrieving - RetrievalOutputMessageFilter = messages => messages.Where(m => m.Role == ChatRole.User) - }; + using var provider = new CosmosChatHistoryProvider( + this._connectionString, + s_testDatabaseId, + TestContainerId, + _ => new CosmosChatHistoryProvider.State(conversationId), + provideOutputMessageFilter: messages => messages.Where(m => m.Role == ChatRole.User)); var requestMessages = new[] { @@ -943,7 +943,7 @@ public sealed class CosmosChatHistoryProviderTests : IAsyncLifetime, IDisposable var invokingContext = new ChatHistoryProvider.InvokingContext(s_mockAgent, session, []); var messages = (await provider.InvokingAsync(invokingContext)).ToList(); - // Assert - Only User messages returned (System and Assistant filtered by RetrievalOutputMessageFilter) + // Assert - Only User messages returned (System and Assistant filtered by ProvideOutputMessageFilter) Assert.Single(messages); Assert.Equal("User message", messages[0].Text); Assert.Equal(ChatRole.User, messages[0].Role); diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentOptionsTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentOptionsTests.cs index f69fb3d636..b8e3c57af6 100644 --- a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentOptionsTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentOptionsTests.cs @@ -115,8 +115,8 @@ public class ChatClientAgentOptionsTests const string Description = "Test description"; var tools = new List { AIFunctionFactory.Create(() => "test") }; - var mockChatHistoryProvider = new Mock().Object; - var mockAIContextProvider = new Mock().Object; + var mockChatHistoryProvider = new Mock(null, null).Object; + var mockAIContextProvider = new Mock(null, null).Object; var original = new ChatClientAgentOptions() { @@ -149,8 +149,8 @@ public class ChatClientAgentOptionsTests public void Clone_WithoutProvidingChatOptions_ClonesCorrectly() { // Arrange - var mockChatHistoryProvider = new Mock().Object; - var mockAIContextProvider = new Mock().Object; + var mockChatHistoryProvider = new Mock(null, null).Object; + var mockAIContextProvider = new Mock(null, null).Object; var original = new ChatClientAgentOptions { diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentTests.cs index d90fe5153a..12446c89c0 100644 --- a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentTests.cs @@ -488,7 +488,7 @@ public partial class ChatClientAgentTests }) .ReturnsAsync(new ChatResponse(responseMessages)); - var mockProvider = new Mock(); + var mockProvider = new Mock(null, null); mockProvider .Protected() .Setup>("InvokingCoreAsync", ItExpr.IsAny(), ItExpr.IsAny()) @@ -559,7 +559,7 @@ public partial class ChatClientAgentTests It.IsAny())) .Throws(new InvalidOperationException("downstream failure")); - var mockProvider = new Mock(); + var mockProvider = new Mock(null, null); mockProvider .Protected() .Setup>("InvokingCoreAsync", ItExpr.IsAny(), ItExpr.IsAny()) @@ -617,7 +617,7 @@ public partial class ChatClientAgentTests }) .ReturnsAsync(new ChatResponse([new(ChatRole.Assistant, "response")])); - var mockProvider = new Mock(); + var mockProvider = new Mock(null, null); mockProvider .Protected() .Setup>("InvokingCoreAsync", ItExpr.IsAny(), ItExpr.IsAny()) @@ -677,7 +677,7 @@ public partial class ChatClientAgentTests .ReturnsAsync(new ChatResponse(responseMessages)); // Provider 1: adds a system message and a tool - var mockProvider1 = new Mock(); + var mockProvider1 = new Mock(null, null); mockProvider1.SetupGet(p => p.StateKey).Returns("Provider1"); mockProvider1 .Protected() @@ -696,7 +696,7 @@ public partial class ChatClientAgentTests // Provider 2: adds another system message and verifies it receives accumulated context from provider 1 AIContext? provider2ReceivedContext = null; - var mockProvider2 = new Mock(); + var mockProvider2 = new Mock(null, null); mockProvider2.SetupGet(p => p.StateKey).Returns("Provider2"); mockProvider2 .Protected() @@ -784,7 +784,7 @@ public partial class ChatClientAgentTests It.IsAny())) .ThrowsAsync(new InvalidOperationException("downstream failure")); - var mockProvider1 = new Mock(); + var mockProvider1 = new Mock(null, null); mockProvider1.SetupGet(p => p.StateKey).Returns("Provider1"); mockProvider1 .Protected() @@ -801,7 +801,7 @@ public partial class ChatClientAgentTests .Setup("InvokedCoreAsync", ItExpr.IsAny(), ItExpr.IsAny()) .Returns(new ValueTask()); - var mockProvider2 = new Mock(); + var mockProvider2 = new Mock(null, null); mockProvider2.SetupGet(p => p.StateKey).Returns("Provider2"); mockProvider2 .Protected() @@ -869,7 +869,7 @@ public partial class ChatClientAgentTests }) .Returns(ToAsyncEnumerableAsync(responseUpdates)); - var mockProvider1 = new Mock(); + var mockProvider1 = new Mock(null, null); mockProvider1.SetupGet(p => p.StateKey).Returns("Provider1"); mockProvider1 .Protected() @@ -886,7 +886,7 @@ public partial class ChatClientAgentTests .Setup("InvokedCoreAsync", ItExpr.IsAny(), ItExpr.IsAny()) .Returns(new ValueTask()); - var mockProvider2 = new Mock(); + var mockProvider2 = new Mock(null, null); mockProvider2.SetupGet(p => p.StateKey).Returns("Provider2"); mockProvider2 .Protected() @@ -1828,7 +1828,7 @@ public partial class ChatClientAgentTests }) .Returns(ToAsyncEnumerableAsync(responseUpdates)); - var mockProvider = new Mock(); + var mockProvider = new Mock(null, null); mockProvider .Protected() .Setup>("InvokingCoreAsync", ItExpr.IsAny(), ItExpr.IsAny()) @@ -1907,7 +1907,7 @@ public partial class ChatClientAgentTests It.IsAny())) .Throws(new InvalidOperationException("downstream failure")); - var mockProvider = new Mock(); + var mockProvider = new Mock(null, null); mockProvider .Protected() .Setup>("InvokingCoreAsync", ItExpr.IsAny(), ItExpr.IsAny()) diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_BackgroundResponsesTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_BackgroundResponsesTests.cs index 84285ff9c4..64835f2b2f 100644 --- a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_BackgroundResponsesTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_BackgroundResponsesTests.cs @@ -338,7 +338,7 @@ public class ChatClientAgent_BackgroundResponsesTests List capturedMessages = []; // Create a mock chat history provider that would normally provide messages - var mockChatHistoryProvider = new Mock(); + var mockChatHistoryProvider = new Mock(null, null); mockChatHistoryProvider.SetupGet(p => p.StateKey).Returns("ChatHistoryProvider"); mockChatHistoryProvider .Protected() @@ -346,7 +346,7 @@ public class ChatClientAgent_BackgroundResponsesTests .ReturnsAsync([new(ChatRole.User, "Message from chat history provider")]); // Create a mock AI context provider that would normally provide context - var mockContextProvider = new Mock(); + var mockContextProvider = new Mock(null, null); mockContextProvider.SetupGet(p => p.StateKey).Returns("Provider1"); mockContextProvider .Protected() @@ -407,7 +407,7 @@ public class ChatClientAgent_BackgroundResponsesTests List capturedMessages = []; // Create a mock chat history provider that would normally provide messages - var mockChatHistoryProvider = new Mock(); + var mockChatHistoryProvider = new Mock(null, null); mockChatHistoryProvider.SetupGet(p => p.StateKey).Returns("ChatHistoryProvider"); mockChatHistoryProvider .Protected() @@ -415,7 +415,7 @@ public class ChatClientAgent_BackgroundResponsesTests .ReturnsAsync([new(ChatRole.User, "Message from chat history provider")]); // Create a mock AI context provider that would normally provide context - var mockContextProvider = new Mock(); + var mockContextProvider = new Mock(null, null); mockContextProvider.SetupGet(p => p.StateKey).Returns("Provider1"); mockContextProvider .Protected() @@ -638,7 +638,7 @@ public class ChatClientAgent_BackgroundResponsesTests .Returns(ToAsyncEnumerableAsync(returnUpdates)); List capturedMessagesAddedToProvider = []; - var mockChatHistoryProvider = new Mock(); + var mockChatHistoryProvider = new Mock(null, null); mockChatHistoryProvider.SetupGet(p => p.StateKey).Returns("ChatHistoryProvider"); mockChatHistoryProvider .Protected() @@ -647,7 +647,7 @@ public class ChatClientAgent_BackgroundResponsesTests .Returns(new ValueTask()); AIContextProvider.InvokedContext? capturedInvokedContext = null; - var mockContextProvider = new Mock(); + var mockContextProvider = new Mock(null, null); mockContextProvider.SetupGet(p => p.StateKey).Returns("Provider1"); mockContextProvider .Protected() @@ -702,7 +702,7 @@ public class ChatClientAgent_BackgroundResponsesTests .Returns(ToAsyncEnumerableAsync(Array.Empty())); List capturedMessagesAddedToProvider = []; - var mockChatHistoryProvider = new Mock(); + var mockChatHistoryProvider = new Mock(null, null); mockChatHistoryProvider.SetupGet(p => p.StateKey).Returns("ChatHistoryProvider"); mockChatHistoryProvider .Protected() @@ -711,7 +711,7 @@ public class ChatClientAgent_BackgroundResponsesTests .Returns(new ValueTask()); AIContextProvider.InvokedContext? capturedInvokedContext = null; - var mockContextProvider = new Mock(); + var mockContextProvider = new Mock(null, null); mockContextProvider.SetupGet(p => p.StateKey).Returns("Provider1"); mockContextProvider .Protected() diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_ChatHistoryManagementTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_ChatHistoryManagementTests.cs index 4731805af7..423c867abd 100644 --- a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_ChatHistoryManagementTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_ChatHistoryManagementTests.cs @@ -185,7 +185,7 @@ public class ChatClientAgent_ChatHistoryManagementTests It.IsAny(), It.IsAny())).ReturnsAsync(new ChatResponse([new(ChatRole.Assistant, "response")])); - Mock mockChatHistoryProvider = new(); + Mock mockChatHistoryProvider = new(null, null); mockChatHistoryProvider .Protected() .Setup>>("InvokingCoreAsync", ItExpr.IsAny(), ItExpr.IsAny()) @@ -240,7 +240,7 @@ public class ChatClientAgent_ChatHistoryManagementTests It.IsAny(), It.IsAny())).Throws(new InvalidOperationException("Test Error")); - Mock mockChatHistoryProvider = new(); + Mock mockChatHistoryProvider = new(null, null); mockChatHistoryProvider .Protected() .Setup>>("InvokingCoreAsync", ItExpr.IsAny(), ItExpr.IsAny()) @@ -311,7 +311,7 @@ public class ChatClientAgent_ChatHistoryManagementTests It.IsAny())).ReturnsAsync(new ChatResponse([new(ChatRole.Assistant, "response")])); // Arrange a chat history provider to override the factory provided one. - Mock mockOverrideChatHistoryProvider = new(); + Mock mockOverrideChatHistoryProvider = new(null, null); mockOverrideChatHistoryProvider .Protected() .Setup>>("InvokingCoreAsync", ItExpr.IsAny(), ItExpr.IsAny()) @@ -324,7 +324,7 @@ public class ChatClientAgent_ChatHistoryManagementTests // Arrange a chat history provider to provide to the agent at construction time. // This one shouldn't be used since it is being overridden. - Mock mockAgentOptionsChatHistoryProvider = new(); + Mock mockAgentOptionsChatHistoryProvider = new(null, null); mockAgentOptionsChatHistoryProvider .Protected() .Setup>>("InvokingCoreAsync", ItExpr.IsAny(), ItExpr.IsAny()) From 8015e00f566f06cc374da3db03f77219f45db612 Mon Sep 17 00:00:00 2001 From: westey <164392973+westey-m@users.noreply.github.com> Date: Tue, 17 Feb 2026 12:53:33 +0000 Subject: [PATCH 17/22] .NET: Surface downstream experimental flags and remove unnecessary suppressions (#3968) * Surface downstream experimental flags and remove unecessary suppressions * Fix format error * Fix file encoding. * Address PR comments. --- dotnet/agent-framework-dotnet.slnx | 4 +++ dotnet/eng/MSBuild/Shared.props | 3 ++ .../AgentWithHostedMCP.csproj | 1 - .../ExperimentalAttribute.cs | 4 +-- .../AgentResponse.cs | 2 ++ .../AgentRunOptions.cs | 3 ++ .../Microsoft.Agents.AI.Abstractions.csproj | 2 ++ .../AzureAIProjectChatClient.cs | 5 +-- .../AzureAIProjectChatClientExtensions.cs | 6 ++-- .../Microsoft.Agents.AI.AzureAI.csproj | 5 +++ .../Microsoft.Agents.AI.CosmosNoSql.csproj | 1 - .../Microsoft.Agents.AI.DurableTask.csproj | 3 +- .../Microsoft.Agents.AI.Hosting.OpenAI.csproj | 2 +- ...StreamingResponseUpdateCollectionResult.cs | 3 ++ .../Extensions/AIAgentWithOpenAIExtensions.cs | 3 ++ .../Extensions/AgentResponseExtensions.cs | 3 ++ .../OpenAIAssistantClientExtensions.cs | 3 ++ .../OpenAIResponseClientExtensions.cs | 3 ++ .../Microsoft.Agents.AI.OpenAI.csproj | 6 +++- .../Shared/DiagnosticIds/DiagnosticsIds.cs | 33 +++++++++++++++++++ dotnet/src/Shared/DiagnosticIds/README.md | 11 +++++++ ...oft.Agents.AI.CosmosNoSql.UnitTests.csproj | 1 - 22 files changed, 93 insertions(+), 14 deletions(-) create mode 100644 dotnet/src/Shared/DiagnosticIds/DiagnosticsIds.cs create mode 100644 dotnet/src/Shared/DiagnosticIds/README.md diff --git a/dotnet/agent-framework-dotnet.slnx b/dotnet/agent-framework-dotnet.slnx index e592e80803..3a3c6bba9a 100644 --- a/dotnet/agent-framework-dotnet.slnx +++ b/dotnet/agent-framework-dotnet.slnx @@ -371,6 +371,10 @@ + + + + diff --git a/dotnet/eng/MSBuild/Shared.props b/dotnet/eng/MSBuild/Shared.props index 95eb5e13da..9b4771a64e 100644 --- a/dotnet/eng/MSBuild/Shared.props +++ b/dotnet/eng/MSBuild/Shared.props @@ -23,4 +23,7 @@ + + + diff --git a/dotnet/samples/HostedAgents/AgentWithHostedMCP/AgentWithHostedMCP.csproj b/dotnet/samples/HostedAgents/AgentWithHostedMCP/AgentWithHostedMCP.csproj index 1244b81542..b00e4a48de 100644 --- a/dotnet/samples/HostedAgents/AgentWithHostedMCP/AgentWithHostedMCP.csproj +++ b/dotnet/samples/HostedAgents/AgentWithHostedMCP/AgentWithHostedMCP.csproj @@ -15,7 +15,6 @@ and cannot access parent folders where Directory.Packages.props resides. --> false - $(NoWarn);MEAI001;OPENAI001 - - $(NoWarn);CA2007;MEAI001 + $(NoWarn);CA2007 diff --git a/dotnet/src/Microsoft.Agents.AI.Hosting.OpenAI/Microsoft.Agents.AI.Hosting.OpenAI.csproj b/dotnet/src/Microsoft.Agents.AI.Hosting.OpenAI/Microsoft.Agents.AI.Hosting.OpenAI.csproj index 923f8e3eb6..78364f4a30 100644 --- a/dotnet/src/Microsoft.Agents.AI.Hosting.OpenAI/Microsoft.Agents.AI.Hosting.OpenAI.csproj +++ b/dotnet/src/Microsoft.Agents.AI.Hosting.OpenAI/Microsoft.Agents.AI.Hosting.OpenAI.csproj @@ -2,7 +2,7 @@ $(TargetFrameworksCore) - $(NoWarn);OPENAI001;MEAI001 + $(NoWarn);MEAI001 Microsoft.Agents.AI.Hosting.OpenAI alpha $(InterceptorsNamespaces);Microsoft.AspNetCore.Http.Generated diff --git a/dotnet/src/Microsoft.Agents.AI.OpenAI/ChatClient/AsyncStreamingResponseUpdateCollectionResult.cs b/dotnet/src/Microsoft.Agents.AI.OpenAI/ChatClient/AsyncStreamingResponseUpdateCollectionResult.cs index 77400b3377..a7e0aa5b67 100644 --- a/dotnet/src/Microsoft.Agents.AI.OpenAI/ChatClient/AsyncStreamingResponseUpdateCollectionResult.cs +++ b/dotnet/src/Microsoft.Agents.AI.OpenAI/ChatClient/AsyncStreamingResponseUpdateCollectionResult.cs @@ -1,10 +1,13 @@ // Copyright (c) Microsoft. All rights reserved. using System.ClientModel; +using System.Diagnostics.CodeAnalysis; +using Microsoft.Shared.DiagnosticIds; using OpenAI.Responses; namespace Microsoft.Agents.AI.OpenAI; +[Experimental(DiagnosticIds.Experiments.AIOpenAIResponses)] internal sealed class AsyncStreamingResponseUpdateCollectionResult : AsyncCollectionResult { private readonly IAsyncEnumerable _updates; diff --git a/dotnet/src/Microsoft.Agents.AI.OpenAI/Extensions/AIAgentWithOpenAIExtensions.cs b/dotnet/src/Microsoft.Agents.AI.OpenAI/Extensions/AIAgentWithOpenAIExtensions.cs index a3598c0abc..5dc0b372ac 100644 --- a/dotnet/src/Microsoft.Agents.AI.OpenAI/Extensions/AIAgentWithOpenAIExtensions.cs +++ b/dotnet/src/Microsoft.Agents.AI.OpenAI/Extensions/AIAgentWithOpenAIExtensions.cs @@ -1,7 +1,9 @@ // Copyright (c) Microsoft. All rights reserved. using System.ClientModel; +using System.Diagnostics.CodeAnalysis; using Microsoft.Agents.AI.OpenAI; +using Microsoft.Shared.DiagnosticIds; using Microsoft.Shared.Diagnostics; using OpenAI.Chat; using OpenAI.Responses; @@ -18,6 +20,7 @@ namespace Microsoft.Agents.AI; /// The methods handle the conversion between OpenAI chat message types and Microsoft Extensions AI types, /// and return OpenAI objects directly from the agent's . /// +[Experimental(DiagnosticIds.Experiments.AIOpenAIResponses)] public static class AIAgentWithOpenAIExtensions { /// diff --git a/dotnet/src/Microsoft.Agents.AI.OpenAI/Extensions/AgentResponseExtensions.cs b/dotnet/src/Microsoft.Agents.AI.OpenAI/Extensions/AgentResponseExtensions.cs index e855aaef56..f9a247832a 100644 --- a/dotnet/src/Microsoft.Agents.AI.OpenAI/Extensions/AgentResponseExtensions.cs +++ b/dotnet/src/Microsoft.Agents.AI.OpenAI/Extensions/AgentResponseExtensions.cs @@ -1,6 +1,8 @@ // Copyright (c) Microsoft. All rights reserved. +using System.Diagnostics.CodeAnalysis; using Microsoft.Agents.AI; +using Microsoft.Shared.DiagnosticIds; using Microsoft.Shared.Diagnostics; using OpenAI.Chat; using OpenAI.Responses; @@ -11,6 +13,7 @@ namespace Microsoft.Agents.AI; /// Provides extension methods for and instances to /// create or extract native OpenAI response objects from the Microsoft Agent Framework responses. /// +[Experimental(DiagnosticIds.Experiments.AIOpenAIResponses)] public static class AgentResponseExtensions { /// diff --git a/dotnet/src/Microsoft.Agents.AI.OpenAI/Extensions/OpenAIAssistantClientExtensions.cs b/dotnet/src/Microsoft.Agents.AI.OpenAI/Extensions/OpenAIAssistantClientExtensions.cs index c56a63c76e..a1f083ae06 100644 --- a/dotnet/src/Microsoft.Agents.AI.OpenAI/Extensions/OpenAIAssistantClientExtensions.cs +++ b/dotnet/src/Microsoft.Agents.AI.OpenAI/Extensions/OpenAIAssistantClientExtensions.cs @@ -1,9 +1,11 @@ // Copyright (c) Microsoft. All rights reserved. using System.ClientModel; +using System.Diagnostics.CodeAnalysis; using Microsoft.Agents.AI; using Microsoft.Extensions.AI; using Microsoft.Extensions.Logging; +using Microsoft.Shared.DiagnosticIds; using Microsoft.Shared.Diagnostics; namespace OpenAI.Assistants; @@ -18,6 +20,7 @@ namespace OpenAI.Assistants; /// The methods handle the conversion from OpenAI clients to instances and then wrap them /// in objects that implement the interface. /// +[Experimental(DiagnosticIds.Experiments.AIOpenAIAssistants)] public static class OpenAIAssistantClientExtensions { /// diff --git a/dotnet/src/Microsoft.Agents.AI.OpenAI/Extensions/OpenAIResponseClientExtensions.cs b/dotnet/src/Microsoft.Agents.AI.OpenAI/Extensions/OpenAIResponseClientExtensions.cs index bc9f28c37d..1bef6cab4d 100644 --- a/dotnet/src/Microsoft.Agents.AI.OpenAI/Extensions/OpenAIResponseClientExtensions.cs +++ b/dotnet/src/Microsoft.Agents.AI.OpenAI/Extensions/OpenAIResponseClientExtensions.cs @@ -1,8 +1,10 @@ // Copyright (c) Microsoft. All rights reserved. +using System.Diagnostics.CodeAnalysis; using Microsoft.Agents.AI; using Microsoft.Extensions.AI; using Microsoft.Extensions.Logging; +using Microsoft.Shared.DiagnosticIds; using Microsoft.Shared.Diagnostics; namespace OpenAI.Responses; @@ -17,6 +19,7 @@ namespace OpenAI.Responses; /// The methods handle the conversion from OpenAI clients to instances and then wrap them /// in objects that implement the interface. /// +[Experimental(DiagnosticIds.Experiments.AIOpenAIResponses)] public static class OpenAIResponseClientExtensions { /// diff --git a/dotnet/src/Microsoft.Agents.AI.OpenAI/Microsoft.Agents.AI.OpenAI.csproj b/dotnet/src/Microsoft.Agents.AI.OpenAI/Microsoft.Agents.AI.OpenAI.csproj index 3de68137ba..6e1947418b 100644 --- a/dotnet/src/Microsoft.Agents.AI.OpenAI/Microsoft.Agents.AI.OpenAI.csproj +++ b/dotnet/src/Microsoft.Agents.AI.OpenAI/Microsoft.Agents.AI.OpenAI.csproj @@ -2,13 +2,17 @@ preview - $(NoWarn);OPENAI001; enable true + + true + true + + diff --git a/dotnet/src/Shared/DiagnosticIds/DiagnosticsIds.cs b/dotnet/src/Shared/DiagnosticIds/DiagnosticsIds.cs new file mode 100644 index 0000000000..6316c6f607 --- /dev/null +++ b/dotnet/src/Shared/DiagnosticIds/DiagnosticsIds.cs @@ -0,0 +1,33 @@ +// Copyright (c) Microsoft. All rights reserved. + +namespace Microsoft.Shared.DiagnosticIds; + +/// +/// Various diagnostic IDs reported by this repo. +/// +internal static class DiagnosticIds +{ + /// + /// Experiments supported by this repo. + /// + internal static class Experiments + { + // This experiment ID is used for all experimental features in the Microsoft Agent Framework. + internal const string AgentsAIExperiments = "MAAI001"; + + // These diagnostic IDs are defined by the MEAI package for its experimental APIs. + // We use the same IDs so consumers do not need to suppress additional diagnostics + // when using the experimental MEAI APIs. + internal const string AIResponseContinuations = MEAIExperiments; + internal const string AIMcpServers = MEAIExperiments; + internal const string AIFunctionApprovals = MEAIExperiments; + + // These diagnostic IDs are defined by the OpenAI package for its experimental APIs. + // We use the same IDs so consumers do not need to suppress additional diagnostics + // when using the experimental OpenAI APIs. + internal const string AIOpenAIResponses = "OPENAI001"; + internal const string AIOpenAIAssistants = "OPENAI001"; + + private const string MEAIExperiments = "MEAI001"; + } +} diff --git a/dotnet/src/Shared/DiagnosticIds/README.md b/dotnet/src/Shared/DiagnosticIds/README.md new file mode 100644 index 0000000000..6035dd8cc7 --- /dev/null +++ b/dotnet/src/Shared/DiagnosticIds/README.md @@ -0,0 +1,11 @@ +# Diagnostic IDs + +Defines various diagnostic IDs reported by this repo. + +To use this in your project, add the following to your `.csproj` file: + +```xml + + true + +``` \ No newline at end of file diff --git a/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/Microsoft.Agents.AI.CosmosNoSql.UnitTests.csproj b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/Microsoft.Agents.AI.CosmosNoSql.UnitTests.csproj index d60418ee2c..78072b8b6a 100644 --- a/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/Microsoft.Agents.AI.CosmosNoSql.UnitTests.csproj +++ b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/Microsoft.Agents.AI.CosmosNoSql.UnitTests.csproj @@ -2,7 +2,6 @@ net10.0;net9.0 - $(NoWarn);MEAI001 From cd4e36ebf7f685c8d3c1bb1dc1b7a9b617f49f4a Mon Sep 17 00:00:00 2001 From: westey <164392973+westey-m@users.noreply.github.com> Date: Tue, 17 Feb 2026 15:06:43 +0000 Subject: [PATCH 18/22] Replace Typed Base Providers with Composition (#3988) --- .../Program.cs | 20 +- .../Program.cs | 16 +- .../AIContextProvider{TState}.cs | 87 -------- .../ChatHistoryProvider{TState}.cs | 87 -------- .../InMemoryChatHistoryProvider.cs | 22 +- .../ProviderSessionState{TState}.cs | 86 ++++++++ .../CosmosChatHistoryProvider.cs | 19 +- .../Microsoft.Agents.AI.Mem0/Mem0Provider.cs | 18 +- .../WorkflowChatHistoryProvider.cs | 23 +- .../Memory/ChatHistoryMemoryProvider.cs | 17 +- .../Microsoft.Agents.AI/TextSearchProvider.cs | 18 +- .../AIContextProviderTStateTests.cs | 199 ------------------ .../ChatHistoryProviderTStateTests.cs | 195 ----------------- .../ProviderSessionStateTests.cs | 189 +++++++++++++++++ 14 files changed, 381 insertions(+), 615 deletions(-) delete mode 100644 dotnet/src/Microsoft.Agents.AI.Abstractions/AIContextProvider{TState}.cs delete mode 100644 dotnet/src/Microsoft.Agents.AI.Abstractions/ChatHistoryProvider{TState}.cs create mode 100644 dotnet/src/Microsoft.Agents.AI.Abstractions/ProviderSessionState{TState}.cs delete mode 100644 dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/AIContextProviderTStateTests.cs delete mode 100644 dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/ChatHistoryProviderTStateTests.cs create mode 100644 dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/ProviderSessionStateTests.cs diff --git a/dotnet/samples/GettingStarted/AgentWithMemory/AgentWithMemory_Step03_CustomMemory/Program.cs b/dotnet/samples/GettingStarted/AgentWithMemory/AgentWithMemory_Step03_CustomMemory/Program.cs index 9311e56c88..fa6940f5fd 100644 --- a/dotnet/samples/GettingStarted/AgentWithMemory/AgentWithMemory_Step03_CustomMemory/Program.cs +++ b/dotnet/samples/GettingStarted/AgentWithMemory/AgentWithMemory_Step03_CustomMemory/Program.cs @@ -86,25 +86,31 @@ namespace SampleApp /// /// Sample memory component that can remember a user's name and age. /// - internal sealed class UserInfoMemory : AIContextProvider + internal sealed class UserInfoMemory : AIContextProvider { + private readonly ProviderSessionState _sessionState; private readonly IChatClient _chatClient; public UserInfoMemory(IChatClient chatClient, Func? stateInitializer = null) - : base(stateInitializer ?? (_ => new UserInfo()), null, null, null, null) + : base(null, null) { + this._sessionState = new ProviderSessionState( + stateInitializer ?? (_ => new UserInfo()), + this.GetType().Name); this._chatClient = chatClient; } + public override string StateKey => this._sessionState.StateKey; + public UserInfo GetUserInfo(AgentSession session) - => this.GetOrInitializeState(session); + => this._sessionState.GetOrInitializeState(session); public void SetUserInfo(AgentSession session, UserInfo userInfo) - => this.SaveState(session, userInfo); + => this._sessionState.SaveState(session, userInfo); protected override async ValueTask StoreAIContextAsync(InvokedContext context, CancellationToken cancellationToken = default) { - var userInfo = this.GetOrInitializeState(context.Session); + var userInfo = this._sessionState.GetOrInitializeState(context.Session); // Try and extract the user name and age from the message if we don't have it already and it's a user message. if ((userInfo.UserName is null || userInfo.UserAge is null) && context.RequestMessages.Any(x => x.Role == ChatRole.User)) @@ -121,12 +127,12 @@ namespace SampleApp userInfo.UserAge ??= result.Result.UserAge; } - this.SaveState(context.Session, userInfo); + this._sessionState.SaveState(context.Session, userInfo); } protected override ValueTask ProvideAIContextAsync(InvokingContext context, CancellationToken cancellationToken = default) { - var userInfo = this.GetOrInitializeState(context.Session); + var userInfo = this._sessionState.GetOrInitializeState(context.Session); StringBuilder instructions = new(); diff --git a/dotnet/samples/GettingStarted/Agents/Agent_Step07_3rdPartyChatHistoryStorage/Program.cs b/dotnet/samples/GettingStarted/Agents/Agent_Step07_3rdPartyChatHistoryStorage/Program.cs index 1c25c36d22..cbcf14157e 100644 --- a/dotnet/samples/GettingStarted/Agents/Agent_Step07_3rdPartyChatHistoryStorage/Program.cs +++ b/dotnet/samples/GettingStarted/Agents/Agent_Step07_3rdPartyChatHistoryStorage/Program.cs @@ -76,25 +76,31 @@ namespace SampleApp /// State (the session DB key) is stored in the so it roundtrips /// automatically with session serialization. /// - internal sealed class VectorChatHistoryProvider : ChatHistoryProvider + internal sealed class VectorChatHistoryProvider : ChatHistoryProvider { + private readonly ProviderSessionState _sessionState; private readonly VectorStore _vectorStore; public VectorChatHistoryProvider( VectorStore vectorStore, Func? stateInitializer = null, string? stateKey = null) - : base(stateInitializer: stateInitializer ?? (_ => new State(Guid.NewGuid().ToString("N"))), stateKey: stateKey, jsonSerializerOptions: null, provideOutputMessageFilter: null, storeInputMessageFilter: null) + : base(provideOutputMessageFilter: null, storeInputMessageFilter: null) { + this._sessionState = new ProviderSessionState( + stateInitializer ?? (_ => new State(Guid.NewGuid().ToString("N"))), + stateKey ?? this.GetType().Name); this._vectorStore = vectorStore ?? throw new ArgumentNullException(nameof(vectorStore)); } + public override string StateKey => this._sessionState.StateKey; + public string GetSessionDbKey(AgentSession session) - => this.GetOrInitializeState(session).SessionDbKey; + => this._sessionState.GetOrInitializeState(session).SessionDbKey; protected override async ValueTask> ProvideChatHistoryAsync(InvokingContext context, CancellationToken cancellationToken = default) { - var state = this.GetOrInitializeState(context.Session); + var state = this._sessionState.GetOrInitializeState(context.Session); var collection = this._vectorStore.GetCollection("ChatHistory"); await collection.EnsureCollectionExistsAsync(cancellationToken); @@ -112,7 +118,7 @@ namespace SampleApp protected override async ValueTask StoreChatHistoryAsync(InvokedContext context, CancellationToken cancellationToken = default) { - var state = this.GetOrInitializeState(context.Session); + var state = this._sessionState.GetOrInitializeState(context.Session); var collection = this._vectorStore.GetCollection("ChatHistory"); await collection.EnsureCollectionExistsAsync(cancellationToken); diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions/AIContextProvider{TState}.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions/AIContextProvider{TState}.cs deleted file mode 100644 index 136724f44b..0000000000 --- a/dotnet/src/Microsoft.Agents.AI.Abstractions/AIContextProvider{TState}.cs +++ /dev/null @@ -1,87 +0,0 @@ -// Copyright (c) Microsoft. All rights reserved. - -using System; -using System.Collections.Generic; -using System.Text.Json; -using Microsoft.Extensions.AI; - -namespace Microsoft.Agents.AI; - -/// -/// Provides an abstract base class for components that enhance AI context during agent invocations with support for maintaining provider state of type . -/// -/// The type of the state to be maintained by the context provider. Must be a reference type. -/// -/// This class extends by introducing a strongly-typed state management mechanism, allowing derived classes to maintain and persist custom state information across invocations. -/// The state is stored in the session's StateBag using a configurable key and JSON serialization options, enabling seamless integration with the agent session lifecycle. -/// -public abstract class AIContextProvider : AIContextProvider - where TState : class -{ - private readonly Func _stateInitializer; - private readonly string _stateKey; - private readonly JsonSerializerOptions _jsonSerializerOptions; - - /// - /// Initializes a new instance of the class. - /// - /// A function to initialize the state for the context provider. - /// The key used to store the state in the session's StateBag. - /// Options for JSON serialization and deserialization of the state. - /// An optional filter function to apply to input messages before providing context. If not set, defaults to including only messages. - /// An optional filter function to apply to request messages before storing context. If not set, defaults to including only messages. - protected AIContextProvider( - Func stateInitializer, - string? stateKey, - JsonSerializerOptions? jsonSerializerOptions, - Func, IEnumerable>? provideInputMessageFilter, - Func, IEnumerable>? storeInputMessageFilter) - : base(provideInputMessageFilter, storeInputMessageFilter) - { - this._stateInitializer = stateInitializer; - this._jsonSerializerOptions = jsonSerializerOptions ?? AgentAbstractionsJsonUtilities.DefaultOptions; - this._stateKey = stateKey ?? this.GetType().Name; - } - - /// - public override string StateKey => this._stateKey; - - /// - /// Gets the state from the session's StateBag, or initializes it using the state initializer if not present. - /// - /// The agent session containing the StateBag. - /// The provider state. - protected virtual TState GetOrInitializeState(AgentSession? session) - { - if (session?.StateBag.TryGetValue(this._stateKey, out var state, this._jsonSerializerOptions) is true && state is not null) - { - return state; - } - - state = this._stateInitializer(session); - if (session is not null) - { - session.StateBag.SetValue(this._stateKey, state, this._jsonSerializerOptions); - } - - return state; - } - - /// - /// Saves the specified state to the session's StateBag using the configured state key and JSON serializer options. - /// If the session is null, this method does nothing. - /// - /// - /// This method provides a convenient way for derived classes to persist state changes back to the session after processing. - /// It abstracts away the details of how state is stored in the session, allowing derived classes to focus on their specific logic. - /// - /// The agent session containing the StateBag. - /// The state to be saved. - protected virtual void SaveState(AgentSession? session, TState state) - { - if (session is not null) - { - session.StateBag.SetValue(this._stateKey, state, this._jsonSerializerOptions); - } - } -} diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions/ChatHistoryProvider{TState}.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions/ChatHistoryProvider{TState}.cs deleted file mode 100644 index ea1280eb1e..0000000000 --- a/dotnet/src/Microsoft.Agents.AI.Abstractions/ChatHistoryProvider{TState}.cs +++ /dev/null @@ -1,87 +0,0 @@ -// Copyright (c) Microsoft. All rights reserved. - -using System; -using System.Collections.Generic; -using System.Text.Json; -using Microsoft.Extensions.AI; - -namespace Microsoft.Agents.AI; - -/// -/// Provides an abstract base class for fetching chat messages from, and adding chat messages to, chat history for the purposes of agent execution with support for maintaining provider state of type . -/// -/// The type of the state to be maintained by the chat history provider. Must be a reference type. -/// -/// This class extends by introducing a strongly-typed state management mechanism, allowing derived classes to maintain and persist custom state information across invocations. -/// The state is stored in the session's StateBag using a configurable key and JSON serialization options, enabling seamless integration with the agent session lifecycle. -/// -public abstract class ChatHistoryProvider : ChatHistoryProvider - where TState : class -{ - private readonly Func _stateInitializer; - private readonly string _stateKey; - private readonly JsonSerializerOptions _jsonSerializerOptions; - - /// - /// Initializes a new instance of the class. - /// - /// A function to initialize the state for the chat history provider. - /// The key used to store the state in the session's StateBag. - /// Options for JSON serialization and deserialization of the state. - /// A filter function to apply to messages when retrieving them from the chat history. - /// A filter function to apply to messages before storing them in the chat history. - protected ChatHistoryProvider( - Func stateInitializer, - string? stateKey, - JsonSerializerOptions? jsonSerializerOptions, - Func, IEnumerable>? provideOutputMessageFilter, - Func, IEnumerable>? storeInputMessageFilter) - : base(provideOutputMessageFilter, storeInputMessageFilter) - { - this._stateInitializer = stateInitializer; - this._jsonSerializerOptions = jsonSerializerOptions ?? AgentAbstractionsJsonUtilities.DefaultOptions; - this._stateKey = stateKey ?? this.GetType().Name; - } - - /// - public override string StateKey => this._stateKey; - - /// - /// Gets the state from the session's StateBag, or initializes it using the state initializer if not present. - /// - /// The agent session containing the StateBag. - /// The provider state, or null if no session is available. - protected virtual TState GetOrInitializeState(AgentSession? session) - { - if (session?.StateBag.TryGetValue(this._stateKey, out var state, this._jsonSerializerOptions) is true && state is not null) - { - return state; - } - - state = this._stateInitializer(session); - if (session is not null) - { - session.StateBag.SetValue(this._stateKey, state, this._jsonSerializerOptions); - } - - return state; - } - - /// - /// Saves the specified state to the session's StateBag using the configured state key and JSON serializer options. - /// If the session is null, this method does nothing. - /// - /// - /// This method provides a convenient way for derived classes to persist state changes back to the session after processing. - /// It abstracts away the details of how state is stored in the session, allowing derived classes to focus on their specific logic. - /// - /// The agent session containing the StateBag. - /// The state to be saved. - protected virtual void SaveState(AgentSession? session, TState state) - { - if (session is not null) - { - session.StateBag.SetValue(this._stateKey, state, this._jsonSerializerOptions); - } - } -} diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions/InMemoryChatHistoryProvider.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions/InMemoryChatHistoryProvider.cs index 2517586c17..12e935b23e 100644 --- a/dotnet/src/Microsoft.Agents.AI.Abstractions/InMemoryChatHistoryProvider.cs +++ b/dotnet/src/Microsoft.Agents.AI.Abstractions/InMemoryChatHistoryProvider.cs @@ -24,8 +24,10 @@ namespace Microsoft.Agents.AI; /// message reduction strategies or alternative storage implementations. /// /// -public sealed class InMemoryChatHistoryProvider : ChatHistoryProvider +public sealed class InMemoryChatHistoryProvider : ChatHistoryProvider { + private readonly ProviderSessionState _sessionState; + /// /// Initializes a new instance of the class. /// @@ -35,16 +37,20 @@ public sealed class InMemoryChatHistoryProvider : ChatHistoryProvider public InMemoryChatHistoryProvider(InMemoryChatHistoryProviderOptions? options = null) : base( - options?.StateInitializer ?? (_ => new State()), - options?.StateKey, - options?.JsonSerializerOptions, options?.ProvideOutputMessageFilter, options?.StorageInputMessageFilter) { + this._sessionState = new ProviderSessionState( + options?.StateInitializer ?? (_ => new State()), + options?.StateKey ?? this.GetType().Name, + options?.JsonSerializerOptions); this.ChatReducer = options?.ChatReducer; this.ReducerTriggerEvent = options?.ReducerTriggerEvent ?? InMemoryChatHistoryProviderOptions.ChatReducerTriggerEvent.BeforeMessagesRetrieval; } + /// + public override string StateKey => this._sessionState.StateKey; + /// /// Gets the chat reducer used to process or reduce chat messages. If null, no reduction logic will be applied. /// @@ -61,7 +67,7 @@ public sealed class InMemoryChatHistoryProvider : ChatHistoryProviderThe agent session containing the state. /// A list of chat messages, or an empty list if no state is found. public List GetMessages(AgentSession? session) - => this.GetOrInitializeState(session).Messages; + => this._sessionState.GetOrInitializeState(session).Messages; /// /// Sets the chat messages for the specified session. @@ -73,14 +79,14 @@ public sealed class InMemoryChatHistoryProvider : ChatHistoryProvider protected override async ValueTask> ProvideChatHistoryAsync(InvokingContext context, CancellationToken cancellationToken = default) { - var state = this.GetOrInitializeState(context.Session); + var state = this._sessionState.GetOrInitializeState(context.Session); if (this.ReducerTriggerEvent is InMemoryChatHistoryProviderOptions.ChatReducerTriggerEvent.BeforeMessagesRetrieval && this.ChatReducer is not null) { @@ -93,7 +99,7 @@ public sealed class InMemoryChatHistoryProvider : ChatHistoryProvider protected override async ValueTask StoreChatHistoryAsync(InvokedContext context, CancellationToken cancellationToken = default) { - var state = this.GetOrInitializeState(context.Session); + var state = this._sessionState.GetOrInitializeState(context.Session); // Add request and response messages to the provider var allNewMessages = context.RequestMessages.Concat(context.ResponseMessages ?? []); diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions/ProviderSessionState{TState}.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions/ProviderSessionState{TState}.cs new file mode 100644 index 0000000000..b88e996644 --- /dev/null +++ b/dotnet/src/Microsoft.Agents.AI.Abstractions/ProviderSessionState{TState}.cs @@ -0,0 +1,86 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Text.Json; + +namespace Microsoft.Agents.AI; + +/// +/// Provides strongly-typed state management for providers, enabling reading and writing of provider-specific state +/// to and from an 's . +/// +/// The type of the state to be maintained. Must be a reference type. +/// +/// +/// This class encapsulates the logic for initializing, retrieving, and persisting provider state in the session's StateBag +/// using a configurable key and JSON serialization options. It is intended to be used as a composed field within provider +/// implementations (e.g., or subclasses) to avoid +/// duplicating state management logic across provider type hierarchies. +/// +/// +/// State is stored in the using the property as the key, +/// enabling multiple providers to maintain independent state within the same session. +/// +/// +public class ProviderSessionState + where TState : class +{ + private readonly Func _stateInitializer; + private readonly JsonSerializerOptions _jsonSerializerOptions; + + /// + /// Initializes a new instance of the class. + /// + /// A function to initialize the state when it is not yet present in the session's StateBag. + /// The key used to store the state in the session's StateBag. + /// Options for JSON serialization and deserialization of the state. + public ProviderSessionState( + Func stateInitializer, + string stateKey, + JsonSerializerOptions? jsonSerializerOptions = null) + { + this._stateInitializer = stateInitializer; + this.StateKey = stateKey; + this._jsonSerializerOptions = jsonSerializerOptions ?? AgentAbstractionsJsonUtilities.DefaultOptions; + } + + /// + /// Gets the key used to store the provider state in the . + /// + public string StateKey { get; } + + /// + /// Gets the state from the session's StateBag, or initializes it using the state initializer if not present. + /// + /// The agent session containing the StateBag. + /// The provider state. + public TState GetOrInitializeState(AgentSession? session) + { + if (session?.StateBag.TryGetValue(this.StateKey, out var state, this._jsonSerializerOptions) is true && state is not null) + { + return state; + } + + state = this._stateInitializer(session); + if (session is not null) + { + session.StateBag.SetValue(this.StateKey, state, this._jsonSerializerOptions); + } + + return state; + } + + /// + /// Saves the specified state to the session's StateBag using the configured state key and JSON serializer options. + /// If the session is null, this method does nothing. + /// + /// The agent session containing the StateBag. + /// The state to be saved. + public void SaveState(AgentSession? session, TState state) + { + if (session is not null) + { + session.StateBag.SetValue(this.StateKey, state, this._jsonSerializerOptions); + } + } +} diff --git a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatHistoryProvider.cs b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatHistoryProvider.cs index b12a465905..c4e13fae29 100644 --- a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatHistoryProvider.cs +++ b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatHistoryProvider.cs @@ -19,8 +19,9 @@ namespace Microsoft.Agents.AI; /// [RequiresUnreferencedCode("The CosmosChatHistoryProvider uses JSON serialization which is incompatible with trimming.")] [RequiresDynamicCode("The CosmosChatHistoryProvider uses JSON serialization which is incompatible with NativeAOT.")] -public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable +public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable { + private readonly ProviderSessionState _sessionState; private readonly CosmosClient _cosmosClient; private readonly Container _container; private readonly bool _ownsClient; @@ -98,8 +99,11 @@ public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IEnumerable>? provideOutputMessageFilter = null, Func, IEnumerable>? storeInputMessageFilter = null) - : base(Throw.IfNull(stateInitializer), stateKey, null, provideOutputMessageFilter, storeInputMessageFilter) + : base(provideOutputMessageFilter, storeInputMessageFilter) { + this._sessionState = new ProviderSessionState( + Throw.IfNull(stateInitializer), + stateKey ?? this.GetType().Name); this._cosmosClient = Throw.IfNull(cosmosClient); this.DatabaseId = Throw.IfNullOrWhitespace(databaseId); this.ContainerId = Throw.IfNullOrWhitespace(containerId); @@ -107,6 +111,9 @@ public sealed class CosmosChatHistoryProvider : ChatHistoryProvider + public override string StateKey => this._sessionState.StateKey; + /// /// Initializes a new instance of the class using a connection string. /// @@ -190,7 +197,7 @@ public sealed class CosmosChatHistoryProvider : ChatHistoryProvider -public sealed class Mem0Provider : AIContextProvider +public sealed class Mem0Provider : AIContextProvider { private const string DefaultContextPrompt = "## Memories\nConsider the following memories when answering user questions:"; + private readonly ProviderSessionState _sessionState; private readonly string _contextPrompt; private readonly bool _enableSensitiveTelemetryData; @@ -51,8 +52,12 @@ public sealed class Mem0Provider : AIContextProvider /// /// public Mem0Provider(HttpClient httpClient, Func stateInitializer, Mem0ProviderOptions? options = null, ILoggerFactory? loggerFactory = null) - : base(ValidateStateInitializer(Throw.IfNull(stateInitializer)), options?.StateKey, Mem0JsonUtilities.DefaultOptions, options?.SearchInputMessageFilter, options?.StorageInputMessageFilter) + : base(options?.SearchInputMessageFilter, options?.StorageInputMessageFilter) { + this._sessionState = new ProviderSessionState( + ValidateStateInitializer(Throw.IfNull(stateInitializer)), + options?.StateKey ?? this.GetType().Name, + Mem0JsonUtilities.DefaultOptions); Throw.IfNull(httpClient); if (string.IsNullOrWhiteSpace(httpClient.BaseAddress?.AbsoluteUri)) { @@ -66,6 +71,9 @@ public sealed class Mem0Provider : AIContextProvider this._enableSensitiveTelemetryData = options?.EnableSensitiveTelemetryData ?? false; } + /// + public override string StateKey => this._sessionState.StateKey; + private static Func ValidateStateInitializer(Func stateInitializer) => session => { @@ -88,7 +96,7 @@ public sealed class Mem0Provider : AIContextProvider { Throw.IfNull(context); - var state = this.GetOrInitializeState(context.Session); + var state = this._sessionState.GetOrInitializeState(context.Session); var searchScope = state.SearchScope; string queryText = string.Join( @@ -165,7 +173,7 @@ public sealed class Mem0Provider : AIContextProvider /// protected override async ValueTask StoreAIContextAsync(InvokedContext context, CancellationToken cancellationToken = default) { - var state = this.GetOrInitializeState(context.Session); + var state = this._sessionState.GetOrInitializeState(context.Session); var storageScope = state.StorageScope; try @@ -200,7 +208,7 @@ public sealed class Mem0Provider : AIContextProvider public Task ClearStoredMemoriesAsync(AgentSession session, CancellationToken cancellationToken = default) { Throw.IfNull(session); - var state = this.GetOrInitializeState(session); + var state = this._sessionState.GetOrInitializeState(session); var storageScope = state.StorageScope; return this._client.ClearMemoryAsync( diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowChatHistoryProvider.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowChatHistoryProvider.cs index 74e2fae3fa..b9d5f3ae49 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowChatHistoryProvider.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowChatHistoryProvider.cs @@ -9,8 +9,10 @@ using Microsoft.Extensions.AI; namespace Microsoft.Agents.AI.Workflows; -internal sealed class WorkflowChatHistoryProvider : ChatHistoryProvider +internal sealed class WorkflowChatHistoryProvider : ChatHistoryProvider { + private readonly ProviderSessionState _sessionState; + /// /// Initializes a new instance of the class. /// @@ -20,10 +22,17 @@ internal sealed class WorkflowChatHistoryProvider : ChatHistoryProvider public WorkflowChatHistoryProvider(JsonSerializerOptions? jsonSerializerOptions = null) - : base(stateInitializer: _ => new StoreState(), stateKey: null, jsonSerializerOptions: jsonSerializerOptions, provideOutputMessageFilter: null, storeInputMessageFilter: null) + : base(provideOutputMessageFilter: null, storeInputMessageFilter: null) { + this._sessionState = new ProviderSessionState( + _ => new StoreState(), + this.GetType().Name, + jsonSerializerOptions); } + /// + public override string StateKey => this._sessionState.StateKey; + internal sealed class StoreState { public int Bookmark { get; set; } @@ -31,21 +40,21 @@ internal sealed class WorkflowChatHistoryProvider : ChatHistoryProvider messages) - => this.GetOrInitializeState(session).Messages.AddRange(messages); + => this._sessionState.GetOrInitializeState(session).Messages.AddRange(messages); protected override ValueTask> ProvideChatHistoryAsync(InvokingContext context, CancellationToken cancellationToken = default) - => new(this.GetOrInitializeState(context.Session).Messages); + => new(this._sessionState.GetOrInitializeState(context.Session).Messages); protected override ValueTask StoreChatHistoryAsync(InvokedContext context, CancellationToken cancellationToken = default) { var allNewMessages = context.RequestMessages.Concat(context.ResponseMessages ?? []); - this.GetOrInitializeState(context.Session).Messages.AddRange(allNewMessages); + this._sessionState.GetOrInitializeState(context.Session).Messages.AddRange(allNewMessages); return default; } public IEnumerable GetFromBookmark(AgentSession session) { - var state = this.GetOrInitializeState(session); + var state = this._sessionState.GetOrInitializeState(session); for (int i = state.Bookmark; i < state.Messages.Count; i++) { @@ -55,7 +64,7 @@ internal sealed class WorkflowChatHistoryProvider : ChatHistoryProvider /// -public sealed class ChatHistoryMemoryProvider : AIContextProvider, IDisposable +public sealed class ChatHistoryMemoryProvider : AIContextProvider, IDisposable { private const string DefaultContextPrompt = "## Memories\nConsider the following memories when answering user questions:"; private const int DefaultMaxResults = 3; private const string DefaultFunctionToolName = "Search"; private const string DefaultFunctionToolDescription = "Allows searching for related previous chat history to help answer the user question."; + private readonly ProviderSessionState _sessionState; + #pragma warning disable CA2213 // VectorStore is not owned by this class - caller is responsible for disposal private readonly VectorStore _vectorStore; #pragma warning restore CA2213 @@ -74,8 +76,12 @@ public sealed class ChatHistoryMemoryProvider : AIContextProvider stateInitializer, ChatHistoryMemoryProviderOptions? options = null, ILoggerFactory? loggerFactory = null) - : base(Throw.IfNull(stateInitializer), options?.StateKey, AgentJsonUtilities.DefaultOptions, options?.SearchInputMessageFilter, options?.StorageInputMessageFilter) + : base(options?.SearchInputMessageFilter, options?.StorageInputMessageFilter) { + this._sessionState = new ProviderSessionState( + Throw.IfNull(stateInitializer), + options?.StateKey ?? this.GetType().Name, + AgentJsonUtilities.DefaultOptions); this._vectorStore = Throw.IfNull(vectorStore); options ??= new ChatHistoryMemoryProviderOptions(); @@ -109,12 +115,15 @@ public sealed class ChatHistoryMemoryProvider : AIContextProvider + public override string StateKey => this._sessionState.StateKey; + /// protected override async ValueTask ProvideAIContextAsync(InvokingContext context, CancellationToken cancellationToken = default) { _ = Throw.IfNull(context); - var state = this.GetOrInitializeState(context.Session); + var state = this._sessionState.GetOrInitializeState(context.Session); var searchScope = state.SearchScope; if (this._searchTime == ChatHistoryMemoryProviderOptions.SearchBehavior.OnDemandFunctionCalling) @@ -186,7 +195,7 @@ public sealed class ChatHistoryMemoryProvider : AIContextProvider /// -public sealed class TextSearchProvider : AIContextProvider +public sealed class TextSearchProvider : AIContextProvider { private const string DefaultPluginSearchFunctionName = "Search"; private const string DefaultPluginSearchFunctionDescription = "Allows searching for additional information to help answer the user question."; private const string DefaultContextPrompt = "## Additional Context\nConsider the following information from source documents when responding to the user:"; private const string DefaultCitationsPrompt = "Include citations to the source document with document name and link if document name and link is available."; + private readonly ProviderSessionState _sessionState; private readonly Func>> _searchAsync; private readonly ILogger? _logger; private readonly AITool[] _tools; @@ -60,8 +61,12 @@ public sealed class TextSearchProvider : AIContextProvider>> searchAsync, TextSearchProviderOptions? options = null, ILoggerFactory? loggerFactory = null) - : base(_ => new TextSearchProviderState(), options?.StateKey, AgentJsonUtilities.DefaultOptions, options?.SearchInputMessageFilter, options?.StorageInputMessageFilter) + : base(options?.SearchInputMessageFilter, options?.StorageInputMessageFilter) { + this._sessionState = new ProviderSessionState( + _ => new TextSearchProviderState(), + options?.StateKey ?? this.GetType().Name, + AgentJsonUtilities.DefaultOptions); // Validate and assign parameters this._searchAsync = Throw.IfNull(searchAsync); this._logger = loggerFactory?.CreateLogger(); @@ -82,6 +87,9 @@ public sealed class TextSearchProvider : AIContextProvider + public override string StateKey => this._sessionState.StateKey; + /// protected override async ValueTask ProvideAIContextAsync(InvokingContext context, CancellationToken cancellationToken = default) { @@ -95,7 +103,7 @@ public sealed class TextSearchProvider : AIContextProvider -/// Contains tests for the class. -/// -public class AIContextProviderTStateTests -{ - private static readonly AIAgent s_mockAgent = new Mock().Object; - - #region GetOrInitializeState Tests - - [Fact] - public void GetOrInitializeState_InitializesFromStateInitializerOnFirstCall() - { - // Arrange - var expectedState = new TestState { Value = "initialized" }; - var provider = new TestAIContextProvider(_ => expectedState); - var session = new TestAgentSession(); - - // Act - var state = provider.GetState(session); - - // Assert - Assert.Same(expectedState, state); - } - - [Fact] - public void GetOrInitializeState_ReturnsCachedStateFromStateBagOnSecondCall() - { - // Arrange - var callCount = 0; - var provider = new TestAIContextProvider(_ => - { - callCount++; - return new TestState { Value = $"init-{callCount}" }; - }); - var session = new TestAgentSession(); - - // Act - var state1 = provider.GetState(session); - var state2 = provider.GetState(session); - - // Assert - initializer called only once; second call reads from StateBag - Assert.Equal(1, callCount); - Assert.Equal("init-1", state1.Value); - Assert.Equal("init-1", state2.Value); - } - - [Fact] - public void GetOrInitializeState_WorksWhenSessionIsNull() - { - // Arrange - var provider = new TestAIContextProvider(_ => new TestState { Value = "no-session" }); - - // Act - var state = provider.GetState(null); - - // Assert - Assert.Equal("no-session", state.Value); - } - - [Fact] - public void GetOrInitializeState_ReInitializesWhenSessionIsNull() - { - // Arrange - without a session, state can't be cached in StateBag - var callCount = 0; - var provider = new TestAIContextProvider(_ => - { - callCount++; - return new TestState { Value = $"init-{callCount}" }; - }); - - // Act - provider.GetState(null); - provider.GetState(null); - - // Assert - initializer called each time since there's no session to cache in - Assert.Equal(2, callCount); - } - - #endregion - - #region SaveState Tests - - [Fact] - public void SaveState_SavesToStateBag() - { - // Arrange - var provider = new TestAIContextProvider(_ => new TestState()); - var session = new TestAgentSession(); - var state = new TestState { Value = "saved" }; - - // Act - provider.DoSaveState(session, state); - var retrieved = provider.GetState(session); - - // Assert - Assert.Equal("saved", retrieved.Value); - } - - [Fact] - public void SaveState_NoOpWhenSessionIsNull() - { - // Arrange - var provider = new TestAIContextProvider(_ => new TestState { Value = "default" }); - - // Act - should not throw - provider.DoSaveState(null, new TestState { Value = "saved" }); - - // Assert - no exception; can't verify further without a session - } - - #endregion - - #region StateKey Tests - - [Fact] - public void StateKey_DefaultsToTypeName() - { - // Arrange - var provider = new TestAIContextProvider(_ => new TestState()); - - // Act & Assert - Assert.Equal(nameof(TestAIContextProvider), provider.StateKey); - } - - [Fact] - public void StateKey_UsesCustomKeyWhenProvided() - { - // Arrange - var provider = new TestAIContextProvider(_ => new TestState(), stateKey: "custom-key"); - - // Act & Assert - Assert.Equal("custom-key", provider.StateKey); - } - - #endregion - - #region Integration Tests - - [Fact] - public async Task InvokingCoreAsync_CanUseStateInProvideAIContextAsync() - { - // Arrange - var provider = new TestAIContextProvider(_ => new TestState { Value = "state-value" }); - var session = new TestAgentSession(); - var inputContext = new AIContext { Messages = [new ChatMessage(ChatRole.User, "Hi")] }; - var context = new AIContextProvider.InvokingContext(s_mockAgent, session, inputContext); - - // Act - var result = await provider.InvokingAsync(context); - - // Assert - the provider uses state to produce context messages - var messages = result.Messages!.ToList(); - Assert.Equal(2, messages.Count); - Assert.Contains("state-value", messages[1].Text); - } - - #endregion - - public sealed class TestState - { - public string Value { get; set; } = string.Empty; - } - - private sealed class TestAIContextProvider : AIContextProvider - { - public TestAIContextProvider( - Func stateInitializer, - string? stateKey = null) - : base(stateInitializer, stateKey, null, null, null) - { - } - - public TestState GetState(AgentSession? session) => this.GetOrInitializeState(session); - - public void DoSaveState(AgentSession? session, TestState state) => this.SaveState(session, state); - - protected override ValueTask ProvideAIContextAsync(InvokingContext context, CancellationToken cancellationToken = default) - { - var state = this.GetOrInitializeState(context.Session); - return new(new AIContext - { - Messages = [new ChatMessage(ChatRole.System, $"Context from state: {state.Value}")] - }); - } - } - - private sealed class TestAgentSession : AgentSession; -} diff --git a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/ChatHistoryProviderTStateTests.cs b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/ChatHistoryProviderTStateTests.cs deleted file mode 100644 index e3c0507e8d..0000000000 --- a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/ChatHistoryProviderTStateTests.cs +++ /dev/null @@ -1,195 +0,0 @@ -// Copyright (c) Microsoft. All rights reserved. - -using System; -using System.Collections.Generic; -using System.Linq; -using System.Threading; -using System.Threading.Tasks; -using Microsoft.Extensions.AI; -using Moq; - -namespace Microsoft.Agents.AI.Abstractions.UnitTests; - -/// -/// Contains tests for the class. -/// -public class ChatHistoryProviderTStateTests -{ - private static readonly AIAgent s_mockAgent = new Mock().Object; - - #region GetOrInitializeState Tests - - [Fact] - public void GetOrInitializeState_InitializesFromStateInitializerOnFirstCall() - { - // Arrange - var expectedState = new TestState { Value = "initialized" }; - var provider = new TestChatHistoryProvider(_ => expectedState); - var session = new TestAgentSession(); - - // Act - var state = provider.GetState(session); - - // Assert - Assert.Same(expectedState, state); - } - - [Fact] - public void GetOrInitializeState_ReturnsCachedStateFromStateBagOnSecondCall() - { - // Arrange - var callCount = 0; - var provider = new TestChatHistoryProvider(_ => - { - callCount++; - return new TestState { Value = $"init-{callCount}" }; - }); - var session = new TestAgentSession(); - - // Act - var state1 = provider.GetState(session); - var state2 = provider.GetState(session); - - // Assert - initializer called only once; second call reads from StateBag - Assert.Equal(1, callCount); - Assert.Equal("init-1", state1.Value); - Assert.Equal("init-1", state2.Value); - } - - [Fact] - public void GetOrInitializeState_WorksWhenSessionIsNull() - { - // Arrange - var provider = new TestChatHistoryProvider(_ => new TestState { Value = "no-session" }); - - // Act - var state = provider.GetState(null); - - // Assert - Assert.Equal("no-session", state.Value); - } - - [Fact] - public void GetOrInitializeState_ReInitializesWhenSessionIsNull() - { - // Arrange - without a session, state can't be cached in StateBag - var callCount = 0; - var provider = new TestChatHistoryProvider(_ => - { - callCount++; - return new TestState { Value = $"init-{callCount}" }; - }); - - // Act - _ = provider.GetState(null); - provider.GetState(null); - - // Assert - initializer called each time since there's no session to cache in - Assert.Equal(2, callCount); - } - - #endregion - - #region SaveState Tests - - [Fact] - public void SaveState_SavesToStateBag() - { - // Arrange - var provider = new TestChatHistoryProvider(_ => new TestState()); - var session = new TestAgentSession(); - var state = new TestState { Value = "saved" }; - - // Act - provider.DoSaveState(session, state); - var retrieved = provider.GetState(session); - - // Assert - Assert.Equal("saved", retrieved.Value); - } - - [Fact] - public void SaveState_NoOpWhenSessionIsNull() - { - // Arrange - var provider = new TestChatHistoryProvider(_ => new TestState { Value = "default" }); - - // Act - should not throw - provider.DoSaveState(null, new TestState { Value = "saved" }); - - // Assert - no exception; can't verify further without a session - } - - #endregion - - #region StateKey Tests - - [Fact] - public void StateKey_DefaultsToTypeName() - { - // Arrange - var provider = new TestChatHistoryProvider(_ => new TestState()); - - // Act & Assert - Assert.Equal(nameof(TestChatHistoryProvider), provider.StateKey); - } - - [Fact] - public void StateKey_UsesCustomKeyWhenProvided() - { - // Arrange - var provider = new TestChatHistoryProvider(_ => new TestState(), stateKey: "custom-key"); - - // Act & Assert - Assert.Equal("custom-key", provider.StateKey); - } - - #endregion - - #region Integration Tests - - [Fact] - public async Task InvokingCoreAsync_CanUseStateInProvideChatHistoryAsync() - { - // Arrange - var provider = new TestChatHistoryProvider(_ => new TestState { Value = "state-value" }); - var session = new TestAgentSession(); - var context = new ChatHistoryProvider.InvokingContext(s_mockAgent, session, [new ChatMessage(ChatRole.User, "Hi")]); - - // Act - var result = (await provider.InvokingAsync(context)).ToList(); - - // Assert - the provider uses state to produce history messages - Assert.Equal(2, result.Count); - Assert.Contains("state-value", result[0].Text); - } - - #endregion - - public sealed class TestState - { - public string Value { get; set; } = string.Empty; - } - - private sealed class TestChatHistoryProvider : ChatHistoryProvider - { - public TestChatHistoryProvider( - Func stateInitializer, - string? stateKey = null) - : base(stateInitializer, stateKey, null, null, null) - { - } - - public TestState GetState(AgentSession? session) => this.GetOrInitializeState(session); - - public void DoSaveState(AgentSession? session, TestState state) => this.SaveState(session, state); - - protected override ValueTask> ProvideChatHistoryAsync(InvokingContext context, CancellationToken cancellationToken = default) - { - var state = this.GetOrInitializeState(context.Session); - return new(new[] { new ChatMessage(ChatRole.System, $"History from state: {state.Value}") }); - } - } - - private sealed class TestAgentSession : AgentSession; -} diff --git a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/ProviderSessionStateTests.cs b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/ProviderSessionStateTests.cs new file mode 100644 index 0000000000..da3992bf7f --- /dev/null +++ b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/ProviderSessionStateTests.cs @@ -0,0 +1,189 @@ +// Copyright (c) Microsoft. All rights reserved. + +namespace Microsoft.Agents.AI.Abstractions.UnitTests; + +/// +/// Contains tests for the class. +/// +public class ProviderSessionStateTests +{ + #region GetOrInitializeState Tests + + [Fact] + public void GetOrInitializeState_InitializesFromStateInitializerOnFirstCall() + { + // Arrange + var expectedState = new TestState { Value = "initialized" }; + var sessionState = new ProviderSessionState(_ => expectedState, "test-key"); + var session = new TestAgentSession(); + + // Act + var state = sessionState.GetOrInitializeState(session); + + // Assert + Assert.Same(expectedState, state); + } + + [Fact] + public void GetOrInitializeState_ReturnsCachedStateFromStateBagOnSecondCall() + { + // Arrange + var callCount = 0; + var sessionState = new ProviderSessionState(_ => + { + callCount++; + return new TestState { Value = $"init-{callCount}" }; + }, "test-key"); + var session = new TestAgentSession(); + + // Act + var state1 = sessionState.GetOrInitializeState(session); + var state2 = sessionState.GetOrInitializeState(session); + + // Assert - initializer called only once; second call reads from StateBag + Assert.Equal(1, callCount); + Assert.Equal("init-1", state1.Value); + Assert.Equal("init-1", state2.Value); + } + + [Fact] + public void GetOrInitializeState_WorksWhenSessionIsNull() + { + // Arrange + var sessionState = new ProviderSessionState(_ => new TestState { Value = "no-session" }, "test-key"); + + // Act + var state = sessionState.GetOrInitializeState(null); + + // Assert + Assert.Equal("no-session", state.Value); + } + + [Fact] + public void GetOrInitializeState_ReInitializesWhenSessionIsNull() + { + // Arrange - without a session, state can't be cached in StateBag + var callCount = 0; + var sessionState = new ProviderSessionState(_ => + { + callCount++; + return new TestState { Value = $"init-{callCount}" }; + }, "test-key"); + + // Act + sessionState.GetOrInitializeState(null); + sessionState.GetOrInitializeState(null); + + // Assert - initializer called each time since there's no session to cache in + Assert.Equal(2, callCount); + } + + #endregion + + #region SaveState Tests + + [Fact] + public void SaveState_SavesToStateBag() + { + // Arrange + var sessionState = new ProviderSessionState(_ => new TestState(), "test-key"); + var session = new TestAgentSession(); + var state = new TestState { Value = "saved" }; + + // Act + sessionState.SaveState(session, state); + var retrieved = sessionState.GetOrInitializeState(session); + + // Assert + Assert.Equal("saved", retrieved.Value); + } + + [Fact] + public void SaveState_NoOpWhenSessionIsNull() + { + // Arrange + var sessionState = new ProviderSessionState(_ => new TestState { Value = "default" }, "test-key"); + + // Act - should not throw + sessionState.SaveState(null, new TestState { Value = "saved" }); + + // Assert - no exception; can't verify further without a session + } + + #endregion + + #region StateKey Tests + + [Fact] + public void StateKey_UsesProvidedKey() + { + // Arrange + var sessionState = new ProviderSessionState(_ => new TestState(), "my-provider-key"); + + // Act & Assert + Assert.Equal("my-provider-key", sessionState.StateKey); + } + + [Fact] + public void StateKey_UsesCustomKeyWhenProvided() + { + // Arrange + var sessionState = new ProviderSessionState(_ => new TestState(), "custom-key"); + + // Act & Assert + Assert.Equal("custom-key", sessionState.StateKey); + } + + #endregion + + #region Isolation Tests + + [Fact] + public void GetOrInitializeState_IsolatesStateBetweenDifferentKeys() + { + // Arrange + var sessionState1 = new ProviderSessionState(_ => new TestState { Value = "state-1" }, "key-1"); + var sessionState2 = new ProviderSessionState(_ => new TestState { Value = "state-2" }, "key-2"); + var session = new TestAgentSession(); + + // Act + var state1 = sessionState1.GetOrInitializeState(session); + var state2 = sessionState2.GetOrInitializeState(session); + + // Assert - each key maintains independent state + Assert.Equal("state-1", state1.Value); + Assert.Equal("state-2", state2.Value); + } + + [Fact] + public void GetOrInitializeState_IsolatesStateBetweenDifferentSessions() + { + // Arrange + var callCount = 0; + var sessionState = new ProviderSessionState(_ => + { + callCount++; + return new TestState { Value = $"init-{callCount}" }; + }, "test-key"); + var session1 = new TestAgentSession(); + var session2 = new TestAgentSession(); + + // Act + var state1 = sessionState.GetOrInitializeState(session1); + var state2 = sessionState.GetOrInitializeState(session2); + + // Assert - each session gets its own state + Assert.Equal(2, callCount); + Assert.Equal("init-1", state1.Value); + Assert.Equal("init-2", state2.Value); + } + + #endregion + + public sealed class TestState + { + public string Value { get; set; } = string.Empty; + } + + private sealed class TestAgentSession : AgentSession; +} From dc6d0bc58b5bc7b4ab455ab21b938232406f9e28 Mon Sep 17 00:00:00 2001 From: westey <164392973+westey-m@users.noreply.github.com> Date: Tue, 17 Feb 2026 15:10:47 +0000 Subject: [PATCH 19/22] Fix sample resource path resolution (#3952) --- .../Workflows/ConditionalEdges/01_EdgeCondition/Resources.cs | 2 +- .../Workflows/ConditionalEdges/02_SwitchCase/Resources.cs | 2 +- .../Workflows/ConditionalEdges/03_MultiSelection/Resources.cs | 2 +- .../samples/GettingStarted/Workflows/SharedStates/Resources.cs | 2 +- 4 files changed, 4 insertions(+), 4 deletions(-) diff --git a/dotnet/samples/GettingStarted/Workflows/ConditionalEdges/01_EdgeCondition/Resources.cs b/dotnet/samples/GettingStarted/Workflows/ConditionalEdges/01_EdgeCondition/Resources.cs index 7a0d0ea2bd..4ac35cfbec 100644 --- a/dotnet/samples/GettingStarted/Workflows/ConditionalEdges/01_EdgeCondition/Resources.cs +++ b/dotnet/samples/GettingStarted/Workflows/ConditionalEdges/01_EdgeCondition/Resources.cs @@ -9,5 +9,5 @@ internal static class Resources { private const string ResourceFolder = "Resources"; - public static string Read(string fileName) => File.ReadAllText($"{ResourceFolder}/{fileName}"); + public static string Read(string fileName) => File.ReadAllText(Path.Combine(AppContext.BaseDirectory, ResourceFolder, fileName)); } diff --git a/dotnet/samples/GettingStarted/Workflows/ConditionalEdges/02_SwitchCase/Resources.cs b/dotnet/samples/GettingStarted/Workflows/ConditionalEdges/02_SwitchCase/Resources.cs index 415a3820a1..236f3a425a 100644 --- a/dotnet/samples/GettingStarted/Workflows/ConditionalEdges/02_SwitchCase/Resources.cs +++ b/dotnet/samples/GettingStarted/Workflows/ConditionalEdges/02_SwitchCase/Resources.cs @@ -9,5 +9,5 @@ internal static class Resources { private const string ResourceFolder = "Resources"; - public static string Read(string fileName) => File.ReadAllText($"{ResourceFolder}/{fileName}"); + public static string Read(string fileName) => File.ReadAllText(Path.Combine(AppContext.BaseDirectory, ResourceFolder, fileName)); } diff --git a/dotnet/samples/GettingStarted/Workflows/ConditionalEdges/03_MultiSelection/Resources.cs b/dotnet/samples/GettingStarted/Workflows/ConditionalEdges/03_MultiSelection/Resources.cs index d04a7c8a5f..d1494b7109 100644 --- a/dotnet/samples/GettingStarted/Workflows/ConditionalEdges/03_MultiSelection/Resources.cs +++ b/dotnet/samples/GettingStarted/Workflows/ConditionalEdges/03_MultiSelection/Resources.cs @@ -9,5 +9,5 @@ internal static class Resources { private const string ResourceFolder = "Resources"; - public static string Read(string fileName) => File.ReadAllText($"{ResourceFolder}/{fileName}"); + public static string Read(string fileName) => File.ReadAllText(Path.Combine(AppContext.BaseDirectory, ResourceFolder, fileName)); } diff --git a/dotnet/samples/GettingStarted/Workflows/SharedStates/Resources.cs b/dotnet/samples/GettingStarted/Workflows/SharedStates/Resources.cs index a831387050..4bdca21dda 100644 --- a/dotnet/samples/GettingStarted/Workflows/SharedStates/Resources.cs +++ b/dotnet/samples/GettingStarted/Workflows/SharedStates/Resources.cs @@ -9,5 +9,5 @@ internal static class Resources { private const string ResourceFolder = "Resources"; - public static string Read(string fileName) => File.ReadAllText($"{ResourceFolder}/{fileName}"); + public static string Read(string fileName) => File.ReadAllText(Path.Combine(AppContext.BaseDirectory, ResourceFolder, fileName)); } From bf7056a131a477adaa1c30ee4d9554e90d9f01d6 Mon Sep 17 00:00:00 2001 From: westey <164392973+westey-m@users.noreply.github.com> Date: Tue, 17 Feb 2026 15:10:53 +0000 Subject: [PATCH 20/22] Add bugbash fixes. (#3961) --- .../Program.cs | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/dotnet/samples/GettingStarted/AgentWithRAG/AgentWithRAG_Step02_CustomVectorStoreRAG/Program.cs b/dotnet/samples/GettingStarted/AgentWithRAG/AgentWithRAG_Step02_CustomVectorStoreRAG/Program.cs index f06fda6a5b..0c299a1445 100644 --- a/dotnet/samples/GettingStarted/AgentWithRAG/AgentWithRAG_Step02_CustomVectorStoreRAG/Program.cs +++ b/dotnet/samples/GettingStarted/AgentWithRAG/AgentWithRAG_Step02_CustomVectorStoreRAG/Program.cs @@ -62,7 +62,7 @@ TextSearchProviderOptions textSearchOptions = new() { // Run the search prior to every model invocation. SearchTime = TextSearchProviderOptions.TextSearchBehavior.BeforeAIInvoke, - // Use up to 4 recent messages when searching so that searches + // Use up to 5 recent messages when searching so that searches // still produce valuable results even when the user is referring // back to previous messages in their request. RecentMessageMemoryLimit = 5 @@ -74,7 +74,14 @@ AIAgent agent = azureOpenAIClient .AsAIAgent(new ChatClientAgentOptions { ChatOptions = new() { Instructions = "You are a helpful support specialist for the Microsoft Agent Framework. Answer questions using the provided context and cite the source document when available. Keep responses brief." }, - AIContextProviders = [new TextSearchProvider(SearchAdapter, textSearchOptions)] + AIContextProviders = [new TextSearchProvider(SearchAdapter, textSearchOptions)], + // Configure a filter on the InMemoryChatHistoryProvider so that we don't persist the messages produced by the TextSearchProvider in chat history. + // The default is to persist all messages except those that came from chat history in the first place. + // You may choose to persist the TextSearchProvider messages, if you want the search output to be provided to the model in future interactions as well. + ChatHistoryProvider = new InMemoryChatHistoryProvider(new InMemoryChatHistoryProviderOptions() + { + StorageInputMessageFilter = msgs => msgs.Where(m => m.GetAgentRequestMessageSourceType() != AgentRequestMessageSourceType.ChatHistory && m.GetAgentRequestMessageSourceType() != AgentRequestMessageSourceType.AIContextProvider) + }) }); AgentSession session = await agent.CreateSessionAsync(); From 36d52a1f9ff0766a01cf033eb34f4f30bc7a8962 Mon Sep 17 00:00:00 2001 From: Pete Roden Date: Tue, 17 Feb 2026 10:22:04 -0500 Subject: [PATCH 21/22] =?UTF-8?q?Python:=20feature:=20Inject=20OpenTelemet?= =?UTF-8?q?ry=20trace=20context=20into=20MCP=20requests=20and=20update?= =?UTF-8?q?=E2=80=A6=20(#3780)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat: Inject OpenTelemetry trace context into MCP requests and update documentation * Update python/samples/getting_started/observability/README.md Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * Update python/packages/core/tests/core/test_mcp.py Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * refactor: move opentelemetry import to module level OpenTelemetry is a hard dependency of agent-framework-core (per pyproject.toml), so the try/except ImportError guard was dead code. Move the import to the top of the file to fail fast on missing dependencies instead of silently hiding installation issues. --------- Co-authored-by: Pete Roden Co-authored-by: Mark Wallace <127216156+markwallace-microsoft@users.noreply.github.com> Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- python/packages/core/agent_framework/_mcp.py | 22 +++++- python/packages/core/tests/core/test_mcp.py | 67 +++++++++++++++++++ .../samples/02-agents/observability/README.md | 4 ++ 3 files changed, 92 insertions(+), 1 deletion(-) diff --git a/python/packages/core/agent_framework/_mcp.py b/python/packages/core/agent_framework/_mcp.py index 2be5e452ce..0c241cb89a 100644 --- a/python/packages/core/agent_framework/_mcp.py +++ b/python/packages/core/agent_framework/_mcp.py @@ -24,6 +24,7 @@ from mcp.client.websocket import websocket_client from mcp.shared.context import RequestContext from mcp.shared.exceptions import McpError from mcp.shared.session import RequestResponder +from opentelemetry import propagate from ._tools import ( FunctionTool, @@ -380,6 +381,22 @@ def _normalize_mcp_name(name: str) -> str: return re.sub(r"[^A-Za-z0-9_.-]", "-", name) +def _inject_otel_into_mcp_meta(meta: dict[str, Any] | None = None) -> dict[str, Any] | None: + """Inject OpenTelemetry trace context into MCP request _meta via the global propagator(s).""" + carrier: dict[str, str] = {} + propagate.inject(carrier) + if not carrier: + return meta + + if meta is None: + meta = {} + for key, value in carrier.items(): + if key not in meta: + meta[key] = value + + return meta + + # region: MCP Plugin @@ -875,12 +892,15 @@ class MCPTool: } } + # Inject OpenTelemetry trace context into MCP _meta for distributed tracing. + otel_meta = _inject_otel_into_mcp_meta() + parser = self.parse_tool_results or _parse_tool_result_from_mcp # Try the operation, reconnecting once if the connection is closed for attempt in range(2): try: - result = await self.session.call_tool(tool_name, arguments=filtered_kwargs) # type: ignore + result = await self.session.call_tool(tool_name, arguments=filtered_kwargs, meta=otel_meta) # type: ignore return parser(result) except ClosedResourceError as cl_ex: if attempt == 0: diff --git a/python/packages/core/tests/core/test_mcp.py b/python/packages/core/tests/core/test_mcp.py index 8b213476aa..09d5b11754 100644 --- a/python/packages/core/tests/core/test_mcp.py +++ b/python/packages/core/tests/core/test_mcp.py @@ -2591,3 +2591,70 @@ async def test_mcp_tool_filters_framework_kwargs(): assert "thread" not in arguments assert "conversation_id" not in arguments assert "options" not in arguments + + +# region: OTel trace context propagation via _meta + + +@pytest.mark.parametrize( + "use_span,expect_traceparent", + [ + (True, True), + (False, False), + ], +) +async def test_mcp_tool_call_tool_otel_meta(use_span, expect_traceparent, span_exporter): + """call_tool propagates OTel trace context via meta only when a span is active.""" + from opentelemetry import trace + + class TestServer(MCPTool): + async def connect(self): + self.session = Mock(spec=ClientSession) + self.session.list_tools = AsyncMock( + return_value=types.ListToolsResult( + tools=[ + types.Tool( + name="test_tool", + description="Test tool", + inputSchema={ + "type": "object", + "properties": {"param": {"type": "string"}}, + "required": ["param"], + }, + ) + ] + ) + ) + self.session.call_tool = AsyncMock( + return_value=types.CallToolResult(content=[types.TextContent(type="text", text="result")]) + ) + + def get_mcp_client(self) -> _AsyncGeneratorContextManager[Any, None]: + return None + + server = TestServer(name="test_server") + async with server: + await server.load_tools() + + if use_span: + tracer = trace.get_tracer("test") + with tracer.start_as_current_span("test_span"): + await server.functions[0].invoke(param="test_value") + else: + # Use an invalid span to ensure no trace context is injected; + # call server.call_tool directly to bypass FunctionTool.invoke's own span. + with trace.use_span(trace.NonRecordingSpan(trace.INVALID_SPAN_CONTEXT)): + await server.call_tool("test_tool", param="test_value") + + meta = server.session.call_tool.call_args.kwargs.get("meta") + if expect_traceparent: + # When a valid span is active, we expect some propagation fields to be injected, + # but we do not assume any specific header name to keep this test propagator-agnostic. + assert meta is not None + assert isinstance(meta, dict) + assert len(meta) > 0 + else: + assert meta is None + + +# endregion diff --git a/python/samples/02-agents/observability/README.md b/python/samples/02-agents/observability/README.md index fa6ba0fec1..b311138714 100644 --- a/python/samples/02-agents/observability/README.md +++ b/python/samples/02-agents/observability/README.md @@ -22,6 +22,10 @@ The Agent Framework Python SDK is designed to efficiently generate comprehensive Next to what happens in the code when you run, we also make setting up observability as easy as possible. By calling a single function `configure_otel_providers()` from the `agent_framework.observability` module, you can enable telemetry for traces, logs, and metrics. The function automatically reads standard OpenTelemetry environment variables to configure exporters and providers, making it simple to get started. +### MCP trace propagation + +Whenever there is an active OpenTelemetry span context, Agent Framework automatically propagates trace context to MCP servers via the `params._meta` field of `tools/call` requests. It uses the globally-configured OpenTelemetry propagator(s) (W3C Trace Context by default, producing `traceparent` and `tracestate`), so custom propagators (B3, Jaeger, etc.) are also supported. This enables distributed tracing across agent-to-MCP-server boundaries for all transports (stdio, HTTP, WebSocket), compliant with the [MCP `_meta` specification](https://modelcontextprotocol.io/specification/2025-11-25/basic#_meta). + ### Five patterns for configuring observability We've identified multiple ways to configure observability in your application, depending on your needs: From 08275f657b1b473144f1bf07e2948672a23ca845 Mon Sep 17 00:00:00 2001 From: westey <164392973+westey-m@users.noreply.github.com> Date: Tue, 17 Feb 2026 15:51:30 +0000 Subject: [PATCH 22/22] .NET: Improve session cast error message quality and consistency (#3973) * Improve session cast error messge consistency * Update changelog --- dotnet/src/Microsoft.Agents.AI.A2A/A2AAgent.cs | 4 ++-- .../Microsoft.Agents.AI.CopilotStudio/CopilotStudioAgent.cs | 6 +++--- dotnet/src/Microsoft.Agents.AI.DurableTask/CHANGELOG.md | 1 + .../src/Microsoft.Agents.AI.DurableTask/DurableAIAgent.cs | 2 +- .../Microsoft.Agents.AI.DurableTask/DurableAIAgentProxy.cs | 2 +- .../GitHubCopilotAgent.cs | 4 ++-- .../src/Microsoft.Agents.AI.Workflows/WorkflowHostAgent.cs | 2 +- .../src/Microsoft.Agents.AI/ChatClient/ChatClientAgent.cs | 4 ++-- .../BasicStreamingTests.cs | 2 +- .../ForwardedPropertiesTests.cs | 2 +- .../SharedStateTests.cs | 2 +- .../AGUIEndpointRouteBuilderExtensionsTests.cs | 4 ++-- .../TestEchoAgent.cs | 2 +- 13 files changed, 19 insertions(+), 18 deletions(-) diff --git a/dotnet/src/Microsoft.Agents.AI.A2A/A2AAgent.cs b/dotnet/src/Microsoft.Agents.AI.A2A/A2AAgent.cs index 18d8bfacdd..2393f59202 100644 --- a/dotnet/src/Microsoft.Agents.AI.A2A/A2AAgent.cs +++ b/dotnet/src/Microsoft.Agents.AI.A2A/A2AAgent.cs @@ -81,7 +81,7 @@ public sealed class A2AAgent : AIAgent if (session is not A2AAgentSession typedSession) { - throw new InvalidOperationException("The provided session is not compatible with the agent. Only sessions created by the agent can be serialized."); + throw new InvalidOperationException($"The provided session type '{session.GetType().Name}' is not compatible with this agent. Only sessions of type '{nameof(A2AAgentSession)}' can be serialized by this agent."); } return new(typedSession.Serialize(jsonSerializerOptions)); @@ -256,7 +256,7 @@ public sealed class A2AAgent : AIAgent if (session is not A2AAgentSession typedSession) { - throw new InvalidOperationException($"The provided session type {session.GetType()} is not compatible with the agent. Only A2A agent created sessions are supported."); + throw new InvalidOperationException($"The provided session type '{session.GetType().Name}' is not compatible with this agent. Only sessions of type '{nameof(A2AAgentSession)}' can be used by this agent."); } return typedSession; diff --git a/dotnet/src/Microsoft.Agents.AI.CopilotStudio/CopilotStudioAgent.cs b/dotnet/src/Microsoft.Agents.AI.CopilotStudio/CopilotStudioAgent.cs index bb53f94b50..5485d08a3f 100644 --- a/dotnet/src/Microsoft.Agents.AI.CopilotStudio/CopilotStudioAgent.cs +++ b/dotnet/src/Microsoft.Agents.AI.CopilotStudio/CopilotStudioAgent.cs @@ -60,7 +60,7 @@ public class CopilotStudioAgent : AIAgent if (session is not CopilotStudioAgentSession typedSession) { - throw new InvalidOperationException("The provided session is not compatible with the agent. Only sessions created by the agent can be serialized."); + throw new InvalidOperationException($"The provided session type '{session.GetType().Name}' is not compatible with this agent. Only sessions of type '{nameof(CopilotStudioAgentSession)}' can be serialized by this agent."); } return new(typedSession.Serialize(jsonSerializerOptions)); @@ -84,7 +84,7 @@ public class CopilotStudioAgent : AIAgent session ??= await this.CreateSessionAsync(cancellationToken).ConfigureAwait(false); if (session is not CopilotStudioAgentSession typedSession) { - throw new InvalidOperationException("The provided session is not compatible with the agent. Only sessions created by the agent can be used."); + throw new InvalidOperationException($"The provided session type '{session.GetType().Name}' is not compatible with this agent. Only sessions of type '{nameof(CopilotStudioAgentSession)}' can be used by this agent."); } typedSession.ConversationId ??= await this.StartNewConversationAsync(cancellationToken).ConfigureAwait(false); @@ -123,7 +123,7 @@ public class CopilotStudioAgent : AIAgent session ??= await this.CreateSessionAsync(cancellationToken).ConfigureAwait(false); if (session is not CopilotStudioAgentSession typedSession) { - throw new InvalidOperationException("The provided session is not compatible with the agent. Only sessions created by the agent can be used."); + throw new InvalidOperationException($"The provided session type '{session.GetType().Name}' is not compatible with this agent. Only sessions of type '{nameof(CopilotStudioAgentSession)}' can be used by this agent."); } typedSession.ConversationId ??= await this.StartNewConversationAsync(cancellationToken).ConfigureAwait(false); diff --git a/dotnet/src/Microsoft.Agents.AI.DurableTask/CHANGELOG.md b/dotnet/src/Microsoft.Agents.AI.DurableTask/CHANGELOG.md index b0124d3de7..e3e90fdae0 100644 --- a/dotnet/src/Microsoft.Agents.AI.DurableTask/CHANGELOG.md +++ b/dotnet/src/Microsoft.Agents.AI.DurableTask/CHANGELOG.md @@ -14,6 +14,7 @@ - Changed AIAgent.SerializeSession to AIAgent.SerializeSessionAsync ([#3879](https://github.com/microsoft/agent-framework/pull/3879)) - Changed ChatHistory and AIContext Providers to have pipeline semantics ([#3806](https://github.com/microsoft/agent-framework/pull/3806)) - Marked all `RunAsync` overloads as `new`, added missing ones, and added support for primitives and arrays ([#3803](https://github.com/microsoft/agent-framework/pull/3803)) +- Improve session cast error message quality and consistency ([#3973](https://github.com/microsoft/agent-framework/pull/3973)) ## v1.0.0-preview.251204.1 diff --git a/dotnet/src/Microsoft.Agents.AI.DurableTask/DurableAIAgent.cs b/dotnet/src/Microsoft.Agents.AI.DurableTask/DurableAIAgent.cs index ee0c457deb..599ea3703f 100644 --- a/dotnet/src/Microsoft.Agents.AI.DurableTask/DurableAIAgent.cs +++ b/dotnet/src/Microsoft.Agents.AI.DurableTask/DurableAIAgent.cs @@ -55,7 +55,7 @@ public sealed class DurableAIAgent : AIAgent if (session is not DurableAgentSession durableSession) { - throw new InvalidOperationException("The provided session is not compatible with the agent. Only sessions created by the agent can be serialized."); + throw new InvalidOperationException($"The provided session type '{session.GetType().Name}' is not compatible with this agent. Only sessions of type '{nameof(DurableAgentSession)}' can be serialized by this agent."); } return new(durableSession.Serialize(jsonSerializerOptions)); diff --git a/dotnet/src/Microsoft.Agents.AI.DurableTask/DurableAIAgentProxy.cs b/dotnet/src/Microsoft.Agents.AI.DurableTask/DurableAIAgentProxy.cs index 8987343c60..36a9336c36 100644 --- a/dotnet/src/Microsoft.Agents.AI.DurableTask/DurableAIAgentProxy.cs +++ b/dotnet/src/Microsoft.Agents.AI.DurableTask/DurableAIAgentProxy.cs @@ -20,7 +20,7 @@ internal class DurableAIAgentProxy(string name, IDurableAgentClient agentClient) if (session is not DurableAgentSession durableSession) { - throw new InvalidOperationException("The provided session is not compatible with the agent. Only sessions created by the agent can be serialized."); + throw new InvalidOperationException($"The provided session type '{session.GetType().Name}' is not compatible with this agent. Only sessions of type '{nameof(DurableAgentSession)}' can be serialized by this agent."); } return new(durableSession.Serialize(jsonSerializerOptions)); diff --git a/dotnet/src/Microsoft.Agents.AI.GitHub.Copilot/GitHubCopilotAgent.cs b/dotnet/src/Microsoft.Agents.AI.GitHub.Copilot/GitHubCopilotAgent.cs index f0533c8461..c966f591fc 100644 --- a/dotnet/src/Microsoft.Agents.AI.GitHub.Copilot/GitHubCopilotAgent.cs +++ b/dotnet/src/Microsoft.Agents.AI.GitHub.Copilot/GitHubCopilotAgent.cs @@ -104,7 +104,7 @@ public sealed class GitHubCopilotAgent : AIAgent, IAsyncDisposable if (session is not GitHubCopilotAgentSession typedSession) { - throw new InvalidOperationException("The provided session is not compatible with the agent. Only sessions created by the agent can be serialized."); + throw new InvalidOperationException($"The provided session type '{session.GetType().Name}' is not compatible with this agent. Only sessions of type '{nameof(GitHubCopilotAgentSession)}' can be serialized by this agent."); } return new(typedSession.Serialize(jsonSerializerOptions)); @@ -139,7 +139,7 @@ public sealed class GitHubCopilotAgent : AIAgent, IAsyncDisposable if (session is not GitHubCopilotAgentSession typedSession) { throw new InvalidOperationException( - $"The provided session type {session.GetType()} is not compatible with the agent. Only GitHub Copilot agent created sessions are supported."); + $"The provided session type '{session.GetType().Name}' is not compatible with this agent. Only sessions of type '{nameof(GitHubCopilotAgentSession)}' can be used by this agent."); } // Ensure the client is started diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowHostAgent.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowHostAgent.cs index 7438ce1a34..fb4ce5c9ab 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowHostAgent.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowHostAgent.cs @@ -74,7 +74,7 @@ internal sealed class WorkflowHostAgent : AIAgent if (session is not WorkflowSession workflowSession) { - throw new InvalidOperationException("The provided session is not compatible with the agent. Only sessions created by the agent can be serialized."); + throw new InvalidOperationException($"The provided session type '{session.GetType().Name}' is not compatible with this agent. Only sessions of type '{nameof(WorkflowSession)}' can be serialized by this agent."); } return new(workflowSession.Serialize(jsonSerializerOptions)); diff --git a/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgent.cs b/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgent.cs index acf90fc16d..604f92ee76 100644 --- a/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgent.cs +++ b/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgent.cs @@ -366,7 +366,7 @@ public sealed partial class ChatClientAgent : AIAgent if (session is not ChatClientAgentSession typedSession) { - throw new InvalidOperationException("The provided session is not compatible with the agent. Only sessions created by the agent can be serialized."); + throw new InvalidOperationException($"The provided session type '{session.GetType().Name}' is not compatible with this agent. Only sessions of type '{nameof(ChatClientAgentSession)}' can be serialized by this agent."); } return new(typedSession.Serialize(jsonSerializerOptions)); @@ -674,7 +674,7 @@ public sealed partial class ChatClientAgent : AIAgent session ??= await this.CreateSessionAsync(cancellationToken).ConfigureAwait(false); if (session is not ChatClientAgentSession typedSession) { - throw new InvalidOperationException("The provided session is not compatible with the agent. Only sessions created by the agent can be used."); + throw new InvalidOperationException($"The provided session type '{session.GetType().Name}' is not compatible with this agent. Only sessions of type '{nameof(ChatClientAgentSession)}' can be used by this agent."); } // Supplying messages when continuing a background response is not allowed. diff --git a/dotnet/tests/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore.IntegrationTests/BasicStreamingTests.cs b/dotnet/tests/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore.IntegrationTests/BasicStreamingTests.cs index 1be1eddef1..d94e520420 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore.IntegrationTests/BasicStreamingTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore.IntegrationTests/BasicStreamingTests.cs @@ -357,7 +357,7 @@ internal sealed class FakeMultiMessageAgent : AIAgent { if (session is not FakeAgentSession fakeSession) { - throw new InvalidOperationException("The provided session is not compatible with the agent. Only sessions created by the agent can be serialized."); + throw new InvalidOperationException($"The provided session type '{session.GetType().Name}' is not compatible with this agent. Only sessions of type '{nameof(FakeAgentSession)}' can be serialized by this agent."); } return new(JsonSerializer.SerializeToElement(fakeSession, jsonSerializerOptions)); diff --git a/dotnet/tests/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore.IntegrationTests/ForwardedPropertiesTests.cs b/dotnet/tests/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore.IntegrationTests/ForwardedPropertiesTests.cs index 844900372b..60d430d23c 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore.IntegrationTests/ForwardedPropertiesTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore.IntegrationTests/ForwardedPropertiesTests.cs @@ -345,7 +345,7 @@ internal sealed class FakeForwardedPropsAgent : AIAgent { if (session is not FakeAgentSession fakeSession) { - throw new InvalidOperationException("The provided session is not compatible with the agent. Only sessions created by the agent can be serialized."); + throw new InvalidOperationException($"The provided session type '{session.GetType().Name}' is not compatible with this agent. Only sessions of type '{nameof(FakeAgentSession)}' can be serialized by this agent."); } return new(JsonSerializer.SerializeToElement(fakeSession, jsonSerializerOptions)); diff --git a/dotnet/tests/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore.IntegrationTests/SharedStateTests.cs b/dotnet/tests/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore.IntegrationTests/SharedStateTests.cs index 78441254d9..cc9c9ce8ef 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore.IntegrationTests/SharedStateTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore.IntegrationTests/SharedStateTests.cs @@ -428,7 +428,7 @@ internal sealed class FakeStateAgent : AIAgent { if (session is not FakeAgentSession fakeSession) { - throw new InvalidOperationException("The provided session is not compatible with the agent. Only sessions created by the agent can be serialized."); + throw new InvalidOperationException($"The provided session type '{session.GetType().Name}' is not compatible with this agent. Only sessions of type '{nameof(FakeAgentSession)}' can be serialized by this agent."); } return new(JsonSerializer.SerializeToElement(fakeSession, jsonSerializerOptions)); diff --git a/dotnet/tests/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore.UnitTests/AGUIEndpointRouteBuilderExtensionsTests.cs b/dotnet/tests/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore.UnitTests/AGUIEndpointRouteBuilderExtensionsTests.cs index 383a7d1062..84a20e1938 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore.UnitTests/AGUIEndpointRouteBuilderExtensionsTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Hosting.AGUI.AspNetCore.UnitTests/AGUIEndpointRouteBuilderExtensionsTests.cs @@ -436,7 +436,7 @@ public sealed class AGUIEndpointRouteBuilderExtensionsTests { if (session is not TestAgentSession testSession) { - throw new InvalidOperationException("The provided session is not compatible with the agent. Only sessions created by the agent can be serialized."); + throw new InvalidOperationException($"The provided session type '{session.GetType().Name}' is not compatible with this agent. Only sessions of type '{nameof(TestAgentSession)}' can be serialized by this agent."); } return new(JsonSerializer.SerializeToElement(testSession, jsonSerializerOptions)); @@ -535,7 +535,7 @@ public sealed class AGUIEndpointRouteBuilderExtensionsTests { if (session is not TestAgentSession testSession) { - throw new InvalidOperationException("The provided session is not compatible with the agent. Only sessions created by the agent can be serialized."); + throw new InvalidOperationException($"The provided session type '{session.GetType().Name}' is not compatible with this agent. Only sessions of type '{nameof(TestAgentSession)}' can be serialized by this agent."); } return new(JsonSerializer.SerializeToElement(testSession, jsonSerializerOptions)); diff --git a/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/TestEchoAgent.cs b/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/TestEchoAgent.cs index b1c7c8cbe3..0d32a6dbaf 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/TestEchoAgent.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/TestEchoAgent.cs @@ -28,7 +28,7 @@ internal class TestEchoAgent(string? id = null, string? name = null, string? pre { if (session is not EchoAgentSession typedSession) { - throw new InvalidOperationException("The provided session is not compatible with the agent. Only sessions created by the agent can be serialized."); + throw new InvalidOperationException($"The provided session type '{session.GetType().Name}' is not compatible with this agent. Only sessions of type '{nameof(EchoAgentSession)}' can be serialized by this agent."); } return new(JsonSerializer.SerializeToElement(typedSession, jsonSerializerOptions));