Update pipeline/clients.py

pipeline/clients.py

@@ -272,9 +272,7 @@ def restore_placeholders(source: str, translated: str) -> str:
     return result
 
 
-@_gpu
-def translate(text: str) -> str:
-    """Translate one string with TranslateGemma (greedy)."""
+def _translate_text(text: str) -> str:
     messages = [
         {
             "role": "user",
@@ -297,8 +295,7 @@ def translate(text: str) -> str:
     return restore_placeholders(text, reply)
 
 
-@_gpu
-def adjust_tone(
+def _adjust_tone_text(
     source: str,
     draft: str,
     wiki: str,
@@ -338,6 +335,65 @@ def adjust_tone(
     return restore_placeholders(source, reply) if reply else draft
 
 
+@_gpu
+def translate(text: str) -> str:
+    """Translate one string with TranslateGemma (greedy)."""
+    return _translate_text(text)
+
+
+@_gpu
+def adjust_tone(
+    source: str,
+    draft: str,
+    wiki: str,
+    char_name: str,
+    context_lines: list[str] | None = None,
+) -> str:
+    """Rewrite a draft translation in the character's voice via Gemma 4."""
+    return _adjust_tone_text(source, draft, wiki, char_name, context_lines)
+
+
+@_gpu
+def translate_and_tone_items(
+    items: list[dict],
+    cache: dict,
+    wiki: str | None,
+    char_name: str | None,
+    context_window: int = 2,
+) -> tuple[list[dict], dict, dict]:
+    """Translate and tone a full review record in one GPU allocation."""
+    cache = dict(cache)
+    translated = 0
+    toned = 0
+
+    for item in items:
+        if item.get("mt") is not None:
+            continue
+        source = item["source"]
+        cached = cache.get(source)
+        if cached is not None:
+            item["mt"] = cached
+        else:
+            item["mt"] = _translate_text(source)
+            cache[source] = item["mt"]
+            translated += 1
+
+    if wiki and char_name:
+        sources = [item["source"] for item in items]
+        for i, item in enumerate(items):
+            if item.get("kind") != "dialogue" or item.get("toned") is not None:
+                continue
+            if item.get("mt") is None:
+                continue
+            context = sources[max(0, i - context_window) : i]
+            item["toned"] = _adjust_tone_text(
+                item["source"], item["mt"], wiki, char_name, context
+            )
+            toned += 1
+
+    return items, cache, {"translated": translated, "toned": toned}
+
+
 @_gpu
 def update_character_wiki(system_prompt: str, user_msg: str) -> str:
     """Update a character wiki via the same in-process Gemma runtime as tone pass."""