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

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())
 

nat/eval/tunable_rag_evaluator/evaluate.py

@@ -13,9 +13,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import asyncio
 import logging
-from typing import Callable
+from collections.abc import Callable
 
 from langchain.output_parsers import ResponseSchema
 from langchain.output_parsers import StructuredOutputParser
@@ -23,7 +22,6 @@ from langchain.schema import HumanMessage
 from langchain.schema import SystemMessage
 from langchain_core.language_models import BaseChatModel
 from langchain_core.runnables import RunnableLambda
-from tqdm import tqdm
 
 from nat.eval.evaluator.base_evaluator import BaseEvaluator
 from nat.eval.evaluator.evaluator_model import EvalInputItem
@@ -31,7 +29,6 @@ from nat.eval.evaluator.evaluator_model import EvalOutputItem
 
 logger = logging.getLogger(__name__)
 
-# pylint: disable=line-too-long
 # flake8: noqa: E501
 
 
@@ -185,8 +182,8 @@ class TunableRagEvaluator(BaseEvaluator):
                     relevance_score = parsed_response["relevance_score"]
                     reasoning = parsed_response["reasoning"]
                 except KeyError as e:
-                    logger.error("Missing required keys in default scoring response: %s",
-                                 ", ".join(str(arg) for arg in e.args))
+                    logger.exception("Missing required keys in default scoring response: %s",
+                                     ", ".join(str(arg) for arg in e.args))
                     reasoning = f"Error in evaluator from parsing judge LLM response. Missing required key(s): {', '.join(str(arg) for arg in e.args)}"
 
                 coverage_weight = self.default_score_weights.get("coverage", 1 / 3)
@@ -218,7 +215,7 @@ class TunableRagEvaluator(BaseEvaluator):
                     reasoning = f"Error in evaluator from parsing judge LLM response. Missing required key(s): {', '.join(str(arg) for arg in e.args)}"
                     raise
         except (KeyError, ValueError) as e:
-            logger.error("Error parsing judge LLM response: %s", e)
+            logger.exception("Error parsing judge LLM response: %s", e)
             score = 0.0
             reasoning = "Error in evaluator from parsing judge LLM response."
 

nat/eval/utils/eval_trace_ctx.py

@@ -0,0 +1,89 @@
+# 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 logging
+from collections.abc import Callable
+from contextlib import contextmanager
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+# Type alias for evaluation call objects that have an optional 'id' attribute
+EvalCallType = Any  # Could be Weave Call object or other tracing framework objects
+
+
+class EvalTraceContext:
+    """
+    Evaluation trace context manager for coordinating traces.
+
+    This class provides a framework-agnostic way to:
+    1. Track evaluation calls/contexts
+    2. Ensure proper parent-child relationships in traces
+    """
+
+    def __init__(self):
+        self.eval_call: EvalCallType | None = None  # Store the evaluation call/context for propagation
+
+    def set_eval_call(self, eval_call: EvalCallType | None) -> None:
+        """Set the evaluation call/context for propagation to traces."""
+        self.eval_call = eval_call
+        if eval_call:
+            logger.debug("Set evaluation call context: %s", getattr(eval_call, 'id', str(eval_call)))
+
+    def get_eval_call(self) -> EvalCallType | None:
+        """Get the current evaluation call/context."""
+        return self.eval_call
+
+    @contextmanager
+    def evaluation_context(self):
+        """
+        Context manager that can be overridden by framework-specific implementations.
+        Default implementation is a no-op.
+        """
+        yield
+
+
+class WeaveEvalTraceContext(EvalTraceContext):
+    """
+    Weave-specific implementation of evaluation trace context.
+    """
+
+    def __init__(self):
+        super().__init__()
+        self.available = False
+        self.set_call_stack: Callable[[list[EvalCallType]], Any] | None = None
+
+        try:
+            from weave.trace.context.call_context import set_call_stack
+            self.set_call_stack = set_call_stack
+            self.available = True
+        except ImportError:
+            self.available = False
+            logger.debug("Weave not available for trace context")
+
+    @contextmanager
+    def evaluation_context(self):
+        """Set the evaluation call as active context for Weave traces."""
+        if self.available and self.eval_call and self.set_call_stack:
+            try:
+                with self.set_call_stack([self.eval_call]):
+                    logger.debug("Set Weave evaluation call context: %s",
+                                 getattr(self.eval_call, 'id', str(self.eval_call)))
+                    yield
+            except Exception as e:
+                logger.warning("Failed to set Weave evaluation call context: %s", e)
+                yield
+        else:
+            yield

nat/eval/utils/weave_eval.py

@@ -15,6 +15,7 @@
 
 import asyncio
 import logging
+from typing import TYPE_CHECKING
 from typing import Any
 
 from nat.eval.evaluator.evaluator_model import EvalInput
@@ -24,26 +25,28 @@ from nat.eval.usage_stats import UsageStats
 from nat.eval.usage_stats import UsageStatsItem
 from nat.profiler.data_models import ProfilerResults
 
+if TYPE_CHECKING:
+    from nat.eval.utils.eval_trace_ctx import EvalTraceContext
+
 logger = logging.getLogger(__name__)
 
 
-class WeaveEvaluationIntegration:  # pylint: disable=too-many-public-methods
+class WeaveEvaluationIntegration:
     """
     Class to handle all Weave integration functionality.
     """
 
-    def __init__(self):
+    def __init__(self, eval_trace_context: "EvalTraceContext"):
         self.available = False
         self.client = None
         self.eval_logger = None
         self.pred_loggers = {}
+        self.eval_trace_context = eval_trace_context
 
         try:
-            from weave.flow.eval_imperative import EvaluationLogger
-            from weave.flow.eval_imperative import ScoreLogger
+            from weave import EvaluationLogger
             from weave.trace.context import weave_client_context
-            self.EvaluationLogger = EvaluationLogger
-            self.ScoreLogger = ScoreLogger
+            self.evaluation_logger_cls = EvaluationLogger
             self.weave_client_context = weave_client_context
             self.available = True
         except ImportError:
@@ -89,9 +92,15 @@ class WeaveEvaluationIntegration:  # pylint: disable=too-many-public-methods
             weave_dataset = self._get_weave_dataset(eval_input)
             config_dict = config.model_dump(mode="json")
             config_dict["name"] = workflow_alias
-            self.eval_logger = self.EvaluationLogger(model=config_dict, dataset=weave_dataset)
+            self.eval_logger = self.evaluation_logger_cls(model=config_dict,
+                                                          dataset=weave_dataset,
+                                                          name=workflow_alias,
+                                                          eval_attributes={})
             self.pred_loggers = {}
 
+            # Capture the current evaluation call for context propagation
+            self.eval_trace_context.set_eval_call(self.eval_logger._evaluate_call)
+
             return True
         except Exception as e:
             self.eval_logger = None
@@ -137,7 +146,7 @@ class WeaveEvaluationIntegration:  # pylint: disable=too-many-public-methods
             await asyncio.gather(*coros)
 
     async def afinish_loggers(self):
-        """Finish all prediction loggers."""
+        """Finish all prediction loggers and wait for exports."""
         if not self.eval_logger:
             return
 
@@ -157,7 +166,6 @@ class WeaveEvaluationIntegration:  # pylint: disable=too-many-public-methods
         if profiler_results.workflow_runtime_metrics:
             profile_metrics["wf_runtime_p95"] = profiler_results.workflow_runtime_metrics.p95
 
-        # TODO:get the LLM tokens from the usage stats and log them
         profile_metrics["total_runtime"] = usage_stats.total_runtime
 
         return profile_metrics
@@ -182,3 +190,4 @@ class WeaveEvaluationIntegration:  # pylint: disable=too-many-public-methods
         # Log the summary to finish the evaluation, disable auto-summarize
         # as we will be adding profiler metrics to the summary
         self.eval_logger.log_summary(summary, auto_summarize=False)
+        logger.info("Logged Evaluation Summary to Weave")

nat/experimental/decorators/experimental_warning_decorator.py

@@ -16,7 +16,12 @@
 import functools
 import inspect
 import logging
+from collections.abc import AsyncGenerator
+from collections.abc import Callable
+from collections.abc import Generator
 from typing import Any
+from typing import TypeVar
+from typing import overload
 
 logger = logging.getLogger(__name__)
 
@@ -25,6 +30,9 @@ BASE_WARNING_MESSAGE = ("is experimental and the API may change in future releas
 
 _warning_issued = set()
 
+# Type variables for overloads
+F = TypeVar('F', bound=Callable[..., Any])
+
 
 def issue_experimental_warning(function_name: str,
                                feature_name: str | None = None,
@@ -53,7 +61,20 @@ def issue_experimental_warning(function_name: str,
         _warning_issued.add(function_name)
 
 
-def experimental(func: Any = None, *, feature_name: str | None = None, metadata: dict[str, Any] | None = None):
+# Overloads for different function types
+@overload
+def experimental(func: F, *, feature_name: str | None = None, metadata: dict[str, Any] | None = None) -> F:
+    """Overload for when a function is passed directly."""
+    ...
+
+
+@overload
+def experimental(*, feature_name: str | None = None, metadata: dict[str, Any] | None = None) -> Callable[[F], F]:
+    """Overload for decorator factory usage (when called with parentheses)."""
+    ...
+
+
+def experimental(func: Any = None, *, feature_name: str | None = None, metadata: dict[str, Any] | None = None) -> Any:
     """
     Decorator that can wrap any type of function (sync, async, generator,
     async generator) and logs a warning that the function is experimental.
@@ -90,7 +111,7 @@ def experimental(func: Any = None, *, feature_name: str | None = None, metadata:
         # ---------------------
 
         @functools.wraps(func)
-        async def async_gen_wrapper(*args, **kwargs):
+        async def async_gen_wrapper(*args, **kwargs) -> AsyncGenerator[Any, Any]:
             issue_experimental_warning(function_name, feature_name, metadata)
             async for item in func(*args, **kwargs):
                 yield item  # yield the original item
@@ -102,7 +123,7 @@ def experimental(func: Any = None, *, feature_name: str | None = None, metadata:
         # ASYNC FUNCTION
         # ---------------------
         @functools.wraps(func)
-        async def async_wrapper(*args, **kwargs):
+        async def async_wrapper(*args, **kwargs) -> Any:
             issue_experimental_warning(function_name, feature_name, metadata)
             result = await func(*args, **kwargs)
             return result
@@ -114,15 +135,14 @@ def experimental(func: Any = None, *, feature_name: str | None = None, metadata:
         # SYNC GENERATOR
         # ---------------------
         @functools.wraps(func)
-        def sync_gen_wrapper(*args, **kwargs):
+        def sync_gen_wrapper(*args, **kwargs) -> Generator[Any, Any, Any]:
             issue_experimental_warning(function_name, feature_name, metadata)
-            for item in func(*args, **kwargs):
-                yield item  # yield the original item
+            yield from func(*args, **kwargs)  # yield the original item
 
         return sync_gen_wrapper
 
     @functools.wraps(func)
-    def sync_wrapper(*args, **kwargs):
+    def sync_wrapper(*args, **kwargs) -> Any:
         issue_experimental_warning(function_name, feature_name, metadata)
         result = func(*args, **kwargs)
         return result

nat/experimental/test_time_compute/functions/plan_select_execute_function.py

@@ -86,7 +86,7 @@ async def plan_select_execute_function(config: PlanSelectExecuteFunctionConfig,
                           "This error can be resolved by installing nvidia-nat-langchain.")
 
     # Get the augmented function's description
-    augmented_function = builder.get_function(config.augmented_fn)
+    augmented_function = await builder.get_function(config.augmented_fn)
 
     # For now, we rely on runtime checking for type conversion
 
@@ -97,11 +97,15 @@ async def plan_select_execute_function(config: PlanSelectExecuteFunctionConfig,
                          f"function without a description.")
 
     # Get the function dependencies of the augmented function
-    function_used_tools = builder.get_function_dependencies(config.augmented_fn).functions
+    function_dependencies = builder.get_function_dependencies(config.augmented_fn)
+    function_used_tools = set(function_dependencies.functions)
+    for function_group in function_dependencies.function_groups:
+        function_used_tools.update(builder.get_function_group_dependencies(function_group).functions)
+
     tool_list = "Tool: Description\n"
 
     for tool in function_used_tools:
-        tool_impl = builder.get_function(tool)
+        tool_impl = await builder.get_function(tool)
         tool_list += f"- {tool}: {tool_impl.description if hasattr(tool_impl, 'description') else ''}\n"
 
     # Draft the reasoning prompt for the augmented function