summarizer.py

from __future__ import annotations

import asyncio
import json
import logging
from pathlib import Path
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from agent.config import Config

logger = logging.getLogger(__name__)

_Q_SYSTEM = (
    "Extract the core intent of the user message as a short phrase. "
    "State WHAT they want, not how they said it. "
    "Reply in the SAME language as the input. "
    "≤15 words. No labels, no trailing punctuation."
)
_A_SYSTEM = (
    "Extract the key result, finding, or decision from the agent response. "
    "Skip all preamble, praise, and process description — only the conclusion matters. "
    "State WHAT was determined or delivered, not how. "
    "Reply in the SAME language as the input. "
    "≤15 words. No labels, no trailing punctuation."
)


def _pick_summarizer_entry(config: "Config", content: str) -> tuple:
    """Smart summarizer routing: GPU when free+fits, CPU fallback otherwise.

    Returns (entry, used_gpu) where entry is the chosen ModelEntry and
    used_gpu indicates whether it's the GPU pool model.
    """
    from agent.config import make_registry
    from agent.core.model_status import get_counts

    # Explicit summarizer pin wins over the background free-cloud offload:
    # without this a `/model summarizer=<entry>` (or [models.summarizer]) pin is
    # silently shadowed whenever the mode allows a free cloud offload, so the
    # user "cannot set it as summarizer". Unpinned → background offload as before.
    reg = make_registry(config)
    cpu_entry = reg.for_role("summarizer") or reg.background
    gpu_pool = config.concurrency.gpu_pool

    # If no GPU pool configured, always use CPU
    if not gpu_pool:
        return cpu_entry, False

    # Find the default model entry — it's the primary GPU model
    default_name = config.model_roles.get("default", "default")
    if default_name not in gpu_pool:
        return cpu_entry, False  # default is not GPU, nothing to route

    default_entry = config.model_entries.get(default_name)
    if default_entry is None:
        return cpu_entry, False

    # Check if GPU is currently busy (main role is running)
    counts = get_counts()
    if counts.get("main", 0) > 0:
        return cpu_entry, False  # GPU busy with main chat

    # Estimate whether the content fits in GPU context (leave 25% headroom)
    from agent.memory.compactor import _count_tokens_approx
    total_tokens = _count_tokens_approx([
        {"role": "system", "content": _Q_SYSTEM},
        {"role": "user", "content": content[:4000]},
    ])
    # Add some margin for the second summary if both run in parallel
    total_tokens = total_tokens * 2 + 200
    ctx = default_entry.ctx_window
    if ctx and total_tokens < ctx * 0.75:
        return default_entry, True  # fits on GPU

    return cpu_entry, False


from contextlib import asynccontextmanager


def _sanitize_summary(text: str) -> str:
    """Truncate at language-switch artifacts (e.g. CJK chars injected into Latin text).

    Local multilingual models sometimes code-switch mid-output. We detect the
    dominant script of the first ~20 chars and truncate at the first run of a
    substantially different script class.
    """
    import unicodedata
    if not text:
        return text

    def _script(ch: str) -> str:
        cp = ord(ch)
        if 0x4E00 <= cp <= 0x9FFF or 0x3400 <= cp <= 0x4DBF or 0xF900 <= cp <= 0xFAFF:
            return "cjk"
        if 0x3040 <= cp <= 0x30FF:  # hiragana/katakana
            return "cjk"
        if 0xAC00 <= cp <= 0xD7AF:  # hangul
            return "cjk"
        if 0x0400 <= cp <= 0x04FF:  # cyrillic
            return "cyrillic"
        if ch.isalpha():
            return "latin"
        return "other"

    # Detect dominant script from first 30 non-other chars.
    scripts: list[str] = [_script(c) for c in text[:60] if _script(c) != "other"]
    if not scripts:
        return text
    dominant = max(set(scripts), key=scripts.count)
    if dominant == "other":
        return text

    # Truncate at first character of a different (non-other, non-dominant) script.
    for i, ch in enumerate(text):
        s = _script(ch)
        if s != "other" and s != dominant:
            return text[:i].rstrip(" ,;:—-")
    return text


@asynccontextmanager
async def _noop():
    yield


async def _consume_stream(stream) -> tuple[list[str], list[str]]:
    content_parts: list[str] = []
    reasoning_parts: list[str] = []
    async for chunk in stream:
        delta = chunk.choices[0].delta if chunk.choices else None
        if delta:
            if delta.content:
                content_parts.append(delta.content)
            if getattr(delta, "reasoning_content", None):
                reasoning_parts.append(delta.reasoning_content)
    return content_parts, reasoning_parts


async def _try_primary(config: "Config", entry, used_gpu: bool, messages: list) -> tuple[list[str], list[str]]:
    """GPU-aware primary attempt — unchanged behavior/semaphore from before W7."""
    from agent.core.llm_client import make_llm_client
    from agent.core.model_status import _inc as _ms_inc, _dec as _ms_dec, gpu_slot as _gpu_slot, provider_label

    client = make_llm_client(config, base_url=entry.base_url, api_key=entry.api_key)
    try:
        from agent.metrics import model_calls
        model_calls.record_entry(entry, role="summarizer")
    except Exception:
        pass
    _ep = provider_label(entry.base_url)
    _role = "sum" if not used_gpu else "main"
    _ms_inc(_role, _ep, entry.model)
    try:
        async with _gpu_slot() if used_gpu else _noop():
            stream = await client.chat.completions.create(
                model=entry.model, messages=messages, stream=True,
            )
            return await _consume_stream(stream)
    finally:
        _ms_dec(_role, _ep, entry.model)
        await client.close()


async def _try_fallback(config: "Config", messages: list) -> tuple[list[str], list[str]]:
    """W7: role-based failover when the GPU-aware primary pick fails outright."""
    from agent.core.llm_retry import open_stream_with_failover
    from agent.core.model_status import _inc as _ms_inc, _dec as _ms_dec, provider_label
    from agent.security.airgap import is_enabled as _airgap_enabled

    stream, _name, entry, client = await open_stream_with_failover(
        config, "background", messages=messages, metrics_role="summarizer",
        local_only=_airgap_enabled(config),
    )
    _ep = provider_label(entry.base_url)
    # Record the entry that actually answered: after failover this is not the
    # primary pick, and the models panel keys its live marker by model.
    _ms_inc("sum", _ep, entry.model)
    try:
        return await _consume_stream(stream)
    finally:
        _ms_dec("sum", _ep, entry.model)
        await client.close()


async def _call_llm_one_line(
    config: "Config",
    system_prompt: str,
    content: str,
) -> str:
    """Stream a one-line summary. Tries the GPU-aware summarizer pick first;
    on any failure falls through to role-based failover (W7) instead of
    giving up outright."""
    entry, used_gpu = _pick_summarizer_entry(config, content)
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": content[:4000]},
    ]
    try:
        content_parts, reasoning_parts = await _try_primary(config, entry, used_gpu, messages)
    except Exception as e:
        logger.info("summarizer: primary entry failed (%s: %s), falling back to failover",
                    type(e).__name__, e)
        content_parts, reasoning_parts = await _try_fallback(config, messages)

    from agent.core.streaming import _clean_output

    raw_content = "".join(content_parts)
    full = _clean_output(raw_content)
    if not full:
        full = _clean_output("".join(reasoning_parts))
    if not full:
        full = raw_content.strip()
    return _sanitize_summary(full)


def _update_json_file(path: Path, updates: dict) -> None:
    try:
        data = json.loads(path.read_text(encoding="utf-8"))
        data.update(updates)
        path.write_text(json.dumps(data, indent=2), encoding="utf-8")
    except Exception as e:
        logger.warning("summarizer: failed to update %s: %s", path, e)


async def summarize_turn_background(
    config: "Config",
    q_path: Path,
    a_path: Path,
) -> bool:
    """Background task: add summary_q / summary_a to the Q and A JSON files.

    Returns True if at least one summary was written, False otherwise.
    Runs in isolation — never touches the agent's primary conversation context.
    Errors are silently logged and never propagate to the caller.
    """
    if not config.ui.q_summaries:
        return False
    try:
        q_content = json.loads(q_path.read_text(encoding="utf-8")).get("content", "")
        a_content = json.loads(a_path.read_text(encoding="utf-8")).get("content", "")

        if not q_content and not a_content:
            return False

        summary_q, summary_a = await asyncio.gather(
            _call_llm_one_line(config, _Q_SYSTEM, q_content) if q_content else asyncio.sleep(0, result=""),
            _call_llm_one_line(config, _A_SYSTEM, a_content) if a_content else asyncio.sleep(0, result=""),
        )

        written = False
        if summary_q:
            await asyncio.to_thread(_update_json_file, q_path, {"summary_q": summary_q})
            written = True
        if summary_a:
            await asyncio.to_thread(_update_json_file, a_path, {"summary_a": summary_a})
            written = True

        logger.debug(
            "summarize_turn_background: done — Q=%r A=%r",
            summary_q[:60] if summary_q else "",
            summary_a[:60] if summary_a else "",
        )
        return written
    except asyncio.CancelledError:
        logger.debug("summarize_turn_background: cancelled")
        raise
    except Exception as _exc:
        logger.warning("summarize_turn_background: error (ignored): %s", _exc)
        return False


_STALE_THRESHOLD = 100  # summaries longer than this were generated by old prompts


async def resummarize_session(
    config: "Config",
    session_id: str,
    force: bool = False,
) -> tuple[int, int]:
    """Re-summarize Q/A entries whose summaries are missing or stale (too long).

    Returns (updated, skipped) counts.
    force=True re-summarizes all entries regardless of length.
    """
    from agent.memory.qa_log import QALogger
    import asyncio as _asyncio

    logger_obj = QALogger(session_id)
    q_dir = logger_obj._get_q_dir()
    a_dir = logger_obj._get_a_dir()

    if not q_dir.exists() and not a_dir.exists():
        return 0, 0

    def _needs_resummarize(summary: str) -> bool:
        if not summary:
            return True
        return force or len(summary) > _STALE_THRESHOLD

    # Collect file pairs by turn_id
    pairs: dict[int, dict] = {}
    for q_file in sorted(q_dir.glob("Q-*.json")) if q_dir.exists() else []:
        try:
            data = json.loads(q_file.read_text(encoding="utf-8"))
            tid = data.get("turn_id")
            if tid is not None:
                pairs.setdefault(tid, {})["q_path"] = q_file
                pairs[tid]["q_data"] = data
        except Exception:
            continue
    for a_file in sorted(a_dir.glob("A-*.json")) if a_dir.exists() else []:
        try:
            data = json.loads(a_file.read_text(encoding="utf-8"))
            tid = data.get("turn_id")
            if tid is not None:
                pairs.setdefault(tid, {})["a_path"] = a_file
                pairs[tid]["a_data"] = data
        except Exception:
            continue

    updated = 0
    skipped = 0
    for tid in sorted(pairs):
        entry = pairs[tid]
        q_data = entry.get("q_data", {})
        a_data = entry.get("a_data", {})
        q_path = entry.get("q_path")
        a_path = entry.get("a_path")

        q_needs = q_path and _needs_resummarize(q_data.get("summary_q", ""))
        a_needs = a_path and _needs_resummarize(a_data.get("summary_a", ""))

        if not q_needs and not a_needs:
            skipped += 1
            continue

        try:
            q_content = q_data.get("content", "") if q_needs else ""
            a_content = a_data.get("content", "") if a_needs else ""

            summary_q, summary_a = await _asyncio.gather(
                _call_llm_one_line(config, _Q_SYSTEM, q_content) if q_content else _asyncio.sleep(0, result=""),
                _call_llm_one_line(config, _A_SYSTEM, a_content) if a_content else _asyncio.sleep(0, result=""),
            )

            if summary_q and q_path:
                await _asyncio.to_thread(_update_json_file, q_path, {"summary_q": summary_q})
            if summary_a and a_path:
                await _asyncio.to_thread(_update_json_file, a_path, {"summary_a": summary_a})
            updated += 1
        except _asyncio.CancelledError:
            raise
        except Exception as exc:
            logger.warning("resummarize_session: turn %s failed (skipped): %s", tid, exc)
            skipped += 1

    return updated, skipped