Update chain.py

chain.py

@@ -5,6 +5,9 @@ from langchain_core.prompts import PromptTemplate
 from langchain_huggingface import HuggingFaceEndpoint
 from langchain_core.output_parsers import StrOutputParser
 
+import logging
+logger = logging.getLogger(__name__)
+
 
 def get_chain(
     vectordb,
@@ -26,8 +29,11 @@ Always say "Thanks for asking!" at the end of the answer.
 Question: {question}
 Helpful Answer:"""
 ):
-    retriever = vectordb.as_retriever(search_type=search_type, search_kwargs={"k": k, "fetch_k": fetch_k})
+    search_kwargs = {"k": k, "fetch_k": fetch_k}
+    logger.info(f'Setting up vectordb retriever with search_type={search_type} and search_kwargs={search_kwargs}')
+    retriever = vectordb.as_retriever(search_type=search_type, search_kwargs=search_kwargs)
 
+    logger.info('Setting up retrieval runnable')
     retrieval = RunnableParallel(
         {
             "context": RunnableLambda(lambda x: retriever.invoke(x["question"])),
@@ -35,8 +41,10 @@ Helpful Answer:"""
         }
     )
 
+    logger.info(f'Setting up prompt from the template:\n{template}')
     prompt = PromptTemplate(input_variables=["context", "question"], template=template)
 
+    logger.info(f'Instantiating llm with repo_id={repo_id}, task={task}, max_new_tokens={max_new_tokens}, top_k={top_k}, temperature={temperature} and repetition_penalty={repetition_penalty}')
     llm = HuggingFaceEndpoint(
         repo_id=repo_id,
         task=task,
@@ -46,5 +54,6 @@ Helpful Answer:"""
         repetition_penalty=repetition_penalty,
     )
 
+    logger.info('Instantiating and returning chain = retrieval | prompt | llm | StrOutputParser()')
     return retrieval | prompt | llm | StrOutputParser()