Spaces:
Sleeping
Sleeping
| """ | |
| SQLite layer. Idempotente: cada corrida hace upsert por arxiv_id sin duplicar. | |
| Decisiones de diseño: | |
| - Listas (autores, categorias, topics, code_urls) se serializan a JSON. | |
| Para volumen <1M de papers, esto es mas simple que tablas hijas y | |
| mantiene la query principal en una sola tabla. | |
| - Indices solo en columnas usadas para filtrar/ordenar: published, primary_cat, score. | |
| - updated_at se setea automaticamente en cada upsert para detectar staleness. | |
| """ | |
| from __future__ import annotations | |
| import json | |
| import sqlite3 | |
| from contextlib import contextmanager | |
| from datetime import datetime, date | |
| from pathlib import Path | |
| from typing import Any, Iterable, Iterator | |
| SCHEMA = """ | |
| CREATE TABLE IF NOT EXISTS papers ( | |
| arxiv_id TEXT PRIMARY KEY, | |
| title TEXT NOT NULL, | |
| abstract TEXT, | |
| authors_json TEXT, -- JSON array of strings | |
| categories_json TEXT, -- JSON array of strings | |
| primary_cat TEXT, | |
| published TEXT, -- ISO date YYYY-MM-DD | |
| updated TEXT, -- ISO date YYYY-MM-DD | |
| doi TEXT, | |
| pdf_url TEXT, | |
| comment TEXT, -- arXiv "comments" field (puede traer github) | |
| journal_ref TEXT, | |
| -- OpenAlex enrichment | |
| openalex_id TEXT, | |
| cited_by_count INTEGER, | |
| fwci REAL, | |
| topics_json TEXT, -- JSON array | |
| publication_year INTEGER, -- ano de publicacion oficial (journal o arXiv) | |
| openalex_abstract TEXT, -- abstract reconstruido del inverted index | |
| -- Semantic Scholar enrichment (opcional) | |
| ss_id TEXT, | |
| influential_citation_count INTEGER, | |
| -- Detected artifacts | |
| has_code INTEGER, -- 0/1 | |
| has_data INTEGER, -- 0/1 | |
| code_urls_json TEXT, -- JSON array | |
| data_urls_json TEXT, -- JSON array | |
| -- Ranking | |
| score REAL, | |
| -- Metadata | |
| fetched_at TEXT NOT NULL, -- ISO datetime | |
| updated_at TEXT NOT NULL -- ISO datetime | |
| ); | |
| CREATE INDEX IF NOT EXISTS idx_published ON papers(published); | |
| CREATE INDEX IF NOT EXISTS idx_primary_cat ON papers(primary_cat); | |
| CREATE INDEX IF NOT EXISTS idx_score ON papers(score DESC); | |
| CREATE INDEX IF NOT EXISTS idx_cited_by ON papers(cited_by_count DESC); | |
| -- Embeddings vectoriales para busqueda semantica. | |
| -- Tabla separada (no en papers) para mantener exports CSV/Parquet lean. | |
| CREATE TABLE IF NOT EXISTS embeddings ( | |
| arxiv_id TEXT PRIMARY KEY, | |
| model TEXT NOT NULL, -- nombre del modelo (ej. allenai-specter) | |
| dim INTEGER NOT NULL, | |
| vector BLOB NOT NULL, -- np.float32 bytes, ya normalizado L2 | |
| computed_at TEXT NOT NULL, | |
| FOREIGN KEY (arxiv_id) REFERENCES papers(arxiv_id) | |
| ); | |
| CREATE INDEX IF NOT EXISTS idx_emb_model ON embeddings(model); | |
| -- Tabla auxiliar: log de corridas para auditoria y debugging | |
| CREATE TABLE IF NOT EXISTS fetch_runs ( | |
| run_id INTEGER PRIMARY KEY AUTOINCREMENT, | |
| started_at TEXT NOT NULL, | |
| finished_at TEXT, | |
| category TEXT, | |
| n_fetched INTEGER, | |
| n_new INTEGER, | |
| n_updated INTEGER, | |
| status TEXT, -- 'ok' | 'partial' | 'error' | |
| error_msg TEXT | |
| ); | |
| """ | |
| # ------------------------------------------------------------ | |
| # Connection helpers | |
| # ------------------------------------------------------------ | |
| def connect(db_path: str | Path) -> sqlite3.Connection: | |
| """Abre conexion con pragmas razonables para uso single-writer.""" | |
| db_path = Path(db_path) | |
| db_path.parent.mkdir(parents=True, exist_ok=True) | |
| conn = sqlite3.connect(str(db_path)) | |
| conn.row_factory = sqlite3.Row | |
| # WAL = mejor concurrencia lectura/escritura, menos lock contention | |
| conn.execute("PRAGMA journal_mode = WAL") | |
| conn.execute("PRAGMA synchronous = NORMAL") | |
| conn.execute("PRAGMA foreign_keys = ON") | |
| return conn | |
| def init_schema(conn: sqlite3.Connection) -> None: | |
| conn.executescript(SCHEMA) | |
| migrate_schema(conn) | |
| conn.commit() | |
| def migrate_schema(conn: sqlite3.Connection) -> None: | |
| """ | |
| Migraciones idempotentes para DBs creadas con versiones viejas del schema. | |
| SQLite no soporta `ALTER TABLE ADD COLUMN IF NOT EXISTS`, asi que | |
| consultamos `PRAGMA table_info` y agregamos solo lo que falta. | |
| """ | |
| existing = {row[1] for row in conn.execute("PRAGMA table_info(papers)")} | |
| additions = { | |
| "publication_year": "INTEGER", | |
| "openalex_abstract": "TEXT", | |
| "user_rating": "INTEGER", # -1 dislike, 0/NULL neutral, +1 like | |
| "user_note": "TEXT", # nota libre del usuario | |
| "rated_at": "TEXT", # ISO datetime del ultimo rating | |
| "methodology_json": "TEXT", # JSON: [{"family","sem","lex"}] metodologias detectadas | |
| } | |
| for col, coltype in additions.items(): | |
| if col not in existing: | |
| conn.execute(f"ALTER TABLE papers ADD COLUMN {col} {coltype}") | |
| # Indice para queries de "mi biblioteca" rapidas | |
| conn.execute("CREATE INDEX IF NOT EXISTS idx_user_rating ON papers(user_rating)") | |
| def session(db_path: str | Path) -> Iterator[sqlite3.Connection]: | |
| conn = connect(db_path) | |
| try: | |
| init_schema(conn) | |
| yield conn | |
| conn.commit() | |
| except Exception: | |
| conn.rollback() | |
| raise | |
| finally: | |
| conn.close() | |
| # ------------------------------------------------------------ | |
| # Serialization helpers | |
| # ------------------------------------------------------------ | |
| def _to_iso(d: Any) -> str | None: | |
| if d is None: | |
| return None | |
| if isinstance(d, (datetime, date)): | |
| return d.isoformat() if isinstance(d, datetime) else d.isoformat() | |
| return str(d) | |
| def _jdump(x: Any) -> str | None: | |
| if x is None: | |
| return None | |
| return json.dumps(x, ensure_ascii=False) | |
| def _jload(s: str | None) -> Any: | |
| if s is None or s == "": | |
| return None | |
| return json.loads(s) | |
| # ------------------------------------------------------------ | |
| # Upserts | |
| # ------------------------------------------------------------ | |
| # Campos provistos por el fetch inicial de arXiv (sin enrichment ni score). | |
| ARXIV_FIELDS = [ | |
| "arxiv_id", "title", "abstract", "authors_json", "categories_json", | |
| "primary_cat", "published", "updated", "doi", "pdf_url", "comment", | |
| "journal_ref", | |
| ] | |
| def upsert_arxiv_paper(conn: sqlite3.Connection, paper: dict) -> str: | |
| """ | |
| Inserta o actualiza un paper con datos crudos de arXiv. | |
| No toca campos de enrichment ni score. | |
| Returns: 'inserted' | 'updated' | |
| """ | |
| now = datetime.utcnow().isoformat(timespec="seconds") | |
| row = { | |
| "arxiv_id": paper["arxiv_id"], | |
| "title": paper["title"], | |
| "abstract": paper.get("abstract"), | |
| "authors_json": _jdump(paper.get("authors") or []), | |
| "categories_json": _jdump(paper.get("categories") or []), | |
| "primary_cat": paper.get("primary_cat"), | |
| "published": _to_iso(paper.get("published")), | |
| "updated": _to_iso(paper.get("updated")), | |
| "doi": paper.get("doi"), | |
| "pdf_url": paper.get("pdf_url"), | |
| "comment": paper.get("comment"), | |
| "journal_ref": paper.get("journal_ref"), | |
| "fetched_at": now, | |
| "updated_at": now, | |
| } | |
| cur = conn.execute( | |
| "SELECT 1 FROM papers WHERE arxiv_id = ?", (row["arxiv_id"],) | |
| ) | |
| exists = cur.fetchone() is not None | |
| if exists: | |
| # UPDATE solo de campos de arXiv; preserva enrichment y score. | |
| conn.execute( | |
| """ | |
| UPDATE papers SET | |
| title=:title, abstract=:abstract, authors_json=:authors_json, | |
| categories_json=:categories_json, primary_cat=:primary_cat, | |
| published=:published, updated=:updated, doi=:doi, | |
| pdf_url=:pdf_url, comment=:comment, journal_ref=:journal_ref, | |
| updated_at=:updated_at | |
| WHERE arxiv_id=:arxiv_id | |
| """, | |
| row, | |
| ) | |
| return "updated" | |
| else: | |
| conn.execute( | |
| """ | |
| INSERT INTO papers ( | |
| arxiv_id, title, abstract, authors_json, categories_json, | |
| primary_cat, published, updated, doi, pdf_url, comment, | |
| journal_ref, fetched_at, updated_at | |
| ) VALUES ( | |
| :arxiv_id, :title, :abstract, :authors_json, :categories_json, | |
| :primary_cat, :published, :updated, :doi, :pdf_url, :comment, | |
| :journal_ref, :fetched_at, :updated_at | |
| ) | |
| """, | |
| row, | |
| ) | |
| return "inserted" | |
| def upsert_many(conn: sqlite3.Connection, papers: Iterable[dict]) -> dict: | |
| """Bulk upsert. Devuelve contadores.""" | |
| counts = {"inserted": 0, "updated": 0} | |
| for p in papers: | |
| result = upsert_arxiv_paper(conn, p) | |
| counts[result] += 1 | |
| conn.commit() | |
| return counts | |
| # ------------------------------------------------------------ | |
| # Run log | |
| # ------------------------------------------------------------ | |
| def start_run(conn: sqlite3.Connection, category: str) -> int: | |
| cur = conn.execute( | |
| "INSERT INTO fetch_runs (started_at, category, status) VALUES (?, ?, 'running')", | |
| (datetime.utcnow().isoformat(timespec="seconds"), category), | |
| ) | |
| conn.commit() | |
| return cur.lastrowid | |
| def finish_run( | |
| conn: sqlite3.Connection, | |
| run_id: int, | |
| n_fetched: int, | |
| n_new: int, | |
| n_updated: int, | |
| status: str = "ok", | |
| error_msg: str | None = None, | |
| ) -> None: | |
| conn.execute( | |
| """ | |
| UPDATE fetch_runs SET | |
| finished_at = ?, n_fetched = ?, n_new = ?, n_updated = ?, | |
| status = ?, error_msg = ? | |
| WHERE run_id = ? | |
| """, | |
| ( | |
| datetime.utcnow().isoformat(timespec="seconds"), | |
| n_fetched, n_new, n_updated, status, error_msg, run_id, | |
| ), | |
| ) | |
| conn.commit() | |
| # ------------------------------------------------------------ | |
| # Read helpers (utiles para los siguientes modulos) | |
| # ------------------------------------------------------------ | |
| def count_papers(conn: sqlite3.Connection) -> int: | |
| return conn.execute("SELECT COUNT(*) FROM papers").fetchone()[0] | |
| def papers_missing_openalex(conn: sqlite3.Connection, limit: int | None = None) -> list[sqlite3.Row]: | |
| sql = ( | |
| "SELECT arxiv_id, doi, title FROM papers " | |
| "WHERE openalex_id IS NULL ORDER BY published DESC" | |
| ) | |
| if limit: | |
| sql += f" LIMIT {int(limit)}" | |
| return conn.execute(sql).fetchall() | |
| # ------------------------------------------------------------ | |
| # Enrichment upserts | |
| # ------------------------------------------------------------ | |
| def update_openalex( | |
| conn: sqlite3.Connection, | |
| arxiv_id: str, | |
| openalex_id: str, | |
| cited_by_count: int | None, | |
| fwci: float | None, | |
| topics: list[dict] | None, | |
| publication_year: int | None = None, | |
| openalex_abstract: str | None = None, | |
| ) -> None: | |
| """Rellena los campos OpenAlex de un paper ya existente.""" | |
| conn.execute( | |
| """ | |
| UPDATE papers SET | |
| openalex_id = ?, | |
| cited_by_count = ?, | |
| fwci = ?, | |
| topics_json = ?, | |
| publication_year = ?, | |
| openalex_abstract = ?, | |
| updated_at = ? | |
| WHERE arxiv_id = ? | |
| """, | |
| ( | |
| openalex_id, | |
| cited_by_count, | |
| fwci, | |
| _jdump(topics) if topics else None, | |
| publication_year, | |
| openalex_abstract, | |
| datetime.utcnow().isoformat(timespec="seconds"), | |
| arxiv_id, | |
| ), | |
| ) | |
| def update_methodology( | |
| conn: sqlite3.Connection, | |
| arxiv_id: str, | |
| tags: list[dict] | None, | |
| ) -> None: | |
| """ | |
| Persiste las metodologias detectadas para un paper. | |
| tags = [{"family": str, "sem": float, "lex": int}, ...] ordenadas por confianza. | |
| No hace commit (el caller comitea por chunk). | |
| """ | |
| conn.execute( | |
| "UPDATE papers SET methodology_json = ?, updated_at = ? WHERE arxiv_id = ?", | |
| ( | |
| _jdump(tags) if tags else None, | |
| datetime.utcnow().isoformat(timespec="seconds"), | |
| arxiv_id, | |
| ), | |
| ) | |
| def all_methodology_families(conn: sqlite3.Connection) -> list[str]: | |
| """Lista deduplicada de familias presentes en el corpus (para poblar filtros).""" | |
| fams: set[str] = set() | |
| for (mj,) in conn.execute( | |
| "SELECT methodology_json FROM papers WHERE methodology_json IS NOT NULL" | |
| ): | |
| try: | |
| for t in json.loads(mj): | |
| f = t.get("family") | |
| if f: | |
| fams.add(f) | |
| except Exception: | |
| pass | |
| return sorted(fams) | |
| def update_artifacts( | |
| conn: sqlite3.Connection, | |
| arxiv_id: str, | |
| has_code: bool, | |
| has_data: bool, | |
| code_urls: list[str], | |
| data_urls: list[str], | |
| ) -> None: | |
| """Persiste el resultado del detector de artefactos.""" | |
| conn.execute( | |
| """ | |
| UPDATE papers SET | |
| has_code = ?, | |
| has_data = ?, | |
| code_urls_json = ?, | |
| data_urls_json = ?, | |
| updated_at = ? | |
| WHERE arxiv_id = ? | |
| """, | |
| ( | |
| int(has_code), | |
| int(has_data), | |
| _jdump(code_urls) if code_urls else None, | |
| _jdump(data_urls) if data_urls else None, | |
| datetime.utcnow().isoformat(timespec="seconds"), | |
| arxiv_id, | |
| ), | |
| ) | |
| # ------------------------------------------------------------ | |
| # User ratings (likes / dislikes / notes) | |
| # ------------------------------------------------------------ | |
| def rate_paper( | |
| conn: sqlite3.Connection, | |
| arxiv_id: str, | |
| rating: int, # -1 dislike, 0 limpia, +1 like | |
| note: str | None = None, | |
| ) -> None: | |
| """Setea/limpia el rating del usuario. rating=0 limpia.""" | |
| now = datetime.utcnow().isoformat(timespec="seconds") | |
| if rating == 0: | |
| conn.execute( | |
| "UPDATE papers SET user_rating = NULL, rated_at = ? WHERE arxiv_id = ?", | |
| (now, arxiv_id), | |
| ) | |
| else: | |
| conn.execute( | |
| "UPDATE papers SET user_rating = ?, user_note = COALESCE(?, user_note), " | |
| "rated_at = ? WHERE arxiv_id = ?", | |
| (rating, note, now, arxiv_id), | |
| ) | |
| conn.commit() | |
| def set_note(conn: sqlite3.Connection, arxiv_id: str, note: str) -> None: | |
| """Guarda solo la nota libre, sin tocar el rating.""" | |
| now = datetime.utcnow().isoformat(timespec="seconds") | |
| conn.execute( | |
| "UPDATE papers SET user_note = ?, rated_at = ? WHERE arxiv_id = ?", | |
| (note or None, now, arxiv_id), | |
| ) | |
| conn.commit() | |
| # ------------------------------------------------------------ | |
| # Embeddings | |
| # ------------------------------------------------------------ | |
| def upsert_embedding( | |
| conn: sqlite3.Connection, | |
| arxiv_id: str, | |
| model: str, | |
| dim: int, | |
| vector_bytes: bytes, | |
| ) -> None: | |
| """Persiste un embedding ya normalizado L2 como BLOB float32.""" | |
| conn.execute( | |
| """ | |
| INSERT INTO embeddings (arxiv_id, model, dim, vector, computed_at) | |
| VALUES (?, ?, ?, ?, ?) | |
| ON CONFLICT(arxiv_id) DO UPDATE SET | |
| model = excluded.model, | |
| dim = excluded.dim, | |
| vector = excluded.vector, | |
| computed_at = excluded.computed_at | |
| """, | |
| ( | |
| arxiv_id, model, dim, vector_bytes, | |
| datetime.utcnow().isoformat(timespec="seconds"), | |
| ), | |
| ) | |
| def papers_missing_embedding( | |
| conn: sqlite3.Connection, | |
| model: str, | |
| ) -> list[sqlite3.Row]: | |
| """Papers que aun no tienen embedding para ese modelo.""" | |
| return conn.execute( | |
| """ | |
| SELECT p.arxiv_id, p.title, p.abstract, p.openalex_abstract | |
| FROM papers p | |
| WHERE p.arxiv_id NOT IN ( | |
| SELECT arxiv_id FROM embeddings WHERE model = ? | |
| ) | |
| AND (p.title IS NOT NULL OR p.abstract IS NOT NULL) | |
| ORDER BY p.published DESC | |
| """, | |
| (model,), | |
| ).fetchall() | |
| def load_all_embeddings( | |
| conn: sqlite3.Connection, | |
| model: str, | |
| ) -> list[sqlite3.Row]: | |
| """Carga todos los embeddings de un modelo (arxiv_id, vector blob).""" | |
| return conn.execute( | |
| "SELECT arxiv_id, vector FROM embeddings WHERE model = ?", | |
| (model,), | |
| ).fetchall() | |
| def count_embeddings(conn: sqlite3.Connection, model: str | None = None) -> int: | |
| if model: | |
| return conn.execute( | |
| "SELECT COUNT(*) FROM embeddings WHERE model = ?", (model,) | |
| ).fetchone()[0] | |
| return conn.execute("SELECT COUNT(*) FROM embeddings").fetchone()[0] | |
| def reset_openalex(conn: sqlite3.Connection) -> int: | |
| """ | |
| Limpia todos los campos OpenAlex para forzar re-enrichment | |
| (util cuando se agregan campos nuevos al pipeline). | |
| Devuelve numero de papers reseteados. | |
| """ | |
| cur = conn.execute( | |
| """ | |
| UPDATE papers SET | |
| openalex_id = NULL, | |
| cited_by_count = NULL, | |
| fwci = NULL, | |
| topics_json = NULL, | |
| publication_year = NULL, | |
| openalex_abstract = NULL | |
| """ | |
| ) | |
| conn.commit() | |
| return cur.rowcount | |
| def mark_openalex_unmatched(conn: sqlite3.Connection, arxiv_id: str) -> None: | |
| """ | |
| Marca explicitamente un paper como buscado-pero-no-encontrado en OpenAlex. | |
| Usa openalex_id = 'NOT_FOUND' como sentinela para evitar re-querying en | |
| futuras corridas (no son free). | |
| """ | |
| conn.execute( | |
| "UPDATE papers SET openalex_id = 'NOT_FOUND', updated_at = ? " | |
| "WHERE arxiv_id = ?", | |
| (datetime.utcnow().isoformat(timespec="seconds"), arxiv_id), | |
| ) | |