refactor: move server/ to tools/server/
This commit is contained in:
51
tools/server/__init__.py
Normal file
51
tools/server/__init__.py
Normal file
@@ -0,0 +1,51 @@
|
||||
# ai_mouse/server/__init__.py
|
||||
"""AI Mouse server package — FastAPI app factory."""
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from pathlib import Path
|
||||
|
||||
from fastapi import FastAPI
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
from fastapi.responses import FileResponse
|
||||
from fastapi.staticfiles import StaticFiles
|
||||
|
||||
from .routes_collect import router as collect_router
|
||||
from .routes_scroll import router as scroll_router
|
||||
from .routes_train import router as train_router
|
||||
from .routes_verify import router as verify_router
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_HERE = Path(__file__).resolve().parent
|
||||
_STATIC_DIR = _HERE.parent.parent / "static"
|
||||
|
||||
|
||||
def create_app() -> FastAPI:
|
||||
"""Create and configure the FastAPI application."""
|
||||
app = FastAPI(title="AI Mouse Trajectory Generator")
|
||||
|
||||
# CORS — allow all origins (local development tool)
|
||||
app.add_middleware(
|
||||
CORSMiddleware,
|
||||
allow_origins=["*"],
|
||||
allow_credentials=True,
|
||||
allow_methods=["*"],
|
||||
allow_headers=["*"],
|
||||
)
|
||||
|
||||
# Routers
|
||||
app.include_router(collect_router, prefix="/api")
|
||||
app.include_router(train_router, prefix="/api")
|
||||
app.include_router(verify_router, prefix="/api")
|
||||
app.include_router(scroll_router, prefix="/api/scroll")
|
||||
|
||||
# Static files
|
||||
app.mount("/static", StaticFiles(directory=str(_STATIC_DIR)), name="static")
|
||||
|
||||
# Serve index.html at root
|
||||
@app.get("/")
|
||||
def index() -> FileResponse:
|
||||
return FileResponse(str(_STATIC_DIR / "index.html"))
|
||||
|
||||
return app
|
||||
49
tools/server/deps.py
Normal file
49
tools/server/deps.py
Normal file
@@ -0,0 +1,49 @@
|
||||
# ai_mouse/server/deps.py
|
||||
"""Shared dependencies for the ai-mouse server package."""
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Data directory
|
||||
# ---------------------------------------------------------------------------
|
||||
_HERE = Path(__file__).resolve().parent
|
||||
_DATA_DIR = _HERE.parent.parent / "data"
|
||||
|
||||
|
||||
def get_data_dir() -> Path:
|
||||
"""Return the resolved data directory (sites/ai_mouse/data/)."""
|
||||
return _DATA_DIR
|
||||
|
||||
|
||||
def validate_path(path: Path, base: Path) -> Path:
|
||||
"""Resolve *path* and ensure it lives under *base*. Raises ValueError on traversal."""
|
||||
resolved = path.resolve()
|
||||
base_resolved = base.resolve()
|
||||
if not str(resolved).startswith(str(base_resolved)):
|
||||
raise ValueError(f"Path traversal detected: {path}")
|
||||
return resolved
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Session state
|
||||
# ---------------------------------------------------------------------------
|
||||
@dataclass
|
||||
class SessionState:
|
||||
"""Mutable singleton holding collectors initialised at runtime."""
|
||||
|
||||
collector: Optional[object] = field(default=None)
|
||||
scroll_collector: Optional[object] = field(default=None)
|
||||
|
||||
|
||||
_state = SessionState()
|
||||
|
||||
|
||||
def get_state() -> SessionState:
|
||||
"""Return the module-level session state singleton."""
|
||||
return _state
|
||||
106
tools/server/routes_collect.py
Normal file
106
tools/server/routes_collect.py
Normal file
@@ -0,0 +1,106 @@
|
||||
# ai_mouse/server/routes_collect.py
|
||||
"""Collection routes: start, trace, skip."""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException
|
||||
from pydantic import BaseModel
|
||||
|
||||
from tools.collector import Collector
|
||||
|
||||
from .deps import SessionState, get_data_dir, get_state
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Request models
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class CollectStartRequest(BaseModel):
|
||||
count: int = 100
|
||||
dist_min: int = 50
|
||||
dist_max: int = 800
|
||||
|
||||
|
||||
class TraceRequest(BaseModel):
|
||||
meta: dict
|
||||
events: list[dict]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Routes
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@router.post("/collect/start")
|
||||
def collect_start(
|
||||
req: CollectStartRequest,
|
||||
state: SessionState = Depends(get_state),
|
||||
) -> dict:
|
||||
traces_path = get_data_dir() / "traces.jsonl"
|
||||
collector = Collector(
|
||||
count=req.count,
|
||||
dist_min=req.dist_min,
|
||||
dist_max=req.dist_max,
|
||||
output_path=traces_path,
|
||||
)
|
||||
state.collector = collector
|
||||
return {"a": list(collector.a_pos), "b": list(collector.b_pos)}
|
||||
|
||||
|
||||
@router.post("/collect/trace")
|
||||
def collect_trace(
|
||||
trace: TraceRequest,
|
||||
state: SessionState = Depends(get_state),
|
||||
) -> dict:
|
||||
if state.collector is None:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail="Collector not started. Call /api/collect/start first.",
|
||||
)
|
||||
|
||||
traces_path = get_data_dir() / "traces.jsonl"
|
||||
traces_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
record = {"meta": trace.meta, "events": trace.events}
|
||||
with traces_path.open("a", encoding="utf-8") as f:
|
||||
f.write(json.dumps(record, ensure_ascii=False) + "\n")
|
||||
|
||||
collector = state.collector
|
||||
collector.collected += 1
|
||||
remaining = collector.count - collector.collected
|
||||
|
||||
if remaining > 0:
|
||||
collector.a_pos, collector.b_pos = collector._new_ab()
|
||||
return {
|
||||
"collected": collector.collected,
|
||||
"remaining": remaining,
|
||||
"a": list(collector.a_pos),
|
||||
"b": list(collector.b_pos),
|
||||
}
|
||||
else:
|
||||
return {
|
||||
"collected": collector.collected,
|
||||
"remaining": 0,
|
||||
"a": None,
|
||||
"b": None,
|
||||
}
|
||||
|
||||
|
||||
@router.post("/collect/skip")
|
||||
def collect_skip(
|
||||
state: SessionState = Depends(get_state),
|
||||
) -> dict:
|
||||
if state.collector is None:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail="Collector not started. Call /api/collect/start first.",
|
||||
)
|
||||
|
||||
collector = state.collector
|
||||
collector.a_pos, collector.b_pos = collector._new_ab()
|
||||
return {"a": list(collector.a_pos), "b": list(collector.b_pos)}
|
||||
202
tools/server/routes_scroll.py
Normal file
202
tools/server/routes_scroll.py
Normal file
@@ -0,0 +1,202 @@
|
||||
# ai_mouse/server/routes_scroll.py
|
||||
"""Scroll collection, training, and verification routes."""
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
from pathlib import Path
|
||||
from typing import AsyncGenerator
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException
|
||||
from fastapi.responses import StreamingResponse
|
||||
from pydantic import BaseModel
|
||||
|
||||
from .deps import SessionState, get_data_dir, get_state
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _paths() -> tuple[Path, Path]:
|
||||
data_dir = get_data_dir()
|
||||
return data_dir / "scroll_traces.jsonl", data_dir / "scroll_models"
|
||||
|
||||
|
||||
def _scroll_trace_count() -> int:
|
||||
traces_path, _ = _paths()
|
||||
if not traces_path.exists():
|
||||
return 0
|
||||
return sum(
|
||||
1
|
||||
for line in traces_path.read_text(encoding="utf-8").splitlines()
|
||||
if line.strip()
|
||||
)
|
||||
|
||||
|
||||
def _scroll_model_trained() -> bool:
|
||||
_, models_dir = _paths()
|
||||
return (models_dir / "scroll_model.pt").exists()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Request models
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class ScrollStartRequest(BaseModel):
|
||||
mode: str = "target"
|
||||
count: int = 50
|
||||
viewport_height: int = 900
|
||||
|
||||
|
||||
class ScrollTraceRequest(BaseModel):
|
||||
meta: dict
|
||||
events: list[dict]
|
||||
|
||||
|
||||
class ScrollSkipRequest(BaseModel):
|
||||
current_scrollY: int = 0
|
||||
|
||||
|
||||
class ScrollTrainRequest(BaseModel):
|
||||
epochs: int = 100
|
||||
|
||||
|
||||
class ScrollVerifyRequest(BaseModel):
|
||||
start_scrollY: int = 1000
|
||||
target_scrollY: int = 3000
|
||||
mode: str = "target"
|
||||
n_paths: int = 5
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Routes
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@router.post("/start")
|
||||
def scroll_start(
|
||||
req: ScrollStartRequest,
|
||||
state: SessionState = Depends(get_state),
|
||||
) -> dict:
|
||||
from tools.scroll.collector import ScrollCollector
|
||||
|
||||
scroll_collector = ScrollCollector(
|
||||
mode=req.mode, count=req.count, viewport_height=req.viewport_height
|
||||
)
|
||||
state.scroll_collector = scroll_collector
|
||||
target = scroll_collector.next_target(current_scrollY=0)
|
||||
return {
|
||||
"success_radius": scroll_collector.success_radius,
|
||||
**target,
|
||||
}
|
||||
|
||||
|
||||
@router.post("/trace")
|
||||
def scroll_trace(
|
||||
trace: ScrollTraceRequest,
|
||||
state: SessionState = Depends(get_state),
|
||||
) -> dict:
|
||||
if state.scroll_collector is None:
|
||||
raise HTTPException(400, "Scroll collector not started")
|
||||
|
||||
traces_path, _ = _paths()
|
||||
traces_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
record = {"meta": trace.meta, "events": trace.events}
|
||||
with traces_path.open("a", encoding="utf-8") as f:
|
||||
f.write(json.dumps(record, ensure_ascii=False) + "\n")
|
||||
|
||||
scroll_collector = state.scroll_collector
|
||||
scroll_collector.collected += 1
|
||||
remaining = scroll_collector.count - scroll_collector.collected
|
||||
|
||||
if remaining > 0:
|
||||
target = scroll_collector.next_target(trace.meta.get("end_scrollY", 0))
|
||||
return {"collected": scroll_collector.collected, "remaining": remaining, **target}
|
||||
return {"collected": scroll_collector.collected, "remaining": 0, "target_scrollY": None}
|
||||
|
||||
|
||||
@router.post("/skip")
|
||||
def scroll_skip(
|
||||
req: ScrollSkipRequest,
|
||||
state: SessionState = Depends(get_state),
|
||||
) -> dict:
|
||||
if state.scroll_collector is None:
|
||||
raise HTTPException(400, "Scroll collector not started")
|
||||
target = state.scroll_collector.next_target(current_scrollY=req.current_scrollY)
|
||||
return target
|
||||
|
||||
|
||||
@router.get("/status")
|
||||
def scroll_status() -> dict:
|
||||
return {"trace_count": _scroll_trace_count(), "model_trained": _scroll_model_trained()}
|
||||
|
||||
|
||||
async def _scroll_train_sse(req: ScrollTrainRequest) -> AsyncGenerator[str, None]:
|
||||
"""Run scroll training in a thread, yield SSE events via asyncio.Queue."""
|
||||
queue: asyncio.Queue[dict] = asyncio.Queue()
|
||||
|
||||
def callback(msg: dict) -> None:
|
||||
queue.put_nowait(msg)
|
||||
|
||||
async def run() -> None:
|
||||
from tools.scroll.trainer import train_scroll
|
||||
|
||||
traces_path, models_dir = _paths()
|
||||
try:
|
||||
await asyncio.to_thread(
|
||||
train_scroll,
|
||||
data_path=traces_path,
|
||||
output_dir=models_dir,
|
||||
epochs=req.epochs,
|
||||
progress_callback=callback,
|
||||
)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
queue.put_nowait({"error": str(exc)})
|
||||
|
||||
task = asyncio.create_task(run())
|
||||
|
||||
while True:
|
||||
msg = await queue.get()
|
||||
yield f"data: {json.dumps(msg)}\n\n"
|
||||
if msg.get("done") or msg.get("error"):
|
||||
break
|
||||
|
||||
await task
|
||||
|
||||
|
||||
@router.post("/train")
|
||||
async def scroll_train(req: ScrollTrainRequest) -> StreamingResponse:
|
||||
return StreamingResponse(
|
||||
_scroll_train_sse(req),
|
||||
media_type="text/event-stream",
|
||||
headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
|
||||
)
|
||||
|
||||
|
||||
@router.post("/verify")
|
||||
def scroll_verify(req: ScrollVerifyRequest) -> dict:
|
||||
from ai_mouse.scroll.generator import generate_scroll
|
||||
|
||||
_, models_dir = _paths()
|
||||
if not (models_dir / "scroll_model.pt").exists():
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail="滚轮模型尚未训练，请先在「训练模型 → 滚轮模型」中完成训练。",
|
||||
)
|
||||
paths = []
|
||||
for _ in range(min(req.n_paths, 12)):
|
||||
events = generate_scroll(
|
||||
req.start_scrollY,
|
||||
req.target_scrollY,
|
||||
mode=req.mode,
|
||||
model_dir=str(models_dir),
|
||||
)
|
||||
paths.append(events)
|
||||
return {"paths": paths}
|
||||
131
tools/server/routes_train.py
Normal file
131
tools/server/routes_train.py
Normal file
@@ -0,0 +1,131 @@
|
||||
# ai_mouse/server/routes_train.py
|
||||
"""Training and status routes."""
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
from pathlib import Path
|
||||
from typing import AsyncGenerator
|
||||
|
||||
from fastapi import APIRouter
|
||||
from fastapi.responses import StreamingResponse
|
||||
from pydantic import BaseModel
|
||||
|
||||
from .deps import get_data_dir
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _paths() -> tuple[Path, Path, Path]:
|
||||
data_dir = get_data_dir()
|
||||
return (
|
||||
data_dir / "traces.jsonl",
|
||||
data_dir / "models_v2",
|
||||
data_dir / "models_v2_pretrained",
|
||||
)
|
||||
|
||||
|
||||
def _trace_count() -> int:
|
||||
traces_path, _, _ = _paths()
|
||||
if not traces_path.exists():
|
||||
return 0
|
||||
return sum(
|
||||
1
|
||||
for line in traces_path.read_text(encoding="utf-8").splitlines()
|
||||
if line.strip()
|
||||
)
|
||||
|
||||
|
||||
def _model_trained() -> bool:
|
||||
_, models_dir, _ = _paths()
|
||||
return (models_dir / "flow_model.pt").exists()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Request models
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TrainRequest(BaseModel):
|
||||
epochs: int = 200
|
||||
data_path: str | None = None
|
||||
output_dir: str | None = None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Routes
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@router.get("/status")
|
||||
def get_status() -> dict:
|
||||
return {"trace_count": _trace_count(), "model_trained": _model_trained()}
|
||||
|
||||
|
||||
async def _train_sse_generator(req: TrainRequest) -> AsyncGenerator[str, None]:
|
||||
"""Run training in a thread via asyncio.to_thread, yield SSE events via asyncio.Queue."""
|
||||
queue: asyncio.Queue[dict] = asyncio.Queue()
|
||||
|
||||
def callback(msg: dict) -> None:
|
||||
queue.put_nowait(msg)
|
||||
|
||||
async def run_training_async() -> None:
|
||||
from tools.trainer import train
|
||||
|
||||
traces_path, models_dir, pretrained_dir = _paths()
|
||||
data_path = Path(req.data_path) if req.data_path else traces_path
|
||||
output_dir = Path(req.output_dir) if req.output_dir else models_dir
|
||||
|
||||
# Auto-detect pretrained checkpoint and switch to fine-tune mode
|
||||
resume_from: Path | None = None
|
||||
effective_lr = 3e-4
|
||||
if (pretrained_dir / "flow_model.pt").exists():
|
||||
resume_from = pretrained_dir
|
||||
effective_lr = 1e-5 # fine-tune lr
|
||||
logger.info("Detected pretrained checkpoint, fine-tuning at lr=%g", effective_lr)
|
||||
queue.put_nowait({
|
||||
"info": f"Detected pretrained checkpoint at {pretrained_dir.name}, "
|
||||
f"running fine-tune at lr={effective_lr}",
|
||||
})
|
||||
|
||||
try:
|
||||
await asyncio.to_thread(
|
||||
train,
|
||||
data_path=data_path,
|
||||
output_dir=output_dir,
|
||||
epochs=req.epochs,
|
||||
lr=effective_lr,
|
||||
progress_callback=callback,
|
||||
resume_from=resume_from,
|
||||
)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
queue.put_nowait({"error": str(exc)})
|
||||
|
||||
task = asyncio.create_task(run_training_async())
|
||||
|
||||
while True:
|
||||
msg = await queue.get()
|
||||
yield f"data: {json.dumps(msg)}\n\n"
|
||||
if msg.get("done") or msg.get("error"):
|
||||
break
|
||||
|
||||
await task
|
||||
|
||||
|
||||
@router.post("/train")
|
||||
async def train_model(req: TrainRequest) -> StreamingResponse:
|
||||
return StreamingResponse(
|
||||
_train_sse_generator(req),
|
||||
media_type="text/event-stream",
|
||||
headers={
|
||||
"Cache-Control": "no-cache",
|
||||
"X-Accel-Buffering": "no",
|
||||
},
|
||||
)
|
||||
51
tools/server/routes_verify.py
Normal file
51
tools/server/routes_verify.py
Normal file
@@ -0,0 +1,51 @@
|
||||
# ai_mouse/server/routes_verify.py
|
||||
"""Verification route: generate trajectories."""
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
|
||||
from fastapi import APIRouter, HTTPException
|
||||
from pydantic import BaseModel
|
||||
|
||||
from .deps import get_data_dir
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Request models
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class VerifyRequest(BaseModel):
|
||||
start: list[int]
|
||||
end: list[int]
|
||||
n_paths: int = 5
|
||||
model_dir: str | None = None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Routes
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@router.post("/verify")
|
||||
def verify(req: VerifyRequest) -> dict:
|
||||
from ai_mouse.generator import generate
|
||||
|
||||
n = max(1, min(req.n_paths, 12))
|
||||
models_dir = get_data_dir() / "models_v2"
|
||||
model_dir_arg = req.model_dir if req.model_dir else str(models_dir)
|
||||
start = tuple(req.start) # type: ignore[arg-type]
|
||||
end = tuple(req.end) # type: ignore[arg-type]
|
||||
|
||||
paths = []
|
||||
try:
|
||||
for _ in range(n):
|
||||
pts = generate(start=start, end=end, model_dir=model_dir_arg)
|
||||
paths.append([[x, y, t] for x, y, t in pts])
|
||||
except FileNotFoundError as exc:
|
||||
raise HTTPException(status_code=404, detail=str(exc)) from exc
|
||||
|
||||
return {"paths": paths}
|
||||
Reference in New Issue
Block a user