service-forge 0.1.11__py3-none-any.whl

service_forge/workflow/triggers/fast_api_trigger.py

@@ -0,0 +1,169 @@
+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, Request
+from fastapi.responses import StreamingResponse
+from service_forge.workflow.port import Port
+from service_forge.utils.default_type_converter import type_converter
+from service_forge.api.routers.websocket.websocket_manager import websocket_manager
+from fastapi import HTTPException
+from google.protobuf.message import Message
+from google.protobuf.json_format import MessageToJson
+
+class FastAPITrigger(Trigger):
+    DEFAULT_INPUT_PORTS = [
+        Port("app", FastAPI),
+        Port("path", str),
+        Port("method", str),
+        Port("data_type", type),
+        Port("is_stream", bool),
+    ]
+
+    DEFAULT_OUTPUT_PORTS = [
+        Port("trigger", bool),
+        Port("user_id", int),
+        Port("data", Any),
+    ]
+
+    def __init__(self, name: str):
+        super().__init__(name)
+        self.events = {}
+        self.is_setup_route = False
+
+    @staticmethod
+    def serialize_result(result: Any):
+        if isinstance(result, Message):
+            return MessageToJson(
+                result,
+                preserving_proto_field_name=True
+            )
+        return result
+
+    async def handle_request(
+        self,
+        request: Request,
+        data_type: type,
+        extract_data_fn: callable[[Request], dict],
+        is_stream: bool,
+    ):
+        task_id = uuid.uuid4()
+        self.result_queues[task_id] = asyncio.Queue()
+
+        body_data = await extract_data_fn(request)
+        converted_data = data_type(**body_data)
+
+        client_id = (
+            body_data.get("client_id")
+            or request.query_params.get("client_id")
+            or request.headers.get("X-Client-ID")
+        )
+        if client_id:
+            workflow_name = getattr(self.workflow, "name", "Unknown")
+            steps = len(self.workflow.nodes) if hasattr(self.workflow, "nodes") else 1
+            websocket_manager.create_task_with_client(task_id, client_id, workflow_name, steps)
+
+        self.trigger_queue.put_nowait({
+            "id": task_id,
+            "user_id": getattr(request.state, "user_id", None),
+            "data": converted_data,
+        })
+
+        if is_stream:
+            self.stream_queues[task_id] = asyncio.Queue()
+
+            async def generate_sse():
+                try:
+                    while True:
+                        item = await self.stream_queues[task_id].get()
+
+                        if item.is_error:
+                            yield f"event: error\ndata: {json.dumps({'detail': str(item.result)})}\n\n"
+                            break
+                        
+                        if item.is_end:
+                            # TODO: send the result?
+                            break
+
+                        # TODO: modify
+                        serialized = self.serialize_result(item.result)
+                        if isinstance(serialized, str):
+                            data = serialized
+                        else:
+                            data = json.dumps(serialized)
+                        
+                        yield f"data: {data}\n\n"
+                    
+                except Exception as e:
+                    yield f"event: error\ndata: {json.dumps({'detail': str(e)})}\n\n"
+                finally:
+                    if task_id in self.stream_queues:
+                        del self.stream_queues[task_id]
+                    if task_id in self.result_queues:
+                        del self.result_queues[task_id]
+            
+            return StreamingResponse(
+                generate_sse(),
+                media_type="text/event-stream",
+                headers={
+                    "Cache-Control": "no-cache",
+                    "Connection": "keep-alive",
+                    "X-Accel-Buffering": "no",
+                }
+            )
+        else:
+            result = await self.result_queues[task_id].get()
+            del self.result_queues[task_id]
+
+            if result.is_error:
+                if isinstance(result.result, HTTPException):
+                    raise result.result
+                else:
+                    raise HTTPException(status_code=500, detail=str(result.result))
+
+            return self.serialize_result(result.result)
+
+    def _setup_route(self, app: FastAPI, path: str, method: str, data_type: type, is_stream: bool) -> None:
+        async def get_data(request: Request) -> dict:
+            return dict(request.query_params)
+
+        async def body_data(request: Request) -> dict:
+            raw = await request.body()
+            if not raw:
+                return {}
+            return json.loads(raw.decode("utf-8"))
+
+        extractor = get_data if method == "GET" else body_data
+
+        async def handler(request: Request):
+            return await self.handle_request(request, data_type, extractor, is_stream)
+
+        if method == "GET":
+            app.get(path)(handler)
+        elif method == "POST":
+            app.post(path)(handler)
+        elif method == "PUT":
+            app.put(path)(handler)
+        elif method == "DELETE":
+            app.delete(path)(handler)
+        else:
+            raise ValueError(f"Invalid method {method}")
+
+    async def _run(self, app: FastAPI, path: str, method: str, data_type: type, is_stream: bool = False) -> AsyncIterator[bool]:
+        if not self.is_setup_route:
+            self._setup_route(app, path, method, data_type, is_stream)
+            self.is_setup_route = True
+
+        while True:
+            try:
+                trigger = await self.trigger_queue.get()
+                self.prepare_output_edges(self.get_output_port_by_name('user_id'), trigger['user_id'])
+                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 FastAPITrigger._run: {e}")
+                continue
+

service_forge/workflow/triggers/kafka_api_trigger.py

@@ -0,0 +1,44 @@
+from __future__ import annotations
+import uuid
+from typing import Any
+from service_forge.workflow.trigger import Trigger
+from typing import AsyncIterator
+from service_forge.workflow.port import Port
+from service_forge.api.kafka_api import KafkaApp
+
+class KafkaAPITrigger(Trigger):
+    DEFAULT_INPUT_PORTS = [
+        Port("app", KafkaApp),
+        Port("topic", str),
+        Port("data_type", type),
+        Port("group_id", str),
+    ]
+
+    DEFAULT_OUTPUT_PORTS = [
+        Port("trigger", bool),
+        Port("data", Any),
+    ]
+    
+    def __init__(self, name: str):
+        super().__init__(name)
+        self.events = {}
+        self.is_setup_kafka_input = False
+
+    def _setup_kafka_input(self, app: KafkaApp, topic: str, data_type: type, group_id: str) -> None:
+        @app.kafka_input(topic, data_type, group_id)
+        async def handle_message(data):
+            task_id = uuid.uuid4()
+            self.trigger_queue.put_nowait({
+                "id": task_id,
+                "data": data,
+            })
+
+    async def _run(self, app: KafkaApp, topic: str, data_type: type, group_id: str) -> AsyncIterator[bool]:
+        if not self.is_setup_kafka_input:
+            self._setup_kafka_input(app, topic, data_type, group_id)
+            self.is_setup_kafka_input = True
+
+        while True:
+            trigger = await self.trigger_queue.get()
+            self.prepare_output_edges(self.get_output_port_by_name('data'), trigger['data'])
+            yield self.trigger(trigger['id'])

service_forge/workflow/triggers/once_trigger.py

@@ -0,0 +1,20 @@
+from __future__ import annotations
+from service_forge.workflow.node import Node
+from service_forge.workflow.port import Port
+from service_forge.workflow.trigger import Trigger
+from typing import AsyncIterator
+import uuid
+
+class OnceTrigger(Trigger):
+    DEFAULT_INPUT_PORTS = [
+    ]
+
+    DEFAULT_OUTPUT_PORTS = [
+        Port("trigger", bool),
+    ]
+
+    def __init__(self, name: str):
+        super().__init__(name)
+
+    async def _run(self) -> AsyncIterator[bool]:
+        yield self.trigger(uuid.uuid4())

service_forge/workflow/triggers/period_trigger.py

@@ -0,0 +1,26 @@
+from __future__ import annotations
+import asyncio
+from service_forge.workflow.node import Node
+from service_forge.workflow.port import Port
+from service_forge.workflow.trigger import Trigger
+from typing import AsyncIterator
+import uuid
+
+class PeriodTrigger(Trigger):
+    DEFAULT_INPUT_PORTS = [
+        Port("TRIGGER", bool),
+        Port("period", float),
+        Port("times", int),
+    ]
+
+    DEFAULT_OUTPUT_PORTS = [
+        Port("trigger", bool),
+    ]
+
+    def __init__(self, name: str):
+        super().__init__(name)
+
+    async def _run(self, times: int, period: float) -> AsyncIterator[bool]:
+        for _ in range(times):
+            await asyncio.sleep(period)
+            yield uuid.uuid4()

service_forge/workflow/triggers/websocket_api_trigger.py

@@ -0,0 +1,184 @@
+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()
+
+        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
+

service_forge/workflow/workflow.py

@@ -0,0 +1,210 @@
+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
+
+class Workflow:
+    def __init__(
+        self,
+        name: str,
+        description: str,
+        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.name = name
+        self.description = description
+        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()
+
+    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)
+
+    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)