recovery.sh migrated data out of ArangoDB. This project abandoned Arango before the extraction, and no one cloning plenum has an Arango server to migrate from — it was carried over by inertia. scripts/deploy_debate_embeddings.sh and deploy_summary_embeddings.sh were one-off rollout scripts for features whose columns are now part of the base schema, so a fresh install already has them. scripts/correct_arguments.py was a single corrective pass over a specific extraction error made by a specific small model on Swedish speeches from 2002 onward. Not reusable, and not something a new deployment should run. _postgres/rename_map.py and scripts/generate_rename_migration.py produced the Swedish-to-English migration. That migration is written, applied, and its rollback verified; the generator has no remaining use and the map described a schema that no longer exists. Root __init__.py was empty and made the repository look like an importable package, which it is not. scripts/test_auth_e2e.py moved to tests/ — it is a test, and putting it in scripts/ implied it was an operational tool. Kept deliberately: the eval harness and its analysis scripts, which are a documented feature rather than leftovers, and the embed/summarise/debate pipeline scripts, which any deployment needs after an ingest run.main
parent
238b46e277
commit
65bb842cc2
10 changed files with 5 additions and 1134 deletions
@ -1,250 +0,0 @@ |
||||
"""The Swedish -> English rename, as data. |
||||
|
||||
One map drives three things that would otherwise drift: the SQL migration, the |
||||
mechanical code rewrite, and the compatibility shim that rewrites SQL replayed from |
||||
saved chat snapshots. |
||||
|
||||
Design notes worth knowing before editing: |
||||
|
||||
* Column renames are **per table**. `dok_id` means different things in different |
||||
tables — in `talks` it identifies the protocol document the speech appeared in, in |
||||
`motions` it is the document's own primary key — so they map to different names. |
||||
* Values are never translated. `Bifall` / `Avslag` stay as published; parliament.yaml |
||||
glosses them. A research tool must not silently rewrite the record. |
||||
* Concepts with no clean cross-country equivalent (yrkande, riksmöte, replik) get |
||||
neutral column names here while the country's own word lives in |
||||
parliament.yaml `vocabulary:`. |
||||
""" |
||||
from __future__ import annotations |
||||
|
||||
# --- tables ------------------------------------------------------------------- |
||||
# "talks" reads as conference talks to everyone outside this project; "motions" is |
||||
# valid parliamentary English but the table really holds member-submitted documents |
||||
# of which a motion is one type — hence the new doc_type column. |
||||
TABLES: dict[str, str] = { |
||||
"talks": "speeches", |
||||
"chunks": "speech_chunks", |
||||
"motions": "documents", |
||||
"motion_chunks": "document_chunks", |
||||
"motion_authors": "document_authors", |
||||
"motion_yrkanden": "document_proposals", |
||||
} |
||||
|
||||
# --- columns, per table (keyed by the ORIGINAL table name) -------------------- |
||||
COLUMNS: dict[str, dict[str, str]] = { |
||||
"people": { |
||||
"intressent_id": "person_id", |
||||
"hangar_id": "source_record_id", |
||||
"hangar_guid": "source_record_guid", |
||||
"sourceid": "source_id", |
||||
"fodd_ar": "birth_year", |
||||
"kon": "gender", |
||||
"efternamn": "last_name", |
||||
# "tilltalsnamn" is the name someone is addressed by, which English has no |
||||
# single word for. "first_name" loses that nuance and is worth the trade. |
||||
"tilltalsnamn": "first_name", |
||||
"sorteringsnamn": "sort_name", |
||||
"iort": "home_town", |
||||
"parti": "party", |
||||
# Meaning varies by country: single-member seat in the UK, multi-member |
||||
# district in Sweden. Nothing in the code assumes either. |
||||
"valkrets": "constituency", |
||||
"person_url_xml": "source_url", |
||||
"bild_url_80": "image_url_small", |
||||
"bild_url_192": "image_url_medium", |
||||
"bild_url_max": "image_url_large", |
||||
"personuppdrag": "assignments", |
||||
"personuppgift": "contact_details", |
||||
"namn": "name", |
||||
"aktiv": "active", |
||||
}, |
||||
"talks": { |
||||
"anforande_id": "source_speech_id", |
||||
"anforandetext": "text", |
||||
"avsnittsrubrik": "section_title", |
||||
# Position within the debate. Not called "number" to avoid confusion with |
||||
# the 0-based `ordinal` used in the child tables. |
||||
"anforande_nummer": "sequence", |
||||
"kammaraktivitet": "activity_type", |
||||
# English "Speaker" is the presiding officer, so a bare `speaker` would be |
||||
# ambiguous in a parliamentary schema. The API keeps `speaker`, which is |
||||
# already its public contract. |
||||
"talare": "speaker_name", |
||||
"parti": "party", |
||||
"intressent_id": "person_id", |
||||
"datum": "date", |
||||
"dok_datum": "source_datetime", |
||||
"period": "session_year", |
||||
"rel_dok_id": "related_doc_id", |
||||
"dok_nummer": "source_doc_number", |
||||
"hangar_id": "source_record_id", |
||||
"dok_id": "source_doc_id", |
||||
"titel": "title", |
||||
"debate": "debate_id", |
||||
# A short right-of-reply intervention. `is_reply` loses the procedural |
||||
# standing the Swedish term carries; nothing downstream depends on it. |
||||
"replik": "is_reply", |
||||
"debateurl": "url_video", |
||||
"audiofileurl": "url_audio_file", |
||||
"startpos": "audio_start_seconds", |
||||
}, |
||||
"chunks": {"talk_id": "speech_id"}, |
||||
"debates": {"debate": "id", "datum": "date"}, |
||||
"motions": { |
||||
"dok_id": "doc_id", |
||||
"hangar_id": "source_record_id", |
||||
# Annual session, e.g. "2022/23". Deliberately not "term": the European |
||||
# Parliament uses that for its five-year cycle. |
||||
"rm": "session_label", |
||||
"beteckning": "designation", |
||||
"subtyp": "subtype", |
||||
# In `motions` this is always a committee, though the source uses `organ` |
||||
# more broadly elsewhere. |
||||
"organ": "committee", |
||||
"datum": "date", |
||||
"systemdatum": "source_updated_at", |
||||
"publicerad": "published_at", |
||||
"year": "session_year", |
||||
"titel": "title", |
||||
"undertitel": "subtitle", |
||||
"forslag_text": "proposals_text", |
||||
"dokument_url_text": "url_text", |
||||
"dokument_url_html": "url_html", |
||||
"pdf_url": "url_pdf", |
||||
"forslag": "proposals_raw", |
||||
"bilagor": "attachments", |
||||
"num_yrkanden": "num_proposals", |
||||
}, |
||||
"motion_authors": { |
||||
"dok_id": "doc_id", |
||||
"intressent_id": "person_id", |
||||
"namn": "name", |
||||
"partibet": "party", |
||||
"roll": "role", |
||||
}, |
||||
"motion_chunks": {"motion_id": "doc_id"}, |
||||
"motion_yrkanden": { |
||||
"dok_id": "doc_id", |
||||
"nummer": "number", |
||||
# The operative demand itself — the thing people search and cite. |
||||
"lydelse": "text", |
||||
"utskottet": "committee_recommendation", |
||||
"kammaren": "chamber_decision", |
||||
"behandlas_i": "handled_in", |
||||
}, |
||||
# session_type values stay 'general' / 'mp'. Rewriting live rows and a CHECK |
||||
# constraint buys nothing: 'mp' reads as "member" generically, and the display |
||||
# label comes from parliament.yaml vocabulary. |
||||
"chat_sessions": {"intressent_id": "person_id", "initial_talk_id": "initial_speech_id"}, |
||||
"chat_snapshots": {"intressent_id": "person_id", "initial_talk_id": "initial_speech_id"}, |
||||
} |
||||
|
||||
# --- indexes ------------------------------------------------------------------ |
||||
INDEXES: dict[str, str] = { |
||||
"talks_search_idx": "speeches_search_idx", |
||||
"talks_debate_idx": "speeches_debate_idx", |
||||
"talks_parti_idx": "speeches_party_idx", |
||||
"talks_datum_idx": "speeches_date_idx", |
||||
"talks_year_idx": "speeches_year_idx", |
||||
"talks_intressent_idx": "speeches_person_idx", |
||||
"talks_talare_idx": "speeches_speaker_idx", |
||||
"talks_dok_id_idx": "speeches_source_doc_idx", |
||||
"talks_summary_embedding_idx": "speeches_summary_embedding_idx", |
||||
"chunks_talk_idx": "speech_chunks_speech_idx", |
||||
"chunks_embedding_idx": "speech_chunks_embedding_idx", |
||||
"debates_datum_idx": "debates_date_idx", |
||||
"motions_search_idx": "documents_search_idx", |
||||
"motions_datum_idx": "documents_date_idx", |
||||
"motions_year_idx": "documents_session_year_idx", |
||||
"motions_organ_idx": "documents_committee_idx", |
||||
"motions_parties_idx": "documents_parties_idx", |
||||
"motion_authors_intressent_idx": "document_authors_person_idx", |
||||
"motion_chunks_motion_idx": "document_chunks_doc_idx", |
||||
"motion_chunks_embedding_idx": "document_chunks_embedding_idx", |
||||
"motion_yrkanden_dok_idx": "document_proposals_doc_idx", |
||||
"motion_yrkanden_embedding_idx": "document_proposals_embedding_idx", |
||||
} |
||||
|
||||
# --- identifiers safe to rewrite globally in code ----------------------------- |
||||
# Names that resolve to the same new name in every table they appear in. Most |
||||
# Swedish column names qualify: `datum` becomes `date` in talks, motions and |
||||
# debates alike, so a global rewrite is correct. |
||||
# |
||||
# Two are deliberately absent because they genuinely differ per table and must be |
||||
# read in context: |
||||
# dok_id -> source_doc_id in talks, but doc_id in motions and its child tables |
||||
# year -> stays `year` in talks (calendar year), but becomes session_year |
||||
# in motions (derived from the session label) |
||||
GLOBAL_IDENTIFIERS: dict[str, str] = { |
||||
# same target in every table |
||||
"datum": "date", |
||||
"titel": "title", |
||||
"namn": "name", |
||||
"parti": "party", |
||||
"talk_id": "speech_id", |
||||
"motion_id": "doc_id", |
||||
"replik": "is_reply", |
||||
"organ": "committee", |
||||
"rm": "session_label", |
||||
"lydelse": "text", |
||||
"subtyp": "subtype", |
||||
"utskottet": "committee_recommendation", |
||||
"kammaren": "chamber_decision", |
||||
"roll": "role", |
||||
"nummer": "number", |
||||
"period": "session_year", |
||||
"forslag": "proposals_raw", |
||||
"aktiv": "active", |
||||
"kon": "gender", |
||||
"iort": "home_town", |
||||
"sourceid": "source_id", |
||||
"hangar_guid": "source_record_guid", |
||||
"hangar_id": "source_record_id", |
||||
"intressent_id": "person_id", |
||||
"intressent_ids": "person_ids", |
||||
"anforandetext": "text", |
||||
"anforande_nummer": "sequence", |
||||
"anforande_id": "source_speech_id", |
||||
"avsnittsrubrik": "section_title", |
||||
"kammaraktivitet": "activity_type", |
||||
"talare": "speaker_name", |
||||
"valkrets": "constituency", |
||||
"efternamn": "last_name", |
||||
"tilltalsnamn": "first_name", |
||||
"sorteringsnamn": "sort_name", |
||||
"fodd_ar": "birth_year", |
||||
"personuppdrag": "assignments", |
||||
"personuppgift": "contact_details", |
||||
"partibet": "party", |
||||
"num_yrkanden": "num_proposals", |
||||
"motion_yrkanden": "document_proposals", |
||||
"motion_authors": "document_authors", |
||||
"motion_chunks": "document_chunks", |
||||
"forslag_text": "proposals_text", |
||||
"dokument_url_text": "url_text", |
||||
"dokument_url_html": "url_html", |
||||
"undertitel": "subtitle", |
||||
"beteckning": "designation", |
||||
"systemdatum": "source_updated_at", |
||||
"publicerad": "published_at", |
||||
"bilagor": "attachments", |
||||
"behandlas_i": "handled_in", |
||||
"debateurl": "url_video", |
||||
"audiofileurl": "url_audio_file", |
||||
"startpos": "audio_start_seconds", |
||||
"bild_url_80": "image_url_small", |
||||
"bild_url_192": "image_url_medium", |
||||
"bild_url_max": "image_url_large", |
||||
"person_url_xml": "source_url", |
||||
} |
||||
|
||||
|
||||
def legacy_sql_map() -> dict[str, str]: |
||||
"""Old identifier -> new, for rewriting SQL replayed from saved snapshots. |
||||
|
||||
Table names come first so `motion_yrkanden` is not partially rewritten by a |
||||
column rule before its table rule fires. |
||||
""" |
||||
out: dict[str, str] = dict(TABLES) |
||||
out.update(GLOBAL_IDENTIFIERS) |
||||
return out |
||||
@ -1,148 +0,0 @@ |
||||
#!/usr/bin/env bash |
||||
# ============================================================================= |
||||
# recovery.sh — full data recovery pipeline |
||||
# |
||||
# Run inside a screen session: |
||||
# screen -S recovery |
||||
# bash recovery.sh 2>&1 | tee logs/recovery.log |
||||
# |
||||
# Steps (ordered, each depends on the previous): |
||||
# 1. Load raw speeches from disk JSON files (ingest.cli load) |
||||
# 2. Enrich with summaries+tags from talks_training in ArangoDB |
||||
# 3. Assign debate IDs to talks that don't have one |
||||
# 4. Migrate debates from ArangoDB |
||||
# 5. Migrate chunks (2.3M rows with embeddings) from ArangoDB ← hours |
||||
# 6. Launch summarize_and_tag.py in the background ← days |
||||
# ============================================================================= |
||||
|
||||
set -euo pipefail |
||||
|
||||
REPO="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" |
||||
LOG_DIR="$REPO/logs" |
||||
PYTHON="python" |
||||
|
||||
cd "$REPO" |
||||
mkdir -p "$LOG_DIR" |
||||
|
||||
# ── colour helpers ──────────────────────────────────────────────────────────── |
||||
GREEN='\033[0;32m'; YELLOW='\033[1;33m'; RED='\033[0;31m'; NC='\033[0m' |
||||
ok() { echo -e "${GREEN}[✓] $*${NC}"; } |
||||
warn() { echo -e "${YELLOW}[!] $*${NC}"; } |
||||
fail() { echo -e "${RED}[✗] $*${NC}"; exit 1; } |
||||
step() { echo; echo -e "${YELLOW}══════════════════════════════════════════${NC}"; \ |
||||
echo -e "${YELLOW} $*${NC}"; \ |
||||
echo -e "${YELLOW}══════════════════════════════════════════${NC}"; } |
||||
|
||||
# ── sanity checks ───────────────────────────────────────────────────────────── |
||||
step "Pre-flight checks" |
||||
|
||||
[[ -f "$REPO/.env" ]] || warn ".env not found — relying on shell environment" |
||||
[[ -d "$REPO/talks" ]] || fail "talks/ directory not found" |
||||
|
||||
$PYTHON -c "import psycopg2, pgvector" 2>/dev/null || fail "psycopg2 / pgvector not installed" |
||||
$PYTHON -c "from postgres_client import pg; rows = pg.execute('SELECT 1'); print('PostgreSQL: OK')" \ |
||||
|| fail "Cannot connect to PostgreSQL" |
||||
|
||||
ok "Pre-flight passed" |
||||
|
||||
# ============================================================================= |
||||
# STEP 1 — Load all talks from disk JSON files |
||||
# ============================================================================= |
||||
step "Step 1 — Load speeches from disk (ingest.cli load)" |
||||
|
||||
$PYTHON -m ingest.cli load --source speeches \ |
||||
2>&1 | tee "$LOG_DIR/step1_documents_to_postgres.log" |
||||
|
||||
ok "Step 1 done" |
||||
|
||||
# ============================================================================= |
||||
# STEP 2 — Enrich with summaries + tags from ArangoDB talks_training |
||||
# ============================================================================= |
||||
step "Step 2 — Import talks_training from ArangoDB" |
||||
|
||||
$PYTHON scripts/migrate_arango_to_postgres.py --collection talks_training \ |
||||
2>&1 | tee "$LOG_DIR/step2_talks_training.log" |
||||
|
||||
ok "Step 2 done" |
||||
|
||||
# ============================================================================= |
||||
# STEP 3 — Assign debate IDs |
||||
# ============================================================================= |
||||
step "Step 3 — Assign debate IDs" |
||||
|
||||
$PYTHON - <<'PYEOF' 2>&1 | tee "$LOG_DIR/step3_debate_ids.log" |
||||
import sys, pathlib |
||||
sys.path.insert(0, str(pathlib.Path.cwd())) |
||||
from scripts.debates import make_debate_ids |
||||
make_debate_ids() |
||||
print("Debate ID assignment complete.") |
||||
PYEOF |
||||
|
||||
ok "Step 3 done" |
||||
|
||||
# ============================================================================= |
||||
# STEP 4 — Migrate debates from ArangoDB |
||||
# ============================================================================= |
||||
step "Step 4 — Migrate debates from ArangoDB" |
||||
|
||||
$PYTHON scripts/migrate_arango_to_postgres.py --collection debates \ |
||||
2>&1 | tee "$LOG_DIR/step4_debates.log" |
||||
|
||||
ok "Step 4 done" |
||||
|
||||
# ============================================================================= |
||||
# STEP 5 — Migrate chunks from ArangoDB (slow — embeddings are large) |
||||
# ============================================================================= |
||||
step "Step 5 — Migrate chunks from ArangoDB (this will take several hours)" |
||||
echo "You can detach the screen session (Ctrl-A D) and come back later." |
||||
echo "Progress is printed to this terminal and logged to $LOG_DIR/step5_chunks.log" |
||||
|
||||
$PYTHON scripts/migrate_arango_to_postgres.py --collection chunks \ |
||||
2>&1 | tee "$LOG_DIR/step5_chunks.log" |
||||
|
||||
ok "Step 5 done" |
||||
|
||||
# ============================================================================= |
||||
# STEP 6 — Launch summarize_and_tag.py in the background |
||||
# ============================================================================= |
||||
step "Step 6 — Launch summarize_and_tag.py (background, may take days)" |
||||
|
||||
SATLOG="$LOG_DIR/summarize_and_tag.log" |
||||
SATPID="$LOG_DIR/summarize_and_tag.pid" |
||||
|
||||
# Kill any previous instance |
||||
if [[ -f "$SATPID" ]]; then |
||||
OLD_PID=$(cat "$SATPID") |
||||
if kill -0 "$OLD_PID" 2>/dev/null; then |
||||
warn "Killing previous summarize_and_tag.py (PID $OLD_PID)" |
||||
kill "$OLD_PID" |
||||
sleep 2 |
||||
fi |
||||
fi |
||||
|
||||
nohup $PYTHON scripts/summarize_and_tag.py >> "$SATLOG" 2>&1 & |
||||
SAT_PID=$! |
||||
echo "$SAT_PID" > "$SATPID" |
||||
|
||||
ok "summarize_and_tag.py launched (PID $SAT_PID)" |
||||
echo " Log: $SATLOG" |
||||
echo " PID file: $SATPID" |
||||
echo " Monitor: tail -f $SATLOG" |
||||
|
||||
# ============================================================================= |
||||
# DONE |
||||
# ============================================================================= |
||||
step "Recovery pipeline complete" |
||||
echo "" |
||||
echo "Verification queries to run in psql:" |
||||
echo " SELECT COUNT(*) FROM talks; -- should be ~450k" |
||||
echo " SELECT COUNT(*) FROM talks WHERE summary IS NOT NULL; -- should be ~11k+" |
||||
echo " SELECT COUNT(*) FROM chunks; -- should be ~2.3M" |
||||
echo " SELECT COUNT(*) FROM debates; -- should be ~17k" |
||||
echo "" |
||||
echo "FK orphan check:" |
||||
echo " SELECT COUNT(*) FROM chunks c" |
||||
echo " WHERE NOT EXISTS (SELECT 1 FROM talks t WHERE t.id = c.talk_id);" |
||||
echo "" |
||||
echo "Summarize progress (run any time):" |
||||
echo " SELECT COUNT(*) FROM talks WHERE summary IS NOT NULL AND tags IS NOT NULL;" |
||||
@ -1,367 +0,0 @@ |
||||
""" |
||||
Language correction pipeline for riksdag talk arguments. |
||||
|
||||
Reads speeches 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 |
||||
""" |
||||
from pathlib import Path |
||||
|
||||
import json |
||||
import logging |
||||
import os |
||||
import re |
||||
import sys |
||||
import threading |
||||
import time |
||||
from concurrent.futures import ThreadPoolExecutor, as_completed |
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[1])) |
||||
import bootstrap # noqa: E402,F401 — sets cwd and sys.path to the project root |
||||
|
||||
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: |
||||
speaker_name = (talk.get("speaker_name") or "Okänd").strip() |
||||
party = (talk.get("party") or "").strip() |
||||
text = talk.get("text", "") |
||||
return ( |
||||
f"Dessa argument var obegripliga. Här är det fullständiga anförandet som referens:\n\n" |
||||
f"Talare: {speaker_name} ({party})\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 speeches SET arguments_corrected = TRUE WHERE id = %s", |
||||
(talk["id"],), |
||||
) |
||||
return True, False |
||||
|
||||
pg.execute_void( |
||||
"UPDATE speeches 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, text, speaker_name, party, arguments |
||||
FROM speeches |
||||
WHERE array_length(arguments, 1) > 0 |
||||
AND date >= '2002-01-01' |
||||
AND arguments_corrected IS NOT TRUE |
||||
AND text IS NOT NULL |
||||
ORDER BY date DESC NULLS LAST |
||||
LIMIT %s |
||||
""", |
||||
(BATCH_SIZE,), |
||||
) |
||||
|
||||
|
||||
def ensure_schema(): |
||||
try: |
||||
pg.execute_void( |
||||
"ALTER TABLE speeches 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 speeches |
||||
WHERE array_length(arguments, 1) > 0 |
||||
AND date >= '2002-01-01' |
||||
AND text 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)} speeches → {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 speeches 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() |
||||
@ -1,93 +0,0 @@ |
||||
#!/bin/bash |
||||
# Deploy debate summary embeddings feature. |
||||
# |
||||
# Steps: |
||||
# 1. Add summary_embedding column + HNSW index to debates table |
||||
# 2. Fix broken talk_ids in old ArangoDB-migrated debate rows |
||||
# 3. Generate debate summaries for debates whose talks are all summarized |
||||
# (runs debates.py once, then exits — use screen for continuous mode) |
||||
# 4. Backfill summary_embedding for all debates with a summary |
||||
# |
||||
# Usage: |
||||
# bash scripts/deploy_debate_embeddings.sh |
||||
|
||||
set -e |
||||
cd "$(dirname "$0")/.." |
||||
|
||||
# Load PG credentials from .env (skip lines that aren't simple KEY=VALUE) |
||||
while IFS='=' read -r key value; do |
||||
[[ "$key" =~ ^[[:space:]]*# ]] && continue |
||||
[[ -z "$key" ]] && continue |
||||
[[ "$key" =~ ^PG_ ]] && export "$key=$value" |
||||
done < .env |
||||
export PGPASSWORD="${PG_PASSWORD}" |
||||
|
||||
PSQL="psql -h ${PG_HOST:-localhost} -U ${PG_USER:-riksdagen} -d ${PG_DB:-riksdagen}" |
||||
|
||||
echo "=== Step 1: Add summary_embedding column to debates ===" |
||||
$PSQL < _postgres/migrations/add_debate_summary_embedding.sql |
||||
echo "Migration applied." |
||||
|
||||
echo "" |
||||
echo "=== Step 2: Fix broken talk_ids in existing debate rows ===" |
||||
$PSQL < _postgres/migrations/fix_debate_talk_ids.sql |
||||
echo "talk_ids fixed." |
||||
|
||||
echo "" |
||||
echo "=== Step 3: Generate missing debate summaries (one pass) ===" |
||||
# Run debates.py in a subprocess that exits after one full pass. |
||||
# For continuous summarization, run: screen -S debates python scripts/debates.py |
||||
python - <<'PYEOF' |
||||
import sys, pathlib |
||||
sys.path.insert(0, str(pathlib.Path.cwd())) |
||||
|
||||
from concurrent.futures import ProcessPoolExecutor, as_completed |
||||
from scripts.debates import process_ready_debate, process_debate_date |
||||
from postgres_client import pg |
||||
|
||||
system_message = """Din uppgift är att sammanfatta debatter i Sveriges riksdag. |
||||
Du kommer först att få enskilda tal som du ska sammanfatta var för sig, efter det ska du sammanfatta hela debatten. |
||||
Sammanfattningarna ska vara på svenska och vara koncisa och informativa. |
||||
Det är viktigt att du förstår vad som är kärnan i varje tal och debatt, fokusera därför på de argument och sakförhållanden som framförs. |
||||
""" |
||||
|
||||
ready = pg.execute(""" |
||||
SELECT t.debate |
||||
FROM talks t |
||||
LEFT JOIN debates d ON t.debate = d.debate |
||||
WHERE t.debate IS NOT NULL AND d.debate IS NULL |
||||
GROUP BY t.debate |
||||
HAVING COUNT(t.id) = COUNT(t.summary) |
||||
AND COUNT(t.id) > 1 |
||||
ORDER BY t.debate |
||||
""") |
||||
ready_ids = [row["debate"] for row in ready] |
||||
print(f"Found {len(ready_ids)} ready debates to summarize.") |
||||
|
||||
with ProcessPoolExecutor(max_workers=4) as executor: |
||||
futures = { |
||||
executor.submit(process_ready_debate, did, system_message): did |
||||
for did in ready_ids |
||||
} |
||||
done = 0 |
||||
for future in as_completed(futures): |
||||
did = futures[future] |
||||
try: |
||||
future.result() |
||||
done += 1 |
||||
if done % 500 == 0: |
||||
print(f" {done}/{len(ready_ids)} done") |
||||
except Exception as exc: |
||||
print(f"Error on {did}: {exc}") |
||||
print(f"Done. Processed {done} debates.") |
||||
PYEOF |
||||
|
||||
echo "" |
||||
echo "=== Step 4: Backfill debate summary embeddings ===" |
||||
python scripts/embed_debate_summaries.py |
||||
|
||||
echo "" |
||||
echo "=== Done! ===" |
||||
echo "Verify with:" |
||||
echo " psql ... -c \"SELECT COUNT(*) FROM debates WHERE summary_embedding IS NOT NULL\"" |
||||
echo " psql ... -c \"SELECT COUNT(*) FROM debates\"" |
||||
@ -1,43 +0,0 @@ |
||||
#!/bin/bash |
||||
# Deploy summary embeddings feature. |
||||
# |
||||
# Steps: |
||||
# 1. Apply DB migration (adds summary_embedding column + HNSW index) |
||||
# 2. Restart summarize_and_tag.py in the 'recovery' screen session |
||||
# 3. Run the backfill script (embeds existing summaries) |
||||
# |
||||
# Usage: |
||||
# bash scripts/deploy_summary_embeddings.sh |
||||
|
||||
set -e |
||||
cd "$(dirname "$0")/.." |
||||
|
||||
# Load PG credentials from .env (skip lines that aren't simple KEY=VALUE) |
||||
while IFS='=' read -r key value; do |
||||
[[ "$key" =~ ^[[:space:]]*# ]] && continue |
||||
[[ -z "$key" ]] && continue |
||||
[[ "$key" =~ ^PG_ ]] && export "$key=$value" |
||||
done < .env |
||||
export PGPASSWORD="${PG_PASSWORD}" |
||||
|
||||
echo "=== Step 1: Apply DB migration ===" |
||||
PSQL="psql -h ${PG_HOST:-localhost} -U ${PG_USER:-riksdagen} -d ${PG_DB:-riksdagen}" |
||||
$PSQL < _postgres/migrations/add_summary_embedding.sql |
||||
echo "Migration applied." |
||||
|
||||
echo "" |
||||
echo "=== Step 2: Restart summarize_and_tag in screen 'recovery' ===" |
||||
screen -S recovery -X stuff $'\009' # send Ctrl+C to stop the running script |
||||
sleep 3 |
||||
screen -S recovery -X stuff "python scripts/summarize_and_tag.py\n" |
||||
echo "summarize_and_tag.py restarted in screen 'recovery'." |
||||
|
||||
echo "" |
||||
echo "=== Step 3: Backfill existing summaries ===" |
||||
python scripts/embed_summaries.py |
||||
|
||||
echo "" |
||||
echo "=== Done! ===" |
||||
echo "Verify with:" |
||||
echo " psql -h \${PG_HOST:-localhost} -U \${PG_USER:-riksdagen} -d \${PG_DB:-riksdagen} \\" |
||||
echo " -c \"SELECT COUNT(*) FROM talks WHERE summary_embedding IS NOT NULL\"" |
||||
@ -1,228 +0,0 @@ |
||||
"""Generate the Swedish -> English migration and its rollback from rename_map.py. |
||||
|
||||
Writing these by hand invites a typo in one direction that the other does not |
||||
mirror. Generating both from one map means the rollback is correct by construction. |
||||
|
||||
python scripts/generate_rename_migration.py |
||||
""" |
||||
from __future__ import annotations |
||||
|
||||
import sys |
||||
from pathlib import Path |
||||
|
||||
ROOT = Path(__file__).resolve().parents[1] |
||||
sys.path.insert(0, str(ROOT)) |
||||
|
||||
from _postgres.rename_map import COLUMNS, INDEXES, TABLES # noqa: E402 |
||||
|
||||
OUT = ROOT / "_postgres" / "migrations" |
||||
STAMP = "20260803_01" |
||||
|
||||
HEADER = """\ |
||||
-- {title} |
||||
-- |
||||
-- Generated by scripts/generate_rename_migration.py from _postgres/rename_map.py. |
||||
-- Do not edit by hand; edit the map and regenerate. |
||||
-- |
||||
-- ALTER TABLE ... RENAME is catalog-only in PostgreSQL: no table rewrite, no data |
||||
-- movement, no index rebuild. Wall time is milliseconds regardless of row count, |
||||
-- and the HNSW vector indexes are not touched. |
||||
-- |
||||
-- Guarded by an existence check, so this is a no-op on a database created from the |
||||
-- current schema.sql and does the work on one created before the rename. One file |
||||
-- serves both a fresh install and an existing deployment. |
||||
|
||||
BEGIN; |
||||
|
||||
SET LOCAL lock_timeout = '5s'; |
||||
SET LOCAL statement_timeout = 0; |
||||
|
||||
DO $rename$ |
||||
BEGIN |
||||
IF {guard} THEN |
||||
""" |
||||
|
||||
FOOTER = """\ |
||||
END IF; |
||||
END |
||||
$rename$; |
||||
|
||||
{triggers} |
||||
COMMIT; |
||||
""" |
||||
|
||||
# Trigger function bodies are stored as opaque text and do NOT follow renames. |
||||
# They would compile fine and then fail at the next INSERT with |
||||
# 'record "new" has no field "anforandetext"'. Recreated unconditionally. |
||||
TRIGGERS_FORWARD = """\ |
||||
-- Recreate the search-vector triggers against the new column names. |
||||
-- |
||||
-- The text-search configuration is read from a database setting rather than |
||||
-- hardcoded, so a deployment in another language does not have to patch DDL: |
||||
-- ALTER DATABASE <db> SET app.fts_config = 'swedish'; |
||||
-- It only takes effect for new sessions, so restart the API afterwards. |
||||
|
||||
DROP TRIGGER IF EXISTS talks_search_vector_trigger ON speeches; |
||||
DROP TRIGGER IF EXISTS speeches_search_vector_trigger ON speeches; |
||||
DROP FUNCTION IF EXISTS talks_search_vector_update(); |
||||
|
||||
CREATE OR REPLACE FUNCTION speeches_search_vector_update() RETURNS TRIGGER AS $fn$ |
||||
BEGIN |
||||
NEW.search_vector := to_tsvector( |
||||
COALESCE(current_setting('app.fts_config', true), 'simple')::regconfig, |
||||
coalesce(NEW.text, '')); |
||||
RETURN NEW; |
||||
END; |
||||
$fn$ LANGUAGE plpgsql; |
||||
|
||||
CREATE TRIGGER speeches_search_vector_trigger |
||||
BEFORE INSERT OR UPDATE OF text ON speeches |
||||
FOR EACH ROW EXECUTE FUNCTION speeches_search_vector_update(); |
||||
|
||||
DROP TRIGGER IF EXISTS motions_search_vector_trigger ON documents; |
||||
DROP TRIGGER IF EXISTS documents_search_vector_trigger ON documents; |
||||
DROP FUNCTION IF EXISTS motions_search_vector_update(); |
||||
|
||||
CREATE OR REPLACE FUNCTION documents_search_vector_update() RETURNS TRIGGER AS $fn$ |
||||
DECLARE cfg regconfig := COALESCE(current_setting('app.fts_config', true), 'simple')::regconfig; |
||||
BEGIN |
||||
NEW.search_vector := |
||||
setweight(to_tsvector(cfg, coalesce(NEW.title, '')), 'A') || |
||||
setweight(to_tsvector(cfg, coalesce(NEW.proposals_text, '')), 'B') || |
||||
setweight(to_tsvector(cfg, coalesce(NEW.subtitle, '')), 'C') || |
||||
setweight(to_tsvector(cfg, coalesce(NEW.text, '')), 'D'); |
||||
RETURN NEW; |
||||
END; |
||||
$fn$ LANGUAGE plpgsql; |
||||
|
||||
CREATE TRIGGER documents_search_vector_trigger |
||||
BEFORE INSERT OR UPDATE OF title, subtitle, text, proposals_text ON documents |
||||
FOR EACH ROW EXECUTE FUNCTION documents_search_vector_update(); |
||||
""" |
||||
|
||||
TRIGGERS_ROLLBACK = """\ |
||||
-- Restore the original trigger functions. |
||||
DROP TRIGGER IF EXISTS speeches_search_vector_trigger ON talks; |
||||
DROP FUNCTION IF EXISTS speeches_search_vector_update(); |
||||
|
||||
CREATE OR REPLACE FUNCTION talks_search_vector_update() RETURNS TRIGGER AS $fn$ |
||||
BEGIN |
||||
NEW.search_vector := to_tsvector('swedish', coalesce(NEW.anforandetext, '')); |
||||
RETURN NEW; |
||||
END; |
||||
$fn$ LANGUAGE plpgsql; |
||||
|
||||
DROP TRIGGER IF EXISTS talks_search_vector_trigger ON talks; |
||||
CREATE TRIGGER talks_search_vector_trigger |
||||
BEFORE INSERT OR UPDATE OF anforandetext ON talks |
||||
FOR EACH ROW EXECUTE FUNCTION talks_search_vector_update(); |
||||
|
||||
DROP TRIGGER IF EXISTS documents_search_vector_trigger ON motions; |
||||
DROP FUNCTION IF EXISTS documents_search_vector_update(); |
||||
|
||||
CREATE OR REPLACE FUNCTION motions_search_vector_update() RETURNS TRIGGER AS $fn$ |
||||
BEGIN |
||||
NEW.search_vector := |
||||
setweight(to_tsvector('swedish', coalesce(NEW.titel, '')), 'A') || |
||||
setweight(to_tsvector('swedish', coalesce(NEW.forslag_text, '')), 'B') || |
||||
setweight(to_tsvector('swedish', coalesce(NEW.undertitel, '')), 'C') || |
||||
setweight(to_tsvector('swedish', coalesce(NEW.text, '')), 'D'); |
||||
RETURN NEW; |
||||
END; |
||||
$fn$ LANGUAGE plpgsql; |
||||
|
||||
DROP TRIGGER IF EXISTS motions_search_vector_trigger ON motions; |
||||
CREATE TRIGGER motions_search_vector_trigger |
||||
BEFORE INSERT OR UPDATE OF titel, undertitel, text, forslag_text ON motions |
||||
FOR EACH ROW EXECUTE FUNCTION motions_search_vector_update(); |
||||
""" |
||||
|
||||
|
||||
def build(forward: bool) -> str: |
||||
"""Emit the migration in one direction. |
||||
|
||||
Column renames run before table renames going forward (the columns are still on |
||||
the old table), and after them going back. |
||||
""" |
||||
lines: list[str] = [] |
||||
ind = " " |
||||
|
||||
def col_statements() -> list[str]: |
||||
# Both directions address the table by its ORIGINAL name. Going forward, |
||||
# columns are renamed before the table is; going back, the table has already |
||||
# been renamed to its original name by the time these run. Getting this |
||||
# wrong produces a rollback that fails halfway, mid-transaction. |
||||
out = [] |
||||
for table, cols in COLUMNS.items(): |
||||
for old, new in cols.items(): |
||||
a, b = (old, new) if forward else (new, old) |
||||
out.append(f"{ind}ALTER TABLE {table} RENAME COLUMN {a} TO {b};") |
||||
return out |
||||
|
||||
def table_statements() -> list[str]: |
||||
out = [] |
||||
for old, new in TABLES.items(): |
||||
a, b = (old, new) if forward else (new, old) |
||||
out.append(f"{ind}ALTER TABLE {a} RENAME TO {b};") |
||||
return out |
||||
|
||||
if forward: |
||||
lines += [" -- columns (still on their original tables)"] |
||||
lines += col_statements() |
||||
lines += ["", " -- tables"] |
||||
lines += table_statements() |
||||
lines += [ |
||||
"", |
||||
" -- New: lets bills, written questions and committee reports share this", |
||||
" -- table instead of needing one table per document type. ADD COLUMN with a", |
||||
" -- constant default is metadata-only in PostgreSQL 11+.", |
||||
f"{ind}ALTER TABLE documents ADD COLUMN IF NOT EXISTS doc_type TEXT NOT NULL DEFAULT 'motion';", |
||||
] |
||||
lines += ["", " -- indexes (cosmetic, but keeps pg_dump diffs readable)"] |
||||
for old, new in INDEXES.items(): |
||||
lines.append(f"{ind}ALTER INDEX IF EXISTS {old} RENAME TO {new};") |
||||
else: |
||||
lines += [" -- indexes"] |
||||
for old, new in INDEXES.items(): |
||||
lines.append(f"{ind}ALTER INDEX IF EXISTS {new} RENAME TO {old};") |
||||
lines += ["", " -- tables", f"{ind}ALTER TABLE documents DROP COLUMN IF EXISTS doc_type;"] |
||||
lines += table_statements() |
||||
lines += ["", " -- columns"] |
||||
lines += col_statements() |
||||
|
||||
guard = ( |
||||
"EXISTS (SELECT 1 FROM information_schema.columns\n" |
||||
" WHERE table_schema = 'public' AND table_name = 'talks'\n" |
||||
" AND column_name = 'anforandetext')" |
||||
if forward |
||||
else "EXISTS (SELECT 1 FROM information_schema.columns\n" |
||||
" WHERE table_schema = 'public' AND table_name = 'speeches'\n" |
||||
" AND column_name = 'text')" |
||||
) |
||||
title = ( |
||||
"Rename the Swedish schema to neutral English" |
||||
if forward |
||||
else "ROLLBACK: restore the Swedish schema" |
||||
) |
||||
return ( |
||||
HEADER.format(title=title, guard=guard) |
||||
+ "\n".join(lines) |
||||
+ "\n" |
||||
+ FOOTER.format(triggers=TRIGGERS_FORWARD if forward else TRIGGERS_ROLLBACK) |
||||
) |
||||
|
||||
|
||||
def main() -> int: |
||||
OUT.mkdir(parents=True, exist_ok=True) |
||||
for forward, name in ((True, f"{STAMP}_rename_to_english.sql"), |
||||
(False, f"{STAMP}_rename_to_english_ROLLBACK.sql")): |
||||
path = OUT / name |
||||
path.write_text(build(forward), encoding="utf-8") |
||||
print(f" wrote {path.relative_to(ROOT)} ({len(path.read_text().splitlines())} lines)") |
||||
total = sum(len(c) for c in COLUMNS.values()) |
||||
print(f"\n {total} column renames, {len(TABLES)} table renames, {len(INDEXES)} index renames") |
||||
return 0 |
||||
|
||||
|
||||
if __name__ == "__main__": |
||||
sys.exit(main()) |
||||
Loading…
Reference in new issue