Update app.py

app.py

@@ -256,6 +256,92 @@ def completions(request: Request, body: ChatRequest):
         "iom_receipt_id": receipt["receipt_id"]
     }
 
+# ========= Session & Merkle helpers =========
+SESSIONS_DIR = pathlib.Path("/tmp/sessions"); SESSIONS_DIR.mkdir(exist_ok=True)
+
+def _session_path(sid: str) -> pathlib.Path:
+    return SESSIONS_DIR / f"{sid}.json"
+
+def merkle_root_from_leaves(leaves: list[str]) -> str:
+    """Compute a simple Merkle root from hex-string leaves.
+    Parent = sha256_hex(left + right) (deterministic demo)."""
+    if not leaves:
+        return ""
+    level = leaves[:]
+    while len(level) > 1:
+        if len(level) % 2 == 1:
+            level.append(level[-1])  # duplicate last if odd
+        nxt = []
+        for i in range(0, len(level), 2):
+            nxt.append(sha256_hex(level[i] + level[i+1]))
+        level = nxt
+    return level[0]
+
+def leaf_from_receipt(r: dict) -> str:
+    """Leaf = sha256( prompt_sha256 + response_sha256 )."""
+    p = r.get("prompt_sha256", "")
+    a = r.get("response_sha256", "")
+    return sha256_hex(p + a)
+
+def session_new() -> dict:
+    sid = str(uuid.uuid4())
+    doc = {"id": sid, "created": datetime.utcnow().isoformat()+"Z",
+           "receipts": [], "leaf_hashes": [], "policy_ids": [], "merkle_root": ""}
+    _session_path(sid).write_text(json.dumps(doc, indent=2))
+    return doc
+
+def session_load(sid: str) -> dict:
+    p = _session_path(sid)
+    if not p.exists():
+        raise FileNotFoundError("session not found")
+    return json.loads(p.read_text())
+
+def session_save(doc: dict):
+    _session_path(doc["id"]).write_text(json.dumps(doc, indent=2))
+
+def session_add_receipt(sid: str, rid: str) -> dict:
+    doc = session_load(sid)
+    rp = RECEIPTS_DIR / f"{rid}.json"
+    if not rp.exists():
+        raise FileNotFoundError("receipt not found")
+    rj = json.loads(rp.read_text())
+    leaf = leaf_from_receipt(rj)
+    doc["receipts"].append(rid)
+    doc["leaf_hashes"].append(leaf)
+    pid = rj.get("policy_id")
+    if pid and pid not in doc["policy_ids"]:
+        doc["policy_ids"].append(pid)
+    doc["merkle_root"] = ""  # will be set on finish
+    session_save(doc)
+    return doc
+
+def session_finish(sid: str) -> dict:
+    doc = session_load(sid)
+    doc["merkle_root"] = merkle_root_from_leaves(doc["leaf_hashes"])
+    doc["finished"] = datetime.utcnow().isoformat()+"Z"
+    session_save(doc)
+    return doc
+
+# ========= Session API =========
+@api.post("/session/new")
+def api_session_new():
+    doc = session_new()
+    return {"session_id": doc["id"]}
+
+@api.post("/session/{sid}/add/{rid}")
+def api_session_add(sid: str, rid: str):
+    doc = session_add_receipt(sid, rid)
+    return {"session_id": doc["id"], "count": len(doc["receipts"])}
+
+@api.post("/session/{sid}/finish")
+def api_session_finish(sid: str):
+    doc = session_finish(sid)
+    return {"session_id": doc["id"], "merkle_root": doc["merkle_root"]}
+
+@api.get("/session/{sid}")
+def api_session_get(sid: str):
+    return session_load(sid)
+
 # ========= UI helpers =========
 def _load_events_df():
     if not os.path.exists(EVENTS_PATH):
@@ -351,7 +437,7 @@ with api_ui:
     gr.Markdown("# IoM-TR Receipts — Portable Governance (NIM + HF Space)")
 
     with gr.Tab("Ask & Verify (in this Space)"):
-        gr.Markdown("Type a prompt, click **Ask & Get Receipt**. You’ll see the answer, a Receipt ID, the full receipt JSON, and a built-in verification result (SHA/HMAC).")
+        gr.Markdown("Type a prompt, click **Ask & Get Receipt**. The UI shows the answer, a Receipt ID, the full receipt JSON, and local verification (SHA/HMAC).")
         p2 = gr.Textbox(label="Prompt", value="In one sentence, what are portable AI governance receipts?")
         ask = gr.Button("Ask & Get Receipt")
         ans = gr.Textbox(label="Answer")
@@ -377,7 +463,45 @@ with api_ui:
         verdict2 = gr.Textbox(label="Result")
         check.click(verify_local, inputs=[p, a, rj, pepper], outputs=verdict2)
 
+    with gr.Tab("Session Root"):
+        gr.Markdown("**Create a session → add last receipt(s) → finish → get one Merkle root + proof JSON.**")
+
+        sid_tb = gr.Textbox(label="Session ID", interactive=False)
+        sess_status = gr.Textbox(label="Status", interactive=False)
+        count_tb = gr.Number(label="Receipts in session", value=0, interactive=False)
+        root_tb = gr.Textbox(label="Merkle root", interactive=False)
+        proof_code = gr.Code(label="Session proof JSON (snapshot)", language="json")
+        proof_link = gr.Markdown()
+
+        def ui_session_new():
+            d = session_new()
+            return d["id"], "Session started", 0, "", "", f"[Open proof](/session/{d['id']})"
+
+        def ui_session_add_last(sid: str, last_rid: str):
+            if not sid:
+                return "", "Start a session first", 0, "", "", ""
+            if not last_rid:
+                return sid, "Run 'Ask & Get Receipt' first (to produce a receipt)", 0, "", "", f"[Open proof](/session/{sid})"
+            d = session_add_receipt(sid, last_rid)
+            snap = json.dumps(d, indent=2)
+            return d["id"], f"Added receipt {last_rid}", len(d["receipts"]), d.get("merkle_root",""), snap, f"[Open proof](/session/{d['id']})"
+
+        def ui_session_finish(sid: str):
+            if not sid:
+                return "", "Start a session first", 0, "", "", ""
+            d = session_finish(sid)
+            snap = json.dumps(d, indent=2)
+            return d["id"], "Session finished", len(d["receipts"]), d["merkle_root"], snap, f"[Open proof](/session/{d['id']})"
+
+        start_btn = gr.Button("Start new session")
+        add_btn   = gr.Button("Add last receipt from 'Ask & Verify'")
+        finish_btn= gr.Button("Finish session (compute root)")
+
+        start_btn.click(ui_session_new, outputs=[sid_tb, sess_status, count_tb, root_tb, proof_code, proof_link])
+        # reuse the 'rid' Textbox from the Ask & Verify tab as input here
+        add_btn.click(ui_session_add_last, inputs=[sid_tb, rid], outputs=[sid_tb, sess_status, count_tb, root_tb, proof_code, proof_link])
+        finish_btn.click(ui_session_finish, inputs=[sid_tb], outputs=[sid_tb, sess_status, count_tb, root_tb, proof_code, proof_link])
+
 # Expose the ASGI app for Spaces (Docker or Gradio SDK)
-# Optional: exposing demo can help the Gradio runner, but not required for Docker
 demo = api_ui
-app = gr.mount_gradio_app(api, api_ui, path="/")
+app  = gr.mount_gradio_app(api, api_ui, path="/")