Add synchronized lego-style tag toggles and prompt builder UI

app.py

@@ -6,7 +6,7 @@ import json
 from datetime import datetime
 from PIL import Image
 from pathlib import Path
-from typing import List
+from typing import Any, Dict, List, Set
 from concurrent.futures import ThreadPoolExecutor, TimeoutError as FutureTimeoutError
 
 from psq_rag.pipeline.preproc import extract_user_provided_tags_upto_3_words
@@ -14,7 +14,7 @@ from psq_rag.llm.rewrite import llm_rewrite_prompt
 from psq_rag.retrieval.psq_retrieval import psq_candidates_from_rewrite_phrases, _norm_tag_for_lookup
 from psq_rag.llm.select import llm_select_indices, llm_infer_structural_tags, llm_infer_probe_tags
 from psq_rag.retrieval.state import expand_tags_via_implications
-from psq_rag.ui.group_ranked_display import render_group_rankings_markdown
+from psq_rag.ui.group_ranked_display import rank_groups_from_tfidf, _load_enabled_groups
 
 
 def _split_prompt_commas(s: str) -> List[str]:
@@ -40,6 +40,181 @@ def compose_final_prompt(rewritten_prompt: str, selected_tags: List[str]) -> str
     return ", ".join(out)
 
 
+def _display_tag_text(tag: str) -> str:
+    return tag.replace("_", " ")
+
+
+def _escape_prompt_tag(tag: str) -> str:
+    return (
+        tag.replace("_", " ")
+        .replace("(", "\\(")
+        .replace(")", "\\)")
+    )
+
+
+def _ordered_selected_for_prompt(selected: Set[str], row_defs: List[Dict[str, Any]]) -> List[str]:
+    out: List[str] = []
+    seen: Set[str] = set()
+    for row in row_defs:
+        for tag in row.get("tags", []):
+            if tag in selected and tag not in seen:
+                out.append(tag)
+                seen.add(tag)
+    # Fallback for any selected tags not present in current rows.
+    for tag in sorted(selected):
+        if tag not in seen:
+            out.append(tag)
+            seen.add(tag)
+    return out
+
+
+def _compose_toggle_prompt_text(selected_tags: List[str], row_defs: List[Dict[str, Any]]) -> str:
+    selected = {t for t in (selected_tags or []) if t}
+    ordered = _ordered_selected_for_prompt(selected, row_defs or [])
+    return ", ".join(_escape_prompt_tag(t) for t in ordered)
+
+
+def _build_toggle_rows(
+    *,
+    seed_terms: List[str],
+    llm_selected_tags: List[str],
+    top_groups: int,
+    top_tags_per_group: int,
+    group_rank_top_k: int,
+) -> List[Dict[str, Any]]:
+    ranked_rows = rank_groups_from_tfidf(
+        seed_terms=seed_terms,
+        top_groups=max(1, int(top_groups)),
+        top_tags_per_group=max(1, int(top_tags_per_group)),
+        group_rank_top_k=max(1, int(group_rank_top_k)),
+    )
+    groups_map = _load_enabled_groups()
+    llm_selected = list(dict.fromkeys(_norm_tag_for_lookup(t) for t in llm_selected_tags if t))
+
+    row_defs: List[Dict[str, Any]] = []
+    displayed_group_names = [r.group_name for r in ranked_rows]
+    displayed_group_tag_sets: Dict[str, Set[str]] = {
+        name: set(groups_map.get(name, [])) for name in displayed_group_names
+    }
+    tags_in_any_displayed_group: Set[str] = set()
+    for tag_set in displayed_group_tag_sets.values():
+        tags_in_any_displayed_group.update(tag_set)
+
+    llm_other = [t for t in llm_selected if t not in tags_in_any_displayed_group]
+    row_defs.append(
+        {
+            "name": "llm_selected_other",
+            "label": "LLM Selected (Other)",
+            "tags": llm_other,
+        }
+    )
+
+    for row in ranked_rows:
+        group_name = row.group_name
+        group_tag_set = displayed_group_tag_sets.get(group_name, set())
+        selected_in_group = [t for t in llm_selected if t in group_tag_set]
+        ranked_tags = [t for t, _ in row.tags]
+        merged = selected_in_group + [t for t in ranked_tags if t not in selected_in_group]
+        keep_n = max(max(1, int(top_tags_per_group)), len(selected_in_group))
+        merged = merged[:keep_n]
+        row_defs.append(
+            {
+                "name": group_name,
+                "label": f"{group_name} (E={row.expected_count:.2f})",
+                "tags": merged,
+            }
+        )
+
+    return row_defs
+
+
+def _build_row_component_updates(
+    row_defs: List[Dict[str, Any]],
+    selected_tags: List[str],
+    max_rows: int,
+):
+    selected = {t for t in (selected_tags or []) if t}
+    row_values_state: List[List[str]] = []
+    header_updates = []
+    checkbox_updates = []
+
+    for idx in range(max_rows):
+        if idx < len(row_defs):
+            row = row_defs[idx]
+            tags = list(dict.fromkeys(row.get("tags", [])))
+            values = [t for t in tags if t in selected]
+            row_values_state.append(values)
+            visible = bool(tags)
+            header_updates.append(gr.update(value=f"**{row.get('label', '')}**", visible=visible))
+            choices = [(_display_tag_text(t), t) for t in tags]
+            checkbox_updates.append(
+                gr.update(
+                    choices=choices,
+                    value=values,
+                    visible=visible,
+                )
+            )
+        else:
+            header_updates.append(gr.update(value="", visible=False))
+            checkbox_updates.append(gr.update(choices=[], value=[], visible=False))
+
+    prompt_text = _compose_toggle_prompt_text(list(selected), row_defs)
+    return prompt_text, row_values_state, header_updates, checkbox_updates
+
+
+def _on_toggle_row(
+    row_idx: int,
+    changed_values: List[str],
+    selected_tags_state: List[str],
+    row_defs_state: List[Dict[str, Any]],
+    row_values_state: List[List[str]],
+    max_rows: int,
+):
+    row_defs = row_defs_state or []
+    selected = set(selected_tags_state or [])
+    prev_values = list(row_values_state or [])
+
+    while len(prev_values) < len(row_defs):
+        prev_values.append([])
+
+    prev_set = set(prev_values[row_idx]) if row_idx < len(prev_values) else set()
+    new_set = set(changed_values or [])
+    selected.update(new_set - prev_set)
+    selected.difference_update(prev_set - new_set)
+
+    prompt_text, new_row_values_state, _header_updates, checkbox_updates = _build_row_component_updates(
+        row_defs=row_defs,
+        selected_tags=list(selected),
+        max_rows=max_rows,
+    )
+
+    return [sorted(selected), new_row_values_state, prompt_text, *checkbox_updates]
+
+
+def _build_ui_payload(
+    *,
+    console_text: str,
+    legacy_prompt_text: str,
+    row_defs: List[Dict[str, Any]],
+    selected_tags: List[str],
+):
+    prompt_text, row_values_state, header_updates, checkbox_updates = _build_row_component_updates(
+        row_defs=row_defs,
+        selected_tags=selected_tags,
+        max_rows=display_max_rows_default,
+    )
+    return [
+        console_text,
+        legacy_prompt_text,
+        prompt_text,
+        sorted(set(selected_tags or [])),
+        row_defs,
+        row_values_state,
+        *header_updates,
+        *checkbox_updates,
+    ]
+
+
 def _build_selection_query(
     prompt_in: str,
     rewritten: str,
@@ -152,6 +327,7 @@ enable_probe_tags = os.environ.get("PSQ_ENABLE_PROBE", "1").strip() not in {"0",
 display_top_groups_default = int(os.environ.get("PSQ_DISPLAY_TOP_GROUPS", "10"))
 display_top_tags_per_group_default = int(os.environ.get("PSQ_DISPLAY_TOP_TAGS_PER_GROUP", "5"))
 display_rank_top_k_default = int(os.environ.get("PSQ_DISPLAY_GROUP_RANK_TOP_K", "5"))
+display_max_rows_default = int(os.environ.get("PSQ_DISPLAY_MAX_ROWS", "14"))
 retrieval_global_k = int(os.environ.get("PSQ_RETRIEVAL_GLOBAL_K", "300"))
 retrieval_per_phrase_k = int(os.environ.get("PSQ_RETRIEVAL_PER_PHRASE_K", "10"))
 retrieval_per_phrase_final_k = int(os.environ.get("PSQ_RETRIEVAL_PER_PHRASE_FINAL_K", "1"))
@@ -189,6 +365,42 @@ css = """
 .pane-right .scrollable-content {
   max-height: 610px;                /* was 420px; tweak to taste */
 }
+
+.lego-tags .gr-checkboxgroup {
+  display: flex;
+  flex-wrap: wrap;
+  gap: 8px;
+}
+
+.lego-tags .gr-checkboxgroup label {
+  margin: 0;
+  padding: 0;
+}
+
+.lego-tags .gr-checkboxgroup input[type="checkbox"] {
+  display: none;
+}
+
+.lego-tags .gr-checkboxgroup span {
+  display: inline-block;
+  padding: 7px 12px;
+  border: 1px solid #8a8a8a;
+  border-radius: 10px;
+  background: #f4f4f4;
+  color: #222;
+  font-size: 0.95rem;
+  line-height: 1.2;
+  cursor: pointer;
+  user-select: none;
+  box-shadow: 0 1px 0 rgba(0,0,0,0.12), inset 0 1px 0 rgba(255,255,255,0.7);
+}
+
+.lego-tags .gr-checkboxgroup input[type="checkbox"]:checked + span {
+  background: #ffd86a;
+  border-color: #c49a00;
+  box-shadow: 0 2px 0 #a98000, inset 0 1px 0 rgba(255,255,255,0.65);
+  transform: translateY(1px);
+}
 """
 
 
@@ -275,7 +487,12 @@ def rag_pipeline_ui(
         log("Start: received prompt")
         prompt_in = (user_prompt or "").strip()
         if not prompt_in:
-            return "Error: empty prompt", "", ""
+            return _build_ui_payload(
+                console_text="Error: empty prompt",
+                legacy_prompt_text="",
+                row_defs=[],
+                selected_tags=[],
+            )
             
         log("Input:")
         log(prompt_in)
@@ -439,6 +656,8 @@ def rag_pipeline_ui(
         elif enable_probe_tags:
             log("  No probe tags inferred")
 
+        llm_selected_tags = list(dict.fromkeys(selected_tags))
+
         log("Step 3c: Expand via tag implications")
         t0 = time.perf_counter()
         tag_set = set(selected_tags)
@@ -469,25 +688,36 @@ def rag_pipeline_ui(
         seed_terms.extend(selected_tags)
         seed_terms = list(dict.fromkeys(seed_terms))
 
-        groups_md = render_group_rankings_markdown(
+        toggle_rows = _build_toggle_rows(
             seed_terms=seed_terms,
+            llm_selected_tags=llm_selected_tags,
             top_groups=max(1, int(display_top_groups)),
             top_tags_per_group=max(1, int(display_top_tags_per_group)),
             group_rank_top_k=max(1, int(display_rank_top_k)),
         )
         dt = time.perf_counter()-t0
         _record_timing("group_display", dt)
-        log(f"Ranked group display: {dt:.2f}s")
+        log(f"Ranked group display: {dt:.2f}s ({len(toggle_rows)} rows)")
 
         total_dt = time.perf_counter()-t_total0
         _emit_timing_summary(total_dt)
         _append_timing_jsonl(total_dt)
         log("Done: final prompt ready")
-        return "\n".join(logs), final_prompt, groups_md
+        return _build_ui_payload(
+            console_text="\n".join(logs),
+            legacy_prompt_text=final_prompt,
+            row_defs=toggle_rows,
+            selected_tags=llm_selected_tags,
+        )
 
     except Exception as e:
         log(f"Error: {type(e).__name__}: {e}")
-        return "\n".join(logs), "", ""
+        return _build_ui_payload(
+            console_text="\n".join(logs),
+            legacy_prompt_text="",
+            row_defs=[],
+            selected_tags=[],
+        )
 
     
 
@@ -529,11 +759,44 @@ then returns a cleaned, model-friendly prompt.
         placeholder="Progress logs will appear here."
     )
 
-    final_prompt = gr.Textbox(
-        label="Final Prompt",
+    suggested_prompt = gr.Textbox(
+        label="Suggested Prompt (From Toggled Tags)",
         lines=3,
         interactive=False,
-        placeholder="Your optimized prompt will appear here."
+        show_copy_button=True,
+        placeholder="Comma-separated tags selected in the rows below."
+    )
+
+    with gr.Accordion("Legacy Pipeline Prompt (for reference)", open=False):
+        legacy_final_prompt = gr.Textbox(
+            label="Legacy Final Prompt",
+            lines=3,
+            interactive=False,
+            show_copy_button=True,
+        )
+
+    selected_tags_state = gr.State([])
+    row_defs_state = gr.State([])
+    row_values_state = gr.State([])
+
+    gr.Markdown("### Toggle Tag Rows")
+    row_headers: List[gr.Markdown] = []
+    row_checkboxes: List[gr.CheckboxGroup] = []
+    for _ in range(display_max_rows_default):
+        row_headers.append(gr.Markdown(value="", visible=False))
+        row_checkboxes.append(
+            gr.CheckboxGroup(
+                choices=[],
+                value=[],
+                visible=False,
+                interactive=True,
+                container=False,
+                elem_classes=["lego-tags"],
+            )
+        )
+
+    gr.Markdown(
+        "Toggling a tag in any row toggles it everywhere else that tag appears."
     )
 
     with gr.Accordion("Display Settings", open=False):
@@ -557,22 +820,42 @@ then returns a cleaned, model-friendly prompt.
                 minimum=1,
             )
 
-    group_rankings_md = gr.Markdown(
-        label="Ranked Group/Category Tag Suggestions",
-        value="",
-    )
+    run_outputs = [
+        console,
+        legacy_final_prompt,
+        suggested_prompt,
+        selected_tags_state,
+        row_defs_state,
+        row_values_state,
+        *row_headers,
+        *row_checkboxes,
+    ]
 
     submit_button.click(
         rag_pipeline_ui,
         inputs=[image_tags, display_top_groups, display_top_tags_per_group, display_rank_top_k],
-        outputs=[console, final_prompt, group_rankings_md]
+        outputs=run_outputs
     )
 
     image_tags.submit(
         rag_pipeline_ui,
         inputs=[image_tags, display_top_groups, display_top_tags_per_group, display_rank_top_k],
-        outputs=[console, final_prompt, group_rankings_md]
+        outputs=run_outputs
     )
 
+    for idx, row_cb in enumerate(row_checkboxes):
+        row_cb.change(
+            fn=lambda changed_values, selected_state, row_defs, row_values, i=idx: _on_toggle_row(
+                i,
+                changed_values,
+                selected_state,
+                row_defs,
+                row_values,
+                display_max_rows_default,
+            ),
+            inputs=[row_cb, selected_tags_state, row_defs_state, row_values_state],
+            outputs=[selected_tags_state, row_values_state, suggested_prompt, *row_checkboxes],
+        )
+
 if __name__ == "__main__":
     app.queue().launch(allowed_paths=[str(MASCOT_DIR)])