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

service_forge/workflow/node.py

@@ -1,50 +1,78 @@
 from __future__ import annotations
 
-from typing import Any, AsyncIterator, Union, TYPE_CHECKING, Callable, Awaitable
-from abc import ABC, abstractmethod
+import inspect
+import json
 import uuid
+from abc import ABC, abstractmethod
+from dataclasses import asdict, is_dataclass
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    AsyncIterator,
+    Awaitable,
+    Callable,
+    Optional,
+    Union,
+    cast,
+)
+
 from loguru import logger
-from .edge import Edge
-from .port import Port
-from .context import Context
-from ..utils.register import Register
+from opentelemetry import context as otel_context_api
+from opentelemetry import trace
+from opentelemetry.trace import SpanKind
+from pydantic import BaseModel
+
 from ..db.database import DatabaseManager, PostgresDatabase, MongoDatabase, RedisDatabase
 from ..utils.workflow_clone import node_clone
+from ..execution_context import (
+    ExecutionContext,
+    get_current_context,
+    reset_current_context,
+    set_current_context,
+)
+from ..utils.register import Register
+from .context import Context
+from .edge import Edge
+from .port import Port
+from .workflow_callback import CallbackEvent
 
 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']
+    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,
+        context: Optional[Context] = None,
+        input_edges: Optional[list[Edge]] = None,
+        output_edges: Optional[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,
+        query_user: Optional[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.workflow: Optional[Workflow] = None
         self.query_user = query_user
-        self.sub_workflows: WorkflowGroup = None
+        self.sub_workflows: Optional[WorkflowGroup] = None
 
         # runtime variables
         self.context = context
         self.input_variables: dict[Port, Any] = {}
         self.num_activated_input_edges = 0
+        self._tracer = trace.get_tracer("service_forge.node")
 
     @property
     def default_postgres_database(self) -> PostgresDatabase | None:
@@ -60,28 +88,39 @@ class Node(ABC):
     
     @property
     def database_manager(self) -> DatabaseManager:
+        assert self.workflow
         return self.workflow.database_manager
 
+    @property
+    def global_context(self) -> Context:
+        return self.workflow.global_context
+
     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_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.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:
+            # TODO: Register currently stores class objects; clarify Register typing vs instance usage.
             node_register.register(cls.__name__, cls)
         return super().__init_subclass__()
 
-    def _query_user(self, prompt: str) -> Callable[[str, str], Awaitable[str]]:
+    def _query_user(self, prompt: str) -> Awaitable[str]:
+        assert self.query_user
         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():
@@ -94,6 +133,7 @@ class Node(ABC):
 
     def is_trigger(self) -> bool:
         from .trigger import Trigger
+
         return isinstance(self, Trigger)
 
     # TODO: maybe add a function before the run function?
@@ -101,23 +141,219 @@ class Node(ABC):
     @abstractmethod
     async def _run(self, **kwargs) -> Union[None, AsyncIterator]:
         ...
+
+    async def clear(self) -> None:
+        ...
     
     def run(self) -> Union[None, AsyncIterator]:
+        task_id: uuid.UUID | None = None
         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)
+        if task_id is not None and "task_id" in self._run.__code__.co_varnames:
+            params["task_id"] = task_id
+        base_context = get_current_context()
+        parent_context = (
+            base_context.trace_context
+            if base_context and base_context.trace_context
+            else otel_context_api.get_current()
+        )
+        span_name = f"Node {self.name}"
 
-    def get_input_port_by_name(self, name: str) -> Port:
-        # TODO: add warning if port is extended
+        if inspect.isasyncgenfunction(self._run):
+            return self._run_async_generator(
+                params, task_id, base_context, parent_context, span_name
+            )
+        if inspect.iscoroutinefunction(self._run):
+            return self._run_async(
+                params, task_id, base_context, parent_context, span_name
+            )
+
+        return self._run_sync(params, task_id, base_context, parent_context, span_name)
+
+    def _build_execution_context(
+        self, base_context: ExecutionContext | None, span: trace.Span
+    ) -> ExecutionContext:
+        return ExecutionContext(
+            trace_context=otel_context_api.get_current(),
+            span=span,
+            state=base_context.state if base_context else {},
+            metadata={
+                **(base_context.metadata if base_context else {}),
+                "node": self.name,
+                "workflow_name": getattr(self.workflow, "name", None),
+            },
+        )
+
+    @staticmethod
+    def _serialize_for_trace(value: Any, max_length: int = 4000) -> tuple[str, bool]:
+        def _normalize(val: Any) -> Any:
+            if isinstance(val, BaseModel):
+                return val.model_dump()
+            if hasattr(val, "model_dump"):
+                try:
+                    return val.model_dump()
+                except Exception:
+                    pass
+            if hasattr(val, "dict"):
+                try:
+                    return val.dict()
+                except Exception:
+                    pass
+            if is_dataclass is not None and is_dataclass(val):
+                return asdict(val) if asdict else val
+            if isinstance(val, dict):
+                return {k: _normalize(v) for k, v in val.items()}
+            if isinstance(val, (list, tuple)):
+                return [_normalize(v) for v in val]
+            return val
+
+        normalized_value = _normalize(value)
+        serialized = json.dumps(normalized_value, ensure_ascii=False, default=str)
+        if len(serialized) > max_length:
+            return serialized[:max_length], True
+        return serialized, False
+
+    def _set_span_attributes(
+        self, span: trace.Span, params: dict[str, Any], task_id: uuid.UUID | None
+    ) -> None:
+        span.set_attribute("node.name", self.name)
+        if self.workflow is not None:
+            span.set_attribute("workflow.name", self.workflow.name)
+        if task_id is not None:
+            span.set_attribute("workflow.task_id", str(task_id))
+        span.set_attribute("node.input_keys", ",".join(params.keys()))
+        serialized_inputs, inputs_truncated = self._serialize_for_trace(params)
+        span.set_attribute("node.inputs", serialized_inputs)
+        if inputs_truncated:
+            span.set_attribute("node.inputs_truncated", True)
+
+    def _record_output(self, span: trace.Span, output: Any) -> None:
+        span.set_attribute(
+            "node.output_type", type(output).__name__ if output is not None else "None"
+        )
+        serialized_output, output_truncated = self._serialize_for_trace(output)
+        span.set_attribute("node.output", serialized_output)
+        if output_truncated:
+            span.set_attribute("node.output_truncated", True)
+
+    async def _run_async(
+        self,
+        params: dict[str, Any],
+        task_id: uuid.UUID | None,
+        base_context: ExecutionContext | None,
+        parent_context: otel_context_api.Context,
+        span_name: str,
+    ) -> Any:
+        with self._tracer.start_as_current_span(
+            span_name,
+            context=parent_context,
+            kind=SpanKind.INTERNAL,
+        ) as span:
+            self._set_span_attributes(span, params, task_id)
+            exec_ctx = self._build_execution_context(base_context, span)
+            token = set_current_context(exec_ctx)
+            try:
+                result = self._run(**params)
+                if inspect.isawaitable(result):
+                    result = await result
+                self._record_output(span, result)
+                return result
+            finally:
+                reset_current_context(token)
+
+    async def _run_async_generator(
+        self,
+        params: dict[str, Any],
+        task_id: uuid.UUID | None,
+        base_context: ExecutionContext | None,
+        parent_context: otel_context_api.Context,
+        span_name: str,
+    ) -> AsyncIterator[Any]:
+        # Trigger 节点是长期运行的 async generator，这里为每次触发单独生成/关闭 span，避免一个 span 挂载所有请求。
+        if self.is_trigger():
+            async for item in self._run(**params):
+                trigger_parent_context = parent_context
+                if hasattr(self, "task_contexts") and isinstance(item, uuid.UUID):
+                    trigger_ctx = getattr(self, "task_contexts").get(item)
+                    if trigger_ctx is not None:
+                        trigger_parent_context = trigger_ctx
+                with self._tracer.start_as_current_span(
+                    span_name,
+                    context=trigger_parent_context,
+                    kind=SpanKind.INTERNAL,
+                ) as span:
+                    self._set_span_attributes(span, params, task_id)
+                    span.set_attribute("node.output_type", "async_generator")
+                    serialized_item, item_truncated = self._serialize_for_trace(item)
+                    span.add_event(
+                        "node.output_item",
+                        {
+                            "value": serialized_item,
+                            "truncated": item_truncated,
+                        },
+                    )
+                    exec_ctx = self._build_execution_context(base_context, span)
+                    token = set_current_context(exec_ctx)
+                    try:
+                        yield item
+                    finally:
+                        reset_current_context(token)
+        else:
+            with self._tracer.start_as_current_span(
+                span_name,
+                context=parent_context,
+                kind=SpanKind.INTERNAL,
+            ) as span:
+                self._set_span_attributes(span, params, task_id)
+                span.set_attribute("node.output_type", "async_generator")
+                exec_ctx = self._build_execution_context(base_context, span)
+                token = set_current_context(exec_ctx)
+                try:
+                    async for item in self._run(**params):
+                        serialized_item, item_truncated = self._serialize_for_trace(item)
+                        span.add_event(
+                            "node.output_item",
+                            {
+                                "value": serialized_item,
+                                "truncated": item_truncated,
+                            },
+                        )
+                        yield item
+                finally:
+                    reset_current_context(token)
+
+    def _run_sync(
+        self,
+        params: dict[str, Any],
+        task_id: uuid.UUID | None,
+        base_context: ExecutionContext | None,
+        parent_context: otel_context_api.Context,
+        span_name: str,
+    ) -> Any:
+        with self._tracer.start_as_current_span(
+            span_name,
+            context=parent_context,
+            kind=SpanKind.INTERNAL,
+        ) as span:
+            self._set_span_attributes(span, params, task_id)
+            exec_ctx = self._build_execution_context(base_context, span)
+            token = set_current_context(exec_ctx)
+            try:
+                result = self._run(**params)
+                self._record_output(span, result)
+                return result
+            finally:
+                reset_current_context(token)
+
+    def get_input_port_by_name(self, name: str) -> Optional[Port]:
         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
+    def get_output_port_by_name(self, name: str) -> Optional[Port]:
         for port in self.output_ports:
             if port.name == name:
                 return port
@@ -143,9 +379,9 @@ class Node(ABC):
         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.')
+            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)
 
@@ -163,7 +399,7 @@ class Node(ABC):
         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)
@@ -173,7 +409,14 @@ class Node(ABC):
 
     # 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)
+        if self.workflow and hasattr(self.workflow, "_handle_workflow_output"):
+            handler = cast(
+                Callable[[str, Any], None],
+                getattr(self.workflow, "_handle_workflow_output"),
+            )
+            handler(self.name, data)
+        else:
+            logger.warning("Workflow output handler not set; skipping output dispatch.")
 
     def extended_output_name(self, name: str, index: int) -> str:
         return name + '_' + str(index)
@@ -181,4 +424,8 @@ class Node(ABC):
     def _clone(self, context: Context) -> Node:
         return node_clone(self, context)
 
-node_register = Register[Node]()
+    async def stream_output(self, data: Any) -> None:
+         await self.workflow.call_callbacks(CallbackEvent.ON_NODE_STREAM_OUTPUT, node=self, output=data)
+
+
+node_register = Register[Node]()

service_forge/workflow/triggers/fast_api_trigger.py

@@ -4,7 +4,7 @@ import asyncio
 import json
 from loguru import logger
 from service_forge.workflow.trigger import Trigger
-from typing import AsyncIterator, Any
+from typing import AsyncIterator, Any, Optional
 from fastapi import FastAPI, Request
 from fastapi.responses import StreamingResponse
 from service_forge.workflow.port import Port
@@ -13,6 +13,7 @@ from service_forge.api.routers.websocket.websocket_manager import websocket_mana
 from fastapi import HTTPException
 from google.protobuf.message import Message
 from google.protobuf.json_format import MessageToJson
+from opentelemetry import context as otel_context_api
 
 class FastAPITrigger(Trigger):
     DEFAULT_INPUT_PORTS = [
@@ -26,13 +27,16 @@ class FastAPITrigger(Trigger):
     DEFAULT_OUTPUT_PORTS = [
         Port("trigger", bool),
         Port("user_id", int),
+        Port("token", str),
         Port("data", Any),
+        Port("path_params", dict),
     ]
 
     def __init__(self, name: str):
         super().__init__(name)
         self.events = {}
         self.is_setup_route = False
+        self.task_contexts: dict[uuid.UUID, otel_context_api.Context] = {}
         self.app = None
         self.route_path = None
         self.route_method = None
@@ -40,23 +44,44 @@ class FastAPITrigger(Trigger):
     @staticmethod
     def serialize_result(result: Any):
         if isinstance(result, Message):
-            return MessageToJson(
-                result,
-                preserving_proto_field_name=True
-            )
+            return MessageToJson(result, preserving_proto_field_name=True)
         return result
 
+    def _normalize_result_or_raise(self, result: Any):
+      # TODO: 检查合并
+        if hasattr(result, "is_error") and hasattr(result, "result"):
+            if result.is_error:
+                if isinstance(result.result, HTTPException):
+                    raise result.result
+                raise HTTPException(status_code=500, detail=str(result.result))
+            return self.serialize_result(result.result)
+
+        if isinstance(result, HTTPException):
+            raise result
+        if isinstance(result, Exception):
+            raise HTTPException(status_code=500, detail=str(result))
+
+        return self.serialize_result(result)
+
     async def handle_request(
         self,
         request: Request,
         data_type: type,
         extract_data_fn: callable[[Request], dict],
         is_stream: bool,
+        path_params: Optional[dict] = None,
     ):
         task_id = uuid.uuid4()
         self.result_queues[task_id] = asyncio.Queue()
 
+        # parse trace context
+        trace_ctx = otel_context_api.get_current()
+        self.task_contexts[task_id] = trace_ctx
+
         body_data = await extract_data_fn(request)
+        # Merge path parameters into body_data (path params take precedence)
+        if path_params:
+            body_data = {**body_data, **path_params}
         converted_data = data_type(**body_data)
 
         client_id = (
@@ -69,10 +94,16 @@ class FastAPITrigger(Trigger):
             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)
 
+        # trigger_queue with trace_context, used in _run()
+        logger.info(f'user_id {getattr(request.state, "user_id", None)} token {getattr(request.state, "auth_token", None)}')
+
         self.trigger_queue.put_nowait({
             "id": task_id,
             "user_id": getattr(request.state, "user_id", None),
+            "token": getattr(request.state, "auth_token", None),
             "data": converted_data,
+            "trace_context": trace_ctx,
+            "path_params": path_params,
         })
 
         if is_stream:
@@ -83,26 +114,23 @@ class FastAPITrigger(Trigger):
                     while True:
                         item = await self.stream_queues[task_id].get()
 
-                        if item.is_error:
+                        if getattr(item, "is_error", False):
                             yield f"event: error\ndata: {json.dumps({'detail': str(item.result)})}\n\n"
                             break
-                        
-                        if item.is_end:
-                            # TODO: send the result?
+
+                        if getattr(item, "is_end", False):
                             break
 
                         # TODO: modify
                         serialized = self.serialize_result(item.result)
-                        if isinstance(serialized, str):
-                            data = serialized
-                        else:
-                            data = json.dumps(serialized)
-                        
+                        data = serialized if isinstance(serialized, str) else 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:
+                    self.stream_queues.pop(task_id, None)
+
                     if task_id in self.stream_queues:
                         del self.stream_queues[task_id]
                     if task_id in self.result_queues:
@@ -115,19 +143,13 @@ class FastAPITrigger(Trigger):
                     "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)
+        # 非流式：等待结果
+        result = await self.result_queues[task_id].get()
+        self.result_queues.pop(task_id, None)
+        return self._normalize_result_or_raise(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:
@@ -142,7 +164,10 @@ class FastAPITrigger(Trigger):
         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)
+            # Get path parameters from FastAPI request
+            # request.path_params is always available in FastAPI and contains path parameters
+            path_params = dict(request.path_params)
+            return await self.handle_request(request, data_type, extractor, is_stream, path_params)
 
         # Save route information for cleanup
         self.app = app
@@ -160,7 +185,14 @@ class FastAPITrigger(Trigger):
         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]:
+    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
@@ -168,8 +200,12 @@ class FastAPITrigger(Trigger):
         while True:
             try:
                 trigger = await self.trigger_queue.get()
+                if trace_ctx := trigger.get("trace_context"):
+                    self.task_contexts[trigger["id"]] = trace_ctx
                 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('path_params'), trigger['path_params'])
                 yield self.trigger(trigger['id'])
             except Exception as e:
                 logger.error(f"Error in FastAPITrigger._run: {e}")