#!/usr/bin/env python3 """ Stage 1: Extract estimated dates for documents. Sources (in priority order): 1. Filename parsing (congress session, year folders, JFK doc IDs) 2. DATE entities already in the entities table (most frequent date per doc) 3. Regex patterns in OCR text (fallback) Populates: document_dates table """ import re import logging import sys from datetime import date, datetime from collections import Counter import psycopg2 import psycopg2.extras from config import CONGRESS_DATES, BATCH_SIZE from db import get_conn, fetch_all, fetch_one logging.basicConfig( level=logging.INFO, format="%(asctime)s %(levelname)-8s %(message)s", handlers=[logging.StreamHandler(sys.stdout)], ) log = logging.getLogger(__name__) def parse_congress_from_path(file_path: str) -> int | None: """Extract congress session number from file path or filename.""" # Match patterns like congress_118, BILLS-118hr, congress_103 m = re.search(r'congress_(\d{2,3})', file_path) if m: return int(m.group(1)) m = re.search(r'BILLS-(\d{2,3})', file_path) if m: return int(m.group(1)) # Congressional Record with ordinal congress m = re.search(r'(\d{2,3})(st|nd|rd|th)\s+Congress', file_path, re.IGNORECASE) if m: return int(m.group(1)) return None def parse_year_from_path(file_path: str) -> int | None: """Extract a year from folder structure like /2021/ or /2017-2018/.""" # Folder-based year m = re.search(r'/(\d{4})(?:[_/-](\d{4}))?/', file_path) if m: return int(m.group(1)) # Year in filename m = re.search(r'[_-](\d{4})[_.-]', file_path) if m: yr = int(m.group(1)) if 1800 <= yr <= 2030: return yr return None def congress_to_date_range(session: int) -> tuple[date | None, date | None]: """Convert congress session to a date range.""" if session in CONGRESS_DATES: s, e = CONGRESS_DATES[session] return date.fromisoformat(s), date.fromisoformat(e) # Approximate: each congress starts Jan 3 of odd year # Congress 1 started 1789, session N starts 1789 + (N-1)*2 start_year = 1789 + (session - 1) * 2 if 1789 <= start_year <= 2030: return date(start_year, 1, 3), date(start_year + 2, 1, 3) return None, None def parse_date_entities(doc_id: int, conn) -> tuple[date | None, float]: """ Find the most common parseable date from DATE entities for a document. Returns (estimated_date, confidence). """ with conn.cursor() as cur: cur.execute( "SELECT entity_text FROM entities " "WHERE document_id = %s AND entity_type = 'DATE'", (doc_id,) ) rows = cur.fetchall() if not rows: return None, 0.0 year_counts = Counter() full_dates = [] for (text,) in rows: text = text.strip() # Try full date patterns for fmt in ("%B %d, %Y", "%b %d, %Y", "%m/%d/%Y", "%Y-%m-%d", "%d %B %Y"): try: dt = datetime.strptime(text, fmt).date() if 1800 <= dt.year <= 2030: full_dates.append(dt) year_counts[dt.year] += 1 break except ValueError: continue else: # Try just year m = re.search(r'\b(1[89]\d{2}|20[0-2]\d)\b', text) if m: year_counts[int(m.group(1))] += 1 if full_dates: # Return most common full date date_counts = Counter(full_dates) best_date, count = date_counts.most_common(1)[0] confidence = min(count / len(rows), 1.0) return best_date, confidence if year_counts: best_year, count = year_counts.most_common(1)[0] confidence = min(count / len(rows) * 0.5, 0.8) # lower confidence for year-only return date(best_year, 7, 1), confidence # midpoint of year return None, 0.0 def process_documents(): """Main processing loop.""" conn = get_conn() conn.autocommit = False # Get documents that don't have dates yet with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur: cur.execute(""" SELECT d.id, d.file_path, d.source_section FROM documents d LEFT JOIN document_dates dd ON dd.document_id = d.id WHERE dd.document_id IS NULL ORDER BY d.id """) docs = cur.fetchall() total = len(docs) log.info(f"Processing {total} documents for date extraction") batch = [] processed = 0 for doc in docs: doc_id = doc["id"] path = doc["file_path"] section = doc["source_section"] estimated_date = None date_source = None date_confidence = 0.0 date_range_start = None date_range_end = None congress_session = None # Priority 1: Congress session from filename congress = parse_congress_from_path(path) if congress: congress_session = congress start, end = congress_to_date_range(congress) if start and end: date_range_start = start date_range_end = end # Midpoint as estimate mid = start.toordinal() + (end.toordinal() - start.toordinal()) // 2 estimated_date = date.fromordinal(mid) date_source = "filename_congress" date_confidence = 0.7 # Priority 2: Year from folder/filename if not estimated_date: year = parse_year_from_path(path) if year: estimated_date = date(year, 7, 1) date_range_start = date(year, 1, 1) date_range_end = date(year, 12, 31) date_source = "filename_year" date_confidence = 0.6 # Priority 3: DATE entities from NER if not estimated_date: ner_date, ner_conf = parse_date_entities(doc_id, conn) if ner_date: estimated_date = ner_date date_source = "ner_entities" date_confidence = ner_conf # Priority 4: Collection-level defaults if not estimated_date: defaults = { "cia_mkultra": (date(1963, 1, 1), "collection_default", 0.3, date(1953, 1, 1), date(1973, 12, 31)), "cia_stargate": (date(1986, 1, 1), "collection_default", 0.3, date(1978, 1, 1), date(1995, 12, 31)), "lincoln_archives": (date(1865, 1, 1), "collection_default", 0.3, date(1860, 1, 1), date(1877, 12, 31)), } if section in defaults: d = defaults[section] estimated_date = d[0] date_source = d[1] date_confidence = d[2] date_range_start = d[3] date_range_end = d[4] batch.append(( doc_id, estimated_date, date_source, date_confidence, date_range_start, date_range_end, congress_session, )) if len(batch) >= BATCH_SIZE: _flush_batch(conn, batch) processed += len(batch) log.info(f"Progress: {processed}/{total} ({processed*100//total}%)") batch = [] if batch: _flush_batch(conn, batch) processed += len(batch) conn.close() log.info(f"Done. Processed {processed} documents.") # Stats stats = fetch_all(""" SELECT date_source, COUNT(*) as cnt, ROUND(AVG(date_confidence)::numeric, 2) as avg_conf FROM document_dates GROUP BY date_source ORDER BY cnt DESC """) log.info("Date extraction stats:") for row in stats: log.info(f" {row['date_source'] or 'no_date'}: {row['cnt']} docs (avg conf: {row['avg_conf']})") def _flush_batch(conn, batch): with conn.cursor() as cur: psycopg2.extras.execute_batch( cur, """INSERT INTO document_dates (document_id, estimated_date, date_source, date_confidence, date_range_start, date_range_end, congress_session) VALUES (%s, %s, %s, %s, %s, %s, %s) ON CONFLICT (document_id) DO UPDATE SET estimated_date = EXCLUDED.estimated_date, date_source = EXCLUDED.date_source, date_confidence = EXCLUDED.date_confidence, date_range_start = EXCLUDED.date_range_start, date_range_end = EXCLUDED.date_range_end, congress_session = EXCLUDED.congress_session, created_at = NOW() """, batch, page_size=500, ) conn.commit() if __name__ == "__main__": process_documents()