File size: 6,723 Bytes
a6e70b1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 | """Run controller for VentureForge UI.
Provides a thin abstraction over the LangGraph pipeline so that the
Gradio UI can:
* start a new run (returns run_id)
* poll the latest state for a run_id
* request cancellation (interrupts the LangGraph execution)
NOTE: The actual persistence is handled by the LangGraph SQLite
checkpointer configured in ``src.graph``. ``GRAPH.get_state`` is used to
recover the latest VentureForgeState snapshot for a given run_id.
"""
from __future__ import annotations
import logging
import threading
from typing import Optional
from pydantic import ValidationError
from src.graph import GRAPH
from src.main import make_initial_state
from src.state.schema import VentureForgeState
logger = logging.getLogger(__name__)
# Single active run tracking (hackathon scope: one run at a time)
_active_thread: Optional[threading.Thread] = None
_active_run_id: Optional[str] = None
_cancel_requested: set[str] = set()
class RunInProgressError(RuntimeError):
"""Raised when attempting to start a new run while one is active."""
def start_run(
domain: str,
*,
max_pain_points: int | None = None,
ideas_per_run: int | None = None,
) -> str:
"""Start a new pipeline run in a background thread.
Returns the newly created ``run_id``. Raises ``RunInProgressError``
if a previous run is still active.
"""
global _active_thread, _active_run_id
if _active_thread is not None and _active_thread.is_alive():
raise RunInProgressError("A run is already in progress")
state = make_initial_state(
domain,
max_pain_points=max_pain_points,
ideas_per_run=ideas_per_run,
)
run_id = state.run_id
def _worker() -> None:
global _active_thread, _active_run_id
import time
last_step_time = time.time()
timeout_seconds = 300 # 5 minutes per step
try:
logger.info(f"[run_controller] Starting pipeline run {run_id}")
# Drive the pipeline via LangGraph streaming so we can
# respond to cancellation requests between steps. The
# checkpointer persists intermediate states; ``poll_state``
# reads them for the UI.
step_count = 0
for chunk in GRAPH.stream(
state,
config={
"recursion_limit": 150, # Increased from 80 to handle retry loops
"configurable": {"thread_id": run_id},
},
):
step_count += 1
current_time = time.time()
elapsed = current_time - last_step_time
# Log each step for debugging
if chunk:
node_name = list(chunk.keys())[0] if chunk else "unknown"
logger.info(f"[run_controller] Step {step_count}: {node_name} (took {elapsed:.1f}s)")
# Warn if step took too long
if elapsed > timeout_seconds:
logger.warning(
f"[run_controller] Step {step_count} ({node_name}) took {elapsed:.1f}s "
f"(timeout threshold: {timeout_seconds}s)"
)
last_step_time = current_time
# Check for cancellation between steps
if is_cancel_requested(run_id):
logger.info(f"[run_controller] Cancellation requested for {run_id}, stopping")
# Update state to mark as cancelled
current_state = poll_state(run_id)
if current_state:
GRAPH.update_state(
config={"configurable": {"thread_id": run_id}},
values=current_state.mark_cancelled("User requested stop"),
)
break
logger.info(f"[run_controller] Pipeline run {run_id} completed after {step_count} steps")
except Exception as exc: # pragma: no cover - surfaced via UI logs
# Surface worker errors to stdout so they are visible when
# running the Gradio UI. Without this, exceptions inside
# the background thread can make the UI hang without any
# clear error message.
logger.error(f"[run_controller] Pipeline worker crashed: {exc!r}", exc_info=True)
print(f"[run_controller] Pipeline worker crashed: {exc!r}", flush=True)
finally:
# Mark thread as finished; the last checkpoint remains
# available for inspection.
logger.info(f"[run_controller] Worker thread for {run_id} finished")
_active_thread = None
_active_run_id = None
t = threading.Thread(target=_worker, daemon=True)
_active_thread = t
_active_run_id = run_id
t.start()
return run_id
def poll_state(run_id: str) -> VentureForgeState | None:
"""Return the latest VentureForgeState for ``run_id`` from checkpoints.
Returns ``None`` if no state has been persisted yet.
"""
stored = GRAPH.get_state(config={"configurable": {"thread_id": run_id}})
# ``stored`` is a LangGraph StoredState; ``values`` holds the latest
# state dict that was checkpointed.
if stored is None or not stored.values:
# No meaningful state yet (LangGraph may persist an empty dict
# before the first node runs). Treat this the same as "no
# checkpoints" so the UI keeps waiting instead of raising a
# validation error about missing required fields like ``domain``.
return None
try:
return VentureForgeState.model_validate(stored.values)
except ValidationError:
# If the stored values are incomplete or from an older schema,
# fail gracefully and let the UI keep polling.
return None
def request_cancel(run_id: str) -> None:
"""Record that the user requested cancellation for ``run_id``.
The cancellation will take effect between LangGraph steps. If an agent
is currently executing (e.g., waiting for an LLM response), the
cancellation will be processed after that agent completes.
"""
logger.info(f"[run_controller] Cancellation requested for run {run_id}")
_cancel_requested.add(run_id)
def is_cancel_requested(run_id: str) -> bool:
"""Return True if cancellation was requested for this run_id."""
return run_id in _cancel_requested
def is_run_active() -> bool:
"""Return True if a pipeline run is currently active."""
return _active_thread is not None and _active_thread.is_alive()
|