Upload alpha_factory/run.py

alpha_factory/run.py

@@ -1,5 +1,6 @@
 """
-Alpha Factory — Entry Point v2
+Alpha Factory — Entry Point v3
+Single asyncio event loop, no session leaks.
 Run: python -m alpha_factory.run [--dry-run] [--batch-size N] [--interactive]
                                  [--proven] [--enable-brain]
 """
@@ -40,6 +41,39 @@ async def setup_models(interactive: bool = False, hf_token: str = None) -> Model
     return manager
 
 
+async def run_pipeline(config, manager: ModelManager | None = None):
+    """Run the full pipeline under a single event loop."""
+    # Update LLM config with Ollama URL
+    config.llm.base_url = f"{args.ollama_url}/v1"
+
+    # Create pipeline
+    pipeline = AlphaPipeline(config)
+    if manager:
+        pipeline.llm = LLMClient(config.llm, model_manager=manager)
+
+    # Initialize BRAIN client if enabled
+    if config.enable_brain_client:
+        try:
+            import aiohttp
+            async with aiohttp.ClientSession() as session:
+                await pipeline.init_brain_client(session)
+                result = await pipeline.run_batch(config.batch_size)
+        except ImportError:
+            console.print("[red]aiohttp required for BRAIN client. pip install aiohttp[/]")
+            config.enable_brain_client = False
+            result = await pipeline.run_batch(config.batch_size)
+        except Exception as e:
+            console.print(f"[red]BRAIN client error: {e}[/]")
+            config.enable_brain_client = False
+            result = await pipeline.run_batch(config.batch_size)
+    else:
+        result = await pipeline.run_batch(config.batch_size)
+
+    console.print(f"\n[bold]Final: {result}[/]")
+    pipeline.close()
+    return result
+
+
 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")
@@ -77,45 +111,25 @@ def main():
   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.[/]
+  [yellow]NOTE: BRAIN integration requires a valid BRAIN_SESSION_TOKEN.[/]
+  [yellow]      See .env.example for setup.[/]
 """)
 
-    # 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)
-    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}[/]")
+        if args.proven:
+            console.print("  [green]Proven template mode — no LLM model discovery needed[/]")
+            asyncio.run(run_pipeline(config, manager=None))
+        else:
+            # Discover and select models, then run pipeline
+            async def _discover_and_run():
+                manager = await setup_models(
+                    interactive=args.interactive,
+                    hf_token=hf_token,
+                )
+                return await run_pipeline(config, manager=manager)
+            asyncio.run(_discover_and_run())
     except KeyboardInterrupt:
         console.print("\n[yellow]Interrupted by user[/]")
-    finally:
-        pipeline.close()
 
 
 if __name__ == "__main__":