core/main_agent.py

"""Main Socratic bot agent.

Receives the manager's verdict + discussion plan + chat history. Streams a
visible response to the participant, with a trailing
`<scratch>...</scratch>` block parsed by the host and logged.

Prompt lives in `prompts/main_prompt.md` and is loaded via
`core.config_loader.MAIN_PROMPT` at module import time.
"""

import re
import time
from dataclasses import dataclass, field
from core.config_loader import (
    client as _default_client, MODEL, REASONING_EFFORT, VERBOSITY, MAIN_PROMPT,
)
from core import plan_parsing


SCRATCH_RE = re.compile(r"<scratch>(.*?)</scratch>\s*", re.DOTALL | re.IGNORECASE)
SCRATCH_OPEN_RE = re.compile(r"<scratch>", re.IGNORECASE)


@dataclass
class BotResponse:
    """One main-bot reply. `scratch` is a parsed dict from the trailing
    <scratch>...</scratch> block; missing fields become None.

    `visible` is the text shown to the participant (scratch stripped).
    `raw` is the full bot output including the <scratch> block."""
    visible: str = ""
    scratch: dict = field(default_factory=dict)
    raw: str = ""
    latency_ms: int = 0
    ttft_ms: int | None = None


def parse_scratch(raw):
    """Extract the scratch fields from a <scratch>...</scratch> block.
    Returns {} if the block is missing or unparseable. Expects YAML-style
    `key: value` lines inside the tags."""
    if not raw:
        return {}
    m = SCRATCH_RE.search(raw)
    if not m:
        return {}
    body = m.group(1).strip()
    out = {}
    for line in body.splitlines():
        line = line.strip()
        if not line or ":" not in line:
            continue
        k, _, v = line.partition(":")
        out[k.strip().lower()] = v.strip()
    # Coerce known integer fields.
    if "chosen_target" in out:
        try:
            out["chosen_target"] = int(out["chosen_target"])
        except (ValueError, TypeError):
            pass
    return out


def split_visible_and_scratch(raw):
    """Split the raw bot output into (visible_text, scratch_dict). The
    scratch block is stripped from visible regardless of whether parsing
    succeeded."""
    scratch = parse_scratch(raw)
    visible = SCRATCH_RE.sub("", raw or "", count=1).strip()
    # If a stray opening tag remains (unmatched), still strip from that point.
    m = SCRATCH_OPEN_RE.search(visible)
    if m:
        visible = visible[: m.start()].rstrip()
    return visible, scratch


def build_plan_control(verdict, discussion_plan_subprobes):
    """Render the per-turn PLAN CONTROL developer message from the manager's
    verdict and the parsed discussion plan."""
    live = plan_parsing.by_id(discussion_plan_subprobes, verdict.live_subprobe_id)
    if live is None:
        live_block = f"(sub-probe id={verdict.live_subprobe_id} not found in plan)"
    else:
        live_block = (
            f"id: {live['id']}\n"
            f"component: {live['component']}\n"
            f"gap: {live['gap']}\n"
            f"rationale: {live['rationale']}\n"
            f"hook: {live.get('hook','')}"
        )

    qset = verdict.question_set or []
    if qset:
        qset_lines = []
        # Display indices as 1-based ([1], [2], ...) so the bot reads them
        # in the same numbering the prompt examples and the manager's
        # `reason` text use. Internal storage stays a 0-based list.
        for i, item in enumerate(qset):
            qset_lines.append(
                f"  [{i + 1}] ({item.get('status','?')}) {item.get('text','')}"
            )
        qset_block = "\n".join(qset_lines)
        # Highlight the first-pending — that's the question the bot must ask.
        first_pending = next(
            (i for i, item in enumerate(qset)
             if item.get("status") == "pending"),
            None,
        )
        if first_pending is not None:
            qset_block += (
                f"\n  ----\n"
                f"  >>> ASK THIS TURN: index [{first_pending + 1}] = "
                f"{qset[first_pending].get('text','')}"
            )
        else:
            qset_block += "\n  (no pending items — should not happen on a question-asking turn)"
    else:
        qset_block = "  (empty — no questions for this turn; expected for plan-complete)"

    return (
        "=== PLAN CONTROL (host-managed, authoritative) ===\n"
        f"turn_type: {verdict.turn_type}\n"
        f"live_subprobe_id: {verdict.live_subprobe_id}\n"
        f"advanced_this_turn: {verdict.advanced_this_turn}\n"
        "\n"
        "LIVE SUB-PROBE (work this one this turn):\n"
        f"{live_block}\n"
        "\n"
        "QUESTION_SET (ordered list; ask the first item with status=pending):\n"
        f"{qset_block}\n"
        "\n"
        "DIRECTIVE for this turn (framing / anchor / recap instructions):\n"
        f"{verdict.directive}\n"
        "=== END PLAN CONTROL ==="
    )


def build_messages(history, case_text, discussion_plan_text, plan_control_text):
    """Assemble the developer + chat messages for the main bot call.

    Layout (oldest to newest):
      developer: MAIN_PROMPT (the bot's system prompt)
      developer: case text
      developer: full discussion plan (for reference; PLAN CONTROL names the live one)
      developer: PLAN CONTROL block (this turn's directive)
      user/assistant: chat history (first 'assistant' = case text, skipped)
    """
    msgs = [
        {"role": "developer", "content": MAIN_PROMPT},
        {"role": "developer", "content": f"The case being discussed:\n\n{case_text}"},
        {
            "role": "developer",
            "content": (
                "=== DISCUSSION PLAN ===\n"
                f"{discussion_plan_text}\n"
                "=== END DISCUSSION PLAN ==="
            ),
        },
        {"role": "developer", "content": plan_control_text},
    ]

    skip_first_assistant = True
    for m in history:
        role = m.get("role", "user")
        content = m.get("content", "")
        if isinstance(content, list):
            content = " ".join(
                b.get("text", "") if isinstance(b, dict) else str(b)
                for b in content
            )
        if skip_first_assistant and role == "assistant":
            skip_first_assistant = False
            continue
        msgs.append({"role": role, "content": str(content)})
    return msgs


def stream(messages, client=None):
    """Single-pass main bot reply. The bot now does both reasoning AND
    rendering (short-question pattern, sentence shape) in one call, so
    no separate paraphrase pass is needed.

    Non-streaming for simplicity. Emits one delta with the full visible
    text, then a done event.

    Generator contract preserved:
      ("delta", str)       — the visible text (yielded once)
      ("done", BotResponse) — final summary; ALWAYS the last yield
    """
    api = client if client is not None else _default_client
    t0 = time.perf_counter()

    bot_resp = api.responses.create(
        model=MODEL,
        reasoning={"effort": REASONING_EFFORT},
        text={"verbosity": VERBOSITY},
        input=messages,
    )
    raw = bot_resp.output_text or ""
    visible, scratch = split_visible_and_scratch(raw)
    done_t = time.perf_counter()

    ttft_ms = int((done_t - t0) * 1000)
    if visible:
        yield ("delta", visible)

    yield (
        "done",
        BotResponse(
            visible=visible,
            scratch=scratch,
            raw=raw,
            latency_ms=int((done_t - t0) * 1000),
            ttft_ms=ttft_ms,
        ),
    )