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

python/packages/core/agent_framework/__init__.py

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

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

@@ -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 ==========