from __future__ import annotations import os from contextlib import asynccontextmanager from html import escape from pathlib import Path from typing import Literal from urllib.parse import quote, urlsplit from uuid import uuid4 from fastapi import Depends, FastAPI, Request, Response from fastapi.exceptions import RequestValidationError from fastapi.responses import FileResponse, HTMLResponse, JSONResponse from pydantic import BaseModel, ConfigDict, Field from .auth import Auth from .db import StoreError, configured_store from .learning import LearningService from .worker import BrowserWorker ROOT = Path(__file__).resolve().parents[2] class Input(BaseModel): model_config = ConfigDict(extra="forbid") class Credentials(Input): username: str = Field(min_length=1, max_length=64) password: str = Field(min_length=1, max_length=256) class NewTask(Input): name: str = Field(min_length=1, max_length=120) platform: Literal["demo", "xiaohongshu", "douyin", "other"] entry_url: str = Field(min_length=1, max_length=2000) goal: str = Field(default="", max_length=4000) account_key: str = Field(min_length=2, max_length=80) parameters: dict[str, str] = Field(default_factory=dict) class Ready(Input): identity_label: str = Field(min_length=1, max_length=120) class Capture(Input): note: str = Field(default="", max_length=1000) class Action(Input): kind: Literal["fill", "click", "read"] selector: str = Field(min_length=1, max_length=1000) value: str | None = Field(default=None, max_length=2000) parameter: str | None = Field(default=None, max_length=64) output_name: str | None = Field(default=None, max_length=64) class Candidate(Input): name: str | None = Field(default=None, max_length=120) class Verify(Input): parameters: dict[str, str] = Field(default_factory=dict) expected_outputs: dict[str, str | list[str]] = Field(default_factory=dict) sample_name: str = Field(default="变参试跑", max_length=120) def create_app(home: Path | None = None, worker=None, *, bootstrap_admin: bool = True, database_url: str | None = None, database_schema: str | None = None) -> FastAPI: home = (home or Path(os.environ.get("ASSISTANT_LEARNING_HOME", ROOT / ".runtime" / "learning"))).resolve() store = configured_store(home, database_url, schema=database_schema) auth = Auth(store) if bootstrap_admin: auth.bootstrap_admin() browser = worker or BrowserWorker(ROOT, home) learning = LearningService(store, browser) instance_id = os.environ.get("ASSISTANT_INSTANCE_ID") or uuid4().hex @asynccontextmanager async def lifespan(application): yield await browser.close() app = FastAPI(title="Assistant 学习维护后台", version="0.1.0", lifespan=lifespan, docs_url=None, redoc_url=None, openapi_url=None) app.state.store, app.state.auth, app.state.learning, app.state.browser = store, auth, learning, browser extra_origins = {value.strip().rstrip("/") for value in os.environ.get("ASSISTANT_ALLOWED_ORIGINS", "").split(",") if value.strip()} extra_origins.update({"http://127.0.0.1:5179", "http://localhost:5179"}) @app.middleware("http") async def local_boundary(request: Request, call_next): try: host = urlsplit("http://" + request.headers.get("host", "")).hostname except ValueError: host = None if host not in {"127.0.0.1", "localhost", "::1"}: return JSONResponse({"detail": "学习维护后台仅接受本机访问"}, status_code=400) origin = request.headers.get("origin") own_origin = str(request.base_url).rstrip("/") if origin and origin.rstrip("/") not in extra_origins | {own_origin}: return JSONResponse({"detail": "请求来源不在维护后台允许范围内"}, status_code=403) response = await call_next(request) if request.url.path.startswith("/api/"): response.headers["Cache-Control"] = "no-store" response.headers["X-Content-Type-Options"] = "nosniff" response.headers["Referrer-Policy"] = "same-origin" return response @app.exception_handler(RequestValidationError) async def validation_error(request, error): fields = ", ".join(".".join(str(part) for part in item["loc"] if part != "body") for item in error.errors()) return JSONResponse({"detail": "输入格式不正确，请检查：" + fields}, status_code=422) @app.exception_handler(StoreError) async def database_error(request, error: StoreError): return JSONResponse({"detail": {"code": error.code, "message": error.message}}, status_code=error.status) def operator(request: Request) -> str: return auth.operator(request, mutate=request.method not in {"GET", "HEAD", "OPTIONS"}) @app.get("/api/health") def health(): return {"status": "ok", "service": "assistant-learning", "version": "0.1.0", "mode": "local_maintainer", "model_connected": False, "instance_id": instance_id, "database_backend": store.backend} @app.get("/api/auth/status") def auth_status(request: Request, response: Response): return auth.status(request, response) @app.post("/api/auth/setup") def setup(body: Credentials, request: Request, response: Response): return auth.setup(request, response, body.username, body.password) @app.post("/api/auth/login") def login(body: Credentials, request: Request, response: Response): return auth.login(request, response, body.username, body.password) @app.post("/api/auth/logout") def logout(request: Request, response: Response): return auth.logout(request, response) @app.get("/api/learning/overview") def overview(owner: str = Depends(operator)): return learning.overview(owner) @app.get("/api/learning/tasks") async def tasks(owner: str = Depends(operator)): for task_id in list(learning.sessions): if store.one("SELECT id FROM learning_tasks WHERE id=$1 AND owner_id=$2", (task_id, owner)): await learning.refresh_browser(task_id, owner) items = learning.tasks(owner) return {"items": items, "total": len(items)} @app.post("/api/learning/tasks", status_code=201) def new_task(body: NewTask, owner: str = Depends(operator)): return learning.create_task(owner, body.model_dump()) @app.get("/api/learning/tasks/{task_id}") async def task_detail(task_id: str, owner: str = Depends(operator)): await learning.refresh_browser(task_id, owner) return learning.detail(task_id, owner) @app.post("/api/learning/tasks/{task_id}/browser/open") async def open_browser(task_id: str, owner: str = Depends(operator)): return await learning.open(task_id, owner) @app.post("/api/learning/tasks/{task_id}/browser/ready") async def ready_browser(task_id: str, body: Ready, owner: str = Depends(operator)): return await learning.ready(task_id, owner, body.identity_label) @app.post("/api/learning/tasks/{task_id}/browser/capture") async def capture_browser(task_id: str, body: Capture | None = None, owner: str = Depends(operator)): return await learning.capture(task_id, owner, body.note if body else "") @app.post("/api/learning/tasks/{task_id}/browser/action") async def action_browser(task_id: str, body: Action, owner: str = Depends(operator)): return await learning.action(task_id, owner, body.model_dump(exclude_none=True)) @app.post("/api/learning/tasks/{task_id}/browser/close") async def close_browser(task_id: str, owner: str = Depends(operator)): return await learning.close(task_id, owner) @app.post("/api/learning/tasks/{task_id}/candidates", status_code=201) def new_candidate(task_id: str, body: Candidate | None = None, owner: str = Depends(operator)): return learning.create_candidate(task_id, owner, body.name if body else None) @app.get("/api/learning/candidates") def candidates(owner: str = Depends(operator)): items = learning.candidates(owner) return {"items": items, "total": len(items)} @app.get("/api/learning/candidates/{candidate_id}") def candidate_detail(candidate_id: str, owner: str = Depends(operator)): return {**learning.candidate(candidate_id, owner), "validations": learning.validations(owner, candidate_id=candidate_id)} @app.post("/api/learning/candidates/{candidate_id}/check") def check_candidate(candidate_id: str, owner: str = Depends(operator)): return learning.check(candidate_id, owner) @app.post("/api/learning/candidates/{candidate_id}/verify") async def verify_candidate(candidate_id: str, body: Verify, owner: str = Depends(operator)): return await learning.verify(candidate_id, owner, body.model_dump()) @app.post("/api/learning/candidates/{candidate_id}/publish-check") def publish_check(candidate_id: str, owner: str = Depends(operator)): return learning.publish_check(candidate_id, owner) @app.get("/api/learning/validations") def validations(owner: str = Depends(operator)): items = learning.validations(owner) return {"items": items, "total": len(items)} @app.get("/api/learning/artifacts/{artifact_id}") def artifact(artifact_id: str, owner: str = Depends(operator)): return FileResponse(learning.artifact(artifact_id, owner), media_type="image/png") # Fixture pages intentionally contain synthetic public content, no operator or platform data. @app.get("/api/learning/demo-page", response_class=HTMLResponse) def demo_page(q: str = ""): term = q.strip()[:120] or "建材" safe = escape(term, quote=True) cards = "" for i in range(1, 4): detail_attr = 'data-testid="detail-link"' if i == 1 else "" cards += f'

{safe}采购需求 {i}

本机演示内容，用于验证关键词与读取对象是否对应。

查看详情 {i}
' return HTMLResponse(f'本机链路验证页面

本机搜索链路

这是本项目的演示页面，不连接小红书或抖音，也不产生真实客户线索。

{cards}
') @app.get("/api/learning/demo-detail/{item}", response_class=HTMLResponse) def demo_detail(item: int, q: str = "建材"): safe = escape(q[:120], quote=True) return HTMLResponse(f'演示详情与评论

{safe}采购需求 {item}

本机演示作者，不是真实平台账号。

评论

演示评论甲：正在了解{safe}，需要比较规格。
演示评论乙：{safe}适用场景有哪些？
返回搜索') @app.post("/api/learning/demo") async def run_demo(request: Request, owner: str = Depends(operator)): return await learning.demo(owner, str(request.base_url)) return app app = create_app()