tatlock/src/responses/context.py

"""
Context window management and token counting.

Handles:
- Token counting (approximate for now, exact in future)
- Context trimming to fit model limits
- Reserve tokens for output generation
"""

from typing import Any


class ContextWindow:
    """
    Manage token limits and context trimming.

    Real production context window management:
    - Approximate token counting (character-based)
    - Context trimming (keep most recent within limits)
    - Reserve tokens for model output
    - Future: Integration with tiktoken for exact counts
    """

    def __init__(self, max_tokens: int = 4096):
        """
        Initialize context window manager.

        Args:
            max_tokens: Maximum context window size in tokens
        """
        self.max_tokens = max_tokens

    async def count_tokens(self, items: list[Any]) -> int:
        """
        Count tokens in items.

        Current implementation: Approximate via character count
        Future: Use tiktoken or similar for exact token counts

        Approximation: ~4 characters per token (common for English text)

        Args:
            items: List of items to count (messages, strings, dicts, etc.)

        Returns:
            int: Approximate token count
        """
        total_chars = 0

        for item in items:
            if isinstance(item, str):
                total_chars += len(item)
            elif isinstance(item, dict):
                # Count all values in dict
                total_chars += len(str(item))
            else:
                # Fallback: convert to string
                total_chars += len(str(item))

        # Approximate: 4 characters per token
        return total_chars // 4

    async def trim_to_fit(self, items: list[Any], reserve_tokens: int = 512) -> list[Any]:
        """
        Trim items to fit within context window.

        Strategy:
        1. Reserve tokens for output generation
        2. Keep most recent items that fit
        3. Drop oldest items first

        Args:
            items: List of items to trim
            reserve_tokens: Tokens to reserve for output (default 512)

        Returns:
            list: Trimmed list of items that fit in context
        """
        available_tokens = self.max_tokens - reserve_tokens

        if available_tokens <= 0:
            # Edge case: reserve too large
            return []

        # Start from most recent, work backwards
        kept_items: list[Any] = []
        current_tokens = 0

        for item in reversed(items):
            item_tokens = await self.count_tokens([item])

            if current_tokens + item_tokens <= available_tokens:
                # Fits - keep this item
                kept_items.insert(0, item)
                current_tokens += item_tokens
            else:
                # Doesn't fit - stop here
                break

        return kept_items

    async def fits_in_context(self, items: list[Any], reserve_tokens: int = 512) -> bool:
        """
        Check if items fit within context window.

        Args:
            items: List of items to check
            reserve_tokens: Tokens to reserve for output

        Returns:
            bool: True if items fit in context
        """
        total_tokens = await self.count_tokens(items)
        available_tokens = self.max_tokens - reserve_tokens
        return total_tokens <= available_tokens

    async def get_usage_stats(self, items: list[Any], reserve_tokens: int = 512) -> dict:
        """
        Get context window usage statistics.

        Args:
            items: List of items to analyze
            reserve_tokens: Tokens reserved for output

        Returns:
            dict: Usage statistics with keys:
                - total_tokens: Total tokens in items
                - max_tokens: Maximum context window
                - reserved_tokens: Tokens reserved for output
                - available_tokens: Tokens available for input
                - usage_percent: Percentage of context used
                - fits: Whether items fit in context
        """
        total_tokens = await self.count_tokens(items)
        available_tokens = self.max_tokens - reserve_tokens

        usage_percent = (total_tokens / available_tokens * 100) if available_tokens > 0 else 100.0

        return {
            "total_tokens": total_tokens,
            "max_tokens": self.max_tokens,
            "reserved_tokens": reserve_tokens,
            "available_tokens": available_tokens,
            "usage_percent": round(usage_percent, 2),
            "fits": total_tokens <= available_tokens,
        }

    # ========================================================================
    # Future: Exact Token Counting
    # ========================================================================

    async def count_tokens_exact(self, text: str, model: str = "gpt-3.5-turbo") -> int:
        """
        Count exact tokens using tiktoken.

        TODO: Integrate tiktoken library for exact token counting
        TODO: Support multiple model encodings (GPT-4, Claude, etc.)

        Args:
            text: Text to count
            model: Model name for encoding

        Returns:
            int: Exact token count (placeholder - returns approximate for now)
        """
        # Placeholder - would use tiktoken here
        return await self.count_tokens([text])