diff --git a/app/pipeline/prompts/citation_prompt.txt b/app/pipeline/prompts/citation_prompt.txt
index e62c2aca..4c143351 100644
--- a/app/pipeline/prompts/citation_prompt.txt
+++ b/app/pipeline/prompts/citation_prompt.txt
@@ -1,6 +1,6 @@
 In the paragraphs below you are provided with an answer to a question. Underneath the answer you will find the paragraphs that the answer was based on.
 Add citations of the paragraphs to the answer. Cite the paragraphs in brackets after the sentence where the information is used in the answer.
-At the end of the answer list each source with its corresponding number and provide the Lecture Title,as well as the page number in this format "[1] Lecture title, page number".
+At the end of the answer list each source with its corresponding number and provide the Lecture Title, Unit Name, as well as the page number in this format "[1] Lecture title, Unit name, page number".
 Do not Include the Actual paragraphs, only the citations at the end.
 Only include the paragraphs that are relevant to the answer.
 If the answer actually does not contain any information from the paragraphs, please do not include any citations and return '!NONE!'.
@@ -10,8 +10,8 @@ Here is an example how to rewrite the answer with citations (note the empty line
 "
 Lorem ipsum dolor sit amet, consectetur adipiscing elit [1]. Ded do eiusmod tempor incididunt ut labore et dolore magna aliqua [2].
 
-[1] Lecture 1, page 2.
-[2] Lecture 2, page 5.
+[1] Lecture 1, Unit A, page 2.
+[2] Lecture 2, Unit B, page 5.
 "
 
 Here are the answer and the paragraphs:
diff --git a/app/pipeline/shared/citation_pipeline.py b/app/pipeline/shared/citation_pipeline.py
index be73b981..6a4aab38 100644
--- a/app/pipeline/shared/citation_pipeline.py
+++ b/app/pipeline/shared/citation_pipeline.py
@@ -51,8 +51,9 @@ def create_formatted_string(self, paragraphs):
         """
         formatted_string = ""
         for i, paragraph in enumerate(paragraphs):
-            lct = "Lecture: {}, Page: {}\nContent:\n---{}---\n\n".format(
+            lct = "Lecture: {}, Unit: {}, Page: {}\nContent:\n---{}---\n\n".format(
                 paragraph.get(LectureSchema.LECTURE_NAME.value),
+                paragraph.get(LectureSchema.LECTURE_UNIT_NAME.value),
                 paragraph.get(LectureSchema.PAGE_NUMBER.value),
                 paragraph.get(LectureSchema.PAGE_TEXT_CONTENT.value),
             )