context.py

Go to the documentation of this file.

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# oc_chat/core/context.py
 
"""
Context-window compression logic.
 
Purpose
-------
Phase 3: :func:`maybe_compress` triggers from the last API
last API ``prompt_tokens`` recorded on the session (or 0 on the first turn), then summarises older *conversation turns* while
keeping the last few turns raw. Summarisation calls are routed to the
provider's ``summary`` tier model.
 
Responsibilities
----------------
- :func:`maybe_compress` — the public entry point; return ``True`` when
  history was rewritten, ``False`` otherwise.
- :func:`_conversation_turns` — split non-system messages into
  user-rooted turn lists.
- :func:`_abbreviate_turn_for_summary` — strip bulky successful tool
  payloads before serialising for the summariser; keep errors verbatim.
- :func:`_effective_prompt_trigger` — uses API-reported prompt size from last turn;
  returns 0 on first turn when no history exists.
 
Diagnostics
-----------
Domain: OC_CHAT.CTX
Levels:
    L2 — compress START / compressed (rebuilt size).
    L3 — threshold comparison, turn-count check, summary model id,
         per-call timing via ``dbg_timed`` label ``summarise``.
    L4 — deep tracing (per-turn abbreviation decisions).
 
Compression-failure paths log at L1 with ``override=True``.
 
Contracts
---------
- ``maybe_compress`` mutates ``session.messages`` only when it
  returns ``True``; on failure the session is left untouched.
- The system prompt at index 0 is preserved verbatim across compression.
- The last ``CONTEXT_KEEP_RECENT_TURNS`` user-rooted turns survive raw.
- ``session.last_turn_max_prompt_tokens`` is reset to 0 on success so
  the next turn re-evaluates the threshold from scratch.
 
Dependencies
------------
- oc_diagnostics (structured debug logging via _dbg)
- oc_chat.core.config (shared config constants and ANSI styles)
- oc_chat.core.client_interface (lazy; ClientInterface type of the summarising client)
- oc_chat.core.session (lazy; ChatSession type of the session being compressed)
- oc_chat.core.tokenizer (lazy; message token counting for the estimate helper)
- oc_ui.presenter_base (lazy; Presenter type for the compression notice)
"""
from __future__ import annotations
 
import copy
import json
from typing import TYPE_CHECKING, Any, List, Dict, Union
 
from oc_diagnostics import _dbg, dbg_timed
 
if TYPE_CHECKING:
    from oc_chat.core.client_interface import ClientInterface
    from oc_chat.core.session import ChatSession
    from oc_ui.presenter_base import Presenter
 
from oc_chat.core.config import (
    CONTEXT_COMPRESS_THRESHOLD,
    CONTEXT_KEEP_RECENT_TURNS,
    MODEL_MAX_TOKENS,
    get_max_context_tokens,
)
 
 
def _estimate_tokens(messages: list[dict[str, Any]]) -> int:
    """Token count estimation for messages.
    
    Uses tiktoken (cl100k_base) when available for accurate counts.
    Falls back to ~4 chars/token heuristic if tiktoken not installed.
    
    This wraps count_message_tokens from the tokenizer module.
    """
    from oc_chat.core.tokenizer import count_message_tokens
    return count_message_tokens(messages)
 
 
 
def _effective_prompt_trigger(session, messages: list) -> int:
    """Use API-reported prompt size from the last turn; 0 on first turn."""
    return int(getattr(session, "last_turn_max_prompt_tokens", 0) or 0)
 
 
def _conversation_turns(messages: list[dict[str, Any]]) -> list[list[dict[str, Any]]]:
    """Split non-system messages into user-rooted turns (each starts with a user message)."""
    if not messages:
        return []
    body = messages[1:] if messages[0].get("role") == "system" else list(messages)
    if not body:
        return []
    turns: list[list[dict[str, Any]]] = []
    cur: list[dict[str, Any]] = []
    for m in body:
        if m.get("role") == "user" and cur:
            turns.append(cur)
            cur = []
        cur.append(m)
    if cur:
        turns.append(cur)
    return turns
 
 
def _abbreviate_tool_for_summary(m: dict[str, Any]) -> dict[str, Any]:
    """Drop bulky successful tool payloads; keep errors verbatim for the summariser."""
    if m.get("role") != "tool":
        return m
    content = m.get("content") or ""
    try:
        obj = json.loads(content)
    except (json.JSONDecodeError, TypeError):
        return m
    if isinstance(obj, dict) and "error" in obj:
        return m
    m2 = dict(m)
    m2["content"] = f"[Successful tool output omitted; {len(content)} chars]"
    return m2
 
 
def _abbreviate_turn_for_summary(turn: list[dict[str, Any]]) -> list[dict[str, Any]]:
    """Aggressively trim old turns before JSON serialisation for the summariser."""
    out: list[dict[str, Any]] = []
    for m in turn:
        role = m.get("role")
        if role == "tool":
            out.append(_abbreviate_tool_for_summary(m))
            continue
        m2 = copy.deepcopy(m)
        if role == "assistant":
            c = m2.get("content")
            if isinstance(c, str) and len(c) > 1200:
                m2["content"] = c[:1200] + "\n…[truncated for summariser]"
            rc = m2.get("reasoning_content")
            if isinstance(rc, str) and len(rc) > 800:
                m2["reasoning_content"] = rc[:800] + "\n…[truncated]"
        out.append(m2)
    return out
 
 
def maybe_compress(session: ChatSession, client: ClientInterface, presenter: Presenter | None = None) -> bool:
    """Compress history if the context budget is exceeded.
 
    Uses API-reported ``prompt_tokens`` from the last turn against the
    session model's context window. Older conversation turns are passed to
    :meth:`DeepSeekClient.one_shot` on ``CONTEXT_COMPRESS_SUMMARY_MODEL`` with
    successful tool bodies stripped; the last ``CONTEXT_KEEP_RECENT_TURNS``
    turns stay verbatim in ``session.messages``.
    """
    func = "maybe_compress"
    messages = session.messages
    max_tokens = get_max_context_tokens(session.model)
    threshold = int(max_tokens * CONTEXT_COMPRESS_THRESHOLD)
    effective = _effective_prompt_trigger(session, messages)
 
    _dbg("OC_CHAT.CTX", func,
         f"effective={effective} threshold={threshold} max_tokens={max_tokens}", 3)
 
    if effective < threshold:
        _dbg("OC_CHAT.CTX", func, "below threshold — skipping", 3)
        return False
 
    turns = _conversation_turns(messages)
    if len(turns) <= CONTEXT_KEEP_RECENT_TURNS:
        _dbg("OC_CHAT.CTX", func,
             f"only {len(turns)} turns <= {CONTEXT_KEEP_RECENT_TURNS} — skipping", 3)
        return False
 
    system_msg = messages[0] if messages and messages[0].get("role") == "system" else None
    if system_msg is None:
        return False
 
    keep_turns = turns[-CONTEXT_KEEP_RECENT_TURNS :]
    old_turns = turns[: -CONTEXT_KEEP_RECENT_TURNS]
 
    _dbg("OC_CHAT.CTX", func,
         f"compressing: {len(old_turns)} old turns, keeping {len(keep_turns)} recent", 2)
 
    abbreviated_old: list[dict[str, Any]] = []
    for t in old_turns:
        for m in _abbreviate_turn_for_summary(t):
            abbreviated_old.append(m)
 
    summary_prompt = [
        {
            "role": "system",
            "content": (
                "Summarise the following older conversation turns in 4-8 short paragraphs. "
                "Preserve critical user goals, file paths, errors, and conclusions. "
                "Omit verbose tool dumps — they were already stripped to placeholders."
            ),
        },
        {
            "role": "user",
            "content": json.dumps(abbreviated_old, ensure_ascii=False),
        },
    ]
 
    summary_model = client.provider.tier("summary") or session.model
    _dbg("OC_CHAT.CTX", func, f"summary model={summary_model}", 3)
    with dbg_timed("OC_CHAT.CTX", func, label="summarise", level=3):
        try:
            summary = client.one_shot(
                summary_prompt,
                summary_model,
                effort="low",
                max_tokens=4096,
            )
        except Exception as exc:
            _dbg("OC_CHAT.CTX", func,
                 f"compression failed: {type(exc).__name__}: {exc}", 1, override=True)
            return False
 
    rebuilt: list[dict[str, Any]] = [system_msg]
    rebuilt.append(
        {"role": "assistant", "content": f"[Earlier conversation summary: {summary}]"}
    )
    for t in keep_turns:
        rebuilt.extend(t)
 
    session.messages = rebuilt
    session.last_turn_max_prompt_tokens = 0
    session.dirty = True
 
    _dbg("OC_CHAT.CTX", func,
         f"compressed: {len(old_turns)} old turns -> summary, messages={len(rebuilt)}", 2)
 
    if presenter:
        presenter.show_compressed()
 
    return True