diff --git a/netra/simulation/api.py b/netra/simulation/api.py index d775a5c..f4d607e 100644 --- a/netra/simulation/api.py +++ b/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 diff --git a/netra/simulation/client.py b/netra/simulation/client.py index e0eda00..da21d0f 100644 --- a/netra/simulation/client.py +++ b/netra/simulation/client.py @@ -10,7 +10,8 @@ LOG_PREFIX, TELEMETRY_SUFFIX, URL_AGENT_RESPONSE, - URL_CREATE_RUN, + URL_FIRST_TURN, + URL_INITIALIZE_RUN, URL_RUN_ITEM_STATUS, URL_RUN_STATUS, ) @@ -109,68 +110,94 @@ def _build_headers(self, config: Config) -> dict[str, str]: headers["x-api-key"] = config.api_key return headers - def create_run( + def initialize_run( self, name: str, dataset_id: str, context: Optional[dict[str, Any]] = None, hooks_meta: Optional[dict[str, Any]] = None, ) -> Optional[dict[str, Any]]: - """Create a new simulation run for the specified dataset. + """Create run + items without generating first user messages. - Args: - name: Name of the simulation run. - dataset_id: Identifier of the dataset to simulate. - context: Optional context data for the simulation. - hooks_meta: Optional metadata describing configured hooks, stored - by the backend for UI display purposes. - - Returns: - Dictionary containing run_id and simulation_items, or None on failure. + Returns dict with run_id and items (testRunItemId + datasetItemId), + or None on failure. """ if not self._ensure_client(): return None response: Optional[httpx.Response] = None try: - url = URL_CREATE_RUN + url = URL_INITIALIZE_RUN payload: dict[str, Any] = { "name": name, "datasetId": dataset_id, "context": context or {}, } if hooks_meta: - # Run-level beforeAll/afterAll + item-level before/after (under items). payload["lifecycleHooks"] = hooks_meta response = self._client.post(url, json=payload) # type:ignore[union-attr] response.raise_for_status() data = response.json() response_data = data.get("data", {}) - user_messages = response_data.get("userMessages", []) - if not user_messages: - logger.warning("%s: No user messages returned from create_run", LOG_PREFIX) + run_id = response_data.get("id", "") + items = response_data.get("items", []) + if not items: + logger.warning("%s: No items returned from initialize_run", LOG_PREFIX) return None - run_id = response_data.get("id", "") - simulation_items = [ - SimulationItem( - run_item_id=msg.get("testRunItemId", ""), - dataset_item_id=msg.get("datasetItemId", ""), - message=msg.get("userMessage", ""), - turn_id=msg.get("turnId", ""), - files=self._parse_files(msg.get("attachments")), - ) - for msg in user_messages - ] return { "run_id": run_id, - "simulation_items": simulation_items, + "dataset_id": response_data.get("datasetId", ""), + "dataset_name": response_data.get("datasetName", ""), + "items": [ + { + "test_run_item_id": item.get("testRunItemId", ""), + "dataset_item_id": item.get("datasetItemId", ""), + } + for item in items + ], } + except Exception as exc: + error_msg = self._extract_error_message(response, exc) + logger.error("%s: Failed to initialize run: %s", LOG_PREFIX, error_msg) + return None + + def generate_first_turn( + self, + run_id: str, + run_item_id: str, + ) -> Optional[SimulationItem]: + """Generate the first user message for a single test run item. + Returns a SimulationItem with message/turnId populated, or None on failure. + """ + if not self._ensure_client(): + return None + + response: Optional[httpx.Response] = None + try: + url = URL_FIRST_TURN.format(run_id=run_id, run_item_id=run_item_id) + response = self._client.post(url, json={}) # type:ignore[union-attr] + response.raise_for_status() + data = response.json() + + resp = data.get("data", {}) + return SimulationItem( + run_item_id=run_item_id, + dataset_item_id=resp.get("datasetItemId", ""), + message=resp.get("userMessage", ""), + turn_id=resp.get("turnId", ""), + files=self._parse_files(resp.get("attachments")), + ) except Exception as exc: error_msg = self._extract_error_message(response, exc) - logger.error("%s: Failed to create simulation run: %s", LOG_PREFIX, error_msg) + logger.error( + "%s: Failed to generate first turn for item %s: %s", + LOG_PREFIX, + run_item_id, + error_msg, + ) return None def trigger_conversation( diff --git a/netra/simulation/constants.py b/netra/simulation/constants.py index 2057947..da30f7f 100644 --- a/netra/simulation/constants.py +++ b/netra/simulation/constants.py @@ -18,7 +18,8 @@ # --------------------------------------------------------------------------- # API endpoints # --------------------------------------------------------------------------- -URL_CREATE_RUN = "/evaluations/test_run/multi-turn" +URL_INITIALIZE_RUN = "/evaluations/test_run/multi-turn/initialize" +URL_FIRST_TURN = "/evaluations/run/{run_id}/item/{run_item_id}/first-turn" URL_AGENT_RESPONSE = "/evaluations/turn/agent-response" URL_RUN_ITEM_STATUS = "/evaluations/run/{run_id}/item/{run_item_id}/status" URL_RUN_STATUS = "/evaluations/run/{run_id}/status" diff --git a/tests/test_simulation.py b/tests/test_simulation.py index 0c93f24..8d165f7 100644 --- a/tests/test_simulation.py +++ b/tests/test_simulation.py @@ -379,58 +379,6 @@ def test_ensure_client_returns_none_when_not_initialized(self) -> None: client = SimulationHttpClient(self._make_config(endpoint="")) assert client._ensure_client() is None - def test_create_run_returns_none_without_client(self) -> None: - from netra.simulation.client import SimulationHttpClient - - client = SimulationHttpClient(self._make_config(endpoint="")) - assert client.create_run(name="test", dataset_id="ds-1") is None - - @patch("netra.simulation.client.httpx.Client") - def test_create_run_success(self, mock_client_cls: MagicMock) -> None: - from netra.simulation.client import SimulationHttpClient - - mock_response = MagicMock(spec=httpx.Response) - mock_response.json.return_value = { - "data": { - "id": "run-1", - "userMessages": [ - { - "testRunItemId": "item-1", - "userMessage": "hello", - "turnId": "turn-1", - "attachments": None, - } - ], - } - } - mock_response.raise_for_status = MagicMock() - - mock_instance = MagicMock() - mock_instance.post.return_value = mock_response - mock_client_cls.return_value = mock_instance - - client = SimulationHttpClient(self._make_config()) - result = client.create_run(name="test", dataset_id="ds-1") - - assert result is not None - assert result["run_id"] == "run-1" - assert len(result["simulation_items"]) == 1 - assert result["simulation_items"][0].message == "hello" - - @patch("netra.simulation.client.httpx.Client") - def test_create_run_returns_none_on_http_error(self, mock_client_cls: MagicMock) -> None: - from netra.simulation.client import SimulationHttpClient - - mock_instance = MagicMock() - mock_instance.post.side_effect = httpx.HTTPStatusError( - "Server Error", request=MagicMock(), response=MagicMock() - ) - mock_client_cls.return_value = mock_instance - - client = SimulationHttpClient(self._make_config()) - result = client.create_run(name="test", dataset_id="ds-1") - assert result is None - @patch("netra.simulation.client.httpx.Client") def test_trigger_conversation_stop(self, mock_client_cls: MagicMock) -> None: from netra.simulation.client import SimulationHttpClient @@ -608,10 +556,10 @@ def test_run_simulation_returns_none_on_invalid_inputs(self, mock_client_cls: Ma assert result is None @patch("netra.simulation.api.SimulationHttpClient") - def test_run_simulation_returns_none_when_create_run_fails(self, mock_client_cls: MagicMock) -> None: + def test_run_simulation_returns_none_when_initialize_run_fails(self, mock_client_cls: MagicMock) -> None: from netra.simulation.api import Simulation - mock_client_cls.return_value.create_run.return_value = None + mock_client_cls.return_value.initialize_run.return_value = None sim = Simulation(self._make_config()) result = sim.run_simulation(name="test", dataset_id="ds-1", task=SyncTask()) assert result is None @@ -633,12 +581,18 @@ def test_run_simulation_success(self, mock_client_cls: MagicMock, mock_span_wrap ) mock_client = MagicMock() - mock_client.create_run.return_value = { + mock_client.initialize_run.return_value = { "run_id": "run-1", - "simulation_items": [ - SimulationItem(run_item_id="item-1", dataset_item_id="ds-item-1", message="hello", turn_id="turn-1"), + "items": [ + {"test_run_item_id": "item-1", "dataset_item_id": "ds-item-1"}, ], } + mock_client.generate_first_turn.return_value = SimulationItem( + run_item_id="item-1", + dataset_item_id="ds-item-1", + message="hello", + turn_id="turn-1", + ) mock_client.trigger_conversation.return_value = stop_response mock_client.post_run_status.return_value = {"status": "completed"} mock_client_cls.return_value = mock_client @@ -665,12 +619,18 @@ def test_run_simulation_marks_failed_on_exception( mock_span_wrapper.return_value = mock_span mock_client = MagicMock() - mock_client.create_run.return_value = { + mock_client.initialize_run.return_value = { "run_id": "run-1", - "simulation_items": [ - SimulationItem(run_item_id="item-1", dataset_item_id="ds-item-1", message="hello", turn_id="turn-1"), + "items": [ + {"test_run_item_id": "item-1", "dataset_item_id": "ds-item-1"}, ], } + mock_client.generate_first_turn.return_value = SimulationItem( + run_item_id="item-1", + dataset_item_id="ds-item-1", + message="hello", + turn_id="turn-1", + ) mock_client.trigger_conversation.side_effect = RuntimeError("backend down") mock_client.post_run_status.return_value = {} mock_client_cls.return_value = mock_client @@ -700,12 +660,18 @@ def test_max_turns_guard(self, mock_client_cls: MagicMock, mock_span_wrapper: Ma ) mock_client = MagicMock() - mock_client.create_run.return_value = { + mock_client.initialize_run.return_value = { "run_id": "run-1", - "simulation_items": [ - SimulationItem(run_item_id="item-1", dataset_item_id="ds-item-1", message="hello", turn_id="turn-1"), + "items": [ + {"test_run_item_id": "item-1", "dataset_item_id": "ds-item-1"}, ], } + mock_client.generate_first_turn.return_value = SimulationItem( + run_item_id="item-1", + dataset_item_id="ds-item-1", + message="hello", + turn_id="turn-1", + ) mock_client.trigger_conversation.return_value = continue_response mock_client.post_run_status.return_value = {} mock_client_cls.return_value = mock_client @@ -740,12 +706,18 @@ def test_trigger_conversation_none_response(self, mock_client_cls: MagicMock, mo mock_span_wrapper.return_value = mock_span mock_client = MagicMock() - mock_client.create_run.return_value = { + mock_client.initialize_run.return_value = { "run_id": "run-1", - "simulation_items": [ - SimulationItem(run_item_id="item-1", dataset_item_id="ds-item-1", message="hello", turn_id="turn-1"), + "items": [ + {"test_run_item_id": "item-1", "dataset_item_id": "ds-item-1"}, ], } + mock_client.generate_first_turn.return_value = SimulationItem( + run_item_id="item-1", + dataset_item_id="ds-item-1", + message="hello", + turn_id="turn-1", + ) mock_client.trigger_conversation.return_value = None mock_client.post_run_status.return_value = {} mock_client_cls.return_value = mock_client