diff --git a/README.md b/README.md
index 7589c17..99e6906 100755
--- a/README.md
+++ b/README.md
@@ -62,7 +62,13 @@ git clone <repository-url>
 
 You're ready to start developing your chat function. Head over to the [Development](#development) section to learn more.
 
-#### 4. Update the README
+#### 4. Deploy the chat function
+
+You will have to add your API key and LLM model name into the Github repo settings. Under `Secrets and variables/Actions`: the API key must be added as a secret and the LLM model must be added as a variable.
+
+You must ensure the same namings as in your `.env` file. So, make sure to update the `.github/{dev and main}.yml` files with the correct parameter names. 
+
+#### 5. Update the README
 
 In the `README.md` file, change the title and description so it fits the purpose of your chat function.
 
diff --git a/index.py b/index.py
index 5ef3bb8..cf071dc 100644
--- a/index.py
+++ b/index.py
@@ -1,42 +1,38 @@
 import json
 try:
     from .src.module import chat_module
-    from .src.agents.utils.types import JsonType
 except ImportError:
     from src.module import chat_module
-    from src.agents.utils.types import JsonType
 
-def handler(event: JsonType, context):
+def handler(event, context):
     """
     Lambda handler function
-    Args:
-        event (JsonType): The AWS Lambda event received by the gateway.
-        context (Any): The AWS Lambda context object.
-
     """
     # Log the input event for debugging purposes
-    print("Received event:", json.dumps(event, indent=2))
+    # print("Received event:", " ".join(json.dumps(event, indent=2).splitlines()))
 
-    if "body" not in event:
-        return {
-            "statusCode": 400,
-            "body": "Missing 'body' key in event. Please confirm the key in the json body."
-        }
-    body = json.loads(event["body"])
-        
-    if "message" not in body:
+    if "body" in event:
+        try:
+            event = json.loads(event["body"])
+        except json.JSONDecodeError:
+            return {
+                "statusCode": 400,
+                "body": "Invalid JSON format in the body or body not found. Please check the input."
+            }
+
+    if "message" not in event:
         return {
             "statusCode": 400,
             "body": "Missing 'message' key in event. Please confirm the key in the json body."
         }
-    if "params" not in body:
+    if "params" not in event:
         return {
             "statusCode": 400,
             "body": "Missing 'params' key in event. Please confirm the key in the json body. Make sure it contains the necessary conversation_id."
         }
     
-    message = body["message"]
-    params = body["params"]
+    message = event.get("message", None)
+    params = event.get("params", None)
 
     try:
         chatbot_response = chat_module(message, params)
@@ -49,7 +45,10 @@ def handler(event: JsonType, context):
     # Create a response
     response = {
         "statusCode": 200,
-        "body": json.dumps(chatbot_response)
+        "body": chatbot_response
     }
 
+    # Log the response for debugging purposes
+    print("Returning response:", " ".join(json.dumps(response, indent=2).splitlines()))
+
     return response
\ No newline at end of file
diff --git a/src/agents/utils/prompt_context_templates.py b/src/agents/utils/prompt_context_templates.py
index 9eb175b..5d3d0af 100644
--- a/src/agents/utils/prompt_context_templates.py
+++ b/src/agents/utils/prompt_context_templates.py
@@ -73,8 +73,9 @@ def format_question_header(
 - Description: {content}
 {duration_text}
 
-> Note: Mathematical equations are in KaTeX format, preserve them the same. Use British English spellings.
-
+> Note: Mathematical equations are in KaTeX format, preserve them the same. Ensure mathematical equations are surrounded by one '$' for in-line equations and '$$' for block equations. 
+Example: '$E=mc^2$' or '$$E=mc^2$$'.
+Use British English spellings.
 ---
 """
 
diff --git a/src/module.py b/src/module.py
index 2b7b2b4..02c6f66 100755
--- a/src/module.py
+++ b/src/module.py
@@ -1,6 +1,7 @@
 import time
 from typing import Any
-from lf_toolkit.chat import ChatResult as Result, ChatParams as Params
+from lf_toolkit.chat.result import ChatResult as Result
+from lf_toolkit.chat.params import ChatParams as Params
 
 try:
     from .agents.utils.parse_json_context_to_prompt import parse_json_to_prompt