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

* 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
@@ -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
@@ -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)