Update app.py

app.py

@@ -10,7 +10,7 @@ from duckduckgo_search import DDGS
 
 app = FastAPI()
 
-# CORS Fix
+# 1. CORS Configuration - Frontend එක සමඟ සම්බන්ධ වීමට
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -18,81 +18,127 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# --- Neural Vault Setup ---
-API_KEYS_DB = {"ELE-PRIME-ADMIN-SYS": {"credits": 999999, "status": "active"}}
+# 2. Database & Storage
+# පද්ධතියේ පවතින Keys සහ ඒවායේ Limits ගබඩා කිරීම
+API_KEYS_DB = {
+    "ELE-PRIME-ADMIN-SYS": {"limit": 10000, "used": 0, "status": "active"}
+}
 ADMIN_SECRET = "MINZO-SECRET-2026"
 LEARNING_VAULT_PATH = "neural_learning_data.jsonl"
 
-# --- New Model Setup (CPU Optimized) ---
-# Qwen 1.5B එක Mistral වලට වඩා 5 ගුණයකට වඩා සැහැල්ලුයි
-model_id = "Qwen/Qwen2.5-1.5B-Instruct" 
+# 3. Light-Weight AI Model Setup (CPU Optimized)
+# Qwen 1.5B මොඩලය 16GB RAM එකක CPU මත ඉතා වේගවත් වේ
+model_id = "Qwen/Qwen2.5-1.5B-Instruct"
+HF_TOKEN = os.getenv("HF_TOKEN")
 
-print("🐘 Elephant Light-Engine Loading on CPU...")
+print("🐘 Elephant Engine Loading on CPU Mode...")
 tokenizer = AutoTokenizer.from_pretrained(model_id)
 model = AutoModelForCausalLM.from_pretrained(
     model_id,
-    torch_dtype="auto", # CPU එකට ගැලපෙන පරිදි auto-select වේ
-    device_map="cpu"    # GPU නැති නිසා අනිවාර්යයෙන් CPU ලෙස දිය යුතුයි
+    torch_dtype="auto",
+    device_map="cpu"  # CPU මත පමණක් ධාවනයට ස්ථාවර කර ඇත
 )
 
-# --- Helpers ---
-class NewKeyRequest(BaseModel):
+# 4. Data Models
+class KeyRequest(BaseModel):
     admin_pass: str
     new_key: str
+    limit: int = 50
 
-def get_live_data(query):
+# 5. Helper Functions
+def get_live_web_data(query):
+    """2026 තොරතුරු සඳහා DuckDuckGo හරහා සෙවීම"""
     try:
         with DDGS() as ddgs:
             results = [r['body'] for r in ddgs.text(query, max_results=3)]
             return "\n".join(results)
-    except: return ""
+    except Exception as e:
+        print(f"Search Error: {e}")
+        return ""
 
-def capture_learning_data(query, context, response, key_id):
+def log_to_learning_vault(query, context, response, key):
+    """පද්ධතිය දියුණු වීමට අවශ්‍ය දත්ත ගබඩා කිරීම"""
     entry = {
-        "timestamp": str(datetime.datetime.now()),
-        "key": key_id,
-        "q": query,
-        "ctx": context,
-        "ans": response
+        "ts": str(datetime.datetime.now()),
+        "key": key,
+        "input": query,
+        "web_ctx": context,
+        "ai_output": response
     }
-    with open(LEARNING_VAULT_PATH, "a") as f:
+    with open(LEARNING_VAULT_PATH, "a", encoding="utf-8") as f:
         f.write(json.dumps(entry) + "\n")
 
-# --- Endpoints ---
+# --- API Endpoints ---
 
 @app.get("/")
-def status():
-    return {"status": "Elephant-Qwen Node Active", "engine": "Qwen-2.5-1.5B"}
+def health_check():
+    return {
+        "status": "Elephant Node v3.6 Active",
+        "engine": "Qwen-2.5-1.5B-Instruct",
+        "keys_loaded": len(API_KEYS_DB)
+    }
 
 @app.post("/admin/add-key")
-async def add_key(data: NewKeyRequest):
-    if data.admin_pass != ADMIN_SECRET: raise HTTPException(status_code=401)
-    API_KEYS_DB[data.new_key] = {"credits": 5000, "status": "active"}
-    return {"message": "Key Registered"}
+async def register_new_key(data: KeyRequest):
+    """නව API Key එකක් පද්ධතියට ලියාපදිංචි කිරීම"""
+    if data.admin_pass != ADMIN_SECRET:
+        raise HTTPException(status_code=401, detail="Unauthorized")
+    
+    API_KEYS_DB[data.new_key] = {
+        "limit": data.limit,
+        "used": 0,
+        "status": "active"
+    }
+    return {"message": f"Token {data.new_key} activated with limit {data.limit}"}
 
 @app.post("/v1/chat")
-async def chat(message: dict, x_api_key: str = Header(None)):
-    if x_api_key not in API_KEYS_DB: raise HTTPException(status_code=403)
+async def chat_service(message: dict, x_api_key: str = Header(None)):
+    """ප්‍රධාන Chat සේවාව"""
+    # Key එක පරීක්ෂා කිරීම
+    if not x_api_key or x_api_key not in API_KEYS_DB:
+        raise HTTPException(status_code=403, detail="Invalid API Key")
     
-    query = message.get("query", "")
-    context = get_live_data(query) if any(w in query.lower() for w in ["now", "2026", "today"]) else ""
+    # Daily Limit පරීක්ෂා කිරීම
+    key_data = API_KEYS_DB[x_api_key]
+    if key_data["used"] >= key_data["limit"]:
+        raise HTTPException(status_code=429, detail="Daily request limit reached")
 
-    # Qwen Chat Template එක භාවිතය
+    user_query = message.get("query", "")
+    if not user_query:
+        return {"reply": "Please provide a query."}
+
+    # 2026 Live Data අවශ්‍යදැයි බැලීම
+    context = ""
+    trigger_words = ["now", "today", "2026", "news", "price", "current"]
+    if any(w in user_query.lower() for w in trigger_words):
+        context = get_live_web_data(user_query)
+
+    # Prompt එක සැකසීම
     messages = [
-        {"role": "system", "content": f"You are Elephant AI. Current year: 2026. Context: {context}"},
-        {"role": "user", "content": query}
+        {"role": "system", "content": f"You are Elephant AI. Year: 2026. Context: {context}"},
+        {"role": "user", "content": user_query}
     ]
-    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-    model_inputs = tokenizer([text], return_tensors="pt").to("cpu")
+    
+    # AI Inference
+    input_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    inputs = tokenizer([input_text], return_tensors="pt").to("cpu")
 
     with torch.no_grad():
-        generated_ids = model.generate(model_inputs.input_ids, max_new_tokens=512)
-        response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
-    
-    # පද්ධතියේ පිළිතුර පිරිසිදු කිරීම
-    clean_response = response.split("assistant")[-1].strip()
-    
-    capture_learning_data(query, context, clean_response, x_api_key)
-    return {"reply": clean_response, "key_id": x_api_key}
+        generated_ids = model.generate(inputs.input_ids, max_new_tokens=512, do_sample=True, temperature=0.7)
+        raw_response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
+
+    # Response එක පිරිසිදු කිරීම
+    final_reply = raw_response.split("assistant")[-1].strip()
+
+    # Usage එක යාවත්කාලීන කිරීම සහ ඉගෙනුම් දත්ත ගබඩා කිරීම
+    API_KEYS_DB[x_api_key]["used"] += 1
+    log_to_learning_vault(user_query, context, final_reply, x_api_key)
+
+    return {
+        "reply": final_reply,
+        "usage": f"{API_KEYS_DB[x_api_key]['used']}/{API_KEYS_DB[x_api_key]['limit']}",
+        "timestamp": "2026-04-27"
+    }
 
+# Entry point for HF Spaces
 main = app