97 lines
3.0 KiB
Python
97 lines
3.0 KiB
Python
import base64
|
|
import json
|
|
from collections.abc import AsyncIterator
|
|
from typing import Annotated
|
|
|
|
from fastapi import Body, FastAPI, HTTPException, Path, Response
|
|
from fastapi.responses import StreamingResponse
|
|
from pydantic import BaseModel
|
|
|
|
from ipython_shell.jupyter import InputRequest, JupyterApp, ParsedJupyterMessage
|
|
|
|
app = FastAPI(title="IPython Jupyter Web App", version="0.1.0")
|
|
shell_app = JupyterApp()
|
|
|
|
|
|
class InputReply(BaseModel):
|
|
"""Value supplied by a client answering a kernel input request."""
|
|
|
|
value: str
|
|
|
|
|
|
def serialize_event(event: ParsedJupyterMessage | InputRequest) -> dict[str, object]:
|
|
"""Turn a shell event into JSON-safe data for a web client."""
|
|
if isinstance(event, InputRequest):
|
|
return {
|
|
"type": "input_request",
|
|
"call_id": event.call_id,
|
|
"prompt": event.prompt,
|
|
"password": event.password,
|
|
}
|
|
|
|
# Some Jupyter messages keep binary payloads in message buffers. Base64
|
|
# makes those buffers safe to carry beside JSON MIME data in the stream.
|
|
content = event.content
|
|
if isinstance(content, BaseModel):
|
|
content = content.model_dump(mode="json")
|
|
|
|
return {
|
|
"type": "kernel_message",
|
|
"msg_type": event.msg_type,
|
|
"parent_id": event.parent_id,
|
|
"content": content,
|
|
"buffers": [
|
|
base64.b64encode(buffer).decode("ascii") for buffer in event.buffers
|
|
],
|
|
}
|
|
|
|
|
|
async def _stream_events(code: str, shell_name: str) -> AsyncIterator[str]:
|
|
"""Encode one execution's events as newline-delimited JSON."""
|
|
async for event in shell_app.run_code_stream(code, shell_name):
|
|
yield json.dumps(serialize_event(event), ensure_ascii=False) + "\n"
|
|
|
|
|
|
@app.post("/shells/{shell_name}/run")
|
|
async def run_code(
|
|
shell_name: Annotated[
|
|
str,
|
|
Path(
|
|
...,
|
|
description="The Jupyter shell to run code in. Existing shell guess. 'new' creates a new shell, 'last' uses the most recent shell.",
|
|
),
|
|
],
|
|
code: Annotated[
|
|
str,
|
|
Body(..., description="The code to execute", media_type="text/plain"),
|
|
],
|
|
) -> StreamingResponse:
|
|
"""Run code and stream each Jupyter event as one JSON line."""
|
|
return StreamingResponse(
|
|
_stream_events(code, shell_name),
|
|
media_type="application/x-ndjson",
|
|
)
|
|
|
|
|
|
@app.post("/calls/{call_id}/input", status_code=204)
|
|
async def reply_to_input(call_id: str, reply: InputReply) -> Response:
|
|
"""Answer an input request belonging to a currently running call."""
|
|
try:
|
|
await shell_app.reply_to_input(call_id, reply.value)
|
|
except KeyError as error:
|
|
raise HTTPException(status_code=404, detail=str(error)) from error
|
|
except RuntimeError as error:
|
|
raise HTTPException(status_code=409, detail=str(error)) from error
|
|
return Response(status_code=204)
|
|
|
|
|
|
def main() -> None:
|
|
"""Run the Jupyter-backed HTTP API with Uvicorn."""
|
|
import uvicorn
|
|
|
|
uvicorn.run(app, host="::", port=8001, log_level="info")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|