Upload app.py with huggingface_hub

app.py

@@ -31,6 +31,9 @@ except ImportError:
 import json
 import logging
 import os
+import csv
+import tempfile
+import time
 from importlib import metadata, util
 from typing import Any
 
@@ -41,6 +44,7 @@ from src.veris_classifier.classifier import (
     answer_question,
     classify_incident,
 )
+from src.veris_classifier.validator import validate_classification
 
 load_dotenv()
 logging.basicConfig(level=logging.INFO)
@@ -276,6 +280,17 @@ textarea:focus {
     text-decoration: none;
 }
 
+.status-card {
+    border: 1px solid #334155;
+    background: rgba(15, 23, 42, 0.6);
+    border-radius: 10px;
+    padding: 8px 12px;
+}
+
+#table-controls .wrap {
+    align-items: end;
+}
+
 /* Mobile */
 @media (max-width: 900px) {
     .hero-section {
@@ -299,6 +314,12 @@ textarea:focus {
     .hero-badges {
         gap: 8px;
     }
+    .primary-btn {
+        width: 100% !important;
+    }
+    #table-controls .wrap {
+        gap: 8px !important;
+    }
 }
 """
 
@@ -330,6 +351,8 @@ EXAMPLES_QA = [
 ZEROGPU_QUEUE_HINT = "No GPU was available after"
 SPACES_PAGE_URL = "https://huggingface.co/spaces/vibesecurityguy/veris-classifier"
 SPACE_HOST_URL = "https://vibesecurityguy-veris-classifier.hf.space"
+ZEROGPU_RETRY_ATTEMPTS = 2
+ZEROGPU_RETRY_DELAY_SECONDS = 3
 
 
 def _is_zerogpu_queue_timeout(err: Exception) -> bool:
@@ -355,6 +378,57 @@ def _spaces_user_logged_in(
     return False
 
 
+def _session_status_markdown(
+    request: gr.Request | None = None,
+    profile: gr.OAuthProfile | None = None,
+) -> str:
+    """Render current Spaces auth status for the user."""
+    if not IS_SPACES:
+        return ""
+
+    if _spaces_user_logged_in(request, profile):
+        username = None
+        if profile is not None:
+            username = (
+                getattr(profile, "preferred_username", None)
+                or getattr(profile, "name", None)
+            )
+        if not username and request is not None:
+            username = getattr(request, "username", None)
+        if username:
+            return (
+                f"**Session status:** Logged in as `{username}`. "
+                "ZeroGPU requests will use your account quota."
+            )
+        return "**Session status:** Logged in. ZeroGPU requests will use your account quota."
+
+    return (
+        "**Session status:** Not logged in. Click sign in to attach this browser session "
+        "to your Hugging Face quota."
+    )
+
+
+def _run_with_zerogpu_retry(call):
+    """Retry queue-timeout failures once before returning an error."""
+    last_error = None
+    for attempt in range(1, ZEROGPU_RETRY_ATTEMPTS + 1):
+        try:
+            return call()
+        except Exception as e:
+            last_error = e
+            if _is_zerogpu_queue_timeout(e) and attempt < ZEROGPU_RETRY_ATTEMPTS:
+                logger.warning(
+                    "ZeroGPU queue timeout (attempt %d/%d). Retrying in %ss.",
+                    attempt,
+                    ZEROGPU_RETRY_ATTEMPTS,
+                    ZEROGPU_RETRY_DELAY_SECONDS,
+                )
+                time.sleep(ZEROGPU_RETRY_DELAY_SECONDS)
+                continue
+            raise
+    raise last_error
+
+
 
 def _use_hf_model() -> bool:
     """Check if we should use the fine-tuned HF model."""
@@ -393,7 +467,7 @@ def classify(
 
     if use_hf:
         try:
-            result = _classify_gpu(description)
+            result = _run_with_zerogpu_retry(lambda: _classify_gpu(description))
             return json.dumps(result, indent=2)
         except Exception as e:
             logger.error(f"HF model error: {e}")
@@ -445,7 +519,7 @@ def ask(
 
     if use_hf:
         try:
-            return _ask_gpu(question)
+            return _run_with_zerogpu_retry(lambda: _ask_gpu(question))
         except Exception as e:
             logger.error(f"HF model error: {e}")
             if _is_zerogpu_queue_timeout(e):
@@ -515,31 +589,150 @@ def _classification_rows_from_json(raw_json: str) -> list[list[str]]:
     return rows
 
 
-def _render_classification_output(raw_json: str, output_format: str):
-    """Render classification as JSON code or flattened table."""
-    if output_format == "Table":
-        table_rows = _classification_rows_from_json(raw_json)
+def _validation_summary_markdown(raw_json: str) -> str:
+    """Build validation summary for the classification output."""
+    try:
+        parsed = json.loads(raw_json)
+    except Exception:
+        return ""
+
+    if not isinstance(parsed, dict) or parsed.get("error"):
+        return "**Validation:** Skipped."
+
+    result = validate_classification(parsed)
+    lines = [f"**Validation:** {'Passed' if result.valid else 'Issues found'}"]
+    if result.errors:
+        lines.append("**Errors**")
+        lines.extend(f"- {err}" for err in result.errors[:8])
+        if len(result.errors) > 8:
+            lines.append(f"- ... {len(result.errors) - 8} more")
+    if result.warnings:
+        lines.append("**Warnings**")
+        lines.extend(f"- {warn}" for warn in result.warnings[:8])
+        if len(result.warnings) > 8:
+            lines.append(f"- ... {len(result.warnings) - 8} more")
+    return "\n".join(lines)
+
+
+def _filter_classification_rows(
+    rows: list[list[str]],
+    dimension_filter: str,
+    errors_only: bool,
+) -> list[list[str]]:
+    """Filter table rows by dimension and optionally error-only rows."""
+    filtered: list[list[str]] = []
+    for row in rows:
+        if len(row) != 3:
+            continue
+        dimension, field, value = row
+
+        if dimension_filter != "All" and dimension != dimension_filter:
+            continue
+
+        if errors_only:
+            blob = f"{dimension} {field} {value}".lower()
+            if "error" not in blob:
+                continue
+
+        filtered.append(row)
+    return filtered
+
+
+def _render_classification_output(
+    raw_json: str,
+    output_format: str,
+    all_rows: list[list[str]],
+    dimension_filter: str,
+    errors_only: bool,
+):
+    """Render classification as JSON code or filtered table."""
+    filtered_rows = _filter_classification_rows(all_rows, dimension_filter, errors_only)
+    show_table = output_format == "Table"
+
+    if show_table:
         return (
             gr.update(value=raw_json, visible=False),
-            gr.update(value=table_rows, visible=True),
+            gr.update(value=filtered_rows, visible=True),
+            gr.update(visible=True),
+            gr.update(visible=True, interactive=bool(filtered_rows)),
         )
+
     return (
         gr.update(value=raw_json, visible=True),
         gr.update(value=[], visible=False),
+        gr.update(visible=False),
+        gr.update(visible=False, interactive=False),
     )
 
 
+def _apply_table_filters(
+    all_rows: list[list[str]],
+    dimension_filter: str,
+    errors_only: bool,
+):
+    """Apply table-only filters without re-running inference."""
+    filtered_rows = _filter_classification_rows(all_rows, dimension_filter, errors_only)
+    return (
+        gr.update(value=filtered_rows),
+        gr.update(interactive=bool(filtered_rows)),
+    )
+
+
+def _build_filtered_csv(
+    all_rows: list[list[str]],
+    dimension_filter: str,
+    errors_only: bool,
+):
+    """Create downloadable CSV file for filtered rows."""
+    filtered_rows = _filter_classification_rows(all_rows, dimension_filter, errors_only)
+    if not filtered_rows:
+        return gr.update(value=None, visible=False)
+
+    with tempfile.NamedTemporaryFile(
+        mode="w",
+        suffix=".csv",
+        delete=False,
+        newline="",
+        encoding="utf-8",
+    ) as tmp:
+        writer = csv.writer(tmp)
+        writer.writerow(["Dimension", "Field", "Value"])
+        writer.writerows(filtered_rows)
+        csv_path = tmp.name
+
+    return gr.update(value=csv_path, visible=True)
+
+
 def classify_and_render(
     description: str,
     api_key: str,
     output_format: str,
+    dimension_filter: str,
+    errors_only: bool,
     request: gr.Request | None = None,
     profile: gr.OAuthProfile | None = None,
 ):
     """Run classification and return display-ready outputs."""
     raw_json = classify(description, api_key, request=request, profile=profile)
-    code_update, table_update = _render_classification_output(raw_json, output_format)
-    return raw_json, code_update, table_update
+    all_rows = _classification_rows_from_json(raw_json)
+    validation_md = _validation_summary_markdown(raw_json)
+    code_update, table_update, controls_update, export_btn_update = _render_classification_output(
+        raw_json,
+        output_format,
+        all_rows,
+        dimension_filter,
+        errors_only,
+    )
+    return (
+        raw_json,
+        all_rows,
+        validation_md,
+        code_update,
+        table_update,
+        controls_update,
+        export_btn_update,
+        gr.update(value=None, visible=False),
+    )
 
 
 # ---------------------------------------------------------------------------
@@ -622,6 +815,7 @@ def build_app() -> gr.Blocks:
         theme=THEME,
         css=CUSTOM_CSS,
     ) as app:
+        session_status = None
 
         # --- Hero Header ---
         gr.HTML("""
@@ -702,6 +896,10 @@ def build_app() -> gr.Blocks:
                     'border:1px solid #334155;color:#cbd5e1;text-decoration:none;font-weight:600;">'
                     "Direct sign-in (if button refreshes)</a>"
                 )
+            session_status = gr.Markdown(
+                value="**Session status:** Checking...",
+                elem_classes=["status-card"],
+            )
         else:
             gr.HTML("""
             <div class="model-banner">
@@ -760,12 +958,38 @@ def build_app() -> gr.Blocks:
                             info="Switch between raw JSON and a flattened table view.",
                         )
                         last_classification_raw = gr.State("")
+                        classification_rows = gr.State([])
+                        validation_output = gr.Markdown(
+                            label="Validation",
+                            value="*Validation summary will appear after classification.*",
+                        )
                         classification_output = gr.Code(
                             label="VERIS Classification (JSON)",
                             language="json",
                             lines=20,
                             elem_classes=["code-output"],
                         )
+                        with gr.Row(visible=False, elem_id="table-controls") as table_controls:
+                            dimension_filter = gr.Dropdown(
+                                choices=["All", "Actor", "Action", "Asset", "Attribute", "Error", "General"],
+                                value="All",
+                                label="Filter Dimension",
+                            )
+                            errors_only = gr.Checkbox(
+                                value=False,
+                                label="Errors Only",
+                            )
+                            export_csv_btn = gr.Button(
+                                "Generate CSV",
+                                size="sm",
+                                interactive=False,
+                                visible=False,
+                            )
+                        csv_file = gr.File(
+                            label="Download Filtered CSV",
+                            visible=False,
+                            interactive=False,
+                        )
                         classification_table = gr.Dataframe(
                             headers=["Dimension", "Field", "Value"],
                             datatype=["str", "str", "str"],
@@ -774,6 +998,7 @@ def build_app() -> gr.Blocks:
                             visible=False,
                             interactive=False,
                             wrap=True,
+                            max_height=500,
                             label="VERIS Classification (Table)",
                         )
 
@@ -787,13 +1012,43 @@ def build_app() -> gr.Blocks:
 
                 classify_btn.click(
                     fn=classify_and_render,
-                    inputs=[incident_input, api_key, output_format],
-                    outputs=[last_classification_raw, classification_output, classification_table],
+                    inputs=[incident_input, api_key, output_format, dimension_filter, errors_only],
+                    outputs=[
+                        last_classification_raw,
+                        classification_rows,
+                        validation_output,
+                        classification_output,
+                        classification_table,
+                        table_controls,
+                        export_csv_btn,
+                        csv_file,
+                    ],
                 )
                 output_format.change(
                     fn=_render_classification_output,
-                    inputs=[last_classification_raw, output_format],
-                    outputs=[classification_output, classification_table],
+                    inputs=[
+                        last_classification_raw,
+                        output_format,
+                        classification_rows,
+                        dimension_filter,
+                        errors_only,
+                    ],
+                    outputs=[classification_output, classification_table, table_controls, export_csv_btn],
+                )
+                dimension_filter.change(
+                    fn=_apply_table_filters,
+                    inputs=[classification_rows, dimension_filter, errors_only],
+                    outputs=[classification_table, export_csv_btn],
+                )
+                errors_only.change(
+                    fn=_apply_table_filters,
+                    inputs=[classification_rows, dimension_filter, errors_only],
+                    outputs=[classification_table, export_csv_btn],
+                )
+                export_csv_btn.click(
+                    fn=_build_filtered_csv,
+                    inputs=[classification_rows, dimension_filter, errors_only],
+                    outputs=[csv_file],
                 )
 
             # ---- TAB 2: Q&A ----
@@ -946,6 +1201,13 @@ def build_app() -> gr.Blocks:
         </div>
         """)
 
+        if IS_SPACES and session_status is not None:
+            app.load(
+                fn=_session_status_markdown,
+                outputs=[session_status],
+                queue=False,
+            )
+
     return app