FEAT: Tool Use + MCP

pyproject.toml

@@ -46,6 +46,7 @@ dependencies = [
     "fastapi>=0.133.0",
     "httpx[http2]>=0.27.2",
     "jinja2>=3.1.6",
+    "mcp>=1.0,<2",
     "numpy>=1.26.0; python_version < '3.14'",
     "numpy>=2.3.0; python_version >= '3.14'",
     "openai>=2.2.0",

pyrit/exceptions/__init__.py

@@ -10,6 +10,8 @@
     MissingPromptPlaceholderException,
     PyritException,
     RateLimitException,
+    ToolCallLoopLimitExceeded,
+    ToolCallNotSupported,
     get_retry_max_num_attempts,
     handle_bad_request_exception,
     pyrit_custom_result_retry,
@@ -59,4 +61,6 @@
     "set_execution_context",
     "set_retry_collector",
     "execution_context",
+    "ToolCallLoopLimitExceeded",
+    "ToolCallNotSupported",
 ]

pyrit/exceptions/exception_classes.py

@@ -233,6 +233,70 @@ def __init__(self, *, message: str = "No prompt placeholder") -> None:
         super().__init__(message=message)
 
 
+class ToolCallNotSupported(PyritException):
+    """
+    Raised when a target produces a tool call that the configured
+    :class:`~pyrit.tools.ToolEventPolicy` does not permit to execute
+    (``ToolEventBehavior.RAISE``, or ``EXECUTE`` without a backend).
+
+    The ``partial_conversation`` attribute carries every message produced
+    up to and including the assistant turn that contained the offending
+    tool call(s). Consumers can inspect it to log the surfaced tool-use
+    attempt.
+    """
+
+    def __init__(
+        self,
+        *,
+        message: str = "Tool call not supported by configured policy.",
+        partial_conversation: Optional[list["Message"]] = None,
+    ) -> None:
+        """
+        Initialize the exception.
+
+        Args:
+            message (str): Human-readable error description.
+            partial_conversation (Optional[list[Message]]): Messages produced by
+                the target up to (and including) the assistant turn that
+                contained the disallowed tool call(s).
+        """
+        super().__init__(status_code=400, message=message)
+        self.partial_conversation: list[Message] = (
+            list(partial_conversation) if partial_conversation is not None else []
+        )
+
+
+class ToolCallLoopLimitExceeded(PyritException):
+    """
+    Raised when the tool-use loop runs for more than
+    ``ToolEventPolicy.max_tool_iterations`` iterations without the model
+    producing a stop response.
+
+    The ``partial_conversation`` attribute carries every message produced
+    across all completed iterations. Consumers can inspect it to debug
+    runaway agentic behavior.
+    """
+
+    def __init__(
+        self,
+        *,
+        message: str = "Tool loop exceeded max_tool_iterations without a stop response.",
+        partial_conversation: Optional[list["Message"]] = None,
+    ) -> None:
+        """
+        Initialize the exception.
+
+        Args:
+            message (str): Human-readable error description.
+            partial_conversation (Optional[list[Message]]): Messages produced by
+                the target across every completed iteration of the tool loop.
+        """
+        super().__init__(status_code=400, message=message)
+        self.partial_conversation: list[Message] = (
+            list(partial_conversation) if partial_conversation is not None else []
+        )
+
+
 def pyrit_custom_result_retry(
     retry_function: Callable[..., bool], retry_max_num_attempts: Optional[int] = None
 ) -> Callable[..., Any]:

pyrit/message_normalizer/chat_message_normalizer.py

@@ -4,6 +4,7 @@
 import base64
 import json
 import os
+from collections.abc import Sequence
 from typing import TYPE_CHECKING, Any, Union
 
 from pyrit.common.data_url_converter import convert_local_image_to_data_url_async
@@ -14,6 +15,7 @@
     apply_system_message_behavior,
 )
 from pyrit.models import ChatMessage, DataTypeSerializer, Message
+from pyrit.models.chat_message import ToolCall, ToolCallFunction
 from pyrit.models.message_piece import MessagePiece
 
 if TYPE_CHECKING:
@@ -83,6 +85,11 @@ async def normalize_async(self, messages: list[Message]) -> list[ChatMessage]:
         chat_messages: list[ChatMessage] = []
         for message in processed_messages:
             pieces = message.message_pieces
+            tool_message = self._try_build_tool_message(pieces=pieces)
+            if tool_message is not None:
+                chat_messages.append(tool_message)
+                continue
+
             role: ChatMessageRole = pieces[0].api_role
 
             # Translate system -> developer for newer OpenAI models
@@ -99,6 +106,89 @@ async def normalize_async(self, messages: list[Message]) -> list[ChatMessage]:
 
         return chat_messages
 
+    def _try_build_tool_message(self, *, pieces: Sequence[MessagePiece]) -> ChatMessage | None:
+        """
+        Build an OpenAI Chat Completions tool message when ``pieces`` carries tool data.
+
+        Returns a populated ``ChatMessage`` when the pieces are tool-call
+        envelopes (``function_call`` or ``function_call_output`` data type),
+        or ``None`` when the pieces are ordinary text / multimodal content.
+
+        ``function_call`` pieces produce a single ``role="assistant"`` message
+        with ``content=None`` and one or more entries in ``tool_calls``.
+        ``function_call_output`` pieces produce a single ``role="tool"``
+        message whose ``content`` is the output payload and whose
+        ``tool_call_id`` matches the originating call.
+
+        Args:
+            pieces (list[MessagePiece]): The pieces making up one PyRIT message.
+
+        Returns:
+            ChatMessage | None: ``None`` when no tool envelopes are present,
+                otherwise the converted tool message.
+        """
+        if not pieces:
+            return None
+        data_types = {p.converted_value_data_type or p.original_value_data_type for p in pieces}
+        if data_types == {"function_call"}:
+            return ChatMessage(
+                role="assistant",
+                content=None,
+                tool_calls=[self._piece_to_tool_call(piece) for piece in pieces],
+            )
+        if data_types == {"function_call_output"}:
+            # A single message carries one or more function_call_output pieces
+            # in declaration order; the OpenAI wire shape sends each as its
+            # own role="tool" message. For multi-piece tool messages, we
+            # surface the first piece here and let the caller emit additional
+            # messages — but in practice tool_loop emits one message per
+            # iteration with multiple pieces, and OpenAI accepts a single
+            # tool message per call_id. Emit the first envelope; warn if
+            # multiple are present.
+            envelope = self._decode_envelope(pieces[0])
+            return ChatMessage(
+                role="tool",
+                content=str(envelope.get("output", "")),
+                tool_call_id=str(envelope["call_id"]),
+            )
+        return None
+
+    @staticmethod
+    def _decode_envelope(piece: MessagePiece) -> dict[str, Any]:
+        """
+        Decode the canonical-envelope JSON carried in a tool piece.
+
+        Args:
+            piece (MessagePiece): A piece whose ``converted_value`` is the
+                canonical-envelope JSON string.
+
+        Returns:
+            dict[str, Any]: The parsed envelope.
+        """
+        return json.loads(piece.converted_value)
+
+    @classmethod
+    def _piece_to_tool_call(cls, piece: MessagePiece) -> ToolCall:
+        """
+        Convert one canonical ``function_call`` piece into an OpenAI ToolCall.
+
+        Args:
+            piece (MessagePiece): A piece carrying a canonical ``function_call``
+                envelope.
+
+        Returns:
+            ToolCall: The corresponding OpenAI Chat Completions tool call.
+        """
+        envelope = cls._decode_envelope(piece)
+        return ToolCall(
+            id=str(envelope["call_id"]),
+            type="function",
+            function=ToolCallFunction(
+                name=str(envelope["name"]),
+                arguments=str(envelope["arguments"]),
+            ),
+        )
+
     async def normalize_string_async(self, messages: list[Message]) -> str:
         """
         Convert a list of Messages to a JSON string representation.

pyrit/models/chat_message.py

@@ -10,13 +10,28 @@
 ALLOWED_CHAT_MESSAGE_ROLES = ["system", "user", "assistant", "simulated_assistant", "tool", "developer"]
 
 
+class ToolCallFunction(BaseModel):
+    """The ``function`` payload of an OpenAI Chat Completions tool call."""
+
+    model_config = ConfigDict(extra="forbid")
+    name: str
+    arguments: str
+
+
 class ToolCall(BaseModel):
-    """Represents a tool invocation requested by the assistant."""
+    """
+    Represents a tool invocation requested by the assistant.
+
+    Matches the OpenAI Chat Completions API ``tool_calls`` shape: each entry
+    has a provider-issued ``id``, a ``type`` string (currently always
+    ``"function"``), and a nested ``function`` object carrying the tool
+    ``name`` and JSON-encoded ``arguments``.
+    """
 
     model_config = ConfigDict(extra="forbid")
     id: str
     type: str
-    function: str
+    function: ToolCallFunction
 
 
 class ChatMessage(BaseModel):
@@ -26,11 +41,12 @@ class ChatMessage(BaseModel):
     The content field can be:
     - A simple string for single-part text messages
     - A list of dicts for multipart messages (e.g., text + images)
+    - ``None`` for assistant messages whose payload is a tool-call only
     """
 
     model_config = ConfigDict(extra="forbid")
     role: ChatMessageRole
-    content: Union[str, list[dict[str, Any]]]
+    content: Optional[Union[str, list[dict[str, Any]]]] = None
     name: Optional[str] = None
     tool_calls: Optional[list[ToolCall]] = None
     tool_call_id: Optional[str] = None