service-forge 0.1.18__py3-none-any.whl

service_forge/workflow/triggers/websocket_api_trigger.py

@@ -0,0 +1,189 @@
+from __future__ import annotations
+import uuid
+import asyncio
+import json
+from loguru import logger
+from service_forge.workflow.trigger import Trigger
+from typing import AsyncIterator, Any
+from fastapi import FastAPI, WebSocket, WebSocketDisconnect
+from service_forge.workflow.port import Port
+from google.protobuf.message import Message
+from google.protobuf.json_format import MessageToJson
+
+class WebSocketAPITrigger(Trigger):
+    DEFAULT_INPUT_PORTS = [
+        Port("app", FastAPI),
+        Port("path", str),
+        Port("data_type", type),
+    ]
+
+    DEFAULT_OUTPUT_PORTS = [
+        Port("trigger", bool),
+        Port("data", Any),
+    ]
+
+    def __init__(self, name: str):
+        super().__init__(name)
+        self.events = {}
+        self.is_setup_websocket = False
+
+    @staticmethod
+    def serialize_result(result: Any):
+        if isinstance(result, Message):
+            return MessageToJson(
+                result,
+                preserving_proto_field_name=True
+            )
+        return result
+
+    async def handle_stream_output(
+        self,
+        websocket: WebSocket,
+        task_id: uuid.UUID,
+    ):
+        try:
+            while True:
+                item = await self.stream_queues[task_id].get()
+
+                if item.is_error:
+                    error_response = {
+                        "type": "stream_error",
+                        "task_id": str(task_id),
+                        "detail": str(item.result)
+                    }
+                    await websocket.send_text(json.dumps(error_response))
+                    break
+                
+                if item.is_end:
+                    # Stream ended, send final result if available
+                    if item.result is not None:
+                        serialized = self.serialize_result(item.result)
+                        if isinstance(serialized, str):
+                            try:
+                                data = json.loads(serialized)
+                            except json.JSONDecodeError:
+                                data = serialized
+                        else:
+                            data = serialized
+                        
+                        end_response = {
+                            "type": "stream_end",
+                            "task_id": str(task_id),
+                            "data": data
+                        }
+                    else:
+                        end_response = {
+                            "type": "stream_end",
+                            "task_id": str(task_id)
+                        }
+                    await websocket.send_text(json.dumps(end_response))
+                    break
+
+                # Send stream data
+                serialized = self.serialize_result(item.result)
+                if isinstance(serialized, str):
+                    try:
+                        data = json.loads(serialized)
+                    except json.JSONDecodeError:
+                        data = serialized
+                else:
+                    data = serialized
+                
+                stream_response = {
+                    "type": "stream",
+                    "task_id": str(task_id),
+                    "data": data
+                }
+                await websocket.send_text(json.dumps(stream_response))
+        except Exception as e:
+            logger.error(f"Error handling stream output for task {task_id}: {e}")
+            error_response = {
+                "type": "stream_error",
+                "task_id": str(task_id),
+                "detail": str(e)
+            }
+            try:
+                await websocket.send_text(json.dumps(error_response))
+            except Exception:
+                pass
+        finally:
+            if task_id in self.stream_queues:
+                del self.stream_queues[task_id]
+
+    async def handle_websocket_message(
+        self,
+        websocket: WebSocket,
+        data_type: type,
+        message_data: dict,
+    ):
+        task_id = uuid.uuid4()
+        self.result_queues[task_id] = asyncio.Queue()
+        self.stream_queues[task_id] = asyncio.Queue()
+
+        if data_type is Any:
+            converted_data = message_data
+        else:
+            try:
+                converted_data = data_type(**message_data)
+            except Exception as e:
+                error_msg = {"error": f"Failed to convert data: {str(e)}"}
+                await websocket.send_text(json.dumps(error_msg))
+                return
+
+        # Always start background task to handle stream output
+        asyncio.create_task(self.handle_stream_output(websocket, task_id))
+
+        self.trigger_queue.put_nowait({
+            "id": task_id,
+            "data": converted_data,
+        })
+
+        # The stream handler will send all messages including stream_end when workflow completes
+
+    def _setup_websocket(self, app: FastAPI, path: str, data_type: type) -> None:
+        async def websocket_handler(websocket: WebSocket):
+            await websocket.accept()
+            
+            try:
+                while True:
+                    # Receive message from client
+                    data = await websocket.receive_text()
+                    try:
+                        message = json.loads(data)
+                        
+                        # Handle the message and trigger workflow
+                        await self.handle_websocket_message(
+                            websocket,
+                            data_type,
+                            message
+                        )
+                    except json.JSONDecodeError:
+                        error_msg = {"error": "Invalid JSON format"}
+                        await websocket.send_text(json.dumps(error_msg))
+                    except Exception as e:
+                        logger.error(f"Error handling websocket message: {e}")
+                        error_msg = {"error": str(e)}
+                        await websocket.send_text(json.dumps(error_msg))
+            except WebSocketDisconnect:
+                logger.info("WebSocket client disconnected")
+            except Exception as e:
+                logger.error(f"WebSocket connection error: {e}")
+
+        app.websocket(path)(websocket_handler)
+
+    async def _run(self, app: FastAPI, path: str, data_type: type) -> AsyncIterator[bool]:
+        if not self.is_setup_websocket:
+            self._setup_websocket(app, path, data_type)
+            self.is_setup_websocket = True
+
+        while True:
+            try:
+                trigger = await self.trigger_queue.get()
+                self.prepare_output_edges(self.get_output_port_by_name('data'), trigger['data'])
+                yield self.trigger(trigger['id'])
+            except Exception as e:
+                logger.error(f"Error in WebSocketAPITrigger._run: {e}")
+                continue
+
+    async def _stop(self) -> AsyncIterator[bool]:
+        pass

service_forge/workflow/workflow.py

@@ -0,0 +1,227 @@
+from __future__ import annotations
+import traceback
+import asyncio
+import uuid
+from typing import AsyncIterator, Awaitable, Callable, Any
+from loguru import logger
+from copy import deepcopy
+from .node import Node
+from .port import Port
+from .trigger import Trigger
+from .edge import Edge
+from ..db.database import DatabaseManager
+from ..utils.workflow_clone import workflow_clone
+from .workflow_callback import WorkflowCallback, BuiltinWorkflowCallback, CallbackEvent
+from .workflow_config import WorkflowConfig
+
+class Workflow:
+    def __init__(
+        self,
+        id: uuid.UUID,
+        config: WorkflowConfig,
+        nodes: list[Node],
+        input_ports: list[Port],
+        output_ports: list[Port],
+        _handle_stream_output: Callable[[str, AsyncIterator[str]], Awaitable[None]] = None, # deprecated
+        _handle_query_user: Callable[[str, str], Awaitable[str]] = None,
+        database_manager: DatabaseManager = None,
+        max_concurrent_runs: int = 10,
+        callbacks: list[WorkflowCallback] = [],
+
+        # for run
+        task_id: uuid.UUID = None,
+        real_trigger_node: Trigger = None,
+    ) -> None:
+        self.id = id
+        self.config = config
+        self.nodes = nodes
+        self.ready_nodes: list[Node] = []
+        self.input_ports = input_ports
+        self.output_ports = output_ports
+        self._handle_stream_output = _handle_stream_output
+        self._handle_query_user = _handle_query_user
+        self.after_trigger_workflow = None
+        self.result_port = Port("result", Any)
+        self.database_manager = database_manager
+        self.max_concurrent_runs = max_concurrent_runs
+        self.run_semaphore = asyncio.Semaphore(max_concurrent_runs)
+        self.callbacks = callbacks
+        self.task_id = task_id
+        self.real_trigger_node = real_trigger_node
+        self._validate()
+    
+    @property
+    def name(self) -> str:
+        return self.config.name
+    
+    @property
+    def version(self) -> str:
+        return self.config.version
+    
+    @property
+    def description(self) -> str:
+        return self.config.description
+
+    def register_callback(self, callback: WorkflowCallback) -> None:
+        self.callbacks.append(callback)
+
+    def unregister_callback(self, callback: WorkflowCallback) -> None:
+        self.callbacks.remove(callback)
+
+    async def call_callbacks(self, callback_type: CallbackEvent, *args, **kwargs) -> None:
+        for callback in self.callbacks:
+            if callback_type == CallbackEvent.ON_WORKFLOW_START:
+                await callback.on_workflow_start(*args, **kwargs)
+            elif callback_type == CallbackEvent.ON_WORKFLOW_END:
+                await callback.on_workflow_end(*args, **kwargs)
+            elif callback_type == CallbackEvent.ON_NODE_START:
+                await callback.on_node_start(*args, **kwargs)
+            elif callback_type == CallbackEvent.ON_NODE_END:
+                await callback.on_node_end(*args, **kwargs)
+            elif callback_type == CallbackEvent.ON_NODE_STREAM_OUTPUT:
+                await callback.on_node_stream_output(*args, **kwargs)
+
+    def add_nodes(self, nodes: list[Node]) -> None:
+        for node in nodes:
+            node.workflow = self
+        self.nodes.extend(nodes)
+
+    def remove_nodes(self, nodes: list[Node]) -> None:
+        for node in nodes:
+            self.nodes.remove(node)
+
+    def load_config(self) -> None:
+        ...
+
+    def _validate(self) -> None:
+        # DAG
+        ...
+
+    def get_input_port_by_name(self, name: str) -> Port:
+        for input_port in self.input_ports:
+            if input_port.name == name:
+                return input_port
+        return None
+
+    def get_output_port_by_name(self, name: str) -> Port:
+        for output_port in self.output_ports:
+            if output_port.name == name:
+                return output_port
+        return None
+
+    def get_trigger_node(self) -> Trigger:
+        trigger_nodes = [node for node in self.nodes if isinstance(node, Trigger)]
+        if not trigger_nodes:
+            raise ValueError("No trigger nodes found in workflow.")
+        if len(trigger_nodes) > 1:
+            raise ValueError("Multiple trigger nodes found in workflow.")
+        return trigger_nodes[0]
+
+    async def _run_node_with_callbacks(self, node: Node) -> None:
+        await self.call_callbacks(CallbackEvent.ON_NODE_START, node=node)
+        
+        try:
+            result = node.run()
+            if hasattr(result, '__anext__'):
+                await self.handle_node_stream_output(node, result)
+            elif asyncio.iscoroutine(result):
+                await result
+        finally:
+            await self.call_callbacks(CallbackEvent.ON_NODE_END, node=node)
+
+    async def run_after_trigger(self) -> Any:
+        logger.info(f"Running workflow: {self.name}")
+
+        await self.call_callbacks(CallbackEvent.ON_WORKFLOW_START, workflow=self)
+
+        self.ready_nodes = []
+        for edge in self.get_trigger_node().output_edges:
+            edge.end_port.trigger()
+
+        try:
+            for input_port in self.input_ports:
+                if input_port.value is not None:
+                    input_port.port.node.fill_input(input_port.port, input_port.value)
+
+            for node in self.nodes:
+                for key in node.AUTO_FILL_INPUT_PORTS:
+                    if key[0] not in [edge.end_port.name for edge in node.input_edges]:
+                        node.fill_input_by_name(key[0], key[1])
+
+            while self.ready_nodes:
+                nodes = self.ready_nodes.copy()
+                self.ready_nodes = []
+
+                tasks = []
+                for node in nodes:
+                    tasks.append(asyncio.create_task(self._run_node_with_callbacks(node)))
+
+                await asyncio.gather(*tasks)
+
+        except Exception as e:
+            error_msg = f"Error in run_after_trigger: {str(e)}"
+            logger.error(error_msg)
+            raise e
+
+        if len(self.output_ports) > 0:
+            if len(self.output_ports) == 1:
+                if self.output_ports[0].is_prepared:
+                    result = self.output_ports[0].value
+                else:
+                    result = None
+            else:
+                result = {}
+                for port in self.output_ports:
+                    if port.is_prepared:
+                        result[port.name] = port.value
+            await self.call_callbacks(CallbackEvent.ON_WORKFLOW_END, workflow=self, output=result)
+        else:
+            await self.call_callbacks(CallbackEvent.ON_WORKFLOW_END, workflow=self, output=None)
+
+    async def _run(self, task_id: uuid.UUID, trigger_node: Trigger) -> None:
+        async with self.run_semaphore:
+            try:
+                new_workflow = self._clone(task_id, trigger_node)
+                await new_workflow.run_after_trigger()
+                # TODO: clear new_workflow
+
+            except Exception as e:
+                error_msg = f"Error running workflow: {str(e)}, {traceback.format_exc()}"
+                logger.error(error_msg)
+
+    async def run(self):
+        tasks = []
+        trigger = self.get_trigger_node()
+
+        async for task_id in trigger.run():
+            tasks.append(asyncio.create_task(self._run(task_id, trigger)))
+
+        if tasks:
+            await asyncio.gather(*tasks)
+
+    async def stop(self):
+        trigger = self.get_trigger_node()
+        await trigger._stop()
+
+    def trigger(self, trigger_name: str, **kwargs) -> uuid.UUID:
+        trigger = self.get_trigger_node()
+        task_id = uuid.uuid4()
+        for key, value in kwargs.items():
+            trigger.prepare_output_edges(key, value)
+        task = asyncio.create_task(self._run(task_id, trigger))
+        return task_id
+
+    async def handle_node_stream_output(
+        self,
+        node: Node,
+        stream: AsyncIterator[Any],
+    ) -> None:
+        async for data in stream:
+            await self.call_callbacks(CallbackEvent.ON_NODE_STREAM_OUTPUT, node=node, output=data)
+
+    # TODO: refactor this
+    async def handle_query_user(self, node_name: str, prompt: str) -> Awaitable[str]:
+        return await asyncio.to_thread(input, f"[{node_name}] {prompt}: ")
+    
+    def _clone(self, task_id: uuid.UUID, trigger_node: Trigger) -> Workflow:
+        return workflow_clone(self, task_id, trigger_node)

service_forge/workflow/workflow_callback.py

@@ -0,0 +1,141 @@
+from __future__ import annotations
+
+import json
+from abc import abstractmethod
+from typing import TYPE_CHECKING, Any
+from typing_extensions import override
+from enum import Enum
+from .workflow_event import WorkflowResult
+from loguru import logger
+
+if TYPE_CHECKING:
+    from .node import Node
+    from .workflow import Workflow
+
+class CallbackEvent(Enum):
+    ON_WORKFLOW_START = "on_workflow_start"
+    ON_WORKFLOW_END = "on_workflow_end"
+    ON_WORKFLOW_ERROR = "on_workflow_error"
+    ON_NODE_START = "on_node_start"
+    ON_NODE_END = "on_node_end"
+    ON_NODE_OUTPUT = "on_node_output"
+    ON_NODE_STREAM_OUTPUT = "on_node_stream_output"
+
+class WorkflowCallback:
+    @abstractmethod
+    async def on_workflow_start(self, workflow: Workflow) -> None:
+        ...
+
+    @abstractmethod
+    async def on_workflow_end(self, workflow: Workflow, output: Any) -> None:
+        pass
+
+    @abstractmethod
+    async def on_workflow_error(self, workflow: Workflow, error: Any) -> None:
+        pass
+
+    @abstractmethod
+    async def on_node_start(self, node: Node) -> None:
+        ...
+
+    @abstractmethod
+    async def on_node_end(self, node: Node) -> None:
+        ...
+
+    @abstractmethod
+    async def on_node_stream_output(self, node: Node, output: Any) -> None:
+        ...
+
+class BuiltinWorkflowCallback(WorkflowCallback):
+    def __init__(self):
+        self._websocket_manager = None
+    
+    def _get_websocket_manager(self):
+        if self._websocket_manager is None:
+            from service_forge.api.routers.websocket.websocket_manager import websocket_manager
+            self._websocket_manager = websocket_manager
+        return self._websocket_manager
+    
+    def _serialize_result(self, result: Any) -> Any:
+        try:
+            json.dumps(result)
+            return result
+        except (TypeError, ValueError):
+            return str(result)
+    
+    @override
+    async def on_workflow_start(self, workflow: Workflow) -> None:
+        ...
+
+    @override
+    async def on_workflow_end(self, workflow: Workflow, output: Any) -> None:
+        workflow_result = WorkflowResult(result=output, is_end=True, is_error=False)
+        
+        if workflow.task_id in workflow.real_trigger_node.result_queues:
+            workflow.real_trigger_node.result_queues[workflow.task_id].put_nowait(workflow_result)
+        if workflow.task_id in workflow.real_trigger_node.stream_queues:
+            workflow.real_trigger_node.stream_queues[workflow.task_id].put_nowait(workflow_result)
+        
+        try:
+            manager = self._get_websocket_manager()
+            message = {
+                "type": "workflow_end",
+                "task_id": str(workflow.task_id),
+                "result": self._serialize_result(output),
+                "is_end": True,
+                "is_error": False
+            }
+            await manager.send_to_task(workflow.task_id, message)
+        except Exception as e:
+            logger.error(f"发送 workflow_end 消息到 websocket 失败: {e}")
+
+    @override
+    async def on_workflow_error(self, workflow: Workflow, error: Any) -> None:
+        workflow_result = WorkflowResult(result=error, is_end=False, is_error=True)
+        
+        if workflow.task_id in workflow.real_trigger_node.result_queues:
+            workflow.real_trigger_node.result_queues[workflow.task_id].put_nowait(workflow_result)
+        if workflow.task_id in workflow.real_trigger_node.stream_queues:
+            workflow.real_trigger_node.stream_queues[workflow.task_id].put_nowait(workflow_result)
+        
+        try:
+            manager = self._get_websocket_manager()
+            message = {
+                "type": "workflow_error",
+                "task_id": str(workflow.task_id),
+                "error": self._serialize_result(error),
+                "is_end": False,
+                "is_error": True
+            }
+            await manager.send_to_task(workflow.task_id, message)
+        except Exception as e:
+            logger.error(f"发送 workflow_error 消息到 websocket 失败: {e}")
+
+    @override
+    async def on_node_start(self, node: Node) -> None:
+        ...
+
+    @override
+    async def on_node_end(self, node: Node) -> None:
+        ...
+
+    @override
+    async def on_node_stream_output(self, node: Node, output: Any) -> None:
+        workflow_result = WorkflowResult(result=output, is_end=False, is_error=False)
+        
+        if node.workflow.task_id in node.workflow.real_trigger_node.stream_queues:
+            node.workflow.real_trigger_node.stream_queues[node.workflow.task_id].put_nowait(workflow_result)
+        
+        try:
+            manager = self._get_websocket_manager()
+            message = {
+                "type": "node_stream_output",
+                "task_id": str(node.workflow.task_id),
+                "node": node.name,
+                "output": self._serialize_result(output),
+                "is_end": False,
+                "is_error": False
+            }
+            await manager.send_to_task(node.workflow.task_id, message)
+        except Exception as e:
+            logger.error(f"发送 node_stream_output 消息到 websocket 失败: {e}")

service_forge/workflow/workflow_config.py

@@ -0,0 +1,66 @@
+from __future__ import annotations
+
+import yaml
+from typing import Any
+from pydantic import BaseModel
+
+DEFAULT_WORKFLOW_VERSION = "0"
+
+class WorkflowNodeOutputConfig(BaseModel):
+    name: str
+    port: str
+
+class WorkflowNodeArgConfig(BaseModel):
+    name: str
+    value: Any | None = None
+
+class WorkflowInputPortConfig(BaseModel):
+    name: str
+    port: str
+    value: Any | None = None
+
+class WorkflowOutputPortConfig(BaseModel):
+    name: str
+    port: str
+
+class WorkflowNodeSubWorkflowConfig(BaseModel):
+    name: str
+    version: str = DEFAULT_WORKFLOW_VERSION
+
+class WorkflowNodeSubWorkflowInputPortConfig(BaseModel):
+    name: str
+    port: str
+    value: Any | None = None
+
+class WorkflowNodeConfig(BaseModel):
+    name: str
+    type: str
+    args: dict | None = None
+    outputs: dict | None = None
+    sub_workflows: list[WorkflowNodeSubWorkflowConfig] | None = None
+    sub_workflows_input_ports: list[WorkflowNodeSubWorkflowInputPortConfig] | None = None
+
+class WorkflowConfig(BaseModel):
+    name: str
+    version: str = DEFAULT_WORKFLOW_VERSION
+    description: str
+    nodes: list[WorkflowNodeConfig]
+    inputs: list[WorkflowInputPortConfig] | None = None
+    outputs: list[WorkflowOutputPortConfig] | None = None
+
+    @classmethod
+    def from_yaml_file(cls, filepath: str) -> WorkflowConfig:
+        with open(filepath, 'r', encoding='utf-8') as f:
+            data = yaml.safe_load(f)
+        return cls(**data)
+
+class WorkflowGroupConfig(BaseModel):
+    workflows: list[WorkflowConfig] | None
+    main_workflow_name: str | None
+    main_workflow_version: str | None = DEFAULT_WORKFLOW_VERSION
+
+    @classmethod
+    def from_yaml_file(cls, filepath: str) -> WorkflowGroupConfig:
+        with open(filepath, 'r', encoding='utf-8') as f:
+            data = yaml.safe_load(f)
+        return cls(**data)

service_forge/workflow/workflow_event.py

@@ -0,0 +1,15 @@
+from __future__ import annotations
+
+from typing import Any
+
+class WorkflowResult:
+    def __init__(
+        self,
+        result: Any | None,
+        is_end: bool,
+        is_error: bool,
+    ) -> None:
+        # when is_end is True, result is from the output port of the workflow
+        self.result = result
+        self.is_end = is_end
+        self.is_error = is_error