From 4e1d81c9f8ed6c758f51ddde11694d02fc34ef83 Mon Sep 17 00:00:00 2001
From: Dong_master <2213070223@qq.com>
Date: Mon, 14 Jul 2025 00:40:02 +0800
Subject: [PATCH] feat:add dify _agent_chat_message streaming

---
 pkg/provider/runners/difysvapi.py | 77 ++++++++++++++++++++++---------
 1 file changed, 56 insertions(+), 21 deletions(-)

diff --git a/pkg/provider/runners/difysvapi.py b/pkg/provider/runners/difysvapi.py
index 1dfde547..566dc0f8 100644
--- a/pkg/provider/runners/difysvapi.py
+++ b/pkg/provider/runners/difysvapi.py
@@ -251,11 +251,26 @@ class DifyServiceAPIRunner(runner.RequestRunner):
 
             if chunk['event'] in ignored_events:
                 continue
+            batch_pending_index += 1
+
+            if chunk['event'] == 'agent_message' or chunk['event'] == 'message_end':
+                if chunk['event'] == 'message_end':
+                    print(chunk['event'])
+                    # break
+                    is_final = True
+                else:
+                    is_final = False
+                    pending_agent_message += chunk['answer']
+                if is_stream:
+                    if batch_pending_index % 64 == 0 or is_final:
+                        yield llm_entities.MessageChunk(
+                            role='assistant',
+                            content=self._try_convert_thinking(pending_agent_message),
+                            is_final=is_final,
+                        )
 
-            if chunk['event'] == 'agent_message':
-                pending_agent_message += chunk['answer']
             else:
-                if pending_agent_message.strip() != '':
+                if pending_agent_message.strip() != '' and not is_stream:
                     pending_agent_message = pending_agent_message.replace('</details>Action:', '</details>')
                     yield llm_entities.Message(
                         role='assistant',
@@ -268,19 +283,34 @@ class DifyServiceAPIRunner(runner.RequestRunner):
                         continue
 
                     if chunk['tool']:
-                        msg = llm_entities.Message(
-                            role='assistant',
-                            tool_calls=[
-                                llm_entities.ToolCall(
-                                    id=chunk['id'],
-                                    type='function',
-                                    function=llm_entities.FunctionCall(
-                                        name=chunk['tool'],
-                                        arguments=json.dumps({}),
-                                    ),
-                                )
-                            ],
-                        )
+                        if is_stream:
+                            msg = llm_entities.MessageChunk(
+                                role='assistant',
+                                tool_calls=[
+                                    llm_entities.ToolCall(
+                                        id=chunk['id'],
+                                        type='function',
+                                        function=llm_entities.FunctionCall(
+                                            name=chunk['tool'],
+                                            arguments=json.dumps({}),
+                                        ),
+                                    )
+                                ],
+                            )
+                        else:
+                            msg = llm_entities.Message(
+                                role='assistant',
+                                tool_calls=[
+                                    llm_entities.ToolCall(
+                                        id=chunk['id'],
+                                        type='function',
+                                        function=llm_entities.FunctionCall(
+                                            name=chunk['tool'],
+                                            arguments=json.dumps({}),
+                                        ),
+                                    )
+                                ],
+                            )
                         yield msg
                 if chunk['event'] == 'message_file':
                     if chunk['type'] == 'image' and chunk['belongs_to'] == 'assistant':
@@ -290,11 +320,16 @@ class DifyServiceAPIRunner(runner.RequestRunner):
                             base_url = base_url[:-3]
 
                         image_url = base_url + chunk['url']
-
-                        yield llm_entities.Message(
-                            role='assistant',
-                            content=[llm_entities.ContentElement.from_image_url(image_url)],
-                        )
+                        if is_stream:
+                            yield llm_entities.MessageChunk(
+                                role='assistant',
+                                content=[llm_entities.ContentElement.from_image_url(image_url)],
+                            )
+                        else:
+                            yield llm_entities.Message(
+                                role='assistant',
+                                content=[llm_entities.ContentElement.from_image_url(image_url)],
+                            )
                 if chunk['event'] == 'error':
                     raise errors.DifyAPIError('dify 服务错误: ' + chunk['message'])