research/activity_narrator.py

"""Rolling fast-model narration for research activity."""

from __future__ import annotations

import asyncio
import json
import re
from collections.abc import Awaitable, Callable, Mapping
from contextvars import ContextVar
from dataclasses import dataclass
from time import monotonic
from typing import Any

from fast_agent.hooks import HookContext
from fast_agent.types.llm_stop_reason import LlmStopReason

from .app_jobs import ResearchJob

SummaryFunction = Callable[[str], Awaitable[str]]
Clock = Callable[[], float]

_SENSITIVE_KEY = re.compile(
    r"(?:authorization|bearer|token|secret|password|api[_-]?key)",
    re.IGNORECASE,
)
_LARGE_VALUE_KEY = re.compile(
    r"(?:content|contents|blob|base64|data|file_data)",
    re.IGNORECASE,
)


@dataclass(frozen=True, slots=True)
class ActivityBatch:
    iteration: int
    stop_reason: str
    reasoning: str
    visible_text: str
    tool_calls: tuple[str, ...]

    @property
    def is_final(self) -> bool:
        return self.stop_reason != LlmStopReason.TOOL_USE.value


current_activity_narrator: ContextVar[ActivityNarrator | None] = ContextVar(
    "current_activity_narrator",
    default=None,
)


class ActivityNarrator:
    """Coalesce Tool Runner observations into a rolling user-facing narrative."""

    def __init__(
        self,
        job: ResearchJob,
        summarize: SummaryFunction,
        *,
        every_n_steps: int = 3,
        max_summary_age: float = 30,
        poll_interval: float = 5,
        timeout: float = 10,
        clock: Clock = monotonic,
    ) -> None:
        self.job = job
        self._summarize = summarize
        self._every_n_steps = max(1, every_n_steps)
        self._max_summary_age = max_summary_age
        self._poll_interval = poll_interval
        self._timeout = timeout
        self._clock = clock

        self._latest: ActivityBatch | None = None
        self._revision = 0
        self._requested_revision = 0
        self._summarized_revision = 0
        self._steps_since_request = 0
        self._last_requested_at = clock()
        self._request_event = asyncio.Event()
        self._idle_event = asyncio.Event()
        self._idle_event.set()
        self._worker: asyncio.Task[None] | None = None
        self._timer: asyncio.Task[None] | None = None

    async def start(self) -> None:
        if self._worker is not None:
            return
        self._worker = asyncio.create_task(self._run(), name=f"{self.job.id}-narrator")
        self._timer = asyncio.create_task(
            self._run_timer(), name=f"{self.job.id}-narrator-timer"
        )

    def observe(self, batch: ActivityBatch) -> None:
        self.job.record_llm_step()
        if self.job.phase in {"reporting", "wrapping_up"}:
            return
        self._latest = batch
        self._revision += 1
        self._steps_since_request += 1

        first = self._summarized_revision == 0 and self._requested_revision == 0
        due_by_steps = self._steps_since_request >= self._every_n_steps
        if first or due_by_steps or batch.is_final:
            self.request_summary()

    def poll(self) -> None:
        """Request a summary when unsummarized activity has aged past the deadline."""
        has_unrequested_activity = self._revision > self._requested_revision
        due = self._clock() - self._last_requested_at >= self._max_summary_age
        if has_unrequested_activity and due:
            self.request_summary()

    def request_summary(self) -> None:
        if self._latest is None:
            return
        self._requested_revision = self._revision
        self._steps_since_request = 0
        self._last_requested_at = self._clock()
        self._idle_event.clear()
        self._request_event.set()

    async def close(self, *, flush_timeout: float = 12) -> None:
        if self._worker is None:
            return
        if self._revision > self._requested_revision:
            self.request_summary()
        if self._requested_revision > self._summarized_revision:
            try:
                await asyncio.wait_for(self._idle_event.wait(), timeout=flush_timeout)
            except TimeoutError:
                pass
        await self._cancel_tasks()

    async def _run(self) -> None:
        while True:
            await self._request_event.wait()
            self._request_event.clear()

            revision = self._requested_revision
            batch = self._latest
            if batch is None:
                self._idle_event.set()
                continue

            prompt = build_summary_prompt(
                topic=self.job.topic,
                previous_summary=self.job.activity_summary,
                batch=batch,
            )
            try:
                summary = await asyncio.wait_for(
                    self._summarize(prompt), timeout=self._timeout
                )
                summary = _clean_summary(summary)
            except asyncio.CancelledError:
                raise
            except Exception:
                if revision == self._requested_revision:
                    self._requested_revision = self._summarized_revision
                self._idle_event.set()
                continue

            if summary and revision >= self._summarized_revision:
                self._summarized_revision = revision
                self.job.set_activity_summary(self.job.narrative_for_phase(summary))

            if self._requested_revision > revision:
                self._request_event.set()
            else:
                self._idle_event.set()

    async def _run_timer(self) -> None:
        while True:
            await asyncio.sleep(self._poll_interval)
            self.poll()

    async def _cancel_tasks(self) -> None:
        tasks = [task for task in (self._timer, self._worker) if task is not None]
        for task in tasks:
            task.cancel()
        await asyncio.gather(*tasks, return_exceptions=True)
        self._timer = None
        self._worker = None


def activity_batch_from_hook(ctx: HookContext) -> ActivityBatch:
    message = ctx.message
    stop_reason = (
        message.stop_reason.value if message.stop_reason is not None else "unknown"
    )
    return ActivityBatch(
        iteration=ctx.iteration,
        stop_reason=stop_reason,
        reasoning=_channel_text(message.channels, "reasoning", limit=1400),
        visible_text=_truncate(message.all_text().strip(), 700),
        tool_calls=tuple(_flatten_tool_calls(message.tool_calls)),
    )


def build_summary_prompt(
    *,
    topic: str,
    previous_summary: str,
    batch: ActivityBatch,
) -> str:
    tools = "\n".join(f"- {call}" for call in batch.tool_calls) or "- none"
    return "\n\n".join(
        [
            f"RESEARCH TASK\n{_truncate(topic, 600)}",
            f"PREVIOUS NARRATIVE\n{_truncate(previous_summary, 900) or '(none yet)'}",
            "\n".join(
                [
                    "LATEST ASSISTANT STEP",
                    f"Iteration: {batch.iteration}",
                    f"Stop reason: {batch.stop_reason}",
                    f"Exposed reasoning:\n{batch.reasoning or '(none)'}",
                    f"Visible response:\n{batch.visible_text or '(none)'}",
                    f"Tool calls:\n{tools}",
                ]
            ),
            (
                "Write the updated progress narrative now. Return only one or two "
                "short sentences. Describe what has been established and what is "
                "currently happening. Preserve uncertainty and tense; planned work "
                "is not completed work. Do not mention internal iterations, JSON, "
                "hooks, or framework details."
            ),
        ]
    )


def _channel_text(
    channels: Mapping[str, Any] | None,
    name: str,
    *,
    limit: int,
) -> str:
    blocks = (channels or {}).get(name) or []
    text = "\n".join(
        str(value).strip()
        for block in blocks
        if (value := getattr(block, "text", None))
    )
    return _truncate(text, limit)


def _flatten_tool_calls(tool_calls: Mapping[str, Any] | None) -> list[str]:
    flattened: list[str] = []
    for call in (tool_calls or {}).values():
        params = getattr(call, "params", None)
        name = str(getattr(params, "name", "tool"))
        arguments = _sanitize(getattr(params, "arguments", None) or {})
        encoded = json.dumps(arguments, ensure_ascii=False, sort_keys=True)
        flattened.append(f"{name}: {_truncate(encoded, 700)}")
    return flattened[:8]


def _sanitize(value: Any, *, depth: int = 0) -> Any:
    if depth >= 4:
        return "…"
    if isinstance(value, Mapping):
        result: dict[str, Any] = {}
        for raw_key, item in list(value.items())[:20]:
            key = str(raw_key)
            if _SENSITIVE_KEY.search(key):
                result[key] = "[redacted]"
            elif _LARGE_VALUE_KEY.fullmatch(key):
                result[key] = "[omitted]"
            else:
                result[key] = _sanitize(item, depth=depth + 1)
        return result
    if isinstance(value, (list, tuple)):
        return [_sanitize(item, depth=depth + 1) for item in value[:12]]
    if isinstance(value, str):
        return _truncate(value, 400)
    if value is None or isinstance(value, (bool, int, float)):
        return value
    return _truncate(str(value), 400)


def _truncate(text: str, limit: int) -> str:
    text = text.strip()
    if len(text) <= limit:
        return text
    marker = "\n…\n"
    half = (limit - len(marker)) // 2
    return f"{text[:half]}{marker}{text[-half:]}"


def _clean_summary(text: str) -> str:
    text = text.strip()
    if text.startswith("```") and text.endswith("```"):
        text = text.strip("`").strip()
    return _truncate(" ".join(text.split()), 500)