backend/app/services/orchestrator_speed.py

"""Speed-oriented orchestration helpers.

Parallel participant turns, compact orchestrator context, and fast-
model routing for lightweight classifier calls. Keeps the same
visible message count while shortening wall-clock time.
"""
from __future__ import annotations

import asyncio
import logging
import re
from dataclasses import dataclass, field
from typing import Any, AsyncIterator, Awaitable, Callable, TYPE_CHECKING

from app.services.resilience import ResilientTurnResult, run_resilient_turn

if TYPE_CHECKING:
    from app.services.models import Participant, Phase, Session

LOG = logging.getLogger(__name__)

# Rough char budget for orchestrator-side prompts that would otherwise
# resend the entire transcript every few turns.
ORCHESTRATOR_TRANSCRIPT_CHAR_BUDGET = 14_000
_RECENT_TAIL_MESSAGES = 24

CallParticipantFn = Callable[..., Awaitable[tuple[str, float, bool, str]]]


@dataclass
class _AiTurnSpec:
    participant: "Participant"
    user_prompt: str
    label: str
    max_tokens: int


@dataclass
class _AiTurnResult:
    participant: "Participant"
    turn: ResilientTurnResult
    pending: list[tuple[str, str, str]] = field(default_factory=list)


def orchestrator_fast_model_id(session: "Session") -> str:
    """Model for lightweight orchestrator classifiers (addressed-to, status)."""
    from app.config import settings
    from app.services.orchestrator import _orchestrator_model_id

    fast = (getattr(settings, "orchestrator_fast_model", None) or "").strip()
    if fast and settings.resolve_model(fast):
        return fast
    return _orchestrator_model_id(session)


def _format_history_for_orchestrator(
    messages: list[dict[str, Any]],
    *,
    include_orchestrator: bool = True,
) -> str:
    from app.services.orchestrator import _format_history

    return _format_history(messages, include_orchestrator=include_orchestrator)


async def compact_transcript_for_orchestrator(
    session: "Session",
    *,
    orchestrator_model_id: str,
) -> str:
    """Return a transcript block sized for orchestrator judge prompts.

    Uses a rolling summary + recent tail when the full history exceeds
    the char budget. Summaries are built lazily (one extra orchestrator
    call) and cached on the session.
    """
    from app.services.json_calls import orchestrator_call
    from app.services.orchestrator import _bump_orchestrator_count

    messages = session.messages
    full = _format_history_for_orchestrator(messages)
    if len(full) <= ORCHESTRATOR_TRANSCRIPT_CHAR_BUDGET:
        return full

    tail = messages[-_RECENT_TAIL_MESSAGES:]
    tail_text = _format_history_for_orchestrator(tail)
    through = len(messages) - len(tail)
    if (
        session.orchestrator_context_summary
        and session.orchestrator_context_through_idx >= through - 2
    ):
        return (
            "[Earlier discussion summary]\n"
            f"{session.orchestrator_context_summary}\n\n"
            "[Recent messages]\n"
            f"{tail_text}"
        )

    prompt = (
        "Summarize the following group discussion for an orchestrator that "
        "will judge consensus and open questions. Preserve names, stances, "
        "and unresolved disagreements. Be concise (under 400 words).\n\n"
        f"{full}"
    )
    raw, _ = await orchestrator_call(
        orchestrator_model_id=orchestrator_model_id,
        user_prompt=prompt,
        label="orchestrator_transcript_summary",
        api_log=session.api_log,
        expect_json=False,
        max_tokens=700,
        temperature=0.2,
    )
    _bump_orchestrator_count(session)
    summary = (raw or "").strip() or full[-ORCHESTRATOR_TRANSCRIPT_CHAR_BUDGET:]
    session.orchestrator_context_summary = summary
    session.orchestrator_context_through_idx = through
    return (
        "[Earlier discussion summary]\n"
        f"{summary}\n\n"
        "[Recent messages]\n"
        f"{tail_text}"
    )


async def _execute_ai_turn(
    session: "Session",
    spec: _AiTurnSpec,
    call_participant: CallParticipantFn,
) -> _AiTurnResult:
    from app.services.orchestrator import _pending_addressed_for

    pending = _pending_addressed_for(session, spec.participant)
    turn = await run_resilient_turn(
        session=session,
        participant=spec.participant,
        user_prompt=spec.user_prompt,
        label=spec.label,
        max_tokens=spec.max_tokens,
        call_participant=call_participant,
    )
    return _AiTurnResult(
        participant=spec.participant,
        turn=turn,
        pending=pending,
    )


PostProcessFn = Callable[
    [_AiTurnResult],
    Awaitable[dict[str, Any] | None],
]


async def run_roster_ai_turns_parallel(
    session: "Session",
    actives: list["Participant"],
    *,
    phase: "Phase",
    build_spec: Callable[["Participant"], _AiTurnSpec | None],
    call_participant: CallParticipantFn,
    on_human_turn: Callable[
        ["Participant"],
        AsyncIterator[str],
    ],
    post_process: PostProcessFn | None = None,
) -> AsyncIterator[str]:
    """Run participant turns: humans sequentially, AI in parallel batches.

    Walks `actives` in roster order. Consecutive AI participants are
    executed with ``asyncio.gather``; results are applied in roster
    order so the message log stays deterministic. Humans are awaited
    one at a time via ``on_human_turn``.

    Yields orchestrator SSE strings (status, message, errors, etc.).
    """
    from app.services.orchestrator import (
        _msg_payload,
        _participant_msg_cap_hit,
        _participant_turn_failure_sse,
        _sse,
        _wait_for_continue,
    )

    ai_batch: list[_AiTurnSpec] = []

    async def flush_ai_batch() -> AsyncIterator[str]:
        nonlocal ai_batch
        if not ai_batch:
            return
        specs = ai_batch
        ai_batch = []
        results = await asyncio.gather(
            *[
                _execute_ai_turn(session, spec, call_participant)
                for spec in specs
            ],
            return_exceptions=True,
        )
        for spec, item in zip(specs, results):
            if isinstance(item, BaseException):
                LOG.exception(
                    "Parallel turn failed for %s: %s",
                    spec.participant.participant_id,
                    item,
                )
                yield _sse("participant_error", {
                    "participant_id": spec.participant.participant_id,
                    "name": spec.participant.name,
                    "phase": phase.value,
                })
                continue
            extra: dict[str, Any] | None = None
            if post_process is not None:
                extra = await post_process(item) or {}
            async for chunk in _emit_ai_turn_result(
                session, item, phase=phase, extra=extra,
            ):
                yield chunk
            if _participant_msg_cap_hit(session):
                async for chunk in _wait_for_continue(session, "messages"):
                    yield chunk

    for p in actives:
        if p.kind == "human":
            async for chunk in flush_ai_batch():
                yield chunk
            async for chunk in on_human_turn(p):
                yield chunk
            continue

        spec = build_spec(p)
        if spec is None:
            continue
        ai_batch.append(spec)

    async for chunk in flush_ai_batch():
        yield chunk


async def run_initial_opinions_roster(
    session: "Session",
    actives: list["Participant"],
    *,
    build_spec: Callable[["Participant"], _AiTurnSpec | None],
    call_participant: CallParticipantFn,
    on_human_turn: Callable[
        ["Participant"],
        AsyncIterator[str],
    ],
    post_process: PostProcessFn | None = None,
) -> AsyncIterator[str]:
    """Phase-1 roster walk with human-aware AI prefetch.

    When a human is in the roster, every AI participant's initial-
    opinion call is fired immediately (in parallel) so answers are ready
    while the human types. SSE ``message`` events are still emitted in
    roster order: any LLMs listed before the human appear as soon as
    their prefetch completes, and LLMs after the human stay hidden until
    the human submits (or skips).
    """
    from app.services.models import Phase
    from app.services.orchestrator import (
        _participant_msg_cap_hit,
        _sse,
        _wait_for_continue,
    )

    phase = Phase.INITIAL_OPINIONS
    has_human = any(p.kind == "human" for p in actives)
    if not has_human:
        async for chunk in run_roster_ai_turns_parallel(
            session,
            actives,
            phase=phase,
            build_spec=build_spec,
            call_participant=call_participant,
            on_human_turn=on_human_turn,
            post_process=post_process,
        ):
            yield chunk
        return

    pending: dict[str, asyncio.Task[_AiTurnResult]] = {}
    for p in actives:
        if p.kind == "human":
            continue
        spec = build_spec(p)
        if spec is None:
            continue
        pending[p.participant_id] = asyncio.create_task(
            _execute_ai_turn(session, spec, call_participant),
            name=f"prefetch_initial:{p.participant_id}",
        )

    for p in actives:
        if p.kind == "human":
            async for chunk in on_human_turn(p):
                yield chunk
            if _participant_msg_cap_hit(session):
                async for chunk in _wait_for_continue(session, "messages"):
                    yield chunk
            continue

        task = pending.pop(p.participant_id, None)
        if task is None:
            continue
        try:
            result = await task
        except BaseException as exc:
            LOG.exception(
                "Prefetched initial opinion failed for %s: %s",
                p.participant_id,
                exc,
            )
            yield _sse("participant_error", {
                "participant_id": p.participant_id,
                "name": p.name,
                "phase": phase.value,
            })
            continue

        extra: dict[str, Any] | None = None
        if post_process is not None:
            extra = await post_process(result) or {}
        async for chunk in _emit_ai_turn_result(
            session, result, phase=phase, extra=extra,
        ):
            yield chunk
        if _participant_msg_cap_hit(session):
            async for chunk in _wait_for_continue(session, "messages"):
                yield chunk

    for pid, task in pending.items():
        if not task.done():
            task.cancel()


async def _emit_ai_turn_result(
    session: "Session",
    result: _AiTurnResult,
    *,
    phase: "Phase",
    extra: dict[str, Any] | None = None,
) -> AsyncIterator[str]:
    """Apply a completed AI turn to the session and yield SSE."""
    from app.services.orchestrator import (
        _add_participant_message,
        _msg_payload,
        _orchestrator_cap_hit,
        _participant_turn_failure_sse,
        _replying_to_ids,
        _sse,
        _wait_for_continue,
    )

    p = result.participant
    turn = result.turn
    substituted = False
    for ev in turn.sse_events:
        if "participant_substituted" in ev:
            substituted = True
        yield ev
    if not turn.ok:
        for chunk in _participant_turn_failure_sse(session, p):
            yield chunk
        return

    speaker = turn.speaker
    meta = extra or {}
    msg = _add_participant_message(
        session,
        speaker,
        turn.text,
        phase=phase,
        elapsed=turn.elapsed,
        addressed_to=meta.get("addressed_to"),
        replying_to=meta.get(
            "replying_to",
            _replying_to_ids(result.pending),
        ),
        message_id=meta.get("message_id"),
    )
    yield _sse("message", _msg_payload(msg))

    if substituted:
        from app.services.orchestrator import (
            _rebuild_participant_credential_on_model_change,
        )
        if await _rebuild_participant_credential_on_model_change(
            session, speaker,
        ):
            yield _sse("credentials_updated", {
                "stage": "model_changed",
                "credentials": session.credential_summary,
            })

    if _orchestrator_cap_hit(session):
        async for chunk in _wait_for_continue(session, "orchestrator"):
            yield chunk


async def run_parallel_coroutines(
    coros: list[Awaitable[Any]],
) -> list[Any]:
    """Gather with exception isolation (failed tasks become exceptions)."""
    return await asyncio.gather(*coros, return_exceptions=True)