Spaces:
Paused
Paused
| #!/usr/bin/env python3 | |
| """ | |
| aichat 訊息中繼（REST）。自架，不經過任何第三方轉發站。 | |
| POST /api/aichat/put 送一則訊息 | |
| GET /api/aichat/poll?channel=&since=&wait= 長輪詢取回後續 | |
| GET /api/aichat/peek?channel=&limit= 不阻塞地抓最近幾則（測試／除錯用） | |
| GET /api/aichat/health 存活與統計 | |
| 為什麼要有這個 | |
| -------------- | |
| 原本訊息是經過 ntfy.sh 之類的**公開 anyone-can-read 轉發站**：內文明文、topic 誰都能猜， | |
| 而且免費站台會限速、會把跑太多的來源 IP 直接封掉且不會自己恢復。訊息走別人的機器， | |
| 也就等於把對話交出去。 | |
| 現在自己架一個：資料只落在這個 Space（SQLite），走 HTTPS，需要 HF token 才能進。 | |
| 「不會漏掉」也不再靠對方的善意： | |
| * 每則訊息有單調遞增的 ``id``。接收端記住最後收到的 ``id``，重連帶 ``since=<id>``， | |
| 伺服器精確地從那裡接續 —— 不是「重送我沒收到的」，是「從沒斷過的那裡接」。 | |
| * 送出端（bridge / 網頁）自己寫磁碟 outbox 才送出，失敗只是延後。 | |
| * 就算游標太舊被修剪，bridge 端還有整篇快照（history.json）可以補齊。 | |
| ## 路由 | |
| 挂在 app.py 的 FastAPI 上（``app.include_router(relay.router)``），不動 app.py 既有的 | |
| ``/``、``/files``、``/download``、``/ws``。前綴 ``/api/aichat`` 跟那些都不撞。 | |
| ## 授權 | |
| Space 本身是 private，HF 的閘門已經擋掉沒 token 的人。這個 Space 的 AICHAT_TOKEN 有設定 | |
| 時，還要再對一次 ``X-Aichat-Token``（雙層）。 | |
| """ | |
| from __future__ import annotations | |
| import asyncio | |
| import json | |
| import os | |
| import re | |
| import sqlite3 | |
| import time | |
| from pathlib import Path | |
| from fastapi import APIRouter, Header, HTTPException, Query, Request | |
| from fastapi.middleware.cors import CORSMiddleware | |
| from fastapi.responses import JSONResponse | |
| router = APIRouter(prefix="/api/aichat", tags=["aichat"]) | |
| VERSION = 2 | |
| # ---------------------------------------------------------------- CORS | |
| # 網頁在 aichatv（static Space）、中繼在 testzerogpu2，兩者不同 host → 瀏覽器一定 | |
| # 會先送 OPTIONS 預檢。這裡放行指定來源，**不是**萬用 *（萬用 * 不能帶 | |
| # Authorization 標頭，實際等於不能用）。 | |
| ALLOW_ORIGINS = [ | |
| o.strip() for o in os.environ.get( | |
| "AICHAT_ORIGINS", | |
| "https://hellosun-aichatv.static.hf.space," | |
| "https://huggingface.co," | |
| "http://localhost:8000,http://127.0.0.1:8000" | |
| ).split(",") if o.strip() | |
| ] | |
| def install_cors(app) -> None: | |
| app.add_middleware( | |
| CORSMiddleware, | |
| allow_origins=ALLOW_ORIGINS, | |
| allow_credentials=False, | |
| allow_methods=["GET", "POST", "DELETE", "OPTIONS"], | |
| allow_headers=["Authorization", "Content-Type", "X-Aichat-Token"], | |
| max_age=600, | |
| ) | |
| # ---------------------------------------------------------------- 設定 | |
| DB_PATH = Path(os.environ.get("AICHAT_DB", "/home/user/data/aichat_relay.db")) | |
| SHARED_TOKEN = os.environ.get("AICHAT_TOKEN", "").strip() | |
| KEEP_PER_CHANNEL = int(os.environ.get("AICHAT_KEEP", "2000")) # 每個 channel 留幾則 | |
| CHANNEL_MAX = 64 | |
| CHANNEL_RE = re.compile(r"^[A-Za-z0-9._~-]{1,%d}$" % CHANNEL_MAX) | |
| PAYLOAD_MAX = 200_000 # 單則最多 200 KB（附件切片後仍在範圍內） | |
| MAX_WAIT = 25 # 長輪詢最多吊幾秒：HF 的 proxy 逾時前一定要回 | |
| # ---------------------------------------------------------------- 儲存 | |
| def _connect() -> sqlite3.Connection: | |
| DB_PATH.parent.mkdir(parents=True, exist_ok=True) | |
| # Space 重啟後檔案還在（重建才會沒），所以 WAL + 自動建表 | |
| con = sqlite3.connect(DB_PATH, timeout=10, check_same_thread=False) | |
| con.execute("PRAGMA journal_mode=WAL") | |
| con.execute("PRAGMA synchronous=NORMAL") | |
| con.execute( | |
| "CREATE TABLE IF NOT EXISTS msgs (" | |
| " id INTEGER PRIMARY KEY AUTOINCREMENT," | |
| " channel TEXT NOT NULL," | |
| " payload TEXT NOT NULL," | |
| " ts REAL NOT NULL)" | |
| ) | |
| con.execute("CREATE INDEX IF NOT EXISTS msgs_ch ON msgs(channel, id)") | |
| return con | |
| _con: sqlite3.Connection | None = None | |
| _lock = asyncio.Lock() | |
| # channel → 一組等待中的未來；有新訊息就把它們叫醒 | |
| _waiters: dict[str, set[asyncio.Future]] = {} | |
| def db() -> sqlite3.Connection: | |
| global _con | |
| if _con is None: | |
| _con = _connect() | |
| return _con | |
| def check_channel(c: str) -> str: | |
| if not c or not CHANNEL_RE.match(c): | |
| raise HTTPException(400, f"channel 只能是 1..{CHANNEL_MAX} 個字元，且只限 A-Za-z0-9._~-") | |
| return c | |
| def prune(channel: str) -> None: | |
| con = db() | |
| row = con.execute("SELECT COUNT(*) FROM msgs WHERE channel=?", (channel,)).fetchone() | |
| n = row[0] if row else 0 | |
| if n > KEEP_PER_CHANNEL: | |
| con.execute( | |
| "DELETE FROM msgs WHERE channel=? AND id <= (" | |
| " SELECT id FROM msgs WHERE channel=? ORDER BY id DESC LIMIT 1 OFFSET ?)", | |
| (channel, channel, KEEP_PER_CHANNEL - 1), | |
| ) | |
| con.commit() | |
| # ---------------------------------------------------------------- 授權 | |
| def authorize(x_aichat_token: str | None) -> None: | |
| """Space 已經是 private（HF 擋掉沒 token 的）。這裡是第二層，設定了才啟用。""" | |
| if not SHARED_TOKEN: | |
| return | |
| if not x_aichat_token or x_aichat_token != SHARED_TOKEN: | |
| raise HTTPException(401, "X-Aichat-Token 不對") | |
| # ---------------------------------------------------------------- 路由 | |
| async def health() -> JSONResponse: | |
| con = db() | |
| total = con.execute("SELECT COUNT(*) FROM msgs").fetchone()[0] | |
| chans = con.execute("SELECT COUNT(DISTINCT channel) FROM msgs").fetchone()[0] | |
| oldest = con.execute("SELECT MIN(id) FROM msgs").fetchone()[0] | |
| newest = con.execute("SELECT MAX(id) FROM msgs").fetchone()[0] | |
| return JSONResponse({ | |
| "ok": True, "version": VERSION, | |
| "hub": "self-hosted (no third party)", | |
| "messages": total, "channels": chans, | |
| "oldest_id": oldest, "newest_id": newest, | |
| "shared_token": bool(SHARED_TOKEN), | |
| "waiters": sum(len(v) for v in _waiters.values()), | |
| }) | |
| async def put(request: Request, x_aichat_token: str | None = Header(None)) -> JSONResponse: | |
| authorize(x_aichat_token) | |
| try: | |
| body = await request.json() | |
| except Exception: | |
| raise HTTPException(400, "body 不是合法 JSON") | |
| channel = check_channel(str(body.get("channel") or "")) | |
| payload = body.get("payload") | |
| if not isinstance(payload, str): | |
| raise HTTPException(400, "payload 必須是字串") | |
| if len(payload.encode("utf-8")) > PAYLOAD_MAX: | |
| raise HTTPException(413, f"payload 超過 {PAYLOAD_MAX} bytes") | |
| con = db() | |
| cur = con.execute("INSERT INTO msgs(channel, payload, ts) VALUES(?,?,?)", | |
| (channel, payload, time.time())) | |
| con.commit() | |
| mid = cur.lastrowid | |
| prune(channel) | |
| # 叫醒在長輪詢的人 | |
| for fut in list(_waiters.get(channel, ())): | |
| if not fut.done(): | |
| fut.set_result(True) | |
| return JSONResponse({"ok": True, "id": mid, "ts": time.time()}) | |
| async def peek(channel: str = Query(...), limit: int = Query(50, ge=1, le=500), | |
| x_aichat_token: str | None = Header(None)) -> JSONResponse: | |
| """不阻塞：抓最近 limit 則（給測試與除錯用）。""" | |
| authorize(x_aichat_token) | |
| channel = check_channel(channel) | |
| con = db() | |
| rows = con.execute( | |
| "SELECT id, payload, ts FROM msgs WHERE channel=? ORDER BY id DESC LIMIT ?", | |
| (channel, limit)).fetchall() | |
| items = [{"id": r[0], "payload": r[1], "ts": r[2]} for r in reversed(rows)] | |
| return JSONResponse({"items": items, "cursor": items[-1]["id"] if items else 0}) | |
| async def poll(channel: str = Query(...), | |
| since: int = Query(0, ge=0), | |
| wait: int = Query(20, ge=0, le=MAX_WAIT), | |
| x_aichat_token: str | None = Header(None)) -> JSONResponse: | |
| """長輪詢：等到有新訊息或逾時。回傳的 cursor 請記住，下次帶回來。""" | |
| authorize(x_aichat_token) | |
| channel = check_channel(channel) | |
| wait = max(0, min(wait, MAX_WAIT)) | |
| def fetch() -> list: | |
| con = db() | |
| rows = con.execute( | |
| "SELECT id, payload, ts FROM msgs WHERE channel=? AND id>? ORDER BY id ASC LIMIT 500", | |
| (channel, since)).fetchall() | |
| return [{"id": r[0], "payload": r[1], "ts": r[2]} for r in rows] | |
| items = fetch() | |
| if items or wait == 0: | |
| return JSONResponse({"items": items, | |
| "cursor": items[-1]["id"] if items else since}) | |
| # 沒有東西 → 掛一個 future，有人送訊息時會被 set_result 叫醒 | |
| fut: asyncio.Future = asyncio.get_running_loop().create_future() | |
| _waiters.setdefault(channel, set()).add(fut) | |
| try: | |
| await asyncio.wait_for(fut, timeout=wait) | |
| except asyncio.TimeoutError: | |
| pass | |
| finally: | |
| _waiters.get(channel, set()).discard(fut) | |
| items = fetch() | |
| return JSONResponse({"items": items, | |
| "cursor": items[-1]["id"] if items else since}) | |
| async def wipe(channel: str, x_aichat_token: str | None = Header(None)) -> JSONResponse: | |
| """清掉某個 channel 的全部訊息（測試重置用）。""" | |
| authorize(x_aichat_token) | |
| channel = check_channel(channel) | |
| con = db() | |
| con.execute("DELETE FROM msgs WHERE channel=?", (channel,)) | |
| con.commit() | |
| return JSONResponse({"ok": True, "channel": channel}) | |
| async def _startup() -> None: | |
| try: | |
| db() | |
| except Exception as e: # 磁碟壞掉不該讓整個 Space 掛掉 | |
| print(f"[relay] 資料庫開不起來：{e!r}") | |
| print(f"[relay] v{VERSION} 就緒 db={DB_PATH} 共享 token={'有' if SHARED_TOKEN else '無'}") |