diff --git a/app.py b/app.py
index 1306d1d..6c17910 100644
--- a/app.py
+++ b/app.py
@@ -169,9 +169,8 @@ def save_listening():
             file_url = upload_file_firebase_get_url(FIREBASE_BUCKET, firebase_file_path, sound_file_path)
             template["parts"][i]["audio"]["source"] = file_url
             template["parts"][i]["exercises"] = part["exercises"]
-        (result, id) = save_to_db("listening", template)
-        if result:
-            return {**template, "id": id}
+        if save_to_db("listening", template):
+            return template
         else:
             raise Exception("Failed to save question: " + parts)
     except Exception as e:
@@ -280,9 +279,8 @@ def save_writing_task():
         template = getWritingTemplate()
         for i, exercise in enumerate(exercises, start=0):
             template["exercises"][i]["prompt"] = exercise
-        (result, id) = save_to_db("writing", template)
-        if result:
-            return {**template, "id": id}
+        if save_to_db("writing", template):
+            return template
         else:
             raise Exception("Failed to save writing: " + template)
     except Exception as e:
@@ -508,9 +506,8 @@ def save_speaking():
         template["exercises"][2]["prompts"] = sp3_questions
         template["exercises"][2]["title"] = exercises[2]["topic"]
 
-        (result, id) = save_to_db("speaking", template)
-        if result:
-            return {**template, "id": id}
+        if save_to_db("speaking", template):
+            return template
         else:
             raise Exception("Failed to save speaking: " + template)
     except Exception as e:
@@ -608,9 +605,8 @@ def save_reading_passage():
         parts = data.get('parts')
         template = getReadingTemplate()
         template["parts"] = parts
-        (result, id) = save_to_db("reading", template)
-        if result:
-            return {**template, "id": id}
+        if save_to_db("reading", template):
+            return template
         else:
             raise Exception("Failed to save reading: " + template)
     except Exception as e:
diff --git a/helper/firebase_helper.py b/helper/firebase_helper.py
index 05fbd89..d169167 100644
--- a/helper/firebase_helper.py
+++ b/helper/firebase_helper.py
@@ -51,10 +51,10 @@ def upload_file_firebase_get_url(bucket_name, destination_blob_name, source_file
 def save_to_db(collection: str, item):
     db = firestore.client()
     collection_ref = db.collection(collection)
-    document_ref = collection_ref.add(item)
+    (update_time, document_ref) = collection_ref.add(item)
     if document_ref:
-        print(f"Document added with ID: {document_ref}")
-        return (True, document_ref)
+        print(f"Document added with ID: {document_ref.id}")
+        return (True, document_ref.id)
     else:
         return (False, None)