alpha_factory/ui.py

"""
Alpha Factory — Gradio UI
View generated alphas, copy expressions, run new batches.

Run: uv run python -m alpha_factory.ui
"""
import os
import sys
import subprocess
import duckdb
import gradio as gr
from pathlib import Path

try:
    from dotenv import load_dotenv
    load_dotenv()
except ImportError:
    pass

DB_PATH = Path("factor_store/alphas.duckdb")


def get_alphas_from_db(limit=50):
    if not DB_PATH.exists():
        return []
    conn = duckdb.connect(str(DB_PATH), read_only=True)
    try:
        rows = conn.execute(f"""
            SELECT 
                alpha_id,
                submitted_at,
                expression,
                theme,
                archetype,
                anomaly_tag,
                neutralization,
                decay,
                fields_used,
                verdict
            FROM alphas 
            ORDER BY submitted_at DESC 
            LIMIT {limit}
        """).fetchall()
        return rows
    except Exception:
        return []
    finally:
        conn.close()


def get_alpha_cards():
    rows = get_alphas_from_db()
    if not rows:
        return [["—", "—", "—", "—", "—", "—", "—", "Run a batch first"]]
    data = []
    for row in rows:
        alpha_id, submitted_at, expression, theme, archetype, tag, neutral, decay, fields, verdict = row
        timestamp = submitted_at.strftime("%Y-%m-%d %H:%M") if submitted_at else "?"
        verdict_str = {"promote": "PASS", "iterate": "PENDING", "kill": "FAIL"}.get(verdict or "", "NEW")
        expr_preview = (expression[:80] + "...") if expression and len(expression) > 80 else (expression or "")
        data.append([timestamp, alpha_id[:10], theme or "", archetype or "", tag or "", str(decay or 0), verdict_str, expr_preview])
    return data


def get_full_expression(evt: gr.SelectData):
    rows = get_alphas_from_db()
    if not rows or evt.index is None:
        return "Click a row above to see the full expression"
    row_idx = evt.index[0] if isinstance(evt.index, (list, tuple)) else evt.index
    if row_idx < len(rows):
        return rows[row_idx][2] or ""
    return ""


def run_batch(batch_size):
    """Run pipeline as subprocess with forced UTF-8 to avoid Windows encoding crash."""
    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

    try:
        result = subprocess.run(
            [sys.executable, "-m", "alpha_factory.run", "--dry-run", "--batch-size", str(int(batch_size))],
            capture_output=True,
            env=env,
            timeout=180,
            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:]
        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."
    except Exception as e:
        return f"ERROR: {str(e)}"


def generate_and_refresh(batch_size):
    log = run_batch(batch_size)
    table = get_alpha_cards()
    return table, log


def build_ui():
    with gr.Blocks(title="Alpha Factory") as app:
        gr.Markdown("""
        # Alpha Factory — Generated Alphas
        View, copy, and manage alphas generated by the pipeline.
        """)

        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,
        )

        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=15, interactive=False)

        # 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],
        )

    return app


if __name__ == "__main__":
    app = build_ui()
    app.launch(share=False)