server.py

import os
import re
import math
import json
import uuid
import logging
from pathlib import Path
from typing import Optional
from collections import defaultdict

import anthropic
from fastapi import FastAPI, UploadFile, File, HTTPException, Form
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse, StreamingResponse
from pydantic import BaseModel

import openai
import google.generativeai as genai
import cohere
import numpy as np
from rank_bm25 import BM25Okapi


# ─── Logging ──────────────────────────────────────────────────────────────────
logging.basicConfig(level=logging.INFO, format="%(asctime)s  %(levelname)s  %(message)s")
log = logging.getLogger(__name__)

# ─── App ──────────────────────────────────────────────────────────────────────
app = FastAPI(title="RAG System", version="2.0.0")
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"], allow_methods=["*"], allow_headers=["*"],
)

# ─── In-memory store ──────────────────────────────────────────────────────────
# { doc_id: { id, name, ext, chunks: [str], char_count: int } }
DOCUMENTS: dict[str, dict] = {}

# ─── Config ───────────────────────────────────────────────────────────────────
CHUNK_SIZE    = int(os.getenv("CHUNK_SIZE",    "600"))
CHUNK_OVERLAP = int(os.getenv("CHUNK_OVERLAP", "100"))
TOP_K         = int(os.getenv("TOP_K",         "4"))
MODEL         = os.getenv("MODEL",             "claude-3-5-sonnet-20240620")
ANTHROPIC_KEY = os.getenv("ANTHROPIC_API_KEY", "")
OPENAI_KEY    = os.getenv("OPENAI_API_KEY",    "")
GOOGLE_KEY    = os.getenv("GOOGLE_API_KEY",    "")
COHERE_KEY    = os.getenv("COHERE_API_KEY",    "")

# ─── Advanced Retrieval Config ────────────────────────────────────────────────
HYBRID_WEIGHT_BM25 = 0.5
HYBRID_WEIGHT_SEMANTIC = 0.5
RRF_K = 60 # Constant for RRF
USE_RERANK = True

# ─── Text extraction ──────────────────────────────────────────────────────────
def extract_text(filename: str, raw: bytes) -> str:
    ext = Path(filename).suffix.lower()

    # PDF
    if ext == ".pdf":
        try:
            import pypdf
            from io import BytesIO
            reader = pypdf.PdfReader(BytesIO(raw))
            return "\n\n".join(p.extract_text() or "" for p in reader.pages)
        except ImportError:
            raise HTTPException(400, "pypdf not installed. Run: pip install pypdf")

    # DOCX
    if ext == ".docx":
        try:
            from docx import Document
            from io import BytesIO
            doc = Document(BytesIO(raw))
            return "\n".join(p.text for p in doc.paragraphs)
        except ImportError:
            raise HTTPException(400, "python-docx not installed. Run: pip install python-docx")

    # Plain-text formats
    if ext in {".txt", ".md", ".py", ".js", ".ts", ".jsx", ".tsx",
               ".json", ".csv", ".html", ".htm", ".xml", ".yaml", ".yml",
               ".toml", ".ini", ".sh", ".bash", ".sql", ".rst", ".tex"}:
        for enc in ("utf-8", "utf-8-sig", "tis-620", "cp874", "latin-1"):
            try:
                return raw.decode(enc)
            except UnicodeDecodeError:
                continue
        return raw.decode("utf-8", errors="replace")

    raise HTTPException(400, f"Unsupported file type: {ext}")


# ─── Chunking ─────────────────────────────────────────────────────────────────
def chunk_text(text: str, size: int = CHUNK_SIZE, overlap: int = CHUNK_OVERLAP) -> list[str]:
    text = re.sub(r"\n{3,}", "\n\n", text.strip())
    chunks, start = [], 0
    while start < len(text):
        end = min(start + size, len(text))
        chunk = text[start:end].strip()
        if len(chunk) > 30:
            chunks.append(chunk)
        start += size - overlap
    return chunks


# ─── TF-IDF Retrieval ─────────────────────────────────────────────────────────
def tokenize(text: str) -> list[str]:
    # ตัดคำทั้ง EN และ TH อย่างง่าย
    tokens = re.findall(r"[ก-๙]+|[a-zA-Z0-9]+", text.lower())
    return [t for t in tokens if len(t) > 1]


def build_idf(all_chunks: list[str]) -> dict[str, float]:
    N = len(all_chunks)
    df: dict[str, int] = defaultdict(int)
    for chunk in all_chunks:
        for tok in set(tokenize(chunk)):
            df[tok] += 1
    return {tok: math.log((N + 1) / (cnt + 1)) + 1 for tok, cnt in df.items()}


def tfidf_score(query: str, chunk: str, idf: dict[str, float]) -> float:
    q_tokens = tokenize(query)
    c_tokens = tokenize(chunk)
    if not c_tokens:
        return 0.0

    freq: dict[str, int] = defaultdict(int)
    for t in c_tokens:
        freq[t] += 1

    score = 0.0
    for qt in q_tokens:
        tf = freq.get(qt, 0) / len(c_tokens)
        score += tf * idf.get(qt, 1.0)

    # Bonus: exact query substring in chunk
    if query.lower() in chunk.lower():
        score += 2.0

    # Bigram bonus
    for i in range(len(q_tokens) - 1):
        bigram = q_tokens[i] + " " + q_tokens[i + 1]
        if bigram in chunk.lower():
            score += 0.5

    return score


def retrieve(query: str, top_k: int = TOP_K):
    if not DOCUMENTS:
        return []

    all_chunks: list[dict] = []
    for doc_id, doc in DOCUMENTS.items():
        for i, chunk in enumerate(doc["chunks"]):
            all_chunks.append({
                "text":      chunk,
                "doc_id":    doc_id,
                "doc_name":  doc["name"],
                "chunk_idx": i,
            })

    if not all_chunks:
        return []

    idf = build_idf([c["text"] for c in all_chunks])
    scored = [
        {**c, "score": tfidf_score(query, c["text"], idf)}
        for c in all_chunks
    ]
    scored.sort(key=lambda x: x["score"], reverse=True)
    return scored[:top_k]


# ─── Advanced Retrieval (Hybrid + RRF + Rerank) ───────────────────────────────
def get_embeddings(texts: list[str], client: openai.OpenAI, model: str = "text-embedding-3-small"):
    res = client.embeddings.create(input=texts, model=model)
    return [d.embedding for d in res.data]


def cosine_similarity(a, b):
    return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))


def hybrid_retrieve(query: str, top_k: int = TOP_K, rerank_top_n: int = 10):
    if not DOCUMENTS:
        return []

    all_chunks: list[dict] = []
    for doc_id, doc in DOCUMENTS.items():
        for i, chunk in enumerate(doc["chunks"]):
            all_chunks.append({
                "text":      chunk,
                "doc_id":    doc_id,
                "doc_name":  doc["name"],
                "chunk_idx": i,
            })

    if not all_chunks:
        return []

    # 1. BM25 Search
    tokenized_corpus = [tokenize(c["text"]) for c in all_chunks]
    bm25 = BM25Okapi(tokenized_corpus)
    tokenized_query = tokenize(query)
    bm25_scores = bm25.get_scores(tokenized_query)
    
    # Get BM25 rankings
    bm25_indices = np.argsort(bm25_scores)[::-1]
    bm25_rank = {idx: rank + 1 for rank, idx in enumerate(bm25_indices)}

    # 2. Semantic Search (Simplified manual cosine similarity for in-memory)
    # Note: ในโปรดักชั่นควรใช้ Vector Database เช่น Pinecone, Qdrant หรือ FAISS
    # เราจะใช้ OpenAI Embeddings
    api_key = os.getenv("OPENAI_API_KEY", OPENAI_KEY)
    if not api_key:
        # Fallback to TF-IDF if no OpenAI key for embeddings
        return retrieve(query, top_k)
    
    client = openai.OpenAI(api_key=api_key)
    query_embedding = get_embeddings([query], client)[0]
    
    # For efficiency we might want to cache embeddings, but for now we re-embed or skip
    # Since we can't easily cache in current DOCUMENTS structure without more logic, 
    # let's assume small doc set or that we have embeddings.
    # To keep it simple and working, I'll only embed the query and compare with pre-stored embeddings if they existed.
    # But since they don't, I'll have to embed all chunks (caution: slow for many documents).
    
    # [TODO] In a real app, chunks would be embedded once during upload.
    # For this demo, let's embed them now if they don't have embeddings.
    chunks_to_embed = [c["text"] for c in all_chunks]
    chunk_embeddings = get_embeddings(chunks_to_embed, client)
    
    semantic_scores = [cosine_similarity(query_embedding, ce) for ce in chunk_embeddings]
    semantic_indices = np.argsort(semantic_scores)[::-1]
    semantic_rank = {idx: rank + 1 for rank, idx in enumerate(semantic_indices)}

    # 3. Reciprocal Rank Fusion (RRF)
    rrf_scores = []
    for i in range(len(all_chunks)):
        score = 0
        if i in bm25_rank:
            score += 1.0 / (RRF_K + bm25_rank[i])
        if i in semantic_rank:
            score += 1.0 / (RRF_K + semantic_rank[i])
        rrf_scores.append(score)

    combined_indices = np.argsort(rrf_scores)[::-1]
    top_indices = combined_indices[:rerank_top_n]
    
    results = []
    for idx in top_indices:
        results.append({
            **all_chunks[idx],
            "score": float(rrf_scores[idx]),
            "bm25_score": float(bm25_scores[idx]),
            "semantic_score": float(semantic_scores[idx])
        })

    # 4. Cohere Re-ranking (Optional)
    cohere_key = os.getenv("COHERE_API_KEY", COHERE_KEY)
    if cohere_key and USE_RERANK:
        co = cohere.Client(cohere_key)
        response = co.rerank(
            model="rerank-english-v3.0", # or rerank-multilingual-v3.0
            query=query,
            documents=[r["text"] for r in results],
            top_n=top_k
        )
        reranked_results = []
        for r in response.results:
            orig = results[r.index]
            orig["re_rank_score"] = r.relevance_score
            reranked_results.append(orig)
        return reranked_results

    return results[:top_k]


# ─── Pydantic models ──────────────────────────────────────────────────────────
class QueryRequest(BaseModel):
    query: str
    history: list[dict] = []
    top_k: int = TOP_K
    api_key: Optional[str] = None
    provider: str = "anthropic" # anthropic, openai, google
    model: Optional[str] = None
    max_tokens: int = 2048
    stream: bool = False
    advanced: bool = False # New flag for Hybrid Search

class AddTextRequest(BaseModel):
    name: str
    content: str
    api_key: Optional[str] = None

#API
# ─── Routes ───────────────────────────────────────────────────────────────────
@app.get("/")
def root():
    index = Path("index.html")
    if index.exists():
        return FileResponse("index.html")
    return {"status": "ok", "docs": "GET /docs"}


@app.get("/api/documents")
def list_documents():
    return [
        {
            "id":         doc_id,
            "name":       doc["name"],
            "ext":        doc["ext"],
            "char_count": doc["char_count"],
            "chunks":     len(doc["chunks"]),
        }
        for doc_id, doc in DOCUMENTS.items()
    ]


@app.post("/api/documents/upload")
async def upload_document(file: UploadFile = File(...)):
    raw = await file.read()
    text = extract_text(file.filename, raw)
    if not text.strip():
        raise HTTPException(400, "Could not extract any text from the file.")

    doc_id = str(uuid.uuid4())
    chunks = chunk_text(text)
    DOCUMENTS[doc_id] = {
        "id":         doc_id,
        "name":       file.filename,
        "ext":        Path(file.filename).suffix.lower(),
        "char_count": len(text),
        "chunks":     chunks,
    }
    log.info(f"Uploaded: {file.filename} → {len(chunks)} chunks")
    return {"id": doc_id, "name": file.filename, "chunks": len(chunks)}


@app.post("/api/documents/text")
def add_text_document(req: AddTextRequest):
    if not req.content.strip():
        raise HTTPException(400, "Content is empty.")
    doc_id = str(uuid.uuid4())
    chunks = chunk_text(req.content)
    DOCUMENTS[doc_id] = {
        "id":         doc_id,
        "name":       req.name or f"text_{doc_id[:8]}",
        "ext":        ".txt",
        "char_count": len(req.content),
        "chunks":     chunks,
    }
    return {"id": doc_id, "name": req.name, "chunks": len(chunks)}


@app.delete("/api/documents/{doc_id}")
def delete_document(doc_id: str):
    if doc_id not in DOCUMENTS:
        raise HTTPException(404, "Document not found.")
    name = DOCUMENTS.pop(doc_id)["name"]
    log.info(f"Deleted: {name}")
    return {"deleted": doc_id}


@app.delete("/api/documents")
def clear_all_documents():
    DOCUMENTS.clear()
    return {"cleared": True}


@app.post("/api/query")
def query(req: QueryRequest):
    provider = req.provider.lower()
    
    # 1. API Key Resolution
    if provider == "anthropic":
        api_key = req.api_key or ANTHROPIC_KEY
        if not api_key: raise HTTPException(400, "ANTHROPIC_API_KEY not set.")
        client = anthropic.Anthropic(api_key=api_key)
        default_model = "claude-3-5-sonnet-20240620"
    elif provider == "openai":
        api_key = req.api_key or OPENAI_KEY
        if not api_key: raise HTTPException(400, "OPENAI_API_KEY not set.")
        client = openai.OpenAI(api_key=api_key)
        default_model = "gpt-4o"
    elif provider == "google":
        api_key = req.api_key or GOOGLE_KEY
        if not api_key: raise HTTPException(400, "GOOGLE_API_KEY not set.")
        genai.configure(api_key=api_key)
        default_model = "gemini-1.5-flash"
    else:
        raise HTTPException(400, f"Unsupported provider: {provider}")

    model = req.model or default_model

    # 2. Retrieval
    if req.advanced:
        retrieved = hybrid_retrieve(req.query, req.top_k)
    else:
        retrieved = retrieve(req.query, req.top_k)
    if retrieved:
        context_parts = [
            f"[Source {i+1}: {c['doc_name']} — chunk {c['chunk_idx']+1}]\n{c['text']}"
            for i, c in enumerate(retrieved)
        ]
        context = "\n\n---\n\n".join(context_parts)
    else:
        context = "ไม่พบเอกสารใน Knowledge Base กรุณาอัปโหลดเอกสารก่อน"

    system_prompt = f"""You are a helpful RAG assistant with access to a knowledge base.
Answer the user's question using the context below.
If the context contains relevant information, use it and cite the source(s).
If the context is insufficient, say so clearly and answer from general knowledge.
Respond in the same language as the user's question (Thai or English).

=== CONTEXT ===
{context}
==============="""

    # 3. Message Prep
    messages = []
    # Convert history to provider-specific format if needed, but here we assume common user/assistant roles
    for m in req.history[-10:]:
        messages.append({"role": m["role"], "content": m["content"]})
    messages.append({"role": "user", "content": req.query})

    # 4. Provider Logic
    if req.stream:
        def generate():
            if provider == "anthropic":
                with client.messages.stream(
                    model=model, max_tokens=req.max_tokens or 2048, system=system_prompt, messages=messages
                ) as stream:
                    for text in stream.text_stream:
                        yield f"data: {json.dumps({'text': text})}\n\n"
            
            elif provider == "openai":
                # OpenAI uses system as a message role
                oa_messages = [{"role": "system", "content": system_prompt}] + messages
                stream = client.chat.completions.create(
                    model=model, messages=oa_messages, stream=True, max_tokens=req.max_tokens or 2048
                )
                for chunk in stream:
                    delta = chunk.choices[0].delta.content
                    if delta:
                        yield f"data: {json.dumps({'text': delta})}\n\n"
            
            elif provider == "google":
                # Gemini
                gemini_model = genai.GenerativeModel(model_name=model, system_instruction=system_prompt)
                # Google expected history differently, but for simplicity we'll send as one-off or use chat session
                content_history = []
                for m in messages[:-1]:
                    role = "user" if m["role"] == "user" else "model"
                    content_history.append({"role": role, "parts": [m["content"]]})
                
                chat = gemini_model.start_chat(history=content_history)
                response = chat.send_message(messages[-1]["content"], stream=True)
                for chunk in response:
                    yield f"data: {json.dumps({'text': chunk.text})}\n\n"

            yield f"data: {json.dumps({'done': True, 'sources': retrieved, 'model': model})}\n\n"

        return StreamingResponse(generate(), media_type="text/event-stream")

    # Non-streaming
    if provider == "anthropic":
        res = client.messages.create(model=model, max_tokens=req.max_tokens or 2048, system=system_prompt, messages=messages)
        answer = res.content[0].text
        usage = {"input_tokens": res.usage.input_tokens, "output_tokens": res.usage.output_tokens}
    
    elif provider == "openai":
        oa_messages = [{"role": "system", "content": system_prompt}] + messages
        res = client.chat.completions.create(model=model, messages=oa_messages, max_tokens=req.max_tokens or 2048)
        answer = res.choices[0].message.content
        usage = {"input_tokens": res.usage.prompt_tokens, "output_tokens": res.usage.completion_tokens}
    
    elif provider == "google":
        gemini_model = genai.GenerativeModel(model_name=model, system_instruction=system_prompt)
        content_history = []
        for m in messages[:-1]:
            role = "user" if m["role"] == "user" else "model"
            content_history.append({"role": role, "parts": [m["content"]]})
        chat = gemini_model.start_chat(history=content_history)
        res = chat.send_message(messages[-1]["content"])
        answer = res.text
        usage = {"input_tokens": 0, "output_tokens": 0} # Google doesn't return usage in the same way easily here

    return {
        "answer":  answer,
        "sources": retrieved,
        "model":   model,
        "usage":   usage,
    }


@app.get("/api/health")
def health():
    return {
        "status":    "ok",
        "documents": len(DOCUMENTS),
        "chunks":    sum(len(d["chunks"]) for d in DOCUMENTS.values()),
        "model":     MODEL,
    }


# ─── Dev runner ───────────────────────────────────────────────────────────────
if __name__ == "__main__":
    import uvicorn
    print("\n RAG System starting...")
    print(" API docs: http://localhost:8000/docs")
    print(" Web UI:   http://localhost:8000\n")
    uvicorn.run("server:app", host="0.0.0.0", port=8000, reload=True)