Translator client: full-URL arg, capability Hello, one-job-at-a-time

scripts/translator.py

@@ -12,19 +12,20 @@
 # ///
 """Pagerite translator service: translate site content with Seed-X-PPO-7B.
 
-Connects to a Pagerite server's translator WebSocket (``/_translate/<key>`` —
-deliberately outside /_api, the key is the access control; find it in the
-site settings, GET /_api/settings -> ``translate_key``), announces the
-languages it handles and translates whatever the server pushes: everything
-pending on connect, then deltas as the content changes. Results go back per
-job and the server stores them (docs/localization.md).
+Connects to a Pagerite server's translator WebSocket — the full URL
+including the access key (the admin finds it in the site settings,
+GET /_api/settings -> ``translate_key``) — and announces the languages the
+model CAN translate (capabilities). The server dispatches one single-item
+job at a time per connection, offered only in its configured target
+languages (``Data.translate_langs``) ∩ the announced capabilities; a
+dropped connection's in-flight item is simply re-offered
+(docs/localization.md). For parallelism, run multiple instances.
 
-Seed-X-PPO-7B (bf16, ~15 GB) is the only supported model; its 28 languages
-are the ceiling of what a site can announce through this script.
+Seed-X-PPO-7B (bf16, ~15 GB) is the only supported model.
 
 Usage:
-    uv run scripts/translator.py ws://localhost:8080 --key "..." --to finnish
-    uv run scripts/translator.py wss://example.com --key "..." --to fi german es
+    uv run scripts/translator.py ws://localhost:8410/_translate/KEY
+    uv run scripts/translator.py wss://example.com/_translate/KEY
 """
 
 import argparse
@@ -58,45 +59,37 @@ SEED_X_NAMES = {v: k for k, v in SEED_X_TAGS.items()}
 NOTE = ", preserving all Markdown formatting, URLs and code exactly unchanged"
 
 
-# The message structs below duplicate pagerite/translate.py 1:1: this script
-# runs in its own uv environment and cannot import the server package. The
+# The wire structs below duplicate pagerite/translate.py: this script runs
+# in its own uv environment and cannot import the server package. The
 # "type" tag selects the frame; bytes fields ride as base64.
 class Hello(msgspec.Struct, tag="hello"):
-    """Client greeting on connect: the target languages it handles."""
+    """Client greeting on connect: the language codes its model CAN produce
+    (capabilities). The server offers jobs only in the intersection with
+    its wanted target languages."""
 
     langs: list[str]
 
 
-class TransItem(msgspec.Struct):
-    """One fragment to translate: original Markdown (or a node title)."""
+class Job(msgspec.Struct, tag="job"):
+    """Server push: ONE fragment to translate. Exactly one job is in flight
+    per connection — the next arrives only after this one's Result."""
 
+    lang: str
     key: bytes  #: 9-byte chunk hash (base64 in the JSON frame)
     text: str
     path: str  #: article it came from ("" = front page), no leading slash
     kind: str  #: "chunk" | "title"
 
 
-class Job(msgspec.Struct, tag="job"):
-    """Server push: pending items for one language."""
+class Result(msgspec.Struct, tag="result"):
+    """Client reply: the translation of the connection's current Job
+    (must match its lang and key exactly)."""
 
     lang: str
-    items: list[TransItem]
-
-
-class TransResult(msgspec.Struct):
-    """One translated fragment."""
-
     key: bytes
     text: str
 
 
-class Result(msgspec.Struct, tag="result"):
-    """Client reply: a batch of translations for one language."""
-
-    lang: str
-    items: list[TransResult]
-
-
 def load_seed_x():
     t0 = time.monotonic()
     tokenizer = AutoTokenizer.from_pretrained(SEED_X)
@@ -119,34 +112,31 @@ def seed_x_chunk(tokenizer, model, text: str, target_lang: str, tag: str, source
 
 
 async def do_job(ws, job: Job, tokenizer, model) -> None:
-    """Translate every item of one job and send the results back as a batch."""
+    """Translate the job's one fragment and send the result back."""
     lang_name = SEED_X_NAMES[job.lang].capitalize()
-    results = []
-    for i, item in enumerate(job.items, 1):
-        # Deliberately blocking: nothing else needs the loop while a job is
+    # Deliberately blocking: nothing else needs the loop while the job is
     # being answered, and the reconnect loop recovers a dropped connection
-        # (results already stored stay stored; only the still-pending items
-        # are re-pushed).
-        text, tokens, dt = seed_x_chunk(tokenizer, model, item.text, lang_name, job.lang, note=NOTE)
-        results.append(TransResult(key=item.key, text=text))
-        print(f"[{job.lang} {i}/{len(job.items)} {item.kind} {item.path or '/'}: "
+    # (the in-flight item is simply re-offered).
+    text, tokens, dt = seed_x_chunk(tokenizer, model, job.text, lang_name, job.lang, note=NOTE)
+    print(f"[{job.lang} {job.kind} {job.path or '/'}: "
           f"{tokens} tokens in {dt:.1f}s = {tokens / dt:.1f} tok/s]", file=sys.stderr)
-    await ws.send(msgspec.json.encode(Result(lang=job.lang, items=results)).decode())
+    await ws.send(msgspec.json.encode(Result(lang=job.lang, key=job.key, text=text)).decode())
 
 
-async def serve(url: str, codes: list[str], tokenizer, model) -> None:
-    """Connect, announce languages, answer jobs; reconnect with backoff."""
+async def serve(url: str, tokenizer, model) -> None:
+    """Connect, announce capabilities, answer jobs; reconnect with backoff."""
     backoff = 1
     while True:
         try:
             async with websockets.connect(url) as ws:
                 backoff = 1
-                await ws.send(msgspec.json.encode(Hello(langs=codes)).decode())
-                print(f"[connected; translating: {', '.join(codes)}]", file=sys.stderr)
+                await ws.send(msgspec.json.encode(Hello(langs=sorted(SEED_X_NAMES))).decode())
+                print(f"[connected; announced {len(SEED_X_NAMES)} language capabilities]",
+                      file=sys.stderr)
                 async for raw in ws:
                     await do_job(ws, msgspec.json.decode(raw, type=Job), tokenizer, model)
         except websockets.exceptions.InvalidHandshake:
-            sys.exit("handshake rejected; check --key and the server URL")
+            sys.exit("handshake rejected; check the URL (including the key)")
         except (OSError, websockets.exceptions.ConnectionClosed) as e:
             print(f"[connection lost ({e}); reconnecting in {backoff}s]", file=sys.stderr)
             await asyncio.sleep(backoff)
@@ -155,24 +145,14 @@ async def serve(url: str, codes: list[str], tokenizer, model) -> None:
 
 def main():
     p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
-    p.add_argument("server", help="server WebSocket base URL, e.g. ws://localhost:8080")
-    p.add_argument("--key", required=True, help="translator API key (site settings: translate_key)")
-    p.add_argument("--to", required=True, nargs="+",
-                   help="target language(s): Seed-X names or codes, e.g. finnish de es")
+    p.add_argument("url", help="full translator WebSocket URL including the key, "
+                               "e.g. ws://localhost:8410/_translate/KEY")
     args = p.parse_args()
-
-    codes = []
-    for lang in args.to:
-        low = lang.lower()
-        code = SEED_X_TAGS.get(low) or (low if low in SEED_X_NAMES else None)
-        if code is None:
-            p.error(f"unknown language {lang!r}; supported: {', '.join(sorted(SEED_X_TAGS))}")
-        if code not in codes:
-            codes.append(code)
+    if not args.url.startswith(("ws://", "wss://")):
+        p.error("url must start with ws:// or wss://")
 
     tokenizer, model = load_seed_x()  # once, before the (re)connect loop
-    url = f"{args.server.rstrip('/')}/_translate/{args.key}"
-    asyncio.run(serve(url, codes, tokenizer, model))
+    asyncio.run(serve(args.url, tokenizer, model))
 
 
 if __name__ == "__main__":