supports batch storage of streaming messages.

Author: JDOxygen

CHANGELOG_zh.md

@@ -9,7 +9,18 @@
 
 ---
 
-## [1.0.9.2] - 2025-12-04
+## [1.0.9.2] - 2025-12-09
+
+### Added
+- 新增流式消息结束标识的stream_end消息
+- stream消息支持分批存储
+
+### Changed
+- 修改message表结构，新增字段
+
+---
+
+## [1.0.9.1] - 2025-12-04
 
 ### Added
 - 打印payload日志

config.json

@@ -33,6 +33,7 @@
             "is_send_think": true,
             "is_send_answer": true,
             "is_stored": false,
+            "stream_batch_size": 256,
             "is_show_in_terminal": false,
             "is_send_full_arguments": false
         },

oxygent/config.py

@@ -68,6 +68,7 @@ class Config:
             "is_send_think": True,
             "is_send_answer": True,
             "is_stored": False,
+            "stream_batch_size": 256,
             "is_show_in_terminal": False,
             "is_send_full_arguments": False,
         },
@@ -359,6 +360,14 @@ def set_message_is_stored(cls, is_stored=True):
     def get_message_is_stored(cls):
         return cls.get_module_config("message", "is_stored")
 
+    @classmethod
+    def set_message_stream_batch_size(cls, stream_batch_size=128):
+        cls.set_module_config("message", "stream_batch_size", stream_batch_size)
+
+    @classmethod
+    def get_message_stream_batch_size(cls):
+        return cls.get_module_config("message", "stream_batch_size")
+
     @classmethod
     def set_message_is_show_in_terminal(cls, is_show_in_terminal=True):
         cls.set_module_config("message", "is_show_in_terminal", is_show_in_terminal)

oxygent/mas.py

@@ -46,6 +46,7 @@
 from .utils.common_utils import (
     generate_uuid,
     get_format_time,
+    get_timestamp,
     msgpack_preprocess,
     print_tree,
     to_json,
@@ -100,6 +101,8 @@ class MAS(BaseModel):
     routers: list = Field(default_factory=list)
     middlewares: list = Field(default_factory=list)
 
+    stream_dict: dict[str, list] = Field(default_factory=dict)
+
     def __init__(self, **kwargs):
         """Construct a new :class:`MAS`.
 
@@ -284,9 +287,14 @@ async def init_db(self):
                     "mappings": {
                         "properties": {
                             "message_id": {"type": "keyword"},
+                            "group_id": {"type": "keyword"},
                             "trace_id": {"type": "keyword"},
+                            "node_id": {"type": "keyword"},
+                            "node_name": {"type": "keyword"},
                             "message": {"type": "text"},
                             "message_type": {"type": "keyword"},
+                            "message_event": {"type": "keyword"},
+                            "message_timestamp": {"type": "long"},
                             "create_time": {
                                 "format": "yyyy-MM-dd HH:mm:ss.SSSSSSSSS",
                                 "type": "date",
@@ -566,7 +574,9 @@ async def call(self, callee, arguments, **kwargs):
         oxy_response = await oxy.execute(oxy_request)
         return oxy_response.output
 
-    async def send_message(self, sse_message: SSEMessage, redis_key: str):
+    async def send_message(
+        self, sse_message: SSEMessage, redis_key: str, group_id: str = ""
+    ):
         """Push *message* onto a capped Redis list.
 
         The data is MsgPack‑encoded before being stored.  At most **10** items
@@ -598,18 +608,78 @@ async def send_message(self, sse_message: SSEMessage, redis_key: str):
             parts = redis_key.split(":")
             current_trace_id = parts[-1] if len(parts) >= 3 else ""
 
-            # Insert into Elasticsearch
-            await self.es_client.index(
-                Config.get_app_name() + "_message",
-                doc_id=sse_message.id,
-                body={
-                    "message_id": sse_message.id,
-                    "trace_id": current_trace_id,
-                    "message": to_json(message),
-                    "message_type": message_type,
-                    "create_time": get_format_time(),
-                },
-            )
+            # 考虑 message 是 str 的情况
+            node_id = ""
+            node_name = ""
+            message_timestamp = get_timestamp()
+            if isinstance(message, dict):
+                message_timestamp = message.get("timestamp", get_timestamp())
+                if isinstance(message.get("content"), dict):
+                    node_id = message.get("content", {}).get("node_id", "")
+                    node_name = message.get("content", {}).get("agent", "")
+
+            if message_type in ["stream", "stream_end"]:
+                # 排队
+                if message_type == "stream":
+                    delta = message.get("content", {}).get("delta", "")
+                    if node_id not in self.stream_dict:
+                        self.stream_dict[node_id] = []
+                    self.stream_dict[node_id].append(delta)
+                if message_type == "stream_end" or (
+                    self.stream_dict[node_id]
+                    and len(self.stream_dict[node_id])
+                    % Config.get_message_stream_batch_size()
+                    == 0
+                ):
+                    message_id = generate_uuid()
+                    merged_type = "merged_stream"
+                    save_message_task = asyncio.create_task(
+                        self.es_client.index(
+                            Config.get_app_name() + "_message",
+                            doc_id=message_id,
+                            body={
+                                "message_id": message_id,
+                                "group_id": group_id,
+                                "trace_id": current_trace_id,
+                                "node_id": node_id,
+                                "node_name": node_name,
+                                "message": to_json(
+                                    {
+                                        "type": merged_type,
+                                        "content": "".join(self.stream_dict[node_id]),
+                                    }
+                                ),
+                                "message_type": merged_type,
+                                "message_event": sse_message.event,
+                                "message_timestamp": message_timestamp,
+                                "create_time": get_format_time(),
+                            },
+                        )
+                    )
+                    save_message_task.add_done_callback(self.background_tasks.discard)
+                    self.background_tasks.add(save_message_task)
+                    self.stream_dict[node_id].clear()
+            else:
+                save_message_task = asyncio.create_task(
+                    self.es_client.index(
+                        Config.get_app_name() + "_message",
+                        doc_id=sse_message.id,
+                        body={
+                            "message_id": sse_message.id,
+                            "group_id": group_id,
+                            "trace_id": current_trace_id,
+                            "node_id": node_id,
+                            "node_name": node_name,
+                            "message": to_json(message),
+                            "message_type": message_type,
+                            "message_event": sse_message.event,
+                            "message_timestamp": message_timestamp,
+                            "create_time": get_format_time(),
+                        },
+                    )
+                )
+                save_message_task.add_done_callback(self.background_tasks.discard)
+                self.background_tasks.add(save_message_task)
         if message_is_send:
             bytes_msg = msgpack.packb(msgpack_preprocess(sse_message.to_sse()))
             await self.redis_client.lpush(redis_key, bytes_msg)
@@ -745,7 +815,9 @@ async def chat_with_agent(
 
             if send_msg_key:
                 await self.send_message(
-                    SSEMessage(event="close", data="done"), send_msg_key
+                    SSEMessage(event="close", data="done"),
+                    send_msg_key,
+                    group_id=oxy_request.group_id,
                 )
             return oxy_response
         except Exception:

oxygent/oxy/agents/sse_oxy_agent.py

@@ -45,6 +45,8 @@ async def _execute(self, oxy_request: OxyRequest) -> OxyResponse:
             del payload["node_id_stack"]
             payload["caller"] = "user"
         del payload["arguments"]
+        if "shared_data" in payload and "_headers" in payload["shared_data"]:
+            del payload["shared_data"]["_headers"]
 
         url = build_url(self.server_url, "/sse/chat")
         answer = ""

oxygent/oxy/llms/http_llm.py

@@ -147,9 +147,18 @@ async def _execute(self, oxy_request: OxyRequest) -> OxyResponse:
                                         "agent": oxy_request.caller,
                                         "node_id": oxy_request.node_id,
                                     },
-                                    "_is_stored": False,
                                 }
                             )
+                await oxy_request.send_message(
+                    {
+                        "type": "stream_end",
+                        "content": {
+                            "delta": "",
+                            "agent": oxy_request.caller,
+                            "node_id": oxy_request.node_id,
+                        },
+                    }
+                )
             result = "".join(result_parts)
             return OxyResponse(state=OxyState.COMPLETED, output=result)
 

oxygent/oxy/llms/openai_llm.py

@@ -83,7 +83,6 @@ async def _execute(self, oxy_request: OxyRequest) -> OxyResponse:
                                     "agent": oxy_request.caller,
                                     "node_id": oxy_request.node_id,
                                 },
-                                "_is_stored": False,
                             }
                         )
                         answer += "<think>"
@@ -100,7 +99,6 @@ async def _execute(self, oxy_request: OxyRequest) -> OxyResponse:
                                     "agent": oxy_request.caller,
                                     "node_id": oxy_request.node_id,
                                 },
-                                "_is_stored": False,
                             }
                         )
                         answer += "</think>"
@@ -116,9 +114,18 @@ async def _execute(self, oxy_request: OxyRequest) -> OxyResponse:
                                 "agent": oxy_request.caller,
                                 "node_id": oxy_request.node_id,
                             },
-                            "_is_stored": False,
                         }
                     )
+            await oxy_request.send_message(
+                {
+                    "type": "stream_end",
+                    "content": {
+                        "delta": "",
+                        "agent": oxy_request.caller,
+                        "node_id": oxy_request.node_id,
+                    },
+                }
+            )
             return OxyResponse(state=OxyState.COMPLETED, output=answer)
         else:
             return OxyResponse(

oxygent/schemas/oxy.py

@@ -361,7 +361,7 @@ async def send_message(self, message=None, event=None, id=None):
             redis_key = (
                 f"{self.mas.message_prefix}:{self.mas.name}:{self.current_trace_id}"
             )
-            await self.mas.send_message(sse_message, redis_key)
+            await self.mas.send_message(sse_message, redis_key, group_id=self.group_id)
 
     def set_query(self, query, master_level=False):
         if master_level:

oxygent/utils/common_utils.py

@@ -39,10 +39,14 @@ def get_mac_address():
     return mac_address
 
 
-def get_timestamp():
+def get_timestamp_str():
     return str(datetime.now().timestamp())
 
 
+def get_timestamp():
+    return datetime.now().timestamp() * 1000
+
+
 def get_format_time():
     """Yyyy-MM-dd HH:mm:ss."""
     """yyyy-MM-dd HH:mm:ss.SSS"""

test/unittest/test_oxy_request.py

@@ -23,7 +23,7 @@ def __init__(self):
         self.message_prefix = "msg"
         self.name = "ut_mas"
 
-    async def send_message(self, message, redis_key):
+    async def send_message(self, message, redis_key, group_id=""):
         self.last_msg = (redis_key, message)