AWSTemplateFormatVersion: 2010-09-09
Description: QnABot on AWS LLM Plugin for AI21 - v0.1.0

Parameters:
  APIKey:
    Type: String
    Description: AI21 API Key (stored in Secrets Manager - see stack Outputs)
    Default: ''
    NoEcho: true

  LLMModelType:
    Type: String
    Default: j2-mid
    AllowedValues:
      - j2-light
      - j2-mid
      - j2-ultra
    Description: AI21 LLM Model Type

Resources:
  ApiKeySecret:
    Type: AWS::SecretsManager::Secret
    Properties: 
      Description: API Key
      Name: !Ref AWS::StackName
      SecretString: !Ref APIKey

  LambdaFunctionRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal:
              Service: lambda.amazonaws.com
            Action: sts:AssumeRole
      ManagedPolicyArns:
        - arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole
      Policies:
        - PolicyDocument:
            Version: 2012-10-17
            Statement:
              - Effect: Allow
                Action:
                  - 'secretsmanager:GetResourcePolicy'
                  - 'secretsmanager:GetSecretValue'
                  - 'secretsmanager:DescribeSecret'
                  - 'secretsmanager:ListSecretVersionIds'
                Resource: !Ref ApiKeySecret
          PolicyName: SecretsManagerPolicy

  LambdaFunction:
    Type: AWS::Lambda::Function
    Properties:
      Handler: "llm.lambda_handler"
      Role: !GetAtt 'LambdaFunctionRole.Arn'
      MemorySize: 128
      Timeout: 60
      Runtime: python3.10       
      Environment:
        Variables:
          API_KEY_SECRET_NAME: !Ref AWS::StackName
      Code: ./src

  OutputSettingsFunctionRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal:
              Service: lambda.amazonaws.com
            Action: sts:AssumeRole
      ManagedPolicyArns:
        - arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole
  
  OutputSettingsFunction:
    Type: AWS::Lambda::Function
    Properties:
      Handler: settings.lambda_handler
      Role: !GetAtt 'OutputSettingsFunctionRole.Arn'
      Runtime: python3.10
      Timeout: 10
      MemorySize: 128
      Code: ./src

  OutputSettings:
    Type: Custom::OutputSettings
    Properties:
      ServiceToken: !GetAtt OutputSettingsFunction.Arn
      ModelType: !Ref LLMModelType

Outputs:
  APIKeySecret:
    Description: Link to Secrets Manager console to input API Key
    Value: !Sub "https://${AWS::Region}.console.aws.amazon.com/secretsmanager/secret?region=${AWS::Region}&name=${AWS::StackName}"
    
  LLMLambdaArn:
    Description: Lambda function ARN (use for QnABot param "LLMLambdaArn")
    Value: !GetAtt LambdaFunction.Arn
  
  QnABotSettingGenerateQueryPromptTemplate:
    Description: QnABot Designer Setting "LLM_GENERATE_QUERY_PROMPT_TEMPLATE"
    Value: !GetAtt OutputSettings.LLM_GENERATE_QUERY_PROMPT_TEMPLATE

  QnABotSettingGenerateQueryModelParams:
    Description: QnABot Designer Setting "LLM_GENERATE_QUERY_MODEL_PARAMS"
    Value: !GetAtt OutputSettings.LLM_GENERATE_QUERY_MODEL_PARAMS

  QnABotSettingQAPromptTemplate:
    Description: QnABot Designer Setting "LLM_QA_PROMPT_TEMPLATE"
    Value: !GetAtt OutputSettings.LLM_QA_PROMPT_TEMPLATE

  QnABotSettingQAModelParams:
    Description: QnABot Designer Setting "LLM_QA_MODEL_PARAMS"
    Value: !GetAtt OutputSettings.LLM_QA_MODEL_PARAMS