Upload alpha_factory/run.py

alpha_factory/run.py

@@ -1,135 +1,124 @@
 """
-Alpha Factory — Entry Point
-Run: python -m alpha_factory.run [--dry-run] [--proven] [--batch-size N]
+Alpha Factory — Entry Point v2
+Run: python -m alpha_factory.run [--dry-run] [--batch-size N] [--interactive]
+                                 [--proven] [--enable-brain]
 """
 import os
 import asyncio
 import argparse
 
+# Load .env file FIRST before anything else reads env vars
 try:
     from dotenv import load_dotenv
-    load_dotenv()
+    load_dotenv()  # Reads .env from current directory
 except ImportError:
-    pass
+    pass  # python-dotenv not installed; rely on system env vars
 
 from rich.console import Console
 from .config import load_config
-from .infra import ModelManager, LLMClient
+from .infra import ModelManager, interactive_model_select, LLMClient
 from .orchestration import AlphaPipeline
 
 console = Console()
 
 
 async def setup_models(interactive: bool = False, hf_token: str = None) -> ModelManager:
+    """Discover models and optionally let user pick interactively."""
     manager = ModelManager(hf_token=hf_token)
-    console.print("\n[bold]Discovering available models...[/]")
+
+    console.print("\n[bold]🔍 Discovering available models...[/]")
     await manager.discover_all()
-    manager.auto_assign_defaults()
-    manager.print_status()
-    return manager
 
+    if interactive:
+        selections = interactive_model_select(manager)
+        for tier, model in selections.items():
+            manager.select_model(tier, model)
+    else:
+        manager.auto_assign_defaults()
 
-def run_proven_mode(batch_size: int):
-    """Generate alphas using proven templates only — no LLM needed."""
-    from .deterministic.proven_templates import generate_batch_from_proven_templates
-    from .deterministic.lint import lint
-    from .infra import FactorStore
-    from pathlib import Path
-
-    console.print(f"\n[bold green]--- Proven Template Mode (no LLM) ---[/]")
-    console.print(f"Generating {batch_size} alphas from Alpha 15/6 structures\n")
-
-    store = FactorStore(Path("factor_store/alphas.duckdb"))
-    existing_hashes = store.get_expression_hashes()
-
-    batch = generate_batch_from_proven_templates(count=batch_size)
-    passed = 0
-    stored = 0
-
-    for i, alpha in enumerate(batch, 1):
-        expr = alpha["expression"]
-        field = alpha["field_id"]
-        template = alpha["template"]
-        ac = alpha["field_ac"]
-        group = alpha["group_key"]
-
-        result = lint(expr)
-        status = "[green]PASS[/]" if result.passed else "[red]FAIL[/]"
-        console.print(f"  {i}. {status} [{template}] field={field} (AC={ac}) group={group}")
-
-        if result.passed:
-            passed += 1
-            # Store
-            from .deterministic.lint import quick_dedup_hash
-            alpha_id = quick_dedup_hash(expr, alpha["neutralization"], alpha["decay"])
-            if alpha_id not in existing_hashes:
-                store.insert_alpha(
-                    alpha_id=alpha_id,
-                    expression=expr,
-                    neutralization=alpha["neutralization"],
-                    decay=alpha["decay"],
-                    fields_used=[alpha["field_id"]],
-                    operators_used=["ts_decay_linear", "group_neutralize", "rank", "zscore", "ts_rank"],
-                    archetype=alpha["archetype"],
-                    theme=alpha["theme"],
-                    anomaly_tag="value",
-                    academic_anchor=None,
-                )
-                existing_hashes.add(alpha_id)
-                stored += 1
-                console.print(f"      {expr[:90]}...")
-            else:
-                console.print(f"      [yellow]DEDUP: already exists[/]")
-        else:
-            console.print(f"      [red]Errors: {result.errors}[/]")
-
-    console.print(f"\n[bold]Results: {passed}/{len(batch)} passed lint, {stored} stored[/]")
-    console.print("[dim]Copy expressions from UI or DuckDB. Paste into BRAIN with settings: USA/TOP3000/D1/Decay=5[/]")
-    store.close()
+    manager.print_status()
+    return manager
 
 
 def main():
-    parser = argparse.ArgumentParser(description="Alpha Factory")
-    parser.add_argument("--dry-run", action="store_true", help="No BRAIN submissions")
-    parser.add_argument("--proven", action="store_true", help="Use proven templates only (no LLM, recommended)")
-    parser.add_argument("--batch-size", type=int, default=10, help="Candidates per batch")
-    parser.add_argument("--interactive", action="store_true", help="Select models interactively")
-    parser.add_argument("--hf-token", type=str, default=None, help="HuggingFace API token")
-    parser.add_argument("--ollama-url", type=str, default="http://localhost:11434", help="Ollama URL")
+    parser = argparse.ArgumentParser(description="Alpha Factory — LLM-Driven Alpha Generation Pipeline")
+    parser.add_argument("--dry-run", action="store_true", help="Run without BRAIN submissions")
+    parser.add_argument("--batch-size", type=int, default=10, help="Number of candidates per batch")
+    parser.add_argument("--interactive", action="store_true", help="Interactively select models")
+    parser.add_argument("--hf-token", type=str, default=None, help="HuggingFace API token (or set HF_TOKEN env)")
+    parser.add_argument("--ollama-url", type=str, default="http://localhost:11434", help="Ollama server URL")
+    parser.add_argument("--proven", action="store_true", help="Use proven templates (no LLM, deterministic generation)")
+    parser.add_argument("--enable-brain", action="store_true", help="Enable live BRAIN submission (requires BRAIN_SESSION_TOKEN)")
+    parser.add_argument("--local-sim", action="store_true", help="Enable local BRAIN simulation (requires yfinance)")
     args = parser.parse_args()
 
     config = load_config()
     config.batch_size = args.batch_size
+    config.use_proven_templates = args.proven
+    config.enable_brain_client = args.enable_brain
+    config.enable_local_sim = args.local_sim
+
+    if args.dry_run:
+        config.enable_brain_client = False
+
+    # Resolve HF token: CLI arg > env var (loaded from .env)
     hf_token = args.hf_token or os.getenv("HF_TOKEN")
 
+    mode_str = "PROVEN TEMPLATES" if args.proven else "LLM GENERATION"
+    brain_str = "LIVE (BRAIN submissions)" if config.enable_brain_client else "DRY RUN"
+
     console.print(f"""
-[bold green]Alpha Factory v0.2.0[/]
+[bold green]━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━[/]
+[bold green]  ALPHA FACTORY v0.2.0[/]
+[bold green]  Open-Source LLM-Driven Pipeline for WorldQuant BRAIN[/]
+[bold green]━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━[/]
 
-  Mode:       {"PROVEN TEMPLATES" if args.proven else "LLM-ASSISTED (dry-run)" if args.dry_run else "LLM-ASSISTED"}
+  Mode:       {mode_str}
+  Brain:      {brain_str}
   Batch size: {args.batch_size}
-  HF Token:   {"Set" if hf_token else "Not set"}
+  Ollama:     {args.ollama_url}
+  HF Token:   {"✓ Set" if hf_token else "✗ Not set (cloud models unavailable)"}
+  
+  [yellow]NOTE: This is v0.2.0 with real personas wired, but BRAIN integration
+  requires a valid BRAIN_SESSION_TOKEN. See .env.example for setup.[/]
 """)
 
-    # PROVEN MODE — no LLM needed, guaranteed valid expressions
-    if args.proven:
-        run_proven_mode(args.batch_size)
-        return
-
-    # LLM MODE — requires HF token or Ollama
-    manager = asyncio.run(setup_models(interactive=args.interactive, hf_token=hf_token))
+    # Discover and select models (only needed for LLM mode)
+    if not args.proven:
+        manager = asyncio.run(setup_models(
+            interactive=args.interactive,
+            hf_token=hf_token,
+        ))
+    else:
+        manager = None
+        console.print("  [green]Proven template mode — no LLM model discovery needed[/]")
+
+    # Update LLM config with Ollama URL
     config.llm.base_url = f"{args.ollama_url}/v1"
 
+    # Create pipeline
     pipeline = AlphaPipeline(config)
-    pipeline.llm = LLMClient(config.llm, model_manager=manager)
+    if manager:
+        pipeline.llm = LLMClient(config.llm, model_manager=manager)
+
+    # Initialize BRAIN client if enabled
+    if config.enable_brain_client:
+        try:
+            import aiohttp
+            session = aiohttp.ClientSession()
+            asyncio.run(pipeline.init_brain_client(session))
+        except ImportError:
+            console.print("[red]aiohttp required for BRAIN client. pip install aiohttp[/]")
+            config.enable_brain_client = False
 
     try:
         result = asyncio.run(pipeline.run_batch(args.batch_size))
         console.print(f"\n[bold]Final: {result}[/]")
     except KeyboardInterrupt:
-        console.print("\n[yellow]Interrupted[/]")
+        console.print("\n[yellow]Interrupted by user[/]")
     finally:
         pipeline.close()
 
 
 if __name__ == "__main__":
-    main()
+    main()