Manok45/jawi-chatbot-api

1from flask import Flask, request, jsonify
2from sentence_transformers import SentenceTransformer
3import faiss
4import numpy as np
5import json
6import google.generativeai as genai
7import os # Baru: Import modul os
8
9app = Flask(__name__)
10
11# --- KONFIGURASI ---
12# Baru: Mengambil API key dari variabel lingkungan
13GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY")
14
15# Pastikan API Key tersedia
16if not GEMINI_API_KEY:
17    raise ValueError("GEMINI_API_KEY is not set as an environment variable.")
18
19genai.configure(api_key=GEMINI_API_KEY)
20
21# --- MUAT MODEL & DATABASE SAAT SERVER START ---
22print("Memuat model dan database...")
23retriever_model = SentenceTransformer('all-MiniLM-L6-v2')
24jawi_index = faiss.read_index('jawi_index.faiss')
25with open('documents.json', 'r', encoding='utf-8') as f:
26    documents = json.load(f)
27
28# Menggunakan model Gemini 1.5 Flash yang lebih cepat dan efisien
29generation_model = genai.GenerativeModel('gemini-1.5-flash-latest') 
30print("Server siap dengan model Gemini 1.5 Flash!")
31
32@app.route('/chat', methods=['POST'])
33def chat():
34    data = request.json
35    user_query = data.get('query')
36
37    if not user_query:
38        return jsonify({"error": "Query tidak ditemukan"}), 400
39
40    # 1. Retrieval
41    query_embedding = retriever_model.encode([user_query])
42    D, I = jawi_index.search(np.array(query_embedding).astype('float32'), k=2)
43    retrieved_context = "\n\n".join([documents[i] for i in I[0]])
44
45    # 2. Augmentation
46    prompt = f"""
47    Anda adalah seorang ahli aksara Jawi yang ramah. Jawab pertanyaan berikut HANYA berdasarkan konteks yang diberikan.
48    Konteks:
49    ---
50    {retrieved_context}
51    ---
52    Pertanyaan: {user_query}
53    Jawaban:
54    """
55
56    # 3. Generation
57    try:
58        response = generation_model.generate_content(prompt)
59        return jsonify({"response": response.text})
60    except Exception as e:
61        print(f"Error: {e}")
62        return jsonify({"error": str(e)}), 500
63
64if __name__ == '__main__':
65    app.run(host='0.0.0.0', port=5000)