add time measure

climateqa/engine/chains/answer_rag.py

@@ -7,6 +7,7 @@ from langchain_core.prompts.base import format_document
 
 from climateqa.engine.chains.prompts import answer_prompt_template,answer_prompt_without_docs_template,answer_prompt_images_template
 from climateqa.engine.chains.prompts import papers_prompt_template
+import time
 from ..utils import rename_chain, pass_values
 
 
@@ -42,6 +43,7 @@ def make_rag_chain(llm):
     prompt = ChatPromptTemplate.from_template(answer_prompt_template)
     chain = ({
         "context":lambda x : _combine_documents(x["documents"]),
+        "context_length":lambda x : print("CONTEXT LENGTH : " , len(_combine_documents(x["documents"]))),
         "query":itemgetter("query"),
         "language":itemgetter("language"),
         "audience":itemgetter("audience"),
@@ -53,7 +55,6 @@ def make_rag_chain_without_docs(llm):
     chain = prompt | llm | StrOutputParser()
     return chain
 
-
 def make_rag_node(llm,with_docs = True):
 
     if with_docs:
@@ -63,9 +64,16 @@ def make_rag_node(llm,with_docs = True):
 
     async def answer_rag(state,config):
         print("---- Answer RAG ----")
+        start_time = time.time()
 
         answer = await rag_chain.ainvoke(state,config)
-        print(f"\n\nAnswer:\n{answer}")
+    
+        end_time = time.time()
+        elapsed_time = end_time - start_time
+        print("RAG elapsed time: ", elapsed_time)
+        print("Answer size : ", len(answer))
+        # print(f"\n\nAnswer:\n{answer}")
+        
         return {"answer":answer}
 
     return answer_rag