Spaces:
Sleeping
Sleeping
File size: 5,313 Bytes
9ad4546 b9f4b13 9ad4546 b9f4b13 9ad4546 b9f4b13 9ad4546 b9f4b13 9ad4546 b9f4b13 9ad4546 b9f4b13 9ad4546 b9f4b13 9ad4546 b9f4b13 9ad4546 b9f4b13 9ad4546 b9f4b13 9ad4546 | 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 | """Run controller for VentureForge UI.
Provides a thin abstraction over the LangGraph pipeline so that the
Gradio UI and API consumers can:
* start a new run (returns run_id)
* poll the latest state for a run_id via SQLite persistence
* stream real-time events and step transitions
* request cancellation (interrupts the LangGraph execution)
"""
from __future__ import annotations
import logging
import threading
import time
from typing import Any, Generator, 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 stream_run_events(
state: VentureForgeState,
*,
recursion_limit: int = 150,
) -> Generator[tuple[str, dict[str, Any]], None, None]:
"""Yield (node_name, chunk) tuples in real-time as the graph executes."""
run_id = state.run_id
config = {
"recursion_limit": recursion_limit,
"configurable": {"thread_id": run_id},
}
for chunk in GRAPH.stream(state, config=config):
if is_cancel_requested(run_id):
logger.info(f"[run_controller] Cancellation requested for {run_id}, stopping stream")
break
if chunk:
for node_name, node_output in chunk.items():
yield node_name, node_output
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
last_step_time = time.time()
timeout_seconds = 300 # 5 minutes per step
try:
logger.info(f"[run_controller] Starting pipeline run {run_id}")
step_count = 0
for node_name, node_output in stream_run_events(state):
step_count += 1
current_time = time.time()
elapsed = current_time - last_step_time
logger.info(
f"[run_controller] Step {step_count}: {node_name} (took {elapsed:.1f}s)"
)
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
# If cancelled, write cancellation state
if is_cancel_requested(run_id):
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"),
)
logger.info(
f"[run_controller] Pipeline run {run_id} completed after {step_count} steps"
)
except Exception as exc: # pragma: no cover
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:
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}})
if stored is None or not stored.values:
return None
try:
return VentureForgeState.model_validate(stored.values)
except ValidationError:
return None
def request_cancel(run_id: str) -> None:
"""Record that the user requested cancellation for ``run_id``."""
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()
|