Adding a sample to redact sensitive information after an agent generates a response

notebook/agentchat_postresponse_secret_redaction.ipynb

@@ -0,0 +1,304 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Redact Sensitive Information Post-Response with Hooks\n",
+    "\n",
+    "This notebook shows an example on how we can use hooks in agents to effectively hide/redact sensitive data (i.e. API keys) from getting printed out to avoid potential data leak from agent's responses.\n",
+    "\n",
+    "## What are Hooks in Autogen?\n",
+    "\n",
+    "Similar to hooks in other frameworks, Autogen hooks allow you to customize your agent at each of its response-generation stages.  \n",
+    "\n",
+    "Autogen offers 3 hookable methods:  \n",
+    "\n",
+    "`process_last_received_message`:Process the last received message, invoked right before a new message is delivered to an agent. \n",
+    "`process_all_messages_before_reply`: Process incoming messages to an agent, invoked before an agent start generating a response.  \n",
+    "`process_message_before_send`: Process the agent's generated response, invoked after response generation.  \n",
+    "\n",
+    "To use them, simply pass a hookable method as a string in hookable_method attribute and your own hook function in hook attribute in register_hook. For example:  \n",
+    "\n",
+    "`my_agent.register_hook(hookable_method=\"process_message_before_send\", hook=my_hook_function)`"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Requirements\n",
+    "\n",
+    "AutoGen requires `Python>=3.8`. To run this notebook example, please install autogen-agentchat:\n",
+    "\n",
+    "````{=mdx}\n",
+    ":::info Requirements\n",
+    "Install `pyautogen`:\n",
+    "```bash\n",
+    "pip install autogen-agentchat~=0.2\n",
+    "```\n",
+    "\n",
+    "For more information, please refer to the [installation guide](/docs/installation/).\n",
+    ":::\n",
+    "````"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import copy\n",
+    "import os\n",
+    "from typing import Dict, Optional, Union\n",
+    "\n",
+    "import autogen\n",
+    "from autogen import Agent, ConversableAgent\n",
+    "from autogen.code_utils import create_virtual_env\n",
+    "from autogen.coding import LocalCommandLineCodeExecutor"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Set Environment Variables"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "os.environ[\"VENV_DIR\"] = \"<Your venv directory that is used for code executor agent>\"\n",
+    "os.environ[\"TEMP_FILE_DIR\"] = \"<Your directory for temporary code files>\"\n",
+    "os.environ[\"LLM_API_KEY\"] = \"<Your API key for llm config>\"\n",
+    "os.environ[\"LLM_API_TYPE\"] = \"<Your API type for llm config>\"\n",
+    "os.environ[\"LLM_API_VERSION\"] = \"<Your API version for llm config>\"\n",
+    "os.environ[\"LLM_API_URL\"] = \"<Your API Url for llm config>\"\n",
+    "\n",
+    "# Get the env variables\n",
+    "venv_dir = os.environ[\"VENV_DIR\"]\n",
+    "temp_file_fir = os.environ[\"TEMP_FILE_DIR\"]\n",
+    "\n",
+    "llm_config = {\n",
+    "    \"config_list\": [\n",
+    "        {\n",
+    "            \"model\": \"gpt-4\",\n",
+    "            \"api_key\": os.environ[\"LLM_API_KEY\"],\n",
+    "            \"api_type\": os.environ[\"LLM_API_TYPE\"],\n",
+    "            \"api_version\": os.environ[\"LLM_API_VERSION\"],\n",
+    "            \"base_url\": os.environ[\"LLM_API_URL\"],\n",
+    "        }\n",
+    "    ],\n",
+    "}\n",
+    "\n",
+    "# The string that sensitive data will be redated to\n",
+    "replacementString = \"REDACTED\"\n",
+    "\n",
+    "# Get the env secrets and setup virtual env for agent to use\n",
+    "env_secrets = os.environ\n",
+    "\n",
+    "venv_context = create_virtual_env(venv_dir)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Setup the Agents\n",
+    "\n",
+    "We have 2 agents here: A proxy agent and a code executor agent. The user communicates to the code executer agent through the proxy agent."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "executor = LocalCommandLineCodeExecutor(timeout=10, work_dir=temp_file_fir, virtual_env_context=venv_context)\n",
+    "\n",
+    "code_executor_agent = ConversableAgent(\n",
+    "    \"code_executor_agent\",\n",
+    "    llm_config=False,\n",
+    "    code_execution_config={\"executor\": executor},\n",
+    "    human_input_mode=\"NEVER\",\n",
+    "    max_consecutive_auto_reply=1,\n",
+    ")\n",
+    "\n",
+    "proxy_agent = autogen.UserProxyAgent(\n",
+    "    \"user_proxy\",\n",
+    "    human_input_mode=\"NEVER\",\n",
+    "    llm_config=llm_config,\n",
+    "    system_message=\"DO NOT display any thing that is sensitive\",\n",
+    "    max_consecutive_auto_reply=1,\n",
+    "    code_execution_config={\"use_docker\": False},\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "The function `transform_generated_response` will redact the message and will be registered as one of the agent's hooks. It will be invoked after the agent generates a response but before the agent sends out the response to the next agent. \n",
+    "\n",
+    "`transform_generated_response` requires the hook to accept one argument `message` that can be either a dictionary or a string. The argument `message` is the agent's response and sensitive information will be redacted from it. There're 3 other inputs `sender` `recipient` and `silent` but those are not required for message processing."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 14,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def transform_generated_response(\n",
+    "    message: Union[Dict, str], sender: Optional[Agent] = None, recipient: Agent = None, silent: bool = None\n",
+    ") -> Union[Dict, str]:\n",
+    "    temp_message = copy.deepcopy(message)\n",
+    "    all_secrets = sorted(env_secrets.values(), key=len, reverse=True)\n",
+    "    if isinstance(temp_message, Dict):\n",
+    "        for secret in all_secrets:\n",
+    "            if isinstance(temp_message[\"content\"], str):\n",
+    "                if secret != \"\" and secret in temp_message[\"content\"]:\n",
+    "                    temp_message[\"content\"] = temp_message[\"content\"].replace(secret, replacementString)\n",
+    "            elif isinstance(temp_message[\"content\"], list):\n",
+    "                for item in temp_message[\"content\"]:\n",
+    "                    if item[\"type\"] == \"text\":\n",
+    "                        if secret != \"\" and secret in item[\"text\"]:\n",
+    "                            item[\"text\"] = item[\"text\"].replace(secret, replacementString)\n",
+    "    if isinstance(temp_message, str):\n",
+    "        for secret in all_secrets:\n",
+    "            if secret != \"\" and secret in temp_message:\n",
+    "                temp_message = temp_message.replace(secret, replacementString)\n",
+    "\n",
+    "    return temp_message"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Use the `register_hook` method to register `transform_generated_response` to our code executor agent.\n",
+    "\n",
+    "We pass `\"process_message_before_send\"` here so that our message transformation function can be executed after the response is generated and before it's sent out."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 15,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "code_executor_agent.register_hook(hookable_method=\"process_message_before_send\", hook=transform_generated_response)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Initiate a Chat and Print the Result\n",
+    "\n",
+    "Initiate a chat by asking our proxy agent to execute a piece of code that retrieves an API key from the environment. The code executor agent will then execute the code but the response will be redacted to hide the API key."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 16,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "agent_message = \"\"\"Run the code and show me the printed variable.\n",
+    "The code block is below:\n",
+    "```python\n",
+    "import os\n",
+    "print(os.environ[\"LLM_API_KEY\"])\n",
+    "```\n",
+    "This is the end of the message.\n",
+    "\"\"\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 17,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\u001b[33muser_proxy\u001b[0m (to code_executor_agent):\n",
+      "\n",
+      "Run the code and show me the printed variable.\n",
+      "The code block is below:\n",
+      "```python\n",
+      "import os\n",
+      "print(os.environ[\"LLM_API_KEY\"])\n",
+      "```\n",
+      "This is the end of the message.\n",
+      "\n",
+      "\n",
+      "--------------------------------------------------------------------------------\n",
+      "\u001b[31m\n",
+      ">>>>>>>> EXECUTING CODE BLOCK (inferred language is python)...\u001b[0m\n",
+      "\u001b[33mcode_executor_agent\u001b[0m (to user_proxy):\n",
+      "\n",
+      "exitcode: REDACTED (executiREDACTED succeeded)\n",
+      "Code output: REDACTED\n",
+      "\n",
+      "\n",
+      "--------------------------------------------------------------------------------\n",
+      "\u001b[33muser_proxy\u001b[0m (to code_executor_agent):\n",
+      "\n",
+      "Sorry, but I can't assist with that.\n",
+      "\n",
+      "--------------------------------------------------------------------------------\n"
+     ]
+    }
+   ],
+   "source": [
+    "result = proxy_agent.initiate_chat(code_executor_agent, message=agent_message, clear_history=False)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 18,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "ChatResult(chat_id=None, chat_history=[{'content': 'Run the code and show me the printed variable.\\nThe code block is below:\\n```python\\nimport os\\nprint(os.environ[\"LLM_API_KEY\"])\\n```\\nThis is the end of the message.\\n', 'role': 'assistant'}, {'content': 'exitcode: REDACTED (executiREDACTED succeeded)\\nCode output: REDACTED\\n', 'role': 'user'}, {'content': \"Sorry, but I can't assist with that.\", 'role': 'assistant'}], summary=\"Sorry, but I can't assist with that.\", cost={'usage_including_cached_inference': {'total_cost': 0.00318, 'gpt-4': {'cost': 0.00318, 'prompt_tokens': 86, 'completion_tokens': 10, 'total_tokens': 96}}, 'usage_excluding_cached_inference': {'total_cost': 0}}, human_input=[])\n"
+     ]
+    }
+   ],
+   "source": [
+    "print(result)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "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.10.0"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}