ai_mouse/tools/server/routes_scroll.py

# 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:
    # Uses the bundled ONNX scroll model exposed via the public ai_mouse API.
    # The legacy scroll_model.pt path is no longer wired in; export a fresh
    # scroll_decoder.onnx via `python -m tools.export_onnx` to update.
    from ai_mouse import generate_scroll

    paths = []
    for _ in range(min(req.n_paths, 12)):
        events = generate_scroll(
            req.start_scrollY,
            req.target_scrollY,
            mode=req.mode,
        )
        paths.append(events)
    return {"paths": paths}