feat: FastAPI app with auth, AG-UI endpoint, session proxy, artifact API
This commit is contained in:
Ken
+374
@@ -0,0 +1,374 @@
|
||||
"""FastAPI application for research workbench backend.
|
||||
|
||||
Provides auth, AG-UI (CopilotKit) endpoint, session proxy to amplifierd,
|
||||
artifact API, MCP App serving, and static frontend.
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
import uuid
|
||||
from contextlib import asynccontextmanager
|
||||
from pathlib import Path
|
||||
from typing import Any, AsyncGenerator
|
||||
|
||||
import httpx
|
||||
from fastapi import Cookie, Depends, FastAPI, HTTPException, Request, Response
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
from fastapi.responses import HTMLResponse, JSONResponse, StreamingResponse
|
||||
from fastapi.staticfiles import StaticFiles
|
||||
|
||||
from agui_adapter import AmplifierdEventTranslator
|
||||
from artifacts import get_artifact, list_artifacts
|
||||
from auth import (
|
||||
create_session_token,
|
||||
invalidate_session_token,
|
||||
validate_session_token,
|
||||
verify_password,
|
||||
)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Module constants
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
AMPLIFIERD_URL: str = os.environ.get("AMPLIFIERD_URL", "http://localhost:8410")
|
||||
AUTH_USER: str = os.environ.get("AUTH_USER", "admin@localhost")
|
||||
BUNDLE_NAME: str = os.environ.get("BUNDLE_NAME", "research-workbench")
|
||||
APPS_DIR: Path = Path(__file__).parent.parent / "bundle" / "apps"
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Lifespan — shared HTTP client
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
|
||||
"""Create and teardown the shared httpx.AsyncClient for amplifierd."""
|
||||
app.state.http_client = httpx.AsyncClient(
|
||||
base_url=AMPLIFIERD_URL,
|
||||
timeout=30.0,
|
||||
)
|
||||
yield
|
||||
await app.state.http_client.aclose()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Application
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
app = FastAPI(lifespan=lifespan)
|
||||
|
||||
app.add_middleware(
|
||||
CORSMiddleware,
|
||||
allow_origins=["*"],
|
||||
allow_methods=["*"],
|
||||
allow_headers=["*"],
|
||||
allow_credentials=True,
|
||||
)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Auth helper
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _require_auth(session: str | None = Cookie(None)) -> str:
|
||||
"""Validate the session cookie and return the email, or raise HTTP 401."""
|
||||
email = validate_session_token(session or "")
|
||||
if not email:
|
||||
raise HTTPException(status_code=401, detail="Not authenticated")
|
||||
return email
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# SSE helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _sse(event: Any) -> str:
|
||||
"""Format an AG-UI event object as an SSE data line."""
|
||||
return f"data: {event.model_dump_json()}\n\n"
|
||||
|
||||
|
||||
def _parse_sse(line: str, state: dict[str, Any]) -> dict[str, Any] | None:
|
||||
"""Parse one SSE line, accumulate state, return event dict when data arrives.
|
||||
|
||||
Args:
|
||||
line: A single line from an SSE stream.
|
||||
state: Mutable dict holding the current event type (reused across calls).
|
||||
|
||||
Returns:
|
||||
A dict ``{event_type, data}`` when a data line is parsed, else ``None``.
|
||||
"""
|
||||
if line.startswith("event:"):
|
||||
state["event_type"] = line[len("event:") :].strip()
|
||||
return None
|
||||
if line.startswith("data:"):
|
||||
raw = line[len("data:") :].strip()
|
||||
try:
|
||||
data = json.loads(raw)
|
||||
except (json.JSONDecodeError, ValueError):
|
||||
data = raw
|
||||
return {"event_type": state.get("event_type", ""), "data": data}
|
||||
return None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Health
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@app.get("/api/health")
|
||||
async def health() -> dict[str, str]:
|
||||
"""Return service health status."""
|
||||
return {"status": "ok"}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Auth endpoints
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@app.post("/api/auth/login")
|
||||
async def login(request: Request) -> JSONResponse:
|
||||
"""Validate credentials, create session token, set session cookie."""
|
||||
body = await request.json()
|
||||
email: str = body.get("email", "")
|
||||
password: str = body.get("password", "")
|
||||
|
||||
if email != AUTH_USER or not verify_password(password):
|
||||
raise HTTPException(status_code=401, detail="Invalid credentials")
|
||||
|
||||
token = create_session_token(email)
|
||||
|
||||
response = JSONResponse({"token": token, "email": email})
|
||||
response.set_cookie(
|
||||
key="session",
|
||||
value=token,
|
||||
httponly=True,
|
||||
samesite="lax",
|
||||
max_age=86400 * 7,
|
||||
)
|
||||
return response
|
||||
|
||||
|
||||
@app.get("/api/auth/me")
|
||||
async def me(email: str = Depends(_require_auth)) -> dict[str, str]:
|
||||
"""Return the authenticated user's email."""
|
||||
return {"email": email}
|
||||
|
||||
|
||||
@app.post("/api/auth/logout")
|
||||
async def logout(
|
||||
response: Response,
|
||||
session: str | None = Cookie(None),
|
||||
) -> dict[str, str]:
|
||||
"""Invalidate the session token and clear the session cookie."""
|
||||
if session:
|
||||
invalidate_session_token(session)
|
||||
response.delete_cookie("session")
|
||||
return {"status": "ok"}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Session proxy (forward to amplifierd)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@app.get("/api/sessions")
|
||||
async def get_sessions(
|
||||
request: Request,
|
||||
email: str = Depends(_require_auth),
|
||||
) -> JSONResponse:
|
||||
"""Proxy GET /sessions to amplifierd."""
|
||||
resp = await request.app.state.http_client.get("/sessions")
|
||||
return JSONResponse(content=resp.json(), status_code=resp.status_code)
|
||||
|
||||
|
||||
@app.post("/api/sessions")
|
||||
async def create_session(
|
||||
request: Request,
|
||||
email: str = Depends(_require_auth),
|
||||
) -> JSONResponse:
|
||||
"""Proxy POST /sessions to amplifierd, injecting bundle_name if absent."""
|
||||
body: dict[str, Any] = await request.json()
|
||||
body.setdefault("bundle_name", BUNDLE_NAME)
|
||||
resp = await request.app.state.http_client.post("/sessions", json=body)
|
||||
return JSONResponse(content=resp.json(), status_code=resp.status_code)
|
||||
|
||||
|
||||
@app.get("/api/sessions/{session_id}")
|
||||
async def get_session(
|
||||
session_id: str,
|
||||
request: Request,
|
||||
email: str = Depends(_require_auth),
|
||||
) -> JSONResponse:
|
||||
"""Proxy GET /sessions/{id} to amplifierd."""
|
||||
resp = await request.app.state.http_client.get(f"/sessions/{session_id}")
|
||||
return JSONResponse(content=resp.json(), status_code=resp.status_code)
|
||||
|
||||
|
||||
@app.delete("/api/sessions/{session_id}")
|
||||
async def delete_session(
|
||||
session_id: str,
|
||||
request: Request,
|
||||
email: str = Depends(_require_auth),
|
||||
) -> JSONResponse:
|
||||
"""Proxy DELETE /sessions/{id} to amplifierd."""
|
||||
resp = await request.app.state.http_client.delete(f"/sessions/{session_id}")
|
||||
return JSONResponse(content=resp.json(), status_code=resp.status_code)
|
||||
|
||||
|
||||
@app.get("/api/sessions/{session_id}/transcript")
|
||||
async def get_transcript(
|
||||
session_id: str,
|
||||
request: Request,
|
||||
email: str = Depends(_require_auth),
|
||||
) -> JSONResponse:
|
||||
"""Proxy GET /sessions/{id}/transcript to amplifierd."""
|
||||
resp = await request.app.state.http_client.get(f"/sessions/{session_id}/transcript")
|
||||
return JSONResponse(content=resp.json(), status_code=resp.status_code)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# AG-UI / CopilotKit streaming endpoint
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@app.post("/api/copilotkit")
|
||||
async def copilotkit(request: Request) -> StreamingResponse:
|
||||
"""AG-UI streaming endpoint.
|
||||
|
||||
Accepts CopilotKit protocol payload, executes via amplifierd, and streams
|
||||
AG-UI SSE events back to the client.
|
||||
"""
|
||||
from ag_ui.core import RunStartedEvent
|
||||
|
||||
body = await request.json()
|
||||
thread_id: str = body.get("threadId", str(uuid.uuid4()))
|
||||
run_id: str = body.get("runId", str(uuid.uuid4()))
|
||||
messages: list[dict[str, Any]] = body.get("messages", [])
|
||||
|
||||
# Extract last user message as prompt
|
||||
prompt: str | None = None
|
||||
for msg in reversed(messages):
|
||||
if msg.get("role") == "user":
|
||||
content = msg.get("content", "")
|
||||
if isinstance(content, str):
|
||||
prompt = content
|
||||
elif isinstance(content, list):
|
||||
for block in content:
|
||||
if isinstance(block, dict) and block.get("type") == "text":
|
||||
prompt = block.get("text", "")
|
||||
break
|
||||
break
|
||||
|
||||
if not prompt:
|
||||
raise HTTPException(status_code=400, detail="Missing user message")
|
||||
|
||||
session_id = thread_id
|
||||
client = request.app.state.http_client
|
||||
|
||||
# Create session if it doesn't exist
|
||||
check_resp = await client.get(f"/sessions/{session_id}")
|
||||
if check_resp.status_code == 404:
|
||||
await client.post(
|
||||
"/sessions",
|
||||
json={"session_id": session_id, "bundle_name": BUNDLE_NAME},
|
||||
)
|
||||
|
||||
# Start streaming execution
|
||||
exec_resp = await client.post(
|
||||
f"/sessions/{session_id}/execute/stream",
|
||||
json={"prompt": prompt},
|
||||
)
|
||||
if exec_resp.status_code != 202:
|
||||
raise HTTPException(status_code=502, detail="Failed to start execution")
|
||||
|
||||
translator = AmplifierdEventTranslator(run_id=run_id, thread_id=thread_id)
|
||||
|
||||
async def event_stream() -> AsyncGenerator[str, None]:
|
||||
# Yield RUN_STARTED first
|
||||
run_started = RunStartedEvent(thread_id=thread_id, run_id=run_id)
|
||||
yield _sse(run_started)
|
||||
|
||||
sse_state: dict[str, Any] = {}
|
||||
async with client.stream("GET", f"/events?session={session_id}") as stream_resp:
|
||||
async for line in stream_resp.aiter_lines():
|
||||
parsed = _parse_sse(line, sse_state)
|
||||
if parsed:
|
||||
sse_state = {} # Reset state for next event
|
||||
ag_events = translator.translate(parsed)
|
||||
for ag_event in ag_events:
|
||||
yield _sse(ag_event)
|
||||
if parsed.get("event_type") == "orchestrator:complete":
|
||||
break
|
||||
|
||||
return StreamingResponse(
|
||||
event_stream(),
|
||||
media_type="text/event-stream",
|
||||
headers={
|
||||
"Cache-Control": "no-cache",
|
||||
"X-Accel-Buffering": "no",
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Artifact endpoints
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@app.get("/api/artifacts/{session_id}")
|
||||
async def list_session_artifacts(
|
||||
session_id: str,
|
||||
email: str = Depends(_require_auth),
|
||||
) -> list[dict[str, Any]]:
|
||||
"""Return sorted list of artifacts for a session."""
|
||||
return list_artifacts(session_id)
|
||||
|
||||
|
||||
@app.get("/api/artifacts/{session_id}/{name}")
|
||||
async def get_session_artifact(
|
||||
session_id: str,
|
||||
name: str,
|
||||
email: str = Depends(_require_auth),
|
||||
) -> JSONResponse:
|
||||
"""Return the content of a specific artifact."""
|
||||
content = get_artifact(session_id, name)
|
||||
if content is None:
|
||||
raise HTTPException(status_code=404, detail="Artifact not found")
|
||||
return JSONResponse({"content": content})
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# MCP App serving
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@app.get("/api/apps/{app_name}")
|
||||
async def serve_mcp_app(
|
||||
app_name: str,
|
||||
email: str = Depends(_require_auth),
|
||||
) -> HTMLResponse:
|
||||
"""Serve an MCP App HTML file from the bundle/apps directory."""
|
||||
# Normalize: add .html extension if missing
|
||||
if not app_name.endswith(".html"):
|
||||
app_name = app_name + ".html"
|
||||
|
||||
app_path = APPS_DIR / app_name
|
||||
if not app_path.exists():
|
||||
raise HTTPException(status_code=404, detail="App not found")
|
||||
|
||||
return HTMLResponse(content=app_path.read_text(encoding="utf-8"))
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Static frontend — MUST be last (catch-all)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_FRONTEND_DIST = Path(__file__).parent.parent / "frontend" / "dist"
|
||||
if _FRONTEND_DIST.exists():
|
||||
app.mount(
|
||||
"/",
|
||||
StaticFiles(directory=str(_FRONTEND_DIST), html=True),
|
||||
name="static",
|
||||
)
|
||||
Reference in New Issue
Block a user