Upload alpha_factory/ui.py

alpha_factory/ui.py

@@ -1,15 +1,19 @@
 """
-Alpha Factory — Gradio UI
-View generated alphas, copy expressions, run new batches.
+Alpha Factory — Gradio UI v2
+View generated alphas, copy expressions, run new batches,
+and SELECT per-tier models from discovered Ollama + HuggingFace models.
 
 Run: uv run python -m alpha_factory.ui
 """
 import os
 import sys
 import subprocess
+import json
+import asyncio
 import duckdb
 import gradio as gr
 from pathlib import Path
+from typing import Optional
 
 try:
     from dotenv import load_dotenv
@@ -17,8 +21,49 @@ try:
 except ImportError:
     pass
 
+from .config import load_config
+from .infra.model_manager import ModelManager, ModelInfo, ModelProvider
+
 DB_PATH = Path("factor_store/alphas.duckdb")
 
+# ── Globals (shared across Gradio sessions) ──────────────────────────────────
+_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}"
+
+
+def _discover_models_sync(
+    ollama_url: str = "http://localhost:11434",
+    hf_token: Optional[str] = None,
+) -> list[ModelInfo]:
+    """Synchronous wrapper around async model discovery."""
+    global _LAST_DISCOVERED_MODELS
+
+    # Resolve HF token
+    token = hf_token or os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACE_TOKEN", "")
+
+    manager = ModelManager(ollama_url=ollama_url, hf_token=token)
+    try:
+        asyncio.run(manager.discover_all())
+    except Exception as e:
+        print(f"Model discovery error: {e}")
+
+    _LAST_DISCOVERED_MODELS = manager.get_all_models()
+    return _LAST_DISCOVERED_MODELS
+
+
+def _get_dropdown_choices(models: list[ModelInfo]) -> list[str]:
+    """Build dropdown choices: [Use Default (auto-assign)] + discovered models."""
+    choices = ["Use Default (auto-assign)"]
+    choices.extend([_model_choice_name(m) for m in models])
+    return choices
+
+
+# ── DB helpers (unchanged) ──────────────────────────────────────────────────
 
 def get_alphas_from_db(limit=50):
     if not DB_PATH.exists():
@@ -72,96 +117,308 @@ def get_full_expression(evt: gr.SelectData):
     return ""
 
 
-def run_batch(batch_size):
-    """Run pipeline as subprocess with forced UTF-8 to avoid Windows encoding crash."""
+# ── Pipeline runner (now with per-tier model overrides) ───────────────────────
+
+def _run_pipeline_subprocess(
+    batch_size: int,
+    proven_mode: bool,
+    enable_brain: bool,
+    ollama_url: str,
+    microfish: str,
+    tinyfish: str,
+    mediumfish: str,
+    bigfish: str,
+) -> str:
+    """Run the pipeline as a subprocess with the selected configuration."""
     env = os.environ.copy()
-    # Force UTF-8 output — prevents Rich/Windows cp1252 crash
     env["PYTHONIOENCODING"] = "utf-8"
     env["PYTHONLEGACYWINDOWSSTDIO"] = "utf-8"
-    # Disable Rich color/formatting when piped (cleaner output)
     env["NO_COLOR"] = "1"
     env["TERM"] = "dumb"
-    # Ensure HF_TOKEN passes through
-    if "HF_TOKEN" not in env:
-        token = os.getenv("HF_TOKEN", "")
-        if token:
-            env["HF_TOKEN"] = token
+
+    # Build CLI args
+    cmd = [
+        sys.executable, "-m", "alpha_factory.run",
+        "--batch-size", str(int(batch_size)),
+        "--ollama-url", ollama_url,
+    ]
+
+    if proven_mode:
+        cmd.append("--proven")
+    if enable_brain:
+        cmd.append("--enable-brain")
+
+    # Only pass per-tier overrides if user selected something other than default
+    def _extract_model_name(choice: str) -> Optional[str]:
+        if not choice or choice == "Use Default (auto-assign)":
+            return None
+        # Strip the [PROVIDER] prefix and size/quant suffix
+        if "]" in choice:
+            return choice.split("]", 1)[1].strip().split(" (")[0].split(" [")[0].strip()
+        return choice.strip()
+
+    mf = _extract_model_name(microfish)
+    tf = _extract_model_name(tinyfish)
+    mmf = _extract_model_name(mediumfish)
+    bf = _extract_model_name(bigfish)
+
+    if mf:
+        cmd.extend(["--microfish", mf])
+    if tf:
+        cmd.extend(["--tinyfish", tf])
+    if mmf:
+        cmd.extend(["--mediumfish", mmf])
+    if bf:
+        cmd.extend(["--bigfish", bf])
+
+    # Log the command for debugging
+    print(f"Running: {' '.join(cmd)}")
 
     try:
         result = subprocess.run(
-            [sys.executable, "-m", "alpha_factory.run", "--dry-run", "--batch-size", str(int(batch_size))],
+            cmd,
             capture_output=True,
             env=env,
-            timeout=180,
+            timeout=300,
             cwd=str(Path.cwd()),
         )
-        # Decode with utf-8, replace errors
         stdout = result.stdout.decode("utf-8", errors="replace") if result.stdout else ""
         stderr = result.stderr.decode("utf-8", errors="replace") if result.stderr else ""
-        
+
         log = ""
         if stdout:
-            log = stdout[-3000:]
+            log = stdout[-4000:]
         if result.returncode != 0 and stderr:
             log += "\n\n--- ERRORS ---\n" + stderr[-2000:]
         if not log.strip():
             log = f"Process exited with code {result.returncode}"
         return log
     except subprocess.TimeoutExpired:
-        return "ERROR: Pipeline timed out after 180 seconds. Try smaller batch size."
+        return "ERROR: Pipeline timed out after 300 seconds. Try smaller batch size."
     except Exception as e:
         return f"ERROR: {str(e)}"
 
 
-def generate_and_refresh(batch_size):
-    log = run_batch(batch_size)
+def generate_and_refresh(
+    batch_size,
+    proven_mode,
+    enable_brain,
+    ollama_url,
+    microfish,
+    tinyfish,
+    mediumfish,
+    bigfish,
+):
+    log = _run_pipeline_subprocess(
+        batch_size, proven_mode, enable_brain, ollama_url,
+        microfish, tinyfish, mediumfish, bigfish,
+    )
     table = get_alpha_cards()
     return table, log
 
 
+# ── Model discovery refresh ──────────────────────────────────────────────────
+
+def refresh_model_list(ollama_url: str, hf_token: str) -> tuple[str, list[str]]:
+    """Discover models and return (status_msg, dropdown_choices)."""
+    models = _discover_models_sync(ollama_url=ollama_url, hf_token=hf_token)
+
+    if not models:
+        return "No models found. Is Ollama running? Is HF_TOKEN set?", ["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} HF models"
+    choices = _get_dropdown_choices(models)
+    return msg, choices
+
+
+# ── UI Builder ──────────────────────────────────────────────────────────────
+
 def build_ui():
-    with gr.Blocks(title="Alpha Factory") as app:
+    with gr.Blocks(title="Alpha Factory v0.2.0") as app:
         gr.Markdown("""
-        # Alpha Factory — Generated Alphas
-        View, copy, and manage alphas generated by the pipeline.
+        # Alpha Factory — LLM-Driven Alpha Generation
+        Generate and manage equity alpha expressions for WorldQuant BRAIN.
         """)
 
-        with gr.Row():
-            with gr.Column(scale=1):
-                batch_size_input = gr.Number(value=3, label="Batch Size", minimum=1, maximum=20)
-                generate_btn = gr.Button("Generate New Batch", variant="primary")
-                refresh_btn = gr.Button("Refresh Table")
-                gr.Markdown("*Dry run mode — no BRAIN submissions*")
-            with gr.Column(scale=3):
-                stats_md = gr.Markdown(f"**Alphas in store:** {len(get_alphas_from_db())}")
-
-        gr.Markdown("### Click any row to see full expression")
-
-        alpha_table = gr.Dataframe(
-            value=get_alpha_cards(),
-            headers=["Time", "ID", "Theme", "Archetype", "Tag", "Decay", "Status", "Expression"],
-            interactive=False,
-            wrap=True,
-        )
+        # ─── SETTINGS TAB ─────────────────────────────────────────────────────
+        with gr.Tab("⚙️ Settings"):
+            with gr.Row():
+                with gr.Column(scale=1):
+                    gr.Markdown("### Connection")
+                    ollama_url_input = gr.Textbox(
+                        value="http://localhost:11434",
+                        label="Ollama URL",
+                    )
+                    hf_token_input = gr.Textbox(
+                        value=os.getenv("HF_TOKEN", ""),
+                        label="HF Token (optional)",
+                        type="password",
+                    )
+                    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",
+                        interactive=False,
+                    )
 
-        gr.Markdown("### Full Expression — Ctrl+A then Ctrl+C to copy")
-        full_expr = gr.Textbox(
-            label="Full Expression",
-            lines=6,
-            interactive=True,
-        )
+                with gr.Column(scale=2):
+                    gr.Markdown("### Model Selection — One Per Tier")
+                    gr.Markdown("""
+                    | Tier | Role | Typical Size |
+                    |------|------|-------------|
+                    | **Microfish** | Hypothesis generation (bulk) | 1.5B-3B |
+                    | **Tinyfish** | Expression compilation | 3B-7B |
+                    | **Mediumfish** | Crowd scout + Performance surgeon | 7B-14B |
+                    | **Bigfish** | Gatekeeper (final memo) | 14B-72B |
+                    """)
 
-        gr.Markdown("### Pipeline Log")
-        pipeline_log = gr.Textbox(label="Output", lines=15, interactive=False)
+                    # Initial choices: just default until discovery
+                    default_choices = ["Use Default (auto-assign)"]
 
-        # Events
-        alpha_table.select(get_full_expression, outputs=[full_expr])
-        refresh_btn.click(get_alpha_cards, outputs=[alpha_table])
-        generate_btn.click(
-            generate_and_refresh,
-            inputs=[batch_size_input],
-            outputs=[alpha_table, pipeline_log],
-        )
+                    microfish_dropdown = gr.Dropdown(
+                        choices=default_choices,
+                        value="Use Default (auto-assign)",
+                        label="Microfish — Hypothesis Generation",
+                    )
+                    tinyfish_dropdown = gr.Dropdown(
+                        choices=default_choices,
+                        value="Use Default (auto-assign)",
+                        label="Tinyfish — Expression Compilation",
+                    )
+                    mediumfish_dropdown = gr.Dropdown(
+                        choices=default_choices,
+                        value="Use Default (auto-assign)",
+                        label="Mediumfish — Critique & Diagnosis",
+                    )
+                    bigfish_dropdown = gr.Dropdown(
+                        choices=default_choices,
+                        value="Use Default (auto-assign)",
+                        label="Bigfish — Final Gatekeeper",
+                    )
+
+            # When refresh is clicked, update all dropdowns
+            refresh_models_btn.click(
+                fn=refresh_model_list,
+                inputs=[ollama_url_input, hf_token_input],
+                outputs=[discovery_status, microfish_dropdown],
+            ).then(
+                lambda choices: gr.Dropdown(choices=choices),
+                inputs=microfish_dropdown,
+                outputs=tinyfish_dropdown,
+            ).then(
+                lambda choices: gr.Dropdown(choices=choices),
+                inputs=microfish_dropdown,
+                outputs=mediumfish_dropdown,
+            ).then(
+                lambda choices: gr.Dropdown(choices=choices),
+                inputs=microfish_dropdown,
+                outputs=bigfish_dropdown,
+            )
+            # Actually the proper way: refresh returns one choices list,
+            # then update all 4 dropdowns with that same list
+            def _update_all_dropdowns(status, choices):
+                return status, choices, choices, choices, choices
+
+            refresh_models_btn.click(
+                fn=lambda url, token: _update_all_dropdowns(*refresh_model_list(url, token)),
+                inputs=[ollama_url_input, hf_token_input],
+                outputs=[
+                    discovery_status,
+                    microfish_dropdown,
+                    tinyfish_dropdown,
+                    mediumfish_dropdown,
+                    bigfish_dropdown,
+                ],
+            )
+
+        # ─── GENERATION TAB ───────────────────────────────────────────────────
+        with gr.Tab("🚀 Generate Alphas"):
+            with gr.Row():
+                with gr.Column(scale=1):
+                    batch_size_input = gr.Number(
+                        value=3, label="Batch Size", minimum=1, maximum=20,
+                    )
+                    proven_mode_cb = gr.Checkbox(
+                        value=False, label="Proven Templates (no LLM)",
+                    )
+                    enable_brain_cb = gr.Checkbox(
+                        value=False, label="Enable BRAIN Submission (needs token)",
+                    )
+                    gr.Markdown("---")
+                    gr.Markdown("*Selected models carry over from the Settings tab*")
+                    generate_btn = gr.Button("Generate New Batch", variant="primary")
+                    refresh_table_btn = gr.Button("Refresh Table Only")
+
+                with gr.Column(scale=3):
+                    stats_md = gr.Markdown(
+                        f"**Alphas in store:** {len(get_alphas_from_db())}"
+                    )
+
+            gr.Markdown("### Click any row to see the full expression")
+            alpha_table = gr.Dataframe(
+                value=get_alpha_cards(),
+                headers=["Time", "ID", "Theme", "Archetype", "Tag", "Decay", "Status", "Expression"],
+                interactive=False,
+                wrap=True,
+            )
+
+            gr.Markdown("### Full Expression — Ctrl+A then Ctrl+C to copy")
+            full_expr = gr.Textbox(
+                label="Full Expression",
+                lines=6,
+                interactive=True,
+            )
+
+            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(
+                fn=generate_and_refresh,
+                inputs=[
+                    batch_size_input,
+                    proven_mode_cb,
+                    enable_brain_cb,
+                    ollama_url_input,
+                    microfish_dropdown,
+                    tinyfish_dropdown,
+                    mediumfish_dropdown,
+                    bigfish_dropdown,
+                ],
+                outputs=[alpha_table, pipeline_log],
+            )
+
+        # ─── ABOUT TAB ──────────────────────────────────────────────────────
+        with gr.Tab("📖 About"):
+            gr.Markdown("""
+            **Alpha Factory v0.2.0** — Open-source LLM-driven pipeline for WorldQuant BRAIN.
+
+            ### How it works
+            1. **Microfish** generates alpha hypotheses (ideas)
+            2. **Tinyfish** compiles the idea into a BRAIN expression
+            3. **Mediumfish** critiques and diagnoses performance
+            4. **Bigfish** makes the final go/no-go decision
+
+            ### Modes
+            - **Proven Templates**: Deterministic, no LLM needed, guaranteed valid expressions
+            - **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
+            - Select which model to use for each tier, or leave as "Use Default"
+
+            ### BRAIN Integration
+            - Requires `BRAIN_SESSION_TOKEN` from browser devtools
+            - Enable "BRAIN Submission" checkbox (disabled by default for safety)
+
+            [GitHub / HuggingFace](https://huggingface.co/gaurv007/alpha-factory)
+            """)
 
     return app