edith-pc 0.3.0__py3-none-any.whl

edith/body/cli.py

@@ -0,0 +1,224 @@
+"""The `edith` command: this PC as a body of your EDITH server, installed with pip.
+
+    pip install edith-pc
+    edith link        link this PC with a code from Google Authenticator; it then runs in
+                      the background and starts with Windows
+    edith status      what this PC is linked to, and whether EDITH is running
+    edith start       start the background body      edith stop     stop it
+    edith app         the EDITH window (the same one the installer's app shows)
+    edith run         the body in this terminal (Ctrl+C stops it)
+    edith update      install the newest EDITH (from PyPI, or wherever pip got it) and restart
+    edith unlink      unlink this PC and stop starting with Windows
+
+No token, key or .env to copy: the code from the authenticator is the only thing typed.
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import subprocess
+import sys
+import time
+
+from .. import __version__
+from . import desktop
+from . import link as linking
+
+
+def _short(server: str) -> str:
+    return server.replace("https://", "")
+
+
+def _wait_linked(seconds: float = 20.0) -> dict:
+    """The background body's status once it says linked (or gives up), for a one-line report."""
+    end = time.time() + seconds
+    st: dict = {}
+    while time.time() < end:
+        st = desktop.read_status()
+        if desktop.body_running() and st.get("state") in ("linked", "refused", "unlinked"):
+            break
+        time.sleep(0.5)
+    return st
+
+
+def _report_start() -> int:
+    st = _wait_linked()
+    if st.get("state") == "linked":
+        print(f"EDITH is running in the background: {st.get('text')}")
+        return 0
+    print(f"EDITH started, but it says: {st.get('text') or 'nothing yet'}. Check with: edith status")
+    return 1
+
+
+def cmd_link(args) -> int:
+    default = linking.default_server()
+    server = args.server
+    while not server:   # the address first: no point asking for a code that expires meanwhile
+        typed = input(f"EDITH server [{_short(default)}]: " if default else "EDITH server (e.g. edith.example.com): ").strip()
+        try:
+            server = linking.normalize_server(typed or default)
+        except linking.LinkError as exc:
+            print(exc)
+    code = args.code or input("Code from Google Authenticator (6 digits): ").strip()
+    desktop.stop_background()   # one still running must not overwrite the new key when it stops
+    try:
+        made = linking.link(server, code)
+    except linking.LinkError as exc:
+        print(f"Not linked: {exc}")
+        return 1
+    print(f"Linked to {_short(made['server'])} as {made.get('name') or made['device_id']}.")
+    if not args.no_autostart:
+        desktop.set_autostart(True)
+        print("EDITH starts with Windows from now on." if desktop.autostart() else
+              "EDITH couldn't be set to start with Windows; start it with: edith start")
+    desktop.start_background()
+    return _report_start()
+
+
+def cmd_status(args) -> int:
+    cfg = linking.load()
+    if not linking.linked(cfg):
+        print("This PC is not linked. Link it with: edith link")
+        return 1
+    running, st = desktop.body_running(), desktop.read_status()
+    print(f"Linked to {_short(str(cfg.get('server')))} as {cfg.get('name') or cfg.get('device_id')}.")
+    print(f"Background body: {'running, ' + str(st.get('state')) if running else 'not running (start it: edith start)'}")
+    if running and st.get("text"):
+        print(f"  {st['text']}")
+    command = desktop.autostart()
+    started_by = "no" if not command else "yes" if desktop.autostart_is_python(command) else "yes, by the EDITH for Windows app"
+    print(f"Starts with Windows: {started_by}")
+    return 0
+
+
+def cmd_start(args) -> int:
+    if not linking.linked():
+        print("This PC is not linked. Link it with: edith link")
+        return 1
+    if desktop.body_running():
+        print("EDITH is already running in the background.")
+        return 0
+    desktop.start_background()
+    return _report_start()
+
+
+def cmd_stop(args) -> int:
+    if not desktop.body_running():
+        print("EDITH is not running.")
+        return 0
+    ok = desktop.stop_background()
+    print("Stopped." if ok else "It didn't stop; end the pythonw process in Task Manager.")
+    return 0 if ok else 1
+
+
+def cmd_unlink(args) -> int:
+    desktop.stop_background()
+    had = linking.linked()
+    dropped = linking.unlink()
+    desktop.set_autostart(False)
+    if not had:
+        print("This PC wasn't linked.")
+    elif dropped:
+        print("Unlinked.")
+    else:
+        print("Unlinked here, but the server couldn't be told: remove this PC in EDITH's HUD (Bodies).")
+    return 0
+
+
+DIST = "edith-pc"   # the name on PyPI ("edith" there is someone else's project)
+
+
+def _source() -> str | None:
+    """What to give pip to get the newest of this EDITH: the package's name when it came from
+    PyPI, else where pip got it (PEP 610's direct_url.json); None for a checkout."""
+    from importlib import metadata
+
+    try:
+        raw = metadata.distribution(DIST).read_text("direct_url.json")
+    except metadata.PackageNotFoundError:
+        return None
+    if not raw:
+        return DIST   # installed from PyPI
+    try:
+        info = json.loads(raw)
+    except ValueError:
+        return None
+    url = str(info.get("url") or "")
+    if "vcs_info" in info and info["vcs_info"].get("vcs") == "git":
+        return f"{DIST} @ git+{url}"
+    if url.startswith(("http://", "https://")):
+        return f"{DIST} @ {url}"
+    return None   # a local folder or an editable checkout
+
+
+def cmd_update(args) -> int:
+    source = _source()
+    if source is None:
+        print("This EDITH wasn't installed from a link pip can fetch again (a checkout, or a local "
+              "folder). Update it there: git pull.")
+        return 1
+    was_running = desktop.body_running()
+    if was_running:
+        desktop.stop_background()
+    pip = [sys.executable, "-m", "pip", "install", "--disable-pip-version-check"]
+    if source == DIST:
+        print("Updating from PyPI ...")
+        steps = (pip + ["--upgrade", DIST],)
+    else:
+        print(f"Updating from {source.split('@', 1)[1].strip()} ...")
+        # the newest code even when its version number didn't change, then anything new it needs
+        steps = (pip + ["--upgrade", "--force-reinstall", "--no-deps", source], pip + [source])
+    for step in steps:
+        if subprocess.call(step) != 0:
+            print("The update failed; the installed EDITH is unchanged or partly updated. Try again.")
+            if was_running:
+                desktop.start_background()
+            return 1
+    print("Updated.")
+    if was_running:
+        desktop.start_background()
+        return _report_start()
+    return 0
+
+
+def cmd_app(args) -> int:
+    return desktop.main([])
+
+
+def cmd_run(args) -> int:
+    from .__main__ import main as body
+
+    return body([])
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = argparse.ArgumentParser(prog="edith", description="EDITH on this PC: a body of your EDITH server.",
+                                     epilog="Start with: edith link")
+    parser.add_argument("--version", action="version", version=f"EDITH {__version__}")
+    sub = parser.add_subparsers(dest="command")
+    p = sub.add_parser("link", help="link this PC with a code from Google Authenticator")
+    p.add_argument("--server", help="the EDITH server's address (default: the one EDITH ships with)")
+    p.add_argument("--code", help="the 6-digit code (default: asked for)")
+    p.add_argument("--no-autostart", action="store_true", help="don't start EDITH with Windows")
+    p.set_defaults(func=cmd_link)
+    for name, func, text in (("status", cmd_status, "what this PC is linked to, and whether EDITH runs"),
+                             ("start", cmd_start, "start the background body"),
+                             ("stop", cmd_stop, "stop the background body"),
+                             ("unlink", cmd_unlink, "unlink this PC and stop starting with Windows"),
+                             ("update", cmd_update, "install the newest EDITH and restart it"),
+                             ("app", cmd_app, "open the EDITH window"),
+                             ("run", cmd_run, "run the body in this terminal")):
+        sub.add_parser(name, help=text).set_defaults(func=func)
+    args = parser.parse_args(sys.argv[1:] if argv is None else argv)
+    if not getattr(args, "func", None):
+        parser.print_help()
+        return 0 if linking.linked() else 1
+    try:
+        return int(args.func(args) or 0)
+    except (KeyboardInterrupt, EOFError):
+        print()
+        return 130
+
+
+if __name__ == "__main__":
+    sys.exit(main())

edith/body/client.py

@@ -0,0 +1,371 @@
+"""The body client: one WebSocket to the server's /ws, tool calls in, results out."""
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+import os
+import re
+import time
+from pathlib import Path
+from typing import Callable
+
+from websockets.asyncio.client import connect
+from websockets.exceptions import ConnectionClosed
+
+from ..tools.registry import SERVER_TOOLS, run_tool, tool_schemas
+from .identity import device_info
+
+_log = logging.getLogger("edith.body")
+MAX_MESSAGE = 8 * 1024 * 1024
+MAX_PARALLEL_TOOLS = 4
+REPLACED = 4409  # the server closed us because a newer connection of this PC arrived
+UNLINKED = 4403  # the server closed us because this PC was removed (its own key stopped working)
+
+
+def point_model_calls_at_server(server_url: str, token: str) -> None:
+    """Every tool builds its OpenAI client from the environment: send those calls to the
+    server's /v1 with the pairing token (or this PC's own key), so the real key never has
+    to be on this PC."""
+    os.environ["OPENAI_BASE_URL"] = server_url.rstrip("/") + "/v1"
+    os.environ["OPENAI_API_KEY"] = token
+
+
+def pairing_token(settings) -> str:
+    """EDITH_SERVICE_TOKEN, else ~/.edith_service_token: the token the server also has."""
+    tok = (getattr(settings, "service_token", None) or "").strip()
+    if tok:
+        return tok
+    try:
+        return (Path.home() / ".edith_service_token").read_text(encoding="utf-8").strip()
+    except OSError:
+        return ""
+
+
+class VoiceTokenError(RuntimeError):
+    """The server would not mint a voice session (402: its daily budget is used up)."""
+
+    def __init__(self, status: int, detail: str) -> None:
+        super().__init__(detail)
+        self.status = status
+        self.detail = detail
+
+
+def _auth(token: str) -> dict:
+    return {"Authorization": f"Bearer {token}", "X-Requested-With": "edith"}
+
+
+async def mint_voice_token(server_url: str, token: str, voice: str | None = None,
+                           transport=None) -> dict:
+    """POST /token on the server: a short-lived Realtime client secret for this PC's voice
+    HUD, so no OpenAI key is needed here. Raises VoiceTokenError with the server's reason."""
+    import httpx
+
+    url = server_url.rstrip("/") + "/token"
+    try:
+        async with httpx.AsyncClient(timeout=20.0, trust_env=False, transport=transport) as http:
+            r = await http.post(url, json={"voice": voice, "body": "pc"}, headers=_auth(token))
+    except httpx.HTTPError as exc:  # unreachable, timed out: a reason to retry, not to crash
+        raise VoiceTokenError(0, f"the server is unreachable ({type(exc).__name__})") from exc
+    if r.status_code != 200:
+        try:
+            detail = (r.json() or {}).get("detail") or r.text
+        except ValueError:
+            detail = r.text
+        raise VoiceTokenError(r.status_code, str(detail)[:200] or f"HTTP {r.status_code}")
+    try:
+        data = r.json()
+    except ValueError as exc:
+        raise VoiceTokenError(502, "the server's answer was not JSON") from exc
+    if not isinstance(data, dict) or not data.get("value"):
+        raise VoiceTokenError(502, "the server returned no secret")
+    return data
+
+
+async def report_voice_usage(server_url: str, token: str, model: str, usage: dict,
+                             transport=None) -> dict | None:
+    """POST /usage: the server books this session's usage on its daily budget and answers
+    with its spend {session, day, cap}; None when it could not be told."""
+    import httpx
+
+    url = server_url.rstrip("/") + "/usage"
+    try:
+        async with httpx.AsyncClient(timeout=15.0, trust_env=False, transport=transport) as http:
+            r = await http.post(url, json={"model": model, "usage": usage}, headers=_auth(token))
+    except httpx.HTTPError:
+        return None
+    if r.status_code != 200:
+        return None
+    try:
+        spend = (r.json() or {}).get("spend")
+    except ValueError:
+        return None
+    return spend if isinstance(spend, dict) else None
+
+
+_BACKUP_NAME = re.compile(r"^edith-backup-\d{4}-\d{2}-\d{2}T\d{6}Z\.tar\.gz$")
+
+
+async def fetch_server_backup(server_url: str, token: str, dest: Path, keep: int = 14,
+                              transport=None) -> Path | None:
+    """The server brain's newest backup, into ``dest``: an off-server copy on this PC. None
+    when the server has none yet or can't be reached. The newest ``keep`` copies are kept."""
+    import httpx
+
+    dest.mkdir(parents=True, exist_ok=True)
+    url = server_url.rstrip("/") + "/backup/latest"
+    target = None
+    try:
+        async with httpx.AsyncClient(timeout=120.0, trust_env=False, transport=transport) as http:
+            async with http.stream("GET", url, headers=_auth(token)) as r:
+                if r.status_code != 200:
+                    return None
+                m = re.search(r'filename="?([^";]+)"?', r.headers.get("content-disposition", ""))
+                name = Path(m.group(1)).name if m else ""
+                if not _BACKUP_NAME.match(name):   # never a path or a name the server made up
+                    name = time.strftime("edith-backup-%Y-%m-%dT%H%M%SZ.tar.gz", time.gmtime())
+                target = dest / name
+                if not target.exists():
+                    part = target.with_name(name + ".part")
+                    with open(part, "wb") as f:
+                        async for chunk in r.aiter_bytes():
+                            f.write(chunk)
+                    part.replace(target)
+    except (httpx.HTTPError, OSError) as exc:
+        _log.warning("backup copy failed: %s", exc)
+        return None
+    for old in sorted(dest.glob("edith-backup-*.tar.gz"), key=lambda p: p.name, reverse=True)[keep:]:
+        try:
+            old.unlink()
+        except OSError:
+            pass
+    return target
+
+
+def ws_url(server_url: str) -> str:
+    url = server_url.strip().rstrip("/")
+    if url.startswith("https://"):
+        url = "wss://" + url[len("https://"):]
+    elif url.startswith("http://"):
+        url = "ws://" + url[len("http://"):]
+    elif not url.startswith(("ws://", "wss://")):
+        url = "wss://" + url
+    return url + "/ws"
+
+
+def pc_tool_names() -> list[str]:
+    """What this PC offers: every PC tool the server does not run itself."""
+    return [t["name"] for t in tool_schemas("pc") if t["name"] not in SERVER_TOOLS]
+
+
+class Body:
+    def __init__(self, server_url: str, token: str, *, say: Callable[[str], None] = print,
+                 run=None, backup_dir: Path | None = None, pc_backup: bool = False,
+                 status_path: Path | None = None, on_unlinked: Callable[[], None] | None = None) -> None:
+        self.server_url = server_url
+        self.url = ws_url(server_url)
+        self.device_id = device_info()["id"]   # until the server's welcome says otherwise
+        # which standard folders the HUD chose to back up (None: not told yet, so nothing)
+        self.backup_choice: list[str] | None = None
+        self.backup_settle = 300.0
+        self._backup_wake = asyncio.Event()
+        # the app's window shows what this body is doing: {state, text, at, pid, server, device}
+        self.status_path = status_path
+        # this PC was removed in the HUD (its own key stopped working): told once, then we stop
+        self.on_unlinked = on_unlinked
+        self._unlinked = False
+        # where to keep the daily copy of the server brain's backup (None: don't fetch)
+        self.backup_dir = backup_dir
+        self._backup_task: asyncio.Task | None = None
+        # back up this PC's folders into the server's vault, nightly (edith/body/pc_backup.py)
+        self.pc_backup = pc_backup
+        self._pc_backup_task: asyncio.Task | None = None
+        self.token = token
+        self._say = say
+        self._run = run or (lambda name, args: run_tool(name, args, role="pc"))
+        self._said = ""
+        self._limit = asyncio.Semaphore(MAX_PARALLEL_TOOLS)
+        # After being replaced (a second python -m edith.body on this PC?) this copy yields:
+        # every replaced_wait it asks for the link WITHOUT taking it, and gets it back once
+        # the other copy has stopped. A fresh start (or a link that came back) takes over,
+        # which is what replaces a dead connection the server has not noticed yet.
+        self.replaced_wait = 60.0
+        self._yield = False
+
+    def _status(self, text: str, state: str | None = None) -> None:
+        if text != self._said:
+            self._said = text
+            self._say(f"  body: {text}")
+            _log.info("body: %s", text)
+        if state is not None:
+            self._state = state
+        if self.status_path is not None:
+            try:
+                tmp = self.status_path.with_name(self.status_path.name + ".tmp")
+                tmp.write_text(json.dumps({"state": getattr(self, "_state", "starting"), "text": text,
+                                           "at": int(time.time()), "pid": os.getpid(),
+                                           "server": self.server_url, "device": self.device_id}), encoding="utf-8")
+                os.replace(tmp, self.status_path)
+            except OSError:
+                pass
+
+    def _apply_config(self, config: dict | None) -> None:
+        """The server's word on this PC's backup: which standard folders go into the vault."""
+        folders = ((config or {}).get("backup") or {}).get("folders") if isinstance(config, dict) else None
+        if not isinstance(folders, list):
+            return
+        choice = [f for f in folders if isinstance(f, str)]
+        if choice != self.backup_choice:
+            self.backup_choice = choice
+            _log.info("backup: %s", ", ".join(choice) + " every night" if choice else "off for the standard folders")
+            self._backup_wake.set()
+
+    def _gone(self) -> None:
+        self._unlinked = True
+        self._status("this PC was removed from EDITH (in the HUD, or unlinked here); "
+                     "open EDITH on this PC and link it again with a code", "unlinked")
+        if self.on_unlinked is not None:
+            try:
+                self.on_unlinked()
+            except Exception as exc:
+                _log.warning("forgetting the link failed: %s", exc)
+
+    async def run_forever(self) -> None:
+        """Link and serve until stopped; returns only when this PC was removed (unlinked)."""
+        backoff = 1.0
+        self._status(f"connecting to {self.server_url}", "connecting")
+        while True:
+            try:
+                async with connect(self.url, max_size=MAX_MESSAGE, ping_interval=20, ping_timeout=20,
+                                   open_timeout=15, proxy=None) as ws:
+                    await ws.send(json.dumps({"type": "hello", "token": self.token, "kind": "pc",
+                                              "device": device_info(), "tools": pc_tool_names(),
+                                              "takeover": not self._yield}))
+                    first = json.loads(await asyncio.wait_for(ws.recv(), 15))
+                    if first.get("code") == "unlinked":
+                        self._gone()
+                        return
+                    if first.get("code") == "in_use":
+                        self._status("another copy of EDITH's body holds this PC's link; this copy waits "
+                                     f"and asks again every {int(self.replaced_wait)} s (stop one of them)", "waiting")
+                        await asyncio.sleep(self.replaced_wait)
+                        continue
+                    if first.get("type") != "welcome":
+                        self._status(f"the server refused this PC: {first.get('message', first.get('type'))}", "refused")
+                        if "token" in str(first.get("message", "")):
+                            await asyncio.sleep(60)  # a wrong token won't fix itself in a second
+                        continue
+                    backoff = 1.0
+                    self._yield = False
+                    if isinstance(first.get("device_id"), str) and first["device_id"]:
+                        self.device_id = first["device_id"]   # the id the server knows this PC by
+                    self._apply_config(first)
+                    if self.backup_dir is not None and self._backup_task is None:
+                        self._backup_task = asyncio.create_task(self._backup_copies())
+                    if self.pc_backup and self._pc_backup_task is None:
+                        self._pc_backup_task = asyncio.create_task(self._pc_backups())
+                    self._status(f"linked to the server as {self.device_id} "
+                                 f"({len(pc_tool_names())} PC tools offered)", "linked")
+                    code = await self._serve(ws)
+                    if code == UNLINKED or self._unlinked:
+                        self._gone()
+                        return
+                    if code == REPLACED:
+                        self._yield = True
+                        self._status("another connection took this PC's place (a second copy of EDITH's "
+                                     "body?); this copy waits until it stops, asking every "
+                                     f"{int(self.replaced_wait)} s", "waiting")
+                        await asyncio.sleep(self.replaced_wait)
+                        continue
+            except asyncio.CancelledError:
+                raise
+            except Exception as exc:
+                self._status(f"server unreachable ({exc or type(exc).__name__}); retrying", "unreachable")
+            await asyncio.sleep(backoff)
+            backoff = min(backoff * 2, 30.0)
+
+    async def _backup_copies(self) -> None:
+        """Once a day, the server's newest backup into backup_dir (a copy off the server)."""
+        while True:
+            newest = max((p.stat().st_mtime for p in self.backup_dir.glob("edith-backup-*.tar.gz")), default=0)
+            if time.time() - newest > 20 * 3600:
+                got = await fetch_server_backup(self.server_url, self.token, self.backup_dir)
+                if got is not None:
+                    _log.info("backup: server copy saved as %s", got)
+            await asyncio.sleep(6 * 3600)
+
+    async def _pc_backups(self) -> None:
+        """This PC's folders into the server's vault, once the HUD switched some on (or
+        EDITH_BACKUP_FOLDERS names some): at night, or within a day and a half when the PC
+        was off at night. The work runs in a thread; tool calls keep flowing."""
+        from . import pc_backup
+
+        await asyncio.sleep(self.backup_settle)   # let a fresh start settle
+        while True:
+            folders = pc_backup.chosen_folders(self.backup_choice)
+            if folders and pc_backup.due():
+                self._status(f"backup: backing up {len(folders)} folder(s) of this PC to the server")
+                try:
+                    report = await asyncio.to_thread(pc_backup.run_backup, self.server_url, self.token,
+                                                     folders=folders, device_id=self.device_id)
+                    done = sum(1 for s in report["sources"].values() if not s.get("error"))
+                    self._status(f"backup: {'done' if report.get('ok') else 'finished with problems'} "
+                                 f"({done}/{len(report['sources'])} sources)")
+                except Exception as exc:
+                    self._status(f"backup failed: {exc}")
+            try:   # every half hour, or at once when the HUD changes the choice
+                await asyncio.wait_for(self._backup_wake.wait(), timeout=1800)
+            except asyncio.TimeoutError:
+                pass
+            self._backup_wake.clear()
+
+    async def _serve(self, ws) -> int | None:
+        """Serve tool calls until the link closes; the close code, if the server sent one."""
+        tasks: set[asyncio.Task] = set()
+        lock = asyncio.Lock()
+
+        async def send(msg: dict) -> None:
+            async with lock:
+                await ws.send(json.dumps(msg, ensure_ascii=False, default=str))
+
+        async def handle_call(msg: dict) -> None:
+            name, args = str(msg.get("name", "")), msg.get("arguments", "{}")
+            async with self._limit:
+                try:
+                    output = await asyncio.to_thread(self._run, name, args if isinstance(args, str) else json.dumps(args))
+                except Exception as exc:  # never let a tool kill the link
+                    output = json.dumps({"error": f"{name} failed on the PC: {exc}"})
+            try:
+                await send({"type": "tool_result", "id": msg.get("id"), "output": output})
+            except ConnectionClosed:
+                pass
+
+        try:
+            async for raw in ws:
+                try:
+                    msg = json.loads(raw)
+                except ValueError:
+                    continue
+                if not isinstance(msg, dict):
+                    continue
+                mtype = msg.get("type")
+                if mtype == "tool_call":
+                    task = asyncio.create_task(handle_call(msg))
+                    tasks.add(task)
+                    task.add_done_callback(tasks.discard)
+                elif mtype == "ping":
+                    await send({"type": "pong"})
+                elif mtype == "config":
+                    self._apply_config(msg)
+                elif mtype == "error":
+                    if msg.get("code") == "unlinked":
+                        self._unlinked = True
+                    self._status(f"server says: {msg.get('message')}")
+        except ConnectionClosed as exc:
+            code = exc.rcvd.code if exc.rcvd else None
+            self._status(f"link closed ({code if code is not None else 'no close frame'}); reconnecting")
+            return code
+        finally:
+            for t in tasks:
+                t.cancel()
+        return None