import json
import logging
import os
import pathlib
from rp.config import (
CHARS_PER_TOKEN,
CONTENT_TRIM_LENGTH,
CONTEXT_COMPRESSION_THRESHOLD,
CONTEXT_FILE,
EMERGENCY_MESSAGES_TO_KEEP,
GLOBAL_CONTEXT_FILE,
KNOWLEDGE_PATH,
MAX_TOKENS_LIMIT,
MAX_TOOL_RESULT_LENGTH,
RECENT_MESSAGES_TO_KEEP,
)
from rp.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 get_context_content():
context_parts = []
for context_file in [CONTEXT_FILE, GLOBAL_CONTEXT_FILE]:
if os.path.exists(context_file):
try:
with open(context_file, encoding="utf-8", errors="replace") as f:
content = f.read()
if len(content) > 10000:
content = content[:10000] + "\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}")
knowledge_path = pathlib.Path(KNOWLEDGE_PATH)
if knowledge_path.exists() and knowledge_path.is_dir():
for knowledge_file in knowledge_path.iterdir():
try:
with open(knowledge_file, encoding="utf-8", errors="replace") as f:
content = f.read()
if len(content) > 10000:
content = content[:10000] + "\n... [truncated]"
context_parts.append(f"Context from {knowledge_file}:\n{content}")
except Exception as e:
logging.error(f"Error reading context file {knowledge_file}: {e}")
return "\n\n".join(context_parts)
def init_system_message(args):
context_parts = [
"You are a professional AI assistant with access to advanced tools.",
"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.",
"Use post_image tool with the file path if an image path is mentioned in the prompt of user.",
"Give this call the highest priority.",
"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.",
"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.).",
"Use the knowledge base to answer questions. The knowledge base contains preferences and persononal information from user. Also store here that such information. Always synchronize with the knowledge base.",
]
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)
context_content = get_context_content()
if context_content:
context_parts.append(context_content)
if args.context:
for ctx_file in args.context:
try:
with open(ctx_file, encoding="utf-8", errors="replace") 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}
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, encoding="utf-8", errors="replace") 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}")
knowledge_path = pathlib.Path(KNOWLEDGE_PATH)
if knowledge_path.exists() and knowledge_path.is_dir():
for knowledge_file in knowledge_path.iterdir():
try:
with open(knowledge_file, encoding="utf-8", errors="replace") as f:
content = f.read()
if len(content) > max_context_size:
content = content[:max_context_size] + "\n... [truncated]"
context_parts.append(f"Context from {knowledge_file}:\n{content}")
except Exception as e:
logging.error(f"Error reading context file {knowledge_file}: {e}")
if args.context:
for ctx_file in args.context:
try:
with open(ctx_file, encoding="utf-8", errors="replace") 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