167 lines
5.6 KiB
Python
167 lines
5.6 KiB
Python
"""Run orchestration: our DB rows + LFX execution + terminal-state mapping."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from collections.abc import AsyncIterator
|
|
from uuid import UUID
|
|
|
|
import httpx
|
|
from sqlmodel import Session
|
|
|
|
from app.db import get_engine
|
|
from app.models import FlowRun, FlowVersion
|
|
|
|
from . import input_schema as form_schema
|
|
from . import store
|
|
from .lfx_client import LfxClient, LfxNotFound, LfxRunFailed
|
|
from .store import InvalidError
|
|
|
|
|
|
def build_envelope(
|
|
*,
|
|
payload: dict,
|
|
run_id: UUID,
|
|
callback_base_url: str,
|
|
continuation_flow_id: UUID | None,
|
|
callback_token: str,
|
|
fields: list[str] | None = None,
|
|
) -> str:
|
|
"""Run envelope for LFX. ``fields`` (asked form-field names, from the
|
|
version's cached input schema) lets form components emit defaults only
|
|
for asked fields; omitted (None) preserves the legacy shape."""
|
|
body: dict = {
|
|
"payload": payload,
|
|
"run_id": str(run_id),
|
|
"callback_url": callback_base_url.rstrip("/"),
|
|
"continuation_flow_id": (str(continuation_flow_id) if continuation_flow_id else None),
|
|
"callback_token": callback_token,
|
|
}
|
|
if fields is not None:
|
|
body["fields"] = list(fields)
|
|
return json.dumps(body)
|
|
|
|
|
|
def _task_exists(session: Session, run_id: UUID) -> bool:
|
|
return store.has_open_task(session, run_id)
|
|
|
|
|
|
async def execute_run(
|
|
session: Session,
|
|
lfx: LfxClient,
|
|
version: FlowVersion,
|
|
*,
|
|
payload: dict,
|
|
tenant_id: str = "default",
|
|
session_id: str | None = None,
|
|
callback_base_url: str,
|
|
continuation_flow_id: UUID | None = None,
|
|
callback_token: str = "",
|
|
) -> FlowRun:
|
|
"""Run a deployed version to a terminal state (no interrupts).
|
|
|
|
The flow always terminates: either it completes, or its approval branch
|
|
POSTs back to /api/approvals (creating a task) and terminates with a
|
|
pending marker. Status is derived from that side effect, not from
|
|
result-text parsing.
|
|
"""
|
|
if not version.lfx_flow_id:
|
|
raise InvalidError(f"version {version.id} is not deployed to LFX")
|
|
run = store.create_run(
|
|
session, version, input=payload, tenant_id=tenant_id, session_id=session_id
|
|
)
|
|
cached = version.input_schema
|
|
envelope = build_envelope(
|
|
payload=payload,
|
|
run_id=run.id,
|
|
callback_base_url=callback_base_url,
|
|
continuation_flow_id=continuation_flow_id,
|
|
callback_token=callback_token,
|
|
fields=form_schema.schema_field_names(cached),
|
|
)
|
|
try:
|
|
result = await lfx.run_flow(
|
|
version.lfx_flow_id,
|
|
input_value=envelope,
|
|
session_id=session_id or str(run.id),
|
|
)
|
|
except LfxRunFailed as exc:
|
|
return store.set_run_result(session, run, "failed", {"error": str(exc)})
|
|
output = {
|
|
"result": result.result,
|
|
"component": result.component,
|
|
"type": result.type,
|
|
}
|
|
if _task_exists(session, run.id):
|
|
return store.set_run_result(session, run, "pending_approval", output)
|
|
return store.set_run_result(session, run, "completed", output)
|
|
|
|
|
|
async def stream_run(
|
|
session: Session,
|
|
lfx: LfxClient,
|
|
version: FlowVersion,
|
|
*,
|
|
payload: dict,
|
|
tenant_id: str,
|
|
session_id: str | None,
|
|
callback_base_url: str,
|
|
continuation_flow_id: UUID | None,
|
|
callback_token: str,
|
|
) -> tuple[FlowRun, str, str, str]:
|
|
"""Create a run + envelope for SSE streaming. Returns (run, envelope, lfx_id, sid).
|
|
|
|
The actual SSE bytes are produced by the API layer; finalization uses a
|
|
fresh session (the request session may be closed when the generator ends).
|
|
"""
|
|
if not version.lfx_flow_id:
|
|
raise InvalidError(f"version {version.id} is not deployed to LFX")
|
|
run = store.create_run(
|
|
session, version, input=payload, tenant_id=tenant_id, session_id=session_id
|
|
)
|
|
cached = version.input_schema
|
|
envelope = build_envelope(
|
|
payload=payload,
|
|
run_id=run.id,
|
|
callback_base_url=callback_base_url,
|
|
continuation_flow_id=continuation_flow_id,
|
|
callback_token=callback_token,
|
|
fields=form_schema.schema_field_names(cached),
|
|
)
|
|
return run, envelope, version.lfx_flow_id, session_id or str(run.id)
|
|
|
|
|
|
async def stream_lfx_chunks(
|
|
lfx: LfxClient, lfx_flow_id: str, envelope: str, sid: str
|
|
) -> AsyncIterator[str]:
|
|
"""Yield raw SSE chunks from LFX, mapping transport errors to error events."""
|
|
try:
|
|
async with (
|
|
httpx.AsyncClient(
|
|
base_url=lfx.base_url, timeout=lfx.timeout, headers=lfx._headers()
|
|
) as client,
|
|
client.stream(
|
|
"POST",
|
|
f"/flows/{lfx_flow_id}/stream",
|
|
json={"input_value": envelope, "session_id": sid},
|
|
) as response,
|
|
):
|
|
if response.status_code == 404:
|
|
raise LfxNotFound(f"flow {lfx_flow_id} unknown to LFX")
|
|
response.raise_for_status()
|
|
async for chunk in response.aiter_text():
|
|
yield chunk
|
|
except (LfxNotFound, httpx.HTTPError) as exc:
|
|
yield f"event: error\ndata: {json.dumps({'error': str(exc)})}\n\n"
|
|
|
|
|
|
def finalize_streamed_run(run_id: UUID) -> str:
|
|
"""Mark a streamed run completed/pending_approval from a fresh session."""
|
|
with Session(get_engine()) as fresh:
|
|
try:
|
|
fresh_run = store.get_run(fresh, run_id)
|
|
status = "pending_approval" if store.has_open_task(fresh, run_id) else "completed"
|
|
store.set_run_result(fresh, fresh_run, status, {"streamed": True})
|
|
except store.NotFoundError:
|
|
status = "failed"
|
|
return status
|