149 lines
5.9 KiB
Python
149 lines
5.9 KiB
Python
"""
|
|
RAG Microservice untuk CP 707 Knowledge Base
|
|
Berjalan di port 5002, dipanggil oleh backend Node.js
|
|
Menggunakan ChromaDB + SentenceTransformers (offline mode)
|
|
"""
|
|
|
|
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:5001", "http://127.0.0.1:5001"],
|
|
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")
|
|
|
|
class QueryResponse(BaseModel):
|
|
success: bool
|
|
chunks: list
|
|
sources: 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()
|
|
|
|
results = collection.query(
|
|
query_embeddings=[query_embedding],
|
|
n_results=min(req.n_results, collection.count()),
|
|
)
|
|
|
|
chunks = results["documents"][0] if results["documents"] else []
|
|
metadatas = results["metadatas"][0] if results["metadatas"] else []
|
|
|
|
sources = [
|
|
f"{m.get('source', 'unknown')} (chunk {m.get('chunk_index', '?')})"
|
|
for m in metadatas
|
|
]
|
|
|
|
return QueryResponse(
|
|
success=True,
|
|
chunks=chunks,
|
|
sources=sources,
|
|
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")
|