Upload 22 files

app.py

@@ -44,38 +44,103 @@ def reset_kb():
 #  TAB 1 — UPLOAD & EXTRACT
 # ═══════════════════════════════════════════════════
 
-def run_extraction(pdf_files, progress=gr.Progress()):
-    if not pdf_files: return "⚠️ No PDFs uploaded.", ""
-    if not cfg.ANTHROPIC_API_KEY: return "❌ ANTHROPIC_API_KEY secret not set.", ""
-    if not cfg.HF_DATASET_REPO:   return "❌ HF_DATASET_REPO secret not set.", ""
-
+def _save_and_resolve_pdfs(pdf_files) -> list:
+    """
+    Gradio 6 passes uploaded files as plain string paths into a temp dir
+    that may be cleaned up before or during processing.
+    
+    This function:
+    1. Immediately copies every uploaded file to /tmp/quant/pdfs/ (persistent for session)
+    2. Uploads each to HuggingFace dataset pdfs/ folder (persistent across restarts)
+    3. Returns stable local Path objects ready for processing
+    """
+    import shutil
+    PDF_DIR = cfg.TMP / "pdfs"
+    PDF_DIR.mkdir(parents=True, exist_ok=True)
+    resolved = []
+    for f in (pdf_files or []):
+        try:
+            # Gradio 6: f is a str path; Gradio 5: f has .name attribute
+            src = Path(f.name if hasattr(f, "name") else f)
+            if not src.exists():
+                logger.warning(f"Uploaded path does not exist: {src}")
+                continue
+            dst = PDF_DIR / src.name
+            if not dst.exists():
+                shutil.copy2(str(src), str(dst))
+            resolved.append(dst)
+            # Persist to HuggingFace
+            if cfg.HF_TOKEN and cfg.HF_DATASET_REPO:
+                hf.pdf_upload(dst)
+        except Exception as e:
+            logger.error(f"Failed to resolve upload {f}: {e}")
+    return resolved
+
+
+def load_pdfs_from_hf() -> list:
+    """List PDFs previously uploaded to HuggingFace dataset."""
+    try:
+        from huggingface_hub import list_repo_files
+        files = list(list_repo_files(
+            repo_id=cfg.HF_DATASET_REPO,
+            repo_type="dataset",
+            token=cfg.HF_TOKEN,
+        ))
+        return sorted([f for f in files if f.startswith("pdfs/") and f.endswith(".pdf")])
+    except Exception as e:
+        logger.warning(f"Could not list HF PDFs: {e}")
+        return []
+
+
+def download_pdf_from_hf(remote_path: str) -> Path | None:
+    """Download a PDF from HuggingFace to local cache."""
+    try:
+        from huggingface_hub import hf_hub_download
+        PDF_DIR = cfg.TMP / "pdfs"
+        PDF_DIR.mkdir(parents=True, exist_ok=True)
+        local = hf_hub_download(
+            repo_id=cfg.HF_DATASET_REPO,
+            filename=remote_path,
+            repo_type="dataset",
+            token=cfg.HF_TOKEN,
+            local_dir=str(PDF_DIR),
+            force_download=False,
+        )
+        return Path(local)
+    except Exception as e:
+        logger.warning(f"Failed to download {remote_path}: {e}")
+        return None
+
+
+def _extract_paths(paths: list, log: list, totals: dict, progress, kb: dict):
+    """Core extraction loop — shared by new upload and re-process from HF."""
     proc  = PDFProcessor()
     ai    = AIExtractor()
     dedup = Deduplicator()
-    kb    = get_kb()
-    log   = []
-    totals = {k:{"added":0,"merged":0,"skipped":0} for k in ("strategies","formulas","systems")}
     hf_files = []
 
-    for i, pdf_file in enumerate(pdf_files):
-        path = Path(pdf_file.name if hasattr(pdf_file, "name") else pdf_file)
-        progress(i/len(pdf_files), desc=f"{path.name}")
-        log.append(f"\n📖 [{i+1}/{len(pdf_files)}] {path.name}")
+    for i, path in enumerate(paths):
+        progress((i + 1) / max(len(paths), 1), desc=f"{path.name}")
+        log.append(f"\n📖 [{i+1}/{len(paths)}] {path.name}")
         try:
             chunks = list(proc.process(path))
-            log.append(f"  → {len(chunks)} chunks")
+            log.append(f"  → {len(chunks)} chunks extracted")
         except Exception as e:
-            log.append(f"  ❌ {e}"); continue
+            log.append(f"  ❌ PDF read error: {e}")
+            continue
 
         for chunk in chunks:
-            extracted = ai.extract(chunk)
-            stats     = dedup.process(extracted, kb)
-            for kind in ("strategies","formulas","systems"):
-                for act in ("added","merged","skipped"):
-                    totals[kind][act] += stats[kind][act]
-
-        log.append(f"  → New: {totals['strategies']['added']} strats, {totals['formulas']['added']} formulas")
-        if cfg.HF_TOKEN: hf.pdf_upload(path)
+            try:
+                extracted = ai.extract(chunk)
+                stats     = dedup.process(extracted, kb)
+                for kind in ("strategies", "formulas", "systems"):
+                    for act in ("added", "merged", "skipped"):
+                        totals[kind][act] += stats[kind][act]
+            except Exception as e:
+                log.append(f"  ⚠️ Chunk error: {e}")
+
+        log.append(f"  → New: {totals['strategies']['added']} strats, "
+                   f"{totals['formulas']['added']} formulas")
 
     for cid, rec in kb["strategies"].items():
         hf_files.append((f"extracted/strategies/{slugify(rec.get('name',''))}.md",
@@ -84,24 +149,82 @@ def run_extraction(pdf_files, progress=gr.Progress()):
         hf_files.append((f"extracted/formulas/{slugify(rec.get('name',''))}.md",
                          formula_md(rec).encode()))
 
-    progress(0.9, desc="Saving to HuggingFace…")
+    progress(0.95, desc="Saving to HuggingFace…")
     hf.kb_save(kb)
     if hf_files and cfg.HF_TOKEN:
         pushed = hf.push_batch(hf_files, "Update extracted knowledge")
-        log.append(f"\n☁️ Pushed {pushed} files to HuggingFace")
+        log.append(f"\n☁️ Pushed {pushed} markdown files to HuggingFace")
     reset_kb()
+    return ai.tokens_used
 
-    counts = {k: len(kb[k]) for k in kb}
+
+def run_extraction(pdf_files, progress=gr.Progress()):
+    if not cfg.ANTHROPIC_API_KEY: return "❌ ANTHROPIC_API_KEY secret not set.", ""
+    if not cfg.HF_DATASET_REPO:   return "❌ HF_DATASET_REPO secret not set.", ""
+
+    # Step 1: resolve uploads → stable local paths + upload to HF
+    progress(0.0, desc="Saving uploads to HuggingFace…")
+    paths = _save_and_resolve_pdfs(pdf_files)
+
+    if not paths:
+        return ("⚠️ No valid PDFs found. Upload files above, "
+                "or use 'Re-process from HF' to reprocess previously uploaded PDFs."), ""
+
+    kb  = get_kb()
+    log = []
+    totals = {k: {"added":0,"merged":0,"skipped":0}
+              for k in ("strategies","formulas","systems")}
+
+    tokens = _extract_paths(paths, log, totals, progress, kb)
+
+    counts  = {k: len(kb[k]) for k in kb}
     summary = f"""✅ Extraction Complete
 
-PDFs processed: {len(pdf_files)}
+PDFs processed : {len(paths)}
 Strategies  — added: {totals['strategies']['added']}  merged: {totals['strategies']['merged']}  skipped: {totals['strategies']['skipped']}
 Formulas    — added: {totals['formulas']['added']}  merged: {totals['formulas']['merged']}  skipped: {totals['formulas']['skipped']}
 Systems     — added: {totals['systems']['added']}  merged: {totals['systems']['merged']}  skipped: {totals['systems']['skipped']}
 
-KB totals: {counts['strategies']} strategies · {counts['formulas']} formulas · {counts['systems']} systems
-Tokens used: {ai.tokens_used:,}"""
-    return summary, "\n".join(log[-40:])
+KB totals   : {counts['strategies']} strategies · {counts['formulas']} formulas · {counts['systems']} systems
+Tokens used : {tokens:,}
+PDFs stored : HuggingFace → {cfg.HF_DATASET_REPO}/pdfs/"""
+    return summary, "\n".join(log[-50:])
+
+
+def reprocess_from_hf(selected_pdfs, progress=gr.Progress()):
+    """Download selected PDFs from HF and re-extract."""
+    if not cfg.ANTHROPIC_API_KEY: return "❌ ANTHROPIC_API_KEY secret not set.", ""
+    if not cfg.HF_DATASET_REPO:   return "❌ HF_DATASET_REPO secret not set.", ""
+    if not selected_pdfs:         return "⚠️ No PDFs selected.", ""
+
+    progress(0.0, desc="Downloading from HuggingFace…")
+    paths = []
+    for remote in selected_pdfs:
+        p = download_pdf_from_hf(remote)
+        if p: paths.append(p)
+
+    if not paths:
+        return "❌ Could not download any PDFs from HuggingFace.", ""
+
+    kb  = get_kb()
+    log = [f"Re-processing {len(paths)} PDF(s) from HuggingFace\n"]
+    totals = {k: {"added":0,"merged":0,"skipped":0}
+              for k in ("strategies","formulas","systems")}
+
+    tokens = _extract_paths(paths, log, totals, progress, kb)
+    counts = {k: len(kb[k]) for k in kb}
+    return (f"✅ Re-extraction complete\n"
+            f"PDFs: {len(paths)} · "
+            f"Strategies: +{totals['strategies']['added']} · "
+            f"Formulas: +{totals['formulas']['added']}\n"
+            f"KB totals: {counts['strategies']} strategies · "
+            f"{counts['formulas']} formulas\n"
+            f"Tokens: {tokens:,}"),  "\n".join(log[-50:])
+
+
+def refresh_hf_pdf_list():
+    pdfs = load_pdfs_from_hf()
+    return gr.update(choices=pdfs, value=[])
 
 
 # ═══════════════════════════════════════════════════
@@ -364,15 +487,39 @@ with gr.Blocks(title="Quant Knowledge Extractor — Julia Engine") as demo:
 
         # Tab 1 — Extract
         with gr.Tab("📤 Upload & Extract"):
-            gr.Markdown("### Upload algorithmic trading PDFs — OCR applied automatically")
+            gr.Markdown("""### Upload algorithmic trading PDFs
+PDFs are **saved to HuggingFace** (`pdfs/` folder) so you can re-process them anytime without re-uploading.
+OCR is applied automatically to scanned pages.""")
             with gr.Row():
                 with gr.Column(scale=2):
-                    pdf_in  = gr.File(label="Drop PDFs here", file_count="multiple", file_types=[".pdf"])
-                    ext_btn = gr.Button("🚀 Extract Knowledge", variant="primary", size="lg")
+                    pdf_in  = gr.File(label="Drop PDFs here", file_count="multiple",
+                                      file_types=[".pdf"])
+                    ext_btn = gr.Button("🚀 Upload + Extract", variant="primary", size="lg")
                 with gr.Column(scale=1):
-                    ext_out = gr.Textbox(label="Result", lines=14, interactive=False, elem_classes=["status-box"])
-            ext_log = gr.Textbox(label="Log", lines=8, interactive=False, elem_classes=["status-box"])
+                    ext_out = gr.Textbox(label="Result", lines=14, interactive=False,
+                                        elem_classes=["status-box"])
+            ext_log = gr.Textbox(label="Log", lines=8, interactive=False,
+                                 elem_classes=["status-box"])
+
+            gr.Markdown("---\n### Re-process PDFs already on HuggingFace")
+            gr.Markdown("*Use this if the container restarted and lost your session, "
+                        "or to re-extract with updated prompts.*")
+            with gr.Row():
+                hf_refresh  = gr.Button("🔄 Refresh HF PDF list")
+                hf_pdf_list = gr.CheckboxGroup(label="PDFs stored on HuggingFace",
+                                               choices=[], value=[])
+            rep_btn = gr.Button("♻️ Re-process selected PDFs from HuggingFace",
+                                variant="secondary")
+            rep_out = gr.Textbox(label="Re-process result", lines=6, interactive=False,
+                                 elem_classes=["status-box"])
+            rep_log = gr.Textbox(label="Re-process log", lines=6, interactive=False,
+                                 elem_classes=["status-box"])
+
             ext_btn.click(fn=run_extraction, inputs=[pdf_in], outputs=[ext_out, ext_log])
+            hf_refresh.click(fn=refresh_hf_pdf_list, outputs=[hf_pdf_list])
+            rep_btn.click(fn=reprocess_from_hf, inputs=[hf_pdf_list],
+                          outputs=[rep_out, rep_log])
+            demo.load(fn=refresh_hf_pdf_list, outputs=[hf_pdf_list])
 
         # Tab 2 — Browse
         with gr.Tab("📚 Knowledge Base"):