dashboard-cpsp/ai-insight/rag_service.py

"""
RAG microservice for the CP 707 knowledge base (ChromaDB + SentenceTransformers).

Runtime entrypoint used by Django `insight_service.fetch_rag_chunks`
via HTTP `RAG_SERVICE_URL` (default compose :5002; NUC AI lab often :5102).
Offline HuggingFace mode — do not confuse with `query_rag.py` (CLI/dev only).
"""

import os
import sys
from pathlib import Path
from dotenv import load_dotenv

# Set offline agar tidak download dari HuggingFace
os.environ["HF_HUB_OFFLINE"] = "1"
os.environ["TRANSFORMERS_OFFLINE"] = "1"

# Load .env dari direktori yang sama dengan script ini
script_dir = Path(__file__).parent
load_dotenv(script_dir / ".env")

CHROMA_DB_DIR = str(script_dir / os.getenv("CHROMA_DB_DIR", "chroma_db"))
EMBEDDING_MODEL_NAME = os.getenv("EMBEDDING_MODEL_NAME", "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2")
RAG_PORT = int(os.getenv("RAG_PORT", "5002"))
COLLECTION_NAME = "company_sop"

print(f"[RAG Service] ChromaDB path: {CHROMA_DB_DIR}")
print(f"[RAG Service] Embedding model: {EMBEDDING_MODEL_NAME}")

# Import setelah env set
import chromadb
from sentence_transformers import SentenceTransformer
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field
import uvicorn

# ─── Init ChromaDB & Embedding Model ─────────────────────────────────────────
print("[RAG Service] Memuat ChromaDB...")
try:
    chroma_client = chromadb.PersistentClient(path=CHROMA_DB_DIR)
    collection = chroma_client.get_collection(name=COLLECTION_NAME)
    total_chunks = collection.count()
    print(f"[RAG Service] ChromaDB loaded. Total chunks: {total_chunks}")
except Exception as e:
    print(f"[RAG Service] ERROR: Gagal load ChromaDB: {e}")
    sys.exit(1)

print(f"[RAG Service] Memuat embedding model '{EMBEDDING_MODEL_NAME}'...")
try:
    embedding_model = SentenceTransformer(EMBEDDING_MODEL_NAME)
    print("[RAG Service] Embedding model berhasil dimuat.")
except Exception as e:
    print(f"[RAG Service] ERROR: Gagal load embedding model: {e}")
    print("[RAG Service] Pastikan model sudah didownload. Jalankan populate_db.py terlebih dahulu.")
    sys.exit(1)

# ─── FastAPI App ──────────────────────────────────────────────────────────────
app = FastAPI(
    title="CP707 RAG Service",
    description="Retrieval-Augmented Generation service untuk buku Manajemen Broiler CP 707",
    version="1.0.0"
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=[
        "http://localhost:3000",
        "http://localhost:3001",
        "http://localhost:5001",
        "http://127.0.0.1:5001",
        "http://localhost:8000",
        "http://127.0.0.1:8000",
    ],
    allow_methods=["GET", "POST"],
    allow_headers=["*"],
)

# ─── Request/Response Models ─────────────────────────────────────────────────
class QueryRequest(BaseModel):
    query: str = Field(..., min_length=1, description="Query text untuk mencari chunk CP707 relevan")
    n_results: int = Field(default=4, ge=1, le=10, description="Jumlah chunk yang dikembalikan")
    topic: str = Field(default="", description="Topic insight: berat_ayam, fcr, iot_panel, dll")
    # Prefer prose SOP; table rows without headers are dangerous for the LLM.
    tipe: str = Field(default="prosa", description="Filter metadata tipe: prosa | tabel | any")

class QueryResponse(BaseModel):
    success: bool
    chunks: list
    sources: list
    metadatas: list
    total_found: int
    query_used: str

class HealthResponse(BaseModel):
    status: str
    total_chunks: int
    embedding_model: str

# ─── Topic → Query enhancement mapping ──────────────────────────────────────
# Tambahkan keyword relevan per topic agar embedding search lebih tepat sasaran
TOPIC_QUERY_HINTS = {
    "berat_ayam":    "berat badan target bobot ADG pertumbuhan standar mingguan ayam broiler",
    "fcr":           "FCR feed conversion ratio konsumsi pakan efisiensi standar broiler",
    "iot_panel":     "suhu kandang kelembapan amonia CO2 ventilasi lingkungan pemeliharaan broiler",
    "eef":           "EEF indeks performa IP efisiensi produksi siklus broiler",
    "hitung_ayam":   "mortalitas deplesi kematian afkir populasi standar toleransi broiler",
    "hitung_karung": "pakan karung konsumsi harian feed intake standar broiler",
}

# ─── Endpoints ───────────────────────────────────────────────────────────────
@app.get("/health", response_model=HealthResponse)
def health_check():
    return HealthResponse(
        status="ok",
        total_chunks=collection.count(),
        embedding_model=EMBEDDING_MODEL_NAME,
    )

@app.post("/query", response_model=QueryResponse)
def query_cp707(req: QueryRequest):
    """
    Cari chunk CP707 yang relevan berdasarkan query.
    Jika topic disediakan, tambahkan hint keyword agar hasil lebih relevan.
    """
    enhanced_query = req.query
    if req.topic and req.topic in TOPIC_QUERY_HINTS:
        enhanced_query = f"{req.query} {TOPIC_QUERY_HINTS[req.topic]}"

    try:
        query_embedding = embedding_model.encode([enhanced_query])[0].tolist()

        # Over-fetch then filter by tipe so prosa chunks win when metadata exists.
        fetch_n = min(max(req.n_results * 3, req.n_results), max(collection.count(), 1))
        where = None
        if req.tipe and req.tipe != "any":
            where = {"tipe": req.tipe}

        try:
            results = collection.query(
                query_embeddings=[query_embedding],
                n_results=fetch_n,
                where=where,
            )
        except Exception:
            # Older indexes may lack tipe metadata — fall back unfiltered.
            results = collection.query(
                query_embeddings=[query_embedding],
                n_results=fetch_n,
            )

        chunks = results["documents"][0] if results["documents"] else []
        metadatas = results["metadatas"][0] if results["metadatas"] else []

        # If unfiltered fallback returned tables, drop them when prosa was requested.
        if req.tipe == "prosa" and metadatas:
            paired = [
                (c, m)
                for c, m in zip(chunks, metadatas)
                if (m or {}).get("tipe", "prosa") != "tabel"
            ]
            if paired:
                chunks, metadatas = [list(x) for x in zip(*paired)]
            else:
                # Keep original if everything was tabel (better something than nothing;
                # Django stripHeaderlessTables still cleans numeric rows).
                pass

        chunks = chunks[: req.n_results]
        metadatas = metadatas[: req.n_results]

        sources = []
        for m in metadatas:
            bab = (m or {}).get("bab") or ""
            src = (m or {}).get("source", "unknown")
            idx = (m or {}).get("chunk_index", "?")
            label = f"{src}"
            if bab:
                label += f" / {bab}"
            label += f" (chunk {idx})"
            sources.append(label)

        return QueryResponse(
            success=True,
            chunks=chunks,
            sources=sources,
            metadatas=metadatas,
            total_found=len(chunks),
            query_used=enhanced_query,
        )

    except Exception as e:
        raise HTTPException(status_code=500, detail=f"RAG query error: {str(e)}")


if __name__ == "__main__":
    print(f"[RAG Service] Starting on http://0.0.0.0:{RAG_PORT}")
    uvicorn.run(app, host="0.0.0.0", port=RAG_PORT, log_level="warning")