WIN-JHFT4D3SIVT\caoxiaozhu
5667190abe
- Phase 1: Infrastructure (state, prompts, registry) - Phase 2: Execution engine (AI adapters, security classifier, executors) - Phase 3: Agent integration (graph nodes, routing) - Phase 4: Streaming interaction (PTY terminal, WebSocket) - Phase 5: Frontend integration (Vue components)
87 lines
2.3 KiB
Python
87 lines
2.3 KiB
Python
"""
|
|
StreamOutput - 流式输出封装
|
|
"""
|
|
|
|
import json
|
|
from datetime import datetime, timezone
|
|
from typing import Any, AsyncGenerator, Callable
|
|
|
|
from dataclasses import dataclass
|
|
|
|
|
|
@dataclass
|
|
class StreamEvent:
|
|
"""流式事件"""
|
|
|
|
type: str # "output" | "error" | "status" | "complete"
|
|
session_id: str
|
|
data: str
|
|
timestamp: str
|
|
|
|
|
|
class StreamOutput:
|
|
"""流式输出处理器"""
|
|
|
|
def __init__(
|
|
self,
|
|
session_id: str,
|
|
websocket_sender: Callable[[str, str], Any] | None = None,
|
|
):
|
|
self.session_id = session_id
|
|
self.websocket_sender = websocket_sender
|
|
self._listeners: list[Callable[[StreamEvent], Any]] = []
|
|
|
|
def add_listener(self, listener: Callable[[StreamEvent], Any]):
|
|
"""添加事件监听器"""
|
|
self._listeners.append(listener)
|
|
|
|
def remove_listener(self, listener: Callable[[StreamEvent], Any]):
|
|
"""移除事件监听器"""
|
|
if listener in self._listeners:
|
|
self._listeners.remove(listener)
|
|
|
|
async def push(self, event_type: str, data: str):
|
|
"""推送事件到 WebSocket 和监听器"""
|
|
event = StreamEvent(
|
|
type=event_type,
|
|
session_id=self.session_id,
|
|
data=data,
|
|
timestamp=datetime.now(timezone.utc).isoformat(),
|
|
)
|
|
|
|
# 发送到 WebSocket
|
|
if self.websocket_sender:
|
|
try:
|
|
await self.websocket_sender(self.session_id, json.dumps(event.__dict__))
|
|
except Exception:
|
|
pass
|
|
|
|
# 通知监听器
|
|
for listener in self._listeners:
|
|
try:
|
|
result = listener(event)
|
|
if hasattr(result, "__await__"):
|
|
await result
|
|
except Exception:
|
|
pass
|
|
|
|
async def stream_execution(
|
|
self,
|
|
executor,
|
|
prompt: str,
|
|
) -> AsyncGenerator[str, None]:
|
|
"""包装执行器，实现流式输出"""
|
|
async for line in executor.execute(prompt):
|
|
await self.push("output", line)
|
|
yield line
|
|
|
|
await self.push("complete", "")
|
|
|
|
async def push_status(self, status: str):
|
|
"""推送状态消息"""
|
|
await self.push("status", status)
|
|
|
|
async def push_error(self, error: str):
|
|
"""推送错误消息"""
|
|
await self.push("error", error)
|