diff --git a/lambdas/inference/src/index.py b/lambdas/inference/src/index.py
index c82bdce..026b50f 100644
--- a/lambdas/inference/src/index.py
+++ b/lambdas/inference/src/index.py
@@ -7,7 +7,7 @@
 
 def prepare_source_prompt(source: str):
     if source == "email":
-        return "You are currently answering an email. After you finish answering the initial query anticipate the user's follow-up questions and answer it too up to 4 questions."
+        return "You are currently answering an email so your answer can be more detailed. After you finish answering the initial query generate follow-up questions and answer it too up to 4 questions."
     elif source == "call":
         return "Make your answer short and concise."
     else:
@@ -40,9 +40,9 @@ def prepare_prompt(query: str, docs: list, history: list, source: str):
     history_prompt = prepare_history_prompt(history)
     source_prompt = prepare_source_prompt(source)
     final_prompt = f"""{basic_prompt}\n
+    {source_prompt}\n
     {document_prompt}\n
     {history_prompt}\n
-    {source_prompt}\n
     {ENV_VARS['system_prompt']}\n
     \nAssistant:"""
     print(final_prompt)
diff --git a/lambdas/inference/src/retrieval.py b/lambdas/inference/src/retrieval.py
index 1b1b123..b72040e 100644
--- a/lambdas/inference/src/retrieval.py
+++ b/lambdas/inference/src/retrieval.py
@@ -50,7 +50,8 @@ def __init__(
     def fetch_documents(self, query: str, top_k: int = 10):
         try:
             docs = self._vector_store.similarity_search_with_relevance_scores(
-                query=query, k=top_k
+                query=query,
+                k=top_k,
             )
             return [x for x in docs if x[1] > self._relevance_treshold]
         except Exception as e: