diff --git a/.aws/task_definition_preview.json b/.aws/task_definition_preview.json
index 950af7d5a..c5f1df54f 100644
--- a/.aws/task_definition_preview.json
+++ b/.aws/task_definition_preview.json
@@ -64,18 +64,72 @@
       }
     },
     {
-      "name": "datadog-agent",
-      "image": "datadog/agent:latest",
-      "cpu": 10,
-      "memory": 256,
-      "essential": true,
-      "portMappings": [
+      "containerDefinitions": [
         {
-          "hostPort": 8126,
-          "protocol": "tcp",
-          "containerPort": 8126
+          "name": "datadog-agent",
+          "image": "public.ecr.aws/datadog/agent:latest",
+          "cpu": 100,
+          "memory": 512,
+          "essential": true,
+          "portMappings": [
+            {
+              "hostPort": 8126,
+              "protocol": "tcp",
+              "containerPort": 8126
+            }
+          ],
+          "mountPoints": [
+            {
+              "containerPath": "/var/run/docker.sock",
+              "sourceVolume": "docker_sock",
+              "readOnly": null
+            },
+            {
+              "containerPath": "/host/sys/fs/cgroup",
+              "sourceVolume": "cgroup",
+              "readOnly": null
+            },
+            {
+              "containerPath": "/host/proc",
+              "sourceVolume": "proc",
+              "readOnly": null
+            }
+          ],
+          "environment": [
+            {
+              "name": "DD_SITE",
+              "value": "datadoghq.com"
+            }
+          ],
+          "environmentFiles": [
+            {
+              "value": "arn:aws:s3:::quivr-env-variables/preview.env",
+              "type": "s3"
+            }
+          ]
         }
-      ]
+      ],
+      "volumes": [
+        {
+          "host": {
+            "sourcePath": "/var/run/docker.sock"
+          },
+          "name": "docker_sock"
+        },
+        {
+          "host": {
+            "sourcePath": "/proc/"
+          },
+          "name": "proc"
+        },
+        {
+          "host": {
+            "sourcePath": "/sys/fs/cgroup/"
+          },
+          "name": "cgroup"
+        }
+      ],
+      "family": "datadog-agent-task"
     }
   ],
   "family": "quivr-preview",
diff --git a/backend/llm/composite_brain_qa.py b/backend/llm/composite_brain_qa.py
index 2225c999f..05bf6b73c 100644
--- a/backend/llm/composite_brain_qa.py
+++ b/backend/llm/composite_brain_qa.py
@@ -289,7 +289,6 @@ class CompositeBrainQA(
                 )
 
                 print("querying brain", function_name)
-                # TODO: extract chat_id from generate_answer function of XBrainQA
                 function_response = function_to_call(
                     chat_id=chat_id,
                     question=question,