1008 lines
39 KiB
Python
1008 lines
39 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
Rewrite Tab for DazedMTLTool GUI
|
||
|
||
Fixes over-limit game dialogue via LLM API.
|
||
Scans translated/*.json for messages exceeding 3 lines / 60 chars per line,
|
||
then rewrites them via the configured LLM while preserving all meaning.
|
||
"""
|
||
|
||
import os
|
||
import json
|
||
import time
|
||
import traceback
|
||
import threading
|
||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||
from pathlib import Path
|
||
from dotenv import load_dotenv
|
||
|
||
# Optional fast token counter — falls back to char/4 heuristic if unavailable
|
||
try:
|
||
import tiktoken as _tiktoken
|
||
_TIKTOKEN_OK = True
|
||
except ImportError:
|
||
_TIKTOKEN_OK = False
|
||
|
||
from PyQt5.QtWidgets import (
|
||
QWidget, QVBoxLayout, QHBoxLayout, QLabel, QPushButton, QGroupBox,
|
||
QTextEdit, QMessageBox, QListWidget, QListWidgetItem,
|
||
QSplitter, QFileDialog, QProgressBar, QFrame, QSpinBox, QLineEdit,
|
||
QSizePolicy,
|
||
)
|
||
from PyQt5.QtCore import Qt, pyqtSignal, QThread, QMutex
|
||
from PyQt5.QtGui import QFont, QColor
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Constants
|
||
# ---------------------------------------------------------------------------
|
||
|
||
MAX_LINES = 3
|
||
MAX_LINE_LEN = 60
|
||
NEWLINE = "\r\n"
|
||
|
||
SYSTEM_PROMPT_TEMPLATE = """\
|
||
You are an editor for a visual novel translation. Rewrite dialogue so it fits
|
||
the game's text box.
|
||
|
||
Hard limits (non-negotiable):
|
||
- Maximum 3 lines
|
||
- Maximum 60 characters per line (every character counts)
|
||
- Break lines on spaces only — never mid-word
|
||
- Put each line on its own line (use a newline character to separate lines)
|
||
|
||
Rewriting rules:
|
||
1. Preserve ALL meaning — every fact and clause must survive in condensed form
|
||
2. NEVER truncate — do not drop the last sentence or clause to make it fit;
|
||
rephrase earlier parts more tightly instead
|
||
3. If a line ends a sentence, the next line begins with a capital letter
|
||
4. Preserve these terms exactly: {terms}
|
||
5. Use only plain ASCII punctuation — NO em dashes (—), en dashes (–),
|
||
ellipsis (…), or curly quotes. Use a hyphen (-) instead of a dash.
|
||
|
||
Reply with ONLY the rewritten text, each line on its own line.
|
||
No explanation, no quotes, no markdown."""
|
||
|
||
DEFAULT_TERMS = (
|
||
"Settler, Assimilation, Synchronization, Those Things, We, Inspector, "
|
||
"Suppressant, Brain Dive, Mutant Form, Mimicry, Resource War, "
|
||
"Resource Crisis, Quarantine, Debris Field"
|
||
)
|
||
|
||
USER_PROMPT_TEMPLATE = """\
|
||
Rewrite this to fit in 3 lines of max 60 characters each.
|
||
Keep all meaning. Do not drop the ending.
|
||
|
||
Original:
|
||
{original_flat}"""
|
||
|
||
RETRY_PROMPT_TEMPLATE = """\
|
||
VALIDATION FAILED: {issues}
|
||
|
||
Rewrite more tightly. Every character counts — including spaces and punctuation.
|
||
Do NOT drop any meaning from the original. Rephrase earlier parts to make room."""
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Helpers
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def is_violation(text: str) -> bool:
|
||
lines = text.split(NEWLINE)
|
||
if len(lines) > MAX_LINES:
|
||
return True
|
||
return any(len(line) > MAX_LINE_LEN for line in lines)
|
||
|
||
|
||
def is_valid(text: str) -> list:
|
||
issues = []
|
||
lines = text.split("\r\n")
|
||
if len(lines) > 3:
|
||
issues.append(f"{len(lines)} lines (max 3)")
|
||
for i, line in enumerate(lines, 1):
|
||
if len(line) > 60:
|
||
issues.append(f"line {i} is {len(line)} chars (max 60)")
|
||
return issues
|
||
|
||
|
||
def scan_folder(folder: Path) -> list:
|
||
violations = []
|
||
for f in sorted(folder.glob("*.json")):
|
||
try:
|
||
data = json.loads(f.read_text(encoding="utf-8"))
|
||
except Exception:
|
||
continue
|
||
if not isinstance(data, list):
|
||
continue
|
||
for i, entry in enumerate(data):
|
||
if not isinstance(entry, dict):
|
||
continue
|
||
m = entry.get("message")
|
||
if isinstance(m, str) and is_violation(m):
|
||
violations.append({"file": f.name, "index": i, "original": m})
|
||
return violations
|
||
|
||
|
||
def load_vocab_terms(project_root: Path) -> str:
|
||
"""Load preserved terms from vocab.txt if present, else use defaults."""
|
||
vocab_path = project_root / "vocab.txt"
|
||
if vocab_path.exists():
|
||
try:
|
||
lines = [ln.strip() for ln in vocab_path.read_text(encoding="utf-8").splitlines() if ln.strip()]
|
||
if lines:
|
||
return ", ".join(lines)
|
||
except Exception:
|
||
pass
|
||
return DEFAULT_TERMS
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Cost estimation helpers
|
||
# ---------------------------------------------------------------------------
|
||
|
||
# Approximate output tokens per rewrite (used for pre-run estimate)
|
||
_EST_OUTPUT_TOKENS = 150
|
||
|
||
|
||
def _count_tokens(text: str, model: str) -> int:
|
||
"""Fast token count using tiktoken; falls back to char/4 heuristic."""
|
||
if _TIKTOKEN_OK:
|
||
try:
|
||
enc = _tiktoken.encoding_for_model(model)
|
||
except Exception:
|
||
try:
|
||
enc = _tiktoken.get_encoding("cl100k_base")
|
||
except Exception:
|
||
return len(text) // 4
|
||
return len(enc.encode(text))
|
||
return len(text) // 4
|
||
|
||
|
||
def estimate_run_cost(violations: list, project_root: Path) -> str:
|
||
"""Return a human-readable cost estimate for a full rewrite run."""
|
||
try:
|
||
load_dotenv()
|
||
model = os.getenv("model", "gpt-4.1").strip()
|
||
from util.translation import getPricingConfig
|
||
pricing = getPricingConfig(model)
|
||
in_rate = pricing["inputAPICost"] # $ per million tokens
|
||
out_rate = pricing["outputAPICost"]
|
||
|
||
terms = load_vocab_terms(project_root)
|
||
system_prompt = SYSTEM_PROMPT_TEMPLATE.format(terms=terms)
|
||
sys_tokens = _count_tokens(system_prompt, model)
|
||
|
||
total_in = 0
|
||
total_out = 0
|
||
for v in violations:
|
||
original_flat = v["original"].replace(NEWLINE, "\n")
|
||
user_msg = USER_PROMPT_TEMPLATE.format(original_flat=original_flat)
|
||
total_in += sys_tokens + _count_tokens(user_msg, model)
|
||
total_out += _EST_OUTPUT_TOKENS
|
||
|
||
cost = (total_in / 1_000_000) * in_rate + (total_out / 1_000_000) * out_rate
|
||
return (
|
||
f"~${cost:.4f} "
|
||
f"({total_in:,} in / ~{total_out:,} out est. | "
|
||
f"${in_rate:.2f}/${out_rate:.2f} per M)"
|
||
)
|
||
except Exception as exc:
|
||
return f"(estimate unavailable: {exc})"
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# API caller (OpenAI-compat + Anthropic)
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def _normalise_line_endings(text: str) -> str:
|
||
"""Normalise any line-ending variant the LLM might emit to actual \\r\\n."""
|
||
# Handle escaped sequences the model may output literally (e.g. backslash-r-backslash-n)
|
||
text = text.replace("\\r\\n", "\n").replace("\\n", "\n")
|
||
# Now normalise all real CR/LF variants
|
||
text = text.replace("\r\n", "\n").replace("\r", "\n")
|
||
return text.replace("\n", "\r\n")
|
||
|
||
|
||
# Characters the game engine cannot render — map to ASCII-safe equivalents
|
||
_CHAR_SUBSTITUTIONS = [
|
||
("\u2014", " - "), # em dash — → hyphen
|
||
("\u2013", " - "), # en dash – → hyphen
|
||
("\u2026", "..."), # ellipsis … → three dots
|
||
("\u201C", '"'), # left double quote "
|
||
("\u201D", '"'), # right double quote "
|
||
("\u2018", "'"), # left single quote '
|
||
("\u2019", "'"), # right single quote '
|
||
]
|
||
|
||
|
||
def _sanitise_output(text: str) -> str:
|
||
"""Replace characters the game engine can't render with ASCII equivalents."""
|
||
for src, dst in _CHAR_SUBSTITUTIONS:
|
||
text = text.replace(src, dst)
|
||
return text
|
||
|
||
|
||
def _extract_claude_text(response) -> str:
|
||
"""Collect text from all text blocks; skip thinking/tool blocks."""
|
||
parts = []
|
||
for block in getattr(response, "content", None) or []:
|
||
if getattr(block, "type", None) == "text":
|
||
t = getattr(block, "text", None)
|
||
if t:
|
||
parts.append(t)
|
||
elif hasattr(block, "text") and block.text:
|
||
# Older SDK blocks without explicit type
|
||
parts.append(block.text)
|
||
text = "".join(parts).strip()
|
||
if not text:
|
||
stop = getattr(response, "stop_reason", None)
|
||
raise ValueError(
|
||
f"Empty response from API (stop_reason={stop!r}). "
|
||
"Try again or use a different model."
|
||
)
|
||
return text
|
||
|
||
|
||
def _extract_openai_text(response) -> str:
|
||
"""Extract assistant message text from an OpenAI-compatible completion."""
|
||
choices = getattr(response, "choices", None) or []
|
||
if not choices:
|
||
raise ValueError("Empty response from API (no choices returned).")
|
||
message = choices[0].message
|
||
text = getattr(message, "content", None)
|
||
if text is None and hasattr(message, "refusal") and message.refusal:
|
||
raise ValueError(f"Model refused: {message.refusal}")
|
||
if not text or not str(text).strip():
|
||
finish = getattr(choices[0], "finish_reason", None)
|
||
raise ValueError(
|
||
f"Empty response from API (finish_reason={finish!r}). "
|
||
"Try again or use a different model."
|
||
)
|
||
return str(text).strip()
|
||
|
||
|
||
def _call_llm(messages: list, model: str, api_url: str, api_key: str,
|
||
is_claude: bool, timeout: int) -> tuple:
|
||
"""Call the LLM. Returns (text, input_tokens, output_tokens)."""
|
||
if is_claude:
|
||
import anthropic
|
||
client = anthropic.Anthropic(api_key=api_key)
|
||
system_text = next((m["content"] for m in messages if m["role"] == "system"), "")
|
||
user_msgs = [m for m in messages if m["role"] != "system"]
|
||
response = client.messages.create(
|
||
model=model,
|
||
max_tokens=512,
|
||
system=system_text,
|
||
messages=user_msgs,
|
||
timeout=timeout,
|
||
)
|
||
text = _extract_claude_text(response)
|
||
in_tok = getattr(response.usage, "input_tokens", 0) or 0
|
||
out_tok = getattr(response.usage, "output_tokens", 0) or 0
|
||
return text, in_tok, out_tok
|
||
else:
|
||
from util.translation import isMistralAPI, callMistral
|
||
if isMistralAPI():
|
||
# Mistral budgets requests AND tokens per minute — go through the
|
||
# shared adaptive limiter (same one the translation runs use).
|
||
response = callMistral({
|
||
"model": model,
|
||
"messages": messages,
|
||
"max_tokens": 512,
|
||
"temperature": 0,
|
||
"timeout": timeout,
|
||
})
|
||
else:
|
||
import openai as _openai
|
||
client_kwargs = {"api_key": api_key}
|
||
if api_url:
|
||
client_kwargs["base_url"] = api_url
|
||
client = _openai.OpenAI(**client_kwargs)
|
||
response = client.chat.completions.create(
|
||
model=model,
|
||
messages=messages,
|
||
max_tokens=512,
|
||
temperature=0,
|
||
timeout=timeout,
|
||
)
|
||
text = _extract_openai_text(response)
|
||
in_tok = getattr(response.usage, "prompt_tokens", 0) or 0
|
||
out_tok = getattr(response.usage, "completion_tokens", 0) or 0
|
||
return text, in_tok, out_tok
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Worker thread
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class RewriteWorker(QThread):
|
||
log_signal = pyqtSignal(str)
|
||
file_status_signal = pyqtSignal(str, str) # filename, status
|
||
progress_signal = pyqtSignal(int, int) # current, total
|
||
cost_signal = pyqtSignal(float, int, int) # running_cost, total_in_tok, total_out_tok
|
||
finished_signal = pyqtSignal(bool, str)
|
||
|
||
def __init__(self, violations: list, output_folder: Path, max_retries: int,
|
||
file_threads: int, project_root: Path):
|
||
super().__init__()
|
||
self.violations = violations
|
||
self.output_folder = output_folder
|
||
self.max_retries = max_retries
|
||
self.file_threads = file_threads
|
||
self.project_root = project_root
|
||
self.should_stop = False
|
||
self._mutex = QMutex()
|
||
|
||
def stop(self):
|
||
self._mutex.lock()
|
||
self.should_stop = True
|
||
self._mutex.unlock()
|
||
|
||
def _is_stopped(self):
|
||
self._mutex.lock()
|
||
v = self.should_stop
|
||
self._mutex.unlock()
|
||
return v
|
||
|
||
def run(self):
|
||
try:
|
||
self._run()
|
||
except Exception as exc:
|
||
self.log_signal.emit(f"[ERROR] Unexpected error: {exc}")
|
||
self.log_signal.emit(traceback.format_exc())
|
||
self.finished_signal.emit(False, str(exc))
|
||
|
||
def _run(self):
|
||
load_dotenv()
|
||
|
||
api_key = os.getenv("key", "").strip()
|
||
api_url = os.getenv("api", "").strip()
|
||
model = os.getenv("model", "gpt-4.1").strip()
|
||
provider = os.getenv("API_PROVIDER", "openai").strip().lower()
|
||
timeout = int(os.getenv("timeout", "120"))
|
||
|
||
if provider == "gemini":
|
||
api_url = api_url or "https://generativelanguage.googleapis.com/v1beta/openai/"
|
||
elif provider == "mistral":
|
||
api_url = api_url or "https://api.mistral.ai/v1/"
|
||
|
||
is_claude = (
|
||
any(x in model.lower() for x in ("claude", "sonnet", "haiku", "opus"))
|
||
and (not api_url or "anthropic" in api_url.lower())
|
||
)
|
||
|
||
if not api_key:
|
||
self.log_signal.emit("[ERROR] No API key configured. Set 'key' in .env")
|
||
self.finished_signal.emit(False, "No API key")
|
||
return
|
||
|
||
from util.translation import getPricingConfig
|
||
pricing = getPricingConfig(model)
|
||
input_cost_per_m = pricing["inputAPICost"]
|
||
output_cost_per_m = pricing["outputAPICost"]
|
||
|
||
terms = load_vocab_terms(self.project_root)
|
||
system_prompt = SYSTEM_PROMPT_TEMPLATE.format(terms=terms)
|
||
|
||
self.log_signal.emit(f"Model: {model} | Provider: {'Claude' if is_claude else provider}")
|
||
self.log_signal.emit(
|
||
f"Pricing: ${input_cost_per_m:.4f}/M in ${output_cost_per_m:.4f}/M out"
|
||
)
|
||
self.log_signal.emit(f"Output folder: {self.output_folder}")
|
||
self.log_signal.emit(
|
||
f"Total violations: {len(self.violations)} | Parallel files: {self.file_threads}\n"
|
||
)
|
||
|
||
# Shared counters — all accesses protected by _counter_lock
|
||
_counter_lock = threading.Lock()
|
||
_totals = {"done": 0, "failed": 0, "in_tok": 0, "out_tok": 0}
|
||
total = len(self.violations)
|
||
|
||
# Group violations by file (preserving order)
|
||
from collections import OrderedDict
|
||
file_groups: dict = OrderedDict()
|
||
for idx, v in enumerate(self.violations):
|
||
file_groups.setdefault(v["file"], []).append((idx, v))
|
||
|
||
# ------------------------------------------------------------------ #
|
||
# Per-file worker — runs in a thread-pool thread
|
||
# ------------------------------------------------------------------ #
|
||
def process_file(filename: str, items: list):
|
||
if self._is_stopped():
|
||
return
|
||
|
||
self.file_status_signal.emit(filename, "processing")
|
||
file_done = 0
|
||
file_failed = 0
|
||
|
||
json_path = self.output_folder / filename
|
||
if not json_path.exists():
|
||
self.log_signal.emit(f"\n[SKIP] {filename} — not found in output folder")
|
||
self.file_status_signal.emit(filename, "skipped")
|
||
with _counter_lock:
|
||
_totals["failed"] += len(items)
|
||
self.progress_signal.emit(_totals["done"] + _totals["failed"], total)
|
||
return
|
||
|
||
try:
|
||
file_data = json.loads(json_path.read_text(encoding="utf-8"))
|
||
except Exception as exc:
|
||
self.log_signal.emit(f"\n[ERROR] Could not read {filename}: {exc}")
|
||
self.file_status_signal.emit(filename, "failed")
|
||
with _counter_lock:
|
||
_totals["failed"] += len(items)
|
||
self.progress_signal.emit(_totals["done"] + _totals["failed"], total)
|
||
return
|
||
|
||
self.log_signal.emit(
|
||
f"\n{'─'*50}\nFile: {filename} ({len(items)} violation{'s' if len(items) != 1 else ''})"
|
||
)
|
||
|
||
for idx, violation in items:
|
||
if self._is_stopped():
|
||
break
|
||
|
||
original = violation["original"]
|
||
original_flat = original.replace(NEWLINE, "\n")
|
||
entry_index = violation["index"]
|
||
|
||
with _counter_lock:
|
||
pos = _totals["done"] + _totals["failed"] + 1
|
||
self.log_signal.emit(f"\n [{pos}/{total}] {filename} entry {entry_index}")
|
||
self.log_signal.emit(f" Original: {repr(original_flat)}")
|
||
|
||
messages = [
|
||
{"role": "system", "content": system_prompt},
|
||
{"role": "user", "content": USER_PROMPT_TEMPLATE.format(original_flat=original_flat)},
|
||
]
|
||
|
||
rewrite = None
|
||
for attempt in range(1, self.max_retries + 1):
|
||
if self._is_stopped():
|
||
break
|
||
try:
|
||
result, in_tok, out_tok = _call_llm(
|
||
messages, model, api_url, api_key, is_claude, timeout
|
||
)
|
||
with _counter_lock:
|
||
_totals["in_tok"] += in_tok
|
||
_totals["out_tok"] += out_tok
|
||
running_cost = (
|
||
(_totals["in_tok"] / 1_000_000) * input_cost_per_m +
|
||
(_totals["out_tok"] / 1_000_000) * output_cost_per_m
|
||
)
|
||
self.cost_signal.emit(running_cost, _totals["in_tok"], _totals["out_tok"])
|
||
|
||
result = _normalise_line_endings(result)
|
||
result = _sanitise_output(result)
|
||
issues = is_valid(result)
|
||
if not issues:
|
||
rewrite = result
|
||
self.log_signal.emit(f" -> OK (attempt {attempt}): {repr(result)}")
|
||
break
|
||
else:
|
||
self.log_signal.emit(f" -> Attempt {attempt} invalid: {issues}")
|
||
messages.append({"role": "assistant", "content": result})
|
||
messages.append({"role": "user", "content": RETRY_PROMPT_TEMPLATE.format(issues="; ".join(issues))})
|
||
except Exception as exc:
|
||
self.log_signal.emit(f" -> Attempt {attempt} API error: {exc}")
|
||
if attempt < self.max_retries:
|
||
time.sleep(2)
|
||
|
||
if rewrite is not None and not self._is_stopped():
|
||
entry = file_data[entry_index]
|
||
if isinstance(entry, dict):
|
||
entry["message"] = rewrite
|
||
try:
|
||
json_path.write_text(
|
||
json.dumps(file_data, ensure_ascii=False, indent=4) + "\n",
|
||
encoding="utf-8",
|
||
)
|
||
except Exception as exc:
|
||
self.log_signal.emit(f" [ERROR] Could not save {filename}: {exc}")
|
||
file_done += 1
|
||
with _counter_lock:
|
||
_totals["done"] += 1
|
||
self.progress_signal.emit(_totals["done"] + _totals["failed"], total)
|
||
else:
|
||
file_failed += 1
|
||
with _counter_lock:
|
||
_totals["failed"] += 1
|
||
self.progress_signal.emit(_totals["done"] + _totals["failed"], total)
|
||
if rewrite is None:
|
||
self.log_signal.emit(f" -> FAILED after {self.max_retries} attempts")
|
||
|
||
if not self._is_stopped():
|
||
status = "done" if file_failed == 0 else ("partial" if file_done > 0 else "failed")
|
||
self.file_status_signal.emit(filename, status)
|
||
self.log_signal.emit(
|
||
f" [FILE DONE] {filename} ✓ {file_done} ✗ {file_failed}"
|
||
)
|
||
|
||
# ------------------------------------------------------------------ #
|
||
# Dispatch files to thread pool
|
||
# ------------------------------------------------------------------ #
|
||
with ThreadPoolExecutor(max_workers=self.file_threads) as executor:
|
||
futures = {
|
||
executor.submit(process_file, fn, items): fn
|
||
for fn, items in file_groups.items()
|
||
}
|
||
for future in as_completed(futures):
|
||
if self._is_stopped():
|
||
# Cancel queued (not yet started) futures
|
||
for f in futures:
|
||
f.cancel()
|
||
break
|
||
try:
|
||
future.result()
|
||
except Exception as exc:
|
||
fn = futures[future]
|
||
self.log_signal.emit(f"[ERROR] {fn}: {exc}")
|
||
|
||
stopped = self._is_stopped()
|
||
summary = (
|
||
f"Done: {_totals['done']} Failed: {_totals['failed']} Total: {total}"
|
||
+ (" (stopped early)" if stopped else "")
|
||
)
|
||
self.log_signal.emit(f"\n{'='*50}\n{summary}")
|
||
self.finished_signal.emit(not stopped and _totals["failed"] == 0, summary)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# UI helpers
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def _make_header(title: str) -> QLabel:
|
||
lbl = QLabel(title)
|
||
lbl.setStyleSheet("""
|
||
QLabel {
|
||
font-size: 13px;
|
||
font-weight: bold;
|
||
color: #007acc;
|
||
padding: 8px 0px 5px 0px;
|
||
background-color: transparent;
|
||
}
|
||
""")
|
||
return lbl
|
||
|
||
|
||
def _make_hline() -> QFrame:
|
||
line = QFrame()
|
||
line.setFrameShape(QFrame.HLine)
|
||
line.setFrameShadow(QFrame.Sunken)
|
||
line.setStyleSheet("QFrame { color: #555555; margin: 5px 0px; }")
|
||
return line
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Main tab widget
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class RewriteTab(QWidget):
|
||
"""Tab for scanning and fixing over-limit dialogue messages."""
|
||
|
||
def __init__(self, parent=None):
|
||
super().__init__(parent)
|
||
self._project_root = Path(__file__).resolve().parent.parent
|
||
load_dotenv()
|
||
self._violations: list = []
|
||
self._worker: RewriteWorker | None = None
|
||
# filename → row index in _file_list for fast status updates
|
||
self._file_row: dict = {}
|
||
self._setup_ui()
|
||
|
||
# ------------------------------------------------------------------
|
||
# UI construction
|
||
# ------------------------------------------------------------------
|
||
|
||
def _setup_ui(self):
|
||
root_layout = QHBoxLayout(self)
|
||
root_layout.setContentsMargins(0, 0, 0, 0)
|
||
root_layout.setSpacing(0)
|
||
|
||
splitter = QSplitter(Qt.Horizontal)
|
||
splitter.setHandleWidth(4)
|
||
|
||
splitter.addWidget(self._build_left_panel())
|
||
splitter.addWidget(self._build_right_panel())
|
||
splitter.setSizes([380, 800])
|
||
splitter.setStretchFactor(0, 0)
|
||
splitter.setStretchFactor(1, 1)
|
||
|
||
root_layout.addWidget(splitter)
|
||
|
||
def _build_left_panel(self) -> QWidget:
|
||
panel = QWidget()
|
||
panel.setMinimumWidth(300)
|
||
panel.setMaximumWidth(460)
|
||
layout = QVBoxLayout(panel)
|
||
layout.setContentsMargins(10, 10, 10, 10)
|
||
layout.setSpacing(6)
|
||
|
||
# ---- Folder settings ----
|
||
layout.addWidget(_make_header("Output Folder"))
|
||
|
||
folder_row = QHBoxLayout()
|
||
self._folder_edit = QLineEdit()
|
||
self._folder_edit.setPlaceholderText("translated/")
|
||
self._folder_edit.setText(str(self._project_root / "translated"))
|
||
self._folder_edit.setReadOnly(False)
|
||
folder_row.addWidget(self._folder_edit)
|
||
btn_browse_folder = QPushButton("Browse")
|
||
btn_browse_folder.setFixedWidth(72)
|
||
btn_browse_folder.clicked.connect(self._browse_folder)
|
||
folder_row.addWidget(btn_browse_folder)
|
||
layout.addLayout(folder_row)
|
||
|
||
layout.addWidget(_make_hline())
|
||
|
||
# ---- Violations ----
|
||
layout.addWidget(_make_header("Violations"))
|
||
|
||
self._violations_label = QLabel("No violations loaded.")
|
||
self._violations_label.setWordWrap(True)
|
||
self._violations_label.setStyleSheet("color: #aaaaaa; font-size: 12px;")
|
||
layout.addWidget(self._violations_label)
|
||
|
||
self._estimate_label = QLabel("")
|
||
self._estimate_label.setWordWrap(True)
|
||
self._estimate_label.setStyleSheet("color: #9cdcfe; font-size: 11px;")
|
||
layout.addWidget(self._estimate_label)
|
||
|
||
scan_row = QHBoxLayout()
|
||
self._btn_scan = QPushButton("Scan Folder")
|
||
self._btn_scan.setToolTip("Scan the output folder and load violations")
|
||
self._btn_scan.clicked.connect(self._scan_folder)
|
||
scan_row.addWidget(self._btn_scan)
|
||
|
||
self._btn_load_violations = QPushButton("Load JSON")
|
||
self._btn_load_violations.setToolTip("Load an existing violations.json file")
|
||
self._btn_load_violations.clicked.connect(self._load_violations_file)
|
||
scan_row.addWidget(self._btn_load_violations)
|
||
layout.addLayout(scan_row)
|
||
|
||
layout.addWidget(_make_hline())
|
||
|
||
# ---- Rewrite settings ----
|
||
layout.addWidget(_make_header("Settings"))
|
||
|
||
retry_row = QHBoxLayout()
|
||
retry_row.addWidget(QLabel("Max retries per message:"))
|
||
self._retry_spin = QSpinBox()
|
||
self._retry_spin.setRange(1, 10)
|
||
self._retry_spin.setValue(3)
|
||
self._retry_spin.setFixedWidth(60)
|
||
retry_row.addWidget(self._retry_spin)
|
||
retry_row.addStretch()
|
||
layout.addLayout(retry_row)
|
||
|
||
threads_row = QHBoxLayout()
|
||
threads_row.addWidget(QLabel("Parallel files:"))
|
||
self._threads_spin = QSpinBox()
|
||
self._threads_spin.setRange(1, 20)
|
||
default_threads = int(os.getenv("fileThreads", "1"))
|
||
self._threads_spin.setValue(max(1, default_threads))
|
||
self._threads_spin.setFixedWidth(60)
|
||
self._threads_spin.setToolTip(
|
||
"Number of files to process simultaneously.\n"
|
||
"Defaults to fileThreads from .env. Keep low for free-tier APIs."
|
||
)
|
||
threads_row.addWidget(self._threads_spin)
|
||
threads_row.addStretch()
|
||
layout.addLayout(threads_row)
|
||
|
||
layout.addWidget(_make_hline())
|
||
|
||
# ---- Progress ----
|
||
layout.addWidget(_make_header("Progress"))
|
||
|
||
self._progress_bar = QProgressBar()
|
||
self._progress_bar.setRange(0, 100)
|
||
self._progress_bar.setValue(0)
|
||
self._progress_bar.setTextVisible(True)
|
||
layout.addWidget(self._progress_bar)
|
||
|
||
self._progress_label = QLabel("Idle")
|
||
self._progress_label.setStyleSheet("color: #aaaaaa; font-size: 11px;")
|
||
layout.addWidget(self._progress_label)
|
||
|
||
self._cost_label = QLabel("Cost: $0.0000")
|
||
self._cost_label.setStyleSheet("color: #aaaaaa; font-size: 11px;")
|
||
layout.addWidget(self._cost_label)
|
||
|
||
layout.addWidget(_make_hline())
|
||
|
||
# ---- Action buttons ----
|
||
btn_row = QHBoxLayout()
|
||
self._btn_start = QPushButton("Start Rewrite")
|
||
self._btn_start.setEnabled(False)
|
||
self._btn_start.setStyleSheet("""
|
||
QPushButton { background-color: #2ea043; }
|
||
QPushButton:hover { background-color: #3fb950; }
|
||
QPushButton:disabled { background-color: #404040; color: #888888; }
|
||
""")
|
||
self._btn_start.clicked.connect(self._start_rewrite)
|
||
btn_row.addWidget(self._btn_start)
|
||
|
||
self._btn_stop = QPushButton("Stop")
|
||
self._btn_stop.setEnabled(False)
|
||
self._btn_stop.setStyleSheet("""
|
||
QPushButton { background-color: #da3633; }
|
||
QPushButton:hover { background-color: #f85149; }
|
||
QPushButton:disabled { background-color: #404040; color: #888888; }
|
||
""")
|
||
self._btn_stop.clicked.connect(self._stop_rewrite)
|
||
btn_row.addWidget(self._btn_stop)
|
||
layout.addLayout(btn_row)
|
||
|
||
layout.addStretch()
|
||
return panel
|
||
|
||
def _build_right_panel(self) -> QWidget:
|
||
panel = QWidget()
|
||
layout = QVBoxLayout(panel)
|
||
layout.setContentsMargins(6, 10, 10, 10)
|
||
layout.setSpacing(6)
|
||
|
||
splitter = QSplitter(Qt.Vertical)
|
||
splitter.setHandleWidth(4)
|
||
|
||
# File checklist
|
||
file_container = QWidget()
|
||
file_layout = QVBoxLayout(file_container)
|
||
file_layout.setContentsMargins(0, 0, 0, 0)
|
||
file_layout.setSpacing(4)
|
||
|
||
file_header_row = QHBoxLayout()
|
||
file_header_row.addWidget(_make_header("Files"))
|
||
file_header_row.addStretch()
|
||
btn_all = QPushButton("All")
|
||
btn_all.setToolTip("Select all files")
|
||
btn_all.clicked.connect(lambda: self._set_all_checked(True))
|
||
btn_none = QPushButton("None")
|
||
btn_none.setToolTip("Deselect all files")
|
||
btn_none.clicked.connect(lambda: self._set_all_checked(False))
|
||
file_header_row.addWidget(btn_all)
|
||
file_header_row.addWidget(btn_none)
|
||
file_layout.addLayout(file_header_row)
|
||
|
||
self._file_list = QListWidget()
|
||
self._file_list.setFont(QFont("Consolas", 9))
|
||
file_layout.addWidget(self._file_list)
|
||
splitter.addWidget(file_container)
|
||
|
||
# Log area
|
||
log_container = QWidget()
|
||
log_layout = QVBoxLayout(log_container)
|
||
log_layout.setContentsMargins(0, 0, 0, 0)
|
||
log_layout.addWidget(_make_header("Log"))
|
||
self._log_text = QTextEdit()
|
||
self._log_text.setReadOnly(True)
|
||
self._log_text.setFont(QFont("Consolas", 9))
|
||
log_layout.addWidget(self._log_text)
|
||
splitter.addWidget(log_container)
|
||
|
||
splitter.setSizes([300, 400])
|
||
layout.addWidget(splitter)
|
||
return panel
|
||
|
||
# ------------------------------------------------------------------
|
||
# Actions
|
||
# ------------------------------------------------------------------
|
||
|
||
def _browse_folder(self):
|
||
folder = QFileDialog.getExistingDirectory(
|
||
self, "Select Output Folder", self._folder_edit.text()
|
||
)
|
||
if folder:
|
||
self._folder_edit.setText(folder)
|
||
|
||
def _scan_folder(self):
|
||
folder = Path(self._folder_edit.text().strip())
|
||
if not folder.exists():
|
||
QMessageBox.warning(self, "Folder Not Found", f"Folder does not exist:\n{folder}")
|
||
return
|
||
|
||
self._log("Scanning folder for violations...")
|
||
violations = scan_folder(folder)
|
||
|
||
# Save violations.json next to the project root
|
||
out_path = self._project_root / "violations.json"
|
||
try:
|
||
out_path.write_text(
|
||
json.dumps(violations, ensure_ascii=False, indent=2), encoding="utf-8"
|
||
)
|
||
self._log(f"Saved {len(violations)} violations → {out_path}")
|
||
except Exception as exc:
|
||
self._log(f"[WARNING] Could not save violations.json: {exc}")
|
||
|
||
self._load_violations(violations)
|
||
|
||
def _load_violations_file(self):
|
||
path, _ = QFileDialog.getOpenFileName(
|
||
self, "Load Violations File",
|
||
str(self._project_root),
|
||
"JSON Files (*.json);;All Files (*)"
|
||
)
|
||
if not path:
|
||
return
|
||
try:
|
||
data = json.loads(Path(path).read_text(encoding="utf-8"))
|
||
if not isinstance(data, list):
|
||
raise ValueError("Expected a JSON array")
|
||
self._load_violations(data)
|
||
self._log(f"Loaded {len(data)} violations from {path}")
|
||
except Exception as exc:
|
||
QMessageBox.critical(self, "Load Error", f"Failed to load violations:\n{exc}")
|
||
|
||
def _load_violations(self, violations: list):
|
||
self._violations = violations
|
||
self._file_list.clear()
|
||
self._file_row.clear()
|
||
|
||
# Group by file, preserving order
|
||
from collections import defaultdict, OrderedDict
|
||
file_groups: dict = OrderedDict()
|
||
for v in violations:
|
||
file_groups.setdefault(v["file"], []).append(v)
|
||
|
||
for row, (filename, items) in enumerate(file_groups.items()):
|
||
n = len(items)
|
||
label = f"{filename} ({n} violation{'s' if n != 1 else ''})"
|
||
item = QListWidgetItem(label)
|
||
item.setFlags(item.flags() | Qt.ItemIsUserCheckable)
|
||
item.setCheckState(Qt.Checked)
|
||
item.setForeground(QColor("#cccccc"))
|
||
item.setData(Qt.UserRole, filename)
|
||
self._file_list.addItem(item)
|
||
self._file_row[filename] = row
|
||
|
||
n_files = len(file_groups)
|
||
count = len(violations)
|
||
self._violations_label.setText(
|
||
f"{count} violation{'s' if count != 1 else ''} across {n_files} file{'s' if n_files != 1 else ''}."
|
||
)
|
||
self._estimate_label.setText("Estimating cost...")
|
||
self._btn_start.setEnabled(count > 0)
|
||
|
||
if count > 0:
|
||
self._log(f"Loaded {count} violations across {n_files} files.")
|
||
self._run_cost_estimate(violations)
|
||
else:
|
||
self._log("No violations found.")
|
||
self._estimate_label.setText("")
|
||
|
||
def _set_all_checked(self, checked: bool):
|
||
state = Qt.Checked if checked else Qt.Unchecked
|
||
for i in range(self._file_list.count()):
|
||
self._file_list.item(i).setCheckState(state)
|
||
|
||
def _checked_violations(self) -> list:
|
||
"""Return only violations belonging to checked files."""
|
||
checked_files = set()
|
||
for i in range(self._file_list.count()):
|
||
item = self._file_list.item(i)
|
||
if item.checkState() == Qt.Checked:
|
||
checked_files.add(item.data(Qt.UserRole))
|
||
return [v for v in self._violations if v["file"] in checked_files]
|
||
|
||
def _run_cost_estimate(self, violations: list):
|
||
"""Compute cost estimate in a background thread and update the label."""
|
||
project_root = self._project_root
|
||
|
||
class _EstimateThread(QThread):
|
||
done = pyqtSignal(str)
|
||
def run(self_t):
|
||
self_t.done.emit(estimate_run_cost(violations, project_root))
|
||
|
||
self._estimate_thread = _EstimateThread(self)
|
||
self._estimate_thread.done.connect(
|
||
lambda text: self._estimate_label.setText(f"Est. cost: {text}")
|
||
)
|
||
self._estimate_thread.done.connect(
|
||
lambda text: self._log(f"Cost estimate: {text}")
|
||
)
|
||
self._estimate_thread.start()
|
||
|
||
def _start_rewrite(self):
|
||
if not self._violations:
|
||
QMessageBox.information(self, "No Violations", "Scan or load violations first.")
|
||
return
|
||
|
||
violations = self._checked_violations()
|
||
if not violations:
|
||
QMessageBox.information(self, "No Files Selected",
|
||
"Check at least one file in the file list.")
|
||
return
|
||
|
||
folder = Path(self._folder_edit.text().strip())
|
||
if not folder.exists():
|
||
QMessageBox.warning(self, "Folder Not Found", f"Output folder not found:\n{folder}")
|
||
return
|
||
|
||
self._btn_start.setEnabled(False)
|
||
self._btn_stop.setEnabled(True)
|
||
self._btn_scan.setEnabled(False)
|
||
self._btn_load_violations.setEnabled(False)
|
||
self._progress_bar.setValue(0)
|
||
self._progress_label.setText("Starting...")
|
||
self._cost_label.setText("Cost: $0.0000")
|
||
self._log("\n" + "="*50)
|
||
n_files = len({v["file"] for v in violations})
|
||
self._log(f"Starting rewrite: {len(violations)} violations across {n_files} files...")
|
||
|
||
# Reset file list colours for checked items only
|
||
for i in range(self._file_list.count()):
|
||
item = self._file_list.item(i)
|
||
if item.checkState() == Qt.Checked:
|
||
item.setForeground(QColor("#cccccc"))
|
||
|
||
self._worker = RewriteWorker(
|
||
violations=violations,
|
||
output_folder=folder,
|
||
max_retries=self._retry_spin.value(),
|
||
file_threads=self._threads_spin.value(),
|
||
project_root=self._project_root,
|
||
)
|
||
self._worker.log_signal.connect(self._log)
|
||
self._worker.file_status_signal.connect(self._on_file_status)
|
||
self._worker.progress_signal.connect(self._on_progress)
|
||
self._worker.cost_signal.connect(self._on_cost)
|
||
self._worker.finished_signal.connect(self._on_finished)
|
||
self._worker.start()
|
||
|
||
def _stop_rewrite(self):
|
||
if self._worker and self._worker.isRunning():
|
||
self._worker.stop()
|
||
self._btn_stop.setEnabled(False)
|
||
self._progress_label.setText("Stopping...")
|
||
self._log("Stop requested...")
|
||
|
||
# ------------------------------------------------------------------
|
||
# Signals
|
||
# ------------------------------------------------------------------
|
||
|
||
def _on_file_status(self, filename: str, status: str):
|
||
row = self._file_row.get(filename)
|
||
if row is None:
|
||
return
|
||
item = self._file_list.item(row)
|
||
if item is None:
|
||
return
|
||
colors = {
|
||
"processing": "#007acc",
|
||
"done": "#2ea043",
|
||
"partial": "#e3b341", # some succeeded, some failed
|
||
"failed": "#da3633",
|
||
"skipped": "#888888",
|
||
}
|
||
item.setForeground(QColor(colors.get(status, "#cccccc")))
|
||
|
||
def _on_progress(self, current: int, total: int):
|
||
if total > 0:
|
||
pct = int(current / total * 100)
|
||
self._progress_bar.setValue(pct)
|
||
self._progress_label.setText(f"{current} / {total} ({pct}%)")
|
||
|
||
def _on_cost(self, cost: float, in_tok: int, out_tok: int):
|
||
self._cost_label.setText(
|
||
f"Cost: ${cost:.4f} ({in_tok:,} in / {out_tok:,} out tokens)"
|
||
)
|
||
|
||
def _on_finished(self, success: bool, message: str):
|
||
self._btn_start.setEnabled(bool(self._violations))
|
||
self._btn_stop.setEnabled(False)
|
||
self._btn_scan.setEnabled(True)
|
||
self._btn_load_violations.setEnabled(True)
|
||
self._progress_label.setText("Finished")
|
||
self._worker = None
|
||
|
||
# ------------------------------------------------------------------
|
||
# Log helper
|
||
# ------------------------------------------------------------------
|
||
|
||
def _log(self, text: str):
|
||
self._log_text.append(text)
|
||
# Auto-scroll
|
||
sb = self._log_text.verticalScrollBar()
|
||
sb.setValue(sb.maximum())
|