Incorporate AgentCore deployment to Terraform

chat_to_agentcore.ipynb

@@ -0,0 +1,361 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "71db7e7e",
+   "metadata": {},
+   "source": [
+    "# Test your agent from a Python Notebook\n",
+    "\n",
+    "This interactive notebook provides an alternative way to test your deployed agent, besides the Streamlit UI app."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "cbced266",
+   "metadata": {},
+   "source": [
+    "## Kernel selection and prerequisites\n",
+    "\n",
+    "You can use the `cx-agent-backend/.venv` as a kernel. If this is not set up already, run the following from your terminal:\n",
+    "\n",
+    "```bash\n",
+    "cd cx-agent-backend\n",
+    "uv venv\n",
+    "uv sync --all-extras --frozen\n",
+    "```\n",
+    "\n",
+    "This notebook assumes:\n",
+    "1. You've already deployed the main solution as described in [README.md](./README.md)\n",
+    "2. Your Python kernel is already configured with AWS credentials and target AWS Region (for example via environment variables, potentially set via a `.env` file as documented [here for VSCode](https://code.visualstudio.com/docs/python/environments#_environment-variables)).\n",
+    "    - Note that the AWS SDK for Python, `boto3`, [expects](https://boto3.amazonaws.com/v1/documentation/api/latest/guide/configuration.html#using-environment-variables) an `AWS_DEFAULT_REGION` environment variable rather than `AWS_REGION`."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "b51a5349",
+   "metadata": {},
+   "source": [
+    "## Dependencies and setup\n",
+    "\n",
+    "First we'll import the necessary libraries, and initialize clients for AWS Services, and define some utility functions that'll be used later:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "93b97307",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Python Built-Ins:\n",
+    "import base64\n",
+    "import json\n",
+    "import getpass\n",
+    "import hashlib\n",
+    "import hmac\n",
+    "import os\n",
+    "import uuid\n",
+    "import secrets\n",
+    "import string\n",
+    "import urllib.parse\n",
+    "\n",
+    "# External Libraries:\n",
+    "import boto3  # AWS SDK for Python\n",
+    "import requests  # For making raw HTTP(S) API calls\n",
+    "\n",
+    "# AWS Service Clients:\n",
+    "botosess = boto3.Session()  # You could set `region_name` here explicitly if wanted\n",
+    "cognito_client = botosess.client(\"cognito-idp\")  # Cognito (Identity Provider)\n",
+    "\n",
+    "\n",
+    "def _set_if_undefined(var: str, name: str | None = None) -> str:\n",
+    "    \"\"\"Utility to prompt user once for a value, and cache it in environment variable\"\"\"\n",
+    "    if not os.environ.get(var):\n",
+    "        os.environ[var] = getpass.getpass(f\"Please provide your {name or var}:\")\n",
+    "    return os.environ[var]\n",
+    "\n",
+    "\n",
+    "def calculate_secret_hash(username, client_id, client_secret):\n",
+    "    \"\"\"Utility to hash a username + client ID + client secret for Cognito login\"\"\"\n",
+    "    message = username + client_id\n",
+    "    return base64.b64encode(\n",
+    "        hmac.new(\n",
+    "            client_secret.encode(\"utf-8\"),\n",
+    "            message.encode(\"utf-8\"),\n",
+    "            hashlib.sha256\n",
+    "        ).digest()\n",
+    "    ).decode(\"utf-8\")\n",
+    "\n",
+    "\n",
+    "def invoke_agent(\n",
+    "    message,\n",
+    "    agent_arn,\n",
+    "    auth_token,\n",
+    "    session_id,\n",
+    "    qualifier=\"DEFAULT\",\n",
+    "    region=botosess.region_name,\n",
+    "):\n",
+    "    \"\"\"Invoke Bedrock AgentCore runtime with a message.\"\"\"\n",
+    "    escaped_agent_arn = urllib.parse.quote(agent_arn, safe='')\n",
+    "    response = requests.post(\n",
+    "        f\"https://bedrock-agentcore.{region}.amazonaws.com/runtimes/{escaped_agent_arn}/invocations?qualifier={qualifier}\",\n",
+    "        headers={\n",
+    "            \"Authorization\": f\"Bearer {auth_token}\",\n",
+    "            \"Content-Type\": \"application/json\",\n",
+    "            \"X-Amzn-Bedrock-AgentCore-Runtime-Session-Id\": session_id\n",
+    "        },\n",
+    "        data=json.dumps({\"input\": {\"prompt\": message, \"conversation_id\": session_id}}),\n",
+    "        timeout=61,\n",
+    "    )\n",
+    "    \n",
+    "    print(f\"Status Code: {response.status_code}\")\n",
+    "    \n",
+    "    if response.status_code == 200:\n",
+    "        return response.json()\n",
+    "    else:\n",
+    "        raise ValueError(f\"HTTP {response.status_code}: {response.text}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "b1f027bf",
+   "metadata": {},
+   "source": [
+    "## Fetch access token from Amazon Cognito\n",
+    "\n",
+    "To talk to the AgentCore-deployed agent, we'll need to log in to Amazon Cognito to fetch a session token.\n",
+    "\n",
+    "You'll need to fetch your Cognito user_pool_id and client_id in the cell below, which you can view by running the `terraform output` command in your terminal:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "09e8aa2d",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "user_pool_id = TODO # E.g. run `terraform output -raw user_pool_id`\n",
+    "client_id = TODO # E.g. run `terraform output -raw client_id`\n",
+    "\n",
+    "# From these we should be able to look up the client secret automatically:\n",
+    "client_secret = cognito_client.describe_user_pool_client(\n",
+    "    UserPoolId=user_pool_id,\n",
+    "    ClientId=client_id\n",
+    ")[\"UserPoolClient\"][\"ClientSecret\"]"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "66238615",
+   "metadata": {},
+   "source": [
+    "The next cell will prompt you for your Cognito username (email address) and password, or re-use the existing one if you run the cell again without restarting the notebook:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "00938dba",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "username = _set_if_undefined(\"COGNITO_USERNAME\", \"Cognito user name (email address)\")\n",
+    "password = _set_if_undefined(\"COGNITO_PASSWORD\", \"Cognito password\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "e00e6221",
+   "metadata": {},
+   "source": [
+    "The deployment steps in [README.md](./README.md) guide you through setting up your Cognito user from the AWS CLI, but you could instead un-comment and run the below to achieve the same effect from Python:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "fde3f16b",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "## Create a user (with temporary password)\n",
+    "# create_user_resp = cognito_client.admin_create_user(\n",
+    "#     UserPoolId=user_pool_id,\n",
+    "#     Username=username,\n",
+    "#     # Temp password is randomized here because we'll never use it:\n",
+    "#     TemporaryPassword=\"\".join((\n",
+    "#         secrets.choice(\n",
+    "#             string.ascii_uppercase + string.ascii_lowercase + string.digits +\n",
+    "#             \"^$*.[]{}()?-'\\\"!@#%&/\\\\,><':;|_~`+=\"\n",
+    "#         ) for i in range(20)\n",
+    "#     )),\n",
+    "#     MessageAction=\"SUPPRESS\"\n",
+    "# )\n",
+    "# print(f\"User created: {create_user_resp['User']['Username']}\")\n",
+    "\n",
+    "## Override the password to the given value (permanently)\n",
+    "# set_password_resp = cognito_client.admin_set_user_password(\n",
+    "#     UserPoolId=user_pool_id,\n",
+    "#     Username=username,\n",
+    "#     Password=password,\n",
+    "#     Permanent=True,\n",
+    "# )\n",
+    "# print(\"Password set successfully\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "ac6d3aff",
+   "metadata": {},
+   "source": [
+    "With the configuration set up, we're ready to request an access token from Cognito:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "33858093",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "auth_resp = cognito_client.initiate_auth(\n",
+    "    ClientId=client_id,\n",
+    "    AuthFlow=\"USER_PASSWORD_AUTH\",\n",
+    "    AuthParameters={\n",
+    "        \"USERNAME\": username,\n",
+    "        \"PASSWORD\": password,\n",
+    "        \"SECRET_HASH\": calculate_secret_hash(username, client_id, client_secret),\n",
+    "    }\n",
+    ")\n",
+    "\n",
+    "access_token = auth_resp[\"AuthenticationResult\"][\"AccessToken\"]\n",
+    "print(\"Access token fetched\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "05e0c2e6",
+   "metadata": {},
+   "source": [
+    "## Invoke the agent\n",
+    "\n",
+    "With the access token ready, we're almost ready to invoke our AgentCore Agent. First though, you'll need to:\n",
+    "1. Look up the deployed AgentRuntime ARN from the terraform, and\n",
+    "2. Choose a session ID (we'll randomize this automatically)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6754ec78",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "agent_arn = TODO  # E.g. run `terraform output -raw agent_runtime_arn`\n",
+    "\n",
+    "session_id = str(uuid.uuid4())  # Can auto-generate this"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "9cf235a1",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "result = invoke_agent(\"Hello, can you help me resetting my router?\", agent_arn, access_token, session_id)\n",
+    "if result:\n",
+    "    print(json.dumps(result, indent=2))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "583c6b54",
+   "metadata": {},
+   "source": [
+    "### Testing math functionality"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "60914206",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "math_test = \"What is 25 + 17?\"\n",
+    "print(f\"Testing math with: {math_test}\")\n",
+    "result = invoke_agent(math_test, agent_arn, access_token, session_id)\n",
+    "if result:\n",
+    "    print(json.dumps(result, indent=2))\n",
+    "print(\"\\n\" + \"=\"*50 + \"\\n\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "ceac8013",
+   "metadata": {},
+   "source": [
+    "### Testing memory persistence"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "68de9444",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "message = \"Add 10 to the result\"\n",
+    "print(message)\n",
+    "result = invoke_agent(message, agent_arn, access_token, session_id)\n",
+    "if result:\n",
+    "    print(json.dumps(result, indent=2))\n",
+    "print(\"\\n\" + \"=\"*50 + \"\\n\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "5caf85c9",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "result = invoke_agent(\"What device did I want to reset?\", agent_arn, access_token, session_id)\n",
+    "if result:\n",
+    "    print(json.dumps(result, indent=2))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "646cef68",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": ".venv",
+   "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.12.8"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}