Streaming RAG

Overview

StreamingRAG wraps a backend to provide token-level streaming, with both a plain-text stream and a richer typed-event stream that surfaces lifecycle metadata (retrieval done, generation started, errors) alongside text chunks.

EventType (str, Enum)

Defined in streaming_enum. Referenced values in streaming_rag include RETRIEVAL_DONE and CHUNK (checked via StreamEvent.is_chunk) — consult streaming_enum directly for the complete member list.

StreamConfig

@dataclass
class StreamConfig:
    chunk_size: int = 10
    sim_delay: float = 0.04
    max_context_docs: int = 5
    max_doc_chars: int = 500
    emit_metadata: bool = True
    retrieval_timeout: float = 10.0
    generation_timeout: float = 60.0
    prompt_template: str = "You are a precise RAG assistant.\n\n..."

Protocols (protocols.py)

@runtime_checkable
class Retrievable(Protocol):
    def retrieve(self, query: str) -> Any: ...

@runtime_checkable
class Queryable(Protocol):
    def query(self, query: str, context: Optional[Dict] = None) -> Any: ...

@runtime_checkable
class SyncStreamableLLM(Protocol):
    def stream(self, prompt: str) -> Iterator[str]: ...

@runtime_checkable
class AsyncStreamableLLM(Protocol):
    async def astream(self, prompt: str) -> AsyncIterator[str]: ...

These are runtime_checkable Protocols — you can isinstance()-check any object against them (e.g. isinstance(my_llm, SyncStreamableLLM)) without subclassing, since Python structurally matches the required methods.

StreamEvent

@dataclass(frozen=True)
class StreamEvent:
    type: EventType
    data: Any = None
    latency: float = 0.0

    @property
    def is_chunk(self) -> bool: ...

Immutable (frozen=True). data holds a text fragment for CHUNK events or a metadata payload otherwise; latency is seconds since the stream started, set by StreamingRAG.

StreamingRAG

Constructor

def __init__(
    self,
    rag_system: Any,
    llm: Any = None,
    config: Optional[StreamConfig] = None,
    *,
    chunk_size: Optional[int] = None,
    on_chunk: Optional[Callable[[str], None]] = None,
    on_event: Optional[Callable[[StreamEvent], None]] = None,
)

Streaming methods

_normalise_doc(doc) (staticmethod) — accepted input shapes

dict                       -> returned as-is
tuple (text,)               -> {"text": text}
tuple (text, score)          -> {"text": text, "score": score}
tuple (text, score, meta)    -> {"text": text, "score": score, **meta}  # meta: dict or any str-able value
anything else                -> {"text": str(doc)}

This means StreamingRAG tolerates almost any retriever output shape without extra glue code.

_generate(prompt, query, context, docs) — generation strategy priority

1. Async LLM astream(), if the llm supports it (AsyncStreamableLLM).
2. Sync LLM stream(), if supported (SyncStreamableLLM).
3. Fallback: call rag_system.generate() for the full answer, then word-level-simulate streaming using config.chunk_size and config.sim_delay.

stats (property)

Runtime statistics (e.g. chunks emitted, latencies) — exact schema not independently specified beyond what's derivable from StreamEvent.latency usage.

Complete example

from zerda.rag.types.streaming_rag import StreamingRAG, StreamConfig, EventType

config = StreamConfig(chunk_size=6, sim_delay=0.02, max_context_docs=4)

def log_event(event):
    if not event.is_chunk:
        print(f"[event] {event.type} @ {event.latency:.2f}s")

streamer = StreamingRAG(backend, llm=my_llm, config=config, on_event=log_event)

# Simple text-only sync iteration:
for chunk in streamer.stream("What is quantum computing?"):
    print(chunk, end="", flush=True)

# Rich async iteration with lifecycle events:
async def run():
    async for event in streamer.astream_events("Explain RAG"):
        if event.is_chunk:
            print(event.data, end="", flush=True)
        elif event.type == EventType.RETRIEVAL_DONE:
            print(f"\nRetrieved {event.data.get('num_docs')} docs")

Best practices

• Provide an llm implementing SyncStreamableLLM/AsyncStreamableLLM whenever possible — the word-level simulation fallback adds artificial sim_delay per chunk purely for UX pacing, not real generation latency.
• Use astream_events/stream_events (not the plain-text variants) when your UI needs to show retrieval progress or handle mid-stream errors distinctly from normal text.

Common mistakes

• Assuming stream() exposes retrieval/generation lifecycle info — it intentionally filters those out; switch to stream_events() for that.
• Forgetting StreamEvent is frozen — attempting to mutate a StreamEvent instance after creation raises dataclasses.FrozenInstanceError.

Related APIs

• Federated RAG — has its own independent streaming implementation (StreamChunk) rather than reusing StreamEvent.
• Domain-Specific RAG's astream — a simpler streaming path that this module's design pattern (fallback to word-level simulation) closely resembles.