From 64826b8f5650aaf9997b9bcd3c5f917714c9529a Mon Sep 17 00:00:00 2001 From: Rishabh Chawla Date: Thu, 6 Nov 2025 23:43:22 -0800 Subject: [PATCH 1/5] 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 2/5] .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 3/5] .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 4/5] .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 5/5] 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