orchestrator/provenance.py

# SQLAlchemy models for provenance
# File: orchestrator/provenance.py
from sqlalchemy import Column, Integer, String, DateTime, ForeignKey, create_engine
from sqlalchemy.orm import declarative_base, relationship, sessionmaker
from datetime import datetime

Base = declarative_base()

class Paper(Base):
    __tablename__ = 'papers'
    id = Column(String, primary_key=True)
    title = Column(String)
    authors = Column(String)
    abstract = Column(String)
    fetched_at = Column(DateTime, default=datetime.utcnow)
    runs = relationship("Run", back_populates="paper")

class Run(Base):
    __tablename__ = 'runs'
    id = Column(Integer, primary_key=True, autoincrement=True)
    paper_id = Column(String, ForeignKey('papers.id'))
    cell_index = Column(Integer)
    output = Column(String)
    executed_at = Column(DateTime, default=datetime.utcnow)
    paper = relationship("Paper", back_populates="runs")

# Utility to initialize and get a session

def init_db(db_url: str):
    engine = create_engine(db_url)
    Base.metadata.create_all(engine)
    return sessionmaker(bind=engine)