feat: Add log retention setting, use JSON logger and add logs to the …

…dashboard (#552)

* feat: Added logs to dashboard + format them in JSON

* chore: Use the logger instead of print

cli/magic-config.ts

@@ -202,6 +202,7 @@ const embeddingModels = [
       // Advanced settings
 
       options.createVpcEndpoints = config.vpc?.createVpcEndpoints;
+      options.logRetention = config.logRetention;
       options.privateWebsite = config.privateWebsite;
       options.certificate = config.certificate;
       options.domain = config.domain;
@@ -808,6 +809,24 @@ async function processCreateOptions(options: any): Promise<void> {
   const models: any = await enquirer.prompt(modelsPrompts);
 
   const advancedSettingsPrompts = [
+    {
+      type: "input",
+      name: "logRetention",
+      message: "For how long do you want to store the logs (in days)?",
+      initial: options.logRetention ? String(options.logRetention) : "7",
+      validate(value: string) {
+        // https://docs.aws.amazon.com/AWSCloudFormation/latest/UserGuide/aws-resource-logs-loggroup.html#cfn-logs-loggroup-retentionindays
+        const allowed = [
+          1, 3, 5, 7, 14, 30, 60, 90, 120, 150, 180, 365, 400, 545, 731, 1096,
+          1827, 2192, 2557, 2922, 3288, 3653,
+        ];
+        if (allowed.includes(Number(value))) {
+          return true;
+        } else {
+          return "Allowed values are: " + allowed.join(", ");
+        }
+      },
+    },
     {
       type: "confirm",
       name: "createVpcEndpoints",
@@ -1087,6 +1106,9 @@ async function processCreateOptions(options: any): Promise<void> {
         }
       : undefined,
     privateWebsite: advancedSettings.privateWebsite,
+    logRetention: advancedSettings.logRetention
+      ? Number(advancedSettings.logRetention)
+      : undefined,
     certificate: advancedSettings.certificate,
     domain: advancedSettings.domain,
     cognitoFederation: advancedSettings.cognitoFederationEnabled

integtests/conftest.py

@@ -65,11 +65,9 @@ def react_url():
 @pytest.fixture(scope="class")
 def selenium_driver(react_url):
     options = webdriver.FirefoxOptions()
-    if os.environ["HEADLESS"]:
+    if os.getenv("HEADLESS"):
         options.add_argument("--headless")
-    driver = webdriver.Remote(
-        command_executor="http://127.0.0.1:4444", options=options
-    )
+    driver = webdriver.Remote(command_executor="http://127.0.0.1:4444", options=options)
     driver.set_window_size(1600, 800)
     driver.get(react_url)
     yield driver

lib/authentication/index.ts

@@ -111,8 +111,10 @@ export class Authentication extends Construct {
           code: lambda.Code.fromAsset(
             "lib/authentication/lambda/updateUserPoolClient"
           ),
+          description: "Updates the user pool client",
           role: lambdaRoleUpdateClient,
-          logRetention: logs.RetentionDays.ONE_WEEK,
+          logRetention: config.logRetention ?? logs.RetentionDays.ONE_WEEK,
+          loggingFormat: lambda.LoggingFormat.JSON,
           environment: {
             USER_POOL_ID: userPool.userPoolId,
             USER_POOL_CLIENT_ID: userPoolClient.userPoolClientId,
@@ -136,7 +138,7 @@ export class Authentication extends Construct {
         );
         this.customOidcProvider = customProvider;
 
-        // Unfortunately the above function does not support SecretValue for Client Secrets so updating with lamda
+        // Unfortunately the above function does not support SecretValue for Client Secrets so updating with lambda
         // Create an IAM Role for the Lambda function
         const lambdaRoleUpdateOidcSecret = new iam.Role(
           this,
@@ -187,8 +189,10 @@ export class Authentication extends Construct {
             code: lambda.Code.fromAsset(
               "lib/authentication/lambda/updateOidcSecret"
             ),
+            description: "Updates OIDC secret",
             role: lambdaRoleUpdateOidcSecret,
-            logRetention: logs.RetentionDays.ONE_WEEK,
+            logRetention: config.logRetention ?? logs.RetentionDays.ONE_WEEK,
+            loggingFormat: lambda.LoggingFormat.JSON,
             environment: {
               USER_POOL_ID: userPool.userPoolId,
               USER_POOL_CLIENT_ID: userPoolClient.userPoolClientId,

lib/aws-genai-llm-chatbot-stack.ts

@@ -15,6 +15,7 @@ import * as sns from "aws-cdk-lib/aws-sns";
 import * as iam from "aws-cdk-lib/aws-iam";
 import * as cr from "aws-cdk-lib/custom-resources";
 import { NagSuppressions } from "cdk-nag";
+import { LogGroup } from "aws-cdk-lib/aws-logs";
 
 export interface AwsGenAILLMChatbotStackProps extends cdk.StackProps {
   readonly config: SystemConfig;
@@ -220,6 +221,13 @@ export class AwsGenAILLMChatbotStack extends cdk.Stack {
     new Monitoring(monitoringStack, "Monitoring", {
       prefix: props.config.prefix,
       appsycnApi: chatBotApi.graphqlApi,
+      appsyncResolversLogGroups: chatBotApi.resolvers.map((r) => {
+        return LogGroup.fromLogGroupName(
+          monitoringStack,
+          "Log" + r.node.id,
+          "/aws/lambda/" + r.functionName
+        );
+      }),
       cognito: {
         userPoolId: authentication.userPool.userPoolId,
         clientId: authentication.userPoolClient.userPoolClientId,

lib/chatbot-api/appsync-ws.ts

@@ -4,6 +4,7 @@ import {
   Code,
   Function as LambdaFunction,
   LayerVersion,
+  LoggingFormat,
   Runtime,
 } from "aws-cdk-lib/aws-lambda";
 import { SqsEventSource } from "aws-cdk-lib/aws-lambda-event-sources";
@@ -21,9 +22,11 @@ interface RealtimeResolversProps {
   readonly userPool: UserPool;
   readonly shared: Shared;
   readonly api: appsync.GraphqlApi;
+  readonly logRetention?: number;
 }
 
 export class RealtimeResolvers extends Construct {
+  public readonly sendQueryHandler: LambdaFunction;
   public readonly outgoingMessageHandler: LambdaFunction;
 
   constructor(scope: Construct, id: string, props: RealtimeResolversProps) {
@@ -42,10 +45,13 @@ export class RealtimeResolvers extends Construct {
         "./lib/chatbot-api/functions/resolvers/send-query-lambda-resolver"
       ),
       handler: "index.handler",
+      description: "Appsync resolver handling LLM Queries",
       runtime: Runtime.PYTHON_3_11,
       environment: {
         SNS_TOPIC_ARN: props.topic.topicArn,
       },
+      logRetention: props.logRetention,
+      loggingFormat: LoggingFormat.JSON,
       layers: [props.shared.powerToolsLayer],
       vpc: props.shared.vpc,
     });
@@ -61,6 +67,7 @@ export class RealtimeResolvers extends Construct {
         layers: [powertoolsLayerJS],
         handler: "index.handler",
         runtime: Runtime.NODEJS_18_X,
+        loggingFormat: LoggingFormat.JSON,
         environment: {
           GRAPHQL_ENDPOINT: props.api.graphqlUrl,
         },
@@ -106,6 +113,7 @@ export class RealtimeResolvers extends Construct {
       dataSource: noneDataSource,
     });
 
+    this.sendQueryHandler = resolverFunction;
     this.outgoingMessageHandler = outgoingMessageHandler;
   }
 }

lib/chatbot-api/functions/api-handler/index.py

@@ -40,11 +40,16 @@
 
 
 @logger.inject_lambda_context(
-    log_event=True, correlation_id_path=correlation_paths.APPSYNC_RESOLVER
+    log_event=False, correlation_id_path=correlation_paths.APPSYNC_RESOLVER
 )
 @tracer.capture_lambda_handler
 def handler(event: dict, context: LambdaContext) -> dict:
     try:
+        logger.info(
+            "Incoming request for " + event["info"]["fieldName"],
+            arguments=event["arguments"],
+            identify=event["identity"],
+        )
         return app.resolve(event, context)
     except ValidationError as e:
         logger.warning(e.errors())

lib/chatbot-api/functions/api-handler/routes/documents.py

@@ -141,7 +141,7 @@ def file_upload(input: dict):
         request.workspaceId, request.fileName
     )
 
-    print(result)
+    logger.info("Generated pre-signed for " + request.fileName)
     return result
 
 

lib/chatbot-api/functions/resolvers/send-query-lambda-resolver/index.py

@@ -5,6 +5,7 @@
 from datetime import datetime
 from aws_lambda_powertools import Logger, Tracer
 from aws_lambda_powertools.utilities.typing import LambdaContext
+from aws_lambda_powertools.logging import correlation_paths
 from pydantic import BaseModel, Field
 
 tracer = Tracer()
@@ -48,10 +49,15 @@ class InputValidation(BaseModel):
 
 
 @tracer.capture_lambda_handler
-@logger.inject_lambda_context(log_event=True)
+@logger.inject_lambda_context(
+    log_event=False, correlation_id_path=correlation_paths.APPSYNC_RESOLVER
+)
 def handler(event, context: LambdaContext):
-    print(event["arguments"]["data"])
-    print(event["identity"])
+    logger.info(
+        "Incoming request for " + event["info"]["fieldName"],
+        arguments=event["arguments"],
+        identify=event["identity"],
+    )
     request = json.loads(event["arguments"]["data"])
     message = {
         "action": request["action"],
@@ -61,7 +67,6 @@ def handler(event, context: LambdaContext):
         "userId": event["identity"]["sub"],
         "data": request.get("data", {}),
     }
-    print(message)
     InputValidation(**message)
 
     try:

lib/chatbot-api/index.ts

@@ -16,6 +16,7 @@ import { ChatBotS3Buckets } from "./chatbot-s3-buckets";
 import { ApiResolvers } from "./rest-api";
 import { RealtimeGraphqlApiBackend } from "./websocket-api";
 import * as appsync from "aws-cdk-lib/aws-appsync";
+import * as lambda from "aws-cdk-lib/aws-lambda";
 import { RetentionDays } from "aws-cdk-lib/aws-logs";
 import { NagSuppressions } from "cdk-nag";
 
@@ -36,6 +37,7 @@ export class ChatBotApi extends Construct {
   public readonly filesBucket: s3.Bucket;
   public readonly userFeedbackBucket: s3.Bucket;
   public readonly graphqlApi: appsync.GraphqlApi;
+  public readonly resolvers: lambda.Function[] = [];
 
   constructor(scope: Construct, id: string, props: ChatBotApiProps) {
     super(scope, id);
@@ -87,19 +89,25 @@ export class ChatBotApi extends Construct {
         : appsync.Visibility.GLOBAL,
     });
 
-    new ApiResolvers(this, "RestApi", {
+    const apiResolvers = new ApiResolvers(this, "RestApi", {
       ...props,
       sessionsTable: chatTables.sessionsTable,
       byUserIdIndex: chatTables.byUserIdIndex,
       api,
       userFeedbackBucket: chatBuckets.userFeedbackBucket,
     });
 
+    this.resolvers.push(apiResolvers.appSyncLambdaResolver);
+
     const realtimeBackend = new RealtimeGraphqlApiBackend(this, "Realtime", {
       ...props,
       api,
+      logRetention: props.config.logRetention,
     });
 
+    this.resolvers.push(realtimeBackend.resolvers.sendQueryHandler);
+    this.resolvers.push(realtimeBackend.resolvers.outgoingMessageHandler);
+
     realtimeBackend.resolvers.outgoingMessageHandler.addEnvironment(
       "GRAPHQL_ENDPOINT",
       api.graphqlUrl

lib/chatbot-api/rest-api.ts

@@ -30,6 +30,7 @@ export interface ApiResolversProps {
 }
 
 export class ApiResolvers extends Construct {
+  readonly appSyncLambdaResolver: lambda.Function;
   constructor(scope: Construct, id: string, props: ApiResolversProps) {
     super(scope, id);
 
@@ -45,12 +46,14 @@ export class ApiResolvers extends Construct {
           path.join(__dirname, "./functions/api-handler")
         ),
         handler: "index.handler",
+        description: "Main Appsync resolver",
         runtime: props.shared.pythonRuntime,
         architecture: props.shared.lambdaArchitecture,
         timeout: cdk.Duration.minutes(10),
         memorySize: 512,
         tracing: lambda.Tracing.ACTIVE,
-        logRetention: logs.RetentionDays.ONE_WEEK,
+        logRetention: props.config.logRetention ?? logs.RetentionDays.ONE_WEEK,
+        loggingFormat: lambda.LoggingFormat.JSON,
         layers: [props.shared.powerToolsLayer, props.shared.commonLayer],
         vpc: props.shared.vpc,
         securityGroups: [apiSecurityGroup],
@@ -117,6 +120,7 @@ export class ApiResolvers extends Construct {
         },
       }
     );
+    this.appSyncLambdaResolver = appSyncLambdaResolver;
 
     function addPermissions(apiHandler: lambda.Function) {
       if (props.ragEngines?.workspacesTable) {

lib/chatbot-api/websocket-api.ts

@@ -17,6 +17,7 @@ interface RealtimeGraphqlApiBackendProps {
   readonly shared: Shared;
   readonly userPool: UserPool;
   readonly api: appsync.GraphqlApi;
+  readonly logRetention?: number;
 }
 
 export class RealtimeGraphqlApiBackend extends Construct {
@@ -64,6 +65,7 @@ export class RealtimeGraphqlApiBackend extends Construct {
       userPool: props.userPool,
       shared: props.shared,
       api: props.api,
+      logRetention: props.logRetention,
     });
 
     // Route all outgoing messages to the websocket interface queue

lib/model-interfaces/idefics/functions/request-handler/adapters/claude.py

@@ -1,3 +1,4 @@
+from aws_lambda_powertools import Logger
 from .base import MultiModalModelBase
 from genai_core.types import ChatbotMessageType
 from urllib.parse import urljoin
@@ -8,6 +9,8 @@
 from base64 import b64encode
 from genai_core.registry import registry
 
+logger = Logger()
+
 
 def get_image_message(
     file: dict,
@@ -69,7 +72,7 @@ def format_prompt(self, prompt: str, messages: list, files: list) -> str:
         )
 
     def handle_run(self, prompt: str, model_kwargs: dict):
-        print(model_kwargs)
+        logger.info("Incoming request for claude", model_kwargs=model_kwargs)
         body = json.loads(prompt)
 
         if "temperature" in model_kwargs:

lib/model-interfaces/idefics/functions/request-handler/adapters/idefics.py

@@ -1,3 +1,4 @@
+from aws_lambda_powertools import Logger
 from .base import MultiModalModelBase
 from genai_core.types import ChatbotMessageType
 from urllib.parse import urljoin
@@ -6,6 +7,8 @@
 from content_handler import ContentHandler
 from genai_core.registry import registry
 
+logger = Logger()
+
 
 class Idefics(MultiModalModelBase):
     model_id: str
@@ -53,11 +56,11 @@ def format_prompt(self, prompt: str, messages: list, files: list) -> str:
         prompts.append("<end_of_utterance>\nAssistant:")
 
         prompt_template = "".join(prompts)
-        print(prompt_template)
+        logger.info(prompt_template)
         return prompt_template
 
     def handle_run(self, prompt: str, model_kwargs: dict):
-        print(model_kwargs)
+        logger.info("Incoming request for idefics", model_kwargs=model_kwargs)
         params = {
             "do_sample": True,
             "top_p": 0.2,

lib/model-interfaces/idefics/functions/request-handler/index.py

@@ -21,7 +21,7 @@
 
 
 def handle_run(record):
-    print(record)
+    logger.info("Incoming request", record=record)
     user_id = record["userId"]
     data = record["data"]
     provider = data["provider"]

lib/model-interfaces/idefics/index.ts

@@ -191,14 +191,16 @@ export class IdeficsInterface extends Construct {
         code: props.shared.sharedCode.bundleWithLambdaAsset(
           path.join(__dirname, "./functions/request-handler")
         ),
+        description: "Idefics request handler",
         runtime: props.shared.pythonRuntime,
         handler: "index.handler",
         layers: [props.shared.powerToolsLayer, props.shared.commonLayer],
         architecture: props.shared.lambdaArchitecture,
         tracing: lambda.Tracing.ACTIVE,
         timeout: cdk.Duration.minutes(lambdaDurationInMinutes),
         memorySize: 1024,
-        logRetention: logs.RetentionDays.ONE_WEEK,
+        logRetention: props.config.logRetention ?? logs.RetentionDays.ONE_WEEK,
+        loggingFormat: lambda.LoggingFormat.JSON,
         environment: {
           ...props.shared.defaultEnvironmentVariables,
           CONFIG_PARAMETER_NAME: props.shared.configParameter.parameterName,