aip-agents-binary 0.0.0__py3-none-any.whl

aip_agents/context/budget.py

@@ -0,0 +1,237 @@
+"""Context budget and compaction decision helpers.
+
+Authors:
+    Raymond Christopher (raymond.christopher@gdplabs.id)
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from dataclasses import dataclass
+from typing import Any, Literal
+
+from aip_agents.context.state import SummaryInsertionPoint
+
+_CHARS_PER_TOKEN = 4
+_BASE_TOKENS_PER_MESSAGE = 4
+_IMAGE_PART_TOKEN_ESTIMATE = 256
+
+CompactionReason = Literal[
+    "below_threshold",
+    "proactive_threshold",
+    "fallback_threshold",
+    "manual_focus",
+    "provider_overflow",
+    "anti_thrash_skip",
+]
+
+
+@dataclass(frozen=True)
+class ContextBudgetConfig:
+    """Configuration for context budget decisions.
+
+    Attributes:
+        threshold_ratio: Fraction of the model context that triggers proactive compaction.
+        fallback_threshold_tokens: Absolute threshold when the model context window is unknown.
+        min_tail_tokens: Minimum recent-token budget preserved after compaction.
+        target_tail_ratio: Fraction of estimated prompt tokens preserved as recent tail.
+        min_head_messages: Number of initial messages always preserved.
+        anti_thrash_min_savings_ratio: Minimum useful savings before repeated compaction.
+        max_retry_after_overflow: Runtime retry cap consumed by overflow handling.
+    """
+
+    threshold_ratio: float = 0.85
+    fallback_threshold_tokens: int = 170_000
+    min_tail_tokens: int = 6_000
+    target_tail_ratio: float = 0.10
+    min_head_messages: int = 2
+    anti_thrash_min_savings_ratio: float = 0.10
+    max_retry_after_overflow: int = 1
+
+
+@dataclass(frozen=True)
+class CompactionDecision:
+    """Result of pure compaction decision-making.
+
+    Attributes:
+        should_compact: Whether the current input should enter the compaction path.
+        reason: Stable reason code for logging, tests, and retry behavior.
+        estimated_prompt_tokens: Estimated prompt size before compaction.
+        model_context_tokens: Model context window used for thresholding, if known.
+        threshold_tokens: Effective token threshold used by this decision.
+        protected_head_indexes: Initial message indexes excluded from replacement.
+        protected_tail_start_index: First recent-tail message index to preserve.
+        focus_topic: Optional user-provided focus topic for manual compaction.
+        summary_insertion_point: Replaceable range where a summary can be inserted.
+    """
+
+    should_compact: bool
+    reason: CompactionReason
+    estimated_prompt_tokens: int
+    model_context_tokens: int | None
+    threshold_tokens: int
+    protected_head_indexes: tuple[int, ...]
+    protected_tail_start_index: int
+    focus_topic: str | None = None
+    summary_insertion_point: SummaryInsertionPoint | None = None
+
+
+def estimate_message_token_counts(messages: Sequence[Any]) -> list[int]:
+    """Estimate token counts per message with conservative heuristics."""
+    return [_estimate_single_message_tokens(message) for message in messages]
+
+
+def estimate_prompt_tokens(messages: Sequence[Any]) -> int:
+    """Estimate total prompt tokens from message sequence."""
+    return sum(estimate_message_token_counts(messages))
+
+
+def compute_protected_regions(
+    token_counts: Sequence[int],
+    config: ContextBudgetConfig | None = None,
+) -> tuple[tuple[int, ...], int]:
+    """Compute protected head indexes and token-budget tail start index."""
+    cfg = config or ContextBudgetConfig()
+    if not token_counts:
+        return (), 0
+
+    message_count = len(token_counts)
+    head_count = min(cfg.min_head_messages, message_count)
+    protected_head_indexes = tuple(range(head_count))
+
+    total_tokens = sum(token_counts)
+    tail_target_tokens = max(cfg.min_tail_tokens, int(total_tokens * cfg.target_tail_ratio))
+
+    running = 0
+    tail_start = message_count
+    for index in range(message_count - 1, -1, -1):
+        running += token_counts[index]
+        tail_start = index
+        if running >= tail_target_tokens:
+            break
+
+    if tail_start < head_count:
+        tail_start = head_count
+
+    return protected_head_indexes, tail_start
+
+
+def classify_provider_overflow_error(error: Exception | str) -> bool:
+    """Return whether an error appears to be a context-window overflow."""
+    text = str(error).lower()
+    patterns = (
+        "context window",
+        "maximum context length",
+        "too many tokens",
+        "prompt is too long",
+        "context_length_exceeded",
+        "input length exceeds",
+        "token limit",
+        "context limit",
+    )
+    return any(pattern in text for pattern in patterns)
+
+
+def decide_compaction(
+    messages: Sequence[Any],
+    *,
+    config: ContextBudgetConfig | None = None,
+    model_context_tokens: int | None = None,
+    focus_topic: str | None = None,
+    provider_overflow: bool = False,
+    recent_savings_ratio: float | None = None,
+) -> CompactionDecision:
+    """Return a pure compaction decision from budget inputs.
+
+    This function does not call models and does not mutate runtime state.
+    """
+    cfg = config or ContextBudgetConfig()
+    token_counts = estimate_message_token_counts(messages)
+    estimated_tokens = sum(token_counts)
+    threshold_tokens = (
+        max(1, int(model_context_tokens * cfg.threshold_ratio))
+        if model_context_tokens is not None
+        else cfg.fallback_threshold_tokens
+    )
+
+    protected_head_indexes, protected_tail_start_index = compute_protected_regions(token_counts, config=cfg)
+    replace_start_index = len(protected_head_indexes)
+    replace_end_index = protected_tail_start_index - 1
+    insertion_point = None
+    if replace_start_index <= replace_end_index:
+        insertion_point = SummaryInsertionPoint(
+            protected_head_indexes=protected_head_indexes,
+            protected_tail_start_index=protected_tail_start_index,
+            replace_start_index=replace_start_index,
+            replace_end_index=replace_end_index,
+        )
+
+    reason: CompactionReason
+    should_compact = False
+
+    normalized_focus_topic = (focus_topic or "").strip() or None
+    if normalized_focus_topic is not None:
+        reason = "manual_focus"
+        should_compact = True
+    elif provider_overflow:
+        reason = "provider_overflow"
+        should_compact = True
+    elif estimated_tokens > threshold_tokens:
+        reason = "proactive_threshold" if model_context_tokens is not None else "fallback_threshold"
+        should_compact = True
+    elif recent_savings_ratio is not None and recent_savings_ratio < cfg.anti_thrash_min_savings_ratio:
+        reason = "anti_thrash_skip"
+        should_compact = False
+    else:
+        reason = "below_threshold"
+
+    return CompactionDecision(
+        should_compact=should_compact,
+        reason=reason,
+        estimated_prompt_tokens=estimated_tokens,
+        model_context_tokens=model_context_tokens,
+        threshold_tokens=threshold_tokens,
+        protected_head_indexes=protected_head_indexes,
+        protected_tail_start_index=protected_tail_start_index,
+        focus_topic=normalized_focus_topic,
+        summary_insertion_point=insertion_point,
+    )
+
+
+def _estimate_single_message_tokens(message: Any) -> int:
+    if isinstance(message, str):
+        content = message
+    elif isinstance(message, dict) and "content" in message:
+        content = message["content"]
+    else:
+        content = getattr(message, "content", message)
+
+    if isinstance(content, list):
+        return _BASE_TOKENS_PER_MESSAGE + sum(_estimate_content_part_tokens(part) for part in content)
+
+    text = str(content or "")
+    return _BASE_TOKENS_PER_MESSAGE + _estimate_text_tokens(text, _CHARS_PER_TOKEN)
+
+
+def _estimate_content_part_tokens(part: Any) -> int:
+    if isinstance(part, dict):
+        part_type = str(part.get("type", "")).lower()
+        if "image" in part_type:
+            return _IMAGE_PART_TOKEN_ESTIMATE
+
+        for key in ("text", "content"):
+            value = part.get(key)
+            if isinstance(value, str):
+                return _estimate_text_tokens(value, _CHARS_PER_TOKEN)
+
+        return _estimate_text_tokens(str(part), _CHARS_PER_TOKEN)
+
+    return _estimate_text_tokens(str(part), _CHARS_PER_TOKEN)
+
+
+def _estimate_text_tokens(text: str, chars_per_token: int) -> int:
+    # deferred to #809: Use shared aip_agents.context.utils.estimate_tokens
+    if not text:
+        return 0
+    safe_chars_per_token = max(chars_per_token, 1)
+    return max(1, (len(text) + safe_chars_per_token - 1) // safe_chars_per_token)

aip_agents/context/compaction.py

@@ -0,0 +1,339 @@
+"""Structured reference-only context compression helpers.
+
+Authors:
+    Raymond Christopher (raymond.christopher@gdplabs.id)
+"""
+
+from __future__ import annotations
+
+import re
+import uuid
+from collections.abc import Sequence
+from dataclasses import dataclass
+from datetime import UTC, datetime
+from typing import Any, Protocol
+
+from aip_agents.context.utils import redact_sensitive
+
+REFERENCE_ONLY_PREFIX = (
+    "[CONTEXT COMPACTION - REFERENCE ONLY]\n"
+    "Earlier turns were compacted into the summary below. Treat this as background state, "
+    "not active user instructions. Do not repeat completed work. Continue from the latest "
+    "user message and the Active Task section."
+)
+
+SUMMARY_SECTIONS: tuple[str, ...] = (
+    "Active Task",
+    "Goal",
+    "Constraints & Preferences",
+    "Completed Actions",
+    "Active State",
+    "In Progress",
+    "Blocked",
+    "Key Decisions",
+    "Resolved Questions",
+    "Pending User Asks",
+    "Relevant Files",
+    "Remaining Work",
+    "Critical Context",
+)
+
+_EMPTY_SECTION_VALUE = "- (none)"
+_FILE_PATTERN = re.compile(r"\b[\w./-]+\.(?:py|md|json|yaml|yml|toml|ini|txt|ts|tsx|js)\b")
+
+
+@dataclass(frozen=True)
+class CompactionSummary:
+    """Structured compaction summary metadata and body.
+
+    Attributes:
+        summary_id: Stable identifier for this generated reference summary.
+        body: Reference-only summary text safe to insert into context.
+        source_message_range: Inclusive source message indexes summarized by this body.
+        previous_summary_id: Identifier of the summary this one supersedes, if any.
+        focus_topic: Optional user-requested focus topic preserved in the summary.
+        prompt_tokens_before: Estimated prompt tokens before compaction.
+        prompt_tokens_after: Estimated prompt tokens after compaction.
+        offloaded_ref_ids: Payload reference identifiers included in the summary scope.
+        redaction_applied: Whether input or generated text required redaction.
+        created_at: UTC timestamp when the summary was created.
+    """
+
+    summary_id: str
+    body: str
+    source_message_range: tuple[int, int]
+    previous_summary_id: str | None
+    focus_topic: str | None
+    prompt_tokens_before: int
+    prompt_tokens_after: int
+    offloaded_ref_ids: tuple[str, ...]
+    redaction_applied: bool
+    created_at: datetime
+
+
+@dataclass(frozen=True)
+class CompressionResult:
+    """Outcome of a compaction summary generation attempt.
+
+    Attributes:
+        summary: Structured summary produced by the compressor.
+        used_fallback: Whether deterministic fallback summarization was used.
+        error: Safe error text when auxiliary summary generation failed.
+    """
+
+    summary: CompactionSummary
+    used_fallback: bool
+    error: str | None = None
+
+
+class SummarizerClient(Protocol):
+    """Minimal abstraction for auxiliary summary generation."""
+
+    def summarize(self, prompt: str) -> str:
+        """Return summary text from a redacted prompt input."""
+
+
+class ContextCompressor:
+    """Generate structured reference-only summaries from prior conversation turns."""
+
+    def __init__(self, summarizer: SummarizerClient | None = None) -> None:
+        """Initialize compressor with optional auxiliary summarizer client."""
+        self._summarizer = summarizer
+
+    def compress(
+        self,
+        messages: Sequence[Any],
+        *,
+        previous_summary: CompactionSummary | None = None,
+        focus_topic: str | None = None,
+        offloaded_ref_ids: Sequence[str] = (),
+        source_message_range: tuple[int, int] | None = None,
+        prompt_tokens_before: int = 0,
+        prompt_tokens_after: int = 0,
+    ) -> CompressionResult:
+        """Create a structured summary with safe fallback behavior.
+
+        Usage::
+
+            from aip_agents.context.compaction import ContextCompressor
+
+            compressor = ContextCompressor()  # uses fallback summarization
+            result = compressor.compress(messages)
+            # result.summary.body contains the structured summary
+            # result.used_fallback is True when no external summarizer was set
+        """
+        redacted_message_texts, message_roles, input_redacted = _prepare_redacted_inputs(messages, previous_summary)
+        generated, used_fallback, generation_error, output_redacted = self._generate_summary_text(
+            redacted_message_texts
+        )
+        input_redacted = input_redacted or output_redacted
+
+        sections = _extract_sections(
+            messages=redacted_message_texts,
+            message_roles=message_roles,
+            generated_summary=generated,
+            focus_topic=focus_topic,
+            previous_summary=previous_summary,
+        )
+        body = _render_summary_body(sections)
+
+        summary_id = f"summary_{uuid.uuid4().hex[:12]}"
+        message_range = _resolve_message_range(messages, source_message_range)
+        summary = CompactionSummary(
+            summary_id=summary_id,
+            body=body,
+            source_message_range=message_range,
+            previous_summary_id=previous_summary.summary_id if previous_summary else None,
+            focus_topic=(focus_topic or "").strip() or None,
+            prompt_tokens_before=prompt_tokens_before,
+            prompt_tokens_after=prompt_tokens_after,
+            offloaded_ref_ids=tuple(offloaded_ref_ids),
+            redaction_applied=input_redacted,
+            created_at=datetime.now(UTC),
+        )
+        return CompressionResult(summary=summary, used_fallback=used_fallback, error=generation_error)
+
+    def _generate_summary_text(self, redacted_message_texts: Sequence[str]) -> tuple[str, bool, str | None, bool]:
+        redacted_input = "\n".join(redacted_message_texts)
+        if self._summarizer is None:
+            return "", True, None, False
+        try:
+            generated_raw = self._summarizer.summarize(redacted_input)
+            generated, output_redacted = redact_sensitive(str(generated_raw))
+            return generated, False, None, output_redacted
+        except Exception as error:  # pragma: no cover - behavior exercised by tests
+            return "", True, str(error), False
+
+
+def _extract_sections(
+    *,
+    messages: Sequence[str],
+    message_roles: Sequence[str] = (),
+    generated_summary: str,
+    focus_topic: str | None,
+    previous_summary: CompactionSummary | None,
+) -> dict[str, str]:
+    """Extract section content from messages and optional generated summary."""
+    latest_user = _latest_user_message(messages, message_roles)
+    decision_lines = _collect_lines(messages, keywords=("decision", "decided", "chose"))
+    blocker_lines = _collect_lines(messages, keywords=("block", "blocked", "waiting", "dependency"))
+    done_lines = _collect_lines(messages, keywords=("done", "completed", "implemented", "finished"))
+    pending_lines = _collect_lines(messages, keywords=("todo", "next", "remaining", "pending"))
+    question_lines = _collect_lines(messages, keywords=("?", "question", "clarify"))
+    file_hits = sorted({match for text in messages for match in _FILE_PATTERN.findall(text)})
+
+    sections: dict[str, str] = dict.fromkeys(SUMMARY_SECTIONS, _EMPTY_SECTION_VALUE)
+    sections["Active Task"] = f"- {latest_user.strip()[:500]}" if latest_user else _EMPTY_SECTION_VALUE
+    sections["Goal"] = f"- Focus on {focus_topic}" if focus_topic else "- Continue current requested scope"
+    sections["Constraints & Preferences"] = "- Preserve prior constraints and avoid re-running completed work"
+    sections["Completed Actions"] = _format_lines(done_lines)
+    sections["Active State"] = "- Compacted reference summary prepared"
+    sections["In Progress"] = _format_lines(pending_lines)
+    sections["Blocked"] = _format_lines(blocker_lines)
+    sections["Key Decisions"] = _format_lines(decision_lines)
+    sections["Resolved Questions"] = _EMPTY_SECTION_VALUE
+    sections["Pending User Asks"] = _format_lines(question_lines)
+    sections["Relevant Files"] = _format_lines([f"`{path}`" for path in file_hits])
+    sections["Remaining Work"] = _format_lines(pending_lines)
+
+    _merge_generated_sections(sections, generated_summary)
+
+    critical_lines: list[str] = []
+    if focus_topic:
+        critical_lines.append(f"Focus topic: {focus_topic}")
+    if previous_summary is not None:
+        critical_lines.append(
+            f"Previous summary id: {previous_summary.summary_id}"
+            if previous_summary is not None
+            else "Previous summary id: none"
+        )
+    if generated_summary:
+        critical_lines.append(generated_summary[:1_000])
+    sections["Critical Context"] = _format_lines(critical_lines)
+
+    return sections
+
+
+def _collect_lines(messages: Sequence[str], *, keywords: Sequence[str]) -> list[str]:
+    matched: list[str] = []
+    keyword_set = tuple(keyword.lower() for keyword in keywords)
+    for text in messages:
+        for line in text.splitlines():
+            lowered = line.lower()
+            if any(keyword in lowered for keyword in keyword_set):
+                cleaned = line.strip()
+                if cleaned and cleaned not in matched:
+                    matched.append(cleaned)
+    return matched
+
+
+def _render_summary_body(sections: dict[str, str]) -> str:
+    lines = [REFERENCE_ONLY_PREFIX, ""]
+    for section in SUMMARY_SECTIONS:
+        lines.append(f"## {section}")
+        lines.append(sections.get(section, _EMPTY_SECTION_VALUE))
+        lines.append("")
+    return "\n".join(lines).strip()
+
+
+def _prepare_redacted_inputs(
+    messages: Sequence[Any], previous_summary: CompactionSummary | None
+) -> tuple[list[str], list[str], bool]:
+    redacted_message_texts: list[str] = []
+    message_roles: list[str] = []
+    input_redacted = False
+    for message in messages:
+        redacted_text, redacted = redact_sensitive(_message_to_text(message))
+        redacted_message_texts.append(redacted_text)
+        input_redacted = input_redacted or redacted
+        message_roles.append(str(message.get("role", "")) if isinstance(message, dict) else "")
+
+    if previous_summary is not None:
+        previous_body, prev_redacted = redact_sensitive(previous_summary.body)
+        redacted_message_texts = [previous_body] + redacted_message_texts
+        message_roles = [""] + message_roles
+        input_redacted = input_redacted or prev_redacted
+
+    return redacted_message_texts, message_roles, input_redacted
+
+
+def _resolve_message_range(messages: Sequence[Any], source_message_range: tuple[int, int] | None) -> tuple[int, int]:
+    if source_message_range is not None:
+        return source_message_range
+    if messages:
+        return (0, len(messages) - 1)
+    return (-1, -1)
+
+
+def _merge_generated_sections(sections: dict[str, str], generated_summary: str) -> None:
+    if not generated_summary:
+        return
+    parsed = _parse_generated_sections(generated_summary)
+    for section in SUMMARY_SECTIONS:
+        content = parsed.get(section)
+        if content:
+            sections[section] = content
+
+
+def _parse_generated_sections(generated_summary: str) -> dict[str, str]:
+    """Parse a generated summary for structured sections matching SUMMARY_SECTIONS."""
+    parsed: dict[str, str] = {}
+    current_section: str | None = None
+    current_lines: list[str] = []
+    for line in generated_summary.splitlines():
+        heading_match = re.match(r"^##\s+(.+)", line.strip())
+        if heading_match:
+            _flush_parsed_section(parsed, current_section, current_lines)
+            heading = heading_match.group(1).strip()
+            current_section = heading if heading in SUMMARY_SECTIONS else None
+            current_lines = []
+        elif current_section:
+            stripped = line.strip()
+            if stripped:
+                cleaned = stripped.removeprefix("- ").strip()
+                current_lines.append(cleaned or stripped)
+    _flush_parsed_section(parsed, current_section, current_lines)
+    return parsed
+
+
+def _flush_parsed_section(parsed: dict[str, str], current_section: str | None, current_lines: list[str]) -> None:
+    if not current_section or not current_lines:
+        return
+    content = _format_lines(current_lines)
+    if content != _EMPTY_SECTION_VALUE:
+        parsed[current_section] = content
+
+
+def _latest_user_message(messages: Sequence[str], message_roles: Sequence[str] = ()) -> str:
+    """Return the latest user/human message text from the given sequence."""
+    if message_roles:
+        for text, role in zip(reversed(messages), reversed(message_roles)):
+            if role.lower() in {"user", "human"} and text.strip():
+                return text.strip()
+    for text in reversed(messages):
+        if text.strip():
+            return text.strip()
+    return ""
+
+
+def _format_lines(lines: Sequence[str]) -> str:
+    if not lines:
+        return _EMPTY_SECTION_VALUE
+    unique = []
+    for line in lines:
+        normalized = line.strip()
+        if normalized and normalized not in unique:
+            unique.append(normalized)
+    if not unique:
+        return _EMPTY_SECTION_VALUE
+    return "\n".join(f"- {line}" for line in unique)
+
+
+def _message_to_text(message: Any) -> str:
+    if isinstance(message, str):
+        return message
+    if isinstance(message, dict):
+        content = message.get("content")
+        return str(content or "")
+    content = getattr(message, "content", "")
+    return str(content or "")