Update app.py

app.py

@@ -1,133 +1,105 @@
+# app.py
+
 import os
+import time
 import gradio as gr
-from openai import OpenAI
+import importlib.util
+
+from huggingface_hub import hf_hub_download
 
 # ----------------------------------------------------------------------
-# Helper to read a secret (fallback is useful when you run locally)
+# Helper to read secrets from the HF Space environment
 # ----------------------------------------------------------------------
-def _secret(key: str, fallback: str = "") -> str:
-    """Return the value of a secret or the supplied fallback."""
-    return os.getenv(key, fallback)
-
+def _secret(key: str, fallback: str = None) -> str:
+    val = os.getenv(key)
+    if val is not None: return val
+    if fallback is not None: return fallback
+    raise RuntimeError(f"Secret '{key}' not found. Please add it to your Space secrets.")
 
 # ----------------------------------------------------------------------
-# Core chat logic – system prompt comes from the secret `prec_chat`
+# 1. Configuration & Constants
 # ----------------------------------------------------------------------
-def respond(
-    message: str,
-    history: list[dict[str, str]],
-    # model_name: str,  
-):
-    """
-    Generate a response using OpenRouter API via OpenAI client.
-    
-    * System prompt = secret `prec_chat`
-    * OpenRouter API key = secret `OPENROUTER_API_KEY`
-    """
-    # 1️⃣ Load the system prompt (fallback = generic assistant)
-    # system_message = _secret("prec_chat", "You are a helpful assistant.")
+# The private repo containing the vector DB and the logic script
+REPO_ID = _secret("REPO_ID") 
+# Files to download from the repo
+FILES_TO_DOWNLOAD = ["index.faiss", "index.pkl", "agent_logic.py"]
+# A local directory to store all downloaded assets
+LOCAL_DOWNLOAD_DIR = "downloaded_assets"
+EMBEDDING_MODEL_NAME = "google/embeddinggemma-300m" 
 
-    # 2️⃣ Load the OpenRouter API key
-    openrouter_api_key = _secret("OPENROUTER_API_KEY")
-    if not openrouter_api_key:
-        raise RuntimeError(
-            "OPENROUTER_API_KEY not found in secrets. Add it to secrets.toml (or via the Space UI)."
+# ----------------------------------------------------------------------
+# 2. Bootstrap Phase: Download assets and initialize the engine
+#    (This code runs only once when the Space starts up)
+# ----------------------------------------------------------------------
+print("--- [UI App] Starting bootstrap process ---")
+os.makedirs(LOCAL_DOWNLOAD_DIR, exist_ok=True)
+hf_token = _secret("HF_TOKEN") # A read-access token is required for private repos
+
+for filename in FILES_TO_DOWNLOAD:
+    print(f"--- [UI App] Downloading '{filename}'... ---")
+    try:
+        hf_hub_download(
+            repo_id=REPO_ID, filename=filename, repo_type="dataset",
+            local_dir=LOCAL_DOWNLOAD_DIR, token=hf_token,
         )
-
-    # 3️⃣ Initialize OpenAI client with OpenRouter configuration
-    # client = OpenAI(
-    #     base_url="https://openrouter.ai/api/v1",        
-    #     api_key=openrouter_api_key,
-    # )
-
-    client = OpenAI(
-    base_url="https://openrouter.ai/api/v1",
-    api_key=openrouter_api_key,
-    default_headers={
-        "X-Title": "hf_precise",        # Optional. Set your app name for tracking.
-    },
+    except Exception as e:
+        raise RuntimeError(f"Failed to download '{filename}'. Check repo/file names and HF_TOKEN. Error: {e}")
+
+# Dynamically import the RAG_Engine class from the downloaded script
+logic_script_path = os.path.join(LOCAL_DOWNLOAD_DIR, "agent_logic.py")
+spec = importlib.util.spec_from_file_location("agent_logic", logic_script_path)
+agent_logic_module = importlib.util.module_from_spec(spec)
+spec.loader.exec_module(agent_logic_module)
+print("--- [UI App] Agent logic module imported successfully. ---")
+
+# Instantiate the engine. This single line triggers all the complex setup
+# defined in the private_logic.py file.
+engine = agent_logic_module.RAG_Engine(
+    local_download_dir=LOCAL_DOWNLOAD_DIR,
+    embedding_model_name=EMBEDDING_MODEL_NAME
 )
+print("--- [UI App] Bootstrap complete. Gradio UI is starting. ---")
 
-    # 4️⃣ Build the message list for the chat completion
-    messages = []
-    messages.extend(history)                     # previous conversation turns
-    messages.append({"role": "user", "content": message})  # current user query
-
-
-    MODEL_NAME = _secret("MODEL_NAME")
+# ----------------------------------------------------------------------
+# 3. Core Gradio Chat Logic (Now a simple wrapper)
+# ----------------------------------------------------------------------
+def respond(message: str, history: list[dict[str, str]]):
+    """
+    This function is called by Gradio for each user message.
+    It passes the inputs to the RAG engine and streams the output.
+    """
+    final_response = engine.get_response(message, history)
     
-    # 5️⃣ Stream the response back to the UI
+    # Stream the response back to the UI for a "typing" effect
     response = ""
-    stream = client.chat.completions.create(
-        # model=model_name,
-        model = MODEL_NAME,
-        messages = messages,
-        reasoning_effort= "high",  # Use high reasoning effort
-        max_tokens=8096,
-        stream=True
-    )
-    
-    for chunk in stream:
-        if chunk.choices[0].delta.content is not None:
-            token = chunk.choices[0].delta.content
-            response += token
-            yield response
-
-
-# # ──────────────────────────────────────────────────────────────────────
-# # List of models available through OpenRouter
-# # ──────────────────────────────────────────────────────────────────────
-AVAILABLE_MODELS = [
-]
+    for char in final_response:
+        response += char
+        time.sleep(0.01)
+        yield response
 
 # ----------------------------------------------------------------------
-# UI – the system‑prompt textbox has been removed.
+# 4. UI Layout and Launch
 # ----------------------------------------------------------------------
 chatbot = gr.ChatInterface(
     respond,
     type="messages",
-    # additional_inputs=[
-    #     gr.Dropdown(
-    #         choices=AVAILABLE_MODELS,
-    #         value=AVAILABLE_MODELS[0],
-    #         label="Model",
-    #         interactive=True,
-    #     ),        
-    # ],
+    title="PRECISE RAG Agent",
+    description="Silakan bertanya tentang PRECISE.",
+    examples=[
+        ["Apa rumus untuk menghitung PVR?"],
+        ["Apa tujuan pengadaan PRECISE?"],
+    ],
+    theme=gr.themes.Soft(),
 )
 
-# ----------------------------------------------------------------------
-# Assemble the Blocks layout (no LoginButton – we use basic auth)
-# ----------------------------------------------------------------------
 with gr.Blocks() as demo:
     chatbot.render()
 
-
-# ----------------------------------------------------------------------
-# Launch – protect the UI with the credentials from secrets.
-# ----------------------------------------------------------------------
 if __name__ == "__main__":
-    # ------------------------------------------------------------------
-    # 1️⃣ Pull the allowed credentials from secrets (fail fast if missing)
-    # ------------------------------------------------------------------
     allowed_user = _secret("CHAT_USER")
     allowed_pass = _secret("CHAT_PASS")
-
-    if not allowed_user or not allowed_pass:
-        raise RuntimeError(
-            "Authentication credentials not found in secrets. "
-            "Add CHAT_USER and CHAT_PASS to secrets.toml (or via the HF Spaces UI)."
-        )
-
-    # ------------------------------------------------------------------
-    # 2️⃣ Launch
-    # ------------------------------------------------------------------
     demo.launch(
-        auth=(allowed_user, allowed_pass),   # <-- Gradio's built‑in basic auth
-        ssr_mode=False,                     # <-- avoids the i18n locale error
-        # In a Space we **must not** set share=True (Spaces already give a public URL)
-        # If you run locally and want a shareable link, add share=True here.
-        server_name="0.0.0.0",              # listen on all interfaces (needed in containers)
-        # Optional: give the app a nice title
-        # title="Secure Chatbot",
+        auth=(allowed_user, allowed_pass),
+        server_name="0.0.0.0",
+        server_port=7860
     )