225 lines
8.8 KiB
Python
225 lines
8.8 KiB
Python
|
"""Startup, CLI parsing, banner, REPL, and main entry."""
|
|||
|
|
|
|||
|
|
import argparse
|
|||
|
|
import os
|
|||
|
|
import sqlite3
|
|||
|
|
import sys
|
|||
|
|
import threading
|
|||
|
|
|
|||
|
|
from tai_pkg.agent import Agent
|
|||
|
|
from tai_pkg.agents import REGISTRY
|
|||
|
|
from tai_pkg.commands import Commands
|
|||
|
|
from tai_pkg.compat import readline
|
|||
|
|
from tai_pkg.config import Config
|
|||
|
|
from tai_pkg.consts import Consts
|
|||
|
|
from tai_pkg.errors import BackendError
|
|||
|
|
from tai_pkg.install import Install
|
|||
|
|
from tai_pkg.models import ModelScoring
|
|||
|
|
from tai_pkg.plexer import NativePlexerBackend
|
|||
|
|
from tai_pkg.release import Release
|
|||
|
|
from tai_pkg.render import Ansi
|
|||
|
|
from tai_pkg.render import Render
|
|||
|
|
from tai_pkg.sandbox import Sandbox
|
|||
|
|
from tai_pkg.scheduler import Scheduler
|
|||
|
|
from tai_pkg.seal import Seal
|
|||
|
|
from tai_pkg.seal import SealError
|
|||
|
|
from tai_pkg.seal import SealManager
|
|||
|
|
from tai_pkg.store import Store
|
|||
|
|
from tai_pkg.telegram import TelegramBot
|
|||
|
|
from tai_pkg.text import Text
|
|||
|
|
|
|||
|
|
__all__ = ["App"]
|
|||
|
|
|
|||
|
|
|
|||
|
|
class App:
|
|||
|
|
"""App."""
|
|||
|
|
|
|||
|
|
@staticmethod
|
|||
|
|
def boot(args):
|
|||
|
|
args.yes = args.yes or args.yolo or args.auto
|
|||
|
|
config = Config(args)
|
|||
|
|
passphrase, default_key = SealManager.resolve_passphrase()
|
|||
|
|
try:
|
|||
|
|
seal = Seal(config.home, passphrase)
|
|||
|
|
except SealError as exc:
|
|||
|
|
if default_key:
|
|||
|
|
raise
|
|||
|
|
try:
|
|||
|
|
old = Seal(config.home, Consts.DEFAULT_PASSPHRASE)
|
|||
|
|
except SealError:
|
|||
|
|
raise exc
|
|||
|
|
try:
|
|||
|
|
seal = SealManager.rotate_seal(config, old, passphrase)
|
|||
|
|
except (OSError, sqlite3.Error) as fail:
|
|||
|
|
raise SealError(str(fail))
|
|||
|
|
print("seal upgraded from default key to TAI_PASSPHRASE")
|
|||
|
|
seal.default_key = default_key and seal.enabled
|
|||
|
|
store = Store(config, seal)
|
|||
|
|
if store.schema_notes:
|
|||
|
|
print("vault schema upgraded: %s" % ", ".join(store.schema_notes))
|
|||
|
|
try:
|
|||
|
|
backed_up = Release.ensure_self_backup(config, store)
|
|||
|
|
except (OSError, sqlite3.Error) as exc:
|
|||
|
|
print("warning: self-backup failed: %s" % Text.short_error(exc))
|
|||
|
|
backed_up = None
|
|||
|
|
if backed_up is not None:
|
|||
|
|
print("self-backup: %s" % os.path.basename(backed_up))
|
|||
|
|
return Agent(config, store, yolo=args.yolo or args.auto, auto=args.auto)
|
|||
|
|
|
|||
|
|
@staticmethod
|
|||
|
|
def build_parser():
|
|||
|
|
parser = argparse.ArgumentParser(prog="tai", description="single-file autonomous agent, standard library only")
|
|||
|
|
parser.add_argument("--profile", default=None)
|
|||
|
|
parser.add_argument("--yes", action="store_true")
|
|||
|
|
parser.add_argument("--yolo", action="store_true")
|
|||
|
|
parser.add_argument("--auto", action="store_true")
|
|||
|
|
parser.add_argument("--continue", dest="resume", action="store_true")
|
|||
|
|
parser.add_argument("--env", default=None)
|
|||
|
|
parser.add_argument("--version", action="store_true")
|
|||
|
|
parser.add_argument("--telegram", action="store_true")
|
|||
|
|
parser.add_argument("--scheduler", action="store_true")
|
|||
|
|
parser.add_argument("--install", action="store_true")
|
|||
|
|
parser.add_argument("--install-telegram", action="store_true")
|
|||
|
|
parser.add_argument("--uninstall-telegram", action="store_true")
|
|||
|
|
parser.add_argument("--plexer-wrap", action="store_true", help="internal: invoked by the bashrc plexer hook, not meant to be run by hand")
|
|||
|
|
parser.add_argument("prompt", nargs="*")
|
|||
|
|
return parser
|
|||
|
|
|
|||
|
|
@staticmethod
|
|||
|
|
def banner(agent):
|
|||
|
|
print(Render.paint("tai %s · profile %s" % (Consts.VERSION, agent.profile), Ansi.BOLD, Ansi.CYAN))
|
|||
|
|
if agent.yolo:
|
|||
|
|
print(Render.paint("YOLO MODE: all approvals off", Ansi.BOLD, Ansi.RED))
|
|||
|
|
if agent.auto:
|
|||
|
|
print(Render.paint("AUTO MODE: never asks, researches instead", Ansi.BOLD, Ansi.YELLOW))
|
|||
|
|
seal_state = "off"
|
|||
|
|
if agent.store.seal.enabled:
|
|||
|
|
seal_state = "on (default key)" if agent.store.seal.default_key else "on (TAI_PASSPHRASE)"
|
|||
|
|
print(" seal: %s" % seal_state)
|
|||
|
|
for label, _chat_url, models_url, key, _model in agent.chat.backends():
|
|||
|
|
print(" %s: %s" % (label, ModelScoring.probe_backend(models_url, key)))
|
|||
|
|
if agent.store.load_turn_state(agent.profile, agent.bot).get("open"):
|
|||
|
|
print(Render.paint("interrupted turn available, restart with --continue to pick it up", Ansi.YELLOW))
|
|||
|
|
print(Render.paint("type /help for commands", Ansi.DIM))
|
|||
|
|
|
|||
|
|
@staticmethod
|
|||
|
|
def repl(agent):
|
|||
|
|
if readline is not None:
|
|||
|
|
try:
|
|||
|
|
readline.read_history_file(agent.config.history_path)
|
|||
|
|
except OSError:
|
|||
|
|
pass
|
|||
|
|
readline.set_completer(Commands.complete_command)
|
|||
|
|
readline.parse_and_bind("tab: complete")
|
|||
|
|
scheduler_stop = Scheduler.start_scheduler(agent.config, agent.store.seal)
|
|||
|
|
try:
|
|||
|
|
while True:
|
|||
|
|
try:
|
|||
|
|
busy = REGISTRY.running_agents()
|
|||
|
|
base = agent.profile + ("!" if agent.yolo else "")
|
|||
|
|
tag = "%s+%d" % (base, busy) if busy else base
|
|||
|
|
line = input(Render.paint("tai[%s]› " % tag, Ansi.BOLD, Ansi.GREEN))
|
|||
|
|
except EOFError:
|
|||
|
|
print()
|
|||
|
|
break
|
|||
|
|
except KeyboardInterrupt:
|
|||
|
|
print()
|
|||
|
|
continue
|
|||
|
|
text = line.strip()
|
|||
|
|
if not text:
|
|||
|
|
continue
|
|||
|
|
if text.startswith("/"):
|
|||
|
|
if not Commands.handle_command(agent, text):
|
|||
|
|
break
|
|||
|
|
continue
|
|||
|
|
try:
|
|||
|
|
agent.run_turn(text)
|
|||
|
|
except KeyboardInterrupt:
|
|||
|
|
print(Render.paint("\ninterrupted", Ansi.YELLOW))
|
|||
|
|
agent.checkpoint()
|
|||
|
|
finally:
|
|||
|
|
scheduler_stop.set()
|
|||
|
|
if readline is not None:
|
|||
|
|
try:
|
|||
|
|
readline.write_history_file(agent.config.history_path)
|
|||
|
|
except OSError:
|
|||
|
|
pass
|
|||
|
|
agent.store.save_session(agent.profile, agent.messages, agent.bot)
|
|||
|
|
|
|||
|
|
@staticmethod
|
|||
|
|
def main(argv=None):
|
|||
|
|
args = App.build_parser().parse_args(argv)
|
|||
|
|
if args.install:
|
|||
|
|
return Install.install_self()
|
|||
|
|
if args.install_telegram:
|
|||
|
|
return Install.install_telegram()
|
|||
|
|
if args.uninstall_telegram:
|
|||
|
|
return Install.uninstall_telegram()
|
|||
|
|
if args.version:
|
|||
|
|
print("tai %s" % Consts.VERSION)
|
|||
|
|
return 0
|
|||
|
|
if args.plexer_wrap:
|
|||
|
|
home = os.environ.get("TAI_HOME") or os.path.join(os.path.expanduser("~"), ".tai")
|
|||
|
|
profile = args.profile or "default"
|
|||
|
|
backend = NativePlexerBackend(home, profile)
|
|||
|
|
return backend.run()
|
|||
|
|
if args.profile is None:
|
|||
|
|
args.profile = "telegram" if args.telegram else "default"
|
|||
|
|
if not Consts.PROFILE_RE.match(args.profile):
|
|||
|
|
print("invalid profile name", file=sys.stderr)
|
|||
|
|
return 1
|
|||
|
|
try:
|
|||
|
|
agent = App.boot(args)
|
|||
|
|
except KeyboardInterrupt:
|
|||
|
|
print(Render.paint("\ninterrupted", Ansi.YELLOW))
|
|||
|
|
return 130
|
|||
|
|
except SealError as exc:
|
|||
|
|
print("seal error: %s" % exc, file=sys.stderr)
|
|||
|
|
return 2
|
|||
|
|
if args.env is not None:
|
|||
|
|
if args.env not in ("home", "sandbox"):
|
|||
|
|
print("use --env home or --env sandbox", file=sys.stderr)
|
|||
|
|
return 1
|
|||
|
|
if args.env == "sandbox":
|
|||
|
|
try:
|
|||
|
|
Sandbox.ensure_box()
|
|||
|
|
except BackendError as exc:
|
|||
|
|
print("sandbox unavailable: %s" % exc, file=sys.stderr)
|
|||
|
|
return 2
|
|||
|
|
agent.env = args.env
|
|||
|
|
if args.scheduler:
|
|||
|
|
print("scheduler online, press Ctrl-C to stop")
|
|||
|
|
try:
|
|||
|
|
Scheduler.run_scheduler_loop(agent.config, agent.store.seal, threading.Event())
|
|||
|
|
except KeyboardInterrupt:
|
|||
|
|
print()
|
|||
|
|
finally:
|
|||
|
|
agent.store.close()
|
|||
|
|
return 0
|
|||
|
|
if args.telegram:
|
|||
|
|
return TelegramBot.run_telegram_bot(agent)
|
|||
|
|
if args.prompt:
|
|||
|
|
try:
|
|||
|
|
agent.run_turn(" ".join(args.prompt))
|
|||
|
|
except KeyboardInterrupt:
|
|||
|
|
print(Render.paint("\ninterrupted", Ansi.YELLOW))
|
|||
|
|
return 130
|
|||
|
|
finally:
|
|||
|
|
try:
|
|||
|
|
agent.checkpoint()
|
|||
|
|
except Exception:
|
|||
|
|
pass
|
|||
|
|
try:
|
|||
|
|
agent.store.close()
|
|||
|
|
except Exception:
|
|||
|
|
pass
|
|||
|
|
return 0
|
|||
|
|
if args.resume and not args.telegram and not args.scheduler:
|
|||
|
|
Agent.resume_turn(agent)
|
|||
|
|
App.banner(agent)
|
|||
|
|
try:
|
|||
|
|
App.repl(agent)
|
|||
|
|
finally:
|
|||
|
|
agent.store.close()
|
|||
|
|
return 0
|