""" Language correction pipeline for riksdag talk arguments. Reads talks from 2002 onwards that have arguments (extracted by a small 9b model with sometimes poor Swedish), sends them to the big-smart LLM for language correction, and writes corrected arguments back. Multi-turn strategy (mirrors summarize_and_tag.py): - Turn 1: send arguments as a keyed dict {argument_1: ..., argument_2: ...}, ask for language correction. guided_json guarantees valid JSON output. Any argument the LLM can't parse is returned as null. - Turn 2 (if any nulls): same conversation, now also includes the full speech text so the LLM can re-derive meaning from source. - Any key still null after turn 2 falls back to the original text. Tracks completion via `arguments_corrected` column (resumable). nohup python scripts/correct_arguments.py >> logs/correct_arguments.log 2>&1 & echo $! > logs/correct_arguments.pid """ import json import logging import os import re import sys import threading import time from concurrent.futures import ThreadPoolExecutor, as_completed os.chdir("/home/lasse/riksdagen") sys.path.append("/home/lasse/riksdagen") from dotenv import load_dotenv load_dotenv() logging.basicConfig( level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s", handlers=[logging.StreamHandler(sys.stdout)], ) logger = logging.getLogger(__name__) logging.getLogger("httpx").setLevel(logging.WARNING) logging.getLogger("requests").setLevel(logging.WARNING) logging.getLogger("openai").setLevel(logging.WARNING) from packages.llm import LLM from postgres_client import pg # ───────────────────────────────────────────────────────────────────────────── # guided_json schema # Each key maps to a corrected string, or null if the argument was incomprehensible. # ───────────────────────────────────────────────────────────────────────────── _GUIDED_NULLABLE_DICT = { "type": "object", "additionalProperties": {"type": ["string", "null"]}, } # ───────────────────────────────────────────────────────────────────────────── # Helpers # ───────────────────────────────────────────────────────────────────────────── def _to_keyed(arguments: list[str]) -> dict[str, str]: return {f"argument_{i+1}": arg for i, arg in enumerate(arguments)} def _from_keyed(keyed: dict, arguments: list[str]) -> list[str | None]: """Map keyed dict back to a list, preserving original order. Missing keys → None.""" return [keyed.get(f"argument_{i+1}") for i in range(len(arguments))] # ───────────────────────────────────────────────────────────────────────────── # Prompts # ───────────────────────────────────────────────────────────────────────────── SYSTEM_PROMPT = """\ Du är ett språkgranskningsverktyg för svenska riksdagsanföranden. Du får en dict med extraherade politiska argument skrivna av en liten AI-modell \ som ibland producerar bristfällig svenska (grammatikfel, knaggliga meningar, konstiga ordval). Regler: - Ändra ENBART språket: grammatik, ordval, meningsbyggnad, stavning. - Ändra INTE innebörden, ståndpunkten eller det politiska innehållet. - Behåll ungefär samma längd och form på varje argument. - Om ett argument är knaggligt men begripligt: rätta språket och behåll innebörden exakt. - Om ett argument är så korrumperat att du inte kan avgöra vad som menas: sätt värdet till null. - Svara BARA med ett JSON-objekt med exakt samma nycklar som indata. """ def _turn1_user(arguments: list[str]) -> str: keyed = _to_keyed(arguments) return ( "Rätta språket i dessa argument:\n\n" + json.dumps(keyed, ensure_ascii=False, indent=2) + "\n\nSvara BARA med ett JSON-objekt med samma nycklar." ) def _turn2_user(unclear_keyed: dict[str, str], talk: dict) -> str: talare = (talk.get("talare") or "Okänd").strip() parti = (talk.get("parti") or "").strip() text = talk.get("anforandetext", "") return ( f"Dessa argument var obegripliga. Här är det fullständiga anförandet som referens:\n\n" f"Talare: {talare} ({parti})\n\n" f"Anförande:\n{text}\n\n" "---\n" "Rätta nu språket i dessa argument med hjälp av anförandet. " "Om ett argument är alltför korrumperat, re-extrahera det korrekt från anförandet. " "Sätt inget värde till null — anförandet ger dig tillräcklig kontext.\n\n" + json.dumps(unclear_keyed, ensure_ascii=False, indent=2) + "\n\nSvara BARA med ett JSON-objekt med samma nycklar." ) # ───────────────────────────────────────────────────────────────────────────── # LLM # ───────────────────────────────────────────────────────────────────────────── _thread_local = threading.local() def get_worker_llm() -> LLM: if not hasattr(_thread_local, "llm"): _thread_local.llm = LLM( model="big-smart", think=False, temperature=0.1, chat=False, silent=True, ) return _thread_local.llm def _parse_dict(content: str) -> dict | None: """Extract a JSON object from model output as robustly as possible.""" text = content.strip() m = re.search(r"```(?:json)?\s*([\s\S]*?)\s*```", text) if m: text = m.group(1).strip() try: val, _ = json.JSONDecoder().raw_decode(text) if isinstance(val, dict): return val except (json.JSONDecodeError, ValueError): pass m = re.search(r"\{[\s\S]*?\}", text) if m: try: val = json.loads(m.group(0)) if isinstance(val, dict): return val except (json.JSONDecodeError, ValueError): pass return None def _call_structured(llm: LLM, messages: list[dict]) -> tuple[dict | None, str]: """Call LLM with guided_json dict schema. Returns (parsed_dict, raw_content).""" result = llm.generate( messages=messages, temperature=0.1, think=False, model="big-smart", extra_body={"guided_json": _GUIDED_NULLABLE_DICT}, ) if isinstance(result, str): raise RuntimeError(f"LLM API error: {result}") content = getattr(result, "content", "") or "" return _parse_dict(content), content # ───────────────────────────────────────────────────────────────────────────── # Core correction logic # ───────────────────────────────────────────────────────────────────────────── def correct_arguments(llm: LLM, talk: dict) -> tuple[list[str] | None, bool]: """ Multi-turn language correction for one talk's arguments. Returns (corrected_list, used_full_text). corrected_list is None when there are no arguments to process. used_full_text is True when turn 2 (full speech context) was needed. """ arguments = talk.get("arguments") if not arguments: return None, False messages = [{"role": "system", "content": SYSTEM_PROMPT}] # Turn 1: correct arguments without full speech text messages.append({"role": "user", "content": _turn1_user(arguments)}) result1, content1 = _call_structured(llm, messages) if result1 is None: raise RuntimeError(f"Unparseable turn-1 response: {content1[:300]}") messages.append({"role": "assistant", "content": content1}) corrected = _from_keyed(result1, arguments) # Turn 2: retry unclear ones with full speech text unclear_indices = [i for i, v in enumerate(corrected) if v is None] used_full_text = bool(unclear_indices) if unclear_indices: logger.info(f"Talk {talk['id']}: {len(unclear_indices)} unclear argument(s), fetching full text") unclear_keyed = {f"argument_{i+1}": arguments[i] for i in unclear_indices} messages.append({"role": "user", "content": _turn2_user(unclear_keyed, talk)}) result2, content2 = _call_structured(llm, messages) if result2 is not None: for i in unclear_indices: key = f"argument_{i+1}" val = result2.get(key) corrected[i] = val if isinstance(val, str) else arguments[i] else: logger.warning(f"Talk {talk['id']}: turn-2 response unusable — keeping originals") for i in unclear_indices: corrected[i] = arguments[i] # Final safety: replace any remaining nulls with originals corrected = [v if isinstance(v, str) else arguments[i] for i, v in enumerate(corrected)] return corrected, used_full_text # ───────────────────────────────────────────────────────────────────────────── # Worker # ───────────────────────────────────────────────────────────────────────────── def process_talk(talk: dict) -> tuple[bool, bool]: """Returns (success, used_full_text).""" try: llm = get_worker_llm() corrected, used_full_text = correct_arguments(llm, talk) if corrected is None: pg.execute_void( "UPDATE talks SET arguments_corrected = TRUE WHERE id = %s", (talk["id"],), ) return True, False pg.execute_void( "UPDATE talks SET arguments = %s, arguments_corrected = TRUE WHERE id = %s", (corrected, talk["id"]), ) return True, used_full_text except Exception as e: logger.error(f"Error processing talk {talk.get('id')}: {e}") return False, False # ───────────────────────────────────────────────────────────────────────────── # Main loop # ───────────────────────────────────────────────────────────────────────────── WORKERS = 5 BATCH_SIZE = 100 def fetch_batch() -> list[dict]: return pg.execute( """ SELECT id, anforandetext, talare, parti, arguments FROM talks WHERE array_length(arguments, 1) > 0 AND datum >= '2002-01-01' AND arguments_corrected IS NOT TRUE AND anforandetext IS NOT NULL ORDER BY datum DESC NULLS LAST LIMIT %s """, (BATCH_SIZE,), ) def ensure_schema(): try: pg.execute_void( "ALTER TABLE talks ADD COLUMN IF NOT EXISTS arguments_corrected BOOLEAN DEFAULT FALSE" ) except Exception as e: logger.warning(f"Could not apply schema change: {e}") def backup_arguments(): """ Dump all original arguments (2002+) to a JSON file before any corrections. Skipped if the backup file already exists. """ path = "logs/arguments_backup.json" if os.path.exists(path): logger.info(f"Backup already exists at {path}, skipping.") return logger.info("Creating arguments backup …") rows = pg.execute( """ SELECT id, arguments FROM talks WHERE array_length(arguments, 1) > 0 AND datum >= '2002-01-01' AND anforandetext IS NOT NULL ORDER BY id """ ) backup = {str(row["id"]): row["arguments"] for row in rows} with open(path, "w", encoding="utf-8") as f: json.dump(backup, f, ensure_ascii=False, indent=2) logger.info(f"Backed up {len(backup)} talks → {path}") def main(): os.makedirs("logs", exist_ok=True) ensure_schema() backup_arguments() logger.info("Starting argument language correction pipeline …") total = 0 errors = 0 full_text_lookups = 0 start = time.time() while True: batch = fetch_batch() if not batch: logger.info("No more talks to correct. All done.") break with ThreadPoolExecutor(max_workers=WORKERS) as executor: futures = {executor.submit(process_talk, talk): talk for talk in batch} for future in as_completed(futures): ok, used_full_text = future.result() total += 1 if not ok: errors += 1 if used_full_text: full_text_lookups += 1 if total % 20 == 0: elapsed = time.time() - start rate = total / (elapsed / 60) if elapsed > 0 else 0 pct = full_text_lookups / total * 100 print( f"Processed: {total} | Errors: {errors} " f"| Full-text lookups: {full_text_lookups} ({pct:.1f}%) " f"| Rate: {rate:.1f}/min" ) elapsed = time.time() - start rate = total / (elapsed / 60) if elapsed > 0 else 0 pct = full_text_lookups / total * 100 if total else 0 logger.info( f"=== DONE: {total} processed, {errors} errors, " f"{full_text_lookups} full-text lookups ({pct:.1f}%), " f"{elapsed/3600:.1f}h total, {rate:.1f}/min ===" ) if __name__ == "__main__": main()