Update orchestrator/provenance.py

orchestrator/provenance.py

@@ -1,6 +1,5 @@
-# SQLAlchemy models for provenance
 # File: orchestrator/provenance.py
-from sqlalchemy import Column, Integer, String, DateTime, ForeignKey, create_engine
+from sqlalchemy import Column, String, Integer, DateTime, ForeignKey, create_engine
 from sqlalchemy.orm import declarative_base, relationship, sessionmaker
 from datetime import datetime
 
@@ -13,7 +12,7 @@ class Paper(Base):
     authors = Column(String)
     abstract = Column(String)
     fetched_at = Column(DateTime, default=datetime.utcnow)
-    runs = relationship("Run", back_populates="paper")
+    runs = relationship('Run', back_populates='paper')
 
 class Run(Base):
     __tablename__ = 'runs'
@@ -22,11 +21,46 @@ class Run(Base):
     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
+    paper = relationship('Paper', back_populates='runs')
 
 def init_db(db_url: str):
     engine = create_engine(db_url)
     Base.metadata.create_all(engine)
-    return sessionmaker(bind=engine)
+    return sessionmaker(bind=engine)
+
+# File: scripts/ingest.py
+import sys
+import yaml
+from orchestrator.client import MCPClient
+
+"""
+Usage:
+    python ingest.py "search query"
+"""
+if __name__ == '__main__':
+    if len(sys.argv) < 2:
+        print('Please provide a search query.')
+        sys.exit(1)
+    query = sys.argv[1]
+    cfg = yaml.safe_load(open('config.yaml'))
+    web = MCPClient(cfg['mcp_servers']['web_search'])
+    pubmed = MCPClient(cfg['mcp_servers']['pubmed'])
+    chroma = MCPClient(cfg['mcp_servers']['chroma'])
+
+    print(f'Ingesting papers for query: {query}')
+    papers = []
+    try:
+        papers += web.call('web_search.search', {'q': query}) or []
+    except Exception as e:
+        print('Web search error:', e)
+    try:
+        papers += pubmed.call('metatool.query', {'source': 'PubMed', 'q': query}) or []
+    except Exception as e:
+        print('PubMed error:', e)
+
+    for paper in papers:
+        pid = paper.get('id')
+        txt = paper.get('abstract', '')
+        meta = {'title': paper.get('title'), 'authors': ','.join(paper.get('authors', []))}
+        chroma.call('chroma.insert', {'id': pid, 'text': txt, 'metadata': meta})
+    print('Done ingesting!')