Update app.py

app.py

@@ -2,98 +2,70 @@ import os
 import time
 import streamlit as st
 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.vectorstores import FAISS
+from langchain.embeddings import HuggingFaceEmbeddings
+from langchain.chains import RetrievalQA
+from langchain.llms import OpenAI
 
 # ----------------- تنظیمات صفحه -----------------
-st.set_page_config(page_title="چت‌بات ارتش - Powered by Groq", page_icon="🪖", layout="wide")
-
-# استایل فارسی و بک‌گراند
-st.markdown("""
-    <style>
-    @import url('https://fonts.googleapis.com/css2?family=Vazirmatn:wght@400;700&display=swap');
-    html, body, [class*="css"] {
-        font-family: 'Vazirmatn', Tahoma, sans-serif;
-        direction: rtl;
-        text-align: right;
-    }
-    .stApp {
-        background: url("./military_bg.jpeg") no-repeat center center fixed;
-        background-size: cover;
-        backdrop-filter: blur(2px);
-    }
-    .stChatMessage {
-        background-color: rgba(255,255,255,0.8);
-        border: 1px solid #4e8a3e;
-        border-radius: 12px;
-        padding: 16px;
-        margin-bottom: 15px;
-        box-shadow: 0 4px 10px rgba(0,0,0,0.2);
-        animation: fadeIn 0.4s ease-in-out;
-    }
-    .stTextInput > div > input, .stTextArea textarea {
-        background-color: rgba(255,255,255,0.9) !important;
-        border-radius: 8px !important;
-        direction: rtl;
-        text-align: right;
-        font-family: 'Vazirmatn', Tahoma;
-    }
-    .stButton>button {
-        background-color: #4e8a3e !important;
-        color: white !important;
-        font-weight: bold;
-        border-radius: 10px;
-        padding: 8px 20px;
-        transition: 0.3s;
-    }
-    .stButton>button:hover {
-        background-color: #3c6d30 !important;
-    }
-    .header-text {
-        text-align: center;
-        margin-top: 20px;
-        margin-bottom: 40px;
-        background-color: rgba(255, 255, 255, 0.75);
-        padding: 20px;
-        border-radius: 20px;
-        box-shadow: 0 4px 12px rgba(0,0,0,0.2);
-    }
-    .header-text h1 {
-        font-size: 42px;
-        color: #2c3e50;
-        margin: 0;
-        font-weight: bold;
-    }
-    .subtitle {
-        font-size: 18px;
-        color: #34495e;
-        margin-top: 8px;
-    }
-    @keyframes fadeIn {
-        from { opacity: 0; transform: translateY(10px); }
-        to { opacity: 1; transform: translateY(0); }
-    }
-    </style>
-""", unsafe_allow_html=True)
-
-# ----------------- لوگو و عنوان -----------------
-col1, col2, col3 = st.columns([1, 1, 1])
-with col2:
-    st.image("army.png", width=240)
-
-st.markdown("""
-    <div class="header-text">
-        <h1>چت‌بات ارتش</h1>
-        <div class="subtitle">دستیار هوشمند میدان نبرد - Powered by Groq</div>
-    </div>
-""", unsafe_allow_html=True)
-
-# ----------------- اتصال به Groq -----------------
-api_key = "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"
-
-client = Groq(api_key=api_key)
-
-selected_model = "llama3-70b-8192"  # بهترین مدل Groq
-
-# ----------------- استیت ذخیره‌ی پیام‌ها -----------------
+st.set_page_config(page_title="چت‌بات ارتش - فقط از PDF", page_icon="🪖", layout="wide")
+
+# استایل فارسی و بک‌گراند (مثل قبل...)
+
+# ----------------- تعریف کلید API -----------------
+groq_api_key = "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"
+
+# ----------------- لود 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]
+
+        # استفاده از HuggingFaceEmbedding محلی برای FAISS
+        embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
+
+        vectordb = FAISS.from_documents(documents, embedding=embeddings)
+
+        return vectordb
+
+# ----------------- ساختن Index از PDF -----------------
+index = 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="llama3-70b-8192")
+
+# ----------------- Retrieval Chain -----------------
+chain = RetrievalQA.from_chain_type(
+    llm=llm,
+    retriever=index.as_retriever(),
+    chain_type="stuff",
+    input_key="question"
+)
+
+# ----------------- استیت برای چت -----------------
 if 'messages' not in st.session_state:
     st.session_state.messages = []
 
@@ -106,28 +78,23 @@ for msg in st.session_state.messages:
         st.markdown(f"🗨️ {msg['content']}", unsafe_allow_html=True)
 
 # ----------------- ورودی چت -----------------
-prompt = st.chat_input("چطور می‌تونم کمک کنم؟")
+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("🤖 در حال فکر کردن...")
+        thinking.markdown("🤖 در حال فکر کردن از روی PDF...")
 
         try:
-            chat_completion = client.chat.completions.create(
-                messages=[
-                    {"role": "system", "content": "پاسخ را همیشه رسمی و فارسی بده."},
-                    {"role": "user", "content": st.session_state.pending_prompt}
-                ],
-                model=selected_model,
-            )
-            answer = chat_completion.choices[0].message.content.strip()
+            # گرفتن جواب فقط از PDF
+            response = chain.run(f"سوال: {st.session_state.pending_prompt}")
+            answer = response.strip()
 
         except Exception as e:
             answer = f"خطا در پاسخ‌دهی: {str(e)}"