AWSTemplateFormatVersion: 2010-09-09
Description: |
  Primary QnABot ChatGPT Plugin CloudFormation template (v1.0.0)
  This template deploys resources to allow QnABot to use OpenAI ChatGPT to answer questions. Presently,
  the implementation supports a configurable CHATGPT_MODEL which defaults to "gpt-3.5-turbo"
  It deploys:
      - A Lambda Layer
      - A Lambda function
      - A DynamoDB table
      - IAM Role / Policy used by the Lambda function

Parameters:
  ChatGptModel:
    Description: >
      The ChatGpt Model used for API requests
    Type: String
    Default: "gpt-3.5-turbo"

  EmbeddingModel:
    Description: >
      The OpenAI Model to use for embeddings
    Type: String
    Default: "text-embedding-ada-002"

  MessageCacheExpirationInHours:
    Description: >
      Time in hours that user message history should be preserved
    Type: Number
    Default: 8

Resources:
  ChatGPTRouterRole:
    Type: AWS::IAM::Role
    Properties:
      Path: /
      AssumeRolePolicyDocument:
        Version: 2012-10-17
        Statement:
          - Principal:
              Service:
                - lambda.amazonaws.com
            Effect: Allow
            Action:
              - sts:AssumeRole
      Policies:
        - PolicyName: CloudWatchLogs
          PolicyDocument:
            Version: 2012-10-17
            Statement:
              - Effect: Allow
                Action:
                  - logs:CreateLogGroup
                  - logs:CreateLogStream
                  - logs:PutLogEvents
                Resource: "*"
        - PolicyName: DynamoDB
          PolicyDocument:
            Version: 2012-10-17
            Statement:
              - Effect: Allow
                Action:
                  - dynamodb:PutItem
                  - dynamodb:Query
                Resource:
                  - !Sub "arn:aws:dynamodb:${AWS::Region}:${AWS::AccountId}:table/QNA-chatgpt-user-cache"
        - PolicyName: SecretsManager
          PolicyDocument:
            Version: 2012-10-17
            Statement:
              - Effect: Allow
                Action:
                  - secretsmanager:GetSecretValue
                Resource:
                  - !Ref OpenAiApiKeySecret

  OpenAIEmbeddingRole:
    Type: AWS::IAM::Role
    Properties:
      Path: /
      AssumeRolePolicyDocument:
        Version: 2012-10-17
        Statement:
          - Principal:
              Service:
                - lambda.amazonaws.com
            Effect: Allow
            Action:
              - sts:AssumeRole
      Policies:
        - PolicyName: CloudWatchLogs
          PolicyDocument:
            Version: 2012-10-17
            Statement:
              - Effect: Allow
                Action:
                  - logs:CreateLogGroup
                  - logs:CreateLogStream
                  - logs:PutLogEvents
                Resource: "*"
        - PolicyName: SecretsManager
          PolicyDocument:
            Version: 2012-10-17
            Statement:
              - Effect: Allow
                Action:
                  - secretsmanager:GetSecretValue
                Resource:
                  - !Ref OpenAiApiKeySecret

  DynamoDBUserMessageCacheTable:
    Type: "AWS::DynamoDB::Table"
    Properties:
      AttributeDefinitions:
        - AttributeName: "messageTimeout"
          AttributeType: "N"
        - AttributeName: "userid"
          AttributeType: "S"
      BillingMode: "PAY_PER_REQUEST"
      TableName: "QNA-chatgpt-user-cache"
      KeySchema:
        - AttributeName: "userid"
          KeyType: "HASH"
        - AttributeName: "messageTimeout"
          KeyType: "RANGE"
      TimeToLiveSpecification:
        AttributeName: "messageTimeout"
        Enabled: true

  OpenAiApiKeySecret:
    Type: AWS::SecretsManager::Secret
    Properties:
      Description: Resource that holds the openai api key to use for calling openai
      Name: "QNA-chatgpt-openai-api-key"

  LambdaLayerVersion:
    Type: "AWS::Lambda::LayerVersion"
    Properties:
      Description: ""
      CompatibleRuntimes:
        - "python3.9"
      LayerName: "QNA-chatgpt-qnabot-plugin-layer"
      Content:
        S3Bucket: !Sub "tioth-chatgpt-plugin-resources-${AWS::Region}"
        S3Key: "plugin-layer.zip"

  LambdaFunction:
    Type: "AWS::Lambda::Function"
    Properties:
      Description: ""
      Environment:
        Variables:
          DYNAMODB_USER_MESSAGE_CACHE: !Ref DynamoDBUserMessageCacheTable
          OPENAI_API_KEY_SECRET_ID: "QNA-chatgpt-openai-api-key"
          MESSAGE_CACHE_EXPIRATION_IN_HOURS: !Ref MessageCacheExpirationInHours
          CHATGPT_MODEL: !Ref ChatGptModel
      FunctionName: "QNA-chatgpt-qnabot-router"
      Handler: "chatgpt-router.lambda_handler"
      Architectures:
        - "x86_64"
      Code:
        S3Bucket: !Sub "tioth-chatgpt-plugin-resources-${AWS::Region}"
        S3Key: "plugin.zip"
      MemorySize: 256
      Role: !GetAtt ChatGPTRouterRole.Arn
      Runtime: "python3.9"
      Timeout: 120
      TracingConfig:
        Mode: "PassThrough"
      Layers:
        - !Ref LambdaLayerVersion
      EphemeralStorage:
        Size: 512

  LambdaFunctionEmbedding:
    Type: "AWS::Lambda::Function"
    Properties:
      Description: ""
      Environment:
        Variables:
          OPENAI_API_KEY_SECRET_ID: "QNA-chatgpt-openai-api-key"
          EMBEDDING_MODEL: !Ref EmbeddingModel
      FunctionName: "QNA-chatgpt-qnabot-embeddings"
      Handler: "chatgpt-embeddings.lambda_handler"
      Architectures:
        - "x86_64"
      Code:
        S3Bucket: !Sub "tioth-chatgpt-plugin-resources-${AWS::Region}"
        S3Key: "plugin.zip"
      MemorySize: 256
      Role: !GetAtt OpenAIEmbeddingRole.Arn
      Runtime: "python3.9"
      Timeout: 120
      TracingConfig:
        Mode: "PassThrough"
      Layers:
        - !Ref LambdaLayerVersion
      EphemeralStorage:
        Size: 512

Outputs:
  BotRouterLambdaFunction:
    Description: Lambda BotRouter Function
    Value: !Ref LambdaFunction

  BotRouterLambdaFunctionArn:
    Description: Lambda BotRouter Function Arn
    Value: !GetAtt LambdaFunction.Arn

  EmbeddingsLambdaFunction:
    Description: Lambda Embeddings Function
    Value: !Ref LambdaFunctionEmbedding

  EmbeddingsLambdaFunctionArn:
    Description: Lambda Embeddings Function Arn
    Value: !GetAtt LambdaFunctionEmbedding.Arn