import os
import json
import logging
from pr.config import (CONTEXT_FILE, GLOBAL_CONTEXT_FILE, CONTEXT_COMPRESSION_THRESHOLD,
RECENT_MESSAGES_TO_KEEP, MAX_TOKENS_LIMIT, CHARS_PER_TOKEN,
EMERGENCY_MESSAGES_TO_KEEP, CONTENT_TRIM_LENGTH, MAX_TOOL_RESULT_LENGTH)
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, 'r') 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, 'r') 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
Reference in New Issue View Git Blame Copy Permalink