88 lines
2.5 KiB
Python
88 lines
2.5 KiB
Python
# SPDX-License-Identifier: MIT
|
|
# Copyright (c) 2025 LlamaIndex Inc.
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import time
|
|
from typing import AsyncGenerator, Callable
|
|
|
|
from workflows.events import Event, StopEvent
|
|
from workflows.runtime.types.plugin import Plugin, SnapshottableRuntime, WorkflowRuntime
|
|
from workflows.runtime.types.step_function import StepWorkerFunction
|
|
from workflows.runtime.types.ticks import WorkflowTick
|
|
from workflows.decorators import P, R
|
|
from workflows.workflow import Workflow
|
|
|
|
|
|
class BasicRuntime:
|
|
def register(
|
|
self,
|
|
workflow: Workflow,
|
|
workflow_function: Callable[P, R],
|
|
steps: dict[str, StepWorkerFunction[R]],
|
|
) -> None:
|
|
return
|
|
|
|
def new_runtime(self, run_id: str) -> WorkflowRuntime:
|
|
snapshottable: SnapshottableRuntime = AsyncioWorkflowRuntime(run_id)
|
|
return snapshottable
|
|
|
|
|
|
basic_runtime: Plugin = BasicRuntime()
|
|
|
|
|
|
class AsyncioWorkflowRuntime:
|
|
"""
|
|
A plugin interface to switch out a broker runtime (external library or service that manages durable/distributed step execution)
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
run_id: str,
|
|
) -> None:
|
|
self.run_id = run_id
|
|
self.receive_queue: asyncio.Queue[WorkflowTick] = asyncio.Queue()
|
|
self.publish_queue: asyncio.Queue[Event] = asyncio.Queue()
|
|
self.ticks: list[WorkflowTick] = []
|
|
|
|
def on_tick(self, tick: WorkflowTick) -> None:
|
|
self.ticks.append(tick)
|
|
|
|
def replay(self) -> list[WorkflowTick]:
|
|
return self.ticks
|
|
|
|
async def wait_receive(self) -> WorkflowTick:
|
|
return await self.receive_queue.get()
|
|
|
|
async def write_to_event_stream(self, event: Event) -> None:
|
|
self.publish_queue.put_nowait(event)
|
|
|
|
async def stream_published_events(self) -> AsyncGenerator[Event, None]:
|
|
while True:
|
|
item = await self.publish_queue.get()
|
|
yield item
|
|
if isinstance(item, StopEvent):
|
|
break
|
|
|
|
async def send_event(self, tick: WorkflowTick) -> None:
|
|
self.receive_queue.put_nowait(tick)
|
|
|
|
async def register_step_worker(
|
|
self, step_name: str, step_worker: StepWorkerFunction[R]
|
|
) -> StepWorkerFunction[R]:
|
|
return step_worker
|
|
|
|
async def register_workflow_function(
|
|
self, workflow_function: Callable[P, R]
|
|
) -> Callable[P, R]:
|
|
return workflow_function
|
|
|
|
async def get_now(self) -> float:
|
|
return time.monotonic()
|
|
|
|
async def sleep(self, seconds: float) -> None:
|
|
await asyncio.sleep(seconds)
|
|
|
|
async def close(self) -> None:
|
|
pass
|