import re
from typing import AsyncGenerator, Dict, Any, Optional
class StreamProcessor:
"""Handles logical processing of LLM streams: thinking tags and content routing."""
def __init__(self, profile: Any):
self._in_thinking_tag = False
self.tag_buffer = ""
self.prefix_buffer = "" # Accumulates start of content to catch split headers
self.header_sent = False
self.profile = profile
def reset_turn(self):
self.header_sent = False
self.prefix_buffer = ""
async def process_chunk(self, content_chunk: str, turn: int) -> AsyncGenerator[Dict[str, Any], None]:
"""Processes a raw content chunk, yields UI events (content, reasoning)."""
self.tag_buffer += content_chunk
turn_header = f"---\n### 🛰️ **[Turn {turn}] Thinking...**\n"
while self.tag_buffer:
if not self._in_thinking_tag:
lower_buf = self.tag_buffer.lower()
start_tag_idx = lower_buf.find("<thinking>")
if start_tag_idx != -1:
strategy_part = self.tag_buffer[:start_tag_idx]
if strategy_part:
# Flush prefix buffer if any before tag
async for event in self._flush_prefix(strategy_part, turn_header):
yield event
self._in_thinking_tag = True
self.tag_buffer = self.tag_buffer[start_tag_idx + len("<thinking>"):]
continue
potential_idx = self.tag_buffer.find("<")
if potential_idx != -1 and "<thinking>".startswith(self.tag_buffer[potential_idx:].lower()):
strategy_part = self.tag_buffer[:potential_idx]
if strategy_part:
async for event in self._flush_prefix(strategy_part, turn_header):
yield event
self.tag_buffer = self.tag_buffer[potential_idx:]
break
else:
# Not in thinking tag, accumulate in prefix buffer if header not sent
if not self.header_sent:
self.prefix_buffer += self.tag_buffer
self.tag_buffer = ""
# If buffer is large enough, or we have a reason to flush
if len(self.prefix_buffer) > 200:
async for event in self._flush_prefix("", turn_header):
yield event
else:
strategy_part = self._apply_turn_header(self.tag_buffer, turn_header)
if strategy_part:
yield {"type": "content", "content": strategy_part}
self.tag_buffer = ""
else:
# Inside thinking tag - reasoning should not be prefix-buffered
# But if we were buffering prefix, flush it now
if not self.header_sent and self.prefix_buffer:
async for event in self._flush_prefix("", turn_header):
yield event
lower_buf = self.tag_buffer.lower()
end_tag_idx = lower_buf.find("</thinking>")
if end_tag_idx != -1:
reasoning_part = self.tag_buffer[:end_tag_idx]
if reasoning_part:
yield {"type": "reasoning", "content": reasoning_part}
self._in_thinking_tag = False
self.tag_buffer = self.tag_buffer[end_tag_idx + len("</thinking>"):]
continue
potential_idx = self.tag_buffer.find("<")
if potential_idx != -1 and "</thinking>".startswith(self.tag_buffer[potential_idx:].lower()):
reasoning_part = self.tag_buffer[:potential_idx]
if reasoning_part:
yield {"type": "reasoning", "content": reasoning_part}
self.tag_buffer = self.tag_buffer[potential_idx:]
break
else:
yield {"type": "reasoning", "content": self.tag_buffer}
self.tag_buffer = ""
async def _flush_prefix(self, extra_text: str, header: str) -> AsyncGenerator[Dict[str, Any], None]:
full_text = self.prefix_buffer + extra_text
self.prefix_buffer = "" # Clear it
processed = self._apply_turn_header(full_text, header)
if processed:
yield {"type": "content", "content": processed}
def _apply_turn_header(self, text: str, header: str) -> Optional[str]:
# List of patterns to strip (hallucinated headers from various LLM generations)
strip_patterns = [
r"(?i)^.*\[Turn\s*\d+\].*$",
r"(?i)Turn\s*\d+:\s*architecting\s*next\s*step\.*",
r"(?i)🏗️\s*BRIDGE\s*ANALYSIS:?",
r"(?i)Analysis: ",
r"(?i)---"
]
if self.profile.silent_stream:
# Aggressively strip any hallucinated headers defined in the profile
for pattern in self.profile.strip_headers:
text = re.sub(pattern, "", text, flags=re.IGNORECASE | re.MULTILINE)
return text
# Even in non-silent mode, strip the patterns from the LLM's raw text to avoid duplication
for pattern in strip_patterns:
text = re.sub(pattern, "", text, flags=re.IGNORECASE | re.MULTILINE)
# Dynamic Title Extraction
if not self.header_sent:
# 1. Primary: Look for "Title: [Title Text]"
title_match = re.search(r"(?i)Title:\s*(.*?)\n", text)
if title_match:
custom_title = title_match.group(1).strip()
header = re.sub(r"Thinking\.\.\.", custom_title, header)
text = text[:title_match.start()] + text[title_match.end():]
else:
# 2. Secondary: If AI uses a Markdown header like "### 🚀 My Title"
md_header_match = re.search(r"(?i)^###\s*.*?\s*(.*?)\n", text, re.MULTILINE)
if md_header_match:
custom_title = md_header_match.group(1).strip()
header = re.sub(r"Thinking\.\.\.", custom_title, header)
text = text[:md_header_match.start()] + text[md_header_match.end():]
if not text.strip():
# If after stripping the text is empty, don't send anything yet (unless it's the very first chunk)
if not self.header_sent:
return None
return ""
if not self.header_sent:
self.header_sent = True
# Prepend the system's authoritative header
return header + text.lstrip()
return text
async def end_stream(self, turn: int) -> AsyncGenerator[Dict[str, Any], None]:
"""Flushes any remaining buffered text at the very end of the stream."""
turn_header = f"---\n### 🛰️ **[Turn {turn}] Thinking...**\n"
# 1. Flush prefix buffer if we haven't sent the header yet
if not self.header_sent and self.prefix_buffer:
async for event in self._flush_prefix("", turn_header):
yield event
# 2. Flush any leftover tag buffer (rarely happens if LLM cuts off)
if self.tag_buffer:
if self._in_thinking_tag:
yield {"type": "reasoning", "content": self.tag_buffer}
else:
processed = self._apply_turn_header(self.tag_buffer, turn_header)
if processed:
yield {"type": "content", "content": processed}
self.tag_buffer = ""
