gns3-server/gns3server/agent/gns3_copilot/agent/context_manager.py

# SPDX-License-Identifier: AGPL-3.0-or-later
#
# Context Window Manager for GNS3-Copilot
#
# This module manages context window limits for different LLM models,
# implementing intelligent message trimming and token counting strategies.

"""
Context Window Manager for GNS3-Copilot

This module provides context window management for different LLM models,
including:
- Model-specific context window limits
- Token counting for messages
- Message trimming strategies
- System message preservation
"""

import json
import logging
from typing import Any, Literal

from langchain_core.messages import (
    AIMessage,
    HumanMessage,
    SystemMessage,
    ToolMessage,
)
from langchain_core.messages.utils import trim_messages

logger = logging.getLogger(__name__)


# ============================================================================
# Token Counting - Using tiktoken for accuracy
# ============================================================================

# Global tiktoken encoding cache (lazy loading)
_tiktoken_encoding = None


def _get_tiktoken_encoding():
    """
    Get tiktoken encoding instance (cached).

    Uses cl100k_base encoding (GPT-4) which is a good approximation
    for most modern LLMs including OpenAI, Anthropic, and DeepSeek.

    Returns:
        Encoding object

    Raises:
        ImportError: If tiktoken is not installed
    """
    global _tiktoken_encoding
    if _tiktoken_encoding is None:
        try:
            import tiktoken
            _tiktoken_encoding = tiktoken.get_encoding("cl100k_base")
            logger.debug("Using tiktoken (cl100k_base) for accurate token counting")
        except ImportError:
            raise ImportError(
                "tiktoken is required for accurate token counting. "
                "Please install it with: pip install tiktoken>=0.8.0"
            )
    return _tiktoken_encoding


def count_tokens_accurately(text: str) -> int:
    """
    Count tokens in text accurately using tiktoken.

    Args:
        text: The text to count tokens for

    Returns:
        Number of tokens

    Raises:
        ImportError: If tiktoken is not installed
    """
    if not text:
        return 0

    encoding = _get_tiktoken_encoding()
    try:
        return len(encoding.encode(text))
    except Exception as e:
        logger.error("tiktoken encoding failed: %s", e)
        raise


def count_messages_tokens(messages: list[Any]) -> int:
    """
    Count total tokens in a list of messages accurately.

    Args:
        messages: List of LangChain messages

    Returns:
        Total token count

    Raises:
        ImportError: If tiktoken is not installed
    """
    total = 0
    for msg in messages:
        if hasattr(msg, 'content') and msg.content:
            # Handle both string and complex content
            content = str(msg.content)
            total += count_tokens_accurately(content)
    return total


# ============================================================================
# Tool Definition Token Estimation
# ============================================================================

def estimate_tool_tokens(tools: list[Any]) -> int:
    """
    Estimate the token count of tool definitions.

    When tools are bound to an LLM, their schemas (name, description, parameters)
    are converted to JSON and sent to the LLM. This function estimates how many
    tokens those definitions will consume.

    Args:
        tools: List of LangChain BaseTool instances

    Returns:
        Estimated token count for all tool definitions

    Raises:
        ImportError: If tiktoken is not installed
    """
    if not tools:
        return 0

    total_tokens = 0
    encoding = _get_tiktoken_encoding()

    for tool in tools:
        try:
            # Build tool schema as it would be sent to LLM
            tool_schema = {
                "type": "function",
                "function": {
                    "name": tool.name,
                    "description": tool.description if hasattr(tool, 'description') else "",
                }
            }

            # Add parameters schema if available
            if hasattr(tool, 'args_schema') and tool.args_schema:
                try:
                    tool_schema["function"]["parameters"] = tool.args_schema.schema()
                except Exception:
                    # If schema generation fails, use empty object
                    tool_schema["function"]["parameters"] = {"type": "object"}

            # Count tokens in the schema
            schema_str = json.dumps(tool_schema, ensure_ascii=False)
            tool_tokens = len(encoding.encode(schema_str))
            total_tokens += tool_tokens

            logger.debug(
                "Tool '%s': ~%d tokens (schema size: %d chars)",
                tool.name, tool_tokens, len(schema_str)
            )

        except Exception as e:
            logger.error("Failed to estimate tokens for tool '%s': %s", tool.name, e)
            raise

    logger.info("Tool definitions estimated at ~%d total tokens (%d tools)", total_tokens, len(tools))
    return total_tokens


# NOTE: Built-in model context limits have been removed.
# Model providers frequently update context limits, and maintaining this list is not sustainable.
# Users MUST configure context_limit explicitly in their LLM model configurations.
#
# For reference, common model context limits as of 2025:
# - OpenAI GPT-4o: 128K tokens
# - OpenAI GPT-4 Turbo: 128K tokens
# - OpenAI GPT-3.5 Turbo: 16K tokens
# - Anthropic Claude 3.5 Sonnet: 200K tokens
# - Google Gemini 1.5 Pro: 2.8M tokens
# - DeepSeek Chat: 128K tokens
#
# Always verify current limits from official provider documentation.


def get_model_context_limit(
    model_name: str,
    llm_config: dict[str, Any] | None = None
) -> int:
    """
    Get the context window limit for a given model.

    IMPORTANT: context_limit MUST be provided in llm_config.
    Model providers frequently update context limits, so built-in defaults are NOT used.
    Users must configure this value explicitly.

    NOTE: context_limit unit is K tokens (1 K = 1000 tokens).
    Example: 128 means 128K = 128,000 tokens.

    Args:
        model_name: Name of the model (e.g., "gpt-4o", "deepseek-chat")
        llm_config: LLM config dict from database (must contain context_limit in K tokens)

    Returns:
        int: Maximum context window size in tokens (actual number, not K)

    Raises:
        ValueError: If context_limit is not provided or invalid
    """
    # Check database config for context_limit
    if llm_config and "context_limit" in llm_config:
        db_limit_k = llm_config["context_limit"]
        if isinstance(db_limit_k, int) and db_limit_k > 0:
            # Convert K tokens to actual tokens
            actual_tokens = db_limit_k * 1000
            logger.debug(
                "Using database config context limit: %dK tokens (%d tokens) for model '%s'",
                db_limit_k, actual_tokens, model_name
            )
            return actual_tokens
        else:
            raise ValueError(
                f"Invalid context_limit in database config: {db_limit_k} "
                f"(type={type(db_limit_k).__name__}, expected positive integer in K tokens)"
            )

    # No context_limit provided - this is a configuration error
    raise ValueError(
        f"context_limit is required but not provided for model '{model_name}'. "
        f"Please configure context_limit in your LLM model configuration (unit: K tokens). "
        f"Example: 128 means 128K = 128,000 tokens. "
        f"Refer to the model provider's documentation for the current context window size."
    )


def calculate_max_tokens(
    model_limit: int,
    strategy: Literal["conservative", "balanced", "aggressive"] = "balanced"
) -> int:
    """
    Calculate the maximum tokens to use, reserving space for output.

    Args:
        model_limit: Model's context window limit
        strategy: How aggressively to use the context window
            - "conservative": Use 60% of limit (safer, more reserved for output)
            - "balanced": Use 75% of limit (default)
            - "aggressive": Use 85% of limit (maximize input, minimal output reserve)

    Returns:
        int: Maximum tokens for input messages
    """
    ratios = {
        "conservative": 0.60,
        "balanced": 0.75,
        "aggressive": 0.85,
    }

    ratio = ratios.get(strategy, 0.75)
    max_tokens = int(model_limit * ratio)

    logger.debug(
        "Context limit: model=%d, strategy=%s, usable=%d tokens",
        model_limit, strategy, max_tokens
    )

    return max_tokens


def trim_messages_for_context(
    messages: list[Any],
    model_name: str,
    llm_config: dict[str, Any] | None = None,
    strategy: Literal["conservative", "balanced", "aggressive"] = "balanced",
    preserve_system: bool = True,
    tool_tokens: int = 0,
) -> list[Any]:
    """
    Trim messages to fit within model's context window.

    This function uses tiktoken for accurate token counting and intelligently
    reduces message history while preserving conversation flow.

    IMPORTANT: Tool definitions are sent separately by LangChain and count towards
    the context limit. This function accounts for tool tokens when making
    trimming decisions.

    Args:
        messages: List of LangChain messages (HumanMessage, AIMessage, etc.)
        model_name: Name of the LLM model being used
        llm_config: Optional LLM config dict from database (may contain context_limit)
        strategy: How aggressively to use the context window
        preserve_system: Whether to always preserve system messages
        tool_tokens: Token count for tool definitions (these are sent separately by LangChain)

    Returns:
        list: Trimmed list of messages that fit within context limit

    Examples:
        >>> messages = [HumanMessage("Hello"), AIMessage("Hi there!")]
        >>> trimmed = trim_messages_for_context(messages, "gpt-4o", tool_tokens=1000)
        >>> len(trimmed) <= len(messages)
        True
    """
    if not messages:
        return messages

    # Get model's context limit (from database or built-in defaults)
    model_limit = get_model_context_limit(model_name, llm_config)

    # Calculate usable tokens (reserve space for output)
    max_tokens = calculate_max_tokens(model_limit, strategy)

    # Account for tool tokens - these are sent separately by LangChain
    # and count towards the context limit
    available_for_messages = max_tokens - tool_tokens

    if available_for_messages < 0:
        logger.warning(
            "Tool definitions (%d tokens) exceed input budget (%d tokens). "
            "Consider reducing context_limit or using fewer tools.",
            tool_tokens, max_tokens
        )
        available_for_messages = 0

    # Check if trimming is needed
    current_tokens = count_messages_tokens(messages)

    if current_tokens <= available_for_messages:
        logger.debug(
            "Messages fit in context: %d / %d tokens (available: %d, tools: %d)",
            current_tokens, max_tokens, available_for_messages, tool_tokens
        )
        return messages

    logger.info(
        "Trimming messages: %d → %d tokens (budget: %d, tools: %d)",
        current_tokens, available_for_messages, max_tokens, tool_tokens
    )

    # Manually separate and trim to ensure system messages are preserved
    # This is more reliable than using trim_messages with include_system
    system_msgs = [m for m in messages if isinstance(m, SystemMessage)]
    other_msgs = [m for m in messages if not isinstance(m, SystemMessage)]

    # Calculate tokens for system messages (these will always be preserved)
    system_tokens = count_messages_tokens(system_msgs)

    # Calculate available tokens for non-system messages (after tools and system)
    available_for_other = available_for_messages - system_tokens

    if available_for_other <= 0:
        # Not enough space for system messages - keep only system messages
        logger.warning(
            "System messages (%d tokens) exceed available space (%d tokens), truncating to system only",
            system_tokens, available_for_messages
        )
        return system_msgs[:1] if system_msgs else messages[-1:]

    # Trim non-system messages to fit available space
    trimmed_other = _trim_to_token_limit(other_msgs, available_for_other)

    # Combine system messages with trimmed conversation
    trimmed = system_msgs + trimmed_other

    logger.info(
        "Trimmed %d → %d messages (system: %d, history: %d → %d)",
        len(messages), len(trimmed),
        len(system_msgs), len(other_msgs), len(trimmed_other)
    )

    return trimmed


def _trim_to_token_limit(messages: list[Any], max_tokens: int) -> list[Any]:
    """
    Trim messages to fit within token limit using tiktoken.

    Iteratively removes oldest messages until under token limit.
    Always keeps at least the most recent message.

    Args:
        messages: List of messages to trim
        max_tokens: Maximum tokens allowed

    Returns:
        Trimmed list of messages
    """
    if not messages:
        return messages

    current_tokens = count_messages_tokens(messages)

    if current_tokens <= max_tokens:
        return messages

    # Iteratively remove oldest messages
    trimmed = list(messages)
    while trimmed and count_messages_tokens(trimmed) > max_tokens:
        trimmed.pop(0)

    # Ensure at least one message remains
    if not trimmed and messages:
        trimmed = [messages[-1]]

    return trimmed


def get_token_usage_summary(
    messages: list[Any],
    model_name: str,
    llm_config: dict[str, Any] | None = None,
    tool_tokens: int = 0,
) -> dict[str, Any]:
    """
    Get a summary of token usage for the given messages.

    Args:
        messages: List of LangChain messages
        model_name: Name of the LLM model
        llm_config: Optional LLM config dict from database (may contain context_limit in K tokens)
        tool_tokens: Optional token count for tool definitions

    Returns:
        dict: Token usage summary including:
            - estimated_tokens: Estimated total tokens (messages only)
            - tool_tokens: Token count for tool definitions
            - total_tokens: Sum of messages and tools
            - model_limit_k: Model's context window limit in K tokens
            - model_limit_tokens: Model's context window limit in actual tokens
            - usage_percentage: Percentage of context used (excluding tools)
            - total_usage_percentage: Percentage including tools
            - message_count: Number of messages
            - needs_trimming: Whether messages exceed 80% of limit
    """
    try:
        estimated_tokens = count_messages_tokens(messages)
    except Exception as e:
        logger.warning("Failed to count tokens: %s", e)
        estimated_tokens = 0

    model_limit_tokens = get_model_context_limit(model_name, llm_config)
    model_limit_k = model_limit_tokens // 1000
    usage_percentage = (estimated_tokens / model_limit_tokens * 100) if model_limit_tokens > 0 else 0
    total_tokens = estimated_tokens + tool_tokens
    total_usage_percentage = (total_tokens / model_limit_tokens * 100) if model_limit_tokens > 0 else 0

    return {
        "estimated_tokens": estimated_tokens,
        "tool_tokens": tool_tokens,
        "total_tokens": total_tokens,
        "model_limit_k": model_limit_k,
        "model_limit_tokens": model_limit_tokens,
        "usage_percentage": round(usage_percentage, 2),
        "total_usage_percentage": round(total_usage_percentage, 2),
        "message_count": len(messages),
        "needs_trimming": usage_percentage > 80,
    }


# Convenience function for GNS3-Copilot integration
def prepare_context_messages(
    state_messages: list[Any],
    system_prompt: str,
    topology_context: str | None,
    model_name: str,
    llm_config: dict[str, Any] | None = None,
    tools: list[Any] | None = None,
) -> list[Any]:
    """
    Prepare full context messages for LLM call with automatic trimming.

    This is the main entry point for GNS3-Copilot to prepare messages
    before calling the LLM.

    Args:
        state_messages: Message history from conversation state
        system_prompt: System prompt text
        topology_context: Optional topology information string
        model_name: Name of the LLM model
        llm_config: Optional LLM config dict from database (may contain context_limit and context_strategy)
        tools: Optional list of LangChain tools (for token estimation)

    Returns:
        list: Prepared messages ready for LLM invocation

    Examples:
        >>> messages = prepare_context_messages(
        ...     state_messages=[HumanMessage("Help me")],
        ...     system_prompt="You are a helpful assistant",
        ...     topology_context=None,
        ...     model_name="gpt-4o"
        ... )
        >>> len(messages)
        2  # System message + Human message
    """
    # Get trimming strategy from config (default: "balanced")
    trim_strategy = "balanced"
    if llm_config and "context_strategy" in llm_config:
        strategy = llm_config["context_strategy"]
        if strategy in ["conservative", "balanced", "aggressive"]:
            trim_strategy = strategy
            logger.debug("Using context_strategy from config: %s", trim_strategy)
        else:
            logger.warning("Invalid context_strategy '%s', using 'balanced'", strategy)

    # Estimate tool tokens (tools are sent with each LLM call)
    tool_tokens = estimate_tool_tokens(tools) if tools else 0

    # Inject topology info into system prompt using template variable
    # The system_prompt contains {{topology_info}} placeholder
    if topology_context:
        topology_formatted = f"Current Topology:\n{topology_context}"
        formatted_prompt = system_prompt.replace("{{topology_info}}", topology_formatted)
    else:
        # If no topology, remove the placeholder
        formatted_prompt = system_prompt.replace("{{topology_info}}", "(No topology information available)")

    # Calculate token breakdown
    system_prompt_tokens = count_tokens_accurately(system_prompt)
    topology_tokens = count_tokens_accurately(topology_formatted) if topology_context else 0
    formatted_tokens = count_tokens_accurately(formatted_prompt)

    # Build context messages (single system message with topology injected)
    context_messages = [SystemMessage(content=formatted_prompt)]

    # Calculate conversation history tokens
    history_tokens = count_messages_tokens(state_messages)
    total_context_tokens = system_prompt_tokens + topology_tokens + history_tokens

    # Combine with conversation history
    full_messages = context_messages + state_messages

    # Trim if needed
    trimmed_messages = trim_messages_for_context(
        full_messages,
        model_name=model_name,
        llm_config=llm_config,
        strategy=trim_strategy,
        preserve_system=True,  # Always keep system prompts
        tool_tokens=tool_tokens,  # Account for tool definitions
    )

    # Recalculate after trimming
    trimmed_history_tokens = count_messages_tokens([m for m in trimmed_messages if not isinstance(m, SystemMessage)])
    trimmed_total_tokens = system_prompt_tokens + topology_tokens + trimmed_history_tokens

    # Log summary with detailed breakdown
    model_limit_k = get_model_context_limit(model_name, llm_config) // 1000

    if trimmed_history_tokens < history_tokens:
        # Trimming happened
        logger.info(
            "Context prepared (trimmed): system=%d (base=%d + topology=%d) + history=%d→%d + tools=%d = %d total / %dK limit (%.1f%%), strategy=%s",
            formatted_tokens,
            system_prompt_tokens,
            topology_tokens,
            history_tokens,
            trimmed_history_tokens,
            tool_tokens,
            formatted_tokens + trimmed_history_tokens + tool_tokens,
            model_limit_k,
            ((formatted_tokens + trimmed_history_tokens) / (model_limit_k * 1000)) * 100,
            trim_strategy
        )
    else:
        # No trimming
        logger.info(
            "Context prepared: system=%d (base=%d + topology=%d) + history=%d + tools=%d = %d total / %dK limit (%.1f%%), strategy=%s",
            formatted_tokens,
            system_prompt_tokens,
            topology_tokens,
            history_tokens,
            tool_tokens,
            formatted_tokens + history_tokens + tool_tokens,
            model_limit_k,
            (formatted_tokens / (model_limit_k * 1000)) * 100,
            trim_strategy
        )

    return trimmed_messages


if __name__ == "__main__":
    # Simple test
    test_messages = [
        HumanMessage(f"Message {i}") for i in range(100)
    ]

    result = trim_messages_for_context(
        test_messages,
        model_name="gpt-4o",
        strategy="balanced"
    )

    print(f"Original: {len(test_messages)} messages")
    print(f"Trimmed: {len(result)} messages")