Update app.py

app.py

@@ -1,89 +1,93 @@
 import gradio as gr
-from langchain_community.document_loaders import PyPDFLoader
-from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain_community.document_loaders import PyPDFLoader
-from langchain_community.embeddings import GPT4AllEmbeddings
-from langchain_community.vectorstores import FAISS
-from langchain.schema.runnable import RunnablePassthrough
-# from langchain.prompts import ChatPromptTemplate
-# from langchain_community.chat_models import ChatOllama
-from prompt_template import *
+from langchain_community.utilities import DuckDuckGoSearchAPIWrapper
+from langchain_community.tools import DuckDuckGoSearchRun
+from langchain.prompts import PromptTemplate, MessagesPlaceholder
+from langchain_core.output_parsers import JsonOutputParser, StrOutputParser
+from typing_extensions import TypedDict
+from langchain_core.prompts import ChatPromptTemplate
+import pickle
+from langchain_core.messages import HumanMessage, AIMessage
 from langgraph.graph import END, StateGraph
-from langchain_community.llms import LlamaCpp
-
-# local_llm = 'aleni_ox'
+from huggingface_hub import hf_hub_download
 
-# llm = ChatOllama(model=local_llm,
-#                  keep_alive="3h", 
-#                  max_tokens=512,  
-#                  temperature=0,
-#                 # callbacks=[StreamingStdOutCallbackHandler()]
-#                 )
+from langchain_community.llms import LlamaCpp
+wrapper = DuckDuckGoSearchAPIWrapper(max_results=5)
+web_search_tool = DuckDuckGoSearchRun(api_wrapper=wrapper)
 
 llm = LlamaCpp(
     model_path="Llama-3.1-8B-Instruct.Q5_K_M.gguf",
     temperature=0,
     max_tokens=512,
+    n_ctx = 8000,
     top_p=1,
     # callback_manager=callback_manager,
     verbose=True,  # Verbose is required to pass to the callback manager
 )
+chat_history = list()
+try:
+    with open("template.pkl", 'rb') as file:
+        template_abox = pickle.load(file)
+except:
+    hf_hub_download(repo_id="linl03/dataAboxChat",local_dir="./", filename="template.pkl", repo_type="dataset")
+    with open("./template.pkl", 'rb') as file:
+        template_abox = pickle.load(file)
+
+router_prompt = PromptTemplate(
+    template=template_abox["router_template"],
+    input_variables=["question"],
+)
+generate_prompt = ChatPromptTemplate.from_messages(
+    [
+        (
+            "system",
+            template_abox["system_prompt"],
+        ),
+        MessagesPlaceholder(variable_name="chat_history"),
+        ("human", "{question}"),
+    ]
+)
+query_prompt = PromptTemplate(
+    template=template_abox["query_template"],
+    input_variables=["question"],
+)
+
+remind_prompt = PromptTemplate(
+    template=template_abox["schedule_template"],
+    input_variables=["time"],
+)
 
 question_router = router_prompt | llm | JsonOutputParser()
 generate_chain = generate_prompt | llm | StrOutputParser()
 query_chain = query_prompt | llm | JsonOutputParser()
-llm_chain = nomalqa_prompt | llm | StrOutputParser()
+# llm_chain = nomalqa_prompt | llm | StrOutputParser()
 
-def generate(state):
-    """
-    Generate answer
 
-    Args:
-        state (dict): The current graph state
+class State(TypedDict):
 
-    Returns:
-        state (dict): New key added to state, generation, that contains LLM generation
-    """
+    question : str
+    generation : str
+    search_query : str
+    context : str
     
-    print("Step: Đang tạo câu trả lời từ những gì tìm được")
+def generate(state):
+
+    print("Step: Đang tạo câu trả lời")
     question = state["question"]
     context = state["context"]
     # return question, context
     return {'question': question, 'context': context}
 
-    # respon=''
-    
-    # for chunk in generate_chain.stream({"context": context, "question": question}):
-    #     respon += chunk
-    #     print(chunk, end="", flush=True)
 
 def transform_query(state):
-    """
-    Transform user question to web search
-
-    Args:
-        state (dict): The current graph state
 
-    Returns:
-        state (dict): Appended search query
-    """
-    
     print("Step: Tối ưu câu hỏi của người dùng")
     question = state['question']
     gen_query = query_chain.invoke({"question": question})
+    print(gen_query)
     search_query = gen_query["query"]
     return {"search_query": search_query}
 
 def web_search(state):
-    """
-    Web search based on the question
-
-    Args:
-        state (dict): The current graph state
-
-    Returns:
-        state (dict): Appended web results to context
-    """
 
     search_query = state['search_query']
     print(f'Step: Đang tìm kiếm web cho: "{search_query}"')
@@ -94,112 +98,72 @@ def web_search(state):
     return {"context": search_result}
 
 def route_question(state):
-    """
-    route question to web search or generation.
-
-    Args:
-        state (dict): The current graph state
-
-    Returns:
-        str: Next node to call
-    """
 
     print("Step: Routing Query")
     question = state['question']
     output = question_router.invoke({"question": question})
     print('Lựa ch���n của AI là: ', output)
-    if output == "web_search":
+    if output['choice'] == "web_search":
         # print("Step: Routing Query to Web Search")
         return "websearch"
-    elif output == 'generate':
+    elif output['choice'] == 'generate':
         # print("Step: Routing Query to Generation")
         return "generate"
 
-workflow = StateGraph(State)
-workflow.add_node("websearch", web_search)
-workflow.add_node("transform_query", transform_query)
-workflow.add_node("generate", generate)
-
-# Build the edges
-workflow.set_conditional_entry_point(
-    route_question,
-    {
-        "websearch": "transform_query",
-        "generate": "generate",
-    },
-)
-workflow.add_edge("transform_query", "websearch")
-workflow.add_edge("websearch", "generate")
-workflow.add_edge("generate", END)
+def Agents():
+    workflow = StateGraph(State)
+    workflow.add_node("websearch", web_search)
+    workflow.add_node("transform_query", transform_query)
+    workflow.add_node("generate", generate)
 
-# Compile the workflow
-local_agent = workflow.compile()
+    # Build the edges
+    workflow.set_conditional_entry_point(
+        route_question,
+        {
+            "websearch": "transform_query",
+            "generate": "generate",
+        },
+    )
+    workflow.add_edge("transform_query", "websearch")
+    workflow.add_edge("websearch", "generate")
+    workflow.add_edge("generate", END)
 
-def run_agent(query):
-    local_agent.invoke({"question": query})
-    print("=======")
+    # Compile the workflow
+    return workflow.compile()
 
-def QA(question: str, history: list, type: str):
-    if 'Agent' in type:
+def QA(question: str, history: list):
+    # print(question.text, question.files, history, type)
+    if not question.files:
+        local_agent = Agents()
         gr.Info("Đang tạo câu trả lời!")
-        respon = ''
-        # print(question)
-        output = local_agent.invoke({"question": question})
-        # print(output)
+        response = ''
+        output = local_agent.invoke({"question": question.text})
         context = output['context']
         questions = output['question']
-        for chunk in generate_chain.stream({"context": context, "question": questions}):
-            respon += chunk
-            print(chunk, end="", flush=True)
-            yield respon
-    else:
-        gr.Info("Đang tạo câu trả lời!")
-        print(question, history)
-        respon = ''
-        for chunk in llm_chain.stream(question):
-            respon += chunk
+        for chunk in generate_chain.stream({"context": context, "question": questions, "chat_history": chat_history}):
+            response += chunk
             print(chunk, end="", flush=True)
-            yield respon
-
-def create_db(doc: str) -> str:
-    loader = PyPDFLoader(doc)
-    text_splitter = RecursiveCharacterTextSplitter(chunk_size=200, chunk_overlap=40)
-
-    chunked_documents = loader.load_and_split(text_splitter)
-    embedding_model = GPT4AllEmbeddings(model_name="all-MiniLM-L6-v2.gguf2.f16.gguf", gpt4all_kwargs={'allow_download': 'True'})
-    db = FAISS.from_documents(chunked_documents, embedding_model)
-    gr.Info("Đã tải lên dữ liệu từ PDF!")
-    
-    retriever = db.as_retriever(
-        search_type="similarity",
-        search_kwargs= {"k": 3}
-    )
-    llm_chain = (
-        {
-            "context": retriever, 
-            "question": RunnablePassthrough()}
-        | nomaldoc_prompt
-        | llm
-    )
-
+            yield response
+                
+        chat_history.append(HumanMessage(content=question.text))
+        chat_history.append(AIMessage(content=response))
 
-with gr.Blocks(fill_height=True) as demo:
-    with gr.Row(equal_height=True):
+    else:
+        print(question.files)
+        # all_loaders = []
+        # for pdf in question.files:
+        #     pdf_loader = PyPDFLoader("data/Đắc Nhân Tâm.pdf")
         
-        with gr.Column(scale=1):
-
-            democ2 = gr.Interface(
-                create_db,
-                [gr.File(file_count='single')],
-                None,
-                )
-        with gr.Column(scale=2):
-            democ1 = gr.ChatInterface(
-                QA,
-                additional_inputs=[gr.Dropdown(["None", "Agent", "Doc"], label="Type", info="Chọn một kiểu chat!"),]
-                
-                
-            )
+
+democ1 = gr.ChatInterface(
+    QA,
+    additional_inputs=[
+        ],
+    fill_height=True,
+    multimodal=True,
+    title="Box Chat(Agent)",
+    type="tuples"
+)
 
 if __name__ == "__main__":
-    demo.launch()
+    democ1.launch()