Upload alpha_factory/run.py

alpha_factory/run.py

@@ -1,8 +1,11 @@
 """
-Alpha Factory — Entry Point v3
+Alpha Factory — Entry Point v4
 Single asyncio event loop, no session leaks.
+Per-tier model selection via CLI or UI.
 Run: python -m alpha_factory.run [--dry-run] [--batch-size N] [--interactive]
                                  [--proven] [--enable-brain]
+                                 [--microfish MODEL] [--tinyfish MODEL]
+                                 [--mediumfish MODEL] [--bigfish MODEL]
 """
 import os
 import asyncio
@@ -18,14 +21,27 @@ except ImportError:
 from rich.console import Console
 from .config import load_config
 from .infra import ModelManager, interactive_model_select, LLMClient
+from .infra.model_manager import ModelInfo, ModelProvider
 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)
+async def setup_models(
+    interactive: bool = False,
+    hf_token: str = None,
+    ollama_url: str = "http://localhost:11434",
+    per_tier: dict[str, str] | None = None,
+) -> ModelManager:
+    """Discover models and optionally let user pick interactively.
+
+    Args:
+        interactive: Open CLI prompt to pick per-tier models.
+        hf_token: HuggingFace token for cloud model validation.
+        ollama_url: Ollama server URL for local model discovery.
+        per_tier: Pre-selected model names (e.g. {"microfish": "qwen2.5:1.5b"}).
+    """
+    manager = ModelManager(ollama_url=ollama_url, hf_token=hf_token)
 
     console.print("\n[bold]🔍 Discovering available models...[/]")
     await manager.discover_all()
@@ -34,6 +50,24 @@ async def setup_models(interactive: bool = False, hf_token: str = None) -> Model
         selections = interactive_model_select(manager)
         for tier, model in selections.items():
             manager.select_model(tier, model)
+    elif per_tier:
+        # Resolve per-tier model names to ModelInfo objects
+        all_models = manager.get_all_models()
+        for tier, model_name in per_tier.items():
+            if not model_name:
+                continue
+            # Try exact match first
+            matched = [m for m in all_models if m.name == model_name]
+            if not matched:
+                # Try case-insensitive partial match
+                matched = [m for m in all_models if model_name.lower() in m.name.lower()]
+            if matched:
+                manager.select_model(tier, matched[0])
+                console.print(f"  [green]{tier}[/]: {matched[0].display_name()}")
+            else:
+                console.print(f"  [yellow]{tier}[/]: '{model_name}' not found — using default")
+        # Fill remaining tiers with defaults
+        manager.auto_assign_defaults()
     else:
         manager.auto_assign_defaults()
 
@@ -41,9 +75,13 @@ async def setup_models(interactive: bool = False, hf_token: str = None) -> Model
     return manager
 
 
-async def run_pipeline(config, ollama_url: str = "http://localhost:11434", manager: ModelManager | None = None):
+async def run_pipeline(
+    config,
+    ollama_url: str = "http://localhost:11434",
+    manager: ModelManager | None = None,
+):
     """Run the full pipeline under a single event loop.
-    
+
     Args:
         config: Pipeline configuration
         ollama_url: Base URL for Ollama (default: http://localhost:11434)
@@ -84,11 +122,16 @@ def main():
     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("--interactive", action="store_true", help="Interactively select models (CLI)")
     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)")
+    # Per-tier model overrides
+    parser.add_argument("--microfish", type=str, default=None, help="Model for hypothesis generation (bulk)")
+    parser.add_argument("--tinyfish", type=str, default=None, help="Model for expression compilation")
+    parser.add_argument("--mediumfish", type=str, default=None, help="Model for crowd scout + performance surgeon")
+    parser.add_argument("--bigfish", type=str, default=None, help="Model for gatekeeper (final memo)")
     args = parser.parse_args()
 
     config = load_config()
@@ -102,6 +145,15 @@ def main():
     # Resolve HF token: CLI arg > env var (loaded from .env)
     hf_token = args.hf_token or os.getenv("HF_TOKEN")
 
+    # Collect per-tier selections from CLI
+    per_tier = {
+        "microfish": args.microfish,
+        "tinyfish": args.tinyfish,
+        "mediumfish": args.mediumfish,
+        "bigfish": args.bigfish,
+    }
+    has_per_tier = any(v is not None for v in per_tier.values())
+
     mode_str = "PROVEN TEMPLATES" if args.proven else "LLM GENERATION"
     brain_str = "LIVE (BRAIN submissions)" if config.enable_brain_client else "DRY RUN"
 
@@ -131,6 +183,8 @@ def main():
                 manager = await setup_models(
                     interactive=args.interactive,
                     hf_token=hf_token,
+                    ollama_url=args.ollama_url,
+                    per_tier=per_tier if has_per_tier else None,
                 )
                 return await run_pipeline(config, ollama_url=args.ollama_url, manager=manager)
             asyncio.run(_discover_and_run())