chore: standardize string quotes and fix import ordering across multiple modules

This commit is contained in:
2025-11-04 07:09:12 +00:00
parent ea29bdc403
commit e9ced4a493
82 changed files with 4963 additions and 3094 deletions
+8 -2
View File
@@ -1,5 +1,11 @@
from pr.core.assistant import Assistant
from pr.core.api import call_api, list_models
from pr.core.assistant import Assistant
from pr.core.context import init_system_message, manage_context_window
__all__ = ['Assistant', 'call_api', 'list_models', 'init_system_message', 'manage_context_window']
__all__ = [
"Assistant",
"call_api",
"list_models",
"init_system_message",
"manage_context_window",
]
+30 -28
View File
@@ -1,20 +1,20 @@
import re
import math
from typing import List, Dict, Any
from collections import Counter
from typing import Any, Dict, List
class AdvancedContextManager:
def __init__(self, knowledge_store=None, conversation_memory=None):
self.knowledge_store = knowledge_store
self.conversation_memory = conversation_memory
def adaptive_context_window(self, messages: List[Dict[str, Any]],
task_complexity: str = 'medium') -> int:
def adaptive_context_window(
self, messages: List[Dict[str, Any]], task_complexity: str = "medium"
) -> int:
complexity_thresholds = {
'simple': 10,
'medium': 20,
'complex': 35,
'very_complex': 50
"simple": 10,
"medium": 20,
"complex": 35,
"very_complex": 50,
}
base_threshold = complexity_thresholds.get(task_complexity, 20)
@@ -31,17 +31,19 @@ class AdvancedContextManager:
return max(base_threshold, adjusted)
def _analyze_message_complexity(self, messages: List[Dict[str, Any]]) -> float:
total_length = sum(len(msg.get('content', '')) for msg in messages)
total_length = sum(len(msg.get("content", "")) for msg in messages)
avg_length = total_length / len(messages) if messages else 0
unique_words = set()
for msg in messages:
content = msg.get('content', '')
words = re.findall(r'\b\w+\b', content.lower())
content = msg.get("content", "")
words = re.findall(r"\b\w+\b", content.lower())
unique_words.update(words)
vocabulary_richness = len(unique_words) / total_length if total_length > 0 else 0
vocabulary_richness = (
len(unique_words) / total_length if total_length > 0 else 0
)
# Simple complexity score based on length and richness
complexity = min(1.0, (avg_length / 100) + vocabulary_richness)
return complexity
@@ -49,10 +51,10 @@ class AdvancedContextManager:
def extract_key_sentences(self, text: str, top_k: int = 5) -> List[str]:
if not text.strip():
return []
sentences = re.split(r'(?<=[.!?])\s+', text)
sentences = re.split(r"(?<=[.!?])\s+", text)
if not sentences:
return []
# Simple scoring based on length and position
scored_sentences = []
for i, sentence in enumerate(sentences):
@@ -60,25 +62,25 @@ class AdvancedContextManager:
position_score = 1.0 if i == 0 else 0.8 if i < len(sentences) / 2 else 0.6
score = (length_score + position_score) / 2
scored_sentences.append((sentence, score))
scored_sentences.sort(key=lambda x: x[1], reverse=True)
return [s[0] for s in scored_sentences[:top_k]]
def advanced_summarize_messages(self, messages: List[Dict[str, Any]]) -> str:
all_content = ' '.join([msg.get('content', '') for msg in messages])
all_content = " ".join([msg.get("content", "") for msg in messages])
key_sentences = self.extract_key_sentences(all_content, top_k=3)
summary = ' '.join(key_sentences)
summary = " ".join(key_sentences)
return summary if summary else "No content to summarize."
def score_message_relevance(self, message: Dict[str, Any], context: str) -> float:
content = message.get('content', '')
content_words = set(re.findall(r'\b\w+\b', content.lower()))
context_words = set(re.findall(r'\b\w+\b', context.lower()))
content = message.get("content", "")
content_words = set(re.findall(r"\b\w+\b", content.lower()))
context_words = set(re.findall(r"\b\w+\b", context.lower()))
intersection = content_words & context_words
union = content_words | context_words
if not union:
return 0.0
return len(intersection) / len(union)
return len(intersection) / len(union)
+39 -33
View File
@@ -1,13 +1,17 @@
import json
import urllib.request
import urllib.error
import logging
from pr.config import DEFAULT_TEMPERATURE, DEFAULT_MAX_TOKENS
import urllib.error
import urllib.request
from pr.config import DEFAULT_MAX_TOKENS, DEFAULT_TEMPERATURE
from pr.core.context import auto_slim_messages
logger = logging.getLogger('pr')
logger = logging.getLogger("pr")
def call_api(messages, model, api_url, api_key, use_tools, tools_definition, verbose=False):
def call_api(
messages, model, api_url, api_key, use_tools, tools_definition, verbose=False
):
try:
messages = auto_slim_messages(messages, verbose=verbose)
@@ -17,62 +21,63 @@ def call_api(messages, model, api_url, api_key, use_tools, tools_definition, ver
logger.debug(f"Use tools: {use_tools}")
logger.debug(f"Message count: {len(messages)}")
headers = {
'Content-Type': 'application/json',
"Content-Type": "application/json",
}
if api_key:
headers['Authorization'] = f'Bearer {api_key}'
headers["Authorization"] = f"Bearer {api_key}"
data = {
'model': model,
'messages': messages,
'temperature': DEFAULT_TEMPERATURE,
'max_tokens': DEFAULT_MAX_TOKENS
"model": model,
"messages": messages,
"temperature": DEFAULT_TEMPERATURE,
"max_tokens": DEFAULT_MAX_TOKENS,
}
if "gpt-5" in model:
del data['temperature']
del data['max_tokens']
del data["temperature"]
del data["max_tokens"]
logger.debug("GPT-5 detected: removed temperature and max_tokens")
if use_tools:
data['tools'] = tools_definition
data['tool_choice'] = 'auto'
data["tools"] = tools_definition
data["tool_choice"] = "auto"
logger.debug(f"Tool calling enabled with {len(tools_definition)} tools")
request_json = json.dumps(data)
logger.debug(f"Request payload size: {len(request_json)} bytes")
req = urllib.request.Request(
api_url,
data=request_json.encode('utf-8'),
headers=headers,
method='POST'
api_url, data=request_json.encode("utf-8"), headers=headers, method="POST"
)
logger.debug("Sending HTTP request...")
with urllib.request.urlopen(req) as response:
response_data = response.read().decode('utf-8')
response_data = response.read().decode("utf-8")
logger.debug(f"Response received: {len(response_data)} bytes")
result = json.loads(response_data)
if 'usage' in result:
if "usage" in result:
logger.debug(f"Token usage: {result['usage']}")
if 'choices' in result and result['choices']:
choice = result['choices'][0]
if 'message' in choice:
msg = choice['message']
if "choices" in result and result["choices"]:
choice = result["choices"][0]
if "message" in choice:
msg = choice["message"]
logger.debug(f"Response role: {msg.get('role', 'N/A')}")
if 'content' in msg and msg['content']:
logger.debug(f"Response content length: {len(msg['content'])} chars")
if 'tool_calls' in msg:
logger.debug(f"Response contains {len(msg['tool_calls'])} tool call(s)")
if "content" in msg and msg["content"]:
logger.debug(
f"Response content length: {len(msg['content'])} chars"
)
if "tool_calls" in msg:
logger.debug(
f"Response contains {len(msg['tool_calls'])} tool call(s)"
)
logger.debug("=== API CALL END ===")
return result
except urllib.error.HTTPError as e:
error_body = e.read().decode('utf-8')
error_body = e.read().decode("utf-8")
logger.error(f"API HTTP Error: {e.code} - {error_body}")
logger.debug("=== API CALL FAILED ===")
return {"error": f"API Error: {e.code}", "message": error_body}
@@ -81,15 +86,16 @@ def call_api(messages, model, api_url, api_key, use_tools, tools_definition, ver
logger.debug("=== API CALL FAILED ===")
return {"error": str(e)}
def list_models(model_list_url, api_key):
try:
req = urllib.request.Request(model_list_url)
if api_key:
req.add_header('Authorization', f'Bearer {api_key}')
req.add_header("Authorization", f"Bearer {api_key}")
with urllib.request.urlopen(req) as response:
data = json.loads(response.read().decode('utf-8'))
data = json.loads(response.read().decode("utf-8"))
return data.get('data', [])
return data.get("data", [])
except Exception as e:
return {"error": str(e)}
+276 -147
View File
@@ -1,63 +1,111 @@
import os
import sys
import json
import sqlite3
import signal
import logging
import traceback
import readline
import glob as glob_module
import json
import logging
import os
import readline
import signal
import sqlite3
import sys
import traceback
from concurrent.futures import ThreadPoolExecutor
from pr.config import DB_PATH, LOG_FILE, DEFAULT_MODEL, DEFAULT_API_URL, MODEL_LIST_URL, HISTORY_FILE
from pr.ui import Colors, render_markdown
from pr.core.context import init_system_message, truncate_tool_result
from pr.commands import handle_command
from pr.config import (
DB_PATH,
DEFAULT_API_URL,
DEFAULT_MODEL,
HISTORY_FILE,
LOG_FILE,
MODEL_LIST_URL,
)
from pr.core.api import call_api
from pr.core.autonomous_interactions import (
get_global_autonomous,
stop_global_autonomous,
)
from pr.core.background_monitor import (
get_global_monitor,
start_global_monitor,
stop_global_monitor,
)
from pr.core.context import init_system_message, truncate_tool_result
from pr.tools import (
http_fetch, run_command, run_command_interactive, read_file, write_file,
list_directory, mkdir, chdir, getpwd, db_set, db_get, db_query,
web_search, web_search_news, python_exec, index_source_directory,
open_editor, editor_insert_text, editor_replace_text, editor_search,
search_replace,close_editor,create_diff,apply_patch,
tail_process, kill_process
apply_patch,
chdir,
close_editor,
create_diff,
db_get,
db_query,
db_set,
editor_insert_text,
editor_replace_text,
editor_search,
getpwd,
http_fetch,
index_source_directory,
kill_process,
list_directory,
mkdir,
open_editor,
python_exec,
read_file,
run_command,
search_replace,
tail_process,
web_search,
web_search_news,
write_file,
)
from pr.tools.base import get_tools_definition
from pr.tools.filesystem import (
clear_edit_tracker,
display_edit_summary,
display_edit_timeline,
)
from pr.tools.interactive_control import (
start_interactive_session, send_input_to_session, read_session_output,
list_active_sessions, close_interactive_session
close_interactive_session,
list_active_sessions,
read_session_output,
send_input_to_session,
start_interactive_session,
)
from pr.tools.patch import display_file_diff
from pr.tools.filesystem import display_edit_summary, display_edit_timeline, clear_edit_tracker
from pr.tools.base import get_tools_definition
from pr.commands import handle_command
from pr.core.background_monitor import start_global_monitor, stop_global_monitor, get_global_monitor
from pr.core.autonomous_interactions import start_global_autonomous, stop_global_autonomous, get_global_autonomous
from pr.ui import Colors, render_markdown
logger = logging.getLogger('pr')
logger = logging.getLogger("pr")
logger.setLevel(logging.DEBUG)
file_handler = logging.FileHandler(LOG_FILE)
file_handler.setFormatter(logging.Formatter('%(asctime)s - %(levelname)s - %(message)s'))
file_handler.setFormatter(
logging.Formatter("%(asctime)s - %(levelname)s - %(message)s")
)
logger.addHandler(file_handler)
class Assistant:
def __init__(self, args):
self.args = args
self.messages = []
self.verbose = args.verbose
self.debug = getattr(args, 'debug', False)
self.debug = getattr(args, "debug", False)
self.syntax_highlighting = not args.no_syntax
if self.debug:
console_handler = logging.StreamHandler()
console_handler.setLevel(logging.DEBUG)
console_handler.setFormatter(logging.Formatter('%(levelname)s: %(message)s'))
console_handler.setFormatter(
logging.Formatter("%(levelname)s: %(message)s")
)
logger.addHandler(console_handler)
logger.debug("Debug mode enabled")
self.api_key = os.environ.get('OPENROUTER_API_KEY', '')
self.model = args.model or os.environ.get('AI_MODEL', DEFAULT_MODEL)
self.api_url = args.api_url or os.environ.get('API_URL', DEFAULT_API_URL)
self.model_list_url = args.model_list_url or os.environ.get('MODEL_LIST_URL', MODEL_LIST_URL)
self.use_tools = os.environ.get('USE_TOOLS', '1') == '1'
self.strict_mode = os.environ.get('STRICT_MODE', '0') == '1'
self.api_key = os.environ.get("OPENROUTER_API_KEY", "")
self.model = args.model or os.environ.get("AI_MODEL", DEFAULT_MODEL)
self.api_url = args.api_url or os.environ.get("API_URL", DEFAULT_API_URL)
self.model_list_url = args.model_list_url or os.environ.get(
"MODEL_LIST_URL", MODEL_LIST_URL
)
self.use_tools = os.environ.get("USE_TOOLS", "1") == "1"
self.strict_mode = os.environ.get("STRICT_MODE", "0") == "1"
self.interrupt_count = 0
self.python_globals = {}
self.db_conn = None
@@ -69,6 +117,7 @@ class Assistant:
try:
from pr.core.enhanced_assistant import EnhancedAssistant
self.enhanced = EnhancedAssistant(self)
if self.debug:
logger.debug("Enhanced assistant features initialized")
@@ -94,13 +143,17 @@ class Assistant:
self.db_conn = sqlite3.connect(DB_PATH, check_same_thread=False)
cursor = self.db_conn.cursor()
cursor.execute('''CREATE TABLE IF NOT EXISTS kv_store
(key TEXT PRIMARY KEY, value TEXT, timestamp REAL)''')
cursor.execute(
"""CREATE TABLE IF NOT EXISTS kv_store
(key TEXT PRIMARY KEY, value TEXT, timestamp REAL)"""
)
cursor.execute('''CREATE TABLE IF NOT EXISTS file_versions
cursor.execute(
"""CREATE TABLE IF NOT EXISTS file_versions
(id INTEGER PRIMARY KEY AUTOINCREMENT,
filepath TEXT, content TEXT, hash TEXT,
timestamp REAL, version INTEGER)''')
timestamp REAL, version INTEGER)"""
)
self.db_conn.commit()
logger.debug("Database initialized successfully")
@@ -110,7 +163,7 @@ class Assistant:
def _handle_background_updates(self, updates):
"""Handle background session updates by injecting them into the conversation."""
if not updates or not updates.get('sessions'):
if not updates or not updates.get("sessions"):
return
# Format the update as a system message
@@ -118,10 +171,12 @@ class Assistant:
# Inject into current conversation if we're in an active session
if self.messages and len(self.messages) > 0:
self.messages.append({
"role": "system",
"content": f"Background session updates: {update_message}"
})
self.messages.append(
{
"role": "system",
"content": f"Background session updates: {update_message}",
}
)
if self.verbose:
print(f"{Colors.CYAN}Background update: {update_message}{Colors.RESET}")
@@ -130,8 +185,8 @@ class Assistant:
"""Format background updates for LLM consumption."""
session_summaries = []
for session_name, session_info in updates.get('sessions', {}).items():
summary = session_info.get('summary', f'Session {session_name}')
for session_name, session_info in updates.get("sessions", {}).items():
summary = session_info.get("summary", f"Session {session_name}")
session_summaries.append(f"{session_name}: {summary}")
if session_summaries:
@@ -151,30 +206,44 @@ class Assistant:
if events:
print(f"\n{Colors.CYAN}Background Events:{Colors.RESET}")
for event in events:
event_type = event.get('type', 'unknown')
session_name = event.get('session_name', 'unknown')
event_type = event.get("type", "unknown")
session_name = event.get("session_name", "unknown")
if event_type == 'session_started':
print(f" {Colors.GREEN}{Colors.RESET} Session '{session_name}' started")
elif event_type == 'session_ended':
print(f" {Colors.YELLOW}{Colors.RESET} Session '{session_name}' ended")
elif event_type == 'output_received':
lines = len(event.get('new_output', {}).get('stdout', []))
print(f" {Colors.BLUE}📝{Colors.RESET} Session '{session_name}' produced {lines} lines of output")
elif event_type == 'possible_input_needed':
print(f" {Colors.RED}{Colors.RESET} Session '{session_name}' may need input")
elif event_type == 'high_output_volume':
total = event.get('total_lines', 0)
print(f" {Colors.YELLOW}📊{Colors.RESET} Session '{session_name}' has high output volume ({total} lines)")
elif event_type == 'inactive_session':
inactive_time = event.get('inactive_seconds', 0)
print(f" {Colors.GRAY}{Colors.RESET} Session '{session_name}' inactive for {inactive_time:.0f}s")
if event_type == "session_started":
print(
f" {Colors.GREEN}{Colors.RESET} Session '{session_name}' started"
)
elif event_type == "session_ended":
print(
f" {Colors.YELLOW}{Colors.RESET} Session '{session_name}' ended"
)
elif event_type == "output_received":
lines = len(event.get("new_output", {}).get("stdout", []))
print(
f" {Colors.BLUE}📝{Colors.RESET} Session '{session_name}' produced {lines} lines of output"
)
elif event_type == "possible_input_needed":
print(
f" {Colors.RED}{Colors.RESET} Session '{session_name}' may need input"
)
elif event_type == "high_output_volume":
total = event.get("total_lines", 0)
print(
f" {Colors.YELLOW}📊{Colors.RESET} Session '{session_name}' has high output volume ({total} lines)"
)
elif event_type == "inactive_session":
inactive_time = event.get("inactive_seconds", 0)
print(
f" {Colors.GRAY}{Colors.RESET} Session '{session_name}' inactive for {inactive_time:.0f}s"
)
print() # Add blank line after events
except Exception as e:
if self.debug:
print(f"{Colors.RED}Error checking background updates: {e}{Colors.RESET}")
print(
f"{Colors.RED}Error checking background updates: {e}{Colors.RESET}"
)
def execute_tool_calls(self, tool_calls):
results = []
@@ -185,114 +254,147 @@ class Assistant:
futures = []
for tool_call in tool_calls:
func_name = tool_call['function']['name']
arguments = json.loads(tool_call['function']['arguments'])
func_name = tool_call["function"]["name"]
arguments = json.loads(tool_call["function"]["arguments"])
logger.debug(f"Tool call: {func_name} with arguments: {arguments}")
func_map = {
'http_fetch': lambda **kw: http_fetch(**kw),
'run_command': lambda **kw: run_command(**kw),
'tail_process': lambda **kw: tail_process(**kw),
'kill_process': lambda **kw: kill_process(**kw),
'start_interactive_session': lambda **kw: start_interactive_session(**kw),
'send_input_to_session': lambda **kw: send_input_to_session(**kw),
'read_session_output': lambda **kw: read_session_output(**kw),
'close_interactive_session': lambda **kw: close_interactive_session(**kw),
'read_file': lambda **kw: read_file(**kw, db_conn=self.db_conn),
'write_file': lambda **kw: write_file(**kw, db_conn=self.db_conn),
'list_directory': lambda **kw: list_directory(**kw),
'mkdir': lambda **kw: mkdir(**kw),
'chdir': lambda **kw: chdir(**kw),
'getpwd': lambda **kw: getpwd(**kw),
'db_set': lambda **kw: db_set(**kw, db_conn=self.db_conn),
'db_get': lambda **kw: db_get(**kw, db_conn=self.db_conn),
'db_query': lambda **kw: db_query(**kw, db_conn=self.db_conn),
'web_search': lambda **kw: web_search(**kw),
'web_search_news': lambda **kw: web_search_news(**kw),
'python_exec': lambda **kw: python_exec(**kw, python_globals=self.python_globals),
'index_source_directory': lambda **kw: index_source_directory(**kw),
'search_replace': lambda **kw: search_replace(**kw, db_conn=self.db_conn),
'open_editor': lambda **kw: open_editor(**kw),
'editor_insert_text': lambda **kw: editor_insert_text(**kw, db_conn=self.db_conn),
'editor_replace_text': lambda **kw: editor_replace_text(**kw, db_conn=self.db_conn),
'editor_search': lambda **kw: editor_search(**kw),
'close_editor': lambda **kw: close_editor(**kw),
'create_diff': lambda **kw: create_diff(**kw),
'apply_patch': lambda **kw: apply_patch(**kw, db_conn=self.db_conn),
'display_file_diff': lambda **kw: display_file_diff(**kw),
'display_edit_summary': lambda **kw: display_edit_summary(),
'display_edit_timeline': lambda **kw: display_edit_timeline(**kw),
'clear_edit_tracker': lambda **kw: clear_edit_tracker(),
'start_interactive_session': lambda **kw: start_interactive_session(**kw),
'send_input_to_session': lambda **kw: send_input_to_session(**kw),
'read_session_output': lambda **kw: read_session_output(**kw),
'list_active_sessions': lambda **kw: list_active_sessions(**kw),
'close_interactive_session': lambda **kw: close_interactive_session(**kw),
'create_agent': lambda **kw: create_agent(**kw),
'list_agents': lambda **kw: list_agents(**kw),
'execute_agent_task': lambda **kw: execute_agent_task(**kw),
'remove_agent': lambda **kw: remove_agent(**kw),
'collaborate_agents': lambda **kw: collaborate_agents(**kw),
'add_knowledge_entry': lambda **kw: add_knowledge_entry(**kw),
'get_knowledge_entry': lambda **kw: get_knowledge_entry(**kw),
'search_knowledge': lambda **kw: search_knowledge(**kw),
'get_knowledge_by_category': lambda **kw: get_knowledge_by_category(**kw),
'update_knowledge_importance': lambda **kw: update_knowledge_importance(**kw),
'delete_knowledge_entry': lambda **kw: delete_knowledge_entry(**kw),
'get_knowledge_statistics': lambda **kw: get_knowledge_statistics(**kw),
"http_fetch": lambda **kw: http_fetch(**kw),
"run_command": lambda **kw: run_command(**kw),
"tail_process": lambda **kw: tail_process(**kw),
"kill_process": lambda **kw: kill_process(**kw),
"start_interactive_session": lambda **kw: start_interactive_session(
**kw
),
"send_input_to_session": lambda **kw: send_input_to_session(**kw),
"read_session_output": lambda **kw: read_session_output(**kw),
"close_interactive_session": lambda **kw: close_interactive_session(
**kw
),
"read_file": lambda **kw: read_file(**kw, db_conn=self.db_conn),
"write_file": lambda **kw: write_file(**kw, db_conn=self.db_conn),
"list_directory": lambda **kw: list_directory(**kw),
"mkdir": lambda **kw: mkdir(**kw),
"chdir": lambda **kw: chdir(**kw),
"getpwd": lambda **kw: getpwd(**kw),
"db_set": lambda **kw: db_set(**kw, db_conn=self.db_conn),
"db_get": lambda **kw: db_get(**kw, db_conn=self.db_conn),
"db_query": lambda **kw: db_query(**kw, db_conn=self.db_conn),
"web_search": lambda **kw: web_search(**kw),
"web_search_news": lambda **kw: web_search_news(**kw),
"python_exec": lambda **kw: python_exec(
**kw, python_globals=self.python_globals
),
"index_source_directory": lambda **kw: index_source_directory(**kw),
"search_replace": lambda **kw: search_replace(
**kw, db_conn=self.db_conn
),
"open_editor": lambda **kw: open_editor(**kw),
"editor_insert_text": lambda **kw: editor_insert_text(
**kw, db_conn=self.db_conn
),
"editor_replace_text": lambda **kw: editor_replace_text(
**kw, db_conn=self.db_conn
),
"editor_search": lambda **kw: editor_search(**kw),
"close_editor": lambda **kw: close_editor(**kw),
"create_diff": lambda **kw: create_diff(**kw),
"apply_patch": lambda **kw: apply_patch(**kw, db_conn=self.db_conn),
"display_file_diff": lambda **kw: display_file_diff(**kw),
"display_edit_summary": lambda **kw: display_edit_summary(),
"display_edit_timeline": lambda **kw: display_edit_timeline(**kw),
"clear_edit_tracker": lambda **kw: clear_edit_tracker(),
"start_interactive_session": lambda **kw: start_interactive_session(
**kw
),
"send_input_to_session": lambda **kw: send_input_to_session(**kw),
"read_session_output": lambda **kw: read_session_output(**kw),
"list_active_sessions": lambda **kw: list_active_sessions(**kw),
"close_interactive_session": lambda **kw: close_interactive_session(
**kw
),
"create_agent": lambda **kw: create_agent(**kw),
"list_agents": lambda **kw: list_agents(**kw),
"execute_agent_task": lambda **kw: execute_agent_task(**kw),
"remove_agent": lambda **kw: remove_agent(**kw),
"collaborate_agents": lambda **kw: collaborate_agents(**kw),
"add_knowledge_entry": lambda **kw: add_knowledge_entry(**kw),
"get_knowledge_entry": lambda **kw: get_knowledge_entry(**kw),
"search_knowledge": lambda **kw: search_knowledge(**kw),
"get_knowledge_by_category": lambda **kw: get_knowledge_by_category(
**kw
),
"update_knowledge_importance": lambda **kw: update_knowledge_importance(
**kw
),
"delete_knowledge_entry": lambda **kw: delete_knowledge_entry(**kw),
"get_knowledge_statistics": lambda **kw: get_knowledge_statistics(
**kw
),
}
if func_name in func_map:
future = executor.submit(func_map[func_name], **arguments)
futures.append((tool_call['id'], future))
futures.append((tool_call["id"], future))
for tool_id, future in futures:
try:
result = future.result(timeout=30)
result = truncate_tool_result(result)
logger.debug(f"Tool result for {tool_id}: {str(result)[:200]}...")
results.append({
"tool_call_id": tool_id,
"role": "tool",
"content": json.dumps(result)
})
results.append(
{
"tool_call_id": tool_id,
"role": "tool",
"content": json.dumps(result),
}
)
except Exception as e:
logger.debug(f"Tool error for {tool_id}: {str(e)}")
error_msg = str(e)[:200] if len(str(e)) > 200 else str(e)
results.append({
"tool_call_id": tool_id,
"role": "tool",
"content": json.dumps({"status": "error", "error": error_msg})
})
results.append(
{
"tool_call_id": tool_id,
"role": "tool",
"content": json.dumps(
{"status": "error", "error": error_msg}
),
}
)
return results
def process_response(self, response):
if 'error' in response:
if "error" in response:
return f"Error: {response['error']}"
if 'choices' not in response or not response['choices']:
if "choices" not in response or not response["choices"]:
return "No response from API"
message = response['choices'][0]['message']
message = response["choices"][0]["message"]
self.messages.append(message)
if 'tool_calls' in message and message['tool_calls']:
if "tool_calls" in message and message["tool_calls"]:
if self.verbose:
print(f"{Colors.YELLOW}Executing tool calls...{Colors.RESET}")
tool_results = self.execute_tool_calls(message['tool_calls'])
tool_results = self.execute_tool_calls(message["tool_calls"])
for result in tool_results:
self.messages.append(result)
follow_up = call_api(
self.messages, self.model, self.api_url, self.api_key,
self.use_tools, get_tools_definition(), verbose=self.verbose
self.messages,
self.model,
self.api_url,
self.api_key,
self.use_tools,
get_tools_definition(),
verbose=self.verbose,
)
return self.process_response(follow_up)
content = message.get('content', '')
content = message.get("content", "")
return render_markdown(content, self.syntax_highlighting)
def signal_handler(self, signum, frame):
@@ -303,7 +405,9 @@ class Assistant:
self.autonomous_mode = False
sys.exit(0)
else:
print(f"\n{Colors.YELLOW}Press Ctrl+C again to force exit{Colors.RESET}")
print(
f"\n{Colors.YELLOW}Press Ctrl+C again to force exit{Colors.RESET}"
)
return
self.interrupt_count += 1
@@ -323,21 +427,34 @@ class Assistant:
readline.set_history_length(1000)
import atexit
atexit.register(readline.write_history_file, HISTORY_FILE)
commands = ['exit', 'quit', 'help', 'reset', 'dump', 'verbose',
'models', 'tools', 'review', 'refactor', 'obfuscate', '/auto']
commands = [
"exit",
"quit",
"help",
"reset",
"dump",
"verbose",
"models",
"tools",
"review",
"refactor",
"obfuscate",
"/auto",
]
def completer(text, state):
options = [cmd for cmd in commands if cmd.startswith(text)]
glob_pattern = os.path.expanduser(text) + '*'
glob_pattern = os.path.expanduser(text) + "*"
path_options = glob_module.glob(glob_pattern)
path_options = [p + os.sep if os.path.isdir(p) else p for p in path_options]
combined_options = sorted(list(set(options + path_options)))
#combined_options.extend(self.commands)
# combined_options.extend(self.commands)
if state < len(combined_options):
return combined_options[state]
@@ -345,10 +462,10 @@ class Assistant:
return None
delims = readline.get_completer_delims()
readline.set_completer_delims(delims.replace('/', ''))
readline.set_completer_delims(delims.replace("/", ""))
readline.set_completer(completer)
readline.parse_and_bind('tab: complete')
readline.parse_and_bind("tab: complete")
def run_repl(self):
self.setup_readline()
@@ -368,8 +485,11 @@ class Assistant:
if self.background_monitoring:
try:
from pr.multiplexer import get_all_sessions
sessions = get_all_sessions()
active_count = sum(1 for s in sessions.values() if s.get('status') == 'running')
active_count = sum(
1 for s in sessions.values() if s.get("status") == "running"
)
if active_count > 0:
prompt += f"[{active_count}bg]"
except:
@@ -405,10 +525,11 @@ class Assistant:
message = sys.stdin.read()
from pr.autonomous.mode import run_autonomous_mode
run_autonomous_mode(self, message)
def cleanup(self):
if hasattr(self, 'enhanced') and self.enhanced:
if hasattr(self, "enhanced") and self.enhanced:
try:
self.enhanced.cleanup()
except Exception as e:
@@ -424,6 +545,7 @@ class Assistant:
try:
from pr.multiplexer import cleanup_all_multiplexers
cleanup_all_multiplexers()
except Exception as e:
logger.error(f"Error cleaning up multiplexers: {e}")
@@ -433,7 +555,9 @@ class Assistant:
def run(self):
try:
print(f"DEBUG: interactive={self.args.interactive}, message={self.args.message}, isatty={sys.stdin.isatty()}")
print(
f"DEBUG: interactive={self.args.interactive}, message={self.args.message}, isatty={sys.stdin.isatty()}"
)
if self.args.interactive or (not self.args.message and sys.stdin.isatty()):
print("DEBUG: calling run_repl")
self.run_repl()
@@ -443,6 +567,7 @@ class Assistant:
finally:
self.cleanup()
def process_message(assistant, message):
assistant.messages.append({"role": "user", "content": message})
@@ -453,9 +578,13 @@ def process_message(assistant, message):
print(f"{Colors.GRAY}Sending request to API...{Colors.RESET}")
response = call_api(
assistant.messages, assistant.model, assistant.api_url,
assistant.api_key, assistant.use_tools, get_tools_definition(),
verbose=assistant.verbose
assistant.messages,
assistant.model,
assistant.api_url,
assistant.api_key,
assistant.use_tools,
get_tools_definition(),
verbose=assistant.verbose,
)
result = assistant.process_response(response)
+49 -30
View File
@@ -1,7 +1,12 @@
import time
import threading
from pr.core.background_monitor import get_global_monitor
from pr.tools.interactive_control import list_active_sessions, get_session_status, read_session_output
import time
from pr.tools.interactive_control import (
get_session_status,
list_active_sessions,
read_session_output,
)
class AutonomousInteractions:
def __init__(self, interaction_interval=10.0):
@@ -16,7 +21,9 @@ class AutonomousInteractions:
self.llm_callback = llm_callback
if self.interaction_thread is None:
self.active = True
self.interaction_thread = threading.Thread(target=self._interaction_loop, daemon=True)
self.interaction_thread = threading.Thread(
target=self._interaction_loop, daemon=True
)
self.interaction_thread.start()
def stop(self):
@@ -48,7 +55,9 @@ class AutonomousInteractions:
if not sessions:
return # No active sessions
sessions_needing_attention = self._identify_sessions_needing_attention(sessions)
sessions_needing_attention = self._identify_sessions_needing_attention(
sessions
)
if sessions_needing_attention and self.llm_callback:
# Format session updates for LLM
@@ -63,26 +72,30 @@ class AutonomousInteractions:
needing_attention = []
for session_name, session_data in sessions.items():
metadata = session_data['metadata']
output_summary = session_data['output_summary']
metadata = session_data["metadata"]
output_summary = session_data["output_summary"]
# Criteria for needing attention:
# 1. Recent output activity
time_since_activity = time.time() - metadata.get('last_activity', 0)
time_since_activity = time.time() - metadata.get("last_activity", 0)
if time_since_activity < 30: # Activity in last 30 seconds
needing_attention.append(session_name)
continue
# 2. High output volume (potential completion or error)
total_lines = output_summary['stdout_lines'] + output_summary['stderr_lines']
total_lines = (
output_summary["stdout_lines"] + output_summary["stderr_lines"]
)
if total_lines > 50: # Arbitrary threshold
needing_attention.append(session_name)
continue
# 3. Long-running sessions that might need intervention
session_age = time.time() - metadata.get('start_time', 0)
if session_age > 300 and time_since_activity > 60: # 5+ minutes old, inactive for 1+ minute
session_age = time.time() - metadata.get("start_time", 0)
if (
session_age > 300 and time_since_activity > 60
): # 5+ minutes old, inactive for 1+ minute
needing_attention.append(session_name)
continue
@@ -95,18 +108,18 @@ class AutonomousInteractions:
def _session_looks_stuck(self, session_name, session_data):
"""Determine if a session appears to be stuck waiting for input."""
metadata = session_data['metadata']
metadata = session_data["metadata"]
# Check if process is still running
status = get_session_status(session_name)
if not status or not status.get('is_active', False):
if not status or not status.get("is_active", False):
return False
time_since_activity = time.time() - metadata.get('last_activity', 0)
interaction_count = metadata.get('interaction_count', 0)
time_since_activity = time.time() - metadata.get("last_activity", 0)
interaction_count = metadata.get("interaction_count", 0)
# If running for a while but no interactions, might be waiting
session_age = time.time() - metadata.get('start_time', 0)
session_age = time.time() - metadata.get("start_time", 0)
if session_age > 60 and interaction_count == 0 and time_since_activity > 30:
return True
@@ -119,9 +132,9 @@ class AutonomousInteractions:
def _format_session_updates(self, session_names):
"""Format session information for LLM consumption."""
updates = {
'type': 'background_session_updates',
'timestamp': time.time(),
'sessions': {}
"type": "background_session_updates",
"timestamp": time.time(),
"sessions": {},
}
for session_name in session_names:
@@ -131,12 +144,12 @@ class AutonomousInteractions:
try:
recent_output = read_session_output(session_name, lines=20)
except:
recent_output = {'stdout': '', 'stderr': ''}
recent_output = {"stdout": "", "stderr": ""}
updates['sessions'][session_name] = {
'status': status,
'recent_output': recent_output,
'summary': self._create_session_summary(status, recent_output)
updates["sessions"][session_name] = {
"status": status,
"recent_output": recent_output,
"summary": self._create_session_summary(status, recent_output),
}
return updates
@@ -145,34 +158,39 @@ class AutonomousInteractions:
"""Create a human-readable summary of session status."""
summary_parts = []
process_type = status.get('metadata', {}).get('process_type', 'unknown')
process_type = status.get("metadata", {}).get("process_type", "unknown")
summary_parts.append(f"Type: {process_type}")
is_active = status.get('is_active', False)
is_active = status.get("is_active", False)
summary_parts.append(f"Status: {'Active' if is_active else 'Inactive'}")
if is_active and 'pid' in status:
if is_active and "pid" in status:
summary_parts.append(f"PID: {status['pid']}")
age = time.time() - status.get('metadata', {}).get('start_time', 0)
age = time.time() - status.get("metadata", {}).get("start_time", 0)
summary_parts.append(f"Age: {age:.1f}s")
output_lines = len(recent_output.get('stdout', '').split('\n')) + len(recent_output.get('stderr', '').split('\n'))
output_lines = len(recent_output.get("stdout", "").split("\n")) + len(
recent_output.get("stderr", "").split("\n")
)
summary_parts.append(f"Recent output: {output_lines} lines")
interaction_count = status.get('metadata', {}).get('interaction_count', 0)
interaction_count = status.get("metadata", {}).get("interaction_count", 0)
summary_parts.append(f"Interactions: {interaction_count}")
return " | ".join(summary_parts)
# Global autonomous interactions instance
_global_autonomous = None
def get_global_autonomous():
"""Get the global autonomous interactions instance."""
global _global_autonomous
return _global_autonomous
def start_global_autonomous(llm_callback=None):
"""Start global autonomous interactions."""
global _global_autonomous
@@ -181,6 +199,7 @@ def start_global_autonomous(llm_callback=None):
_global_autonomous.start(llm_callback)
return _global_autonomous
def stop_global_autonomous():
"""Stop global autonomous interactions."""
global _global_autonomous
+94 -64
View File
@@ -1,8 +1,9 @@
import queue
import threading
import time
import queue
from pr.multiplexer import get_all_multiplexer_states, get_multiplexer
from pr.tools.interactive_control import get_session_status
class BackgroundMonitor:
def __init__(self, check_interval=5.0):
@@ -17,7 +18,9 @@ class BackgroundMonitor:
"""Start the background monitoring thread."""
if self.monitor_thread is None:
self.active = True
self.monitor_thread = threading.Thread(target=self._monitor_loop, daemon=True)
self.monitor_thread = threading.Thread(
target=self._monitor_loop, daemon=True
)
self.monitor_thread.start()
def stop(self):
@@ -78,19 +81,18 @@ class BackgroundMonitor:
# Check for new sessions
for session_name in new_states:
if session_name not in old_states:
events.append({
'type': 'session_started',
'session_name': session_name,
'metadata': new_states[session_name]['metadata']
})
events.append(
{
"type": "session_started",
"session_name": session_name,
"metadata": new_states[session_name]["metadata"],
}
)
# Check for ended sessions
for session_name in old_states:
if session_name not in new_states:
events.append({
'type': 'session_ended',
'session_name': session_name
})
events.append({"type": "session_ended", "session_name": session_name})
# Check for activity in existing sessions
for session_name, new_state in new_states.items():
@@ -98,92 +100,112 @@ class BackgroundMonitor:
old_state = old_states[session_name]
# Check for output changes
old_stdout_lines = old_state['output_summary']['stdout_lines']
new_stdout_lines = new_state['output_summary']['stdout_lines']
old_stderr_lines = old_state['output_summary']['stderr_lines']
new_stderr_lines = new_state['output_summary']['stderr_lines']
old_stdout_lines = old_state["output_summary"]["stdout_lines"]
new_stdout_lines = new_state["output_summary"]["stdout_lines"]
old_stderr_lines = old_state["output_summary"]["stderr_lines"]
new_stderr_lines = new_state["output_summary"]["stderr_lines"]
if new_stdout_lines > old_stdout_lines or new_stderr_lines > old_stderr_lines:
if (
new_stdout_lines > old_stdout_lines
or new_stderr_lines > old_stderr_lines
):
# Get the new output
mux = get_multiplexer(session_name)
if mux:
all_output = mux.get_all_output()
new_output = {
'stdout': all_output['stdout'].split('\n')[old_stdout_lines:],
'stderr': all_output['stderr'].split('\n')[old_stderr_lines:]
"stdout": all_output["stdout"].split("\n")[
old_stdout_lines:
],
"stderr": all_output["stderr"].split("\n")[
old_stderr_lines:
],
}
events.append({
'type': 'output_received',
'session_name': session_name,
'new_output': new_output,
'total_lines': {
'stdout': new_stdout_lines,
'stderr': new_stderr_lines
events.append(
{
"type": "output_received",
"session_name": session_name,
"new_output": new_output,
"total_lines": {
"stdout": new_stdout_lines,
"stderr": new_stderr_lines,
},
}
})
)
# Check for state changes
old_metadata = old_state['metadata']
new_metadata = new_state['metadata']
old_metadata = old_state["metadata"]
new_metadata = new_state["metadata"]
if old_metadata.get('state') != new_metadata.get('state'):
events.append({
'type': 'state_changed',
'session_name': session_name,
'old_state': old_metadata.get('state'),
'new_state': new_metadata.get('state')
})
if old_metadata.get("state") != new_metadata.get("state"):
events.append(
{
"type": "state_changed",
"session_name": session_name,
"old_state": old_metadata.get("state"),
"new_state": new_metadata.get("state"),
}
)
# Check for process type identification
if (old_metadata.get('process_type') == 'unknown' and
new_metadata.get('process_type') != 'unknown'):
events.append({
'type': 'process_identified',
'session_name': session_name,
'process_type': new_metadata.get('process_type')
})
if (
old_metadata.get("process_type") == "unknown"
and new_metadata.get("process_type") != "unknown"
):
events.append(
{
"type": "process_identified",
"session_name": session_name,
"process_type": new_metadata.get("process_type"),
}
)
# Check for sessions needing attention (based on heuristics)
for session_name, state in new_states.items():
metadata = state['metadata']
output_summary = state['output_summary']
metadata = state["metadata"]
output_summary = state["output_summary"]
# Heuristic: High output volume might indicate completion or error
total_lines = output_summary['stdout_lines'] + output_summary['stderr_lines']
total_lines = (
output_summary["stdout_lines"] + output_summary["stderr_lines"]
)
if total_lines > 100: # Arbitrary threshold
events.append({
'type': 'high_output_volume',
'session_name': session_name,
'total_lines': total_lines
})
events.append(
{
"type": "high_output_volume",
"session_name": session_name,
"total_lines": total_lines,
}
)
# Heuristic: Long-running session without recent activity
time_since_activity = time.time() - metadata.get('last_activity', 0)
time_since_activity = time.time() - metadata.get("last_activity", 0)
if time_since_activity > 300: # 5 minutes
events.append({
'type': 'inactive_session',
'session_name': session_name,
'inactive_seconds': time_since_activity
})
events.append(
{
"type": "inactive_session",
"session_name": session_name,
"inactive_seconds": time_since_activity,
}
)
# Heuristic: Sessions that might be waiting for input
# This would be enhanced with prompt detection in later phases
if self._might_be_waiting_for_input(session_name, state):
events.append({
'type': 'possible_input_needed',
'session_name': session_name
})
events.append(
{"type": "possible_input_needed", "session_name": session_name}
)
return events
def _might_be_waiting_for_input(self, session_name, state):
"""Heuristic to detect if a session might be waiting for input."""
metadata = state['metadata']
process_type = metadata.get('process_type', 'unknown')
metadata = state["metadata"]
metadata.get("process_type", "unknown")
# Simple heuristics based on process type and recent activity
time_since_activity = time.time() - metadata.get('last_activity', 0)
time_since_activity = time.time() - metadata.get("last_activity", 0)
# If it's been more than 10 seconds since last activity, might be waiting
if time_since_activity > 10:
@@ -191,9 +213,11 @@ class BackgroundMonitor:
return False
# Global monitor instance
_global_monitor = None
def get_global_monitor():
"""Get the global background monitor instance."""
global _global_monitor
@@ -201,20 +225,24 @@ def get_global_monitor():
_global_monitor = BackgroundMonitor()
return _global_monitor
def start_global_monitor():
"""Start the global background monitor."""
monitor = get_global_monitor()
monitor.start()
def stop_global_monitor():
"""Stop the global background monitor."""
global _global_monitor
if _global_monitor:
_global_monitor.stop()
# Global monitor instance
_global_monitor = None
def start_global_monitor():
"""Start the global background monitor."""
global _global_monitor
@@ -223,6 +251,7 @@ def start_global_monitor():
_global_monitor.start()
return _global_monitor
def stop_global_monitor():
"""Stop the global background monitor."""
global _global_monitor
@@ -230,6 +259,7 @@ def stop_global_monitor():
_global_monitor.stop()
_global_monitor = None
def get_global_monitor():
"""Get the global background monitor instance."""
global _global_monitor
+8 -13
View File
@@ -1,22 +1,17 @@
import os
import configparser
from typing import Dict, Any
import os
from typing import Any, Dict
from pr.core.logging import get_logger
logger = get_logger('config')
logger = get_logger("config")
CONFIG_FILE = os.path.expanduser("~/.prrc")
LOCAL_CONFIG_FILE = ".prrc"
def load_config() -> Dict[str, Any]:
config = {
'api': {},
'autonomous': {},
'ui': {},
'output': {},
'session': {}
}
config = {"api": {}, "autonomous": {}, "ui": {}, "output": {}, "session": {}}
global_config = _load_config_file(CONFIG_FILE)
local_config = _load_config_file(LOCAL_CONFIG_FILE)
@@ -55,9 +50,9 @@ def _load_config_file(filepath: str) -> Dict[str, Dict[str, Any]]:
def _parse_value(value: str) -> Any:
value = value.strip()
if value.lower() == 'true':
if value.lower() == "true":
return True
if value.lower() == 'false':
if value.lower() == "false":
return False
if value.isdigit():
@@ -99,7 +94,7 @@ max_history = 1000
"""
try:
with open(filepath, 'w') as f:
with open(filepath, "w") as f:
f.write(default_config)
logger.info(f"Created default configuration at {filepath}")
return True
+125 -38
View File
@@ -1,11 +1,21 @@
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)
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
@@ -17,24 +27,36 @@ def truncate_tool_result(result, max_length=None):
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]"
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]"
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]"
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]"
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.
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
@@ -51,14 +73,15 @@ Process Management:
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."]
- 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('_'):
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]"
@@ -67,7 +90,7 @@ Shell Commands:
for context_file in [CONTEXT_FILE, GLOBAL_CONTEXT_FILE]:
if os.path.exists(context_file):
try:
with open(context_file, 'r') as f:
with open(context_file) as f:
content = f.read()
if len(content) > max_context_size:
content = content[:max_context_size] + "\n... [truncated]"
@@ -78,7 +101,7 @@ Shell Commands:
if args.context:
for ctx_file in args.context:
try:
with open(ctx_file, 'r') as f:
with open(ctx_file) as f:
content = f.read()
if len(content) > max_context_size:
content = content[:max_context_size] + "\n... [truncated]"
@@ -88,22 +111,29 @@ Shell Commands:
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]"
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}")
print(
f"{Colors.YELLOW}📄 Managing context window (current: {len(messages)} messages)...{Colors.RESET}"
)
system_message = messages[0]
recent_messages = messages[-RECENT_MESSAGES_TO_KEEP:]
@@ -113,18 +143,21 @@ def manage_context_window(messages, verbose):
summary = summarize_messages(middle_messages)
summary_message = {
"role": "system",
"content": f"[Previous conversation summary: {summary}]"
"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}")
print(
f"{Colors.GREEN}✓ Context compressed to {len(new_messages)} messages{Colors.RESET}"
)
return new_messages
return messages
def summarize_messages(messages):
summary_parts = []
@@ -142,6 +175,7 @@ def summarize_messages(messages):
return " | ".join(summary_parts[:10])
def estimate_tokens(messages):
total_chars = 0
@@ -155,6 +189,7 @@ def estimate_tokens(messages):
return int(estimated_tokens * overhead_multiplier)
def trim_message_content(message, max_length):
trimmed_msg = message.copy()
@@ -162,14 +197,22 @@ def trim_message_content(message, max_length):
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]"
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]"
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)
@@ -179,35 +222,61 @@ def trim_message_content(message, max_length):
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]"
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]"
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
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 []
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))
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:
@@ -233,6 +302,7 @@ def intelligently_trim_messages(messages, target_tokens, keep_recent=3):
return ([system_msg] if system_msg else []) + messages[-1:]
def auto_slim_messages(messages, verbose=False):
estimated_tokens = estimate_tokens(messages)
@@ -240,29 +310,46 @@ def auto_slim_messages(messages, verbose=False):
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}")
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)
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}")
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}")
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}")
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
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
+112 -82
View File
@@ -1,22 +1,29 @@
import logging
import json
import logging
import uuid
from typing import Optional, Dict, Any, List
from pr.config import (
DB_PATH, CACHE_ENABLED, API_CACHE_TTL, TOOL_CACHE_TTL,
WORKFLOW_EXECUTOR_MAX_WORKERS, AGENT_MAX_WORKERS,
KNOWLEDGE_SEARCH_LIMIT, ADVANCED_CONTEXT_ENABLED,
MEMORY_AUTO_SUMMARIZE, CONVERSATION_SUMMARY_THRESHOLD
)
from pr.cache import APICache, ToolCache
from pr.workflows import WorkflowEngine, WorkflowStorage
from typing import Any, Dict, List, Optional
from pr.agents import AgentManager
from pr.memory import KnowledgeStore, ConversationMemory, FactExtractor
from pr.cache import APICache, ToolCache
from pr.config import (
ADVANCED_CONTEXT_ENABLED,
API_CACHE_TTL,
CACHE_ENABLED,
CONVERSATION_SUMMARY_THRESHOLD,
DB_PATH,
KNOWLEDGE_SEARCH_LIMIT,
MEMORY_AUTO_SUMMARIZE,
TOOL_CACHE_TTL,
WORKFLOW_EXECUTOR_MAX_WORKERS,
)
from pr.core.advanced_context import AdvancedContextManager
from pr.core.api import call_api
from pr.memory import ConversationMemory, FactExtractor, KnowledgeStore
from pr.tools.base import get_tools_definition
from pr.workflows import WorkflowEngine, WorkflowStorage
logger = logging.getLogger("pr")
logger = logging.getLogger('pr')
class EnhancedAssistant:
def __init__(self, base_assistant):
@@ -32,7 +39,7 @@ class EnhancedAssistant:
self.workflow_storage = WorkflowStorage(DB_PATH)
self.workflow_engine = WorkflowEngine(
tool_executor=self._execute_tool_for_workflow,
max_workers=WORKFLOW_EXECUTOR_MAX_WORKERS
max_workers=WORKFLOW_EXECUTOR_MAX_WORKERS,
)
self.agent_manager = AgentManager(DB_PATH, self._api_caller_for_agent)
@@ -44,20 +51,21 @@ class EnhancedAssistant:
if ADVANCED_CONTEXT_ENABLED:
self.context_manager = AdvancedContextManager(
knowledge_store=self.knowledge_store,
conversation_memory=self.conversation_memory
conversation_memory=self.conversation_memory,
)
else:
self.context_manager = None
self.current_conversation_id = str(uuid.uuid4())[:16]
self.conversation_memory.create_conversation(
self.current_conversation_id,
session_id=str(uuid.uuid4())[:16]
self.current_conversation_id, session_id=str(uuid.uuid4())[:16]
)
logger.info("Enhanced Assistant initialized with all features")
def _execute_tool_for_workflow(self, tool_name: str, arguments: Dict[str, Any]) -> Any:
def _execute_tool_for_workflow(
self, tool_name: str, arguments: Dict[str, Any]
) -> Any:
if self.tool_cache:
cached_result = self.tool_cache.get(tool_name, arguments)
if cached_result is not None:
@@ -65,41 +73,66 @@ class EnhancedAssistant:
return cached_result
func_map = {
'read_file': lambda **kw: self.base.execute_tool_calls([{
'id': 'temp',
'function': {'name': 'read_file', 'arguments': json.dumps(kw)}
}])[0],
'write_file': lambda **kw: self.base.execute_tool_calls([{
'id': 'temp',
'function': {'name': 'write_file', 'arguments': json.dumps(kw)}
}])[0],
'list_directory': lambda **kw: self.base.execute_tool_calls([{
'id': 'temp',
'function': {'name': 'list_directory', 'arguments': json.dumps(kw)}
}])[0],
'run_command': lambda **kw: self.base.execute_tool_calls([{
'id': 'temp',
'function': {'name': 'run_command', 'arguments': json.dumps(kw)}
}])[0],
"read_file": lambda **kw: self.base.execute_tool_calls(
[
{
"id": "temp",
"function": {"name": "read_file", "arguments": json.dumps(kw)},
}
]
)[0],
"write_file": lambda **kw: self.base.execute_tool_calls(
[
{
"id": "temp",
"function": {"name": "write_file", "arguments": json.dumps(kw)},
}
]
)[0],
"list_directory": lambda **kw: self.base.execute_tool_calls(
[
{
"id": "temp",
"function": {
"name": "list_directory",
"arguments": json.dumps(kw),
},
}
]
)[0],
"run_command": lambda **kw: self.base.execute_tool_calls(
[
{
"id": "temp",
"function": {
"name": "run_command",
"arguments": json.dumps(kw),
},
}
]
)[0],
}
if tool_name in func_map:
result = func_map[tool_name](**arguments)
if self.tool_cache:
content = result.get('content', '')
content = result.get("content", "")
try:
parsed_content = json.loads(content) if isinstance(content, str) else content
parsed_content = (
json.loads(content) if isinstance(content, str) else content
)
self.tool_cache.set(tool_name, arguments, parsed_content)
except Exception:
pass
return result
return {'error': f'Unknown tool: {tool_name}'}
return {"error": f"Unknown tool: {tool_name}"}
def _api_caller_for_agent(self, messages: List[Dict[str, Any]],
temperature: float, max_tokens: int) -> Dict[str, Any]:
def _api_caller_for_agent(
self, messages: List[Dict[str, Any]], temperature: float, max_tokens: int
) -> Dict[str, Any]:
return call_api(
messages,
self.base.model,
@@ -109,15 +142,12 @@ class EnhancedAssistant:
tools=None,
temperature=temperature,
max_tokens=max_tokens,
verbose=self.base.verbose
verbose=self.base.verbose,
)
def enhanced_call_api(self, messages: List[Dict[str, Any]]) -> Dict[str, Any]:
if self.api_cache and CACHE_ENABLED:
cached_response = self.api_cache.get(
self.base.model, messages,
0.7, 4096
)
cached_response = self.api_cache.get(self.base.model, messages, 0.7, 4096)
if cached_response:
logger.debug("API cache hit")
return cached_response
@@ -129,15 +159,13 @@ class EnhancedAssistant:
self.base.api_key,
self.base.use_tools,
get_tools_definition(),
verbose=self.base.verbose
verbose=self.base.verbose,
)
if self.api_cache and CACHE_ENABLED and 'error' not in response:
token_count = response.get('usage', {}).get('total_tokens', 0)
if self.api_cache and CACHE_ENABLED and "error" not in response:
token_count = response.get("usage", {}).get("total_tokens", 0)
self.api_cache.set(
self.base.model, messages,
0.7, 4096,
response, token_count
self.base.model, messages, 0.7, 4096, response, token_count
)
return response
@@ -146,35 +174,33 @@ class EnhancedAssistant:
self.base.messages.append({"role": "user", "content": user_message})
self.conversation_memory.add_message(
self.current_conversation_id,
str(uuid.uuid4())[:16],
'user',
user_message
self.current_conversation_id, str(uuid.uuid4())[:16], "user", user_message
)
if MEMORY_AUTO_SUMMARIZE and len(self.base.messages) % 5 == 0:
facts = self.fact_extractor.extract_facts(user_message)
for fact in facts[:3]:
entry_id = str(uuid.uuid4())[:16]
from pr.memory import KnowledgeEntry
import time
categories = self.fact_extractor.categorize_content(fact['text'])
from pr.memory import KnowledgeEntry
categories = self.fact_extractor.categorize_content(fact["text"])
entry = KnowledgeEntry(
entry_id=entry_id,
category=categories[0] if categories else 'general',
content=fact['text'],
metadata={'type': fact['type'], 'confidence': fact['confidence']},
category=categories[0] if categories else "general",
content=fact["text"],
metadata={"type": fact["type"], "confidence": fact["confidence"]},
created_at=time.time(),
updated_at=time.time()
updated_at=time.time(),
)
self.knowledge_store.add_entry(entry)
if self.context_manager and ADVANCED_CONTEXT_ENABLED:
enhanced_messages, context_info = self.context_manager.create_enhanced_context(
self.base.messages,
user_message,
include_knowledge=True
enhanced_messages, context_info = (
self.context_manager.create_enhanced_context(
self.base.messages, user_message, include_knowledge=True
)
)
if self.base.verbose:
@@ -189,38 +215,40 @@ class EnhancedAssistant:
result = self.base.process_response(response)
if len(self.base.messages) >= CONVERSATION_SUMMARY_THRESHOLD:
summary = self.context_manager.advanced_summarize_messages(
self.base.messages[-CONVERSATION_SUMMARY_THRESHOLD:]
) if self.context_manager else "Conversation in progress"
summary = (
self.context_manager.advanced_summarize_messages(
self.base.messages[-CONVERSATION_SUMMARY_THRESHOLD:]
)
if self.context_manager
else "Conversation in progress"
)
topics = self.fact_extractor.categorize_content(summary)
self.conversation_memory.update_conversation_summary(
self.current_conversation_id,
summary,
topics
self.current_conversation_id, summary, topics
)
return result
def execute_workflow(self, workflow_name: str,
initial_variables: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
def execute_workflow(
self, workflow_name: str, initial_variables: Optional[Dict[str, Any]] = None
) -> Dict[str, Any]:
workflow = self.workflow_storage.load_workflow_by_name(workflow_name)
if not workflow:
return {'error': f'Workflow "{workflow_name}" not found'}
return {"error": f'Workflow "{workflow_name}" not found'}
context = self.workflow_engine.execute_workflow(workflow, initial_variables)
execution_id = self.workflow_storage.save_execution(
self.workflow_storage.load_workflow_by_name(workflow_name).name,
context
self.workflow_storage.load_workflow_by_name(workflow_name).name, context
)
return {
'success': True,
'execution_id': execution_id,
'results': context.step_results,
'execution_log': context.execution_log
"success": True,
"execution_id": execution_id,
"results": context.step_results,
"execution_log": context.execution_log,
}
def create_agent(self, role_name: str, agent_id: Optional[str] = None) -> str:
@@ -230,20 +258,22 @@ class EnhancedAssistant:
return self.agent_manager.execute_agent_task(agent_id, task)
def collaborate_agents(self, task: str, agent_roles: List[str]) -> Dict[str, Any]:
orchestrator_id = self.agent_manager.create_agent('orchestrator')
orchestrator_id = self.agent_manager.create_agent("orchestrator")
return self.agent_manager.collaborate_agents(orchestrator_id, task, agent_roles)
def search_knowledge(self, query: str, limit: int = KNOWLEDGE_SEARCH_LIMIT) -> List[Any]:
def search_knowledge(
self, query: str, limit: int = KNOWLEDGE_SEARCH_LIMIT
) -> List[Any]:
return self.knowledge_store.search_entries(query, top_k=limit)
def get_cache_statistics(self) -> Dict[str, Any]:
stats = {}
if self.api_cache:
stats['api_cache'] = self.api_cache.get_statistics()
stats["api_cache"] = self.api_cache.get_statistics()
if self.tool_cache:
stats['tool_cache'] = self.tool_cache.get_statistics()
stats["tool_cache"] = self.tool_cache.get_statistics()
return stats
+8 -11
View File
@@ -1,6 +1,7 @@
import logging
import os
from logging.handlers import RotatingFileHandler
from pr.config import LOG_FILE
@@ -9,21 +10,19 @@ def setup_logging(verbose=False):
if log_dir and not os.path.exists(log_dir):
os.makedirs(log_dir, exist_ok=True)
logger = logging.getLogger('pr')
logger = logging.getLogger("pr")
logger.setLevel(logging.DEBUG if verbose else logging.INFO)
if logger.handlers:
logger.handlers.clear()
file_handler = RotatingFileHandler(
LOG_FILE,
maxBytes=10 * 1024 * 1024,
backupCount=5
LOG_FILE, maxBytes=10 * 1024 * 1024, backupCount=5
)
file_handler.setLevel(logging.DEBUG)
file_formatter = logging.Formatter(
'%(asctime)s - %(name)s - %(levelname)s - %(filename)s:%(lineno)d - %(message)s',
datefmt='%Y-%m-%d %H:%M:%S'
"%(asctime)s - %(name)s - %(levelname)s - %(filename)s:%(lineno)d - %(message)s",
datefmt="%Y-%m-%d %H:%M:%S",
)
file_handler.setFormatter(file_formatter)
logger.addHandler(file_handler)
@@ -31,9 +30,7 @@ def setup_logging(verbose=False):
if verbose:
console_handler = logging.StreamHandler()
console_handler.setLevel(logging.INFO)
console_formatter = logging.Formatter(
'%(levelname)s: %(message)s'
)
console_formatter = logging.Formatter("%(levelname)s: %(message)s")
console_handler.setFormatter(console_formatter)
logger.addHandler(console_handler)
@@ -42,5 +39,5 @@ def setup_logging(verbose=False):
def get_logger(name=None):
if name:
return logging.getLogger(f'pr.{name}')
return logging.getLogger('pr')
return logging.getLogger(f"pr.{name}")
return logging.getLogger("pr")
+35 -30
View File
@@ -2,9 +2,10 @@ import json
import os
from datetime import datetime
from typing import Dict, List, Optional
from pr.core.logging import get_logger
logger = get_logger('session')
logger = get_logger("session")
SESSIONS_DIR = os.path.expanduser("~/.assistant_sessions")
@@ -14,18 +15,20 @@ class SessionManager:
def __init__(self):
os.makedirs(SESSIONS_DIR, exist_ok=True)
def save_session(self, name: str, messages: List[Dict], metadata: Optional[Dict] = None) -> bool:
def save_session(
self, name: str, messages: List[Dict], metadata: Optional[Dict] = None
) -> bool:
try:
session_file = os.path.join(SESSIONS_DIR, f"{name}.json")
session_data = {
'name': name,
'created_at': datetime.now().isoformat(),
'messages': messages,
'metadata': metadata or {}
"name": name,
"created_at": datetime.now().isoformat(),
"messages": messages,
"metadata": metadata or {},
}
with open(session_file, 'w') as f:
with open(session_file, "w") as f:
json.dump(session_data, f, indent=2)
logger.info(f"Session saved: {name}")
@@ -43,7 +46,7 @@ class SessionManager:
logger.warning(f"Session not found: {name}")
return None
with open(session_file, 'r') as f:
with open(session_file) as f:
session_data = json.load(f)
logger.info(f"Session loaded: {name}")
@@ -58,22 +61,24 @@ class SessionManager:
try:
for filename in os.listdir(SESSIONS_DIR):
if filename.endswith('.json'):
if filename.endswith(".json"):
filepath = os.path.join(SESSIONS_DIR, filename)
try:
with open(filepath, 'r') as f:
with open(filepath) as f:
data = json.load(f)
sessions.append({
'name': data.get('name', filename[:-5]),
'created_at': data.get('created_at', 'unknown'),
'message_count': len(data.get('messages', [])),
'metadata': data.get('metadata', {})
})
sessions.append(
{
"name": data.get("name", filename[:-5]),
"created_at": data.get("created_at", "unknown"),
"message_count": len(data.get("messages", [])),
"metadata": data.get("metadata", {}),
}
)
except Exception as e:
logger.warning(f"Error reading session file {filename}: {e}")
sessions.sort(key=lambda x: x['created_at'], reverse=True)
sessions.sort(key=lambda x: x["created_at"], reverse=True)
except Exception as e:
logger.error(f"Error listing sessions: {e}")
@@ -96,39 +101,39 @@ class SessionManager:
logger.error(f"Error deleting session {name}: {e}")
return False
def export_session(self, name: str, output_path: str, format: str = 'json') -> bool:
def export_session(self, name: str, output_path: str, format: str = "json") -> bool:
session_data = self.load_session(name)
if not session_data:
return False
try:
if format == 'json':
with open(output_path, 'w') as f:
if format == "json":
with open(output_path, "w") as f:
json.dump(session_data, f, indent=2)
elif format == 'markdown':
with open(output_path, 'w') as f:
elif format == "markdown":
with open(output_path, "w") as f:
f.write(f"# Session: {name}\n\n")
f.write(f"Created: {session_data['created_at']}\n\n")
f.write("---\n\n")
for msg in session_data['messages']:
role = msg.get('role', 'unknown')
content = msg.get('content', '')
for msg in session_data["messages"]:
role = msg.get("role", "unknown")
content = msg.get("content", "")
f.write(f"## {role.capitalize()}\n\n")
f.write(f"{content}\n\n")
f.write("---\n\n")
elif format == 'txt':
with open(output_path, 'w') as f:
elif format == "txt":
with open(output_path, "w") as f:
f.write(f"Session: {name}\n")
f.write(f"Created: {session_data['created_at']}\n")
f.write("=" * 80 + "\n\n")
for msg in session_data['messages']:
role = msg.get('role', 'unknown')
content = msg.get('content', '')
for msg in session_data["messages"]:
role = msg.get("role", "unknown")
content = msg.get("content", "")
f.write(f"[{role.upper()}]\n")
f.write(f"{content}\n")
+62 -63
View File
@@ -2,20 +2,21 @@ import json
import os
from datetime import datetime
from typing import Dict, Optional
from pr.core.logging import get_logger
logger = get_logger('usage')
logger = get_logger("usage")
USAGE_DB_FILE = os.path.expanduser("~/.assistant_usage.json")
MODEL_COSTS = {
'x-ai/grok-code-fast-1': {'input': 0.0, 'output': 0.0},
'gpt-4': {'input': 0.03, 'output': 0.06},
'gpt-4-turbo': {'input': 0.01, 'output': 0.03},
'gpt-3.5-turbo': {'input': 0.0005, 'output': 0.0015},
'claude-3-opus': {'input': 0.015, 'output': 0.075},
'claude-3-sonnet': {'input': 0.003, 'output': 0.015},
'claude-3-haiku': {'input': 0.00025, 'output': 0.00125},
"x-ai/grok-code-fast-1": {"input": 0.0, "output": 0.0},
"gpt-4": {"input": 0.03, "output": 0.06},
"gpt-4-turbo": {"input": 0.01, "output": 0.03},
"gpt-3.5-turbo": {"input": 0.0005, "output": 0.0015},
"claude-3-opus": {"input": 0.015, "output": 0.075},
"claude-3-sonnet": {"input": 0.003, "output": 0.015},
"claude-3-haiku": {"input": 0.00025, "output": 0.00125},
}
@@ -23,12 +24,12 @@ class UsageTracker:
def __init__(self):
self.session_usage = {
'requests': 0,
'total_tokens': 0,
'input_tokens': 0,
'output_tokens': 0,
'estimated_cost': 0.0,
'models_used': {}
"requests": 0,
"total_tokens": 0,
"input_tokens": 0,
"output_tokens": 0,
"estimated_cost": 0.0,
"models_used": {},
}
def track_request(
@@ -36,30 +37,30 @@ class UsageTracker:
model: str,
input_tokens: int,
output_tokens: int,
total_tokens: Optional[int] = None
total_tokens: Optional[int] = None,
):
if total_tokens is None:
total_tokens = input_tokens + output_tokens
self.session_usage['requests'] += 1
self.session_usage['total_tokens'] += total_tokens
self.session_usage['input_tokens'] += input_tokens
self.session_usage['output_tokens'] += output_tokens
self.session_usage["requests"] += 1
self.session_usage["total_tokens"] += total_tokens
self.session_usage["input_tokens"] += input_tokens
self.session_usage["output_tokens"] += output_tokens
if model not in self.session_usage['models_used']:
self.session_usage['models_used'][model] = {
'requests': 0,
'tokens': 0,
'cost': 0.0
if model not in self.session_usage["models_used"]:
self.session_usage["models_used"][model] = {
"requests": 0,
"tokens": 0,
"cost": 0.0,
}
model_usage = self.session_usage['models_used'][model]
model_usage['requests'] += 1
model_usage['tokens'] += total_tokens
model_usage = self.session_usage["models_used"][model]
model_usage["requests"] += 1
model_usage["tokens"] += total_tokens
cost = self._calculate_cost(model, input_tokens, output_tokens)
model_usage['cost'] += cost
self.session_usage['estimated_cost'] += cost
model_usage["cost"] += cost
self.session_usage["estimated_cost"] += cost
self._save_to_history(model, input_tokens, output_tokens, cost)
@@ -67,9 +68,11 @@ class UsageTracker:
f"Tracked request: {model}, tokens: {total_tokens}, cost: ${cost:.4f}"
)
def _calculate_cost(self, model: str, input_tokens: int, output_tokens: int) -> float:
def _calculate_cost(
self, model: str, input_tokens: int, output_tokens: int
) -> float:
if model not in MODEL_COSTS:
base_model = model.split('/')[0] if '/' in model else model
base_model = model.split("/")[0] if "/" in model else model
if base_model not in MODEL_COSTS:
logger.warning(f"Unknown model for cost calculation: {model}")
return 0.0
@@ -77,31 +80,35 @@ class UsageTracker:
else:
costs = MODEL_COSTS[model]
input_cost = (input_tokens / 1000) * costs['input']
output_cost = (output_tokens / 1000) * costs['output']
input_cost = (input_tokens / 1000) * costs["input"]
output_cost = (output_tokens / 1000) * costs["output"]
return input_cost + output_cost
def _save_to_history(self, model: str, input_tokens: int, output_tokens: int, cost: float):
def _save_to_history(
self, model: str, input_tokens: int, output_tokens: int, cost: float
):
try:
history = []
if os.path.exists(USAGE_DB_FILE):
with open(USAGE_DB_FILE, 'r') as f:
with open(USAGE_DB_FILE) as f:
history = json.load(f)
history.append({
'timestamp': datetime.now().isoformat(),
'model': model,
'input_tokens': input_tokens,
'output_tokens': output_tokens,
'total_tokens': input_tokens + output_tokens,
'cost': cost
})
history.append(
{
"timestamp": datetime.now().isoformat(),
"model": model,
"input_tokens": input_tokens,
"output_tokens": output_tokens,
"total_tokens": input_tokens + output_tokens,
"cost": cost,
}
)
if len(history) > 10000:
history = history[-10000:]
with open(USAGE_DB_FILE, 'w') as f:
with open(USAGE_DB_FILE, "w") as f:
json.dump(history, f, indent=2)
except Exception as e:
@@ -121,42 +128,34 @@ class UsageTracker:
f"Estimated Cost: ${usage['estimated_cost']:.4f}",
]
if usage['models_used']:
if usage["models_used"]:
lines.append("\nModels Used:")
for model, stats in usage['models_used'].items():
for model, stats in usage["models_used"].items():
lines.append(
f" {model}: {stats['requests']} requests, "
f"{stats['tokens']:,} tokens, ${stats['cost']:.4f}"
)
return '\n'.join(lines)
return "\n".join(lines)
@staticmethod
def get_total_usage() -> Dict:
if not os.path.exists(USAGE_DB_FILE):
return {
'total_requests': 0,
'total_tokens': 0,
'total_cost': 0.0
}
return {"total_requests": 0, "total_tokens": 0, "total_cost": 0.0}
try:
with open(USAGE_DB_FILE, 'r') as f:
with open(USAGE_DB_FILE) as f:
history = json.load(f)
total_tokens = sum(entry['total_tokens'] for entry in history)
total_cost = sum(entry['cost'] for entry in history)
total_tokens = sum(entry["total_tokens"] for entry in history)
total_cost = sum(entry["cost"] for entry in history)
return {
'total_requests': len(history),
'total_tokens': total_tokens,
'total_cost': total_cost
"total_requests": len(history),
"total_tokens": total_tokens,
"total_cost": total_cost,
}
except Exception as e:
logger.error(f"Error loading usage history: {e}")
return {
'total_requests': 0,
'total_tokens': 0,
'total_cost': 0.0
}
return {"total_requests": 0, "total_tokens": 0, "total_cost": 0.0}
+6 -4
View File
@@ -1,5 +1,5 @@
import os
from typing import Optional
from pr.core.exceptions import ValidationError
@@ -16,7 +16,9 @@ def validate_file_path(path: str, must_exist: bool = False) -> str:
return os.path.abspath(path)
def validate_directory_path(path: str, must_exist: bool = False, create: bool = False) -> str:
def validate_directory_path(
path: str, must_exist: bool = False, create: bool = False
) -> str:
if not path:
raise ValidationError("Directory path cannot be empty")
@@ -48,7 +50,7 @@ def validate_api_url(url: str) -> str:
if not url:
raise ValidationError("API URL cannot be empty")
if not url.startswith(('http://', 'https://')):
if not url.startswith(("http://", "https://")):
raise ValidationError("API URL must start with http:// or https://")
return url
@@ -58,7 +60,7 @@ def validate_session_name(name: str) -> str:
if not name:
raise ValidationError("Session name cannot be empty")
invalid_chars = ['/', '\\', ':', '*', '?', '"', '<', '>', '|']
invalid_chars = ["/", "\\", ":", "*", "?", '"', "<", ">", "|"]
for char in invalid_chars:
if char in name:
raise ValidationError(f"Session name contains invalid character: {char}")