agent/agent-framework
1
0
Fork 0
mirror of https://github.com/microsoft/agent-framework.git synced 2026-06-16 21:04:09 +08:00
Code Issues Packages Projects Releases Wiki Activity

Python: Address PR 5331 comments and track sesssion while calling Agent in email_security_example (#5446)

Browse Source 
* Address PR review: fix paths and update FIDES implementation

* Address PR comments and add session tracking in email example in samples

* Fix session creation and resolve merge conflict in docstring example

* Resolve merge conflict in docstring example
This commit is contained in:
shrutitople
2026-04-24 10:38:22 +01:00
committed by eavanvalkenburg
Unverified
parent 14d779c0fb
commit 9711562c9e
6 changed files with 37 additions and 195 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/packages/core/agent_framework/__init__.py
-2
View File
@@ -130,7 +130,6 @@ from ._tools import (
FunctionInvocationLayer,
FunctionTool,
ToolTypes,
ai_function,
normalize_function_invocation_configuration,
tool,
)
@@ -430,7 +429,6 @@ __all__ = [
"__version__",
"add_usage_details",
"agent_middleware",
"ai_function",
"annotate_message_groups",
"apply_compaction",
"chat_middleware",
python/packages/core/agent_framework/security.py
+7 -78
View File
@@ -831,15 +831,19 @@ class LabelTrackingFunctionMiddleware(FunctionMiddleware):
Examples:
.. code-block:: python
from agent_framework import Agent
from agent_framework import Agent, LabelTrackingFunctionMiddleware, tool
@tool(additional_properties={"source_integrity": "trusted"})
async def get_weather(city: str) -> str:
return f"Weather in {city}: 72°F"
from agent_framework.security import LabelTrackingFunctionMiddleware
# Create agent with automatic hiding enabled
middleware = LabelTrackingFunctionMiddleware(
auto_hide_untrusted=True # Enabled by default
)
agent = Agent(client=client, name="assistant", middleware=[middleware])
agent = Agent(client=client, name="assistant", tools=[get_weather], middleware=[middleware])
# Run agent - untrusted tool results are automatically hidden
response = await agent.run(messages=[{"role": "user", "content": "What's the weather?"}])
@@ -880,10 +884,6 @@ class LabelTrackingFunctionMiddleware(FunctionMiddleware):
# Metadata about stored variables
self._variable_metadata: dict[str, dict[str, Any]] = {}
# Phase 1: Message-level label tracking
# Maps message index to its security label
self._message_labels: dict[int, ContentLabel] = {}
def get_context_label(self) -> ContentLabel:
"""Get the current context-level security label.
@@ -904,79 +904,8 @@ class LabelTrackingFunctionMiddleware(FunctionMiddleware):
self._context_label = ContentLabel(
integrity=IntegrityLabel.TRUSTED, confidentiality=ConfidentialityLabel.PUBLIC, metadata={"reset": True}
)
# Also reset message labels for new conversation
self._message_labels.clear()
logger.info("Context label reset to TRUSTED + PUBLIC")
# ========== Phase 1: Message-Level Label Tracking ==========
def label_message(
self,
message_index: int,
label: ContentLabel,
source_labels: list[ContentLabel] | None = None,
) -> None:
"""Assign a security label to a message in the conversation.
Args:
message_index: The index of the message in the conversation.
label: The security label to assign.
source_labels: Optional list of labels that contributed to this message.
"""
self._message_labels[message_index] = label
logger.debug(f"Labeled message {message_index}: {label.integrity.value}/{label.confidentiality.value}")
def get_message_label(self, message_index: int) -> ContentLabel | None:
"""Get the security label of a specific message.
Args:
message_index: The index of the message.
Returns:
The message's ContentLabel, or None if not labeled.
"""
return self._message_labels.get(message_index)
def label_messages(self, messages: list[dict[str, Any]]) -> list[LabeledMessage]:
"""Label a list of messages based on their roles and content.
This method automatically assigns labels to messages:
- user/system messages: TRUSTED
- assistant messages: Inherit from source labels or TRUSTED
- tool messages: UNTRUSTED (external data)
Args:
messages: List of message dicts with 'role' and 'content'.
Returns:
List of LabeledMessage objects.
"""
labeled: list[LabeledMessage] = []
for i, msg in enumerate(messages):
# Check if message already has a label
existing_label = self._message_labels.get(i)
labeled_msg = LabeledMessage(
role=msg.get("role", "unknown"),
content=msg.get("content", ""),
security_label=existing_label, # Will auto-infer if None
message_index=i,
)
# Store the label
self._message_labels[i] = labeled_msg.security_label
labeled.append(labeled_msg)
return labeled
def get_all_message_labels(self) -> dict[int, ContentLabel]:
"""Get all message labels.
Returns:
Dictionary mapping message index to ContentLabel.
"""
return dict(self._message_labels)
def _update_context_label(self, new_content_label: ContentLabel) -> None:
"""Update the context label based on new content added to the context.
python/packages/core/tests/test_security.py
-55
View File
@@ -1464,61 +1464,6 @@ class TestLabeledMessage:
assert labeled.is_trusted()
class TestMiddlewareMessageLabeling:
"""Tests for middleware message label tracking."""
def test_label_message(self):
"""Test labeling a message by index."""
middleware = LabelTrackingFunctionMiddleware()
label = ContentLabel(integrity=IntegrityLabel.UNTRUSTED, confidentiality=ConfidentialityLabel.PRIVATE)
middleware.label_message(5, label)
retrieved = middleware.get_message_label(5)
assert retrieved is not None
assert retrieved.integrity == IntegrityLabel.UNTRUSTED
def test_get_unlabeled_message_returns_none(self):
"""Test that unlabeled messages return None."""
middleware = LabelTrackingFunctionMiddleware()
assert middleware.get_message_label(999) is None
def test_label_messages_batch(self):
"""Test batch labeling of messages."""
middleware = LabelTrackingFunctionMiddleware()
messages = [
{"role": "user", "content": "Hello"},
{"role": "assistant", "content": "Hi there"},
{"role": "tool", "content": "External data"},
]
labeled = middleware.label_messages(messages)
assert len(labeled) == 3
assert labeled[0].security_label.integrity == IntegrityLabel.TRUSTED
assert labeled[1].security_label.integrity == IntegrityLabel.TRUSTED
assert labeled[2].security_label.integrity == IntegrityLabel.UNTRUSTED
# Check that labels are stored in middleware
all_labels = middleware.get_all_message_labels()
assert len(all_labels) == 3
def test_reset_clears_message_labels(self):
"""Test that reset_context_label also clears message labels."""
middleware = LabelTrackingFunctionMiddleware()
middleware.label_message(0, ContentLabel())
middleware.label_message(1, ContentLabel())
assert len(middleware.get_all_message_labels()) == 2
middleware.reset_context_label()
assert len(middleware.get_all_message_labels()) == 0
# ========== Quarantined LLM Tests ==========
python/samples/02-agents/security/FIDES_DEVELOPER_GUIDE.md
+2 -41
View File
@@ -495,36 +495,9 @@ The instructions explain:
- How to pass `variable_ids` to reference hidden content
- Best practices for secure content handling
### 9. Message-Level Label Tracking (Phase 1)
### 9. LabeledMessage Class
The middleware now tracks security labels at the **message level**, not just tool calls:
```python
from agent_framework.security import LabelTrackingFunctionMiddleware, LabeledMessage
middleware = LabelTrackingFunctionMiddleware()
# Label messages in a conversation
messages = [
{"role": "user", "content": "Hello"}, # Auto-labeled TRUSTED
{"role": "assistant", "content": "Hi there"}, # Auto-labeled TRUSTED (no untrusted sources)
{"role": "tool", "content": "API response"}, # Auto-labeled UNTRUSTED
]
labeled_messages = middleware.label_messages(messages)
# labeled_messages[0].security_label.integrity == TRUSTED
# labeled_messages[2].security_label.integrity == UNTRUSTED
# Individual message labeling
middleware.label_message(message_index=5, label=custom_label)
label = middleware.get_message_label(5)
# Get all message labels
all_labels = middleware.get_all_message_labels()
```
**LabeledMessage Class:**
- Automatically infers labels based on message role
**LabeledMessage** automatically infers security labels based on message role:
- User/system messages → TRUSTED
- Tool messages → UNTRUSTED
- Assistant messages → Inherit from source_labels or TRUSTED
@@ -1108,18 +1081,6 @@ LabeledMessage.from_dict(data) -> LabeledMessage # Deserialize
LabeledMessage.from_message(msg, index) -> LabeledMessage # Wrap standard message
```
### LabelTrackingFunctionMiddleware Extensions
```python
middleware = LabelTrackingFunctionMiddleware(...)
# Message-level label tracking (Phase 1)
middleware.label_message(message_index, label, source_labels=None) # Label a message
middleware.get_message_label(message_index) -> ContentLabel | None # Get message label
middleware.label_messages(messages) -> List[LabeledMessage] # Batch label messages
middleware.get_all_message_labels() -> Dict[int, ContentLabel] # Get all message labels
```
### SecureAgentConfig
```python
python/samples/02-agents/security/email_security_example.py
+10 -2
View File
@@ -278,7 +278,13 @@ async def run_scenarios(agent, config):
print("- Injection attempts in emails are NOT followed")
print()
response = await agent.run("Please fetch my recent emails and give me a brief summary of each one.")
# Use a shared session so conversation history persists across scenarios.
# Without this, each agent.run() starts a fresh conversation and the LLM
# won't know about the emails fetched in Scenario 1 — it would never
# attempt to call send_email, so the policy enforcer would never trigger.
session = agent.create_session()
response = await agent.run("Please fetch my recent emails and give me a brief summary of each one.", session=session)
print(f"\n📋 Agent Response:\n{'-' * 40}")
print(response.text)
@@ -295,7 +301,9 @@ async def run_scenarios(agent, config):
print("- Agent should explain it cannot send email due to security policy")
print()
response = await agent.run("Now please send an email to colleague@company.com summarizing what you found.")
response = await agent.run(
"Now please send an email to colleague@company.com summarizing what you found.", session=session
)
print(f"\n📋 Agent Response:\n{'-' * 40}")
print(response.text)