coze-coding-utils 0.2.3a2__py3-none-any.whl → 0.2.3a3__py3-none-any.whl

coze_coding_utils/helper/stream_runner.py

@@ -7,24 +7,30 @@ from abc import ABC, abstractmethod
 from typing import Any, Dict, Iterator, AsyncIterable
 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__)
 
 TIMEOUT_SECONDS = 900
 PING_INTERVAL_SECONDS = 30
 
-
 class WorkflowEventType:
     WORKFLOW_START = "workflow_start"
     WORKFLOW_END = "workflow_end"
@@ -450,6 +456,106 @@ 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, need_detail=False):
+            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],
+    workflow_debug: bool = False,
+) -> AsyncGenerator[str, None]:
+    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, need_detail=workflow_debug):
+            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-0.2.3a2.dist-info → coze_coding_utils-0.2.3a3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: coze-coding-utils
-Version: 0.2.3a2
+Version: 0.2.3a3
 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.3a3.dist-info}/RECORD

@@ -10,7 +10,7 @@ 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=f66n6QJ3zCakhk7Fe4Vz9vTZ2KJuM9v9UJfqX5S3nDA,24050
 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
@@ -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.3a3.dist-info/METADATA,sha256=BpT4ybyeGIjxSW4G62Z56WZ05ITL24mGEkKh-zDG47E,979
+coze_coding_utils-0.2.3a3.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+coze_coding_utils-0.2.3a3.dist-info/licenses/LICENSE,sha256=lzckZhAjHlpSJcWvppoST095IHFpBwKiB2pKcBv7vP4,1078
+coze_coding_utils-0.2.3a3.dist-info/RECORD,,