Update app.py

app.py

@@ -1,13 +1,15 @@
 import os
 import torch
+import uuid
+import re
 from fastapi import FastAPI, Header, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from duckduckgo_search import DDGS
-import re
 
 app = FastAPI()
 
+# 🔱 CORS Setup
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -15,66 +17,103 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# --- Specialist DB ---
+# --- 🔱 Specialist DB (Memory Based) ---
+# සටහන: සර්වර් එක Restart වූ විට මේවා මැකේ. ස්ථිර කිරීමට DB එකක් අවශ්‍යයි.
 API_KEYS_DB = {
-    "ELE-PRIME-ADMIN-SYS": {"limit": 100000, "used": 0, "status": "active"},
-    "ELE-PRIME-VOID-X": {"limit": 50000, "used": 0, "status": "active"}
+    "ELE-PRIME-ADMIN-SYS": {"limit": 100000, "used": 0, "status": "active", "owner": "MINZO-PRIME"},
 }
 
-# 🔱 Model Configuration (CPU Friendly Load)
+# --- 🔱 Model Configuration (CPU Stable Engine) ---
 MODEL_ID = "google/gemma-3-270m"
-print(f"🔱 INACHI-CORE: Launching Gemma-4 on CPU Engine...")
+print(f"🔱 INACHI-CORE: Launching Gemma-3 on CPU Engine...")
 
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
 
-# GPU නැති නිසා 4-bit quantization වෙනුවට CPU Memory optimization පාවිච්චි කරමු
+# float16 නිසා එන 'NaN' error එක වැළැක්වීමට float32 පාවිච්චි කරමු
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_ID,
-    torch_dtype=torch.float16,
+    torch_dtype=torch.float32, 
     low_cpu_mem_usage=True,
     device_map="cpu" 
 )
 
+# --- 🔱 Web Context Retrieval ---
 def get_web_context(query: str):
     try:
         with DDGS() as ddgs:
+            # නව ddgs version එකට ගැලපෙන පරිදි update කර ඇත
             results = [r['body'] for r in ddgs.text(query, max_results=3)]
             return "\n".join(results)
-    except:
+    except Exception as e:
+        print(f"Search Error: {e}")
         return ""
 
+# --- 🔱 Admin Routes ---
+@app.get("/sys/generate-key")
+async def create_key(admin_key: str = Header(None)):
+    """අලුත් API Keys සාදා ගැනීමට: Header එකේ 'admin-key' ලෙස ELE-PRIME-ADMIN-SYS ලබා දෙන්න."""
+    if admin_key != "ELE-PRIME-ADMIN-SYS":
+        raise HTTPException(status_code=403, detail="Unauthorized Specialist Access")
+    
+    new_key = f"ELE-PRIME-{uuid.uuid4().hex[:8].upper()}"
+    API_KEYS_DB[new_key] = {"limit": 5000, "used": 0, "status": "active", "owner": "Specialist"}
+    return {"status": "success", "new_key": new_key}
+
+# --- 🔱 Chat Endpoint ---
 @app.post("/v1/chat")
 async def chat(message: dict, x_api_key: str = Header(None)):
+    # 1. API Key Validation
     if not x_api_key or x_api_key not in API_KEYS_DB:
         raise HTTPException(status_code=403, detail="Invalid Specialist Key")
-        
+    
+    key_info = API_KEYS_DB[x_api_key]
+    if key_info["used"] >= key_info["limit"]:
+        raise HTTPException(status_code=429, detail="API Limit Reached for this Key")
+
     user_query = message.get("query", "")
     web_data = get_web_context(user_query)
     
+    # 2. Inachi System Prompt
     system_prompt = (
         "You are Inachi-Prime, a multimodal AI developed by Specialist MINZO-PRIME. "
-        "Respond directly without showing your internal thought process. "
-        f"\nWeb Context: {web_data}"
+        "Respond directly without internal thought process. "
+        f"\nContext: {web_data}"
     )
 
     full_prompt = f"<start_of_turn>system\n{system_prompt}<end_of_turn>\n<start_of_turn>user\n{user_query}<end_of_turn>\n<start_of_turn>model\n"
     
     inputs = tokenizer(full_prompt, return_tensors="pt")
     
+    # 3. Generation Logic (Stable Parameters)
     with torch.no_grad():
-        outputs = model.generate(
-            **inputs,
-            max_new_tokens=512,
-            temperature=0.7,
-            do_sample=True
-        )
-        
-        full_response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        # Thinking කොටස ඉවත් කිරීම
-        final_reply = full_response.split("model\n")[-1].strip()
-        final_reply = re.sub(r'<thought>.*?</thought>', '', final_reply, flags=re.DOTALL).strip()
+        try:
+            outputs = model.generate(
+                **inputs,
+                max_new_tokens=512,
+                temperature=0.7,
+                do_sample=True,
+                renormalize_logits=True # Probability Error එක වැළැක්වීමට
+            )
+            
+            full_response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+            
+            # Response එක පිරිසිදු කිරීම
+            final_reply = full_response.split("model\n")[-1].strip()
+            final_reply = re.sub(r'<thought>.*?</thought>', '', final_reply, flags=re.DOTALL).strip()
+
+            # 4. Usage tracking
+            API_KEYS_DB[x_api_key]["used"] += 1
+            
+            return {
+                "reply": final_reply,
+                "usage": f"{API_KEYS_DB[x_api_key]['used']}/{API_KEYS_DB[x_api_key]['limit']}"
+            }
 
-    API_KEYS_DB[x_api_key]["used"] += 1
-    return {"reply": final_reply}
+        except RuntimeError as e:
+            print(f"Generation Error: {e}")
+            return {"reply": "Core engine destabilized. Retrying process recommended."}
 
-main = app 
+if __name__ == "__main__":
+    import uvicorn
+    # Hugging Face Space සඳහා Port 7860 අනිවාර්යයි
+    uvicorn.run(app, host="0.0.0.0", port=7860)