mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
Python: [Purview] Update CorrelationId (#3745)
This commit is contained in:
committed by
GitHub
Unverified
parent
f106a1a2b1
commit
7e7d72275d
@@ -21,12 +21,14 @@ internal abstract class ProcessContentMetadataBase : GraphDataTypeBase
|
||||
/// <param name="identifier">The unique identifier for the content.</param>
|
||||
/// <param name="isTruncated">Indicates if the content is truncated.</param>
|
||||
/// <param name="name">The name of the content.</param>
|
||||
protected ProcessContentMetadataBase(ContentBase content, string identifier, bool isTruncated, string name) : base(ProcessConversationMetadataDataType)
|
||||
/// <param name="correlationId">The correlation ID for the content.</param>
|
||||
protected ProcessContentMetadataBase(ContentBase content, string identifier, bool isTruncated, string name, string correlationId) : base(ProcessConversationMetadataDataType)
|
||||
{
|
||||
this.Identifier = identifier;
|
||||
this.IsTruncated = isTruncated;
|
||||
this.Content = content;
|
||||
this.Name = name;
|
||||
this.CorrelationId = correlationId;
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
@@ -55,7 +57,7 @@ internal abstract class ProcessContentMetadataBase : GraphDataTypeBase
|
||||
/// Identifier to group multiple contents.
|
||||
/// </summary>
|
||||
[JsonPropertyName("correlationId")]
|
||||
public string? CorrelationId { get; set; }
|
||||
public string CorrelationId { get; set; }
|
||||
|
||||
/// <summary>
|
||||
/// Gets or sets the sequenceNumber.
|
||||
|
||||
@@ -15,7 +15,7 @@ internal sealed class ProcessConversationMetadata : ProcessContentMetadataBase
|
||||
/// <summary>
|
||||
/// Initializes a new instance of the <see cref="ProcessConversationMetadata"/> class.
|
||||
/// </summary>
|
||||
public ProcessConversationMetadata(ContentBase contentBase, string identifier, bool isTruncated, string name) : base(contentBase, identifier, isTruncated, name)
|
||||
public ProcessConversationMetadata(ContentBase contentBase, string identifier, bool isTruncated, string name, string correlationId) : base(contentBase, identifier, isTruncated, name, correlationId)
|
||||
{
|
||||
this.DataType = ProcessConversationMetadataDataType;
|
||||
}
|
||||
|
||||
@@ -14,7 +14,7 @@ internal sealed class ProcessFileMetadata : ProcessContentMetadataBase
|
||||
/// <summary>
|
||||
/// Initializes a new instance of the <see cref="ProcessFileMetadata"/> class.
|
||||
/// </summary>
|
||||
public ProcessFileMetadata(ContentBase contentBase, string identifier, bool isTruncated, string name) : base(contentBase, identifier, isTruncated, name)
|
||||
public ProcessFileMetadata(ContentBase contentBase, string identifier, bool isTruncated, string name, string correlationId) : base(contentBase, identifier, isTruncated, name, correlationId)
|
||||
{
|
||||
this.DataType = ProcessFileMetadataDataType;
|
||||
}
|
||||
|
||||
@@ -19,7 +19,7 @@ public class PurviewSettings
|
||||
/// <param name="appName">The publicly visible name of the application.</param>
|
||||
public PurviewSettings(string appName)
|
||||
{
|
||||
this.AppName = appName;
|
||||
this.AppName = string.IsNullOrWhiteSpace(appName) ? throw new ArgumentException("AppName cannot be null or whitespace.", nameof(appName)) : appName;
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
|
||||
@@ -53,7 +53,7 @@ internal sealed class PurviewWrapper : IDisposable
|
||||
}
|
||||
}
|
||||
|
||||
return Guid.NewGuid().ToString();
|
||||
return string.Empty;
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
@@ -136,12 +136,15 @@ internal sealed class PurviewWrapper : IDisposable
|
||||
/// <returns>The agent's response. This could be the response from the agent or a message indicating that Purview has blocked the prompt or response.</returns>
|
||||
public async Task<AgentResponse> ProcessAgentContentAsync(IEnumerable<ChatMessage> messages, AgentSession? session, AgentRunOptions? options, AIAgent innerAgent, CancellationToken cancellationToken)
|
||||
{
|
||||
string sessionId = GetSessionIdFromAgentSession(session, messages);
|
||||
|
||||
string? resolvedUserId = null;
|
||||
|
||||
string sessionId = string.Empty;
|
||||
try
|
||||
{
|
||||
sessionId = GetSessionIdFromAgentSession(session, messages);
|
||||
if (string.IsNullOrEmpty(sessionId))
|
||||
{
|
||||
sessionId = Guid.NewGuid().ToString();
|
||||
}
|
||||
(bool shouldBlockPrompt, resolvedUserId) = await this._scopedProcessor.ProcessMessagesAsync(messages, sessionId, Activity.UploadText, this._purviewSettings, null, cancellationToken).ConfigureAwait(false);
|
||||
|
||||
if (shouldBlockPrompt)
|
||||
@@ -171,7 +174,19 @@ internal sealed class PurviewWrapper : IDisposable
|
||||
|
||||
try
|
||||
{
|
||||
(bool shouldBlockResponse, _) = await this._scopedProcessor.ProcessMessagesAsync(response.Messages, sessionId, Activity.UploadText, this._purviewSettings, resolvedUserId, cancellationToken).ConfigureAwait(false);
|
||||
string sessionIdResponse = GetSessionIdFromAgentSession(session, messages);
|
||||
if (string.IsNullOrEmpty(sessionIdResponse))
|
||||
{
|
||||
if (string.IsNullOrEmpty(sessionId))
|
||||
{
|
||||
sessionIdResponse = Guid.NewGuid().ToString();
|
||||
}
|
||||
else
|
||||
{
|
||||
sessionIdResponse = sessionId;
|
||||
}
|
||||
}
|
||||
(bool shouldBlockResponse, _) = await this._scopedProcessor.ProcessMessagesAsync(response.Messages, sessionIdResponse, Activity.UploadText, this._purviewSettings, resolvedUserId, cancellationToken).ConfigureAwait(false);
|
||||
|
||||
if (shouldBlockResponse)
|
||||
{
|
||||
|
||||
@@ -121,9 +121,10 @@ internal sealed class ScopedContentProcessor : IScopedContentProcessor
|
||||
{
|
||||
string messageId = message.MessageId ?? Guid.NewGuid().ToString();
|
||||
ContentBase content = new PurviewTextContent(message.Text);
|
||||
ProcessConversationMetadata conversationmetadata = new(content, messageId, false, $"Agent Framework Message {messageId}")
|
||||
string correlationId = (sessionId ?? Guid.NewGuid().ToString()) + "@AF";
|
||||
ProcessConversationMetadata conversationMetadata = new(content, messageId, false, $"Agent Framework Message {messageId}", correlationId)
|
||||
{
|
||||
CorrelationId = sessionId ?? Guid.NewGuid().ToString()
|
||||
SequenceNumber = DateTime.UtcNow.Ticks,
|
||||
};
|
||||
ActivityMetadata activityMetadata = new(activity);
|
||||
PolicyLocation policyLocation;
|
||||
@@ -162,7 +163,7 @@ internal sealed class ScopedContentProcessor : IScopedContentProcessor
|
||||
OperatingSystemVersion = "Unknown"
|
||||
}
|
||||
};
|
||||
ContentToProcess contentToProcess = new([conversationmetadata], activityMetadata, deviceMetadata, integratedAppMetadata, protectedAppMetadata);
|
||||
ContentToProcess contentToProcess = new([conversationMetadata], activityMetadata, deviceMetadata, integratedAppMetadata, protectedAppMetadata);
|
||||
|
||||
if (userId == null &&
|
||||
tokenInfo?.UserId != null)
|
||||
|
||||
@@ -478,7 +478,7 @@ public sealed class PurviewClientTests : IDisposable
|
||||
private static ContentToProcess CreateValidContentToProcess()
|
||||
{
|
||||
var content = new PurviewTextContent("Test content");
|
||||
var metadata = new ProcessConversationMetadata(content, "msg-123", false, "Test message");
|
||||
var metadata = new ProcessConversationMetadata(content, "msg-123", false, "Test message", "test-correlation-id");
|
||||
var activityMetadata = new ActivityMetadata(Activity.UploadText);
|
||||
var deviceMetadata = new DeviceMetadata
|
||||
{
|
||||
|
||||
@@ -45,6 +45,25 @@ class PurviewPolicyMiddleware(AgentMiddleware):
|
||||
self._processor = ScopedContentProcessor(self._client, settings, cache_provider)
|
||||
self._settings = settings
|
||||
|
||||
@staticmethod
|
||||
def _get_agent_session_id(context: AgentContext) -> str | None:
|
||||
"""Resolve a session/conversation id from the agent run context.
|
||||
|
||||
Resolution order:
|
||||
1. thread.service_thread_id
|
||||
2. First message whose additional_properties contains 'conversation_id'
|
||||
3. None: the downstream processor will generate a new UUID
|
||||
"""
|
||||
if context.thread and context.thread.service_thread_id:
|
||||
return context.thread.service_thread_id
|
||||
|
||||
for message in context.messages:
|
||||
conversation_id = message.additional_properties.get("conversation_id")
|
||||
if conversation_id is not None:
|
||||
return str(conversation_id)
|
||||
|
||||
return None
|
||||
|
||||
async def process(
|
||||
self,
|
||||
context: AgentContext,
|
||||
@@ -53,8 +72,9 @@ class PurviewPolicyMiddleware(AgentMiddleware):
|
||||
resolved_user_id: str | None = None
|
||||
try:
|
||||
# Pre (prompt) check
|
||||
session_id = self._get_agent_session_id(context)
|
||||
should_block_prompt, resolved_user_id = await self._processor.process_messages(
|
||||
context.messages, Activity.UPLOAD_TEXT
|
||||
context.messages, Activity.UPLOAD_TEXT, session_id=session_id
|
||||
)
|
||||
if should_block_prompt:
|
||||
from agent_framework import AgentResponse, ChatMessage
|
||||
@@ -79,10 +99,14 @@ class PurviewPolicyMiddleware(AgentMiddleware):
|
||||
try:
|
||||
# Post (response) check only if we have a normal AgentResponse
|
||||
# Use the same user_id from the request for the response evaluation
|
||||
session_id_response = self._get_agent_session_id(context)
|
||||
if session_id_response is None:
|
||||
session_id_response = session_id
|
||||
if context.result and not context.stream:
|
||||
should_block_response, _ = await self._processor.process_messages(
|
||||
context.result.messages, # type: ignore[union-attr]
|
||||
Activity.UPLOAD_TEXT,
|
||||
session_id=session_id,
|
||||
user_id=resolved_user_id,
|
||||
)
|
||||
if should_block_response:
|
||||
@@ -144,8 +168,9 @@ class PurviewChatPolicyMiddleware(ChatMiddleware):
|
||||
) -> None: # type: ignore[override]
|
||||
resolved_user_id: str | None = None
|
||||
try:
|
||||
session_id = context.options.get("conversation_id") if context.options else None
|
||||
should_block_prompt, resolved_user_id = await self._processor.process_messages(
|
||||
context.messages, Activity.UPLOAD_TEXT
|
||||
context.messages, Activity.UPLOAD_TEXT, session_id=session_id
|
||||
)
|
||||
if should_block_prompt:
|
||||
from agent_framework import ChatMessage, ChatResponse
|
||||
@@ -169,12 +194,15 @@ class PurviewChatPolicyMiddleware(ChatMiddleware):
|
||||
try:
|
||||
# Post (response) evaluation only if non-streaming and we have messages result shape
|
||||
# Use the same user_id from the request for the response evaluation
|
||||
session_id_response = context.options.get("conversation_id") if context.options else None
|
||||
if session_id_response is None:
|
||||
session_id_response = session_id
|
||||
if context.result and not context.stream:
|
||||
result_obj = context.result
|
||||
messages = getattr(result_obj, "messages", None)
|
||||
if messages:
|
||||
should_block_response, _ = await self._processor.process_messages(
|
||||
messages, Activity.UPLOAD_TEXT, user_id=resolved_user_id
|
||||
messages, Activity.UPLOAD_TEXT, session_id=session_id_response, user_id=resolved_user_id
|
||||
)
|
||||
if should_block_response:
|
||||
from agent_framework import ChatMessage, ChatResponse
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import asyncio
|
||||
import time
|
||||
import uuid
|
||||
from collections.abc import Iterable, MutableMapping
|
||||
from typing import Any
|
||||
@@ -62,13 +63,18 @@ class ScopedContentProcessor:
|
||||
self._background_tasks: set[asyncio.Task[Any]] = set()
|
||||
|
||||
async def process_messages(
|
||||
self, messages: Iterable[ChatMessage], activity: Activity, user_id: str | None = None
|
||||
self,
|
||||
messages: Iterable[ChatMessage],
|
||||
activity: Activity,
|
||||
session_id: str | None = None,
|
||||
user_id: str | None = None,
|
||||
) -> tuple[bool, str | None]:
|
||||
"""Process messages for policy evaluation.
|
||||
|
||||
Args:
|
||||
messages: The messages to process
|
||||
activity: The activity type (e.g., UPLOAD_TEXT)
|
||||
session_id: Optional session/conversation id. Else, a new GUID is generated.
|
||||
user_id: Optional user_id to use for all messages. If provided, this is the fallback.
|
||||
|
||||
Returns:
|
||||
@@ -76,7 +82,7 @@ class ScopedContentProcessor:
|
||||
The resolved_user_id can be stored and passed back when processing the response
|
||||
to ensure the same user context is maintained throughout the request/response cycle.
|
||||
"""
|
||||
pc_requests, resolved_user_id = await self._map_messages(messages, activity, user_id)
|
||||
pc_requests, resolved_user_id = await self._map_messages(messages, activity, session_id, user_id)
|
||||
should_block = False
|
||||
for req in pc_requests:
|
||||
resp = await self._process_with_scopes(req)
|
||||
@@ -90,13 +96,18 @@ class ScopedContentProcessor:
|
||||
return should_block, resolved_user_id
|
||||
|
||||
async def _map_messages(
|
||||
self, messages: Iterable[ChatMessage], activity: Activity, provided_user_id: str | None = None
|
||||
self,
|
||||
messages: Iterable[ChatMessage],
|
||||
activity: Activity,
|
||||
session_id: str | None = None,
|
||||
provided_user_id: str | None = None,
|
||||
) -> tuple[list[ProcessContentRequest], str | None]:
|
||||
"""Map messages to ProcessContentRequests.
|
||||
|
||||
Args:
|
||||
messages: The messages to map
|
||||
activity: The activity type
|
||||
session_id: Optional session/conversation id to use for correlation
|
||||
provided_user_id: Optional user_id to use. If provided, this is the fallback.
|
||||
|
||||
Returns:
|
||||
@@ -137,12 +148,14 @@ class ScopedContentProcessor:
|
||||
for m in messages:
|
||||
message_id = m.message_id or str(uuid.uuid4())
|
||||
content = PurviewTextContent(data=m.text or "")
|
||||
correlation_id = (session_id or str(uuid.uuid4())) + "@AF"
|
||||
meta = ProcessConversationMetadata(
|
||||
identifier=message_id,
|
||||
content=content,
|
||||
name=f"Agent Framework Message {message_id}",
|
||||
is_truncated=False,
|
||||
correlation_id=str(uuid.uuid4()),
|
||||
correlation_id=correlation_id,
|
||||
sequence_number=time.time_ns(),
|
||||
)
|
||||
activity_meta = ActivityMetadata(activity=activity)
|
||||
|
||||
@@ -159,12 +172,13 @@ class ScopedContentProcessor:
|
||||
else:
|
||||
raise ValueError("App location not provided or inferable")
|
||||
|
||||
app_version = self._settings.app_version or "Unknown"
|
||||
protected_app = ProtectedAppMetadata(
|
||||
name=self._settings.app_name,
|
||||
version="1.0",
|
||||
version=app_version,
|
||||
application_location=policy_location,
|
||||
)
|
||||
integrated_app = IntegratedAppMetadata(name=self._settings.app_name, version="1.0")
|
||||
integrated_app = IntegratedAppMetadata(name=self._settings.app_name, version=app_version)
|
||||
device_meta = DeviceMetadata(
|
||||
operating_system_specifications=OperatingSystemSpecifications(
|
||||
operating_system_platform="Unknown", operating_system_version="Unknown"
|
||||
|
||||
@@ -35,7 +35,7 @@ class PurviewAppLocation(BaseModel):
|
||||
|
||||
|
||||
class PurviewSettings(AFBaseSettings):
|
||||
"""Settings for Purview integration mirroring .NET PurviewSettings.
|
||||
"""Settings for Purview integration.
|
||||
|
||||
Attributes:
|
||||
app_name: Public app name.
|
||||
|
||||
@@ -9,6 +9,7 @@ from agent_framework import ChatContext, ChatMessage, MiddlewareTermination
|
||||
from azure.core.credentials import AccessToken
|
||||
|
||||
from agent_framework_purview import PurviewChatPolicyMiddleware, PurviewSettings
|
||||
from agent_framework_purview._models import Activity
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -82,7 +83,7 @@ class TestPurviewChatPolicyMiddleware:
|
||||
async def test_blocks_response(self, middleware: PurviewChatPolicyMiddleware, chat_context: ChatContext) -> None:
|
||||
call_state = {"count": 0}
|
||||
|
||||
async def side_effect(messages, activity, user_id=None):
|
||||
async def side_effect(messages, activity, session_id=None, user_id=None):
|
||||
call_state["count"] += 1
|
||||
should_block = call_state["count"] == 2
|
||||
return (should_block, "user-123")
|
||||
@@ -157,7 +158,7 @@ class TestPurviewChatPolicyMiddleware:
|
||||
"""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):
|
||||
async def mock_process_messages(messages, activity, session_id=None, user_id=None):
|
||||
captured_user_ids.append(user_id)
|
||||
return (False, "resolved-user-123")
|
||||
|
||||
@@ -362,3 +363,67 @@ class TestPurviewChatPolicyMiddleware:
|
||||
|
||||
with pytest.raises(ValueError, match="post"):
|
||||
await middleware.process(context, mock_next)
|
||||
|
||||
async def test_chat_middleware_uses_conversation_id_from_options(
|
||||
self, middleware: PurviewChatPolicyMiddleware
|
||||
) -> None:
|
||||
"""Test that session_id is extracted from context.options['conversation_id']."""
|
||||
chat_client = DummyChatClient()
|
||||
messages = [ChatMessage(role="user", text="Hello")]
|
||||
options = {"conversation_id": "conv-123", "model": "test-model"}
|
||||
context = ChatContext(chat_client=chat_client, messages=messages, options=options)
|
||||
|
||||
with patch.object(middleware._processor, "process_messages", return_value=(False, "user-123")) as mock_proc:
|
||||
|
||||
async def mock_next(ctx: ChatContext) -> None:
|
||||
result = MagicMock()
|
||||
result.messages = [ChatMessage(role="assistant", text="Hi")]
|
||||
ctx.result = result
|
||||
|
||||
await middleware.process(context, mock_next)
|
||||
|
||||
# Verify session_id is passed to both pre-check and post-check
|
||||
assert mock_proc.call_count == 2
|
||||
mock_proc.assert_any_call(messages, Activity.UPLOAD_TEXT, session_id="conv-123")
|
||||
|
||||
async def test_chat_middleware_passes_none_session_id_when_options_missing(
|
||||
self, middleware: PurviewChatPolicyMiddleware
|
||||
) -> None:
|
||||
"""Test that session_id is None when options don't contain conversation_id."""
|
||||
chat_client = DummyChatClient()
|
||||
messages = [ChatMessage(role="user", text="Hello")]
|
||||
context = ChatContext(chat_client=chat_client, messages=messages, options=None)
|
||||
|
||||
with patch.object(middleware._processor, "process_messages", return_value=(False, "user-123")) as mock_proc:
|
||||
|
||||
async def mock_next(ctx: ChatContext) -> None:
|
||||
result = MagicMock()
|
||||
result.messages = [ChatMessage(role="assistant", text="Hi")]
|
||||
ctx.result = result
|
||||
|
||||
await middleware.process(context, mock_next)
|
||||
|
||||
# Verify session_id=None is passed
|
||||
mock_proc.assert_any_call(messages, Activity.UPLOAD_TEXT, session_id=None)
|
||||
|
||||
async def test_chat_middleware_session_id_used_in_post_check(self, middleware: PurviewChatPolicyMiddleware) -> None:
|
||||
"""Test that session_id is passed to post-check process_messages call."""
|
||||
chat_client = DummyChatClient()
|
||||
messages = [ChatMessage(role="user", text="Hello")]
|
||||
options = {"conversation_id": "conv-999"}
|
||||
context = ChatContext(chat_client=chat_client, messages=messages, options=options)
|
||||
|
||||
with patch.object(middleware._processor, "process_messages", return_value=(False, "user-123")) as mock_proc:
|
||||
|
||||
async def mock_next(ctx: ChatContext) -> None:
|
||||
result = MagicMock()
|
||||
result.messages = [ChatMessage(role="assistant", text="Response")]
|
||||
ctx.result = result
|
||||
|
||||
await middleware.process(context, mock_next)
|
||||
|
||||
# Verify both calls include session_id
|
||||
assert mock_proc.call_count == 2
|
||||
# Check post-check call includes session_id
|
||||
post_check_call = mock_proc.call_args_list[1]
|
||||
assert post_check_call[1]["session_id"] == "conv-999"
|
||||
|
||||
@@ -5,10 +5,11 @@
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from agent_framework import AgentContext, AgentResponse, ChatMessage, MiddlewareTermination
|
||||
from agent_framework import AgentContext, AgentResponse, AgentThread, ChatMessage, MiddlewareTermination
|
||||
from azure.core.credentials import AccessToken
|
||||
|
||||
from agent_framework_purview import PurviewPolicyMiddleware, PurviewSettings
|
||||
from agent_framework_purview._models import Activity
|
||||
|
||||
|
||||
class TestPurviewPolicyMiddleware:
|
||||
@@ -92,7 +93,7 @@ class TestPurviewPolicyMiddleware:
|
||||
|
||||
call_count = 0
|
||||
|
||||
async def mock_process_messages(messages, activity, user_id=None):
|
||||
async def mock_process_messages(messages, activity, session_id=None, user_id=None):
|
||||
nonlocal call_count
|
||||
call_count += 1
|
||||
should_block = call_count != 1
|
||||
@@ -335,3 +336,93 @@ class TestPurviewPolicyMiddleware:
|
||||
# Should raise the exception
|
||||
with pytest.raises(ValueError, match="Test error"):
|
||||
await middleware.process(context, mock_next)
|
||||
|
||||
async def test_middleware_uses_thread_service_thread_id_as_session_id(
|
||||
self, middleware: PurviewPolicyMiddleware, mock_agent: MagicMock
|
||||
) -> None:
|
||||
"""Test that session_id is extracted from thread.service_thread_id."""
|
||||
thread = AgentThread(service_thread_id="thread-123")
|
||||
context = AgentContext(agent=mock_agent, messages=[ChatMessage(role="user", text="Hello")], thread=thread)
|
||||
|
||||
with patch.object(middleware._processor, "process_messages", return_value=(False, "user-123")) as mock_proc:
|
||||
|
||||
async def mock_next(ctx: AgentContext) -> None:
|
||||
ctx.result = AgentResponse(messages=[ChatMessage(role="assistant", text="Hi")])
|
||||
|
||||
await middleware.process(context, mock_next)
|
||||
|
||||
# Verify session_id is passed to both pre-check and post-check
|
||||
assert mock_proc.call_count == 2
|
||||
mock_proc.assert_any_call(context.messages, Activity.UPLOAD_TEXT, session_id="thread-123")
|
||||
|
||||
async def test_middleware_uses_message_conversation_id_as_session_id(
|
||||
self, middleware: PurviewPolicyMiddleware, mock_agent: MagicMock
|
||||
) -> None:
|
||||
"""Test that session_id is extracted from message.additional_properties['conversation_id']."""
|
||||
messages = [ChatMessage(role="user", text="Hello", additional_properties={"conversation_id": "conv-456"})]
|
||||
context = AgentContext(agent=mock_agent, messages=messages)
|
||||
|
||||
with patch.object(middleware._processor, "process_messages", return_value=(False, "user-123")) as mock_proc:
|
||||
|
||||
async def mock_next(ctx: AgentContext) -> None:
|
||||
ctx.result = AgentResponse(messages=[ChatMessage(role="assistant", text="Hi")])
|
||||
|
||||
await middleware.process(context, mock_next)
|
||||
|
||||
# Verify session_id is passed to both pre-check and post-check
|
||||
assert mock_proc.call_count == 2
|
||||
mock_proc.assert_any_call(messages, Activity.UPLOAD_TEXT, session_id="conv-456")
|
||||
|
||||
async def test_middleware_thread_id_takes_precedence_over_message_conversation_id(
|
||||
self, middleware: PurviewPolicyMiddleware, mock_agent: MagicMock
|
||||
) -> None:
|
||||
"""Test that thread.service_thread_id takes precedence over message conversation_id."""
|
||||
thread = AgentThread(service_thread_id="thread-789")
|
||||
messages = [ChatMessage(role="user", text="Hello", additional_properties={"conversation_id": "conv-456"})]
|
||||
context = AgentContext(agent=mock_agent, messages=messages, thread=thread)
|
||||
|
||||
with patch.object(middleware._processor, "process_messages", return_value=(False, "user-123")) as mock_proc:
|
||||
|
||||
async def mock_next(ctx: AgentContext) -> None:
|
||||
ctx.result = AgentResponse(messages=[ChatMessage(role="assistant", text="Hi")])
|
||||
|
||||
await middleware.process(context, mock_next)
|
||||
|
||||
# Verify thread ID is used, not message conversation_id
|
||||
mock_proc.assert_any_call(messages, Activity.UPLOAD_TEXT, session_id="thread-789")
|
||||
|
||||
async def test_middleware_passes_none_session_id_when_not_available(
|
||||
self, middleware: PurviewPolicyMiddleware, mock_agent: MagicMock
|
||||
) -> None:
|
||||
"""Test that session_id is None when no thread or conversation_id is available."""
|
||||
context = AgentContext(agent=mock_agent, messages=[ChatMessage(role="user", text="Hello")])
|
||||
|
||||
with patch.object(middleware._processor, "process_messages", return_value=(False, "user-123")) as mock_proc:
|
||||
|
||||
async def mock_next(ctx: AgentContext) -> None:
|
||||
ctx.result = AgentResponse(messages=[ChatMessage(role="assistant", text="Hi")])
|
||||
|
||||
await middleware.process(context, mock_next)
|
||||
|
||||
# Verify session_id=None is passed
|
||||
mock_proc.assert_any_call(context.messages, Activity.UPLOAD_TEXT, session_id=None)
|
||||
|
||||
async def test_middleware_session_id_used_in_post_check(
|
||||
self, middleware: PurviewPolicyMiddleware, mock_agent: MagicMock
|
||||
) -> None:
|
||||
"""Test that session_id is passed to post-check process_messages call."""
|
||||
thread = AgentThread(service_thread_id="thread-999")
|
||||
context = AgentContext(agent=mock_agent, messages=[ChatMessage(role="user", text="Hello")], thread=thread)
|
||||
|
||||
with patch.object(middleware._processor, "process_messages", return_value=(False, "user-123")) as mock_proc:
|
||||
|
||||
async def mock_next(ctx: AgentContext) -> None:
|
||||
ctx.result = AgentResponse(messages=[ChatMessage(role="assistant", text="Response")])
|
||||
|
||||
await middleware.process(context, mock_next)
|
||||
|
||||
# Verify both calls include session_id
|
||||
assert mock_proc.call_count == 2
|
||||
# Check post-check call includes session_id
|
||||
post_check_call = mock_proc.call_args_list[1]
|
||||
assert post_check_call[1]["session_id"] == "thread-999"
|
||||
|
||||
@@ -92,7 +92,7 @@ class TestScopedContentProcessor:
|
||||
|
||||
assert should_block is False
|
||||
assert user_id is None
|
||||
mock_map.assert_called_once_with(messages, Activity.UPLOAD_TEXT, None)
|
||||
mock_map.assert_called_once_with(messages, Activity.UPLOAD_TEXT, None, None)
|
||||
|
||||
async def test_process_messages_blocks_content(
|
||||
self, processor: ScopedContentProcessor, process_content_request_factory
|
||||
|
||||
Reference in New Issue
Block a user