Update app.py

app.py

@@ -1,110 +1,118 @@
 import os
-import shutil
-import logging
 import gradio as gr
-from langchain_community.document_loaders import PyPDFLoader
+import logging
+from langchain.vectorstores import FAISS
+from langchain_community.storage import InMemoryDocstore
+from langchain_community.docstore import Document
 from langchain_openai import OpenAIEmbeddings
-from langchain_community.vectorstores import FAISS
-from langchain_community.chat_models import ChatOpenAI
+from langchain.schema import Document as LangchainDocument
+from langchain.chat_models import ChatOpenAI
 
 # تنظیمات لاگ‌گیری
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# مسیرهای مهم
-UPLOAD_FOLDER = "uploaded_files"
+# مسیر دیتابیس برداری و پوشه آپلود
 DATABASE_FILE = "/home/user/app/vector_database"
+UPLOAD_FOLDER = "/home/user/app/uploaded_files"
+
+# ایجاد پوشه در صورت عدم وجود
 os.makedirs(UPLOAD_FOLDER, exist_ok=True)
 
-# بارگذاری دیتابیس برداری
+# بارگذاری دیتابیس
 def load_database():
-    embeddings = OpenAIEmbeddings(openai_api_key=os.getenv("My_huggingface_key"))
-    if os.path.exists(DATABASE_FILE):
-        logger.info(f"✅ دیتابیس بارگذاری شد: {DATABASE_FILE}")
-        return FAISS.load_local(DATABASE_FILE, embeddings)
-    else:
-        logger.warning(f"❌ دیتابیس وجود ندارد. ایجاد دیتابیس جدید...")
-        return FAISS.from_texts([""], embeddings)
-
-def save_database(vector_db):
-    vector_db.save_local(DATABASE_FILE)
-    logger.info(f"✅ دیتابیس ذخیره شد: {DATABASE_FILE}")
-
-global_vector_db = load_database()
-
-# پردازش PDF و اضافه به دیتابیس
-def process_and_store_pdfs(file_paths):
-    texts, metadata = [], []
-    for file_path in file_paths:
-        if file_path.endswith(".pdf"):
-            loader = PyPDFLoader(file_path)
-            documents = loader.load()
-            for doc in documents:
-                texts.append(doc.page_content)
-                metadata.append({"source": file_path, "page": doc.metadata["page_number"]})
-            logger.info(f"✅ پردازش {len(documents)} صفحه از {file_path}")
-
-    if texts:
-        embeddings = OpenAIEmbeddings(openai_api_key=os.getenv("My_huggingface_key"))
-        vector_db = FAISS.from_texts(texts, embeddings, metadatas=metadata)
-        return vector_db
-    return None
-
-# پاسخ‌گویی با نمایش منبع
-def chat_with_doc(query):
-    global global_vector_db
-    if not global_vector_db:
-        return "<p style='color:red;'>❌ لطفاً یک فایل مرتبط آپلود کنید.</p>"
-    
-    retriever = global_vector_db.as_retriever(search_kwargs={"k": 5})
-    docs = retriever.get_relevant_documents(query)
-
-    if not docs:
-        return "<p style='color:red;'>هیچ اطلاعات مرتبطی یافت نشد.</p>"
-    
-    context = "\n\n".join([f"[صفحه {doc.metadata['page']}] {doc.page_content}" for doc in docs])
-    sources = "\n".join([f"📄 صفحه {doc.metadata['page']} - {doc.metadata['source']}" for doc in docs])
-    
-    llm = ChatOpenAI(model_name="gpt-3.5-turbo", openai_api_key=os.getenv("My_huggingface_key"))
-    prompt = f"متن زیر از اسناد بازیابی شده است و باید فقط از آن برای پاسخ به سوال استفاده شود:\n\n{context}\n\nسوال: {query}\n\nلطفاً به سوال بر اساس اطلاعات موجود پاسخ دهید و منابع را نادیده نگیرید."
-    response = llm.predict(prompt)
-
-    return f"""
-    <div style='max-height: 400px; overflow-y: auto; border: 1px solid #ddd; padding: 10px;'>
-        <h3 style='color:blue;'>پاسخ:</h3>
-        <p>{response}</p>
-        <h4 style='color:green;'>📌 منابع:</h4>
-        <pre>{sources}</pre>
-    </div>
-    """
-
-# ذخیره فایل و به‌روزرسانی دیتابیس
+    try:
+        if os.path.exists(DATABASE_FILE):
+            embeddings = OpenAIEmbeddings(openai_api_key=os.getenv("My_huggingface_key"))
+            vector_db = FAISS.load_local(DATABASE_FILE, embeddings)
+            logger.info(f"✅ دیتابیس بارگذاری شد: {DATABASE_FILE}")
+            return vector_db
+        else:
+            logger.warning(f"❌ دیتابیس در مسیر `{DATABASE_FILE}` وجود ندارد. ایجاد دیتابیس جدید...")
+
+            embeddings = OpenAIEmbeddings(openai_api_key=os.getenv("My_huggingface_key"))
+            index = FAISS.new_index(embeddings)
+            docstore = InMemoryDocstore()
+            index_to_docstore_id = {}
+
+            return FAISS(index, docstore, index_to_docstore_id, embeddings)
+
+    except Exception as e:
+        logger.error(f"❌ خطا در بارگذاری دیتابیس: {e}")
+        return None
+
+# ذخیره فایل و پردازش آن
 def save_uploaded_file(file):
-    file_name = os.path.basename(file)
-    file_path = os.path.join(UPLOAD_FOLDER, file_name)
-    shutil.copy(file, file_path)
-
-    global global_vector_db
-    new_vector_db = process_and_store_pdfs([file_path])
-    if new_vector_db:
-        global_vector_db = new_vector_db
-        save_database(global_vector_db)
-        return f"✅ فایل ذخیره شد و به دیتابیس اضافه شد: {file_path}"
-    return "❌ خطا در پردازش فایل."
-
-# رابط کاربری Gradio
+    try:
+        if not file:
+            return "❌ فایلی انتخاب نشده است."
+        
+        file_name = os.path.basename(file.name)
+        file_path = os.path.join(UPLOAD_FOLDER, file_name)
+
+        with open(file_path, "wb") as f:
+            f.write(file.read())
+        
+        logger.info(f"✅ فایل ذخیره شد: {file_path}")
+
+        global global_vector_db
+        new_vector_db = process_and_store_pdfs([file_path])
+        if new_vector_db:
+            global_vector_db = new_vector_db
+            save_database(global_vector_db)
+            return f"✅ فایل پردازش و ذخیره شد: {file_path}"
+        else:
+            return "❌ خطا در پردازش فایل."
+
+    except Exception as e:
+        logger.error(f"❌ خطا در ذخیره فایل: {e}")
+        return f"❌ خطایی رخ داد: {e}"
+
+# دریافت پاسخ از دیتابیس
+def chat_with_doc(query):
+    try:
+        global global_vector_db
+        if not global_vector_db:
+            return "❌ لطفاً یک فایل مرتبط آپلود کنید."
+
+        if not query.strip():
+            return "❌ لطفاً سوال خود را وارد کنید."
+
+        retriever = global_vector_db.as_retriever(search_kwargs={"k": 5})
+        docs = retriever.get_relevant_documents(query)
+
+        logger.info(f"📄 تعداد اسناد بازیابی شده: {len(docs)}")
+
+        context = "\n\n".join([f"📜 {doc.page_content[:500]}...\n🔗 **منبع:** {doc.metadata.get('source', 'نامشخص')}, صفحه {doc.metadata.get('page', 'نامشخص')}" for doc in docs])
+
+        if not context:
+            return "هیچ اطلاعات مرتبطی یافت نشد."
+
+        llm = ChatOpenAI(model_name="gpt-3.5-turbo", openai_api_key=os.getenv("My_huggingface_key"))
+
+        prompt = f"""### سوال: {query}\n\n**🔍 اطلاعات مرتبط:**\n{context}\n\n**🤖 پاسخ:**"""
+        response = llm.predict(prompt)
+
+        return f"**📢 پاسخ:**\n{response}\n\n**📚 منابع:**\n" + "\n".join([f"- {doc.metadata.get('source', 'نامشخص')}, صفحه {doc.metadata.get('page', 'نامشخص')}" for doc in docs])
+
+    except Exception as e:
+        logger.error(f"❌ خطا در پاسخ‌گویی: {e}")
+        return f"❌ خطایی رخ داده است: {e}"
+
+# رابط کاربری با Gradio
 with gr.Blocks() as demo:
-    gr.Markdown("# 🏭 هوش مصنوعی همراه کارشناسان توزیع برق ایران")
+    gr.Markdown("### 🎛 **ابزار تحلیل اسناد توزیع برق**")
 
-    query = gr.Textbox(label="🔍 سوال خود را بپرسید", lines=2)
-    response = gr.HTML()
-    submit_btn = gr.Button("ارسال سوال")
+    upload_button = gr.File(label="آپلود فایل PDF", type="file")
+    upload_output = gr.Textbox(label="نتیجه آپلود", interactive=False)
 
-    file = gr.File(label="📎 آپلود فایل", file_types=[".pdf"])
-    upload_status = gr.Textbox(label="وضعیت آپلود", interactive=False)
-    upload_btn = gr.Button("ارسال فایل")
+    query_input = gr.Textbox(label="سوال خود را وارد کنید")
+    response_output = gr.HTML(label="📢 پاسخ", interactive=False)
 
-    upload_btn.click(save_uploaded_file, inputs=[file], outputs=[upload_status])
-    submit_btn.click(chat_with_doc, inputs=[query], outputs=[response])
+    upload_button.change(save_uploaded_file, inputs=upload_button, outputs=upload_output)
+    query_input.submit(chat_with_doc, inputs=query_input, outputs=response_output)
 
-demo.launch()
+# اجرای برنامه
+if __name__ == "__main__":
+    global_vector_db = load_database()
+    demo.launch()