coze-coding-utils 0.2.3a2__tar.gz → 0.2.3a3__tar.gz

{coze_coding_utils-0.2.3a2 → coze_coding_utils-0.2.3a3}/PKG-INFO

@@ -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 → coze_coding_utils-0.2.3a3}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "coze-coding-utils"
-version = "0.2.3a2"
+version = "0.2.3a3"
 description = "Utilities for Coze coding client runtime context and helpers."
 readme = "README.md"
 requires-python = ">=3.10"

{coze_coding_utils-0.2.3a2 → coze_coding_utils-0.2.3a3}/src/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()