chore: scaffold project with editorconfig, ci workflows, gitignore, pre-commit, changelog, contributing guide, license, and makefile
This commit is contained in:
@@ -0,0 +1,5 @@
|
||||
from pr.core.assistant import Assistant
|
||||
from pr.core.api import call_api, list_models
|
||||
from pr.core.context import init_system_message, manage_context_window
|
||||
|
||||
__all__ = ['Assistant', 'call_api', 'list_models', 'init_system_message', 'manage_context_window']
|
||||
@@ -0,0 +1,82 @@
|
||||
import re
|
||||
import math
|
||||
from typing import List, Dict, Any
|
||||
from collections import Counter
|
||||
|
||||
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:
|
||||
complexity_thresholds = {
|
||||
'simple': 10,
|
||||
'medium': 20,
|
||||
'complex': 35,
|
||||
'very_complex': 50
|
||||
}
|
||||
|
||||
base_threshold = complexity_thresholds.get(task_complexity, 20)
|
||||
|
||||
message_complexity_score = self._analyze_message_complexity(messages)
|
||||
|
||||
if message_complexity_score > 0.7:
|
||||
adjusted = int(base_threshold * 1.5)
|
||||
elif message_complexity_score < 0.3:
|
||||
adjusted = int(base_threshold * 0.7)
|
||||
else:
|
||||
adjusted = base_threshold
|
||||
|
||||
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)
|
||||
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())
|
||||
unique_words.update(words)
|
||||
|
||||
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
|
||||
|
||||
def extract_key_sentences(self, text: str, top_k: int = 5) -> List[str]:
|
||||
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):
|
||||
length_score = min(1.0, len(sentence) / 50)
|
||||
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])
|
||||
key_sentences = self.extract_key_sentences(all_content, top_k=3)
|
||||
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()))
|
||||
|
||||
intersection = content_words & context_words
|
||||
union = content_words | context_words
|
||||
|
||||
if not union:
|
||||
return 0.0
|
||||
|
||||
return len(intersection) / len(union)
|
||||
@@ -0,0 +1,95 @@
|
||||
import json
|
||||
import urllib.request
|
||||
import urllib.error
|
||||
import logging
|
||||
from pr.config import DEFAULT_TEMPERATURE, DEFAULT_MAX_TOKENS
|
||||
from pr.core.context import auto_slim_messages
|
||||
|
||||
logger = logging.getLogger('pr')
|
||||
|
||||
def call_api(messages, model, api_url, api_key, use_tools, tools_definition, verbose=False):
|
||||
try:
|
||||
messages = auto_slim_messages(messages, verbose=verbose)
|
||||
|
||||
logger.debug(f"=== API CALL START ===")
|
||||
logger.debug(f"Model: {model}")
|
||||
logger.debug(f"API URL: {api_url}")
|
||||
logger.debug(f"Use tools: {use_tools}")
|
||||
logger.debug(f"Message count: {len(messages)}")
|
||||
headers = {
|
||||
'Content-Type': 'application/json',
|
||||
}
|
||||
|
||||
if api_key:
|
||||
headers['Authorization'] = f'Bearer {api_key}'
|
||||
|
||||
data = {
|
||||
'model': model,
|
||||
'messages': messages,
|
||||
'temperature': DEFAULT_TEMPERATURE,
|
||||
'max_tokens': DEFAULT_MAX_TOKENS
|
||||
}
|
||||
|
||||
if "gpt-5" in model:
|
||||
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'
|
||||
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'
|
||||
)
|
||||
|
||||
logger.debug("Sending HTTP request...")
|
||||
with urllib.request.urlopen(req) as response:
|
||||
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:
|
||||
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']
|
||||
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)")
|
||||
|
||||
logger.debug("=== API CALL END ===")
|
||||
return result
|
||||
|
||||
except urllib.error.HTTPError as e:
|
||||
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}
|
||||
except Exception as e:
|
||||
logger.error(f"API call failed: {e}")
|
||||
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}')
|
||||
|
||||
with urllib.request.urlopen(req) as response:
|
||||
data = json.loads(response.read().decode('utf-8'))
|
||||
|
||||
return data.get('data', [])
|
||||
except Exception as e:
|
||||
return {"error": str(e)}
|
||||
@@ -0,0 +1,325 @@
|
||||
import os
|
||||
import sys
|
||||
import json
|
||||
import sqlite3
|
||||
import signal
|
||||
import logging
|
||||
import traceback
|
||||
import readline
|
||||
import glob as glob_module
|
||||
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.core.api import call_api
|
||||
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
|
||||
)
|
||||
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
|
||||
|
||||
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'))
|
||||
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.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'))
|
||||
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.interrupt_count = 0
|
||||
self.python_globals = {}
|
||||
self.db_conn = None
|
||||
self.autonomous_mode = False
|
||||
self.autonomous_iterations = 0
|
||||
self.init_database()
|
||||
self.messages.append(init_system_message(args))
|
||||
|
||||
try:
|
||||
from pr.core.enhanced_assistant import EnhancedAssistant
|
||||
self.enhanced = EnhancedAssistant(self)
|
||||
if self.debug:
|
||||
logger.debug("Enhanced assistant features initialized")
|
||||
except Exception as e:
|
||||
logger.warning(f"Could not initialize enhanced features: {e}")
|
||||
self.enhanced = None
|
||||
|
||||
def init_database(self):
|
||||
try:
|
||||
logger.debug(f"Initializing database at {DB_PATH}")
|
||||
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 file_versions
|
||||
(id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
filepath TEXT, content TEXT, hash TEXT,
|
||||
timestamp REAL, version INTEGER)''')
|
||||
|
||||
self.db_conn.commit()
|
||||
logger.debug("Database initialized successfully")
|
||||
except Exception as e:
|
||||
logger.error(f"Database initialization error: {e}")
|
||||
self.db_conn = None
|
||||
|
||||
def execute_tool_calls(self, tool_calls):
|
||||
results = []
|
||||
|
||||
logger.debug(f"Executing {len(tool_calls)} tool call(s)")
|
||||
|
||||
with ThreadPoolExecutor(max_workers=5) as executor:
|
||||
futures = []
|
||||
|
||||
for tool_call in tool_calls:
|
||||
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),
|
||||
'run_command_interactive': lambda **kw: run_command_interactive(**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(),
|
||||
}
|
||||
|
||||
if func_name in func_map:
|
||||
future = executor.submit(func_map[func_name], **arguments)
|
||||
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)
|
||||
})
|
||||
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})
|
||||
})
|
||||
|
||||
return results
|
||||
|
||||
def process_response(self, response):
|
||||
if 'error' in response:
|
||||
return f"Error: {response['error']}"
|
||||
|
||||
if 'choices' not in response or not response['choices']:
|
||||
return "No response from API"
|
||||
|
||||
message = response['choices'][0]['message']
|
||||
self.messages.append(message)
|
||||
|
||||
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'])
|
||||
|
||||
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
|
||||
)
|
||||
return self.process_response(follow_up)
|
||||
|
||||
content = message.get('content', '')
|
||||
return render_markdown(content, self.syntax_highlighting)
|
||||
|
||||
def signal_handler(self, signum, frame):
|
||||
if self.autonomous_mode:
|
||||
self.interrupt_count += 1
|
||||
if self.interrupt_count >= 2:
|
||||
print(f"\n{Colors.RED}Force exiting autonomous mode...{Colors.RESET}")
|
||||
self.autonomous_mode = False
|
||||
sys.exit(0)
|
||||
else:
|
||||
print(f"\n{Colors.YELLOW}Press Ctrl+C again to force exit{Colors.RESET}")
|
||||
return
|
||||
|
||||
self.interrupt_count += 1
|
||||
if self.interrupt_count >= 2:
|
||||
print(f"\n{Colors.RED}Exiting...{Colors.RESET}")
|
||||
self.cleanup()
|
||||
sys.exit(0)
|
||||
else:
|
||||
print(f"\n{Colors.YELLOW}Press Ctrl+C again to exit{Colors.RESET}")
|
||||
|
||||
def setup_readline(self):
|
||||
try:
|
||||
readline.read_history_file(HISTORY_FILE)
|
||||
except FileNotFoundError:
|
||||
pass
|
||||
|
||||
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']
|
||||
|
||||
def completer(text, state):
|
||||
options = [cmd for cmd in commands if cmd.startswith(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)))
|
||||
|
||||
if state < len(combined_options):
|
||||
return combined_options[state]
|
||||
|
||||
return None
|
||||
|
||||
delims = readline.get_completer_delims()
|
||||
readline.set_completer_delims(delims.replace('/', ''))
|
||||
|
||||
readline.set_completer(completer)
|
||||
readline.parse_and_bind('tab: complete')
|
||||
|
||||
def run_repl(self):
|
||||
self.setup_readline()
|
||||
signal.signal(signal.SIGINT, self.signal_handler)
|
||||
|
||||
print(f"{Colors.BOLD}r{Colors.RESET}")
|
||||
print(f"Type 'help' for commands or start chatting")
|
||||
|
||||
while True:
|
||||
try:
|
||||
user_input = input(f"{Colors.BLUE}You>{Colors.RESET} ").strip()
|
||||
|
||||
if not user_input:
|
||||
continue
|
||||
|
||||
cmd_result = handle_command(self, user_input)
|
||||
|
||||
if cmd_result is False:
|
||||
break
|
||||
elif cmd_result is True:
|
||||
continue
|
||||
|
||||
process_message(self, user_input)
|
||||
|
||||
except EOFError:
|
||||
break
|
||||
except KeyboardInterrupt:
|
||||
self.signal_handler(None, None)
|
||||
except Exception as e:
|
||||
print(f"{Colors.RED}Error: {e}{Colors.RESET}")
|
||||
logging.error(f"REPL error: {e}\n{traceback.format_exc()}")
|
||||
|
||||
def run_single(self):
|
||||
if self.args.message:
|
||||
message = self.args.message
|
||||
else:
|
||||
message = sys.stdin.read()
|
||||
|
||||
process_message(self, message)
|
||||
|
||||
def cleanup(self):
|
||||
if hasattr(self, 'enhanced') and self.enhanced:
|
||||
try:
|
||||
self.enhanced.cleanup()
|
||||
except Exception as e:
|
||||
logger.error(f"Error cleaning up enhanced features: {e}")
|
||||
|
||||
try:
|
||||
from pr.multiplexer import cleanup_all_multiplexers
|
||||
cleanup_all_multiplexers()
|
||||
except Exception as e:
|
||||
logger.error(f"Error cleaning up multiplexers: {e}")
|
||||
|
||||
if self.db_conn:
|
||||
self.db_conn.close()
|
||||
|
||||
def run(self):
|
||||
try:
|
||||
if self.args.interactive or (not self.args.message and sys.stdin.isatty()):
|
||||
self.run_repl()
|
||||
else:
|
||||
self.run_single()
|
||||
finally:
|
||||
self.cleanup()
|
||||
|
||||
def process_message(assistant, message):
|
||||
assistant.messages.append({"role": "user", "content": message})
|
||||
|
||||
logger.debug(f"Processing user message: {message[:100]}...")
|
||||
logger.debug(f"Current message count: {len(assistant.messages)}")
|
||||
|
||||
if assistant.verbose:
|
||||
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
|
||||
)
|
||||
result = assistant.process_response(response)
|
||||
|
||||
print(f"\n{Colors.GREEN}r:{Colors.RESET} {result}\n")
|
||||
@@ -0,0 +1,108 @@
|
||||
import os
|
||||
import configparser
|
||||
from typing import Dict, Any
|
||||
from pr.core.logging import get_logger
|
||||
|
||||
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': {}
|
||||
}
|
||||
|
||||
global_config = _load_config_file(CONFIG_FILE)
|
||||
local_config = _load_config_file(LOCAL_CONFIG_FILE)
|
||||
|
||||
for section in config.keys():
|
||||
if section in global_config:
|
||||
config[section].update(global_config[section])
|
||||
if section in local_config:
|
||||
config[section].update(local_config[section])
|
||||
|
||||
return config
|
||||
|
||||
|
||||
def _load_config_file(filepath: str) -> Dict[str, Dict[str, Any]]:
|
||||
if not os.path.exists(filepath):
|
||||
return {}
|
||||
|
||||
try:
|
||||
parser = configparser.ConfigParser()
|
||||
parser.read(filepath)
|
||||
|
||||
config = {}
|
||||
for section in parser.sections():
|
||||
config[section] = {}
|
||||
for key, value in parser.items(section):
|
||||
config[section][key] = _parse_value(value)
|
||||
|
||||
logger.debug(f"Loaded configuration from {filepath}")
|
||||
return config
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error loading config from {filepath}: {e}")
|
||||
return {}
|
||||
|
||||
|
||||
def _parse_value(value: str) -> Any:
|
||||
value = value.strip()
|
||||
|
||||
if value.lower() == 'true':
|
||||
return True
|
||||
if value.lower() == 'false':
|
||||
return False
|
||||
|
||||
if value.isdigit():
|
||||
return int(value)
|
||||
|
||||
try:
|
||||
return float(value)
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
return value
|
||||
|
||||
|
||||
def create_default_config(filepath: str = CONFIG_FILE):
|
||||
default_config = """[api]
|
||||
default_model = x-ai/grok-code-fast-1
|
||||
timeout = 30
|
||||
temperature = 0.7
|
||||
max_tokens = 8096
|
||||
|
||||
[autonomous]
|
||||
max_iterations = 50
|
||||
context_threshold = 30
|
||||
recent_messages_to_keep = 10
|
||||
|
||||
[ui]
|
||||
syntax_highlighting = true
|
||||
show_timestamps = false
|
||||
color_output = true
|
||||
|
||||
[output]
|
||||
format = text
|
||||
verbose = false
|
||||
quiet = false
|
||||
|
||||
[session]
|
||||
auto_save = false
|
||||
max_history = 1000
|
||||
"""
|
||||
|
||||
try:
|
||||
with open(filepath, 'w') as f:
|
||||
f.write(default_config)
|
||||
logger.info(f"Created default configuration at {filepath}")
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.error(f"Error creating config file: {e}")
|
||||
return False
|
||||
@@ -0,0 +1,289 @@
|
||||
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
|
||||
@@ -0,0 +1,278 @@
|
||||
import logging
|
||||
import json
|
||||
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 pr.agents import AgentManager
|
||||
from pr.memory import KnowledgeStore, ConversationMemory, FactExtractor
|
||||
from pr.core.advanced_context import AdvancedContextManager
|
||||
from pr.core.api import call_api
|
||||
from pr.tools.base import get_tools_definition
|
||||
|
||||
logger = logging.getLogger('pr')
|
||||
|
||||
class EnhancedAssistant:
|
||||
def __init__(self, base_assistant):
|
||||
self.base = base_assistant
|
||||
|
||||
if CACHE_ENABLED:
|
||||
self.api_cache = APICache(DB_PATH, API_CACHE_TTL)
|
||||
self.tool_cache = ToolCache(DB_PATH, TOOL_CACHE_TTL)
|
||||
else:
|
||||
self.api_cache = None
|
||||
self.tool_cache = None
|
||||
|
||||
self.workflow_storage = WorkflowStorage(DB_PATH)
|
||||
self.workflow_engine = WorkflowEngine(
|
||||
tool_executor=self._execute_tool_for_workflow,
|
||||
max_workers=WORKFLOW_EXECUTOR_MAX_WORKERS
|
||||
)
|
||||
|
||||
self.agent_manager = AgentManager(DB_PATH, self._api_caller_for_agent)
|
||||
|
||||
self.knowledge_store = KnowledgeStore(DB_PATH)
|
||||
self.conversation_memory = ConversationMemory(DB_PATH)
|
||||
self.fact_extractor = FactExtractor()
|
||||
|
||||
if ADVANCED_CONTEXT_ENABLED:
|
||||
self.context_manager = AdvancedContextManager(
|
||||
knowledge_store=self.knowledge_store,
|
||||
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]
|
||||
)
|
||||
|
||||
logger.info("Enhanced Assistant initialized with all features")
|
||||
|
||||
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:
|
||||
logger.debug(f"Tool cache hit for {tool_name}")
|
||||
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],
|
||||
}
|
||||
|
||||
if tool_name in func_map:
|
||||
result = func_map[tool_name](**arguments)
|
||||
|
||||
if self.tool_cache:
|
||||
content = result.get('content', '')
|
||||
try:
|
||||
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}'}
|
||||
|
||||
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,
|
||||
self.base.api_url,
|
||||
self.base.api_key,
|
||||
use_tools=False,
|
||||
tools=None,
|
||||
temperature=temperature,
|
||||
max_tokens=max_tokens,
|
||||
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
|
||||
)
|
||||
if cached_response:
|
||||
logger.debug("API cache hit")
|
||||
return cached_response
|
||||
|
||||
response = call_api(
|
||||
messages,
|
||||
self.base.model,
|
||||
self.base.api_url,
|
||||
self.base.api_key,
|
||||
self.base.use_tools,
|
||||
get_tools_definition(),
|
||||
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)
|
||||
self.api_cache.set(
|
||||
self.base.model, messages,
|
||||
0.7, 4096,
|
||||
response, token_count
|
||||
)
|
||||
|
||||
return response
|
||||
|
||||
def process_with_enhanced_context(self, user_message: str) -> str:
|
||||
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
|
||||
)
|
||||
|
||||
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'])
|
||||
entry = KnowledgeEntry(
|
||||
entry_id=entry_id,
|
||||
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()
|
||||
)
|
||||
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
|
||||
)
|
||||
|
||||
if self.base.verbose:
|
||||
logger.info(f"Enhanced context: {context_info}")
|
||||
|
||||
working_messages = enhanced_messages
|
||||
else:
|
||||
working_messages = self.base.messages
|
||||
|
||||
response = self.enhanced_call_api(working_messages)
|
||||
|
||||
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"
|
||||
|
||||
topics = self.fact_extractor.categorize_content(summary)
|
||||
self.conversation_memory.update_conversation_summary(
|
||||
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]:
|
||||
workflow = self.workflow_storage.load_workflow_by_name(workflow_name)
|
||||
|
||||
if not workflow:
|
||||
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
|
||||
)
|
||||
|
||||
return {
|
||||
'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:
|
||||
return self.agent_manager.create_agent(role_name, agent_id)
|
||||
|
||||
def agent_task(self, agent_id: str, task: str) -> Dict[str, Any]:
|
||||
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')
|
||||
return self.agent_manager.collaborate_agents(orchestrator_id, task, agent_roles)
|
||||
|
||||
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()
|
||||
|
||||
if self.tool_cache:
|
||||
stats['tool_cache'] = self.tool_cache.get_statistics()
|
||||
|
||||
return stats
|
||||
|
||||
def get_workflow_list(self) -> List[Dict[str, Any]]:
|
||||
return self.workflow_storage.list_workflows()
|
||||
|
||||
def get_agent_summary(self) -> Dict[str, Any]:
|
||||
return self.agent_manager.get_session_summary()
|
||||
|
||||
def get_knowledge_statistics(self) -> Dict[str, Any]:
|
||||
return self.knowledge_store.get_statistics()
|
||||
|
||||
def get_conversation_history(self, limit: int = 10) -> List[Dict[str, Any]]:
|
||||
return self.conversation_memory.get_recent_conversations(limit=limit)
|
||||
|
||||
def clear_caches(self):
|
||||
if self.api_cache:
|
||||
self.api_cache.clear_all()
|
||||
|
||||
if self.tool_cache:
|
||||
self.tool_cache.clear_all()
|
||||
|
||||
logger.info("All caches cleared")
|
||||
|
||||
def cleanup(self):
|
||||
if self.api_cache:
|
||||
self.api_cache.clear_expired()
|
||||
|
||||
if self.tool_cache:
|
||||
self.tool_cache.clear_expired()
|
||||
|
||||
self.agent_manager.clear_session()
|
||||
@@ -0,0 +1,44 @@
|
||||
class PRException(Exception):
|
||||
pass
|
||||
|
||||
|
||||
class APIException(PRException):
|
||||
pass
|
||||
|
||||
|
||||
class APIConnectionError(APIException):
|
||||
pass
|
||||
|
||||
|
||||
class APITimeoutError(APIException):
|
||||
pass
|
||||
|
||||
|
||||
class APIResponseError(APIException):
|
||||
pass
|
||||
|
||||
|
||||
class ConfigurationError(PRException):
|
||||
pass
|
||||
|
||||
|
||||
class ToolExecutionError(PRException):
|
||||
def __init__(self, tool_name: str, message: str):
|
||||
self.tool_name = tool_name
|
||||
super().__init__(f"Error executing tool '{tool_name}': {message}")
|
||||
|
||||
|
||||
class FileSystemError(PRException):
|
||||
pass
|
||||
|
||||
|
||||
class SessionError(PRException):
|
||||
pass
|
||||
|
||||
|
||||
class ContextError(PRException):
|
||||
pass
|
||||
|
||||
|
||||
class ValidationError(PRException):
|
||||
pass
|
||||
@@ -0,0 +1,46 @@
|
||||
import logging
|
||||
import os
|
||||
from logging.handlers import RotatingFileHandler
|
||||
from pr.config import LOG_FILE
|
||||
|
||||
|
||||
def setup_logging(verbose=False):
|
||||
log_dir = os.path.dirname(LOG_FILE)
|
||||
if log_dir and not os.path.exists(log_dir):
|
||||
os.makedirs(log_dir, exist_ok=True)
|
||||
|
||||
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
|
||||
)
|
||||
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'
|
||||
)
|
||||
file_handler.setFormatter(file_formatter)
|
||||
logger.addHandler(file_handler)
|
||||
|
||||
if verbose:
|
||||
console_handler = logging.StreamHandler()
|
||||
console_handler.setLevel(logging.INFO)
|
||||
console_formatter = logging.Formatter(
|
||||
'%(levelname)s: %(message)s'
|
||||
)
|
||||
console_handler.setFormatter(console_formatter)
|
||||
logger.addHandler(console_handler)
|
||||
|
||||
return logger
|
||||
|
||||
|
||||
def get_logger(name=None):
|
||||
if name:
|
||||
return logging.getLogger(f'pr.{name}')
|
||||
return logging.getLogger('pr')
|
||||
@@ -0,0 +1,146 @@
|
||||
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')
|
||||
|
||||
SESSIONS_DIR = os.path.expanduser("~/.assistant_sessions")
|
||||
|
||||
|
||||
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:
|
||||
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 {}
|
||||
}
|
||||
|
||||
with open(session_file, 'w') as f:
|
||||
json.dump(session_data, f, indent=2)
|
||||
|
||||
logger.info(f"Session saved: {name}")
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error saving session {name}: {e}")
|
||||
return False
|
||||
|
||||
def load_session(self, name: str) -> Optional[Dict]:
|
||||
try:
|
||||
session_file = os.path.join(SESSIONS_DIR, f"{name}.json")
|
||||
|
||||
if not os.path.exists(session_file):
|
||||
logger.warning(f"Session not found: {name}")
|
||||
return None
|
||||
|
||||
with open(session_file, 'r') as f:
|
||||
session_data = json.load(f)
|
||||
|
||||
logger.info(f"Session loaded: {name}")
|
||||
return session_data
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error loading session {name}: {e}")
|
||||
return None
|
||||
|
||||
def list_sessions(self) -> List[Dict]:
|
||||
sessions = []
|
||||
|
||||
try:
|
||||
for filename in os.listdir(SESSIONS_DIR):
|
||||
if filename.endswith('.json'):
|
||||
filepath = os.path.join(SESSIONS_DIR, filename)
|
||||
try:
|
||||
with open(filepath, 'r') 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', {})
|
||||
})
|
||||
except Exception as e:
|
||||
logger.warning(f"Error reading session file {filename}: {e}")
|
||||
|
||||
sessions.sort(key=lambda x: x['created_at'], reverse=True)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error listing sessions: {e}")
|
||||
|
||||
return sessions
|
||||
|
||||
def delete_session(self, name: str) -> bool:
|
||||
try:
|
||||
session_file = os.path.join(SESSIONS_DIR, f"{name}.json")
|
||||
|
||||
if not os.path.exists(session_file):
|
||||
logger.warning(f"Session not found: {name}")
|
||||
return False
|
||||
|
||||
os.remove(session_file)
|
||||
logger.info(f"Session deleted: {name}")
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error deleting session {name}: {e}")
|
||||
return False
|
||||
|
||||
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:
|
||||
json.dump(session_data, f, indent=2)
|
||||
|
||||
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', '')
|
||||
|
||||
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:
|
||||
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', '')
|
||||
|
||||
f.write(f"[{role.upper()}]\n")
|
||||
f.write(f"{content}\n")
|
||||
f.write("-" * 80 + "\n\n")
|
||||
|
||||
else:
|
||||
logger.error(f"Unsupported export format: {format}")
|
||||
return False
|
||||
|
||||
logger.info(f"Session exported to {output_path}")
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error exporting session: {e}")
|
||||
return False
|
||||
@@ -0,0 +1,162 @@
|
||||
import json
|
||||
import os
|
||||
from datetime import datetime
|
||||
from typing import Dict, Optional
|
||||
from pr.core.logging import get_logger
|
||||
|
||||
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},
|
||||
}
|
||||
|
||||
|
||||
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': {}
|
||||
}
|
||||
|
||||
def track_request(
|
||||
self,
|
||||
model: str,
|
||||
input_tokens: int,
|
||||
output_tokens: int,
|
||||
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
|
||||
|
||||
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
|
||||
|
||||
cost = self._calculate_cost(model, input_tokens, output_tokens)
|
||||
model_usage['cost'] += cost
|
||||
self.session_usage['estimated_cost'] += cost
|
||||
|
||||
self._save_to_history(model, input_tokens, output_tokens, cost)
|
||||
|
||||
logger.debug(
|
||||
f"Tracked request: {model}, tokens: {total_tokens}, cost: ${cost:.4f}"
|
||||
)
|
||||
|
||||
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
|
||||
if base_model not in MODEL_COSTS:
|
||||
logger.warning(f"Unknown model for cost calculation: {model}")
|
||||
return 0.0
|
||||
costs = MODEL_COSTS[base_model]
|
||||
else:
|
||||
costs = MODEL_COSTS[model]
|
||||
|
||||
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):
|
||||
try:
|
||||
history = []
|
||||
if os.path.exists(USAGE_DB_FILE):
|
||||
with open(USAGE_DB_FILE, 'r') 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
|
||||
})
|
||||
|
||||
if len(history) > 10000:
|
||||
history = history[-10000:]
|
||||
|
||||
with open(USAGE_DB_FILE, 'w') as f:
|
||||
json.dump(history, f, indent=2)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error saving usage history: {e}")
|
||||
|
||||
def get_session_summary(self) -> Dict:
|
||||
return self.session_usage.copy()
|
||||
|
||||
def get_formatted_summary(self) -> str:
|
||||
usage = self.session_usage
|
||||
lines = [
|
||||
"\n=== Session Usage Summary ===",
|
||||
f"Total Requests: {usage['requests']}",
|
||||
f"Total Tokens: {usage['total_tokens']:,}",
|
||||
f" Input: {usage['input_tokens']:,}",
|
||||
f" Output: {usage['output_tokens']:,}",
|
||||
f"Estimated Cost: ${usage['estimated_cost']:.4f}",
|
||||
]
|
||||
|
||||
if usage['models_used']:
|
||||
lines.append("\nModels Used:")
|
||||
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)
|
||||
|
||||
@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
|
||||
}
|
||||
|
||||
try:
|
||||
with open(USAGE_DB_FILE, 'r') 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)
|
||||
|
||||
return {
|
||||
'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
|
||||
}
|
||||
@@ -0,0 +1,86 @@
|
||||
import os
|
||||
from typing import Optional
|
||||
from pr.core.exceptions import ValidationError
|
||||
|
||||
|
||||
def validate_file_path(path: str, must_exist: bool = False) -> str:
|
||||
if not path:
|
||||
raise ValidationError("File path cannot be empty")
|
||||
|
||||
if must_exist and not os.path.exists(path):
|
||||
raise ValidationError(f"File does not exist: {path}")
|
||||
|
||||
if must_exist and os.path.isdir(path):
|
||||
raise ValidationError(f"Path is a directory, not a file: {path}")
|
||||
|
||||
return os.path.abspath(path)
|
||||
|
||||
|
||||
def validate_directory_path(path: str, must_exist: bool = False, create: bool = False) -> str:
|
||||
if not path:
|
||||
raise ValidationError("Directory path cannot be empty")
|
||||
|
||||
abs_path = os.path.abspath(path)
|
||||
|
||||
if must_exist and not os.path.exists(abs_path):
|
||||
if create:
|
||||
os.makedirs(abs_path, exist_ok=True)
|
||||
else:
|
||||
raise ValidationError(f"Directory does not exist: {abs_path}")
|
||||
|
||||
if os.path.exists(abs_path) and not os.path.isdir(abs_path):
|
||||
raise ValidationError(f"Path is not a directory: {abs_path}")
|
||||
|
||||
return abs_path
|
||||
|
||||
|
||||
def validate_model_name(model: str) -> str:
|
||||
if not model:
|
||||
raise ValidationError("Model name cannot be empty")
|
||||
|
||||
if len(model) < 2:
|
||||
raise ValidationError("Model name too short")
|
||||
|
||||
return model
|
||||
|
||||
|
||||
def validate_api_url(url: str) -> str:
|
||||
if not url:
|
||||
raise ValidationError("API URL cannot be empty")
|
||||
|
||||
if not url.startswith(('http://', 'https://')):
|
||||
raise ValidationError("API URL must start with http:// or https://")
|
||||
|
||||
return url
|
||||
|
||||
|
||||
def validate_session_name(name: str) -> str:
|
||||
if not name:
|
||||
raise ValidationError("Session name cannot be empty")
|
||||
|
||||
invalid_chars = ['/', '\\', ':', '*', '?', '"', '<', '>', '|']
|
||||
for char in invalid_chars:
|
||||
if char in name:
|
||||
raise ValidationError(f"Session name contains invalid character: {char}")
|
||||
|
||||
if len(name) > 255:
|
||||
raise ValidationError("Session name too long (max 255 characters)")
|
||||
|
||||
return name
|
||||
|
||||
|
||||
def validate_temperature(temp: float) -> float:
|
||||
if not 0.0 <= temp <= 2.0:
|
||||
raise ValidationError("Temperature must be between 0.0 and 2.0")
|
||||
|
||||
return temp
|
||||
|
||||
|
||||
def validate_max_tokens(tokens: int) -> int:
|
||||
if tokens < 1:
|
||||
raise ValidationError("Max tokens must be at least 1")
|
||||
|
||||
if tokens > 100000:
|
||||
raise ValidationError("Max tokens too high (max 100000)")
|
||||
|
||||
return tokens
|
||||
Reference in New Issue
Block a user