File size: 5,069 Bytes
b2c86fd | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 | """
StreamingCallback：拦截 agent 的输出，转成结构化 SSE 事件
后端通过 patch agent 的 print 输出方法捕获步骤事件，并转成 SSE。
"""
import io
import queue
import re
import sys
class StreamingCallback:
"""
拦截 agent.go() 的 stdout 输出，解析成结构化事件。
Biomni 的 agent 通过 print() 输出中间步骤：
- "Thought: ..." → thinking 事件
- "```python\n...\n```" → code 事件
- "Observation: ..." → observation 事件
- 最终答案 → result 事件
"""
def __init__(self, event_queue: queue.Queue):
self.q = event_queue
self._original_stdout = None
self._buf = io.StringIO()
self._in_solution = False
self._solution_lines: list[str] = []
def attach(self, _agent):
"""patch sys.stdout 来捕获 agent 的 print 输出"""
self._original_stdout = sys.stdout
sys.stdout = _InterceptWriter(self._parse_and_emit, self._original_stdout)
def detach(self, _agent):
# 兜底：若流结束时仍在 solution 块内，仍然输出 result
self._flush_solution_if_any()
if self._original_stdout:
sys.stdout = self._original_stdout
def _parse_and_emit(self, text: str):
"""解析 agent 输出文本，转成结构化事件"""
raw_text = text.rstrip("\n")
if self._handle_solution_markup(raw_text):
return
text = raw_text.strip()
if not text:
return
# 匹配 Biomni A1 的输出格式
if text.startswith("Thought:") or text.startswith("[Thinking]"):
self.q.put({"type": "thinking", "content": text})
elif "```python" in text or "```r" in text or "```bash" in text:
# 提取代码块
match = re.search(r"```(\w+)\n(.*?)```", text, re.DOTALL)
if match:
lang, code = match.group(1), match.group(2)
self.q.put({"type": "code", "lang": lang, "content": code})
else:
self.q.put({"type": "thinking", "content": text})
elif text.startswith("Observation:") or text.startswith("[Output]"):
self.q.put({"type": "observation", "content": text})
elif text.startswith("Action:") or "tool" in text.lower()[:30]:
# 尝试提取工具名
tool_match = re.search(r"(\w+)\(", text)
tool_name = tool_match.group(1) if tool_match else "tool"
self.q.put({"type": "tool_use", "tool": tool_name, "content": text})
elif text.startswith("Final Answer:") or text.startswith("Answer:"):
self.q.put({"type": "result", "content": text})
else:
# 其他输出作为 thinking
self.q.put({"type": "thinking", "content": text})
def _handle_solution_markup(self, raw_text: str) -> bool:
"""
解析 <solution>...</solution> 多行块。
命中时输出 type=result，并返回 True 表示该行已被消费。
"""
if self._in_solution:
if "</solution>" in raw_text:
before_end, _after_end = raw_text.split("</solution>", 1)
if before_end:
self._solution_lines.append(before_end)
self._flush_solution_if_any()
self._in_solution = False
else:
self._solution_lines.append(raw_text)
return True
if "<solution>" in raw_text:
_before_start, after_start = raw_text.split("<solution>", 1)
self._in_solution = True
self._solution_lines = []
if "</solution>" in after_start:
content, _after_end = after_start.split("</solution>", 1)
if content:
self._solution_lines.append(content)
self._flush_solution_if_any()
self._in_solution = False
else:
if after_start:
self._solution_lines.append(after_start)
return True
return False
def _flush_solution_if_any(self):
content = "\n".join(self._solution_lines).strip()
self._solution_lines = []
if content:
self.q.put({"type": "result", "content": content})
class _InterceptWriter:
"""替换 sys.stdout，拦截写入"""
def __init__(self, callback, original):
self._cb = callback
self._orig = original
self._line_buf = ""
def write(self, text):
self._orig.write(text) # 同时写到真实 stdout（方便调试）
self._line_buf += text
if "\n" in self._line_buf:
lines = self._line_buf.split("\n")
for line in lines[:-1]:
self._cb(line)
self._line_buf = lines[-1]
def flush(self):
if self._line_buf:
self._cb(self._line_buf)
self._line_buf = ""
self._orig.flush()
def __getattr__(self, name):
return getattr(self._orig, name)
|