chatbot/retrieval.py

from langchain.tools import SerpAPIWrapper
from langchain.retrievers import WikipediaRetriever
from chatbot.llm import gemini_llm  
from chatbot.memory import memory  
from chatbot.prompts import chat_prompt  
from langchain.chains import ConversationalRetrievalChain

# Option 1: Use SerpAPI for Internet Search
search = SerpAPIWrapper()
retriever = search  # Acts as the retriever

# Option 2: Use WikipediaRetriever
# retriever = WikipediaRetriever()

qa_chain = ConversationalRetrievalChain.from_llm(
    llm=gemini_llm, 
    retriever=retriever, 
    memory=memory,  
    return_source_documents=False,  
    combine_docs_chain_kwargs={"prompt": chat_prompt},
    output_key="result"
)

def get_chat_response(user_input: str) -> str:
    response = qa_chain(user_input)

    # Lưu vào bộ nhớ hội thoại
    memory.save_context({"input": user_input}, {"output": response["result"]})

    return response["result"]