api/ml/rag_engine.py

import os
import faiss
import numpy as np
from sentence_transformers import SentenceTransformer

import json

# Paths for persistence
INDEX_PATH = os.path.join(os.path.dirname(__file__), "faiss_index.bin")
STORE_PATH = os.path.join(os.path.dirname(__file__), "doc_store.json")

# Initialize model
# all-MiniLM-L6-v2 is small, fast, and great for sentence embeddings
model = SentenceTransformer('all-MiniLM-L6-v2')

dimension = 384

# Map vector ID to text snippet
document_store = {}
current_id = 0

if os.path.exists(INDEX_PATH):
    index = faiss.read_index(INDEX_PATH)
else:
    index = faiss.IndexFlatL2(dimension)

if os.path.exists(STORE_PATH):
    with open(STORE_PATH, "r", encoding="utf-8") as f:
        stored_data = json.load(f)
        # Convert string keys back to int
        document_store = {int(k): v for k, v in stored_data.items()}
        current_id = max(document_store.keys()) + 1 if document_store else 0

def save_state():
    faiss.write_index(index, INDEX_PATH)
    with open(STORE_PATH, "w", encoding="utf-8") as f:
        json.dump(document_store, f, ensure_ascii=False, indent=2)

def add_documents(texts: list[str]):
    global current_id
    if not texts:
        return
        
    embeddings = model.encode(texts)
    index.add(np.array(embeddings).astype('float32'))
    
    for i, text in enumerate(texts):
        document_store[current_id + i] = text
        
    current_id += len(texts)
    
    # Save to disk after adding
    save_state()

def search_library(query: str, top_k: int = 3) -> list[str]:
    if index.ntotal == 0:
        return ["Perpustakaan Anda masih kosong. Tidak ada data jurnal yang bisa dicari."]
        
    query_vector = model.encode([query])
    distances, indices = index.search(np.array(query_vector).astype('float32'), top_k)
    
    results = []
    for idx in indices[0]:
        if idx != -1 and idx in document_store:
            results.append(document_store[idx])
            
    return results