devplacepy/devplacepy/services/devii/agentic/compaction.py

# retoor <retoor@molodetz.nl>

from __future__ import annotations

import json
import logging
from typing import Any

from ..errors import LLMError
from ..text import normalize_newlines

logger = logging.getLogger("devii.agentic.compaction")

SUMMARY_INPUT_CAP = 600_000
CONTEXT_COMPACT_TARGET_RATIO = 0.4
CONTEXT_LENGTH_ERROR_CODES = {"context_length_exceeded"}
CONTEXT_LENGTH_ERROR_PHRASES = (
    "maximum context length",
    "context length exceeded",
    "reduce the length",
    "context_length_exceeded",
)
SUMMARY_PROMPT = (
    "Summarize the following assistant conversation segment as a concise factual log of "
    "actions taken, tools called, entities created or changed, conclusions reached, and "
    "outstanding work. Keep identifiers, slugs, uids, and decisions verbatim. "
    "Write plain markdown with real line breaks (never the two-character sequence \\n). "
    "Use headings and bullet lists where helpful. Maximum 800 words.\n\n"
    "---\n\n"
)


def context_size(messages: list[dict[str, Any]]) -> int:
    return len(json.dumps(messages, default=str))


def shrink_large_messages(
    messages: list[dict[str, Any]], max_message_chars: int, skip_leading: int = 1
) -> bool:
    shrunk = False
    for message in messages[skip_leading:]:
        content = message.get("content")
        if isinstance(content, str) and len(content) > max_message_chars:
            removed = len(content) - max_message_chars
            message["content"] = (
                content[:max_message_chars]
                + f"\n...[truncated {removed} more chars to fit the model's context window]"
            )
            shrunk = True
    return shrunk


def is_context_length_error(exc: LLMError) -> bool:
    if exc.details.get("status") != 400:
        return False
    body = str(exc.details.get("body") or "")
    haystack = f"{exc.message} {body}".lower()
    if any(phrase in haystack for phrase in CONTEXT_LENGTH_ERROR_PHRASES):
        return True
    try:
        parsed = json.loads(body)
    except (ValueError, TypeError):
        return False
    error = parsed.get("error") if isinstance(parsed, dict) else None
    code = str((error or {}).get("code") or "").lower() if isinstance(error, dict) else ""
    return code in CONTEXT_LENGTH_ERROR_CODES


def find_compaction_split(messages: list[dict[str, Any]], keep_tail: int) -> int:
    if len(messages) <= keep_tail:
        return 1
    candidate = len(messages) - keep_tail
    fallback = 0
    while candidate > 1:
        role = messages[candidate].get("role")
        if role == "user":
            return candidate
        if role != "tool" and fallback == 0:
            fallback = candidate
        candidate -= 1
    return fallback or 1


def _segment_plain(messages: list[dict[str, Any]]) -> str:
    parts: list[str] = []
    for message in messages:
        role = str(message.get("role") or "unknown")
        content = message.get("content")
        if isinstance(content, str) and content.strip():
            parts.append(f"{role}:\n{normalize_newlines(content)}")
            continue
        if isinstance(content, list):
            chunks: list[str] = []
            for part in content:
                if isinstance(part, dict) and part.get("type") == "text":
                    chunks.append(str(part.get("text") or ""))
                elif isinstance(part, str):
                    chunks.append(part)
            text = normalize_newlines("\n".join(c for c in chunks if c))
            if text.strip():
                parts.append(f"{role}:\n{text}")
            continue
        tool_calls = message.get("tool_calls")
        if tool_calls:
            names = []
            for call in tool_calls:
                fn = (call or {}).get("function") or {}
                name = fn.get("name") or "tool"
                names.append(str(name))
            if names:
                parts.append(f"{role}: called {', '.join(names)}")
    return "\n\n".join(parts)


async def compact_messages(
    llm: Any,
    messages: list[dict[str, Any]],
    keep_tail: int,
    max_summary_chars: int = SUMMARY_INPUT_CAP,
) -> list[dict[str, Any]]:
    if len(messages) < keep_tail + 3:
        return messages
    split = find_compaction_split(messages, keep_tail)
    if split <= 1:
        return messages
    system_message = messages[0]
    middle = messages[1:split]
    tail = messages[split:]
    if not middle:
        return messages

    segment = _segment_plain(middle)[:max_summary_chars]
    if not segment.strip():
        segment = json.dumps(middle, default=str)[:max_summary_chars]
    try:
        summary = await llm.summarize(SUMMARY_PROMPT + segment)
    except Exception:  # noqa: BLE001 - compaction must never break the loop
        logger.exception("Compaction summary failed; keeping full context")
        return messages

    summary = normalize_newlines(summary or "").strip()
    if not summary:
        return messages
    logger.info("Compacted %d messages into a summary", len(middle))
    return [
        system_message,
        {
            "role": "assistant",
            "content": f"[compacted earlier turns]\n\n{summary}",
        },
        *tail,
    ]


async def shrink_and_compact(
    llm: Any,
    messages: list[dict[str, Any]],
    keep_tail: int,
    message_cap: int,
    summary_max_chars: int,
) -> list[dict[str, Any]]:
    shrink_large_messages(messages, message_cap)
    return await compact_messages(llm, messages, keep_tail, summary_max_chars)