123 lines
4.7 KiB
Python
123 lines
4.7 KiB
Python
import json
|
|
import re
|
|
|
|
from colorama import Fore, Style, init
|
|
|
|
init(autoreset=True)
|
|
|
|
# DEBUG = True # puts full message outputs in the terminal
|
|
DEBUG = False # only dumps important messages in the terminal
|
|
|
|
|
|
async def internal_monologue(msg):
|
|
# ANSI escape code for italic is '\x1B[3m'
|
|
print(f'\x1B[3m{Fore.LIGHTBLACK_EX}💭 {msg}{Style.RESET_ALL}')
|
|
|
|
async def assistant_message(msg):
|
|
print(f'{Fore.YELLOW}{Style.BRIGHT}🤖 {Fore.YELLOW}{msg}{Style.RESET_ALL}')
|
|
|
|
async def memory_message(msg):
|
|
print(f'{Fore.LIGHTMAGENTA_EX}{Style.BRIGHT}🧠 {Fore.LIGHTMAGENTA_EX}{msg}{Style.RESET_ALL}')
|
|
|
|
async def system_message(msg):
|
|
print(f'{Fore.MAGENTA}{Style.BRIGHT}🖥️ [system] {Fore.MAGENTA}{msg}{Style.RESET_ALL}')
|
|
|
|
async def user_message(msg, raw=False):
|
|
if isinstance(msg, str):
|
|
if raw:
|
|
print(f'{Fore.GREEN}{Style.BRIGHT}🧑 {Fore.GREEN}{msg}{Style.RESET_ALL}')
|
|
return
|
|
else:
|
|
try:
|
|
msg_json = json.loads(msg)
|
|
except:
|
|
print(f"Warning: failed to parse user message into json")
|
|
print(f'{Fore.GREEN}{Style.BRIGHT}🧑 {Fore.GREEN}{msg}{Style.RESET_ALL}')
|
|
return
|
|
|
|
if msg_json['type'] == 'user_message':
|
|
msg_json.pop('type')
|
|
print(f'{Fore.GREEN}{Style.BRIGHT}🧑 {Fore.GREEN}{msg_json}{Style.RESET_ALL}')
|
|
elif msg_json['type'] == 'heartbeat':
|
|
if DEBUG:
|
|
msg_json.pop('type')
|
|
print(f'{Fore.GREEN}{Style.BRIGHT}💓 {Fore.GREEN}{msg_json}{Style.RESET_ALL}')
|
|
elif msg_json['type'] == 'system_message':
|
|
msg_json.pop('type')
|
|
print(f'{Fore.GREEN}{Style.BRIGHT}🖥️ {Fore.GREEN}{msg_json}{Style.RESET_ALL}')
|
|
else:
|
|
print(f'{Fore.GREEN}{Style.BRIGHT}🧑 {Fore.GREEN}{msg_json}{Style.RESET_ALL}')
|
|
|
|
async def function_message(msg):
|
|
|
|
if isinstance(msg, dict):
|
|
print(f'{Fore.RED}{Style.BRIGHT}⚡ [function] {Fore.RED}{msg}{Style.RESET_ALL}')
|
|
return
|
|
|
|
if msg.startswith('Success: '):
|
|
if DEBUG:
|
|
print(f'{Fore.RED}{Style.BRIGHT}⚡🟢 [function] {Fore.RED}{msg}{Style.RESET_ALL}')
|
|
elif msg.startswith('Error: '):
|
|
print(f'{Fore.RED}{Style.BRIGHT}⚡🔴 [function] {Fore.RED}{msg}{Style.RESET_ALL}')
|
|
elif msg.startswith('Running '):
|
|
if DEBUG:
|
|
print(f'{Fore.RED}{Style.BRIGHT}⚡ [function] {Fore.RED}{msg}{Style.RESET_ALL}')
|
|
else:
|
|
if 'memory' in msg:
|
|
match = re.search(r'Running (\w+)\(', msg)
|
|
if match:
|
|
function_name = match.group(1)
|
|
print(f'{Fore.RED}{Style.BRIGHT}⚡🧠 [function] {Fore.RED}updating memory with {function_name}{Style.RESET_ALL}')
|
|
else:
|
|
print(f"Warning: did not recognize function message")
|
|
print(f'{Fore.RED}{Style.BRIGHT}⚡ [function] {Fore.RED}{msg}{Style.RESET_ALL}')
|
|
elif 'send_message' in msg:
|
|
# ignore in debug mode
|
|
pass
|
|
else:
|
|
print(f'{Fore.RED}{Style.BRIGHT}⚡ [function] {Fore.RED}{msg}{Style.RESET_ALL}')
|
|
else:
|
|
print(f"Warning: did not recognize function message")
|
|
print(f'{Fore.RED}{Style.BRIGHT}⚡ [function] {Fore.RED}{msg}{Style.RESET_ALL}')
|
|
|
|
async def print_messages(message_sequence):
|
|
for msg in message_sequence:
|
|
role = msg['role']
|
|
content = msg['content']
|
|
|
|
if role == 'system':
|
|
await system_message(content)
|
|
elif role == 'assistant':
|
|
# Differentiate between internal monologue, function calls, and messages
|
|
if msg.get('function_call'):
|
|
if content is not None:
|
|
await internal_monologue(content)
|
|
await function_message(msg['function_call'])
|
|
# assistant_message(content)
|
|
else:
|
|
await internal_monologue(content)
|
|
elif role == 'user':
|
|
await user_message(content)
|
|
elif role == 'function':
|
|
await function_message(content)
|
|
else:
|
|
print(f'Unknown role: {content}')
|
|
|
|
async def print_messages_simple(message_sequence):
|
|
for msg in message_sequence:
|
|
role = msg['role']
|
|
content = msg['content']
|
|
|
|
if role == 'system':
|
|
await system_message(content)
|
|
elif role == 'assistant':
|
|
await assistant_message(content)
|
|
elif role == 'user':
|
|
await user_message(content, raw=True)
|
|
else:
|
|
print(f'Unknown role: {content}')
|
|
|
|
async def print_messages_raw(message_sequence):
|
|
for msg in message_sequence:
|
|
print(msg)
|