""" URAL AI — Worker Space App Nhẹ, nhanh, chỉ làm đúng 1 việc: nhận task từ main space và xử lý AI. ENV cần set trong HF Secrets của mỗi worker: WORKER_SECRET = (giống main space) WORKER_ROLE = strong | code | creative | backup MAIN_SPACE_URL = https://bachdathan-ural-ai.hf.space (URL main space) """ import warnings # Cohere SDK (5.11.4) touches Pydantic v1-style `__fields__` internally, # spamming PydanticDeprecatedSince20 warnings on import. Silence centrally, # before any module that imports cohere gets loaded (mirrors main app.py). warnings.filterwarnings("ignore", category=DeprecationWarning, module=r"pydantic.*") warnings.filterwarnings("ignore", message=r".*__fields__.*attribute is deprecated.*") try: from pydantic import PydanticDeprecatedSince20 warnings.filterwarnings("ignore", category=PydanticDeprecatedSince20) except ImportError: pass import os, time, random, asyncio import requests as _req from fastapi import FastAPI, Request from fastapi.responses import JSONResponse # 🛡️ (2026-07-22) Bảo mật worker — trước đây worker KHÔNG có bất kỳ lớp # phòng thủ nào (không rate-limit, không chặn path-traversal, /docs mở # công khai lộ toàn bộ shape API) dù nhận traffic public y hệt space # chính. Route xử lý dữ liệu (/api/worker, /api/hist-*...) đã có # WORKER_SECRET check nên KHÔNG bị lộ dữ liệu — nhưng vẫn cần thêm 2 lớp # này để chặn spam/DoS và giảm bề mặt dò quét, giống bản đã vá ở app.py # chính. _EXPOSE_API_DOCS = os.getenv("EXPOSE_API_DOCS", "0") == "1" app = FastAPI( title="Yui Worker", docs_url=("/docs" if _EXPOSE_API_DOCS else None), redoc_url=("/redoc" if _EXPOSE_API_DOCS else None), openapi_url=("/openapi.json" if _EXPOSE_API_DOCS else None), ) from starlette.middleware.base import BaseHTTPMiddleware _SEC_RATE_LIMIT_MAX = 30 # request / phút / IP — worker chỉ nhận _SEC_RATE_LIMIT_WINDOW = 60 # traffic từ main space + kẻ dò quét, _SEC_BLOCK_SECONDS = 300 # không cần whitelist route nội bộ như _SEC_IP_REQUESTS: dict = {} # app.py chính (main space luôn có secret _SEC_IP_BLOCK: dict = {} # đúng nên request hợp lệ không bao giờ _SEC_SUSPICIOUS_PATTERNS = ( # đủ dồn dập để chạm ngưỡng 30/phút). "../", "..\\", "%2e%2e", "%252e%252e", "%2f..", "%252f..", "/proc/self", "/.git/", "/.env", ) def _sec_get_real_ip(request) -> str: xff = request.headers.get("x-forwarded-for", "") if xff: return xff.split(",")[0].strip() return request.client.host if request.client else "unknown" def _sec_check_rate_limit(ip: str) -> bool: now = time.time() if ip in _SEC_IP_BLOCK and now < _SEC_IP_BLOCK[ip]: return False reqs = [t for t in _SEC_IP_REQUESTS.get(ip, []) if now - t < _SEC_RATE_LIMIT_WINDOW] if len(reqs) >= _SEC_RATE_LIMIT_MAX: _SEC_IP_BLOCK[ip] = now + _SEC_BLOCK_SECONDS print(f"[WORKER-SECURITY] 🚫 Block IP {ip} trong {_SEC_BLOCK_SECONDS}s") return False reqs.append(now) _SEC_IP_REQUESTS[ip] = reqs return True class WorkerSecurityMiddleware(BaseHTTPMiddleware): async def dispatch(self, request, call_next): path = request.url.path raw = str(request.url).lower() if any(pat in raw for pat in _SEC_SUSPICIOUS_PATTERNS): real_ip = _sec_get_real_ip(request) print(f"[WORKER-SECURITY] 🚨 Path traversal từ {real_ip}: {path[:80]}") return JSONResponse({"detail": "Forbidden"}, status_code=403) # 🐛 FIX (2026-09-01) — Boss phát hiện qua log thực tế: middleware # này kế thừa BaseHTTPMiddleware — cách nó "đọc trước rồi nạp lại" # request.body() (xem đoạn `await request.body()` + # `request._receive = _receive` bên dưới) là ANTI-PATTERN đã biết # của Starlette, XUNG ĐỘT với StreamingResponse khi client (Space # chính) đóng kết nối giữa lúc route đang stream — 2 task cùng # tranh đọc receive() channel, gây lỗi "Unexpected message # received: http.request" trong anyio.TaskGroup (crash 500, không # phải lỗi logic gate — Aren build response đúng, chỉ là tầng # middleware/framework xung đột khi kết hợp với streaming). # Sửa: BỎ QUA hoàn toàn bước "đọc body để check secret" cho route # streaming (/api/local-llm/stream) — route đó đã tự đọc và kiểm # tra "secret" trong JSON body của chính nó ở tầng handler rồi # (xem handle_stream_request() → local_llm_stream_endpoint()), nên # middleware không cần (và không nên) đọc trước ở đây nữa. Route # vẫn được rate-limit như request không trusted khác nếu cần (an # toàn tương đương — chỉ mất phần "miễn rate-limit sớm" cho riêng # route stream, không mất kiểm tra secret thật, vì secret vẫn được # route tự xác thực). _is_streaming_route = path == "/api/local-llm/stream" # Request có đúng WORKER_SECRET (tức từ main space thật) được # miễn rate-limit — main space có thể gọi dồn dập khi nhiều # người chat cùng lúc, đó là traffic hợp lệ chứ không phải bot # dò quét. Chỉ đọc body cho POST — không ảnh hưởng GET (/health). # 🐛 FIX (2026-09-02) — Boss phát hiện ca thực tế: GET # /api/local-llm/status từ gate Qwen (app.py._qwen_stream_gate_sse) # gửi secret qua header "X-Worker-Secret", nhưng middleware ở đây # TRƯỚC ĐÂY chỉ kiểm tra "X-Local-Model-Secret" — mismatch tên # header khiến middleware LUÔN coi request /status hợp lệ này là # "không đáng tin" (_is_trusted=False), dồn nó vào rate-limit check # thay vì miễn — có thể dẫn tới bị chặn/chậm dù request hoàn toàn # hợp lệ. Route handler (local_llm_status_endpoint) đã đúng từ đầu # — chấp nhận CẢ 2 tên header — nay middleware cũng được sửa khớp, # để nhất quán 2 tầng. _is_trusted = False if WORKER_SECRET and ( request.headers.get("X-Local-Model-Secret") == WORKER_SECRET or request.headers.get("X-Worker-Secret") == WORKER_SECRET ): _is_trusted = True if WORKER_SECRET and request.method == "POST" and not _is_trusted and not _is_streaming_route: try: _body_bytes = await request.body() if _body_bytes: import json as _sec_json _peek = _sec_json.loads(_body_bytes) if _peek.get("secret") == WORKER_SECRET: _is_trusted = True # Request body đã được đọc — nạp lại để route xử lý bình thường async def _receive(): return {"type": "http.request", "body": _body_bytes, "more_body": False} request._receive = _receive except Exception: pass # 🐛 (2026-09-01) Route streaming: bỏ qua rate-limit ở tầng # middleware luôn (không chỉ bỏ qua đọc body) — vì _is_trusted sẽ # LUÔN False cho route này nếu request không có header # X-Local-Model-Secret (Space chính hiện gửi secret trong JSON # body, không phải header, cho route stream — xem # _try_local_model_stream() phía Space chính) — nếu không loại # trừ ở đây, MỌI request /api/local-llm/stream hợp lệ từ Space # chính sẽ luôn bị tính vào rate-limit như 1 IP lạ, có thể bị # 429 sai khi nhiều request dồn dập (dù đã đúng thiết kế "bỏ qua # gate nếu Aren đang bận" ở tầng router — nhưng rate-limit theo # IP là lớp riêng, không liên quan trạng thái bận/rảnh của model). if not _is_trusted and not _is_streaming_route: real_ip = _sec_get_real_ip(request) if not _sec_check_rate_limit(real_ip): return JSONResponse({"detail": "Too Many Requests"}, status_code=429) return await call_next(request) app.add_middleware(WorkerSecurityMiddleware) # ── Config ──────────────────────────────────────────────────────── WORKER_SECRET = os.environ.get("WORKER_SECRET", "") WORKER_ROLE = os.environ.get("WORKER_ROLE", "backup") MAIN_SPACE_URL = os.environ.get("MAIN_SPACE_URL", "https://bachdathan-ural-ai.hf.space") # ── Key cache ───────────────────────────────────────────────────── _keys: dict = {"gemini": [], "cerebras": [], "openrouter": [], "cohere": [], "mistral": []} _keys_at: float = 0.0 _KEY_TTL = 1500 # 25 phút def _fetch_keys_sync(): """Lấy API keys từ main space (sync, dùng khi startup).""" global _keys, _keys_at try: r = _req.post( f"{MAIN_SPACE_URL}/api/get-keys", json={"secret": WORKER_SECRET, "role": WORKER_ROLE}, timeout=20 ) if r.status_code == 200: _keys = r.json() _keys_at = time.time() g = len(_keys.get("gemini", [])) cb = len(_keys.get("cerebras", [])) co = len(_keys.get("cohere", [])) print(f"[WORKER-{WORKER_ROLE.upper()}] ✅ Keys: G×{g} CB×{cb} CO×{co}") else: print(f"[WORKER] ⚠️ get-keys trả {r.status_code}") except Exception as e: print(f"[WORKER] ⚠️ get-keys lỗi: {e}") def _get_keys() -> dict: if time.time() - _keys_at > _KEY_TTL: _fetch_keys_sync() return _keys # ── Helpers gọi AI ──────────────────────────────────────────────── def _build_msgs(system: str, history: list, message: str) -> list: msgs = [] if system: # 🆕 (2026-07-11) BUG NGHIÊM TRỌNG ĐÃ VÁ: system[:1500] cắt mất phần # lớn system prompt mà app.py chính build (soul header + task rules + # style guide cấm "Ồ/Ừ" đầu câu + teencode rules...). Vì các phần này # được ghép NỐI TIẾP nhau và có thể vượt 1500 ký tự, phần nằm sau vị # trí cắt bị mất HOÀN TOÀN trước khi tới Gemini/Cohere — model tự viết # theo thói quen huấn luyện (chèn "Ồ, chào Boss") vì rule cấm điều đó # không còn trong prompt nữa. Đây là lý do các fix văn phong ở app.py # chính KHÔNG có tác dụng khi câu trả lời đến từ Worker (khác hẳn # app.py chính — worker chỉ forward "dumb relay", không có logic # riêng). Tăng giới hạn lên 6000 ký tự — đủ chứa toàn bộ system prompt # điển hình của Yui mà vẫn tránh trường hợp cực đoan (system quá dài # làm tốn quota/context của Gemini Flash, Cohere). msgs.append({"role": "system", "content": system[:6000]}) for turn in history[-6:]: if isinstance(turn, (list, tuple)) and len(turn) >= 2: msgs.append({"role": "user", "content": str(turn[0])[:800]}) msgs.append({"role": "assistant", "content": str(turn[1])[:800]}) msgs.append({"role": "user", "content": message}) return msgs def _call_gemini(keys: list, msgs: list, system: str, max_tok: int = 1200) -> str: for key in random.sample(keys, min(3, len(keys))): try: contents = [{"role": "user" if m["role"] == "user" else "model", "parts": [{"text": m["content"]}]} for m in msgs if m["role"] != "system"] if not contents: continue payload: dict = { "contents": contents, "generationConfig": {"maxOutputTokens": max_tok, "temperature": 0.7} } if system: # 🆕 (2026-07-11) Cùng bug với _build_msgs — system[:1000] cắt # mất phần cấm "Ồ/Ừ"/teencode rules nằm sau vị trí 1000. Tăng # lên 6000, khớp với _build_msgs ở trên. payload["system_instruction"] = {"parts": [{"text": system[:6000]}]} r = _req.post( f"https://generativelanguage.googleapis.com/v1beta/models/" f"gemini-2.5-flash:generateContent?key={key}", json=payload, timeout=30 ) if r.status_code == 200: cands = r.json().get("candidates", []) if cands: text = (cands[0].get("content", {}) .get("parts", [{}])[0].get("text", "")) if text and len(text) > 5: return text except Exception as e: print(f"[WORKER] Gemini key fail: {type(e).__name__}") return "" # [DỌN 2026-07-17] Đã bỏ hẳn _call_cerebras() — Cerebras free tier đóng # cửa 17/08/2026, không còn được gọi trong cascade (xem worker_chat() # bên dưới: Gemini→Cohere→OpenRouter). def _call_cohere(keys: list, msgs: list, message: str, system: str, model: str = "", v2_models: set = None, history_raw: list = None, system_raw: str = "") -> str: """ 🆕 (2026-07-11) Trước đây hardcode model="command-r-08-2024" (model cũ, yếu) — không bao giờ tận dụng được command-a-plus mà Cohere Hunter bên app.py chính tìm ra. Giờ nhận `model` từ request; command-a* cần ClientV2 (messages format) khác hẳn command-r* (Client v1, preamble/chat_history) nên phải tự chọn đúng client theo model. 🆕 (2026-07-11) TẬN DỤNG QUOTA: Cohere tính theo LƯỢT REQUEST, không theo token — mỗi lần gọi tốn 1 lượt dù context/output nhỏ hay lớn. command-a-plus hỗ trợ 256K context / 8K output, nhưng `msgs` truyền vào đây đã bị _build_msgs() cắt ngắn (system[:6000], history mỗi turn [:800]) để an toàn cho các model context nhỏ hơn (Gemini Flash free tier, Cerebras). Nếu dùng command-a-plus (v2), dùng `history_raw`/`system_raw` (CHƯA CẮT) thay vì `msgs` đã cắt — tận dụng tối đa context/output mỗi lượt gọi thay vì lãng phí quota cho 1 request với input/output nhỏ hơn khả năng thật. """ v2_models = v2_models or set() _model = model or "command-r-08-2024" _needs_v2 = _model in v2_models or _model.lower().startswith("command-a") try: import cohere as _co api_key = random.choice(keys) if _needs_v2: # ── ClientV2: messages array — TẬN DỤNG TỐI ĐA 256K ctx/8K out ── co2 = _co.ClientV2(api_key=api_key) v2_msgs = [] _full_system = system_raw or system if _full_system: # 256K context → 40000 ký tự system prompt vẫn còn rất xa giới # hạn thật (an toàn hơn hẳn 6000 trước đây, tận dụng đúng quota) v2_msgs.append({"role": "system", "content": _full_system[:40000]}) _full_history = history_raw if history_raw is not None else None if _full_history: # Dùng nhiều turn hơn (12 thay vì 6) + mỗi turn giữ nguyên dài # hơn hẳn (2000 thay vì 800) — vẫn rất xa 256K nên an toàn. for turn in _full_history[-12:]: if isinstance(turn, (list, tuple)) and len(turn) >= 2: v2_msgs.append({"role": "user", "content": str(turn[0])[:2000]}) v2_msgs.append({"role": "assistant", "content": str(turn[1])[:2000]}) else: # Fallback: không có history_raw → dùng msgs đã cắt như cũ for m in msgs[:-1]: if m["role"] == "system": continue v2_msgs.append({ "role": "user" if m["role"] == "user" else "assistant", "content": m["content"], }) v2_msgs.append({"role": "user", "content": message}) r = co2.chat( model=_model, messages=v2_msgs, temperature=0.7, max_tokens=8192, # ↑ command-a-plus hỗ trợ 8K output — tận dụng hết ) # Trích text an toàn từ response ClientV2 try: for item in r.message.content: if getattr(item, "type", "") == "text": return item.text or "" except Exception: pass return "" # ── Client v1: preamble + chat_history (command-r*, context nhỏ hơn # command-a nhưng vẫn hỗ trợ 128K context / 4K output — TẬN DỤNG # TỐI ĐA giống hệt Client v2 ở trên và giống Space chính, vì # Cohere tính quota theo LƯỢT REQUEST chứ không theo token) ── client = _co.Client(api_key=api_key) # 🆕 (2026-07-14) Đồng bộ với Space chính: dùng history_raw (CHƯA CẮT) # thay vì msgs đã bị _build_msgs() cắt ngắn (mỗi turn [:800], chỉ 6 # lượt) — trước đây còn cắt thêm 1 lần nữa xuống [:400], lãng phí # quota cho 1 request nhỏ hơn nhiều so với khả năng thật của model. _full_history_v1 = history_raw if history_raw is not None else None hist = [] if _full_history_v1: # Giống Space chính: 20 lượt gần nhất, không cắt ngắn nội dung for turn in _full_history_v1[-20:]: if isinstance(turn, (list, tuple)) and len(turn) >= 2: hist.append({"role": "USER", "message": str(turn[0])}) hist.append({"role": "CHATBOT", "message": str(turn[1])}) else: # Fallback: không có history_raw → dùng msgs đã cắt như cũ for m in msgs[:-1]: if m["role"] == "system": continue hist.append({"role": "USER" if m["role"] == "user" else "CHATBOT", "message": m["content"][:400]}) # 🆕 (2026-07-14) Dùng system_raw (CHƯA CẮT) thay vì system[:6000] — # đồng bộ với Space chính (preamble không giới hạn độ dài). _full_system_v1 = system_raw or system r = client.chat( model=_model, message=message, chat_history=hist or None, preamble=_full_system_v1 if _full_system_v1 else None, # 🆕 (2026-07-14) Nâng từ 600 → 1200 → giờ 8192, khớp đúng Space # chính. Cohere tính quota theo request, không theo token — mỗi # lượt gọi đã tốn 1 request rồi, nên phải tận dụng hết output tối # đa model hỗ trợ, không có lý do gì giữ giới hạn thấp. max_tokens=8192, ) return r.text or "" except Exception as e: print(f"[WORKER] Cohere fail: {type(e).__name__}") return "" def _call_openrouter(keys: list, msgs: list, task: str = "chat") -> str: try: import openai as _oai model = ("qwen/qwen3-coder:free" if task in ("code", "math", "debug") else "openrouter/owl-alpha") client = _oai.OpenAI(api_key=random.choice(keys), base_url="https://openrouter.ai/api/v1") r = client.chat.completions.create( model=model, messages=msgs, # 🆕 (2026-07-14) Nâng từ 600 lên 1200 — đồng bộ với Cohere ở # trên, tránh phản hồi bị cắt ngắn thiếu tự nhiên. max_tokens=1200, timeout=20 ) return r.choices[0].message.content or "" except Exception as e: print(f"[WORKER] OpenRouter fail: {type(e).__name__}") return "" # ── Endpoints ───────────────────────────────────────────────────── @app.get("/") @app.head("/") async def root(): return {"yui": "worker", "role": WORKER_ROLE, "status": "ready"} @app.get("/health") @app.head("/health") async def health(): return {"status": "ok", "role": WORKER_ROLE, "has_keys": bool(_keys.get("gemini"))} @app.post("/api/worker") async def worker_chat(request: Request): """Endpoint chính — nhận task từ main space và trả kết quả.""" t0 = time.time() # Parse body try: body = await request.json() except Exception: return JSONResponse({"ok": False, "error": "invalid json"}, status_code=400) # Auth if not WORKER_SECRET or body.get("secret") != WORKER_SECRET: return JSONResponse({"ok": False, "error": "unauthorized"}, status_code=403) message = body.get("message", "").strip() system = body.get("system", "") history = body.get("history", []) task = body.get("task", "chat") # 🆕 (2026-07-11) Nhận model Cohere tốt nhất đã hunt được từ app.py chính # — trước đây worker luôn hardcode command-r-08-2024 (model cũ/yếu), # không bao giờ tận dụng được command-a-plus mà Cohere Hunter tìm ra. _chat_model = body.get("chat_model", "") or "command-r-08-2024" _v2_models = set(body.get("v2_models", []) or []) if not message: return JSONResponse({"ok": False, "error": "empty message"}, status_code=400) # Lấy keys keys = _get_keys() msgs = _build_msgs(system, history, message) reply = "" # Chọn model theo role và task gemini = keys.get("gemini", []) cerebras = keys.get("cerebras", []) cohere = keys.get("cohere", []) openrouter = keys.get("openrouter", []) # Gemini Flash — ưu tiên hàng đầu if gemini and not reply: max_tok = 1500 if task in ("code", "math", "debug") else 800 reply = _call_gemini(gemini, msgs, system, max_tok) # 🔄 MIGRATION (2026-07-17): Đã bỏ bước Cerebras khỏi cascade — Cerebras # free tier đóng cửa 17/08/2026 (bắt buộc thẻ thanh toán). Worker này là # process độc lập (không import được yui_model_router của Space chính), # nên thay vì thêm provider mới ở đây, đơn giản bỏ bước này — cascade # Gemini→Cohere→OpenRouter vẫn đủ 3 lớp fallback để hoạt động bình thường. # (Biến `cerebras` giữ nguyên phía trên — không dùng nữa nhưng vô hại, # không cần sửa _get_keys()/_fetch_keys_sync() ở phần đầu file.) # Cohere — fallback trung bình if cohere and not reply: reply = _call_cohere(cohere, msgs, message, system, model=_chat_model, v2_models=_v2_models, history_raw=history, system_raw=system) # OpenRouter — fallback cuối if openrouter and not reply: reply = _call_openrouter(openrouter, msgs, task) elapsed = round(time.time() - t0, 2) if reply: return JSONResponse({ "ok": True, "reply": reply, "role": WORKER_ROLE, "elapsed": elapsed }) else: return JSONResponse({ "ok": False, "error": "all models failed", "role": WORKER_ROLE, "elapsed": elapsed }, status_code=503) @app.post("/api/tts") async def worker_tts(request: Request): """TTS endpoint — edge-tts.""" try: body = await request.json() if body.get("secret") != WORKER_SECRET: return JSONResponse({"ok": False}, status_code=403) text = body.get("text", "").strip() if not text: return JSONResponse({"ok": False, "error": "no text"}) import edge_tts, base64 voice = "vi-VN-HoaiMyNeural" comm = edge_tts.Communicate(text[:600], voice) chunks = [] async for chunk in comm.stream(): if chunk["type"] == "audio": chunks.append(chunk["data"]) audio_b64 = base64.b64encode(b"".join(chunks)).decode() return JSONResponse({"ok": True, "audio_b64": audio_b64}) except Exception as e: return JSONResponse({"ok": False, "error": str(e)}) @app.post("/api/search") async def worker_search(request: Request): """Search endpoint — Tavily.""" try: body = await request.json() if body.get("secret") != WORKER_SECRET: return JSONResponse({"ok": False}, status_code=403) query = body.get("query", "") tavily = os.environ.get("TAVILY_KEY", "") if not tavily or not query: return JSONResponse({"ok": False, "results": []}) r = _req.post("https://api.tavily.com/search", json={"api_key": tavily, "query": query, "max_results": 3}, timeout=10) return JSONResponse({"ok": True, "results": r.json().get("results", [])}) except Exception as e: return JSONResponse({"ok": False, "results": [], "error": str(e)}) @app.post("/api/backup-brain") async def worker_backup_brain(request: Request): """Backup brain — forward to Cohere/OpenRouter khi main space quá tải.""" return await worker_chat(request) # ══════════════════════════════════════════════════════════════════════════════ # 🖥️ LOCAL MODEL (Qwen3-4B-Instruct qua llama.cpp) — GATE TRƯỚC KHI GỌI API # (admin/friend: câu đơn giản; guest: luôn thử) + LƯỚI AN TOÀN CUỐI khi # mọi API free đều fail. Aren dùng chung /data với Space chính (persistent # storage share) nên model file chỉ tải 1 lần, không cần thao tác tay (HF # Space free tier không có terminal — package cài qua requirements.txt, # model tự tải bằng huggingface_hub ngay trong code, xem # yui_local_model_server.py). 3 route: /api/local-llm (tác vụ nền, cũ), # /api/local-llm/chat (gate admin/friend, tool-calling), /api/local-llm/ # stream (gate guest, SSE). # ══════════════════════════════════════════════════════════════════════════════ from fastapi.responses import StreamingResponse import threading as _local_llm_threading try: from yui_local_model_server import ( handle_request as _local_llm_handle, handle_chat_request as _local_llm_chat_handle, handle_stream_request as _local_llm_stream_handle, init_local_model_background as _local_llm_init_bg, status as _local_llm_status, ) _LOCAL_MODEL_MODULE_OK = True except Exception as _lme: _LOCAL_MODEL_MODULE_OK = False print(f"[LOCAL-MODEL] ⚠️ Module yui_local_model_server chưa sẵn sàng: {_lme}") def _local_llm_check_secret(body: dict) -> bool: """Kiểm tra secret trong body JSON — dùng chung cho 3 route local-llm (giữ đúng cách handle_request() cũ tự kiểm tra secret bên trong nó; 2 route MỚI /chat và /stream cần route tự kiểm tra TRƯỚC vì handle_stream_request() là generator — không có chỗ nào "return lỗi" giữa chừng generator một cách gọn gàng như handle_request()/ handle_chat_request() làm được).""" if not WORKER_SECRET: return True return body.get("secret") == WORKER_SECRET @app.post("/api/local-llm") async def local_llm_endpoint(request: Request): """Endpoint để Space chính gọi model local trên Aren (fallback cuối khi mọi API free đều fail — xem yui_model_router.py phía Space chính, hàm _try_local_model_fallback()).""" if not _LOCAL_MODEL_MODULE_OK: return JSONResponse({"ok": False, "reply": "", "error": "Module chưa load được"}, status_code=503) try: body = await request.json() except Exception: return JSONResponse({"ok": False, "reply": "", "error": "invalid json"}, status_code=400) # handle_request() tự kiểm tra secret bên trong — chạy đồng bộ (llama.cpp # không async-native), bọc to_thread để không block event loop chính # trong lúc CPU đang bận generate (có thể 10-25s). result = await asyncio.to_thread(_local_llm_handle, body) return JSONResponse(result) # 🆕 (2026-09-01) Route MỚI — tool-calling, dùng cho gate admin/friend # (yui_model_router._try_local_model_chat, gọi từ app.py._qwen_local_gate # ở Space chính). Cùng pattern to_thread như /api/local-llm ở trên — chạy # đồng bộ (create_chat_completion không async-native), không chặn loop. @app.post("/api/local-llm/chat") async def local_llm_chat_endpoint(request: Request): """Endpoint tool-calling cho gate admin/friend — xem yui_local_model_server.generate_chat() phía Aren và yui_model_router._try_local_model_chat() phía Space chính.""" if not _LOCAL_MODEL_MODULE_OK: return JSONResponse({"ok": False, "reply": "", "tool_calls": None, "error": "Module chưa load được"}, status_code=503) try: body = await request.json() except Exception: return JSONResponse({"ok": False, "reply": "", "tool_calls": None, "error": "invalid json"}, status_code=400) result = await asyncio.to_thread(_local_llm_chat_handle, body) return JSONResponse(result) # 🆕 (2026-09-01) Route MỚI — streaming SSE, dùng cho gate guest # (yui_model_router._try_local_model_stream, gọi từ app.py._qwen_local_gate # ở Space chính). handle_stream_request() là generator SYNC (giữ _lock # llama.cpp suốt thời gian sinh, giống handle_request()) — bọc chạy trong # thread riêng qua 1 hàng đợi (queue) để StreamingResponse (async) có thể # yield từng đoạn NGAY khi generator sync sinh ra, không chặn event loop # chính trong lúc CPU đang generate. @app.post("/api/local-llm/stream") async def local_llm_stream_endpoint(request: Request): """Endpoint streaming SSE cho gate guest — xem yui_local_model_server.generate_stream() phía Aren và yui_model_router._try_local_model_stream() phía Space chính.""" if not _LOCAL_MODEL_MODULE_OK: return JSONResponse({"ok": False, "error": "Module chưa load được"}, status_code=503) try: body = await request.json() except Exception: return JSONResponse({"ok": False, "error": "invalid json"}, status_code=400) if not _local_llm_check_secret(body): return JSONResponse({"ok": False, "error": "Sai secret"}, status_code=403) async def _sse(): # Generator sync (handle_stream_request) chạy trong thread riêng, # đẩy từng chunk qua queue — event loop chính chỉ await queue, # không bao giờ bị CPU-bound generate() làm nghẽn các request khác # (Discord bot, health-check...) đang chạy song song trên Aren. _queue: asyncio.Queue = asyncio.Queue() _SENTINEL = object() _loop = asyncio.get_event_loop() # lấy TRƯỚC khi thread khởi chạy — _producer (closure) cần biến này sẵn có def _producer(): try: for chunk in _local_llm_stream_handle(body): asyncio.run_coroutine_threadsafe(_queue.put(chunk), _loop).result() except Exception as _e: print(f"[LOCAL-MODEL-STREAM] ⚠️ Producer lỗi: {_e.__class__.__name__}: {str(_e)[:150]}") finally: asyncio.run_coroutine_threadsafe(_queue.put(_SENTINEL), _loop).result() _thread = _local_llm_threading.Thread(target=_producer, daemon=True, name="local-llm-stream-producer") _thread.start() while True: chunk = await _queue.get() if chunk is _SENTINEL: break yield f"data: {chunk}\n\n" yield "data: [DONE]\n\n" return StreamingResponse(_sse(), media_type="text/event-stream") @app.get("/api/local-llm/status") async def local_llm_status_endpoint(request: Request): """Status có auth; trả heartbeat/progress thực tế của request Aren.""" if not _LOCAL_MODEL_MODULE_OK: return JSONResponse({"loaded": False, "error": "Module chưa load được"}, status_code=503) _status_secret = ( request.headers.get("X-Local-Model-Secret", "") or request.headers.get("X-Worker-Secret", "") ) if WORKER_SECRET and _status_secret != WORKER_SECRET: return JSONResponse({"loaded": False, "error": "Forbidden"}, status_code=403) return JSONResponse(_local_llm_status()) # ══════════════════════════════════════════════════════════════════════════════ # 📦 WORKER COLD HISTORY — Lưu lịch sử chat Guest/Friend từ main space # Main space push cold history về đây để giảm RAM chính # Tự xóa history cũ > 7 ngày | Auth: WORKER_SECRET # Sharding: main space tự chọn worker theo hash(user_id) # ══════════════════════════════════════════════════════════════════════════════ import sqlite3 as _hist_sq import json as _hist_js import time as _hist_tm import hmac as _hist_hm import hashlib as _hist_hs _HIST_DB_PATH = "/data/worker_cold_history.db" _HIST_TTL_DAYS = 7 _HIST_MAX_TURNS = 30 _HIST_SECRET = WORKER_SECRET # Dùng lại WORKER_SECRET đã có sẵn def _hist_db_init(): try: with _hist_sq.connect(_HIST_DB_PATH, timeout=10) as _c: _c.execute(""" CREATE TABLE IF NOT EXISTS cold_history ( user_hash TEXT PRIMARY KEY, role TEXT DEFAULT 'guest', history_json TEXT NOT NULL, updated_at REAL NOT NULL, turn_count INTEGER DEFAULT 0 ) """) _c.execute("CREATE INDEX IF NOT EXISTS idx_upd ON cold_history(updated_at)") print(f"[HIST-WORKER-{WORKER_ROLE.upper()}] ✅ Cold History DB ready: {_HIST_DB_PATH}") except Exception as _e: print(f"[HIST-WORKER] ⚠️ DB init: {_e}") _hist_db_init() def _hist_auth_ok(secret: str) -> bool: if not _HIST_SECRET or len(_HIST_SECRET) < 8: return False try: return _hist_hm.compare_digest( secret.encode("utf-8"), _HIST_SECRET.encode("utf-8") ) except Exception: return False def _hist_cleanup_ttl(): try: _cut = _hist_tm.time() - (_HIST_TTL_DAYS * 86400) with _hist_sq.connect(_HIST_DB_PATH, timeout=5) as _c: _c.execute("DELETE FROM cold_history WHERE updated_at < ?", (_cut,)) except Exception: pass @app.post("/api/hist-save") async def hist_save_endpoint(request: Request): """Nhận cold history từ main space → lưu vào SQLite worker.""" try: _data = await request.json() except Exception: return JSONResponse({"ok": False, "error": "Bad JSON"}, status_code=400) if not _hist_auth_ok(_data.get("secret", "")): return JSONResponse({"ok": False, "error": "Unauthorized"}, status_code=401) _user_hash = str(_data.get("user_hash", ""))[:64] _role = str(_data.get("role", "guest"))[:10] _new_turns = _data.get("new_turns", []) _action = _data.get("action", "append") if not _user_hash or not isinstance(_new_turns, list): return JSONResponse({"ok": False, "error": "Missing fields"}, status_code=400) _clean = [[str(t[0])[:4000], str(t[1])[:4000]] for t in _new_turns if isinstance(t, (list, tuple)) and len(t) >= 2] try: with _hist_sq.connect(_HIST_DB_PATH, timeout=10) as _c: _row = _c.execute( "SELECT history_json FROM cold_history WHERE user_hash=?", (_user_hash,) ).fetchone() if _action == "append" and _row: _merged = _hist_js.loads(_row[0]) + _clean else: _merged = _clean _merged = _merged[-_HIST_MAX_TURNS:] _c.execute( "INSERT OR REPLACE INTO cold_history VALUES (?,?,?,?,?)", (_user_hash, _role, _hist_js.dumps(_merged, ensure_ascii=False), _hist_tm.time(), len(_merged)) ) import random as _rh if _rh.random() < 0.02: _hist_cleanup_ttl() return JSONResponse({"ok": True, "turns": len(_merged)}) except Exception as _e: return JSONResponse({"ok": False, "error": str(_e)[:80]}, status_code=500) @app.post("/api/hist-load") async def hist_load_endpoint(request: Request): """Trả cold history về main space khi local session hết hạn.""" try: _data = await request.json() except Exception: return JSONResponse({"ok": False, "error": "Bad JSON"}, status_code=400) if not _hist_auth_ok(_data.get("secret", "")): return JSONResponse({"ok": False, "error": "Unauthorized"}, status_code=401) _user_hash = str(_data.get("user_hash", ""))[:64] if not _user_hash: return JSONResponse({"ok": False, "error": "Missing user_hash"}, status_code=400) try: with _hist_sq.connect(_HIST_DB_PATH, timeout=5) as _c: _row = _c.execute( "SELECT history_json, updated_at FROM cold_history WHERE user_hash=?", (_user_hash,) ).fetchone() if _row: return JSONResponse({ "ok": True, "history": _hist_js.loads(_row[0]), "turns": len(_hist_js.loads(_row[0])), "age_h": round((_hist_tm.time() - _row[1]) / 3600, 1) }) return JSONResponse({"ok": True, "history": [], "turns": 0}) except Exception as _e: return JSONResponse({"ok": False, "error": str(_e)[:80]}, status_code=500) @app.get("/api/hist-stats") async def hist_stats_endpoint(request: Request): """Admin xem thống kê cold history của worker này.""" if not _hist_auth_ok(request.headers.get("X-Worker-Secret", "")): return JSONResponse({"ok": False, "error": "Unauthorized"}, status_code=401) try: with _hist_sq.connect(_HIST_DB_PATH, timeout=5) as _c: _total = _c.execute("SELECT COUNT(*) FROM cold_history").fetchone()[0] _turns = _c.execute("SELECT SUM(turn_count) FROM cold_history").fetchone()[0] or 0 _roles = _c.execute("SELECT role, COUNT(*) FROM cold_history GROUP BY role").fetchall() return JSONResponse({"ok": True, "role": WORKER_ROLE, "total_users": _total, "total_turns": _turns, "by_role": dict(_roles)}) except Exception as _e: return JSONResponse({"ok": False, "error": str(_e)[:80]}) print(f"[HIST-WORKER-{WORKER_ROLE.upper()}] ✅ 3 endpoints: hist-save | hist-load | hist-stats") # ── Startup ─────────────────────────────────────────────────────── @app.on_event("startup") async def startup(): import threading print(f"[WORKER-{WORKER_ROLE.upper()}] 🚀 Khởi động...") # Fetch keys ngay lúc startup (sync trong thread để không block) threading.Thread(target=_fetch_keys_sync, daemon=True).start() # 🆕 (2026-08-09) Khởi động model local ở NỀN — dùng bản _background # (không phải init_local_model() đồng bộ) để KHÔNG chặn app khởi động. # Lần đầu tải model (~4.4GB) + build có thể mất vài phút — nếu chặn ở # đây, HF health-check có thể coi Aren "không phản hồi". Model tải/ # load ngầm phía sau, trong lúc đó /api/local-llm trả lỗi rõ ràng # "chưa sẵn sàng" thay vì treo — Space chính tự hiểu là fail bình # thường, không ảnh hưởng gì tới các endpoint khác của worker. if _LOCAL_MODEL_MODULE_OK: _local_llm_init_bg() # Background: refresh keys mỗi 25 phút async def _refresh_loop(): while True: await asyncio.sleep(1500) _fetch_keys_sync() asyncio.create_task(_refresh_loop()) print(f"[WORKER-{WORKER_ROLE.upper()}] ✅ Sẵn sàng! Role={WORKER_ROLE}") # 🆕 (2026-07-13) Đã bỏ khối "Embedded Discord Bot" từng gắn ở đây — # bot Discord giờ chạy độc lập trên Wispbyte (xem README_WISPBYTE.md), # gọi API sang main Space (bachdathan-ural-ai) trực tiếp qua HF_API_URL. # Worker này quay lại đúng vai trò gốc: model-router phụ, không còn # chạy bot Discord chung nữa.