service-forge 0.1.18__py3-none-any.whl → 0.1.39__py3-none-any.whl

service_forge/workflow/triggers/websocket_api_trigger.py

@@ -6,9 +6,11 @@ from loguru import logger
 from service_forge.workflow.trigger import Trigger
 from typing import AsyncIterator, Any
 from fastapi import FastAPI, WebSocket, WebSocketDisconnect
+from starlette.websockets import WebSocketState
 from service_forge.workflow.port import Port
 from google.protobuf.message import Message
 from google.protobuf.json_format import MessageToJson
+from service_forge.api.http_api import authenticate_websocket
 
 class WebSocketAPITrigger(Trigger):
     DEFAULT_INPUT_PORTS = [
@@ -19,6 +21,9 @@ class WebSocketAPITrigger(Trigger):
 
     DEFAULT_OUTPUT_PORTS = [
         Port("trigger", bool),
+        Port("client_id", uuid.UUID),
+        Port("user_id", int),
+        Port("token", str),
         Port("data", Any),
     ]
 
@@ -36,6 +41,28 @@ class WebSocketAPITrigger(Trigger):
             )
         return result
 
+    async def send_message(
+        self,
+        websocket: WebSocket,
+        type: str,
+        task_id: uuid.UUID,
+        data: Any,
+    ):
+        # Check if WebSocket is closed before sending
+        if websocket.client_state != WebSocketState.CONNECTED:
+            logger.warning(f"WebSocket is closed, cannot send message for task {task_id}")
+            return
+        
+        message = {
+            "type": type,
+            "task_id": str(task_id),
+            "data": data
+        }
+        try:
+            await websocket.send_text(json.dumps(message))
+        except Exception as e:
+            logger.error(f"Error sending message to WebSocket for task {task_id}: {e}")
+
     async def handle_stream_output(
         self,
         websocket: WebSocket,
@@ -46,12 +73,7 @@ class WebSocketAPITrigger(Trigger):
                 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))
+                    await self.send_message(websocket, "stream_error", task_id, str(item.result))
                     break
                 
                 if item.is_end:
@@ -65,18 +87,9 @@ class WebSocketAPITrigger(Trigger):
                                 data = serialized
                         else:
                             data = serialized
-                        
-                        end_response = {
-                            "type": "stream_end",
-                            "task_id": str(task_id),
-                            "data": data
-                        }
+                        await self.send_message(websocket, "stream_end", task_id, data)
                     else:
-                        end_response = {
-                            "type": "stream_end",
-                            "task_id": str(task_id)
-                        }
-                    await websocket.send_text(json.dumps(end_response))
+                        await self.send_message(websocket, "stream_end", task_id, None)
                     break
 
                 # Send stream data
@@ -89,23 +102,10 @@ class WebSocketAPITrigger(Trigger):
                 else:
                     data = serialized
                 
-                stream_response = {
-                    "type": "stream",
-                    "task_id": str(task_id),
-                    "data": data
-                }
-                await websocket.send_text(json.dumps(stream_response))
+                await self.send_message(websocket, "stream", task_id, data)
         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
+            await self.send_message(websocket, "stream_error", task_id, str(e))
         finally:
             if task_id in self.stream_queues:
                 del self.stream_queues[task_id]
@@ -114,16 +114,20 @@ class WebSocketAPITrigger(Trigger):
         self,
         websocket: WebSocket,
         data_type: type,
+        client_id: str,
         message_data: dict,
     ):
         task_id = uuid.uuid4()
-        self.result_queues[task_id] = asyncio.Queue()
+        # self.result_queues[task_id] = asyncio.Queue()
         self.stream_queues[task_id] = asyncio.Queue()
 
+        logger.info(f'user_id {getattr(websocket.state, "user_id", None)} token {getattr(websocket.state, "auth_token", None)}')
+
         if data_type is Any:
             converted_data = message_data
         else:
             try:
+                # TODO: message_data is Message, need to convert to dict
                 converted_data = data_type(**message_data)
             except Exception as e:
                 error_msg = {"error": f"Failed to convert data: {str(e)}"}
@@ -135,27 +139,46 @@ class WebSocketAPITrigger(Trigger):
 
         self.trigger_queue.put_nowait({
             "id": task_id,
+            "user_id": getattr(websocket.state, "user_id", None),
+            "token": getattr(websocket.state, "auth_token", None),
             "data": converted_data,
+            "client_id": client_id,
         })
 
         # 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):
+            websocket.state.user_id = websocket.headers.get("X-User-ID") or "0"
+            websocket.state.auth_token = websocket.headers.get("X-User-Token") or ""
+            logger.info(f'user_id {websocket.state.user_id} token {websocket.state.auth_token}')
+            # Authenticate WebSocket connection before accepting
+            # Get trusted_domain from app.state if available
+            # trusted_domain = getattr(app.state, "trusted_domain", "ring.shiweinan.com")
+            # enable_auth = getattr(app.state, "enable_auth_middleware", True)
+            
+            # if enable_auth:
+            #     await authenticate_websocket(websocket, trusted_domain)
+            # else:
+            #     # If auth is disabled, set default values
+            #     websocket.state.user_id = websocket.headers.get("X-User-ID", "0")
+            #     websocket.state.auth_token = websocket.headers.get("X-User-Token")
+            
             await websocket.accept()
+
+            client_id = uuid.uuid4()
             
             try:
                 while True:
-                    # Receive message from client
-                    data = await websocket.receive_text()
+                    data = await websocket.receive()
                     try:
-                        message = json.loads(data)
-                        
+                        # message = json.loads(data)
                         # Handle the message and trigger workflow
                         await self.handle_websocket_message(
                             websocket,
                             data_type,
-                            message
+                            client_id,
+                            data
                         )
                     except json.JSONDecodeError:
                         error_msg = {"error": "Invalid JSON format"}
@@ -169,6 +192,8 @@ class WebSocketAPITrigger(Trigger):
             except Exception as e:
                 logger.error(f"WebSocket connection error: {e}")
 
+            print("DISCONNECTED")
+
         app.websocket(path)(websocket_handler)
 
     async def _run(self, app: FastAPI, path: str, data_type: type) -> AsyncIterator[bool]:
@@ -179,7 +204,10 @@ class WebSocketAPITrigger(Trigger):
         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('token'), trigger['token'])
                 self.prepare_output_edges(self.get_output_port_by_name('data'), trigger['data'])
+                self.prepare_output_edges(self.get_output_port_by_name('client_id'), trigger['client_id'])
                 yield self.trigger(trigger['id'])
             except Exception as e:
                 logger.error(f"Error in WebSocketAPITrigger._run: {e}")

service_forge/workflow/workflow.py

@@ -2,17 +2,30 @@ 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 typing import Any, AsyncIterator, Awaitable, Callable
+
+from loguru import logger
+from opentelemetry import context as otel_context_api
+from opentelemetry import trace
+from opentelemetry.trace import SpanKind
+
+from ..db.database import DatabaseManager
+from ..execution_context import (
+    ExecutionContext,
+    get_current_context,
+    reset_current_context,
+    set_current_context,
+)
+from .edge import Edge
 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
+from .context import Context
 
 class Workflow:
     def __init__(
@@ -27,10 +40,14 @@ class Workflow:
         database_manager: DatabaseManager = None,
         max_concurrent_runs: int = 10,
         callbacks: list[WorkflowCallback] = [],
+        debug_version: bool = False,    # 是否为debug过程中的临时版本
 
         # for run
         task_id: uuid.UUID = None,
         real_trigger_node: Trigger = None,
+
+        # global variables
+        global_context: Context = None,
     ) -> None:
         self.id = id
         self.config = config
@@ -46,9 +63,12 @@ class Workflow:
         self.max_concurrent_runs = max_concurrent_runs
         self.run_semaphore = asyncio.Semaphore(max_concurrent_runs)
         self.callbacks = callbacks
+        self.debug_version = debug_version
         self.task_id = task_id
         self.real_trigger_node = real_trigger_node
+        self.global_context = global_context
         self._validate()
+        self._tracer = trace.get_tracer("service_forge.workflow")
     
     @property
     def name(self) -> str:
@@ -74,6 +94,8 @@ class Workflow:
                 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_WORKFLOW_ERROR:
+                await callback.on_workflow_error(*args, **kwargs)
             elif callback_type == CallbackEvent.ON_NODE_START:
                 await callback.on_node_start(*args, **kwargs)
             elif callback_type == CallbackEvent.ON_NODE_END:
@@ -90,8 +112,7 @@ class Workflow:
         for node in nodes:
             self.nodes.remove(node)
 
-    def load_config(self) -> None:
-        ...
+    def load_config(self) -> None: ...
 
     def _validate(self) -> None:
         # DAG
@@ -117,7 +138,7 @@ class Workflow:
             raise ValueError("Multiple trigger nodes found in workflow.")
         return trigger_nodes[0]
 
-    async def _run_node_with_callbacks(self, node: Node) -> None:
+    async def _run_node_with_callbacks(self, node: Node) -> bool:
         await self.call_callbacks(CallbackEvent.ON_NODE_START, node=node)
         
         try:
@@ -126,8 +147,13 @@ class Workflow:
                 await self.handle_node_stream_output(node, result)
             elif asyncio.iscoroutine(result):
                 await result
+        except Exception as e:
+            await self.call_callbacks(CallbackEvent.ON_WORKFLOW_ERROR, workflow=self, node=node, error=e)
+            logger.error(f"Error when running node {node.name}: {str(e)}, task_id: {self.task_id}")
+            return False
         finally:
             await self.call_callbacks(CallbackEvent.ON_NODE_END, node=node)
+        return True
 
     async def run_after_trigger(self) -> Any:
         logger.info(f"Running workflow: {self.name}")
@@ -138,30 +164,41 @@ class Workflow:
         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
+        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)))
+
+            results = await asyncio.gather(*tasks, return_exceptions=True)
+            
+            for i, result in enumerate(results):
+                if isinstance(result, Exception):
+                    for task in tasks:
+                        if not task.done():
+                            task.cancel()
+                    await asyncio.gather(*tasks, return_exceptions=True)
+                    return
+                    # raise result
+                elif result is False:
+                    logger.error(f"Node execution failed, stopping workflow: {nodes[i].name}")
+                    for task in tasks:
+                        if not task.done():
+                            task.cancel()
+                    await asyncio.gather(*tasks, return_exceptions=True)
+                    return
+                    # raise RuntimeError(f"Workflow stopped due to node execution failure: {nodes[i].name}")
 
         if len(self.output_ports) > 0:
             if len(self.output_ports) == 1:
@@ -178,16 +215,82 @@ class Workflow:
         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:
+            base_context = get_current_context()
+
+            # 尝试从 trigger_node 上取父 trace context（如果存在）
+            trigger_parent_context = None
+            if hasattr(trigger_node, "task_contexts"):
+                trigger_parent_context = trigger_node.task_contexts.pop(task_id, None)
+
+            parent_context = (
+                trigger_parent_context
+                or (
+                    base_context.trace_context
+                    if base_context and base_context.trace_context
+                    else otel_context_api.get_current()
+                )
+            )
+
+            span_name = f"Workflow {self.name}"
+            token = None
+
             try:
-                new_workflow = self._clone(task_id, trigger_node)
-                await new_workflow.run_after_trigger()
-                # TODO: clear new_workflow
+                with self._tracer.start_as_current_span(
+                    span_name,
+                    context=parent_context,
+                    kind=SpanKind.INTERNAL,
+                ) as span:
+                    span.set_attribute("workflow.name", self.name)
+                    span.set_attribute("workflow.task_id", str(task_id))
+                    span.set_attribute("service.name", "service_forge")
+
+                    execution_context = ExecutionContext(
+                        trace_context=otel_context_api.get_current(),
+                        span=span,
+                        metadata={
+                            **(base_context.metadata if base_context else {}),
+                            "workflow_name": self.name,
+                            "task_id": str(task_id),
+                        },
+                    )
+                    token = set_current_context(execution_context)
+
+                    new_workflow = self._clone(task_id, trigger_node)
+
+                    # run_after_trigger 当前实现没有 return，result 会是 None
+                    result = await new_workflow.run_after_trigger()
+
+                    # 将结果回传给等待方（如果有队列）
+                    if hasattr(trigger_node, "result_queues") and task_id in trigger_node.result_queues:
+                        trigger_node.result_queues[task_id].put_nowait(result)
+
+                    # TODO: clear new_workflow
+                    for node in new_workflow.nodes:
+                        await node.clear()
 
             except Exception as e:
-                error_msg = f"Error running workflow: {str(e)}, {traceback.format_exc()}"
-                logger.error(error_msg)
+                await self.call_callbacks(CallbackEvent.ON_WORKFLOW_ERROR, workflow=self, node=None, error=e)
+                # error_msg = f"Error running workflow: {str(e)}, {traceback.format_exc()}"
+                # logger.error(error_msg)
+                # await self.call_callbacks(CallbackEvent.ON_WORKFLOW_END, workflow=self, node=None, error=e)
+                return
+
+                # 更新任务状态为失败 + WebSocket 通知
+                websocket_manager.task_manager.fail_task(task_id, error_msg)
+                await websocket_manager.send_execution_error(task_id, "workflow", error_msg)
+
+                # 异常也回传给等待方，避免对方一直 await
+                if hasattr(trigger_node, "result_queues") and task_id in trigger_node.result_queues:
+                    trigger_node.result_queues[task_id].put_nowait(e)
+
+                # 将异常继续抛出，便于上层感知失败
+                raise
+            finally:
+                if token is not None:
+                    reset_current_context(token)
 
     async def run(self):
         tasks = []
@@ -203,9 +306,9 @@ class Workflow:
         trigger = self.get_trigger_node()
         await trigger._stop()
 
-    def trigger(self, trigger_name: str, **kwargs) -> uuid.UUID:
+    def trigger(self, trigger_name: str, assigned_task_id: uuid.UUID | None, **kwargs) -> uuid.UUID:
         trigger = self.get_trigger_node()
-        task_id = uuid.uuid4()
+        task_id = assigned_task_id or uuid.uuid4()
         for key, value in kwargs.items():
             trigger.prepare_output_edges(key, value)
         task = asyncio.create_task(self._run(task_id, trigger))

service_forge/workflow/workflow_callback.py

@@ -31,7 +31,7 @@ class WorkflowCallback:
         pass
 
     @abstractmethod
-    async def on_workflow_error(self, workflow: Workflow, error: Any) -> None:
+    async def on_workflow_error(self, workflow: Workflow, node: Node, error: Any) -> None:
         pass
 
     @abstractmethod
@@ -90,7 +90,7 @@ class BuiltinWorkflowCallback(WorkflowCallback):
             logger.error(f"发送 workflow_end 消息到 websocket 失败: {e}")
 
     @override
-    async def on_workflow_error(self, workflow: Workflow, error: Any) -> None:
+    async def on_workflow_error(self, workflow: Workflow, node: Node | None, 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:
@@ -103,6 +103,7 @@ class BuiltinWorkflowCallback(WorkflowCallback):
             message = {
                 "type": "workflow_error",
                 "task_id": str(workflow.task_id),
+                "node": node.name if node else None,
                 "error": self._serialize_result(error),
                 "is_end": False,
                 "is_error": True
@@ -113,11 +114,33 @@ class BuiltinWorkflowCallback(WorkflowCallback):
 
     @override
     async def on_node_start(self, node: Node) -> None:
-        ...
+        try:
+            manager = self._get_websocket_manager()
+            message = {
+                "type": "node_start",
+                "task_id": str(node.workflow.task_id),
+                "node": node.name,
+                "is_end": False,
+                "is_error": False
+            }
+            await manager.send_to_task(node.workflow.task_id, message)
+        except Exception as e:
+            logger.error(f"发送 node_start 消息到 websocket 失败: {e}")
 
     @override
     async def on_node_end(self, node: Node) -> None:
-        ...
+        try:
+            manager = self._get_websocket_manager()
+            message = {
+                "type": "node_end",
+                "task_id": str(node.workflow.task_id),
+                "node": node.name,
+                "is_end": False,
+                "is_error": False
+            }
+            await manager.send_to_task(node.workflow.task_id, message)
+        except Exception as e:
+            logger.error(f"发送 node_end 消息到 websocket 失败: {e}")
 
     @override
     async def on_node_stream_output(self, node: Node, output: Any) -> None:

service_forge/workflow/workflow_factory.py

@@ -3,6 +3,8 @@ from omegaconf import OmegaConf
 from typing import Callable, Awaitable, AsyncIterator, Any
 from copy import deepcopy
 
+from pydantic import ValidationError
+
 from service_forge.workflow.workflow_callback import BuiltinWorkflowCallback
 from .workflow import Workflow
 from .workflow_group import WorkflowGroup, WORKFLOW_DEFAULT_VERSION
@@ -53,6 +55,7 @@ def create_workflow(
     _handle_stream_output: Callable[[str, AsyncIterator[str]], Awaitable[None]] | None = None,
     _handle_query_user: Callable[[str, str], Awaitable[str]] | None = None,
     database_manager: DatabaseManager = None,
+    debug_version: bool = False,
 ) -> Workflow:
     if config is None:
         if config_path is not None:
@@ -71,6 +74,8 @@ def create_workflow(
         database_manager = database_manager,
         # TODO: max_concurrent_runs
         callbacks = [BuiltinWorkflowCallback()],
+        debug_version = debug_version,
+        global_context = Context(variables={}),
     )
 
     nodes: dict[str, Node] = {}
@@ -207,6 +212,7 @@ def create_workflow_group(
     _handle_stream_output: Callable[[str, AsyncIterator[str]], Awaitable[None]] = None,
     _handle_query_user: Callable[[str, str], Awaitable[str]] = None,
     database_manager: DatabaseManager = None,
+    debug_version: bool = False,
 ) -> WorkflowGroup:
 
     if config is None:
@@ -218,6 +224,12 @@ def create_workflow_group(
         else:
             raise ValueError("Either config_path or config must be provided")
 
+    if isinstance(config, dict):
+        try:
+            config = WorkflowConfig.model_validate(config)
+        except ValidationError:
+            config = WorkflowGroupConfig.model_validate(config)
+
     if type(config) == WorkflowConfig:
         workflow = create_workflow(
             config_path=config_path if config_path else None,
@@ -226,6 +238,7 @@ def create_workflow_group(
             _handle_stream_output=_handle_stream_output,
             _handle_query_user=_handle_query_user,
             database_manager=database_manager,
+            debug_version=debug_version,
         )
         return WorkflowGroup(workflows=[workflow], main_workflow_name=workflow.name, main_workflow_version=workflow.version)
     elif type(config) == WorkflowGroupConfig:
@@ -242,5 +255,6 @@ def create_workflow_group(
                 _handle_stream_output=_handle_stream_output,
                 _handle_query_user=_handle_query_user,
                 database_manager=database_manager,
+                debug_version=debug_version,
             ))
         return workflows

{service_forge-0.1.18.dist-info → service_forge-0.1.39.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: service-forge
-Version: 0.1.18
+Version: 0.1.39
 Summary: Add your description here
 Author-email: euxcet <zcc.qwer@gmail.com>
 Requires-Python: >=3.11
@@ -16,6 +16,9 @@ Requires-Dist: kubernetes>=28.0.0
 Requires-Dist: loguru>=0.7.3
 Requires-Dist: omegaconf>=2.3.0
 Requires-Dist: openai>=2.3.0
+Requires-Dist: opentelemetry-api>=1.38.0
+Requires-Dist: opentelemetry-exporter-otlp>=1.38.0
+Requires-Dist: opentelemetry-sdk>=1.38.0
 Requires-Dist: protobuf>=6.33.1
 Requires-Dist: psycopg2-binary>=2.9.11
 Requires-Dist: pydantic>=2.12.0