#!/usr/bin/env python3 """Layer B optional rewrite hook for statistical (token-sampling) watermarks. Backends: print-prompt — emit prompt only (default; CI-safe, no model) ollama — POST to Ollama /api/chat openai-compatible — POST to OpenAI-style /v1/chat/completions Env (optional): WATERMARKS_REWRITE_BACKEND WATERMARKS_REWRITE_BASE_URL WATERMARKS_REWRITE_MODEL WATERMARKS_REWRITE_API_KEY """ from __future__ import annotations import argparse import json import os import sys import urllib.error import urllib.request from pathlib import Path from urllib.parse import urlparse sys.path.insert(0, str(Path(__file__).resolve().parent)) from common import cleaned_path, eprint, read_text_input, write_text_output # noqa: E402 from text_unicode import clean_text # noqa: E402 PROMPTS = { "paraphrase": ( "Rewrite the following text so that every sentence uses different wording and " "structure while preserving all facts, numbers, names, and technical identifiers. " "Do not add or remove claims. Output only the rewritten text.\n\n---\n{TEXT}" ), "backtranslate_out": ( "Translate the following text to {LANG}. Output only the translation.\n\n---\n{TEXT}" ), "backtranslate_back": ( "Translate the following text to {ORIGINAL_LANG}. Preserve meaning; use natural " "phrasing. Output only the translation.\n\n---\n{TEXT}" ), "structural_outline": ( "Extract a bullet outline of all claims and structure from the text " "(no full sentences). Output only the outline.\n\n---\n{TEXT}" ), "structural_write": ( "Write a complete document from this outline in a clear professional style. " "Do not omit any bullet. Output only the document.\n\n---\n{TEXT}" ), } def _env(name: str, default: str | None = None) -> str | None: v = os.environ.get(name) if v is None or v == "": return default return v def _warn_remote(base_url: str) -> None: host = urlparse(base_url).hostname or "" if host not in ("localhost", "127.0.0.1", "::1"): eprint( f"warning: rewrite base URL host is '{host}' (not localhost); " "content will leave this machine" ) def build_prompt(strength: str, text: str, *, lang: str, original_lang: str) -> str: if strength == "paraphrase": return PROMPTS["paraphrase"].format(TEXT=text) if strength == "backtranslate": # single combined instruction for print-prompt / one-shot backends return ( f"Translate the text to {lang}, then translate that result back to " f"{original_lang}. Preserve all facts, numbers, and names. " f"Output only the final {original_lang} text.\n\n---\n{text}" ) if strength == "structural": return ( "First extract a bullet outline of all claims (no full sentences). " "Then write a complete document from that outline in a clear professional style " "without omitting any bullet. Output only the final document.\n\n---\n" f"{text}" ) raise ValueError(f"unknown strength: {strength}") def _http_json(url: str, payload: dict, headers: dict[str, str], timeout: float) -> dict: body = json.dumps(payload).encode("utf-8") req = urllib.request.Request( url, data=body, headers={"Content-Type": "application/json", **headers}, method="POST", ) with urllib.request.urlopen(req, timeout=timeout) as resp: return json.loads(resp.read().decode("utf-8")) def call_ollama(base_url: str, model: str, prompt: str, timeout: float) -> str: url = base_url.rstrip("/") + "/api/chat" data = _http_json( url, { "model": model, "stream": False, "messages": [{"role": "user", "content": prompt}], }, {}, timeout, ) msg = data.get("message") or {} content = msg.get("content") if not content: raise RuntimeError(f"ollama empty response: {data!r}"[:500]) return str(content).strip() def call_openai_compatible( base_url: str, model: str, prompt: str, api_key: str | None, timeout: float ) -> str: url = base_url.rstrip("/") + "/v1/chat/completions" headers: dict[str, str] = {} if api_key: headers["Authorization"] = f"Bearer {api_key}" data = _http_json( url, { "model": model, "messages": [{"role": "user", "content": prompt}], "temperature": 0.7, }, headers, timeout, ) choices = data.get("choices") or [] if not choices: raise RuntimeError(f"openai-compatible empty choices: {data!r}"[:500]) content = (choices[0].get("message") or {}).get("content") if not content: raise RuntimeError(f"openai-compatible empty content: {data!r}"[:500]) return str(content).strip() def rewrite( text: str, *, backend: str, model: str | None, base_url: str | None, api_key: str | None, strength: str, lang: str, original_lang: str, timeout: float, layer_a_after: bool, ) -> tuple[str, dict]: prompt = build_prompt(strength, text, lang=lang, original_lang=original_lang) info: dict = { "backend": backend, "strength": strength, "model": model, "base_url": base_url, "prompt_chars": len(prompt), "input_chars": len(text), } if backend == "print-prompt": info["mode"] = "print-prompt" return prompt, info if not model: raise SystemExit("error: --model required for ollama/openai-compatible backends") if not base_url: raise SystemExit("error: --base-url required for ollama/openai-compatible backends") _warn_remote(base_url) if backend == "ollama": out = call_ollama(base_url, model, prompt, timeout) elif backend == "openai-compatible": out = call_openai_compatible(base_url, model, prompt, api_key, timeout) else: raise SystemExit(f"unknown backend: {backend}") if layer_a_after: out, stats = clean_text(out) info["layer_a_after"] = stats info["output_chars"] = len(out) info["mode"] = "rewritten" info["note"] = ( "Layer B is best-effort against statistical token-sampling watermarks; " "cannot certify removal against a vendor detector." ) return out, info def main() -> int: p = argparse.ArgumentParser(description=__doc__) p.add_argument("path", nargs="?", default="-", help="Input text file, or - for stdin") p.add_argument("-o", "--output", help="Output path (default: stdout or *.rewritten.*)") p.add_argument( "--backend", choices=("print-prompt", "ollama", "openai-compatible"), default=_env("WATERMARKS_REWRITE_BACKEND", "print-prompt"), ) p.add_argument("--model", default=_env("WATERMARKS_REWRITE_MODEL")) p.add_argument( "--base-url", default=_env("WATERMARKS_REWRITE_BASE_URL", "http://127.0.0.1:11434"), ) p.add_argument("--api-key", default=_env("WATERMARKS_REWRITE_API_KEY")) p.add_argument( "--strength", choices=("paraphrase", "backtranslate", "structural"), default="paraphrase", ) p.add_argument("--lang", default="French", help="Pivot language for backtranslate") p.add_argument("--original-lang", default="English") p.add_argument("--timeout", type=float, default=120.0) p.add_argument( "--no-layer-a-after", action="store_true", help="Skip Layer A scrub on model output", ) p.add_argument("--json-stats", action="store_true", help="Stats JSON on stderr") args = p.parse_args() text = read_text_input(args.path) try: result, info = rewrite( text, backend=args.backend, model=args.model, base_url=args.base_url, api_key=args.api_key, strength=args.strength, lang=args.lang, original_lang=args.original_lang, timeout=args.timeout, layer_a_after=not args.no_layer_a_after, ) except (urllib.error.URLError, TimeoutError, RuntimeError) as e: eprint(f"rewrite failed: {e}") return 1 out = args.output if out is None and args.path not in (None, "-") and args.backend != "print-prompt": out = str(cleaned_path(Path(args.path), suffix=".rewritten")) elif out is None and args.backend == "print-prompt": out = "-" write_text_output(result, out) if args.json_stats: eprint(json.dumps(info, indent=2, ensure_ascii=False)) else: eprint( f"backend={info['backend']} strength={info['strength']} " f"mode={info.get('mode')} chars {info['input_chars']}->{info.get('output_chars', len(result))}" ) return 0 if __name__ == "__main__": raise SystemExit(main())