sensai ¶

@router.get("/config")
async def get_config() -> SensAISettings:
    """Returns the current SensAI settings."""
    return MujocoMojoSettings().dojo.sensai

@router.post("/config")
async def post_config(body: SensAISettings) -> SensAISettings:
    """Persists updated SensAI settings to disk."""
    settings = MujocoMojoSettings()
    updated = settings.model_copy(
        update={"dojo": settings.dojo.model_copy(update={"sensai": body})}
    )
    updated.save()
    return updated.dojo.sensai

@router.post("/chat")
async def post_chat(body: ChatRequest) -> EventSourceResponse:
    """Runs the SensAI agents and streams the response as SSE."""
    settings = MujocoMojoSettings()

    if not settings.dojo.sensai.enabled:
        raise HTTPException(
            status_code=status.HTTP_403_FORBIDDEN,
            detail="SensAI is not enabled. Set dojo.sensai.enabled = true in settings.",
        )

    current_plot_config: PlotConfig | None = None
    if body.current_plot_config_json:
        try:
            current_plot_config = PlotConfig.model_validate_json(
                body.current_plot_config_json
            )
        except Exception:
            logger.warning(
                "Failed to parse current_plot_config_json; proceeding without it."
            )
    column_manifest: ColumnManifest = {
        "all": body.all_columns,
        "rotatable_vectors": body.rotatable_vectors,
        "available_quats": body.available_quats,
        "column_metadata": body.column_metadata,
    }  # pyright: ignore[reportAssignmentType]

    deps = SensAIDeps(
        job_status=shared.CURRENT_JOB,
        column_manifest=column_manifest,
        current_plot_config=current_plot_config,
    )

    model = build_fallback_model()

    message_history = _to_model_messages(body.message_history)
    context_block = _build_context_block(deps)
    augmented_message = f"{context_block}\n\n{body.message}"

    async def event_stream() -> AsyncIterator[dict]:
        # undo path: restore previous plot config via client-side history
        if _is_undo_intent(body.message):
            yield {
                "event": "result",
                "data": json.dumps(
                    {
                        "message": "Undid the last plot change.",
                        "plot_config_update": None,
                        "routed_to": "undo",
                    }
                ),
            }
            return

        # tool_call events (fired the moment a delegation tool starts, well before
        # its result/the reply text) and text_delta events (from the plain text
        # loop below) both need to reach the client in the order they actually
        # happen - a shared queue lets _ToolResultCapture push the former in
        # concurrently with the latter, rather than only being able to report
        # tool calls after the fact once the whole run finishes.
        queue: asyncio.Queue[dict[str, str] | None] = asyncio.Queue()
        capture = _ToolResultCapture(queue)
        accumulated = ""
        stream_exc: Exception | None = None

        async def _produce_text() -> None:
            nonlocal accumulated, stream_exc
            try:
                async with chat_agent.run_stream(
                    augmented_message,
                    model=model,
                    deps=deps,
                    message_history=message_history,
                    event_stream_handler=capture,
                ) as streamed:
                    async for delta in streamed.stream_text(
                        delta=True, debounce_by=None
                    ):
                        accumulated += delta
                        await queue.put(
                            {
                                "event": "text_delta",
                                "data": json.dumps({"delta": delta}),
                            }
                        )
            except Exception as exc:
                stream_exc = exc
                logger.warning("SensAI chat run_stream failed", exc_info=True)
            finally:
                await queue.put(None)

        producer = asyncio.create_task(_produce_text())
        try:
            while True:
                item = await queue.get()
                if item is None:
                    break
                yield item
        finally:
            if not producer.done():
                producer.cancel()
            with contextlib.suppress(asyncio.CancelledError):
                await producer

        if not accumulated:
            detail = (
                f"No response from model: {stream_exc}"
                if stream_exc
                else "No response from model."
            )
            yield {"event": "error", "data": json.dumps({"detail": detail})}
            return

        yield {
            "event": "result",
            "data": json.dumps(
                {
                    "message": accumulated,
                    "plot_config_update": (
                        capture.plot_config_update.model_dump()
                        if capture.plot_config_update is not None
                        else None
                    ),
                    "routed_to": capture.routed_to,
                }
            ),
        }

    return EventSourceResponse(event_stream())