app.py

import os
import time
import streamlit as st
import spacy
from groq import Groq
from langchain.document_loaders import PyPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.schema import Document as LangchainDocument
from langchain.chains import RetrievalQA
from langchain.llms import OpenAI

# ----------------- تنظیمات صفحه -----------------
st.set_page_config(page_title="چت‌بات ارتش - فقط از PDF", page_icon="🪖", layout="wide")

# استایل فارسی و بک‌گراند (مثل قبل...)

# ----------------- تعریف کلید API -----------------
groq_api_key = "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"

# ----------------- بارگذاری مدل spaCy برای امبدینگ -----------------
nlp = spacy.load("en_core_web_md")  # یا 'en_core_web_lg' برای مدل بزرگتر

# ----------------- لود PDF و ساخت ایندکس -----------------
@st.cache_resource
def build_pdf_index():
    with st.spinner('📄 در حال پردازش فایل PDF...'):
        loader = PyPDFLoader("test1.pdf")
        pages = loader.load()

        # تکه‌تکه کردن متن PDF
        splitter = RecursiveCharacterTextSplitter(
            chunk_size=500,
            chunk_overlap=50
        )

        texts = []
        for page in pages:
            texts.extend(splitter.split_text(page.page_content))

        # تبدیل به Document
        documents = [LangchainDocument(page_content=t) for t in texts]

        # استفاده از spaCy برای تولید امبدینگ
        embeddings = []
        for doc in documents:
            spacy_doc = nlp(doc.page_content)  # تبدیل متن به شیء spaCy
            embeddings.append(spacy_doc.vector)  # استخراج امبدینگ از spaCy

        # به جای FAISS، فقط لیست امبدینگ‌ها را برمی‌گردانیم
        return documents, embeddings

# ----------------- ساختن Index از PDF -----------------
documents, embeddings = build_pdf_index()

# ----------------- تعریف LLM Groq -----------------
client = Groq(api_key=groq_api_key)

class GroqLLM(OpenAI):
    def __init__(self, api_key, model_name):
        super().__init__(openai_api_key=api_key, model_name=model_name, base_url="https://api.groq.com/openai/v1")

llm = GroqLLM(api_key=groq_api_key, model_name="deepseek-r1-distill-llama-70b")

# ----------------- Retrieval Chain -----------------
# به جای استفاده از FAISS، اکنون از لیست امبدینگ‌ها استفاده می‌کنیم
class SimpleRetriever:
    def __init__(self, documents, embeddings):
        self.documents = documents
        self.embeddings = embeddings

    def retrieve(self, query, top_k=1):
        query_embedding = nlp(query).vector  # تبدیل سوال به امبدینگ با استفاده از spaCy
        similarities = []
        for doc_embedding in self.embeddings:
            similarity = query_embedding.dot(doc_embedding)  # محاسبه شباهت بین امبدینگ‌ها
            similarities.append(similarity)
        
        # یافتن مستندات مشابه بر اساس بیشترین شباهت
        ranked_docs = sorted(zip(similarities, self.documents), reverse=True)
        return [doc for _, doc in ranked_docs[:top_k]]

# استفاده از retriever ساده برای جستجو در امبدینگ‌ها
retriever = SimpleRetriever(documents, embeddings)

# ----------------- Retrieval Chain -----------------
chain = RetrievalQA.from_chain_type(
    llm=llm,
    retriever=retriever.retrieve,  # استفاده از متد retrieve
    chain_type="stuff",
    input_key="question"
)

# ----------------- استیت برای چت -----------------
if 'messages' not in st.session_state:
    st.session_state.messages = []

if 'pending_prompt' not in st.session_state:
    st.session_state.pending_prompt = None

# ----------------- نمایش پیام‌های قبلی -----------------
for msg in st.session_state.messages:
    with st.chat_message(msg['role']):
        st.markdown(f"🗨️ {msg['content']}", unsafe_allow_html=True)

# ----------------- ورودی چت -----------------
prompt = st.chat_input("سوالی در مورد فایل بپرس...")

if prompt:
    st.session_state.messages.append({'role': 'user', 'content': prompt})
    st.session_state.pending_prompt = prompt
    st.rerun()

# ----------------- پاسخ مدل فقط از روی PDF -----------------
if st.session_state.pending_prompt:
    with st.chat_message('ai'):
        thinking = st.empty()
        thinking.markdown("🤖 در حال فکر کردن از روی PDF...")

        try:
            # گرفتن جواب فقط از PDF
            response = chain.run(f"سوال: {st.session_state.pending_prompt}")
            answer = response.strip()

        except Exception as e:
            answer = f"خطا در پاسخ‌دهی: {str(e)}"

        thinking.empty()

        # انیمیشن تایپ پاسخ
        full_response = ""
        placeholder = st.empty()
        for word in answer.split():
            full_response += word + " "
            placeholder.markdown(full_response + "▌")
            time.sleep(0.03)

        placeholder.markdown(full_response)
        st.session_state.messages.append({'role': 'ai', 'content': full_response})
        st.session_state.pending_prompt = None