Merge pull request #82 from YangZhiBoGreenHand/yzb/fix/sub-stream

fix: sub stream with submit tool outputs

Author: liuooo

app/api/v1/runs.py

@@ -121,7 +121,6 @@ async def list_run_steps(
     return page.model_dump(by_alias=True)
 
 
-
 @router.get(
     "/{thread_id}/runs/{run_id}/steps/{step_id}",
     response_model=RunStepRead,
@@ -160,7 +159,7 @@ async def submit_tool_outputs_to_run(
     db_run = await RunService.submit_tool_outputs_to_run(session=session, thread_id=thread_id, run_id=run_id, body=body)
     # Resume async task
     if db_run.status == "queued":
-        run_task.apply_async(args=(db_run.id,))
+        run_task.apply_async(args=(db_run.id, body.stream))
 
     if body.stream:
         return pub_handler.sub_stream(db_run.id, request)

app/core/runner/pub_handler.py

@@ -49,6 +49,24 @@ def read_event(channel: str, x_index: str = None) -> Tuple[Optional[str], Option
     return stream_id, event
 
 
+def save_last_stream_id(run_id: str, stream_id: str):
+    """
+    保存当前 run_id 对应的最新 stream_id
+    :param run_id: 当前的运行 ID
+    :param stream_id: 最新的 stream_id
+    """
+    redis_client.set(f"run:{run_id}:last_stream_id", stream_id, 10 * 60)
+
+
+def get_last_stream_id(run_id: str) -> str:
+    """
+    获取上次保存的 stream_id
+    :param run_id: 当前的运行 ID
+    :return: 上次的 stream_id 或 None
+    """
+    return redis_client.get(f"run:{run_id}:last_stream_id")
+
+
 def _data_adjust_tools(tools: List[dict]) -> List[dict]:
     def _adjust_tool(tool: dict):
         if tool["type"] not in {"code_interpreter", "file_search", "function"}:
@@ -108,7 +126,8 @@ async def _stream():
         for event in prefix_events:
             yield event
 
-        x_index = None
+        last_index = get_last_stream_id(run_id)  # 获取上次的 stream_id
+        x_index = last_index or None
         while True:
             if await request.is_disconnected():
                 break
@@ -120,12 +139,15 @@ async def _stream():
                 break
 
             if data["event"] == "done":
+                save_last_stream_id(run_id, x_index)  # 记录最新的 stream_id
                 break
 
             if data["event"] == "error":
+                save_last_stream_id(run_id, x_index)  # 记录最新的 stream_id
                 raise InternalServerError(data["data"])
 
             yield data
+            save_last_stream_id(run_id, x_index)  # 记录最新的 stream_id
 
         for event in suffix_events:
             yield event

tests/e2e/sub_stream_test.py

@@ -0,0 +1,88 @@
+import json
+
+from openai import AssistantEventHandler
+
+
+def test_sub_stream_with_submit_tool_outputs_stream(client):
+
+    def get_current_weather(location):
+        return f"{location}今天是雨天。 "
+
+    assistant = client.beta.assistants.create(
+        name="Assistant Demo",
+        instructions="You are a helpful assistant. When asked a question, use tools wherever possible.",
+        model="gpt-4o",
+        tools=[
+            {
+                "type": "function",
+                "function": {
+                    "name": "get_current_weather",
+                    "description": "当你想查询指定城市的天气时非常有用。",
+                    "parameters": {"type": "object", "properties": {"location": {"type": "string", "description": "城市或县区，比如北京市、杭州市、余杭区等。"}}, "required": ["location"]},  # 查询天气时需要提供位置，因此参数设置为location
+                },
+            }
+        ],
+    )
+    print("=====> : %s\n", assistant)
+
+    thread = client.beta.threads.create()
+    print("=====> : %s\n", thread)
+
+    message = client.beta.threads.messages.create(
+        thread_id=thread.id,
+        role="user",
+        content="北京天气如何？",
+    )
+    print("=====> : %s\n", message)
+
+    funcs = [get_current_weather]
+
+    class EventHandler(AssistantEventHandler):
+
+        def on_event(self, event):
+            print(event.event)
+            if event.event == "thread.run.requires_action":
+                print(event)
+                run_id = event.data.id  # Retrieve the run ID from the event data
+                self.handle_requires_action(event.data, run_id)
+
+        def handle_requires_action(self, data, run_id):
+            tool_outputs = []
+
+            for tool in data.required_action.submit_tool_outputs.tool_calls:
+                func = next(iter([func for func in funcs if func.__name__ == tool.function.name]))
+                try:
+                    output = func(**eval(tool.function.arguments))
+                except Exception as e:
+                    output = "Error: " + str(e)
+
+                tool_outputs.append({"tool_call_id": tool.id, "output": json.dumps(output)})
+
+            print(tool_outputs)
+
+            # Submit all tool_outputs at the same time
+            self.submit_tool_outputs(tool_outputs, run_id)
+
+        def submit_tool_outputs(self, tool_outputs, run_id):
+            # Use the submit_tool_outputs_stream helper
+            with client.beta.threads.runs.submit_tool_outputs_stream(
+                thread_id=self.current_run.thread_id,
+                run_id=self.current_run.id,
+                tool_outputs=tool_outputs,
+                event_handler=EventHandler(),
+            ) as stream:
+                # for text in stream.text_deltas:
+                #     print(text, end="", flush=True)
+                #     print()
+                stream.until_done()
+
+        def on_text_delta(self, delta, snapshot) -> None:
+            print("=====> text delta")
+            print("delta   : %s", delta)
+
+    with client.beta.threads.runs.stream(
+        thread_id=thread.id,
+        assistant_id=assistant.id,
+        event_handler=EventHandler(),
+    ) as stream:
+        stream.until_done()