From ac018f700bc8bd1d6d44f1ad795ef17276b560c9 Mon Sep 17 00:00:00 2001 From: Evan Mattson <35585003+moonbox3@users.noreply.github.com> Date: Thu, 6 Nov 2025 11:31:24 +0900 Subject: [PATCH 01/16] Python: Fix ag-ui examples packaging for PyPI publish (#1953) * Fix ag-ui examples packaging for PyPI publish * Fix markdown links --- python/packages/ag-ui/README.md | 4 ++-- .../.env.example | 0 .../.vscode/settings.json | 0 .../{examples => agent_framework_ag_ui_examples}/README.md | 0 .../{examples => agent_framework_ag_ui_examples}/__init__.py | 0 .../{examples => agent_framework_ag_ui_examples}/__main__.py | 0 .../agents/__init__.py | 0 .../agents/document_writer_agent.py | 0 .../agents/human_in_the_loop_agent.py | 0 .../agents/recipe_agent.py | 0 .../agents/research_assistant_agent.py | 0 .../agents/simple_agent.py | 0 .../agents/task_planner_agent.py | 0 .../agents/task_steps_agent.py | 0 .../agents/ui_generator_agent.py | 0 .../agents/weather_agent.py | 0 .../server/__init__.py | 0 .../server/api/__init__.py | 0 .../server/api/backend_tool_rendering.py | 0 .../server/main.py | 0 python/packages/ag-ui/getting_started/README.md | 2 +- python/packages/ag-ui/pyproject.toml | 5 ++--- python/uv.lock | 2 +- 23 files changed, 6 insertions(+), 7 deletions(-) rename python/packages/ag-ui/{examples => agent_framework_ag_ui_examples}/.env.example (100%) rename python/packages/ag-ui/{examples => agent_framework_ag_ui_examples}/.vscode/settings.json (100%) rename python/packages/ag-ui/{examples => agent_framework_ag_ui_examples}/README.md (100%) rename python/packages/ag-ui/{examples => agent_framework_ag_ui_examples}/__init__.py (100%) rename python/packages/ag-ui/{examples => agent_framework_ag_ui_examples}/__main__.py (100%) rename python/packages/ag-ui/{examples => agent_framework_ag_ui_examples}/agents/__init__.py (100%) rename python/packages/ag-ui/{examples => agent_framework_ag_ui_examples}/agents/document_writer_agent.py (100%) rename python/packages/ag-ui/{examples => agent_framework_ag_ui_examples}/agents/human_in_the_loop_agent.py (100%) rename python/packages/ag-ui/{examples => agent_framework_ag_ui_examples}/agents/recipe_agent.py (100%) rename python/packages/ag-ui/{examples => agent_framework_ag_ui_examples}/agents/research_assistant_agent.py (100%) rename python/packages/ag-ui/{examples => agent_framework_ag_ui_examples}/agents/simple_agent.py (100%) rename python/packages/ag-ui/{examples => agent_framework_ag_ui_examples}/agents/task_planner_agent.py (100%) rename python/packages/ag-ui/{examples => agent_framework_ag_ui_examples}/agents/task_steps_agent.py (100%) rename python/packages/ag-ui/{examples => agent_framework_ag_ui_examples}/agents/ui_generator_agent.py (100%) rename python/packages/ag-ui/{examples => agent_framework_ag_ui_examples}/agents/weather_agent.py (100%) rename python/packages/ag-ui/{examples => agent_framework_ag_ui_examples}/server/__init__.py (100%) rename python/packages/ag-ui/{examples => agent_framework_ag_ui_examples}/server/api/__init__.py (100%) rename python/packages/ag-ui/{examples => agent_framework_ag_ui_examples}/server/api/backend_tool_rendering.py (100%) rename python/packages/ag-ui/{examples => agent_framework_ag_ui_examples}/server/main.py (100%) diff --git a/python/packages/ag-ui/README.md b/python/packages/ag-ui/README.md index 7e0d6b73d9..a7e24b5fbc 100644 --- a/python/packages/ag-ui/README.md +++ b/python/packages/ag-ui/README.md @@ -36,7 +36,7 @@ add_agent_framework_fastapi_endpoint(app, agent, "/") ## Documentation - **[Getting Started Tutorial](getting_started/)** - Step-by-step guide to building your first AG-UI server and client -- **[Examples](examples/)** - Complete examples for AG-UI features +- **[Examples](agent_framework_ag_ui_examples/)** - Complete examples for AG-UI features ## Features @@ -64,7 +64,7 @@ The package uses a clean, orchestrator-based architecture: ## Next Steps 1. **New to AG-UI?** Start with the [Getting Started Tutorial](getting_started/) -2. **Want to see examples?** Check out the [Examples](examples/) for AG-UI features +2. **Want to see examples?** Check out the [Examples](agent_framework_ag_ui_examples/) for AG-UI features ## License diff --git a/python/packages/ag-ui/examples/.env.example b/python/packages/ag-ui/agent_framework_ag_ui_examples/.env.example similarity index 100% rename from python/packages/ag-ui/examples/.env.example rename to python/packages/ag-ui/agent_framework_ag_ui_examples/.env.example diff --git a/python/packages/ag-ui/examples/.vscode/settings.json b/python/packages/ag-ui/agent_framework_ag_ui_examples/.vscode/settings.json similarity index 100% rename from python/packages/ag-ui/examples/.vscode/settings.json rename to python/packages/ag-ui/agent_framework_ag_ui_examples/.vscode/settings.json diff --git a/python/packages/ag-ui/examples/README.md b/python/packages/ag-ui/agent_framework_ag_ui_examples/README.md similarity index 100% rename from python/packages/ag-ui/examples/README.md rename to python/packages/ag-ui/agent_framework_ag_ui_examples/README.md diff --git a/python/packages/ag-ui/examples/__init__.py b/python/packages/ag-ui/agent_framework_ag_ui_examples/__init__.py similarity index 100% rename from python/packages/ag-ui/examples/__init__.py rename to python/packages/ag-ui/agent_framework_ag_ui_examples/__init__.py diff --git a/python/packages/ag-ui/examples/__main__.py b/python/packages/ag-ui/agent_framework_ag_ui_examples/__main__.py similarity index 100% rename from python/packages/ag-ui/examples/__main__.py rename to python/packages/ag-ui/agent_framework_ag_ui_examples/__main__.py diff --git a/python/packages/ag-ui/examples/agents/__init__.py b/python/packages/ag-ui/agent_framework_ag_ui_examples/agents/__init__.py similarity index 100% rename from python/packages/ag-ui/examples/agents/__init__.py rename to python/packages/ag-ui/agent_framework_ag_ui_examples/agents/__init__.py diff --git a/python/packages/ag-ui/examples/agents/document_writer_agent.py b/python/packages/ag-ui/agent_framework_ag_ui_examples/agents/document_writer_agent.py similarity index 100% rename from python/packages/ag-ui/examples/agents/document_writer_agent.py rename to python/packages/ag-ui/agent_framework_ag_ui_examples/agents/document_writer_agent.py diff --git a/python/packages/ag-ui/examples/agents/human_in_the_loop_agent.py b/python/packages/ag-ui/agent_framework_ag_ui_examples/agents/human_in_the_loop_agent.py similarity index 100% rename from python/packages/ag-ui/examples/agents/human_in_the_loop_agent.py rename to python/packages/ag-ui/agent_framework_ag_ui_examples/agents/human_in_the_loop_agent.py diff --git a/python/packages/ag-ui/examples/agents/recipe_agent.py b/python/packages/ag-ui/agent_framework_ag_ui_examples/agents/recipe_agent.py similarity index 100% rename from python/packages/ag-ui/examples/agents/recipe_agent.py rename to python/packages/ag-ui/agent_framework_ag_ui_examples/agents/recipe_agent.py diff --git a/python/packages/ag-ui/examples/agents/research_assistant_agent.py b/python/packages/ag-ui/agent_framework_ag_ui_examples/agents/research_assistant_agent.py similarity index 100% rename from python/packages/ag-ui/examples/agents/research_assistant_agent.py rename to python/packages/ag-ui/agent_framework_ag_ui_examples/agents/research_assistant_agent.py diff --git a/python/packages/ag-ui/examples/agents/simple_agent.py b/python/packages/ag-ui/agent_framework_ag_ui_examples/agents/simple_agent.py similarity index 100% rename from python/packages/ag-ui/examples/agents/simple_agent.py rename to python/packages/ag-ui/agent_framework_ag_ui_examples/agents/simple_agent.py diff --git a/python/packages/ag-ui/examples/agents/task_planner_agent.py b/python/packages/ag-ui/agent_framework_ag_ui_examples/agents/task_planner_agent.py similarity index 100% rename from python/packages/ag-ui/examples/agents/task_planner_agent.py rename to python/packages/ag-ui/agent_framework_ag_ui_examples/agents/task_planner_agent.py diff --git a/python/packages/ag-ui/examples/agents/task_steps_agent.py b/python/packages/ag-ui/agent_framework_ag_ui_examples/agents/task_steps_agent.py similarity index 100% rename from python/packages/ag-ui/examples/agents/task_steps_agent.py rename to python/packages/ag-ui/agent_framework_ag_ui_examples/agents/task_steps_agent.py diff --git a/python/packages/ag-ui/examples/agents/ui_generator_agent.py b/python/packages/ag-ui/agent_framework_ag_ui_examples/agents/ui_generator_agent.py similarity index 100% rename from python/packages/ag-ui/examples/agents/ui_generator_agent.py rename to python/packages/ag-ui/agent_framework_ag_ui_examples/agents/ui_generator_agent.py diff --git a/python/packages/ag-ui/examples/agents/weather_agent.py b/python/packages/ag-ui/agent_framework_ag_ui_examples/agents/weather_agent.py similarity index 100% rename from python/packages/ag-ui/examples/agents/weather_agent.py rename to python/packages/ag-ui/agent_framework_ag_ui_examples/agents/weather_agent.py diff --git a/python/packages/ag-ui/examples/server/__init__.py b/python/packages/ag-ui/agent_framework_ag_ui_examples/server/__init__.py similarity index 100% rename from python/packages/ag-ui/examples/server/__init__.py rename to python/packages/ag-ui/agent_framework_ag_ui_examples/server/__init__.py diff --git a/python/packages/ag-ui/examples/server/api/__init__.py b/python/packages/ag-ui/agent_framework_ag_ui_examples/server/api/__init__.py similarity index 100% rename from python/packages/ag-ui/examples/server/api/__init__.py rename to python/packages/ag-ui/agent_framework_ag_ui_examples/server/api/__init__.py diff --git a/python/packages/ag-ui/examples/server/api/backend_tool_rendering.py b/python/packages/ag-ui/agent_framework_ag_ui_examples/server/api/backend_tool_rendering.py similarity index 100% rename from python/packages/ag-ui/examples/server/api/backend_tool_rendering.py rename to python/packages/ag-ui/agent_framework_ag_ui_examples/server/api/backend_tool_rendering.py diff --git a/python/packages/ag-ui/examples/server/main.py b/python/packages/ag-ui/agent_framework_ag_ui_examples/server/main.py similarity index 100% rename from python/packages/ag-ui/examples/server/main.py rename to python/packages/ag-ui/agent_framework_ag_ui_examples/server/main.py diff --git a/python/packages/ag-ui/getting_started/README.md b/python/packages/ag-ui/getting_started/README.md index 2d1219bed4..f656e49979 100644 --- a/python/packages/ag-ui/getting_started/README.md +++ b/python/packages/ag-ui/getting_started/README.md @@ -629,7 +629,7 @@ Now that you understand the basics of AG-UI, you can: ## Additional Resources -- [AG-UI Examples](../examples/README.md): Complete working examples for all 7 features +- [AG-UI Examples](../agent_framework_ag_ui_examples/README.md): Complete working examples for all 7 features - [Agent Framework Documentation](../../core/README.md): Learn more about creating agents - [AG-UI Protocol Spec](https://docs.ag-ui.com/): Official protocol documentation diff --git a/python/packages/ag-ui/pyproject.toml b/python/packages/ag-ui/pyproject.toml index 85297ac73f..400e5e6e34 100644 --- a/python/packages/ag-ui/pyproject.toml +++ b/python/packages/ag-ui/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "agent-framework-ag-ui" -version = "1.0.0b251106" +version = "1.0.0b251106.post1" description = "AG-UI protocol integration for Agent Framework" readme = "README.md" license-files = ["LICENSE"] @@ -40,8 +40,7 @@ requires = ["hatchling"] build-backend = "hatchling.build" [tool.hatch.build.targets.wheel] -packages = ["agent_framework_ag_ui"] -force-include = { "examples" = "agent_framework_ag_ui_examples" } +packages = ["agent_framework_ag_ui", "agent_framework_ag_ui_examples"] [tool.pytest.ini_options] asyncio_mode = "auto" diff --git a/python/uv.lock b/python/uv.lock index ad924492ac..ddf1ac4205 100644 --- a/python/uv.lock +++ b/python/uv.lock @@ -193,7 +193,7 @@ requires-dist = [ [[package]] name = "agent-framework-ag-ui" -version = "1.0.0b251106" +version = "1.0.0b251106.post1" source = { editable = "packages/ag-ui" } dependencies = [ { name = "ag-ui-protocol", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, From ee1661ecb7d38284098eba97f66ad8befa8b9445 Mon Sep 17 00:00:00 2001 From: Giles Odigwe <79032838+giles17@users.noreply.github.com> Date: Wed, 5 Nov 2025 18:59:22 -0800 Subject: [PATCH 02/16] Python: Thread Samples Fix (#1945) * thread samples fix * custom chat message store fix --- .../samples/getting_started/threads/README.md | 11 +++- .../custom_chat_message_store_thread.py | 29 +++++++-- .../redis_chat_message_store_thread.py | 8 +++ .../threads/suspend_resume_thread.py | 61 ++++++++++++------- 4 files changed, 79 insertions(+), 30 deletions(-) diff --git a/python/samples/getting_started/threads/README.md b/python/samples/getting_started/threads/README.md index d1ece041f7..32c19d537f 100644 --- a/python/samples/getting_started/threads/README.md +++ b/python/samples/getting_started/threads/README.md @@ -7,5 +7,14 @@ This folder contains examples demonstrating different ways to manage conversatio | File | Description | |------|-------------| | [`custom_chat_message_store_thread.py`](custom_chat_message_store_thread.py) | Demonstrates how to implement a custom `ChatMessageStore` for persisting conversation history. Shows how to create a custom store with serialization/deserialization capabilities and integrate it with agents for thread management across multiple sessions. | -| [`suspend_resume_thread.py`](suspend_resume_thread.py) | Shows how to suspend and resume conversation threads, allowing you to save the state of a conversation and continue it later. This is useful for long-running conversations or when you need to persist conversation state across application restarts. | | [`redis_chat_message_store_thread.py`](redis_chat_message_store_thread.py) | Comprehensive examples of using the Redis-backed `RedisChatMessageStore` for persistent conversation storage. Covers basic usage, user session management, conversation persistence across app restarts, thread serialization, and automatic message trimming. Requires Redis server and demonstrates production-ready patterns for scalable chat applications. | +| [`suspend_resume_thread.py`](suspend_resume_thread.py) | Shows how to suspend and resume conversation threads, comparing service-managed threads (Azure AI) with in-memory threads (OpenAI). Demonstrates saving conversation state and continuing it later, useful for long-running conversations or persisting state across application restarts. | + +## Environment Variables + +Make sure to set the following environment variables before running the examples: + +- `OPENAI_API_KEY`: Your OpenAI API key (required for all samples) +- `OPENAI_CHAT_MODEL_ID`: The OpenAI model to use (e.g., `gpt-4o`, `gpt-4o-mini`, `gpt-3.5-turbo`) (required for all samples) +- `AZURE_AI_PROJECT_ENDPOINT`: Azure AI Project endpoint URL (required for service-managed thread examples) +- `AZURE_AI_MODEL_DEPLOYMENT_NAME`: The name of your model deployment (required for service-managed thread examples) diff --git a/python/samples/getting_started/threads/custom_chat_message_store_thread.py b/python/samples/getting_started/threads/custom_chat_message_store_thread.py index f4a9d79112..266bb095f2 100644 --- a/python/samples/getting_started/threads/custom_chat_message_store_thread.py +++ b/python/samples/getting_started/threads/custom_chat_message_store_thread.py @@ -8,6 +8,14 @@ from agent_framework import ChatMessage, ChatMessageStoreProtocol from agent_framework._threads import ChatMessageStoreState from agent_framework.openai import OpenAIChatClient +""" +Custom Chat Message Store Thread Example + +This sample demonstrates how to implement and use a custom chat message store +for thread management, allowing you to persist conversation history in your +preferred storage solution (database, file system, etc.). +""" + class CustomChatMessageStore(ChatMessageStoreProtocol): """Implementation of custom chat message store. @@ -24,13 +32,22 @@ class CustomChatMessageStore(ChatMessageStoreProtocol): async def list_messages(self) -> list[ChatMessage]: return self._messages - async def deserialize_state(self, serialized_store_state: Any, **kwargs: Any) -> None: + @classmethod + async def deserialize(cls, serialized_store_state: Any, **kwargs: Any) -> "CustomChatMessageStore": + """Create a new instance from serialized state.""" + store = cls() + await store.update_from_state(serialized_store_state, **kwargs) + return store + + async def update_from_state(self, serialized_store_state: Any, **kwargs: Any) -> None: + """Update this instance from serialized state.""" if serialized_store_state: state = ChatMessageStoreState.from_dict(serialized_store_state, **kwargs) if state.messages: self._messages.extend(state.messages) - async def serialize_state(self, **kwargs: Any) -> Any: + async def serialize(self, **kwargs: Any) -> Any: + """Serialize this store's state.""" state = ChatMessageStoreState(messages=self._messages) return state.to_dict(**kwargs) @@ -42,8 +59,8 @@ async def main() -> None: # OpenAI Chat Client is used as an example here, # other chat clients can be used as well. agent = OpenAIChatClient().create_agent( - name="Joker", - instructions="You are good at telling jokes.", + name="CustomBot", + instructions="You are a helpful assistant that remembers our conversation.", # Use custom chat message store. # If not provided, the default in-memory store will be used. chat_message_store_factory=CustomChatMessageStore, @@ -53,7 +70,7 @@ async def main() -> None: thread = agent.get_new_thread() # Respond to user input. - query = "Tell me a joke about a pirate." + query = "Hello! My name is Alice and I love pizza." print(f"User: {query}") print(f"Agent: {await agent.run(query, thread=thread)}\n") @@ -67,7 +84,7 @@ async def main() -> None: resumed_thread = await agent.deserialize_thread(serialized_thread) # Respond to user input. - query = "Now tell the same joke in the voice of a pirate, and add some emojis to the joke." + query = "What do you remember about me?" print(f"User: {query}") print(f"Agent: {await agent.run(query, thread=resumed_thread)}\n") diff --git a/python/samples/getting_started/threads/redis_chat_message_store_thread.py b/python/samples/getting_started/threads/redis_chat_message_store_thread.py index 30d647fd40..d4b8f03c26 100644 --- a/python/samples/getting_started/threads/redis_chat_message_store_thread.py +++ b/python/samples/getting_started/threads/redis_chat_message_store_thread.py @@ -8,6 +8,14 @@ from agent_framework import AgentThread from agent_framework.openai import OpenAIChatClient from agent_framework.redis import RedisChatMessageStore +""" +Redis Chat Message Store Thread Example + +This sample demonstrates how to use Redis as a chat message store for thread +management, enabling persistent conversation history storage across sessions +with Redis as the backend data store. +""" + async def example_manual_memory_store() -> None: """Basic example of using Redis chat message store.""" diff --git a/python/samples/getting_started/threads/suspend_resume_thread.py b/python/samples/getting_started/threads/suspend_resume_thread.py index 88065aa987..99880fe8ab 100644 --- a/python/samples/getting_started/threads/suspend_resume_thread.py +++ b/python/samples/getting_started/threads/suspend_resume_thread.py @@ -2,38 +2,51 @@ import asyncio +from agent_framework.azure import AzureAIAgentClient from agent_framework.openai import OpenAIChatClient +from azure.identity.aio import AzureCliCredential + +""" +Thread Suspend and Resume Example + +This sample demonstrates how to suspend and resume conversation threads, comparing +service-managed threads (Azure AI) with in-memory threads (OpenAI) for persistent +conversation state across sessions. +""" async def suspend_resume_service_managed_thread() -> None: """Demonstrates how to suspend and resume a service-managed thread.""" print("=== Suspend-Resume Service-Managed Thread ===") - # OpenAI Chat Client is used as an example here, - # other chat clients can be used as well. - agent = OpenAIChatClient().create_agent(name="Joker", instructions="You are good at telling jokes.") + # AzureAIAgentClient supports service-managed threads. + async with ( + AzureCliCredential() as credential, + AzureAIAgentClient(async_credential=credential).create_agent( + name="MemoryBot", instructions="You are a helpful assistant that remembers our conversation." + ) as agent, + ): + # Start a new thread for the agent conversation. + thread = agent.get_new_thread() - # Start a new thread for the agent conversation. - thread = agent.get_new_thread() + # Respond to user input. + query = "Hello! My name is Alice and I love pizza." + print(f"User: {query}") + print(f"Agent: {await agent.run(query, thread=thread)}\n") - # Respond to user input. - query = "Tell me a joke about a pirate." - print(f"User: {query}") - print(f"Agent: {await agent.run(query, thread=thread)}\n") + # Serialize the thread state, so it can be stored for later use. + serialized_thread = await thread.serialize() - # Serialize the thread state, so it can be stored for later use. - serialized_thread = await thread.serialize() + # The thread can now be saved to a database, file, or any other storage mechanism and loaded again later. + print(f"Serialized thread: {serialized_thread}\n") - # The thread can now be saved to a database, file, or any other storage mechanism and loaded again later. - print(f"Serialized thread: {serialized_thread}\n") + # Deserialize the thread state after loading from storage. + resumed_thread = await agent.deserialize_thread(serialized_thread) - # Deserialize the thread state after loading from storage. - resumed_thread = await agent.deserialize_thread(serialized_thread) - - # Respond to user input. - query = "Now tell the same joke in the voice of a pirate, and add some emojis to the joke." - print(f"User: {query}") - print(f"Agent: {await agent.run(query, thread=resumed_thread)}\n") + # Respond to user input. + query = "What do you remember about me?" + print(f"User: {query}") + print(f"Agent: {await agent.run(query, thread=resumed_thread)}\n") async def suspend_resume_in_memory_thread() -> None: @@ -42,13 +55,15 @@ async def suspend_resume_in_memory_thread() -> None: # OpenAI Chat Client is used as an example here, # other chat clients can be used as well. - agent = OpenAIChatClient().create_agent(name="Joker", instructions="You are good at telling jokes.") + agent = OpenAIChatClient().create_agent( + name="MemoryBot", instructions="You are a helpful assistant that remembers our conversation." + ) # Start a new thread for the agent conversation. thread = agent.get_new_thread() # Respond to user input. - query = "Tell me a joke about a pirate." + query = "Hello! My name is Alice and I love pizza." print(f"User: {query}") print(f"Agent: {await agent.run(query, thread=thread)}\n") @@ -62,7 +77,7 @@ async def suspend_resume_in_memory_thread() -> None: resumed_thread = await agent.deserialize_thread(serialized_thread) # Respond to user input. - query = "Now tell the same joke in the voice of a pirate, and add some emojis to the joke." + query = "What do you remember about me?" print(f"User: {query}") print(f"Agent: {await agent.run(query, thread=resumed_thread)}\n") From 708556e4ee9d4ad3eb1f6241451a224300286004 Mon Sep 17 00:00:00 2001 From: Evan Mattson <35585003+moonbox3@users.noreply.github.com> Date: Thu, 6 Nov 2025 12:21:54 +0900 Subject: [PATCH 03/16] Python: Update changelog with ag-ui changes (#1954) * Update changelog with ag-ui changes * Changed -> Fixed --- python/CHANGELOG.md | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/python/CHANGELOG.md b/python/CHANGELOG.md index 4e1971bab0..b4485b914c 100644 --- a/python/CHANGELOG.md +++ b/python/CHANGELOG.md @@ -7,6 +7,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +## [1.0.0b251106.post1] - 2025-11-06 + +### Fixed + +- **agent-framework-ag-ui**: Fix ag-ui examples packaging for PyPI publish ([#1953](https://github.com/microsoft/agent-framework/pull/1953)) + ## [1.0.0b251106] - 2025-11-06 ### Changed From 6e205445be51c07257e1ac96f4eff2757cc01300 Mon Sep 17 00:00:00 2001 From: westey <164392973+westey-m@users.noreply.github.com> Date: Thu, 6 Nov 2025 14:34:27 +0000 Subject: [PATCH 04/16] .NET: Add tool calling sample with OpenAPI (#1968) * Add tool calling sample with OpenAPI * Address PR comments. * Rename folders and moved literal to inline. * Fix broken link. --- dotnet/Directory.Packages.props | 8 +- dotnet/agent-framework-dotnet.slnx | 7 +- .../Agent_Step03.1_UsingFunctionTools.csproj} | 0 .../Program.cs | 0 ...03.2_UsingFunctionTools_FromOpenAPI.csproj | 28 ++ .../OpenAPISpec.json | 354 ++++++++++++++++++ .../Program.cs | 33 ++ .../samples/GettingStarted/Agents/README.md | 3 +- 8 files changed, 426 insertions(+), 7 deletions(-) rename dotnet/samples/GettingStarted/Agents/{Agent_Step03_UsingFunctionTools/Agent_Step03_UsingFunctionTools.csproj => Agent_Step03.1_UsingFunctionTools/Agent_Step03.1_UsingFunctionTools.csproj} (100%) rename dotnet/samples/GettingStarted/Agents/{Agent_Step03_UsingFunctionTools => Agent_Step03.1_UsingFunctionTools}/Program.cs (100%) create mode 100644 dotnet/samples/GettingStarted/Agents/Agent_Step03.2_UsingFunctionTools_FromOpenAPI/Agent_Step03.2_UsingFunctionTools_FromOpenAPI.csproj create mode 100644 dotnet/samples/GettingStarted/Agents/Agent_Step03.2_UsingFunctionTools_FromOpenAPI/OpenAPISpec.json create mode 100644 dotnet/samples/GettingStarted/Agents/Agent_Step03.2_UsingFunctionTools_FromOpenAPI/Program.cs diff --git a/dotnet/Directory.Packages.props b/dotnet/Directory.Packages.props index 7fb6459906..71eceb7f46 100644 --- a/dotnet/Directory.Packages.props +++ b/dotnet/Directory.Packages.props @@ -68,13 +68,15 @@ - + - - + + + + diff --git a/dotnet/agent-framework-dotnet.slnx b/dotnet/agent-framework-dotnet.slnx index 7cbe76b6fc..03f8a910d3 100644 --- a/dotnet/agent-framework-dotnet.slnx +++ b/dotnet/agent-framework-dotnet.slnx @@ -47,7 +47,8 @@ - + + @@ -154,8 +155,8 @@ - + @@ -309,8 +310,8 @@ - + diff --git a/dotnet/samples/GettingStarted/Agents/Agent_Step03_UsingFunctionTools/Agent_Step03_UsingFunctionTools.csproj b/dotnet/samples/GettingStarted/Agents/Agent_Step03.1_UsingFunctionTools/Agent_Step03.1_UsingFunctionTools.csproj similarity index 100% rename from dotnet/samples/GettingStarted/Agents/Agent_Step03_UsingFunctionTools/Agent_Step03_UsingFunctionTools.csproj rename to dotnet/samples/GettingStarted/Agents/Agent_Step03.1_UsingFunctionTools/Agent_Step03.1_UsingFunctionTools.csproj diff --git a/dotnet/samples/GettingStarted/Agents/Agent_Step03_UsingFunctionTools/Program.cs b/dotnet/samples/GettingStarted/Agents/Agent_Step03.1_UsingFunctionTools/Program.cs similarity index 100% rename from dotnet/samples/GettingStarted/Agents/Agent_Step03_UsingFunctionTools/Program.cs rename to dotnet/samples/GettingStarted/Agents/Agent_Step03.1_UsingFunctionTools/Program.cs diff --git a/dotnet/samples/GettingStarted/Agents/Agent_Step03.2_UsingFunctionTools_FromOpenAPI/Agent_Step03.2_UsingFunctionTools_FromOpenAPI.csproj b/dotnet/samples/GettingStarted/Agents/Agent_Step03.2_UsingFunctionTools_FromOpenAPI/Agent_Step03.2_UsingFunctionTools_FromOpenAPI.csproj new file mode 100644 index 0000000000..e2edbb2f8d --- /dev/null +++ b/dotnet/samples/GettingStarted/Agents/Agent_Step03.2_UsingFunctionTools_FromOpenAPI/Agent_Step03.2_UsingFunctionTools_FromOpenAPI.csproj @@ -0,0 +1,28 @@ + + + + Exe + net9.0 + + enable + enable + + + + + + + + + + + + + + + + PreserveNewest + + + + diff --git a/dotnet/samples/GettingStarted/Agents/Agent_Step03.2_UsingFunctionTools_FromOpenAPI/OpenAPISpec.json b/dotnet/samples/GettingStarted/Agents/Agent_Step03.2_UsingFunctionTools_FromOpenAPI/OpenAPISpec.json new file mode 100644 index 0000000000..84715914da --- /dev/null +++ b/dotnet/samples/GettingStarted/Agents/Agent_Step03.2_UsingFunctionTools_FromOpenAPI/OpenAPISpec.json @@ -0,0 +1,354 @@ +{ + "openapi": "3.0.1", + "info": { + "title": "Github Versions API", + "version": "1.0.0" + }, + "servers": [ + { + "url": "https://api.github.com" + } + ], + "components": { + "schemas": { + "basic-error": { + "title": "Basic Error", + "description": "Basic Error", + "type": "object", + "properties": { + "message": { + "type": "string" + }, + "documentation_url": { + "type": "string" + }, + "url": { + "type": "string" + }, + "status": { + "type": "string" + } + } + }, + "label": { + "title": "Label", + "description": "Color-coded labels help you categorize and filter your issues (just like labels in Gmail).", + "type": "object", + "properties": { + "id": { + "description": "Unique identifier for the label.", + "type": "integer", + "format": "int64", + "example": 208045946 + }, + "node_id": { + "type": "string", + "example": "MDU6TGFiZWwyMDgwNDU5NDY=" + }, + "url": { + "description": "URL for the label", + "example": "https://api.github.com/repositories/42/labels/bug", + "type": "string", + "format": "uri" + }, + "name": { + "description": "The name of the label.", + "example": "bug", + "type": "string" + }, + "description": { + "description": "Optional description of the label, such as its purpose.", + "type": "string", + "example": "Something isn't working", + "nullable": true + }, + "color": { + "description": "6-character hex code, without the leading #, identifying the color", + "example": "FFFFFF", + "type": "string" + }, + "default": { + "description": "Whether this label comes by default in a new repository.", + "type": "boolean", + "example": true + } + }, + "required": [ + "id", + "node_id", + "url", + "name", + "description", + "color", + "default" + ] + }, + "tag": { + "title": "Tag", + "description": "Tag", + "type": "object", + "properties": { + "name": { + "type": "string", + "example": "v0.1" + }, + "commit": { + "type": "object", + "properties": { + "sha": { + "type": "string" + }, + "url": { + "type": "string", + "format": "uri" + } + }, + "required": [ + "sha", + "url" + ] + }, + "zipball_url": { + "type": "string", + "format": "uri", + "example": "https://github.com/octocat/Hello-World/zipball/v0.1" + }, + "tarball_url": { + "type": "string", + "format": "uri", + "example": "https://github.com/octocat/Hello-World/tarball/v0.1" + }, + "node_id": { + "type": "string" + } + }, + "required": [ + "name", + "node_id", + "commit", + "zipball_url", + "tarball_url" + ] + } + }, + "examples": { + "label-items": { + "value": [ + { + "id": 208045946, + "node_id": "MDU6TGFiZWwyMDgwNDU5NDY=", + "url": "https://api.github.com/repos/octocat/Hello-World/labels/bug", + "name": "bug", + "description": "Something isn't working", + "color": "f29513", + "default": true + }, + { + "id": 208045947, + "node_id": "MDU6TGFiZWwyMDgwNDU5NDc=", + "url": "https://api.github.com/repos/octocat/Hello-World/labels/enhancement", + "name": "enhancement", + "description": "New feature or request", + "color": "a2eeef", + "default": false + } + ] + }, + "tag-items": { + "value": [ + { + "name": "v0.1", + "commit": { + "sha": "c5b97d5ae6c19d5c5df71a34c7fbeeda2479ccbc", + "url": "https://api.github.com/repos/octocat/Hello-World/commits/c5b97d5ae6c19d5c5df71a34c7fbeeda2479ccbc" + }, + "zipball_url": "https://github.com/octocat/Hello-World/zipball/v0.1", + "tarball_url": "https://github.com/octocat/Hello-World/tarball/v0.1", + "node_id": "MDQ6VXNlcjE=" + } + ] + } + }, + "parameters": { + "owner": { + "name": "owner", + "description": "The account owner of the repository. The name is not case sensitive.", + "in": "path", + "required": true, + "schema": { + "type": "string" + } + }, + "repo": { + "name": "repo", + "description": "The name of the repository without the `.git` extension. The name is not case sensitive.", + "in": "path", + "required": true, + "schema": { + "type": "string" + } + }, + "per-page": { + "name": "per_page", + "description": "The number of results per page (max 100). For more information, see \"[Using pagination in the REST API](https://docs.github.com/rest/using-the-rest-api/using-pagination-in-the-rest-api).\"", + "in": "query", + "schema": { + "type": "integer", + "default": 30 + } + }, + "page": { + "name": "page", + "description": "The page number of the results to fetch. For more information, see \"[Using pagination in the REST API](https://docs.github.com/rest/using-the-rest-api/using-pagination-in-the-rest-api).\"", + "in": "query", + "schema": { + "type": "integer", + "default": 1 + } + } + }, + "responses": { + "not_found": { + "description": "Resource not found", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/basic-error" + } + } + } + } + }, + "headers": { + "link": { + "example": "; rel=\"next\", ; rel=\"last\"", + "schema": { + "type": "string" + } + } + } + }, + "paths": { + "/repos/{owner}/{repo}/tags": { + "get": { + "summary": "List repository tags", + "description": "", + "tags": [ + "repos" + ], + "operationId": "repos/list-tags", + "externalDocs": { + "description": "API method documentation", + "url": "https://docs.github.com/rest/repos/repos#list-repository-tags" + }, + "parameters": [ + { + "$ref": "#/components/parameters/owner" + }, + { + "$ref": "#/components/parameters/repo" + }, + { + "$ref": "#/components/parameters/per-page" + }, + { + "$ref": "#/components/parameters/page" + } + ], + "responses": { + "200": { + "description": "Response", + "content": { + "application/json": { + "schema": { + "type": "array", + "items": { + "$ref": "#/components/schemas/tag" + } + }, + "examples": { + "default": { + "$ref": "#/components/examples/tag-items" + } + } + } + }, + "headers": { + "Link": { + "$ref": "#/components/headers/link" + } + } + } + }, + "x-github": { + "githubCloudOnly": false, + "enabledForGitHubApps": true, + "category": "repos", + "subcategory": "repos" + } + } + }, + "/repos/{owner}/{repo}/labels": { + "get": { + "summary": "List labels for a repository", + "description": "Lists all labels for a repository.", + "tags": [ + "issues" + ], + "operationId": "issues/list-labels-for-repo", + "externalDocs": { + "description": "API method documentation", + "url": "https://docs.github.com/rest/issues/labels#list-labels-for-a-repository" + }, + "parameters": [ + { + "$ref": "#/components/parameters/owner" + }, + { + "$ref": "#/components/parameters/repo" + }, + { + "$ref": "#/components/parameters/per-page" + }, + { + "$ref": "#/components/parameters/page" + } + ], + "responses": { + "200": { + "description": "Response", + "content": { + "application/json": { + "schema": { + "type": "array", + "items": { + "$ref": "#/components/schemas/label" + } + }, + "examples": { + "default": { + "$ref": "#/components/examples/label-items" + } + } + } + }, + "headers": { + "Link": { + "$ref": "#/components/headers/link" + } + } + }, + "404": { + "$ref": "#/components/responses/not_found" + } + }, + "x-github": { + "githubCloudOnly": false, + "enabledForGitHubApps": true, + "category": "issues", + "subcategory": "labels" + } + } + } + } +} \ No newline at end of file diff --git a/dotnet/samples/GettingStarted/Agents/Agent_Step03.2_UsingFunctionTools_FromOpenAPI/Program.cs b/dotnet/samples/GettingStarted/Agents/Agent_Step03.2_UsingFunctionTools_FromOpenAPI/Program.cs new file mode 100644 index 0000000000..e61c9f845a --- /dev/null +++ b/dotnet/samples/GettingStarted/Agents/Agent_Step03.2_UsingFunctionTools_FromOpenAPI/Program.cs @@ -0,0 +1,33 @@ +// Copyright (c) Microsoft. All rights reserved. + +// This sample demonstrates how to use a ChatClientAgent with function tools provided via an OpenAPI spec. +// It uses functionality from Semantic Kernel to parse the OpenAPI spec and create function tools to use with the Agent Framework Agent. + +using Azure.AI.OpenAI; +using Azure.Identity; +using Microsoft.Agents.AI; +using Microsoft.Extensions.AI; +using Microsoft.SemanticKernel; +using Microsoft.SemanticKernel.Plugins.OpenApi; +using OpenAI; + +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"; + +// Load the OpenAPI Spec from a file. +KernelPlugin plugin = await OpenApiKernelPluginFactory.CreateFromOpenApiAsync("github", "OpenAPISpec.json"); + +// Convert the Semantic Kernel plugin to Agent Framework function tools. +// This requires a dummy Kernel instance, since KernelFunctions cannot execute without one. +Kernel kernel = new(); +List tools = plugin.Select(x => x.WithKernel(kernel)).Cast().ToList(); + +// Create the chat client and agent, and provide the OpenAPI function tools to the agent. +AIAgent agent = new AzureOpenAIClient( + new Uri(endpoint), + new AzureCliCredential()) + .GetChatClient(deploymentName) + .CreateAIAgent(instructions: "You are a helpful assistant", tools: tools); + +// Run the agent with the OpenAPI function tools. +Console.WriteLine(await agent.RunAsync("Please list the names, colors and descriptions of all the labels available in the microsoft/agent-framework repository on github.")); diff --git a/dotnet/samples/GettingStarted/Agents/README.md b/dotnet/samples/GettingStarted/Agents/README.md index bf5c291f9e..562b6b2500 100644 --- a/dotnet/samples/GettingStarted/Agents/README.md +++ b/dotnet/samples/GettingStarted/Agents/README.md @@ -28,7 +28,8 @@ Before you begin, ensure you have the following prerequisites: |---|---| |[Running a simple agent](./Agent_Step01_Running/)|This sample demonstrates how to create and run a basic agent with instructions| |[Multi-turn conversation with a simple agent](./Agent_Step02_MultiturnConversation/)|This sample demonstrates how to implement a multi-turn conversation with a simple agent| -|[Using function tools with a simple agent](./Agent_Step03_UsingFunctionTools/)|This sample demonstrates how to use function tools with a simple agent| +|[Using function tools with a simple agent](./Agent_Step03.1_UsingFunctionTools/)|This sample demonstrates how to use function tools with a simple agent| +|[Using OpenAPI function tools with a simple agent](./Agent_Step03.2_UsingFunctionTools_FromOpenAPI/)|This sample demonstrates how to create function tools from an OpenAPI spec and use them with a simple agent| |[Using function tools with approvals](./Agent_Step04_UsingFunctionToolsWithApprovals/)|This sample demonstrates how to use function tools where approvals require human in the loop approvals before execution| |[Structured output with a simple agent](./Agent_Step05_StructuredOutput/)|This sample demonstrates how to use structured output with a simple agent| |[Persisted conversations with a simple agent](./Agent_Step06_PersistedConversations/)|This sample demonstrates how to persist conversations and reload them later. This is useful for cases where an agent is hosted in a stateless service| From b374ff0e10aca0114277dd6a5467639138d3d223 Mon Sep 17 00:00:00 2001 From: Jeff Handley Date: Thu, 6 Nov 2025 14:11:00 -0800 Subject: [PATCH 05/16] Rename nightly build GitHub source in FAQS.md to be consistent (#1755) --- docs/FAQS.md | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/docs/FAQS.md b/docs/FAQS.md index 8363fd40d1..3ecd5514cc 100644 --- a/docs/FAQS.md +++ b/docs/FAQS.md @@ -23,23 +23,23 @@ To download nightly builds follow the following steps: - + - + - - - - + + + + ``` From 25f405c7ce72bcee2b1ca1aca8ed5788ef92c66d Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Thu, 6 Nov 2025 17:12:45 -0500 Subject: [PATCH 06/16] Bump CommunityToolkit.Aspire.OllamaSharp from 9.8.0 to 9.9.0 (#1961) --- updated-dependencies: - dependency-name: CommunityToolkit.Aspire.OllamaSharp dependency-version: 9.9.0 dependency-type: direct:production update-type: version-update:semver-minor ... Signed-off-by: dependabot[bot] Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> --- dotnet/Directory.Packages.props | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dotnet/Directory.Packages.props b/dotnet/Directory.Packages.props index 71eceb7f46..d54038ec79 100644 --- a/dotnet/Directory.Packages.props +++ b/dotnet/Directory.Packages.props @@ -15,7 +15,7 @@ - + From cb50f3e070d3fe35096570f776562ada42f01126 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Thu, 6 Nov 2025 22:13:25 +0000 Subject: [PATCH 07/16] Bump AWSSDK.Extensions.Bedrock.MEAI from 4.0.4.1 to 4.0.4.2 (#1707) --- updated-dependencies: - dependency-name: AWSSDK.Extensions.Bedrock.MEAI dependency-version: 4.0.4.2 dependency-type: direct:production update-type: version-update:semver-patch ... Signed-off-by: dependabot[bot] Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> Co-authored-by: Chris <66376200+crickman@users.noreply.github.com> --- dotnet/Directory.Packages.props | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dotnet/Directory.Packages.props b/dotnet/Directory.Packages.props index d54038ec79..d110ec4426 100644 --- a/dotnet/Directory.Packages.props +++ b/dotnet/Directory.Packages.props @@ -86,7 +86,7 @@ - + From 820c6afe0991035b835dd1d265554ba6e787c1be Mon Sep 17 00:00:00 2001 From: Stephen Toub Date: Thu, 6 Nov 2025 17:23:37 -0500 Subject: [PATCH 08/16] .NET: Fix the ordering of chained resolvers in JsonSerializerOptions (#1974) * Fix the ordering of chained resolvers in JsonSerializerOptions We want the resolvers from AIJsonUtilities to be used before the ones from the source generator, in case the source generator emits its own copy in that assembly for the M.E.AI types. * Apply suggestions from code review Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * Update dotnet/src/Microsoft.Agents.AI/AgentJsonUtilities.cs * Update dotnet/src/Microsoft.Agents.AI.Mem0/Mem0JsonUtilities.cs Co-authored-by: westey <164392973+westey-m@users.noreply.github.com> * Remove unused using directive in Mem0JsonUtilities Removed unused using directive for Microsoft.Extensions.AI. --------- Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> Co-authored-by: westey <164392973+westey-m@users.noreply.github.com> --- .../AgentAbstractionsJsonUtilities.cs | 8 ++++++-- .../ChatCompletionsJsonSerializerOptions.cs | 6 ++++++ .../OpenAIHostingJsonUtilities.cs | 6 ++++++ .../Mem0JsonUtilities.cs | 9 ++++++--- .../WorkflowsJsonUtilities.cs | 5 ++++- .../Microsoft.Agents.AI/AgentJsonUtilities.cs | 6 +++++- .../InMemoryChatMessageStoreTests.cs | 17 +++++++++++++++++ 7 files changed, 50 insertions(+), 7 deletions(-) diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions/AgentAbstractionsJsonUtilities.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions/AgentAbstractionsJsonUtilities.cs index 8943e29c79..d5003cace0 100644 --- a/dotnet/src/Microsoft.Agents.AI.Abstractions/AgentAbstractionsJsonUtilities.cs +++ b/dotnet/src/Microsoft.Agents.AI.Abstractions/AgentAbstractionsJsonUtilities.cs @@ -50,12 +50,16 @@ public static partial class AgentAbstractionsJsonUtilities Encoder = JavaScriptEncoder.UnsafeRelaxedJsonEscaping, // same as AIJsonUtilities }; - // Chain with all supported types from Microsoft.Extensions.AI.Abstractions. + // Chain in the resolvers from both AIJsonUtilities and our source generated context. + // We want AIJsonUtilities first to ensure any M.E.AI types are handled via its resolver. + options.TypeInfoResolverChain.Clear(); + options.TypeInfoResolverChain.Add(AIJsonUtilities.DefaultOptions.TypeInfoResolver!); + options.TypeInfoResolverChain.Add(JsonContext.Default.Options.TypeInfoResolver!); + // If reflection-based serialization is enabled by default, this includes // the default type info resolver that utilizes reflection, but we need to manually // apply the same converter AIJsonUtilities adds for string-based enum serialization, // as that's not propagated as part of the resolver. - options.TypeInfoResolverChain.Add(AIJsonUtilities.DefaultOptions.TypeInfoResolver!); if (JsonSerializer.IsReflectionEnabledByDefault) { options.Converters.Add(new JsonStringEnumConverter()); diff --git a/dotnet/src/Microsoft.Agents.AI.Hosting.OpenAI/ChatCompletions/ChatCompletionsJsonSerializerOptions.cs b/dotnet/src/Microsoft.Agents.AI.Hosting.OpenAI/ChatCompletions/ChatCompletionsJsonSerializerOptions.cs index b009b82d29..301cae1f8f 100644 --- a/dotnet/src/Microsoft.Agents.AI.Hosting.OpenAI/ChatCompletions/ChatCompletionsJsonSerializerOptions.cs +++ b/dotnet/src/Microsoft.Agents.AI.Hosting.OpenAI/ChatCompletions/ChatCompletionsJsonSerializerOptions.cs @@ -17,7 +17,13 @@ internal static class ChatCompletionsJsonSerializerOptions private static JsonSerializerOptions Create() { JsonSerializerOptions options = new(ChatCompletionsJsonContext.Default.Options); + + // Chain in the resolvers from both AgentAbstractionsJsonUtilities and our source generated context. + // We want AgentAbstractionsJsonUtilities first to ensure any M.E.AI types are handled via its resolver. + options.TypeInfoResolverChain.Clear(); options.TypeInfoResolverChain.Add(AgentAbstractionsJsonUtilities.DefaultOptions.TypeInfoResolver!); + options.TypeInfoResolverChain.Add(ChatCompletionsJsonContext.Default.Options.TypeInfoResolver!); + options.MakeReadOnly(); return options; } diff --git a/dotnet/src/Microsoft.Agents.AI.Hosting.OpenAI/OpenAIHostingJsonUtilities.cs b/dotnet/src/Microsoft.Agents.AI.Hosting.OpenAI/OpenAIHostingJsonUtilities.cs index ceac8b872f..49ceef622a 100644 --- a/dotnet/src/Microsoft.Agents.AI.Hosting.OpenAI/OpenAIHostingJsonUtilities.cs +++ b/dotnet/src/Microsoft.Agents.AI.Hosting.OpenAI/OpenAIHostingJsonUtilities.cs @@ -24,7 +24,13 @@ internal static class OpenAIHostingJsonUtilities private static JsonSerializerOptions CreateDefaultOptions() { JsonSerializerOptions options = new(OpenAIHostingJsonContext.Default.Options); + + // Chain in the resolvers from both AgentAbstractionsJsonUtilities and our source generated context. + // We want AgentAbstractionsJsonUtilities first to ensure any M.E.AI types are handled via its resolver. + options.TypeInfoResolverChain.Clear(); options.TypeInfoResolverChain.Add(AgentAbstractionsJsonUtilities.DefaultOptions.TypeInfoResolver!); + options.TypeInfoResolverChain.Add(OpenAIHostingJsonContext.Default.Options.TypeInfoResolver!); + options.MakeReadOnly(); return options; } diff --git a/dotnet/src/Microsoft.Agents.AI.Mem0/Mem0JsonUtilities.cs b/dotnet/src/Microsoft.Agents.AI.Mem0/Mem0JsonUtilities.cs index eb50d31f70..d139cb0f76 100644 --- a/dotnet/src/Microsoft.Agents.AI.Mem0/Mem0JsonUtilities.cs +++ b/dotnet/src/Microsoft.Agents.AI.Mem0/Mem0JsonUtilities.cs @@ -4,7 +4,6 @@ using System.Diagnostics.CodeAnalysis; using System.Text.Encodings.Web; using System.Text.Json; using System.Text.Json.Serialization; -using Microsoft.Extensions.AI; namespace Microsoft.Agents.AI.Mem0; @@ -44,8 +43,12 @@ public static partial class Mem0JsonUtilities Encoder = JavaScriptEncoder.UnsafeRelaxedJsonEscaping, // same as in AIJsonUtilities }; - // Chain with all supported types from Microsoft.Extensions.AI.Abstractions. - options.TypeInfoResolverChain.Add(AIJsonUtilities.DefaultOptions.TypeInfoResolver!); + // Chain in the resolvers from both AgentAbstractionsJsonUtilities and our source generated context. + // We want AgentAbstractionsJsonUtilities first to ensure any M.E.AI types are handled via its resolver. + options.TypeInfoResolverChain.Clear(); + options.TypeInfoResolverChain.Add(AgentAbstractionsJsonUtilities.DefaultOptions.TypeInfoResolver!); + options.TypeInfoResolverChain.Add(JsonContext.Default.Options.TypeInfoResolver!); + if (JsonSerializer.IsReflectionEnabledByDefault) { options.Converters.Add(new JsonStringEnumConverter()); diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowsJsonUtilities.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowsJsonUtilities.cs index 752bb4bac7..d8241f4681 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowsJsonUtilities.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowsJsonUtilities.cs @@ -50,8 +50,11 @@ internal static partial class WorkflowsJsonUtilities // Copy the configuration from the source generated context. JsonSerializerOptions options = new(JsonContext.Default.Options); - // Chain with all supported types from Microsoft.Extensions.AI.Abstractions and Microsoft.Agents.AI.Abstractions. + // Chain in the resolvers from both AgentAbstractionsJsonUtilities and our source generated context. + // We want AgentAbstractionsJsonUtilities first to ensure any M.E.AI types are handled via its resolver. + options.TypeInfoResolverChain.Clear(); options.TypeInfoResolverChain.Add(AgentAbstractionsJsonUtilities.DefaultOptions.TypeInfoResolver!); + options.TypeInfoResolverChain.Add(JsonContext.Default.Options.TypeInfoResolver!); options.MakeReadOnly(); return options; diff --git a/dotnet/src/Microsoft.Agents.AI/AgentJsonUtilities.cs b/dotnet/src/Microsoft.Agents.AI/AgentJsonUtilities.cs index 5ef6978c00..36bef4a2af 100644 --- a/dotnet/src/Microsoft.Agents.AI/AgentJsonUtilities.cs +++ b/dotnet/src/Microsoft.Agents.AI/AgentJsonUtilities.cs @@ -44,8 +44,12 @@ internal static partial class AgentJsonUtilities Encoder = JavaScriptEncoder.UnsafeRelaxedJsonEscaping, // same as in AgentAbstractionsJsonUtilities and AIJsonUtilities }; - // Chain with all supported types from Microsoft.Agents.AI.Abstractions. + // Chain in the resolvers from both AgentAbstractionsJsonUtilities and our source generated context. + // We want AgentAbstractionsJsonUtilities first to ensure any M.E.AI types are handled via its resolver. + options.TypeInfoResolverChain.Clear(); options.TypeInfoResolverChain.Add(AgentAbstractionsJsonUtilities.DefaultOptions.TypeInfoResolver!); + options.TypeInfoResolverChain.Add(JsonContext.Default.Options.TypeInfoResolver!); + if (JsonSerializer.IsReflectionEnabledByDefault) { options.Converters.Add(new JsonStringEnumConverter()); diff --git a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/InMemoryChatMessageStoreTests.cs b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/InMemoryChatMessageStoreTests.cs index fedd0ce591..4c793d17f4 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/InMemoryChatMessageStoreTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/InMemoryChatMessageStoreTests.cs @@ -114,6 +114,23 @@ public class InMemoryChatMessageStoreTests Assert.Equal("B", newStore[1].Text); } + [Fact] + public async Task SerializeAndDeserializeWorksWithExperimentalContentTypesAsync() + { + var store = new InMemoryChatMessageStore + { + 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"))]) + }; + + var jsonElement = store.Serialize(); + var newStore = new InMemoryChatMessageStore(jsonElement); + + Assert.Equal(2, newStore.Count); + Assert.IsType(newStore[0].Contents[0]); + Assert.IsType(newStore[1].Contents[0]); + } + [Fact] public async Task AddMessagesAsyncWithEmptyMessagesDoesNotChangeStoreAsync() { From 64826b8f5650aaf9997b9bcd3c5f917714c9529a Mon Sep 17 00:00:00 2001 From: Rishabh Chawla Date: Thu, 6 Nov 2025 23:43:22 -0800 Subject: [PATCH 09/16] Python: [Purview] Add Caching and background processing in Python Purview Middleware (#1844) * [PythonPurview] Add Caching and background processing * [PythonPurview] Updates based on comments --- .../agent_framework/microsoft/__init__.py | 2 + .../agent_framework/microsoft/__init__.pyi | 4 + python/packages/purview/README.md | 197 +++++++++--- .../agent_framework_purview/__init__.py | 4 + .../purview/agent_framework_purview/_cache.py | 191 ++++++++++++ .../agent_framework_purview/_client.py | 87 +++++- .../agent_framework_purview/_exceptions.py | 5 + .../agent_framework_purview/_middleware.py | 48 ++- .../agent_framework_purview/_models.py | 43 ++- .../agent_framework_purview/_processor.py | 112 ++++++- .../agent_framework_purview/_settings.py | 23 +- python/packages/purview/tests/test_cache.py | 196 ++++++++++++ .../purview/tests/test_chat_middleware.py | 135 ++++++++- python/packages/purview/tests/test_client.py | 157 ++++++++++ .../packages/purview/tests/test_exceptions.py | 7 + .../packages/purview/tests/test_middleware.py | 59 +++- python/packages/purview/tests/test_models.py | 7 +- .../packages/purview/tests/test_processor.py | 285 +++++++++++++++++- .../packages/purview/tests/test_settings.py | 3 - .../getting_started/purview_agent/README.md | 68 ++++- .../purview_agent/sample_purview_agent.py | 148 ++++++++- 21 files changed, 1657 insertions(+), 124 deletions(-) create mode 100644 python/packages/purview/agent_framework_purview/_cache.py create mode 100644 python/packages/purview/tests/test_cache.py diff --git a/python/packages/core/agent_framework/microsoft/__init__.py b/python/packages/core/agent_framework/microsoft/__init__.py index 2874488829..106facf54c 100644 --- a/python/packages/core/agent_framework/microsoft/__init__.py +++ b/python/packages/core/agent_framework/microsoft/__init__.py @@ -14,9 +14,11 @@ _IMPORTS: dict[str, tuple[str, list[str]]] = { "PurviewAppLocation": ("agent_framework_purview", ["microsoft-purview", "purview"]), "PurviewLocationType": ("agent_framework_purview", ["microsoft-purview", "purview"]), "PurviewAuthenticationError": ("agent_framework_purview", ["microsoft-purview", "purview"]), + "PurviewPaymentRequiredError": ("agent_framework_purview", ["microsoft-purview", "purview"]), "PurviewRateLimitError": ("agent_framework_purview", ["microsoft-purview", "purview"]), "PurviewRequestError": ("agent_framework_purview", ["microsoft-purview", "purview"]), "PurviewServiceError": ("agent_framework_purview", ["microsoft-purview", "purview"]), + "CacheProvider": ("agent_framework_purview", ["microsoft-purview", "purview"]), } diff --git a/python/packages/core/agent_framework/microsoft/__init__.pyi b/python/packages/core/agent_framework/microsoft/__init__.pyi index 99ba2af489..f3c5c27a0e 100644 --- a/python/packages/core/agent_framework/microsoft/__init__.pyi +++ b/python/packages/core/agent_framework/microsoft/__init__.pyi @@ -2,10 +2,12 @@ from agent_framework_copilotstudio import CopilotStudioAgent, __version__, acquire_token from agent_framework_purview import ( + CacheProvider, PurviewAppLocation, PurviewAuthenticationError, PurviewChatPolicyMiddleware, PurviewLocationType, + PurviewPaymentRequiredError, PurviewPolicyMiddleware, PurviewRateLimitError, PurviewRequestError, @@ -14,11 +16,13 @@ from agent_framework_purview import ( ) __all__ = [ + "CacheProvider", "CopilotStudioAgent", "PurviewAppLocation", "PurviewAuthenticationError", "PurviewChatPolicyMiddleware", "PurviewLocationType", + "PurviewPaymentRequiredError", "PurviewPolicyMiddleware", "PurviewRateLimitError", "PurviewRequestError", diff --git a/python/packages/purview/README.md b/python/packages/purview/README.md index b52fdca807..4e6690bf31 100644 --- a/python/packages/purview/README.md +++ b/python/packages/purview/README.md @@ -10,20 +10,45 @@ - Blocks or allows content at both ingress (prompt) and egress (response) - Works with any `ChatAgent` / agent orchestration using the standard Agent Framework middleware pipeline - Supports both synchronous `TokenCredential` and `AsyncTokenCredential` from `azure-identity` -- Simple, typed configuration via `PurviewSettings` / `PurviewAppLocation` -- Two middleware types: - - `PurviewPolicyMiddleware` (Agent pipeline) - - `PurviewChatPolicyMiddleware` (Chat client middleware list) +- Configuration via `PurviewSettings` / `PurviewAppLocation` +- Built-in caching with configurable TTL and size limits for protection scopes in `PurviewSettings` +- Background processing for content activities and offline policy evaluation ### When to Use Add Purview when you need to: +- **Prevent sensitive data leaks**: Inline blocking of sensitive content based on Data Loss Prevention (DLP) policies. +- **Enable governance**: Log AI interactions in Purview for Audit, Communication Compliance, Insider Risk Management, eDiscovery, and Data Lifecycle Management. - Prevent sensitive or disallowed content from being sent to an LLM - Prevent model output containing disallowed data from leaving the system - Apply centrally managed policies without rewriting agent logic --- +## Prerequisites + +- Microsoft Azure subscription with Microsoft Purview configured. +- Microsoft 365 subscription with an E5 license and pay-as-you-go billing setup. + - For testing, you can use a Microsoft 365 Developer Program tenant. For more information, see [Join the Microsoft 365 Developer Program](https://developer.microsoft.com/en-us/microsoft-365/dev-program). + +### Authentication + +`PurviewClient` uses the `azure-identity` library for token acquisition. You can use any `TokenCredential` or `AsyncTokenCredential` implementation. + +- **Entra registration**: Register your agent and add the required Microsoft Graph permissions (`dataSecurityAndGovernance`) to the Service Principal. For more information, see [Register an application in Microsoft Entra ID](https://learn.microsoft.com/en-us/entra/identity-platform/quickstart-register-app) and [dataSecurityAndGovernance resource type](https://learn.microsoft.com/en-us/graph/api/resources/datasecurityandgovernance). You'll need the Microsoft Entra app ID in the next step. + +- **Graph Permissions**: +- ProtectionScopes.Compute.All : [userProtectionScopeContainer](https://learn.microsoft.com/en-us/graph/api/userprotectionscopecontainer-compute) +- Content.Process.All : [processContent](https://learn.microsoft.com/en-us/graph/api/userdatasecurityandgovernance-processcontent) +- ContentActivity.Write : [contentActivity](https://learn.microsoft.com/en-us/graph/api/activitiescontainer-post-contentactivities) + +- **Purview policies**: Configure Purview policies using the Microsoft Entra app ID to enable agent communications data to flow into Purview. For more information, see [Configure Microsoft Purview](https://learn.microsoft.com/purview/developer/configurepurview). + +#### Scopes +`PurviewSettings.get_scopes()` derives the Graph scope list (currently `https://graph.microsoft.com/.default` style). + +--- + ## Quick Start ```python @@ -57,37 +82,75 @@ If a policy violation is detected on the prompt, the middleware terminates the r --- -## Authentication - -`PurviewClient` uses the `azure-identity` library for token acquisition. You can use any `TokenCredential` or `AsyncTokenCredential` implementation. - -The APIs require the following Graph Permissions: -- ProtectionScopes.Compute.All : [userProtectionScopeContainer](https://learn.microsoft.com/en-us/graph/api/userprotectionscopecontainer-compute) -- Content.Process.All : [processContent](https://learn.microsoft.com/en-us/graph/api/userdatasecurityandgovernance-processcontent) -- ContentActivity.Write : [contentActivity](https://learn.microsoft.com/en-us/graph/api/activitiescontainer-post-contentactivities) - -### Scopes -`PurviewSettings.get_scopes()` derives the Graph scope list (currently `https://graph.microsoft.com/.default` style). - -### Tenant Enablement for Purview -- The tenant requires an e5 license and consumptive billing setup. -- There need to be [Data Loss Prevention](https://learn.microsoft.com/en-us/purview/dlp-create-deploy-policy) or [Data Collection Policies](https://learn.microsoft.com/en-us/purview/collection-policies-policy-reference) that apply to the user to call Process Content API else it calls Content Activities API for auditing the message. - ---- - ## Configuration ### `PurviewSettings` ```python PurviewSettings( - app_name="My App", # Display / logical name - tenant_id=None, # Optional – used mainly for auth context - purview_app_location=None, # Optional PurviewAppLocation for scoping + app_name="My App", # Required: Display / logical name + app_version=None, # Optional: Version string of the application + tenant_id=None, # Optional: Tenant id (guid), used mainly for auth context + purview_app_location=None, # Optional: PurviewAppLocation for scoping graph_base_uri="https://graph.microsoft.com/v1.0/", - process_inline=False, # Reserved for future inline processing optimizations - blocked_prompt_message="Prompt blocked by policy", # Custom message for blocked prompts - blocked_response_message="Response blocked by policy" # Custom message for blocked responses + blocked_prompt_message="Prompt blocked by policy", # Custom message for blocked prompts + blocked_response_message="Response blocked by policy", # Custom message for blocked responses + ignore_exceptions=False, # If True, non-payment exceptions are logged but not thrown + ignore_payment_required=False, # If True, 402 payment required errors are logged but not thrown + cache_ttl_seconds=14400, # Cache TTL in seconds (default 4 hours) + max_cache_size_bytes=200 * 1024 * 1024 # Max cache size in bytes (default 200MB) +) +``` + +### Caching + +The Purview integration includes built-in caching for protection scopes responses to improve performance and reduce API calls: + +- **Default TTL**: 4 hours (14400 seconds) +- **Default Cache Size**: 200MB +- **Cache Provider**: `InMemoryCacheProvider` is used by default, but you can provide a custom implementation via the `CacheProvider` protocol +- **Cache Invalidation**: Cache is automatically invalidated when protection scope state is modified +- **Exception Caching**: 402 Payment Required errors are cached to avoid repeated failed API calls + +You can customize caching behavior in `PurviewSettings`: + +```python +from agent_framework.microsoft import PurviewSettings + +settings = PurviewSettings( + app_name="My App", + cache_ttl_seconds=14400, # 4 hours + max_cache_size_bytes=200 * 1024 * 1024 # 200MB +) +``` + +Or provide your own cache provider: + +```python +from typing import Any +from agent_framework.microsoft import PurviewPolicyMiddleware, PurviewSettings, CacheProvider +from azure.identity import DefaultAzureCredential + +class MyCustomCache(CacheProvider): + async def get(self, key: str) -> Any | None: + # Your implementation + pass + + async def set(self, key: str, value: Any, ttl_seconds: int | None = None) -> None: + # Your implementation + pass + + async def remove(self, key: str) -> None: + # Your implementation + pass + +credential = DefaultAzureCredential() +settings = PurviewSettings(app_name="MyApp") + +middleware = PurviewPolicyMiddleware( + credential=credential, + settings=settings, + cache_provider=MyCustomCache() ) ``` @@ -123,11 +186,32 @@ settings = PurviewSettings( ) ``` -This is useful for: -- Providing more user-friendly error messages -- Including support contact information -- Localizing messages for different languages -- Adding branding or specific guidance for your application +### Exception Handling Controls + +The Purview integration provides fine-grained control over exception handling to support graceful degradation scenarios: + +```python +from agent_framework.microsoft import PurviewSettings + +# Ignore all non-payment exceptions (continue execution even if policy check fails) +settings = PurviewSettings( + app_name="My App", + ignore_exceptions=True # Log errors but don't throw +) + +# Ignore only 402 Payment Required errors (useful for tenants without proper licensing) +settings = PurviewSettings( + app_name="My App", + ignore_payment_required=True # Continue even without Purview Consumptive Billing Setup +) + +# Both can be combined +settings = PurviewSettings( + app_name="My App", + ignore_exceptions=True, + ignore_payment_required=True +) +``` ### Selecting Agent vs Chat Middleware @@ -178,12 +262,17 @@ The policy logic is identical; the difference is only the hook point in the pipe ## Middleware Lifecycle -1. Before agent execution (`prompt phase`): all `context.messages` are evaluated. -2. If blocked: `context.result` is replaced with a system message and `context.terminate = True`. -3. After successful agent execution (`response phase`): the produced messages are evaluated. -4. If blocked: result messages are replaced with a blocking notice. +1. **Before agent execution** (`prompt phase`): all `context.messages` are evaluated. + - If no valid user_id is found, processing is skipped (no policy evaluation) + - Protection scopes are retrieved (with caching) + - Applicable scopes are checked to determine execution mode + - In inline mode: content is evaluated immediately + - In offline mode: evaluation is queued in background +2. **If blocked**: `context.result` is replaced with a system message and `context.terminate = True`. +3. **After successful agent execution** (`response phase`): the produced messages are evaluated using the same user_id from the prompt phase. +4. **If blocked**: result messages are replaced with a blocking notice. -When a user identifier is discovered (e.g. in `ChatMessage.additional_properties['user_id']`) during the prompt phase it is reused for the response phase so both evaluations map consistently to the same user. +The user identifier is discovered from `ChatMessage.additional_properties['user_id']` during the prompt phase and reused for the response phase, ensuring both evaluations map consistently to the same user. If no user_id is present, policy evaluation is skipped entirely. You can customize the blocking messages using the `blocked_prompt_message` and `blocked_response_message` fields in `PurviewSettings`. For more advanced scenarios, you can wrap the middleware or post-process `context.result` in later middleware. @@ -193,32 +282,44 @@ You can customize the blocking messages using the `blocked_prompt_message` and ` | Exception | Scenario | |-----------|----------| +| `PurviewPaymentRequiredError` | 402 Payment Required - tenant lacks proper Purview licensing or consumptive billing setup | | `PurviewAuthenticationError` | Token acquisition / validation issues | | `PurviewRateLimitError` | 429 responses from service | | `PurviewRequestError` | 4xx client errors (bad input, unauthorized, forbidden) | | `PurviewServiceError` | 5xx or unexpected service errors | -Catch broadly if you want unified fallback: +### Exception Handling + +All exceptions inherit from `PurviewServiceError`. You can catch specific exceptions or use the base class: ```python from agent_framework.microsoft import ( - PurviewAuthenticationError, PurviewRateLimitError, - PurviewRequestError, PurviewServiceError + PurviewPaymentRequiredError, + PurviewAuthenticationError, + PurviewRateLimitError, + PurviewRequestError, + PurviewServiceError ) try: - ... + # Your code here + pass +except PurviewPaymentRequiredError as ex: + # Handle licensing issues specifically + print(f"Purview licensing required: {ex}") except (PurviewAuthenticationError, PurviewRateLimitError, PurviewRequestError, PurviewServiceError) as ex: - # Log / degrade gracefully - print(f"Purview enforcement skipped: {ex}") + # Handle other errors + print(f"Purview enforcement skipped: {ex}") ``` --- ## Notes -- Provide a `user_id` per request (e.g. in `ChatMessage(..., additional_properties={"user_id": ""})`) when possible for per-user policy scoping; otherwise supply a default via settings or environment. -- Blocking messages can be customized via `blocked_prompt_message` and `blocked_response_message` in `PurviewSettings`. By default, they are "Prompt blocked by policy" and "Response blocked by policy" respectively. -- Streaming responses: post-response policy evaluation presently applies only to non-streaming chat responses. -- Errors during policy checks are logged and do not fail the run; they degrade gracefully. +- **User Identification**: Provide a `user_id` per request (e.g. in `ChatMessage(..., additional_properties={"user_id": ""})`) for per-user policy scoping. If no user_id is provided, policy evaluation is skipped entirely. +- **Blocking Messages**: Can be customized via `blocked_prompt_message` and `blocked_response_message` in `PurviewSettings`. By default, they are "Prompt blocked by policy" and "Response blocked by policy" respectively. +- **Streaming Responses**: Post-response policy evaluation presently applies only to non-streaming chat responses. +- **Error Handling**: Use `ignore_exceptions` and `ignore_payment_required` settings for graceful degradation. When enabled, errors are logged but don't fail the request. +- **Caching**: Protection scopes responses and 402 errors are cached by default with a 4-hour TTL. Cache is automatically invalidated when protection scope state changes. +- **Background Processing**: Content Activities and offline Process Content requests are handled asynchronously using background tasks to avoid blocking the main execution flow. diff --git a/python/packages/purview/agent_framework_purview/__init__.py b/python/packages/purview/agent_framework_purview/__init__.py index 74e09164f7..79722f1b50 100644 --- a/python/packages/purview/agent_framework_purview/__init__.py +++ b/python/packages/purview/agent_framework_purview/__init__.py @@ -1,7 +1,9 @@ # Copyright (c) Microsoft. All rights reserved. +from ._cache import CacheProvider from ._exceptions import ( PurviewAuthenticationError, + PurviewPaymentRequiredError, PurviewRateLimitError, PurviewRequestError, PurviewServiceError, @@ -10,10 +12,12 @@ from ._middleware import PurviewChatPolicyMiddleware, PurviewPolicyMiddleware from ._settings import PurviewAppLocation, PurviewLocationType, PurviewSettings __all__ = [ + "CacheProvider", "PurviewAppLocation", "PurviewAuthenticationError", "PurviewChatPolicyMiddleware", "PurviewLocationType", + "PurviewPaymentRequiredError", "PurviewPolicyMiddleware", "PurviewRateLimitError", "PurviewRequestError", diff --git a/python/packages/purview/agent_framework_purview/_cache.py b/python/packages/purview/agent_framework_purview/_cache.py new file mode 100644 index 0000000000..5f9f6fe5bb --- /dev/null +++ b/python/packages/purview/agent_framework_purview/_cache.py @@ -0,0 +1,191 @@ +# Copyright (c) Microsoft. All rights reserved. +"""Cache provider for Purview data.""" + +from __future__ import annotations + +import hashlib +import heapq +import json +import sys +import time +from typing import Any, Protocol + +from ._models import ProtectionScopesRequest + + +class CacheProvider(Protocol): + """Protocol for cache providers used by Purview integration.""" + + async def get(self, key: str) -> Any | None: + """Get a value from the cache. + + Args: + key: The cache key. + + Returns: + The cached value or None if not found or expired. + """ + ... + + async def set(self, key: str, value: Any, ttl_seconds: int | None = None) -> None: + """Set a value in the cache. + + Args: + key: The cache key. + value: The value to cache. + ttl_seconds: Time to live in seconds. If None, uses provider default. + """ + ... + + async def remove(self, key: str) -> None: + """Remove a value from the cache. + + Args: + key: The cache key. + """ + ... + + +class InMemoryCacheProvider: + """Simple in-memory cache implementation for Purview data. + + This implementation uses a dictionary with expiration tracking and size limits. + """ + + def __init__(self, default_ttl_seconds: int = 1800, max_size_bytes: int = 200 * 1024 * 1024): + """Initialize the in-memory cache. + + Args: + default_ttl_seconds: Default time to live in seconds (default 1800 = 30 minutes). + max_size_bytes: Maximum cache size in bytes (default 200MB). + """ + self._cache: dict[str, tuple[Any, float, int]] = {} # key -> (value, expiry, size) + self._expiry_heap: list[tuple[float, str]] = [] # min-heap of (expiry_time, key) + self._default_ttl = default_ttl_seconds + self._max_size_bytes = max_size_bytes + self._current_size_bytes = 0 + + def _estimate_size(self, value: Any) -> int: + """Estimate the size of a cached value in bytes. + + Args: + value: The value to estimate size for. + + Returns: + Estimated size in bytes. + """ + try: + if hasattr(value, "model_dump_json"): + return len(value.model_dump_json().encode("utf-8")) + + return len(json.dumps(value, default=str).encode("utf-8")) + except Exception: + # Fallback to sys.getsizeof if JSON serialization fails + try: + return sys.getsizeof(value) + except Exception: + # Conservative fallback estimate + return 1024 + + def _evict_if_needed(self, required_size: int) -> None: + """Evict oldest entries if needed to make room for new entry. + + Uses a min-heap to efficiently find and evict entries with earliest expiry times. + Also cleans up stale heap entries for keys that no longer exist in cache. + + Args: + required_size: Size in bytes needed for new entry. + """ + if self._current_size_bytes + required_size <= self._max_size_bytes: + return + + while self._expiry_heap and self._current_size_bytes + required_size > self._max_size_bytes: + expiry_time, key = heapq.heappop(self._expiry_heap) + + if key in self._cache: + _, cached_expiry, size = self._cache[key] + if cached_expiry == expiry_time: + del self._cache[key] + self._current_size_bytes -= size + # else: stale heap entry, already updated/removed, skip it + + async def get(self, key: str) -> Any | None: + """Get a value from the cache. + + Args: + key: The cache key. + + Returns: + The cached value or None if not found or expired. + """ + if key not in self._cache: + return None + + value, expiry, size = self._cache[key] + if time.time() > expiry: + del self._cache[key] + self._current_size_bytes -= size + return None + + return value + + async def set(self, key: str, value: Any, ttl_seconds: int | None = None) -> None: + """Set a value in the cache. + + Args: + key: The cache key. + value: The value to cache. + ttl_seconds: Time to live in seconds. If None, uses default TTL. + """ + ttl = ttl_seconds if ttl_seconds is not None else self._default_ttl + expiry = time.time() + ttl + size = self._estimate_size(value) + + # Remove old entry if exists + if key in self._cache: + old_size = self._cache[key][2] + self._current_size_bytes -= old_size + + # Evict if needed + self._evict_if_needed(size) + + self._cache[key] = (value, expiry, size) + self._current_size_bytes += size + + heapq.heappush(self._expiry_heap, (expiry, key)) + + async def remove(self, key: str) -> None: + """Remove a value from the cache. + + Args: + key: The cache key. + """ + entry = self._cache.pop(key, None) + if entry is not None: + self._current_size_bytes -= entry[2] + self._cache.pop(key, None) + + +def create_protection_scopes_cache_key(request: ProtectionScopesRequest) -> str: + """Create a cache key for a ProtectionScopesRequest. + + The key is based on the serialized request content (excluding correlation_id). + + Args: + request: The protection scopes request. + + Returns: + A string cache key. + """ + data = request.to_dict(exclude_none=True) + + for field in ["correlation_id"]: + data.pop(field, None) + + json_str = json.dumps(data, sort_keys=True) + return f"purview:protection_scopes:{hashlib.sha256(json_str.encode()).hexdigest()}" + + +__all__ = [ + "CacheProvider", +] diff --git a/python/packages/purview/agent_framework_purview/_client.py b/python/packages/purview/agent_framework_purview/_client.py index 8679a3bf84..2ed37722cf 100644 --- a/python/packages/purview/agent_framework_purview/_client.py +++ b/python/packages/purview/agent_framework_purview/_client.py @@ -5,15 +5,19 @@ import base64 import inspect import json 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 +from opentelemetry import trace from ._exceptions import ( PurviewAuthenticationError, + PurviewPaymentRequiredError, PurviewRateLimitError, PurviewRequestError, PurviewServiceError, @@ -28,6 +32,8 @@ from ._models import ( ) from ._settings import PurviewSettings +logger = get_logger("agent_framework.purview") + class PurviewClient: """Async client for calling Graph Purview endpoints. @@ -85,13 +91,39 @@ class PurviewClient: with get_tracer().start_as_current_span("purview.process_content"): token = await self._get_token(tenant_id=request.tenant_id) url = f"{self._graph_uri}/users/{request.user_id}/dataSecurityAndGovernance/processContent" - return cast(ProcessContentResponse, await self._post(url, request, ProcessContentResponse, token)) + headers = {} + # Add If-None-Match header if scope_identifier is present + if hasattr(request, "scope_identifier") and request.scope_identifier: + headers["If-None-Match"] = request.scope_identifier + # Add Prefer: evaluateInline header if process_inline is True + if hasattr(request, "process_inline") and request.process_inline: + headers["Prefer"] = "evaluateInline" + + response = await self._post( + url, request, ProcessContentResponse, token, headers=headers, return_response=True + ) + + if isinstance(response, tuple) and len(response) == 2: + response_obj, _ = response + return cast(ProcessContentResponse, response_obj) + + return cast(ProcessContentResponse, response) async def get_protection_scopes(self, request: ProtectionScopesRequest) -> ProtectionScopesResponse: with get_tracer().start_as_current_span("purview.get_protection_scopes"): token = await self._get_token() url = f"{self._graph_uri}/users/{request.user_id}/dataSecurityAndGovernance/protectionScopes/compute" - return cast(ProtectionScopesResponse, await self._post(url, request, ProtectionScopesResponse, token)) + response = await self._post(url, request, ProtectionScopesResponse, token, return_response=True) + + # Extract etag from response headers + if isinstance(response, tuple) and len(response) == 2: + response_obj, headers = response + if "etag" in headers: + etag_value = headers["etag"].strip('"') + response_obj.scope_identifier = etag_value + return cast(ProtectionScopesResponse, response_obj) + + return cast(ProtectionScopesResponse, response) async def send_content_activities(self, request: ContentActivitiesRequest) -> ContentActivitiesResponse: with get_tracer().start_as_current_span("purview.send_content_activities"): @@ -99,16 +131,44 @@ class PurviewClient: url = f"{self._graph_uri}/users/{request.user_id}/dataSecurityAndGovernance/activities/contentActivities" return cast(ContentActivitiesResponse, await self._post(url, request, ContentActivitiesResponse, token)) - async def _post(self, url: str, model: Any, response_type: type[Any], token: str) -> Any: + async def _post( + self, + url: str, + model: Any, + response_type: type[Any], + token: str, + headers: dict[str, str] | None = None, + return_response: bool = False, + ) -> Any: + if hasattr(model, "correlation_id") and not model.correlation_id: + model.correlation_id = str(uuid4()) + + correlation_id = getattr(model, "correlation_id", None) + if correlation_id: + span = trace.get_current_span() + if span and span.is_recording(): + span.set_attribute("correlation_id", correlation_id) + logger.info(f"Purview request to {url} with correlation_id: {correlation_id}") + payload = model.model_dump(by_alias=True, exclude_none=True, mode="json") - headers = { + request_headers = { "Authorization": f"Bearer {token}", "User-Agent": AGENT_FRAMEWORK_USER_AGENT, "Content-Type": "application/json", } - resp = await self._client.post(url, json=payload, headers=headers) + if correlation_id: + request_headers["client-request-id"] = correlation_id + + if headers: + request_headers.update(headers) + resp = await self._client.post(url, json=payload, headers=request_headers) + if resp.status_code in (401, 403): raise PurviewAuthenticationError(f"Auth failure {resp.status_code}: {resp.text}") + if resp.status_code == 402: + if self._settings.ignore_payment_required: + return response_type() # type: ignore[call-arg, no-any-return] + raise PurviewPaymentRequiredError(f"Payment required {resp.status_code}: {resp.text}") if resp.status_code == 429: raise PurviewRateLimitError(f"Rate limited {resp.status_code}: {resp.text}") if resp.status_code not in (200, 201, 202): @@ -117,10 +177,21 @@ class PurviewClient: data = resp.json() except ValueError: data = {} + try: # Prefer pydantic-style model_validate if present, else fall back to constructor. if hasattr(response_type, "model_validate"): - return response_type.model_validate(data) # type: ignore[no-any-return] - return response_type(**data) # type: ignore[call-arg, no-any-return] - except Exception as ex: # pragma: no cover + response_obj = response_type.model_validate(data) # type: ignore[no-any-return] + else: + response_obj = response_type(**data) # type: ignore[call-arg, no-any-return] + + # Extract correlation_id from response headers if response object supports it + if "client-request-id" in resp.headers and hasattr(response_obj, "correlation_id"): + response_obj.correlation_id = resp.headers["client-request-id"] + logger.info(f"Purview response from {url} with correlation_id: {response_obj.correlation_id}") + + if return_response: + return (response_obj, resp.headers) + return response_obj + except Exception as ex: raise PurviewServiceError(f"Failed to deserialize Purview response: {ex}") from ex diff --git a/python/packages/purview/agent_framework_purview/_exceptions.py b/python/packages/purview/agent_framework_purview/_exceptions.py index 0ce8463504..8ddaf5bc28 100644 --- a/python/packages/purview/agent_framework_purview/_exceptions.py +++ b/python/packages/purview/agent_framework_purview/_exceptions.py @@ -7,6 +7,7 @@ from agent_framework.exceptions import ServiceResponseException __all__ = [ "PurviewAuthenticationError", + "PurviewPaymentRequiredError", "PurviewRateLimitError", "PurviewRequestError", "PurviewServiceError", @@ -21,6 +22,10 @@ class PurviewAuthenticationError(PurviewServiceError): """Authentication / authorization failure (401/403).""" +class PurviewPaymentRequiredError(PurviewServiceError): + """Payment required (402).""" + + class PurviewRateLimitError(PurviewServiceError): """Rate limiting or throttling (429).""" diff --git a/python/packages/purview/agent_framework_purview/_middleware.py b/python/packages/purview/agent_framework_purview/_middleware.py index 4ed21aa1a3..da43d8d9f4 100644 --- a/python/packages/purview/agent_framework_purview/_middleware.py +++ b/python/packages/purview/agent_framework_purview/_middleware.py @@ -8,7 +8,9 @@ from agent_framework._logging import get_logger from azure.core.credentials import TokenCredential from azure.core.credentials_async import AsyncTokenCredential +from ._cache import CacheProvider from ._client import PurviewClient +from ._exceptions import PurviewPaymentRequiredError from ._models import Activity from ._processor import ScopedContentProcessor from ._settings import PurviewSettings @@ -38,9 +40,10 @@ class PurviewPolicyMiddleware(AgentMiddleware): self, credential: TokenCredential | AsyncTokenCredential, settings: PurviewSettings, + cache_provider: CacheProvider | None = None, ) -> None: self._client = PurviewClient(credential, settings) - self._processor = ScopedContentProcessor(self._client, settings) + self._processor = ScopedContentProcessor(self._client, settings, cache_provider) self._settings = settings async def process( @@ -62,9 +65,14 @@ class PurviewPolicyMiddleware(AgentMiddleware): ) context.terminate = True return + except PurviewPaymentRequiredError as ex: + logger.error(f"Purview payment required error in policy pre-check: {ex}") + if not self._settings.ignore_payment_required: + raise except Exception as ex: - # Log and continue if there's an error in the pre-check logger.error(f"Error in Purview policy pre-check: {ex}") + if not self._settings.ignore_exceptions: + raise await next(context) @@ -86,9 +94,14 @@ class PurviewPolicyMiddleware(AgentMiddleware): else: # Streaming responses are not supported for post-checks logger.debug("Streaming responses are not supported for Purview policy post-checks") + except PurviewPaymentRequiredError as ex: + logger.error(f"Purview payment required error in policy post-check: {ex}") + if not self._settings.ignore_payment_required: + raise except Exception as ex: - # Log and continue if there's an error in the post-check logger.error(f"Error in Purview policy post-check: {ex}") + if not self._settings.ignore_exceptions: + raise class PurviewChatPolicyMiddleware(ChatMiddleware): @@ -118,9 +131,10 @@ class PurviewChatPolicyMiddleware(ChatMiddleware): self, credential: TokenCredential | AsyncTokenCredential, settings: PurviewSettings, + cache_provider: CacheProvider | None = None, ) -> None: self._client = PurviewClient(credential, settings) - self._processor = ScopedContentProcessor(self._client, settings) + self._processor = ScopedContentProcessor(self._client, settings, cache_provider) self._settings = settings async def process( @@ -134,15 +148,20 @@ class PurviewChatPolicyMiddleware(ChatMiddleware): context.messages, Activity.UPLOAD_TEXT ) if should_block_prompt: - from agent_framework import ChatMessage + from agent_framework import ChatMessage, ChatResponse - context.result = [ # type: ignore[assignment] - ChatMessage(role="system", text=self._settings.blocked_prompt_message) - ] + blocked_message = ChatMessage(role="system", text=self._settings.blocked_prompt_message) + context.result = ChatResponse(messages=[blocked_message]) context.terminate = True return + except PurviewPaymentRequiredError as ex: + logger.error(f"Purview payment required error in policy pre-check: {ex}") + if not self._settings.ignore_payment_required: + raise except Exception as ex: logger.error(f"Error in Purview policy pre-check: {ex}") + if not self._settings.ignore_exceptions: + raise await next(context) @@ -157,12 +176,17 @@ class PurviewChatPolicyMiddleware(ChatMiddleware): messages, Activity.UPLOAD_TEXT, user_id=resolved_user_id ) if should_block_response: - from agent_framework import ChatMessage + from agent_framework import ChatMessage, ChatResponse - context.result = [ # type: ignore[assignment] - ChatMessage(role="system", text=self._settings.blocked_response_message) - ] + blocked_message = ChatMessage(role="system", text=self._settings.blocked_response_message) + context.result = ChatResponse(messages=[blocked_message]) else: logger.debug("Streaming responses are not supported for Purview policy post-checks") + except PurviewPaymentRequiredError as ex: + logger.error(f"Purview payment required error in policy post-check: {ex}") + if not self._settings.ignore_payment_required: + raise except Exception as ex: logger.error(f"Error in Purview policy post-check: {ex}") + if not self._settings.ignore_exceptions: + raise diff --git a/python/packages/purview/agent_framework_purview/_models.py b/python/packages/purview/agent_framework_purview/_models.py index a629536901..0ee502da1a 100644 --- a/python/packages/purview/agent_framework_purview/_models.py +++ b/python/packages/purview/agent_framework_purview/_models.py @@ -642,7 +642,9 @@ class ContentToProcess(_AliasSerializable): class ProcessContentRequest(_AliasSerializable): _ALIASES: ClassVar[dict[str, str]] = {"content_to_process": "contentToProcess"} - DEFAULT_EXCLUDE: ClassVar[set[str]] = {"user_id", "tenant_id", "correlation_id", "process_inline"} + DEFAULT_EXCLUDE: ClassVar[set[str]] = { + "correlation_id", + } def __init__( self, @@ -651,6 +653,7 @@ class ProcessContentRequest(_AliasSerializable): tenant_id: str, correlation_id: str | None = None, process_inline: bool | None = None, + scope_identifier: str | None = None, **kwargs: Any, ) -> None: # Extract aliased values from kwargs @@ -668,10 +671,11 @@ class ProcessContentRequest(_AliasSerializable): self.tenant_id = tenant_id self.correlation_id = correlation_id self.process_inline = process_inline + self.scope_identifier = scope_identifier class ProtectionScopesRequest(_AliasSerializable): - DEFAULT_EXCLUDE: ClassVar[set[str]] = {"user_id", "tenant_id", "correlation_id", "scope_identifier"} + DEFAULT_EXCLUDE: ClassVar[set[str]] = {"correlation_id"} _ALIASES: ClassVar[dict[str, str]] = { "pivot_on": "pivotOn", "device_metadata": "deviceMetadata", @@ -743,7 +747,7 @@ class ContentActivitiesRequest(_AliasSerializable): "scope_identifier": "scopeIdentifier", "content_to_process": "contentMetadata", } - DEFAULT_EXCLUDE: ClassVar[set[str]] = {"tenant_id", "correlation_id"} + DEFAULT_EXCLUDE: ClassVar[set[str]] = {"correlation_id"} def __init__( self, @@ -800,12 +804,15 @@ class ProcessContentResponse(_AliasSerializable): "protection_scope_state": "protectionScopeState", "policy_actions": "policyActions", "processing_errors": "processingErrors", + "correlation_id": "correlationId", } + DEFAULT_EXCLUDE: ClassVar[set[str]] = {"correlation_id"} id: str | None protection_scope_state: ProtectionScopeState | None policy_actions: list[DlpActionInfo] | None processing_errors: list[ProcessingError] | None + correlation_id: str | None def __init__( self, @@ -813,6 +820,7 @@ class ProcessContentResponse(_AliasSerializable): protection_scope_state: ProtectionScopeState | None = None, policy_actions: list[DlpActionInfo | MutableMapping[str, Any]] | None = None, processing_errors: list[ProcessingError | MutableMapping[str, Any]] | None = None, + correlation_id: str | None = None, **kwargs: Any, ) -> None: # Extract aliased values from kwargs @@ -822,6 +830,8 @@ class ProcessContentResponse(_AliasSerializable): policy_actions = kwargs["policyActions"] if "processingErrors" in kwargs: processing_errors = kwargs["processingErrors"] + if "correlationId" in kwargs: + correlation_id = kwargs["correlationId"] # Convert to objects converted_policy_actions: list[DlpActionInfo] | None = None @@ -838,12 +848,12 @@ class ProcessContentResponse(_AliasSerializable): [pe if isinstance(pe, ProcessingError) else ProcessingError(**pe) for pe in processing_errors], ) - # Call parent without explicit params with aliases super().__init__(**kwargs) self.id = id self.protection_scope_state = protection_scope_state self.policy_actions = converted_policy_actions self.processing_errors = converted_processing_errors + self.correlation_id = correlation_id class PolicyScope(_AliasSerializable): @@ -909,15 +919,22 @@ class PolicyScope(_AliasSerializable): class ProtectionScopesResponse(_AliasSerializable): - _ALIASES: ClassVar[dict[str, str]] = {"scope_identifier": "scopeIdentifier", "scopes": "value"} + _ALIASES: ClassVar[dict[str, str]] = { + "scope_identifier": "scopeIdentifier", + "scopes": "value", + "correlation_id": "correlationId", + } + DEFAULT_EXCLUDE: ClassVar[set[str]] = {"correlation_id"} scope_identifier: str | None scopes: list[PolicyScope] | None + correlation_id: str | None def __init__( self, scope_identifier: str | None = None, scopes: list[PolicyScope | MutableMapping[str, Any]] | None = None, + correlation_id: str | None = None, **kwargs: Any, ) -> None: # Extract aliased values from kwargs before they're normalized by parent @@ -925,6 +942,8 @@ class ProtectionScopesResponse(_AliasSerializable): scope_identifier = kwargs["scopeIdentifier"] if "value" in kwargs: scopes = kwargs["value"] + if "correlationId" in kwargs: + correlation_id = kwargs["correlationId"] converted_scopes: list[PolicyScope] | None = None if scopes is not None: @@ -936,22 +955,32 @@ class ProtectionScopesResponse(_AliasSerializable): super().__init__(**kwargs) self.scope_identifier = scope_identifier self.scopes = converted_scopes + self.correlation_id = correlation_id class ContentActivitiesResponse(_AliasSerializable): - DEFAULT_EXCLUDE: ClassVar[set[str]] = {"status_code"} + DEFAULT_EXCLUDE: ClassVar[set[str]] = {"correlation_id"} + _ALIASES: ClassVar[dict[str, str]] = {"correlation_id": "correlationId"} + + status_code: int | None + error: ErrorDetails | None + correlation_id: str | None def __init__( self, status_code: int | None = None, error: ErrorDetails | MutableMapping[str, Any] | None = None, + correlation_id: str | None = None, **kwargs: Any, ) -> None: + if "correlationId" in kwargs: + correlation_id = kwargs["correlationId"] if isinstance(error, MutableMapping): error = ErrorDetails(**error) - super().__init__(status_code=status_code, error=error, **kwargs) + super().__init__(status_code=status_code, error=error, correlation_id=correlation_id, **kwargs) self.status_code = status_code self.error = error # type: ignore[assignment] + self.correlation_id = correlation_id __all__ = [ diff --git a/python/packages/purview/agent_framework_purview/_processor.py b/python/packages/purview/agent_framework_purview/_processor.py index 81367dd58b..44163dc2f6 100644 --- a/python/packages/purview/agent_framework_purview/_processor.py +++ b/python/packages/purview/agent_framework_purview/_processor.py @@ -1,13 +1,17 @@ # Copyright (c) Microsoft. All rights reserved. from __future__ import annotations +import asyncio import uuid from collections.abc import Iterable, MutableMapping from typing import Any from agent_framework import ChatMessage +from agent_framework._logging import get_logger +from ._cache import CacheProvider, InMemoryCacheProvider, create_protection_scopes_cache_key from ._client import PurviewClient +from ._exceptions import PurviewPaymentRequiredError from ._models import ( Activity, ActivityMetadata, @@ -16,22 +20,25 @@ from ._models import ( DeviceMetadata, DlpAction, DlpActionInfo, + ExecutionMode, IntegratedAppMetadata, OperatingSystemSpecifications, PolicyLocation, ProcessContentRequest, ProcessContentResponse, ProcessConversationMetadata, - ProcessingError, ProtectedAppMetadata, ProtectionScopesRequest, ProtectionScopesResponse, + ProtectionScopeState, PurviewTextContent, RestrictionAction, translate_activity, ) from ._settings import PurviewSettings +logger = get_logger("agent_framework.purview") + def _is_valid_guid(value: str | None) -> bool: """Check if a string is a valid GUID/UUID format using uuid module.""" @@ -47,9 +54,13 @@ def _is_valid_guid(value: str | None) -> bool: class ScopedContentProcessor: """Combine protection scopes, process content, and content activities logic.""" - def __init__(self, client: PurviewClient, settings: PurviewSettings): + def __init__(self, client: PurviewClient, settings: PurviewSettings, cache_provider: CacheProvider | None = None): self._client = client self._settings = settings + self._cache: CacheProvider = cache_provider or InMemoryCacheProvider( + default_ttl_seconds=settings.cache_ttl_seconds, max_size_bytes=settings.max_cache_size_bytes + ) + self._background_tasks: set[asyncio.Task[Any]] = set() async def process_messages( self, messages: Iterable[ChatMessage], activity: Activity, user_id: str | None = None @@ -173,7 +184,7 @@ class ScopedContentProcessor: user_id=resolved_user_id, # Use the resolved user_id for all messages tenant_id=tenant_id, correlation_id=meta.correlation_id, - process_inline=True if self._settings.process_inline else None, + process_inline=None, # Will be set based on execution mode ) results.append(req) return results, resolved_user_id @@ -191,23 +202,86 @@ class ScopedContentProcessor: integrated_app_metadata=pc_request.content_to_process.integrated_app_metadata, correlation_id=pc_request.correlation_id, ) - ps_resp = await self._client.get_protection_scopes(ps_req) - should_process, dlp_actions = self._check_applicable_scopes(pc_request, ps_resp) + + # Check for tenant-level 402 exception cache first + tenant_payment_cache_key = f"purview:payment_required:{pc_request.tenant_id}" + cached_payment_exception = await self._cache.get(tenant_payment_cache_key) + if isinstance(cached_payment_exception, PurviewPaymentRequiredError): + raise cached_payment_exception + + cache_key = create_protection_scopes_cache_key(ps_req) + cached_ps_resp = await self._cache.get(cache_key) + + if cached_ps_resp is not None: + if isinstance(cached_ps_resp, ProtectionScopesResponse): + ps_resp = cached_ps_resp + else: + try: + ps_resp = await self._client.get_protection_scopes(ps_req) + await self._cache.set(cache_key, ps_resp, ttl_seconds=self._settings.cache_ttl_seconds) + except PurviewPaymentRequiredError as ex: + # Cache the exception at tenant level so all subsequent requests for this tenant fail fast + await self._cache.set(tenant_payment_cache_key, ex, ttl_seconds=self._settings.cache_ttl_seconds) + raise + + if ps_resp.scope_identifier: + pc_request.scope_identifier = ps_resp.scope_identifier + + should_process, dlp_actions, execution_mode = self._check_applicable_scopes(pc_request, ps_resp) if should_process: + # Set process_inline based on execution mode + pc_request.process_inline = execution_mode == ExecutionMode.EVALUATE_INLINE + + # If execution mode is offline, queue the PC request in background + if execution_mode != ExecutionMode.EVALUATE_INLINE: + task = asyncio.create_task(self._process_content_background(pc_request, cache_key)) + self._background_tasks.add(task) + task.add_done_callback(self._background_tasks.discard) + return ProcessContentResponse(id="204", correlation_id=pc_request.correlation_id) + pc_resp = await self._client.process_content(pc_request) + + if pc_request.scope_identifier and pc_resp.protection_scope_state == ProtectionScopeState.MODIFIED: + await self._cache.remove(cache_key) + pc_resp.policy_actions = self._combine_policy_actions(pc_resp.policy_actions, dlp_actions) return pc_resp + + # No applicable scopes - send content activities in background ca_req = ContentActivitiesRequest( user_id=pc_request.user_id, tenant_id=pc_request.tenant_id, content_to_process=pc_request.content_to_process, correlation_id=pc_request.correlation_id, ) - ca_resp = await self._client.send_content_activities(ca_req) - if ca_resp.error: - return ProcessContentResponse(processing_errors=[ProcessingError(message=str(ca_resp.error))]) - return ProcessContentResponse() + + task = asyncio.create_task(self._send_content_activities_background(ca_req)) + self._background_tasks.add(task) + task.add_done_callback(self._background_tasks.discard) + # Respond with HttpStatusCode 204(No Content) + return ProcessContentResponse(id="204", correlation_id=pc_request.correlation_id) + + async def _process_content_background(self, pc_request: ProcessContentRequest, cache_key: str) -> None: + """Process content in background for offline execution mode.""" + try: + pc_resp = await self._client.process_content(pc_request) + + # If protection scope state is modified, make another PC request and invalidate cache + if pc_request.scope_identifier and pc_resp.protection_scope_state == ProtectionScopeState.MODIFIED: + await self._cache.remove(cache_key) + await self._client.process_content(pc_request) + except Exception as ex: + # Log errors but don't propagate since this is fire-and-forget + logger.warning(f"Background process content request failed: {ex}") + + async def _send_content_activities_background(self, ca_req: ContentActivitiesRequest) -> None: + """Send content activities in background without blocking.""" + try: + await self._client.send_content_activities(ca_req) + except Exception as ex: + # Log errors but don't propagate since this is fire-and-forget + logger.warning(f"Background content activities request failed: {ex}") @staticmethod def _combine_policy_actions( @@ -225,11 +299,22 @@ class ScopedContentProcessor: @staticmethod def _check_applicable_scopes( pc_request: ProcessContentRequest, ps_response: ProtectionScopesResponse - ) -> tuple[bool, list[DlpActionInfo]]: + ) -> tuple[bool, list[DlpActionInfo], ExecutionMode]: + """Check if any scopes are applicable to the request. + + Args: + pc_request: The process content request + ps_response: The protection scopes response + + Returns: + A tuple of (should_process, dlp_actions, execution_mode) + """ req_activity = translate_activity(pc_request.content_to_process.activity_metadata.activity) location = pc_request.content_to_process.protected_app_metadata.application_location should_process: bool = False dlp_actions: list[DlpActionInfo] = [] + execution_mode: ExecutionMode = ExecutionMode.EVALUATE_OFFLINE # Default to offline + for scope in ps_response.scopes or []: # Check if all activities in req_activity are present in scope.activities using bitwise flags. activity_match = bool(scope.activities and (scope.activities & req_activity) == req_activity) @@ -246,6 +331,11 @@ class ScopedContentProcessor: break if activity_match and location_match: should_process = True + + # If any scope has EvaluateInline, upgrade to inline mode + if scope.execution_mode == ExecutionMode.EVALUATE_INLINE: + execution_mode = ExecutionMode.EVALUATE_INLINE + if scope.policy_actions: dlp_actions.extend(scope.policy_actions) - return should_process, dlp_actions + return should_process, dlp_actions, execution_mode diff --git a/python/packages/purview/agent_framework_purview/_settings.py b/python/packages/purview/agent_framework_purview/_settings.py index 6b7e8d12f4..3e3d13ada6 100644 --- a/python/packages/purview/agent_framework_purview/_settings.py +++ b/python/packages/purview/agent_framework_purview/_settings.py @@ -41,18 +41,23 @@ class PurviewSettings(AFBaseSettings): Attributes: app_name: Public app name. + app_version: Optional version string of the application. tenant_id: Optional tenant id (guid) of the user making the request. purview_app_location: Optional app location for policy evaluation. graph_base_uri: Base URI for Microsoft Graph. blocked_prompt_message: Custom message to return when a prompt is blocked by policy. blocked_response_message: Custom message to return when a response is blocked by policy. + ignore_exceptions: If True, all Purview exceptions will be logged but not thrown in middleware. + ignore_payment_required: If True, 402 payment required errors will be logged but not thrown. + cache_ttl_seconds: Time to live for cache entries in seconds (default 14400 = 4 hours). + max_cache_size_bytes: Maximum cache size in bytes (default 200MB). """ app_name: str = Field(...) + app_version: str | None = Field(default=None) tenant_id: str | None = Field(default=None) purview_app_location: PurviewAppLocation | None = Field(default=None) graph_base_uri: str = Field(default="https://graph.microsoft.com/v1.0/") - process_inline: bool = Field(default=False, description="Process content inline if supported.") blocked_prompt_message: str = Field( default="Prompt blocked by policy", description="Message to return when a prompt is blocked by policy.", @@ -61,6 +66,22 @@ class PurviewSettings(AFBaseSettings): default="Response blocked by policy", description="Message to return when a response is blocked by policy.", ) + ignore_exceptions: bool = Field( + default=False, + description="If True, all Purview exceptions will be logged but not thrown in middleware.", + ) + ignore_payment_required: bool = Field( + default=False, + description="If True, 402 payment required errors will be logged but not thrown.", + ) + cache_ttl_seconds: int = Field( + default=14400, + description="Time to live for cache entries in seconds (default 14400 = 4 hours).", + ) + max_cache_size_bytes: int = Field( + default=200 * 1024 * 1024, + description="Maximum cache size in bytes (default 200MB).", + ) model_config = SettingsConfigDict(populate_by_name=True, validate_assignment=True) diff --git a/python/packages/purview/tests/test_cache.py b/python/packages/purview/tests/test_cache.py new file mode 100644 index 0000000000..2089d9e2e7 --- /dev/null +++ b/python/packages/purview/tests/test_cache.py @@ -0,0 +1,196 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Tests for Purview cache provider.""" + +import asyncio + +from agent_framework_purview._cache import ( + InMemoryCacheProvider, + create_protection_scopes_cache_key, +) +from agent_framework_purview._models import PolicyLocation, ProtectionScopesRequest + + +class TestInMemoryCacheProvider: + """Test InMemoryCacheProvider functionality.""" + + async def test_cache_set_and_get(self) -> None: + """Test basic set and get operations.""" + cache = InMemoryCacheProvider() + + await cache.set("key1", "value1") + result = await cache.get("key1") + + assert result == "value1" + + async def test_cache_get_nonexistent_key(self) -> None: + """Test get returns None for non-existent key.""" + cache = InMemoryCacheProvider() + + result = await cache.get("nonexistent") + + assert result is None + + async def test_cache_expiration(self) -> None: + """Test that cached values expire after TTL.""" + cache = InMemoryCacheProvider(default_ttl_seconds=1) + + await cache.set("key1", "value1") + result = await cache.get("key1") + assert result == "value1" + + await asyncio.sleep(1.1) + result = await cache.get("key1") + assert result is None + + async def test_cache_custom_ttl(self) -> None: + """Test that custom TTL overrides default.""" + cache = InMemoryCacheProvider(default_ttl_seconds=10) + + await cache.set("key1", "value1", ttl_seconds=1) + result = await cache.get("key1") + assert result == "value1" + + await asyncio.sleep(1.1) + result = await cache.get("key1") + assert result is None + + async def test_cache_update_existing_key(self) -> None: + """Test updating an existing cache entry.""" + cache = InMemoryCacheProvider() + + await cache.set("key1", "value1") + await cache.set("key1", "value2") + result = await cache.get("key1") + + assert result == "value2" + + async def test_cache_remove(self) -> None: + """Test removing a cache entry.""" + cache = InMemoryCacheProvider() + + await cache.set("key1", "value1") + await cache.remove("key1") + result = await cache.get("key1") + + assert result is None + + async def test_cache_remove_nonexistent_key(self) -> None: + """Test removing non-existent key does not raise error.""" + cache = InMemoryCacheProvider() + + await cache.remove("nonexistent") + + async def test_cache_size_limit_eviction(self) -> None: + """Test that cache evicts old entries when size limit is reached.""" + cache = InMemoryCacheProvider(max_size_bytes=200) + + await cache.set("key1", "a" * 50) + await cache.set("key2", "b" * 50) + await cache.set("key3", "c" * 50) + + await cache.set("key4", "d" * 100) + + result1 = await cache.get("key1") + assert result1 is None + + async def test_estimate_size_with_pydantic_model(self) -> None: + """Test size estimation with Pydantic models.""" + cache = InMemoryCacheProvider() + + location = PolicyLocation(**{"@odata.type": "microsoft.graph.policyLocationApplication", "value": "app-id"}) + request = ProtectionScopesRequest(user_id="user1", tenant_id="tenant1", locations=[location]) + + await cache.set("key1", request) + result = await cache.get("key1") + + assert result == request + + async def test_estimate_size_fallback(self) -> None: + """Test size estimation fallback for non-serializable objects.""" + cache = InMemoryCacheProvider() + + class CustomObject: + pass + + obj = CustomObject() + await cache.set("key1", obj) + result = await cache.get("key1") + + assert result == obj + + async def test_cache_multiple_updates(self) -> None: + """Test that updating a key multiple times maintains correct size tracking.""" + cache = InMemoryCacheProvider(max_size_bytes=1000) + + await cache.set("key1", "a" * 100) + initial_size = cache._current_size_bytes + + await cache.set("key1", "b" * 200) + + assert cache._current_size_bytes != initial_size + + async def test_eviction_with_stale_heap_entries(self) -> None: + """Test that eviction correctly handles stale heap entries.""" + cache = InMemoryCacheProvider(max_size_bytes=500) + + await cache.set("key1", "a" * 100, ttl_seconds=10) + await cache.set("key2", "b" * 100, ttl_seconds=10) + await cache.set("key1", "c" * 100, ttl_seconds=20) + + await cache.set("key3", "d" * 300) + + result = await cache.get("key1") + assert result is not None + + +class TestCreateProtectionScopesCacheKey: + """Test cache key generation for ProtectionScopesRequest.""" + + def test_cache_key_deterministic(self) -> None: + """Test that same request generates same cache key.""" + location = PolicyLocation(**{"@odata.type": "microsoft.graph.policyLocationApplication", "value": "app-id"}) + request1 = ProtectionScopesRequest(user_id="user1", tenant_id="tenant1", locations=[location]) + request2 = ProtectionScopesRequest(user_id="user1", tenant_id="tenant1", locations=[location]) + + key1 = create_protection_scopes_cache_key(request1) + key2 = create_protection_scopes_cache_key(request2) + + assert key1 == key2 + + def test_cache_key_different_for_different_requests(self) -> None: + """Test that different requests generate different cache keys.""" + location1 = PolicyLocation(**{"@odata.type": "microsoft.graph.policyLocationApplication", "value": "app-id1"}) + location2 = PolicyLocation(**{"@odata.type": "microsoft.graph.policyLocationApplication", "value": "app-id2"}) + request1 = ProtectionScopesRequest(user_id="user1", tenant_id="tenant1", locations=[location1]) + request2 = ProtectionScopesRequest(user_id="user1", tenant_id="tenant1", locations=[location2]) + + key1 = create_protection_scopes_cache_key(request1) + key2 = create_protection_scopes_cache_key(request2) + + assert key1 != key2 + + def test_cache_key_excludes_correlation_id(self) -> None: + """Test that correlation_id is excluded from cache key.""" + location = PolicyLocation(**{"@odata.type": "microsoft.graph.policyLocationApplication", "value": "app-id"}) + request1 = ProtectionScopesRequest( + user_id="user1", tenant_id="tenant1", locations=[location], correlation_id="corr1" + ) + request2 = ProtectionScopesRequest( + user_id="user1", tenant_id="tenant1", locations=[location], correlation_id="corr2" + ) + + key1 = create_protection_scopes_cache_key(request1) + key2 = create_protection_scopes_cache_key(request2) + + assert key1 == key2 + + def test_cache_key_format(self) -> None: + """Test that cache key has expected format.""" + location = PolicyLocation(**{"@odata.type": "microsoft.graph.policyLocationApplication", "value": "app-id"}) + request = ProtectionScopesRequest(user_id="user1", tenant_id="tenant1", locations=[location]) + + key = create_protection_scopes_cache_key(request) + + assert key.startswith("purview:protection_scopes:") + assert len(key) > len("purview:protection_scopes:") diff --git a/python/packages/purview/tests/test_chat_middleware.py b/python/packages/purview/tests/test_chat_middleware.py index 7e4faa0236..5633488a7e 100644 --- a/python/packages/purview/tests/test_chat_middleware.py +++ b/python/packages/purview/tests/test_chat_middleware.py @@ -74,7 +74,8 @@ class TestPurviewChatPolicyMiddleware: await middleware.process(chat_context, mock_next) assert chat_context.terminate assert chat_context.result - msg = chat_context.result[0] # type: ignore[index] + assert hasattr(chat_context.result, "messages") + msg = chat_context.result.messages[0] assert msg.role in ("system", Role.SYSTEM) assert "blocked" in msg.text.lower() @@ -112,7 +113,7 @@ class TestPurviewChatPolicyMiddleware: chat_options=chat_options, is_streaming=True, ) - with patch.object(middleware._processor, "process_messages", return_value=False) as mock_proc: + with patch.object(middleware._processor, "process_messages", return_value=(False, "user-123")) as mock_proc: async def mock_next(ctx: ChatContext) -> None: ctx.result = MagicMock() @@ -123,7 +124,10 @@ class TestPurviewChatPolicyMiddleware: async def test_chat_middleware_handles_post_check_exception( self, middleware: PurviewChatPolicyMiddleware, chat_context: ChatContext ) -> None: - """Test that exceptions in post-check are logged but don't affect result.""" + """Test that exceptions in post-check are logged but don't affect result when ignore_exceptions=True.""" + # Set ignore_exceptions to True to test exception suppression + middleware._settings.ignore_exceptions = True + call_count = 0 async def mock_process_messages(*args, **kwargs): @@ -136,7 +140,6 @@ class TestPurviewChatPolicyMiddleware: with patch.object(middleware._processor, "process_messages", side_effect=mock_process_messages): async def mock_next(ctx: ChatContext) -> None: - # Create a mock result with messages attribute result = MagicMock() result.messages = [ChatMessage(role=Role.ASSISTANT, text="Response")] ctx.result = result @@ -147,3 +150,127 @@ class TestPurviewChatPolicyMiddleware: assert call_count == 2 # Result should still be set assert chat_context.result is not None + + async def test_chat_middleware_uses_consistent_user_id( + self, middleware: PurviewChatPolicyMiddleware, chat_context: ChatContext + ) -> None: + """Test that the same user_id from pre-check is used in post-check.""" + captured_user_ids = [] + + async def mock_process_messages(messages, activity, user_id=None): + captured_user_ids.append(user_id) + return (False, "resolved-user-123") + + with patch.object(middleware._processor, "process_messages", side_effect=mock_process_messages): + + async def mock_next(ctx: ChatContext) -> None: + result = MagicMock() + result.messages = [ChatMessage(role=Role.ASSISTANT, text="Response")] + ctx.result = result + + await middleware.process(chat_context, mock_next) + + # Should have been called twice + assert len(captured_user_ids) == 2 + # First call should have None (no user_id provided yet) + assert captured_user_ids[0] is None + # Second call should have the resolved user_id from first call + assert captured_user_ids[1] == "resolved-user-123" + + async def test_chat_middleware_handles_payment_required_pre_check(self, mock_credential: AsyncMock) -> None: + """Test that 402 in pre-check is handled based on settings.""" + from agent_framework_purview._exceptions import PurviewPaymentRequiredError + + # Test with ignore_payment_required=False + settings = PurviewSettings(app_name="Test App", ignore_payment_required=False) + middleware = PurviewChatPolicyMiddleware(mock_credential, settings) + + chat_client = DummyChatClient() + chat_options = MagicMock() + chat_options.model = "test-model" + context = ChatContext( + chat_client=chat_client, messages=[ChatMessage(role=Role.USER, text="Hello")], chat_options=chat_options + ) + + async def mock_process_messages(*args, **kwargs): + raise PurviewPaymentRequiredError("Payment required") + + with patch.object(middleware._processor, "process_messages", side_effect=mock_process_messages): + + async def mock_next(ctx: ChatContext) -> None: + raise AssertionError("next should not be called") + + # Should raise the exception + with pytest.raises(PurviewPaymentRequiredError): + await middleware.process(context, mock_next) + + async def test_chat_middleware_ignores_payment_required_when_configured(self, mock_credential: AsyncMock) -> None: + """Test that 402 is ignored when ignore_payment_required=True.""" + from agent_framework_purview._exceptions import PurviewPaymentRequiredError + + settings = PurviewSettings(app_name="Test App", ignore_payment_required=True) + middleware = PurviewChatPolicyMiddleware(mock_credential, settings) + + chat_client = DummyChatClient() + chat_options = MagicMock() + chat_options.model = "test-model" + context = ChatContext( + chat_client=chat_client, messages=[ChatMessage(role=Role.USER, text="Hello")], chat_options=chat_options + ) + + async def mock_process_messages(*args, **kwargs): + raise PurviewPaymentRequiredError("Payment required") + + with patch.object(middleware._processor, "process_messages", side_effect=mock_process_messages): + + async def mock_next(ctx: ChatContext) -> None: + result = MagicMock() + result.messages = [ChatMessage(role=Role.ASSISTANT, text="Response")] + context.result = result + + # Should not raise, just log + await middleware.process(context, mock_next) + # Next should have been called + assert context.result is not None + + async def test_chat_middleware_handles_result_without_messages_attribute( + self, middleware: PurviewChatPolicyMiddleware, chat_context: ChatContext + ) -> None: + """Test middleware handles result that doesn't have messages attribute.""" + with patch.object(middleware._processor, "process_messages", return_value=(False, "user-123")): + + async def mock_next(ctx: ChatContext) -> None: + # Set result to something without messages attribute + ctx.result = "Some string result" + + await middleware.process(chat_context, mock_next) + + # Should not crash, result should be unchanged + assert chat_context.result == "Some string result" + + async def test_chat_middleware_with_ignore_exceptions(self, mock_credential: AsyncMock) -> None: + """Test that middleware respects ignore_exceptions setting.""" + settings = PurviewSettings(app_name="Test App", ignore_exceptions=True) + middleware = PurviewChatPolicyMiddleware(mock_credential, settings) + + chat_client = DummyChatClient() + chat_options = MagicMock() + chat_options.model = "test-model" + context = ChatContext( + chat_client=chat_client, messages=[ChatMessage(role=Role.USER, text="Hello")], chat_options=chat_options + ) + + async def mock_process_messages(*args, **kwargs): + raise ValueError("Some error") + + with patch.object(middleware._processor, "process_messages", side_effect=mock_process_messages): + + async def mock_next(ctx: ChatContext) -> None: + result = MagicMock() + result.messages = [ChatMessage(role=Role.ASSISTANT, text="Response")] + context.result = result + + # Should not raise, just log + await middleware.process(context, mock_next) + # Next should have been called + assert context.result is not None diff --git a/python/packages/purview/tests/test_client.py b/python/packages/purview/tests/test_client.py index 3becc9b49f..7953c16b77 100644 --- a/python/packages/purview/tests/test_client.py +++ b/python/packages/purview/tests/test_client.py @@ -12,6 +12,7 @@ from agent_framework_purview import PurviewSettings from agent_framework_purview._client import PurviewClient from agent_framework_purview._exceptions import ( PurviewAuthenticationError, + PurviewPaymentRequiredError, PurviewRateLimitError, PurviewRequestError, PurviewServiceError, @@ -157,6 +158,7 @@ class TestPurviewClient: mock_response = MagicMock(spec=httpx.Response) mock_response.status_code = 200 + mock_response.headers = {} mock_response.json.return_value = {"id": "response-123", "protectionScopeState": "notModified"} with patch.object(client._client, "post", return_value=mock_response): @@ -174,6 +176,7 @@ class TestPurviewClient: mock_response = MagicMock(spec=httpx.Response) mock_response.status_code = 200 + mock_response.headers = {} # Add headers attribute mock_response.json.return_value = {"scopeIdentifier": "scope-123", "value": []} with patch.object(client._client, "post", return_value=mock_response): @@ -236,3 +239,157 @@ class TestPurviewClient: pytest.raises(PurviewRequestError, match="Purview request failed"), ): await client.process_content(request) + + async def test_prefer_header_sent_when_process_inline_true( + self, client: PurviewClient, content_to_process_factory + ) -> None: + """Test that Prefer: evaluateInline header is sent when process_inline is True.""" + content = content_to_process_factory() + request = ProcessContentRequest( + content_to_process=content, + user_id="user-123", + tenant_id="tenant-456", + process_inline=True, + ) + + posted_headers = {} + mock_response = MagicMock(spec=httpx.Response) + mock_response.status_code = 200 + mock_response.headers = {} + mock_response.json.return_value = {} + + async def capture_post(url, json, headers): + posted_headers.update(headers) + return mock_response + + with patch.object(client._client, "post", side_effect=capture_post): + await client.process_content(request) + + assert "Prefer" in posted_headers + assert posted_headers["Prefer"] == "evaluateInline" + + async def test_prefer_header_not_sent_when_process_inline_false( + self, client: PurviewClient, content_to_process_factory + ) -> None: + """Test that Prefer header is not sent when process_inline is False.""" + content = content_to_process_factory() + request = ProcessContentRequest( + content_to_process=content, + user_id="user-123", + tenant_id="tenant-456", + process_inline=False, + ) + + posted_headers = {} + mock_response = MagicMock(spec=httpx.Response) + mock_response.status_code = 200 + mock_response.headers = {} + mock_response.json.return_value = {} + + async def capture_post(url, json, headers): + posted_headers.update(headers) + return mock_response + + with patch.object(client._client, "post", side_effect=capture_post): + await client.process_content(request) + + assert "Prefer" not in posted_headers + + async def test_prefer_header_not_sent_when_process_inline_none( + self, client: PurviewClient, content_to_process_factory + ) -> None: + """Test that Prefer header is not sent when process_inline is None.""" + content = content_to_process_factory() + request = ProcessContentRequest( + content_to_process=content, + user_id="user-123", + tenant_id="tenant-456", + process_inline=None, + ) + + posted_headers = {} + mock_response = MagicMock(spec=httpx.Response) + mock_response.status_code = 200 + mock_response.headers = {} + mock_response.json.return_value = {} + + async def capture_post(url, json, headers): + posted_headers.update(headers) + return mock_response + + with patch.object(client._client, "post", side_effect=capture_post): + await client.process_content(request) + + assert "Prefer" not in posted_headers + + async def test_scope_identifier_extraction_from_etag(self, client: PurviewClient) -> None: + """Test that scope_identifier is extracted from ETag header.""" + mock_response = MagicMock(spec=httpx.Response) + mock_response.status_code = 200 + mock_response.headers = {"etag": '"test-scope-id"'} + mock_response.json.return_value = {"value": []} + + with patch.object(client._client, "post", return_value=mock_response): + req = ProtectionScopesRequest(user_id="user1", tenant_id="tenant1") + response = await client.get_protection_scopes(req) + + assert response.scope_identifier == "test-scope-id" + + async def test_scope_identifier_sent_as_if_none_match_header( + self, client: PurviewClient, content_to_process_factory + ) -> None: + """Test that scope_identifier is sent as If-None-Match header.""" + content = content_to_process_factory() + request = ProcessContentRequest( + content_to_process=content, + user_id="user-123", + tenant_id="tenant-456", + scope_identifier="test-scope-id", + ) + + posted_headers = {} + mock_response = MagicMock(spec=httpx.Response) + mock_response.status_code = 200 + mock_response.headers = {} + mock_response.json.return_value = {} + + async def capture_post(url, json, headers): + posted_headers.update(headers) + return mock_response + + with patch.object(client._client, "post", side_effect=capture_post): + await client.process_content(request) + + assert "If-None-Match" in posted_headers + assert posted_headers["If-None-Match"] == "test-scope-id" + + async def test_402_payment_required_raises_exception_by_default(self, client: PurviewClient) -> None: + """Test that 402 raises exception when ignore_payment_required is False.""" + mock_response = MagicMock(spec=httpx.Response) + mock_response.status_code = 402 + mock_response.text = "Payment required" + + with patch.object(client._client, "post", return_value=mock_response): + req = ProtectionScopesRequest(user_id="user1", tenant_id="tenant1") + + with pytest.raises(PurviewPaymentRequiredError): + await client.get_protection_scopes(req) + + async def test_402_payment_required_returns_empty_when_ignored(self, mock_credential: MagicMock) -> None: + """Test that 402 returns empty response when ignore_payment_required is True.""" + settings = PurviewSettings(app_name="Test App", ignore_payment_required=True) + client = PurviewClient(mock_credential, settings) + + mock_response = MagicMock(spec=httpx.Response) + mock_response.status_code = 402 + mock_response.text = "Payment required" + + with patch.object(client._client, "post", return_value=mock_response): + req = ProtectionScopesRequest(user_id="user1", tenant_id="tenant1") + response = await client.get_protection_scopes(req) + + # Should return empty response without raising + assert response is not None + assert response.scopes is None or response.scopes == [] + + await client.close() diff --git a/python/packages/purview/tests/test_exceptions.py b/python/packages/purview/tests/test_exceptions.py index cb417d028d..adae0ee091 100644 --- a/python/packages/purview/tests/test_exceptions.py +++ b/python/packages/purview/tests/test_exceptions.py @@ -4,6 +4,7 @@ from agent_framework_purview import ( PurviewAuthenticationError, + PurviewPaymentRequiredError, PurviewRateLimitError, PurviewRequestError, PurviewServiceError, @@ -25,6 +26,12 @@ class TestPurviewExceptions: assert str(error) == "Authentication failed" assert isinstance(error, PurviewServiceError) + def test_purview_payment_required_error(self) -> None: + """Test PurviewPaymentRequiredError exception.""" + error = PurviewPaymentRequiredError("Payment required") + assert str(error) == "Payment required" + assert isinstance(error, PurviewServiceError) + def test_purview_rate_limit_error(self) -> None: """Test PurviewRateLimitError exception.""" error = PurviewRateLimitError("Rate limit exceeded") diff --git a/python/packages/purview/tests/test_middleware.py b/python/packages/purview/tests/test_middleware.py index 8d31d25950..0769efab27 100644 --- a/python/packages/purview/tests/test_middleware.py +++ b/python/packages/purview/tests/test_middleware.py @@ -120,6 +120,9 @@ class TestPurviewPolicyMiddleware: self, middleware: PurviewPolicyMiddleware, mock_agent: MagicMock ) -> None: """Test middleware handles result that doesn't have messages attribute.""" + # Set ignore_exceptions to True so AttributeError is caught and logged + middleware._settings.ignore_exceptions = True + context = AgentRunContext(agent=mock_agent, messages=[ChatMessage(role=Role.USER, text="Hello")]) with patch.object(middleware._processor, "process_messages", return_value=(False, "user-123")): @@ -153,7 +156,10 @@ class TestPurviewPolicyMiddleware: async def test_middleware_handles_pre_check_exception( self, middleware: PurviewPolicyMiddleware, mock_agent: MagicMock ) -> None: - """Test that exceptions in pre-check are logged but don't stop processing.""" + """Test that exceptions in pre-check are logged but don't stop processing when ignore_exceptions=True.""" + # Set ignore_exceptions to True + middleware._settings.ignore_exceptions = True + context = AgentRunContext(agent=mock_agent, messages=[ChatMessage(role=Role.USER, text="Test")]) with patch.object( @@ -175,7 +181,10 @@ class TestPurviewPolicyMiddleware: async def test_middleware_handles_post_check_exception( self, middleware: PurviewPolicyMiddleware, mock_agent: MagicMock ) -> None: - """Test that exceptions in post-check are logged but don't affect result.""" + """Test that exceptions in post-check are logged but don't affect result when ignore_exceptions=True.""" + # Set ignore_exceptions to True + middleware._settings.ignore_exceptions = True + context = AgentRunContext(agent=mock_agent, messages=[ChatMessage(role=Role.USER, text="Test")]) call_count = 0 @@ -199,3 +208,49 @@ class TestPurviewPolicyMiddleware: # Result should still be set assert context.result is not None assert hasattr(context.result, "messages") + + async def test_middleware_with_ignore_exceptions_true(self, mock_credential: AsyncMock) -> None: + """Test that middleware logs but doesn't throw when ignore_exceptions is True.""" + settings = PurviewSettings(app_name="Test App", ignore_exceptions=True) + middleware = PurviewPolicyMiddleware(mock_credential, settings) + + mock_agent = MagicMock() + mock_agent.name = "test-agent" + context = AgentRunContext(agent=mock_agent, messages=[ChatMessage(role=Role.USER, text="Test")]) + + # Mock processor to raise an exception + async def mock_process_messages(*args, **kwargs): + raise ValueError("Test error") + + with patch.object(middleware._processor, "process_messages", side_effect=mock_process_messages): + + async def mock_next(ctx): + ctx.result = AgentRunResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="Response")]) + + # Should not raise, just log + await middleware.process(context, mock_next) + + # Result should be set because next was called despite the error + assert context.result is not None + + async def test_middleware_with_ignore_exceptions_false(self, mock_credential: AsyncMock) -> None: + """Test that middleware throws exceptions when ignore_exceptions is False.""" + settings = PurviewSettings(app_name="Test App", ignore_exceptions=False) + middleware = PurviewPolicyMiddleware(mock_credential, settings) + + mock_agent = MagicMock() + mock_agent.name = "test-agent" + context = AgentRunContext(agent=mock_agent, messages=[ChatMessage(role=Role.USER, text="Test")]) + + # Mock processor to raise an exception + async def mock_process_messages(*args, **kwargs): + raise ValueError("Test error") + + with patch.object(middleware._processor, "process_messages", side_effect=mock_process_messages): + + async def mock_next(ctx): + pass + + # Should raise the exception + with pytest.raises(ValueError, match="Test error"): + await middleware.process(context, mock_next) diff --git a/python/packages/purview/tests/test_models.py b/python/packages/purview/tests/test_models.py index aa4f62d79d..d0c968df5c 100644 --- a/python/packages/purview/tests/test_models.py +++ b/python/packages/purview/tests/test_models.py @@ -237,10 +237,7 @@ class TestModelDeserialization: dumped = request.model_dump(by_alias=True, exclude_none=True, mode="json") - # Check that excluded fields are not present - assert "user_id" not in dumped - assert "tenant_id" not in dumped + assert "user_id" in dumped + assert "tenant_id" in dumped assert "correlation_id" not in dumped - - # Check that content is present assert "contentToProcess" in dumped diff --git a/python/packages/purview/tests/test_processor.py b/python/packages/purview/tests/test_processor.py index 96554af126..c517da2459 100644 --- a/python/packages/purview/tests/test_processor.py +++ b/python/packages/purview/tests/test_processor.py @@ -170,7 +170,7 @@ class TestScopedContentProcessor: request = process_content_request_factory() response = ProtectionScopesResponse(**{"value": None}) - should_process, actions = processor._check_applicable_scopes(request, response) + should_process, actions, execution_mode = processor._check_applicable_scopes(request, response) assert should_process is False assert actions == [] @@ -200,7 +200,7 @@ class TestScopedContentProcessor: }) response = ProtectionScopesResponse(**{"value": [scope]}) - should_process, actions = processor._check_applicable_scopes(request, response) + should_process, actions, execution_mode = processor._check_applicable_scopes(request, response) assert should_process is True assert len(actions) == 1 @@ -220,7 +220,7 @@ class TestScopedContentProcessor: async def test_process_with_scopes_calls_client_methods( self, processor: ScopedContentProcessor, mock_client: AsyncMock, process_content_request_factory ) -> None: - """Test _process_with_scopes calls get_protection_scopes and process_content.""" + """Test _process_with_scopes calls get_protection_scopes when scopes response is empty.""" from agent_framework_purview._models import ( ContentActivitiesResponse, ProtectionScopesResponse, @@ -237,9 +237,10 @@ class TestScopedContentProcessor: response = await processor._process_with_scopes(request) mock_client.get_protection_scopes.assert_called_once() + # When no scopes apply, process_content is not called (activities are sent in background) mock_client.process_content.assert_not_called() - mock_client.send_content_activities.assert_called_once() - assert response.id is None + # The response should have id=204 (No Content) when no scopes apply + assert response.id == "204" async def test_map_messages_with_user_id_in_additional_properties(self, mock_client: AsyncMock) -> None: """Test user_id extraction from message additional_properties.""" @@ -308,7 +309,7 @@ class TestScopedContentProcessor: async def test_process_content_sends_activities_when_not_applicable( self, mock_client: AsyncMock, process_content_request_factory ) -> None: - """Test that content activities are sent when scopes don't apply.""" + """Test that response is returned when scopes don't apply (activities sent in background).""" settings = PurviewSettings( app_name="Test App", tenant_id="12345678-1234-1234-1234-123456789012", @@ -325,7 +326,7 @@ class TestScopedContentProcessor: mock_ps_response.scopes = [] mock_client.get_protection_scopes.return_value = mock_ps_response - # Mock send_content_activities to return success + # Mock send_content_activities to return success (called in background) mock_ca_response = MagicMock() mock_ca_response.error = None mock_client.send_content_activities.return_value = mock_ca_response @@ -334,14 +335,13 @@ class TestScopedContentProcessor: mock_client.get_protection_scopes.assert_called_once() mock_client.process_content.assert_not_called() - mock_client.send_content_activities.assert_called_once() - # When content activities succeed, response has no errors (processing_errors can be None or empty) - assert response.processing_errors is None or response.processing_errors == [] + # Response should have id=204 when no scopes apply + assert response.id == "204" async def test_process_content_handles_activities_error( self, mock_client: AsyncMock, process_content_request_factory ) -> None: - """Test error handling when content activities fail.""" + """Test that errors in background activities don't affect the response.""" settings = PurviewSettings( app_name="Test App", tenant_id="12345678-1234-1234-1234-123456789012", @@ -358,12 +358,269 @@ class TestScopedContentProcessor: mock_ps_response.scopes = [] mock_client.get_protection_scopes.return_value = mock_ps_response - # Mock send_content_activities to return error + # Mock send_content_activities to return error (called in background task) mock_ca_response = MagicMock() mock_ca_response.error = "Test error message" mock_client.send_content_activities.return_value = mock_ca_response response = await processor._process_with_scopes(pc_request) - assert len(response.processing_errors) == 1 - assert response.processing_errors[0].message == "Test error message" + # Since activities are sent in background, errors don't affect the response + # Response should have id=204 when no scopes apply + assert response.id == "204" + + +class TestUserIdResolution: + """Test user ID resolution from various sources.""" + + @pytest.fixture + def mock_client(self) -> AsyncMock: + """Create a mock Purview client.""" + client = AsyncMock() + client.get_user_info_from_token = AsyncMock( + return_value={ + "tenant_id": "12345678-1234-1234-1234-123456789012", + "user_id": "11111111-1111-1111-1111-111111111111", + "client_id": "12345678-1234-1234-1234-123456789012", + } + ) + return client + + @pytest.fixture + def settings(self) -> PurviewSettings: + """Create settings.""" + return PurviewSettings( + app_name="Test App", + tenant_id="12345678-1234-1234-1234-123456789012", + purview_app_location=PurviewAppLocation( + location_type=PurviewLocationType.APPLICATION, location_value="app-id" + ), + ) + + async def test_user_id_from_token_when_no_other_source(self, mock_client: AsyncMock) -> None: + """Test user_id is extracted from token when no other source available.""" + settings = PurviewSettings(app_name="Test App") # No tenant_id or app_location + processor = ScopedContentProcessor(mock_client, settings) + + messages = [ChatMessage(role=Role.USER, text="Test")] + + requests, user_id = await processor._map_messages(messages, Activity.UPLOAD_TEXT) + + mock_client.get_user_info_from_token.assert_called_once() + assert user_id == "11111111-1111-1111-1111-111111111111" + + async def test_user_id_from_additional_properties_takes_priority( + self, mock_client: AsyncMock, settings: PurviewSettings + ) -> None: + """Test user_id from additional_properties takes priority over token.""" + processor = ScopedContentProcessor(mock_client, settings) + + messages = [ + ChatMessage( + role=Role.USER, + text="Test", + additional_properties={"user_id": "22222222-2222-2222-2222-222222222222"}, + ) + ] + + requests, user_id = await processor._map_messages(messages, Activity.UPLOAD_TEXT) + + # Token info should not be called since we have user_id in message + mock_client.get_user_info_from_token.assert_not_called() + assert user_id == "22222222-2222-2222-2222-222222222222" + + async def test_user_id_from_author_name_as_fallback( + self, mock_client: AsyncMock, settings: PurviewSettings + ) -> None: + """Test user_id is extracted from author_name when it's a valid GUID.""" + processor = ScopedContentProcessor(mock_client, settings) + + messages = [ + ChatMessage( + role=Role.USER, + text="Test", + author_name="33333333-3333-3333-3333-333333333333", + ) + ] + + requests, user_id = await processor._map_messages(messages, Activity.UPLOAD_TEXT) + + assert user_id == "33333333-3333-3333-3333-333333333333" + + async def test_author_name_ignored_if_not_valid_guid( + self, mock_client: AsyncMock, settings: PurviewSettings + ) -> None: + """Test author_name is ignored if it's not a valid GUID.""" + processor = ScopedContentProcessor(mock_client, settings) + + messages = [ + ChatMessage( + role=Role.USER, + text="Test", + author_name="John Doe", # Not a GUID + ) + ] + + requests, user_id = await processor._map_messages(messages, Activity.UPLOAD_TEXT) + + # Should return empty since author_name is not a valid GUID + assert user_id is None + assert len(requests) == 0 + + async def test_provided_user_id_used_as_last_resort( + self, mock_client: AsyncMock, settings: PurviewSettings + ) -> None: + """Test provided_user_id parameter is used as last resort.""" + processor = ScopedContentProcessor(mock_client, settings) + + messages = [ChatMessage(role=Role.USER, text="Test")] + + requests, user_id = await processor._map_messages( + messages, Activity.UPLOAD_TEXT, provided_user_id="44444444-4444-4444-4444-444444444444" + ) + + assert user_id == "44444444-4444-4444-4444-444444444444" + + async def test_invalid_provided_user_id_ignored(self, mock_client: AsyncMock, settings: PurviewSettings) -> None: + """Test invalid provided_user_id is ignored.""" + processor = ScopedContentProcessor(mock_client, settings) + + messages = [ChatMessage(role=Role.USER, text="Test")] + + requests, user_id = await processor._map_messages(messages, Activity.UPLOAD_TEXT, provided_user_id="not-a-guid") + + assert user_id is None + assert len(requests) == 0 + + async def test_multiple_messages_same_user_id(self, mock_client: AsyncMock, settings: PurviewSettings) -> None: + """Test that all messages use the same resolved user_id.""" + processor = ScopedContentProcessor(mock_client, settings) + + messages = [ + ChatMessage( + role=Role.USER, text="First", additional_properties={"user_id": "55555555-5555-5555-5555-555555555555"} + ), + ChatMessage(role=Role.ASSISTANT, text="Response"), + ChatMessage(role=Role.USER, text="Second"), + ] + + requests, user_id = await processor._map_messages(messages, Activity.UPLOAD_TEXT) + + assert user_id == "55555555-5555-5555-5555-555555555555" + # All requests should have the same user_id + assert all(req.user_id == "55555555-5555-5555-5555-555555555555" for req in requests) + + async def test_first_valid_user_id_in_messages_is_used( + self, mock_client: AsyncMock, settings: PurviewSettings + ) -> None: + """Test that the first valid user_id found in messages is used for all.""" + processor = ScopedContentProcessor(mock_client, settings) + + messages = [ + ChatMessage(role=Role.USER, text="First", author_name="Not a GUID"), + ChatMessage( + role=Role.ASSISTANT, + text="Response", + additional_properties={"user_id": "66666666-6666-6666-6666-666666666666"}, + ), + ChatMessage( + role=Role.USER, text="Third", additional_properties={"user_id": "77777777-7777-7777-7777-777777777777"} + ), + ] + + requests, user_id = await processor._map_messages(messages, Activity.UPLOAD_TEXT) + + # First valid user_id (from second message) should be used + assert user_id == "66666666-6666-6666-6666-666666666666" + assert all(req.user_id == "66666666-6666-6666-6666-666666666666" for req in requests) + + +class TestScopedContentProcessorCaching: + """Test caching functionality in ScopedContentProcessor.""" + + @pytest.fixture + def mock_client(self) -> AsyncMock: + """Create a mock Purview client.""" + client = AsyncMock() + client.get_user_info_from_token = AsyncMock( + return_value={ + "tenant_id": "12345678-1234-1234-1234-123456789012", + "user_id": "12345678-1234-1234-1234-123456789012", + "client_id": "12345678-1234-1234-1234-123456789012", + } + ) + client.get_protection_scopes = AsyncMock() + return client + + @pytest.fixture + def settings(self) -> PurviewSettings: + """Create test settings.""" + location = PurviewAppLocation(location_type=PurviewLocationType.APPLICATION, location_value="app-id") + return PurviewSettings( + app_name="Test App", + tenant_id="12345678-1234-1234-1234-123456789012", + default_user_id="12345678-1234-1234-1234-123456789012", + purview_app_location=location, + ) + + async def test_protection_scopes_cached_on_first_call( + self, mock_client: AsyncMock, settings: PurviewSettings + ) -> None: + """Test that protection scopes response is cached after first call.""" + from agent_framework_purview._cache import InMemoryCacheProvider + from agent_framework_purview._models import ProtectionScopesResponse + + cache_provider = InMemoryCacheProvider() + processor = ScopedContentProcessor(mock_client, settings, cache_provider=cache_provider) + + mock_client.get_protection_scopes.return_value = ProtectionScopesResponse( + scope_identifier="scope-123", scopes=[] + ) + + messages = [ChatMessage(role=Role.USER, text="Test")] + + await processor.process_messages(messages, Activity.UPLOAD_TEXT, user_id="12345678-1234-1234-1234-123456789012") + + mock_client.get_protection_scopes.assert_called_once() + + await processor.process_messages(messages, Activity.UPLOAD_TEXT, user_id="12345678-1234-1234-1234-123456789012") + + mock_client.get_protection_scopes.assert_called_once() + + async def test_payment_required_exception_cached_at_tenant_level( + self, mock_client: AsyncMock, settings: PurviewSettings + ) -> None: + """Test that 402 payment required exceptions are cached at tenant level.""" + from agent_framework_purview._cache import InMemoryCacheProvider + from agent_framework_purview._exceptions import PurviewPaymentRequiredError + + cache_provider = InMemoryCacheProvider() + processor = ScopedContentProcessor(mock_client, settings, cache_provider=cache_provider) + + mock_client.get_protection_scopes.side_effect = PurviewPaymentRequiredError("Payment required") + + messages = [ChatMessage(role=Role.USER, text="Test")] + + with pytest.raises(PurviewPaymentRequiredError): + await processor.process_messages( + messages, Activity.UPLOAD_TEXT, user_id="12345678-1234-1234-1234-123456789012" + ) + + mock_client.get_protection_scopes.assert_called_once() + + with pytest.raises(PurviewPaymentRequiredError): + await processor.process_messages( + messages, Activity.UPLOAD_TEXT, user_id="12345678-1234-1234-1234-123456789012" + ) + + mock_client.get_protection_scopes.assert_called_once() + + async def test_custom_cache_provider_used(self, mock_client: AsyncMock, settings: PurviewSettings) -> None: + """Test that custom cache provider is used when provided.""" + from agent_framework_purview._cache import InMemoryCacheProvider + + custom_cache = InMemoryCacheProvider(default_ttl_seconds=60) + processor = ScopedContentProcessor(mock_client, settings, cache_provider=custom_cache) + + assert processor._cache is custom_cache + assert processor._cache._default_ttl == 60 diff --git a/python/packages/purview/tests/test_settings.py b/python/packages/purview/tests/test_settings.py index dac8d68d63..9abc27aed1 100644 --- a/python/packages/purview/tests/test_settings.py +++ b/python/packages/purview/tests/test_settings.py @@ -18,7 +18,6 @@ class TestPurviewSettings: assert settings.graph_base_uri == "https://graph.microsoft.com/v1.0/" assert settings.tenant_id is None assert settings.purview_app_location is None - assert settings.process_inline is False def test_settings_with_custom_values(self) -> None: """Test PurviewSettings with custom values.""" @@ -28,13 +27,11 @@ class TestPurviewSettings: app_name="Test App", graph_base_uri="https://graph.microsoft-ppe.com", tenant_id="test-tenant-id", - process_inline=True, purview_app_location=app_location, ) assert settings.graph_base_uri == "https://graph.microsoft-ppe.com" assert settings.tenant_id == "test-tenant-id" - assert settings.process_inline is True assert settings.purview_app_location.location_value == "app-123" @pytest.mark.parametrize( diff --git a/python/samples/getting_started/purview_agent/README.md b/python/samples/getting_started/purview_agent/README.md index 59b2ee918b..d4cfeca3df 100644 --- a/python/samples/getting_started/purview_agent/README.md +++ b/python/samples/getting_started/purview_agent/README.md @@ -2,9 +2,14 @@ This getting-started sample shows how to attach Microsoft Purview policy evaluation to an Agent Framework `ChatAgent` using the **middleware** approach. +**What this sample demonstrates:** 1. Configure an Azure OpenAI chat client 2. Add Purview policy enforcement middleware (`PurviewPolicyMiddleware`) -3. Run a short conversation and observe prompt / response blocking behavior +3. Add Purview policy enforcement at the chat client level (`PurviewChatPolicyMiddleware`) +4. Implement a custom cache provider for advanced caching scenarios +5. Run conversations and observe prompt / response blocking behavior + +**Note:** Caching is **automatic** and enabled by default with sensible defaults (30-minute TTL, 200MB max size). --- ## 1. Setup @@ -20,8 +25,6 @@ This getting-started sample shows how to attach Microsoft Purview policy evaluat | `PURVIEW_CERT_PATH` | Yes (when cert auth on) | Path to your .pfx certificate | | `PURVIEW_CERT_PASSWORD` | Optional | Password for encrypted certs | -*A demo default exists in code for illustration only—always set your own value. - ### 2. Auth Modes Supported #### A. Interactive Browser Authentication (default) @@ -42,7 +45,7 @@ $env:PURVIEW_CERT_PATH = "C:\path\to\cert.pfx" $env:PURVIEW_CERT_PASSWORD = "optional-password" ``` -Certificate steps (summary): create / register app, generate certificate, upload public key, export .pfx with private key, grant required Graph / Purview permissions. +Certificate steps (summary): create / register entra app, generate certificate, upload public key, export .pfx with private key, grant required Graph / Purview permissions. --- @@ -61,18 +64,39 @@ If interactive auth is used, a browser window will appear the first time. ## 4. How It Works +The sample demonstrates three different scenarios: + +### A. Agent Middleware (`run_with_agent_middleware`) 1. Builds an Azure OpenAI chat client (using the environment endpoint / deployment) 2. Chooses credential mode (certificate vs interactive) 3. Creates `PurviewPolicyMiddleware` with `PurviewSettings` 4. Injects middleware into the agent at construction 5. Sends two user messages sequentially 6. Prints results (or policy block messages) +7. Uses default caching automatically +### B. Chat Client Middleware (`run_with_chat_middleware`) +1. Creates a chat client with `PurviewChatPolicyMiddleware` attached directly +2. Policy evaluation happens at the chat client level rather than agent level +3. Demonstrates an alternative integration point for Purview policies +4. Uses default caching automatically + +### C. Custom Cache Provider (`run_with_custom_cache_provider`) +1. Implements the `CacheProvider` protocol with a custom class (`SimpleDictCacheProvider`) +2. Shows how to add custom logging and metrics to cache operations +3. The custom provider must implement three async methods: + - `async def get(self, key: str) -> Any | None` + - `async def set(self, key: str, value: Any, ttl_seconds: int | None = None) -> None` + - `async def remove(self, key: str) -> None` + +**Policy Behavior:** Prompt blocks set a system-level message: `Prompt blocked by policy` and terminate the run early. Response blocks rewrite the output to `Response blocked by policy`. --- -## 5. Code Snippet (Middleware Injection) +## 5. Code Snippets + +### Agent Middleware Injection ```python agent = ChatAgent( @@ -80,9 +104,41 @@ agent = ChatAgent( instructions="You are good at telling jokes.", name="Joker", middleware=[ - PurviewPolicyMiddleware(credential, PurviewSettings(app_name="Sample App", default_user_id="")) + PurviewPolicyMiddleware(credential, PurviewSettings(app_name="Sample App")) ], ) ``` +### Custom Cache Provider Implementation + +This is only needed if you want to integrate with external caching systems. + +```python +class SimpleDictCacheProvider: + """Custom cache provider that implements the CacheProvider protocol.""" + + def __init__(self) -> None: + self._cache: dict[str, Any] = {} + + async def get(self, key: str) -> Any | None: + """Get a value from the cache.""" + return self._cache.get(key) + + async def set(self, key: str, value: Any, ttl_seconds: int | None = None) -> None: + """Set a value in the cache.""" + self._cache[key] = value + + async def remove(self, key: str) -> None: + """Remove a value from the cache.""" + self._cache.pop(key, None) + +# Use the custom cache provider +custom_cache = SimpleDictCacheProvider() +middleware = PurviewPolicyMiddleware( + credential, + PurviewSettings(app_name="Sample App"), + cache_provider=custom_cache, +) +``` + --- diff --git a/python/samples/getting_started/purview_agent/sample_purview_agent.py b/python/samples/getting_started/purview_agent/sample_purview_agent.py index 00e596c32f..4d9b0d612c 100644 --- a/python/samples/getting_started/purview_agent/sample_purview_agent.py +++ b/python/samples/getting_started/purview_agent/sample_purview_agent.py @@ -5,7 +5,10 @@ Shows: 1. Creating a basic chat agent 2. Adding Purview policy evaluation via AGENT middleware (agent-level) 3. Adding Purview policy evaluation via CHAT middleware (chat-client level) -4. Running a threaded conversation and printing results +4. Implementing a custom cache provider for advanced caching scenarios +5. Running threaded conversations and printing results + +Note: Caching is automatic and enabled by default. Environment variables: - AZURE_OPENAI_ENDPOINT (required) @@ -31,7 +34,6 @@ from azure.identity import ( InteractiveBrowserCredential, ) -# Purview integration pieces from agent_framework.microsoft import ( PurviewPolicyMiddleware, PurviewChatPolicyMiddleware, @@ -42,6 +44,59 @@ JOKER_NAME = "Joker" JOKER_INSTRUCTIONS = "You are good at telling jokes. Keep responses concise." +# Custom Cache Provider Implementation +class SimpleDictCacheProvider: + """A simple custom cache provider that stores everything in a dictionary. + + This example demonstrates how to implement the CacheProvider protocol. + """ + + def __init__(self) -> None: + """Initialize the simple dictionary cache.""" + self._cache: dict[str, Any] = {} + self._access_count: dict[str, int] = {} + + async def get(self, key: str) -> Any | None: + """Get a value from the cache. + + Args: + key: The cache key. + + Returns: + The cached value or None if not found. + """ + value = self._cache.get(key) + if value is not None: + self._access_count[key] = self._access_count.get(key, 0) + 1 + print(f"[CustomCache] Cache HIT for key: {key[:50]}... (accessed {self._access_count[key]} times)") + else: + print(f"[CustomCache] Cache MISS for key: {key[:50]}...") + return value + + async def set(self, key: str, value: Any, ttl_seconds: int | None = None) -> None: + """Set a value in the cache. + + Args: + key: The cache key. + value: The value to cache. + ttl_seconds: Time to live in seconds (ignored in this simple implementation). + """ + self._cache[key] = value + print(f"[CustomCache] Cached value for key: {key[:50]}... (TTL: {ttl_seconds}s)") + + async def remove(self, key: str) -> None: + """Remove a value from the cache. + + Args: + key: The cache key. + """ + if key in self._cache: + del self._cache[key] + self._access_count.pop(key, None) + print(f"[CustomCache] Removed key: {key[:50]}...") + + + def _get_env(name: str, *, required: bool = True, default: str | None = None) -> str: val = os.environ.get(name, default) if required and not val: @@ -161,9 +216,91 @@ async def run_with_chat_middleware() -> None: ) print("Second response (chat middleware):\n", second) +async def run_with_custom_cache_provider() -> None: + """Demonstrate implementing and using a custom cache provider.""" + endpoint = os.environ.get("AZURE_OPENAI_ENDPOINT") + if not endpoint: + print("Skipping custom cache provider run: AZURE_OPENAI_ENDPOINT not set") + return + + deployment = os.environ.get("AZURE_OPENAI_DEPLOYMENT_NAME", "gpt-4o-mini") + user_id = os.environ.get("PURVIEW_DEFAULT_USER_ID") + chat_client = AzureOpenAIChatClient(deployment_name=deployment, endpoint=endpoint, credential=AzureCliCredential()) + + custom_cache = SimpleDictCacheProvider() + + purview_agent_middleware = PurviewPolicyMiddleware( + build_credential(), + PurviewSettings( + app_name="Agent Framework Sample App (Custom Provider)", + ), + cache_provider=custom_cache, + ) + + agent = ChatAgent( + chat_client=chat_client, + instructions=JOKER_INSTRUCTIONS, + name=JOKER_NAME, + middleware=purview_agent_middleware, + ) + + print("-- Custom Cache Provider Path --") + print("Using SimpleDictCacheProvider") + + first: AgentRunResponse = await agent.run( + ChatMessage(role=Role.USER, text="Tell me a joke about a programmer.", additional_properties={"user_id": user_id}) + ) + print("First response (custom provider):\n", first) + + second: AgentRunResponse = await agent.run( + ChatMessage(role=Role.USER, text="That's hilarious! One more?", additional_properties={"user_id": user_id}) + ) + print("Second response (custom provider):\n", second) + + """Demonstrate using the default built-in cache.""" + endpoint = os.environ.get("AZURE_OPENAI_ENDPOINT") + if not endpoint: + print("Skipping default cache run: AZURE_OPENAI_ENDPOINT not set") + return + + deployment = os.environ.get("AZURE_OPENAI_DEPLOYMENT_NAME", "gpt-4o-mini") + user_id = os.environ.get("PURVIEW_DEFAULT_USER_ID") + chat_client = AzureOpenAIChatClient(deployment_name=deployment, endpoint=endpoint, credential=AzureCliCredential()) + + # No cache_provider specified - uses default InMemoryCacheProvider + purview_agent_middleware = PurviewPolicyMiddleware( + build_credential(), + PurviewSettings( + app_name="Agent Framework Sample App (Default Cache)", + cache_ttl_seconds=3600, + max_cache_size_bytes=100 * 1024 * 1024, # 100MB + ), + ) + + agent = ChatAgent( + chat_client=chat_client, + instructions=JOKER_INSTRUCTIONS, + name=JOKER_NAME, + middleware=purview_agent_middleware, + ) + + print("-- Default Cache Path --") + print("Using default InMemoryCacheProvider with settings-based configuration") + + first: AgentRunResponse = await agent.run( + ChatMessage(role=Role.USER, text="Tell me a joke about AI.", additional_properties={"user_id": user_id}) + ) + print("First response (default cache):\n", first) + + second: AgentRunResponse = await agent.run( + ChatMessage(role=Role.USER, text="Nice! Another AI joke please.", additional_properties={"user_id": user_id}) + ) + print("Second response (default cache):\n", second) + async def main() -> None: - print("== Purview Agent Sample (Agent & Chat Middleware) ==") + print("== Purview Agent Sample (Middleware with Automatic Caching) ==") + try: await run_with_agent_middleware() except Exception as ex: # pragma: no cover - demo resilience @@ -174,6 +311,11 @@ async def main() -> None: except Exception as ex: # pragma: no cover - demo resilience print(f"Chat middleware path failed: {ex}") + try: + await run_with_custom_cache_provider() + except Exception as ex: # pragma: no cover - demo resilience + print(f"Custom cache provider path failed: {ex}") + if __name__ == "__main__": asyncio.run(main()) From e5d9d74c2da0302d5e2ca01e174a7b1cbe4c03e0 Mon Sep 17 00:00:00 2001 From: westey <164392973+westey-m@users.noreply.github.com> Date: Fri, 7 Nov 2025 11:46:35 +0000 Subject: [PATCH 10/16] .NET: Chathistory memory provider add (#1867) * Add ChatHistoryMemoryProvider with unit tests * Set new project to not packable. * Fix bugs * Add serialization support. * Update dotnet/src/Microsoft.Agents.AI.VectorDataMemory/ChatHistoryMemoryProvider.cs Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * Remove unnecessary line * Convert ChatHistoryMemoryProvider to use Dynamic collections. * Sealing options and scope classes. * Add sample, add scope to logs and improve scope validation * Move ChatHistoryMemoryProvider to MAAI project. --------- Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- dotnet/agent-framework-dotnet.slnx | 1 + .../Agent_Step18_TextSearchRag/Program.cs | 2 + ...nt_Step21_ChatHistoryMemoryProvider.csproj | 23 + .../Program.cs | 60 +++ .../Microsoft.Agents.AI.Mem0/Mem0Provider.cs | 8 +- .../Microsoft.Agents.AI/AgentJsonUtilities.cs | 1 + .../ChatClient/ChatClientAgentThread.cs | 4 +- .../Memory/ChatHistoryMemoryProvider.cs | 483 ++++++++++++++++++ .../ChatHistoryMemoryProviderOptions.cs | 56 ++ .../Memory/ChatHistoryMemoryProviderScope.cs | 53 ++ .../Memory/ChatHistoryMemoryProviderTests.cs | 396 ++++++++++++++ 11 files changed, 1081 insertions(+), 6 deletions(-) create mode 100644 dotnet/samples/GettingStarted/Agents/Agent_Step21_ChatHistoryMemoryProvider/Agent_Step21_ChatHistoryMemoryProvider.csproj create mode 100644 dotnet/samples/GettingStarted/Agents/Agent_Step21_ChatHistoryMemoryProvider/Program.cs create mode 100644 dotnet/src/Microsoft.Agents.AI/Memory/ChatHistoryMemoryProvider.cs create mode 100644 dotnet/src/Microsoft.Agents.AI/Memory/ChatHistoryMemoryProviderOptions.cs create mode 100644 dotnet/src/Microsoft.Agents.AI/Memory/ChatHistoryMemoryProviderScope.cs create mode 100644 dotnet/tests/Microsoft.Agents.AI.UnitTests/Memory/ChatHistoryMemoryProviderTests.cs diff --git a/dotnet/agent-framework-dotnet.slnx b/dotnet/agent-framework-dotnet.slnx index 03f8a910d3..56211e457f 100644 --- a/dotnet/agent-framework-dotnet.slnx +++ b/dotnet/agent-framework-dotnet.slnx @@ -66,6 +66,7 @@ + diff --git a/dotnet/samples/GettingStarted/Agents/Agent_Step18_TextSearchRag/Program.cs b/dotnet/samples/GettingStarted/Agents/Agent_Step18_TextSearchRag/Program.cs index 65f3a9e98f..c56be11fc5 100644 --- a/dotnet/samples/GettingStarted/Agents/Agent_Step18_TextSearchRag/Program.cs +++ b/dotnet/samples/GettingStarted/Agents/Agent_Step18_TextSearchRag/Program.cs @@ -4,6 +4,8 @@ // capabilities to an AI agent. The provider runs a search against an external knowledge base // before each model invocation and injects the results into the model context. +// Also see the AgentWithRAG folder for more advanced RAG scenarios. + using Azure.AI.OpenAI; using Azure.Identity; using Microsoft.Agents.AI; diff --git a/dotnet/samples/GettingStarted/Agents/Agent_Step21_ChatHistoryMemoryProvider/Agent_Step21_ChatHistoryMemoryProvider.csproj b/dotnet/samples/GettingStarted/Agents/Agent_Step21_ChatHistoryMemoryProvider/Agent_Step21_ChatHistoryMemoryProvider.csproj new file mode 100644 index 0000000000..1caf270c49 --- /dev/null +++ b/dotnet/samples/GettingStarted/Agents/Agent_Step21_ChatHistoryMemoryProvider/Agent_Step21_ChatHistoryMemoryProvider.csproj @@ -0,0 +1,23 @@ + + + + Exe + net9.0 + + enable + enable + + + + + + + + + + + + + + + diff --git a/dotnet/samples/GettingStarted/Agents/Agent_Step21_ChatHistoryMemoryProvider/Program.cs b/dotnet/samples/GettingStarted/Agents/Agent_Step21_ChatHistoryMemoryProvider/Program.cs new file mode 100644 index 0000000000..9e4c27cebb --- /dev/null +++ b/dotnet/samples/GettingStarted/Agents/Agent_Step21_ChatHistoryMemoryProvider/Program.cs @@ -0,0 +1,60 @@ +// Copyright (c) Microsoft. All rights reserved. + +// This sample shows how to create and use a simple AI agent that stores chat messages in a vector store using the ChatHistoryMemoryProvider. +// It can then use the chat history from prior conversations to inform responses in new conversations. + +using Azure.AI.OpenAI; +using Azure.Identity; +using Microsoft.Agents.AI; +using Microsoft.Extensions.AI; +using Microsoft.Extensions.VectorData; +using Microsoft.SemanticKernel.Connectors.InMemory; +using OpenAI; + +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"; +var embeddingDeploymentName = Environment.GetEnvironmentVariable("AZURE_OPENAI_EMBEDDING_DEPLOYMENT_NAME") ?? "text-embedding-3-large"; + +// Create a vector store to store the chat messages in. +// For demonstration purposes, we are using an in-memory vector store. +// Replace this with a vector store implementation of your choice that can persist the chat history long term. +VectorStore vectorStore = new InMemoryVectorStore(new InMemoryVectorStoreOptions() +{ + EmbeddingGenerator = new AzureOpenAIClient(new Uri(endpoint), new AzureCliCredential()) + .GetEmbeddingClient(embeddingDeploymentName) + .AsIEmbeddingGenerator() +}); + +// Create the agent and add the ChatHistoryMemoryProvider to store chat messages in the vector store. +AIAgent agent = new AzureOpenAIClient( + new Uri(endpoint), + new AzureCliCredential()) + .GetChatClient(deploymentName) + .CreateAIAgent(new ChatClientAgentOptions + { + Instructions = "You are good at telling jokes.", + Name = "Joker", + AIContextProviderFactory = (ctx) => 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 thread ID for each new thread. + storageScope: new() { UserId = "UID1", ThreadId = new Guid().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 threads. + searchScope: new() { UserId = "UID1" }) + }); + +// Start a new thread for the agent conversation. +AgentThread thread = agent.GetNewThread(); + +// Run the agent with the thread that stores conversation history in the vector store. +Console.WriteLine(await agent.RunAsync("I like jokes about Pirates. Tell me a joke about a pirate.", thread)); + +// Start a second thread. Since we configured the search scope to be across all threads for the user, +// the agent should remember that the user likes pirate jokes. +AgentThread thread2 = agent.GetNewThread(); + +// Run the agent with the second thread. +Console.WriteLine(await agent.RunAsync("Tell me a joke that I might like.", thread2)); diff --git a/dotnet/src/Microsoft.Agents.AI.Mem0/Mem0Provider.cs b/dotnet/src/Microsoft.Agents.AI.Mem0/Mem0Provider.cs index 4aae5de59b..d18ed2b460 100644 --- a/dotnet/src/Microsoft.Agents.AI.Mem0/Mem0Provider.cs +++ b/dotnet/src/Microsoft.Agents.AI.Mem0/Mem0Provider.cs @@ -153,7 +153,7 @@ public sealed class Mem0Provider : AIContextProvider if (this._logger is not null) { this._logger.LogInformation( - "Mem0AIContextProvider: Retrieved {Count} memories. ApplicationId: '{ApplicationId}', AgentId: '{AgentId}', ThreadId: '{ThreadId}', UserId: '{UserId}'", + "Mem0AIContextProvider: Retrieved {Count} memories. ApplicationId: '{ApplicationId}', AgentId: '{AgentId}', ThreadId: '{ThreadId}', UserId: '{UserId}'.", memories.Count, this._searchScope.ApplicationId, this._searchScope.AgentId, @@ -162,7 +162,7 @@ public sealed class Mem0Provider : AIContextProvider if (outputMessageText is not null) { this._logger.LogTrace( - "Mem0AIContextProvider: Search Results\nInput:{Input}\nOutput:{MessageText}\nApplicationId: '{ApplicationId}', AgentId: '{AgentId}', ThreadId: '{ThreadId}', UserId: '{UserId}'", + "Mem0AIContextProvider: Search Results\nInput:{Input}\nOutput:{MessageText}\nApplicationId: '{ApplicationId}', AgentId: '{AgentId}', ThreadId: '{ThreadId}', UserId: '{UserId}'.", queryText, outputMessageText, this._searchScope.ApplicationId, @@ -185,7 +185,7 @@ 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}'", + "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, @@ -211,7 +211,7 @@ 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}'", + "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, diff --git a/dotnet/src/Microsoft.Agents.AI/AgentJsonUtilities.cs b/dotnet/src/Microsoft.Agents.AI/AgentJsonUtilities.cs index 36bef4a2af..c400a1cb6c 100644 --- a/dotnet/src/Microsoft.Agents.AI/AgentJsonUtilities.cs +++ b/dotnet/src/Microsoft.Agents.AI/AgentJsonUtilities.cs @@ -68,6 +68,7 @@ internal static partial class AgentJsonUtilities // Agent abstraction types [JsonSerializable(typeof(ChatClientAgentThread.ThreadState))] [JsonSerializable(typeof(TextSearchProvider.TextSearchProviderState))] + [JsonSerializable(typeof(ChatHistoryMemoryProvider.ChatHistoryMemoryProviderState))] [ExcludeFromCodeCoverage] internal sealed partial class JsonContext : JsonSerializerContext; diff --git a/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgentThread.cs b/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgentThread.cs index baa36c0054..ad224e6777 100644 --- a/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgentThread.cs +++ b/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgentThread.cs @@ -166,8 +166,8 @@ public class ChatClientAgentThread : AgentThread var state = new ThreadState { ConversationId = this.ConversationId, - StoreState = storeState, - AIContextProviderState = aiContextProviderState + StoreState = storeState is { ValueKind: not JsonValueKind.Undefined } ? storeState : null, + AIContextProviderState = aiContextProviderState is { ValueKind: not JsonValueKind.Undefined } ? aiContextProviderState : null, }; return JsonSerializer.SerializeToElement(state, AgentJsonUtilities.DefaultOptions.GetTypeInfo(typeof(ThreadState))); diff --git a/dotnet/src/Microsoft.Agents.AI/Memory/ChatHistoryMemoryProvider.cs b/dotnet/src/Microsoft.Agents.AI/Memory/ChatHistoryMemoryProvider.cs new file mode 100644 index 0000000000..6d90c877e8 --- /dev/null +++ b/dotnet/src/Microsoft.Agents.AI/Memory/ChatHistoryMemoryProvider.cs @@ -0,0 +1,483 @@ +// Copyright (c) Microsoft. All rights reserved. + +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; +using Microsoft.Extensions.Logging; +using Microsoft.Extensions.VectorData; +using Microsoft.Shared.Diagnostics; + +namespace Microsoft.Agents.AI; + +/// +/// A context provider that stores all chat history in a vector store and is able to +/// retrieve related chat history later to augment the current conversation. +/// +/// +/// +/// This provider stores chat messages in a vector store and retrieves relevant previous messages +/// to provide as context during agent invocations. It uses the VectorStore and VectorStoreCollection +/// abstractions to work with any compatible vector store implementation. +/// +/// +/// Messages are stored during the method and retrieved during the +/// method using semantic similarity search. +/// +/// +/// Behavior is configurable through . When +/// is selected the provider +/// exposes a function tool that the model can invoke to retrieve relevant memories on demand instead of +/// injecting them automatically on each invocation. +/// +/// +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 VectorStore _vectorStore; + private readonly VectorStoreCollection> _collection; + private readonly int _maxResults; + private readonly string _contextPrompt; + private readonly ChatHistoryMemoryProviderOptions.SearchBehavior _searchTime; + private readonly AITool[] _tools; + private readonly ILogger? _logger; + + private readonly ChatHistoryMemoryProviderScope _storageScope; + private readonly ChatHistoryMemoryProviderScope _searchScope; + + private bool _collectionInitialized; + private readonly SemaphoreSlim _initializationLock = new(1, 1); + private bool _disposedValue; + + /// + /// Initializes a new instance of the class. + /// + /// 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. + /// Optional configuration options. + /// Optional logger factory. + /// Thrown when is . + public ChatHistoryMemoryProvider( + VectorStore vectorStore, + string collectionName, + int vectorDimensions, + ChatHistoryMemoryProviderScope storageScope, + ChatHistoryMemoryProviderScope? searchScope = null, + ChatHistoryMemoryProviderOptions? options = null, + ILoggerFactory? loggerFactory = null) + : this( + vectorStore, + collectionName, + vectorDimensions, + new ChatHistoryMemoryProviderState + { + StorageScope = new(Throw.IfNull(storageScope)), + SearchScope = searchScope ?? new(storageScope), + }, + options, + loggerFactory) + { + } + + /// + /// 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._searchTime = options.SearchTime; + 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) + ]; + + // Create a definition so that we can use the dimensions provided at runtime. + var definition = new VectorStoreCollectionDefinition + { + Properties = new List + { + new VectorStoreKeyProperty("Key", typeof(Guid)), + new VectorStoreDataProperty("Role", typeof(string)) { IsIndexed = true }, + new VectorStoreDataProperty("MessageId", typeof(string)) { IsIndexed = true }, + new VectorStoreDataProperty("AuthorName", typeof(string)), + new VectorStoreDataProperty("ApplicationId", typeof(string)) { IsIndexed = true }, + new VectorStoreDataProperty("AgentId", typeof(string)) { IsIndexed = true }, + new VectorStoreDataProperty("UserId", typeof(string)) { IsIndexed = true }, + new VectorStoreDataProperty("ThreadId", typeof(string)) { IsIndexed = true }, + new VectorStoreDataProperty("Content", typeof(string)) { IsFullTextIndexed = true }, + new VectorStoreDataProperty("CreatedAt", typeof(string)) { IsIndexed = true }, + new VectorStoreVectorProperty("ContentEmbedding", typeof(string), Throw.IfLessThan(vectorDimensions, 1)) + } + }; + + this._collection = this._vectorStore.GetDynamicCollection(Throw.IfNullOrWhitespace(collectionName), definition); + } + + /// + public override async ValueTask InvokingAsync(InvokingContext context, CancellationToken cancellationToken = default) + { + _ = Throw.IfNull(context); + + if (this._searchTime == ChatHistoryMemoryProviderOptions.SearchBehavior.OnDemandFunctionCalling) + { + // Expose search tool for on-demand invocation by the model + return new AIContext { Tools = this._tools }; + } + + try + { + // Get the text from the current request messages + var requestText = string.Join("\n", context.RequestMessages + .Where(m => m != null && !string.IsNullOrWhiteSpace(m.Text)) + .Select(m => m.Text)); + + if (string.IsNullOrWhiteSpace(requestText)) + { + return new AIContext(); + } + + // Search for relevant chat history + var contextText = await this.SearchTextAsync(requestText, cancellationToken).ConfigureAwait(false); + + if (string.IsNullOrWhiteSpace(contextText)) + { + return new AIContext(); + } + + return new AIContext + { + Messages = [new ChatMessage(ChatRole.User, contextText)] + }; + } + catch (Exception ex) + { + this._logger?.LogError( + ex, + "ChatHistoryMemoryProvider: Failed to search for chat history due to error. ApplicationId: '{ApplicationId}', AgentId: '{AgentId}', ThreadId: '{ThreadId}', UserId: '{UserId}'.", + this._searchScope.ApplicationId, + this._searchScope.AgentId, + this._searchScope.ThreadId, + this._searchScope.UserId); + return new AIContext(); + } + } + + /// + public override async ValueTask InvokedAsync(InvokedContext context, CancellationToken cancellationToken = default) + { + _ = Throw.IfNull(context); + + // Only store if invocation was successful + if (context.InvokeException != null) + { + return; + } + + try + { + // Ensure the collection is initialized + var collection = await this.EnsureCollectionExistsAsync(cancellationToken).ConfigureAwait(false); + + List> itemsToStore = context.RequestMessages + .Concat(context.ResponseMessages ?? []) + .Select(message => new Dictionary + { + ["Key"] = Guid.NewGuid(), + ["Role"] = message.Role.ToString(), + ["MessageId"] = message.MessageId, + ["AuthorName"] = message.AuthorName, + ["ApplicationId"] = this._storageScope?.ApplicationId, + ["AgentId"] = this._storageScope?.AgentId, + ["UserId"] = this._storageScope?.UserId, + ["ThreadId"] = this._storageScope?.ThreadId, + ["Content"] = message.Text, + ["CreatedAt"] = message.CreatedAt?.ToString("O") ?? DateTimeOffset.UtcNow.ToString("O"), + ["ContentEmbedding"] = message.Text, + }) + .ToList(); + + if (itemsToStore.Count > 0) + { + await collection.UpsertAsync(itemsToStore, cancellationToken).ConfigureAwait(false); + } + } + catch (Exception ex) + { + this._logger?.LogError( + ex, + "ChatHistoryMemoryProvider: Failed to add messages to chat history vector store due to error. ApplicationId: '{ApplicationId}', AgentId: '{AgentId}', ThreadId: '{ThreadId}', UserId: '{UserId}'.", + this._searchScope.ApplicationId, + this._searchScope.AgentId, + this._searchScope.ThreadId, + this._searchScope.UserId); + } + } + + /// + /// Function callable by the AI model (when enabled) to perform an ad-hoc chat history search. + /// + /// The query text. + /// Cancellation token. + /// Formatted search results (may be empty). + internal async Task SearchTextAsync(string userQuestion, CancellationToken cancellationToken = default) + { + if (string.IsNullOrWhiteSpace(userQuestion)) + { + return string.Empty; + } + + var results = await this.SearchChatHistoryAsync(userQuestion, this._maxResults, cancellationToken).ConfigureAwait(false); + if (!results.Any()) + { + return string.Empty; + } + + // Format the results as a single context message + var outputResultsText = string.Join("\n", results.Select(x => (string?)x["Content"]).Where(c => !string.IsNullOrWhiteSpace(c))); + if (string.IsNullOrWhiteSpace(outputResultsText)) + { + return string.Empty; + } + + var formatted = $"{this._contextPrompt}\n{outputResultsText}"; + + this._logger?.LogTrace( + "ChatHistoryMemoryProvider: Search Results\nInput:{Input}\nOutput:{MessageText}\n ApplicationId: '{ApplicationId}', AgentId: '{AgentId}', ThreadId: '{ThreadId}', UserId: '{UserId}'.", + userQuestion, + formatted, + this._searchScope.ApplicationId, + this._searchScope.AgentId, + this._searchScope.ThreadId, + this._searchScope.UserId); + return formatted; + } + + /// + /// Searches for relevant chat history items based on the provided query text. + /// + /// The text to search for. + /// The maximum number of results to return. + /// The cancellation token. + /// A list of relevant chat history items. + private async Task>> SearchChatHistoryAsync( + string queryText, + int top, + CancellationToken cancellationToken = default) + { + if (string.IsNullOrWhiteSpace(queryText)) + { + return []; + } + + var collection = await this.EnsureCollectionExistsAsync(cancellationToken).ConfigureAwait(false); + + string? applicationId = this._searchScope.ApplicationId; + string? agentId = this._searchScope.AgentId; + string? userId = this._searchScope.UserId; + string? threadId = this._searchScope.ThreadId; + + Expression, bool>>? filter = null; + if (applicationId != null) + { + filter = x => (string?)x["ApplicationId"] == applicationId; + } + + if (agentId != null) + { + Expression, bool>> agentIdFilter = x => (string?)x["AgentId"] == agentId; + filter = filter == null ? agentIdFilter : Expression.Lambda, bool>>( + Expression.AndAlso(filter.Body, agentIdFilter.Body), + filter.Parameters); + } + + if (userId != null) + { + Expression, bool>> userIdFilter = x => (string?)x["UserId"] == userId; + filter = filter == null ? userIdFilter : Expression.Lambda, bool>>( + Expression.AndAlso(filter.Body, userIdFilter.Body), + filter.Parameters); + } + + if (threadId != null) + { + Expression, bool>> threadIdFilter = x => (string?)x["ThreadId"] == threadId; + filter = filter == null ? threadIdFilter : Expression.Lambda, bool>>( + Expression.AndAlso(filter.Body, threadIdFilter.Body), + filter.Parameters); + } + + // Use search to find relevant messages + var searchResults = collection.SearchAsync( + queryText, + top, + options: new() + { + Filter = filter + }, + cancellationToken: cancellationToken); + + var results = new List>(); + await foreach (var result in searchResults.WithCancellation(cancellationToken).ConfigureAwait(false)) + { + results.Add(result.Record); + } + + this._logger?.LogInformation( + "ChatHistoryMemoryProvider: Retrieved {Count} search results. ApplicationId: '{ApplicationId}', AgentId: '{AgentId}', ThreadId: '{ThreadId}', UserId: '{UserId}'.", + results.Count, + this._searchScope.ApplicationId, + this._searchScope.AgentId, + this._searchScope.ThreadId, + this._searchScope.UserId); + + return results; + } + + /// + /// Ensures the collection exists in the vector store, creating it if necessary. + /// + /// The cancellation token. + /// The vector store collection. + private async Task>> EnsureCollectionExistsAsync( + CancellationToken cancellationToken = default) + { + if (this._collectionInitialized) + { + return this._collection; + } + + await this._initializationLock.WaitAsync(cancellationToken).ConfigureAwait(false); + try + { + if (this._collectionInitialized) + { + return this._collection; + } + + await this._collection.EnsureCollectionExistsAsync(cancellationToken).ConfigureAwait(false); + this._collectionInitialized = true; + + return this._collection; + } + finally + { + this._initializationLock.Release(); + } + } + + /// + private void Dispose(bool disposing) + { + if (!this._disposedValue) + { + if (disposing) + { + this._initializationLock.Dispose(); + this._collection?.Dispose(); + } + + this._disposedValue = true; + } + } + + /// + public void Dispose() + { + // Do not change this code. Put cleanup code in 'Dispose(bool disposing)' method + this.Dispose(disposing: true); + 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; + } + + internal sealed class ChatHistoryMemoryProviderState + { + public ChatHistoryMemoryProviderScope? StorageScope { get; set; } + public ChatHistoryMemoryProviderScope? SearchScope { get; set; } + } +} diff --git a/dotnet/src/Microsoft.Agents.AI/Memory/ChatHistoryMemoryProviderOptions.cs b/dotnet/src/Microsoft.Agents.AI/Memory/ChatHistoryMemoryProviderOptions.cs new file mode 100644 index 0000000000..55f06d7429 --- /dev/null +++ b/dotnet/src/Microsoft.Agents.AI/Memory/ChatHistoryMemoryProviderOptions.cs @@ -0,0 +1,56 @@ +// Copyright (c) Microsoft. All rights reserved. + +namespace Microsoft.Agents.AI; + +/// +/// Options controlling the behavior of . +/// +public sealed class ChatHistoryMemoryProviderOptions +{ + /// + /// Gets or sets a value indicating when the search should be executed. + /// + /// by default. + public SearchBehavior SearchTime { get; set; } = SearchBehavior.BeforeAIInvoke; + + /// + /// Gets or sets the name of the exposed search tool when operating in on-demand mode. + /// + /// Defaults to "Search". + public string? FunctionToolName { get; set; } + + /// + /// Gets or sets the description of the exposed search tool when operating in on-demand mode. + /// + /// Defaults to "Allows searching through previous chat history to help answer the user question.". + public string? FunctionToolDescription { get; set; } + + /// + /// Gets or sets the context prompt prefixed to results. + /// + public string? ContextPrompt { get; set; } + + /// + /// Gets or sets the maximum number of results to retrieve from the chat history. + /// + /// + /// Defaults to 3 if not set. + /// + public int? MaxResults { get; set; } + + /// + /// Behavior choices for the provider. + /// + public enum SearchBehavior + { + /// + /// Execute search prior to each invocation and inject results as a message. + /// + BeforeAIInvoke, + + /// + /// Expose a function tool to perform search on-demand via function/tool calling. + /// + OnDemandFunctionCalling + } +} diff --git a/dotnet/src/Microsoft.Agents.AI/Memory/ChatHistoryMemoryProviderScope.cs b/dotnet/src/Microsoft.Agents.AI/Memory/ChatHistoryMemoryProviderScope.cs new file mode 100644 index 0000000000..2715ed2e20 --- /dev/null +++ b/dotnet/src/Microsoft.Agents.AI/Memory/ChatHistoryMemoryProviderScope.cs @@ -0,0 +1,53 @@ +// Copyright (c) Microsoft. All rights reserved. + +using Microsoft.Shared.Diagnostics; + +namespace Microsoft.Agents.AI; + +/// +/// Allows scoping of chat history for the . +/// +public sealed class ChatHistoryMemoryProviderScope +{ + /// + /// Initializes a new instance of the class. + /// + public ChatHistoryMemoryProviderScope() { } + + /// + /// Initializes a new instance of the class by cloning an existing scope. + /// + /// The scope to clone. + public ChatHistoryMemoryProviderScope(ChatHistoryMemoryProviderScope sourceScope) + { + Throw.IfNull(sourceScope); + + this.ApplicationId = sourceScope.ApplicationId; + this.AgentId = sourceScope.AgentId; + this.ThreadId = sourceScope.ThreadId; + this.UserId = sourceScope.UserId; + } + + /// + /// Gets or sets an optional ID for the application to scope chat history to. + /// + /// If not set, the scope of the chat history will span all applications. + public string? ApplicationId { get; set; } + + /// + /// Gets or sets an optional ID for the agent to scope chat history to. + /// + /// If not set, the scope of the chat history will span all agents. + public string? AgentId { get; set; } + + /// + /// Gets or sets an optional ID for the thread to scope chat history to. + /// + public string? ThreadId { get; set; } + + /// + /// Gets or sets an optional ID for the user to scope chat history to. + /// + /// If not set, the scope of the chat history will span all users. + public string? UserId { get; set; } +} diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/Memory/ChatHistoryMemoryProviderTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/Memory/ChatHistoryMemoryProviderTests.cs new file mode 100644 index 0000000000..49e5a5d29c --- /dev/null +++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/Memory/ChatHistoryMemoryProviderTests.cs @@ -0,0 +1,396 @@ +// 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.Extensions.Logging; +using Microsoft.Extensions.VectorData; +using Moq; + +namespace Microsoft.Agents.AI.Memory.UnitTests; + +/// +/// Contains unit tests for the class. +/// +public class ChatHistoryMemoryProviderTests +{ + private readonly Mock> _loggerMock; + private readonly Mock _loggerFactoryMock; + + private readonly Mock _vectorStoreMock; + private readonly Mock>> _vectorStoreCollectionMock; + private const string TestCollectionName = "testcollection"; + + public ChatHistoryMemoryProviderTests() + { + this._loggerMock = new(); + this._loggerFactoryMock = new(); + this._loggerFactoryMock + .Setup(f => f.CreateLogger(It.IsAny())) + .Returns(this._loggerMock.Object); + this._loggerFactoryMock + .Setup(f => f.CreateLogger(typeof(ChatHistoryMemoryProvider).FullName!)) + .Returns(this._loggerMock.Object); + + this._vectorStoreCollectionMock = new(MockBehavior.Strict); + this._vectorStoreMock = new(MockBehavior.Strict); + + this._vectorStoreCollectionMock + .Setup(c => c.EnsureCollectionExistsAsync(It.IsAny())) + .Returns(Task.CompletedTask); + + this._vectorStoreMock + .Setup(vs => vs.GetDynamicCollection( + It.IsAny(), + It.IsAny())) + .Returns(this._vectorStoreCollectionMock.Object); + } + + [Fact] + public void Constructor_Throws_ForNullVectorStore() + { + // Act & Assert + Assert.Throws(() => new ChatHistoryMemoryProvider(null!, "testcollection", 1, 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" })); + } + + [Fact] + public void Constructor_Throws_ForNullStorageScope() + { + // Act & Assert + 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" })); + } + + #region InvokedAsync Tests + + [Fact] + public async Task InvokedAsync_UpsertsMessages_ToCollectionAsync() + { + // 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 storeScope = new ChatHistoryMemoryProviderScope + { + ApplicationId = "app1", + AgentId = "agent1", + ThreadId = "thread1", + UserId = "user1" + }; + + var provider = new ChatHistoryMemoryProvider(this._vectorStoreMock.Object, TestCollectionName, 1, 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([requestMsgWithValues, requestMsgWithNulls], aiContextProviderMessages: null) + { + ResponseMessages = [responseMsg] + }; + + // Act + await provider.InvokedAsync(invokedContext, CancellationToken.None); + + // Assert + this._vectorStoreCollectionMock.Verify( + m => m.EnsureCollectionExistsAsync(It.IsAny()), + Times.Once); + + Assert.Equal(3, stored.Count); + + Assert.Equal("req-1", stored[0]["MessageId"]); + Assert.Equal("request text", stored[0]["Content"]); + Assert.Equal("user1", stored[0]["AuthorName"]); + Assert.Equal(ChatRole.User.ToString(), stored[0]["Role"]); + Assert.Equal("2000-01-01T00:00:00.0000000+00:00", stored[0]["CreatedAt"]); + Assert.Equal("app1", stored[0]["ApplicationId"]); + Assert.Equal("agent1", stored[0]["AgentId"]); + Assert.Equal("thread1", stored[0]["ThreadId"]); + Assert.Equal("user1", stored[0]["UserId"]); + + Assert.Null(stored[1]["MessageId"]); + Assert.Equal("request text nulls", stored[1]["Content"]); + Assert.Null(stored[1]["AuthorName"]); + Assert.Equal(ChatRole.User.ToString(), stored[1]["Role"]); + Assert.Equal("app1", stored[1]["ApplicationId"]); + Assert.Equal("agent1", stored[1]["AgentId"]); + Assert.Equal("thread1", stored[1]["ThreadId"]); + Assert.Equal("user1", stored[1]["UserId"]); + + Assert.Equal("resp-1", stored[2]["MessageId"]); + Assert.Equal("response text", stored[2]["Content"]); + Assert.Equal("assistant", stored[2]["AuthorName"]); + Assert.Equal(ChatRole.Assistant.ToString(), stored[2]["Role"]); + Assert.Equal("app1", stored[2]["ApplicationId"]); + Assert.Equal("agent1", stored[2]["AgentId"]); + Assert.Equal("thread1", stored[2]["ThreadId"]); + Assert.Equal("user1", stored[2]["UserId"]); + } + + [Fact] + public async Task InvokedAsync_DoesNotUpsertMessages_WhenInvokeFailedAsync() + { + // Arrange + this._vectorStoreCollectionMock + .Setup(c => c.UpsertAsync(It.IsAny>>(), It.IsAny())) + .Returns(Task.CompletedTask); + + var provider = new ChatHistoryMemoryProvider( + this._vectorStoreMock.Object, + TestCollectionName, + 1, + new ChatHistoryMemoryProviderScope() { UserId = "UID" }); + var requestMsg = new ChatMessage(ChatRole.User, "request text") { MessageId = "req-1" }; + var invokedContext = new AIContextProvider.InvokedContext([requestMsg], aiContextProviderMessages: null) + { + InvokeException = new InvalidOperationException("Invoke failed") + }; + + // Act + await provider.InvokedAsync(invokedContext, CancellationToken.None); + + // Assert + this._vectorStoreCollectionMock.Verify( + c => c.UpsertAsync(It.IsAny>>(), It.IsAny()), + Times.Never); + } + + [Fact] + public async Task InvokedAsync_DoesNotThrow_WhenUpsertThrowsAsync() + { + // Arrange + this._vectorStoreCollectionMock + .Setup(c => c.UpsertAsync(It.IsAny>>(), It.IsAny())) + .ThrowsAsync(new InvalidOperationException("Upsert failed")); + + var provider = new ChatHistoryMemoryProvider( + this._vectorStoreMock.Object, + TestCollectionName, + 1, + new ChatHistoryMemoryProviderScope() { UserId = "UID" }, + loggerFactory: this._loggerFactoryMock.Object); + var requestMsg = new ChatMessage(ChatRole.User, "request text") { MessageId = "req-1" }; + var invokedContext = new AIContextProvider.InvokedContext([requestMsg], aiContextProviderMessages: null); + + // Act + await provider.InvokedAsync(invokedContext, CancellationToken.None); + + // Assert + this._loggerMock.Verify( + l => l.Log( + LogLevel.Error, + It.IsAny(), + It.Is((v, t) => v.ToString()!.Contains("ChatHistoryMemoryProvider: Failed to add messages to chat history vector store due to error")), + It.IsAny(), + It.IsAny>()), + Times.Once); + } + + #endregion + + #region InvokingAsync Tests + + [Fact] + public async Task InvokedAsync_SearchesVectorStoreAsync() + { + // Arrange + var providerOptions = new ChatHistoryMemoryProviderOptions + { + SearchTime = ChatHistoryMemoryProviderOptions.SearchBehavior.BeforeAIInvoke, + MaxResults = 2, + ContextPrompt = "Here is the relevant chat history:\n" + }; + + var storedItems = new List>> + { + new( + new Dictionary + { + ["MessageId"] = "msg-1", + ["Content"] = "First stored message", + ["Role"] = ChatRole.User.ToString(), + ["CreatedAt"] = "2023-01-01T00:00:00.0000000+00:00" + }, + 0.9f), + new( + new Dictionary + { + ["MessageId"] = "msg-2", + ["Content"] = "Second stored message", + ["Role"] = ChatRole.User.ToString(), + ["CreatedAt"] = "2023-01-02T00:00:00.0000000+00:00" + }, + 0.8f) + }; + + this._vectorStoreCollectionMock + .Setup(c => c.SearchAsync( + It.IsAny(), + It.IsAny(), + It.IsAny>>(), + It.IsAny())) + .Returns(ToAsyncEnumerableAsync(storedItems)); + + var provider = new ChatHistoryMemoryProvider( + this._vectorStoreMock.Object, + TestCollectionName, + 1, + new ChatHistoryMemoryProviderScope() { UserId = "UID" }, + options: providerOptions); + + var requestMsg = new ChatMessage(ChatRole.User, "requesting relevant history"); + var invokingContext = new AIContextProvider.InvokingContext([requestMsg]); + + // Act + await provider.InvokingAsync(invokingContext, CancellationToken.None); + + // Assert + this._vectorStoreCollectionMock.Verify( + c => c.SearchAsync( + It.Is(s => s == "requesting relevant history"), + 2, + It.IsAny>>(), + It.IsAny()), + Times.Once); + } + + [Fact] + public async Task InvokedAsync_CreatesFilter_WhenSearchScopeProvidedAsync() + { + // Arrange + var providerOptions = new ChatHistoryMemoryProviderOptions + { + SearchTime = ChatHistoryMemoryProviderOptions.SearchBehavior.BeforeAIInvoke, + MaxResults = 2, + ContextPrompt = "Here is the relevant chat history:\n" + }; + + var searchScope = new ChatHistoryMemoryProviderScope + { + ApplicationId = "app1", + AgentId = "agent1", + ThreadId = "thread1", + UserId = "user1" + }; + + this._vectorStoreCollectionMock + .Setup(c => c.SearchAsync( + It.IsAny(), + It.IsAny(), + It.IsAny>>(), + It.IsAny())) + .Callback((string query, int maxResults, VectorSearchOptions> options, CancellationToken ct) => + { + // Verify that the filter was created correctly + const string ExpectedFilter = "x => ((((x.ApplicationId == value(Microsoft.Agents.AI.VectorDataMemory.ChatHistoryMemoryProvider+<>c__DisplayClass20_0).applicationId) AndAlso (x.AgentId == value(Microsoft.Agents.AI.VectorDataMemory.ChatHistoryMemoryProvider+<>c__DisplayClass20_0).agentId)) AndAlso (x.UserId == value(Microsoft.Agents.AI.VectorDataMemory.ChatHistoryMemoryProvider+<>c__DisplayClass20_0).userId)) AndAlso (x.ThreadId == value(Microsoft.Agents.AI.VectorDataMemory.ChatHistoryMemoryProvider+<>c__DisplayClass20_0).threadId))"; + Assert.Equal(ExpectedFilter, options.Filter!.ToString()); + }) + .Returns(ToAsyncEnumerableAsync(new List>>())); + + var provider = new ChatHistoryMemoryProvider(this._vectorStoreMock.Object, TestCollectionName, 1, options: providerOptions, storageScope: searchScope, searchScope: searchScope); + + var requestMsg = new ChatMessage(ChatRole.User, "requesting relevant history"); + var invokingContext = new AIContextProvider.InvokingContext([requestMsg]); + + // Act + await provider.InvokingAsync(invokingContext, CancellationToken.None); + + // Assert + this._vectorStoreCollectionMock.Verify( + c => c.SearchAsync( + It.Is(s => s == "requesting relevant history"), + 2, + It.IsAny>>(), + It.IsAny()), + Times.Once); + } + + #endregion + + #region Serialization Tests + + [Fact] + public void Serialize_Deserialize_RoundtripsScopes() + { + // Arrange + var storageScope = new ChatHistoryMemoryProviderScope + { + ApplicationId = "app", + AgentId = "agent", + ThreadId = "thread", + UserId = "user" + }; + + var searchScope = new ChatHistoryMemoryProviderScope + { + ApplicationId = "app2", + AgentId = "agent2", + ThreadId = "thread2", + UserId = "user2" + }; + + var provider = new ChatHistoryMemoryProvider(this._vectorStoreMock.Object, TestCollectionName, 1, storageScope: storageScope, searchScope: searchScope); + + // Act + var stateElement = provider.Serialize(); + + 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("thread", storage.GetProperty("threadId").GetString()); + Assert.Equal("user", storage.GetProperty("userId").GetString()); + + var search = doc.RootElement.GetProperty("searchScope"); + Assert.Equal("app2", search.GetProperty("applicationId").GetString()); + Assert.Equal("agent2", search.GetProperty("agentId").GetString()); + Assert.Equal("thread2", search.GetProperty("threadId").GetString()); + Assert.Equal("user2", search.GetProperty("userId").GetString()); + + // Act - deserialize and serialize again + var provider2 = new ChatHistoryMemoryProvider(this._vectorStoreMock.Object, TestCollectionName, 1, serializedState: stateElement); + var stateElement2 = provider2.Serialize(); + + // Assert - roundtrip the state + Assert.Equal(stateElement.GetRawText(), stateElement2.GetRawText()); + } + + #endregion + + private static async IAsyncEnumerable ToAsyncEnumerableAsync(IEnumerable values) + { + await Task.Yield(); + foreach (var update in values) + { + yield return update; + } + } +} From b7cddaaba870205d3fbaaf58c579c38b41770a57 Mon Sep 17 00:00:00 2001 From: Korolev Dmitry Date: Fri, 7 Nov 2025 12:58:21 +0100 Subject: [PATCH 11/16] .NET: Remove sequential\concurrent workflow extensions (#1731) * remove non-standard workflow extensions, create overloads for IHostedAgentBuilder * rework * rollback extensions in a2a --- .../AgentWebChat.AgentHost/Program.cs | 15 ++- .../DevUI/DevUI_Step01_BasicUsage/Program.cs | 11 ++- ...ostApplicationBuilderWorkflowExtensions.cs | 41 -------- ...plicationBuilderWorkflowExtensionsTests.cs | 96 ------------------- 4 files changed, 20 insertions(+), 143 deletions(-) diff --git a/dotnet/samples/AgentWebChat/AgentWebChat.AgentHost/Program.cs b/dotnet/samples/AgentWebChat/AgentWebChat.AgentHost/Program.cs index d86c53958d..ac905d3aa1 100644 --- a/dotnet/samples/AgentWebChat/AgentWebChat.AgentHost/Program.cs +++ b/dotnet/samples/AgentWebChat/AgentWebChat.AgentHost/Program.cs @@ -78,8 +78,19 @@ var literatureAgent = builder.AddAIAgent("literator", description: "An agent that helps with literature.", chatClientServiceKey: "chat-model"); -builder.AddSequentialWorkflow("science-sequential-workflow", [chemistryAgent, mathsAgent, literatureAgent]).AddAsAIAgent(); -builder.AddConcurrentWorkflow("science-concurrent-workflow", [chemistryAgent, mathsAgent, literatureAgent]).AddAsAIAgent(); +var scienceSequentialWorkflow = builder.AddWorkflow("science-sequential-workflow", (sp, key) => +{ + List usedAgents = [chemistryAgent, mathsAgent, literatureAgent]; + var agents = usedAgents.Select(ab => sp.GetRequiredKeyedService(ab.Name)); + return AgentWorkflowBuilder.BuildSequential(workflowName: key, agents: agents); +}).AddAsAIAgent(); + +var scienceConcurrentWorkflow = builder.AddWorkflow("science-concurrent-workflow", (sp, key) => +{ + List usedAgents = [chemistryAgent, mathsAgent, literatureAgent]; + var agents = usedAgents.Select(ab => sp.GetRequiredKeyedService(ab.Name)); + return AgentWorkflowBuilder.BuildConcurrent(workflowName: key, agents: agents); +}).AddAsAIAgent(); builder.AddOpenAIChatCompletions(); builder.AddOpenAIResponses(); diff --git a/dotnet/samples/GettingStarted/DevUI/DevUI_Step01_BasicUsage/Program.cs b/dotnet/samples/GettingStarted/DevUI/DevUI_Step01_BasicUsage/Program.cs index e2e6e6b727..68a3043f3f 100644 --- a/dotnet/samples/GettingStarted/DevUI/DevUI_Step01_BasicUsage/Program.cs +++ b/dotnet/samples/GettingStarted/DevUI/DevUI_Step01_BasicUsage/Program.cs @@ -4,8 +4,10 @@ using Azure.AI.OpenAI; using Azure.Identity; +using Microsoft.Agents.AI; using Microsoft.Agents.AI.DevUI; using Microsoft.Agents.AI.Hosting; +using Microsoft.Agents.AI.Workflows; using Microsoft.Extensions.AI; namespace DevUI_Step01_BasicUsage; @@ -56,10 +58,11 @@ internal static class Program // Register sample workflows var assistantBuilder = builder.AddAIAgent("workflow-assistant", "You are a helpful assistant in a workflow."); var reviewerBuilder = builder.AddAIAgent("workflow-reviewer", "You are a reviewer. Review and critique the previous response."); - builder.AddSequentialWorkflow( - "review-workflow", - [assistantBuilder, reviewerBuilder]) - .AddAsAIAgent(); + builder.AddWorkflow("review-workflow", (sp, key) => + { + var agents = new List() { assistantBuilder, reviewerBuilder }.Select(ab => sp.GetRequiredKeyedService(ab.Name)); + return AgentWorkflowBuilder.BuildSequential(workflowName: key, agents: agents); + }).AddAsAIAgent(); if (builder.Environment.IsDevelopment()) { diff --git a/dotnet/src/Microsoft.Agents.AI.Hosting/HostApplicationBuilderWorkflowExtensions.cs b/dotnet/src/Microsoft.Agents.AI.Hosting/HostApplicationBuilderWorkflowExtensions.cs index ac78877682..2215a52a69 100644 --- a/dotnet/src/Microsoft.Agents.AI.Hosting/HostApplicationBuilderWorkflowExtensions.cs +++ b/dotnet/src/Microsoft.Agents.AI.Hosting/HostApplicationBuilderWorkflowExtensions.cs @@ -1,7 +1,6 @@ // Copyright (c) Microsoft. All rights reserved. using System; -using System.Collections.Generic; using System.Linq; using Microsoft.Agents.AI.Hosting.Local; using Microsoft.Agents.AI.Workflows; @@ -16,46 +15,6 @@ namespace Microsoft.Agents.AI.Hosting; /// public static class HostApplicationBuilderWorkflowExtensions { - /// - /// Registers a concurrent workflow that executes multiple agents in parallel. - /// - /// The to configure. - /// The unique name for the workflow. - /// A collection of instances representing agents to execute concurrently. - /// An that can be used to further configure the workflow. - /// Thrown when , , or is null. - /// Thrown when or is empty. - public static IHostedWorkflowBuilder AddConcurrentWorkflow(this IHostApplicationBuilder builder, string name, IEnumerable agentBuilders) - { - Throw.IfNullOrEmpty(agentBuilders); - - return builder.AddWorkflow(name, (sp, key) => - { - var agents = agentBuilders.Select(ab => sp.GetRequiredKeyedService(ab.Name)); - return AgentWorkflowBuilder.BuildConcurrent(workflowName: name, agents: agents); - }); - } - - /// - /// Registers a sequential workflow that executes agents in a specific order. - /// - /// The to configure. - /// The unique name for the workflow. - /// A collection of instances representing agents to execute in sequence. - /// An that can be used to further configure the workflow. - /// Thrown when , , or is null. - /// Thrown when or is empty. - public static IHostedWorkflowBuilder AddSequentialWorkflow(this IHostApplicationBuilder builder, string name, IEnumerable agentBuilders) - { - Throw.IfNullOrEmpty(agentBuilders); - - return builder.AddWorkflow(name, (sp, key) => - { - var agents = agentBuilders.Select(ab => sp.GetRequiredKeyedService(ab.Name)); - return AgentWorkflowBuilder.BuildSequential(workflowName: name, agents: agents); - }); - } - /// /// Registers a custom workflow using a factory delegate. /// diff --git a/dotnet/tests/Microsoft.Agents.AI.Hosting.UnitTests/HostApplicationBuilderWorkflowExtensionsTests.cs b/dotnet/tests/Microsoft.Agents.AI.Hosting.UnitTests/HostApplicationBuilderWorkflowExtensionsTests.cs index fc96af2637..a1b7d29f55 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Hosting.UnitTests/HostApplicationBuilderWorkflowExtensionsTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Hosting.UnitTests/HostApplicationBuilderWorkflowExtensionsTests.cs @@ -137,102 +137,6 @@ public class HostApplicationBuilderWorkflowExtensionsTests Assert.NotNull(descriptor); } - /// - /// Verifies that providing a null builder to AddConcurrentWorkflow throws an ArgumentNullException. - /// - [Fact] - public void AddConcurrentWorkflow_NullBuilder_ThrowsArgumentNullException() - { - Assert.Throws(() => - HostApplicationBuilderWorkflowExtensions.AddConcurrentWorkflow(null!, "workflow", [null!])); - } - - /// - /// Verifies that AddConcurrentWorkflow throws ArgumentNullException for null name. - /// - [Fact] - public void AddConcurrentWorkflow_NullName_ThrowsArgumentNullException() - { - var builder = new HostApplicationBuilder(); - - var exception = Assert.Throws(() => - builder.AddConcurrentWorkflow(null!, [new HostedAgentBuilder("test", builder)])); - Assert.Equal("name", exception.ParamName); - } - - /// - /// Verifies that AddConcurrentWorkflow throws ArgumentNullException for null agent builders. - /// - [Fact] - public void AddConcurrentWorkflow_NullAgentBuilders_ThrowsArgumentNullException() - { - var builder = new HostApplicationBuilder(); - - var exception = Assert.Throws(() => - builder.AddConcurrentWorkflow("workflowName", null!)); - Assert.Equal("agentBuilders", exception.ParamName); - } - - /// - /// Verifies that AddConcurrentWorkflow returns IHostWorkflowBuilder instance. - /// - [Fact] - public void AddConcurrentWorkflow_ValidParameters_ReturnsBuilder() - { - var builder = new HostApplicationBuilder(); - - var result = builder.AddConcurrentWorkflow("concurrentWorkflow", [new HostedAgentBuilder("test", builder)]); - - Assert.NotNull(result); - Assert.IsAssignableFrom(result); - } - - /// - /// Verifies that providing a null builder to AddSequentialWorkflow throws an ArgumentNullException. - /// - [Fact] - public void AddSequentialWorkflow_NullBuilder_ThrowsArgumentNullException() - { - Assert.Throws(() => - HostApplicationBuilderWorkflowExtensions.AddSequentialWorkflow(null!, "workflow", [null!])); - } - - /// - /// Verifies that AddSequentialWorkflow throws ArgumentNullException for null name. - /// - [Fact] - public void AddSequentialWorkflow_NullName_ThrowsArgumentNullException() - { - var builder = new HostApplicationBuilder(); - - var exception = Assert.Throws(() => - builder.AddSequentialWorkflow(null!, [new HostedAgentBuilder("test", builder)])); - Assert.Equal("name", exception.ParamName); - } - - /// - /// Verifies that AddSequentialWorkflow throws ArgumentNullException for null agent builders. - /// - [Fact] - public void AddSequentialWorkflow_NullAgentBuilders_ThrowsArgumentNullException() - { - var builder = new HostApplicationBuilder(); - - var exception = Assert.Throws(() => - builder.AddSequentialWorkflow("workflowName", null!)); - Assert.Equal("agentBuilders", exception.ParamName); - } - - [Fact] - public void AddSequentialWorkflow_EmptyAgentBuilders_Throws() - { - var builder = new HostApplicationBuilder(); - - var exception = Assert.Throws(() => - builder.AddSequentialWorkflow("sequentialWorkflow", Array.Empty())); - Assert.Equal("agentBuilders", exception.ParamName); - } - /// /// Verifies that AddAsAIAgent without a name parameter uses the workflow name as the agent name. /// From 32b984b09bf16076e0af1047d72fe0f6ec2e02de Mon Sep 17 00:00:00 2001 From: Korolev Dmitry Date: Fri, 7 Nov 2025 17:28:28 +0100 Subject: [PATCH 12/16] .NET: Add `WithAITool` extensions for Hosting AIAgents (#1990) * add extensions to register tools via fluentAPI * fix --- .../AgentWebChat.AgentHost.csproj | 6 +- .../Custom/CustomAITools.cs | 17 ++ .../AgentWebChat.AgentHost/Program.cs | 3 + ...AgentHostingServiceCollectionExtensions.cs | 22 ++- .../HostedAgentBuilderExtensions.cs | 51 ++++++ .../Local/LocalAgentToolRegistry.cs | 27 ++++ .../Microsoft.Agents.AI.csproj | 4 +- .../HostedAgentBuilderToolsExtensionsTests.cs | 150 ++++++++++++++++++ 8 files changed, 272 insertions(+), 8 deletions(-) create mode 100644 dotnet/samples/AgentWebChat/AgentWebChat.AgentHost/Custom/CustomAITools.cs create mode 100644 dotnet/src/Microsoft.Agents.AI.Hosting/Local/LocalAgentToolRegistry.cs create mode 100644 dotnet/tests/Microsoft.Agents.AI.Hosting.UnitTests/HostedAgentBuilderToolsExtensionsTests.cs diff --git a/dotnet/samples/AgentWebChat/AgentWebChat.AgentHost/AgentWebChat.AgentHost.csproj b/dotnet/samples/AgentWebChat/AgentWebChat.AgentHost/AgentWebChat.AgentHost.csproj index 802c864c1f..53fd4757ee 100644 --- a/dotnet/samples/AgentWebChat/AgentWebChat.AgentHost/AgentWebChat.AgentHost.csproj +++ b/dotnet/samples/AgentWebChat/AgentWebChat.AgentHost/AgentWebChat.AgentHost.csproj @@ -1,4 +1,4 @@ - + net9.0 @@ -13,7 +13,7 @@ - + @@ -37,4 +37,4 @@ - + \ No newline at end of file diff --git a/dotnet/samples/AgentWebChat/AgentWebChat.AgentHost/Custom/CustomAITools.cs b/dotnet/samples/AgentWebChat/AgentWebChat.AgentHost/Custom/CustomAITools.cs new file mode 100644 index 0000000000..d3deb9162c --- /dev/null +++ b/dotnet/samples/AgentWebChat/AgentWebChat.AgentHost/Custom/CustomAITools.cs @@ -0,0 +1,17 @@ +// Copyright (c) Microsoft. All rights reserved. + +using Microsoft.Extensions.AI; + +namespace AgentWebChat.AgentHost.Custom; + +public class CustomAITool : AITool +{ +} + +public class CustomFunctionTool : AIFunction +{ + protected override ValueTask InvokeCoreAsync(AIFunctionArguments arguments, CancellationToken cancellationToken) + { + return new ValueTask(arguments.Context?.Count ?? 0); + } +} diff --git a/dotnet/samples/AgentWebChat/AgentWebChat.AgentHost/Program.cs b/dotnet/samples/AgentWebChat/AgentWebChat.AgentHost/Program.cs index ac905d3aa1..cb1a7e3cd9 100644 --- a/dotnet/samples/AgentWebChat/AgentWebChat.AgentHost/Program.cs +++ b/dotnet/samples/AgentWebChat/AgentWebChat.AgentHost/Program.cs @@ -2,6 +2,7 @@ using A2A.AspNetCore; using AgentWebChat.AgentHost; +using AgentWebChat.AgentHost.Custom; using AgentWebChat.AgentHost.Utilities; using Microsoft.Agents.AI; using Microsoft.Agents.AI.Hosting; @@ -25,6 +26,8 @@ var pirateAgentBuilder = builder.AddAIAgent( instructions: "You are a pirate. Speak like a pirate", description: "An agent that speaks like a pirate.", chatClientServiceKey: "chat-model") + .WithAITool(new CustomAITool()) + .WithAITool(new CustomFunctionTool()) .WithInMemoryThreadStore(); var knightsKnavesAgentBuilder = builder.AddAIAgent("knights-and-knaves", (sp, key) => diff --git a/dotnet/src/Microsoft.Agents.AI.Hosting/AgentHostingServiceCollectionExtensions.cs b/dotnet/src/Microsoft.Agents.AI.Hosting/AgentHostingServiceCollectionExtensions.cs index 3c5a5b84c2..d958fc3578 100644 --- a/dotnet/src/Microsoft.Agents.AI.Hosting/AgentHostingServiceCollectionExtensions.cs +++ b/dotnet/src/Microsoft.Agents.AI.Hosting/AgentHostingServiceCollectionExtensions.cs @@ -1,6 +1,7 @@ // Copyright (c) Microsoft. All rights reserved. using System; +using System.Collections.Generic; using System.Linq; using Microsoft.Agents.AI.Hosting.Local; using Microsoft.Extensions.AI; @@ -29,7 +30,8 @@ public static class AgentHostingServiceCollectionExtensions return services.AddAIAgent(name, (sp, key) => { var chatClient = sp.GetRequiredService(); - return new ChatClientAgent(chatClient, instructions, key); + var tools = GetRegisteredToolsForAgent(sp, name); + return new ChatClientAgent(chatClient, instructions, key, tools: tools); }); } @@ -46,7 +48,11 @@ public static class AgentHostingServiceCollectionExtensions { Throw.IfNull(services); Throw.IfNullOrEmpty(name); - return services.AddAIAgent(name, (sp, key) => new ChatClientAgent(chatClient, instructions, key)); + return services.AddAIAgent(name, (sp, key) => + { + var tools = GetRegisteredToolsForAgent(sp, name); + return new ChatClientAgent(chatClient, instructions, key, tools: tools); + }); } /// @@ -65,7 +71,8 @@ public static class AgentHostingServiceCollectionExtensions return services.AddAIAgent(name, (sp, key) => { var chatClient = chatClientServiceKey is null ? sp.GetRequiredService() : sp.GetRequiredKeyedService(chatClientServiceKey); - return new ChatClientAgent(chatClient, instructions, key); + var tools = GetRegisteredToolsForAgent(sp, name); + return new ChatClientAgent(chatClient, instructions, key, tools: tools); }); } @@ -86,7 +93,8 @@ public static class AgentHostingServiceCollectionExtensions return services.AddAIAgent(name, (sp, key) => { var chatClient = chatClientServiceKey is null ? sp.GetRequiredService() : sp.GetRequiredKeyedService(chatClientServiceKey); - return new ChatClientAgent(chatClient, instructions: instructions, name: key, description: description); + var tools = GetRegisteredToolsForAgent(sp, name); + return new ChatClientAgent(chatClient, instructions: instructions, name: key, description: description, tools: tools); }); } @@ -142,4 +150,10 @@ public static class AgentHostingServiceCollectionExtensions services.Add(ServiceDescriptor.Singleton(agentHostBuilderContext)); services.AddSingleton(); } + + private static IList GetRegisteredToolsForAgent(IServiceProvider serviceProvider, string agentName) + { + var registry = serviceProvider.GetService(); + return registry?.GetTools(agentName) ?? []; + } } diff --git a/dotnet/src/Microsoft.Agents.AI.Hosting/HostedAgentBuilderExtensions.cs b/dotnet/src/Microsoft.Agents.AI.Hosting/HostedAgentBuilderExtensions.cs index 902c54ebe9..e8a55b3baa 100644 --- a/dotnet/src/Microsoft.Agents.AI.Hosting/HostedAgentBuilderExtensions.cs +++ b/dotnet/src/Microsoft.Agents.AI.Hosting/HostedAgentBuilderExtensions.cs @@ -1,6 +1,9 @@ // Copyright (c) Microsoft. All rights reserved. using System; +using System.Linq; +using Microsoft.Agents.AI.Hosting.Local; +using Microsoft.Extensions.AI; using Microsoft.Extensions.DependencyInjection; using Microsoft.Shared.Diagnostics; @@ -59,4 +62,52 @@ public static class HostedAgentBuilderExtensions }); return builder; } + + /// + /// Adds an AI tool to an agent being configured with the service collection. + /// + /// The hosted agent builder. + /// The AI tool to add to the agent. + /// The same instance so that additional calls can be chained. + /// Thrown when or is . + public static IHostedAgentBuilder WithAITool(this IHostedAgentBuilder builder, AITool tool) + { + Throw.IfNull(builder); + Throw.IfNull(tool); + + var agentName = builder.Name; + var services = builder.ServiceCollection; + + // Get or create the agent tool registry + var descriptor = services.FirstOrDefault(sd => !sd.IsKeyedService && sd.ServiceType.Equals(typeof(LocalAgentToolRegistry))); + if (descriptor?.ImplementationInstance is not LocalAgentToolRegistry toolRegistry) + { + toolRegistry = new(); + services.Add(ServiceDescriptor.Singleton(toolRegistry)); + } + + toolRegistry.AddTool(agentName, tool); + + return builder; + } + + /// + /// Adds multiple AI tools to an agent being configured with the service collection. + /// + /// The hosted agent builder. + /// The collection of AI tools to add to the agent. + /// The same instance so that additional calls can be chained. + /// Thrown when or is . + public static IHostedAgentBuilder WithAITools(this IHostedAgentBuilder builder, params AITool[] tools) + { + Throw.IfNull(builder); + Throw.IfNull(tools); + + foreach (var tool in tools) + { + builder.WithAITool(tool); + } + + return builder; + } } diff --git a/dotnet/src/Microsoft.Agents.AI.Hosting/Local/LocalAgentToolRegistry.cs b/dotnet/src/Microsoft.Agents.AI.Hosting/Local/LocalAgentToolRegistry.cs new file mode 100644 index 0000000000..ea8d8ad74e --- /dev/null +++ b/dotnet/src/Microsoft.Agents.AI.Hosting/Local/LocalAgentToolRegistry.cs @@ -0,0 +1,27 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System.Collections.Generic; +using Microsoft.Extensions.AI; + +namespace Microsoft.Agents.AI.Hosting.Local; + +internal sealed class LocalAgentToolRegistry +{ + private readonly Dictionary> _toolsByAgentName = new(); + + public void AddTool(string agentName, AITool tool) + { + if (!this._toolsByAgentName.TryGetValue(agentName, out var tools)) + { + tools = []; + this._toolsByAgentName[agentName] = tools; + } + + tools.Add(tool); + } + + public IList GetTools(string agentName) + { + return this._toolsByAgentName.TryGetValue(agentName, out var tools) ? tools : []; + } +} diff --git a/dotnet/src/Microsoft.Agents.AI/Microsoft.Agents.AI.csproj b/dotnet/src/Microsoft.Agents.AI/Microsoft.Agents.AI.csproj index 59345d21ae..e3d7f00aa1 100644 --- a/dotnet/src/Microsoft.Agents.AI/Microsoft.Agents.AI.csproj +++ b/dotnet/src/Microsoft.Agents.AI/Microsoft.Agents.AI.csproj @@ -31,8 +31,10 @@ - + + + diff --git a/dotnet/tests/Microsoft.Agents.AI.Hosting.UnitTests/HostedAgentBuilderToolsExtensionsTests.cs b/dotnet/tests/Microsoft.Agents.AI.Hosting.UnitTests/HostedAgentBuilderToolsExtensionsTests.cs new file mode 100644 index 0000000000..9993007de1 --- /dev/null +++ b/dotnet/tests/Microsoft.Agents.AI.Hosting.UnitTests/HostedAgentBuilderToolsExtensionsTests.cs @@ -0,0 +1,150 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Collections.Generic; +using System.Threading; +using System.Threading.Tasks; +using Microsoft.Extensions.AI; +using Microsoft.Extensions.DependencyInjection; + +namespace Microsoft.Agents.AI.Hosting.UnitTests; + +/// +/// Unit tests for AI tool registration extensions on . +/// +public sealed class HostedAgentBuilderToolsExtensionsTests +{ + [Fact] + public void WithAITool_ThrowsWhenBuilderIsNull() + { + // Arrange + var tool = new DummyAITool(); + + // Act & Assert + Assert.Throws(() => HostedAgentBuilderExtensions.WithAITool(null!, tool)); + } + + [Fact] + public void WithAITool_ThrowsWhenToolIsNull() + { + // Arrange + var services = new ServiceCollection(); + var builder = services.AddAIAgent("test-agent", "Test instructions"); + + // Act & Assert + Assert.Throws(() => builder.WithAITool(null!)); + } + + [Fact] + public void WithAITools_ThrowsWhenBuilderIsNull() + { + // Arrange + var tools = new[] { new DummyAITool() }; + + // Act & Assert + Assert.Throws(() => HostedAgentBuilderExtensions.WithAITools(null!, tools)); + } + + [Fact] + public void WithAITools_ThrowsWhenToolsArrayIsNull() + { + // Arrange + var services = new ServiceCollection(); + var builder = services.AddAIAgent("test-agent", "Test instructions"); + + // Act & Assert + Assert.Throws(() => builder.WithAITools(null!)); + } + + [Fact] + public void RegisteredTools_ResolvesAllToolsForAgent() + { + // Arrange + var services = new ServiceCollection(); + services.AddSingleton(new MockChatClient()); + + var builder = services.AddAIAgent("test-agent", "Test instructions"); + var tool1 = new DummyAITool(); + var tool2 = new DummyAITool(); + + builder + .WithAITool(tool1) + .WithAITool(tool2); + + var serviceProvider = services.BuildServiceProvider(); + + var agent1Tools = ResolveAgentTools(serviceProvider, "test-agent"); + Assert.Contains(tool1, agent1Tools); + Assert.Contains(tool2, agent1Tools); + } + + [Fact] + public void RegisteredTools_IsolatedPerAgent() + { + var services = new ServiceCollection(); + services.AddSingleton(new MockChatClient()); + + var builder1 = services.AddAIAgent("agent1", "Agent 1 instructions"); + var builder2 = services.AddAIAgent("agent2", "Agent 2 instructions"); + + var tool1 = new DummyAITool(); + var tool2 = new DummyAITool(); + var tool3 = new DummyAITool(); + + builder1 + .WithAITool(tool1) + .WithAITool(tool2); + + builder2 + .WithAITool(tool3); + + var serviceProvider = services.BuildServiceProvider(); + + var agent1Tools = ResolveAgentTools(serviceProvider, "agent1"); + var agent2Tools = ResolveAgentTools(serviceProvider, "agent2"); + + Assert.Contains(tool1, agent1Tools); + Assert.Contains(tool2, agent1Tools); + Assert.Contains(tool3, agent2Tools); + } + + private static IList ResolveAgentTools(IServiceProvider serviceProvider, string name) + { + var agent = serviceProvider.GetRequiredKeyedService(name) as ChatClientAgent; + Assert.NotNull(agent?.ChatOptions?.Tools); + return agent.ChatOptions.Tools; + } + + /// + /// Dummy AITool implementation for testing. + /// + private sealed class DummyAITool : AITool + { + } + + /// + /// Mock chat client for testing. + /// + private sealed class MockChatClient : IChatClient + { + public Task GetResponseAsync(IEnumerable messages, ChatOptions? options = null, CancellationToken cancellationToken = default) + { + throw new NotImplementedException(); + } + + public IAsyncEnumerable GetStreamingResponseAsync(IEnumerable messages, ChatOptions? options = null, CancellationToken cancellationToken = default) + { + throw new NotImplementedException(); + } + + public object? GetService(Type serviceType, object? serviceKey = null) + { + return null; + } + + public void Dispose() + { + throw new NotImplementedException(); + } + } +} From 00b67e191be890f6ef6a09896a50f98b32241dde Mon Sep 17 00:00:00 2001 From: Korolev Dmitry Date: Fri, 7 Nov 2025 17:32:06 +0100 Subject: [PATCH 13/16] move a2a tests; add extensions for a2a+openairesponses (#1992) --- dotnet/agent-framework-dotnet.slnx | 2 +- .../EndpointRouteBuilderExtensions.cs | 78 +++ .../Microsoft.Agents.AI.Hosting.A2A.csproj | 2 +- ...ndpointRouteBuilderExtensions.Responses.cs | 24 + .../Converters/MessageConverterTests.cs | 2 +- .../EndpointRouteA2ABuilderExtensionsTests.cs | 502 ++++++++++++++++++ ...ft.Agents.AI.Hosting.A2A.UnitTests.csproj} | 1 + .../EndpointRouteBuilderExtensionsTests.cs | 170 ++++++ 8 files changed, 778 insertions(+), 3 deletions(-) rename dotnet/tests/{Microsoft.Agents.AI.Hosting.A2A.Tests => Microsoft.Agents.AI.Hosting.A2A.UnitTests}/Converters/MessageConverterTests.cs (97%) create mode 100644 dotnet/tests/Microsoft.Agents.AI.Hosting.A2A.UnitTests/EndpointRouteA2ABuilderExtensionsTests.cs rename dotnet/tests/{Microsoft.Agents.AI.Hosting.A2A.Tests/Microsoft.Agents.AI.Hosting.A2A.Tests.csproj => Microsoft.Agents.AI.Hosting.A2A.UnitTests/Microsoft.Agents.AI.Hosting.A2A.UnitTests.csproj} (82%) diff --git a/dotnet/agent-framework-dotnet.slnx b/dotnet/agent-framework-dotnet.slnx index 56211e457f..3e845c75d9 100644 --- a/dotnet/agent-framework-dotnet.slnx +++ b/dotnet/agent-framework-dotnet.slnx @@ -314,7 +314,7 @@ - + diff --git a/dotnet/src/Microsoft.Agents.AI.Hosting.A2A.AspNetCore/EndpointRouteBuilderExtensions.cs b/dotnet/src/Microsoft.Agents.AI.Hosting.A2A.AspNetCore/EndpointRouteBuilderExtensions.cs index cae9801148..b6b9c4da48 100644 --- a/dotnet/src/Microsoft.Agents.AI.Hosting.A2A.AspNetCore/EndpointRouteBuilderExtensions.cs +++ b/dotnet/src/Microsoft.Agents.AI.Hosting.A2A.AspNetCore/EndpointRouteBuilderExtensions.cs @@ -18,6 +18,21 @@ namespace Microsoft.AspNetCore.Builder; /// public static class MicrosoftAgentAIHostingA2AEndpointRouteBuilderExtensions { + /// + /// Attaches A2A (Agent2Agent) communication capabilities via Message processing to the specified web application. + /// + /// The to add the A2A endpoints to. + /// The configuration builder for . + /// The route group to use for A2A endpoints. + /// Configured for A2A integration. + /// + /// This method can be used to access A2A agents that support the + /// Curated Registries (Catalog-Based Discovery) + /// discovery mechanism. + /// + public static IEndpointConventionBuilder MapA2A(this IEndpointRouteBuilder endpoints, IHostedAgentBuilder agentBuilder, string path) + => endpoints.MapA2A(agentBuilder, path, _ => { }); + /// /// Attaches A2A (Agent2Agent) communication capabilities via Message processing to the specified web application. /// @@ -28,6 +43,25 @@ public static class MicrosoftAgentAIHostingA2AEndpointRouteBuilderExtensions public static IEndpointConventionBuilder MapA2A(this IEndpointRouteBuilder endpoints, string agentName, string path) => endpoints.MapA2A(agentName, path, _ => { }); + /// + /// Attaches A2A (Agent2Agent) communication capabilities via Message processing to the specified web application. + /// + /// The to add the A2A endpoints to. + /// The configuration builder for . + /// The route group to use for A2A endpoints. + /// The callback to configure . + /// Configured for A2A integration. + /// + /// This method can be used to access A2A agents that support the + /// Curated Registries (Catalog-Based Discovery) + /// discovery mechanism. + /// + public static IEndpointConventionBuilder MapA2A(this IEndpointRouteBuilder endpoints, IHostedAgentBuilder agentBuilder, string path, Action configureTaskManager) + { + ArgumentNullException.ThrowIfNull(agentBuilder); + return endpoints.MapA2A(agentBuilder.Name, path, configureTaskManager); + } + /// /// Attaches A2A (Agent2Agent) communication capabilities via Message processing to the specified web application. /// @@ -38,10 +72,27 @@ public static class MicrosoftAgentAIHostingA2AEndpointRouteBuilderExtensions /// Configured for A2A integration. public static IEndpointConventionBuilder MapA2A(this IEndpointRouteBuilder endpoints, string agentName, string path, Action configureTaskManager) { + ArgumentNullException.ThrowIfNull(endpoints); var agent = endpoints.ServiceProvider.GetRequiredKeyedService(agentName); return endpoints.MapA2A(agent, path, configureTaskManager); } + /// + /// Attaches A2A (Agent2Agent) communication capabilities via Message processing to the specified web application. + /// + /// The to add the A2A endpoints to. + /// The configuration builder for . + /// The route group to use for A2A endpoints. + /// Agent card info to return on query. + /// Configured for A2A integration. + /// + /// This method can be used to access A2A agents that support the + /// Curated Registries (Catalog-Based Discovery) + /// discovery mechanism. + /// + public static IEndpointConventionBuilder MapA2A(this IEndpointRouteBuilder endpoints, IHostedAgentBuilder agentBuilder, string path, AgentCard agentCard) + => endpoints.MapA2A(agentBuilder, path, agentCard, _ => { }); + /// /// Attaches A2A (Agent2Agent) communication capabilities via Message processing to the specified web application. /// @@ -58,6 +109,26 @@ public static class MicrosoftAgentAIHostingA2AEndpointRouteBuilderExtensions public static IEndpointConventionBuilder MapA2A(this IEndpointRouteBuilder endpoints, string agentName, string path, AgentCard agentCard) => endpoints.MapA2A(agentName, path, agentCard, _ => { }); + /// + /// Attaches A2A (Agent2Agent) communication capabilities via Message processing to the specified web application. + /// + /// The to add the A2A endpoints to. + /// The configuration builder for . + /// The route group to use for A2A endpoints. + /// Agent card info to return on query. + /// The callback to configure . + /// Configured for A2A integration. + /// + /// This method can be used to access A2A agents that support the + /// Curated Registries (Catalog-Based Discovery) + /// discovery mechanism. + /// + public static IEndpointConventionBuilder MapA2A(this IEndpointRouteBuilder endpoints, IHostedAgentBuilder agentBuilder, string path, AgentCard agentCard, Action configureTaskManager) + { + ArgumentNullException.ThrowIfNull(agentBuilder); + return endpoints.MapA2A(agentBuilder.Name, path, agentCard, configureTaskManager); + } + /// /// Attaches A2A (Agent2Agent) communication capabilities via Message processing to the specified web application. /// @@ -74,6 +145,7 @@ public static class MicrosoftAgentAIHostingA2AEndpointRouteBuilderExtensions /// public static IEndpointConventionBuilder MapA2A(this IEndpointRouteBuilder endpoints, string agentName, string path, AgentCard agentCard, Action configureTaskManager) { + ArgumentNullException.ThrowIfNull(endpoints); var agent = endpoints.ServiceProvider.GetRequiredKeyedService(agentName); return endpoints.MapA2A(agent, path, agentCard, configureTaskManager); } @@ -98,6 +170,9 @@ public static class MicrosoftAgentAIHostingA2AEndpointRouteBuilderExtensions /// Configured for A2A integration. public static IEndpointConventionBuilder MapA2A(this IEndpointRouteBuilder endpoints, AIAgent agent, string path, Action configureTaskManager) { + ArgumentNullException.ThrowIfNull(endpoints); + ArgumentNullException.ThrowIfNull(agent); + var loggerFactory = endpoints.ServiceProvider.GetRequiredService(); var agentThreadStore = endpoints.ServiceProvider.GetKeyedService(agent.Name); var taskManager = agent.MapA2A(loggerFactory: loggerFactory, agentThreadStore: agentThreadStore); @@ -139,6 +214,9 @@ public static class MicrosoftAgentAIHostingA2AEndpointRouteBuilderExtensions /// public static IEndpointConventionBuilder MapA2A(this IEndpointRouteBuilder endpoints, AIAgent agent, string path, AgentCard agentCard, Action configureTaskManager) { + ArgumentNullException.ThrowIfNull(endpoints); + ArgumentNullException.ThrowIfNull(agent); + var loggerFactory = endpoints.ServiceProvider.GetRequiredService(); var agentThreadStore = endpoints.ServiceProvider.GetKeyedService(agent.Name); var taskManager = agent.MapA2A(agentCard: agentCard, agentThreadStore: agentThreadStore, loggerFactory: loggerFactory); diff --git a/dotnet/src/Microsoft.Agents.AI.Hosting.A2A/Microsoft.Agents.AI.Hosting.A2A.csproj b/dotnet/src/Microsoft.Agents.AI.Hosting.A2A/Microsoft.Agents.AI.Hosting.A2A.csproj index 5076e25c05..f300483f63 100644 --- a/dotnet/src/Microsoft.Agents.AI.Hosting.A2A/Microsoft.Agents.AI.Hosting.A2A.csproj +++ b/dotnet/src/Microsoft.Agents.AI.Hosting.A2A/Microsoft.Agents.AI.Hosting.A2A.csproj @@ -29,6 +29,6 @@ - + diff --git a/dotnet/src/Microsoft.Agents.AI.Hosting.OpenAI/EndpointRouteBuilderExtensions.Responses.cs b/dotnet/src/Microsoft.Agents.AI.Hosting.OpenAI/EndpointRouteBuilderExtensions.Responses.cs index f4159011cf..9a395b9b12 100644 --- a/dotnet/src/Microsoft.Agents.AI.Hosting.OpenAI/EndpointRouteBuilderExtensions.Responses.cs +++ b/dotnet/src/Microsoft.Agents.AI.Hosting.OpenAI/EndpointRouteBuilderExtensions.Responses.cs @@ -3,6 +3,7 @@ using System; using System.Diagnostics.CodeAnalysis; using Microsoft.Agents.AI; +using Microsoft.Agents.AI.Hosting; using Microsoft.Agents.AI.Hosting.OpenAI; using Microsoft.Agents.AI.Hosting.OpenAI.Conversations; using Microsoft.Agents.AI.Hosting.OpenAI.Responses; @@ -17,6 +18,29 @@ namespace Microsoft.AspNetCore.Builder; /// public static partial class MicrosoftAgentAIHostingOpenAIEndpointRouteBuilderExtensions { + /// + /// Maps OpenAI Responses API endpoints to the specified for the given . + /// + /// The to add the OpenAI Responses endpoints to. + /// The builder for to map the OpenAI Responses endpoints for. + public static IEndpointConventionBuilder MapOpenAIResponses(this IEndpointRouteBuilder endpoints, IHostedAgentBuilder agentBuilder) + => MapOpenAIResponses(endpoints, agentBuilder, path: null); + + /// + /// Maps OpenAI Responses API endpoints to the specified for the given . + /// + /// The to add the OpenAI Responses endpoints to. + /// The builder for to map the OpenAI Responses endpoints for. + /// Custom route path for the OpenAI Responses endpoint. + public static IEndpointConventionBuilder MapOpenAIResponses(this IEndpointRouteBuilder endpoints, IHostedAgentBuilder agentBuilder, string? path) + { + ArgumentNullException.ThrowIfNull(endpoints); + ArgumentNullException.ThrowIfNull(agentBuilder); + + var agent = endpoints.ServiceProvider.GetRequiredKeyedService(agentBuilder.Name); + return MapOpenAIResponses(endpoints, agent, path); + } + /// /// Maps OpenAI Responses API endpoints to the specified for the given . /// diff --git a/dotnet/tests/Microsoft.Agents.AI.Hosting.A2A.Tests/Converters/MessageConverterTests.cs b/dotnet/tests/Microsoft.Agents.AI.Hosting.A2A.UnitTests/Converters/MessageConverterTests.cs similarity index 97% rename from dotnet/tests/Microsoft.Agents.AI.Hosting.A2A.Tests/Converters/MessageConverterTests.cs rename to dotnet/tests/Microsoft.Agents.AI.Hosting.A2A.UnitTests/Converters/MessageConverterTests.cs index 81ce582870..69eaf3a535 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Hosting.A2A.Tests/Converters/MessageConverterTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Hosting.A2A.UnitTests/Converters/MessageConverterTests.cs @@ -5,7 +5,7 @@ using A2A; using Microsoft.Agents.AI.Hosting.A2A.Converters; using Microsoft.Extensions.AI; -namespace Microsoft.Agents.AI.Hosting.A2A.Tests.Converters; +namespace Microsoft.Agents.AI.Hosting.A2A.UnitTests.Converters; public class MessageConverterTests { diff --git a/dotnet/tests/Microsoft.Agents.AI.Hosting.A2A.UnitTests/EndpointRouteA2ABuilderExtensionsTests.cs b/dotnet/tests/Microsoft.Agents.AI.Hosting.A2A.UnitTests/EndpointRouteA2ABuilderExtensionsTests.cs new file mode 100644 index 0000000000..1ae0dda908 --- /dev/null +++ b/dotnet/tests/Microsoft.Agents.AI.Hosting.A2A.UnitTests/EndpointRouteA2ABuilderExtensionsTests.cs @@ -0,0 +1,502 @@ +// Copyright (c) Microsoft. All rights reserved. + +using System; +using System.Collections.Generic; +using System.Threading; +using System.Threading.Tasks; +using A2A; +using Microsoft.AspNetCore.Builder; +using Microsoft.Extensions.AI; +using Microsoft.Extensions.DependencyInjection; + +namespace Microsoft.Agents.AI.Hosting.A2A.UnitTests; + +/// +/// Tests for MicrosoftAgentAIHostingA2AEndpointRouteBuilderExtensions.MapA2A method. +/// +public sealed class EndpointRouteA2ABuilderExtensionsTests +{ + /// + /// Verifies that MapA2A throws ArgumentNullException for null endpoints. + /// + [Fact] + public void MapA2A_WithAgentBuilder_NullEndpoints_ThrowsArgumentNullException() + { + // Arrange + AspNetCore.Routing.IEndpointRouteBuilder endpoints = null!; + WebApplicationBuilder builder = WebApplication.CreateBuilder(); + IChatClient mockChatClient = new DummyChatClient(); + builder.Services.AddKeyedSingleton("chat-client", mockChatClient); + IHostedAgentBuilder agentBuilder = builder.AddAIAgent("agent", "Instructions", chatClientServiceKey: "chat-client"); + + // Act & Assert + ArgumentNullException exception = Assert.Throws(() => + endpoints.MapA2A(agentBuilder, "/a2a")); + + Assert.Equal("endpoints", exception.ParamName); + } + + /// + /// Verifies that MapA2A throws ArgumentNullException for null agentBuilder. + /// + [Fact] + public void MapA2A_WithAgentBuilder_NullAgentBuilder_ThrowsArgumentNullException() + { + // Arrange + WebApplicationBuilder builder = WebApplication.CreateBuilder(); + IChatClient mockChatClient = new DummyChatClient(); + builder.Services.AddKeyedSingleton("chat-client", mockChatClient); + builder.AddAIAgent("agent", "Instructions", chatClientServiceKey: "chat-client"); + builder.Services.AddLogging(); + using WebApplication app = builder.Build(); + IHostedAgentBuilder agentBuilder = null!; + + // Act & Assert + ArgumentNullException exception = Assert.Throws(() => + app.MapA2A(agentBuilder, "/a2a")); + + Assert.Equal("agentBuilder", exception.ParamName); + } + + /// + /// Verifies that MapA2A with IHostedAgentBuilder correctly maps the agent with default task manager configuration. + /// + [Fact] + public void MapA2A_WithAgentBuilder_DefaultConfiguration_Succeeds() + { + // Arrange + WebApplicationBuilder builder = WebApplication.CreateBuilder(); + IChatClient mockChatClient = new DummyChatClient(); + builder.Services.AddKeyedSingleton("chat-client", mockChatClient); + IHostedAgentBuilder agentBuilder = builder.AddAIAgent("agent", "Instructions", chatClientServiceKey: "chat-client"); + builder.Services.AddLogging(); + using WebApplication app = builder.Build(); + + // Act & Assert - Should not throw + var result = app.MapA2A(agentBuilder, "/a2a"); + Assert.NotNull(result); + Assert.NotNull(app); + } + + /// + /// Verifies that MapA2A with IHostedAgentBuilder and custom task manager configuration succeeds. + /// + [Fact] + public void MapA2A_WithAgentBuilder_CustomTaskManagerConfiguration_Succeeds() + { + // Arrange + WebApplicationBuilder builder = WebApplication.CreateBuilder(); + IChatClient mockChatClient = new DummyChatClient(); + builder.Services.AddKeyedSingleton("chat-client", mockChatClient); + IHostedAgentBuilder agentBuilder = builder.AddAIAgent("agent", "Instructions", chatClientServiceKey: "chat-client"); + builder.Services.AddLogging(); + using WebApplication app = builder.Build(); + + // Act & Assert - Should not throw + var result = app.MapA2A(agentBuilder, "/a2a", taskManager => { }); + Assert.NotNull(result); + Assert.NotNull(app); + } + + /// + /// Verifies that MapA2A with IHostedAgentBuilder and agent card succeeds. + /// + [Fact] + public void MapA2A_WithAgentBuilder_WithAgentCard_Succeeds() + { + // Arrange + WebApplicationBuilder builder = WebApplication.CreateBuilder(); + IChatClient mockChatClient = new DummyChatClient(); + builder.Services.AddKeyedSingleton("chat-client", mockChatClient); + IHostedAgentBuilder agentBuilder = builder.AddAIAgent("agent", "Instructions", chatClientServiceKey: "chat-client"); + builder.Services.AddLogging(); + using WebApplication app = builder.Build(); + + var agentCard = new AgentCard + { + Name = "Test Agent", + Description = "A test agent for A2A communication" + }; + + // Act & Assert - Should not throw + var result = app.MapA2A(agentBuilder, "/a2a", agentCard); + Assert.NotNull(result); + Assert.NotNull(app); + } + + /// + /// Verifies that MapA2A with IHostedAgentBuilder, agent card, and custom task manager configuration succeeds. + /// + [Fact] + public void MapA2A_WithAgentBuilder_WithAgentCardAndCustomConfiguration_Succeeds() + { + // Arrange + WebApplicationBuilder builder = WebApplication.CreateBuilder(); + IChatClient mockChatClient = new DummyChatClient(); + builder.Services.AddKeyedSingleton("chat-client", mockChatClient); + IHostedAgentBuilder agentBuilder = builder.AddAIAgent("agent", "Instructions", chatClientServiceKey: "chat-client"); + builder.Services.AddLogging(); + using WebApplication app = builder.Build(); + + var agentCard = new AgentCard + { + Name = "Test Agent", + Description = "A test agent for A2A communication" + }; + + // Act & Assert - Should not throw + var result = app.MapA2A(agentBuilder, "/a2a", agentCard, taskManager => { }); + Assert.NotNull(result); + Assert.NotNull(app); + } + + /// + /// Verifies that MapA2A throws ArgumentNullException for null endpoints when using string agent name. + /// + [Fact] + public void MapA2A_WithAgentName_NullEndpoints_ThrowsArgumentNullException() + { + // Arrange + AspNetCore.Routing.IEndpointRouteBuilder endpoints = null!; + + // Act & Assert + ArgumentNullException exception = Assert.Throws(() => + endpoints.MapA2A("agent", "/a2a")); + + Assert.Equal("endpoints", exception.ParamName); + } + + /// + /// Verifies that MapA2A with string agent name correctly maps the agent. + /// + [Fact] + public void MapA2A_WithAgentName_DefaultConfiguration_Succeeds() + { + // Arrange + WebApplicationBuilder builder = WebApplication.CreateBuilder(); + IChatClient mockChatClient = new DummyChatClient(); + builder.Services.AddKeyedSingleton("chat-client", mockChatClient); + builder.AddAIAgent("agent", "Instructions", chatClientServiceKey: "chat-client"); + builder.Services.AddLogging(); + using WebApplication app = builder.Build(); + + // Act & Assert - Should not throw + var result = app.MapA2A("agent", "/a2a"); + Assert.NotNull(result); + Assert.NotNull(app); + } + + /// + /// Verifies that MapA2A with string agent name and custom task manager configuration succeeds. + /// + [Fact] + public void MapA2A_WithAgentName_CustomTaskManagerConfiguration_Succeeds() + { + // Arrange + WebApplicationBuilder builder = WebApplication.CreateBuilder(); + IChatClient mockChatClient = new DummyChatClient(); + builder.Services.AddKeyedSingleton("chat-client", mockChatClient); + builder.AddAIAgent("agent", "Instructions", chatClientServiceKey: "chat-client"); + builder.Services.AddLogging(); + using WebApplication app = builder.Build(); + + // Act & Assert - Should not throw + var result = app.MapA2A("agent", "/a2a", taskManager => { }); + Assert.NotNull(result); + Assert.NotNull(app); + } + + /// + /// Verifies that MapA2A with string agent name and agent card succeeds. + /// + [Fact] + public void MapA2A_WithAgentName_WithAgentCard_Succeeds() + { + // Arrange + WebApplicationBuilder builder = WebApplication.CreateBuilder(); + IChatClient mockChatClient = new DummyChatClient(); + builder.Services.AddKeyedSingleton("chat-client", mockChatClient); + builder.AddAIAgent("agent", "Instructions", chatClientServiceKey: "chat-client"); + builder.Services.AddLogging(); + using WebApplication app = builder.Build(); + + var agentCard = new AgentCard + { + Name = "Test Agent", + Description = "A test agent for A2A communication" + }; + + // Act & Assert - Should not throw + var result = app.MapA2A("agent", "/a2a", agentCard); + Assert.NotNull(result); + Assert.NotNull(app); + } + + /// + /// Verifies that MapA2A with string agent name, agent card, and custom task manager configuration succeeds. + /// + [Fact] + public void MapA2A_WithAgentName_WithAgentCardAndCustomConfiguration_Succeeds() + { + // Arrange + WebApplicationBuilder builder = WebApplication.CreateBuilder(); + IChatClient mockChatClient = new DummyChatClient(); + builder.Services.AddKeyedSingleton("chat-client", mockChatClient); + builder.AddAIAgent("agent", "Instructions", chatClientServiceKey: "chat-client"); + builder.Services.AddLogging(); + using WebApplication app = builder.Build(); + + var agentCard = new AgentCard + { + Name = "Test Agent", + Description = "A test agent for A2A communication" + }; + + // Act & Assert - Should not throw + var result = app.MapA2A("agent", "/a2a", agentCard, taskManager => { }); + Assert.NotNull(result); + Assert.NotNull(app); + } + + /// + /// Verifies that MapA2A throws ArgumentNullException for null endpoints when using AIAgent. + /// + [Fact] + public void MapA2A_WithAIAgent_NullEndpoints_ThrowsArgumentNullException() + { + // Arrange + AspNetCore.Routing.IEndpointRouteBuilder endpoints = null!; + + // Act & Assert + ArgumentNullException exception = Assert.Throws(() => + endpoints.MapA2A((AIAgent)null!, "/a2a")); + + Assert.Equal("endpoints", exception.ParamName); + } + + /// + /// Verifies that MapA2A with AIAgent correctly maps the agent. + /// + [Fact] + public void MapA2A_WithAIAgent_DefaultConfiguration_Succeeds() + { + // Arrange + WebApplicationBuilder builder = WebApplication.CreateBuilder(); + IChatClient mockChatClient = new DummyChatClient(); + builder.Services.AddKeyedSingleton("chat-client", mockChatClient); + builder.AddAIAgent("agent", "Instructions", chatClientServiceKey: "chat-client"); + builder.Services.AddLogging(); + using WebApplication app = builder.Build(); + AIAgent agent = app.Services.GetRequiredKeyedService("agent"); + + // Act & Assert - Should not throw + var result = app.MapA2A(agent, "/a2a"); + Assert.NotNull(result); + Assert.NotNull(app); + } + + /// + /// Verifies that MapA2A with AIAgent and custom task manager configuration succeeds. + /// + [Fact] + public void MapA2A_WithAIAgent_CustomTaskManagerConfiguration_Succeeds() + { + // Arrange + WebApplicationBuilder builder = WebApplication.CreateBuilder(); + IChatClient mockChatClient = new DummyChatClient(); + builder.Services.AddKeyedSingleton("chat-client", mockChatClient); + builder.AddAIAgent("agent", "Instructions", chatClientServiceKey: "chat-client"); + builder.Services.AddLogging(); + using WebApplication app = builder.Build(); + AIAgent agent = app.Services.GetRequiredKeyedService("agent"); + + // Act & Assert - Should not throw + var result = app.MapA2A(agent, "/a2a", taskManager => { }); + Assert.NotNull(result); + Assert.NotNull(app); + } + + /// + /// Verifies that MapA2A with AIAgent and agent card succeeds. + /// + [Fact] + public void MapA2A_WithAIAgent_WithAgentCard_Succeeds() + { + // Arrange + WebApplicationBuilder builder = WebApplication.CreateBuilder(); + IChatClient mockChatClient = new DummyChatClient(); + builder.Services.AddKeyedSingleton("chat-client", mockChatClient); + builder.AddAIAgent("agent", "Instructions", chatClientServiceKey: "chat-client"); + builder.Services.AddLogging(); + using WebApplication app = builder.Build(); + AIAgent agent = app.Services.GetRequiredKeyedService("agent"); + + var agentCard = new AgentCard + { + Name = "Test Agent", + Description = "A test agent for A2A communication" + }; + + // Act & Assert - Should not throw + var result = app.MapA2A(agent, "/a2a", agentCard); + Assert.NotNull(result); + Assert.NotNull(app); + } + + /// + /// Verifies that MapA2A with AIAgent, agent card, and custom task manager configuration succeeds. + /// + [Fact] + public void MapA2A_WithAIAgent_WithAgentCardAndCustomConfiguration_Succeeds() + { + // Arrange + WebApplicationBuilder builder = WebApplication.CreateBuilder(); + IChatClient mockChatClient = new DummyChatClient(); + builder.Services.AddKeyedSingleton("chat-client", mockChatClient); + builder.AddAIAgent("agent", "Instructions", chatClientServiceKey: "chat-client"); + builder.Services.AddLogging(); + using WebApplication app = builder.Build(); + AIAgent agent = app.Services.GetRequiredKeyedService("agent"); + + var agentCard = new AgentCard + { + Name = "Test Agent", + Description = "A test agent for A2A communication" + }; + + // Act & Assert - Should not throw + var result = app.MapA2A(agent, "/a2a", agentCard, taskManager => { }); + Assert.NotNull(result); + Assert.NotNull(app); + } + + /// + /// Verifies that MapA2A throws ArgumentNullException for null endpoints when using ITaskManager. + /// + [Fact] + public void MapA2A_WithTaskManager_NullEndpoints_ThrowsArgumentNullException() + { + // Arrange + AspNetCore.Routing.IEndpointRouteBuilder endpoints = null!; + ITaskManager taskManager = null!; + + // Act & Assert + ArgumentNullException exception = Assert.Throws(() => + endpoints.MapA2A(taskManager, "/a2a")); + + Assert.Equal("endpoints", exception.ParamName); + } + + /// + /// Verifies that multiple agents can be mapped to different paths. + /// + [Fact] + public void MapA2A_MultipleAgents_Succeeds() + { + // Arrange + WebApplicationBuilder builder = WebApplication.CreateBuilder(); + IChatClient mockChatClient = new DummyChatClient(); + builder.Services.AddKeyedSingleton("chat-client", mockChatClient); + IHostedAgentBuilder agent1Builder = builder.AddAIAgent("agent1", "Instructions1", chatClientServiceKey: "chat-client"); + IHostedAgentBuilder agent2Builder = builder.AddAIAgent("agent2", "Instructions2", chatClientServiceKey: "chat-client"); + builder.Services.AddLogging(); + using WebApplication app = builder.Build(); + + // Act & Assert - Should not throw + app.MapA2A(agent1Builder, "/a2a/agent1"); + app.MapA2A(agent2Builder, "/a2a/agent2"); + Assert.NotNull(app); + } + + /// + /// Verifies that custom paths can be specified for A2A endpoints. + /// + [Fact] + public void MapA2A_WithCustomPath_AcceptsValidPath() + { + // Arrange + WebApplicationBuilder builder = WebApplication.CreateBuilder(); + IChatClient mockChatClient = new DummyChatClient(); + builder.Services.AddKeyedSingleton("chat-client", mockChatClient); + IHostedAgentBuilder agentBuilder = builder.AddAIAgent("agent", "Instructions", chatClientServiceKey: "chat-client"); + builder.Services.AddLogging(); + using WebApplication app = builder.Build(); + + // Act & Assert - Should not throw + app.MapA2A(agentBuilder, "/custom/a2a/path"); + Assert.NotNull(app); + } + + /// + /// Verifies that task manager configuration callback is invoked correctly. + /// + [Fact] + public void MapA2A_WithAgentBuilder_TaskManagerConfigurationCallbackInvoked() + { + // Arrange + WebApplicationBuilder builder = WebApplication.CreateBuilder(); + IChatClient mockChatClient = new DummyChatClient(); + builder.Services.AddKeyedSingleton("chat-client", mockChatClient); + IHostedAgentBuilder agentBuilder = builder.AddAIAgent("agent", "Instructions", chatClientServiceKey: "chat-client"); + builder.Services.AddLogging(); + using WebApplication app = builder.Build(); + + bool configureCallbackInvoked = false; + + // Act + app.MapA2A(agentBuilder, "/a2a", taskManager => + { + configureCallbackInvoked = true; + Assert.NotNull(taskManager); + }); + + // Assert + Assert.True(configureCallbackInvoked); + } + + /// + /// Verifies that agent card with all properties is accepted. + /// + [Fact] + public void MapA2A_WithAgentBuilder_FullAgentCard_Succeeds() + { + // Arrange + WebApplicationBuilder builder = WebApplication.CreateBuilder(); + IChatClient mockChatClient = new DummyChatClient(); + builder.Services.AddKeyedSingleton("chat-client", mockChatClient); + IHostedAgentBuilder agentBuilder = builder.AddAIAgent("agent", "Instructions", chatClientServiceKey: "chat-client"); + builder.Services.AddLogging(); + using WebApplication app = builder.Build(); + + var agentCard = new AgentCard + { + Name = "Test Agent", + Description = "A comprehensive test agent" + }; + + // Act & Assert - Should not throw + var result = app.MapA2A(agentBuilder, "/a2a", agentCard); + Assert.NotNull(result); + } + + private sealed class DummyChatClient : IChatClient + { + public void Dispose() + { + throw new NotImplementedException(); + } + + public Task GetResponseAsync(IEnumerable messages, ChatOptions? options = null, CancellationToken cancellationToken = default) + { + throw new NotImplementedException(); + } + + public object? GetService(Type serviceType, object? serviceKey = null) => + serviceType.IsInstanceOfType(this) ? this : null; + + public IAsyncEnumerable GetStreamingResponseAsync(IEnumerable messages, ChatOptions? options = null, CancellationToken cancellationToken = default) + { + throw new NotImplementedException(); + } + } +} diff --git a/dotnet/tests/Microsoft.Agents.AI.Hosting.A2A.Tests/Microsoft.Agents.AI.Hosting.A2A.Tests.csproj b/dotnet/tests/Microsoft.Agents.AI.Hosting.A2A.UnitTests/Microsoft.Agents.AI.Hosting.A2A.UnitTests.csproj similarity index 82% rename from dotnet/tests/Microsoft.Agents.AI.Hosting.A2A.Tests/Microsoft.Agents.AI.Hosting.A2A.Tests.csproj rename to dotnet/tests/Microsoft.Agents.AI.Hosting.A2A.UnitTests/Microsoft.Agents.AI.Hosting.A2A.UnitTests.csproj index 9e0a79d646..63387ae458 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Hosting.A2A.Tests/Microsoft.Agents.AI.Hosting.A2A.Tests.csproj +++ b/dotnet/tests/Microsoft.Agents.AI.Hosting.A2A.UnitTests/Microsoft.Agents.AI.Hosting.A2A.UnitTests.csproj @@ -11,6 +11,7 @@ + diff --git a/dotnet/tests/Microsoft.Agents.AI.Hosting.OpenAI.UnitTests/EndpointRouteBuilderExtensionsTests.cs b/dotnet/tests/Microsoft.Agents.AI.Hosting.OpenAI.UnitTests/EndpointRouteBuilderExtensionsTests.cs index 38b13fbfac..e3effac07a 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Hosting.OpenAI.UnitTests/EndpointRouteBuilderExtensionsTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Hosting.OpenAI.UnitTests/EndpointRouteBuilderExtensionsTests.cs @@ -223,4 +223,174 @@ public sealed class EndpointRouteBuilderExtensionsTests app.MapOpenAIResponses(responsesPath: "/custom/path/responses"); Assert.NotNull(app); } + + /// + /// Verifies that MapOpenAIResponses throws ArgumentNullException for null endpoints when using IHostedAgentBuilder. + /// + [Fact] + public void MapOpenAIResponses_WithAgentBuilder_NullEndpoints_ThrowsArgumentNullException() + { + // Arrange + AspNetCore.Routing.IEndpointRouteBuilder endpoints = null!; + WebApplicationBuilder builder = WebApplication.CreateBuilder(); + IChatClient mockChatClient = new TestHelpers.SimpleMockChatClient(); + builder.Services.AddKeyedSingleton("chat-client", mockChatClient); + IHostedAgentBuilder agentBuilder = builder.AddAIAgent("agent", "Instructions", chatClientServiceKey: "chat-client"); + + // Act & Assert + ArgumentNullException exception = Assert.Throws(() => + endpoints.MapOpenAIResponses(agentBuilder)); + + Assert.Equal("endpoints", exception.ParamName); + } + + /// + /// Verifies that MapOpenAIResponses throws ArgumentNullException for null agentBuilder. + /// + [Fact] + public void MapOpenAIResponses_WithAgentBuilder_NullAgentBuilder_ThrowsArgumentNullException() + { + // Arrange + WebApplicationBuilder builder = WebApplication.CreateBuilder(); + IChatClient mockChatClient = new TestHelpers.SimpleMockChatClient(); + builder.Services.AddKeyedSingleton("chat-client", mockChatClient); + builder.AddOpenAIResponses(); + using WebApplication app = builder.Build(); + IHostedAgentBuilder agentBuilder = null!; + + // Act & Assert + ArgumentNullException exception = Assert.Throws(() => + app.MapOpenAIResponses(agentBuilder)); + + Assert.Equal("agentBuilder", exception.ParamName); + } + + /// + /// Verifies that MapOpenAIResponses with IHostedAgentBuilder correctly resolves and maps the agent. + /// + [Fact] + public void MapOpenAIResponses_WithAgentBuilder_Succeeds() + { + // Arrange + WebApplicationBuilder builder = WebApplication.CreateBuilder(); + IChatClient mockChatClient = new TestHelpers.SimpleMockChatClient(); + builder.Services.AddKeyedSingleton("chat-client", mockChatClient); + IHostedAgentBuilder agentBuilder = builder.AddAIAgent("agent", "Instructions", chatClientServiceKey: "chat-client"); + builder.AddOpenAIResponses(); + using WebApplication app = builder.Build(); + + // Act & Assert - Should not throw + app.MapOpenAIResponses(agentBuilder); + Assert.NotNull(app); + } + + /// + /// Verifies that MapOpenAIResponses with IHostedAgentBuilder and custom path works correctly. + /// + [Fact] + public void MapOpenAIResponses_WithAgentBuilder_CustomPath_Succeeds() + { + // Arrange + WebApplicationBuilder builder = WebApplication.CreateBuilder(); + IChatClient mockChatClient = new TestHelpers.SimpleMockChatClient(); + builder.Services.AddKeyedSingleton("chat-client", mockChatClient); + IHostedAgentBuilder agentBuilder = builder.AddAIAgent("my-agent", "Instructions", chatClientServiceKey: "chat-client"); + builder.AddOpenAIResponses(); + using WebApplication app = builder.Build(); + + // Act & Assert - Should not throw + app.MapOpenAIResponses(agentBuilder, path: "/agents/my-agent/responses"); + Assert.NotNull(app); + } + + /// + /// Verifies that multiple agents can be mapped using IHostedAgentBuilder. + /// + [Fact] + public void MapOpenAIResponses_WithAgentBuilder_MultipleAgents_Succeeds() + { + // Arrange + WebApplicationBuilder builder = WebApplication.CreateBuilder(); + IChatClient mockChatClient = new TestHelpers.SimpleMockChatClient(); + builder.Services.AddKeyedSingleton("chat-client", mockChatClient); + IHostedAgentBuilder agent1Builder = builder.AddAIAgent("agent1", "Instructions1", chatClientServiceKey: "chat-client"); + IHostedAgentBuilder agent2Builder = builder.AddAIAgent("agent2", "Instructions2", chatClientServiceKey: "chat-client"); + builder.AddOpenAIResponses(); + using WebApplication app = builder.Build(); + + // Act & Assert - Should not throw + app.MapOpenAIResponses(agent1Builder); + app.MapOpenAIResponses(agent2Builder); + Assert.NotNull(app); + } + + /// + /// Verifies that IHostedAgentBuilder overload validates agent name characters. + /// + [Theory] + [InlineData("agent with spaces")] + [InlineData("agent