rag.py

import os
import logging

from pathlib import Path
from dotenv import load_dotenv

from langchain_community.vectorstores import FAISS
from langchain_community.document_loaders import PyPDFLoader
from langchain_community.embeddings import OpenAIEmbeddings
from langchain_openai import ChatOpenAI
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_core.prompts import ChatPromptTemplate

logger = logging.getLogger(__name__)

BASE_DIR = Path(__file__).resolve().parent
ENV_PATH = BASE_DIR / ".env"
if ENV_PATH.exists():
    load_dotenv(dotenv_path=ENV_PATH, override=False)

# --- Konstanta global ---
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
PDF_DIR = "documents"
VECTOR_DIR = "vector_store"
VECTOR_INDEX_PATH = os.path.join(VECTOR_DIR, "faiss_index")
_CHAT_KW = {"model": "gpt-4o-mini", "api_key": OPENAI_API_KEY}

MASKOT_NAME = "Hadmin AI"
SYSTEM_INSTRUCTION = f"""
Kamu adalah maskot chat AI bernama {MASKOT_NAME}, mewakili perusahaan. 
Kamu ramah, sopan, profesional, dan mudah diajak bicara — seperti customer service yang hangat dan ahli.

Tugas:
1. Pahami konteks dan kebutuhan pengguna dengan cermat.
2. Berikan saran, rekomendasi, dan informasi relevan secara jelas, logis, dan solutif.
3. Jika informasi kurang lengkap atau tidak tersedia:
   - Sampaikan dengan sopan bahwa informasi terkait tidak tersedia.
   - Tawarkan alternatif atau sumber resmi yang relevan (misal website, pendaftaran, layanan terkait).
   - Jelaskan mengapa alternatif tersebut bisa menjadi opsi yang cocok.
4. Informasi yang sudah tersedia **tidak boleh diubah, dikarang, atau disesuaikan** atas permintaan pengguna.
5. Gunakan nama {MASKOT_NAME} **hanya saat menyapa pengguna atau membalas sapaan pertama**.
6. Jangan menganggap diri sebagai manusia; kamu adalah maskot/Asisten Virtual.
7. Jawaban harus empatik, mudah dipahami, natural, profesional, dan nyaman dibaca.

Format HTML:
- Gunakan <p> untuk paragraf.
- Gunakan <ul>/<li> untuk daftar.
- Gunakan <b>/<strong> untuk penekanan.
- Gunakan <table>/<tr>/<td>/<th> untuk tabel.
- Jangan menggunakan <h1>–<h5>, CSS, warna, atau layout kompleks.
- Jika menyertakan link, gunakan <a href="URL" target="_blank">nama link</a>.

Perilaku:
- Saat pengguna menyapa: balas hangat, cukup tanyakan kebutuhan mereka.
- Saat pengguna meminta saran/rekomendasi: berikan beberapa opsi relevan, jelaskan alasannya, dan tawarkan solusi praktis.
- Jika info dari pengguna kurang jelas atau ambigu: ajukan pertanyaan klarifikasi sebelum menjawab.
- Sesuaikan bahasa jawaban dengan bahasa pesan pengguna:
   - Inggris → jawaban dalam bahasa Inggris.
   - Prancis → jawaban dalam bahasa Prancis.
   - Bahasa lain → jawab dengan bahasa yang sama.
   - Pertahankan istilah atau nama properti unik dalam bahasa Indonesia (misal: “BPJS Kesehatan”, nama kota, produk).
- Format HTML, daftar, dan tabel tetap sama.
- Jika ditanya kemampuan bahasa (misal: "bisa kamu berbicara bahasa Indonesia?"):
   - Jawab **singkat dan jelas** sesuai bahasa yang diminta.
   - Tidak perlu menambahkan saran, rekomendasi, atau info lain.
- Jangan menyapa pada setiap jawaban; sapaan hanya pada interaksi pertama atau bila diperlukan.
- Jangan menyatakan diri sebagai AI, menyebut model OpenAI, atau membahas batasan teknologi.
- Jangan memberikan informasi palsu, menyesatkan, atau tindakan ilegal/berbahaya.
- Jangan memberikan nasihat medis, hukum, atau keuangan profesional.
"""

# ------------------ INIT ------------------
def get_embeddings():
    """Inisialisasi embedding OpenAI."""
    try:
        return OpenAIEmbeddings(api_key=OPENAI_API_KEY)
    except Exception as e:
        logger.warning(f"OpenAIEmbeddings init failed: {e}")
        raise RuntimeError("Embeddings gagal diinisialisasi. Pastikan OPENAI_API_KEY sudah benar.")


# ------------------ BUILD ------------------
def build_vector_index(pdf_dir: str = PDF_DIR, vector_path: str = VECTOR_INDEX_PATH):
    """
    Bangun FAISS index dari seluruh PDF di folder pdf_dir.
    Simpan ke vector_path.
    """
    embeddings = get_embeddings()
    os.makedirs(vector_path, exist_ok=True)

    pdf_files = [f for f in os.listdir(pdf_dir) if f.lower().endswith(".pdf")]
    if not pdf_files:
        raise ValueError(f"Tidak ada PDF ditemukan di folder: {pdf_dir}")

    splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
    all_docs = []

    for file in pdf_files:
        path = os.path.join(pdf_dir, file)
        try:
            loader = PyPDFLoader(path)
            docs = loader.load_and_split(text_splitter=splitter)
            for d in docs:
                d.metadata["source"] = path
            all_docs.extend(docs)
            logger.info(f"✅ Berhasil memproses {file}, total potongan: {len(docs)}")
        except Exception as e:
            logger.error(f"❌ Gagal memproses {file}: {e}")

    if not all_docs:
        raise ValueError("Tidak ada teks valid yang berhasil di-load dari PDF mana pun.")

    vectordb = FAISS.from_documents(all_docs, embeddings)
    vectordb.save_local(vector_path)

    return {
        "status": "success",
        "pdf_count": len(pdf_files),
        "chunks_total": len(all_docs),
        "vector_path": vector_path,
    }


# ------------------ SESSION ------------------
# Simpan percakapan per session (untuk prototipe, pakai memory)
chat_sessions = {}  # key = session_id, value = list of {"role": "user"/"assistant", "content": str}

def add_to_session(session_id: str, role: str, content: str):
    if session_id not in chat_sessions:
        chat_sessions[session_id] = []
    chat_sessions[session_id].append({"role": role, "content": content})
    

# ------------------ QUERY ------------------
def query_from_vector(query: str, session_id: str = None, vector_path: str = VECTOR_INDEX_PATH, k: int = 3):
    """
    Query ke FAISS index dan hasilkan jawaban singkat dari LLM.
    session_id -> untuk multi-turn chat
    output_format: "string" | "markdown" | "html"
    """
    embeddings = get_embeddings()
    if not os.path.exists(vector_path):
        raise FileNotFoundError(f"Vector DB belum dibuat di {vector_path}")

    vectordb = FAISS.load_local(vector_path, embeddings, allow_dangerous_deserialization=True)
    docs = vectordb.similarity_search(query, k=k)
    if not docs:
        return {"error": "Tidak ditemukan hasil relevan."}

    context = "\n\n---\n\n".join([d.page_content for d in docs])
    sources = [d.metadata.get("source", "") for d in docs]

    # Ambil history chat sebelumnya
    history_messages = chat_sessions.get(session_id, []) if session_id else []

    # Siapkan prompt dengan history
    messages = [("system", SYSTEM_INSTRUCTION)]
    for msg in history_messages:
        messages.append((msg["role"], msg["content"]))
    messages.append(("human", f"Pertanyaan: {query}\n\nKonteks:\n{context}\n\nJawaban:"))

    prompt = ChatPromptTemplate.from_messages(messages)
    llm = ChatOpenAI(temperature=0, **_CHAT_KW)
    result = (prompt | llm).invoke({"q": query, "ctx": context})

    answer_text = (result.content or "").strip()

    # Simpan percakapan ke session
    if session_id:
        add_to_session(session_id, "user", query)
        add_to_session(session_id, "assistant", answer_text)

    token_usage = getattr(result, "usage_metadata", {})

    return {
        "query": query,
        "answer": answer_text,
        "sources": sources,
        "token_usage": {
            "input_tokens": token_usage.get("input_tokens", 0),
            "output_tokens": token_usage.get("output_tokens", 0),
            "total_tokens": token_usage.get("total_tokens", 0),
        },
    }