Update custom_llm.py

custom_llm.py

@@ -47,11 +47,11 @@ def custom_chain_with_history(llm, memory):
     <|you|>
     """)
     
-    def prompt_memory(_):
+    def prompt_memory(memory):
       t = ""
       # for x in memory.chat_memory.messages:
       for x in memory.messages:
-        t += f"<|assistant|>\n<s>{x.content}</s>\n" if type(x) is AIMessage else f"<|user|>\n{x.content}\n"
+        t += f"<|you|>\n<s>{x.content}</s>\n" if type(x) is AIMessage else f"<|user|>\n{x.content}\n"
       return "" if len(t) == 0 else t
     
     def format_docs(docs):
@@ -66,7 +66,8 @@ def custom_chain_with_history(llm, memory):
     #     ]
     # )
     
-    return {"chat_history":prompt_memory, "context":create_vectorstore().as_retriever(search_type="similarity", search_kwargs={"k": 8}) | format_docs, "question": RunnablePassthrough()} | prompt | llm
+    # return {"chat_history":lambda x:, "context":create_vectorstore().as_retriever(search_type="similarity", search_kwargs={"k": 8}) | format_docs, "question": RunnablePassthrough()} | prompt | llm
+    return {"chat_history":lambda x:prompt_memory(x['memory']), "context":itemgetter("question") | create_vectorstore().as_retriever(search_type="similarity", search_kwargs={"k": 8}) | format_docs, "question": lambda x:x['question']} | prompt | llm
 
 class CustomLLM(LLM):
     repo_id : str