mirror of
https://github.com/allaunthefox/Research-Stack.git
synced 2026-07-30 18:56:16 +00:00
allaun
e02786e310
- distribute_extraction.py: escape strings and batch Gremlin vertex upserts. - populate_ene_tables.py: load extraction JSONs and map concepts into ENE tables on neon-64gb (packages, receipts, sidon_labels, braid_strands, eigensolid_snapshots, gossip_surface_nodes/edges, relations, sessions, ingest_events).
310 lines
12 KiB
Python
310 lines
12 KiB
Python
#!/usr/bin/env -S uv run
|
|
# /// script
|
|
# requires-python = ">=3.11"
|
|
# dependencies = ["gremlinpython>=3.7", "psycopg2-binary", "requests"]
|
|
# ///
|
|
# INFRA:LIVE gremlin
|
|
# INFRA:LIVE postgres
|
|
# INFRA:LIVE contextstream
|
|
"""
|
|
distribute_extraction.py
|
|
|
|
Reads curated extraction/ideas.json and extraction/math.json and distributes to:
|
|
- Gremlin (ideas → graph vertices + edges)
|
|
- PostgreSQL (math → math_objects table on arxiv DB)
|
|
|
|
By default uses the hand-curated files. Use --from-merged to ingest the
|
|
auto-extracted all_concepts_merged.json instead.
|
|
|
|
AppFlowy and ContextStream distribution handled separately
|
|
(AppFlowy needs REST API credentials; ContextStream via MCP session).
|
|
|
|
Usage:
|
|
uv run scripts/distribute_extraction.py [--dry-run] [--from-merged] [--target gremlin|postgres|all]
|
|
"""
|
|
|
|
import json
|
|
import os
|
|
import re
|
|
import sys
|
|
import argparse
|
|
import time
|
|
from pathlib import Path
|
|
|
|
ROOT = Path(__file__).parent.parent
|
|
|
|
|
|
def gremlin_escape(s: str) -> str:
|
|
"""Escape a string for safe interpolation into a Gremlin single-quoted literal."""
|
|
s = s.replace("\\", "\\\\")
|
|
s = s.replace("'", "\\'")
|
|
s = s.replace("\n", " ").replace("\r", "")
|
|
s = re.sub(r"[\x00-\x08\x0b\x0c\x0e-\x1f]", "", s)
|
|
if len(s) > 500:
|
|
s = s[:500]
|
|
return s
|
|
|
|
# ── Load infra.json first — abort if service is DEAD ─────────────────────────
|
|
|
|
def load_infra() -> dict:
|
|
p = ROOT / "infra.json"
|
|
if not p.exists():
|
|
sys.exit("ERROR: infra.json not found. Run from project root.")
|
|
with open(p) as f:
|
|
return json.load(f)
|
|
|
|
def require_live(infra: dict, service: str):
|
|
s = infra["services"].get(service, {})
|
|
status = s.get("status", "UNKNOWN")
|
|
if status != "LIVE":
|
|
sys.exit(f"ERROR: service '{service}' is {status} per infra.json. Aborting.")
|
|
|
|
# ── Gremlin distribution (ideas → vertices + edges) ──────────────────────────
|
|
|
|
def gremlin_submit(gc, query, max_retries=5, base_delay=0.5):
|
|
"""Submit a Gremlin query with exponential backoff for rate limiting."""
|
|
for attempt in range(max_retries):
|
|
try:
|
|
return gc.submit(query).all().result()
|
|
except Exception as exc:
|
|
err = str(exc).lower()
|
|
is_rate_limit = any(tag in err for tag in [
|
|
"requestratetoolarge", "too many requests", "429",
|
|
"request rate is large", "throttled"
|
|
])
|
|
if not is_rate_limit or attempt == max_retries - 1:
|
|
raise
|
|
delay = base_delay * (2 ** attempt)
|
|
print(f" [retry {attempt + 1}/{max_retries}] rate limited, sleeping {delay:.2f}s")
|
|
sys.stdout.flush()
|
|
time.sleep(delay)
|
|
|
|
|
|
def distribute_gremlin(ideas: list, infra: dict, dry_run: bool, batch_size: int = 50):
|
|
require_live(infra, "gremlin")
|
|
creds_file = ROOT / infra["services"]["gremlin"].get("credentials", ".env.gremlin")
|
|
|
|
env = {}
|
|
if creds_file.exists():
|
|
for line in creds_file.read_text().splitlines():
|
|
if "=" in line and not line.startswith("#"):
|
|
k, _, v = line.partition("=")
|
|
env[k.strip()] = v.strip().strip('"')
|
|
|
|
endpoint = env.get("GREMLIN_ENDPOINT") or infra["services"]["gremlin"]["endpoint"]
|
|
key = env.get("GREMLIN_PASSWORD") or env.get("GREMLIN_KEY", "")
|
|
database = infra["services"]["gremlin"].get("database", "research")
|
|
graph = infra["services"]["gremlin"].get("graph", "concepts")
|
|
|
|
print(f"\n[Gremlin] endpoint={endpoint} graph={graph} ideas={len(ideas)}")
|
|
|
|
def vertex_label(idea):
|
|
lbl = idea.get('label', '')
|
|
if lbl and re.fullmatch(r"[A-Za-z_][A-Za-z0-9_]*", lbl) and len(lbl) <= 40:
|
|
return lbl
|
|
return "concept"
|
|
|
|
if dry_run:
|
|
for idea in ideas:
|
|
print(f" DRY-RUN vertex: id={idea['id']} label={vertex_label(idea)}")
|
|
return
|
|
|
|
try:
|
|
from gremlin_python.driver import client as gremlin_client
|
|
from gremlin_python.driver.serializer import GraphSONSerializersV2d0
|
|
except ImportError as e:
|
|
print(f" [skip] gremlinpython import failed: {e}")
|
|
print(" Run: uv pip install gremlinpython")
|
|
return
|
|
|
|
gc = gremlin_client.Client(
|
|
endpoint,
|
|
"g",
|
|
username=f"/dbs/{database}/colls/{graph}",
|
|
password=key,
|
|
message_serializer=GraphSONSerializersV2d0(),
|
|
)
|
|
|
|
def build_vertex_query(idea):
|
|
name = idea.get('name', '')
|
|
title = idea.get('title', '')
|
|
summary = idea.get('summary', '') or idea.get('description', '')
|
|
props = {
|
|
"title": title or name,
|
|
"name": name or title,
|
|
"summary": summary[:500],
|
|
"status": idea.get("status", ""),
|
|
"type": idea.get("type", ""),
|
|
"tags": ",".join(idea.get("tags", [])),
|
|
"source": idea.get("source", idea.get("source_file", "")),
|
|
}
|
|
prop_str = " ".join(f".property('{k}', '{gremlin_escape(v)}')" for k, v in props.items()
|
|
if v)
|
|
label = vertex_label(idea)
|
|
eid = gremlin_escape(idea['id'])
|
|
query = f"g.V().has('id', '{eid}').fold().coalesce(unfold(), addV('{label}').property('id', '{eid}').property('pk', '{eid}')){prop_str}"
|
|
return query, label
|
|
|
|
# Batch async vertex upserts
|
|
total = len(ideas)
|
|
for batch_start in range(0, total, batch_size):
|
|
batch = ideas[batch_start:batch_start + batch_size]
|
|
futures = []
|
|
for idea in batch:
|
|
query, label = build_vertex_query(idea)
|
|
try:
|
|
future = gc.submitAsync(query)
|
|
futures.append((future, idea['id'], label))
|
|
except Exception as exc:
|
|
print(f" ERROR submit {idea['id']} {str(exc)[:120]}")
|
|
sys.stdout.flush()
|
|
|
|
for future, eid, label in futures:
|
|
try:
|
|
future.result(timeout=30)
|
|
print(f" vertex upserted: {eid} ({label[:40]})")
|
|
except Exception as exc:
|
|
print(f" ERROR {eid} {str(exc)[:120]}")
|
|
sys.stdout.flush()
|
|
|
|
if batch_start + batch_size < total:
|
|
time.sleep(0.1)
|
|
|
|
# Add edges
|
|
for idea in ideas:
|
|
src_id = gremlin_escape(idea['id'])
|
|
for edge in idea.get("edges", []):
|
|
rel = gremlin_escape(edge['rel'])
|
|
dst_id = gremlin_escape(edge['to'])
|
|
query = (f"g.V().has('id', '{src_id}').has('pk', '{src_id}')"
|
|
f".coalesce("
|
|
f" outE('{rel}').where(inV().has('id', '{dst_id}').has('pk', '{dst_id}')).fold(),"
|
|
f" addE('{rel}').to(g.V().has('id', '{dst_id}').has('pk', '{dst_id}'))"
|
|
f")")
|
|
try:
|
|
gremlin_submit(gc, query)
|
|
print(f" edge: {idea['id']} --[{edge['rel']}]--> {edge['to']}")
|
|
except Exception as exc:
|
|
print(f" edge skip ({idea['id']}→{edge['to']}): {exc}")
|
|
|
|
gc.close()
|
|
|
|
# ── PostgreSQL distribution (math → math_objects table) ──────────────────────
|
|
|
|
def distribute_postgres(math_items: list, infra: dict, dry_run: bool, sql_only: bool = False):
|
|
require_live(infra, "postgres")
|
|
svc = infra["services"]["postgres"]
|
|
host = svc.get("host", "neon-64gb")
|
|
|
|
print(f"\n[PostgreSQL] host={host} math_items={len(math_items)}")
|
|
|
|
if dry_run:
|
|
for m in math_items:
|
|
print(f" DRY-RUN row: id={m['id']} type={m['type']}")
|
|
return
|
|
|
|
import subprocess, shlex
|
|
|
|
def esc(s):
|
|
return str(s or "").replace("'", "''")
|
|
|
|
# Build complete SQL
|
|
sql_lines = [
|
|
"CREATE TABLE IF NOT EXISTS math_objects (",
|
|
" id TEXT PRIMARY KEY, type TEXT NOT NULL, name TEXT NOT NULL,",
|
|
" expression TEXT, latex TEXT, description TEXT,",
|
|
" source_file TEXT, lean_name TEXT, lean_status TEXT,",
|
|
" tags TEXT[], related TEXT[], created_at TIMESTAMP DEFAULT NOW()",
|
|
");",
|
|
""
|
|
]
|
|
|
|
for m in math_items:
|
|
tags = "{" + ",".join(esc(t) for t in m.get("tags", [])) + "}"
|
|
related = "{" + ",".join(esc(r) for r in m.get("related", [])) + "}"
|
|
sql_lines.append(f"""
|
|
INSERT INTO math_objects
|
|
(id,type,name,expression,latex,description,source_file,lean_name,lean_status,tags,related)
|
|
VALUES (
|
|
'{esc(m["id"])}','{esc(m["type"])}','{esc(m["name"])}',
|
|
'{esc(m.get("expression"))}','{esc(m.get("latex"))}',
|
|
'{esc(m.get("description"))}','{esc(m.get("source_file"))}',
|
|
'{esc(m.get("lean_name"))}','{esc(m.get("lean_status"))}',
|
|
'{tags}','{related}'
|
|
)
|
|
ON CONFLICT (id) DO UPDATE SET
|
|
expression=EXCLUDED.expression, latex=EXCLUDED.latex,
|
|
description=EXCLUDED.description, lean_status=EXCLUDED.lean_status,
|
|
tags=EXCLUDED.tags;
|
|
""".strip())
|
|
|
|
sql_file = ROOT / "extraction" / "math_postgres.sql"
|
|
sql_file.write_text("\n".join(sql_lines))
|
|
print(f" SQL written to {sql_file}")
|
|
|
|
if sql_only:
|
|
print(f" --sql-only mode: file written, not executed. Pipe via SSH:")
|
|
print(f" ssh {host} 'podman exec -i arxiv-pg psql -U postgres -d arxiv' < {sql_file}")
|
|
return
|
|
|
|
# Execute via SSH in one pass
|
|
cmd = ["ssh", host, "podman exec -i arxiv-pg psql -U postgres -d arxiv"]
|
|
proc = subprocess.Popen(cmd, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
|
|
stdout, stderr = proc.communicate(input="\n".join(sql_lines))
|
|
|
|
if proc.returncode != 0:
|
|
print(f" ERROR: {stderr.strip()}")
|
|
return
|
|
|
|
print(" math_objects table ensured")
|
|
for m in math_items:
|
|
print(f" upserted: {m['id']} ({m.get('lean_status','')})")
|
|
|
|
# ── Main ──────────────────────────────────────────────────────────────────────
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser(description="Distribute extraction JSON to live infra")
|
|
ap.add_argument("--dry-run", action="store_true", help="Print actions without executing")
|
|
ap.add_argument("--from-merged", action="store_true", help="Use auto-extracted all_concepts_merged.json instead of curated ideas.json/math.json")
|
|
ap.add_argument("--target", default="all", help="gremlin|postgres|all")
|
|
ap.add_argument("--sql-only", action="store_true", help="Write SQL file but don't execute (postgres only)")
|
|
ap.add_argument("--limit", type=int, default=0, help="Process only the first N items per target (0 = all)")
|
|
ap.add_argument("--offset", type=int, default=0, help="Skip the first N items per target")
|
|
ap.add_argument("--batch-size", type=int, default=50, help="Gremlin async batch size (default 50)")
|
|
args = ap.parse_args()
|
|
|
|
infra = load_infra()
|
|
math_types = {"theorem","definition","structure","type","predicate",
|
|
"operator","invariant","constant","equation","algorithm",
|
|
"dataset","model"}
|
|
|
|
if args.from_merged:
|
|
merged = json.loads((ROOT / "extraction/all_concepts_merged.json").read_text())["math"]
|
|
math_obj = [x for x in merged if x.get("type") in math_types]
|
|
ideas = [x for x in merged if x.get("type") == "concept"]
|
|
else:
|
|
ideas = json.loads((ROOT / "extraction/ideas.json").read_text())["ideas"]
|
|
math_obj = [x for x in json.loads((ROOT / "extraction/math.json").read_text())["math"]
|
|
if x.get("type") in math_types]
|
|
|
|
def slice(items):
|
|
start = args.offset
|
|
end = None if args.limit == 0 else start + args.limit
|
|
return items[start:end]
|
|
|
|
ideas = slice(ideas)
|
|
math_obj = slice(math_obj)
|
|
|
|
print(f"Loaded {len(ideas)} ideas, {len(math_obj)} math objects (offset={args.offset}, limit={args.limit})")
|
|
print(f"Dry run: {args.dry_run} From merged: {args.from_merged} Target: {args.target} SQL-only: {args.sql_only}")
|
|
|
|
if args.target in ("gremlin", "all"):
|
|
distribute_gremlin(ideas, infra, args.dry_run, batch_size=args.batch_size)
|
|
|
|
if args.target in ("postgres", "all"):
|
|
distribute_postgres(math_obj, infra, args.dry_run, args.sql_only)
|
|
|
|
print("\nDone.")
|
|
|
|
if __name__ == "__main__":
|
|
main()
|