nvidia-nat 1.2.0rc5__py3-none-any.whl

aiq/profiler/inference_optimization/prompt_caching.py

@@ -0,0 +1,163 @@
+# 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 aiq.data_models.intermediate_step import IntermediateStep
+from aiq.profiler.inference_optimization.data_models import CommonPrefixesOutput
+from aiq.profiler.inference_optimization.data_models import FrameworkLLMPrefixData
+from aiq.profiler.inference_optimization.data_models import PrefixInfo
+from aiq.profiler.utils import create_standardized_dataframe
+
+
+# -----------------------------------------------------------
+# 1. Helper: Build a prefix trie
+# -----------------------------------------------------------
+def build_prefix_trie(strings: list[str]) -> dict:
+    """
+    Build a trie from a list of strings.
+
+    Returns a nested dictionary with::
+
+        {
+            'count': int,         # number of strings passing through this node
+            'children': dict[str, TrieNode]
+        }
+
+    """
+    root = {'count': 0, 'children': {}}
+    for s in strings:
+        node = root
+        node['count'] += 1  # every string passes through the root
+        for ch in s:
+            if ch not in node['children']:
+                node['children'][ch] = {'count': 0, 'children': {}}
+            node = node['children'][ch]
+            node['count'] += 1
+    return root
+
+
+# -----------------------------------------------------------
+# 2. Helper: Iterative traversal of the trie
+# -----------------------------------------------------------
+def collect_prefixes_iterative(root: dict, total_calls: int) -> list[dict]:
+    """
+    Iteratively traverse the trie to collect prefix statistics,
+    avoiding recursion depth limits.
+
+    :param root: Trie node with 'count' and 'children'
+    :param total_calls: Number of total calls in this group (denominator for percentages)
+    :return: A list of dicts, each dict containing prefix info
+    """
+    results = []
+    # stack holds (node, prefix_so_far)
+    stack = [(root, "")]
+
+    while stack:
+        node, prefix = stack.pop()
+
+        # Skip storing the empty root prefix
+        if prefix:
+            calls_count = node['count']
+            calls_percentage = calls_count / total_calls
+            results.append({
+                'prefix': prefix,
+                'prefix_length': len(prefix),
+                'calls_count': calls_count,
+                'calls_percentage': calls_percentage
+            })
+
+        # Add children to the stack
+        for ch, child_node in node['children'].items():
+            stack.append((child_node, prefix + ch))
+
+    return results
+
+
+# -----------------------------------------------------------
+# 3. Main Function
+# -----------------------------------------------------------
+def get_common_prefixes(all_steps: list[list[IntermediateStep]],
+                        min_call_percentage: float = 0.0) -> CommonPrefixesOutput:
+    """
+    Given a pandas DataFrame with columns 'framework', 'llm_name',
+    and 'llm_text_input', return a Pydantic-validated RootModel
+    keyed by "<llm_name>" with a sorted list of
+    common prefix statistics.
+
+    1) Only includes prefixes with calls_percentage >= `min_call_percentage`.
+    2) Excludes any prefix that is a substring of another (longer) prefix
+       that already meets the threshold and is retained.
+    3) Optionally writes the resulting dictionary to JSON if `output_path` is provided.
+
+    :param all_steps: Intermediate Steps
+    :param min_call_percentage: Exclude prefixes that appear in fewer than this fraction
+                                of total calls. (Default 0.0 = no filtering)
+
+    Sorting: primarily by prefix length (descending),
+             secondarily by frequency (descending).
+    """
+    # Validate necessary columns
+    df = create_standardized_dataframe(all_steps)
+
+    required_cols = {'framework', 'llm_name', 'llm_text_input'}
+    if not required_cols.issubset(df.columns):
+        missing = required_cols - set(df.columns)
+        raise ValueError(f"DataFrame missing required columns: {missing}")
+
+    output_data: dict[str, FrameworkLLMPrefixData] = {}
+
+    # Group DataFrame by (framework, llm_name)
+    grouped = df.groupby(['llm_name'])
+    for llm_name, group_df in grouped:
+        # Unpack llm_name Tuple
+        llm_name = llm_name[0]
+
+        text_inputs = group_df['llm_text_input'].astype(str).tolist()
+        total_calls = len(text_inputs)
+
+        # Build trie for all text inputs
+        trie = build_prefix_trie(text_inputs)
+
+        # Collect prefix info using iterative traversal
+        results = collect_prefixes_iterative(trie, total_calls=total_calls)
+
+        # 1) Filter out prefixes below min_call_percentage
+        results_filtered = [r for r in results if r['calls_percentage'] >= min_call_percentage]
+
+        # 2) Sort results: prefix_length desc, then calls_count desc
+        results_sorted = sorted(results_filtered, key=lambda x: (x['prefix_length'], x['calls_count']), reverse=True)
+
+        # 3) Substring filtering:
+        #    Because results_sorted is in descending length order,
+        #    if we keep a prefix, we exclude any shorter prefix that
+        #    is a substring of that already-kept prefix.
+        final_results = []
+        for r in results_sorted:
+            pfx = r['prefix']
+            # Check if this prefix is contained in any longer prefix we have kept
+            if not any(pfx in kept['prefix'] for kept in final_results):
+                final_results.append(r)
+
+        # Convert each dict to a PrefixInfo model
+        prefix_info_list = [PrefixInfo(**res) for res in final_results]
+
+        # Construct the dictionary key
+        framework_llm_key = f"{llm_name}"
+
+        # Save the data for this group
+        output_data[framework_llm_key] = FrameworkLLMPrefixData(total_calls=total_calls, prefix_info=prefix_info_list)
+
+    # Package the final result in a validated RootModel
+    result_model = CommonPrefixesOutput(root=output_data)
+    return result_model

aiq/profiler/inference_optimization/token_uniqueness.py

@@ -0,0 +1,107 @@
+# 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 re
+
+import numpy as np
+
+from aiq.data_models.intermediate_step import IntermediateStep
+from aiq.profiler.inference_optimization.data_models import LLMUniquenessMetrics
+from aiq.profiler.inference_optimization.data_models import LLMUniquenessMetricsByLLM
+from aiq.profiler.utils import create_standardized_dataframe
+
+
+# ----------------------------------------------------------------
+# 1. Main Function
+# ----------------------------------------------------------------
+def compute_inter_query_token_uniqueness_by_llm(all_steps: list[list[IntermediateStep]]) -> LLMUniquenessMetricsByLLM:
+    """
+    Computes p90, p95, and p99 of 'new words added' between consecutive llm_start events,
+    grouped by (llm_name, example_number).
+
+    Steps:
+
+    1. Filter df to only llm_start events.
+    2. Group first by (llm_name, example_number), then sort by event_timestamp in each group.
+    3. Compare each llm_text_input to the previous one in the same group to find how many new words appear.
+    4. Aggregate all 'new words count' across each llm_name, compute p90/p95/p99 for each LLM.
+    5. Return a Pydantic RootModel containing a dictionary::
+
+         { llm_name -> LLMUniquenessMetrics(p90, p95, p99) }.
+    """
+    df = create_standardized_dataframe(all_steps)
+    # Validate that the necessary columns exist
+    required_cols = {'event_type', 'llm_name', 'example_number', 'event_timestamp', 'llm_text_input'}
+    missing = required_cols - set(df.columns)
+    if missing:
+        raise ValueError(f"DataFrame missing required columns: {missing}")
+
+    # 1) Filter to llm_start events
+    cdf = df[df['event_type'] == 'LLM_START'].copy()
+    if cdf.empty:
+        # Return an empty dictionary if no llm_start events
+        return LLMUniquenessMetricsByLLM(root={})
+
+    # Helper to tokenize text into a set of words
+    def tokenize_to_set(text: str) -> set:
+        if not isinstance(text, str):
+            return set()
+        return set(re.findall(r"\w+", text.lower()))
+
+    # We'll store new_words counts for each llm_name
+    llm_to_counts: dict[str, list[int]] = {}
+
+    # 2) Group by (llm_name, example_number), then sort each group
+    grouped = cdf.groupby(['llm_name', 'example_number'], as_index=False, group_keys=True)
+
+    for (llm, ex_num), group_df in grouped:  # pylint: disable=unused-variable
+        # Sort by event_timestamp
+        group_df = group_df.sort_values('event_timestamp', ascending=True)
+
+        # Shift the llm_text_input to compare consecutive calls
+        group_df['prev_llm_text_input'] = group_df['llm_text_input'].shift(1)
+
+        # Compute new words for each row (excluding the first in the group)
+        def compute_new_words(row):
+            current_tokens = tokenize_to_set(row['llm_text_input'])
+            prev_tokens = tokenize_to_set(row['prev_llm_text_input'])
+            return len(current_tokens - prev_tokens)
+
+        group_df['new_words_count'] = group_df.apply(compute_new_words, axis=1)
+
+        # Drop rows where there's no 'previous' call
+        valid_rows = group_df.dropna(subset=['prev_llm_text_input'])
+
+        # Gather the new_words_count
+        counts = valid_rows['new_words_count'].tolist()
+        if counts:
+            # Accumulate them in llm_to_counts
+            if llm not in llm_to_counts:
+                llm_to_counts[llm] = []
+            llm_to_counts[llm].extend(counts)
+
+    # 4) For each llm_name, compute p90, p95, p99
+    output_dict = {}
+    for llm_name, counts_list in llm_to_counts.items():
+        arr = np.array(counts_list)
+        p90_val = float(np.percentile(arr, 90))
+        p95_val = float(np.percentile(arr, 95))
+        p99_val = float(np.percentile(arr, 99))
+
+        output_dict[llm_name] = LLMUniquenessMetrics(p90=p90_val, p95=p95_val, p99=p99_val)
+
+    ret_val = LLMUniquenessMetricsByLLM(root=output_dict)
+    # Validate & return as a RootModel
+    return ret_val

aiq/profiler/inference_optimization/workflow_runtimes.py

@@ -0,0 +1,72 @@
+# 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 numpy as np
+
+from aiq.data_models.intermediate_step import IntermediateStep
+from aiq.profiler.inference_optimization.data_models import WorkflowRuntimeMetrics
+from aiq.profiler.utils import create_standardized_dataframe
+
+
+def compute_workflow_runtime_metrics(all_steps: list[list[IntermediateStep]]) -> WorkflowRuntimeMetrics:
+    """
+    Computes the p90, p95, and p99 of workflow runtime for each example_number.
+
+    The 'workflow runtime' per example is::
+
+        max(event_timestamp) - min(event_timestamp)
+
+    for that example_number.
+
+    Parameters
+    ----------
+    all_steps : IntermediateStep
+        Must contain at least two columns:
+          - 'example_number'
+          - 'event_timestamp'
+
+    Returns
+    -------
+    WorkflowRuntimeMetrics
+        A Pydantic model with 'p90', 'p95', and 'p99' attributes.
+    """
+    df = create_standardized_dataframe(all_steps)
+    required_cols = {"example_number", "event_timestamp"}
+    missing = required_cols - set(df.columns)
+    if missing:
+        raise ValueError(f"DataFrame is missing required columns: {missing}")
+
+    # Group by example_number, then find min and max timestamp
+    grouped = df.groupby("example_number")["event_timestamp"]
+    min_timestamps = grouped.min()
+    max_timestamps = grouped.max()
+
+    # Workflow runtime is difference between max and min
+    runtimes = max_timestamps - min_timestamps
+
+    # Convert to a NumPy array for percentile calculations
+    runtimes_arr = runtimes.values
+
+    # Edge case: if there's only one example or no data
+    # (NumPy percentile can handle 1-element arrays, but let's guard for empties)
+    if len(runtimes_arr) == 0:
+        return WorkflowRuntimeMetrics(p90=0.0, p95=0.0, p99=0.0)
+
+    # Compute p90, p95, p99
+    p90_val = float(np.percentile(runtimes_arr, 90))
+    p95_val = float(np.percentile(runtimes_arr, 95))
+    p99_val = float(np.percentile(runtimes_arr, 99))
+
+    return WorkflowRuntimeMetrics(p90=p90_val, p95=p95_val, p99=p99_val)

aiq/profiler/intermediate_property_adapter.py

@@ -0,0 +1,102 @@
+# 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 aiq.data_models.intermediate_step import IntermediateStep
+from aiq.data_models.intermediate_step import IntermediateStepType
+from aiq.data_models.intermediate_step import TokenUsageBaseModel
+
+
+class IntermediatePropertyAdaptor(IntermediateStep):
+
+    @classmethod
+    def from_intermediate_step(cls, step: IntermediateStep) -> "IntermediatePropertyAdaptor":
+        """
+        Create an adaptor instance from an existing IntermediateStep.
+        Uses the dict() representation of the instance to initialize the adaptor.
+        """
+        return cls(**step.model_dump())
+
+    @property
+    def token_usage(self) -> TokenUsageBaseModel:
+        return self.payload.usage_info.token_usage if self.payload.usage_info else TokenUsageBaseModel()
+
+    @property
+    def seconds_between_calls(self) -> int:
+        return self.payload.usage_info.seconds_between_calls if self.payload.usage_info else 0
+
+    @property
+    def llm_text_input(self) -> str:
+        ret = ""
+        if self.payload.data and self.event_type == IntermediateStepType.LLM_START:
+            ret = self.payload.data.input
+        return ret
+
+    @property
+    def llm_text_output(self) -> str:
+        ret = ""
+        if self.payload.data and self.event_type == IntermediateStepType.LLM_END:
+            ret = self.payload.data.output
+        return ret
+
+    @property
+    def llm_text_chunk(self) -> str:
+        ret = ""
+        if self.payload.data and self.event_type == IntermediateStepType.LLM_NEW_TOKEN:
+            ret = self.payload.data.chunk
+        return ret
+
+    @property
+    def tool_input(self) -> str:
+        ret = ""
+        if self.payload.data and self.event_type == IntermediateStepType.TOOL_START:
+            ret = self.payload.data.input
+        return ret
+
+    @property
+    def tool_output(self) -> str:
+        ret = ""
+        if self.payload.data and self.event_type == IntermediateStepType.TOOL_END:
+            ret = self.payload.data.output
+        return ret
+
+    @property
+    def llm_name(self) -> str:
+        ret = ""
+        if self.payload.name and self.event_type in [IntermediateStepType.LLM_START, IntermediateStepType.LLM_END]:
+            ret = self.payload.name
+        return ret
+
+    @property
+    def tool_name(self) -> str:
+        ret = ""
+        if self.payload.name and self.event_type in [IntermediateStepType.TOOL_START, IntermediateStepType.TOOL_END]:
+            ret = self.payload.name
+        return ret
+
+    @property
+    def function_name(self) -> str:
+        return self.function_ancestry.function_name
+
+    @property
+    def function_id(self) -> str:
+        return self.function_ancestry.function_id
+
+    @property
+    def parent_function_id(self) -> str:
+        return self.function_ancestry.parent_id
+
+    @property
+    def parent_function_name(self) -> str:
+        return self.function_ancestry.parent_name