app.py

import os
import re
import time
from typing import List, Optional

import faiss
import numpy as np
import pandas as pd
import requests
from datasets import load_dataset
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, Field
from rank_bm25 import BM25Okapi
from sentence_transformers import SentenceTransformer

MODEL_NAME = os.getenv("OPENROUTER_MODEL", "deepseek/deepseek-chat-v3-0324")
OPENROUTER_API_KEY = os.getenv("OPENROUTER_API_KEY", "").strip()

KNOWLEDGE_URL = os.getenv(
    "KNOWLEDGE_URL",
    "https://huggingface.co/datasets/sharshar1/arabic-hr-rag-dataset/resolve/main/knowledge_chunks.jsonl",
)
FAQ_URL = os.getenv(
    "FAQ_URL",
    "https://huggingface.co/datasets/sharshar1/arabic-hr-rag-dataset/resolve/main/faq_pairs.jsonl",
)

app = FastAPI(title="HR Genie RAG API", version="1.0.0")


class ChatRequest(BaseModel):
    message: str = Field(..., min_length=1, max_length=2000)
    session_id: Optional[str] = None
    user_id: Optional[int] = None


class SourceItem(BaseModel):
    title: str
    source_file: str
    score: float


class ChatResponse(BaseModel):
    answer: str
    mode: str
    sources: List[SourceItem]
    latency_ms: int


search_df: Optional[pd.DataFrame] = None
embedding_model: Optional[SentenceTransformer] = None
faiss_index = None
bm25 = None


def normalize_ar(text: str) -> str:
    text = str(text)
    text = re.sub(r"[\u064B-\u0652\u0670\u0640]", "", text)
    text = re.sub(r"[أإآٱ]", "ا", text)
    text = text.replace("ى", "ي").replace("ة", "ه")
    text = re.sub(r"\s+", " ", text).strip()
    return text


def simple_ar_tokenize(text: str) -> List[str]:
    text = normalize_ar(text)
    text = re.sub(r"[^\w\s]", " ", text)
    return text.split()


def detect_response_style(query: str) -> str:
    text = str(query).strip()
    arabic_chars = len(re.findall(r"[\u0600-\u06FF]", text))
    latin_chars = len(re.findall(r"[A-Za-z]", text))

    if latin_chars > arabic_chars and latin_chars > 0:
        return "english"

    if arabic_chars > 0:
        egyptian_markers = ["ازاي", "عاوز", "عايز", "ليه", "عامل", "ازيك", "دلوقتي", "ممكن"]
        lowered = text.lower()
        if any(marker in lowered for marker in egyptian_markers):
            return "egyptian_ar"
        return "fusha_ar"

    return "english"


def style_instruction(style: str) -> str:
    if style == "english":
        return "Answer in clear natural English."
    if style == "egyptian_ar":
        return "أجب باللهجة المصرية بشكل طبيعي وواضح، وبأسلوب شات بوت ودود."
    return "أجب بالعربية الفصحى بشكل واضح ومهذب."


def source_label_for_style(style: str) -> str:
    return "Source" if style == "english" else "المصدر"


def sanitize_model_output(text: str) -> str:
    text = str(text or "")
    text = re.sub(r"[\u4e00-\u9fff]+", " ", text)
    text = re.sub(r"\n{3,}", "\n\n", text)
    return re.sub(r"[ \t]{2,}", " ", text).strip()


def openrouter_chat(messages: List[dict], temperature: float = 0.2, max_tokens: int = 300) -> str:
    if not OPENROUTER_API_KEY:
        raise RuntimeError("OPENROUTER_API_KEY is missing.")

    url = "https://openrouter.ai/api/v1/chat/completions"
    headers = {
        "Authorization": f"Bearer {OPENROUTER_API_KEY}",
        "Content-Type": "application/json",
    }

    payload = {
        "model": MODEL_NAME,
        "messages": messages,
        "temperature": temperature,
        "max_tokens": max_tokens,
    }

    response = requests.post(url, headers=headers, json=payload, timeout=60)
    response.raise_for_status()
    data = response.json()
    return sanitize_model_output(data["choices"][0]["message"]["content"])


def load_data() -> pd.DataFrame:
    knowledge_ds = load_dataset("json", data_files=KNOWLEDGE_URL)["train"]
    faq_ds = load_dataset("json", data_files=FAQ_URL)["train"]

    knowledge_df = pd.DataFrame(knowledge_ds)
    faq_df = pd.DataFrame(faq_ds)

    knowledge_search = knowledge_df.copy()
    knowledge_search["search_text"] = (
        "الفئة: "
        + knowledge_search["category_ar"].fillna("")
        + "\nالعنوان: "
        + knowledge_search["topic"].fillna("")
        + "\nالمحتوى: "
        + knowledge_search["content"].fillna("")
    )
    knowledge_search["answer_text"] = knowledge_search["content"]
    knowledge_search["display_title"] = knowledge_search["topic"]
    knowledge_search["record_type"] = "knowledge"

    faq_search = faq_df.copy()
    faq_search["search_text"] = (
        "الفئة: "
        + faq_search["category_ar"].fillna("")
        + "\nالسؤال: "
        + faq_search["question"].fillna("")
        + "\nالإجابة: "
        + faq_search["answer"].fillna("")
    )
    faq_search["answer_text"] = faq_search["answer"]
    faq_search["display_title"] = faq_search["question"]
    faq_search["record_type"] = "faq"

    common_cols = [
        "category_ar",
        "source_file",
        "record_type",
        "display_title",
        "answer_text",
        "search_text",
    ]

    return pd.concat([knowledge_search[common_cols], faq_search[common_cols]], ignore_index=True)


def build_indexes() -> None:
    global search_df, embedding_model, faiss_index, bm25

    search_df = load_data()

    embedding_model = SentenceTransformer("sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2")

    search_texts = search_df["search_text"].apply(normalize_ar).tolist()
    vectors = embedding_model.encode(
        search_texts,
        batch_size=32,
        convert_to_numpy=True,
        normalize_embeddings=True,
    ).astype("float32")

    faiss_index = faiss.IndexFlatIP(vectors.shape[1])
    faiss_index.add(vectors)

    bm25_corpus = [simple_ar_tokenize(t) for t in search_df["search_text"].tolist()]
    bm25 = BM25Okapi(bm25_corpus)


def retrieve(query: str, top_k: int = 3) -> List[dict]:
    if search_df is None or embedding_model is None or faiss_index is None or bm25 is None:
        raise RuntimeError("Indexes are not initialized.")

    query_norm = normalize_ar(query)

    query_vec = embedding_model.encode(
        [query_norm],
        convert_to_numpy=True,
        normalize_embeddings=True,
    ).astype("float32")

    vector_scores, vector_indices = faiss_index.search(query_vec, 8)

    bm25_scores = bm25.get_scores(simple_ar_tokenize(query_norm))
    bm25_top_idx = np.argsort(bm25_scores)[::-1][:8]

    merged = {}

    for rank, (score, idx) in enumerate(zip(vector_scores[0], vector_indices[0]), start=1):
        merged[int(idx)] = {
            "vector_rank": rank,
            "vector_score": float(score),
            "bm25_rank": None,
        }

    for rank, idx in enumerate(bm25_top_idx, start=1):
        idx = int(idx)
        merged.setdefault(
            idx,
            {
                "vector_rank": None,
                "vector_score": 0.0,
                "bm25_rank": None,
            },
        )
        merged[idx]["bm25_rank"] = rank

    items = []
    for idx, score_obj in merged.items():
        v_rank_score = 1.0 / score_obj["vector_rank"] if score_obj["vector_rank"] else 0.0
        b_rank_score = 1.0 / score_obj["bm25_rank"] if score_obj["bm25_rank"] else 0.0
        hybrid_score = 0.65 * v_rank_score + 0.35 * b_rank_score

        row = search_df.iloc[idx]
        items.append(
            {
                "display_title": str(row.get("display_title", "")),
                "answer_text": str(row.get("answer_text", "")),
                "source_file": str(row.get("source_file", "")),
                "record_type": str(row.get("record_type", "")),
                "category_ar": str(row.get("category_ar", "")),
                "score": float(hybrid_score),
            }
        )

    return sorted(items, key=lambda x: x["score"], reverse=True)[:top_k]


def build_prompt(query: str, chunks: List[dict]) -> str:
    style = detect_response_style(query)
    response_style_instruction = style_instruction(style)

    context = []
    for i, c in enumerate(chunks, start=1):
        context.append(
            "\n".join(
                [
                    f"[Source {i}]",
                    f"Type: {c['record_type']}",
                    f"Category: {c['category_ar']}",
                    f"Title: {c['display_title']}",
                    f"Content: {c['answer_text']}",
                    f"File: {c['source_file']}",
                ]
            )
        )

    context_text = "\n\n".join(context)

    return (
        "You are an HR smart assistant inside a company.\n\n"
        f"{response_style_instruction}\n\n"
        "Use the provided sources when they are relevant and enough.\n"
        "If the sources are not enough, answer naturally from general knowledge.\n"
        "Do not mention retrieval or hidden reasoning.\n"
        "Keep answers short, clear, and practical.\n"
        "If you used sources, add one final source line in the same language style.\n\n"
        f"Sources:\n{context_text}\n\n"
        f"User question:\n{query}\n\n"
        "Final answer:"
    )


def ask_assistant(query: str) -> dict:
    chunks = retrieve(query, top_k=3)
    style = detect_response_style(query)
    source_label = source_label_for_style(style)

    if not chunks:
        prompt = (
            "You are a smart and friendly assistant. "
            f"{style_instruction(style)} "
            "Answer naturally and briefly."
        )
        answer = openrouter_chat(
            [
                {"role": "system", "content": prompt},
                {"role": "user", "content": query},
            ],
            temperature=0.4,
            max_tokens=260,
        )
        return {"answer": answer, "mode": "general_chat", "sources": []}

    prompt = build_prompt(query, chunks)
    answer = openrouter_chat(
        [{"role": "user", "content": prompt}],
        temperature=0.15,
        max_tokens=300,
    )

    answer = sanitize_model_output(answer)

    if source_label not in answer and chunks[0]["source_file"]:
        answer = f"{answer}\n{source_label}\n{chunks[0]['source_file']}"

    sources = [
        {
            "title": c["display_title"],
            "source_file": c["source_file"],
            "score": round(c["score"], 4),
        }
        for c in chunks
    ]

    return {"answer": answer, "mode": "rag", "sources": sources}


@app.on_event("startup")
def on_startup() -> None:
    build_indexes()


@app.get("/")
def root() -> dict:
    return {
        "name": "HR Genie RAG API",
        "status": "running",
        "endpoints": {
            "health": "/health",
            "chat": "/chat",
            "docs": "/docs",
        },
    }


@app.get("/health")
def health() -> dict:
    return {"status": "ok", "model": MODEL_NAME, "has_openrouter_key": bool(OPENROUTER_API_KEY)}


@app.post("/chat", response_model=ChatResponse)
def chat(payload: ChatRequest) -> ChatResponse:
    start = time.time()

    try:
        result = ask_assistant(payload.message.strip())
    except Exception as exc:
        raise HTTPException(status_code=500, detail=str(exc)) from exc

    latency_ms = int((time.time() - start) * 1000)

    sources = [SourceItem(**src) for src in result.get("sources", [])]

    return ChatResponse(
        answer=result.get("answer", ""),
        mode=result.get("mode", "unknown"),
        sources=sources,
        latency_ms=latency_ms,
    )