{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# MessageNormalizer\n",
    "\n",
    "MessageNormalizers convert PyRIT's `Message` format into other formats that specific targets require. Different LLMs and APIs expect messages in different formats:\n",
    "\n",
    "- **OpenAI-style APIs** expect `ChatMessage` objects with `role` and `content` fields\n",
    "- **HuggingFace models** expect specific chat templates (ChatML, Llama, Mistral, etc.)\n",
    "- **Some models** don't support system messages and need them merged into user messages\n",
    "- **Attack components** sometimes need conversation history as a formatted text string\n",
    "\n",
    "The `MessageNormalizer` classes handle these conversions, making it easy to work with any target regardless of its expected input format.\n",
    "\n",
    "## Memory is canonical, the normalized payload is not\n",
    "\n",
    "A normalizer builds a **target-facing view at send time**. It is never written back to memory.\n",
    "\n",
    "This matters most for a target that cannot accept editable history. If you prepend eight\n",
    "structured turns to a conversation, memory keeps eight structured turns, and the UI, scorers,\n",
    "resume, and exports all see eight turns. `HistorySquashNormalizer` flattens those turns into\n",
    "a single prompt only for the wire, then discards the flattened copy. The two representations\n",
    "are expected to differ.\n",
    "\n",
    "Flattening the conversation in memory instead would break scoring, resume, and evaluation for\n",
    "the sake of one target's wire format. If a prepended piece is not text (an image, for example),\n",
    "the flattened view holds a text placeholder and the normalizer logs a warning, because the\n",
    "target receives a description instead of the media.\n",
    "\n",
    "For prepended history on a target without editable history, PyRIT:\n",
    "\n",
    "1. Converts and persists the structured prepended messages.\n",
    "2. Converts the live request.\n",
    "3. Applies the per-send `EDITABLE_HISTORY` normalizer selected by\n",
    "   `PrependedConversationConfig`.\n",
    "4. Applies the target's remaining capability normalizers.\n",
    "5. Serializes the normalized view and invokes the provider.\n",
    "\n",
    "The attack-owned `PrependedHistorySendContext` records the persisted prepended-message boundary.\n",
    "Stateful targets consume it after the first successful target invocation; failed or cancelled\n",
    "invocations retain it for retry. Stateless targets reuse it for each current request. Targets\n",
    "interact with that state only through an internal `TargetSendContext` protocol at the send boundary.\n",
    "\n",
    "## Base Classes\n",
    "\n",
    "There are two base normalizer types:\n",
    "- **`MessageListNormalizer[T]`**: Converts `list[Message]` → `list[T]` (e.g., to `ChatMessage` objects)\n",
    "- **`MessageStringNormalizer`**: Converts `list[Message]` → `str` (e.g., to ChatML format)\n",
    "\n",
    "Some normalizers implement both interfaces."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1",
   "metadata": {},
   "source": [
    "## Two History-Squashing Scopes\n",
    "\n",
    "PyRIT uses `HistorySquashNormalizer` in two places that adapt different target capabilities:\n",
    "\n",
    "- **Multi-turn support** answers whether the target can continue a conversation across sends.\n",
    "- **Editable-history support** answers whether PyRIT can supply or rewrite earlier turns before\n",
    "  sending the current message.\n",
    "\n",
    "| Target capabilities | Prepended-history behavior |\n",
    "|---|---|\n",
    "| Multi-turn with editable history | Send the structured history directly; no history squashing is needed. |\n",
    "| Multi-turn without editable history | A context-scoped `HistorySquashNormalizer` encodes the prepended history and first live message into the initial request. Later turns use the target's own conversation state. |\n",
    "| Single-turn without editable history | The context-scoped `HistorySquashNormalizer` first produces one message. The target's ordinary use of the same normalizer then sees one message and does nothing. |\n",
    "\n",
    "The first-turn distinction comes from the attack-owned prepended-history send context, not from a\n",
    "separate normalizer implementation. The context applies its configured `HistorySquashNormalizer`\n",
    "once to bootstrap prepended history for a target that cannot accept caller-supplied prior turns.\n",
    "Its key use case is a multi-turn, server-managed target without editable history.\n",
    "\n",
    "The target's ordinary capability pipeline independently uses `HistorySquashNormalizer` for a target\n",
    "without multi-turn support. Both scopes preserve original-versus-converted text views and keep\n",
    "non-text pieces from the current request separate. The context-scoped use can supply a custom\n",
    "formatter; the ordinary use defaults to `[Conversation History]` and `[Current Message]` sections."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Sample messages created:\n",
      "  system: You are a helpful assistant....\n",
      "  user: What is the capital of France?...\n",
      "  assistant: The capital of France is Paris....\n",
      "  user: What about Germany?...\n"
     ]
    }
   ],
   "source": [
    "from pyrit.models import Message\n",
    "\n",
    "# Create sample messages for demonstration\n",
    "system_message = Message.from_prompt(prompt=\"You are a helpful assistant.\", role=\"system\")\n",
    "user_message = Message.from_prompt(prompt=\"What is the capital of France?\", role=\"user\")\n",
    "assistant_message = Message.from_prompt(prompt=\"The capital of France is Paris.\", role=\"assistant\")\n",
    "followup_message = Message.from_prompt(prompt=\"What about Germany?\", role=\"user\")\n",
    "\n",
    "messages = [system_message, user_message, assistant_message, followup_message]\n",
    "\n",
    "print(\"Sample messages created:\")\n",
    "for msg in messages:\n",
    "    print(f\"  {msg.api_role}: {msg.get_piece().converted_value[:50]}...\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3",
   "metadata": {},
   "source": [
    "## ChatMessageNormalizer\n",
    "\n",
    "The `ChatMessageNormalizer` converts `Message` objects to `ChatMessage` objects, which are the standard format for OpenAI chat-based API calls. It handles both single-part text messages and multipart messages (with images, audio, etc.).\n",
    "\n",
    "Key features:\n",
    "- Single text pieces become simple string content\n",
    "- Multiple pieces become content arrays with type information\n",
    "- Supports `use_developer_role=True` for newer OpenAI models that use \"developer\" instead of \"system\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ChatMessage output:\n",
      "  Role: system, Content: You are a helpful assistant.\n",
      "  Role: user, Content: What is the capital of France?\n",
      "  Role: assistant, Content: The capital of France is Paris.\n",
      "  Role: user, Content: What about Germany?\n"
     ]
    }
   ],
   "source": [
    "from pyrit.message_normalizer import ChatMessageNormalizer\n",
    "\n",
    "# Standard usage\n",
    "normalizer = ChatMessageNormalizer()\n",
    "chat_messages = await normalizer.normalize_async(messages)  # type: ignore[top-level-await]\n",
    "\n",
    "print(\"ChatMessage output:\")\n",
    "for msg in chat_messages:  # type: ignore[assignment]\n",
    "    print(f\"  Role: {msg.role}, Content: {msg.content}\")  # type: ignore[attr-defined]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ChatMessage with developer role:\n",
      "  Role: developer, Content: You are a helpful assistant.\n",
      "  Role: user, Content: What is the capital of France?\n",
      "  Role: assistant, Content: The capital of France is Paris.\n",
      "  Role: user, Content: What about Germany?\n"
     ]
    }
   ],
   "source": [
    "# With developer role for newer OpenAI models (o1, o3, gpt-4.1+)\n",
    "dev_normalizer = ChatMessageNormalizer(use_developer_role=True)\n",
    "dev_chat_messages = await dev_normalizer.normalize_async(messages)  # type: ignore[top-level-await]\n",
    "\n",
    "print(\"ChatMessage with developer role:\")\n",
    "for msg in dev_chat_messages:  # type: ignore[assignment]\n",
    "    print(f\"  Role: {msg.role}, Content: {msg.content}\")  # type: ignore[attr-defined]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "JSON string output:\n",
      "[\n",
      "  {\n",
      "    \"role\": \"system\",\n",
      "    \"content\": \"You are a helpful assistant.\"\n",
      "  },\n",
      "  {\n",
      "    \"role\": \"user\",\n",
      "    \"content\": \"What is the capital of France?\"\n",
      "  },\n",
      "  {\n",
      "    \"role\": \"assistant\",\n",
      "    \"content\": \"The capital of France is Paris.\"\n",
      "  },\n",
      "  {\n",
      "    \"role\": \"user\",\n",
      "    \"content\": \"What about Germany?\"\n",
      "  }\n",
      "]\n"
     ]
    }
   ],
   "source": [
    "# ChatMessageNormalizer also implements MessageStringNormalizer for JSON output\n",
    "json_output = await normalizer.normalize_string_async(messages)  # type: ignore[top-level-await]\n",
    "print(\"JSON string output:\")\n",
    "print(json_output)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7",
   "metadata": {},
   "source": [
    "## GenericSystemSquashNormalizer\n",
    "\n",
    "Some models don't support system messages. The `GenericSystemSquashNormalizer` combines consecutive\n",
    "system messages in their original order and merges them into the user message immediately following\n",
    "them. If no user immediately follows, it converts the system messages to a user message in their\n",
    "original position.\n",
    "\n",
    "For example, `system: Policy`, `system: Persona`, `user: Question` becomes one user message containing\n",
    "the Policy and Persona instructions followed by the Question.\n",
    "\n",
    "The format is:\n",
    "```\n",
    "### Instructions ###\n",
    "\n",
    "{system_content}\n",
    "\n",
    "######\n",
    "\n",
    "{user_content}\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Original message count: 4\n",
      "Squashed message count: 3\n",
      "\n",
      "First message after squashing:\n",
      "### Instructions ###\n",
      "\n",
      "You are a helpful assistant.\n",
      "\n",
      "######\n",
      "\n",
      "What is the capital of France?\n"
     ]
    }
   ],
   "source": [
    "from pyrit.message_normalizer import GenericSystemSquashNormalizer\n",
    "\n",
    "squash_normalizer = GenericSystemSquashNormalizer()\n",
    "squashed_messages = await squash_normalizer.normalize_async(messages)  # type: ignore[top-level-await]\n",
    "\n",
    "print(f\"Original message count: {len(messages)}\")\n",
    "print(f\"Squashed message count: {len(squashed_messages)}\")\n",
    "print(\"\\nFirst message after squashing:\")\n",
    "print(squashed_messages[0].get_piece().converted_value)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9",
   "metadata": {},
   "source": [
    "## ConversationContextNormalizer\n",
    "\n",
    "The `ConversationContextNormalizer` formats conversation history as a turn-based text string. This is useful for:\n",
    "- Including conversation history in attack prompts\n",
    "- Logging and debugging conversations\n",
    "- Creating context strings for adversarial chat\n",
    "\n",
    "The output format is:\n",
    "```\n",
    "Turn 1:\n",
    "User: <content>\n",
    "Assistant: <content>\n",
    "\n",
    "Turn 2:\n",
    "User: <content>\n",
    "...\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Conversation context format:\n",
      "Turn 1:\n",
      "user: What is the capital of France?\n",
      "assistant: The capital of France is Paris.\n",
      "Turn 2:\n",
      "user: What about Germany?\n"
     ]
    }
   ],
   "source": [
    "from pyrit.message_normalizer import ConversationContextNormalizer\n",
    "\n",
    "context_normalizer = ConversationContextNormalizer()\n",
    "context_string = await context_normalizer.normalize_string_async(messages)  # type: ignore[top-level-await]\n",
    "\n",
    "print(\"Conversation context format:\")\n",
    "print(context_string)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11",
   "metadata": {},
   "source": [
    "## TokenizerTemplateNormalizer\n",
    "\n",
    "The `TokenizerTemplateNormalizer` uses HuggingFace tokenizer chat templates to format messages. This is essential for:\n",
    "- Local LLM inference with proper formatting\n",
    "- Matching the exact prompt format a model was trained with\n",
    "- Working with various open-source models\n",
    "\n",
    "### Using Model Aliases\n",
    "\n",
    "For convenience, common models have aliases that automatically configure the normalizer:\n",
    "\n",
    "| Alias | Model | Notes |\n",
    "|-------|-------|-------|\n",
    "| `chatml` | HuggingFaceH4/zephyr-7b-beta | No auth required |\n",
    "| `phi3` | microsoft/Phi-3-mini-4k-instruct | No auth required |\n",
    "| `qwen` | Qwen/Qwen2-7B-Instruct | No auth required |\n",
    "| `llama3` | meta-llama/Meta-Llama-3-8B-Instruct | Requires HF token |\n",
    "| `gemma` | google/gemma-7b-it | Requires HF token, auto-squashes system |\n",
    "| `mistral` | mistralai/Mistral-7B-Instruct-v0.2 | Requires HF token |"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "No HuggingFace token provided. Gated models may fail to load without authentication.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ChatML formatted output:\n",
      "<|system|>\n",
      "You are a helpful assistant.</s>\n",
      "<|user|>\n",
      "What is the capital of France?</s>\n",
      "<|assistant|>\n",
      "The capital of France is Paris.</s>\n",
      "<|user|>\n",
      "What about Germany?</s>\n",
      "<|assistant|>\n",
      "\n"
     ]
    }
   ],
   "source": [
    "from pyrit.message_normalizer import TokenizerTemplateNormalizer\n",
    "\n",
    "# Using an alias (no auth required for this model)\n",
    "template_normalizer = TokenizerTemplateNormalizer.from_model(\"chatml\")\n",
    "formatted = await template_normalizer.normalize_string_async(messages)  # type: ignore[top-level-await]\n",
    "\n",
    "print(\"ChatML formatted output:\")\n",
    "print(formatted)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13",
   "metadata": {},
   "source": [
    "### System Message Behavior\n",
    "\n",
    "The `TokenizerTemplateNormalizer` supports different strategies for handling system messages:\n",
    "\n",
    "- **`keep`**: Pass system messages as-is (default)\n",
    "- **`squash`**: Merge system messages into the following user message using `GenericSystemSquashNormalizer`\n",
    "- **`ignore`**: Drop system messages entirely\n",
    "- **`developer`**: Change system role to developer role (for newer OpenAI models)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "No HuggingFace token provided. Gated models may fail to load without authentication.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ChatML with squashed system message:\n",
      "<|user|>\n",
      "### Instructions ###\n",
      "\n",
      "You are a helpful assistant.\n",
      "\n",
      "######\n",
      "\n",
      "What is the capital of France?</s>\n",
      "<|assistant|>\n",
      "The capital of France is Paris.</s>\n",
      "<|user|>\n",
      "What about Germany?</s>\n",
      "<|assistant|>\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Using squash behavior for models that don't support system messages\n",
    "squash_template_normalizer = TokenizerTemplateNormalizer.from_model(\"chatml\", system_message_behavior=\"squash\")\n",
    "squashed_formatted = await squash_template_normalizer.normalize_string_async(messages)  # type: ignore[top-level-await]\n",
    "\n",
    "print(\"ChatML with squashed system message:\")\n",
    "print(squashed_formatted)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15",
   "metadata": {},
   "source": [
    "### Using Custom Models\n",
    "\n",
    "You can also use any HuggingFace model with a chat template by providing the full model name."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "No HuggingFace token provided. Gated models may fail to load without authentication.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "TinyLlama formatted output:\n",
      "<|system|>\n",
      "You are a helpful assistant.</s>\n",
      "<|user|>\n",
      "What is the capital of France?</s>\n",
      "<|assistant|>\n",
      "The capital of France is Paris.</s>\n",
      "<|user|>\n",
      "What about Germany?</s>\n",
      "<|assistant|>\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Using a custom HuggingFace model\n",
    "# Note: Some models require authentication via HUGGINGFACE_TOKEN env var or token parameter\n",
    "custom_normalizer = TokenizerTemplateNormalizer.from_model(\"TinyLlama/TinyLlama-1.1B-Chat-v1.0\")\n",
    "custom_formatted = await custom_normalizer.normalize_string_async(messages)  # type: ignore[top-level-await]\n",
    "\n",
    "print(\"TinyLlama formatted output:\")\n",
    "print(custom_formatted)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17",
   "metadata": {},
   "source": [
    "## Creating Custom Normalizers\n",
    "\n",
    "You can create custom normalizers by extending the base classes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Markdown formatted output:\n",
      "**System**: You are a helpful assistant.\n",
      "\n",
      "**User**: What is the capital of France?\n",
      "\n",
      "**Assistant**: The capital of France is Paris.\n",
      "\n",
      "**User**: What about Germany?\n"
     ]
    }
   ],
   "source": [
    "from pyrit.message_normalizer import MessageStringNormalizer\n",
    "from pyrit.models import Message\n",
    "\n",
    "\n",
    "class SimpleMarkdownNormalizer(MessageStringNormalizer):\n",
    "    \"\"\"Custom normalizer that formats messages as Markdown.\"\"\"\n",
    "\n",
    "    async def normalize_string_async(self, messages: list[Message]) -> str:\n",
    "        lines = []\n",
    "        for msg in messages:\n",
    "            piece = msg.get_piece()\n",
    "            role = piece.api_role.capitalize()\n",
    "            content = piece.converted_value\n",
    "            lines.append(f\"**{role}**: {content}\")\n",
    "        return \"\\n\\n\".join(lines)\n",
    "\n",
    "\n",
    "# Use the custom normalizer\n",
    "md_normalizer = SimpleMarkdownNormalizer()\n",
    "md_output = await md_normalizer.normalize_string_async(messages)  # type: ignore[top-level-await]\n",
    "\n",
    "print(\"Markdown formatted output:\")\n",
    "print(md_output)"
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "main_language": "python"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.14.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}