context.py

"""Context management — size estimation, persistence, compaction, and
summarisation for long-running agent conversations."""

import json
import time

import litellm
from litellm.utils import trim_messages

from config import MY_MODEL, TOOL_RESULTS_DIR, TRANSCRIPT_DIR, PERSIST_THRESHOLD, get_logger

logger = get_logger(__name__)


def estimate_size(msgs) -> int:
    """Rough byte-size estimate of serialised messages."""
    return len(str(msgs))


# ── Large-output persistence ──


def persist_large_output(tool_use_id: str, output: str) -> str:
    """Persist large tool outputs to disk, returning a placeholder with preview."""
    if len(output) <= PERSIST_THRESHOLD:
        return output
    TOOL_RESULTS_DIR.mkdir(parents=True, exist_ok=True)
    path = TOOL_RESULTS_DIR / f"{tool_use_id}.txt"
    logger.info(f"Large output detected, persisting to {path}")

    if not path.exists():
        path.write_text(output)

    return (
        f"<persisted-output>\nFull output: {path}\n"
        f"Preview:\n{output[:2000]}\n</persisted-output>"
    )


# ── Compaction strategies ──


def tool_result_compact(messages: list, max_byte: int = 200_000) -> list:
    """Persist consecutive tool-result tail whose total size exceeds *max_byte*."""
    if not messages or len(messages) == 0:
        return messages
    
    tool_msgs = [m for m in messages if m.get("role") == "tool"]
    if len(tool_msgs) == 0:
        return messages
    def get_tool_msg_size(msgs: list):
        return sum([len(str(m.get("content", ""))) for m in msgs])
    
    if get_tool_msg_size(tool_msgs) <= max_byte:
        return messages
    
    for msg in sorted(tool_msgs, key=lambda m: len(str(m.get("content", ""))), reverse=True):
        msg["content"] = persist_large_output(msg.get("tool_call_id", "unknown"), msg.get("content", ""))
        if get_tool_msg_size(tool_msgs) <= max_byte:
            break
    return messages

# ── Transcript & full summarisation ──

def write_transcript(messages: list) -> str:
    """Write the full message history to a timestamped JSONL transcript."""
    TRANSCRIPT_DIR.mkdir(parents=True, exist_ok=True)
    path = TRANSCRIPT_DIR / f"transcript_{int(time.time())}.jsonl"
    with path.open("w") as f:
        for msg in messages:
            f.write(json.dumps(msg, default=str) + "\n")
    return path


def summarize_history(messages: list) -> str:
    """Ask the LLM to summarise the conversation so work can continue."""
    conversation = json.dumps(messages, default=str)[:80000]
    prompt = (
        "Summarize this coding-agent conversation so work can continue.\n"
        "Preserve: 1. current goal, 2. key findings/decisions, 3. files read/changed, "
        "4. remaining work, 5. user constraints.\nBe compact but concrete.\n\n"
        + conversation
    )
    response = litellm.completion(
        messages=[{"role": "user", "content": prompt}],
        model=MY_MODEL,
        max_tokens=2000,
    )
    return response.choices[0].message.content


def compact_context_by_llm(messages: list) -> list:
    """Write a transcript, summarise, and return a single compacted message."""
    # TODO can add compact tool to LLM
    transcript_path = write_transcript(messages[1:])
    logger.info(f"[transcript saved: {transcript_path}]")
    summary = summarize_history(messages[1:])
    return [
        messages[0],
        {"role": "user", "content": f"[Compacted]\n\n{summary}"}
    ]

def compact_context_fast(messages: list) -> list:
    messages[:] = tool_result_compact(messages)
    messages[:] = trim_messages(messages=messages, model=MY_MODEL)
    return messages