add question rephrase

main.py

@@ -26,9 +26,13 @@ RETRIEVER_MAX_DOCS = 50
 RERANKER_MAX_RESULTS = 25
 
 with open("prompt.md") as f:
-    template = f.read()
+    question_template = f.read()
 
-prompt = ChatPromptTemplate.from_template(template)
+with open("question_rewrite_prompt.md") as f:
+    question_rewrite_template = f.read()
+
+question_prompt = ChatPromptTemplate.from_template(question_template)
+question_rewrite_prompt = ChatPromptTemplate.from_template(question_rewrite_template)
 
 
 def format_docs(question: str) -> str:
@@ -53,6 +57,12 @@ def format_docs(question: str) -> str:
     return "\n\n".join(doc.page_content for doc in reranked_docs)
 
 
+def log_rewritten_question(rewritten_question: str) -> str:
+    print("=== REWRITTEN QUESTION ===")
+    print(rewritten_question)
+    return rewritten_question
+
+
 llm = ChatGoogleGenerativeAI(
     model=MODEL,
     project=PROJECT,
@@ -82,31 +92,29 @@ compression_retriever = ContextualCompressionRetriever(
     base_compressor=reranker, base_retriever=base_retriever
 )
 
+question_rewrite_chain = (
+    {"question": RunnablePassthrough()}
+    | question_rewrite_prompt
+    | llm
+    | StrOutputParser()
+    | RunnableLambda(log_rewritten_question)
+)
+
 rag_chain = (
     {"context": RunnableLambda(format_docs), "question": RunnablePassthrough()}
-    | prompt
+    | question_prompt
     | llm
     | StrOutputParser()
 )
 
-
-async def async_invoke(rag_chain, prompt: str):
-    return rag_chain.invoke(prompt)
+full_chain = question_rewrite_chain | rag_chain
 
 
 async def main():
-    (
-        res1,
-        res2,
-    ) = await asyncio.gather(
-        async_invoke(rag_chain, "come si calcola l'angolo di fase?"),
-        async_invoke(rag_chain, "cos'e' la massa magra?"),
+    response = await full_chain.ainvoke(
+        "Buongiorno, non so se è la mail specifica ma volevo se possibile dei chiarimenti per l’interpretazione dei parametri BCM /SMM/ASMM. Mi capita a volte di trovare casi in cui la BCM è aumentata ma allo stesso tempo SMM/ASMM hanno subito una piccola flessione in negativo (o viceversa). Se la parte metabolicamente attiva aumenta perchè può succedere che gli altri compartimenti si riducono?? E allo stesso tempo phA e BCM possono essere inversamente proporzionali?? So che il phA correla con massa e struttura + idratazione."
     )
-
-    print("RES1")
-    print(res1)
-    print("\n\nRES2")
-    print(res2)
+    print(response)
 
 
 if __name__ == "__main__":