Update llm.py

llm.py

@@ -49,13 +49,16 @@ def create_rag_chain():
     llm = OllamaLLM(model="gemma2:2b")
     
     system_prompt = (
-        "You are an assistant for question-answering tasks. "
-        "Use the following pieces of retrieved context to answer "
-        "the question. If you don't know the answer, say that you "
-        "don't know. Provide a clear and concise answer, using "
-        "up to three sentences."
-        "\n\n"
-        "{context}"
+        """
+<|context|>
+Bạn là một AI trợ lý chuyên cung cấp về sản phẩm bảo hiểm của công ty MB Ageas Life tại Việt Nam.
+Hãy trả lời  chuyên nghiệp, chính xác, cung cấp thông tin bao quát trước, đưa ví dụ rồi mới đặt câu hỏi gợi mở nếu chưa rõ. Tất cả các thông tin cung cấp đều trong phạm vi MBAL. Những câu hỏi không có đủ thông tin thì mời khách hàng đăng ký để nhận tư vấn trên https://www.mbageas.life/
+</s>
+<|user|>
+{query}
+</s>
+ <|assistant|>
+"""
     )
     
     prompt = ChatPromptTemplate.from_messages(
@@ -69,7 +72,9 @@ def create_rag_chain():
     return question_answer_chain
 
 def get_response(rag_chain, vectorstore, question):
-    retriever = vectorstore.as_retriever()
+    retriever = vectorstore.as_retriever(
+    search_kwargs={'k': 5}
+)
     retrieval_chain = create_retrieval_chain(retriever, rag_chain)
     response = retrieval_chain.invoke({"input": question})
     return response["answer"]