mirror of
https://github.com/agentscope-ai/ReMe.git
synced 2026-09-15 23:31:05 +00:00
62 lines
2.4 KiB
Python
62 lines
2.4 KiB
Python
"""Streaming chat for test."""
|
|
|
|
from loguru import logger
|
|
|
|
from ..core.enumeration import Role, ChunkEnum
|
|
from ..core.op import BaseTool
|
|
from ..core.schema import Message, ToolCall
|
|
|
|
|
|
class StreamChat(BaseTool):
|
|
"""Streaming chat agent that handles real-time conversation streaming."""
|
|
|
|
def _build_tool_call(self) -> ToolCall:
|
|
return ToolCall(
|
|
**{
|
|
"description": "simple chat agent",
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {
|
|
"query": {
|
|
"type": "string",
|
|
"description": "query",
|
|
},
|
|
"messages": {
|
|
"type": "array",
|
|
"items": {
|
|
"type": "object",
|
|
"properties": {
|
|
"role": {
|
|
"type": "string",
|
|
"description": "role",
|
|
},
|
|
"content": {
|
|
"type": "string",
|
|
"description": "content",
|
|
},
|
|
},
|
|
"required": ["role", "content"],
|
|
},
|
|
},
|
|
},
|
|
"required": [],
|
|
},
|
|
},
|
|
)
|
|
|
|
async def execute(self):
|
|
"""Execute streaming chat operation with query or messages."""
|
|
if "query" in self.context:
|
|
messages = [
|
|
Message(role=Role.SYSTEM, content="You are a helpful assistant."),
|
|
Message(role=Role.USER, content=self.context.query),
|
|
]
|
|
elif "messages" in self.context:
|
|
messages = [Message(**m) if isinstance(m, dict) else m for m in self.context.messages if m]
|
|
else:
|
|
raise ValueError("query or messages must be provided!")
|
|
logger.info(f"messages={messages}")
|
|
|
|
async for stream_chunk in self.llm.stream_chat(messages):
|
|
if stream_chunk.chunk_type in [ChunkEnum.ANSWER, ChunkEnum.THINK, ChunkEnum.ERROR, ChunkEnum.TOOL]:
|
|
await self.context.add_stream_chunk(stream_chunk)
|