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