107 lines
2.9 KiB
Python
107 lines
2.9 KiB
Python
# 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)}
|