nvidia-nat 1.3.0.dev2__py3-none-any.whl → 1.3.0rc1__py3-none-any.whl

nat/eval/evaluate.py

@@ -42,7 +42,7 @@ from nat.runtime.session import SessionManager
 logger = logging.getLogger(__name__)
 
 
-class EvaluationRun:  # pylint: disable=too-many-public-methods
+class EvaluationRun:
     """
     Instantiated for each evaluation run and used to store data for that single run.
 
@@ -63,7 +63,16 @@ class EvaluationRun:  # pylint: disable=too-many-public-methods
 
         # Helpers
         self.intermediate_step_adapter: IntermediateStepAdapter = IntermediateStepAdapter()
-        self.weave_eval: WeaveEvaluationIntegration = WeaveEvaluationIntegration()
+
+        # Create evaluation trace context
+        try:
+            from nat.eval.utils.eval_trace_ctx import WeaveEvalTraceContext
+            self.eval_trace_context = WeaveEvalTraceContext()
+        except Exception:
+            from nat.eval.utils.eval_trace_ctx import EvalTraceContext
+            self.eval_trace_context = EvalTraceContext()
+
+        self.weave_eval: WeaveEvaluationIntegration = WeaveEvaluationIntegration(self.eval_trace_context)
         # Metadata
         self.eval_input: EvalInput | None = None
         self.workflow_interrupted: bool = False
@@ -159,17 +168,17 @@ class EvaluationRun:  # pylint: disable=too-many-public-methods
                 intermediate_future = None
 
                 try:
-
                     # Start usage stats and intermediate steps collection in parallel
                     intermediate_future = pull_intermediate()
                     runner_result = runner.result()
                     base_output = await runner_result
                     intermediate_steps = await intermediate_future
                 except NotImplementedError as e:
+                    logger.error("Failed to run the workflow: %s", e)
                     # raise original error
-                    raise e
+                    raise
                 except Exception as e:
-                    logger.exception("Failed to run the workflow: %s", e, exc_info=True)
+                    logger.exception("Failed to run the workflow: %s", e)
                     # stop processing if a workflow error occurs
                     self.workflow_interrupted = True
 
@@ -308,9 +317,9 @@ class EvaluationRun:  # pylint: disable=too-many-public-methods
                 logger.info("Deleting old job directory: %s", dir_to_delete)
                 shutil.rmtree(dir_to_delete)
             except Exception as e:
-                logger.exception("Failed to delete old job directory: %s: %s", dir_to_delete, e, exc_info=True)
+                logger.exception("Failed to delete old job directory: %s: %s", dir_to_delete, e)
 
-    def write_output(self, dataset_handler: DatasetHandler, profiler_results: ProfilerResults):  # pylint: disable=unused-argument  # noqa: E501
+    def write_output(self, dataset_handler: DatasetHandler, profiler_results: ProfilerResults):
         workflow_output_file = self.eval_config.general.output_dir / "workflow_output.json"
         workflow_output_file.parent.mkdir(parents=True, exist_ok=True)
 
@@ -358,7 +367,7 @@ class EvaluationRun:  # pylint: disable=too-many-public-methods
 
             await self.weave_eval.alog_score(eval_output, evaluator_name)
         except Exception as e:
-            logger.exception("An error occurred while running evaluator %s: %s", evaluator_name, e, exc_info=True)
+            logger.exception("An error occurred while running evaluator %s: %s", evaluator_name, e)
 
     async def run_evaluators(self, evaluators: dict[str, Any]):
         """Run all configured evaluators asynchronously."""
@@ -371,7 +380,7 @@ class EvaluationRun:  # pylint: disable=too-many-public-methods
         try:
             await asyncio.gather(*tasks)
         except Exception as e:
-            logger.exception("An error occurred while running evaluators: %s", e, exc_info=True)
+            logger.error("An error occurred while running evaluators: %s", e)
             raise
         finally:
             # Finish prediction loggers in Weave
@@ -401,6 +410,33 @@ class EvaluationRun:  # pylint: disable=too-many-public-methods
 
         return workflow_type
 
+    async def wait_for_all_export_tasks_local(self, session_manager: SessionManager, timeout: float) -> None:
+        """Wait for all trace export tasks to complete for local workflows.
+
+        This only works for local workflows where we have direct access to the
+        SessionManager and its underlying workflow with exporter manager.
+        """
+        try:
+            workflow = session_manager.workflow
+            all_exporters = await workflow.get_all_exporters()
+            if not all_exporters:
+                logger.debug("No exporters to wait for")
+                return
+
+            logger.info("Waiting for export tasks from %d local exporters (timeout: %ds)", len(all_exporters), timeout)
+
+            for name, exporter in all_exporters.items():
+                try:
+                    await exporter.wait_for_tasks(timeout=timeout)
+                    logger.info("Export tasks completed for exporter: %s", name)
+                except Exception as e:
+                    logger.warning("Error waiting for export tasks from %s: %s", name, e)
+
+            logger.info("All local export task waiting completed")
+
+        except Exception as e:
+            logger.warning("Failed to wait for local export tasks: %s", e)
+
     async def run_and_evaluate(self,
                                session_manager: SessionManager | None = None,
                                job_id: str | None = None) -> EvaluationRunOutput:
@@ -413,10 +449,14 @@ class EvaluationRun:  # pylint: disable=too-many-public-methods
         from nat.runtime.loader import load_config
 
         # Load and override the config
-        if self.config.override:
+        config = None
+        if isinstance(self.config.config_file, BaseModel):
+            config = self.config.config_file
+        elif self.config.override:
             config = self.apply_overrides()
         else:
             config = load_config(self.config.config_file)
+
         self.eval_config = config.eval
         workflow_alias = self._get_workflow_alias(config.workflow.type)
         logger.debug("Loaded %s evaluation configuration: %s", workflow_alias, self.eval_config)
@@ -442,44 +482,59 @@ class EvaluationRun:  # pylint: disable=too-many-public-methods
         dataset_config = self.eval_config.general.dataset  # Currently only one dataset is supported
         if not dataset_config:
             logger.info("No dataset found, nothing to evaluate")
-            return EvaluationRunOutput(
-                workflow_output_file=self.workflow_output_file,
-                evaluator_output_files=self.evaluator_output_files,
-                workflow_interrupted=self.workflow_interrupted,
-            )
-
+            return EvaluationRunOutput(workflow_output_file=self.workflow_output_file,
+                                       evaluator_output_files=self.evaluator_output_files,
+                                       workflow_interrupted=self.workflow_interrupted,
+                                       eval_input=EvalInput(eval_input_items=[]),
+                                       evaluation_results=[],
+                                       usage_stats=UsageStats(),
+                                       profiler_results=ProfilerResults())
+
+        custom_pre_eval_process_function = self.eval_config.general.output.custom_pre_eval_process_function \
+            if self.eval_config.general.output else None
         dataset_handler = DatasetHandler(dataset_config=dataset_config,
                                          reps=self.config.reps,
                                          concurrency=self.eval_config.general.max_concurrency,
                                          num_passes=self.config.num_passes,
-                                         adjust_dataset_size=self.config.adjust_dataset_size)
+                                         adjust_dataset_size=self.config.adjust_dataset_size,
+                                         custom_pre_eval_process_function=custom_pre_eval_process_function)
         self.eval_input = dataset_handler.get_eval_input_from_dataset(self.config.dataset)
         if not self.eval_input.eval_input_items:
             logger.info("Dataset is empty. Nothing to evaluate.")
-            return EvaluationRunOutput(
-                workflow_output_file=self.workflow_output_file,
-                evaluator_output_files=self.evaluator_output_files,
-                workflow_interrupted=self.workflow_interrupted,
-            )
+            return EvaluationRunOutput(workflow_output_file=self.workflow_output_file,
+                                       evaluator_output_files=self.evaluator_output_files,
+                                       workflow_interrupted=self.workflow_interrupted,
+                                       eval_input=self.eval_input,
+                                       evaluation_results=self.evaluation_results,
+                                       usage_stats=self.usage_stats,
+                                       profiler_results=ProfilerResults())
 
         # Run workflow and evaluate
         async with WorkflowEvalBuilder.from_config(config=config) as eval_workflow:
             # Initialize Weave integration
             self.weave_eval.initialize_logger(workflow_alias, self.eval_input, config)
 
-            # Run workflow
-            if self.config.endpoint:
-                await self.run_workflow_remote()
-            else:
-                if not self.config.skip_workflow:
+            with self.eval_trace_context.evaluation_context():
+                # Run workflow
+                if self.config.endpoint:
+                    await self.run_workflow_remote()
+                elif not self.config.skip_workflow:
                     if session_manager is None:
-                        session_manager = SessionManager(eval_workflow.build(),
+                        workflow = await eval_workflow.build()
+                        session_manager = SessionManager(workflow,
                                                          max_concurrency=self.eval_config.general.max_concurrency)
                     await self.run_workflow_local(session_manager)
 
-            # Evaluate
-            evaluators = {name: eval_workflow.get_evaluator(name) for name in self.eval_config.evaluators}
-            await self.run_evaluators(evaluators)
+                # Pre-evaluation process the workflow output
+                self.eval_input = dataset_handler.pre_eval_process_eval_input(self.eval_input)
+
+                # Evaluate
+                evaluators = {name: eval_workflow.get_evaluator(name) for name in self.eval_config.evaluators}
+                await self.run_evaluators(evaluators)
+
+                # Wait for all trace export tasks to complete (local workflows only)
+                if session_manager and not self.config.endpoint:
+                    await self.wait_for_all_export_tasks_local(session_manager, timeout=self.config.export_timeout)
 
         # Profile the workflow
         profiler_results = await self.profile_workflow()

nat/eval/evaluator/base_evaluator.py

@@ -71,7 +71,7 @@ class BaseEvaluator(ABC):
             TqdmPositionRegistry.release(tqdm_position)
 
         # Compute average if possible
-        numeric_scores = [item.score for item in output_items if isinstance(item.score, (int, float))]
+        numeric_scores = [item.score for item in output_items if isinstance(item.score, int | float)]
         avg_score = round(sum(numeric_scores) / len(numeric_scores), 2) if numeric_scores else None
 
         return EvalOutput(average_score=avg_score, eval_output_items=output_items)

nat/eval/evaluator/evaluator_model.py

@@ -29,6 +29,19 @@ class EvalInputItem(BaseModel):
     trajectory: list[IntermediateStep] = []  # populated by the workflow
     full_dataset_entry: typing.Any
 
+    def copy_with_updates(self, **updates) -> "EvalInputItem":
+        """
+        Copy EvalInputItem with optional field updates.
+        """
+        # Get all current fields
+        item_data = self.model_dump()
+
+        # Apply any updates
+        item_data.update(updates)
+
+        # Create new item with all fields
+        return EvalInputItem(**item_data)
+
 
 class EvalInput(BaseModel):
     eval_input_items: list[EvalInputItem]

nat/eval/intermediate_step_adapter.py

@@ -40,7 +40,7 @@ class IntermediateStepAdapter:
             try:
                 validated_steps.append(IntermediateStep.model_validate(step_data))
             except Exception as e:
-                logger.exception("Validation failed for step: %r, Error: %s", step_data, e, exc_info=True)
+                logger.exception("Validation failed for step: %r, Error: %s", step_data, e)
         return validated_steps
 
     def serialize_intermediate_steps(self, intermediate_steps: list[IntermediateStep]) -> list[dict]:

nat/eval/rag_evaluator/evaluate.py

@@ -102,7 +102,7 @@ class RAGEvaluator:
         """Converts the ragas EvaluationResult to nat EvalOutput"""
 
         if not results_dataset:
-            logger.error("Ragas evaluation failed with no results")
+            logger.error("Ragas evaluation failed with no results", exc_info=True)
             return EvalOutput(average_score=0.0, eval_output_items=[])
 
         scores: list[dict[str, float]] = results_dataset.scores
@@ -169,7 +169,7 @@ class RAGEvaluator:
                                              _pbar=pbar)
         except Exception as e:
             # On exception we still continue with other evaluators. Log and return an avg_score of 0.0
-            logger.exception("Error evaluating ragas metric, Error: %s", e, exc_info=True)
+            logger.exception("Error evaluating ragas metric, Error: %s", e)
             results_dataset = None
         finally:
             pbar.close()

nat/eval/rag_evaluator/register.py

@@ -73,7 +73,7 @@ class RagasEvaluatorConfig(EvaluatorBaseConfig, name="ragas"):
         if isinstance(self.metric, str):
             return self.metric
         if isinstance(self.metric, dict) and self.metric:
-            return next(iter(self.metric.keys()))  # pylint: disable=no-member
+            return next(iter(self.metric.keys()))
         return ""
 
     @property
@@ -82,7 +82,7 @@ class RagasEvaluatorConfig(EvaluatorBaseConfig, name="ragas"):
         if isinstance(self.metric, str):
             return RagasMetricConfig()  # Default config when only a metric name is given
         if isinstance(self.metric, dict) and self.metric:
-            return next(iter(self.metric.values()))  # pylint: disable=no-member
+            return next(iter(self.metric.values()))
         return RagasMetricConfig()  # Default config when an invalid type is provided
 
 
@@ -104,7 +104,7 @@ async def register_ragas_evaluator(config: RagasEvaluatorConfig, builder: EvalBu
             raise ValueError(message) from e
         except AttributeError as e:
             message = f"Ragas metric {metric_name} not found {e}."
-            logger.error(message)
+            logger.exception(message)
             return None
 
     async def evaluate_fn(eval_input: EvalInput) -> EvalOutput:

nat/eval/register.py

@@ -14,10 +14,13 @@
 # limitations under the License.
 
 # flake8: noqa
-# pylint: disable=unused-import
 
 # Import evaluators which need to be automatically registered here
 from .rag_evaluator.register import register_ragas_evaluator
+from .runtime_evaluator.register import register_avg_llm_latency_evaluator
+from .runtime_evaluator.register import register_avg_num_llm_calls_evaluator
+from .runtime_evaluator.register import register_avg_tokens_per_llm_end_evaluator
+from .runtime_evaluator.register import register_avg_workflow_runtime_evaluator
 from .swe_bench_evaluator.register import register_swe_bench_evaluator
 from .trajectory_evaluator.register import register_trajectory_evaluator
 from .tunable_rag_evaluator.register import register_tunable_rag_evaluator

nat/eval/remote_workflow.py

@@ -74,7 +74,7 @@ class EvaluationRemoteWorkflowHandler:
                             if chunk_data.get("value"):
                                 final_response = chunk_data.get("value")
                         except json.JSONDecodeError as e:
-                            logger.error("Failed to parse generate response chunk: %s", e)
+                            logger.exception("Failed to parse generate response chunk: %s", e)
                             continue
                     elif line.startswith(INTERMEDIATE_DATA_PREFIX):
                         # This is an intermediate step
@@ -90,12 +90,12 @@ class EvaluationRemoteWorkflowHandler:
                                                                  payload=payload)
                             intermediate_steps.append(intermediate_step)
                         except (json.JSONDecodeError, ValidationError) as e:
-                            logger.error("Failed to parse intermediate step: %s", e)
+                            logger.exception("Failed to parse intermediate step: %s", e)
                             continue
 
         except aiohttp.ClientError as e:
             # Handle connection or HTTP-related errors
-            logger.error("Request failed for question %s: %s", question, e)
+            logger.exception("Request failed for question %s: %s", question, e)
             item.output_obj = None
             item.trajectory = []
             return

nat/eval/runtime_evaluator/__init__.py

@@ -0,0 +1,14 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

nat/eval/runtime_evaluator/evaluate.py

@@ -0,0 +1,123 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from collections import defaultdict
+from dataclasses import dataclass
+
+from nat.data_models.intermediate_step import IntermediateStepType
+from nat.eval.evaluator.base_evaluator import BaseEvaluator
+from nat.eval.evaluator.evaluator_model import EvalInputItem
+from nat.eval.evaluator.evaluator_model import EvalOutputItem
+from nat.profiler.intermediate_property_adapter import IntermediatePropertyAdaptor
+
+
+@dataclass
+class _CallTiming:
+    start_ts: float | None = None
+    end_ts: float | None = None
+
+    @property
+    def latency(self) -> float | None:
+        if self.start_ts is None or self.end_ts is None:
+            return None
+        return max(0.0, self.end_ts - self.start_ts)
+
+
+class AverageLLMLatencyEvaluator(BaseEvaluator):
+    """
+    Mean difference between connected LLM_START and LLM_END events (same UUID).
+    The score is the average latency in seconds for the item. Reasoning contains per-call latencies.
+    """
+
+    def __init__(self, max_concurrency: int = 8):
+        super().__init__(max_concurrency=max_concurrency, tqdm_desc="Evaluating Avg LLM Latency")
+
+    async def evaluate_item(self, item: EvalInputItem) -> EvalOutputItem:  # noqa: D401
+        calls: dict[str, _CallTiming] = defaultdict(_CallTiming)
+
+        for step in (IntermediatePropertyAdaptor.from_intermediate_step(s) for s in item.trajectory):
+            if step.event_type == IntermediateStepType.LLM_START:
+                calls[step.UUID].start_ts = step.event_timestamp
+            elif step.event_type == IntermediateStepType.LLM_END:
+                calls[step.UUID].end_ts = step.event_timestamp
+
+        latencies = [ct.latency for ct in calls.values() if ct.latency is not None]
+        avg_latency = sum(latencies) / len(latencies) if latencies else 0.0
+
+        reasoning = {
+            "num_llm_calls": len(latencies),
+            "latencies": latencies,
+        }
+        return EvalOutputItem(id=item.id, score=round(avg_latency, 4), reasoning=reasoning)
+
+
+class AverageWorkflowRuntimeEvaluator(BaseEvaluator):
+    """
+    Average workflow runtime per item: max(event_timestamp) - min(event_timestamp) across the trajectory.
+    The score is the runtime in seconds for the item.
+    """
+
+    def __init__(self, max_concurrency: int = 8):
+        super().__init__(max_concurrency=max_concurrency, tqdm_desc="Evaluating Avg Workflow Runtime")
+
+    async def evaluate_item(self, item: EvalInputItem) -> EvalOutputItem:  # noqa: D401
+        if not item.trajectory:
+            return EvalOutputItem(id=item.id, score=0.0, reasoning={"note": "no steps"})
+
+        timestamps = [s.event_timestamp for s in item.trajectory]
+        runtime = max(timestamps) - min(timestamps)
+        return EvalOutputItem(id=item.id, score=round(max(0.0, runtime), 4), reasoning={"steps": len(timestamps)})
+
+
+class AverageNumberOfLLMCallsEvaluator(BaseEvaluator):
+    """
+    Average number of LLM calls per item. The score is the count for the item.
+    """
+
+    def __init__(self, max_concurrency: int = 8):
+        super().__init__(max_concurrency=max_concurrency, tqdm_desc="Evaluating Avg # LLM Calls")
+
+    async def evaluate_item(self, item: EvalInputItem) -> EvalOutputItem:  # noqa: D401
+        num_calls = sum(1 for s in item.trajectory if s.event_type == IntermediateStepType.LLM_END)
+        return EvalOutputItem(id=item.id, score=float(num_calls), reasoning={"num_llm_end": num_calls})
+
+
+class AverageTokensPerLLMEndEvaluator(BaseEvaluator):
+    """
+    Average total tokens per LLM_END event: sum of prompt and completion tokens if available.
+    The score is the average tokens per LLM_END for the item (0 if none).
+    """
+
+    def __init__(self, max_concurrency: int = 8):
+        super().__init__(max_concurrency=max_concurrency, tqdm_desc="Evaluating Avg Tokens/LLM_END")
+
+    async def evaluate_item(self, item: EvalInputItem) -> EvalOutputItem:  # noqa: D401
+        totals: list[int] = []
+        for step in (IntermediatePropertyAdaptor.from_intermediate_step(s) for s in item.trajectory):
+            if step.event_type == IntermediateStepType.LLM_END:
+                total_tokens = step.token_usage.total_tokens
+                # If framework doesn't set total, compute from prompt+completion
+                if total_tokens == 0:
+                    total_tokens = step.token_usage.prompt_tokens + step.token_usage.completion_tokens
+                totals.append(total_tokens)
+
+        avg_tokens = (sum(totals) / len(totals)) if totals else 0.0
+        reasoning = {
+            "num_llm_end": len(totals),
+            "totals": totals,
+        }
+        return EvalOutputItem(id=item.id, score=round(avg_tokens, 2), reasoning=reasoning)

nat/eval/runtime_evaluator/register.py

@@ -0,0 +1,100 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from pydantic import Field
+
+from nat.builder.builder import EvalBuilder
+from nat.builder.evaluator import EvaluatorInfo
+from nat.cli.register_workflow import register_evaluator
+from nat.data_models.evaluator import EvaluatorBaseConfig
+from nat.eval.evaluator.evaluator_model import EvalInput
+from nat.eval.evaluator.evaluator_model import EvalOutput
+
+
+class AverageLLMLatencyConfig(EvaluatorBaseConfig, name="avg_llm_latency"):
+    """Mean difference between connected LLM_START and LLM_END events (same UUID)."""
+
+    max_concurrency: int = Field(default=8, description="Max concurrency for evaluation.")
+
+
+class AverageWorkflowRuntimeConfig(EvaluatorBaseConfig, name="avg_workflow_runtime"):
+    """Average workflow runtime per item (max timestamp - min timestamp)."""
+
+    max_concurrency: int = Field(default=8, description="Max concurrency for evaluation.")
+
+
+class AverageNumberOfLLMCallsConfig(EvaluatorBaseConfig, name="avg_num_llm_calls"):
+    """Average number of LLM calls per item (count of LLM_END)."""
+
+    max_concurrency: int = Field(default=8, description="Max concurrency for evaluation.")
+
+
+class AverageTokensPerLLMEndConfig(EvaluatorBaseConfig, name="avg_tokens_per_llm_end"):
+    """Average total tokens per LLM_END event (prompt + completion if available)."""
+
+    max_concurrency: int = Field(default=8, description="Max concurrency for evaluation.")
+
+
+@register_evaluator(config_type=AverageLLMLatencyConfig)
+async def register_avg_llm_latency_evaluator(config: AverageLLMLatencyConfig, builder: EvalBuilder):
+    from .evaluate import AverageLLMLatencyEvaluator
+
+    evaluator = AverageLLMLatencyEvaluator(max_concurrency=config.max_concurrency or builder.get_max_concurrency())
+
+    async def evaluate_fn(eval_input: EvalInput) -> EvalOutput:
+        return await evaluator.evaluate(eval_input)
+
+    yield EvaluatorInfo(config=config,
+                        evaluate_fn=evaluate_fn,
+                        description="Average LLM latency (s) from LLM_START to LLM_END")
+
+
+@register_evaluator(config_type=AverageWorkflowRuntimeConfig)
+async def register_avg_workflow_runtime_evaluator(config: AverageWorkflowRuntimeConfig, builder: EvalBuilder):
+    from .evaluate import AverageWorkflowRuntimeEvaluator
+
+    evaluator = AverageWorkflowRuntimeEvaluator(max_concurrency=config.max_concurrency or builder.get_max_concurrency())
+
+    async def evaluate_fn(eval_input: EvalInput) -> EvalOutput:
+        return await evaluator.evaluate(eval_input)
+
+    yield EvaluatorInfo(config=config, evaluate_fn=evaluate_fn, description="Average workflow runtime (s)")
+
+
+@register_evaluator(config_type=AverageNumberOfLLMCallsConfig)
+async def register_avg_num_llm_calls_evaluator(config: AverageNumberOfLLMCallsConfig, builder: EvalBuilder):
+    from .evaluate import AverageNumberOfLLMCallsEvaluator
+
+    evaluator = AverageNumberOfLLMCallsEvaluator(
+        max_concurrency=config.max_concurrency or builder.get_max_concurrency())
+
+    async def evaluate_fn(eval_input: EvalInput) -> EvalOutput:
+        return await evaluator.evaluate(eval_input)
+
+    yield EvaluatorInfo(config=config, evaluate_fn=evaluate_fn, description="Average number of LLM calls")
+
+
+@register_evaluator(config_type=AverageTokensPerLLMEndConfig)
+async def register_avg_tokens_per_llm_end_evaluator(config: AverageTokensPerLLMEndConfig, builder: EvalBuilder):
+    from .evaluate import AverageTokensPerLLMEndEvaluator
+
+    evaluator = AverageTokensPerLLMEndEvaluator(max_concurrency=config.max_concurrency or builder.get_max_concurrency())
+
+    async def evaluate_fn(eval_input: EvalInput) -> EvalOutput:
+        return await evaluator.evaluate(eval_input)
+
+    yield EvaluatorInfo(config=config,
+                        evaluate_fn=evaluate_fn,
+                        description="Average total tokens per LLM_END (prompt + completion)")

nat/eval/swe_bench_evaluator/evaluate.py

@@ -69,13 +69,13 @@ class SweBenchEvaluator:
         try:
             shutil.move(swe_bench_report_file, report_dir)
         except Exception as e:
-            logger.exception("Error moving report file: %s", e, exc_info=True)
+            logger.exception("Error moving report file: %s", e)
 
         try:
             dest_logs_dir = os.path.join(report_dir, 'logs')
             shutil.move(logs_dir, dest_logs_dir)
         except Exception as e:
-            logger.exception("Error moving logs directory: %s", e, exc_info=True)
+            logger.exception("Error moving logs directory: %s", e)
 
     def is_repo_supported(self, repo: str, version: str) -> bool:
         """Check if the repo is supported by swebench"""
@@ -106,7 +106,7 @@ class SweBenchEvaluator:
                     self._model_name_or_path = swebench_output.model_name_or_path
 
             except Exception as e:
-                logger.exception("Failed to parse EvalInputItem %s: %s", item.id, e, exc_info=True)
+                logger.exception("Failed to parse EvalInputItem %s: %s", item.id, e)
 
         # Filter out repos/version not supported by SWEBench
         supported_inputs = [
@@ -114,7 +114,7 @@ class SweBenchEvaluator:
         ]
 
         if not supported_inputs:
-            logger.error("No supported instances; nothing to evaluate")
+            logger.exception("No supported instances; nothing to evaluate")
             return None, None
 
         if len(supported_inputs) < len(swebench_inputs):
@@ -135,7 +135,7 @@ class SweBenchEvaluator:
         filtered_outputs = [output for output in swebench_outputs if output.instance_id in valid_instance_ids]
 
         if not filtered_outputs:
-            logger.error("No supported outputs; nothing to evaluate")
+            logger.error("No supported outputs; nothing to evaluate", exc_info=True)
             return None, None
 
         # Write SWEBenchOutput to file
@@ -204,7 +204,7 @@ class SweBenchEvaluator:
         # if report file is not present, return empty EvalOutput
         avg_score = 0.0
         if report_file.exists():
-            with open(report_file, "r", encoding="utf-8") as f:
+            with open(report_file, encoding="utf-8") as f:
                 report = json.load(f)
                 resolved_instances = report.get("resolved_instances", 0)
                 total_instances = report.get("total_instances", 0)

nat/eval/trajectory_evaluator/evaluate.py

@@ -65,7 +65,7 @@ class TrajectoryEvaluator(BaseEvaluator):
                 prediction=generated_answer,
             )
         except Exception as e:
-            logger.exception("Error evaluating trajectory for question: %s, Error: %s", question, e, exc_info=True)
+            logger.exception("Error evaluating trajectory for question: %s, Error: %s", question, e)
             return EvalOutputItem(id=item.id, score=0.0, reasoning=f"Error evaluating trajectory: {e}")
 
         reasoning = {

nat/eval/trajectory_evaluator/register.py

@@ -33,7 +33,7 @@ async def register_trajectory_evaluator(config: TrajectoryEvaluatorConfig, build
 
     from .evaluate import TrajectoryEvaluator
     llm = await builder.get_llm(config.llm_name, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
-    tools = builder.get_all_tools(wrapper_type=LLMFrameworkEnum.LANGCHAIN)
+    tools = await builder.get_all_tools(wrapper_type=LLMFrameworkEnum.LANGCHAIN)
 
     _evaluator = TrajectoryEvaluator(llm, tools, builder.get_max_concurrency())