Update app.py

app.py

@@ -6,41 +6,58 @@ from langchain.prompts import PromptTemplate
 from langchain.docstore.document import Document
 
 def extract_text_from_pdf(file):
-    doc = fitz.open(stream=file.read(), filetype="pdf")
+    doc = fitz.open(file.name)  # Fix lỗi: dùng tên file thay vì file.read()
     text = ""
     for page in doc:
         text += page.get_text()
     return text
 
-def answer_question(pdf_file, question):
-    if not pdf_file or not question:
-        return "Vui lòng tải file PDF và nhập câu hỏi."
-
-    content = extract_text_from_pdf(pdf_file)
-
-    # Khởi tạo model Ollama
-    llm = Ollama(model="llama3")  # Bạn có thể dùng 'mistral', 'phi', v.v.
-
-    # Prompt đơn giản
-    prompt_template = PromptTemplate.from_template(
-        "Trả lời câu hỏi sau dựa trên tài liệu:\n\n{context}\n\nCâu hỏi: {question}"
-    )
-
-    chain = load_qa_chain(llm, chain_type="stuff", prompt=prompt_template)
-
-    # LangChain yêu cầu list Document
-    docs = [Document(page_content=content)]
-
-    result = chain.run(input_documents=docs, question=question)
-    return result
-
-gr.Interface(
-    fn=answer_question,
-    inputs=[
-        gr.File(label="Tải lên PDF"),
-        gr.Textbox(label="Câu hỏi về nội dung PDF")
-    ],
-    outputs="text",
-    title="PDF Q&A với Ollama",
-    description="Sử dụng mô hình Ollama (như llama3) để trả lời câu hỏi từ file PDF.",
-).launch()
+# Trạng thái lưu lịch sử chat (nếu không dùng PDF)
+chat_history = []
+
+def answer_question(pdf_file, question, history):
+    if not question:
+        return "Vui lòng nhập câu hỏi.", history
+
+    llm = Ollama(model="llama3")
+
+    if pdf_file:  # Nếu có PDF được tải lên
+        content = extract_text_from_pdf(pdf_file)
+        prompt_template = PromptTemplate.from_template(
+            "Trả lời câu hỏi sau dựa trên tài liệu:\n\n{context}\n\nCâu hỏi: {question}"
+        )
+        chain = load_qa_chain(llm, chain_type="stuff", prompt=prompt_template)
+        docs = [Document(page_content=content)]
+        result = chain.run(input_documents=docs, question=question)
+        return result, history
+    else:
+        # Chat tự do (không có PDF)
+        context = "\n".join([f"User: {q}\nAssistant: {a}" for q, a in history[-5:]])  # nhớ 5 lượt gần nhất
+        prompt = f"{context}\nUser: {question}\nAssistant:"
+
+        answer = llm(prompt)
+        history.append((question, answer.strip()))
+        return answer.strip(), history
+
+# Giao diện Gradio nâng cấp
+chatbot = gr.Chatbot()
+with gr.Blocks() as demo:
+    gr.Markdown("## 🧠 PDF Q&A và Chat tự do với Ollama")
+    with gr.Row():
+        pdf_file = gr.File(label="📄 Tải PDF (tùy chọn)")
+    with gr.Row():
+        question = gr.Textbox(label="💬 Nhập câu hỏi", placeholder="Hỏi gì cũng được...")
+    with gr.Row():
+        submit_btn = gr.Button("Gửi câu hỏi")
+    with gr.Row():
+        chatbot_output = chatbot
+
+    state = gr.State([])  # Lưu lịch sử hội thoại
+
+    def respond(pdf, q, hist):
+        answer, updated_history = answer_question(pdf, q, hist)
+        return updated_history, updated_history
+
+    submit_btn.click(respond, inputs=[pdf_file, question, state], outputs=[chatbot_output, state])
+
+demo.launch()