AI butler Jarvis
This commit is contained in:
fiatrete
@@ -0,0 +1,195 @@
|
||||
import asyncio
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
|
||||
from jarvis import CFG
|
||||
from jarvis.logger import logger
|
||||
from pathlib import Path
|
||||
|
||||
from aiohttp import web
|
||||
import socketio
|
||||
import socketio.exceptions
|
||||
|
||||
import importlib.util
|
||||
from importlib.machinery import SourceFileLoader
|
||||
from jarvis.gateway.session import Session, SioServerConnection, SioClientConnection
|
||||
from jarvis.utils.incoming_chat_message_parser import parse_incoming_chat_message
|
||||
|
||||
|
||||
def _import_external_functions():
|
||||
def import_recursive(path: str):
|
||||
files = os.listdir(path)
|
||||
no_subdir = False
|
||||
for file in files:
|
||||
if file.endswith(".module.py"):
|
||||
# If a module file exists, then it's the only module we are going to load
|
||||
full_path = os.path.join(path, file)
|
||||
# Add the module path
|
||||
sys.path.append(os.path.dirname(full_path))
|
||||
SourceFileLoader(full_path, full_path).load_module()
|
||||
no_subdir = True
|
||||
|
||||
if not no_subdir:
|
||||
# This is not the root of a module, let's dig in
|
||||
for file in files:
|
||||
full_path = os.path.join(path, file)
|
||||
if os.path.isdir(full_path):
|
||||
import_recursive(full_path)
|
||||
|
||||
import_recursive(CFG.external_function_module_dirs)
|
||||
|
||||
|
||||
def _import_functions():
|
||||
py_files = []
|
||||
dir_path = os.path.join(Path(__file__).parent, "functional_modules")
|
||||
for file in os.listdir(dir_path):
|
||||
if file.endswith(".py"):
|
||||
py_files.append(file)
|
||||
|
||||
for file in py_files:
|
||||
if file == "functional_module.py" or file == "caller_context.py":
|
||||
continue
|
||||
SourceFileLoader(file, os.path.join("jarvis/functional_modules", file)).load_module()
|
||||
|
||||
_import_external_functions()
|
||||
|
||||
|
||||
logger.info("Registering functions...")
|
||||
_import_functions()
|
||||
|
||||
|
||||
def run_server_mode():
|
||||
logger.info("Starting server...")
|
||||
|
||||
async def index(request):
|
||||
"""Serve the client-side application."""
|
||||
with open('./TestPage/index.html') as f:
|
||||
return web.Response(text=f.read(), content_type='text/html')
|
||||
|
||||
app = web.Application()
|
||||
session_map = {}
|
||||
|
||||
sio: socketio.AsyncServer = socketio.AsyncServer(
|
||||
max_http_buffer_size=50000000, # 50M
|
||||
)
|
||||
sio.attach(app)
|
||||
|
||||
@sio.event
|
||||
def connect(sid, environ):
|
||||
logger.debug(f"connect {sid}")
|
||||
session_map.update({sid: Session(SioServerConnection(sio, sid), sid)})
|
||||
|
||||
@sio.event
|
||||
async def disconnect(sid):
|
||||
logger.debug(f'disconnect {sid}')
|
||||
session: Session = session_map[sid]
|
||||
session_map.update({sid: None})
|
||||
await session.stop()
|
||||
|
||||
@sio.on('chat_message')
|
||||
async def chat_message(sid, data):
|
||||
logger.debug(f"message {data}")
|
||||
msg = parse_incoming_chat_message(data)
|
||||
if msg is None:
|
||||
return
|
||||
|
||||
session = session_map[sid]
|
||||
if session is None:
|
||||
logger.debug(f"Error: session {sid} not found!")
|
||||
return
|
||||
|
||||
if msg.message_type == 'clear':
|
||||
session.clear_history()
|
||||
elif msg.message_type == 'set_ts_offset':
|
||||
offset = int(msg.message_content)
|
||||
if offset > 12 or offset < -12:
|
||||
logger.error(f"Invalid tz offset: {msg.message_content}")
|
||||
return
|
||||
session.set_tz_offset(offset)
|
||||
elif msg.message_type == 'text':
|
||||
await session.on_chat_message(msg)
|
||||
|
||||
app.router.add_static('/js', './TestPage/js')
|
||||
app.router.add_static('/css', './TestPage/css')
|
||||
app.router.add_get('/', index)
|
||||
web.run_app(app, host='0.0.0.0', port=CFG.server_mode_port)
|
||||
|
||||
|
||||
async def run_client_mode(session_map: dict[str, Session]):
|
||||
sio = socketio.AsyncClient()
|
||||
# The connection is re-established, thus re-set sio of all sessions.
|
||||
for s in session_map.values():
|
||||
s.set_sio(SioClientConnection(sio))
|
||||
|
||||
# @sio.event
|
||||
@sio.on('connect')
|
||||
def connect():
|
||||
logger.debug(f"connected")
|
||||
|
||||
@sio.event
|
||||
def disconnect():
|
||||
logger.debug(f'disconnected')
|
||||
# Do nothing, sessions will not be proactively destoryed in this mode.
|
||||
|
||||
@sio.on('chat_message')
|
||||
async def chat_message(data):
|
||||
logger.debug(f"message {data}")
|
||||
msg = parse_incoming_chat_message(data)
|
||||
if msg is None:
|
||||
return
|
||||
sid = msg.chat_id
|
||||
if sid in session_map.keys():
|
||||
session = session_map[sid]
|
||||
assert session is not None
|
||||
else:
|
||||
session = Session(SioClientConnection(sio), sid)
|
||||
session_map.update({sid: session})
|
||||
|
||||
if msg.message_type == 'clear':
|
||||
session.clear_history()
|
||||
elif msg.message_type == 'set_ts_offset':
|
||||
offset = int(msg.message_content)
|
||||
if offset > 12 or offset < -12:
|
||||
logger.error(f"Invalid tz offset: {msg.message_content}")
|
||||
return
|
||||
session.set_tz_offset(offset)
|
||||
elif msg.message_type == 'text':
|
||||
await session.on_chat_message(msg)
|
||||
|
||||
await sio.connect(CFG.bot_server_url)
|
||||
try:
|
||||
await sio.wait()
|
||||
except:
|
||||
# I don't known why, but if we don't catch here, the logger.debug below will
|
||||
# die when the program is interrupted by SIGINT
|
||||
raise
|
||||
finally:
|
||||
del sio
|
||||
logger.debug("Client mode end")
|
||||
|
||||
|
||||
async def run_client_mode_async(session_map: dict[str, Session]):
|
||||
while True:
|
||||
try:
|
||||
await run_client_mode(session_map)
|
||||
except (InterruptedError, asyncio.CancelledError):
|
||||
logger.info(f"Interrupted, exit...")
|
||||
break
|
||||
except BaseException as e:
|
||||
logger.error(f"Failed to run in client mode, try again 1 seconds later: {str(e)}")
|
||||
await asyncio.sleep(1)
|
||||
|
||||
|
||||
def main():
|
||||
if CFG.is_server_mode:
|
||||
run_server_mode()
|
||||
else:
|
||||
session_map = {}
|
||||
asyncio.run(run_client_mode_async(session_map))
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
|
||||
logger.debug("End jarvis")
|
||||
Reference in New Issue
Block a user