Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

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

Open
wants to merge 12 commits into
base: main
Choose a base branch
from
Open

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

Changes from all commits
Commits
Show all changes
12 commits
Select commit Hold shift + click to select a range
8bca41e
Add files via upload
Git-Noob123 Jun 12, 2024
a319a65
Update agentchat_postresponse_secret_redaction.ipynb, removed error f…
Git-Noob123 Jun 12, 2024
99337b3
Removed cell output
Git-Noob123 Jun 12, 2024
0a28a50
Update agentchat_postresponse_secret_redaction.ipynb
Git-Noob123 Jun 12, 2024
15224f1
Update agentchat_postresponse_secret_redaction.ipynb
Git-Noob123 Jun 12, 2024
53c04e2
Update agentchat_postresponse_secret_redaction.ipynb
Git-Noob123 Jun 12, 2024
229154a
Update agentchat_postresponse_secret_redaction.ipynb
Git-Noob123 Jun 12, 2024
cd70b6d
Update agentchat_postresponse_secret_redaction.ipynb
Git-Noob123 Jun 12, 2024
6e6a276
Update agentchat_postresponse_secret_redaction.ipynb
Git-Noob123 Jun 12, 2024
5fbb833
Update agentchat_postresponse_secret_redaction.ipynb
Git-Noob123 Jun 12, 2024
4b8ab22
Merge branch 'microsoft:main' into post_process_message_redaction_sample
Git-Noob123 Jun 12, 2024
a63db7c
Added changes based on comments in pr
Git-Noob123 Jun 14, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
305 changes: 305 additions & 0 deletions notebook/agentchat_postresponse_secret_redaction.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,305 @@
{
"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 pyautogen:\n",
"\n",
"````{=mdx}\n",
":::info Requirements\n",
"Install `pyautogen`:\n",
"```bash\n",
"pip install pyautogen\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": [
"from typing import Dict, Union, Optional\n",
"import os\n",
"import copy\n",
"\n",
"import autogen\n",
"from autogen import ConversableAgent, Agent\n",
"from autogen.coding import LocalCommandLineCodeExecutor\n",
"from autogen.code_utils import create_virtual_env"
]
},
{
"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",
" \"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",
"# 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(\n",
" timeout=10, \n",
" work_dir=temp_file_fir, \n",
" virtual_env_context=venv_context\n",
")\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(message: Union[Dict, str], sender: Optional[Agent] = None, recipient: Agent = None, silent: bool = None ) -> 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(\n",
" code_executor_agent, message=agent_message, clear_history=False\n",
")"
]
},
{
"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
}