dynamic_context_manager.py

import asyncio
import json
import tiktoken
from dataclasses import dataclass, field
from enum import Enum
from typing import Optional

from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage, SystemMessage, AIMessage


class RetentionPriority(str, Enum):
    CRITICAL = "critical"
    HIGH = "high"
    MEDIUM = "medium"
    LOW = "low"
    DROPPABLE = "droppable"


@dataclass
class ManagedMessage:
    role: str
    content: str
    turn_id: int
    priority: RetentionPriority = RetentionPriority.MEDIUM
    compressed_version: Optional[str] = None
    is_compressed: bool = False
    original_tokens: int = 0
    current_tokens: int = 0


@dataclass
class CompressionDecision:
    turn_id: int
    action: str
    reason: str
    priority: RetentionPriority


SCORER_SYSTEM = """You are a context importance scorer for an AI agent conversation.
Given the current task goal and a set of conversation turns, score each turn's
importance to the ongoing task.

For each turn, output a JSON object with:
- turn_id: the integer turn ID
- priority: one of "critical", "high", "medium", "low", "droppable"
- reason: one sentence explaining the score

Priority definitions:
- critical: Contains the core problem definition, key constraints, agreed solution, or active bug
- high: Contains important decisions, working code, or data that is actively referenced
- medium: Contains useful context that might be referenced later
- low: Contains exploratory discussion, failed attempts, or verbose explanations
- droppable: Contains filler, off-topic content, or complete repetition of other turns

Output a JSON array of these objects, one per turn. Nothing else.
"""


class DynamicContextManager:
    """
    Agent-driven context manager that scores message importance against
    the current task goal and makes intelligent compression decisions.
    """

    def __init__(
        self,
        budget_tokens: int = 8000,
        scorer_model: str = "gpt-4o-mini",
        compressor_model: str = "gpt-4o-mini",
        working_model: str = "gpt-4o",
    ):
        self.budget_tokens = budget_tokens
        self.enc = tiktoken.encoding_for_model(working_model)
        self.scorer = ChatOpenAI(model=scorer_model, temperature=0.0)
        self.compressor = ChatOpenAI(model=compressor_model, temperature=0.0)
        self.working_llm = ChatOpenAI(model=working_model, temperature=0.0)
        self.messages: list[ManagedMessage] = []
        self.task_goal: str = ""
        self._turn_counter = 0

    def count_tokens(self, text: str) -> int:
        return len(self.enc.encode(text))

    def set_task_goal(self, goal: str):
        """Set the current task goal used for importance scoring."""
        self.task_goal = goal
        print(f"[DynamicCtx] Task goal: {goal[:80]}")

    def add_message(self, role: str, content: str) -> ManagedMessage:
        """Add a message and track its token count."""
        tokens = self.count_tokens(content)
        msg = ManagedMessage(
            role=role,
            content=content,
            turn_id=self._turn_counter,
            original_tokens=tokens,
            current_tokens=tokens,
        )
        self.messages.append(msg)
        self._turn_counter += 1
        return msg

    def total_tokens(self) -> int:
        return sum(m.current_tokens for m in self.messages)

    async def score_importance(self, messages: list[ManagedMessage]) -> list[CompressionDecision]:
        """
        Ask the scorer LLM to evaluate each message's importance
        relative to the current task goal.

        Args:
            messages: Messages to score.

        Returns:
            List of CompressionDecisions with priority and reasoning.
        """
        turns_text = "\n\n".join(
            f"Turn {m.turn_id} [{m.role.upper()}]: {m.content[:500]}"
            for m in messages
        )

        prompt = (
            f"Current task goal: {self.task_goal}\n\n"
            f"Conversation turns to score:\n\n{turns_text}"
        )

        response = await self.scorer.ainvoke([
            SystemMessage(content=SCORER_SYSTEM),
            HumanMessage(content=prompt),
        ])

        try:
            decisions_raw = json.loads(response.content)
            return [
                CompressionDecision(
                    turn_id=d["turn_id"],
                    action="score",
                    reason=d["reason"],
                    priority=RetentionPriority(d["priority"]),
                )
                for d in decisions_raw
            ]
        except (json.JSONDecodeError, KeyError, ValueError) as e:
            print(f"[Scorer] Parse error: {e}. Defaulting to MEDIUM priority.")
            return [
                CompressionDecision(
                    turn_id=m.turn_id, action="score",
                    reason="parse error", priority=RetentionPriority.MEDIUM,
                )
                for m in messages
            ]

    async def compress_message(self, msg: ManagedMessage) -> str:
        """
        Compress a single low-priority message using the structured compressor.
        Stores the original for possible re-expansion.
        """
        prompt = (
            f"Compress this conversation turn to its minimum essential content.\n"
            f"Preserve: all code identifiers, numbers, error messages, file names, and decisions.\n"
            f"Remove: prose explanations, pleasantries, verbose context already captured elsewhere.\n"
            f"Output only the compressed content, no preamble.\n\n"
            f"Turn [{msg.role.upper()}]: {msg.content}"
        )

        response = await self.compressor.ainvoke([HumanMessage(content=prompt)])
        return response.content.strip()

    async def optimize(self) -> dict:
        """
        Main optimization loop. If context exceeds budget:
        1. Score all uncompressed messages by importance against the task goal.
        2. Compress low/droppable priority messages first.
        3. Repeat until within budget or only critical/high messages remain.

        Returns:
            Stats dict with before/after token counts and actions taken.
        """
        before = self.total_tokens()
        if before <= self.budget_tokens:
            return {"action": "none", "before": before, "after": before}

        print(f"[DynamicCtx] Optimizing: {before} / {self.budget_tokens} tokens")

        uncompressed = [m for m in self.messages if not m.is_compressed]
        decisions = await self.score_importance(uncompressed)

        priority_map = {d.turn_id: d for d in decisions}
        for msg in self.messages:
            if msg.turn_id in priority_map:
                msg.priority = priority_map[msg.turn_id].priority

        compress_order = [
            RetentionPriority.DROPPABLE,
            RetentionPriority.LOW,
            RetentionPriority.MEDIUM,
        ]

        actions_taken = []

        for priority_level in compress_order:
            if self.total_tokens() <= self.budget_tokens:
                break

            candidates = [
                m for m in self.messages
                if m.priority == priority_level and not m.is_compressed
            ]

            for msg in candidates:
                if self.total_tokens() <= self.budget_tokens:
                    break

                if priority_level == RetentionPriority.DROPPABLE:
                    original_tokens = msg.current_tokens
                    msg.content = f"[DROPPED - {msg.role} turn {msg.turn_id}: low-signal content]"
                    msg.current_tokens = self.count_tokens(msg.content)
                    msg.is_compressed = True
                    actions_taken.append({
                        "turn": msg.turn_id,
                        "action": "dropped",
                        "saved": original_tokens - msg.current_tokens,
                    })
                else:
                    compressed = await self.compress_message(msg)
                    original_tokens = msg.current_tokens
                    msg.compressed_version = msg.content
                    msg.content = compressed
                    msg.current_tokens = self.count_tokens(compressed)
                    msg.is_compressed = True
                    actions_taken.append({
                        "turn": msg.turn_id,
                        "action": "compressed",
                        "priority": priority_level.value,
                        "saved": original_tokens - msg.current_tokens,
                    })

        after = self.total_tokens()
        print(f"[DynamicCtx] Optimized: {before} -> {after} tokens ({len(actions_taken)} actions)")

        return {"action": "optimized", "before": before, "after": after, "actions": actions_taken}

    def expand_message(self, turn_id: int) -> bool:
        """
        Re-expand a previously compressed message back to its original content.
        Returns True if expansion was possible, False if original is unavailable.
        """
        for msg in self.messages:
            if msg.turn_id == turn_id and msg.compressed_version:
                msg.content = msg.compressed_version
                msg.current_tokens = self.count_tokens(msg.content)
                msg.is_compressed = False
                print(f"[DynamicCtx] Expanded turn {turn_id}: {msg.current_tokens} tokens restored")
                return True
        return False

    def build_prompt_messages(self) -> list:
        """Build the final message list for the working LLM."""
        result = []
        for msg in self.messages:
            if msg.role == "user":
                result.append(HumanMessage(content=msg.content))
            elif msg.role == "assistant":
                result.append(AIMessage(content=msg.content))
            else:
                result.append(SystemMessage(content=f"[{msg.role.upper()}]: {msg.content}"))
        return result


async def main():
    manager = DynamicContextManager(budget_tokens=3000)

    manager.set_task_goal(
        "Debug a FastAPI endpoint that returns 422 Unprocessable Entity "
        "when receiving webhook payloads from Stripe."
    )

    turns = [
        ("user", "Let's start building the Stripe webhook handler in FastAPI."),
        ("assistant", "We'll create a POST /webhook endpoint using FastAPI's Request object to read raw body bytes for Stripe signature verification."),
        ("user", "What's the signature verification process?"),
        ("assistant", "Use stripe.Webhook.construct_event(payload_bytes, sig_header, webhook_secret). The sig_header comes from request.headers.get('stripe-signature'). Store webhook_secret in STRIPE_WEBHOOK_SECRET env var."),
        ("user", "I also want to log all events to a PostgreSQL database."),
        ("assistant", "Create a webhook_events table: id (UUID PK), stripe_event_id (TEXT UNIQUE), event_type (TEXT), payload (JSONB), received_at (TIMESTAMPTZ), processed (BOOL DEFAULT FALSE)."),
        ("user", "I deployed it and now I'm getting 422 Unprocessable Entity on every webhook."),
        ("assistant", "422 means FastAPI is parsing the request body as a Pydantic model before your handler sees it. Use `body: bytes = Body(...)` or read with `await request.body()` to get raw bytes."),
        ("user", "I changed it to `async def webhook(request: Request)` and read with `await request.body()`. Now I get 400: No signatures found."),
        ("assistant", "The 400 means raw bytes don't match Stripe's expected signature. Most common cause: a GZip or logging middleware is consuming the body stream before your handler. Audit your middleware stack."),
    ]

    for role, content in turns:
        manager.add_message(role, content)

    print(f"Before optimization: {manager.total_tokens()} tokens")

    stats = await manager.optimize()
    print(f"After optimization: {stats['after']} tokens")
    print(f"\nActions taken:")
    for action in stats.get("actions", []):
        print(f"  Turn {action['turn']}: {action['action']} (saved {action.get('saved', 0)} tokens)")


if __name__ == "__main__":
    asyncio.run(main())