[NET-1341] update: Use new split version of initialse multi turn runs and user message generation

netra/simulation/api.py

@@ -94,36 +94,170 @@ def run_simulation(
         hooks_meta = hooks.describe() if hooks else None
 
         start_time = time.time()
-        run_result = self._client.create_run(
+
+        # --- Phase 1: Initialize run (DB only, no LLM) ---
+        init_result = self._client.initialize_run(
             name=name,
             dataset_id=dataset_id,
             context=context or {},
             hooks_meta=hooks_meta,
         )
-        if not run_result:
+        if init_result is None:
             return None
 
-        run_id = run_result.get("run_id")
-        simulation_items = run_result.get("simulation_items")
-        if not simulation_items:
-            logger.error("%s: No items returned from create_run", LOG_PREFIX)
+        run_id: str = init_result["run_id"]
+        items: list[dict[str, str]] = init_result["items"]
+
+        if not items:
+            logger.error("%s: No items returned from initialize_run", LOG_PREFIX)
             return None
 
+        # --- Phase 2: Run before_all hook (no LLM cost yet) ---
+        shared_context: Optional[dict[str, Any]] = None
+        if hooks and hooks.before_all is not None:
+            try:
+                shared_context = run_async_safely(run_before_all(hooks))
+            except Exception as exc:
+                logger.error(
+                    "%s: before_all hook failed: %s — aborting run (no LLM spent)",
+                    LOG_PREFIX,
+                    exc,
+                    exc_info=True,
+                )
+                for item in items:
+                    self._client.report_failure(
+                        run_id=run_id,
+                        run_item_id=item["test_run_item_id"],
+                        error=f"before_all hook failed: {exc}",
+                        status="prescript_failed",
+                    )
+                self._client.post_run_status(run_id, "completed")
+                return {
+                    "success": False,
+                    "completed": [],
+                    "failed": [
+                        {
+                            "run_item_id": item["test_run_item_id"],
+                            "success": False,
+                            "error": f"before_all hook failed: {exc}",
+                        }
+                        for item in items
+                    ],
+                    "total_items": len(items),
+                }
+
+        # --- Phase 3: Run per-item before hooks + generate first turns (parallel) ---
+        simulation_items: list[SimulationItem] = []
+        setup_contexts: dict[str, Optional[dict[str, Any]]] = {}
+        failed_items: list[dict[str, Any]] = []
+
+        async def run_before_and_generate_first_turn(item: dict[str, str]) -> Optional[SimulationItem]:
+            """Execute the before hook and generate the first turn for a single item.
+
+            Returns:
+                SimulationItem if successful, None if failed (failure recorded in failed_items).
+            """
+            run_item_id = item["test_run_item_id"]
+            dataset_item_id = item["dataset_item_id"]
+
+            has_before_hook = hooks and hooks.before and dataset_item_id in hooks.before
+            if has_before_hook:
+                try:
+                    item_context = await run_before(hooks, dataset_item_id, shared_context)
+                    setup_contexts[run_item_id] = item_context
+                except Exception as exc:
+                    error_msg = f"before hook failed: {exc}"
+                    logger.error(
+                        "%s: %s for run_item_id=%s",
+                        LOG_PREFIX,
+                        error_msg,
+                        run_item_id,
+                        exc_info=True,
+                    )
+                    self._client.report_failure(
+                        run_id=run_id,
+                        run_item_id=run_item_id,
+                        error=error_msg,
+                        status="prescript_failed",
+                    )
+                    item_result = {
+                        "run_item_id": run_item_id,
+                        "success": False,
+                        "error": error_msg,
+                    }
+                    failed_items.append(item_result)
+                    # Call after hook even when before fails (cleanup contract)
+                    await run_after(hooks, dataset_item_id, item_result, shared_context)
+                    return None
+            else:
+                setup_contexts[run_item_id] = shared_context
+
+            sim_item = self._client.generate_first_turn(
+                run_id=run_id,
+                run_item_id=run_item_id,
+            )
+            if sim_item is None:
+                logger.warning(
+                    "%s: Failed to generate first turn for item %s, marking failed",
+                    LOG_PREFIX,
+                    run_item_id,
+                )
+                self._client.report_failure(
+                    run_id=run_id,
+                    run_item_id=run_item_id,
+                    error="Failed to generate first user message",
+                )
+                item_result = {
+                    "run_item_id": run_item_id,
+                    "success": False,
+                    "error": "Failed to generate first user message",
+                }
+                failed_items.append(item_result)
+                # Call after hook when first turn generation fails (cleanup contract)
+                await run_after(hooks, dataset_item_id, item_result, setup_contexts[run_item_id])
+                return None
+            return sim_item
+
+        async def _run_all_before_and_first_turns() -> list[Optional[SimulationItem]]:
+            return list(await asyncio.gather(*[run_before_and_generate_first_turn(item) for item in items]))
+
+        setup_results = run_async_safely(_run_all_before_and_first_turns())
+        simulation_items = [item for item in setup_results if item is not None]
+
+        if not simulation_items:
+            logger.error("%s: All items failed during setup/generation", LOG_PREFIX)
+            self._client.post_run_status(run_id, "completed")
+            return {
+                "success": False,
+                "completed": [],
+                "failed": failed_items,
+                "total_items": len(items),
+            }
+
+        # --- Phase 4: Run conversation loops (before hooks already executed) ---
         logger.info("%s: Starting simulation with %d items", LOG_PREFIX, len(simulation_items))
         try:
             result = run_async_safely(
                 self._run_simulation_async(
-                    run_id, simulation_items, task, max_concurrency, max_turns, hooks  # type:ignore[arg-type]
+                    run_id,
+                    simulation_items,
+                    task,
+                    max_concurrency,
+                    max_turns,
+                    hooks,
+                    shared_context=shared_context,
+                    setup_contexts=setup_contexts,
+                    setup_failed_items=failed_items,
                 )
             )
 
             elapsed_time = time.time() - start_time
             logger.info("%s: Simulation completed in %.2f seconds", LOG_PREFIX, elapsed_time)
-            self._client.post_run_status(run_id, "completed")  # type:ignore[arg-type]
+            self._client.post_run_status(run_id, "completed")
             return result
         except Exception:
             logger.error("%s: Run simulation failed", LOG_PREFIX, exc_info=True)
-            self._client.post_run_status(run_id, "failed")  # type:ignore[arg-type]
+            self._client.post_run_status(run_id, "failed")
             return None
 
     async def _run_simulation_async(
@@ -134,20 +268,19 @@ async def _run_simulation_async(
         max_concurrency: int,
         max_turns: int,
         hooks: Optional[SimulationHooks],
+        shared_context: Optional[dict[str, Any]] = None,
+        setup_contexts: Optional[dict[str, Optional[dict[str, Any]]]] = None,
+        setup_failed_items: Optional[list[dict[str, Any]]] = None,
     ) -> dict[str, Any]:
         """Orchestrate concurrent simulation execution.
 
         Each simulation item is dispatched to a thread via ``run_in_executor``.
         Inside each thread, ``run_async_safely`` creates a **new** event loop
         so that async user tasks (``BaseTask.run``) work correctly without
-        nesting into the orchestrator's loop.  This two-level design lets us
-        honour ``max_concurrency`` while supporting both sync and async tasks
-        transparently.
+        nesting into the orchestrator's loop.
 
-        Executes ``before_all`` first (if configured). If it fails the entire
-        run is aborted.  Individual scenarios run concurrently via a thread
-        pool; each thread gets its own event loop so sync and async tasks work
-        without loop nesting.
+        When ``setup_contexts`` is provided, per-item before hooks have already
+        been executed and the conversation loop skips them.
 
         Args:
             run_id: The simulation run identifier.
@@ -156,6 +289,9 @@ async def _run_simulation_async(
             max_concurrency: Maximum concurrent executions.
             max_turns: Maximum conversation turns per item.
             hooks: Optional lifecycle hooks.
+            shared_context: Pre-computed before_all context.
+            setup_contexts: Pre-computed per-item setup contexts keyed by run_item_id.
+            setup_failed_items: Items that failed during before-hook or first-turn generation.
 
         Returns:
             Dictionary with simulation results.
@@ -167,51 +303,23 @@ async def _run_simulation_async(
             "total_items": len(simulation_items),
         }
 
-        # --- before_all ---
-        shared_context: Optional[dict[str, Any]] = None
-        if hooks and hooks.before_all is not None:
-            try:
-                shared_context = await run_before_all(hooks)
-            except Exception as exc:
-                logger.error(
-                    "%s: before_all hook failed: %s — aborting run",
-                    LOG_PREFIX,
-                    exc,
-                    exc_info=True,
-                )
-                # Mark every item as prescript_failed and abort
-                for item in simulation_items:
-                    self._client.report_failure(
-                        run_id=run_id,
-                        run_item_id=item.run_item_id,
-                        error=f"before_all hook failed: {exc}",
-                        status="prescript_failed",
-                    )
-                results["success"] = False
-                results["failed"] = [
-                    {
-                        "run_item_id": item.run_item_id,
-                        "success": False,
-                        "error": f"before_all hook failed: {exc}",
-                    }
-                    for item in simulation_items
-                ]
-                return results
-
         processed_count = 0
         lock = asyncio.Lock()
         loop = asyncio.get_running_loop()
 
         def run_item_in_thread(run_item: SimulationItem) -> dict[str, Any]:
-            """Run a single simulation item in a dedicated thread/event-loop.
-
-            Args:
-                run_item: The simulation item to run.
-
-            Returns: Dictionary with simulation result.
-            """
+            """Run a single simulation item in a dedicated thread/event-loop."""
+            precomputed = setup_contexts.get(run_item.run_item_id) if setup_contexts is not None else None
             return run_async_safely(
-                self._execute_conversation(run_id, run_item, task, max_turns, hooks, shared_context)
+                self._execute_conversation(
+                    run_id,
+                    run_item,
+                    task,
+                    max_turns,
+                    hooks,
+                    shared_context,
+                    setup_context=precomputed,
+                )
             )
 
         async def process_item(run_item: SimulationItem) -> None:
@@ -244,6 +352,10 @@ async def process_item(run_item: SimulationItem) -> None:
                 await asyncio.gather(*tasks, return_exceptions=True)
                 executor.shutdown(wait=False, cancel_futures=True)
 
+        # Merge setup failures before calling after_all so it sees complete results
+        if setup_failed_items:
+            results.setdefault("failed", []).extend(setup_failed_items)
+
         # --- after_all ---
         await run_after_all(hooks, results, shared_context)
 
@@ -263,11 +375,13 @@ async def _execute_conversation(
         max_turns: int,
         hooks: Optional[SimulationHooks],
         shared_context: Optional[dict[str, Any]],
+        setup_context: Optional[dict[str, Any]] = None,
     ) -> dict[str, Any]:
         """Execute a multi-turn conversation for a single simulation item.
 
-        Runs the ``before`` hook before starting the conversation loop and
-        the ``after`` hook when the conversation ends (success or failure).
+        The per-item ``before`` hook has already been executed by the caller;
+        the merged ``setup_context`` is passed directly to ``BaseTask.run``.
+        The ``after`` hook runs when the conversation ends (success or failure).
 
         Args:
             run_id: The simulation run identifier.
@@ -276,6 +390,7 @@ async def _execute_conversation(
             max_turns: Safety limit on the number of conversation turns.
             hooks: Optional lifecycle hooks.
             shared_context: Context dict returned by the ``before_all`` hook.
+            setup_context: Merged context from ``before_all`` + item ``before``.
 
         Returns:
             Dictionary with execution result including success status.
@@ -287,37 +402,6 @@ async def _execute_conversation(
         raw_files: list[FileData] = run_item.files
         session_id: Optional[str] = None
 
-        # --- before ---
-        setup_context: Optional[dict[str, Any]] = None
-        has_before_hooks = hooks and hooks.before and dataset_item_id in hooks.before
-        if has_before_hooks:
-            try:
-                setup_context = await run_before(hooks, dataset_item_id, shared_context)
-            except Exception as exc:
-                error_msg = f"before hook failed: {exc}"
-                logger.error(
-                    "%s: %s for run_item_id=%s",
-                    LOG_PREFIX,
-                    error_msg,
-                    run_item_id,
-                    exc_info=True,
-                )
-                self._client.report_failure(
-                    run_id=run_id,
-                    run_item_id=run_item_id,
-                    error=error_msg,
-                    status="prescript_failed",
-                )
-                item_result = {
-                    "run_item_id": run_item_id,
-                    "success": False,
-                    "error": error_msg,
-                }
-                await run_after(hooks, dataset_item_id, item_result, shared_context)
-                return item_result
-        else:
-            setup_context = shared_context
-
         for turn_number in range(1, max_turns + 1):
             try:
                 with SpanWrapper(SPAN_NAME, module_name=LOG_PREFIX) as span:
@@ -402,4 +486,5 @@ async def _execute_conversation(
             "error": error_msg,
             "turn_id": turn_id,
         }
+        await run_after(hooks, dataset_item_id, item_result, setup_context)
         return item_result