service-forge 0.1.18__py3-none-any.whl

service_forge/workflow/node.py

@@ -0,0 +1,184 @@
+from __future__ import annotations
+
+from typing import Any, AsyncIterator, Union, TYPE_CHECKING, Callable, Awaitable
+from abc import ABC, abstractmethod
+import uuid
+from loguru import logger
+from .edge import Edge
+from .port import Port
+from .context import Context
+from ..utils.register import Register
+from ..db.database import DatabaseManager, PostgresDatabase, MongoDatabase, RedisDatabase
+from ..utils.workflow_clone import node_clone
+
+if TYPE_CHECKING:
+    from .workflow import Workflow
+
+class Node(ABC):
+    DEFAULT_INPUT_PORTS: list[Port] = []
+    DEFAULT_OUTPUT_PORTS: list[Port] = []
+
+    CLASS_NOT_REQUIRED_TO_REGISTER = ['Node']
+    AUTO_FILL_INPUT_PORTS = []
+
+    def __init__(
+        self,
+        name: str,
+        context: Context = None,
+        input_edges: list[Edge] = None,
+        output_edges: list[Edge] = None,
+        input_ports: list[Port] = DEFAULT_INPUT_PORTS,
+        output_ports: list[Port] = DEFAULT_OUTPUT_PORTS,
+        query_user: Callable[[str, str], Awaitable[str]] = None,
+    ) -> None:
+        from .workflow_group import WorkflowGroup
+        self.name = name
+        self.input_edges = [] if input_edges is None else input_edges
+        self.output_edges = [] if output_edges is None else output_edges
+        self.input_ports = input_ports
+        self.output_ports = output_ports
+        self.workflow: Workflow = None
+        self.query_user = query_user
+        self.sub_workflows: WorkflowGroup = None
+
+        # runtime variables
+        self.context = context
+        self.input_variables: dict[Port, Any] = {}
+        self.num_activated_input_edges = 0
+
+    @property
+    def default_postgres_database(self) -> PostgresDatabase | None:
+        return self.database_manager.get_default_postgres_database()
+
+    @property
+    def default_mongo_database(self) -> MongoDatabase | None:
+        return self.database_manager.get_default_mongo_database()
+
+    @property
+    def default_redis_database(self) -> RedisDatabase | None:
+        return self.database_manager.get_default_redis_database()
+    
+    @property
+    def database_manager(self) -> DatabaseManager:
+        return self.workflow.database_manager
+
+    def backup(self) -> None:
+        # do NOT use deepcopy here
+        # self.bak_context = deepcopy(self.context)
+        # TODO: what if the value changes after backup?
+        self.bak_input_variables = {port: value for port, value in self.input_variables.items()}
+        self.bak_num_activated_input_edges = self.num_activated_input_edges
+
+    def reset(self) -> None:
+        # self.context = deepcopy(self.bak_context)
+        self.input_variables = {port: value for port, value in self.bak_input_variables.items()}
+        self.num_activated_input_edges = self.bak_num_activated_input_edges
+
+    def __init_subclass__(cls) -> None:
+        if cls.__name__ not in Node.CLASS_NOT_REQUIRED_TO_REGISTER:
+            node_register.register(cls.__name__, cls)
+        return super().__init_subclass__()
+
+    def _query_user(self, prompt: str) -> Callable[[str, str], Awaitable[str]]:
+        return self.query_user(self.name, prompt)
+    
+    def variables_to_params(self) -> dict[str, Any]:
+        params = {port.name: self.input_variables[port] for port in self.input_variables.keys() if not port.is_extended_generated}
+        for port in self.input_variables.keys():
+            if port.is_extended_generated:
+                if port.get_extended_name() not in params:
+                    params[port.get_extended_name()] = []
+                params[port.get_extended_name()].append((port.get_extended_index(), self.input_variables[port]))
+                params[port.get_extended_name()].sort()
+        return params
+
+    def is_trigger(self) -> bool:
+        from .trigger import Trigger
+        return isinstance(self, Trigger)
+
+    # TODO: maybe add a function before the run function?
+
+    @abstractmethod
+    async def _run(self, **kwargs) -> Union[None, AsyncIterator]:
+        ...
+    
+    def run(self) -> Union[None, AsyncIterator]:
+        for key in list(self.input_variables.keys()):
+            if key and key.name[0].isupper():
+                del self.input_variables[key]
+        params = self.variables_to_params()
+        return self._run(**params)
+
+    def get_input_port_by_name(self, name: str) -> Port:
+        # TODO: add warning if port is extended
+        for port in self.input_ports:
+            if port.name == name:
+                return port
+        return None
+
+    def get_output_port_by_name(self, name: str) -> Port:
+        # TODO: add warning if port is extended
+        for port in self.output_ports:
+            if port.name == name:
+                return port
+        return None
+
+    def try_create_extended_input_port(self, name: str) -> None:
+        for port in self.input_ports:
+            if port.is_extended and name.startswith(port.name + '_') and name[len(port.name + '_'):].isdigit():
+                self.input_ports.append(Port(name=name, type=port.type, node=port.node, port=port.port, value=port.value, default=port.default, is_extended=False, is_extended_generated=True))
+
+    def try_create_extended_output_port(self, name: str) -> None:
+        for port in self.output_ports:
+            if port.is_extended and name.startswith(port.name + '_') and name[len(port.name + '_'):].isdigit():
+                self.output_ports.append(Port(name=name, type=port.type, node=port.node, port=port.port, value=port.value, default=port.default, is_extended=False, is_extended_generated=True))
+
+    def num_input_ports(self) -> int:
+        return sum(1 for port in self.input_ports if not port.is_extended)
+    
+    def is_ready(self) -> bool:
+        return self.num_activated_input_edges == self.num_input_ports()
+
+    def fill_input_by_name(self, port_name: str, value: Any) -> None:
+        self.try_create_extended_input_port(port_name)
+        port = self.get_input_port_by_name(port_name)
+        if port is None:
+            raise ValueError(f'{port_name} is not a valid input port.')
+        self.fill_input(port, value)
+    
+    def fill_input(self, port: Port, value: Any) -> None:
+        port.activate(value)
+
+    def activate_output_edges(self, port: str | Port, data: Any) -> None:
+        if isinstance(port, str):
+            port = self.get_output_port_by_name(port)
+        for output_edge in self.output_edges:
+            if output_edge.start_port == port:
+                output_edge.end_port.activate(data)
+
+    # for trigger nodes
+    def prepare_output_edges(self, port: Port, data: Any) -> None:
+        if isinstance(port, str):
+            port = self.get_output_port_by_name(port)
+        for output_edge in self.output_edges:
+            if output_edge.start_port == port:
+                output_edge.end_port.prepare(data)
+    
+    def trigger_output_edges(self, port: Port) -> None:
+        if isinstance(port, str):
+            port = self.get_output_port_by_name(port)
+        for output_edge in self.output_edges:
+            if output_edge.start_port == port:
+                output_edge.end_port.trigger()
+
+    # TODO: the result is outputed to the trigger now, maybe we should add a new function to output the result to the workflow
+    def output_to_workflow(self, data: Any) -> None:
+        self.workflow._handle_workflow_output(self.name, data)
+
+    def extended_output_name(self, name: str, index: int) -> str:
+        return name + '_' + str(index)
+
+    def _clone(self, context: Context) -> Node:
+        return node_clone(self, context)
+
+node_register = Register[Node]()

service_forge/workflow/nodes/__init__.py

@@ -0,0 +1,8 @@
+from .output.print_node import PrintNode
+from .output.kafka_output_node import KafkaOutputNode
+from .control.if_node import IfNode
+from .control.switch_node import SwitchNode
+from .llm.query_llm_node import QueryLLMNode
+from .test.if_console_input_node import IfConsoleInputNode
+from .nested.workflow_node import WorkflowNode
+from .test.time_consuming_node import TimeConsumingNode

service_forge/workflow/nodes/control/if_node.py

@@ -0,0 +1,29 @@
+from __future__ import annotations
+from ...node import Node
+from ...port import Port
+
+class IfNode(Node):
+    DEFAULT_INPUT_PORTS = [
+        Port("TRIGGER", bool),
+        Port("condition", str),
+    ]
+
+    DEFAULT_OUTPUT_PORTS = [
+        Port("true", bool),
+        Port("false", bool),
+    ]
+
+    def __init__(
+        self, 
+        name: str,
+    ) -> None:
+        super().__init__(
+            name,
+        )
+
+    async def _run(self, condition: str) -> None:
+        result = eval(condition)
+        if result:
+            self.activate_output_edges('true', True)
+        else:
+            self.activate_output_edges('false', False)

service_forge/workflow/nodes/control/switch_node.py

@@ -0,0 +1,28 @@
+from __future__ import annotations
+from typing import Any
+from ...node import Node
+from ...port import Port
+
+class SwitchNode(Node):
+    DEFAULT_INPUT_PORTS = [
+        Port("TRIGGER", bool),
+        Port("condition", str, is_extended=True),
+    ]
+
+    DEFAULT_OUTPUT_PORTS = [
+        Port("result", Any, is_extended=True),
+    ]
+
+    def __init__(
+        self, 
+        name: str,
+    ) -> None:
+        super().__init__(
+            name,
+        )
+
+    async def _run(self, condition: list[tuple[int, str]]) -> None:
+        for index, cond in condition:
+            if eval(cond):
+                self.activate_output_edges(self.extended_output_name('result', index), str(index))
+                break

service_forge/workflow/nodes/input/console_input_node.py

@@ -0,0 +1,26 @@
+from __future__ import annotations
+from service_forge.workflow.node import Node
+from service_forge.workflow.port import Port
+from service_forge.llm import chat_stream, Model
+
+class ConsoleInputNode(Node):
+    DEFAULT_INPUT_PORTS = [
+        Port("TRIGGER", bool),
+        Port("prompt", str),
+    ]
+
+    DEFAULT_OUTPUT_PORTS = [
+        Port("user_input", str),
+    ]
+
+    AUTO_FILL_INPUT_PORTS = []
+
+    def __init__(
+        self, 
+        name: str,
+    ) -> None:
+        super().__init__(name)
+
+    async def _run(self, prompt: str) -> None:
+        user_input = self._query_user(prompt)
+        self.activate_output_edges(self.get_output_port_by_name('user_input'), user_input)

service_forge/workflow/nodes/llm/query_llm_node.py

@@ -0,0 +1,41 @@
+from __future__ import annotations
+import os
+
+from service_forge.workflow.node import Node
+from service_forge.workflow.port import Port
+from service_forge.llm import chat_stream, Model
+
+class QueryLLMNode(Node):
+    DEFAULT_INPUT_PORTS = [
+        Port("prompt", str),
+        Port("system_prompt", str),
+        Port("temperature", float),
+        Port("TRIGGER", bool),
+    ]
+
+    DEFAULT_OUTPUT_PORTS = [
+        Port("response", str),
+    ]
+
+    AUTO_FILL_INPUT_PORTS = [('TRIGGER', True)]
+
+    def __init__(
+        self, 
+        name: str,
+    ) -> None:
+        super().__init__(
+            name,
+        )
+
+    async def _run(self, prompt: str, system_prompt: str, temperature: float) -> None:
+        if os.path.exists(system_prompt):
+            with open(system_prompt, "r") as f:
+                system_prompt = f.read()
+        if os.path.exists(prompt):
+            with open(prompt, "r") as f:
+                prompt = f.read()
+
+        print(f"prompt: {prompt} temperature: {temperature}")
+        response = chat_stream(prompt, system_prompt, Model.GEMINI, temperature)
+        for chunk in response:
+            yield chunk

service_forge/workflow/nodes/nested/workflow_node.py

@@ -0,0 +1,28 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+from service_forge.workflow.node import Node
+from service_forge.workflow.port import Port
+
+class WorkflowNode(Node):
+    from service_forge.workflow.workflow import Workflow
+    DEFAULT_INPUT_PORTS = [
+        Port("workflow", Workflow),
+    ]
+
+    DEFAULT_OUTPUT_PORTS = [
+    ]
+
+    def __init__(
+        self, 
+        name: str,
+    ) -> None:
+        super().__init__(
+            name,
+        )
+
+    async def _run(self, workflow: Workflow, **kwargs) -> None:
+        for input_port in self.input_ports:
+            if input_port.is_sub_workflow_input_port():
+                input_port.port.node.fill_input(input_port.port, input_port.value)
+        await workflow.run()

service_forge/workflow/nodes/output/kafka_output_node.py

@@ -0,0 +1,27 @@
+from __future__ import annotations
+from service_forge.workflow.node import Node
+from service_forge.workflow.port import Port
+from service_forge.api.kafka_api import KafkaApp
+from typing import Any
+
+class KafkaOutputNode(Node):
+    DEFAULT_INPUT_PORTS = [
+        Port("app", KafkaApp),
+        Port("topic", str),
+        Port("data_type", type),
+        Port("data", Any),
+    ]
+
+    DEFAULT_OUTPUT_PORTS = [
+    ]
+
+    def __init__(
+        self, 
+        name: str,
+    ) -> None:
+        super().__init__(
+            name,
+        )
+
+    async def _run(self, app: KafkaApp, topic: str, data_type: type, data: Any) -> None:
+        await app.send_message(topic, data_type, data)

service_forge/workflow/nodes/output/print_node.py

@@ -0,0 +1,29 @@
+from __future__ import annotations
+import asyncio
+from typing import AsyncIterator
+from service_forge.workflow.node import Node
+from service_forge.workflow.port import Port
+
+class PrintNode(Node):
+    DEFAULT_INPUT_PORTS = [
+        Port("TRIGGER", bool),
+        Port("message", str),
+    ]
+
+    DEFAULT_OUTPUT_PORTS = [
+    ]
+
+    AUTO_FILL_INPUT_PORTS = [('TRIGGER', True)]
+
+    def __init__(
+        self, 
+        name: str,
+    ) -> None:
+        super().__init__(
+            name,
+        )
+
+    async def _run(self, message: str) -> AsyncIterator[str]:
+        for char in str(message):
+            await asyncio.sleep(0.1)
+            yield char

service_forge/workflow/nodes/test/if_console_input_node.py

@@ -0,0 +1,33 @@
+from __future__ import annotations
+from typing import AsyncIterator
+from service_forge.workflow.node import Node
+from service_forge.workflow.port import Port
+
+class IfConsoleInputNode(Node):
+    DEFAULT_INPUT_PORTS = [
+        Port("TRIGGER", bool),
+        Port("condition", str),
+    ]
+
+    DEFAULT_OUTPUT_PORTS = [
+        Port("true", bool),
+        Port("false", bool),
+    ]
+
+    def __init__(
+        self, 
+        name: str,
+    ) -> None:
+        super().__init__(
+            name,
+        )
+
+    async def _run(self, condition: str) -> None:
+        while True:
+            user_input = await self._query_user(condition)
+            if user_input.lower() in ['y', 'yes']:
+                self.activate_output_edges(self.get_output_port_by_name('true'), True)
+                break
+            elif user_input.lower() in ['n', 'no']:
+                self.activate_output_edges(self.get_output_port_by_name('false'), False)
+                break

service_forge/workflow/nodes/test/time_consuming_node.py

@@ -0,0 +1,62 @@
+
+from __future__ import annotations
+import asyncio
+import uuid
+import json
+from typing import Any
+from ...node import Node
+from ...port import Port
+from ....api.routers.websocket.websocket_manager import websocket_manager
+
+# It's deprecated, just for testing
+class TimeConsumingNode(Node):
+    """模拟耗时节点，定期发送进度更新"""
+
+    DEFAULT_INPUT_PORTS = [
+        Port("input", Any)
+    ]
+
+    DEFAULT_OUTPUT_PORTS = [
+        Port("output", Any)
+    ]
+
+    def __init__(self, name: str, duration: float = 2.0):
+        super().__init__(name)
+        self.duration = duration  # 总耗时（秒）
+        self.progress = 0.0
+        self.task_id = None
+
+    async def _run(self, input: Any = None, task_id: uuid.UUID = None) -> str:
+        """执行耗时任务，定期更新进度"""
+        # 保存任务ID（如果有）
+        if task_id is not None:
+            self.task_id = task_id
+
+        total_steps = 10
+        result = f"Completed {self.name} after {self.duration} seconds"
+
+        # 分步骤执行，每步更新进度
+        for i in range(total_steps + 1):
+            # 更新进度
+            self.progress = i / total_steps
+
+            # 发送进度更新
+            if self.task_id:
+                await websocket_manager.send_progress(
+                    self.task_id, 
+                    self.name, 
+                    self.progress
+                )
+
+            # 模拟耗时
+            if i < total_steps:  # 最后一步不需要等待
+                await asyncio.sleep(self.duration / total_steps)
+
+        # 获取输出端口并设置值
+        output_port = self.get_output_port_by_name('output')
+        output_port.prepare(result)
+
+        # 激活输出端口
+        self.activate_output_edges(output_port, result)
+
+        return result

service_forge/workflow/port.py

@@ -0,0 +1,89 @@
+from __future__ import annotations
+from typing import TYPE_CHECKING, Any
+from ..utils.workflow_clone import port_clone
+
+if TYPE_CHECKING:
+    from .node import Node
+
+class Port:
+    def __init__(
+        self,
+        name: str,
+        type: type,
+        node: Node = None,
+        port: Port = None,
+        value: Any = None,
+        default: Any = None,
+        is_extended: bool = False,
+        is_extended_generated: bool = False,
+    ) -> None:
+        self.name = name
+        self.type = type
+        self.node = node
+        self.port = port
+        self.value = value
+        # not used yet
+        self.default = default
+        self.is_prepared = False
+        self.is_extended = is_extended
+        self.is_extended_generated = is_extended_generated
+    
+    def is_sub_workflow_input_port(self) -> bool:
+        return self.port != None
+
+    def prepare(self, data: Any) -> None:
+        from ..utils.default_type_converter import type_converter
+        data = type_converter.convert(data, self.type, node=self.node)
+        self.value = data
+        self.is_prepared = True
+
+    def trigger(self) -> None:
+        if self.node is None:
+            return
+        if self in self.node.input_variables:
+            return
+        self.node.input_variables[self] = self.value
+        self.node.num_activated_input_edges += 1
+        if self.node.is_ready():
+            self.node.workflow.ready_nodes.append(self.node)
+
+    def activate(self, data: Any) -> None:
+        self.prepare(data)
+        self.trigger()
+
+    def get_extended_name(self) -> str:
+        if self.is_extended_generated:
+            return '_'.join(self.name.split('_')[:-1])
+        raise ValueError(f"Port {self.name} is not extended generated.")
+    
+    def get_extended_index(self) -> int:
+        if self.is_extended_generated:
+            return int(self.name.split('_')[-1])
+        raise ValueError(f"Port {self.name} is not extended generated.")
+
+    def _clone(self, node_map: dict[Node, Node]) -> Port:
+        return port_clone(self, node_map)
+
+# node port
+def create_port(name: str, type: type, node: Node = None, value: Any = None, port: Port = None) -> Port:
+    return Port(name, type, node, port, value)
+
+# workflow input port
+def create_workflow_input_port(name: str, port: Port, value: Any = None) -> Port:
+    if value is None:
+        value = port.value
+    return Port(name, port.type, port.node, port, value)
+
+# sub workflow input port
+# node is the node that the sub workflow is running on
+def create_sub_workflow_input_port(name: str, node: Node, port: Port, value: Any = None) -> Port:
+    if value is None:
+        value = port.value
+    return Port(name, port.type, node, port, value)
+
+PORT_DELIMITER = '|'
+
+def parse_port_name(port_name: str) -> tuple[str, str]:
+    if PORT_DELIMITER not in port_name or len(port_name.split(PORT_DELIMITER)) != 2:
+        raise ValueError(f"Invalid port name: {port_name}")
+    return port_name.split(PORT_DELIMITER)

service_forge/workflow/trigger.py

@@ -0,0 +1,28 @@
+from __future__ import annotations
+import asyncio
+from typing import AsyncIterator
+from abc import ABC, abstractmethod
+import uuid
+from .node import Node
+from .workflow_event import WorkflowResult
+
+class Trigger(Node, ABC):
+    def __init__(self, name: str):
+        super().__init__(name)
+        self.trigger_queue = asyncio.Queue()
+        # for workflow result
+        self.result_queues: dict[uuid.UUID, asyncio.Queue[WorkflowResult]] = {}
+        # for node stream output
+        self.stream_queues: dict[uuid.UUID, asyncio.Queue[WorkflowResult]] = {}
+
+    @abstractmethod
+    async def _run(self) -> AsyncIterator[bool]:
+        ...
+
+    @abstractmethod
+    async def _stop(self) -> AsyncIterator[bool]:
+        ...
+
+    def trigger(self, task_id: uuid.UUID) -> bool:
+        self.prepare_output_edges(self.get_output_port_by_name('trigger'), True)
+        return task_id

service_forge/workflow/triggers/__init__.py

@@ -0,0 +1,6 @@
+from .once_trigger import OnceTrigger
+from .period_trigger import PeriodTrigger
+from .fast_api_trigger import FastAPITrigger
+from .kafka_api_trigger import KafkaAPITrigger
+from .websocket_api_trigger import WebSocketAPITrigger
+from .a2a_api_trigger import A2AAPITrigger