Update app.py
Browse files
app.py
CHANGED
|
@@ -1 +1 @@
|
|
| 1 |
-
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(
|
|
|
|
| 1 |
+
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
|