nvidia-nat 1.2.0__py3-none-any.whl

nat/eval/runners/multi_eval_runner.py

@@ -0,0 +1,54 @@
+# 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.
+
+import copy
+import typing
+
+from nat.eval.config import EvaluationRunConfig
+from nat.eval.config import EvaluationRunOutput
+from nat.eval.evaluate import EvaluationRun
+from nat.eval.runners.config import MultiEvaluationRunConfig
+
+
+class MultiEvaluationRunner:
+    """
+    Run a multi-evaluation run.
+    """
+
+    def __init__(self, config: MultiEvaluationRunConfig):
+        """
+        Initialize a multi-evaluation run.
+        """
+        self.config = config
+        self.evaluation_run_outputs: dict[typing.Any, EvaluationRunOutput] = {}
+
+    async def run_all(self):
+        """
+        Run all evaluations defined by the overrides.
+        """
+        for id, config in self.config.configs.items():
+            output = await self.run_single_evaluation(id, config)
+            self.evaluation_run_outputs[id] = output
+
+        return self.evaluation_run_outputs
+
+    async def run_single_evaluation(self, id: typing.Any, config: EvaluationRunConfig) -> EvaluationRunOutput:
+        """
+        Run a single evaluation and return the output.
+        """
+        # copy the config in case the caller is using the same config for multiple evaluations
+        config_copy = copy.deepcopy(config)
+        evaluation_run = EvaluationRun(config_copy)
+        return await evaluation_run.run_and_evaluate()

nat/eval/runtime_event_subscriber.py

@@ -0,0 +1,52 @@
+# 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.
+
+import asyncio
+import logging
+
+from nat.builder.context import Context
+from nat.data_models.intermediate_step import IntermediateStep
+
+logger = logging.getLogger(__name__)
+
+
+def pull_intermediate() -> asyncio.Future[list[dict]]:
+    """
+    Subscribes to the runner's event stream using callbacks.
+    Intermediate steps are collected and, when complete, the future is set
+    with the list of dumped intermediate steps.
+    """
+    future = asyncio.Future()
+    intermediate_steps = []  # We'll store the dumped steps here.
+    context = Context.get()
+
+    def on_next_cb(item: IntermediateStep):
+        # Append each new intermediate step (dumped to dict) to the list.
+        intermediate_steps.append(item.model_dump())
+
+    def on_error_cb(exc: Exception):
+        logger.error("Hit on_error: %s", exc)
+        if not future.done():
+            future.set_exception(exc)
+
+    def on_complete_cb():
+        logger.debug("Completed reading intermediate steps")
+        if not future.done():
+            future.set_result(intermediate_steps)
+
+    # Subscribe with our callbacks.
+    context.intermediate_step_manager.subscribe(on_next=on_next_cb, on_error=on_error_cb, on_complete=on_complete_cb)
+
+    return future

nat/eval/swe_bench_evaluator/evaluate.py

@@ -0,0 +1,215 @@
+# 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.
+
+import json
+import logging
+import os
+import shutil
+from pathlib import Path
+
+from nat.data_models.swe_bench_model import SWEBenchInput
+from nat.data_models.swe_bench_model import SWEBenchOutput
+from nat.eval.evaluator.evaluator_model import EvalInput
+from nat.eval.evaluator.evaluator_model import EvalOutput
+
+try:
+    import swebench.harness.run_evaluation as swebench_eval
+    from swebench.harness.constants import MAP_REPO_VERSION_TO_SPECS
+except ImportError as exc:
+    raise ImportError("Please install swebench to use this evaluator") from exc
+
+logger = logging.getLogger(__name__)
+
+
+class SweBenchEvaluator:
+
+    def __init__(self, run_id: str, max_workers: int, output_dir: Path):
+
+        self.run_id = run_id
+        self.max_workers = max_workers
+        self.output_dir = output_dir
+
+        # metadata
+        self._unsupported_repos = []
+        self._swe_bench_inputs = []
+        self._swe_bench_outputs = []
+        self._model_name_or_path = "no_llm"
+
+    def get_model_name_from_output(self, workflow_output: list[dict]) -> str | None:
+        """Fetch the `model_name_or_path` from the first entry in the list."""
+        return workflow_output[0].get("model_name_or_path") if workflow_output else None
+
+    @staticmethod
+    def empty_report_dir(report_dir: Path):
+        """Remove the current contents of the report directory."""
+        os.makedirs(report_dir, exist_ok=True)
+
+        # Iterate through all files in the directory and remove them
+        for item in report_dir.iterdir():
+            if item.is_file():  # Remove files only
+                item.unlink()
+            elif item.is_dir():  # Remove subdirectories and their contents
+                shutil.rmtree(item)
+
+    @staticmethod
+    def move_report_and_logs(swe_bench_report_file: str, logs_dir: str, report_dir: Path):
+        """ Temorary function to move the report and logs to the output directory"""
+        try:
+            shutil.move(swe_bench_report_file, report_dir)
+        except Exception as e:
+            logger.exception("Error moving report file: %s", e, exc_info=True)
+
+        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)
+
+    def is_repo_supported(self, repo: str, version: str) -> bool:
+        """Check if the repo is supported by swebench"""
+
+        try:
+            _ = MAP_REPO_VERSION_TO_SPECS[repo][str(version)]
+        except KeyError:
+            self._unsupported_repos.append({repo, version})
+            return False
+        return True
+
+    def process_eval_input(self, eval_input: EvalInput) -> tuple[Path, Path]:
+        """Converts EvalInput into lists of SWEBenchInput and SWEBenchOutput models and applies filtering."""
+        # Convert input_obj and output_obj JSON strings to SWEBenchInput and SWEBenchOutput models
+        swebench_inputs = []
+        swebench_outputs = []
+
+        for item in eval_input.eval_input_items:
+            try:
+                swebench_input = SWEBenchInput.model_validate_json(item.input_obj)  # Convert input JSON to model
+                swebench_input.version = str(swebench_input.version)  # Convert version to string
+                swebench_inputs.append(swebench_input)
+
+                if item.output_obj:  # Convert output JSON to model if available
+                    swebench_output = SWEBenchOutput.model_validate_json(item.output_obj)
+                    swebench_outputs.append(swebench_output)
+                    # this is bit of a hack to match the swe_bench harness
+                    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)
+
+        # Filter out repos/version not supported by SWEBench
+        supported_inputs = [
+            swebench for swebench in swebench_inputs if self.is_repo_supported(swebench.repo, swebench.version)
+        ]
+
+        if not supported_inputs:
+            logger.error("No supported instances; nothing to evaluate")
+            return None, None
+
+        if len(supported_inputs) < len(swebench_inputs):
+            logger.warning("The following repos are not supported by SWEBench and were skipped:\n %s",
+                           {s.repo
+                            for s in swebench_inputs if s not in supported_inputs})
+
+        # Write SWEBenchInput to file
+        workflow_input_file = self.output_dir / "nat_workflow_input.json"
+        workflow_input_file.parent.mkdir(parents=True, exist_ok=True)
+        Path(workflow_input_file).write_text(json.dumps([swebench.model_dump() for swebench in supported_inputs],
+                                                        indent=2),
+                                             encoding="utf-8")
+        logger.info("Workflow input written to %s", workflow_input_file)
+
+        # Filter SWEBenchOutput to include only instance_ids present in SWEBenchInput
+        valid_instance_ids = {swebench.instance_id for swebench in supported_inputs}
+        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")
+            return None, None
+
+        # Write SWEBenchOutput to file
+        workflow_output_file = self.output_dir / "nat_workflow_output.json"
+        Path(workflow_output_file).write_text(json.dumps([output.model_dump() for output in filtered_outputs],
+                                                         indent=2),
+                                              encoding="utf-8")
+        logger.info("Workflow output written to %s", workflow_output_file)
+
+        self._swe_bench_inputs = supported_inputs
+        self._swe_bench_outputs = filtered_outputs
+        return workflow_input_file, workflow_output_file
+
+    def build_eval_output(self):
+        """Builds the EvalOutput object from the SWEBenchOutput models and the average score."""
+        # WIP: Build a score based on eval run logs
+        for swebench_output in self._swe_bench_outputs:
+            yield {"id": swebench_output.instance_id, "score": "-", "reasoning": "-"}
+
+    @staticmethod
+    def compute_score(success_cnt: int, total_cnt: int) -> float:
+        if total_cnt == 0:
+            return 0.0
+        score = success_cnt / total_cnt
+        return min(max(score, 0.0), 1.0)
+
+    async def evaluate(self, eval_input: EvalInput) -> EvalOutput:
+        '''Run the swebench evaluation and store the report in the output directory'''
+
+        # Process the EvalInput
+        workflow_input_file, workflow_output_file = self.process_eval_input(eval_input)
+        if not workflow_input_file or not workflow_output_file:
+            # nothing to evaluate
+            return EvalOutput(average_score=0.0, eval_output_items=[])
+
+        report_dir = self.output_dir / "swe_bench_reports"
+        self.empty_report_dir(report_dir)
+
+        logger.info("Starting swe_bench run %s", self.run_id)
+        swebench_eval.main(dataset_name=str(workflow_input_file),
+                           split="dev",
+                           instance_ids=[],
+                           predictions_path=str(workflow_output_file),
+                           max_workers=self.max_workers,
+                           force_rebuild=False,
+                           cache_level="env",
+                           clean=False,
+                           open_file_limit=4096,
+                           run_id=self.run_id,
+                           timeout=1800,
+                           namespace=None,
+                           rewrite_reports=False,
+                           modal=False,
+                           instance_image_tag='latest',
+                           report_dir=str(report_dir))
+        logger.info("Completed swe_bench run %s", self.run_id)
+
+        swe_bench_report_file = f"{self._model_name_or_path}.{self.run_id}.json"
+
+        # There is a bug in swebench because of which report_dir is being ignored. Copy the report to the output dir
+        self.move_report_and_logs(swe_bench_report_file=swe_bench_report_file, logs_dir="logs", report_dir=report_dir)
+        logger.info("SWE_bench report and logs written to %s directory", report_dir)
+
+        # read the swe_bench report file
+        report_file = report_dir / swe_bench_report_file
+        # 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:
+                report = json.load(f)
+                resolved_instances = report.get("resolved_instances", 0)
+                total_instances = report.get("total_instances", 0)
+                avg_score = self.compute_score(resolved_instances, total_instances)
+
+        # Build the EvalOutput from self._swe_bench_outputs and avg_score
+        eval_output_items = list(self.build_eval_output())
+        return EvalOutput(average_score=avg_score, eval_output_items=eval_output_items)

nat/eval/swe_bench_evaluator/register.py

@@ -0,0 +1,36 @@
+# 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
+
+
+class SweBenchEvaluatorConfig(EvaluatorBaseConfig, name="swe_bench"):
+    """Code patch evaluation for SWE Bench problems."""
+
+    run_id: str = Field(description="swe-bench test harness run identifier.")
+
+
+@register_evaluator(config_type=SweBenchEvaluatorConfig)
+async def register_swe_bench_evaluator(config: SweBenchEvaluatorConfig, builder: EvalBuilder):
+
+    from .evaluate import SweBenchEvaluator
+    _evaluator = SweBenchEvaluator(config.run_id, builder.get_max_concurrency(), builder.get_output_dir())
+
+    yield EvaluatorInfo(config=config, evaluate_fn=_evaluator.evaluate, description="SWE Bench Evaluator")

nat/eval/trajectory_evaluator/evaluate.py

@@ -0,0 +1,75 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-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.
+
+import logging
+
+from langchain.evaluation import TrajectoryEvalChain
+from langchain_core.language_models import BaseChatModel
+from langchain_core.tools import BaseTool
+
+from nat.eval.evaluator.base_evaluator import BaseEvaluator
+from nat.eval.evaluator.evaluator_model import EvalInputItem
+from nat.eval.evaluator.evaluator_model import EvalOutputItem
+
+logger = logging.getLogger(__name__)
+
+
+class TrajectoryEvaluator(BaseEvaluator):
+
+    def __init__(
+        self,
+        llm: BaseChatModel,
+        tools: list[BaseTool] | None = None,
+        max_concurrency: int = 8,
+    ):
+        super().__init__(max_concurrency=max_concurrency, tqdm_desc="Evaluating Trajectory")
+        self.llm = llm
+        self.tools = tools
+        # Initialize trajectory evaluation chain
+        self.traj_eval_chain = TrajectoryEvalChain.from_llm(llm=self.llm,
+                                                            tools=self.tools,
+                                                            return_reasoning=True,
+                                                            requires_reference=True)
+        logger.debug("Trajectory evaluation chain initialized.")
+
+    async def evaluate_item(self, item: EvalInputItem) -> EvalOutputItem:
+        """
+        Evaluate a single EvalInputItem and return an EvalOutputItem.
+        """
+        from nat.data_models.intermediate_step import IntermediateStepType
+        from nat.eval.intermediate_step_adapter import IntermediateStepAdapter
+
+        intermediate_step_adapter = IntermediateStepAdapter()
+        event_filter = [IntermediateStepType.LLM_END, IntermediateStepType.TOOL_END]
+
+        question = item.input_obj
+        generated_answer = item.output_obj
+        agent_trajectory = intermediate_step_adapter.get_agent_actions(item.trajectory, event_filter)
+
+        try:
+            eval_result = await self.traj_eval_chain.aevaluate_agent_trajectory(
+                input=question,
+                agent_trajectory=agent_trajectory,
+                prediction=generated_answer,
+            )
+        except Exception as e:
+            logger.exception("Error evaluating trajectory for question: %s, Error: %s", question, e, exc_info=True)
+            return EvalOutputItem(id=item.id, score=0.0, reasoning=f"Error evaluating trajectory: {e}")
+
+        reasoning = {
+            "reasoning": eval_result["reasoning"],
+            "trajectory": [(action.model_dump(), output) for (action, output) in agent_trajectory]
+        }
+        return EvalOutputItem(id=item.id, score=eval_result["score"], reasoning=reasoning)

nat/eval/trajectory_evaluator/register.py

@@ -0,0 +1,40 @@
+# 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
+
+
+class TrajectoryEvaluatorConfig(EvaluatorBaseConfig, name="trajectory"):
+    """Agent Trajectory Evaluation."""
+
+    llm_name: str = Field(description="LLM as a judge.")
+
+
+@register_evaluator(config_type=TrajectoryEvaluatorConfig)
+async def register_trajectory_evaluator(config: TrajectoryEvaluatorConfig, builder: EvalBuilder):
+    from nat.builder.framework_enum import LLMFrameworkEnum
+
+    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)
+
+    _evaluator = TrajectoryEvaluator(llm, tools, builder.get_max_concurrency())
+
+    yield EvaluatorInfo(config=config, evaluate_fn=_evaluator.evaluate, description="Trajectory Evaluator")