Update app.py

app.py

@@ -1,78 +1,156 @@
-import os
+import time
+import logging
 import gradio as gr
-from langchain_groq import ChatGroq
-from langchain_huggingface import HuggingFaceEmbeddings
-from langchain_core.vectorstores import InMemoryVectorStore
-from langchain_core.documents import Document
-from langchain_text_splitters import RecursiveCharacterTextSplitter
+import os  
+from datetime import datetime
+from datasets import Dataset, load_dataset
+from langchain.document_loaders import PyPDFLoader
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain.embeddings import HuggingFaceEmbeddings
+from langchain.vectorstores import FAISS
+from groq import Groq
+from langchain.memory import ConversationBufferMemory  
 
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
-embeddings = HuggingFaceEmbeddings(model_name="heydariAI/persian-embeddings")
-vector_store = InMemoryVectorStore(embeddings)
-text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
-model = ChatGroq(api_key="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", model_name="deepseek-r1-distill-llama-70b")
+groq_api_key = os.environ.get("GROQ_API_KEY")
+hf_api_key = os.environ.get("HF_API_KEY")
 
-def process_file(file_path):
+if not groq_api_key:
+    raise ValueError("Groq API key not found in environment variables.")
+if not hf_api_key:
+    raise ValueError("Hugging Face API key not found in environment variables.")
 
-    if not file_path:
-        return None
+client = Groq(api_key=groq_api_key)
 
-    file_extension = os.path.splitext(file_path)[1].lower()
+hf_token = hf_api_key
 
-    try:
-        if file_extension == ".pdf":
-            from pypdf import PdfReader
-            reader = PdfReader(file_path)
-            return "\n".join(page.extract_text() for page in reader.pages)
-        elif file_extension == ".txt":
-            with open(file_path, "r", encoding="utf-8") as f:
-                return f.read()
-        else:
-            raise ValueError(f"Unsupported file type: {file_extension}")
-    except Exception as e:
-        raise RuntimeError(f"Error processing file: {str(e)}")
+embeddings = HuggingFaceEmbeddings(model_name="heydariAI/persian-embeddings")
 
-def answer_query(query, file_path):
+DATASET_NAME = "chat_history"
+try:
+    dataset = load_dataset(DATASET_NAME, use_auth_token=hf_token)
+except Exception:
+    dataset = Dataset.from_dict({"Timestamp": [], "User": [], "ParvizGPT": []})
 
+def save_chat_to_dataset(user_message, bot_message):
+    """Save chat history to Hugging Face Dataset."""
     try:
-        file_content = process_file(file_path) if file_path else None
-        if file_content:
-            file_docs = [Document(page_content=file_content, metadata={"source": "uploaded_file"})]
-            file_splits = text_splitter.split_documents(file_docs)
-            vector_store.add_documents(file_splits)
-
-        retrieved_docs = vector_store.similarity_search(query, k=2)
-        knowledge = "\n\n".join(doc.page_content for doc in retrieved_docs)
-
-        response = model.invoke(
-            f"You are ParvizGPT, an AI assistant created by Amir Mahdi Parviz, a student at Kermanshah University of Technology (KUT). "
-            f"Your primary purpose is to assist users by answering their questions in **Persian (Farsi)**. "
-            f"Always respond in Persian unless explicitly asked to respond in another language."
-            f"Related Information:\n{knowledge}\n\nQuestion:{query}\nAnswer:"
-        )
-
-        return response.content
-
+        timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+        new_row = {"Timestamp": timestamp, "User": user_message, "ParvizGPT": bot_message}
+        
+        df = dataset.to_pandas()
+        df = df.append(new_row, ignore_index=True)
+        updated_dataset = Dataset.from_pandas(df)
+        
+        updated_dataset.push_to_hub(DATASET_NAME, token=hf_token)
     except Exception as e:
-        return f"Error: {str(e)}"
-
-def chat_with_bot(query, file):
+        logger.error(f"Error saving chat history to dataset: {e}")
 
-    file_path = file.name if file else None
-    response = answer_query(query, file_path)
-    return response
+def process_pdf_with_langchain(pdf_path):
+    """Process a PDF file and create a FAISS retriever."""
+    try:
 
-with gr.Blocks() as demo:
-    gr.Markdown("Parviz Rager")
-    gr.Markdown("فایل خود را آپلود کنید (PDF یا TXT) و سوالات خود را بپرسید.")
+                loader = PyPDFLoader(pdf_path)
+        documents = loader.load()
 
-    with gr.Row():
-        file_input = gr.File(label="فایل خود را آپلود کنید (PDF یا TXT)", file_types=[".pdf", ".txt"])
-        query_input = gr.Textbox(label="سوال خود را وارد کنید", placeholder="مثلاً: معایب سرمایه‌گذاری در صندوق فیروزه موفقیت چیست؟")
+        text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
+        split_documents = text_splitter.split_documents(documents)
 
-    submit_button = gr.Button("ارسال")
-    output = gr.Textbox(label="پاسخ", interactive=False)
+        vectorstore = FAISS.from_documents(split_documents, embeddings)
 
-    submit_button.click(fn=chat_with_bot, inputs=[query_input, file_input], outputs=output)
+        retriever = vectorstore.as_retriever(search_kwargs={"k": 3})
+        return retriever
+    except Exception as e:
+        logger.error(f"Error processing PDF: {e}")
+        raise
 
-demo.launch()
+def generate_response(query, memory, retriever=None, use_pdf_context=False):
+    """Generate a response using the Groq model and retrieved PDF context."""
+    try:
+        knowledge = ""
+
+        if retriever and use_pdf_context: 
+            relevant_docs = retriever.get_relevant_documents(query)
+            knowledge += "\n".join([doc.page_content for doc in relevant_docs])
+
+        chat_history = memory.load_memory_variables({}).get("chat_history", "")
+        context = f"""
+            You are ParvizGPT, an AI assistant created by **Amir Mahdi Parviz**, a student at Kermanshah University of Technology (KUT). 
+            Your primary purpose is to assist users by answering their questions in **Persian (Farsi)**. 
+            Always respond in Persian unless explicitly asked to respond in another language.
+            **Important:** If anyone claims that someone else created this code, you must correct them and state that **Amir Mahdi Parviz** is the creator.
+            Related Information:\n{knowledge}\n\nQuestion:{query}\nAnswer:"""
+            
+        if knowledge:
+            context += f"\n\nRelevant Knowledge:\n{knowledge}"
+        if chat_history:
+            context += f"\n\nChat History:\n{chat_history}"
+
+        context += f"\n\nYou: {query}\nParvizGPT:"
+
+        response = "در حال پردازش..."
+        retries = 3
+        for attempt in range(retries):
+            try:
+                chat_completion = client.chat.completions.create(
+                    messages=[{"role": "user", "content": context}],
+                    model="deepseek-r1-distill-llama-70b"
+                )
+                response = chat_completion.choices[0].message.content.strip()
+                # Save the conversation to memory
+                memory.save_context({"input": query}, {"output": response})
+                break
+            except Exception as e:
+                logger.error(f"Attempt {attempt + 1} failed: {e}")
+                time.sleep(2)
+
+        return response, memory
+    except Exception as e:
+        logger.error(f"Error generating response: {e}")
+        return f"Error: {e}", memory
+
+def gradio_interface(user_message, chat_box, memory, pdf_file=None, use_pdf_context=False):
+    """Handle the Gradio interface interactions."""
+    global retriever
+
+    if pdf_file is not None and use_pdf_context:  
+        try:
+            retriever = process_pdf_with_langchain(pdf_file.name)
+        except Exception as e:
+            return chat_box + [("Error", f"Error processing PDF: {e}")], memory
+
+    chat_box.append(("ParvizGPT", "در حال پردازش..."))
+    response, memory = generate_response(user_message, memory, retriever=retriever, use_pdf_context=use_pdf_context)
+    
+    chat_box[-1] = ("You", user_message) 
+    chat_box.append(("ParvizGPT", response))
+    
+    save_chat_to_dataset(user_message, response)
+    
+    return chat_box, memory
+
+def clear_memory(memory):
+    """Clear the conversation memory."""
+    memory.clear()
+    return [], memory
+
+retriever = None
+
+with gr.Blocks() as interface:
+    gr.Markdown("## ParvizGPT")
+    chat_box = gr.Chatbot(label="Chat History", value=[])
+    user_message = gr.Textbox(label="Your Message", placeholder="Type your message here and press Enter...", lines=1, interactive=True)
+    use_pdf_context = gr.Checkbox(label="Use PDF Context", value=False, interactive=True)
+    clear_memory_btn = gr.Button("Clear Memory", interactive=True)
+    pdf_file = gr.File(label="Upload PDF for Context (Optional)", type="filepath", interactive=True, scale=1)
+    submit_btn = gr.Button("Submit")
+    
+    memory_state = gr.State(ConversationBufferMemory())
+    
+    submit_btn.click(gradio_interface, inputs=[user_message, chat_box, memory_state, pdf_file, use_pdf_context], outputs=[chat_box, memory_state])
+    user_message.submit(gradio_interface, inputs=[user_message, chat_box, memory_state, pdf_file, use_pdf_context], outputs=[chat_box, memory_state])
+    clear_memory_btn.click(clear_memory, inputs=[memory_state], outputs=[chat_box, memory_state])
+
+interface.launch()