Skip to content

Commit

Permalink
one bug further lol
Browse files Browse the repository at this point in the history
  • Loading branch information
shroominic committed Aug 1, 2024
1 parent 718e41f commit bf2cc51
Show file tree
Hide file tree
Showing 4 changed files with 26 additions and 34 deletions.
2 changes: 1 addition & 1 deletion Dockerfile
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
FROM ghcr.io/astral-sh/uv as uv

FROM --platform=arm64 python:3.11 as build
FROM --platform=amd64 python:3.11 as build

ENV VIRTUAL_ENV=/.venv PATH="/.venv/bin:$PATH"

Expand Down
33 changes: 21 additions & 12 deletions src/codeboxapi/api.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,9 +2,10 @@
from contextlib import asynccontextmanager
from datetime import datetime, timedelta
from os import getenv
from tempfile import SpooledTemporaryFile
from typing import AsyncGenerator, Literal

from fastapi import Depends, FastAPI, HTTPException, UploadFile
from fastapi import Body, Depends, FastAPI, HTTPException, UploadFile
from fastapi.responses import StreamingResponse
from pydantic import BaseModel

Expand All @@ -18,8 +19,9 @@
@asynccontextmanager
async def lifespan(_: FastAPI) -> AsyncGenerator[None, None]:
async def timeout():
timeout_secs = float(getenv("CODEBOX_TIMEOUT", "900"))
while last_interaction + timedelta(seconds=timeout_secs) > datetime.utcnow():
if (_timeout := getenv("CODEBOX_TIMEOUT", "90")).lower() == "none":
return
while last_interaction + timedelta(seconds=float(_timeout)) > datetime.utcnow():
await asyncio.sleep(1)
exit(0)

Expand All @@ -28,18 +30,14 @@ async def timeout():
t.cancel()


app = FastAPI(title="Codebox API", lifespan=lifespan)


async def get_codebox() -> AsyncGenerator[LocalBox, None]:
global codebox, last_interaction
last_interaction = datetime.utcnow()
yield codebox


@app.get("/")
async def healthcheck() -> dict[str, str]:
return {"status": "ok"}
app = FastAPI(title="Codebox API", lifespan=lifespan)
app.get("/")(lambda: {"status": "ok"})


class ExecBody(BaseModel):
Expand All @@ -62,7 +60,7 @@ async def event_stream() -> AsyncGenerator[str, None]:
return StreamingResponse(event_stream())


@app.get("/download/{file_name}")
@app.get("/files/download/{file_name}")
async def download(
file_name: str,
timeout: int | None = None,
Expand All @@ -71,21 +69,32 @@ async def download(
return StreamingResponse(codebox.astream_download(file_name, timeout))


@app.post("/upload")
@app.post("/files/upload")
async def upload(
file: UploadFile,
timeout: int | None = None,
codebox: LocalBox = Depends(get_codebox),
) -> "CodeBoxFile":
if not file.filename:
raise HTTPException(status_code=400, detail="A file name is required")
if isinstance(file.file, SpooledTemporaryFile):
file.file = file.file
return await codebox.aupload(file.filename, file.file, timeout)


@app.post("/code/execute")
async def deprecated_exec(
body: dict = Body(), codebox: LocalBox = Depends(get_codebox)
) -> dict:
"""deprecated: use /exec instead"""
ex = await codebox.aexec(body["properties"]["code"])
return {"properties": {"stdout": ex.text, "stderr": ex.errors, "result": ex.text}}


def serve():
import uvicorn

uvicorn.run(app, host="0.0.0.0", port=getenv("CODEBOX_PORT", 8069))
uvicorn.run(app, host="0.0.0.0", port=8000)


if __name__ == "__main__":
Expand Down
2 changes: 1 addition & 1 deletion src/codeboxapi/docker.py
Original file line number Diff line number Diff line change
Expand Up @@ -43,7 +43,7 @@ def __init__(
"-e",
f"CODEBOX_TIMEOUT={timeout}",
"-p",
f"{self.port}:8069",
f"{self.port}:8000",
image,
],
check=True,
Expand Down
23 changes: 3 additions & 20 deletions src/codeboxapi/remote.py
Original file line number Diff line number Diff line change
Expand Up @@ -100,21 +100,6 @@ async def astream_exec(
async for c in self.astream_exec(code, kernel, timeout, cwd):
yield c

def upload(
self,
remote_file_path: str,
content: BinaryIO | bytes | str,
timeout: float | None = None,
) -> CodeBoxFile:
if isinstance(content, str):
content = content.encode("utf-8")
response = self.client.post(
url="/upload",
files={"file": (remote_file_path, content)},
timeout=timeout,
)
return CodeBoxFile(**response.json())

async def aupload(
self,
file_name: str,
Expand All @@ -124,7 +109,7 @@ async def aupload(
if isinstance(content, str):
content = content.encode("utf-8")
response = await self.aclient.post(
url="/upload",
url="/files/upload",
files={"file": (file_name, content)},
timeout=timeout,
)
Expand All @@ -137,9 +122,8 @@ def stream_download(
) -> Generator[bytes, None, None]:
with self.client.stream(
method="GET",
url="/download",
url=f"/files/download/{remote_file_path}",
timeout=timeout,
params={"file_name": remote_file_path},
) as response:
for chunk in response.iter_bytes():
yield chunk
Expand All @@ -151,9 +135,8 @@ async def astream_download(
) -> AsyncGenerator[bytes, None]:
async with self.aclient.stream(
method="GET",
url="/download",
url=f"/files/download/{remote_file_path}",
timeout=timeout,
params={"file_name": remote_file_path},
) as response:
async for chunk in response.aiter_bytes():
yield chunk

0 comments on commit bf2cc51

Please sign in to comment.