Create main.py

main.py

@@ -0,0 +1,458 @@
+import os, json, logging, warnings, time, certifi, pymysql, requests
+from contextlib import contextmanager
+from datetime import date
+from flask import Flask, request, jsonify
+from flask_cors import CORS
+
+# ---- Optional Google GenAI (Gemini) ----
+from google import genai
+from google.genai import types
+
+warnings.filterwarnings("ignore")
+
+# ───────────────────────────────────────────────────────────────────────────────
+# CONFIG
+# ───────────────────────────────────────────────────────────────────────────────
+DB_NAME        = os.getenv("TIDB_DB", "test")
+TIDB_HOST      = os.getenv("TIDB_HOST", "")
+TIDB_PORT      = int(os.getenv("TIDB_PORT", "4000"))
+TIDB_USER      = os.getenv("TIDB_USER", "")
+TIDB_PASS      = os.getenv("TIDB_PASS", "")
+
+VEC_DIM        = int(os.getenv("VEC_DIM", "1536"))
+EMBED_MODEL    = os.getenv("EMBED_MODEL", "sentence-transformers/all-MiniLM-L6-v2")
+USE_GPU        = os.getenv("USE_GPU", "0") == "1"  # Spaces are usually CPU; works either way
+
+# Policy windows (server is single source of truth for the client)
+POLICY_WINDOWS = [
+    {
+        "code": "NAZI_ERA",
+        "label": "Washington Conference Principles (1933–1945)",
+        "from": "1933-01-01",
+        "to":   "1945-12-31",
+        "ref":  "https://www.state.gov/washington-conference-principles-on-nazi-confiscated-art"
+    },
+    {
+        "code": "UNESCO_1970",
+        "label": "UNESCO 1970 Convention",
+        "from": "1970-11-14",
+        "to":   None,
+        "ref":  "https://www.unesco.org/en/legal-affairs/convention-means-prohibiting-and-preventing-illicit-import-export-and-transfer-ownership-cultural"
+    }
+]
+
+# ───────────────────────────────────────────────────────────────────────────────
+# APP + LOGGING
+# ───────────────────────────────────────────────────────────────────────────────
+logging.basicConfig(level=logging.INFO)
+log = logging.getLogger("provenance-api")
+
+app = Flask(__name__)
+CORS(app)
+
+# ───────────────────────────────────────────────────────────────────────────────
+# DB CONNECTION (autocommit + TLS + auto-reconnect)
+# ───────────────────────────────────────────────────────────────────────────────
+_CONN = None
+
+def _connect():
+    """(Re)connect to TiDB with TLS + autocommit; DictCursor for JSON friendliness."""
+    global _CONN
+    try:
+        if _CONN:
+            _CONN.close()
+    except Exception:
+        pass
+    _CONN = pymysql.connect(
+        host=TIDB_HOST,
+        port=TIDB_PORT,
+        user=TIDB_USER,
+        password=TIDB_PASS,
+        database=DB_NAME,
+        ssl={"ca": certifi.where()},
+        ssl_verify_cert=True,
+        ssl_verify_identity=True,
+        autocommit=True,  # ensure object rows are durable before child rows (FKs)
+        charset="utf8mb4",
+        cursorclass=pymysql.cursors.DictCursor,
+    )
+
+def _ensure_conn():
+    global _CONN
+    if _CONN is None:
+        _connect()
+    else:
+        try:
+            _CONN.ping(reconnect=True)
+        except Exception:
+            _connect()
+    return _CONN
+
+@contextmanager
+def cursor():
+    """DictCursor with auto-ping; use in each route."""
+    conn = _ensure_conn()
+    with conn.cursor() as cur:
+        yield cur
+
+# ───────────────────────────────────────────────────────────────────────────────
+# EMBEDDINGS (lazy-load; same model as ingest; pad to 1536)
+# ───────────────────────────────────────────────────────────────────────────────
+_MODEL = None
+_DEVICE_INFO = "cpu"
+
+def _pad(vec, dim=VEC_DIM):
+    return vec[:dim] + [0.0] * max(0, dim - len(vec))
+
+def _load_model():
+    global _MODEL, _DEVICE_INFO
+    if _MODEL is not None:
+        return _MODEL
+    if USE_GPU:
+        try:
+            import torch
+            if torch.cuda.is_available():
+                _DEVICE_INFO = "cuda"
+        except Exception:
+            _DEVICE_INFO = "cpu"
+    from sentence_transformers import SentenceTransformer
+    _MODEL = SentenceTransformer(EMBED_MODEL, device=_DEVICE_INFO)
+    log.info(f"Loaded embedding model on '{_DEVICE_INFO}': {EMBED_MODEL}")
+    return _MODEL
+
+def embed_text_to_vec1536(text: str):
+    model = _load_model()
+    vec = model.encode(
+        [text], batch_size=1, show_progress_bar=False, convert_to_numpy=True
+    )[0].tolist()
+    return _pad(vec, VEC_DIM)
+
+# ───────────────────────────────────────────────────────────────────────────────
+# GEMINI (explanations / descriptions)
+# ───────────────────────────────────────────────────────────────────────────────
+GEMINI_KEY = os.environ.get("Gemini")
+_gclient = None
+
+def _gemini():
+    global _gclient
+    if _gclient is not None:
+        return _gclient
+    if not GEMINI_KEY:
+        return None
+    try:
+        _gclient = genai.Client(api_key=GEMINI_KEY)
+        log.info("Gemini client initialized.")
+        return _gclient
+    except Exception as e:
+        log.warning(f"Gemini init failed: {e}")
+        return None
+
+EXPLAIN_MODEL = "gemini-2.5-flash"
+
+def gemini_explain(prompt: str, sys: str = None, model: str = EXPLAIN_MODEL) -> str:
+    g = _gemini()
+    if g is None:
+        # Graceful fallback so the API still works without a key
+        return "(Gemini not configured) " + prompt[:180]
+    # chat-style to mirror your original pattern
+    chat = g.chats.create(model=model)
+    # Add a light system preamble for style/constraints
+    if sys:
+        chat.send_message(f"[SYSTEM]\n{sys}")
+    resp = chat.send_message(prompt)
+    return getattr(resp, "text", "").strip() or ""
+
+# ───────────────────────────────────────────────────────────────────────────────
+# UTIL: Build graph & timeline from events (+ risk overlays)
+# ───────────────────────────────────────────────────────────────────────────────
+def _policy_hits_for_date(d: str):
+    """Return policy codes a given ISO date falls into."""
+    if not d:
+        return []
+    hits = []
+    for w in POLICY_WINDOWS:
+        start_ok = (d >= w["from"]) if w["from"] else True
+        end_ok   = (d <= w["to"])   if w["to"]   else True
+        if start_ok and end_ok:
+            hits.append(w["code"])
+    return hits
+
+def build_graph_from_events(obj_row, events):
+    """Cytoscape.js-style graph: nodes+edges."""
+    nodes = []
+    edges = []
+
+    # center object node
+    onode = {
+        "id": f"obj:{obj_row['object_id']}",
+        "label": f"{obj_row.get('title') or 'Untitled'} ({obj_row.get('source')})",
+        "type": "object"
+    }
+    nodes_map = {onode["id"]: onode}
+
+    def add_node(kind, label):
+        if not label:
+            return None
+        nid = f"{kind}:{label}"
+        if nid not in nodes_map:
+            nodes_map[nid] = {"id": nid, "label": label, "type": kind}
+        return nid
+
+    for ev in events:
+        actor = ev.get("actor")
+        place = ev.get("place")
+        etype = ev.get("event_type") or "UNKNOWN"
+        d_iso = (ev.get("date_from") or "")[:10] if ev.get("date_from") else None
+
+        actor_id = add_node("actor", actor) if actor else None
+        place_id = add_node("place", place) if place else None
+
+        # Edge semantics: actor -> object; place is context (not endpoint)
+        if actor_id:
+            edges.append({
+                "source": actor_id,
+                "target": onode["id"],
+                "label": etype,
+                "date": d_iso,
+                "weight": 1.0,   # client may recompute with risk overlays
+                "source_ref": ev.get("source_ref"),
+                "policy": _policy_hits_for_date(d_iso)
+            })
+
+        # Optional: object -> place (to visualize locations)
+        if place_id and place:
+            edges.append({
+                "source": onode["id"],
+                "target": place_id,
+                "label": "LOCATED",
+                "date": d_iso,
+                "weight": 0.5,
+                "source_ref": ev.get("source_ref"),
+                "policy": _policy_hits_for_date(d_iso)
+            })
+
+    return {"nodes": list(nodes_map.values()), "edges": edges}
+
+def build_timeline_from_events_and_sentences(events, sentences):
+    """Simple list items for any timeline widget."""
+    items = []
+    s_by_seq = {s["seq"]: s["sentence"] for s in sentences}
+    for ev in events:
+        start = ev.get("date_from")
+        end   = ev.get("date_to")
+        title = ev.get("event_type") or "Event"
+        txt   = None
+        # Try to pull the nearest sentence by seq if present
+        # (ingest stored seq starting at 0)
+        for k in (0, 1, 2, 3):
+            if k in s_by_seq:
+                txt = s_by_seq[k]; break
+        items.append({
+            "title": title,
+            "start_date": start,
+            "end_date": end,
+            "text": txt or "",
+            "source_ref": ev.get("source_ref")
+        })
+    return items
+
+# ───────────────────────────────────────────────────────────────────────────────
+# ROUTES
+# ───────────────────────────────────────────────────────────────────────────────
+
+@app.get("/")
+def root():
+    return jsonify({"ok": True, "service": "provenance-radar-api", "device": _DEVICE_INFO})
+
+@app.get("/api/health")
+def health():
+    try:
+        with cursor() as cur:
+            cur.execute("SELECT COUNT(*) AS c FROM objects");      objects = cur.fetchone()["c"]
+            cur.execute("SELECT COUNT(*) AS c FROM provenance_sentences"); sentences = cur.fetchone()["c"]
+            cur.execute("SELECT COUNT(*) AS c FROM risk_signals"); risks = cur.fetchone()["c"]
+        return jsonify({"ok": True, "device": _DEVICE_INFO, "counts": {
+            "objects": objects, "sentences": sentences, "risk_signals": risks}})
+    except Exception as e:
+        log.exception("health failed")
+        return jsonify({"ok": False, "error": str(e)}), 500
+
+@app.get("/api/policy/windows")
+def policy_windows():
+    return jsonify({"ok": True, "windows": POLICY_WINDOWS})
+
+@app.get("/api/leads")
+def leads():
+    limit = max(1, min(int(request.args.get("limit", 50)), 200))
+    min_score = float(request.args.get("min_score", 0))
+    source = request.args.get("source")
+    sql = (
+        "SELECT object_id, source, title, creator, risk_score, top_signals "
+        "FROM flagged_leads WHERE risk_score >= %s "
+    )
+    args = [min_score]
+    if source:
+        sql += " AND source = %s "
+        args.append(source)
+    sql += " LIMIT %s"
+    args.append(limit)
+    with cursor() as cur:
+        cur.execute(sql, args)
+        rows = cur.fetchall()
+    return jsonify({"ok": True, "data": rows})
+
+@app.get("/api/object/<int:object_id>")
+def object_detail(object_id: int):
+    with cursor() as cur:
+        cur.execute("SELECT * FROM objects WHERE object_id=%s", (object_id,))
+        obj = cur.fetchone()
+        if not obj:
+            return jsonify({"ok": False, "error": "not_found"}), 404
+        cur.execute("SELECT seq, sentence FROM provenance_sentences WHERE object_id=%s ORDER BY seq", (object_id,))
+        sents = cur.fetchall()
+        cur.execute("""SELECT event_type, date_from, date_to, place, actor, method, source_ref
+                       FROM provenance_events WHERE object_id=%s
+                       ORDER BY COALESCE(date_from,'0001-01-01')""", (object_id,))
+        events = cur.fetchall()
+        cur.execute("SELECT code, detail, weight FROM risk_signals WHERE object_id=%s ORDER BY weight DESC", (object_id,))
+        risks = cur.fetchall()
+    return jsonify({"ok": True, "object": obj, "sentences": sents, "events": events, "risks": risks})
+
+@app.get("/api/graph/<int:object_id>")
+def graph(object_id: int):
+    with cursor() as cur:
+        cur.execute("SELECT object_id, source, title FROM objects WHERE object_id=%s", (object_id,))
+        obj = cur.fetchone()
+        if not obj:
+            return jsonify({"ok": False, "error": "not_found"}), 404
+        cur.execute("""SELECT event_type, date_from, date_to, place, actor, source_ref
+                       FROM provenance_events WHERE object_id=%s
+                       ORDER BY COALESCE(date_from,'0001-01-01')""", (object_id,))
+        events = cur.fetchall()
+    return jsonify({"ok": True, **build_graph_from_events(obj, events)})
+
+@app.get("/api/timeline/<int:object_id>")
+def timeline(object_id: int):
+    with cursor() as cur:
+        cur.execute("SELECT seq, sentence FROM provenance_sentences WHERE object_id=%s ORDER BY seq", (object_id,))
+        sents = cur.fetchall()
+        cur.execute("""SELECT event_type, date_from, date_to, place, actor, source_ref
+                       FROM provenance_events WHERE object_id=%s
+                       ORDER BY COALESCE(date_from,'0001-01-01')""", (object_id,))
+        events = cur.fetchall()
+    items = build_timeline_from_events_and_sentences(events, sents)
+    return jsonify({"ok": True, "items": items})
+
+@app.get("/api/keyword")
+def keyword_search():
+    q = (request.args.get("q") or "").strip()
+    limit = max(1, min(int(request.args.get("limit", 50)), 200))
+    if not q:
+        return jsonify({"ok": False, "error": "q required"}), 400
+    like = "%" + q.replace("%","").replace("_","") + "%"
+    with cursor() as cur:
+        cur.execute(
+            """SELECT ps.object_id, ps.seq, ps.sentence, o.source, o.title, o.creator
+               FROM provenance_sentences ps
+               JOIN objects o ON o.object_id = ps.object_id
+               WHERE ps.sentence LIKE %s
+               LIMIT %s""", (like, limit)
+        )
+        rows = cur.fetchall()
+    return jsonify({"ok": True, "query": q, "data": rows})
+
+@app.post("/api/similar")
+def similar_search():
+    payload = request.get_json(force=True) or {}
+    text = (payload.get("text") or "").strip()
+    limit = max(1, min(int(payload.get("limit", 20)), 100))
+    if not text:
+        return jsonify({"ok": False, "error": "text required"}), 400
+    vec = embed_text_to_vec1536(text)
+    vec_json = json.dumps(vec)
+    sql = (
+        "SELECT ps.object_id, ps.seq, ps.sentence, o.source, o.title, o.creator, "
+        f"VEC_COSINE_DISTANCE(ps.embedding, CAST(%s AS VECTOR({VEC_DIM}))) AS distance "
+        "FROM provenance_sentences ps "
+        "JOIN objects o ON o.object_id = ps.object_id "
+        "ORDER BY distance ASC "
+        "LIMIT %s"
+    )
+    with cursor() as cur:
+        cur.execute(sql, (vec_json, limit))
+        rows = cur.fetchall()
+    return jsonify({"ok": True, "device": _DEVICE_INFO, "query": text, "data": rows})
+
+@app.get("/api/vocab")
+def vocab():
+    field = (request.args.get("field") or "").strip().lower()
+    limit = max(1, min(int(request.args.get("limit", 100)), 500))
+    if field not in {"actor", "place", "source", "culture"}:
+        return jsonify({"ok": False, "error": "field must be one of actor|place|source|culture"}), 400
+    if field in {"actor", "place"}:
+        sql = f"SELECT {field} AS v, COUNT(*) AS n FROM provenance_events WHERE {field} IS NOT NULL AND {field}<>'' GROUP BY {field} ORDER BY n DESC LIMIT %s"
+    elif field == "source":
+        sql = "SELECT source AS v, COUNT(*) AS n FROM objects GROUP BY source ORDER BY n DESC LIMIT %s"
+    else:  # culture
+        sql = "SELECT culture AS v, COUNT(*) AS n FROM objects WHERE culture IS NOT NULL AND culture<>'' GROUP BY culture ORDER BY n DESC LIMIT %s"
+    with cursor() as cur:
+        cur.execute(sql, (limit,))
+        rows = cur.fetchall()
+    return jsonify({"ok": True, "field": field, "data": rows})
+
+# ── Gemini-powered explanations ────────────────────────────────────────────────
+
+@app.get("/api/explain/object/<int:object_id>")
+def explain_object(object_id: int):
+    """Generate a concise, policy-aware research note for an object."""
+    with cursor() as cur:
+        cur.execute("SELECT object_id, source, title, creator, date_display, risk_score FROM objects WHERE object_id=%s", (object_id,))
+        obj = cur.fetchone()
+        if not obj:
+            return jsonify({"ok": False, "error": "not_found"}), 404
+        cur.execute("SELECT seq, sentence FROM provenance_sentences WHERE object_id=%s ORDER BY seq", (object_id,))
+        sents = cur.fetchall()
+        cur.execute("SELECT event_type, date_from, date_to, place, actor, source_ref FROM provenance_events WHERE object_id=%s ORDER BY COALESCE(date_from,'0001-01-01')", (object_id,))
+        events = cur.fetchall()
+
+    # Build a compact prompt (few sentences) to keep latency low
+    bullets = []
+    for s in sents[:8]:  # keep prompt small
+        bullets.append(f"- {s['sentence']}")
+    evsumm = []
+    for e in events[:8]:
+        evsumm.append(f"{e.get('event_type')} @ {e.get('place') or '—'} on {e.get('date_from') or '—'} (actor: {e.get('actor') or '—'})")
+
+    sys = ("You are assisting provenance researchers. Write a neutral, concise brief (120–180 words) that:\n"
+           "1) summarizes the chain of custody in plain language; 2) clearly marks any timeline gaps; "
+           "3) calls out potential red flags (e.g., confiscated/looted, sales during 1933–45, exports post-1970) "
+           "without making legal conclusions; 4) ends with a short 'Next leads' list (max 3).")
+    prompt = (
+        f"Object: {obj.get('title') or 'Untitled'} — {obj.get('creator') or ''} (source {obj['source']}). "
+        f"Display date: {obj.get('date_display') or 'n/a'}. Current risk_score={obj.get('risk_score', 0)}.\n\n"
+        f"Provenance sentences:\n" + "\n".join(bullets) + "\n\n"
+        f"Structured events (first 8):\n- " + "\n- ".join(evsumm) + "\n\n"
+        f"Policy windows to consider: Nazi era 1933–1945; UNESCO 1970 onwards."
+    )
+    text = gemini_explain(prompt, sys=sys)
+    return jsonify({"ok": True, "model": EXPLAIN_MODEL, "note": text})
+
+@app.post("/api/explain/text")
+def explain_text():
+    """Explain a specific provenance sentence or user query with policy context."""
+    payload = request.get_json(force=True) or {}
+    sentence = (payload.get("text") or "").strip()
+    if not sentence:
+        return jsonify({"ok": False, "error": "text required"}), 400
+    sys = ("Explain this text as a provenance note for curators. "
+           "Be precise and cautious; highlight possible red flags tied to 1933–1945 and post-1970 export rules.")
+    prompt = f"Explain and contextualize this provenance fragment:\n\n“{sentence}”."
+    text = gemini_explain(prompt, sys=sys)
+    return jsonify({"ok": True, "model": EXPLAIN_MODEL, "explanation": text})
+
+# ───────────────────────────────────────────────────────────────────────────────
+# MAIN (Spaces expects 7860)
+# ───────────────────────────────────────────────────────────────────────────────
+if __name__ == "__main__":
+    port = int(os.environ.get("PORT", "7860"))
+    app.run(host="0.0.0.0", port=port, debug=False)