| """ |
| 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): |
| |
| 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 |
|
|
| |
| 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: |
| |
| 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) |
| 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) |
|
|
