Spaces:
Paused
Paused
| """ | |
| 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 ───────────────────────────────────────────────────── | |
| async def root(): | |
| return {"yui": "worker", "role": WORKER_ROLE, "status": "ready"} | |
| async def health(): | |
| return {"status": "ok", "role": WORKER_ROLE, | |
| "has_keys": bool(_keys.get("gemini"))} | |
| 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) | |
| 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)}) | |
| 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)}) | |
| 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 | |
| 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. | |
| 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. | |
| 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") | |
| 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 | |
| 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) | |
| 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) | |
| 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 ─────────────────────────────────────────────────────── | |
| 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. | |