Refactoring code making lambda more readable and reusable

serverless/common/utils/encoders.py

@@ -1,6 +1,7 @@
 import json
 from decimal import Decimal
 
+
 # Custom JSON Encoder for handling Decimal types from DynamoDB
 class DecimalEncoder(json.JSONEncoder):
     def default(self, obj):
@@ -9,4 +10,4 @@ def default(self, obj):
                 return float(obj)
             else:
                 return int(obj)
-        return super(DecimalEncoder, self).default(obj)
+        return super(DecimalEncoder, self).default(obj)

serverless/common/utils/error_handler.py

@@ -0,0 +1,22 @@
+# error_handlers.py
+import json
+
+
+def error_response(message, status_code=400):
+    return {
+        "statusCode": status_code,
+        "headers": {
+            "Access-Control-Allow-Origin": "*",
+            "Access-Control-Allow-Headers": "*",
+            "Content-Type": "application/json",
+        },
+        "body": json.dumps({"message": message}),
+    }
+
+
+def not_found_error():
+    return error_response("The specified resource was not found.", 404)
+
+
+def internal_server_error():
+    return error_response("Internal server error.", 500)

serverless/common/utils/lambda_utils.py

@@ -1,5 +1,6 @@
 import json
 
+
 def load_body_from_event(event):
     # Attempt to load JSON from the request body
     body = event.get("body", "{}")  # Default to empty JSON string if body is None
@@ -8,9 +9,9 @@ def load_body_from_event(event):
             body = json.loads(body)
         except json.JSONDecodeError:
             body = {}
-            print("Failed to parse JSON from request body")
     return body
 
+
 def load_path_parameter_from_event(event, param_name):
     # Extract a path parameter from the event
-    return event.get("pathParameters", {}).get(param_name)
+    return event.get("pathParameters", {}).get(param_name)

serverless/common/utils/response_utils.py

@@ -0,0 +1,15 @@
+# response_utils.py
+import json
+from common.utils.encoders import DecimalEncoder
+
+
+def success_response(data, status_code=200):
+    return {
+        "statusCode": status_code,
+        "headers": {
+            "Access-Control-Allow-Origin": "*",
+            "Access-Control-Allow-Headers": "*",
+            "Content-Type": "application/json",
+        },
+        "body": json.dumps(data, cls=DecimalEncoder),
+    }

serverless/serverless.yml

@@ -49,6 +49,7 @@ functions:
       patterns:
         - "!**"
         - services/health/**
+        - common/**
   upload_url_s3:
     handler: services/files/get_upload_url_s3.lambda_handler
     description: Get expires URL to upload a file to AWS bucket.
@@ -64,6 +65,7 @@ functions:
       patterns:
         - "!**"
         - services/files/**
+        - common/**
   gpt_ask:
     handler: services/gpt/ask.lambda_handler
     description: Given a set of messages ask to chat gpt and respond with a message.
@@ -115,6 +117,7 @@ functions:
       patterns:
         - "!**"
         - services/memory/**
+        - common/**
   update_session_metadata:
     handler: services/memory/update_session_metadata.lambda_handler
     description: Update session metadata like 'title' or 'user_id'.
@@ -174,6 +177,7 @@ functions:
       patterns:
         - "!**"
         - services/memory/**
+        - common/**
 
 resources:
   Resources:

serverless/services/files/get_upload_url_s3.py

@@ -1,49 +1,50 @@
 import os
-import json
 import datetime
 import boto3
 
+from common.utils.error_handler import error_response, internal_server_error
+from common.utils.response_utils import success_response
+
+s3_client = boto3.client("s3")
+s3_bucket_name = os.getenv("S3_BUCKET_NAME")
 
-def lambda_handler(event, context):
-    if not event.get("queryStringParameters"):
-        return {
-            "statusCode": 400,
-            "body": json.dumps({"message": "Query parameters are missing."}),
-        }
 
-    # Extract parameters from the query string
-    query_params = event["queryStringParameters"]
+def parse_and_validate(event):
+    query_params = event.get("queryStringParameters", {})
     s3_key = query_params.get("key")
     s3_expiresin = query_params.get("expiresIn")
 
-    # Check if required parameters are provided
-    if not s3_key or not s3_expiresin:
-        return {
-            "statusCode": 400,
-            "body": json.dumps(
-                {
-                    "message": "You must include 'key' and 'expiresIn' in your query request."
-                }
-            ),
-        }
+    if not s3_key or s3_expiresin is None:
+        raise ValueError(
+            "Both 'key' and 'expiresIn' parameters are required in the query."
+        )
 
     try:
-        # Ensure expiresIn is an integer
         s3_expiresin = int(s3_expiresin)
     except ValueError:
-        return {
-            "statusCode": 400,
-            "body": json.dumps({"message": "expiresIn must be a valid integer."}),
-        }
+        raise ValueError("'expiresIn' must be a valid integer.")
+
+    if (
+        s3_expiresin <= 0 or s3_expiresin > 86400
+    ):  # Example: limits expiresIn to 24 hours
+        raise ValueError("'expiresIn' must be between 1 and 86400 seconds.")
 
-    # Get the S3 bucket name from environment variables
-    s3_bucket_name = os.getenv("S3_BUCKET_NAME")
+    return s3_key, s3_expiresin
 
-    # Generate the presigned URL
-    s3_client = boto3.client("s3")
-    url = s3_client.generate_presigned_post(
-        Bucket=s3_bucket_name, Key=s3_key, ExpiresIn=s3_expiresin
-    )
+
+def lambda_handler(event, context):
+    try:
+        s3_key, s3_expiresin = parse_and_validate(event)
+    except ValueError as e:
+        return error_response(str(e))
+
+    try:
+        # Generate the presigned URL
+        url = s3_client.generate_presigned_post(
+            Bucket=s3_bucket_name, Key=s3_key, ExpiresIn=s3_expiresin
+        )
+    except Exception as e:
+        return internal_server_error()
 
     body = {
         "url": url,
@@ -52,13 +53,4 @@ def lambda_handler(event, context):
         "created_at": datetime.datetime.now().isoformat(),
     }
 
-    response = {
-        "statusCode": 200,
-        "headers": {
-            "Access-Control-Allow-Origin": "*",
-            "Access-Control-Allow-Headers": "*",
-            "Content-Type": "application/json",
-        },
-        "body": json.dumps(body),
-    }
-    return response
+    return success_response(body)

serverless/services/gpt/ask.py

@@ -1,52 +1,40 @@
 import os
-import json
 import openai
 
-from typing import List
+from typing import List, Dict, Any
 from .src.models import OpenAIModel
+
 from common.utils.lambda_utils import load_body_from_event
+from common.utils.error_handler import error_response, internal_server_error
+from common.utils.response_utils import success_response
 
 
-def lambda_handler(event, context):
-    if not event.get("body"):
-        return {
-            "statusCode": 400,
-            "headers": {"Content-Type": "application/json"},
-            "body": json.dumps(
-                {"message": "Body parameters are missing. You must include 'messages'."}
-            ),
-        }
+API_KEY = os.getenv("OPENAI_API_KEY")
+MODEL_ENGINE = os.getenv("OPENAI_GPTMODEL")
+TEMPERATURE = int(os.getenv("OPENAI_TEMPERATURE", 0))
+MAX_TOKENS = int(os.getenv("OPENAI_TOKENS", 0))
 
-    body = load_body_from_event(event)
-    messages = body.get("messages")
 
-    if len(messages) == 0:
-        return {
-            "statusCode": 400,
-            "headers": {"Content-Type": "application/json"},
-            "body": json.dumps({"message": "Body 'messages' must not be empty."}),
-        }
+def parse_and_validate(event: Dict[str, Any]) -> List[Dict[str, str]]:
+    messages = load_body_from_event(event).get("messages")
 
-    if isinstance(messages, List) is False:
-        return {
-            "statusCode": 400,
-            "headers": {"Content-Type": "application/json"},
-            "body": json.dumps({"message": "Body 'messages' must be a list."}),
-        }
+    if not messages or len(messages) == 0:
+        raise ValueError("Body 'messages' must not be empty.")
+
+    if not isinstance(messages, list):
+        raise ValueError("Body 'messages' must be a list.")
 
     for msg in messages:
         if not msg.get("content") or not msg.get("role"):
-            return {
-                "statusCode": 400,
-                "headers": {"Content-Type": "application/json"},
-                "body": json.dumps(
-                    {
-                        "message": "Messages should not be empty. It should have 'role' and 'content' keys."
-                    }
-                ),
-            }
+            raise ValueError(
+                "Messages should not be empty. It should have 'role' and 'content' keys."
+            )
+    return messages
 
+
+def lambda_handler(event, context):
     try:
+        messages = parse_and_validate(event)
         model = OpenAIModel(
             api_key=os.getenv("OPENAI_API_KEY"),
             model_engine=os.getenv("OPENAI_GPTMODEL"),
@@ -56,43 +44,23 @@ def lambda_handler(event, context):
 
         res = model.chat_completion(messages)
         if res.get("status") != "success":
-            return {
-                "statusCode": 400,
-                "headers": {"Content-Type": "application/json"},
-                "body": json.dumps({"role": res["role"], "content": res["content"]}),
-            }
+            error_content = res["content"]
+            return error_response(error_content)
         else:
             messages.append({"role": res["role"], "content": res["content"]})
-            return {
-                "statusCode": 200,
-                "headers": {"Content-Type": "application/json"},
-                "body": json.dumps(
-                    {
-                        "context": model.__str__(),
-                        "last_message": {
-                            "role": res["role"],
-                            "content": res["content"],
-                        },
-                        "memory": messages,
-                        "memory_count": len(messages),
-                    }
-                ),
+            body = {
+                "context": str(model),
+                "last_message": {
+                    "role": res["role"],
+                    "content": res["content"],
+                },
+                "memory": messages,
+                "memory_count": len(messages),
             }
-    except ValueError as ve:
-        return {
-            "statusCode": 400,
-            "headers": {"Content-Type": "application/json"},
-            "body": json.dumps({"role": ve["role"], "content": ve["content"]}),
-        }
+            return success_response(body)
+    except ValueError as e:
+        return error_response(str(e))
     except openai.OpenAIError as e:
-        return {
-            "statusCode": 400,
-            "headers": {"Content-Type": "application/json"},
-            "body": json.dumps({"role": ve["role"], "content": ve["content"]}),
-        }
+        return error_response(str(e))
     except Exception as e:
-        return {
-            "statusCode": 500,
-            "headers": {"Content-Type": "application/json"},
-            "body": json.dumps({"message": "Internal Server Error."}),
-        }
+        return internal_server_error()

serverless/services/health/health_check.py

@@ -1,26 +1,11 @@
-import json
+from common.utils.response_utils import success_response
 
 
 def lambda_handler(event, context):
-    """Health check function to test if the Lambda function is running.
-
-    Args:
-        event (obj): Event data passed to the function.
-        context (obj): Runtime information provided by AWS Lambda.
-
-    Returns:
-        json: A JSON object containing the health check message.
-    """
     body = {
         "message": "Health check successful",
         "region": context.invoked_function_arn.split(":")[3],
         "function_name": context.function_name,
     }
 
-    return {
-        "statusCode": 200,
-        "headers": {
-            "Content-Type": "application/json",
-        },
-        "body": json.dumps(body),
-    }
+    return success_response(body)