Update app/llm.py

app/llm.py

@@ -22,7 +22,7 @@ from langchain_core.prompts import PromptTemplate
 from langchain import hub
 from langchain_core.runnables import RunnablePassthrough, RunnablePick
 
-rag_prompt = hub.pull("rlm/rag-prompt-llama")
+rag_prompt_llama = hub.pull("rlm/rag-prompt-llama")
 rag_prompt.messages
 
 class RagChat:
@@ -38,8 +38,9 @@ class RagChat:
         return vectorstore
 
     def chat(self, question):
+        retriever = vectorstore.as_retriever()
         chain = (
-            RunnablePassthrough.assign(context=RunnablePick("context") | format_docs)
+            {"context": retriever | format_docs, "question": RunnablePassthrough()}
             | rag_prompt_llama
             | llm
             | StrOutputParser()
@@ -50,10 +51,6 @@ class RagChat:
         docs = vectorstore.similarity_search(question)
         return docs
         
-        
-    
-
-
 
 class GenModel(BaseModel):
     question: str