{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# GPT Action Library: AWS RedShift"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Introduction"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This page provides an instruction & guide for developers building a GPT Action for a specific application. Before you proceed, make sure to first familiarize yourself with the following information: \n",
    "- [Introduction to GPT Actions](https://platform.openai.com/docs/actions)\n",
    "- [Introduction to GPT Actions Library](https://platform.openai.com/docs/actions/actions-library)\n",
    "- [Example of Building a GPT Action from Scratch](https://platform.openai.com/docs/actions/getting-started)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This solution enables a GPT action to retrieve data from Redshift and perform data analysis.It uses AWS Functions, performing every action from AWS ecosystem and network. The middleware (AWS function) will perform the SQL query, wait for its completion and return the data as a file. The code is provided for information purpose only and should be modified to your needs.\n",
    "\n",
    "This solution uses the ability to [retrieve files in Actions](https://platform.openai.com/docs/actions/sending-files) and use them as if you had uploaded them directly to a conversation.\n",
    "\n",
    "This solution highlight a connection to Redshift serverless, the integration with a provisioned Redshift might differ slighltly to retrieve networks and set-up connection, the overall code and (minimal) integration should be similar."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Value & Example Business Use Cases"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Value**: Leverage ChatGPT's natural language capabilities to connect to Redshift's DWH.\n",
    "\n",
    "**Example Use Cases**:\n",
    "- Data scientists can connect to tables and run data analyses using ChatGPT's Data Analysis\n",
    "- Citizen data users can ask basic questions of their transactional data\n",
    "- Users gain more visibility into their data & potential anomalies"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Application Information"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Application Prerequisites"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Before you get started, make sure that:\n",
    "- You have access to a Redshift environment\n",
    "- You have the rights to deploy AWS function in the same VPC (Virtual Private Network)\n",
    "- Your AWS CLI is authenticated"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Middleware Information"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Install required libraries\n",
    "- Install AWS CLI, required for AWS SAM ([docs](https://docs.aws.amazon.com/cli/latest/userguide/getting-started-install.html#getting-started-install-instructions))\n",
    "- Install AWS SAM CLI ([docs](https://docs.aws.amazon.com/serverless-application-model/latest/developerguide/install-sam-cli.html))\n",
    "- Install Python\n",
    "- Install yq [docs](https://github.com/mikefarah/yq?tab=readme-ov-file#install)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Middleware function\n",
    "\n",
    "To create a function, follow the steps in the [AWS Middleware Action cookbook](https://cookbook.openai.com/examples/chatgpt/gpt_actions_library/gpt_middleware_aws_function).\n",
    "\n",
    "To deploy specifically an application that connects to Redshift, use the following code instead of the \"hello-world\" GitHub repository referenced in the Middleware AWS Function cookbook. You can either clone the repository or take the code pasted below and modify it to your needs.\n",
    "\n",
    "> This code is meant to be directional - while it should work out of the box, it is designed to be customized to your needs (see examples towards the end of this document).\n",
    "\n",
    "To get the code, you can clone openai-cookbook repository and navigate to the redshift-middleware directory\n",
    "\n",
    "```\n",
    "git clone https://github.com/pap-openai/redshift-middleware\n",
    "cd redshift-middleware\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import psycopg2\n",
    "import os\n",
    "import base64\n",
    "import tempfile\n",
    "import csv\n",
    "\n",
    "# Fetch Redshift credentials from environment variables\n",
    "host = os.environ['REDSHIFT_HOST']\n",
    "port = os.environ['REDSHIFT_PORT']\n",
    "user = os.environ['REDSHIFT_USER']\n",
    "password = os.environ['REDSHIFT_PASSWORD']\n",
    "database = os.environ['REDSHIFT_DB']\n",
    "\n",
    "def execute_statement(sql_statement):\n",
    "    try:\n",
    "        # Establish connection\n",
    "        conn = psycopg2.connect(\n",
    "            host=host,\n",
    "            port=port,\n",
    "            user=user,\n",
    "            password=password,\n",
    "            dbname=database\n",
    "        )\n",
    "        cur = conn.cursor()\n",
    "        cur.execute(sql_statement)\n",
    "        conn.commit()\n",
    "\n",
    "        # Fetch all results\n",
    "        if cur.description:\n",
    "            columns = [desc[0] for desc in cur.description]\n",
    "            result = cur.fetchall()\n",
    "        else:\n",
    "            columns = []\n",
    "            result = []\n",
    "\n",
    "        cur.close()\n",
    "        conn.close()\n",
    "        return columns, result\n",
    "\n",
    "    except Exception as e:\n",
    "        raise Exception(f\"Database query failed: {str(e)}\")\n",
    "\n",
    "def lambda_handler(event, context):\n",
    "    try:\n",
    "        data = json.loads(event['body'])\n",
    "        sql_statement = data['sql_statement']\n",
    "\n",
    "        # Execute the statement and fetch results\n",
    "        columns, result = execute_statement(sql_statement)\n",
    "        \n",
    "        # Create a temporary file to save the result as CSV\n",
    "        with tempfile.NamedTemporaryFile(delete=False, mode='w', suffix='.csv', newline='') as tmp_file:\n",
    "            csv_writer = csv.writer(tmp_file)\n",
    "            if columns:\n",
    "                csv_writer.writerow(columns)  # Write the header\n",
    "            csv_writer.writerows(result)  # Write all rows\n",
    "            tmp_file_path = tmp_file.name\n",
    "\n",
    "        # Read the file and encode its content to base64\n",
    "        with open(tmp_file_path, 'rb') as f:\n",
    "            file_content = f.read()\n",
    "            encoded_content = base64.b64encode(file_content).decode('utf-8')\n",
    "\n",
    "        response = {\n",
    "            'openaiFileResponse': [\n",
    "                {\n",
    "                    'name': 'query_result.csv',\n",
    "                    'mime_type': 'text/csv',\n",
    "                    'content': encoded_content\n",
    "                }\n",
    "            ]\n",
    "        }\n",
    "\n",
    "        return {\n",
    "            'statusCode': 200,\n",
    "            'headers': {\n",
    "                'Content-Type': 'application/json'\n",
    "            },\n",
    "            'body': json.dumps(response)\n",
    "        }\n",
    "\n",
    "    except Exception as e:\n",
    "        return {\n",
    "            'statusCode': 500,\n",
    "            'body': json.dumps({'error': str(e)})\n",
    "        }\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Retrieve VPC information\n",
    "\n",
    "We will need to connnect our function to Redshift, therefore we need to find the network used by Redshift. You can find this on your Redshift interface the AWS console, under Amazon Redshift Serverless > Workgroup configuration > `your_workgroup` > Data access, or through the CLI:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "aws redshift-serverless get-workgroup --workgroup-name default-workgroup --query 'workgroup.{address: endpoint.address, port: endpoint.port, SecurityGroupIds: securityGroupIds, SubnetIds: subnetIds}'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Set up AWS function\n",
    "\n",
    "Copy `env.sample.yaml` to `env.yaml` and replace with the values obtained above. You will need a Redshift user with access to your DB/schema.\n",
    "\n",
    "```\n",
    "cp env.sample.yaml env.yaml\n",
    "```\n",
    "\n",
    "Fill in `env.yaml` with the values retrieved by the previous command as well as your credentials to Redshift.\n",
    "Alternatively, you can create a file named `env.yaml` manually and fill the following variables:\n",
    "```\n",
    "RedshiftHost: default-workgroup.xxxxx.{region}.redshift-serverless.amazonaws.com\n",
    "RedshiftPort: 5439\n",
    "RedshiftUser: username\n",
    "RedshiftPassword: password\n",
    "RedshiftDb: my-db\n",
    "SecurityGroupId: sg-xx\n",
    "SubnetId1: subnet-xx\n",
    "SubnetId2: subnet-xx\n",
    "SubnetId3: subnet-xx\n",
    "SubnetId4: subnet-xx\n",
    "SubnetId5: subnet-xx\n",
    "SubnetId6: subnet-xx\n",
    "```\n",
    "\n",
    "This file will be used to deploy your function with parameters, as shown below:\n",
    "\n",
    "```\n",
    "PARAM_FILE=\"env.yaml\"\n",
    "PARAMS=$(yq eval -o=json $PARAM_FILE | jq -r 'to_entries | map(\"\\(.key)=\\(.value|tostring)\") | join(\" \")')\n",
    "sam deploy --template-file template.yaml --stack-name redshift-middleware --capabilities CAPABILITY_IAM --parameter-overrides $PARAMS\n",
    "```\n",
    "\n",
    "The template.yaml has the following content:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "AWSTemplateFormatVersion: '2010-09-09'\n",
    "Transform: AWS::Serverless-2016-10-31\n",
    "Description: >\n",
    "  redshift-middleware\n",
    "\n",
    "  Middleware to fetch RedShift data and return it through HTTP as files\n",
    "\n",
    "Globals:\n",
    "  Function:\n",
    "    Timeout: 3\n",
    "\n",
    "Parameters:\n",
    "  RedshiftHost:\n",
    "    Type: String\n",
    "  RedshiftPort:\n",
    "    Type: String\n",
    "  RedshiftUser:\n",
    "    Type: String\n",
    "  RedshiftPassword:\n",
    "    Type: String\n",
    "  RedshiftDb:\n",
    "    Type: String\n",
    "  SecurityGroupId:\n",
    "    Type: String\n",
    "  SubnetId1:\n",
    "    Type: String\n",
    "  SubnetId2:\n",
    "    Type: String\n",
    "  SubnetId3:\n",
    "    Type: String\n",
    "  SubnetId4:\n",
    "    Type: String\n",
    "  SubnetId5:\n",
    "    Type: String\n",
    "  SubnetId6:\n",
    "    Type: String\n",
    "  CognitoUserPoolName:\n",
    "    Type: String\n",
    "    Default: MyCognitoUserPool\n",
    "  CognitoUserPoolClientName:\n",
    "    Type: String\n",
    "    Default: MyCognitoUserPoolClient\n",
    "\n",
    "Resources:\n",
    "  MyCognitoUserPool:\n",
    "    Type: AWS::Cognito::UserPool\n",
    "    Properties:\n",
    "      UserPoolName: !Ref CognitoUserPoolName\n",
    "      Policies:\n",
    "        PasswordPolicy:\n",
    "          MinimumLength: 8\n",
    "      UsernameAttributes:\n",
    "        - email\n",
    "      Schema:\n",
    "        - AttributeDataType: String\n",
    "          Name: email\n",
    "          Required: false\n",
    "\n",
    "  MyCognitoUserPoolClient:\n",
    "    Type: AWS::Cognito::UserPoolClient\n",
    "    Properties:\n",
    "      UserPoolId: !Ref MyCognitoUserPool\n",
    "      ClientName: !Ref CognitoUserPoolClientName\n",
    "      GenerateSecret: true\n",
    "\n",
    "  RedshiftMiddlewareApi:\n",
    "    Type: AWS::Serverless::Api\n",
    "    Properties:\n",
    "      StageName: Prod\n",
    "      Cors: \"'*'\"\n",
    "      Auth:\n",
    "        DefaultAuthorizer: MyCognitoAuthorizer\n",
    "        Authorizers:\n",
    "          MyCognitoAuthorizer:\n",
    "            AuthorizationScopes:\n",
    "              - openid\n",
    "              - email\n",
    "              - profile\n",
    "            UserPoolArn: !GetAtt MyCognitoUserPool.Arn\n",
    "        \n",
    "  RedshiftMiddlewareFunction:\n",
    "    Type: AWS::Serverless::Function\n",
    "    Properties:\n",
    "      CodeUri: redshift-middleware/\n",
    "      Handler: app.lambda_handler\n",
    "      Runtime: python3.11\n",
    "      Timeout: 45\n",
    "      Architectures:\n",
    "        - x86_64\n",
    "      Events:\n",
    "        SqlStatement:\n",
    "          Type: Api\n",
    "          Properties:\n",
    "            Path: /sql_statement\n",
    "            Method: post\n",
    "            RestApiId: !Ref RedshiftMiddlewareApi\n",
    "      Environment:\n",
    "        Variables:\n",
    "          REDSHIFT_HOST: !Ref RedshiftHost\n",
    "          REDSHIFT_PORT: !Ref RedshiftPort\n",
    "          REDSHIFT_USER: !Ref RedshiftUser\n",
    "          REDSHIFT_PASSWORD: !Ref RedshiftPassword\n",
    "          REDSHIFT_DB: !Ref RedshiftDb\n",
    "      VpcConfig:\n",
    "        SecurityGroupIds:\n",
    "          - !Ref SecurityGroupId\n",
    "        SubnetIds:\n",
    "          - !Ref SubnetId1\n",
    "          - !Ref SubnetId2\n",
    "          - !Ref SubnetId3\n",
    "          - !Ref SubnetId4\n",
    "          - !Ref SubnetId5\n",
    "          - !Ref SubnetId6\n",
    "\n",
    "Outputs:\n",
    "  RedshiftMiddlewareApi:\n",
    "    Description: \"API Gateway endpoint URL for Prod stage for SQL Statement function\"\n",
    "    Value: !Sub \"https://${RedshiftMiddlewareApi}.execute-api.${AWS::Region}.amazonaws.com/Prod/sql_statement/\"\n",
    "  RedshiftMiddlewareFunction:\n",
    "    Description: \"SQL Statement Lambda Function ARN\"\n",
    "    Value: !GetAtt RedshiftMiddlewareFunction.Arn\n",
    "  RedshiftMiddlewareFunctionIamRole:\n",
    "    Description: \"Implicit IAM Role created for SQL Statement function\"\n",
    "    Value: !GetAtt RedshiftMiddlewareFunctionRole.Arn\n",
    "  CognitoUserPoolArn:\n",
    "    Description: \"ARN of the Cognito User Pool\"\n",
    "    Value: !GetAtt MyCognitoUserPool.Arn\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Retrieve the URL information from the previous command output, you can then run a cURL request, which should return data in a file format:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "curl -X POST https://<your_url>/Prod/sql_statement/ \\\n",
    "-H \"Content-Type: application/json\" \\\n",
    "-d '{ \"sql_statement\": \"SELECT * FROM customers LIMIT 10\", \"workgroup_name\": \"default-workgroup\", \"database_name\": \"pap-db\" }'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## ChatGPT Steps"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Custom GPT Instructions "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once you've created a Custom GPT, copy the text below in the Instructions panel."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "plaintext"
    }
   },
   "outputs": [],
   "source": [
    "**Context**: You are an expert at writing Redshift SQL queries. You will initially retrieve the table schema that you will use thoroughly. Every attributes, table names or data type will be known by you.\n",
    "\n",
    "**Instructions**:\n",
    "1. No matter the user's question, start by running `runQuery` operation using this query: \"SELECT table_name, column_name FROM INFORMATION_SCHEMA.COLUMNS WHERE table_schema = 'public' ORDER BY table_name, ordinal_position;\"  It will help you understand how to query the data. A CSV will be returned with all the attributes and their table. Make sure to read it fully and understand all available tables & their attributes before querying. You don't have to show this to the user.\n",
    "2. Convert the user's question into a SQL statement that leverages the step above and run the `runQuery` operation on that SQL statement to confirm the query works. Let the user know which table you will use/query.\n",
    "3. Execute the query and show him the data. Show only the first few rows.\n",
    "\n",
    "**Additional Notes**: If the user says \"Let's get started\", explain they can ask a question they want answered about data that we have access to. If the user has no ideas, suggest that we have transactions data they can query - ask if they want you to query that.\n",
    "**Important**: Never make up a table name or table attribute. If you don't know, go back to the data you've retrieved to check what is available. If you think no table or attribute is available, then tell the user you can't perform this query for them."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### OpenAPI Schema "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once you've created a Custom GPT, copy the text below in the Actions panel.\n",
    "\n",
    "This expects a response that matches the file retrieval structure in our doc [here](https://platform.openai.com/docs/actions/sending-files) and passes in a `query` as a parameter to execute.\n",
    "\n",
    "Make sure to follow the steps in the [AWS Middleware cookbook](https://cookbook.openai.com/examples/chatgpt/gpt_actions_library/gpt_middleware_aws_function) to set up authentication.\n",
    "\n",
    "> Make sure to switch the function app name based on your function deployment."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "yaml"
    }
   },
   "outputs": [],
   "source": [
    "openapi: 3.1.0\n",
    "info:\n",
    "  title: SQL Execution API\n",
    "  description: API to execute SQL statements and return results as a file.\n",
    "  version: 1.0.0\n",
    "servers:\n",
    "  - url: {your_function_url}/Prod\n",
    "    description: Production server\n",
    "paths:\n",
    "  /sql_statement:\n",
    "    post:\n",
    "      operationId: executeSqlStatement\n",
    "      summary: Executes a SQL statement and returns the result as a file.\n",
    "      requestBody:\n",
    "        required: true\n",
    "        content:\n",
    "          application/json:\n",
    "            schema:\n",
    "              type: object\n",
    "              properties:\n",
    "                sql_statement:\n",
    "                  type: string\n",
    "                  description: The SQL statement to execute.\n",
    "                  example: SELECT * FROM customers LIMIT 10\n",
    "              required:\n",
    "                - sql_statement\n",
    "      responses:\n",
    "        '200':\n",
    "          description: The SQL query result as a JSON file.\n",
    "          content:\n",
    "            application/json:\n",
    "              schema:\n",
    "                type: object\n",
    "                properties:\n",
    "                  openaiFileResponse:\n",
    "                    type: array\n",
    "                    items:\n",
    "                      type: object\n",
    "                      properties:\n",
    "                        name:\n",
    "                          type: string\n",
    "                          description: The name of the file.\n",
    "                          example: query_result.json\n",
    "                        mime_type:\n",
    "                          type: string\n",
    "                          description: The MIME type of the file.\n",
    "                          example: application/json\n",
    "                        content:\n",
    "                          type: string\n",
    "                          description: The base64 encoded content of the file.\n",
    "                          format: byte\n",
    "                          example: eyJrZXkiOiJ2YWx1ZSJ9\n",
    "        '500':\n",
    "          description: Error response\n",
    "          content:\n",
    "            application/json:\n",
    "              schema:\n",
    "                type: object\n",
    "                properties:\n",
    "                  error:\n",
    "                    type: string\n",
    "                    description: Error message.\n",
    "                    example: Database query failed error details\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Conclusion\n",
    "\n",
    "You now have deployed a GPT that uses a middleware in AWS, in an authenticated manner, that's able to connect to Redsfhit. Users with access (that are in Cognito) can now query your databases to perform data analysis task:\n",
    "\n",
    "![../../../images/redshift_gpt.png](../../../images/redshift_gpt.png)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.11.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}