98 lines
3.0 KiB
Python
98 lines
3.0 KiB
Python
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
|