aiqtoolkit 1.2.0a20250706__py3-none-any.whl → 1.2.0a20250730__py3-none-any.whl

aiq/eval/rag_evaluator/evaluate.py

@@ -14,6 +14,7 @@
 # limitations under the License.
 
 import logging
+import math
 from collections.abc import Sequence
 
 from pydantic import BaseModel
@@ -53,9 +54,9 @@ class RAGEvaluator:
             if self.input_obj_field and hasattr(input_obj, self.input_obj_field):
                 # If input_obj_field is specified, return the value of that field
                 return str(getattr(input_obj, self.input_obj_field, ""))
-            else:
-                # If no input_obj_field is specified, return the string representation of the model
-                return input_obj.model_dump_json()
+
+            # If no input_obj_field is specified, return the string representation of the model
+            return input_obj.model_dump_json()
 
         if isinstance(input_obj, dict):
             # If input_obj is a dict, return the JSON string representation
@@ -105,19 +106,29 @@ class RAGEvaluator:
             return EvalOutput(average_score=0.0, eval_output_items=[])
 
         scores: list[dict[str, float]] = results_dataset.scores
+
+        # If Ragas returned no scores, return empty output to avoid downstream errors
         if not scores:
-            logger.error("Ragas returned empty score list")
+            logger.warning("Ragas returned empty score list")
             return EvalOutput(average_score=0.0, eval_output_items=[])
 
-        # Convert from list of dicts to dict of lists
-        scores_dict = {metric: [score[metric] for score in scores] for metric in scores[0]}
+        def _nan_to_zero(v: float | None) -> float:
+            """Convert NaN or None to 0.0 for safe arithmetic/serialization."""
+            return 0.0 if v is None or (isinstance(v, float) and math.isnan(v)) else v
+
+        # Convert from list of dicts to dict of lists, coercing NaN/None to 0.0
+        scores_dict = {metric: [_nan_to_zero(score.get(metric)) for score in scores] for metric in scores[0]}
+        first_metric_name = list(scores_dict.keys())[0] if scores_dict else None
 
-        # Compute the average of each metric
-        average_scores = {metric: sum(values) / len(values) for metric, values in scores_dict.items()}
+        # Compute the average of each metric, guarding against empty lists
+        average_scores = {
+            metric: (sum(values) / len(values) if values else 0.0)
+            for metric, values in scores_dict.items()
+        }
 
-        # Extract the first (and only) metric's average score
-        first_avg_score = next(iter(average_scores.values()))
-        first_metric_name = list(scores_dict.keys())[0]
+        first_avg_score = average_scores.get(list(scores_dict.keys())[0], 0.0)
+        if isinstance(first_avg_score, float) and math.isnan(first_avg_score):
+            first_avg_score = 0.0
 
         df = results_dataset.to_pandas()
         # Get id from eval_input if df size matches number of eval_input_items
@@ -130,7 +141,7 @@ class RAGEvaluator:
         eval_output_items = [
             EvalOutputItem(
                 id=ids[i],
-                score=getattr(row, first_metric_name, 0.0),
+                score=_nan_to_zero(getattr(row, first_metric_name, 0.0) if first_metric_name else 0.0),
                 reasoning={
                     key:
                         getattr(row, key, None)  # Use getattr to safely access attributes

aiq/eval/remote_workflow.py

@@ -24,6 +24,7 @@ from tqdm import tqdm
 from aiq.data_models.api_server import AIQResponseIntermediateStep
 from aiq.data_models.intermediate_step import IntermediateStep
 from aiq.data_models.intermediate_step import IntermediateStepPayload
+from aiq.data_models.invocation_node import InvocationNode
 from aiq.eval.config import EvaluationRunConfig
 from aiq.eval.evaluator.evaluator_model import EvalInput
 from aiq.eval.evaluator.evaluator_model import EvalInputItem
@@ -81,8 +82,12 @@ class EvaluationRemoteWorkflowHandler:
                             step_data = json.loads(line[len(INTERMEDIATE_DATA_PREFIX):])
                             response_intermediate = AIQResponseIntermediateStep.model_validate(step_data)
                             # The payload is expected to be IntermediateStepPayload
-                            intermediate_step = IntermediateStep(
-                                payload=IntermediateStepPayload.model_validate_json(response_intermediate.payload))
+                            payload = IntermediateStepPayload.model_validate_json(response_intermediate.payload)
+                            intermediate_step = IntermediateStep(parent_id="remote",
+                                                                 function_ancestry=InvocationNode(
+                                                                     function_name=payload.name or "remote_function",
+                                                                     function_id=payload.UUID or "remote_function_id"),
+                                                                 payload=payload)
                             intermediate_steps.append(intermediate_step)
                         except (json.JSONDecodeError, ValidationError) as e:
                             logger.error("Failed to parse intermediate step: %s", e)

aiq/eval/runners/__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.

aiq/eval/runners/config.py

@@ -0,0 +1,39 @@
+# 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 typing
+
+from pydantic import BaseModel
+
+from aiq.eval.config import EvaluationRunConfig
+from aiq.eval.config import EvaluationRunOutput
+
+
+class MultiEvaluationRunConfig(BaseModel):
+    """
+    Parameters used for a multi-evaluation run.
+    This includes a dict of configs. The key is an id of any type.
+    Each pass loads the config, applies the overrides and runs to completion
+    before the next pass starts.
+    """
+    configs: dict[typing.Any, EvaluationRunConfig]
+
+
+class MultiEvaluationRunOutput(BaseModel):
+    """
+    Output of a multi-evaluation run.
+    The results per-pass are accumulated in the evaluation_run_outputs dict.
+    """
+    evaluation_run_outputs: dict[typing.Any, EvaluationRunOutput]

aiq/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 aiq.eval.config import EvaluationRunConfig
+from aiq.eval.config import EvaluationRunOutput
+from aiq.eval.evaluate import EvaluationRun
+from aiq.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()

aiq/eval/usage_stats.py

@@ -28,8 +28,14 @@ class UsageStatsItem(BaseModel):
     usage_stats_per_llm: dict[str, UsageStatsLLM]
     total_tokens: int | None = None
     runtime: float = 0.0
+    min_timestamp: float = 0.0
+    max_timestamp: float = 0.0
+    llm_latency: float = 0.0
 
 
 class UsageStats(BaseModel):
     # key is the id or input_obj from EvalInputItem
+    min_timestamp: float = 0.0
+    max_timestamp: float = 0.0
+    total_runtime: float = 0.0
     usage_stats_items: dict[typing.Any, UsageStatsItem] = {}

aiq/eval/utils/weave_eval.py

@@ -152,10 +152,14 @@ class WeaveEvaluationIntegration:  # pylint: disable=too-many-public-methods
     def _log_profiler_metrics(self, profiler_results: ProfilerResults, usage_stats: UsageStats) -> dict[str, Any]:
         """Log profiler metrics to Weave."""
         profile_metrics = {}
+        if profiler_results.llm_latency_ci:
+            profile_metrics["llm_latency_p95"] = profiler_results.llm_latency_ci.p95
         if profiler_results.workflow_runtime_metrics:
-            profile_metrics["wf_p95_runtime"] = profiler_results.workflow_runtime_metrics.p95
+            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
 
     def log_summary(self,

aiq/experimental/decorators/experimental_warning_decorator.py

@@ -0,0 +1,130 @@
+# 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 functools
+import inspect
+import logging
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+BASE_WARNING_MESSAGE = ("is experimental and the API may change in future releases. "
+                        "Future versions may introduce breaking changes without notice.")
+
+_warning_issued = set()
+
+
+def issue_experimental_warning(function_name: str,
+                               feature_name: str | None = None,
+                               metadata: dict[str, Any] | None = None):
+    """
+    Log a warning message that the function is experimental.
+
+    A warning is emitted only once per function.  When a ``metadata`` dict
+    is supplied, it is appended to the log entry to provide extra context
+    (e.g., version, author, feature flag).
+    """
+    if function_name not in _warning_issued:
+        if (feature_name):
+            warning_message = f"The {feature_name} feature {BASE_WARNING_MESSAGE}"
+        else:
+            warning_message = f"This function {BASE_WARNING_MESSAGE}"
+
+        warning_message += f" Function: {function_name}"
+
+        if (metadata):
+            warning_message += f" | Metadata: {metadata}"
+
+        # Issue warning and save function name to avoid duplicate warnings
+        logger.warning(warning_message)
+
+        _warning_issued.add(function_name)
+
+
+def aiq_experimental(func: Any = None, *, feature_name: str | None = None, metadata: dict[str, Any] | None = None):
+    """
+    Decorator that can wrap any type of function (sync, async, generator,
+    async generator) and logs a warning that the function is experimental.
+
+    Args:
+        func: The function to be decorated.
+        feature_name: Optional name of the feature that is experimental. If provided, the warning will be
+        prefixed with "The <feature_name> feature is experimental".
+        metadata: Optional dictionary of metadata to log with the warning. This can include information
+        like version, author, etc. If provided, the metadata will be
+        logged alongside the experimental warning.
+    """
+    function_name: str = f"{func.__module__}.{func.__qualname__}" if func else "<unknown_function>"
+
+    # If called as @track_function(...) but not immediately passed a function
+    if func is None:
+
+        def decorator_wrapper(actual_func):
+            return aiq_experimental(actual_func, feature_name=feature_name, metadata=metadata)
+
+        return decorator_wrapper
+
+    # --- Validate metadata ---
+    if metadata is not None:
+        if not isinstance(metadata, dict):
+            raise TypeError("metadata must be a dict[str, Any].")
+        if any(not isinstance(k, str) for k in metadata.keys()):
+            raise TypeError("All metadata keys must be strings.")
+
+    # --- Now detect the function type and wrap accordingly ---
+    if inspect.isasyncgenfunction(func):
+        # ---------------------
+        # ASYNC GENERATOR
+        # ---------------------
+
+        @functools.wraps(func)
+        async def async_gen_wrapper(*args, **kwargs):
+            issue_experimental_warning(function_name, feature_name, metadata)
+            async for item in func(*args, **kwargs):
+                yield item  # yield the original item
+
+        return async_gen_wrapper
+
+    if inspect.iscoroutinefunction(func):
+        # ---------------------
+        # ASYNC FUNCTION
+        # ---------------------
+        @functools.wraps(func)
+        async def async_wrapper(*args, **kwargs):
+            issue_experimental_warning(function_name, feature_name, metadata)
+            result = await func(*args, **kwargs)
+            return result
+
+        return async_wrapper
+
+    if inspect.isgeneratorfunction(func):
+        # ---------------------
+        # SYNC GENERATOR
+        # ---------------------
+        @functools.wraps(func)
+        def sync_gen_wrapper(*args, **kwargs):
+            issue_experimental_warning(function_name, feature_name, metadata)
+            for item in func(*args, **kwargs):
+                yield item  # yield the original item
+
+        return sync_gen_wrapper
+
+    @functools.wraps(func)
+    def sync_wrapper(*args, **kwargs):
+        issue_experimental_warning(function_name, feature_name, metadata)
+        result = func(*args, **kwargs)
+        return result
+
+    return sync_wrapper

aiq/experimental/inference_time_scaling/editing/iterative_plan_refinement_editor.py

@@ -0,0 +1,147 @@
+# 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
+import re
+
+from aiq.builder.builder import Builder
+from aiq.builder.framework_enum import LLMFrameworkEnum
+from aiq.cli.register_workflow import register_its_strategy
+from aiq.data_models.its_strategy import ITSStrategyBaseConfig
+from aiq.experimental.inference_time_scaling.models.editor_config import IterativePlanRefinementConfig
+from aiq.experimental.inference_time_scaling.models.its_item import ITSItem
+from aiq.experimental.inference_time_scaling.models.stage_enums import PipelineTypeEnum
+from aiq.experimental.inference_time_scaling.models.stage_enums import StageTypeEnum
+from aiq.experimental.inference_time_scaling.models.strategy_base import StrategyBase
+from aiq.utils.io.model_processing import remove_r1_think_tags
+
+logger = logging.getLogger(__name__)
+
+
+class IterativePlanRefinementEditor(StrategyBase):
+    """
+    A planner that generates an initial plan, then refines it multiple times
+    using the same LLM. Each iteration updates the plan to (hopefully) be better.
+    """
+
+    def __init__(self, config: ITSStrategyBaseConfig) -> None:
+        super().__init__(config)
+        self.llm_bound = None
+
+    def supported_pipeline_types(self) -> [PipelineTypeEnum]:
+        return [PipelineTypeEnum.PLANNING]
+
+    def stage_type(self) -> StageTypeEnum:
+        return StageTypeEnum.EDITING
+
+    async def build_components(self, builder: Builder) -> None:
+        """
+        Build the components required for the iterative planner.
+        """
+        logger.debug("Building components for IterativePlanRefinementEditor")
+        self.llm_bound = await builder.get_llm(self.config.editor_llm, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
+
+    async def refine_single(self, prompt: str, context: str, its_item: ITSItem, prompt_idx: int) -> ITSItem:
+        from langchain_core.language_models import BaseChatModel
+        from langchain_core.prompts import PromptTemplate
+
+        if not isinstance(self.llm_bound, BaseChatModel):
+            raise ValueError("editor_llm must be a BaseChatModel instance for iterative plan refinement.")
+
+        llm: BaseChatModel = self.llm_bound
+
+        # Refinement loop
+        refinement_template = PromptTemplate(
+            template=self.config.refinement_template,
+            input_variables=["current_plan", "context", "original_prompt"],
+            validate_template=True,
+        )
+
+        current_plan = its_item.plan
+        for iteration in range(1, self.config.num_iterations + 1):
+            logger.info("Refinement iteration %d / %d for prompt %d", iteration, self.config.num_iterations, prompt_idx)
+            refine_prompt = (await refinement_template.ainvoke({
+                "current_plan": current_plan, "context": context, "original_prompt": prompt
+            })).to_string()
+
+            refine_response = await llm.ainvoke(refine_prompt)
+            refined_plan = remove_r1_think_tags(
+                refine_response.content if hasattr(refine_response, 'content') else str(refine_response))
+            refined_plan = re.sub(r'(?i)^\s*EDITED PLAN:\s*', '', refined_plan).strip()
+            if refined_plan:
+                current_plan = refined_plan
+            else:
+                logger.warning("Refinement iteration %d for prompt %d produced an empty plan; keeping existing plan.",
+                               iteration,
+                               prompt_idx)
+
+        logger.info("IterativePlanRefinementPlanner produced a final plan after %d iterations.",
+                    self.config.num_iterations)
+
+        its_item.plan = current_plan
+        # Return a single final plan
+        return its_item
+
+    async def ainvoke(self,
+                      items: list[ITSItem],
+                      original_prompt: str | None = None,
+                      agent_context: str | None = None,
+                      **kwargs) -> list[ITSItem]:
+        """
+        Runs the iterative plan refinement process on the provided planning items.
+
+        Each planning item is refined in parallel the configured number of times. Default is 3.
+
+        Args:
+            items (list[ITSItem]): The planning items to refine.
+            original_prompt (str): The original prompt used to generate the plans.
+            agent_context (str): The context for the agent.
+
+        Returns:
+            list[ITSItem]: The refined planning items.
+        """
+
+        if not original_prompt or not agent_context:
+            raise ValueError("Arguments original_prompt and agent_context must be provdied.")
+
+        # Generate feedback for each planning item concurrently
+        tasks = [
+            self.refine_single(prompt=original_prompt, context=agent_context, its_item=item, prompt_idx=i + 1)
+            for i, item in enumerate(items)
+        ]
+
+        # Run the tasks concurrently and gather results
+        refined_planning_items = await asyncio.gather(*tasks)
+
+        return refined_planning_items
+
+
+@register_its_strategy(config_type=IterativePlanRefinementConfig)
+async def register_iterative_plan_refinement_editor(config: IterativePlanRefinementConfig, builder: Builder):
+    """
+    Register the IterativePlanRefinementEditor strategy.
+
+    Args:
+        config (IterativePlanRefinementConfig): The configuration for the strategy.
+
+    Returns:
+        IterativePlanRefinementEditor: The registered strategy instance.
+    """
+
+    editor = IterativePlanRefinementEditor(config)
+    await editor.build_components(builder=builder)
+
+    yield editor