Update app.py

app.py

@@ -1,104 +1 @@
-from fastapi import FastAPI, Header, HTTPException
-from fastapi.middleware.cors import CORSMiddleware
-from pydantic import BaseModel
-import torch
-import os
-from transformers import AutoModelForCausalLM, AutoTokenizer
-from duckduckgo_search import DDGS
-
-# 🔱 Server එකට "app" කියන නමම අවශ්‍යයි
-app = FastAPI()
-
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
-
-# --- Database & Config ---
-API_KEYS_DB = {
-    "ELE-PRIME-ADMIN-SYS": {"limit": 10000, "used": 0, "status": "active"}
-}
-ADMIN_SECRET = "MINZO-SECRET-2026"
-
-# --- AI Model (Gemma-3-1B-it) ---
-model_id = "google/gemma-3-1b-it"
-HF_TOKEN = os.getenv("HF_TOKEN")
-
-print(f"🐘 Elephant Node v3.7 Loading: {model_id}...")
-
-tokenizer = AutoTokenizer.from_pretrained(model_id, token=HF_TOKEN)
-
-# 🔱 [Fix] torch_dtype වෙනුවට dtype පාවිච්චි කරන ලදී
-model = AutoModelForCausalLM.from_pretrained(
-    model_id, 
-    dtype=torch.bfloat16, 
-    device_map="cpu", 
-    token=HF_TOKEN
-)
-
-# --- Data Models ---
-class KeyRequest(BaseModel):
-    admin_pass: str
-    new_key: str
-    limit: int = 100
-
-# --- API Endpoints ---
-@app.get("/")
-def home():
-    return {"status": "Elephant Pro Active", "model": "Gemma-3-1B"}
-
-@app.post("/admin/add-key")
-async def add_key(data: KeyRequest):
-    if data.admin_pass != ADMIN_SECRET:
-        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)):
-    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"]
-    }
-
-@app.post("/v1/chat")
-async def chat(message: dict, x_api_key: str = Header(None)):
-    if not x_api_key or x_api_key not in API_KEYS_DB:
-        raise HTTPException(status_code=403)
-        
-    key_info = API_KEYS_DB[x_api_key]
-    if key_info["used"] >= key_info["limit"]:
-        raise HTTPException(status_code=429, detail="Limit Reached")
-    
-    query = message.get("query", "")
-    
-    # 2026 Web Search Logic
-    context = ""
-    try:
-        if any(w in query.lower() for w in ["today", "now", "2026"]):
-            with DDGS() as ddgs:
-                context = "\n".join([r['body'] for r in ddgs.text(query, max_results=2)])
-    except: pass
-
-    # AI Inference
-    msgs = [
-        {"role": "system", "content": f"Elephant AI by MINZO-PRIME. 2026 Edition. 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():
-        ids = model.generate(inputs.input_ids, max_new_tokens=300, temperature=0.7, do_sample=True)
-        full_ans = tokenizer.batch_decode(ids, skip_special_tokens=True)[0]
-        ans = full_ans.split("model")[-1].strip()
-
-    API_KEYS_DB[x_api_key]["used"] += 1
-    return {"reply": ans, "usage": API_KEYS_DB[x_api_key]["used"]}
-
+from fastapi import FastAPI, Header, HTTPExceptionfrom fastapi.middleware.cors import CORSMiddlewarefrom pydantic import BaseModelimport torchimport osimport jsonimport datetimefrom transformers import AutoModelForCausalLM, AutoTokenizerfrom duckduckgo_search import DDGSapp = FastAPI()# CORS Fix for Dashboard connectivityapp.add_middleware(    CORSMiddleware,    allow_origins=["*"],    allow_methods=["*"],    allow_headers=["*"],)# --- Database & Config ---API_KEYS_DB = {    "ELE-PRIME-ADMIN-SYS": {"limit": 10000, "used": 0, "status": "active"}}ADMIN_SECRET = "MINZO-SECRET-2026"LEARNING_VAULT = "neural_learning_data.jsonl"# --- AI Model (Qwen-2.5-1.5B) ---model_id = "Qwen/Qwen2.5-1.5B-Instruct"print("🐘 Elephant Node v3.7 Loading...")tokenizer = AutoTokenizer.from_pretrained(model_id)model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype="auto", device_map="cpu")# --- Data Models ---class KeyRequest(BaseModel):    admin_pass: str    new_key: str    limit: int = 100# --- API Endpoints ---@app.get("/")def home():    return {"status": "Elephant Pro Active", "keys": len(API_KEYS_DB)}@app.post("/admin/add-key")async def add_key(data: KeyRequest):    if data.admin_pass != ADMIN_SECRET:        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    }@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)        key_info = API_KEYS_DB[x_api_key]    if key_info["used"] >= key_info["limit"]:        raise HTTPException(status_code=429, detail="Limit Reached")    query = message.get("query", "")        # 2026 Web Search Logic    context = ""    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    # AI Inference    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():        ids = model.generate(inputs.input_ids, max_new_tokens=256)        ans = tokenizer.batch_decode(ids, skip_special_tokens=True)[0].split("assistant")[-1].strip()    # Update Stats    API_KEYS_DB[x_api_key]["used"] += 1    return {"reply": ans, "usage": API_KEYS_DB[x_api_key]["used"]}main = app