feat: initial commit - pipeline runner with real-time streaming
This commit is contained in:
@@ -0,0 +1,226 @@
|
||||
import asyncio
|
||||
import json
|
||||
import uuid
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
from fastapi import FastAPI, WebSocket, WebSocketDisconnect, HTTPException
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
from fastapi.staticfiles import StaticFiles
|
||||
from pydantic import BaseModel
|
||||
|
||||
app = FastAPI(title="Pipeline Runner")
|
||||
|
||||
app.add_middleware(
|
||||
CORSMiddleware,
|
||||
allow_origins=["*"],
|
||||
allow_credentials=True,
|
||||
allow_methods=["*"],
|
||||
allow_headers=["*"],
|
||||
)
|
||||
|
||||
BASE_DIR = Path(__file__).parent
|
||||
PIPELINES_DIR = BASE_DIR / "storage" / "pipelines"
|
||||
LOGS_DIR = BASE_DIR / "storage" / "logs"
|
||||
|
||||
PIPELINES_DIR.mkdir(parents=True, exist_ok=True)
|
||||
LOGS_DIR.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
active_connections: dict[str, set[WebSocket]] = {}
|
||||
active_runs: dict[str, asyncio.Task] = {}
|
||||
|
||||
|
||||
class Prompt(BaseModel):
|
||||
id: str
|
||||
text: str
|
||||
title: str
|
||||
order: int
|
||||
|
||||
|
||||
class Pipeline(BaseModel):
|
||||
id: str
|
||||
name: str
|
||||
prompts: list[Prompt]
|
||||
created_at: str
|
||||
updated_at: str
|
||||
|
||||
|
||||
def load_pipelines() -> list[Pipeline]:
|
||||
pipelines = []
|
||||
for f in PIPELINES_DIR.glob("*.json"):
|
||||
try:
|
||||
with open(f) as file:
|
||||
data = json.load(file)
|
||||
pipelines.append(Pipeline(**data))
|
||||
except (json.JSONDecodeError, OSError):
|
||||
continue
|
||||
return sorted(pipelines, key=lambda x: x.updated_at, reverse=True)
|
||||
|
||||
|
||||
def save_pipeline(pipeline: Pipeline):
|
||||
pipeline.updated_at = datetime.now(timezone.utc).isoformat()
|
||||
try:
|
||||
with open(PIPELINES_DIR / f"{pipeline.id}.json", "w") as f:
|
||||
json.dump(pipeline.model_dump(), f, indent=2)
|
||||
except OSError as e:
|
||||
raise HTTPException(status_code=500, detail=f"Failed to save pipeline: {e}")
|
||||
|
||||
|
||||
def delete_pipeline(pipeline_id: str):
|
||||
path = PIPELINES_DIR / f"{pipeline_id}.json"
|
||||
if path.exists():
|
||||
try:
|
||||
path.unlink()
|
||||
except OSError as e:
|
||||
raise HTTPException(status_code=500, detail=f"Failed to delete pipeline: {e}")
|
||||
|
||||
|
||||
@app.get("/api/pipelines", response_model=list[Pipeline])
|
||||
async def list_pipelines():
|
||||
return load_pipelines()
|
||||
|
||||
|
||||
@app.post("/api/pipelines", response_model=Pipeline)
|
||||
async def create_pipeline(pipeline: Pipeline):
|
||||
existing = load_pipelines()
|
||||
if any(p.id == pipeline.id for p in existing):
|
||||
raise HTTPException(status_code=409, detail="Pipeline already exists")
|
||||
save_pipeline(pipeline)
|
||||
return pipeline
|
||||
|
||||
|
||||
@app.delete("/api/pipelines/{pipeline_id}")
|
||||
async def delete(pipeline_id: str):
|
||||
pipelines = load_pipelines()
|
||||
if not any(p.id == pipeline_id for p in pipelines):
|
||||
raise HTTPException(status_code=404, detail="Pipeline not found")
|
||||
delete_pipeline(pipeline_id)
|
||||
return {"deleted": pipeline_id}
|
||||
|
||||
|
||||
@app.post("/api/pipelines/{pipeline_id}/execute")
|
||||
async def execute_pipeline(pipeline_id: str):
|
||||
pipelines = load_pipelines()
|
||||
pipeline = next((p for p in pipelines if p.id == pipeline_id), None)
|
||||
if not pipeline:
|
||||
raise HTTPException(status_code=404, detail="Pipeline not found")
|
||||
|
||||
run_id = str(uuid.uuid4())
|
||||
log_file = LOGS_DIR / f"{run_id}.log"
|
||||
|
||||
active_connections[run_id] = set()
|
||||
task = asyncio.create_task(run_pipeline_task(pipeline, run_id, log_file))
|
||||
active_runs[run_id] = task
|
||||
|
||||
return {"run_id": run_id, "log_file": str(log_file)}
|
||||
|
||||
|
||||
async def run_pipeline_task(pipeline: Pipeline, run_id: str, log_file: Path):
|
||||
async def write_log(line: str):
|
||||
with open(log_file, "a") as f:
|
||||
f.write(f"[{datetime.now(timezone.utc).isoformat()}] {line}\n")
|
||||
for ws in list(active_connections.get(run_id, set())):
|
||||
try:
|
||||
await ws.send_json({"type": "log", "data": line})
|
||||
except:
|
||||
pass
|
||||
|
||||
async def stream_reader(stream, prefix, output_collector):
|
||||
while True:
|
||||
line = await stream.readline()
|
||||
if not line:
|
||||
break
|
||||
text = line.decode("utf-8", errors="replace").rstrip("\n")
|
||||
if text:
|
||||
output_collector.append(text)
|
||||
await write_log(f"{prefix} {text}")
|
||||
|
||||
await write_log(f"Starting pipeline: {pipeline.name} (run_id: {run_id})")
|
||||
|
||||
try:
|
||||
for i, prompt in enumerate(pipeline.prompts, 1):
|
||||
await write_log(f"[Step {i}/{len(pipeline.prompts)}] Starting: {prompt.title}")
|
||||
await write_log(f"[Step {i}/{len(pipeline.prompts)}] Prompt: {prompt.text}")
|
||||
|
||||
for ws in list(active_connections.get(run_id, set())):
|
||||
try:
|
||||
await ws.send_json({
|
||||
"type": "status",
|
||||
"step_index": i - 1,
|
||||
"status": "running",
|
||||
"title": prompt.title
|
||||
})
|
||||
except:
|
||||
pass
|
||||
|
||||
cmd = ["opencode", "run", prompt.text, "--title", prompt.title]
|
||||
process = await asyncio.create_subprocess_exec(
|
||||
*cmd,
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
stderr=asyncio.subprocess.PIPE
|
||||
)
|
||||
|
||||
stdout_lines: list[str] = []
|
||||
stderr_lines: list[str] = []
|
||||
|
||||
await asyncio.gather(
|
||||
stream_reader(process.stdout, f"[Step {i}/{len(pipeline.prompts)}]", stdout_lines),
|
||||
stream_reader(process.stderr, f"[Step {i}/{len(pipeline.prompts)}]", stderr_lines)
|
||||
)
|
||||
|
||||
await process.wait()
|
||||
|
||||
output = "\n".join(stdout_lines)
|
||||
error = "\n".join(stderr_lines)
|
||||
|
||||
status = "completed" if process.returncode == 0 else "error"
|
||||
result = {"returncode": process.returncode, "output": output, "error": error}
|
||||
|
||||
await write_log(f"[Step {i}/{len(pipeline.prompts)}] {status.upper()} (returncode: {process.returncode})")
|
||||
|
||||
for ws in list(active_connections.get(run_id, set())):
|
||||
try:
|
||||
await ws.send_json({
|
||||
"type": "status",
|
||||
"step_index": i - 1,
|
||||
"status": status,
|
||||
"title": prompt.title,
|
||||
"result": result
|
||||
})
|
||||
except:
|
||||
pass
|
||||
|
||||
if process.returncode != 0:
|
||||
await write_log(f"Pipeline failed at step {i}")
|
||||
break
|
||||
|
||||
await write_log("Pipeline finished")
|
||||
|
||||
finally:
|
||||
if run_id in active_runs:
|
||||
del active_runs[run_id]
|
||||
|
||||
|
||||
@app.websocket("/ws/{run_id}")
|
||||
async def websocket_endpoint(websocket: WebSocket, run_id: str):
|
||||
await websocket.accept()
|
||||
if run_id not in active_connections:
|
||||
active_connections[run_id] = set()
|
||||
active_connections[run_id].add(websocket)
|
||||
|
||||
try:
|
||||
while True:
|
||||
await websocket.receive_text()
|
||||
except Exception:
|
||||
active_connections[run_id].discard(websocket)
|
||||
if not active_connections[run_id]:
|
||||
del active_connections[run_id]
|
||||
|
||||
|
||||
if (BASE_DIR / "static").exists():
|
||||
app.mount("/static", StaticFiles(directory="static"), name="static")
|
||||
|
||||
@app.get("/")
|
||||
async def root():
|
||||
return {"message": "Pipeline Runner API", "ws_endpoint": "/ws/{run_id}"}
|
||||
Reference in New Issue
Block a user