flowyml 1.3.0__py3-none-any.whl → 1.5.0__py3-none-any.whl

flowyml/ui/backend/routers/runs.py

@@ -2,19 +2,15 @@ from fastapi import APIRouter, HTTPException
 from pydantic import BaseModel
 from flowyml.storage.metadata import SQLiteMetadataStore
 from flowyml.core.project import ProjectManager
-from typing import Optional
 import json
+from flowyml.ui.backend.dependencies import get_store
 
 router = APIRouter()
 
 
-def get_store():
-    return SQLiteMetadataStore()
-
-
 def _iter_metadata_stores():
     """Yield tuples of (project_name, store) including global and project stores."""
-    stores: list[tuple[Optional[str], SQLiteMetadataStore]] = [(None, SQLiteMetadataStore())]
+    stores: list[tuple[str | None, SQLiteMetadataStore]] = [(None, get_store())]
     try:
         manager = ProjectManager()
         for project_meta in manager.list_projects():
@@ -50,16 +46,43 @@ def _sort_runs(runs):
 
 
 @router.get("/")
-async def list_runs(limit: int = 20, project: str = None):
-    """List all runs, optionally filtered by project."""
+async def list_runs(
+    limit: int = 20,
+    project: str = None,
+    pipeline_name: str = None,
+    status: str = None,
+):
+    """List all runs, optionally filtered by project, pipeline_name, and status."""
     try:
         combined = []
         for project_name, store in _iter_metadata_stores():
             # Skip other projects if filtering by project name
             if project and project_name and project != project_name:
                 continue
-            store_runs = store.list_runs(limit=limit)
+
+            # Use store's query method if available for better performance, or list_runs
+            # SQLMetadataStore has query method.
+            if hasattr(store, "query"):
+                filters = {}
+                if pipeline_name:
+                    filters["pipeline_name"] = pipeline_name
+                if status:
+                    filters["status"] = status
+
+                # We can't pass limit to query easily if it doesn't support it,
+                # but SQLMetadataStore.query usually returns all matching.
+                # We'll slice later.
+                store_runs = store.query(**filters)
+            else:
+                store_runs = store.list_runs(limit=limit)
+
             for run in store_runs:
+                # Apply filters if store didn't (e.g. if we used list_runs or store doesn't support query)
+                if pipeline_name and run.get("pipeline_name") != pipeline_name:
+                    continue
+                if status and run.get("status") != status:
+                    continue
+
                 combined.append((run, project_name))
 
         runs = _deduplicate_runs(combined)
@@ -73,12 +96,74 @@ async def list_runs(limit: int = 20, project: str = None):
         return {"runs": [], "error": str(e)}
 
 
+class RunCreate(BaseModel):
+    run_id: str
+    pipeline_name: str
+    status: str = "pending"
+    start_time: str
+    end_time: str | None = None
+    duration: float | None = None
+    metadata: dict = {}
+    project: str | None = None
+    metrics: dict | None = None
+    parameters: dict | None = None
+
+
+@router.post("/")
+async def create_run(run: RunCreate):
+    """Create or update a run."""
+    try:
+        store = get_store()
+
+        # Prepare metadata dict
+        metadata = run.metadata.copy()
+        metadata.update(
+            {
+                "pipeline_name": run.pipeline_name,
+                "status": run.status,
+                "start_time": run.start_time,
+                "end_time": run.end_time,
+                "duration": run.duration,
+                "project": run.project,
+            },
+        )
+
+        if run.metrics:
+            metadata["metrics"] = run.metrics
+
+        if run.parameters:
+            metadata["parameters"] = run.parameters
+
+        store.save_run(run.run_id, metadata)
+        return {"status": "success", "run_id": run.run_id}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+
 @router.get("/{run_id}")
 async def get_run(run_id: str):
     """Get details for a specific run."""
     run, _ = _find_run(run_id)
     if not run:
         raise HTTPException(status_code=404, detail="Run not found")
+
+    # Mark dead steps
+    dead_steps = _get_dead_steps(run_id)
+    if dead_steps and "steps" in run:
+        for step_name in dead_steps:
+            if step_name in run["steps"]:
+                # Only mark as dead if it was running
+                if run["steps"][step_name].get("status") == "running":
+                    run["steps"][step_name]["status"] = "dead"
+                    run["steps"][step_name]["success"] = False
+
+    # Inject heartbeat timestamps
+    with _heartbeat_lock:
+        if run_id in _heartbeat_timestamps:
+            for step_name, ts in _heartbeat_timestamps[run_id].items():
+                if step_name in run.get("steps", {}):
+                    run["steps"][step_name]["last_heartbeat"] = ts
+
     return run
 
 
@@ -206,3 +291,196 @@ async def get_cloud_status(run_id: str):
         "cloud_status": cloud_status,
         "cloud_error": cloud_error,
     }
+
+
+class HeartbeatRequest(BaseModel):
+    step_name: str
+    status: str = "running"
+
+
+# In-memory storage for heartbeat timestamps
+# Format: {run_id: {step_name: last_heartbeat_timestamp}}
+_heartbeat_timestamps: dict[str, dict[str, float]] = {}
+_heartbeat_lock = __import__("threading").Lock()
+
+# Heartbeat interval in seconds (should match executor's interval)
+HEARTBEAT_INTERVAL = 5
+# Number of missed heartbeats before marking step as dead
+DEAD_THRESHOLD = 3
+
+
+def _record_heartbeat(run_id: str, step_name: str) -> None:
+    """Record heartbeat timestamp for a step."""
+    import time
+
+    with _heartbeat_lock:
+        if run_id not in _heartbeat_timestamps:
+            _heartbeat_timestamps[run_id] = {}
+        _heartbeat_timestamps[run_id][step_name] = time.time()
+
+
+def _get_dead_steps(run_id: str) -> list[str]:
+    """Get list of steps that have missed too many heartbeats."""
+    import time
+
+    dead_steps = []
+    timeout = HEARTBEAT_INTERVAL * DEAD_THRESHOLD
+
+    with _heartbeat_lock:
+        if run_id not in _heartbeat_timestamps:
+            return []
+
+        current_time = time.time()
+        for step_name, last_heartbeat in _heartbeat_timestamps[run_id].items():
+            if current_time - last_heartbeat > timeout:
+                dead_steps.append(step_name)
+
+    return dead_steps
+
+
+def _cleanup_heartbeats(run_id: str) -> None:
+    """Remove heartbeat tracking for a completed run."""
+    with _heartbeat_lock:
+        _heartbeat_timestamps.pop(run_id, None)
+
+
+@router.post("/{run_id}/steps/{step_name}/heartbeat")
+async def step_heartbeat(run_id: str, step_name: str, heartbeat: HeartbeatRequest):
+    """Receive heartbeat from a running step.
+
+    Returns:
+        dict: Instructions for the step (e.g., {"action": "continue"} or {"action": "stop"})
+    """
+    store = _find_store_for_run(run_id)
+
+    # Record heartbeat timestamp
+    _record_heartbeat(run_id, step_name)
+
+    # Check if run is marked for stopping
+    run = store.load_run(run_id)
+    if not run:
+        raise HTTPException(status_code=404, detail="Run not found")
+
+    run_status = run.get("status")
+    if run_status in ["stopping", "stopped", "cancelled", "cancelling"]:
+        return {"action": "stop"}
+
+    return {"action": "continue"}
+
+
+@router.get("/{run_id}/dead-steps")
+async def get_dead_steps(run_id: str):
+    """Get list of steps that appear to be dead (missed heartbeats)."""
+    dead_steps = _get_dead_steps(run_id)
+    return {"dead_steps": dead_steps}
+
+
+@router.post("/{run_id}/stop")
+async def stop_run(run_id: str):
+    """Signal a run to stop."""
+    store = _find_store_for_run(run_id)
+
+    try:
+        # Update run status to STOPPING
+        # This will be picked up by the next heartbeat
+        store.update_run_status(run_id, "stopping")
+        return {"status": "success", "message": "Stop signal sent"}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+class LogChunk(BaseModel):
+    content: str
+    level: str = "INFO"
+    timestamp: str | None = None
+
+
+@router.post("/{run_id}/steps/{step_name}/logs")
+async def post_step_logs(run_id: str, step_name: str, log_chunk: LogChunk):
+    """Receive log chunk from a running step."""
+    import anyio
+
+    from flowyml.utils.config import get_config
+
+    # Store logs in the runs directory
+    runs_dir = get_config().runs_dir
+    log_dir = runs_dir / run_id / "logs"
+    log_dir.mkdir(parents=True, exist_ok=True)
+
+    log_file = log_dir / f"{step_name}.log"
+
+    # Append log content
+    timestamp = log_chunk.timestamp or ""
+    line = f"[{timestamp}] [{log_chunk.level}] {log_chunk.content}\n"
+
+    def write_log():
+        with open(log_file, "a") as f:
+            f.write(line)
+
+    await anyio.to_thread.run_sync(write_log)
+
+    # Broadcast to WebSocket clients
+    try:
+        from flowyml.ui.backend.routers.websocket import manager
+
+        await manager.broadcast_log(run_id, step_name, log_chunk.content)
+    except Exception:
+        pass  # Ignore WebSocket broadcast failures
+
+    return {"status": "success"}
+
+
+@router.get("/{run_id}/steps/{step_name}/logs")
+async def get_step_logs(run_id: str, step_name: str, offset: int = 0):
+    """Get logs for a specific step."""
+    import anyio
+
+    from flowyml.utils.config import get_config
+
+    runs_dir = get_config().runs_dir
+    log_file = runs_dir / run_id / "logs" / f"{step_name}.log"
+
+    if not log_file.exists():
+        return {"logs": "", "offset": 0, "has_more": False}
+
+    def read_log():
+        with open(log_file) as f:
+            return f.read()
+
+    content = await anyio.to_thread.run_sync(read_log)
+
+    # Return content from offset
+    if offset > 0 and offset < len(content):
+        content = content[offset:]
+
+    return {
+        "logs": content,
+        "offset": offset + len(content),
+        "has_more": False,  # For now, always return all available
+    }
+
+
+@router.get("/{run_id}/logs")
+async def get_run_logs(run_id: str):
+    """Get all logs for a run."""
+    import anyio
+
+    from flowyml.utils.config import get_config
+
+    runs_dir = get_config().runs_dir
+    log_dir = runs_dir / run_id / "logs"
+
+    if not log_dir.exists():
+        return {"logs": {}}
+
+    def read_all_logs():
+        logs = {}
+        for log_file in log_dir.glob("*.log"):
+            step_name = log_file.stem
+            with open(log_file) as f:
+                logs[step_name] = f.read()
+        return logs
+
+    logs = await anyio.to_thread.run_sync(read_all_logs)
+
+    return {"logs": logs}

flowyml/ui/backend/routers/schedules.py

@@ -71,16 +71,13 @@ async def create_schedule(schedule: ScheduleRequest):
         else:
             # Check if it's a historical pipeline (in metadata but not registry)
             # This means we can't run it because we don't have the code loaded
-            from flowyml.storage.metadata import SQLiteMetadataStore
+            from flowyml.ui.backend.dependencies import get_store
 
-            store = SQLiteMetadataStore()
+            store = get_store()
             pipelines = store.list_pipelines()
 
             if schedule.pipeline_name in pipelines:
                 # Try to load pipeline definition
-                from flowyml.storage.metadata import SQLiteMetadataStore
-
-                store = SQLiteMetadataStore()
                 pipeline_def = store.get_pipeline_definition(schedule.pipeline_name)
 
                 if pipeline_def:
@@ -188,7 +185,6 @@ async def get_schedule_history(schedule_name: str, limit: int = 50):
 async def list_registered_pipelines(project: str = None):
     """List all pipelines available for scheduling."""
     from flowyml.core.templates import list_templates
-    from flowyml.storage.metadata import SQLiteMetadataStore
 
     registered = pipeline_registry.list_pipelines()
     templates = list_templates()
@@ -196,22 +192,10 @@ async def list_registered_pipelines(project: str = None):
     # Also get pipelines from metadata store (historical runs)
     metadata_pipelines = []
     try:
-        store = SQLiteMetadataStore()
-        import sqlite3
-
-        conn = sqlite3.connect(store.db_path)
-        cursor = conn.cursor()
-
-        if project:
-            cursor.execute(
-                "SELECT DISTINCT pipeline_name FROM runs WHERE project = ? ORDER BY pipeline_name",
-                (project,),
-            )
-        else:
-            cursor.execute("SELECT DISTINCT pipeline_name FROM runs ORDER BY pipeline_name")
+        from flowyml.ui.backend.dependencies import get_store
 
-        metadata_pipelines = [row[0] for row in cursor.fetchall()]
-        conn.close()
+        store = get_store()
+        metadata_pipelines = store.list_pipelines(project=project)
     except Exception as e:
         print(f"Failed to fetch pipelines from metadata store: {e}")
 

flowyml/ui/backend/routers/stats.py

@@ -0,0 +1,14 @@
+from fastapi import APIRouter, HTTPException
+from flowyml.ui.backend.dependencies import get_store
+
+router = APIRouter()
+
+
+@router.get("/")
+async def get_global_stats(project: str | None = None):
+    """Get global statistics."""
+    try:
+        store = get_store()
+        return store.get_statistics(project=project)
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))

flowyml/ui/backend/routers/traces.py

@@ -1,7 +1,6 @@
 from fastapi import APIRouter, HTTPException
-from flowyml.storage.metadata import SQLiteMetadataStore
-import contextlib
-import builtins
+from pydantic import BaseModel
+from flowyml.ui.backend.dependencies import get_store
 
 router = APIRouter()
 
@@ -14,61 +13,14 @@ async def list_traces(
     project: str | None = None,
 ):
     """List traces, optionally filtered by project."""
-    store = SQLiteMetadataStore()
-
-    # We need to implement list_traces in metadata store or query manually
-    # For now, let's query manually via sqlite
-    import sqlite3
-
-    conn = sqlite3.connect(store.db_path)
-    cursor = conn.cursor()
-
-    query = "SELECT * FROM traces"
-    params = []
-    conditions = []
-
-    if trace_id:
-        conditions.append("trace_id = ?")
-        params.append(trace_id)
-
-    if event_type:
-        conditions.append("event_type = ?")
-        params.append(event_type)
-
-    if project:
-        conditions.append("project = ?")
-        params.append(project)
-
-    if conditions:
-        query += " WHERE " + " AND ".join(conditions)
-
-    query += " ORDER BY start_time DESC LIMIT ?"
-    params.append(limit)
-
-    cursor.execute(query, params)
-    columns = [description[0] for description in cursor.description]
-    rows = cursor.fetchall()
-
-    traces = []
-    import json
-
-    for row in rows:
-        trace = dict(zip(columns, row, strict=False))
-        # Parse JSON fields
-        for field in ["inputs", "outputs", "metadata"]:
-            if trace[field]:
-                with contextlib.suppress(builtins.BaseException):
-                    trace[field] = json.loads(trace[field])
-        traces.append(trace)
-
-    conn.close()
-    return traces
+    store = get_store()
+    return store.list_traces(limit=limit, trace_id=trace_id, event_type=event_type, project=project)
 
 
 @router.get("/{trace_id}")
 async def get_trace(trace_id: str):
     """Get a specific trace tree."""
-    store = SQLiteMetadataStore()
+    store = get_store()
     events = store.get_trace(trace_id)
     if not events:
         raise HTTPException(status_code=404, detail="Trace not found")
@@ -82,3 +34,35 @@ async def get_trace(trace_id: str):
         return event
 
     return [build_tree(root) for root in root_events]
+
+
+class TraceEventCreate(BaseModel):
+    event_id: str
+    trace_id: str
+    parent_id: str | None = None
+    event_type: str
+    name: str
+    inputs: dict | None = None
+    outputs: dict | None = None
+    start_time: float | None = None
+    end_time: float | None = None
+    duration: float | None = None
+    status: str | None = None
+    error: str | None = None
+    metadata: dict | None = None
+    prompt_tokens: int | None = None
+    completion_tokens: int | None = None
+    total_tokens: int | None = None
+    cost: float | None = None
+    model: str | None = None
+
+
+@router.post("/")
+async def create_trace_event(event: TraceEventCreate):
+    """Create or update a trace event."""
+    try:
+        store = get_store()
+        store.save_trace_event(event.dict())
+        return {"status": "success", "event_id": event.event_id}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))

flowyml/ui/backend/routers/websocket.py

@@ -0,0 +1,121 @@
+"""WebSocket router for real-time log streaming."""
+
+from fastapi import APIRouter, WebSocket, WebSocketDisconnect
+from datetime import datetime
+import asyncio
+import contextlib
+
+router = APIRouter()
+
+
+class ConnectionManager:
+    """Manage WebSocket connections for log streaming."""
+
+    def __init__(self):
+        # Format: {run_id: {step_name: [websocket, ...]}}
+        self.active_connections: dict[str, dict[str, list[WebSocket]]] = {}
+        self._lock = asyncio.Lock()
+
+    async def connect(self, websocket: WebSocket, run_id: str, step_name: str = "__all__"):
+        """Accept and track a WebSocket connection."""
+        await websocket.accept()
+        async with self._lock:
+            if run_id not in self.active_connections:
+                self.active_connections[run_id] = {}
+            if step_name not in self.active_connections[run_id]:
+                self.active_connections[run_id][step_name] = []
+            self.active_connections[run_id][step_name].append(websocket)
+
+    async def disconnect(self, websocket: WebSocket, run_id: str, step_name: str = "__all__"):
+        """Remove a WebSocket connection."""
+        async with self._lock:
+            if run_id in self.active_connections:
+                if step_name in self.active_connections[run_id]:
+                    with contextlib.suppress(ValueError):
+                        self.active_connections[run_id][step_name].remove(websocket)
+                    if not self.active_connections[run_id][step_name]:
+                        del self.active_connections[run_id][step_name]
+                if not self.active_connections[run_id]:
+                    del self.active_connections[run_id]
+
+    async def broadcast_log(self, run_id: str, step_name: str, log_content: str):
+        """Broadcast log content to all connected clients for a run/step."""
+        async with self._lock:
+            connections_to_notify = []
+
+            if run_id in self.active_connections:
+                # Notify step-specific subscribers
+                if step_name in self.active_connections[run_id]:
+                    connections_to_notify.extend(self.active_connections[run_id][step_name])
+                # Notify all-steps subscribers
+                if "__all__" in self.active_connections[run_id]:
+                    connections_to_notify.extend(self.active_connections[run_id]["__all__"])
+
+        # Send to all relevant connections (outside lock to avoid blocking)
+        message = {
+            "type": "log",
+            "step": step_name,
+            "content": log_content,
+            "timestamp": datetime.now().isoformat(),
+        }
+        for ws in connections_to_notify:
+            with contextlib.suppress(Exception):
+                await ws.send_json(message)
+
+
+# Global connection manager instance
+manager = ConnectionManager()
+
+
+@router.websocket("/ws/runs/{run_id}/logs")
+async def websocket_logs(websocket: WebSocket, run_id: str, step_name: str = "__all__"):
+    """WebSocket endpoint for streaming logs.
+
+    Query params:
+        step_name: Optional. Subscribe to specific step logs only.
+
+    Messages sent:
+        {"type": "log", "step": "step_name", "content": "...", "timestamp": "..."}
+        {"type": "dead_steps", "steps": ["step1", "step2"]}
+    """
+    await manager.connect(websocket, run_id, step_name)
+    try:
+        # Keep connection alive and handle incoming messages
+        while True:
+            try:
+                # Wait for any message (ping/pong or close)
+                data = await asyncio.wait_for(websocket.receive_text(), timeout=30.0)
+                # Could handle client commands here if needed
+                if data == "ping":
+                    await websocket.send_text("pong")
+            except asyncio.TimeoutError:
+                # Send heartbeat to keep connection alive
+                try:
+                    await websocket.send_json({"type": "heartbeat"})
+                except Exception:
+                    break
+    except WebSocketDisconnect:
+        pass
+    finally:
+        await manager.disconnect(websocket, run_id, step_name)
+
+
+@router.websocket("/ws/runs/{run_id}/steps/{step_name}/logs")
+async def websocket_step_logs(websocket: WebSocket, run_id: str, step_name: str):
+    """WebSocket endpoint for streaming logs of a specific step."""
+    await manager.connect(websocket, run_id, step_name)
+    try:
+        while True:
+            try:
+                data = await asyncio.wait_for(websocket.receive_text(), timeout=30.0)
+                if data == "ping":
+                    await websocket.send_text("pong")
+            except asyncio.TimeoutError:
+                try:
+                    await websocket.send_json({"type": "heartbeat"})
+                except Exception:
+                    break
+    except WebSocketDisconnect:
+        pass
+    finally:
+        await manager.disconnect(websocket, run_id, step_name)