bvandeusen
b11a92f32d
- stream_chat: add think=False parameter passed through to Ollama payload. qwen3 models have thinking enabled by default; without this flag the model spends minutes generating internal thinking tokens that stream_chat silently discards, leaving the frontend spinner blank until the SSE connection times out and the widget disappears. - create_assist_buffer: orphan (overwrite) a still-running buffer instead of raising. The old asyncio task holds a direct reference and completes harmlessly against the stale buffer. New requests always win. - assist_route: remove the 409 guard that blocked new requests when a previous generation got stuck. create_assist_buffer now handles this transparently. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
139 lines
4.5 KiB
Python
139 lines
4.5 KiB
Python
"""In-memory generation event buffer for SSE fan-out.
|
|
|
|
Each active generation gets a GenerationBuffer keyed by conversation_id.
|
|
SSE clients tail the buffer and can reconnect at any point without data loss.
|
|
"""
|
|
|
|
import asyncio
|
|
import json
|
|
import logging
|
|
import time
|
|
from dataclasses import dataclass, field
|
|
from enum import Enum
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class GenerationState(str, Enum):
|
|
RUNNING = "running"
|
|
COMPLETED = "completed"
|
|
ERRORED = "errored"
|
|
|
|
|
|
@dataclass
|
|
class SSEEvent:
|
|
index: int
|
|
event_type: str # "context", "chunk", "done", "error"
|
|
data: dict
|
|
|
|
|
|
@dataclass
|
|
class GenerationBuffer:
|
|
conversation_id: int
|
|
assistant_message_id: int
|
|
state: GenerationState = GenerationState.RUNNING
|
|
events: list[SSEEvent] = field(default_factory=list)
|
|
content_so_far: str = ""
|
|
finished_at: float | None = None
|
|
cancel_event: asyncio.Event = field(default_factory=asyncio.Event)
|
|
_notify: asyncio.Event = field(default_factory=asyncio.Event)
|
|
# Tool confirmation — set when a write tool is awaiting user approval
|
|
confirmation_future: asyncio.Future | None = None
|
|
pending_tool: dict | None = None
|
|
|
|
def append_event(self, event_type: str, data: dict) -> SSEEvent:
|
|
event = SSEEvent(index=len(self.events), event_type=event_type, data=data)
|
|
self.events.append(event)
|
|
# Wake all waiting SSE clients, then reset for next wait
|
|
old = self._notify
|
|
self._notify = asyncio.Event()
|
|
old.set()
|
|
return event
|
|
|
|
async def wait_for_event(self, after_index: int, timeout: float = 30.0) -> bool:
|
|
"""Wait until there are events beyond after_index. Returns True if new events exist."""
|
|
if after_index + 1 < len(self.events):
|
|
return True
|
|
try:
|
|
await asyncio.wait_for(self._notify.wait(), timeout=timeout)
|
|
return True
|
|
except asyncio.TimeoutError:
|
|
return False
|
|
|
|
def events_after(self, index: int) -> list[SSEEvent]:
|
|
"""Return events with index > given index (for replay on reconnect)."""
|
|
start = index + 1
|
|
if start >= len(self.events):
|
|
return []
|
|
return self.events[start:]
|
|
|
|
@staticmethod
|
|
def format_sse(event: SSEEvent) -> str:
|
|
return f"id: {event.index}\nevent: {event.event_type}\ndata: {json.dumps(event.data)}\n\n"
|
|
|
|
|
|
# Module-level singleton registry
|
|
_buffers: dict[int | str, GenerationBuffer] = {}
|
|
|
|
_cleanup_task: asyncio.Task | None = None
|
|
_GRACE_PERIOD = 60.0 # seconds to keep completed buffers
|
|
|
|
|
|
def create_buffer(conv_id: int, msg_id: int) -> GenerationBuffer:
|
|
existing = _buffers.get(conv_id)
|
|
if existing and existing.state == GenerationState.RUNNING:
|
|
raise RuntimeError(f"Generation already running for conversation {conv_id}")
|
|
buf = GenerationBuffer(conversation_id=conv_id, assistant_message_id=msg_id)
|
|
_buffers[conv_id] = buf
|
|
return buf
|
|
|
|
|
|
def get_buffer(conv_id: int) -> GenerationBuffer | None:
|
|
return _buffers.get(conv_id)
|
|
|
|
|
|
def remove_buffer(conv_id: int) -> None:
|
|
_buffers.pop(conv_id, None)
|
|
|
|
|
|
def create_assist_buffer(user_id: int) -> GenerationBuffer:
|
|
key = f"assist:{user_id}"
|
|
existing = _buffers.get(key)
|
|
if existing and existing.state == GenerationState.RUNNING:
|
|
# Orphan the old buffer — the background task holds a direct reference
|
|
# and will complete against it harmlessly. A new request always wins.
|
|
logger.warning("Assist generation still running for user %d; orphaning old buffer", user_id)
|
|
buf = GenerationBuffer(conversation_id=0, assistant_message_id=0)
|
|
_buffers[key] = buf
|
|
return buf
|
|
|
|
|
|
def get_assist_buffer(user_id: int) -> GenerationBuffer | None:
|
|
return _buffers.get(f"assist:{user_id}")
|
|
|
|
|
|
def remove_assist_buffer(user_id: int) -> None:
|
|
_buffers.pop(f"assist:{user_id}", None)
|
|
|
|
|
|
async def _cleanup_loop() -> None:
|
|
"""Remove completed/errored buffers after grace period."""
|
|
while True:
|
|
await asyncio.sleep(15)
|
|
now = time.monotonic()
|
|
to_remove = [
|
|
key for key, buf in _buffers.items()
|
|
if buf.state != GenerationState.RUNNING
|
|
and buf.finished_at is not None
|
|
and now - buf.finished_at > _GRACE_PERIOD
|
|
]
|
|
for key in to_remove:
|
|
_buffers.pop(key, None)
|
|
logger.debug("Cleaned up generation buffer for key %s", key)
|
|
|
|
|
|
def start_cleanup_loop() -> None:
|
|
global _cleanup_task
|
|
if _cleanup_task is None or _cleanup_task.done():
|
|
_cleanup_task = asyncio.create_task(_cleanup_loop())
|