Update lambda.tf and index.py files

lambdas/inference/lambda.tf

@@ -61,17 +61,13 @@ module "lambda_function_container_image" {
         "logs:PutLogEvents",
       ]
     }
+
     bedrock_usage = {
       effect = "Allow"
-
-      resources = [
-        "*"
-      ]
-
-      actions = [
-        "bedrock:*"
-      ]
+      resources = ["*"]
+      actions = ["bedrock:*"]
     }
+
     rds_connect_readwrite = {
       effect = "Allow"
 

lambdas/inference/src/index.py

@@ -4,7 +4,11 @@
 from retrieval import Retrieval
 from history import History
 import uuid
+from aws_lambda_powertools import Logger, Metrics, Tracer
 
+tracer = Tracer()
+logger = Logger()
+metrics = Metrics()
 
 HEADERS = {
     "Access-Control-Allow-Origin": "*",
@@ -78,8 +82,6 @@ def invoke_model(
     if source == "call":
         maxtokens //= 2
 
-    messages.append({"role": "user", "content": user_message})
-
     body = json.dumps(
         {
             "anthropic_version": "bedrock-2023-05-31",
@@ -104,6 +106,9 @@ def invoke_model(
         raise e
 
 
+@metrics.log_metrics
+@logger.inject_lambda_context
+@tracer.capture_lambda_handler
 def lambda_handler(event, context):
     response = "this is a dummy response"
     source = event.get("queryStringParameters", {}).get("source", "message")
@@ -127,13 +132,14 @@ def lambda_handler(event, context):
         )
         docs = retrieval.fetch_documents(query=query, top_k=ENV_VARS["top_k"])
 
-
         # prepare the prompt
         system_prompt = prepare_system_prompt(docs, source)
 
         chat_history = history.get(limit=5)
         chat_history.append({"role": "user", "content": query})
 
+        logger.info(f"Chat history: {chat_history}")
+
         response = invoke_model(
             system_prompt=system_prompt,
             source=source,

lambdas/inference/src/requirements.txt

@@ -3,4 +3,5 @@ langchain-community
 langchain
 psycopg2-binary
 pgvector
-boto3
+boto3
+aws-lambda-powertools