template.yml

AWSTemplateFormatVersion: "2010-09-09"
Transform: AWS::Serverless-2016-10-31
Description: >
  Sample SAM Template for bedrockbenchmark

Resources:
  BedrockBenchmarkStateMachine:
    Type: AWS::Serverless::StateMachine
    Properties:
      DefinitionUri: statemachine/bedrock_benchmark.asl.json
      DefinitionSubstitutions:
        AnthropicFunctionArn: !GetAtt AnthropicFunction.Arn
        AmazonFunctionArn: !GetAtt AmazonFunction.Arn
        CohereFunctionArn: !GetAtt CohereFunction.Arn
        Ai21FunctionArn: !GetAtt Ai21Function.Arn
        MetaFunctionArn: !GetAtt MetaFunction.Arn
        MistralFunctionArn: !GetAtt MistralFunction.Arn
      Events:
        DailyBenchmarkingSchedule:
          Type: Schedule
          Properties:
            Description: Schedule to run the bedrock benchmarking state machine every day
            Enabled: True 
            Schedule: "rate(1 day)"
      Policies:
        - LambdaInvokePolicy:
            FunctionName: !Ref AnthropicFunction
        - LambdaInvokePolicy:
            FunctionName: !Ref CohereFunction
        - LambdaInvokePolicy:
            FunctionName: !Ref Ai21Function
        - LambdaInvokePolicy:
            FunctionName: !Ref AmazonFunction
        - LambdaInvokePolicy:
            FunctionName: !Ref MetaFunction
        - LambdaInvokePolicy:
            FunctionName: !Ref MistralFunction

  AnthropicFunction:
    Type: AWS::Serverless::Function
    Properties:
      CodeUri: functions/main/
      Handler: app.lambda_handler
      Runtime: python3.12
      Architectures:
        - x86_64
      Timeout: 900
      MemorySize: 150
      Policies: arn:aws:iam::aws:policy/AdministratorAccess
      Environment:
        Variables:
          model_shape: '{"anthropic_version": "bedrock-2023-05-31", "max_tokens": 4000, "messages": [{"role": "user", "content": ""}], "temperature": 1.0,"stop_sequences": ["."], "top_p": 1.0, "top_k": 250}'
          supported_models: '["anthropic.claude-v2:1", "anthropic.claude-3-sonnet-20240229-v1:0", "anthropic.claude-3-haiku-20240307-v1:0", "anthropic.claude-3-opus-20240229-v1:0"]'
          benchmark_table: !Ref BedrockBenchmarkTable
          prompt_catalog: !Ref BedrockBenchmarkPromptsTable

  CohereFunction:
    Type: AWS::Serverless::Function
    Properties:
      CodeUri: functions/main/
      Handler: app.lambda_handler
      Runtime: python3.12
      Architectures:
        - x86_64
      Timeout: 900
      MemorySize: 150
      Policies: arn:aws:iam::aws:policy/AdministratorAccess
      Environment:
        Variables:
          model_shape: '{"message": "", "temperature": 0.75, "p": 0.01, "k": 0, "stop_sequences": [], "max_tokens": 2048}'
          supported_models: '["cohere.command-r-v1:0", "cohere.command-r-plus-v1:0"]'
          benchmark_table: !Ref BedrockBenchmarkTable
          prompt_catalog: !Ref BedrockBenchmarkPromptsTable

  Ai21Function:
    Type: AWS::Serverless::Function
    Properties:
      CodeUri: functions/main/
      Handler: app.lambda_handler
      Runtime: python3.12
      Architectures:
        - x86_64
      Timeout: 900
      MemorySize: 150 
      Policies: arn:aws:iam::aws:policy/AdministratorAccess
      Environment:
        Variables:
          model_shape:  '{"prompt": "", "temperature": 1.0, "maxTokens": 2048}'
          supported_models: '["ai21.j2-mid", "ai21.j2-ultra"]'
          benchmark_table: !Ref BedrockBenchmarkTable
          prompt_catalog: !Ref BedrockBenchmarkPromptsTable

  AmazonFunction:
    Type: AWS::Serverless::Function
    Properties:
      CodeUri: functions/main/
      Handler: app.lambda_handler
      Runtime: python3.12
      Architectures:
        - x86_64
      Timeout: 900
      MemorySize: 150
      Policies: arn:aws:iam::aws:policy/AdministratorAccess
      Environment:
        Variables:
          model_shape: '{"inputText": "", "textGenerationConfig": {"temperature": 1.0, "topP": 1.0, "maxTokenCount": 4096, "stopSequences": []}}'
          supported_models: '["amazon.titan-tg1-large", "amazon.titan-text-express-v1"]'
          benchmark_table: !Ref BedrockBenchmarkTable
          prompt_catalog: !Ref BedrockBenchmarkPromptsTable

  MetaFunction:
    Type: AWS::Serverless::Function
    Properties:
      CodeUri: functions/main/
      Handler: app.lambda_handler
      Runtime: python3.12
      Architectures:
        - x86_64
      Timeout: 900
      MemorySize: 150
      Policies: arn:aws:iam::aws:policy/AdministratorAccess
      Environment:
        Variables:
          model_shape: '{"prompt": "", "max_gen_len":512, "temperature":0.5, "top_p":0.9}'
          supported_models: '["meta.llama3-8b-instruct-v1:0", "meta.llama3-70b-instruct-v1:0"]'
          benchmark_table: !Ref BedrockBenchmarkTable
          prompt_catalog: !Ref BedrockBenchmarkPromptsTable
  
  MistralFunction:
    Type: AWS::Serverless::Function
    Properties:
      CodeUri: functions/main/
      Handler: app.lambda_handler
      Runtime: python3.12
      Architectures:
        - x86_64
      Timeout: 900
      MemorySize: 250
      Policies: arn:aws:iam::aws:policy/AdministratorAccess
      Environment:
        Variables:
          model_shape: '{"prompt": "", "max_tokens":4000, "temperature":0.5, "top_p":0.9, "top_k":50 }'
          supported_models: '["mistral.mistral-7b-instruct-v0:2", "mistral.mixtral-8x7b-instruct-v0:1"]'
          benchmark_table: !Ref BedrockBenchmarkTable
          prompt_catalog: !Ref BedrockBenchmarkPromptsTable

  BedrockBenchmarkTable:
    Type: 'AWS::DynamoDB::Table'
    Properties:
      TableName: bedrockbenchmark
      AttributeDefinitions:
        - AttributeName: model_prompt_id
          AttributeType: S
        - AttributeName: date
          AttributeType: S
        - AttributeName: prompt_model_id
          AttributeType: S
      KeySchema:
        - AttributeName: model_prompt_id
          KeyType: HASH
        - AttributeName: date
          KeyType: RANGE
      BillingMode: PAY_PER_REQUEST
      GlobalSecondaryIndexes:
        - IndexName: prompt_model_id_gsi
          KeySchema:
            - AttributeName: prompt_model_id
              KeyType: HASH
          Projection:
            ProjectionType: ALL

  BedrockBenchmarkPromptsTable:
    Type: 'AWS::DynamoDB::Table'
    Properties:
      TableName: bedrockbenchmarkprompts
      AttributeDefinitions:
        - AttributeName: id
          AttributeType: S
        - AttributeName: prompt
          AttributeType: S
      KeySchema:
        - AttributeName: id
          KeyType: HASH
        - AttributeName: prompt
          KeyType: RANGE
      BillingMode: PAY_PER_REQUEST

  DynamodbUpsertFunction:
    Type: AWS::Serverless::Function
    Properties:
      CodeUri: functions/dynamodb/
      Handler: app.lambda_handler
      Policies: arn:aws:iam::aws:policy/AdministratorAccess
      Runtime: 'python3.12'
      Timeout: 60
      Architectures:
        - x86_64
  
  DdbInsertCustomFunction:
    Type: Custom::DdbInsertCustomFunction
    Properties:
      ServiceToken: !GetAtt DynamodbUpsertFunction.Arn
      TableName: !Ref BedrockBenchmarkPromptsTable

Outputs:
  benchmarktable:
    Value: !Ref BedrockBenchmarkTable
  promptcatalogtable:
    Value: !Ref BedrockBenchmarkPromptsTable