Add multi-turn conversation feature

chatbot/Custom_chain.py

@@ -0,0 +1,42 @@
+from langchain.chains.base import Chain
+from langchain.schema import BaseRetriever
+from langchain.llms import BaseLLM
+from langchain.prompts import PromptTemplate
+from pydantic import Field
+from typing import Dict, Any
+
+class MyCustomMemoryRetrievalChain(Chain):
+    """
+    Custom chain cho phép truyền question, memory.
+    Lấy docs từ retriever, trộn với prompt, gọi LLM.
+    """
+
+    llm: BaseLLM = Field(...)
+    retriever: BaseRetriever = Field(...)
+    prompt: PromptTemplate = Field(...)
+    output_key: str = "result"  
+
+    @property
+    def input_keys(self) -> list:
+        return ["question", "memory"]
+
+    @property
+    def output_keys(self) -> list:
+        return [self.output_key]
+
+    def _call(self, inputs: Dict[str, Any], run_manager=None) -> Dict[str, Any]:
+        question = inputs["question"]
+        memory = inputs["memory"]
+
+        docs = self.retriever.get_relevant_documents(question)
+        context = "\n".join(doc.page_content for doc in docs)
+
+        final_prompt = self.prompt.format(
+            question=question,
+            memory=memory,
+            context=context
+        )
+
+        answer = self.llm(final_prompt)
+
+        return {self.output_key: answer}

chatbot/core.py

@@ -1,13 +1,13 @@
 from .llm import gemini_llm  
 from .retrieval import load_vectordb
-from .memory import memory
-from .prompts import chat_prompt, classification_prompt, category_tree_json
-from langchain.chains import ConversationalRetrievalChain
+from .prompts import chat_prompt_no_memory, chat_prompt_memory, classification_prompt, category_tree_json
+from langchain.chains import RetrievalQA
 from .metadata_selfquery import metadata_field_info
 from langchain.retrievers.self_query.base import SelfQueryRetriever
 from langchain.retrievers.self_query.qdrant import QdrantTranslator
-
-
+from .memory import ShortTermMemory
+from .Custom_chain import MyCustomMemoryRetrievalChain
+memory = ShortTermMemory()
 vector_store = load_vectordb()
 
 def classify_query(query):
@@ -21,22 +21,49 @@ retriever = SelfQueryRetriever.from_llm(
     metadata_field_info=metadata_field_info,
     structured_query_translator= QdrantTranslator(metadata_key="metadata"),
     search_type="similarity",
-    search_kwargs={"k": 10, "score_threshold": 0.3}
+    search_kwargs={"k": 10}
 )
 
-qa_chain = ConversationalRetrievalChain.from_llm(
+qa_chain1 = RetrievalQA.from_chain_type(
     llm=gemini_llm, 
     retriever=retriever, 
-    memory=memory,  
     return_source_documents= False,  
-    combine_docs_chain_kwargs={"prompt": chat_prompt},
+    chain_type_kwargs={"prompt": chat_prompt_no_memory},
     output_key="result" 
 )
+qa_chain2 = MyCustomMemoryRetrievalChain(
+    llm= gemini_llm,
+    retriever= retriever,
+    prompt= chat_prompt_memory,
+    output_key="result"
+)
+
+# qa_chain2 = RetrievalQA.from_chain_type(
+#     llm=gemini_llm,
+#     retriever=retriever,
+#     return_source_documents=False,
+#     chain_type_kwargs={
+#         "prompt": chat_prompt_memory,
+#         "document_variable_name": "context"
+#     },
+#     output_key="result"
+# )
+
 
 def get_chat_response(user_input: str) -> str:
-    classified_query = classify_query(user_input)
-    response = qa_chain({"question": classified_query})
+    restriction = classify_query(user_input.strip())
 
-    memory.save_context({"input": classified_query}, {"output": response["result"]})
-    print(classified_query)
+    if memory.related_to_cache(user_input):
+        print("Liên quan tới câu trước")
+        response = qa_chain2({"question": user_input + memory.restrict + restriction, "memory": memory.get_memory_text()})
+        
+    else:
+        memory.reset()
+        classified_query = user_input + restriction
+        memory.restrict = restriction
+        print("Không liên quan tới câu trước")
+        response = qa_chain1({"query": classified_query})
+    print(restriction)
+    memory.add(user_input, response["result"])
+    print(memory.restrict)
     return response["result"]

chatbot/memory.py

@@ -1,3 +1,45 @@
-from langchain.memory import ConversationBufferMemory
+from collections import deque
+from difflib import SequenceMatcher
+from langchain_core.prompts import PromptTemplate
+from langchain.chains import LLMChain
+from .llm import gemini_llm
+
+related_prompt = PromptTemplate(
+    input_variables=["q1", "q2"],
+    template="""
+Bạn là một trợ lý thông minh. Nhiệm vụ của bạn là xác định xem hai câu hỏi hiện tại, có phải là nối tiếp hay cùng mục đích, liên quan tới ngữ cảnh trước đó không:
+
+Câu hỏi hiện tại: {q1}
+Hội thoại trước đó: {q2} (nếu không có gì, trả lời là không)
+
+Hai câu hỏi này có liên quan không? Trả lời ngắn gọn: Có hoặc Không.
+"""
+)
+
+check_related_chain = LLMChain(llm=gemini_llm, prompt=related_prompt)
+
+class ShortTermMemory:
+    def __init__(self, maxlen=3):
+        self.cache = deque(maxlen=maxlen)
+        self.restrict = ""
+
+    def is_similar(self, q1, q2):
+        response = check_related_chain.invoke({"q1": q1, "q2": q2})
+        print(response['text']) 
+        return "có" in response['text'].lower()
+
+
+    def related_to_cache(self, query):
+        return self.is_similar(query, self.get_memory_text())
+        
+    def add(self, query, answer):
+        text_entry = f"Người dùng hỏi: {query}\n Hệ thống trả lời: {answer}"
+        self.cache.append(text_entry)
+
+    def reset(self):
+        self.cache.clear()
+        self.restrict = ""
+    def get_memory_text(self):
+        return "\n".join(self.cache)
+    
 
-memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)

chatbot/prompts.py

@@ -3,7 +3,7 @@ from .llm import gemini_llm
 from langchain_core.prompts import ChatPromptTemplate
 import json
 
-chat_prompt = PromptTemplate(
+chat_prompt_no_memory = PromptTemplate(
     input_variables=["context", "question"],
     template="""
 Bạn là trợ lý AI bán hàng của **Rạng Đông Store**, chuyên hỗ trợ khách hàng tìm kiếm và lựa chọn các sản phẩm chiếu sáng và gia dụng chất lượng cao.
@@ -43,7 +43,7 @@ Bạn là trợ lý AI bán hàng của **Rạng Đông Store**, chuyên hỗ tr
   `Giá`, `Công suất`, `Góc chiếu`, `Độ rọi`.
 - Nếu là **bình giữ nhiệt hoặc phích nước**, hãy nêu rõ:  
   `Giá`, `Dung tích`, `Thời gian giữ nhiệt`.
-
+- Ghi đầy đủ các thông tin khác về thông số kỹ thuật, mô tả sản phẩm
 ---
 
 ### Câu hỏi từ khách hàng:
@@ -229,5 +229,29 @@ classification_prompt = ChatPromptTemplate.from_messages([
                "Hãy đọc câu hỏi của khách hàng và xác định danh mục thích hợp L1, L2, L3"
                "Và giá thấp nhất khách hàng mua, giá cao nhất khách hàng mua"),
     ("human", "Câu hỏi: {query}. Hãy trả về danh mục thích hợp."
-    "Trả lời theo định dạng sau: {query}, L1:... - L2:....")
+    "Trả lời theo định dạng sau: L1(Nếu xác định được):... - L2(Nếu xác định được):.... -L3(Nếu xác định được):... - Giá thấp nhất(Nếu xác định được):... - Giá cao nhất(nếu xác định được):..."
+    "Nếu chỉ có thông tin về giá vẫn trả lời Giá thấp nhất:.. - Giá cao nhât:...."
+    "Nếu không tìm được yếu tố nào thì bỏ trống (không cần ghi yếu tố không tìm thấy ví dụ (L1: ... - Giá thấp nhất 100 000))"
+    "Nếu không tìm thấy danh mục nào phù hợp trả về nội dung sau: '...'")
 ])
+
+
+chat_prompt_memory = PromptTemplate(
+    input_variables=["context", "question", "memory"],
+    template="""
+Bạn là một trợ lý AI bán hàng của Rạng Đông Store, chuyên hỗ trợ khách hàng tìm kiếm và lựa chọn các sản phẩm chiếu sáng và gia dụng chất lượng cao.
+
+Dưới đây là các thông tin bạn cần để trả lời khách hàng:
+
+Lịch sử trò chuyện trước đó với khách hàng:
+{memory}
+
+Thông tin sản phẩm liên quan đến câu hỏi:
+{context}
+
+Câu hỏi hiện tại của khách hàng (tập trung trả lời câu hỏi này, quan trọng nhất):
+{question}
+
+Hãy dựa vào các thông tin trên để đưa ra câu trả lời chính xác thân thiện nhằm hỗ trợ khách hàng một cách hiệu quả nhất.
+"""
+)