import os import json import sqlite3 import uuid import logging from typing import Optional from src.config import settings from src.models.paper import Paper from src.models.claim import Claim, Entity, ClaimType, Polarity, StudyDesign from src.models.contradiction import ContradictionPair, ContradictionType logger = logging.getLogger(__name__) def validate_pmids(pmids: list[str]) -> None: """Validate that PMIDs contain only safe characters to prevent SQL injection.""" for pmid in pmids: if not pmid or not isinstance(pmid, str): raise ValueError(f"Invalid PMID: {pmid}") # PMIDs should be alphanumeric with possible hyphens/underscores if not pmid.replace('-', '').replace('_', '').isalnum(): raise ValueError(f"PMID contains invalid characters: {pmid}") def get_connection(db_path: Optional[str] = None) -> sqlite3.Connection: """Establish a connection to the SQLite database and enable foreign keys.""" if db_path is None: db_path = settings.db_path db_dir = os.path.dirname(os.path.abspath(db_path)) if db_dir: os.makedirs(db_dir, exist_ok=True) conn = sqlite3.connect(db_path) conn.execute("PRAGMA foreign_keys = ON;") conn.row_factory = sqlite3.Row return conn def init_db(db_path: Optional[str] = None) -> None: """Initialize the SQLite tables matching the schema.""" if db_path is None: db_path = settings.db_path conn = get_connection(db_path) with conn: conn.execute(""" CREATE TABLE IF NOT EXISTS papers ( pmid TEXT PRIMARY KEY, title TEXT NOT NULL, authors TEXT NOT NULL, -- JSON array of strings year INTEGER NOT NULL, journal TEXT, abstract_text TEXT NOT NULL, full_text TEXT, doi TEXT, created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ); """) conn.execute(""" CREATE TABLE IF NOT EXISTS claims ( id TEXT PRIMARY KEY, -- UUID string paper_id TEXT NOT NULL REFERENCES papers(pmid) ON DELETE CASCADE, text TEXT NOT NULL, normalized_text TEXT, polarity TEXT NOT NULL, population TEXT, context TEXT, section TEXT, -- Section of the paper quote_anchor TEXT NOT NULL, claim_type TEXT NOT NULL, study_design TEXT, confidence_score REAL NOT NULL, is_primary_finding BOOLEAN DEFAULT TRUE, sample_size INTEGER, entities TEXT, -- JSON array of Entity dicts created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ); """) conn.execute(""" CREATE TABLE IF NOT EXISTS contradictions ( id TEXT PRIMARY KEY, -- UUID string claim_a_id TEXT NOT NULL REFERENCES claims(id) ON DELETE CASCADE, claim_b_id TEXT NOT NULL REFERENCES claims(id) ON DELETE CASCADE, contradiction_score REAL NOT NULL, contradiction_type TEXT NOT NULL, explanation TEXT, scope_note TEXT, temporal_resolution TEXT, is_genuine BOOLEAN NOT NULL, created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ); """) conn.execute(""" CREATE TABLE IF NOT EXISTS pipeline_runs ( id TEXT PRIMARY KEY, query TEXT NOT NULL, status TEXT NOT NULL, -- RUNNING | COMPLETED | FAILED papers_fetched INTEGER, claims_extracted INTEGER, contradictions_found INTEGER, started_at TIMESTAMP, completed_at TIMESTAMP, error_message TEXT, pmids TEXT, -- JSON array of PMID strings report_json TEXT, -- Serialized SynthesisReport status_message TEXT, total_papers INTEGER, papers_extracted INTEGER, nli_pairs_total INTEGER, nli_pairs_scored INTEGER, judge_pairs_total INTEGER, judge_pairs_scored INTEGER, seed_claim TEXT, date_from INTEGER, date_to INTEGER, journals TEXT ); """) # Run automatic migrations for existing databases with conn: try: conn.execute("ALTER TABLE pipeline_runs ADD COLUMN pmids TEXT;") except sqlite3.OperationalError: pass try: conn.execute("ALTER TABLE pipeline_runs ADD COLUMN report_json TEXT;") except sqlite3.OperationalError: pass # New columns for progress tracking and seed claims new_cols = [ ("status_message", "TEXT"), ("total_papers", "INTEGER"), ("papers_extracted", "INTEGER"), ("nli_pairs_total", "INTEGER"), ("nli_pairs_scored", "INTEGER"), ("judge_pairs_total", "INTEGER"), ("judge_pairs_scored", "INTEGER"), ("seed_claim", "TEXT"), ("date_from", "INTEGER"), ("date_to", "INTEGER"), ("journals", "TEXT") ] for col_name, col_type in new_cols: try: conn.execute(f"ALTER TABLE pipeline_runs ADD COLUMN {col_name} {col_type};") except sqlite3.OperationalError: pass try: conn.execute("ALTER TABLE claims ADD COLUMN embedding BLOB;") except sqlite3.OperationalError: pass conn.close() logger.info(f"Database initialized at: {db_path}") def save_papers(papers: list[Paper], db_path: Optional[str] = None) -> None: """Save a list of Paper objects to the database.""" if db_path is None: db_path = settings.db_path conn = get_connection(db_path) with conn: for paper in papers: conn.execute(""" INSERT OR REPLACE INTO papers (pmid, title, authors, year, journal, abstract_text, full_text, doi) VALUES (?, ?, ?, ?, ?, ?, ?, ?) """, ( paper.pmid, paper.title, json.dumps(paper.authors), paper.year, paper.journal, paper.abstract_text, paper.full_text, paper.doi )) conn.close() def save_claims(claims: list[Claim], db_path: Optional[str] = None) -> None: """Save a list of Claim objects to the database.""" if db_path is None: db_path = settings.db_path conn = get_connection(db_path) with conn: for claim in claims: emb_bytes = None if claim.embedding is not None: import numpy as np emb_bytes = np.asarray(claim.embedding, dtype=np.float32).tobytes() conn.execute(""" INSERT OR REPLACE INTO claims ( id, paper_id, text, normalized_text, polarity, population, context, section, quote_anchor, claim_type, study_design, confidence_score, is_primary_finding, sample_size, entities, embedding ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """, ( str(claim.id), claim.paper_id, claim.text, claim.normalized_text, claim.polarity.value, claim.population, claim.context, claim.section, claim.quote_anchor, claim.claim_type.value, claim.study_design.value, claim.confidence_score, 1 if claim.is_primary_finding else 0, claim.sample_size, json.dumps([e.model_dump() for e in claim.entities]), emb_bytes )) conn.close() def save_contradictions(contradictions: list[ContradictionPair], db_path: Optional[str] = None) -> list[str]: """Save a list of ContradictionPair objects to the database and return their assigned UUIDs.""" if db_path is None: db_path = settings.db_path conn = get_connection(db_path) assigned_ids = [] with conn: for pair in contradictions: pair_id = str(uuid.uuid4()) assigned_ids.append(pair_id) conn.execute(""" INSERT INTO contradictions ( id, claim_a_id, claim_b_id, contradiction_score, contradiction_type, explanation, scope_note, temporal_resolution, is_genuine ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?) """, ( pair_id, str(pair.claim_a.id), str(pair.claim_b.id), pair.contradiction_score, pair.contradiction_type.value, pair.explanation, pair.scope_note, pair.temporal_resolution, 1 if pair.is_genuine else 0 )) conn.close() return assigned_ids def save_pipeline_run( run_id: str, query: str, status: str, papers_fetched: Optional[int] = None, claims_extracted: Optional[int] = None, contradictions_found: Optional[int] = None, started_at: Optional[str] = None, completed_at: Optional[str] = None, error_message: Optional[str] = None, pmids: Optional[list[str]] = None, report_json: Optional[str] = None, status_message: Optional[str] = None, total_papers: Optional[int] = None, papers_extracted: Optional[int] = None, nli_pairs_total: Optional[int] = None, nli_pairs_scored: Optional[int] = None, judge_pairs_total: Optional[int] = None, judge_pairs_scored: Optional[int] = None, seed_claim: Optional[str] = None, date_from: Optional[int] = None, date_to: Optional[int] = None, journals: Optional[list[str]] = None, db_path: Optional[str] = None ) -> None: """Save or update a pipeline run status.""" if db_path is None: db_path = settings.db_path conn = get_connection(db_path) with conn: row = conn.execute("SELECT 1 FROM pipeline_runs WHERE id = ?", (run_id,)).fetchone() if not row: conn.execute(""" INSERT INTO pipeline_runs ( id, query, status, papers_fetched, claims_extracted, contradictions_found, started_at, completed_at, error_message, pmids, report_json, status_message, total_papers, papers_extracted, nli_pairs_total, nli_pairs_scored, judge_pairs_total, judge_pairs_scored, seed_claim, date_from, date_to, journals ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """, ( run_id, query, status, papers_fetched, claims_extracted, contradictions_found, started_at, completed_at, error_message, json.dumps(pmids) if pmids is not None else None, report_json, status_message, total_papers, papers_extracted, nli_pairs_total, nli_pairs_scored, judge_pairs_total, judge_pairs_scored, seed_claim, date_from, date_to, json.dumps(journals) if journals is not None else None )) else: update_fields = [] params = [] update_fields.append("query = ?") params.append(query) update_fields.append("status = ?") params.append(status) if papers_fetched is not None: update_fields.append("papers_fetched = ?") params.append(papers_fetched) if claims_extracted is not None: update_fields.append("claims_extracted = ?") params.append(claims_extracted) if contradictions_found is not None: update_fields.append("contradictions_found = ?") params.append(contradictions_found) if started_at is not None: update_fields.append("started_at = ?") params.append(started_at) if completed_at is not None: update_fields.append("completed_at = ?") params.append(completed_at) if error_message is not None: update_fields.append("error_message = ?") params.append(error_message) if pmids is not None: update_fields.append("pmids = ?") params.append(json.dumps(pmids)) if report_json is not None: update_fields.append("report_json = ?") params.append(report_json) if status_message is not None: update_fields.append("status_message = ?") params.append(status_message) if total_papers is not None: update_fields.append("total_papers = ?") params.append(total_papers) if papers_extracted is not None: update_fields.append("papers_extracted = ?") params.append(papers_extracted) if nli_pairs_total is not None: update_fields.append("nli_pairs_total = ?") params.append(nli_pairs_total) if nli_pairs_scored is not None: update_fields.append("nli_pairs_scored = ?") params.append(nli_pairs_scored) if judge_pairs_total is not None: update_fields.append("judge_pairs_total = ?") params.append(judge_pairs_total) if judge_pairs_scored is not None: update_fields.append("judge_pairs_scored = ?") params.append(judge_pairs_scored) if seed_claim is not None: update_fields.append("seed_claim = ?") params.append(seed_claim) if date_from is not None: update_fields.append("date_from = ?") params.append(date_from) if date_to is not None: update_fields.append("date_to = ?") params.append(date_to) if journals is not None: update_fields.append("journals = ?") params.append(json.dumps(journals)) params.append(run_id) query_str = f"UPDATE pipeline_runs SET {', '.join(update_fields)} WHERE id = ?" conn.execute(query_str, tuple(params)) conn.close() def get_pipeline_run(run_id: str, db_path: Optional[str] = None) -> Optional[dict]: """Retrieve a pipeline run record from the database by run_id.""" if db_path is None: db_path = settings.db_path conn = get_connection(db_path) try: row = conn.execute("SELECT * FROM pipeline_runs WHERE id = ?", (run_id,)).fetchone() return dict(row) if row else None finally: conn.close() def _row_to_paper(row: sqlite3.Row) -> Paper: return Paper( pmid=row["pmid"], title=row["title"], authors=json.loads(row["authors"]), year=row["year"], journal=row["journal"] or "", abstract_text=row["abstract_text"], full_text=row["full_text"], doi=row["doi"] ) def _row_to_claim(row: sqlite3.Row, conn: sqlite3.Connection) -> Claim: # Use pre-fetched authors and year if they exist in the Row (e.g. from JOIN query) if "authors" in row.keys() and "year" in row.keys(): authors = json.loads(row["authors"]) if row["authors"] else [] year = row["year"] if row["year"] is not None else None else: # Query paper to resolve authors/year since they are dynamic metadata on the Claim model paper_row = conn.execute("SELECT authors, year FROM papers WHERE pmid = ?", (row["paper_id"],)).fetchone() authors = json.loads(paper_row["authors"]) if paper_row else [] year = paper_row["year"] if paper_row else None entities_data = json.loads(row["entities"]) if row["entities"] else [] entities = [Entity(**e) for e in entities_data] embedding = None if "embedding" in row.keys() and row["embedding"] is not None: import numpy as np embedding = np.frombuffer(row["embedding"], dtype=np.float32).tolist() return Claim( id=uuid.UUID(row["id"]), text=row["text"], normalized_text=row["normalized_text"], paper_id=row["paper_id"], section=row["section"] or "Abstract", authors=authors, year=year if (year is not None and year > 0) else 0, # 0 = unknown/unparseable year; excluded from temporal & supersession logic confidence_score=row["confidence_score"], claim_type=ClaimType(row["claim_type"]), polarity=Polarity(row["polarity"]), entities=entities, population=row["population"] or "", context=row["context"] or "", quote_anchor=row["quote_anchor"], sample_size=row["sample_size"], study_design=StudyDesign(row["study_design"]), is_primary_finding=bool(row["is_primary_finding"]), embedding=embedding ) def _get_claim_with_conn(claim_id: str, conn: sqlite3.Connection) -> Optional[Claim]: """Retrieve a single Claim using an existing connection, using a JOIN to get paper metadata.""" row = conn.execute(""" SELECT c.*, p.authors, p.year FROM claims c JOIN papers p ON c.paper_id = p.pmid WHERE c.id = ? """, (claim_id,)).fetchone() if row: return _row_to_claim(row, conn) return None def get_paper(pmid: str, db_path: Optional[str] = None) -> Optional[Paper]: """Retrieve a single Paper from the database by PMID.""" if db_path is None: db_path = settings.db_path conn = get_connection(db_path) try: row = conn.execute("SELECT * FROM papers WHERE pmid = ?", (pmid,)).fetchone() return _row_to_paper(row) if row else None finally: conn.close() def get_claim(claim_id: str, db_path: Optional[str] = None) -> Optional[Claim]: """Retrieve a single Claim from the database by its ID.""" if db_path is None: db_path = settings.db_path conn = get_connection(db_path) try: return _get_claim_with_conn(claim_id, conn) finally: conn.close() def get_claims_for_run(paper_ids: list[str], db_path: Optional[str] = None) -> list[Claim]: """Retrieve all claims associated with a list of PMIDs using a JOIN to fetch paper metadata in one query.""" if not paper_ids: return [] validate_pmids(paper_ids) if db_path is None: db_path = settings.db_path conn = get_connection(db_path) try: placeholders = ",".join("?" for _ in paper_ids) rows = conn.execute(f""" SELECT c.*, p.authors, p.year FROM claims c JOIN papers p ON c.paper_id = p.pmid WHERE c.paper_id IN ({placeholders}) """, paper_ids).fetchall() return [_row_to_claim(r, conn) for r in rows] finally: conn.close() def get_papers_for_run(paper_ids: list[str], db_path: Optional[str] = None) -> list[Paper]: """Retrieve all papers associated with a list of PMIDs.""" if not paper_ids: return [] validate_pmids(paper_ids) if db_path is None: db_path = settings.db_path conn = get_connection(db_path) try: placeholders = ",".join("?" for _ in paper_ids) rows = conn.execute(f"SELECT * FROM papers WHERE pmid IN ({placeholders})", paper_ids).fetchall() return [_row_to_paper(r) for r in rows] finally: conn.close() def get_contradictions_for_run(paper_ids: list[str], db_path: Optional[str] = None) -> list[ContradictionPair]: """Retrieve all contradiction pairs where both claims belong to the set of paper PMIDs.""" if not paper_ids: return [] validate_pmids(paper_ids) if db_path is None: db_path = settings.db_path conn = get_connection(db_path) try: placeholders = ",".join("?" for _ in paper_ids) query = f""" SELECT c.* FROM contradictions c JOIN claims ca ON c.claim_a_id = ca.id JOIN claims cb ON c.claim_b_id = cb.id WHERE ca.paper_id IN ({placeholders}) AND cb.paper_id IN ({placeholders}) """ params = paper_ids + paper_ids rows = conn.execute(query, params).fetchall() contradictions = [] for r in rows: claim_a = _get_claim_with_conn(r["claim_a_id"], conn) claim_b = _get_claim_with_conn(r["claim_b_id"], conn) if claim_a and claim_b: contradictions.append( ContradictionPair( claim_a=claim_a, claim_b=claim_b, contradiction_score=r["contradiction_score"], contradiction_type=ContradictionType(r["contradiction_type"]), explanation=r["explanation"], scope_note=r["scope_note"], temporal_resolution=r["temporal_resolution"], is_genuine=bool(r["is_genuine"]) ) ) return contradictions finally: conn.close()