coze-coding-utils 0.2.3a2__py3-none-any.whl → 0.2.4__py3-none-any.whl

coze_coding_utils/helper/stream_runner.py

@@ -4,20 +4,28 @@ import threading
 import contextvars
 import logging
 from abc import ABC, abstractmethod
-from typing import Any, Dict, Iterator, AsyncIterable
+from dataclasses import dataclass
+from typing import Any, Dict, Iterator, AsyncIterable, Optional, Literal
 from langchain_core.runnables import RunnableConfig
 from langgraph.graph.state import CompiledStateGraph
-from coze_coding_utils.runtime_ctx.context import Context
 from coze_coding_utils.helper.agent_helper import (
     to_stream_input,
-    to_client_message,
     agent_iter_server_messages,
 )
+
+from coze_coding_utils.error import classify_error
+import asyncio
+import time
+import traceback
+from typing import Any, Dict, AsyncGenerator, Callable
+from coze_coding_utils.runtime_ctx.context import Context
 from coze_coding_utils.messages.server import (
+    create_message_end_dict,
+    create_message_error_dict,
     MESSAGE_END_CODE_CANCELED,
-    create_message_end_dict, create_message_error_dict,
 )
-from coze_coding_utils.error import classify_error
+from coze_coding_utils.helper.agent_helper import to_client_message
+from coze_coding_utils.error.classifier import ErrorClassifier
 
 logger = logging.getLogger(__name__)
 
@@ -25,6 +33,14 @@ TIMEOUT_SECONDS = 900
 PING_INTERVAL_SECONDS = 30
 
 
+@dataclass
+class RunOpt:
+    workflow_debug: bool = False
+
+    @property
+    def stream_mode(self) -> Literal["debug", "updates"]:
+        return "debug" if self.workflow_debug else "updates"
+
 class WorkflowEventType:
     WORKFLOW_START = "workflow_start"
     WORKFLOW_END = "workflow_end"
@@ -41,16 +57,16 @@ class WorkflowErrorCode:
 
 class BaseStreamRunner(ABC):
     @abstractmethod
-    def stream(self, payload: Dict[str, Any], graph: CompiledStateGraph, run_config: RunnableConfig, ctx: Context) -> Iterator[Any]:
+    def stream(self, payload: Dict[str, Any], graph: CompiledStateGraph, run_config: RunnableConfig, ctx: Context, run_opt: Optional[RunOpt] = None) -> Iterator[Any]:
         pass
 
     @abstractmethod
-    async def astream(self, payload: Dict[str, Any], graph: CompiledStateGraph, run_config: RunnableConfig, ctx: Context) -> AsyncIterable[Any]:
+    async def astream(self, payload: Dict[str, Any], graph: CompiledStateGraph, run_config: RunnableConfig, ctx: Context, run_opt: Optional[RunOpt] = None) -> AsyncIterable[Any]:
         pass
 
 
 class AgentStreamRunner(BaseStreamRunner):
-    def stream(self, payload: Dict[str, Any], graph: CompiledStateGraph, run_config: RunnableConfig, ctx: Context) -> Iterator[Any]:
+    def stream(self, payload: Dict[str, Any], graph: CompiledStateGraph, run_config: RunnableConfig, ctx: Context, run_opt: Optional[RunOpt] = None) -> Iterator[Any]:
         client_msg, session_id = to_client_message(payload)
         run_config["recursion_limit"] = 100
         run_config["configurable"] = {"thread_id": session_id}
@@ -95,7 +111,7 @@ class AgentStreamRunner(BaseStreamRunner):
             )
             yield end_msg
 
-    async def astream(self, payload: Dict[str, Any], graph: CompiledStateGraph, run_config: RunnableConfig, ctx: Context) -> AsyncIterable[Any]:
+    async def astream(self, payload: Dict[str, Any], graph: CompiledStateGraph, run_config: RunnableConfig, ctx: Context, run_opt: Optional[RunOpt] = None) -> AsyncIterable[Any]:
         client_msg, session_id = to_client_message(payload)
         run_config["recursion_limit"] = 100
         run_config["configurable"] = {"thread_id": session_id}
@@ -188,9 +204,6 @@ class AgentStreamRunner(BaseStreamRunner):
 
 
 class WorkflowStreamRunner(BaseStreamRunner):
-    def __init__(self):
-        self._node_start_times: Dict[str, float] = {}
-
     def _serialize_data(self, data: Any) -> Any:
         if isinstance(data, dict):
             return {k: self._serialize_data(v) for k, v in data.items()}
@@ -215,7 +228,9 @@ class WorkflowStreamRunner(BaseStreamRunner):
         result.update(kwargs)
         return result
 
-    def stream(self, payload: Dict[str, Any], graph: CompiledStateGraph, run_config: RunnableConfig, ctx: Context) -> Iterator[Any]:
+    def stream(self, payload: Dict[str, Any], graph: CompiledStateGraph, run_config: RunnableConfig, ctx: Context, run_opt: Optional[RunOpt] = None) -> Iterator[Any]:
+        if run_opt is None:
+            run_opt = RunOpt()
         run_config["recursion_limit"] = 100
         if "configurable" not in run_config:
             run_config["configurable"] = {}
@@ -226,21 +241,19 @@ class WorkflowStreamRunner(BaseStreamRunner):
         node_start_times: Dict[str, float] = {}
         final_output = {}
         seq = 0
-        is_debug = run_config.get("configurable", {}).get("workflow_debug", False)
-        stream_mode = "debug" if is_debug else "updates"
 
         try:
             seq += 1
             yield (seq, self._build_event(WorkflowEventType.WORKFLOW_START, ctx))
 
-            for event in graph.stream(payload, stream_mode=stream_mode, config=run_config, context=ctx):
+            for event in graph.stream(payload, stream_mode=run_opt.stream_mode, config=run_config, context=ctx):
                 current_time = time.time()
                 if current_time - last_ping_time >= PING_INTERVAL_SECONDS:
                     seq += 1
                     yield (seq, self._build_event(WorkflowEventType.PING, ctx))
                     last_ping_time = current_time
                 
-                if not is_debug:
+                if not run_opt.workflow_debug:
                     if isinstance(event, dict):
                         logger.info(f"Debug event: {event}")
                         for node_name, node_output in event.items():
@@ -305,7 +318,9 @@ class WorkflowStreamRunner(BaseStreamRunner):
             seq += 1
             yield (seq, self._build_event(WorkflowEventType.ERROR, ctx, code=str(err.code), error_msg=err.message))
 
-    async def astream(self, payload: Dict[str, Any], graph: CompiledStateGraph, run_config: RunnableConfig, ctx: Context) -> AsyncIterable[Any]:
+    async def astream(self, payload: Dict[str, Any], graph: CompiledStateGraph, run_config: RunnableConfig, ctx: Context, run_opt: Optional[RunOpt] = None) -> AsyncIterable[Any]:
+        if run_opt is None:
+            run_opt = RunOpt()
         run_config["recursion_limit"] = 100
         if "configurable" not in run_config:
             run_config["configurable"] = {}
@@ -317,9 +332,7 @@ class WorkflowStreamRunner(BaseStreamRunner):
         start_time = time.time()
         cancelled = threading.Event()
         last_ping_time = [start_time]
-        is_debug = run_config.get("configurable", {}).get("workflow_debug", False)
-        stream_mode = "debug" if is_debug else "updates"
-        logger.info(f"Stream mode: {stream_mode}")
+        logger.info(f"Stream mode: {run_opt.stream_mode}")
         seq = [0]
 
         def producer():
@@ -333,7 +346,7 @@ class WorkflowStreamRunner(BaseStreamRunner):
                 seq[0] += 1
                 loop.call_soon_threadsafe(q.put_nowait, (seq[0], self._build_event(WorkflowEventType.WORKFLOW_START, ctx)))
 
-                for event in graph.stream(payload, stream_mode=stream_mode, config=run_config, context=ctx):
+                for event in graph.stream(payload, stream_mode=run_opt.stream_mode, config=run_config, context=ctx):
                     if cancelled.is_set():
                         logger.info(f"Workflow producer cancelled during iteration for run_id: {ctx.run_id}")
                         seq[0] += 1
@@ -352,7 +365,7 @@ class WorkflowStreamRunner(BaseStreamRunner):
                         loop.call_soon_threadsafe(q.put_nowait, (seq[0], self._build_event(WorkflowEventType.PING, ctx)))
                         last_ping_time[0] = current_time
 
-                    if not is_debug:
+                    if not run_opt.workflow_debug:
                         if isinstance(event, dict):
                             for node_name, node_output in event.items():
                                 logger.info(f"Node output: {node_name}")
@@ -450,6 +463,108 @@ class WorkflowStreamRunner(BaseStreamRunner):
                 pass
 
 
+async def agent_stream_handler(
+    payload: Dict[str, Any],
+    ctx: Context,
+    run_id: str,
+    stream_sse_func: Callable,
+    sse_event_func: Callable,
+    error_classifier: ErrorClassifier,
+    register_task_func: Callable[[str, asyncio.Task], None],
+) -> AsyncGenerator[str, None]:
+    task = asyncio.current_task()
+    if task:
+        register_task_func(run_id, task)
+        logger.info(f"Registered agent streaming task for run_id: {run_id}")
+
+    client_msg, _ = to_client_message(payload)
+    t0 = time.time()
+
+    try:
+        async for chunk in stream_sse_func(payload, ctx):
+            yield chunk
+    except asyncio.CancelledError:
+        logger.info(f"Agent stream cancelled for run_id: {run_id}")
+        end_msg = create_message_end_dict(
+            code=MESSAGE_END_CODE_CANCELED,
+            message="Stream cancelled by user",
+            session_id=client_msg.session_id,
+            query_msg_id=client_msg.local_msg_id,
+            log_id=ctx.logid,
+            time_cost_ms=int((time.time() - t0) * 1000),
+            reply_id="",
+            sequence_id=1,
+        )
+        yield sse_event_func(end_msg)
+        raise
+    except Exception as ex:
+        err = error_classifier.classify(ex, {"node_name": "agent_stream", "run_id": run_id})
+        logger.error(
+            f"Unexpected error in agent_stream: [{err.code}] {err.message}, "
+            f"traceback: {traceback.format_exc()}"
+        )
+        error_msg = create_message_error_dict(
+            code=str(err.code),
+            message=str(ex),
+            session_id=client_msg.session_id,
+            query_msg_id=client_msg.local_msg_id,
+            log_id=ctx.logid,
+            reply_id="",
+            sequence_id=1,
+            local_msg_id=client_msg.local_msg_id,
+        )
+        yield sse_event_func(error_msg)
+
+
+async def workflow_stream_handler(
+    payload: Dict[str, Any],
+    ctx: Context,
+    run_id: str,
+    stream_sse_func: Callable,
+    sse_event_func: Callable,
+    error_classifier: ErrorClassifier,
+    register_task_func: Callable[[str, asyncio.Task], None],
+    run_opt: Optional[RunOpt] = None,
+) -> AsyncGenerator[str, None]:
+    if run_opt is None:
+        run_opt = RunOpt()
+    task = asyncio.current_task()
+    if task:
+        register_task_func(run_id, task)
+        logger.info(f"Registered workflow streaming task for run_id: {run_id}")
+
+    try:
+        async for chunk in stream_sse_func(payload, ctx, run_opt):
+            yield chunk
+    except asyncio.CancelledError:
+        logger.info(f"Workflow stream cancelled for run_id: {run_id}")
+        cancel_event = {
+            "type": "error",
+            "timestamp": int(time.time() * 1000),
+            "log_id": ctx.logid,
+            "run_id": run_id,
+            "code": "CANCELED",
+            "error_msg": "Stream cancelled by user",
+        }
+        yield sse_event_func(cancel_event)
+        raise
+    except Exception as ex:
+        err = error_classifier.classify(ex, {"node_name": "workflow_stream", "run_id": run_id})
+        logger.error(
+            f"Unexpected error in workflow_stream: [{err.code}] {err.message}, "
+            f"traceback: {traceback.format_exc()}"
+        )
+        error_event = {
+            "type": "error",
+            "timestamp": int(time.time() * 1000),
+            "log_id": ctx.logid,
+            "run_id": run_id,
+            "code": str(err.code),
+            "error_msg": str(ex),
+        }
+        yield sse_event_func(error_event)
+
+
 def get_stream_runner(is_agent: bool) -> BaseStreamRunner:
     if is_agent:
         return AgentStreamRunner()

coze_coding_utils/log/node_log.py

@@ -1,7 +1,6 @@
 import time
 import logging
 from uuid import UUID
-from openai import BaseModel
 from coze_coding_utils.log.config import LOG_DIR
 from coze_coding_utils.log.common import get_execute_mode, is_prod
 import uuid

{coze_coding_utils-0.2.3a2.dist-info → coze_coding_utils-0.2.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: coze-coding-utils
-Version: 0.2.3a2
+Version: 0.2.4
 Summary: Utilities for Coze coding client runtime context and helpers.
 Project-URL: Homepage, https://code.byted.org/stone/coze-coding-client
 Author: Bytedance Stone Team

{coze_coding_utils-0.2.3a2.dist-info → coze_coding_utils-0.2.4.dist-info}/RECORD

@@ -10,13 +10,13 @@ coze_coding_utils/file/file.py,sha256=fBda18EGSQZ3Xl8OqEaGAb5Rd90_SmhJ1k0jgQk2v7
 coze_coding_utils/helper/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 coze_coding_utils/helper/agent_helper.py,sha256=q1ZM30xLXoW-m0NJmJ_Y0M-kUAQCBstG_j7xkqsyRSU,22546
 coze_coding_utils/helper/graph_helper.py,sha256=UNtqqiQNAQ4319qcC1vHiLYIL2eGzvGQRgXu3mgLq8Y,8893
-coze_coding_utils/helper/stream_runner.py,sha256=q8lLo3i-PMVgbhWiv3X51NoiHkmyhT0-25hI_W1BDi0,20384
+coze_coding_utils/helper/stream_runner.py,sha256=K446dPUdwemA0_FwbtyscQz5hGmhRxUCG0XYUTFVSI8,24329
 coze_coding_utils/log/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 coze_coding_utils/log/common.py,sha256=mUNkCm68oaPaI6-a5UwLf87AfhrMnVPkEuri16guqKc,168
 coze_coding_utils/log/config.py,sha256=Qkw3JRuGUKJ6CBY7WqHJOFeyCU47cArvUtMsSBifFMo,195
 coze_coding_utils/log/err_trace.py,sha256=iwt5g8-AX0N2KuUpuLXjk5PocL6NZdAGddJh27Sxp_o,2954
 coze_coding_utils/log/loop_trace.py,sha256=68sI1AHKd8oeLti-7trBygTZqIiUu_TcHt93DLItkXQ,2129
-coze_coding_utils/log/node_log.py,sha256=AAreNJJc8Lx1QuCIwvss2ty9z5cDuHEzMobA4ZYSYBI,18040
+coze_coding_utils/log/node_log.py,sha256=cqx0l8hgDj8adydj5lYQbpnFwyd3k_QyNeq43Qsvl1I,18011
 coze_coding_utils/log/parser.py,sha256=XuzOkO8SzTbBP_XL1RMFzjqG4IG0iSRStAA2M6latt8,9728
 coze_coding_utils/log/write_log.py,sha256=lDu4_Tjk6eYtWBXI-OOYLqe5ejqTtGnDcrH6M7fMQ2A,7223
 coze_coding_utils/messages/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -32,7 +32,7 @@ coze_coding_utils/openai/types/request.py,sha256=IuNMT2Ce1--_32R30Q2q7Lb2dAwKNy3
 coze_coding_utils/openai/types/response.py,sha256=pjHHVR8LSMVFCc3fGzKqXrdoKDIfSCJEfICd_X9Nohc,4808
 coze_coding_utils/runtime_ctx/__init__.py,sha256=4W8VliAYUP1KY2gLJ_YDy2TmcXYVm-PY7XikQD_bFwA,2
 coze_coding_utils/runtime_ctx/context.py,sha256=G8ld-WnQ1pTJe5OOXC_dTbagXj9IxmpRiPM4X_jWW6o,3992
-coze_coding_utils-0.2.3a2.dist-info/METADATA,sha256=cGQN_B0A1fz_AH5FqH_yjgcCaQkjzURl6bJxQurKRVk,979
-coze_coding_utils-0.2.3a2.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-coze_coding_utils-0.2.3a2.dist-info/licenses/LICENSE,sha256=lzckZhAjHlpSJcWvppoST095IHFpBwKiB2pKcBv7vP4,1078
-coze_coding_utils-0.2.3a2.dist-info/RECORD,,
+coze_coding_utils-0.2.4.dist-info/METADATA,sha256=clx-L3s4LykCyhZxHNxKCc_Ip_u1AY3cpdyxqQkVL74,977
+coze_coding_utils-0.2.4.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+coze_coding_utils-0.2.4.dist-info/licenses/LICENSE,sha256=lzckZhAjHlpSJcWvppoST095IHFpBwKiB2pKcBv7vP4,1078
+coze_coding_utils-0.2.4.dist-info/RECORD,,