Update app.py

app.py

@@ -1,17 +1,31 @@
 from fastapi import FastAPI, Header, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel
 import torch
 import os
 import json
+import datetime
 from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 from duckduckgo_search import DDGS
 
 app = FastAPI()
 
-# 1. API Keys 50 ක ලැයිස්තුව (Hardcoded for now as requested)
-# Format: ELE-PRIME-001, ELE-PRIME-002 ... ELE-PRIME-050
-API_KEYS_DB = {f"ELE-PRIME-{i:03d}": {"credits": 5000, "status": "active"} for i in range(1, 51)}
+# CORS Settings
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# --- දත්ත ගබඩාව සහ මතකය (Learning Path) ---
+API_KEYS_DB = {
+    "ELE-PRIME-ADMIN-SYS": {"credits": 999999, "status": "active"}
+}
+ADMIN_SECRET = "MINZO-SECRET-2026"
+LEARNING_VAULT_PATH = "neural_learning_data.jsonl" # මෙතන තමයි AI එක ඉගෙන ගන්න දත්ත Save වෙන්නේ
 
-# 2. GPU පවතිනවාදැයි පරීක්ෂා කිරීම සහ Quantization සැකසීම
+# --- AI Model සැකසුම් ---
 model_id = "mistralai/Mistral-7B-v0.3"
 HF_TOKEN = os.getenv("HF_TOKEN")
 
@@ -22,8 +36,7 @@ quant_config = BitsAndBytesConfig(
     bnb_4bit_use_double_quant=True,
 )
 
-# 3. මොඩලය Load කිරීම
-print("Loading Elephant Engine (Mistral-7B)...")
+print("🐘 Elephant Learning Engine Loading...")
 tokenizer = AutoTokenizer.from_pretrained(model_id, token=HF_TOKEN)
 model = AutoModelForCausalLM.from_pretrained(
     model_id,
@@ -32,54 +45,84 @@ model = AutoModelForCausalLM.from_pretrained(
     token=HF_TOKEN
 )
 
-# 4. Web Search පහසුකම
+# --- Helpers ---
+class NewKeyRequest(BaseModel):
+    admin_pass: str
+    new_key: str
+
 def get_live_data(query):
     try:
         with DDGS() as ddgs:
             results = [r['body'] for r in ddgs.text(query, max_results=3)]
             return "\n".join(results)
-    except:
-        return ""
+    except: return ""
+
+# 🧠 CONTINUOUS LEARNING FUNCTION
+# පද්ධතිය විසින් අලුත් දැනුම ගබඩා කරගන්නා ආකාරය
+def capture_learning_data(query, context, response, key_id):
+    timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+    learning_entry = {
+        "timestamp": timestamp,
+        "key_node": key_id,
+        "instruction": query,
+        "external_context": context,
+        "trained_output": response
+    }
+    # JSONL ගොනුවකට දත්ත එකතු කිරීම
+    with open(LEARNING_VAULT_PATH, "a", encoding="utf-8") as f:
+        f.write(json.dumps(learning_entry) + "\n")
+    print(f"📊 Neural Data Captured via {key_id}")
+
+# --- Endpoints ---
 
 @app.get("/")
-def health_check():
-    return {"status": "Elephant API Node 2026 is Active", "keys_loaded": len(API_KEYS_DB)}
+def home():
+    # පද්ධතිය කොච්චර ඉගෙන ගෙන තියෙනවද කියලා බලාගන්න
+    learning_count = 0
+    if os.path.exists(LEARNING_VAULT_PATH):
+        with open(LEARNING_VAULT_PATH, "r") as f:
+            learning_count = sum(1 for line in f)
+            
+    return {
+        "status": "Elephant AI Node 2026 Live",
+        "active_keys": len(API_KEYS_DB),
+        "learning_entries_captured": learning_count
+    }
+
+@app.post("/admin/add-key")
+async def register_key(data: NewKeyRequest):
+    if data.admin_pass != ADMIN_SECRET:
+        raise HTTPException(status_code=401, detail="Unauthorized Admin Access")
+    API_KEYS_DB[data.new_key] = {"credits": 5000, "status": "active"}
+    return {"message": f"Key {data.new_key} Registered Successfully"}
 
 @app.post("/v1/chat")
-async def chat_api(message: dict, x_api_key: str = Header(None)):
-    # API Key එක පරීක්ෂා කිරීම
+async def chat_endpoint(message: dict, x_api_key: str = Header(None)):
     if x_api_key not in API_KEYS_DB:
-        raise HTTPException(status_code=403, detail="Unauthorized: Invalid API Key")
+        raise HTTPException(status_code=403, detail="Invalid API Key")
     
     user_query = message.get("query", "")
-    
-    # 2026 දත්ත සඳහා Web Search කිරීම
     context = ""
-    if any(word in user_query.lower() for word in ["today", "now", "2026", "news", "current"]):
+    
+    # 🌐 2026 Live Web Search
+    if any(word in user_query.lower() for word in ["today", "now", "2026", "news", "price"]):
         context = get_live_data(user_query)
 
-    # Prompt එක සැකසීම
-    system_instr = "Current Year: 2026. You are Elephant AI. Use the provided context to answer."
-    full_prompt = f"System: {system_instr}\nContext: {context}\nUser: {user_query}\nAssistant:"
-    
+    full_prompt = f"System: Year 2026. Context: {context}\nUser: {user_query}\nAssistant:"
     inputs = tokenizer(full_prompt, return_tensors="pt").to("cuda")
     
-    # Response එක Generate කිරීම
     with torch.no_grad():
         output_tokens = model.generate(**inputs, max_new_tokens=300, do_sample=True, temperature=0.7)
     
     response = tokenizer.decode(output_tokens[0], skip_special_tokens=True).split("Assistant:")[-1].strip()
 
-    # Learning Loop: පද්ධතිය ඉගෙන ගැනීමට දත්ත ලොග් කිරීම
-    with open("learning_vault.jsonl", "a") as f:
-        log_entry = {"q": user_query, "ctx": context, "ans": response, "key": x_api_key}
-        f.write(json.dumps(log_entry) + "\n")
+    # 🧠 දියුණු වීමට අවශ්‍ය දත්ත ගබඩා කිරීම (Learning Trigger)
+    capture_learning_data(user_query, context, response, x_api_key)
 
     return {
-        "reply": response,
-        "model": "Elephant-Mistral-7B-v0.3",
+        "reply": response, 
         "key_id": x_api_key,
-        "timestamp": "2026-04-27"
+        "learning_status": "synced"
     }
 
 main = app