services/rag_coordinator.py

"""
GovBridge India — RAG Coordination Matrix (Sprint 40)


True zero-allocation async state machine for parsing interleaved LLM output.
Routes prose to UI and structured JSON to OpenFisca eligibility engine.

ARCHITECTURAL CONSTRAINTS:
  - No collections.deque (heap fragmentation in ASGI environments)
  - No regex for hot-path parsing (catastrophic backtracking risk)
  - Pure string mathematical overlap checks for TCP fragmentation
  - O(1) fast-path bypass: if sentinel start char ('<') not in chunk,
    yield immediately with ZERO intermediate allocations
  - JSON accumulated via list[str].append() → single join() on close

Sentinel Protocol:
  - <|param_start|> opens a JSON parameter block
  - <|param_end|>   closes it
  - Everything between sentinels is valid JSON for OpenFisca
"""

from typing import AsyncGenerator, Tuple, AsyncIterator

# ── Sentinel Constants ─────────────────────────────────────────────
PARAM_START = "<|param_start|>"
PARAM_END = "<|param_end|>"
PARAM_START_LEN = len(PARAM_START)  # 15
PARAM_END_LEN = len(PARAM_END)      # 13

# O(1) character probes — avoid .find() entirely on the fast path
_START_CHAR = PARAM_START[0]  # "<"
_END_CHAR = PARAM_END[0]      # "<"


async def parse_interleaved_stream(
    token_stream: AsyncIterator[str],
) -> AsyncGenerator[Tuple[str, str], None]:
    """
    Async generator that consumes raw LLM token chunks and yields
    classified tuples: ("text", prose_string) or ("json", json_string).

    FAST PATH (O(1), zero allocation):
        When state == TEXT, buffer is empty, and the chunk does not
        contain the sentinel start character '<', yield the raw chunk
        pointer directly. No string concatenation, no slicing, no
        intermediate objects.

    SLOW PATH (sentinel detection + TCP fragmentation):
        When '<' IS detected or buffer has residual bytes, concatenate
        buffer + chunk into a working string and scan for full sentinels.
        Uses decreasing suffix-to-prefix overlap to handle sentinels
        split across TCP packet boundaries.

    JSON ACCUMULATION:
        Uses list[str].append() for O(1) amortized accumulation.
        Single "".join() fires only when <|param_end|> is found.

    Args:
        token_stream: AsyncIterator yielding raw string chunks from Groq SDK.

    Yields:
        ("text", str): Prose content for SSE `event: message`
        ("json", str): Complete JSON string for SSE `event: parameter_update`
    """
    state = "TEXT"                       # FSM state: "TEXT" or "JSON"
    buffer = ""                          # Residual bytes from TCP fragmentation
    json_accumulator: list[str] = []     # O(1) amortized JSON token collector

    async for chunk in token_stream:
        if not chunk:
            continue

        # ═══════════════════════════════════════════════════════════
        # O(1) FAST PATH — ZERO ALLOCATION BYPASS
        # ═══════════════════════════════════════════════════════════
        # Preconditions for the fast path:
        #   1. We are in TEXT state (not accumulating JSON)
        #   2. No residual buffer from a previous partial sentinel
        #   3. The chunk does not contain '<' (sentinel start char)
        #
        # When all three hold, the chunk is guaranteed to be pure
        # prose. We yield the raw chunk object directly — CPython
        # passes the existing PyUnicode pointer with zero copy.
        # This covers ~95% of LLM tokens (words, spaces, punctuation).
        # ═══════════════════════════════════════════════════════════
        if state == "TEXT" and not buffer:
            if _START_CHAR not in chunk:
                yield ("text", chunk)
                continue

        # ═══════════════════════════════════════════════════════════
        # SLOW PATH — SENTINEL DETECTION & TCP FRAGMENTATION
        # ═══════════════════════════════════════════════════════════
        # Only reached when:
        #   - Chunk contains '<' (potential sentinel boundary), OR
        #   - Buffer has residual bytes from a split sentinel, OR
        #   - We are in JSON state (accumulating parameter block)
        # ═══════════════════════════════════════════════════════════
        working = buffer + chunk
        buffer = ""

        while working:
            if state == "TEXT":
                # ── TEXT STATE: Scan for <|param_start|> ──────────
                idx = working.find(PARAM_START)

                if idx == -1:
                    # No complete sentinel found.
                    # Check if the tail could be a PARTIAL sentinel
                    # split across TCP packets.
                    overlap = _find_suffix_prefix_overlap(working, PARAM_START)

                    if overlap > 0:
                        # Hold back the ambiguous suffix
                        safe = working[:-overlap]
                        buffer = working[-overlap:]
                    else:
                        safe = working
                        # buffer stays "" (already cleared above)

                    if safe:
                        yield ("text", safe)
                    break  # Consume next chunk from token_stream

                else:
                    # ── Complete sentinel found at position `idx` ──
                    # Yield everything before it as prose
                    if idx > 0:
                        yield ("text", working[:idx])

                    # Transition to JSON state
                    state = "JSON"
                    json_accumulator = []  # Reset for new parameter block
                    # Advance past the sentinel tag
                    working = working[idx + PARAM_START_LEN:]

            elif state == "JSON":
                # ── JSON STATE: Scan for <|param_end|> ────────────
                idx = working.find(PARAM_END)

                if idx == -1:
                    # End sentinel not yet received.
                    # Accumulate into the list (O(1) amortized).
                    overlap = _find_suffix_prefix_overlap(working, PARAM_END)

                    if overlap > 0:
                        # Safe portion goes to accumulator;
                        # ambiguous suffix held in buffer
                        if working[:-overlap]:
                            json_accumulator.append(working[:-overlap])
                        buffer = working[-overlap:]
                    else:
                        json_accumulator.append(working)
                        # buffer stays "" (already cleared above)

                    break  # Consume next chunk from token_stream

                else:
                    # ── End sentinel found ─────────────────────────
                    # Append final fragment before sentinel
                    if idx > 0:
                        json_accumulator.append(working[:idx])

                    # Single join — the ONLY allocation for the
                    # entire JSON block regardless of token count
                    json_content = "".join(json_accumulator).strip()
                    json_accumulator = []  # Reset

                    if json_content:
                        yield ("json", json_content)

                    # Transition back to TEXT state
                    state = "TEXT"
                    working = working[idx + PARAM_END_LEN:]

    # ── Flush remaining state at stream end ────────────────────────
    if state == "TEXT":
        if buffer:
            yield ("text", buffer)
    elif state == "JSON":
        # Incomplete JSON block — LLM was cut off.
        # Append buffer to accumulator and flush as prose (defensive).
        if buffer:
            json_accumulator.append(buffer)
        remaining = "".join(json_accumulator).strip()
        if remaining:
            yield ("text", remaining)


def _find_suffix_prefix_overlap(text: str, sentinel: str) -> int:
    """
    Find the longest suffix of `text` that matches a prefix of `sentinel`.
    Returns the length of the overlap (0 if none).

    This handles TCP fragmentation where a sentinel like '<|param_start|>'
    is split across two chunks: chunk1 ends with '<|para' and chunk2
    starts with 'm_start|>...'.

    Uses decreasing length comparison — O(k²) where k = len(sentinel).
    Since sentinels are 13-15 chars, this is effectively O(1) constant.
    """
    max_overlap = min(len(text), len(sentinel))
    for length in range(max_overlap, 0, -1):
        if text[-length:] == sentinel[:length]:
            return length
    return 0