Update app.py

app.py

@@ -1,183 +1,146 @@
-import time
-import logging
-import gradio as gr
 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__)
-
-groq_api_key = os.environ.get("GROQ_API_KEY")
-hf_api_key = os.environ.get("HF_API_KEY")
-
-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.")
-
-client = Groq(api_key=groq_api_key)
-hf_token = hf_api_key
+import re
+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
+
 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")
 
-DATASET_NAME = "chat_history"
-try:
-    dataset = load_dataset(DATASET_NAME, use_auth_token=hf_token)
-except Exception:
-    dataset = Dataset.from_dict({"Timestamp": [], "User": [], "ParvizGPT": []})
+chat_history = [] 
 
-def save_chat_to_dataset(user_message, bot_message):
-    try:
-        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:
-        logger.error(f"Error saving chat history to dataset: {e}")
+PRICE_PER_TOKEN = 0.00001  
 
-def process_pdf_with_langchain(pdf_path):
-    try:
-        loader = PyPDFLoader(pdf_path)
-        documents = loader.load()
+def count_tokens(text):
+    """تخمین تعداد توکن‌های متن."""
+    return len(text.split())
 
-        text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
-        split_documents = text_splitter.split_documents(documents)
+def calculate_price(input_text, output_text):
+    """محاسبه هزینه بر اساس تعداد توکن‌ها."""
+    input_tokens = count_tokens(input_text)
+    output_tokens = count_tokens(output_text)
+    total_tokens = input_tokens + output_tokens
+    total_price = total_tokens * PRICE_PER_TOKEN
+    return total_tokens, f"{total_price:.6f} دلار"
 
-        vectorstore = FAISS.from_documents(split_documents, embeddings)
-        retriever = vectorstore.as_retriever(search_kwargs={"k": 3})
-        return retriever
-    except Exception as e:
-        logger.error(f"Error processing PDF: {e}")
-        raise
+def process_file(file_path):
+    """پردازش فایل و بازگرداندن محتوای آن."""
+    if not file_path:
+        return None
+
+    file_extension = os.path.splitext(file_path)[1].lower()
 
-def generate_response(query, memory, retriever=None, use_pdf_context=False, tone="friendly"):
     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", "")
-        
-        tone_instruction = ""
-        if tone == "friendly":
-            tone_instruction = "Please respond in a friendly and informal tone."
-        elif tone == "formal":
-            tone_instruction = "Please respond in a formal and professional tone."
-        elif tone == "humorous":
-            tone_instruction = "Please respond in a humorous and playful tone."
-        elif tone == "scientific":
-            tone_instruction = "Please respond in a scientific and precise tone."
-
-        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.
-            {tone_instruction}
-            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 = "Processing..."
-        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()
-                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
+        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"فرمت فایل پشتیبانی نمی‌شود: {file_extension}")
     except Exception as e:
-        logger.error(f"Error generating response: {e}")
-        return f"Error: {e}", memory
+        raise RuntimeError(f"خطا در پردازش فایل: {str(e)}")
+
+
+def remove_think_sections(response_text):
+    """حذف بخش‌های که با <think> شروع و با </think> تمام می‌شوند."""
+
+    cleaned_text = re.sub(r"<think>.*?</think>", "", response_text, flags=re.DOTALL)
+    return cleaned_text
 
-def summarize_chat_history(chat_history):
+def answer_query(query, file_path, summarize, tone):
+    """پاسخ به سوالات کاربر با تنظیم لحن و محاسبه هزینه توکن."""
+    global chat_history
     try:
-        chat_text = "\n".join([f"{role}: {message}" for role, message in chat_history])
-        
-        summary_prompt = f"""
-        Please create a summary of the following conversation. The summary should include key points and details:
-        {chat_text}
-        """
-        
-        chat_completion = client.chat.completions.create(
-            messages=[{"role": "user", "content": summary_prompt}],
-            model="deepseek-r1-distill-llama-70b"
+        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)
+
+        tone_prompts = {
+            "رسمی": "پاسخ را با لحنی رسمی و مودبانه ارائه کن.",
+            "محاوره‌ای": "پاسخ را به صورت دوستانه و غیررسمی ارائه کن.",
+            "علمی": "پاسخ را با ذکر منابع علمی و استدلال‌های منطقی ارائه کن.",
+            "طنزآمیز": "پاسخ را با لحنی طنزآمیز و سرگرم‌کننده ارائه کن.",
+        }
+        tone_instruction = tone_prompts.get(tone, "پاسخ را به زبان فارسی ارائه کن.")
+
+        prompt = (
+            f"شما ParvizGPT هستید، یک دستیار هوش مصنوعی که توسط امیر مهدی پرویز ساخته شده است. "
+            f"همیشه به فارسی پاسخ دهید. {tone_instruction} "
+            f"\n\nاطلاعات مرتبط:\n{knowledge}\n\nسوال: {query}\nپاسخ:"
         )
-        summary = chat_completion.choices[0].message.content.strip()
-        return summary
-    except Exception as e:
-        logger.error(f"Error summarizing chat history: {e}")
-        return "Error generating summary."
 
-def gradio_interface(user_message, chat_box, memory, pdf_file=None, use_pdf_context=False, tone="friendly", summarize_chat=False):
-    global retriever
+        response = model.invoke(prompt)
+        response_text = response.content
 
-    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
+        cleaned_response = remove_think_sections(response_text)
 
-    chat_box.append(("You", user_message))
-    chat_box.append(("ParvizGPT", "Processing..."))
-    response, memory = generate_response(user_message, memory, retriever=retriever, use_pdf_context=use_pdf_context, tone=tone)
-    
-    chat_box[-1] = ("ParvizGPT", response)
-    
-    save_chat_to_dataset(user_message, response)
-    
-    if summarize_chat:
-        summary = summarize_chat_history(chat_box)
-        chat_box.append(("System", f"Summary of the conversation:\n{summary}"))
-    
-    return chat_box, memory
-
-def clear_memory(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)
-    tone = gr.Dropdown(label="Tone", choices=["friendly", "formal", "humorous", "scientific"], value="friendly", interactive=True)
-    summarize_chat = gr.Checkbox(label="Show conversation summary", 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")
+        chat_history.append((query, cleaned_response))  
+
+        total_tokens, price = calculate_price(prompt, cleaned_response)
+
+        summary = summarize_chat() if summarize else "خلاصه‌سازی غیرفعال است."
+
+        return cleaned_response, summary, total_tokens, price
+
+    except Exception as e:
+        return f"خطا: {str(e)}", "", 0, "0 دلار"
+
+def summarize_chat():
+    """خلاصه‌سازی مکالمات اخیر."""
+    chat_text = "\n".join([f"پرسش: {q}\nپاسخ: {a}" for q, a in chat_history])  
+    summary_prompt = f"یک خلاصه کوتاه و دقیق از مکالمه زیر ارائه کن:\n\n{chat_text}\n\nخلاصه:"
+    summary_response = model.invoke(summary_prompt)
+    return summary_response.content
+
+def chat_with_bot(query, file, summarize, tone):
+    """رابط Gradio برای چت."""
+    file_path = file.name if file else None
+    response, summary, total_tokens, price = answer_query(query, file_path, summarize, tone)
+    return response, summary, total_tokens, price
+
+with gr.Blocks() as demo:
+    gr.Markdown("## 🤖 Parviz GPT")
+    gr.Markdown("**یک فایل (PDF یا TXT) آپلود کنید و سوال خود را بپرسید.**")
     
-    memory_state = gr.State(ConversationBufferMemory())
+    with gr.Column():
+
+        chat_output = gr.Textbox(label="📝 تاریخچه چت", interactive=False, lines=10)
+        summary_output = gr.Textbox(label="📌 خلاصه مکالمه", interactive=False)
+
+        query_input = gr.Textbox(label="❓ سوال خود را وارد کنید", placeholder="مثلاً: کی تو را ساخته است؟")
+
+        with gr.Row():
+            summarize_checkbox = gr.Checkbox(label="📌 خلاصه‌ساز را فعال کن")
+            submit_button = gr.Button("🚀 ارسال")
+            tone_dropdown = gr.Dropdown(label="🎭 انتخاب لحن پاسخ", choices=["رسمی", "محاوره‌ای", "علمی", "طنزآمیز"], value="رسمی")
+
+
+    with gr.Row():
+        token_count = gr.Textbox(label="🔢 تعداد توکن‌ها", interactive=False)
+        token_price = gr.Textbox(label="💰 هزینه تخمینی", interactive=False)
+
+    with gr.Row():
+        file_input = gr.File(label="📂 فایل خود را آپلود کنید", file_types=[".pdf", ".txt"])
+
+    query_input.submit(fn=chat_with_bot, 
+                        inputs=[query_input, file_input, summarize_checkbox, tone_dropdown], 
+                        outputs=[chat_output, summary_output, token_count, token_price])
     
-    submit_btn.click(gradio_interface, inputs=[user_message, chat_box, memory_state, pdf_file, use_pdf_context, tone, summarize_chat], outputs=[chat_box, memory_state])
-    user_message.submit(gradio_interface, inputs=[user_message, chat_box, memory_state, pdf_file, use_pdf_context, tone, summarize_chat], outputs=[chat_box, memory_state])
-    clear_memory_btn.click(clear_memory, inputs=[memory_state], outputs=[chat_box, memory_state])
+    submit_button.click(fn=chat_with_bot, 
+                        inputs=[query_input, file_input, summarize_checkbox, tone_dropdown], 
+                        outputs=[chat_output, summary_output, token_count, token_price])
 
-interface.launch()
+demo.launch()