mirror of
https://github.com/imartinez/privateGPT.git
synced 2026-07-17 20:03:12 +00:00
620 lines
23 KiB
Python
620 lines
23 KiB
Python
import asyncio
|
|
import contextlib
|
|
import logging
|
|
from collections import defaultdict
|
|
from collections.abc import AsyncGenerator
|
|
from typing import TYPE_CHECKING, Final
|
|
|
|
from injector import inject, singleton
|
|
from pydantic import BaseModel
|
|
|
|
from private_gpt.components.streaming.providers.models import StreamStatus
|
|
from private_gpt.components.streaming.stream.event_handler import EventHandler
|
|
from private_gpt.components.streaming.stream_component import StreamComponent
|
|
from private_gpt.settings.settings import Settings
|
|
|
|
if TYPE_CHECKING:
|
|
from private_gpt.components.streaming.providers.models import StreamMetadata
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
DEFAULT_BLOCK_MS: Final[int] = 1000
|
|
BATCH_SIZE: Final[int] = 5
|
|
# How often (seconds) to poll for terminal status when no events arrive.
|
|
# Halves Redis metadata round-trips vs. checking every block cycle (1 s).
|
|
STATUS_CHECK_INTERVAL: Final[float] = 2.0
|
|
TERMINAL_STATUSES: Final[set[StreamStatus]] = {
|
|
StreamStatus.COMPLETED,
|
|
StreamStatus.CANCELLED,
|
|
StreamStatus.ERROR,
|
|
}
|
|
|
|
|
|
class StreamBroadcast:
|
|
"""Append-only batch log that broadcasts to N concurrent readers.
|
|
|
|
A single producer calls ``push`` / ``close`` (both synchronous); each
|
|
subscriber holds an integer *cursor* into ``_batches`` and parks on a
|
|
``asyncio.Future`` until new data arrives. One ``push`` resolves *all*
|
|
pending futures in a single pass — no per-subscriber queues.
|
|
|
|
Safety: all operations are synchronous except ``read_from``, so no lock is
|
|
needed (asyncio cooperative threading guarantees no interleaving between
|
|
non-await statements).
|
|
"""
|
|
|
|
def __init__(self) -> None:
|
|
self._batches: list[list[BaseModel] | None] = []
|
|
self._waiters: list[asyncio.Future[None]] = []
|
|
|
|
def _notify(self) -> None:
|
|
waiters, self._waiters = self._waiters, []
|
|
for fut in waiters:
|
|
if not fut.done():
|
|
fut.set_result(None)
|
|
|
|
def push(self, batch: list[BaseModel]) -> None:
|
|
self._batches.append(batch)
|
|
self._notify()
|
|
|
|
def close(self) -> None:
|
|
"""Append the terminal sentinel and wake all parked readers."""
|
|
self._batches.append(None)
|
|
self._notify()
|
|
|
|
def next_cursor(self) -> int:
|
|
"""Cursor for a subscriber joining right now (no backfill)."""
|
|
return len(self._batches)
|
|
|
|
@property
|
|
def is_closed(self) -> bool:
|
|
return bool(self._batches) and self._batches[-1] is None
|
|
|
|
async def read_from(
|
|
self,
|
|
cursor: int,
|
|
stop_event: asyncio.Event | None = None,
|
|
) -> AsyncGenerator[BaseModel, None]:
|
|
"""Yield events starting at *cursor*, blocking until new ones arrive."""
|
|
while True:
|
|
# Drain already-available batches. Safe without a lock: _batches is
|
|
# append-only and list operations are atomic between await points.
|
|
while cursor < len(self._batches):
|
|
batch = self._batches[cursor]
|
|
cursor += 1
|
|
if batch is None:
|
|
return
|
|
for event in batch:
|
|
yield event
|
|
|
|
if stop_event and stop_event.is_set():
|
|
return
|
|
|
|
# Park until the producer pushes more data.
|
|
fut: asyncio.Future[None] = asyncio.get_event_loop().create_future()
|
|
self._waiters.append(fut)
|
|
try:
|
|
await asyncio.wait_for(fut, timeout=1.0)
|
|
except TimeoutError:
|
|
pass # Re-check stop_event and drain any new batches
|
|
finally:
|
|
# Clean up if push() hasn't already swapped us out.
|
|
if fut in self._waiters:
|
|
self._waiters.remove(fut)
|
|
|
|
|
|
class StreamConsumer:
|
|
"""A subscriber on a ``StreamBroadcast`` with an individual read cursor."""
|
|
|
|
def __init__(
|
|
self,
|
|
correlation_id: str,
|
|
event_handler: EventHandler,
|
|
broadcast: StreamBroadcast,
|
|
cursor: int,
|
|
) -> None:
|
|
self.correlation_id = correlation_id
|
|
self.event_handler = event_handler
|
|
self.broadcast = broadcast
|
|
self.cursor = cursor # Index into broadcast._batches
|
|
|
|
|
|
class StreamState:
|
|
"""Read position and broadcast log for a single multiplexed stream."""
|
|
|
|
def __init__(self, last_id: str = "0") -> None:
|
|
self.last_id = last_id
|
|
self.cached_events: list[BaseModel] | None = None
|
|
# 0.0 → check on the first idle cycle (ancient enough to pass the gate)
|
|
self.last_status_check: float = 0.0
|
|
self.broadcast = StreamBroadcast()
|
|
|
|
|
|
@singleton
|
|
class StreamReader:
|
|
"""Reads events from the backend (Redis or in-memory) and deserializes them.
|
|
|
|
On Redis, ``read_events`` with ``block_ms > 0`` issues ``XREAD BLOCK`` which
|
|
parks the connection — the event loop stays free. On the in-memory backend
|
|
the provider uses an ``asyncio.Event`` to avoid busy-waiting.
|
|
``block_ms=None`` is always non-blocking (used for catch-up / drain reads).
|
|
"""
|
|
|
|
@inject
|
|
def __init__(self, stream_component: StreamComponent) -> None:
|
|
self.stream_service = stream_component.stream
|
|
|
|
async def read_events(
|
|
self,
|
|
event_handler: EventHandler,
|
|
correlation_id: str,
|
|
last_id: str = "0",
|
|
count: int | None = None,
|
|
block_ms: int | None = None,
|
|
) -> tuple[list[BaseModel], str]:
|
|
raw_events, next_last_id = await self.stream_service.read_events(
|
|
correlation_id=correlation_id,
|
|
last_id=last_id,
|
|
count=count,
|
|
block_ms=block_ms,
|
|
)
|
|
|
|
events: list[BaseModel] = []
|
|
for raw_data in raw_events:
|
|
try:
|
|
events.append(event_handler.deserialize(raw_data))
|
|
except Exception as e:
|
|
logger.error(f"Error deserializing event: {e}")
|
|
return events, next_last_id
|
|
|
|
async def check_terminal_status(
|
|
self,
|
|
correlation_id: str,
|
|
event_handler: EventHandler,
|
|
cached_events: list[BaseModel] | None,
|
|
) -> bool:
|
|
"""Return True when the stream has reached a terminal status.
|
|
|
|
Tries the last cached event first; only hits Redis when that check is
|
|
inconclusive, avoiding a round-trip when the cached event already
|
|
carries terminal state.
|
|
"""
|
|
if cached_events:
|
|
try:
|
|
status = await event_handler.get_current_status(cached_events[-1])
|
|
if status in TERMINAL_STATUSES:
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f"Error checking cached status for {correlation_id}: {e}")
|
|
|
|
try:
|
|
metadata: (
|
|
StreamMetadata | None
|
|
) = await self.stream_service.get_stream_metadata(correlation_id)
|
|
if metadata and metadata.status in TERMINAL_STATUSES:
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f"Error checking metadata for {correlation_id}: {e}")
|
|
|
|
return False
|
|
|
|
async def stream_events(
|
|
self,
|
|
event_handler: EventHandler,
|
|
correlation_id: str,
|
|
last_id: str = "0",
|
|
block_ms: int | None = None,
|
|
stop_event: asyncio.Event | None = None,
|
|
) -> AsyncGenerator[BaseModel, None]:
|
|
"""Stream events as an async generator (direct, non-multiplexed path).
|
|
|
|
A producer task reads from the backend in blocking mode and feeds an
|
|
internal queue; the generator drains the queue and yields individual
|
|
events. After the terminal status is detected a final non-blocking
|
|
drain read is performed so no event published between the last read
|
|
and the terminal signal is lost.
|
|
"""
|
|
stop_event = stop_event or asyncio.Event()
|
|
block_ms = block_ms if block_ms is not None else DEFAULT_BLOCK_MS
|
|
event_queue: asyncio.Queue[list[BaseModel] | None] = asyncio.Queue()
|
|
|
|
async def produce_events(current_last_id: str) -> None:
|
|
try:
|
|
cached_events: list[BaseModel] | None = None
|
|
last_status_check: float = 0.0 # 0 → check on first idle cycle
|
|
|
|
while not stop_event.is_set():
|
|
events, current_last_id = await self.read_events(
|
|
event_handler=event_handler,
|
|
correlation_id=correlation_id,
|
|
last_id=current_last_id,
|
|
block_ms=block_ms,
|
|
)
|
|
|
|
if events:
|
|
cached_events = events
|
|
await event_queue.put(events)
|
|
else:
|
|
current_time = asyncio.get_event_loop().time()
|
|
if (current_time - last_status_check) >= STATUS_CHECK_INTERVAL:
|
|
if await self.check_terminal_status(
|
|
correlation_id,
|
|
event_handler,
|
|
cached_events,
|
|
):
|
|
break
|
|
last_status_check = current_time
|
|
|
|
events, _ = await self.read_events(
|
|
event_handler=event_handler,
|
|
correlation_id=correlation_id,
|
|
last_id=current_last_id,
|
|
)
|
|
if events:
|
|
await event_queue.put(events)
|
|
|
|
except Exception as e:
|
|
logger.error(f"Producer error for {correlation_id}: {e}", exc_info=True)
|
|
finally:
|
|
await event_queue.put(None)
|
|
|
|
async def get() -> AsyncGenerator[BaseModel, None]:
|
|
producer_task = asyncio.create_task(produce_events(last_id))
|
|
|
|
try:
|
|
while not stop_event.is_set():
|
|
events = await event_queue.get()
|
|
if events is None:
|
|
break
|
|
for event in events:
|
|
yield event
|
|
finally:
|
|
producer_task.cancel()
|
|
with contextlib.suppress(asyncio.CancelledError):
|
|
await producer_task
|
|
|
|
while not event_queue.empty():
|
|
events = event_queue.get_nowait()
|
|
if events is None:
|
|
break
|
|
for event in events:
|
|
yield event
|
|
|
|
return get()
|
|
|
|
|
|
class StreamMultiplexer:
|
|
"""Multiplexes many stream consumers onto a single worker task.
|
|
|
|
Each correlation id has one ``StreamState`` containing a ``StreamBroadcast``
|
|
shared by all subscribers. The worker calls ``_process_stream`` once per
|
|
active stream per cycle: one Redis read, one ``broadcast.push`` — no
|
|
per-consumer queues.
|
|
|
|
Key invariants
|
|
---------------
|
|
* ``broadcast.push`` is synchronous and resolves all parked futures in one
|
|
pass, so all N subscribers wake atomically after a single call.
|
|
* State is updated under the lock; ``broadcast.push`` happens outside it so
|
|
the lock is held for as short a time as possible.
|
|
* ``close`` on the broadcast terminates all ``read_from`` generators, making
|
|
``stop`` and terminal-status cleanup self-contained.
|
|
"""
|
|
|
|
def __init__(self, stream_reader: StreamReader) -> None:
|
|
self.stream_reader = stream_reader
|
|
self.consumers: dict[str, list[StreamConsumer]] = defaultdict(list)
|
|
self.states: dict[str, StreamState] = {}
|
|
self.worker_task: asyncio.Task[None] | None = None
|
|
self.shutdown_event = asyncio.Event()
|
|
self._lock = asyncio.Lock()
|
|
|
|
async def start(self) -> None:
|
|
if self.worker_task is None or self.worker_task.done():
|
|
self.shutdown_event.clear()
|
|
self.worker_task = asyncio.create_task(self._run_worker())
|
|
logger.info("Stream multiplexer started")
|
|
|
|
async def stop(self) -> None:
|
|
"""Cancel the worker, clear all state, and close every broadcast."""
|
|
self.shutdown_event.set()
|
|
task = self.worker_task
|
|
self.worker_task = None
|
|
if task is not None and not task.done():
|
|
task.cancel()
|
|
with contextlib.suppress(asyncio.CancelledError):
|
|
await task
|
|
|
|
async with self._lock:
|
|
states_to_close = list(self.states.values())
|
|
self.consumers.clear()
|
|
self.states.clear()
|
|
|
|
for state in states_to_close:
|
|
state.broadcast.close()
|
|
|
|
async def add_consumer(
|
|
self,
|
|
correlation_id: str,
|
|
event_handler: EventHandler,
|
|
last_id: str = "0",
|
|
) -> StreamConsumer:
|
|
"""Register a new subscriber for *correlation_id*.
|
|
|
|
Each call returns a distinct ``StreamConsumer`` with its own cursor so
|
|
all concurrent subscribers receive every event independently (true
|
|
fan-out). When this is the first consumer for the stream the read
|
|
position is initialised to *last_id*.
|
|
"""
|
|
async with self._lock:
|
|
if correlation_id not in self.states:
|
|
self.states[correlation_id] = StreamState(last_id=last_id)
|
|
|
|
state = self.states[correlation_id]
|
|
# last_id="0" means "from the start" — replay every batch already in
|
|
# the broadcast so a late joiner never misses history (e.g. the
|
|
# content_block_start that arrived before it connected).
|
|
# A specific last_id means the consumer is already caught up to that
|
|
# position, so start at the current end and receive only new batches.
|
|
cursor = 0 if last_id == "0" else state.broadcast.next_cursor()
|
|
consumer = StreamConsumer(
|
|
correlation_id, event_handler, state.broadcast, cursor
|
|
)
|
|
self.consumers[correlation_id].append(consumer)
|
|
return consumer
|
|
|
|
async def remove_consumer(self, consumer: StreamConsumer) -> None:
|
|
broadcast_to_close: StreamBroadcast | None = None
|
|
async with self._lock:
|
|
correlation_id = consumer.correlation_id
|
|
cl = self.consumers.get(correlation_id)
|
|
if cl and consumer in cl:
|
|
cl.remove(consumer)
|
|
if correlation_id in self.consumers and not self.consumers[correlation_id]:
|
|
del self.consumers[correlation_id]
|
|
state = self.states.pop(correlation_id, None)
|
|
if state:
|
|
broadcast_to_close = state.broadcast
|
|
|
|
if broadcast_to_close is not None:
|
|
broadcast_to_close.close()
|
|
|
|
async def _run_worker(self) -> None:
|
|
try:
|
|
while not self.shutdown_event.is_set():
|
|
async with self._lock:
|
|
correlation_ids = list(self.states.keys())
|
|
|
|
if not correlation_ids:
|
|
await asyncio.sleep(0.5)
|
|
continue
|
|
|
|
tasks = [self._process_stream(cid) for cid in correlation_ids]
|
|
await asyncio.gather(*tasks, return_exceptions=True)
|
|
await asyncio.sleep(0)
|
|
except Exception as e:
|
|
logger.error(f"Worker error: {e}", exc_info=True)
|
|
finally:
|
|
async with self._lock:
|
|
states_to_close = list(self.states.values())
|
|
self.consumers.clear()
|
|
self.states.clear()
|
|
for state in states_to_close:
|
|
state.broadcast.close()
|
|
|
|
async def _process_stream(self, correlation_id: str) -> None:
|
|
try:
|
|
async with self._lock:
|
|
if correlation_id not in self.states:
|
|
return
|
|
state = self.states[correlation_id]
|
|
consumer_list = list(self.consumers.get(correlation_id, []))
|
|
if not consumer_list:
|
|
return
|
|
event_handler = consumer_list[0].event_handler
|
|
|
|
events, new_last_id = await self.stream_reader.read_events(
|
|
event_handler=event_handler,
|
|
correlation_id=correlation_id,
|
|
last_id=state.last_id,
|
|
block_ms=DEFAULT_BLOCK_MS,
|
|
)
|
|
|
|
current_time = asyncio.get_event_loop().time()
|
|
|
|
if events:
|
|
await self._dispatch(
|
|
correlation_id, state, events, new_last_id, current_time
|
|
)
|
|
|
|
elif (current_time - state.last_status_check) >= STATUS_CHECK_INTERVAL:
|
|
state.last_status_check = current_time
|
|
if await self.stream_reader.check_terminal_status(
|
|
correlation_id,
|
|
event_handler,
|
|
state.cached_events,
|
|
):
|
|
drain_events, drain_last_id = await self.stream_reader.read_events(
|
|
event_handler=event_handler,
|
|
correlation_id=correlation_id,
|
|
last_id=state.last_id,
|
|
block_ms=None,
|
|
)
|
|
if drain_events:
|
|
await self._dispatch(
|
|
correlation_id,
|
|
state,
|
|
drain_events,
|
|
drain_last_id,
|
|
current_time,
|
|
)
|
|
await self._close_all_consumers(correlation_id)
|
|
logger.info(f"Stream {correlation_id} reached terminal status")
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error processing {correlation_id}: {e}", exc_info=True)
|
|
|
|
async def _dispatch(
|
|
self,
|
|
correlation_id: str,
|
|
state: StreamState,
|
|
events: list[BaseModel],
|
|
new_last_id: str,
|
|
current_time: float,
|
|
) -> None:
|
|
"""Advance stream state then broadcast *events* to all subscribers.
|
|
|
|
State is updated under the lock; ``broadcast.push`` is synchronous and
|
|
happens outside so the lock is held for the minimum time.
|
|
"""
|
|
async with self._lock:
|
|
if correlation_id not in self.states:
|
|
return
|
|
if not self.consumers.get(correlation_id):
|
|
return
|
|
state.last_id = new_last_id
|
|
state.cached_events = events
|
|
state.last_status_check = current_time
|
|
|
|
state.broadcast.push(events)
|
|
|
|
async def _close_all_consumers(self, correlation_id: str) -> None:
|
|
async with self._lock:
|
|
state = self.states.pop(correlation_id, None)
|
|
self.consumers.pop(correlation_id, None)
|
|
if state:
|
|
state.broadcast.close()
|
|
|
|
|
|
@singleton
|
|
class AdaptiveStreamReader:
|
|
"""Routes stream consumers between direct and multiplexed paths.
|
|
|
|
Direct path: one XREAD BLOCK per consumer → one Redis connection held.
|
|
Multiplexed path: one shared worker for all active streams on this process.
|
|
|
|
When concurrent direct-path readers reach ``multiplexing_threshold`` new
|
|
consumers are routed to the multiplexer instead of opening another Redis
|
|
connection. When the multiplexer drains to zero active streams it is
|
|
stopped, so load-shedding is fully reversible — consumers that arrive when
|
|
load is back below the threshold go direct again.
|
|
|
|
``multiplexing_threshold=None`` disables multiplexing entirely.
|
|
``enable_multiplexing=False`` is a hard kill-switch (e.g. for tests).
|
|
"""
|
|
|
|
@inject
|
|
def __init__(
|
|
self,
|
|
settings: Settings,
|
|
stream_reader: StreamReader,
|
|
) -> None:
|
|
self.stream_reader = stream_reader
|
|
self.multiplexer: StreamMultiplexer | None = None
|
|
self._lock = asyncio.Lock()
|
|
self.enable_multiplexing = True
|
|
self.multiplexing_threshold: int | None = settings.chat.multiplexing_threshold
|
|
# Number of consumers currently on the direct path.
|
|
# Each one holds an XREAD BLOCK connection; this is what we throttle.
|
|
self._direct_count = 0
|
|
|
|
def _should_multiplex(self) -> bool:
|
|
return (
|
|
self.enable_multiplexing
|
|
and self.multiplexing_threshold is not None
|
|
and self._direct_count >= self.multiplexing_threshold
|
|
)
|
|
|
|
async def read_events(
|
|
self,
|
|
event_handler: EventHandler,
|
|
correlation_id: str,
|
|
last_id: str = "0",
|
|
count: int | None = None,
|
|
block_ms: int | None = None,
|
|
) -> tuple[list[BaseModel], str]:
|
|
return await self.stream_reader.read_events(
|
|
event_handler=event_handler,
|
|
correlation_id=correlation_id,
|
|
last_id=last_id,
|
|
count=count,
|
|
block_ms=block_ms,
|
|
)
|
|
|
|
async def stream_events(
|
|
self,
|
|
event_handler: EventHandler,
|
|
correlation_id: str,
|
|
last_id: str = "0",
|
|
block_ms: int | None = None,
|
|
stop_event: asyncio.Event | None = None,
|
|
) -> AsyncGenerator[BaseModel, None]:
|
|
async def gen() -> AsyncGenerator[BaseModel, None]:
|
|
# Decide path under the lock so _direct_count is consistent.
|
|
async with self._lock:
|
|
use_mux = self._should_multiplex()
|
|
if not use_mux:
|
|
self._direct_count += 1
|
|
|
|
try:
|
|
if use_mux:
|
|
async for event in await self._stream_multiplexed(
|
|
event_handler, correlation_id, last_id, stop_event
|
|
):
|
|
yield event
|
|
else:
|
|
async for event in await self.stream_reader.stream_events(
|
|
event_handler, correlation_id, last_id, block_ms, stop_event
|
|
):
|
|
yield event
|
|
finally:
|
|
if not use_mux:
|
|
async with self._lock:
|
|
self._direct_count -= 1
|
|
|
|
return gen()
|
|
|
|
async def _stream_multiplexed(
|
|
self,
|
|
event_handler: EventHandler,
|
|
correlation_id: str,
|
|
last_id: str,
|
|
stop_event: asyncio.Event | None,
|
|
) -> AsyncGenerator[BaseModel, None]:
|
|
async def gen() -> AsyncGenerator[BaseModel, None]:
|
|
# Start the multiplexer on first use; capture the reference so the
|
|
# finally block targets the right instance even if it is replaced.
|
|
async with self._lock:
|
|
if self.multiplexer is None:
|
|
self.multiplexer = StreamMultiplexer(self.stream_reader)
|
|
await self.multiplexer.start()
|
|
logger.info("Multiplexer started (threshold reached)")
|
|
mux = self.multiplexer
|
|
|
|
consumer = await mux.add_consumer(
|
|
correlation_id, event_handler, last_id=last_id
|
|
)
|
|
try:
|
|
async for event in consumer.broadcast.read_from(
|
|
cursor=consumer.cursor,
|
|
stop_event=stop_event,
|
|
):
|
|
yield event
|
|
finally:
|
|
await mux.remove_consumer(consumer)
|
|
# Stop the multiplexer when it has no more active streams so the
|
|
# next consumer below the threshold can go direct again.
|
|
async with self._lock:
|
|
if self.multiplexer is mux and not mux.states:
|
|
await mux.stop()
|
|
self.multiplexer = None
|
|
logger.info("Multiplexer stopped (no active streams)")
|
|
|
|
return gen()
|
|
|
|
async def shutdown(self) -> None:
|
|
async with self._lock:
|
|
mux, self.multiplexer = self.multiplexer, None
|
|
if mux:
|
|
await mux.stop()
|