rp/pr/core/context.py

import json
import logging
import os

from pr.config import (
    CHARS_PER_TOKEN,
    CONTENT_TRIM_LENGTH,
    CONTEXT_COMPRESSION_THRESHOLD,
    CONTEXT_FILE,
    EMERGENCY_MESSAGES_TO_KEEP,
    GLOBAL_CONTEXT_FILE,
    MAX_TOKENS_LIMIT,
    MAX_TOOL_RESULT_LENGTH,
    RECENT_MESSAGES_TO_KEEP,
)
from pr.ui import Colors


def truncate_tool_result(result, max_length=None):
    if max_length is None:
        max_length = MAX_TOOL_RESULT_LENGTH

    if not isinstance(result, dict):
        return result

    result_copy = result.copy()

    if "output" in result_copy and isinstance(result_copy["output"], str):
        if len(result_copy["output"]) > max_length:
            result_copy["output"] = (
                result_copy["output"][:max_length]
                + f"\n... [truncated {len(result_copy['output']) - max_length} chars]"
            )

    if "content" in result_copy and isinstance(result_copy["content"], str):
        if len(result_copy["content"]) > max_length:
            result_copy["content"] = (
                result_copy["content"][:max_length]
                + f"\n... [truncated {len(result_copy['content']) - max_length} chars]"
            )

    if "data" in result_copy and isinstance(result_copy["data"], str):
        if len(result_copy["data"]) > max_length:
            result_copy["data"] = (
                result_copy["data"][:max_length] + f"\n... [truncated]"
            )

    if "error" in result_copy and isinstance(result_copy["error"], str):
        if len(result_copy["error"]) > max_length // 2:
            result_copy["error"] = (
                result_copy["error"][: max_length // 2] + "... [truncated]"
            )

    return result_copy


def init_system_message(args):
    context_parts = [
        """You are a professional AI assistant with access to advanced tools.

File Operations:
- Use RPEditor tools (open_editor, editor_insert_text, editor_replace_text, editor_search, close_editor) for precise file modifications
- Always close editor files when finished
- Use write_file for complete file rewrites, search_replace for simple text replacements

Process Management:
- run_command executes shell commands with a timeout (default 30s)
- If a command times out, you receive a PID in the response
- Use tail_process(pid) to monitor running processes
- Use kill_process(pid) to terminate processes
- Manage long-running commands effectively using these tools

Shell Commands:
- Be a shell ninja using native OS tools
- Prefer standard Unix utilities over complex scripts
- Use run_command_interactive for commands requiring user input (vim, nano, etc.)"""
    ]
    # context_parts = ["You are a helpful AI assistant with access to advanced tools, including a powerful built-in editor (RPEditor). For file editing tasks, prefer using the editor-related tools like write_file, search_replace, open_editor, editor_insert_text, editor_replace_text, and editor_search, as they provide advanced editing capabilities with undo/redo, search, and precise text manipulation. The editor is integrated seamlessly and should be your primary tool for modifying files."]
    max_context_size = 10000

    if args.include_env:
        env_context = "Environment Variables:\n"
        for key, value in os.environ.items():
            if not key.startswith("_"):
                env_context += f"{key}={value}\n"
        if len(env_context) > max_context_size:
            env_context = env_context[:max_context_size] + "\n... [truncated]"
        context_parts.append(env_context)

    for context_file in [CONTEXT_FILE, GLOBAL_CONTEXT_FILE]:
        if os.path.exists(context_file):
            try:
                with open(context_file) as f:
                    content = f.read()
                    if len(content) > max_context_size:
                        content = content[:max_context_size] + "\n... [truncated]"
                    context_parts.append(f"Context from {context_file}:\n{content}")
            except Exception as e:
                logging.error(f"Error reading context file {context_file}: {e}")

    if args.context:
        for ctx_file in args.context:
            try:
                with open(ctx_file) as f:
                    content = f.read()
                    if len(content) > max_context_size:
                        content = content[:max_context_size] + "\n... [truncated]"
                    context_parts.append(f"Context from {ctx_file}:\n{content}")
            except Exception as e:
                logging.error(f"Error reading context file {ctx_file}: {e}")

    system_message = "\n\n".join(context_parts)
    if len(system_message) > max_context_size * 3:
        system_message = (
            system_message[: max_context_size * 3] + "\n... [system message truncated]"
        )

    return {"role": "system", "content": system_message}


def should_compress_context(messages):
    return len(messages) > CONTEXT_COMPRESSION_THRESHOLD


def compress_context(messages):
    return manage_context_window(messages, verbose=False)


def manage_context_window(messages, verbose):
    if len(messages) <= CONTEXT_COMPRESSION_THRESHOLD:
        return messages

    if verbose:
        print(
            f"{Colors.YELLOW}📄 Managing context window (current: {len(messages)} messages)...{Colors.RESET}"
        )

    system_message = messages[0]
    recent_messages = messages[-RECENT_MESSAGES_TO_KEEP:]
    middle_messages = messages[1:-RECENT_MESSAGES_TO_KEEP]

    if middle_messages:
        summary = summarize_messages(middle_messages)
        summary_message = {
            "role": "system",
            "content": f"[Previous conversation summary: {summary}]",
        }

        new_messages = [system_message, summary_message] + recent_messages

        if verbose:
            print(
                f"{Colors.GREEN}✓ Context compressed to {len(new_messages)} messages{Colors.RESET}"
            )

        return new_messages

    return messages


def summarize_messages(messages):
    summary_parts = []

    for msg in messages:
        role = msg.get("role", "unknown")
        content = msg.get("content", "")

        if role == "tool":
            continue

        if isinstance(content, str) and len(content) > 200:
            content = content[:200] + "..."

        summary_parts.append(f"{role}: {content}")

    return " | ".join(summary_parts[:10])


def estimate_tokens(messages):
    total_chars = 0

    for msg in messages:
        msg_json = json.dumps(msg)
        total_chars += len(msg_json)

    estimated_tokens = total_chars / CHARS_PER_TOKEN

    overhead_multiplier = 1.3

    return int(estimated_tokens * overhead_multiplier)


def trim_message_content(message, max_length):
    trimmed_msg = message.copy()

    if "content" in trimmed_msg:
        content = trimmed_msg["content"]

        if isinstance(content, str) and len(content) > max_length:
            trimmed_msg["content"] = (
                content[:max_length]
                + f"\n... [trimmed {len(content) - max_length} chars]"
            )
        elif isinstance(content, list):
            trimmed_content = []
            for item in content:
                if isinstance(item, dict):
                    trimmed_item = item.copy()
                    if (
                        "text" in trimmed_item
                        and len(trimmed_item["text"]) > max_length
                    ):
                        trimmed_item["text"] = (
                            trimmed_item["text"][:max_length] + f"\n... [trimmed]"
                        )
                    trimmed_content.append(trimmed_item)
                else:
                    trimmed_content.append(item)
            trimmed_msg["content"] = trimmed_content

    if trimmed_msg.get("role") == "tool":
        if "content" in trimmed_msg and isinstance(trimmed_msg["content"], str):
            content = trimmed_msg["content"]
            if len(content) > MAX_TOOL_RESULT_LENGTH:
                trimmed_msg["content"] = (
                    content[:MAX_TOOL_RESULT_LENGTH]
                    + f"\n... [trimmed {len(content) - MAX_TOOL_RESULT_LENGTH} chars]"
                )

            try:
                parsed = json.loads(content)
                if isinstance(parsed, dict):
                    if (
                        "output" in parsed
                        and isinstance(parsed["output"], str)
                        and len(parsed["output"]) > MAX_TOOL_RESULT_LENGTH // 2
                    ):
                        parsed["output"] = (
                            parsed["output"][: MAX_TOOL_RESULT_LENGTH // 2]
                            + f"\n... [truncated]"
                        )
                    if (
                        "content" in parsed
                        and isinstance(parsed["content"], str)
                        and len(parsed["content"]) > MAX_TOOL_RESULT_LENGTH // 2
                    ):
                        parsed["content"] = (
                            parsed["content"][: MAX_TOOL_RESULT_LENGTH // 2]
                            + f"\n... [truncated]"
                        )
                    trimmed_msg["content"] = json.dumps(parsed)
            except:
                pass

    return trimmed_msg


def intelligently_trim_messages(messages, target_tokens, keep_recent=3):
    if estimate_tokens(messages) <= target_tokens:
        return messages

    system_msg = (
        messages[0] if messages and messages[0].get("role") == "system" else None
    )
    start_idx = 1 if system_msg else 0

    recent_messages = (
        messages[-keep_recent:] if len(messages) > keep_recent else messages[start_idx:]
    )
    middle_messages = (
        messages[start_idx:-keep_recent] if len(messages) > keep_recent else []
    )

    trimmed_middle = []
    for msg in middle_messages:
        if msg.get("role") == "tool":
            trimmed_middle.append(
                trim_message_content(msg, MAX_TOOL_RESULT_LENGTH // 2)
            )
        elif msg.get("role") in ["user", "assistant"]:
            trimmed_middle.append(trim_message_content(msg, CONTENT_TRIM_LENGTH))
        else:
            trimmed_middle.append(msg)

    result = ([system_msg] if system_msg else []) + trimmed_middle + recent_messages

    if estimate_tokens(result) <= target_tokens:
        return result

    step_size = len(trimmed_middle) // 4 if len(trimmed_middle) >= 4 else 1
    while len(trimmed_middle) > 0 and estimate_tokens(result) > target_tokens:
        remove_count = min(step_size, len(trimmed_middle))
        trimmed_middle = trimmed_middle[remove_count:]
        result = ([system_msg] if system_msg else []) + trimmed_middle + recent_messages

    if estimate_tokens(result) <= target_tokens:
        return result

    keep_recent -= 1
    if keep_recent > 0:
        return intelligently_trim_messages(messages, target_tokens, keep_recent)

    return ([system_msg] if system_msg else []) + messages[-1:]


def auto_slim_messages(messages, verbose=False):
    estimated_tokens = estimate_tokens(messages)

    if estimated_tokens <= MAX_TOKENS_LIMIT:
        return messages

    if verbose:
        print(
            f"{Colors.YELLOW}⚠️  Token limit approaching: ~{estimated_tokens} tokens (limit: {MAX_TOKENS_LIMIT}){Colors.RESET}"
        )
        print(
            f"{Colors.YELLOW}🔧 Intelligently trimming message content...{Colors.RESET}"
        )

    result = intelligently_trim_messages(
        messages, MAX_TOKENS_LIMIT, keep_recent=EMERGENCY_MESSAGES_TO_KEEP
    )
    final_tokens = estimate_tokens(result)

    if final_tokens > MAX_TOKENS_LIMIT:
        if verbose:
            print(
                f"{Colors.RED}⚠️  Still over limit after trimming, applying emergency reduction...{Colors.RESET}"
            )
        result = emergency_reduce_messages(result, MAX_TOKENS_LIMIT, verbose)
        final_tokens = estimate_tokens(result)

    if verbose:
        removed_count = len(messages) - len(result)
        print(
            f"{Colors.GREEN}✓ Optimized from {len(messages)} to {len(result)} messages{Colors.RESET}"
        )
        print(
            f"{Colors.GREEN}  Token estimate: {estimated_tokens} → {final_tokens} (~{estimated_tokens - final_tokens} saved){Colors.RESET}"
        )
        if removed_count > 0:
            print(
                f"{Colors.GREEN}  Removed {removed_count} older messages{Colors.RESET}"
            )

    return result


def emergency_reduce_messages(messages, target_tokens, verbose=False):
    system_msg = (
        messages[0] if messages and messages[0].get("role") == "system" else None
    )
    start_idx = 1 if system_msg else 0

    keep_count = 2
    while estimate_tokens(messages) > target_tokens and keep_count >= 1:
        if len(messages[start_idx:]) <= keep_count:
            break

        result = ([system_msg] if system_msg else []) + messages[-keep_count:]

        for i in range(len(result)):
            result[i] = trim_message_content(result[i], CONTENT_TRIM_LENGTH // 2)

        if estimate_tokens(result) <= target_tokens:
            return result

        keep_count -= 1

    final = ([system_msg] if system_msg else []) + messages[-1:]

    for i in range(len(final)):
        if final[i].get("role") != "system":
            final[i] = trim_message_content(final[i], 100)

    return final