Update app.py

app.py

@@ -10,7 +10,7 @@ from duckduckgo_search import DDGS
 
 app = FastAPI()
 
-# 1. CORS Configuration - Frontend එක සමඟ සම්බන්ධ වීමට
+# CORS Fix for Dashboard connectivity
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -18,127 +18,80 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# 2. Database & Storage
-# පද්ධතියේ පවතින Keys සහ ඒවායේ Limits ගබඩා කිරීම
+# --- Database & Config ---
 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"
+LEARNING_VAULT = "neural_learning_data.jsonl"
 
-# 3. Light-Weight AI Model Setup (CPU Optimized)
-# Qwen 1.5B මොඩලය 16GB RAM එකක CPU මත ඉතා වේගවත් වේ
+# --- AI Model (Qwen-2.5-1.5B) ---
 model_id = "Qwen/Qwen2.5-1.5B-Instruct"
-HF_TOKEN = os.getenv("HF_TOKEN")
-
-print("🐘 Elephant Engine Loading on CPU Mode...")
+print("🐘 Elephant Node v3.7 Loading...")
 tokenizer = AutoTokenizer.from_pretrained(model_id)
-model = AutoModelForCausalLM.from_pretrained(
-    model_id,
-    torch_dtype="auto",
-    device_map="cpu"  # CPU මත පමණක් ධාවනයට ස්ථාවර කර ඇත
-)
+model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype="auto", device_map="cpu")
 
-# 4. Data Models
+# --- Data Models ---
 class KeyRequest(BaseModel):
     admin_pass: str
     new_key: str
-    limit: int = 50
-
-# 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 Exception as e:
-        print(f"Search Error: {e}")
-        return ""
-
-def log_to_learning_vault(query, context, response, key):
-    """පද්ධතිය දියුණු වීමට අවශ්‍ය දත්ත ගබඩා කිරීම"""
-    entry = {
-        "ts": str(datetime.datetime.now()),
-        "key": key,
-        "input": query,
-        "web_ctx": context,
-        "ai_output": response
-    }
-    with open(LEARNING_VAULT_PATH, "a", encoding="utf-8") as f:
-        f.write(json.dumps(entry) + "\n")
+    limit: int = 100
 
 # --- API Endpoints ---
 
 @app.get("/")
-def health_check():
-    return {
-        "status": "Elephant Node v3.6 Active",
-        "engine": "Qwen-2.5-1.5B-Instruct",
-        "keys_loaded": len(API_KEYS_DB)
-    }
+def home():
+    return {"status": "Elephant Pro Active", "keys": len(API_KEYS_DB)}
 
 @app.post("/admin/add-key")
-async def register_new_key(data: KeyRequest):
-    """නව API Key එකක් පද්ධතියට ලියාපදිංචි කිරීම"""
+async def add_key(data: KeyRequest):
     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"
+        raise HTTPException(status_code=401)
+    API_KEYS_DB[data.new_key] = {"limit": data.limit, "used": 0, "status": "active"}
+    return {"message": "Key Activated"}
+
+@app.get("/v1/usage")
+async def get_usage(x_api_key: str = Header(None)):
+    """Key එකේ පාවිච්චිය පරීක්ෂා කිරීමේ Endpoint එක"""
+    if not x_api_key or x_api_key not in API_KEYS_DB:
+        raise HTTPException(status_code=403, detail="Invalid Key")
+    info = API_KEYS_DB[x_api_key]
+    return {
+        "used": info["used"],
+        "limit": info["limit"],
+        "percentage": (info["used"] / info["limit"]) * 100 if info["limit"] > 0 else 0
     }
-    return {"message": f"Token {data.new_key} activated with limit {data.limit}"}
 
 @app.post("/v1/chat")
-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")
+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)
     
-    # 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")
-
-    user_query = message.get("query", "")
-    if not user_query:
-        return {"reply": "Please provide a query."}
+    key_info = API_KEYS_DB[x_api_key]
+    if key_info["used"] >= key_info["limit"]:
+        raise HTTPException(status_code=429, detail="Limit Reached")
 
-    # 2026 Live Data අවශ්‍යදැයි බැලීම
+    query = message.get("query", "")
+    
+    # 2026 Web Search Logic
     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)
+    if any(w in query.lower() for w in ["today", "now", "2026"]):
+        try:
+            with DDGS() as ddgs:
+                context = "\n".join([r['body'] for r in ddgs.text(query, max_results=2)])
+        except: pass
 
-    # Prompt එක සැකසීම
-    messages = [
-        {"role": "system", "content": f"You are Elephant AI. Year: 2026. Context: {context}"},
-        {"role": "user", "content": user_query}
-    ]
-    
     # AI Inference
-    input_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-    inputs = tokenizer([input_text], return_tensors="pt").to("cpu")
+    msgs = [{"role": "system", "content": f"Elephant AI. 2026 mode. Context: {context}"}, {"role": "user", "content": query}]
+    text = tokenizer.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)
+    inputs = tokenizer([text], return_tensors="pt").to("cpu")
 
     with torch.no_grad():
-        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]
+        ids = model.generate(inputs.input_ids, max_new_tokens=256)
+        ans = tokenizer.batch_decode(ids, skip_special_tokens=True)[0].split("assistant")[-1].strip()
 
-    # Response එක පිරිසිදු කිරීම
-    final_reply = raw_response.split("assistant")[-1].strip()
-
-    # Usage එක යාවත්කාලීන කිරීම සහ ඉගෙනුම් දත්ත ගබඩා කිරීම
+    # Update Stats
     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"
-    }
+    return {"reply": ans, "usage": API_KEYS_DB[x_api_key]["used"]}
 
-# Entry point for HF Spaces
 main = app