IRI-PowerDistAI / app.py
Reza-galaxy21's picture
Create app.py
1f4020e verified
import os
import shutil
import logging
import gradio as gr
from langchain_community.document_loaders import PyPDFLoader
from langchain_community.embeddings import OpenAIEmbeddings
from langchain_community.vectorstores import FAISS
from langchain_community.chat_models import ChatOpenAI
# تنظیمات لاگ‌گیری
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
# تنظیمات مسیرها
UPLOAD_FOLDER = "uploaded_files"
DATABASE_FILE = "/home/user/app/vector_database"
os.makedirs(UPLOAD_FOLDER, exist_ok=True)
# بررسی و بارگذاری دیتابیس برداری
def load_database():
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}` وجود ندارد.")
return None
except Exception as e:
logger.error(f"❌ خطا در بارگذاری دیتابیس: {e}")
return None
def save_database(vector_db):
try:
vector_db.save_local(DATABASE_FILE)
logger.info(f"✅ دیتابیس ذخیره شد: {DATABASE_FILE}")
except Exception as e:
logger.error(f"❌ خطا در ذخیره دیتابیس: {e}")
raise e
global_vector_db = load_database()
# پردازش و ذخیره‌سازی فایل‌های PDF
def process_and_store_pdfs(file_paths):
texts = []
for file_path in file_paths:
try:
logger.info(f"در حال پردازش فایل: {file_path}")
if file_path.endswith(".pdf"):
loader = PyPDFLoader(file_path)
documents = loader.load()
texts.extend(documents)
logger.info(f"✅ تعداد صفحات پردازش شده: {len(documents)}")
else:
logger.warning(f"فرمت فایل {file_path} پشتیبانی نمی‌شود.")
except Exception as e:
logger.error(f"خطا در پردازش فایل {file_path}: {e}")
try:
logger.info("در حال ایجاد embeddings و ذخیره‌سازی در FAISS...")
embeddings = OpenAIEmbeddings(openai_api_key=os.getenv("My_huggingface_key"))
vector_db = FAISS.from_documents(texts, embeddings)
logger.info(f"✅ ذخیره‌سازی در FAISS با موفقیت انجام شد. تعداد اسناد: {vector_db.index.ntotal}")
return vector_db
except Exception as e:
logger.error(f"خطا در ایجاد embeddings یا ذخیره‌سازی در FAISS: {e}")
return None
# ایجاد لینک منبع
def generate_source_link(source, page):
return f'<a href="{source}#page={page}" target="_blank">{source}, صفحه {page}</a>'
# پاسخ‌گویی بر اساس اسناد بارگذاری‌شده
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": 8})
docs = retriever.get_relevant_documents(query)
if not docs:
return "هیچ اطلاعات مرتبطی یافت نشد."
# نمایش لاگ برای بررسی اسناد
for doc in docs:
logger.info(f"📄 سند: {doc.metadata.get('source', 'نامشخص')} | صفحه {doc.metadata.get('page', 'نامشخص')}")
# تفکیک بخش‌های نقل‌قول شده از متن سند
citations = []
context = ""
for doc in docs:
quoted_text = f"«{doc.page_content.strip()}»"
source_info = generate_source_link(doc.metadata.get("source", "نامشخص"), doc.metadata.get("page", "نامشخص"))
citations.append(f"{quoted_text} ({source_info})")
context += f"{quoted_text}\n\n"
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)
final_response = f"**پاسخ:**\n{response}\n\n**نقل‌قول‌های مستقیم:**\n" + "\n".join(citations)
return final_response
except Exception as e:
logger.error(f"خطا در پاسخ‌گویی بر اساس سند: {e}")
return f"❌ خطایی رخ داده است: {e}"
# فرمت کردن خروجی با HTML برای Gradio
def format_response(response_text):
response_text = response_text.replace("**پاسخ:**", "<h3 style='color:blue;'>پاسخ:</h3>")
response_text = response_text.replace("**نقل‌قول‌های مستقیم:**", "<h4 style='color:green;'>نقل‌قول‌های مستقیم:</h4>")
return response_text
# ذخیره فایل آپلود شده
def save_uploaded_file(file):
try:
file_name = os.path.basename(file)
file_path = os.path.join(UPLOAD_FOLDER, file_name)
shutil.copy(file, file_path)
global global_vector_db
global_vector_db = process_and_store_pdfs([file_path])
if global_vector_db:
save_database(global_vector_db)
return f"✅ فایل ذخیره شد: {file_path}"
except Exception as e:
return f"❌ خطا در ذخیره فایل: {e}"
# رابط کاربری با Gradio
with gr.Blocks() as demo:
gr.Markdown("# هوش مصنوعی همراه کارشناسان توزیع برق ایران")
query = gr.Textbox(label="سوال خود را بپرسید", lines=2)
response = gr.HTML(label="پاسخ", interactive=False) # تغییر از Textbox به HTML
submit_btn = gr.Button("ارسال سوال")
file = gr.File(label="📎 آپلود فایل", file_types=[".pdf"])
upload_status = gr.Textbox(label="وضعیت آپلود", interactive=False)
upload_btn = gr.Button("ارسال فایل")
upload_btn.click(save_uploaded_file, inputs=[file], outputs=[upload_status])
submit_btn.click(lambda q: format_response(chat_with_doc(q)), inputs=[query], outputs=[response])
demo.launch()