chore: address review comments - first run

Author: massi-ang

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

@@ -71,14 +71,19 @@ export class AwsGenAILLMChatbotStack extends cdk.Stack {
       (model) => model.interface === ModelInterface.LangChain
     );
 
+    let bedrockAgentsInterface: BedrockAgentInterface | undefined;
     if (props.config.bedrock?.enabled) {
-      new BedrockAgentInterface(this, "IBedrockAgent", {
-        shared,
-        config: props.config,
-        messagesTopic: chatBotApi.messagesTopic,
-        sessionsTable: chatBotApi.sessionsTable,
-        byUserIdIndex: chatBotApi.byUserIdIndex,
-      });
+      bedrockAgentsInterface = new BedrockAgentInterface(
+        this,
+        "IBedrockAgent",
+        {
+          shared,
+          config: props.config,
+          messagesTopic: chatBotApi.messagesTopic,
+          sessionsTable: chatBotApi.sessionsTable,
+          byUserIdIndex: chatBotApi.byUserIdIndex,
+        }
+      );
     }
 
     // check if any deployed model requires langchain interface or if bedrock is enabled from config
@@ -220,6 +225,9 @@ export class AwsGenAILLMChatbotStack extends cdk.Stack {
         chatBotApi.outBoundQueue,
         ideficsInterface.ingestionQueue,
         ...(langchainInterface ? [langchainInterface.ingestionQueue] : []),
+        ...(bedrockAgentsInterface
+          ? [bedrockAgentsInterface.ingestionQueue]
+          : []),
       ],
       aurora: ragEngines?.auroraPgVector?.database,
       opensearch: ragEngines?.openSearchVector?.openSearchCollection,

lib/model-interfaces/bedrock-agents/functions/request-handler/adapters/bedrock/base.py

@@ -2,6 +2,9 @@
 from typing import Any, Dict, Iterator, Optional
 from pydantic import root_validator, Extra
 from ..base import AgentAdapter
+from aws_lambda_powertools import Logger
+
+logger = Logger()
 
 
 class AgentInputOutputAdapter:
@@ -19,7 +22,7 @@ def prepare_output_stream(cls, response: Any) -> Iterator[str]:
             yield ""
 
         for event in stream:
-            print(event)
+            logger.info("Stream event", event=event)
             chunk = event.get("chunk")
             # chunk obj format varies with provider
             yield chunk["bytes"].decode("utf8")
@@ -117,6 +120,3 @@ def _invoke_agent(
 
         for event in response["completion"]:
             yield event
-
-        # for chunk in AgentInputOutputAdapter.prepare_output_stream(response):
-        #     yield chunk

lib/model-interfaces/bedrock-agents/functions/request-handler/index.py

@@ -18,7 +18,6 @@
 logger = Logger()
 
 AWS_REGION = os.environ["AWS_REGION"]
-API_KEYS_SECRETS_ARN = os.environ["API_KEYS_SECRETS_ARN"]
 CONFIG_PARAMETER_NAME = os.environ["CONFIG_PARAMETER_NAME"]
 
 sequence_number = 0
@@ -49,7 +48,7 @@ def handle_run(record):
         prompt,
     )
 
-    logger.info(response)
+    logger.info("Bedrock Agent response", response=response)
     sequence_number = 0
     run_id = str(uuid.uuid4())
     for r in response:
@@ -109,7 +108,7 @@ def record_handler(record: SQSRecord):
     payload: str = record.body
     message: dict = json.loads(payload)
     detail: dict = json.loads(message["Message"])
-    logger.info(detail)
+    logger.info("Incoming event", detail=detail)
 
     if detail["action"] == ChatbotAction.RUN.value:
         handle_run(detail)
@@ -123,7 +122,7 @@ def handle_failed_records(records):
         payload: str = record.body
         message: dict = json.loads(payload)
         detail: dict = json.loads(message["Message"])
-        logger.info(detail)
+        logger.info("Failed event", detail=detail)
         user_id = detail["userId"]
         data = detail.get("data", {})
         session_id = data.get("sessionId", "")
@@ -137,7 +136,7 @@ def handle_failed_records(records):
                 "timestamp": str(int(round(datetime.now().timestamp()))),
                 "data": {
                     "sessionId": session_id,
-                    "content": str(error),
+                    "content": "Something went wrong",
                     "type": "text",
                 },
             }
@@ -155,7 +154,12 @@ def handler(event, context: LambdaContext):
     except BatchProcessingError as e:
         logger.error(e)
 
-    logger.info(processed_messages)
+    for message in processed_messages:
+        logger.info(
+            "Request complete with status " + message[0],
+            status=message[0],
+            cause=message[1],
+        )
     handle_failed_records(
         message for message in processed_messages if message[0] == "fail"
     )

lib/model-interfaces/bedrock-agents/index.ts

@@ -39,14 +39,14 @@ export class BedrockAgentInterface extends Construct {
       tracing: lambda.Tracing.ACTIVE,
       timeout: cdk.Duration.minutes(15),
       memorySize: 1024,
-      logRetention: logs.RetentionDays.ONE_WEEK,
+      logRetention: props.config.logRetention ?? logs.RetentionDays.ONE_WEEK,
+      loggingFormat: lambda.LoggingFormat.JSON,
       layers: [props.shared.powerToolsLayer, props.shared.commonLayer],
       environment: {
         ...props.shared.defaultEnvironmentVariables,
         CONFIG_PARAMETER_NAME: props.shared.configParameter.parameterName,
         SESSIONS_TABLE_NAME: props.sessionsTable.tableName,
         SESSIONS_BY_USER_ID_INDEX_NAME: props.byUserIdIndex,
-        API_KEYS_SECRETS_ARN: props.shared.apiKeysSecret.secretArn,
         MESSAGES_TOPIC_ARN: props.messagesTopic.topicArn,
       },
     });

lib/shared/layers/python-sdk/python/genai_core/models.py

@@ -70,7 +70,7 @@ def list_bedrock_agents():
 
         return bedrock_agents
     except Exception as e:
-        print(f"Error listing Bedrock agents: {e}")
+        logger.error("Error listing Bedrock agents:", e)
         return None