app.py

import os
import faiss
import numpy as np
import json
from fastapi import FastAPI, HTTPException
from fastapi.responses import HTMLResponse, RedirectResponse
from pydantic import BaseModel
from sentence_transformers import SentenceTransformer
from openai import OpenAI
from typing import Optional
import uvicorn
from rank_bm25 import BM25Okapi

# =====================
# CONFIG
# =====================
HF_TOKEN = os.environ.get("HF_TOKEN", "")
LLM_MODEL = "deepseek-ai/DeepSeek-V4-Flash:novita"
EMBED_MODEL = "intfloat/multilingual-e5-large"
DATA_FILE = "plant_diseases_guide.json"
CHUNK_SIZE = 500
CHUNK_OVERLAP = 50
TOP_K = 6

# =====================
# LOAD & CHUNK JSON
# =====================
def load_and_chunk(filepath: str, chunk_size: int = CHUNK_SIZE, overlap: int = CHUNK_OVERLAP):
    with open(filepath, "r", encoding="utf-8") as f:
        data = json.load(f)

    chunks = []

    if isinstance(data, dict) and "diseases" in data:
        for disease in data["diseases"]:
            chunk_text = f"""
المرض: {disease.get('arabic_name', '')}
النبات: {disease.get('plant', '')}
النوع: {disease.get('type', '')}
مسبب المرض: {disease.get('pathogen', '')}
مستوى الخطورة: {disease.get('danger_level', '')}
الانتشار في مصر: {disease.get('egypt_prevalence', '')}
الأعراض:
{chr(10).join(disease.get('symptoms', []))}
طرق الوقاية:
{chr(10).join(disease.get('prevention', []))}
العلاج:
{chr(10).join(disease.get('treatment', []))}
معلومات إضافية:
{disease.get('egypt_notes', '')}
{disease.get('danger_details', '')}
{disease.get('contagion_method', '')}
الظروف المساعدة: {' | '.join(disease.get('favorable_conditions', []))}
""".strip()
            chunks.append(chunk_text)

        if "general_agriculture_info" in data:
            general = data["general_agriculture_info"]
            for section_key, section in general.get("sections", {}).items():
                section_text = f"""
موضوع: {section.get('title', '')}
{chr(10).join(section.get('content', []))}
""".strip()
                chunks.append(section_text)

    else:
        if isinstance(data, list):
            text = "\n".join([json.dumps(item, ensure_ascii=False) for item in data])
        else:
            text = json.dumps(data, ensure_ascii=False, indent=2)

        start = 0
        while start < len(text):
            end = start + chunk_size
            chunk = text[start:end].strip()
            if chunk:
                chunks.append(chunk)
            start += chunk_size - overlap

    return chunks

# =====================
# BUILD FAISS INDEX — Cosine Similarity
# =====================
def build_index(chunks, embed_model):
    print("جاري بناء الـ index...")
    embeddings = embed_model.encode(chunks, show_progress_bar=True)
    embeddings = np.array(embeddings).astype("float32")
    faiss.normalize_L2(embeddings)
    dim = embeddings.shape[1]
    index = faiss.IndexFlatIP(dim)
    index.add(embeddings)
    print(f"تم بناء الـ index بنجاح - {len(chunks)} chunk")
    return index, embeddings

# =====================
# QUERY EXPANSION
# =====================
def expand_query(query: str) -> list:
    variations = [query]
    replacements = {
        "بيتلون": "تلون تغير لون",
        "بيعفن": "عفن تعفن",
        "بيموت": "يذبل موت ذبول",
        "عليه بقع": "بقع تبقع",
        "اصفرار": "اصفرار أوراق صفراء",
        "بتاع": "خاص بـ",
        "ليه": "لماذا سبب",
        "ازاي": "كيف طريقة",
        "امتا": "متى وقت",
        "علاج": "علاج معالجة مكافحة",
        "وقاية": "وقاية وقاء حماية منع",
        "اعراض": "أعراض علامات",
        "مرض": "مرض إصابة",
        "نبات": "نبات محصول زراعة",
        "فطر": "فطر فطريات مرض فطري",
        "حشرة": "حشرة آفة",
    }
    expanded = query
    for colloquial, formal in replacements.items():
        if colloquial in query:
            expanded = expanded.replace(colloquial, formal)
    if expanded != query:
        variations.append(expanded)
    return variations

# =====================
# QUERY REWRITING — يحل مشكلة الضمائر مع history طويل
# =====================
def rewrite_query(query: str, history: list, llm_client: OpenAI) -> str:
    if not history:
        return query

    reference_words = ["ده", "هو", "هي", "دي", "نفسه", "نفسها", "علاجه",
                       "علاجها", "أعراضه", "أعراضها", "خطير", "منتشر",
                       "سببه", "وقايته", "بتاعه", "بتاعها", "فيه", "عنه", "عنها"]

    if not any(word in query for word in reference_words):
        return query

    # بنبعت كل الـ history عشان يلاقي الموضوع حتى لو بعيد
    history_text = "\n".join([
        f"سؤال: {turn['question']}\nجواب: {turn['answer'][:150]}..."
        for turn in history
    ])

    rewrite_prompt = f"""المحادثة كلها:
{history_text}

السؤال الجديد: "{query}"

بناءً على المحادثة كلها، حدد الموضوع الرئيسي وأعد صياغة السؤال الجديد كسؤال مستقل وكامل بدون ضمائر.
أرجع السؤال المعاد صياغته فقط بدون أي كلام تاني."""

    try:
        response = llm_client.chat.completions.create(
            model=LLM_MODEL,
            messages=[{"role": "user", "content": rewrite_prompt}],
            max_tokens=100,
            temperature=0,
        )
        rewritten = response.choices[0].message.content.strip()
        print(f"Query rewritten: '{query}' → '{rewritten}'")
        return rewritten
    except Exception as e:
        print(f"Rewrite failed, using original: {e}")
        return query

# =====================
# HYBRID RETRIEVE — FAISS (Cosine) + BM25 + RRF
# =====================
def retrieve(query: str, index, chunks, embed_model, bm25: BM25Okapi, top_k: int = TOP_K):
    queries = expand_query(query)
    rrf_scores = {}

    # Dense retrieval (FAISS Cosine)
    for q in queries:
        query_vec = embed_model.encode([q]).astype("float32")
        faiss.normalize_L2(query_vec)
        distances, indices = index.search(query_vec, top_k * 2)
        for rank, idx in enumerate(indices[0]):
            if idx < len(chunks):
                rrf_scores[idx] = rrf_scores.get(idx, 0) + 1 / (rank + 60)

    # Sparse retrieval (BM25)
    for q in queries:
        bm25_scores = bm25.get_scores(q.split())
        top_bm25 = np.argsort(bm25_scores)[::-1][: top_k * 2]
        for rank, idx in enumerate(top_bm25):
            rrf_scores[idx] = rrf_scores.get(idx, 0) + 1 / (rank + 60)

    top_indices = sorted(rrf_scores.keys(), key=lambda x: rrf_scores[x], reverse=True)[:top_k]
    return [{"chunk": chunks[i], "score": rrf_scores[i]} for i in top_indices if i < len(chunks)]

# =====================
# GENERATE
# =====================
def generate_answer(query: str, context_chunks: list, llm_client: OpenAI, history: list = None):
    context = "\n---\n".join([c["chunk"] for c in context_chunks])

    system_prompt = """أنت مساعد ذكي متخصص في أمراض النباتات.
استخدم المعلومات المقدمة في كل رسالة للإجابة على السؤال مع مراعاة سياق المحادثة السابقة كاملاً.
إذا كان السؤال مرتبطاً بموضوع سبق ذكره في المحادثة، استخدم ذلك السياق مباشرةً دون طلب توضيح.
إذا لم تجد الإجابة في المعلومات المتاحة، أجب من معلوماتك العامة بس لازم يبقي سؤال لي علاقه بالزراعه او النباتات، لو بره الحاجات دي قوله مش تخصصي.
للتحيات والأسئلة الاعتيادية أجب بشكل طبيعي.
لا تستخدم أي تنسيق markdown مثل ** أو ## أو * في إجاباتك، اكتب نص عادي فقط."""

    messages = [{"role": "system", "content": system_prompt}]

    if history:
        for turn in history:
            messages.append({"role": "user", "content": turn["question"]})
            messages.append({"role": "assistant", "content": turn["answer"]})

    user_prompt = f"""المعلومات من قاعدة البيانات:
{context}
السؤال الحالي: {query}
الإجابة:"""

    messages.append({"role": "user", "content": user_prompt})

    response = llm_client.chat.completions.create(
        model=LLM_MODEL,
        messages=messages,
        max_tokens=512,
        temperature=0.3,
    )

    return response.choices[0].message.content.strip()

# =====================
# INIT
# =====================
print("جاري تحميل الـ embedding model...")
embed_model = SentenceTransformer(EMBED_MODEL)

print("جاري تحميل الداتا...")
chunks = load_and_chunk(DATA_FILE)
index, _ = build_index(chunks, embed_model)

print("جاري بناء الـ BM25 index...")
tokenized_chunks = [c.split() for c in chunks]
bm25 = BM25Okapi(tokenized_chunks)
print(f"BM25 جاهز - {len(tokenized_chunks)} chunk")

print("جاري الاتصال بـ DeepSeek...")
llm_client = OpenAI(
    base_url="https://router.huggingface.co/v1",
    api_key=HF_TOKEN,
)

app = FastAPI(title="Arabic RAG Chat API")

# =====================
# SCHEMAS
# =====================
class HistoryTurn(BaseModel):
    question: str
    answer: str

class QueryRequest(BaseModel):
    question: str
    top_k: Optional[int] = TOP_K
    history: Optional[list[HistoryTurn]] = []

class QueryResponse(BaseModel):
    answer: str
    sources: list[str]

# =====================
# ENDPOINTS
# =====================
@app.get("/")
def root():
    return RedirectResponse(url="/ui")

@app.get("/ui", response_class=HTMLResponse)
def ui():
    return r"""
<!DOCTYPE html>
<html dir="rtl" lang="ar">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>مساعد أمراض النباتات</title>
    <style>
        :root {
            --green-dark: #1b4332;
            --green-mid: #2d6a4f;
            --green-light: #52b788;
            --green-pale: #d8f3dc;
            --green-faint: #f0faf2;
            --gray-100: #f5f5f5;
            --gray-200: #e8e8e8;
            --gray-400: #aaa;
            --gray-600: #555;
            --gray-800: #222;
            --white: #ffffff;
            --border: #d4e8da;
        }
        * { box-sizing: border-box; margin: 0; padding: 0; }
        body {
            font-family: 'Segoe UI', Tahoma, Arial, sans-serif;
            background: var(--green-faint);
            height: 100dvh;
            display: flex;
            flex-direction: column;
            overflow: hidden;
        }
        header {
            background: var(--green-dark);
            color: white;
            padding: 14px 20px;
            display: flex;
            align-items: center;
            gap: 12px;
            flex-shrink: 0;
            box-shadow: 0 2px 8px rgba(0,0,0,0.2);
        }
        .header-icon {
            width: 40px; height: 40px;
            background: var(--green-light);
            border-radius: 50%;
            display: flex; align-items: center; justify-content: center;
            font-size: 20px;
        }
        .header-text h1 { font-size: 16px; font-weight: 700; }
        .header-text p  { font-size: 12px; opacity: 0.75; margin-top: 1px; }
        .header-actions { margin-right: auto; }
        #clear-btn {
            background: rgba(255,255,255,0.12);
            border: 1px solid rgba(255,255,255,0.25);
            color: white;
            padding: 6px 14px;
            border-radius: 20px;
            font-size: 12px;
            cursor: pointer;
            transition: background .2s;
            font-family: inherit;
        }
        #clear-btn:hover { background: rgba(255,255,255,0.22); }
        #chat-area {
            flex: 1;
            overflow-y: auto;
            padding: 20px 16px;
            display: flex;
            flex-direction: column;
            gap: 16px;
            scroll-behavior: smooth;
        }
        .welcome {
            text-align: center;
            padding: 30px 16px;
            color: var(--gray-600);
        }
        .welcome-icon { font-size: 48px; margin-bottom: 12px; }
        .welcome h2 { font-size: 18px; color: var(--green-dark); margin-bottom: 8px; }
        .welcome p  { font-size: 14px; line-height: 1.7; }
        .suggestions {
            display: flex; flex-wrap: wrap;
            gap: 8px; justify-content: center;
            margin-top: 16px;
        }
        .suggestion-chip {
            background: var(--white);
            border: 1.5px solid var(--border);
            color: var(--green-mid);
            padding: 8px 14px;
            border-radius: 20px;
            font-size: 13px;
            cursor: pointer;
            transition: all .2s;
            font-family: inherit;
        }
        .suggestion-chip:hover { background: var(--green-pale); border-color: var(--green-light); }
        .msg-row {
            display: flex;
            gap: 10px;
            align-items: flex-end;
            animation: fadeUp .25s ease;
        }
        @keyframes fadeUp {
            from { opacity: 0; transform: translateY(8px); }
            to   { opacity: 1; transform: translateY(0); }
        }
        .msg-row.user { flex-direction: row-reverse; }
        .msg-row.bot  { flex-direction: row; }
        .avatar {
            width: 32px; height: 32px;
            border-radius: 50%;
            flex-shrink: 0;
            display: flex; align-items: center; justify-content: center;
            font-size: 15px;
        }
        .avatar.bot  { background: var(--green-light); }
        .avatar.user { background: var(--green-dark); color: white; font-size: 12px; font-weight: 700; }
        .bubble {
            max-width: 72%;
            padding: 12px 16px;
            border-radius: 18px;
            font-size: 15px;
            line-height: 1.75;
            word-break: break-word;
            white-space: pre-wrap;
        }
        .msg-row.user .bubble {
            background: var(--green-mid);
            color: white;
            border-bottom-right-radius: 4px;
        }
        .msg-row.bot .bubble {
            background: var(--white);
            color: var(--gray-800);
            border-bottom-left-radius: 4px;
            border: 1px solid var(--border);
            box-shadow: 0 1px 4px rgba(0,0,0,0.06);
        }
        .typing-bubble {
            background: var(--white);
            border: 1px solid var(--border);
            border-radius: 18px;
            border-bottom-left-radius: 4px;
            padding: 14px 18px;
            display: flex; gap: 5px; align-items: center;
        }
        .dot {
            width: 7px; height: 7px;
            border-radius: 50%;
            background: var(--green-light);
            animation: bounce 1.2s infinite;
        }
        .dot:nth-child(2) { animation-delay: .2s; }
        .dot:nth-child(3) { animation-delay: .4s; }
        @keyframes bounce {
            0%,60%,100% { transform: translateY(0); }
            30%          { transform: translateY(-6px); }
        }
        #input-area {
            padding: 12px 16px 16px;
            background: var(--white);
            border-top: 1px solid var(--border);
            flex-shrink: 0;
        }
        .input-row {
            display: flex;
            gap: 10px;
            align-items: flex-end;
            max-width: 780px;
            margin: 0 auto;
        }
        #question {
            flex: 1;
            padding: 12px 16px;
            font-size: 15px;
            font-family: inherit;
            border: 1.5px solid var(--border);
            border-radius: 24px;
            outline: none;
            resize: none;
            max-height: 130px;
            overflow-y: auto;
            line-height: 1.5;
            transition: border-color .2s;
            background: var(--gray-100);
        }
        #question:focus { border-color: var(--green-light); background: var(--white); }
        #question::placeholder { color: var(--gray-400); }
        #send-btn {
            width: 46px; height: 46px;
            background: var(--green-mid);
            border: none;
            border-radius: 50%;
            color: white;
            font-size: 20px;
            cursor: pointer;
            flex-shrink: 0;
            transition: background .2s, transform .1s;
            display: flex; align-items: center; justify-content: center;
        }
        #send-btn:hover:not(:disabled)  { background: var(--green-dark); }
        #send-btn:active:not(:disabled) { transform: scale(.93); }
        #send-btn:disabled { background: var(--gray-200); color: var(--gray-400); cursor: not-allowed; }
        .char-hint { text-align: center; font-size: 11px; color: var(--gray-400); margin-top: 6px; }
        #chat-area::-webkit-scrollbar { width: 5px; }
        #chat-area::-webkit-scrollbar-thumb { background: var(--green-pale); border-radius: 4px; }
    </style>
</head>
<body>
<header>
    <div class="header-icon">🌿</div>
    <div class="header-text">
        <h1>مساعد أمراض النباتات</h1>
        <p>اسألني عن أي مرض نباتي — أتذكر المحادثة كلها</p>
    </div>
    <div class="header-actions">
        <button id="clear-btn" onclick="clearChat()">محادثة جديدة 🔄</button>
    </div>
</header>
<div id="chat-area">
    <div class="welcome" id="welcome-msg">
        <div class="welcome-icon">🌱</div>
        <h2>أهلاً! أنا مساعدك لأمراض النباتات</h2>
        <p>اسألني عن أي مرض وهجاوبك من قاعدة البيانات.<br>ومش محتاج تكرر اسم المرض في كل سؤال — أنا بتذكر!</p>
        <div class="suggestions">
            <button class="suggestion-chip" onclick="useSuggestion(this)">ما هي أعراض البياض الدقيقي؟</button>
            <button class="suggestion-chip" onclick="useSuggestion(this)">كيف أعالج الصدأ في القمح؟</button>
            <button class="suggestion-chip" onclick="useSuggestion(this)">ما أخطر أمراض الطماطم؟</button>
            <button class="suggestion-chip" onclick="useSuggestion(this)">كيف أقي نباتاتي من الفطريات؟</button>
        </div>
    </div>
</div>
<div id="input-area">
    <div class="input-row">
        <textarea id="question" rows="1" placeholder="اكتب سؤالك هنا..." maxlength="1000"></textarea>
        <button id="send-btn" onclick="sendMessage()" title="إرسال">➤</button>
    </div>
    <p class="char-hint">Enter للإرسال — Shift+Enter لسطر جديد</p>
</div>
<script>
    let chatHistory = [];
    let isLoading = false;
    const chatArea   = document.getElementById('chat-area');
    const questionEl = document.getElementById('question');
    const sendBtn    = document.getElementById('send-btn');

    questionEl.addEventListener('input', () => {
        questionEl.style.height = 'auto';
        questionEl.style.height = Math.min(questionEl.scrollHeight, 130) + 'px';
    });
    questionEl.addEventListener('keydown', (e) => {
        if (e.key === 'Enter' && !e.shiftKey) { e.preventDefault(); sendMessage(); }
    });

    async function sendMessage() {
        const q = questionEl.value.trim();
        if (!q || isLoading) return;
        hideWelcome();
        appendUserBubble(q);
        questionEl.value = '';
        questionEl.style.height = 'auto';
        setLoading(true);
        const typingId = showTyping();
        try {
            const res = await fetch('/query', {
                method: 'POST',
                headers: { 'Content-Type': 'application/json' },
                body: JSON.stringify({ question: q, history: chatHistory })
            });
            removeTyping(typingId);
            if (!res.ok) {
                const err = await res.json().catch(() => ({}));
                appendBotBubble('❌ ' + (err.detail || 'حصل خطأ، حاول تاني.'));
            } else {
                const data = await res.json();
                appendBotBubble(data.answer);
                chatHistory.push({ question: q, answer: data.answer });
                if (chatHistory.length>10){
                chatHistory.shift()
                }
            }
        } catch (e) {
            removeTyping(typingId);
            appendBotBubble('❌ مشكلة في الاتصال بالخادم، حاول مرة تانية.');
        }
        setLoading(false);
    }

    function hideWelcome() {
        const w = document.getElementById('welcome-msg');
        if (w) w.style.display = 'none';
    }
    function appendUserBubble(text) {
        const row = document.createElement('div');
        row.className = 'msg-row user';
        row.innerHTML = `<div class="avatar user">أنت</div><div class="bubble">${escapeHtml(text)}</div>`;
        chatArea.appendChild(row);
        scrollBottom();
    }
    function appendBotBubble(text) {
        const row = document.createElement('div');
        row.className = 'msg-row bot';
        row.innerHTML = `<div class="avatar bot">🌿</div><div class="bubble">${escapeHtml(text)}</div>`;
        chatArea.appendChild(row);
        scrollBottom();
    }
    function showTyping() {
        const id = 'typing-' + Date.now();
        const row = document.createElement('div');
        row.className = 'msg-row bot';
        row.id = id;
        row.innerHTML = `<div class="avatar bot">🌿</div><div class="typing-bubble"><div class="dot"></div><div class="dot"></div><div class="dot"></div></div>`;
        chatArea.appendChild(row);
        scrollBottom();
        return id;
    }
    function removeTyping(id) {
        const el = document.getElementById(id);
        if (el) el.remove();
    }
    function setLoading(state) {
        isLoading = state;
        sendBtn.disabled = state;
        questionEl.disabled = state;
    }
    function scrollBottom() { chatArea.scrollTop = chatArea.scrollHeight; }
    function escapeHtml(str) {
        return str.replace(/&/g,'&amp;').replace(/</g,'&lt;').replace(/>/g,'&gt;').replace(/"/g,'&quot;');
    }
    function useSuggestion(btn) {
        questionEl.value = btn.textContent;
        questionEl.dispatchEvent(new Event('input'));
        sendMessage();
    }
    function clearChat() {
        chatHistory = [];
        chatArea.innerHTML = '';
        const w = document.createElement('div');
        w.className = 'welcome';
        w.id = 'welcome-msg';
        w.innerHTML = `
            <div class="welcome-icon">🌱</div>
            <h2>أهلاً! أنا مساعدك لأمراض النباتات</h2>
            <p>اسألني عن أي مرض وهجاوبك من قاعدة البيانات.<br>ومش محتاج تكرر اسم المرض في كل سؤال — أنا بتذكر!</p>
            <div class="suggestions">
                <button class="suggestion-chip" onclick="useSuggestion(this)">ما هي أعراض البياض الدقيقي؟</button>
                <button class="suggestion-chip" onclick="useSuggestion(this)">كيف أعالج الصدأ في القمح؟</button>
                <button class="suggestion-chip" onclick="useSuggestion(this)">ما أخطر أمراض الطماطم؟</button>
                <button class="suggestion-chip" onclick="useSuggestion(this)">كيف أقي نباتاتي من الفطريات؟</button>
            </div>`;
        chatArea.appendChild(w);
    }
</script>
</body>
</html>
"""

@app.post("/query", response_model=QueryResponse)
def query(req: QueryRequest):
    if not req.question.strip():
        raise HTTPException(status_code=400, detail="السؤال فاضي!")

    history_dicts = [{"question": h.question, "answer": h.answer} for h in (req.history or [])]

    # Query Rewriting — يحل مشكلة الضمائر مع كل الـ history
    resolved_question = rewrite_query(req.question, history_dicts, llm_client)

    # Hybrid Retrieval بالسؤال المعاد صياغته
    results = retrieve(resolved_question, index, chunks, embed_model, bm25, req.top_k)

    if not results:
        raise HTTPException(status_code=404, detail="مفيش نتائج في الداتا")

    # Generate بالسؤال الأصلي + كل الـ history
    answer = generate_answer(req.question, results, llm_client, history=history_dicts)

    return QueryResponse(
        answer=answer,
        sources=[r["chunk"][:200] + "..." for r in results]
    )

@app.get("/health")
def health():
    return {"status": "healthy", "chunks": len(chunks)}

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)