app.py

import os
import re
from pypdf import PdfReader
import gradio as gr

from langchain_groq import ChatGroq
from langchain_huggingface import HuggingFaceEmbeddings
from langchain.vectorstores import Chroma
from langchain_core.documents import Document
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain.utilities import DuckDuckGoSearchAPIWrapper
from langchain_community.tools import DuckDuckGoSearchRun , DuckDuckGoSearchResults

embeddings = HuggingFaceEmbeddings(model_name="heydariAI/persian-embeddings")
vector_store = Chroma(embedding_function=embeddings)
text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)

models = ["deepseek-r1-distill-llama-70b", "llama-3.3-70b-versatile", "gemma2-9b-it"]
default_model = models[0]

model_instance = ChatGroq(
    api_key="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",
    model_name=default_model
)

chat_history = []
PRICE_PER_TOKEN = 0.00001

def summarize_chat(model):
    """Generate a short summary of the conversation so far."""
    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 process_file(file_path):
    """Process the uploaded file (PDF or TXT) and index its content."""
    if not file_path:
        return None
    file_extension = os.path.splitext(file_path)[1].lower()
    try:
        if file_extension == ".pdf":
            reader = PdfReader(file_path)
            file_text = "\n".join(page.extract_text() for page in reader.pages)
        elif file_extension == ".txt":
            with open(file_path, "r", encoding="utf-8") as f:
                file_text = f.read()
        else:
            raise ValueError(f"Unsupported file format: {file_extension}")

        file_docs = [Document(page_content=file_text, metadata={"source": "uploaded_file"})]
        file_splits = text_splitter.split_documents(file_docs)
        vector_store.add_documents(file_splits)
        return file_text
    except Exception as e:
        raise RuntimeError(f"Error processing file: {str(e)}")

def answer_query(query, file_path, summarize, tone, model_name, creativity, keywords, language, response_length, welcome_message, exclusion_words):
    """
    Generate an answer by searching indexed content, applying prompt instructions,
    and querying the language model.
    """
    global chat_history, model_instance

    model_instance = ChatGroq(
        api_key="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",
        model_name=model_name
    )
    try:
        if file_path:
            process_file(file_path)

        search_query = f"{keywords} {query}" if keywords else query
        retrieved_docs = vector_store.similarity_search(search_query, k=3)
        knowledge = "\n\n".join(doc.page_content for doc in retrieved_docs)

        tone_prompts = {
            "رسمی": "پاسخ را با لحنی رسمی و مودبانه ارائه کن.",
            "محاوره‌ای": "پاسخ را به صورت دوستانه ارائه کن.",
            "علمی": "پاسخ را با استدلال‌های منطقی ارائه کن.",
            "طنزآمیز": "پاسخ را با لحنی طنزآمیز ارائه کن.",
        }
        tone_instruction = tone_prompts.get(tone, f"پاسخ را به زبان {language} ارائه کن.")
        language_instruction = (f"پاسخ را فقط به زبان {language} ارائه کن و از زبان دیگری استفاده نکن مگر آنکه بخواهی کد بنویسی که در آن صورت فقط از زبان انگلیسی استفاده کن مگر اینکه کاربر از تو درخواست کند از زبان دیگری استفاده بکنی و از زبان چینی استفاده نکن.") if language else ""
        if response_length == "کوتاه":
            length_instruction = "پاسخ را به صورت مختصر ارائه کن."
        elif response_length == "بلند":
            length_instruction = "پاسخ را به صورت مفصل و جامع ارائه کن."
        else:
            length_instruction = ""
        exclusion_instruction = f"از کلمات زیر در پاسخ استفاده نکن: {exclusion_words}" if exclusion_words else ""

        prompt = (
            f"شما ParvizGPT هستید، یک دستیار هوش مصنوعی ساخته شده توسط امیرمهدی پرویز دانشجو دانشگاه صنعتی کرمانشاه "
            f"{tone_instruction} {language_instruction} {length_instruction} {exclusion_instruction}\n\n"
        )

        if welcome_message and not chat_history:
            prompt = f"{welcome_message}\n\n" + prompt

        if chat_history:
            conversation_history = "\n".join([f"پرسش: {q}\nپاسخ: {a}" for q, a in chat_history])
            prompt = f"{conversation_history}\n\n" + prompt

        prompt += f"اطلاعات مرتبط:\n{knowledge}\n\nسوال: {query}\nپاسخ:"

        response = model_instance.invoke(prompt, temperature=creativity)
        cleaned_response = remove_think_sections(response.content)

        chat_history.append((query, cleaned_response))

        total_tokens, price = calculate_price(prompt, cleaned_response)
        summary = summarize_chat(model_instance) if summarize else "خلاصه‌سازی غیرفعال است."

        return cleaned_response, summary, total_tokens, price
    except Exception as e:
        return f"خطا: {str(e)}", "", 0, "0 دلار"

def count_tokens(text):
    """A simple token counter based on whitespace splitting."""
    return len(text.split())

def calculate_price(input_text, output_text):
    """Estimate the total token count and cost."""
    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} دلار"

def remove_think_sections(response_text):
    """Remove any <think>...</think> sections from the response text."""
    return re.sub(r"<think>.*?</think>", "", response_text, flags=re.DOTALL)

def chat_with_bot(query, file, summarize, tone, model_name, creativity, keywords, language, response_length, welcome_message, exclusion_words, chat_history_input):
    """
    Handle a chat message by processing file (if provided), querying the model,
    and updating the chat history.
    """
    file_path = file.name if file else None
    response, summary, total_tokens, price = answer_query(
        query, file_path, summarize, tone, model_name, creativity, keywords,
        language, response_length, welcome_message, exclusion_words
    )
    chat_history_input.append((query, response))
    return chat_history_input, summary, total_tokens, price

def clear_memory():
    """Clear the conversation history."""
    global chat_history
    chat_history = []
    return [], "", 0, "0 دلار"

def search_and_summarize(query: str) -> str:
    """Perform web and news search, then summarize results using the model."""
    try:
        regular_search = DuckDuckGoSearchRun()
        regular_results = regular_search.invoke(query)
        
        news_search = DuckDuckGoSearchResults(backend="news")
        news_results = news_search.invoke(query)
        
        combined_results = f"نتایج جستجو:\n{regular_results}\n\nنتایج اخبار:\n{news_results}"
        
        summary_prompt = (
            "لطفاً نتایج جستجو و اخبار زیر را به صورت خلاصه و منسجم به زبان فارسی ارائه دهید. "
            "خلاصه باید شامل نکات کلیدی از هر دو بخش باشد و به صورت واضح و مختصر نوشته شود.\n\n"
            f"{combined_results}\n\n"
            "خلاصه:"
        )
        
        # Initialize the model
        model = ChatGroq(
            api_key="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",
            model_name=default_model
        )
        
        summary_response = model.invoke(summary_prompt)
        return f"🔍 نتایج کامل:\n{combined_results}\n\n 📝 خلاصه نتایج:\n{remove_think_sections(summary_response.content)}"
    except Exception as e:
        return f"خطا در پردازش جستجو: {str(e)}"

with gr.Blocks() as demo:
    gr.Markdown("# Parviz Mind")
    gr.Markdown("زنده باد")

    with gr.Tabs():

        with gr.TabItem("Chatbot"):
            gr.Markdown("## 🤖 Parviz GPT - چت بات هوش مصنوعی")
            gr.Markdown("**یک فایل (PDF یا TXT) آپلود کنید و سوال خود را بپرسید.**")

            chatbot = gr.Chatbot(label="💬 تاریخچه چت")
            query_input = gr.Textbox(label="❓ سوال خود را وارد کنید")
            summarize_checkbox = gr.Checkbox(label="📌 خلاصه‌ساز را فعال کن")
            submit_button = gr.Button("🚀 ارسال")
            del_button = gr.Button("🗑 پاک کردن حافظه")
            file_input = gr.File(label="📂 آپلود فایل", file_types=[".pdf", ".txt"])

            with gr.Accordion("تنظیمات پیشرفته", open=False):

                summary_output = gr.Textbox(label="📌 خلاصه مکالمه", interactive=False)
                token_count = gr.Textbox(label="🔢 تعداد توکن‌ها", interactive=False)
                token_price = gr.Textbox(label="💰 هزینه تخمینی", interactive=False)

                with gr.Row():
                    model_dropdown = gr.Dropdown(label="🔍 انتخاب مدل", choices=models, value=default_model)
                    tone_dropdown = gr.Dropdown(label="🎭 لحن پاسخ", choices=["رسمی", "محاوره‌ای", "علمی", "طنزآمیز"], value="رسمی")
                    language_dropdown = gr.Dropdown(label="🌐 زبان چت بات", choices=["فارسی", "انگلیسی", "عربی"], value="فارسی")

                with gr.Row():
                    creativity_slider = gr.Slider(label="🎨 خلاقیت (Temperature)", minimum=0.0, maximum=1.0, step=0.1, value=0.7)
                    response_length_dropdown = gr.Dropdown(label="📏 طول پاسخ", choices=["کوتاه", "بلند"], value="بلند")

                keywords_input = gr.Textbox(label="🔑 کلمات کلیدی (اختیاری)")
                welcome_message_input = gr.Textbox(label="👋 پیام خوش آمدگویی (اختیاری)")
                exclusion_words_input = gr.Textbox(label="🚫 کلمات استثنا (اختیاری)")

            del_button.click(
                clear_memory,
                inputs=[],
                outputs=[chatbot, summary_output, token_count, token_price]
            )

            submit_button.click(
                chat_with_bot,
                inputs=[
                    query_input, file_input, summarize_checkbox, tone_dropdown, model_dropdown,
                    creativity_slider, keywords_input, language_dropdown, response_length_dropdown,
                    welcome_message_input, exclusion_words_input, chatbot
                ],
                outputs=[chatbot, summary_output, token_count, token_price]
            )

            query_input.submit(
                chat_with_bot,
                inputs=[
                    query_input, file_input, summarize_checkbox, tone_dropdown, model_dropdown,
                    creativity_slider, keywords_input, language_dropdown, response_length_dropdown,
                    welcome_message_input, exclusion_words_input, chatbot
                ],
                outputs=[chatbot, summary_output, token_count, token_price]
            )

        with gr.TabItem("Web Search"):
            gr.Markdown("## 🌐 Web Search")
            gr.Markdown("سؤالات خود را تایپ کنید:")
            search_query = gr.Textbox(label="عبارت جستجو", placeholder="متنی برای جستجو وارد کنید...")
            search_button = gr.Button("جستجو")
            web_results_output = gr.Textbox(label="نتایج جستجو و خلاصه", lines=10)
            search_button.click(fn=search_and_summarize, inputs=search_query, outputs=web_results_output)
            search_query.submit(fn=search_and_summarize, inputs=search_query, outputs=web_results_output)

demo.launch()