Update app.py

app.py

@@ -8,6 +8,8 @@ import shutil
 import tempfile
 import uuid
 import time
+import uvicorn
+from contextlib import asynccontextmanager
 
 from fastapi import FastAPI, Request, UploadFile, File
 from fastapi.responses import JSONResponse, StreamingResponse
@@ -17,54 +19,54 @@ from faster_whisper import WhisperModel
 from gtts import gTTS
 
 # ====================================
-# CONFIG
+# CONFIG & MODELS
 # ====================================
-
 TIMEOUT = 120
-
-OPENROUTER_KEYS = os.getenv("OPENROUTER_KEYS", "").split(",")
-
-CF_ACCOUNTS = []
-cf_env = os.getenv("CLOUDFLARE_ACCOUNTS", "")
-
-for acc in cf_env.split(","):
-    if ":" in acc:
-        acc_id, token = acc.split(":")
-        CF_ACCOUNTS.append({"id": acc_id, "token": token})
-
-OLLAMA_URL = os.getenv("OLLAMA_URL", "http://127.0.0.1:11434/api/chat")
-
-# ====================================
-# MODEL ROUTER
-# ====================================
-
 MODEL_ROUTER = {
-
     "deepseek": "deepseek/deepseek-r1",
-
     "gpt": "openai/gpt-4o-mini",
-
     "claude": "anthropic/claude-3-haiku",
-
     "gemini": "google/gemini-2.0-flash"
-
 }
 
-# ====================================
-# WHISPER
-# ====================================
-
+# Inisialisasi Whisper (Pastikan path cache sesuai dengan Docker)
 whisper_model = WhisperModel(
-    "tiny",
-    device="cpu",
-    compute_type="int8"
+    "tiny", 
+    device="cpu", 
+    compute_type="int8",
+    download_root="./models_cache"
 )
 
 # ====================================
-# FASTAPI
+# LIFESPAN & WORKER
 # ====================================
+QUEUE = asyncio.Queue()
 
-app = FastAPI(title="AI Gateway")
+async def worker():
+    while True:
+        messages, model, future = await QUEUE.get()
+        try:
+            result = await parallel_ai(messages, model)
+            future.set_result(result)
+        except Exception as e:
+            future.set_exception(e)
+        finally:
+            QUEUE.task_done()
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    # Startup: Jalankan 4 worker
+    workers = [asyncio.create_task(worker()) for _ in range(4)]
+    print("=== WORKERS READY ===")
+    yield
+    # Shutdown: Batalkan worker jika perlu
+    for w in workers:
+        w.cancel()
+
+# ====================================
+# FASTAPI SETUP
+# ====================================
+app = FastAPI(title="AI Gateway", lifespan=lifespan)
 
 app.add_middleware(
     CORSMiddleware,
@@ -74,360 +76,143 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# ====================================
-# CACHE
-# ====================================
-
 CACHE = {}
 
-def cache_key(messages):
-    return hashlib.md5(json.dumps(messages).encode()).hexdigest()
-
+# --- Helper Functions (Cache & AI Calls) ---
 def get_cache(messages):
-    key = cache_key(messages)
+    key = hashlib.md5(json.dumps(messages).encode()).hexdigest()
     return CACHE.get(key)
 
 def set_cache(messages, data):
-    key = cache_key(messages)
+    key = hashlib.md5(json.dumps(messages).encode()).hexdigest()
     CACHE[key] = data
 
-# ====================================
-# PROVIDERS
-# ====================================
-
 async def call_ollama(messages, model):
-
     try:
-
         async with httpx.AsyncClient(timeout=TIMEOUT) as client:
-
             res = await client.post(
-                OLLAMA_URL,
-                json={
-                    "model": model,
-                    "messages": messages,
-                    "stream": False
-                }
+                os.getenv("OLLAMA_URL", "http://127.0.0.1:11434/api/chat"),
+                json={"model": model, "messages": messages, "stream": False}
             )
-
-            if res.status_code == 200:
-
-                data = res.json()
-
-                return data["message"]["content"]
-
-    except:
-        return None
-
+            return res.json()["message"]["content"] if res.status_code == 200 else None
+    except: return None
 
 async def call_openrouter(messages, model):
-
-    random.shuffle(OPENROUTER_KEYS)
-
+    keys = os.getenv("OPENROUTER_KEYS", "").split(",")
+    random.shuffle(keys)
     async with httpx.AsyncClient(timeout=TIMEOUT) as client:
-
-        for key in OPENROUTER_KEYS:
-
+        for key in keys:
             try:
-
                 res = await client.post(
                     "https://openrouter.ai/api/v1/chat/completions",
-                    headers={
-                        "Authorization": f"Bearer {key}",
-                        "Content-Type": "application/json"
-                    },
-                    json={
-                        "model": model,
-                        "messages": messages
-                    }
+                    headers={"Authorization": f"Bearer {key}"},
+                    json={"model": model, "messages": messages}
                 )
-
-                if res.status_code == 200:
-
-                    data = res.json()
-
-                    return data["choices"][0]["message"]["content"]
-
-            except:
-                continue
-
+                if res.status_code == 200: return res.json()["choices"][0]["message"]["content"]
+            except: continue
     return None
 
-
 async def call_cloudflare(messages):
-
-    random.shuffle(CF_ACCOUNTS)
-
+    cf_env = os.getenv("CLOUDFLARE_ACCOUNTS", "")
+    accounts = [{"id": a.split(":")[0], "token": a.split(":")[1]} for a in cf_env.split(",") if ":" in a]
+    random.shuffle(accounts)
     async with httpx.AsyncClient(timeout=TIMEOUT) as client:
-
-        for acc in CF_ACCOUNTS:
-
+        for acc in accounts:
             try:
-
                 url = f"https://api.cloudflare.com/client/v4/accounts/{acc['id']}/ai/run/@cf/meta/llama-3-8b-instruct"
-
-                res = await client.post(
-                    url,
-                    headers={
-                        "Authorization": f"Bearer {acc['token']}"
-                    },
-                    json={
-                        "messages": messages
-                    }
-                )
-
-                if res.status_code == 200:
-
-                    data = res.json()
-
-                    return data["result"]["response"]
-
-            except:
-                continue
-
+                res = await client.post(url, headers={"Authorization": f"Bearer {acc['token']}"}, json={"messages": messages})
+                if res.status_code == 200: return res.json()["result"]["response"]
+            except: continue
     return None
 
-
-# ====================================
-# PARALLEL ROUTER
-# ====================================
-
 async def parallel_ai(messages, model):
-
     tasks = [
-
         asyncio.create_task(call_ollama(messages, model)),
-
         asyncio.create_task(call_openrouter(messages, model)),
-
         asyncio.create_task(call_cloudflare(messages))
-
     ]
-
-    done, pending = await asyncio.wait(
-        tasks,
-        return_when=asyncio.FIRST_COMPLETED
-    )
-
-    for p in pending:
-        p.cancel()
-
+    done, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
+    for p in pending: p.cancel()
     for d in done:
-
-        result = d.result()
-
-        if result:
-            return result
-
+        res = d.result()
+        if res: return res
     return None
 
-
-# ====================================
-# QUEUE
-# ====================================
-
-QUEUE = asyncio.Queue()
-
-async def worker():
-
-    while True:
-
-        messages, model, future = await QUEUE.get()
-
-        try:
-
-            result = await parallel_ai(messages, model)
-
-            future.set_result(result)
-
-        except Exception as e:
-
-            future.set_exception(e)
-
-        QUEUE.task_done()
-
-
-@app.on_event("startup")
-async def startup():
-
-    for _ in range(4):
-        asyncio.create_task(worker())
-
-
 # ====================================
-# STREAM GENERATOR
+# ENDPOINTS
 # ====================================
 
-async def stream_text(text):
-
-    words = text.split(" ")
-
-    for w in words:
-
-        chunk = {
-            "choices": [
-                {
-                    "delta": {
-                        "content": w + " "
-                    }
-                }
-            ]
-        }
-
-        yield f"data: {json.dumps(chunk)}\n\n"
-
-        await asyncio.sleep(0.02)
-
-    yield "data: [DONE]\n\n"
-
-
-# ====================================
-# CHAT COMPLETION
-# ====================================
+@app.get("/")
+async def root():
+    return {"status": "AI Gateway Running", "engine": "FastAPI Lifespan"}
 
 @app.post("/v1/chat/completions")
 async def chat(request: Request):
-
     body = await request.json()
-
     messages = body.get("messages", [])
-    model = body.get("model", "deepseek")
-
+    model_key = body.get("model", "deepseek")
     stream = body.get("stream", False)
-
-    model = MODEL_ROUTER.get(model, model)
-
-    cached = get_cache(messages)
-
-    if cached and not stream:
-        return cached
-
-    loop = asyncio.get_event_loop()
-    future = loop.create_future()
-
+    
+    model = MODEL_ROUTER.get(model_key, model_key)
+    
+    if not stream:
+        cached = get_cache(messages)
+        if cached: return cached
+
+    future = asyncio.get_event_loop().create_future()
     await QUEUE.put((messages, model, future))
-
     result = await future
 
     if not result:
-        return JSONResponse({"error": "AI failed"}, status_code=500)
+        return JSONResponse({"error": "AI failed to respond"}, status_code=500)
 
     if stream:
-
-        return StreamingResponse(
-            stream_text(result),
-            media_type="text/event-stream"
-        )
+        async def stream_gen():
+            for word in result.split(" "):
+                chunk = {"choices": [{"delta": {"content": word + " "}}]}
+                yield f"data: {json.dumps(chunk)}\n\n"
+                await asyncio.sleep(0.01)
+            yield "data: [DONE]\n\n"
+        return StreamingResponse(stream_gen(), media_type="text/event-stream")
 
     response = {
         "id": str(uuid.uuid4()),
         "object": "chat.completion",
         "created": int(time.time()),
         "model": model,
-        "choices": [
-            {
-                "message": {
-                    "role": "assistant",
-                    "content": result
-                }
-            }
-        ]
+        "choices": [{"message": {"role": "assistant", "content": result}}]
     }
-
     set_cache(messages, response)
-
     return response
 
-
-# ====================================
-# WHISPER STT
-# ====================================
-
 @app.post("/v1/audio/transcriptions")
 async def transcribe(file: UploadFile = File(...)):
-
-    with tempfile.NamedTemporaryFile(delete=False) as tmp:
-
+    with tempfile.NamedTemporaryFile(delete=False, suffix=".tmp") as tmp:
         shutil.copyfileobj(file.file, tmp)
-        path = tmp.name
-
-    segments, info = whisper_model.transcribe(path)
-
-    text = ""
-
-    for s in segments:
-        text += s.text
-
-    os.remove(path)
-
-    return {
-        "text": text,
-        "language": info.language
-    }
-
-
-# ====================================
-# TEXT TO SPEECH
-# ====================================
+        tmp_path = tmp.name
+    
+    try:
+        segments, info = whisper_model.transcribe(tmp_path)
+        text = "".join([s.text for s in segments])
+        return {"text": text, "language": info.language}
+    finally:
+        if os.path.exists(tmp_path): os.remove(tmp_path)
 
 @app.post("/v1/audio/speech")
 async def tts(request: Request):
-
     body = await request.json()
-
-    text = body.get("input", "")
-
-    voice = body.get("voice", "en")
-
-    lang = "en"
-
-    if voice == "id":
-        lang = "id"
-
-    tts = gTTS(text=text, lang=lang)
-
+    text, voice = body.get("input", ""), body.get("voice", "en")
+    lang = "id" if voice == "id" else "en"
+    
     tmp_path = f"/tmp/{uuid.uuid4()}.mp3"
+    gTTS(text=text, lang=lang).save(tmp_path)
 
-    tts.save(tmp_path)
-
-    def stream():
-
-        with open(tmp_path, "rb") as f:
+    def iterfile():
+        with open(tmp_path, mode="rb") as f:
             yield from f
+        if os.path.exists(tmp_path): os.remove(tmp_path)
 
-        os.remove(tmp_path)
-
-    return StreamingResponse(stream(), media_type="audio/mpeg")
-
-
-# ====================================
-# MODELS
-# ====================================
-
-@app.get("/v1/models")
-async def models():
-
-    return {
-        "data": [
-
-            {"id": "deepseek"},
-            {"id": "gpt"},
-            {"id": "claude"},
-            {"id": "gemini"},
-            {"id": "whisper-1"},
-            {"id": "tts-1"}
-
-        ]
-    }
-
-
-# ====================================
-# HEALTH
-# ====================================
-
-@app.get("/")
-async def root():
+    return StreamingResponse(iterfile(), media_type="audio/mpeg")
 
-    return {
-        "status": "AI Gateway Running"
-    }
+if __name__ == "__main__":
+    uvicorn.run(app, host="0.0.0.0", port=7860)