import hashlib import ipaddress import json import os import re import threading import time from typing import Any from urllib.parse import quote from fastapi import APIRouter, Request from fastapi.responses import JSONResponse from . import parse_core, pipeline_dispatch, source_files from .doc_store import parse_queue_docs from .document_identity import content_hash from .parse_core import ParseRequest from .storage_lock import serving_lock router = APIRouter() GITHUB_TOKEN = os.environ.get("BHA_PROOFREAD_GITHUB_TOKEN", "") PIPELINE_REPOSITORY = os.environ.get("BHA_PROOFREAD_PIPELINE_REPOSITORY", "anftm/pipeline") PARSE_ARCHIVE_ID = 25 PARSE_QUEUE_PREFIX = "【文章待拆分】" PARSE_OCR_BASE_URL = os.environ.get( "BHA_PARSE_OCR_BASE_URL", "https://raw.githubusercontent.com/banned-historical-archives/banned-historical-archives25/ocr_cache", ).rstrip("/") PARSE_OCR_MAX_BYTES = 2 * 1024 * 1024 PARSE_MAX_PAYLOAD_CHARACTERS = 2_000_000 PARSE_RATE_LIMIT = 30 PARSE_RATE_WINDOW_SECONDS = 60 * 60 TRUST_PROXY = os.environ.get("BHA_TRUST_PROXY", "1" if os.environ.get("SPACE_ID") else "0") == "1" try: PROXY_NETWORKS = tuple( ipaddress.ip_network(value.strip()) for value in os.environ.get("BHA_TRUST_PROXY_CIDRS", "").split(",") if value.strip() ) except ValueError: PROXY_NETWORKS = () try: PROXY_HOPS = max(1, int(os.environ.get("BHA_TRUST_PROXY_HOPS", "1"))) except ValueError: PROXY_HOPS = 1 parse_requests: dict[str, list[float]] = {} parse_rate_lock = threading.Lock() parse_dispatches: dict[str, dict[str, Any]] = {} parse_dispatch_lock = threading.Lock() async def dispatch_parse(payload: dict[str, Any]) -> None: await pipeline_dispatch.dispatch(payload, "submit-parse", GITHUB_TOKEN, PIPELINE_REPOSITORY) def parse_client_key(request: Request) -> str: peer = request.client.host if request.client else "unknown" trusted_peer = False try: peer_address = ipaddress.ip_address(peer) trusted_peer = TRUST_PROXY and bool(PROXY_NETWORKS) and any(peer_address in network for network in PROXY_NETWORKS) except ValueError: pass forwarded = request.headers.get("x-forwarded-for", "") if trusted_peer else "" candidates = [value.strip() for value in forwarded.split(",") if value.strip()] if candidates and len(candidates) > PROXY_HOPS: candidate = candidates[-PROXY_HOPS - 1] try: ipaddress.ip_address(candidate) return candidate except ValueError: pass return peer def parse_rate_error(request: Request) -> JSONResponse | None: key = parse_client_key(request) now = time.monotonic() with parse_rate_lock: recent = [value for value in parse_requests.get(key, []) if now - value < PARSE_RATE_WINDOW_SECONDS] if len(recent) >= PARSE_RATE_LIMIT: return JSONResponse({"error": "parse rate limit exceeded"}, status_code=429) recent.append(now) parse_requests[key] = recent if len(parse_requests) > 10000: for identity in list(parse_requests): parse_requests[identity] = [value for value in parse_requests[identity] if now - value < PARSE_RATE_WINDOW_SECONDS] if not parse_requests[identity]: parse_requests.pop(identity, None) while len(parse_requests) > 10000: oldest = min(parse_requests, key=lambda identity: min(parse_requests[identity])) parse_requests.pop(oldest, None) request.state.parse_rate_key = key request.state.parse_rate_timestamp = now return None def release_parse_rate(request: Request) -> None: key = getattr(request.state, "parse_rate_key", None) timestamp = getattr(request.state, "parse_rate_timestamp", None) if key is None or timestamp is None: return with parse_rate_lock: values = parse_requests.get(key, []) if timestamp in values: values.remove(timestamp) if values: parse_requests[key] = values else: parse_requests.pop(key, None) def duplicate_parse_response(request_id: str) -> JSONResponse | None: now = time.time() with parse_dispatch_lock: for key, value in list(parse_dispatches.items()): if now - float(value["created_at"]) > 86400: parse_dispatches.pop(key, None) existing = parse_dispatches.get(request_id) if existing is not None: return JSONResponse({ "ok": True, "status": "duplicate", "original_status": existing["status"], "request_id": request_id, "tasks": 1, }) if len(parse_dispatches) >= 1000: oldest = min(parse_dispatches, key=lambda key: parse_dispatches[key]["created_at"]) parse_dispatches.pop(oldest, None) parse_dispatches[request_id] = {"created_at": now, "status": "processing"} return None def parse_queue_source(doc_id: str) -> dict[str, Any] | None: source = source_files.get_source_doc(doc_id) if not parse_core.is_queue_source(source, PARSE_ARCHIVE_ID, PARSE_QUEUE_PREFIX): return None return source @router.get("/api/parse/queue") async def parse_queue(page: int = 1, page_size: int = 20): if page < 1 or page_size < 1 or page_size > 100: return JSONResponse({"error": "invalid pagination"}, status_code=400) with serving_lock(): total, sources = parse_queue_docs(PARSE_ARCHIVE_ID, PARSE_QUEUE_PREFIX, page, page_size) return { "archive_id": PARSE_ARCHIVE_ID, "prefix": PARSE_QUEUE_PREFIX, "total": total, "page": page, "page_size": page_size, "items": [parse_core.queue_item(source) for source in sources], } @router.get("/api/parse/{doc_id}/ocr") async def parse_ocr_page(doc_id: str, page: int = 1): source = parse_queue_source(doc_id) if not source: return JSONResponse({"error": "parse item not found"}, status_code=404) article = source.get("article") or {} page_start = int(article.get("page_start") or 0) page_end = int(article.get("page_end") or 0) if page < page_start or page > page_end: return JSONResponse({"error": "page is outside the source range"}, status_code=404) publication_id = str(source.get("publication_id") or "") if not re.fullmatch(r"[0-9a-f-]{36}", publication_id, re.IGNORECASE): return JSONResponse({"error": "invalid publication id"}, status_code=500) url = f"{PARSE_OCR_BASE_URL}/{quote(publication_id, safe='')}/{page}.json" try: response = await source_files.source_client.get(url, follow_redirects=True) response.raise_for_status() if len(response.content) > PARSE_OCR_MAX_BYTES: return JSONResponse({"error": "OCR page is too large"}, status_code=502) payload = response.json() except Exception as exc: return JSONResponse({"error": f"OCR page unavailable: {exc}"}, status_code=502) if not isinstance(payload, dict) or not isinstance(payload.get("ocr_results"), list): return JSONResponse({"error": "OCR page is invalid"}, status_code=502) return JSONResponse({"publication_id": publication_id, "page": page, **payload}, headers={ "Cache-Control": "public, max-age=300", }) @router.post("/api/parse", status_code=202) async def submit_parse(body: ParseRequest, request: Request): if not GITHUB_TOKEN: return JSONResponse({"error": "parsing is not configured"}, status_code=503) if len(json.dumps(body.model_dump(), ensure_ascii=False)) > PARSE_MAX_PAYLOAD_CHARACTERS: return JSONResponse({"error": "parse payload is too large"}, status_code=413) source = parse_queue_source(body.doc_id) if not source: return JSONResponse({"error": "parse item not found or already resolved"}, status_code=404) current_hash = content_hash(source) if body.base_content_hash and body.base_content_hash != current_hash: return JSONResponse({"error": "parse item changed since preview", "content_hash": current_hash}, status_code=409) articles, error = parse_core.validate_articles(body.articles, source, PARSE_QUEUE_PREFIX) if error: return JSONResponse({"error": error}, status_code=400) payload: dict[str, Any] = { "doc_id": body.doc_id, "archive_id": source.get("archive_id"), "publication_id": source.get("publication_id"), "kind": "parse", "body": { "articles": articles, "locator": { "title": (source.get("article") or {}).get("title"), "page_start": (source.get("article") or {}).get("page_start"), "page_end": (source.get("article") or {}).get("page_end"), }, "source_files": source.get("source_files") or [], "description": body.description, }, } request_id = hashlib.sha256(json.dumps(payload, ensure_ascii=False, sort_keys=True, separators=(",", ":")).encode()).hexdigest() payload["request_id"] = request_id if duplicate := duplicate_parse_response(request_id): return duplicate if error := parse_rate_error(request): with parse_dispatch_lock: parse_dispatches.pop(request_id, None) return error try: await dispatch_parse(payload) except Exception as exc: with parse_dispatch_lock: parse_dispatches.pop(request_id, None) release_parse_rate(request) return JSONResponse({"error": str(exc)}, status_code=502) with parse_dispatch_lock: parse_dispatches[request_id] = {"created_at": time.time(), "status": "accepted"} return {"ok": True, "status": "accepted", "request_id": request_id, "tasks": 1}