aiqtoolkit 1.2.0.dev0__py3-none-any.whl → 1.2.0rc2__py3-none-any.whl

aiq/eval/usage_stats.py

@@ -0,0 +1,41 @@
+# 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
+
+
+class UsageStatsLLM(BaseModel):
+    prompt_tokens: int = 0
+    completion_tokens: int = 0
+    total_tokens: int = 0
+
+
+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/output_uploader.py

@@ -78,9 +78,18 @@ class OutputUploader:
 
         session = aioboto3.Session()
         try:
+            if self.s3_config.endpoint_url:
+                region_name = None
+                endpoint_url = self.s3_config.endpoint_url
+            elif self.s3_config.region_name:
+                region_name = self.s3_config.region_name
+                endpoint_url = None
+            else:
+                raise ValueError("No endpoint_url or region_name provided in the config: eval.general.output.s3")
             async with session.client(
                     "s3",
-                    endpoint_url=self.s3_config.endpoint_url,
+                    endpoint_url=endpoint_url,
+                    region_name=region_name,
                     aws_access_key_id=self.s3_config.access_key,
                     aws_secret_access_key=self.s3_config.secret_key,
             ) as s3_client:

aiq/eval/utils/weave_eval.py

@@ -0,0 +1,184 @@
+# 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 typing import Any
+
+from aiq.eval.evaluator.evaluator_model import EvalInput
+from aiq.eval.evaluator.evaluator_model import EvalInputItem
+from aiq.eval.evaluator.evaluator_model import EvalOutput
+from aiq.eval.usage_stats import UsageStats
+from aiq.eval.usage_stats import UsageStatsItem
+from aiq.profiler.data_models import ProfilerResults
+
+logger = logging.getLogger(__name__)
+
+
+class WeaveEvaluationIntegration:  # pylint: disable=too-many-public-methods
+    """
+    Class to handle all Weave integration functionality.
+    """
+
+    def __init__(self):
+        self.available = False
+        self.client = None
+        self.eval_logger = None
+        self.pred_loggers = {}
+
+        try:
+            from weave.flow.eval_imperative import EvaluationLogger
+            from weave.flow.eval_imperative import ScoreLogger
+            from weave.trace.context import weave_client_context
+            self.EvaluationLogger = EvaluationLogger
+            self.ScoreLogger = ScoreLogger
+            self.weave_client_context = weave_client_context
+            self.available = True
+        except ImportError:
+            self.available = False
+            # we simply don't do anything if weave is not available
+            pass
+
+    def initialize_client(self):
+        """Initialize the Weave client if available."""
+        if not self.available:
+            return False
+
+        try:
+            self.client = self.weave_client_context.require_weave_client()
+            return self.client is not None
+        except Exception:
+            self.client = None
+            return False
+
+    def _get_prediction_inputs(self, item: EvalInputItem):
+        """Get the inputs for displaying in the UI.
+        The following fields are excluded as they are too large to display in the UI:
+        - full_dataset_entry
+        - expected_trajectory
+        - trajectory
+
+        output_obj is excluded because it is displayed separately.
+        """
+        include = {"id", "input_obj", "expected_output_obj"}
+        return item.model_dump(include=include)
+
+    def _get_weave_dataset(self, eval_input: EvalInput):
+        """Get the full dataset for Weave."""
+        return [item.full_dataset_entry for item in eval_input.eval_input_items]
+
+    def initialize_logger(self, workflow_alias: str, eval_input: EvalInput, config: Any):
+        """Initialize the Weave evaluation logger."""
+        if not self.client and not self.initialize_client():
+            # lazy init the client
+            return False
+
+        try:
+            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.pred_loggers = {}
+
+            return True
+        except Exception as e:
+            self.eval_logger = None
+            logger.warning("Failed to initialize Weave `EvaluationLogger`: %s", e)
+
+            return False
+
+    def log_prediction(self, item: EvalInputItem, output: Any):
+        """Log a prediction to Weave."""
+        if not self.eval_logger:
+            return
+
+        pred_logger = self.eval_logger.log_prediction(inputs=self._get_prediction_inputs(item), output=output)
+        self.pred_loggers[item.id] = pred_logger
+
+    async def log_usage_stats(self, item: EvalInputItem, usage_stats_item: UsageStatsItem):
+        """Log usage stats to Weave."""
+        if not self.eval_logger:
+            return
+
+        # log each usage stat as a score
+        await self.pred_loggers[item.id].alog_score(scorer="wf_runtime", score=usage_stats_item.runtime)
+
+        # log the total tokens for this item, per-llm tokens can be exported later if needed
+        await self.pred_loggers[item.id].alog_score(scorer="wf_tokens", score=usage_stats_item.total_tokens)
+
+    async def alog_score(self, eval_output: EvalOutput, evaluator_name: str):
+        """Log scores for evaluation outputs."""
+        if not self.eval_logger:
+            return
+
+        # Create coroutines for all score logging operations
+        coros = []
+        for eval_output_item in eval_output.eval_output_items:
+            if eval_output_item.id in self.pred_loggers:
+                coros.append(self.pred_loggers[eval_output_item.id].alog_score(
+                    scorer=evaluator_name,
+                    score=eval_output_item.score,
+                ))
+
+        # Execute all coroutines concurrently
+        if coros:
+            await asyncio.gather(*coros)
+
+    async def afinish_loggers(self):
+        """Finish all prediction loggers."""
+        if not self.eval_logger:
+            return
+
+        async def _finish_one(pred_logger):
+            if hasattr(pred_logger, '_has_finished') and not pred_logger._has_finished:
+                return
+            # run the *blocking* finish() in a thread so we don't nest loops
+            await asyncio.to_thread(pred_logger.finish)
+
+        await asyncio.gather(*[_finish_one(pl) for pl in self.pred_loggers.values()])
+
+    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_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,
+                    usage_stats: UsageStats,
+                    evaluation_results: list[tuple[str, EvalOutput]],
+                    profiler_results: ProfilerResults):
+        """Log summary statistics to Weave."""
+        if not self.eval_logger:
+            return
+
+        summary = {}
+        # add evaluation results to the summary
+        for evaluator_name, eval_output in evaluation_results:
+            summary[evaluator_name] = eval_output.average_score
+
+        # add profiler metrics to the summary
+        profile_metrics = self._log_profiler_metrics(profiler_results, usage_stats)
+        summary.update(profile_metrics)
+
+        # 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)

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