Upload alpha_factory/ui.py

alpha_factory/ui.py

@@ -1,7 +1,10 @@
 """
-Alpha Factory — Gradio UI v2
+Alpha Factory — Gradio UI v3
 View generated alphas, copy expressions, run new batches,
-and SELECT per-tier models from discovered Ollama + HuggingFace models.
+and SELECT per-tier models from:
+  - Installed Ollama models
+  - Pullable Ollama models (recommended tags)
+  - HuggingFace Inference API (cloud)
 
 Run: uv run python -m alpha_factory.ui
 """
@@ -31,9 +34,7 @@ _LAST_DISCOVERED_MODELS: list[ModelInfo] = []
 
 def _model_choice_name(m: ModelInfo) -> str:
     """Human-readable label for a model in the dropdown."""
-    size = f" ({m.size_gb:.1f}GB)" if m.size_gb else ""
-    quant = f" [{m.quantization}]" if m.quantization else ""
-    return f"[{m.provider.value.upper()}] {m.name}{size}{quant}"
+    return m.display_name()
 
 
 def _discover_models_sync(
@@ -70,15 +71,16 @@ def refresh_all_dropdowns(ollama_url: str, hf_token: str):
         msg = "⚠️ No models found. Is Ollama running? Is HF_TOKEN set?"
         return msg, ["Use Default (auto-assign)"], ["Use Default (auto-assign)"], ["Use Default (auto-assign)"], ["Use Default (auto-assign)"]
 
-    local_count = sum(1 for m in models if m.provider == ModelProvider.OLLAMA)
-    cloud_count = sum(1 for m in models if m.provider == ModelProvider.HUGGINGFACE)
-    msg = f"✅ Found {local_count} Ollama + {cloud_count} HuggingFace models"
+    installed = sum(1 for m in models if m.provider == ModelProvider.OLLAMA and m.is_installed)
+    pullable = sum(1 for m in models if m.provider == ModelProvider.OLLAMA and not m.is_installed)
+    cloud = sum(1 for m in models if m.provider == ModelProvider.HUGGINGFACE)
+    msg = f"✅ Found {installed} Ollama installed + {pullable} pullable + {cloud} HF cloud"
 
     choices = _get_dropdown_choices(models)
     return msg, choices, choices, choices, choices
 
 
-# ── DB helpers (unchanged) ──────────────────────────────────────────────────
+# ── DB helpers ──────────────────────────────────────────────────
 
 def get_alphas_from_db(limit=50):
     if not DB_PATH.exists():
@@ -132,16 +134,19 @@ def get_full_expression(evt: gr.SelectData):
     return ""
 
 
-# ── Pipeline runner (with per-tier model overrides) ───────────────────────────
+# ── Pipeline runner ───────────────────────────
 
 def _extract_model_name(choice: str) -> Optional[str]:
-    """Strip [PROVIDER] prefix and size/quant suffix to get raw model name."""
+    """Strip [PROVIDER] prefix and size/quant/pullable suffix to get raw model name."""
     if not choice or choice == "Use Default (auto-assign)":
         return None
     if "]" in choice:
         raw = choice.split("]", 1)[1].strip()
         # Remove trailing size/quant suffixes: " (4.7GB) [q4_k_m]" etc.
         raw = raw.split(" (")[0].split(" [")[0].strip()
+        # Remove " [PULLABLE — ollama pull ...]" suffix if present
+        if "[PULLABLE" in raw:
+            raw = raw.split(" [PULLABLE")[0].strip()
         return raw
     return choice.strip()
 
@@ -261,9 +266,15 @@ def build_ui():
                     refresh_models_btn = gr.Button("🔍 Refresh Model List", variant="secondary")
                     discovery_status = gr.Textbox(
                         label="Discovery Status",
-                        value="Click 'Refresh Model List' to discover Ollama + HF models",
+                        value="Click 'Refresh Model List' to discover installed + pullable models",
                         interactive=False,
                     )
+                    gr.Markdown("""
+                    **Legend:**
+                    - ✅ `[ollama]` — already installed locally
+                    - ⬇️ `[ollama]` + `[PULLABLE]` — run `ollama pull <name>` first
+                    - ☁️ `[huggingface]` — cloud Inference API (needs HF token)
+                    """)
 
                 with gr.Column(scale=2):
                     gr.Markdown("### Model Selection — One Per Tier")
@@ -283,24 +294,28 @@ def build_ui():
                         value="Use Default (auto-assign)",
                         label="🐟 Microfish — Hypothesis Generation",
                         info="Bulk idea generation (fast, small model)",
+                        allow_custom_value=True,
                     )
                     tinyfish_dropdown = gr.Dropdown(
                         choices=default_choices,
                         value="Use Default (auto-assign)",
                         label="🐠 Tinyfish — Expression Compilation",
                         info="Converts hypothesis to BRAIN expression",
+                        allow_custom_value=True,
                     )
                     mediumfish_dropdown = gr.Dropdown(
                         choices=default_choices,
                         value="Use Default (auto-assign)",
                         label="🦈 Mediumfish — Critique & Diagnosis",
                         info="Novelty check + performance analysis",
+                        allow_custom_value=True,
                     )
                     bigfish_dropdown = gr.Dropdown(
                         choices=default_choices,
                         value="Use Default (auto-assign)",
                         label="🐋 Bigfish — Final Gatekeeper",
                         info="Production go/no-go decision (slowest, most capable)",
+                        allow_custom_value=True,
                     )
 
             refresh_models_btn.click(
@@ -360,7 +375,6 @@ def build_ui():
             gr.Markdown("### Pipeline Log")
             pipeline_log = gr.Textbox(label="Output", lines=20, interactive=False)
 
-            # Events
             alpha_table.select(get_full_expression, outputs=[full_expr])
             refresh_table_btn.click(get_alpha_cards, outputs=[alpha_table])
             generate_btn.click(
@@ -394,8 +408,9 @@ def build_ui():
             - **LLM Mode**: Uses local (Ollama) or cloud (HuggingFace) models
 
             ### Model Discovery
-            - Set your **Ollama URL** and click **Refresh Model List** to find local models
-            - Set your **HF Token** to see HuggingFace Inference API models
+            - **Installed Ollama models** — already pulled, ready to run
+            - **Pullable Ollama models** — recommended tags you can `ollama pull` first
+            - **HuggingFace cloud models** — Inference API, needs HF token
             - Select which model to use for each tier, or leave as "Use Default"
 
             ### BRAIN Integration