diff --git a/backend/package/yuxi/channels/adapters/signal/streaming.py b/backend/package/yuxi/channels/adapters/signal/streaming.py new file mode 100644 index 00000000..831e65b0 --- /dev/null +++ b/backend/package/yuxi/channels/adapters/signal/streaming.py @@ -0,0 +1,97 @@ +from __future__ import annotations + +import logging +import time + +logger = logging.getLogger(__name__) + +STREAM_BUFFER_MAX_ENTRIES = 100 +STREAM_BUFFER_TTL = 300 + + +class StreamBuffer: + def __init__(self, max_entries: int = STREAM_BUFFER_MAX_ENTRIES, ttl: float = STREAM_BUFFER_TTL): + self._buffer: dict[str, tuple[str, float]] = {} + self._max_entries = max_entries + self._ttl = ttl + + def get(self, chat_id: str) -> str: + buffered, _ = self._buffer.get(chat_id, ("", 0)) + return buffered + + def set(self, chat_id: str, content: str) -> None: + self._buffer[chat_id] = (content, time.monotonic()) + + def pop(self, chat_id: str) -> str | None: + entry = self._buffer.pop(chat_id, None) + return entry[0] if entry else None + + def idle_ms(self, chat_id: str) -> float: + _, last_ts = self._buffer.get(chat_id, ("", 0)) + if not last_ts: + return 0 + return (time.monotonic() - last_ts) * 1000 + + def prune(self) -> None: + now = time.monotonic() + expired = [k for k, (_, ts) in self._buffer.items() if now - ts > self._ttl] + for k in expired: + self._buffer.pop(k, None) + + if len(self._buffer) > self._max_entries: + sorted_keys = sorted( + self._buffer.keys(), + key=lambda k: self._buffer[k][1], + ) + excess = len(self._buffer) - self._max_entries + for k in sorted_keys[:excess]: + self._buffer.pop(k, None) + + def clear(self) -> None: + self._buffer.clear() + + +class StreamCoalescer: + def __init__( + self, + buffer: StreamBuffer, + min_chars: int = 1500, + idle_ms: int = 1000, + lane_separator: str = "---", + ): + self._buffer = buffer + self._min_chars = min_chars + self._idle_ms = idle_ms + self._lane_separator = lane_separator + + def should_flush_on_idle(self, chat_id: str) -> bool: + buffered = self._buffer.get(chat_id) + if not buffered: + return False + return self._buffer.idle_ms(chat_id) >= self._idle_ms + + def should_flush_on_size(self, chat_id: str, new_chunk: str = "") -> bool: + accumulated = self._buffer.get(chat_id) + new_chunk + return len(accumulated) >= self._min_chars + + def format_reasoning(self, chunk: str) -> str: + reasoning_content = chunk.replace("lane:reasoning:", "") + sep = self._lane_separator + return f"(thinking)\n{reasoning_content}\n{sep}\n" + + @staticmethod + def is_reasoning_chunk(chunk: str) -> bool: + return "lane:reasoning:" in chunk + + +class ReasoningLaneFormatter: + def __init__(self, lane_separator: str = "---"): + self._lane_separator = lane_separator + + def format(self, chunk: str) -> str: + reasoning_content = chunk.replace("lane:reasoning:", "") + return f"(thinking)\n{reasoning_content}\n{self._lane_separator}\n" + + @staticmethod + def is_reasoning(chunk: str) -> bool: + return "lane:reasoning:" in chunk