hkuds--deeptutor/deeptutor/book/streaming.py

"""
Book Engine streaming helpers
=============================

Thin wrapper around ``StreamBus`` that fixes ``source="book_engine"`` and
defines book-specific event metadata schemas.
"""

from __future__ import annotations

from contextlib import asynccontextmanager
from typing import Any

from deeptutor.core.stream import StreamEvent, StreamEventType
from deeptutor.core.stream_bus import StreamBus

SOURCE = "book_engine"


# Stage names used across the pipeline
STAGE_IDEATION = "ideation"
STAGE_EXPLORATION = "exploration"  # SourceExplorer (Stage 2 prep)
STAGE_SYNTHESIS = "synthesis"  # SpineSynthesizer draft / revise
STAGE_CRITIQUE = "critique"  # SpineSynthesizer critique round
STAGE_OVERVIEW = "overview"  # Engine-injected Overview chapter
STAGE_SPINE = "spine"  # Outer wrapper for the spine stage
STAGE_PAGE_PLAN = "page_plan"
STAGE_COMPILATION = "compilation"
STAGE_BLOCK = "block"
STAGE_INTERACTION = "interaction"


class BookStream:
    """High-level helpers around a ``StreamBus`` for the BookEngine."""

    def __init__(self, bus: StreamBus) -> None:
        self.bus = bus

    @asynccontextmanager
    async def stage(self, name: str, metadata: dict[str, Any] | None = None):
        async with self.bus.stage(name, source=SOURCE, metadata=metadata or {}):
            yield

    async def content(
        self,
        text: str,
        stage: str = "",
        metadata: dict[str, Any] | None = None,
    ) -> None:
        await self.bus.content(text, source=SOURCE, stage=stage, metadata=metadata)

    async def thinking(
        self,
        text: str,
        stage: str = "",
        metadata: dict[str, Any] | None = None,
    ) -> None:
        await self.bus.thinking(text, source=SOURCE, stage=stage, metadata=metadata)

    async def progress(
        self,
        message: str,
        current: int = 0,
        total: int = 0,
        stage: str = "",
        metadata: dict[str, Any] | None = None,
    ) -> None:
        await self.bus.progress(
            message,
            current=current,
            total=total,
            source=SOURCE,
            stage=stage,
            metadata=metadata,
        )

    async def result(self, data: dict[str, Any], metadata: dict[str, Any] | None = None) -> None:
        await self.bus.result(data, source=SOURCE, metadata=metadata)

    async def error(
        self,
        message: str,
        stage: str = "",
        metadata: dict[str, Any] | None = None,
    ) -> None:
        await self.bus.error(message, source=SOURCE, stage=stage, metadata=metadata)

    async def emit(self, event_type: StreamEventType, **kwargs: Any) -> None:
        kwargs.setdefault("source", SOURCE)
        await self.bus.emit(StreamEvent(type=event_type, **kwargs))

    # ── Book-specific events ────────────────────────────────────────────

    async def book_event(
        self,
        kind: str,
        data: dict[str, Any],
        stage: str = "",
    ) -> None:
        """Emit a custom 'book' progress event using the PROGRESS channel.

        Frontend distinguishes by ``metadata.kind``::

            kind ∈ {
              "proposal_ready", "spine_ready", "page_planned",
              "block_ready", "block_error", "page_ready",
              "compilation_complete", ...
            }
        """
        await self.bus.emit(
            StreamEvent(
                type=StreamEventType.PROGRESS,
                source=SOURCE,
                stage=stage,
                content=kind,
                metadata={"kind": kind, **data},
            )
        )


__all__ = [
    "BookStream",
    "SOURCE",
    "STAGE_IDEATION",
    "STAGE_EXPLORATION",
    "STAGE_SYNTHESIS",
    "STAGE_CRITIQUE",
    "STAGE_OVERVIEW",
    "STAGE_SPINE",
    "STAGE_PAGE_PLAN",
    "STAGE_COMPILATION",
    "STAGE_BLOCK",
    "STAGE_INTERACTION",
]