app.py

import time
import streamlit as st
st.set_page_config(page_title="چت بات ارتش", page_icon="🪖", layout="wide")
st.markdown("""
    <style>
    .main {
        background-color: #f4f6f7;
    }

    .stChatMessage {
        background-color: #e8f0fe;
        border-radius: 12px;
        padding: 10px;
        margin-bottom: 10px;
        direction: rtl;
        text-align: right;
        font-family: 'Tahoma', sans-serif;
    }

    .stMarkdown, .stTextInput, .stTextArea {
        direction: rtl !important;
        text-align: right !important;
        font-family: 'Tahoma', sans-serif;
    }

    .center-header {
        display: flex;
        flex-direction: column;
        align-items: center;
        justify-content: center;
        margin-bottom: 30px;
    }

    .center-header img {
        width: 150px; /* اندازه بزرگ‌تر برای لوگو */
        height: auto;
        border-radius: 20px;
        box-shadow: 0 4px 12px rgba(0,0,0,0.15);
        margin-bottom: 10px;
    }

    .center-header h1 {
        font-size: 28px;
        font-weight: bold;
        color: #2c3e50;
        font-family: 'Tahoma', sans-serif;
        text-align: center;
        margin: 0;
    }
    </style>
""", unsafe_allow_html=True)

st.markdown("""
    <div class="center-header">
        <img src="https://your-logo-url.com/logo.png" alt="لوگو">
        <h1>اسم اپلیکیشن یا شرکت شما</h1>
    </div>
""", unsafe_allow_html=True)


from langchain.document_loaders import PyPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.embeddings.base import Embeddings
from langchain.vectorstores import FAISS
from langchain.indexes import VectorstoreIndexCreator
from langchain.chains import RetrievalQA
from langchain.chat_models import ChatOpenAI
from typing import List
from together import Together

class TogetherEmbeddings(Embeddings):
    def __init__(self, model_name: str, api_key: str):
        self.model_name = model_name
        self.client = Together(api_key=api_key)

    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        response = self.client.embeddings.create(
            model=self.model_name,
            input=texts
        )
        return [item.embedding for item in response.data]

    def embed_query(self, text: str) -> List[float]:
        return self.embed_documents([text])[0]

@st.cache_resource
def get_pdf_index():
    with st.spinner('لطفاً لحظه‌ای صبر کنید...'):
        pdf_reader = [PyPDFLoader('test1.pdf')]
        embeddings = TogetherEmbeddings(
            model_name="togethercomputer/m2-bert-80M-8k-retrieval",
            api_key="0291f33aee03412a47fa5d8e562e515182dcc5d9aac5a7fb5eefdd1759005979"
        )
        return VectorstoreIndexCreator(
            embedding=embeddings,
            text_splitter=RecursiveCharacterTextSplitter(chunk_size=300, chunk_overlap=0)
        ).from_loaders(pdf_reader)

index = get_pdf_index()
llm = ChatOpenAI(
    base_url="https://api.together.xyz/v1",
    api_key='0291f33aee03412a47fa5d8e562e515182dcc5d9aac5a7fb5eefdd1759005979',
    model="meta-llama/Llama-3.3-70B-Instruct-Turbo-Free"
)
chain = RetrievalQA.from_chain_type(
    llm=llm,
    chain_type='stuff',
    retriever=index.vectorstore.as_retriever(),
    input_key='question'
)

# --- UI زیباسازی ---

col1, col2 = st.columns([1, 10])
with col1:
    st.image("army.png", width=70)
with col2:
    st.title('🤖 چت‌بات هوشمند ارتش')

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 message in st.session_state.messages:
    with st.chat_message(message['role']):
        st.markdown(f"🗨️ {message['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()

if st.session_state.pending_prompt:
    with st.chat_message('ai'):
        thinking_placeholder = st.empty()
        thinking_placeholder.markdown("🤖 در حال فکر کردن...")

        response = chain.run(f'پاسخ را فقط به زبان فارسی بده. سوال: {st.session_state.pending_prompt}')
        helpful_answer = response.split("Helpful Answer:")[-1]
        if not helpful_answer.strip():
            helpful_answer = "اطلاعات دقیقی در دسترس نیست، اما می‌توانم به شما کمک کنم تا از منابع دیگر بررسی کنید."

        thinking_placeholder.empty()
        full_response = ""
        placeholder = st.empty()
        for chunk in helpful_answer.split():
            full_response += chunk + " "
            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