Initial local translation API with persistent projects and batch jobs

local_translator/engine.py

@@ -0,0 +1,138 @@
+import asyncio
+from contextlib import suppress
+import hashlib
+import logging
+
+from .config import ProviderConfig
+from .models import Project
+from .protection import protect
+from .provider import Provider, ProviderError
+from .storage import encode
+
+log = logging.getLogger(__name__)
+PROMPT_VERSION = 2
+
+
+class Engine:
+    def __init__(self, storage, provider_factory=Provider):
+        self.storage = storage
+        self.provider_factory = provider_factory
+        self.wakeup = asyncio.Event()
+        self.task = None
+
+    def start(self):
+        self.storage.recover()
+        self.task = asyncio.create_task(self.run())
+
+    async def stop(self):
+        if self.task:
+            self.task.cancel()
+            with suppress(asyncio.CancelledError):
+                await self.task
+
+    async def run(self):
+        while True:
+            self.wakeup.clear()
+            identity = self.storage.next_job()
+            if identity:
+                try:
+                    await self.process(identity)
+                except Exception:
+                    log.exception("Translation job %s failed unexpectedly", identity)
+                    self.storage.job_status(identity, "failed", "Internal processing error; see server log")
+            else:
+                await self.wakeup.wait()
+
+    async def process(self, identity):
+        job = self.storage.get_job(identity, internal=True)
+        snapshot, options = job["snapshot"], job["request"]
+        project = Project.model_validate({k: v for k, v in snapshot["project"].items() if k not in {"id", "revision"}})
+        config = ProviderConfig.model_validate(snapshot["provider"])
+        provider = self.provider_factory(config)
+        self.storage.job_status(identity, "running")
+        try:
+            model = job["model"] or await provider.resolve_model(project.model)
+            self.storage.job_status(identity, "running", model=model)
+            fingerprint = hashlib.sha256(encode({"snapshot": snapshot, "model": model, "prompt_version": PROMPT_VERSION}).encode()).hexdigest()
+            pending = []
+            for item in job["items"]:
+                if item["status"] == "completed":
+                    continue
+                if not item["source"].strip():
+                    self.storage.item_result(identity, item["position"], "completed", item["source"])
+                    continue
+                if item["source"] in project.glossary:
+                    # A complete dictionary match is authoritative, even when
+                    # a model would mistake a short UI label for a code token.
+                    translated = project.glossary[item["source"]]
+                    original_tokens = list(protect(item["source"], project).tokens.values())
+                    translated_tokens = list(protect(translated, project).tokens.values())
+                    if not translated.strip() or original_tokens != translated_tokens:
+                        self.storage.item_result(identity, item["position"], "failed", error="Exact glossary entry changes protected syntax or has empty text")
+                    else:
+                        self.storage.item_result(identity, item["position"], "completed", translated)
+                        self.storage.cache_put(fingerprint, item["source"], translated)
+                    continue
+                cached = self.storage.cache_get(fingerprint, item["source"]) if options["use_cache"] else None
+                if cached is not None:
+                    self.storage.item_result(identity, item["position"], "completed", cached, cached=True)
+                else:
+                    pending.append(item)
+            while pending:
+                batch, size = [], 0
+                while pending and len(batch) < options["batch_size"]:
+                    item = pending[0]
+                    if batch and size + len(item["source"]) > options["max_batch_chars"]:
+                        break
+                    batch.append(pending.pop(0))
+                    size += len(item["source"])
+                await self.translate_batch(provider, model, project, job, fingerprint, batch)
+            result = self.storage.get_job(identity)
+            self.storage.job_status(identity, "completed_with_errors" if result["progress"]["failed"] else "completed")
+        except ProviderError as exc:
+            self.storage.job_status(identity, "failed", str(exc))
+        finally:
+            await provider.close()
+
+    async def translate_batch(self, provider, model, project, job, fingerprint, items):
+        identity = job["id"]
+        remaining = items
+        for attempt in range(job["request"]["retries"] + 1):
+            eligible = []
+            for item in remaining:
+                if item["attempts"] >= job["request"]["retries"] + 1:
+                    self.storage.item_result(identity, item["position"], "failed", error="Attempt budget exhausted; retry this job explicitly")
+                else:
+                    eligible.append(item)
+            remaining = eligible
+            if not remaining:
+                break
+            protected = {item["id"]: protect(item["source"], project) for item in remaining}
+            self.storage.attempted(identity, [item["position"] for item in remaining])
+            for item in remaining:
+                item["attempts"] += 1
+            error = "Model omitted the item or returned a duplicate ID"
+            try:
+                output = await provider.translate(model, project, [{"id": item["id"], "text": protected[item["id"]].text} for item in remaining])
+            except ProviderError as exc:
+                output, error = {}, str(exc)
+            failed = []
+            for item in remaining:
+                item_error = error
+                try:
+                    text = output[item["id"]]
+                    if not text.strip():
+                        raise ValueError("Model returned an empty translation")
+                    translated = protected[item["id"]].restore(text)
+                except (KeyError, ValueError) as exc:
+                    if isinstance(exc, ValueError):
+                        item_error = str(exc)
+                    failed.append(item)
+                    # Keep retryable rows pending until the attempt budget ends.
+                    self.storage.item_result(identity, item["position"], "failed" if item["attempts"] >= job["request"]["retries"] + 1 else "pending", error=item_error)
+                else:
+                    self.storage.cache_put(fingerprint, item["source"], translated)
+                    self.storage.item_result(identity, item["position"], "completed", translated)
+            remaining = failed
+            if remaining and attempt < job["request"]["retries"]:
+                await asyncio.sleep(min(2 ** attempt, 8))