aiqtoolkit 1.2.0rc4__py3-none-any.whl → 1.2rc9__py3-none-any.whl

aiq/profiler/calc/plot.py

@@ -1,345 +0,0 @@
-# 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 pathlib import Path
-
-import matplotlib.pyplot as plt
-import numpy as np
-import pandas as pd
-
-from aiq.profiler.calc.data_models import LinearFitResult
-from aiq.profiler.calc.data_models import SizingMetrics
-
-logger = logging.getLogger(__name__)
-
-
-# Plotting constants
-class PlotConfig:
-    # Simple plot settings
-    SIMPLE_FIGSIZE = (12, 6)
-    SIMPLE_LINEWIDTH = 2
-    SIMPLE_DPI = 150
-
-    # Enhanced plot settings
-    ENHANCED_FIGSIZE = (16, 6)
-    ENHANCED_DPI = 300
-
-    # Marker and styling
-    DATA_MARKER = 'o'
-    OUTLIER_MARKER = 'x'
-    OUTLIER_COLOR = 'crimson'
-    TREND_COLOR = 'r'
-    TREND_LINESTYLE = '--'
-    TREND_ALPHA = 0.8
-    TREND_LINEWIDTH = 2.0
-
-    # Colors
-    LLM_LATENCY_COLOR = 'steelblue'
-    RUNTIME_COLOR = 'darkgreen'
-    SLA_COLOR = 'red'
-    NOTE_BOX_COLOR = 'mistyrose'
-    NOTE_TEXT_COLOR = 'crimson'
-    STATS_BOX_COLOR = 'lightblue'
-
-    # Alpha values
-    DATA_ALPHA = 0.7
-    OUTLIER_ALPHA = 0.9
-    GRID_ALPHA = 0.3
-    SLA_ALPHA = 0.7
-    NOTE_BOX_ALPHA = 0.7
-    STATS_BOX_ALPHA = 0.8
-
-    # Sizes
-    DATA_POINT_SIZE = 120
-    OUTLIER_POINT_SIZE = 140
-    DATA_LINEWIDTH = 1
-
-    # Font sizes
-    AXIS_LABEL_FONTSIZE = 12
-    TITLE_FONTSIZE = 14
-    LEGEND_FONTSIZE = 10
-    NOTE_FONTSIZE = 10
-    STATS_FONTSIZE = 10
-
-    # Text positioning
-    NOTE_X_POS = 0.98
-    NOTE_Y_POS = 0.02
-    STATS_X_POS = 0.02
-    STATS_Y_POS = 0.02
-
-    # Box styling
-    NOTE_BOX_PAD = 0.3
-    STATS_BOX_PAD = 0.5
-
-    # Trend line points
-    TREND_LINE_POINTS = 100
-
-    # Font weights
-    AXIS_LABEL_FONTWEIGHT = 'bold'
-    TITLE_FONTWEIGHT = 'bold'
-
-
-def plot_concurrency_vs_time_metrics_simple(df: pd.DataFrame, output_dir: Path) -> None:
-    """
-    Save a simple plot of concurrency vs. p95 LLM latency and workflow runtime.
-    """
-    plt.figure(figsize=PlotConfig.SIMPLE_FIGSIZE)
-    plt.plot(df["concurrency"],
-             df["llm_latency_p95"],
-             label="p95 LLM Latency (s)",
-             marker=PlotConfig.DATA_MARKER,
-             linewidth=PlotConfig.SIMPLE_LINEWIDTH)
-    plt.plot(df["concurrency"],
-             df["workflow_runtime_p95"],
-             label="p95 Workflow Runtime (s)",
-             marker="s",
-             linewidth=PlotConfig.SIMPLE_LINEWIDTH)
-    plt.xlabel("Concurrency")
-    plt.ylabel("Time (seconds)")
-    plt.title("Concurrency vs. p95 LLM Latency and Workflow Runtime")
-    plt.grid(True, alpha=PlotConfig.GRID_ALPHA)
-    plt.legend()
-    plt.tight_layout()
-
-    simple_plot_path = output_dir / "concurrency_vs_p95_simple.png"
-    plt.savefig(simple_plot_path, dpi=PlotConfig.SIMPLE_DPI, bbox_inches='tight')
-    plt.close()
-    logger.info("Simple plot saved to %s", simple_plot_path)
-
-
-def plot_metric_vs_concurrency_with_optional_fit(
-    ax: plt.Axes,
-    x: np.ndarray,
-    y: np.ndarray,
-    metric_name: str,
-    y_label: str,
-    title: str,
-    color: str,
-    sla_value: float = 0.0,
-    sla_label: str = None,
-    fit: LinearFitResult | None = None,
-):
-    """
-    Helper to plot a metric vs concurrency with pre-computed fit, outlier highlighting, and SLA line.
-    Requires pre-computed fit to be provided.
-    """
-    marker = PlotConfig.DATA_MARKER
-    outlier_marker = PlotConfig.OUTLIER_MARKER
-    outlier_color = PlotConfig.OUTLIER_COLOR
-    trend_color = PlotConfig.TREND_COLOR
-    trend_linestyle = PlotConfig.TREND_LINESTYLE
-    trend_alpha = PlotConfig.TREND_ALPHA
-    trend_linewidth = PlotConfig.TREND_LINEWIDTH
-    note_box_color = PlotConfig.NOTE_BOX_COLOR
-    note_text_color = PlotConfig.NOTE_TEXT_COLOR
-    legend_fontsize = PlotConfig.LEGEND_FONTSIZE
-    outliers_x = outliers_y = np.array([])
-    outliers_note = ""
-
-    # Skip analysis plot if no fit is available
-    if not fit:
-        logger.warning(f"No linear fit available for {metric_name}, skipping analysis plot")
-        return False
-
-    if fit.outliers_removed:
-        # Use the concurrencies that were removed to identify outlier points
-        outlier_mask = np.isin(x, fit.outliers_removed)
-        outliers_x = x[outlier_mask]
-        outliers_y = y[outlier_mask]
-        outliers_note = f"Outliers removed: concurrencies {fit.outliers_removed}"
-        # Plot cleaned data (points that weren't removed as outliers)
-        non_outlier_mask = ~np.isin(x, fit.outliers_removed)
-        x_clean = x[non_outlier_mask]
-        y_clean = y[non_outlier_mask]
-        ax.scatter(x_clean,
-                   y_clean,
-                   alpha=PlotConfig.DATA_ALPHA,
-                   s=PlotConfig.DATA_POINT_SIZE,
-                   c=color,
-                   edgecolors='white',
-                   linewidth=PlotConfig.DATA_LINEWIDTH,
-                   marker=marker,
-                   label='Data Points')
-        ax.scatter(outliers_x,
-                   outliers_y,
-                   alpha=PlotConfig.OUTLIER_ALPHA,
-                   s=PlotConfig.OUTLIER_POINT_SIZE,
-                   c=outlier_color,
-                   marker=outlier_marker,
-                   label='Removed Outliers')
-    else:
-        # No outliers plot all data points
-        ax.scatter(x,
-                   y,
-                   alpha=PlotConfig.DATA_ALPHA,
-                   s=PlotConfig.DATA_POINT_SIZE,
-                   c=color,
-                   edgecolors='white',
-                   linewidth=PlotConfig.DATA_LINEWIDTH,
-                   marker=marker,
-                   label='Data Points')
-
-    # Plot trend line using the fit
-    x_fit = np.linspace(x.min(), x.max(), PlotConfig.TREND_LINE_POINTS)
-    y_fit = fit.slope * x_fit + fit.intercept
-    ax.plot(x_fit,
-            y_fit,
-            trend_linestyle,
-            alpha=trend_alpha,
-            linewidth=trend_linewidth,
-            color=trend_color,
-            label=f'Trend (slope={fit.slope:.4f}, R²={fit.r_squared:.3f})')
-
-    if sla_value > 0:
-        ax.axhline(y=sla_value,
-                   color=PlotConfig.SLA_COLOR,
-                   linestyle=':',
-                   alpha=PlotConfig.SLA_ALPHA,
-                   linewidth=2,
-                   label=sla_label or f'SLA Threshold ({sla_value}s)')
-
-    ax.set_xlabel('Concurrency', fontsize=PlotConfig.AXIS_LABEL_FONTSIZE, fontweight=PlotConfig.AXIS_LABEL_FONTWEIGHT)
-    ax.set_ylabel(y_label, fontsize=PlotConfig.AXIS_LABEL_FONTSIZE, fontweight=PlotConfig.AXIS_LABEL_FONTWEIGHT)
-    ax.set_title(title, fontsize=PlotConfig.TITLE_FONTSIZE, fontweight=PlotConfig.TITLE_FONTWEIGHT)
-    ax.grid(True, alpha=PlotConfig.GRID_ALPHA)
-    ax.legend(fontsize=legend_fontsize)
-    if outliers_note:
-        ax.text(PlotConfig.NOTE_X_POS,
-                PlotConfig.NOTE_Y_POS,
-                outliers_note,
-                transform=ax.transAxes,
-                fontsize=PlotConfig.NOTE_FONTSIZE,
-                color=note_text_color,
-                ha='right',
-                va='bottom',
-                bbox=dict(boxstyle=f'round,pad={PlotConfig.NOTE_BOX_PAD}',
-                          facecolor=note_box_color,
-                          alpha=PlotConfig.NOTE_BOX_ALPHA))
-
-    return True
-
-
-def plot_concurrency_vs_time_metrics(metrics_per_concurrency: dict[int, SizingMetrics],
-                                     output_dir: Path,
-                                     target_llm_latency: float = 0.0,
-                                     target_runtime: float = 0.0,
-                                     llm_latency_fit: LinearFitResult | None = None,
-                                     runtime_fit: LinearFitResult | None = None) -> None:
-    """
-    Plot concurrency vs. p95 latency and workflow runtime using metrics_per_concurrency.
-    Enhanced with better styling, trend analysis, and annotations.
-    Only plots valid runs and requires pre-computed fits.
-    """
-    rows = []
-
-    for concurrency, metrics in metrics_per_concurrency.items():
-        llm_latency = metrics.llm_latency_p95
-        workflow_runtime = metrics.workflow_runtime_p95
-
-        rows.append({
-            "concurrency": concurrency, "llm_latency_p95": llm_latency, "workflow_runtime_p95": workflow_runtime
-        })
-
-    if not rows:
-        logger.warning("No valid metrics data available to plot.")
-        return
-
-    plt.style.use('seaborn-v0_8')
-    df = pd.DataFrame(rows).sort_values("concurrency")
-
-    # Always generate simple plot first
-    plot_concurrency_vs_time_metrics_simple(df, output_dir)
-
-    # Check if we have fits available for analysis plots
-    has_llm_latency_fit = llm_latency_fit is not None
-    has_runtime_fit = runtime_fit is not None
-
-    if not has_llm_latency_fit and not has_runtime_fit:
-        logger.warning("No linear fits available for analysis plots, skipping enhanced plot")
-        return
-
-    # Create subplots based on available fits
-    if has_llm_latency_fit and has_runtime_fit:
-        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=PlotConfig.ENHANCED_FIGSIZE)
-    else:
-        fig, ax1 = plt.subplots(1, 1, figsize=(8, 6))
-        ax2 = None
-
-    # Plot llm_latency if fit is available
-    llm_latency_plotted = False
-    if has_llm_latency_fit:
-        llm_latency_plotted = plot_metric_vs_concurrency_with_optional_fit(
-            ax1,
-            df["concurrency"].to_numpy(),
-            df["llm_latency_p95"].to_numpy(),
-            metric_name="llm_latency",
-            y_label='P95 LLM Latency (seconds)',
-            title='Concurrency vs P95 LLM Latency',
-            color=PlotConfig.LLM_LATENCY_COLOR,
-            sla_value=target_llm_latency,
-            sla_label=f'SLA Threshold ({target_llm_latency}s)' if target_llm_latency > 0 else None,
-            fit=llm_latency_fit,
-        )
-
-    # Plot runtime if fit is available
-    runtime_plotted = False
-    if has_runtime_fit and ax2 is not None:
-        runtime_plotted = plot_metric_vs_concurrency_with_optional_fit(
-            ax2,
-            df["concurrency"].to_numpy(),
-            df["workflow_runtime_p95"].to_numpy(),
-            metric_name="runtime",
-            y_label='P95 Workflow Runtime (seconds)',
-            title='Concurrency vs P95 Workflow Runtime',
-            color=PlotConfig.RUNTIME_COLOR,
-            sla_value=target_runtime,
-            sla_label=f'SLA Threshold ({target_runtime}s)' if target_runtime > 0 else None,
-            fit=runtime_fit,
-        )
-
-    # Check if any plots were successfully created
-    plots_created = (llm_latency_plotted or runtime_plotted)
-
-    if not plots_created:
-        logger.warning("No analysis plots could be created, skipping enhanced plot")
-        plt.close(fig)
-        return
-
-    # Add summary statistics
-    stats_text = f'Data Points: {len(df)}\n'
-    stats_text += f'LLM Latency Range: {df["llm_latency_p95"].min():.3f}-{df["llm_latency_p95"].max():.3f}s\n'
-    stats_text += f'WF Runtime Range: {df["workflow_runtime_p95"].min():.3f}-{df["workflow_runtime_p95"].max():.3f}s'
-
-    fig.text(PlotConfig.STATS_X_POS,
-             PlotConfig.STATS_Y_POS,
-             stats_text,
-             fontsize=PlotConfig.STATS_FONTSIZE,
-             bbox=dict(boxstyle=f'round,pad={PlotConfig.STATS_BOX_PAD}',
-                       facecolor=PlotConfig.STATS_BOX_COLOR,
-                       alpha=PlotConfig.STATS_BOX_ALPHA))
-
-    plt.tight_layout()
-    output_dir.mkdir(parents=True, exist_ok=True)
-
-    enhanced_plot_path = output_dir / "concurrency_vs_p95_analysis.png"
-    plt.savefig(enhanced_plot_path,
-                dpi=PlotConfig.ENHANCED_DPI,
-                bbox_inches='tight',
-                facecolor='white',
-                edgecolor='none')
-    plt.close()
-
-    logger.info("Enhanced plot saved to %s", enhanced_plot_path)

aiq/profiler/callbacks/agno_callback_handler.py

@@ -1,295 +0,0 @@
-# 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 logging
-import threading
-import time
-from collections.abc import Callable
-from typing import Any
-from uuid import uuid4
-
-import litellm
-
-from aiq.builder.context import AIQContext
-from aiq.builder.framework_enum import LLMFrameworkEnum
-from aiq.data_models.intermediate_step import IntermediateStepPayload
-from aiq.data_models.intermediate_step import IntermediateStepType
-from aiq.data_models.intermediate_step import StreamEventData
-from aiq.data_models.intermediate_step import TraceMetadata
-from aiq.data_models.intermediate_step import UsageInfo
-from aiq.profiler.callbacks.base_callback_class import BaseProfilerCallback
-from aiq.profiler.callbacks.token_usage_base_model import TokenUsageBaseModel
-
-logger = logging.getLogger(__name__)
-
-
-class AgnoProfilerHandler(BaseProfilerCallback):
-    """
-    A callback manager/handler for Agno that intercepts calls to:
-
-      - Tool execution
-      - LLM Calls
-
-    to collect usage statistics (tokens, inputs, outputs, time intervals, etc.)
-    and store them in AIQ Toolkit's usage_stats queue for subsequent analysis.
-    """
-
-    def __init__(self) -> None:
-        super().__init__()
-        self._lock = threading.Lock()
-        self.last_call_ts = time.time()
-        self.step_manager = AIQContext.get().intermediate_step_manager
-
-        # Original references to Agno methods (for uninstrumenting if needed)
-        self._original_tool_execute = None
-        self._original_llm_call = None
-
-    def instrument(self) -> None:
-        """
-        Monkey-patch the relevant Agno methods with usage-stat collection logic.
-        """
-        # Save the originals and apply patches
-        self._original_llm_call = getattr(litellm, "completion", None)
-
-        # Patch LLM completion if available
-        if self._original_llm_call:
-            litellm.completion = self._llm_call_monkey_patch()
-            logger.debug("AgnoProfilerHandler LLM call instrumentation applied successfully.")
-        else:
-            logger.debug("Could not patch Agno LLM calls: litellm.completion not found")
-
-        # Note: Agno doesn't have a class-based tool structure to patch directly.
-        # Instead, it uses decorators to convert functions to tools.
-        # In AIQ Toolkit, tool executions are captured at the execute_agno_tool level
-        # in packages/aiqtoolkit_agno/src/aiq/plugins/agno/tool_wrapper.py
-
-        # To properly monitor Agno tool executions, we would need to either:
-        # 1. Patch the execute_agno_tool function in tool_wrapper.py
-        # 2. Add explicit instrumentation in that function to push events to the step manager
-        # 3. Or, if Agno updates to have a class-based tool structure, update this handler
-        #    to patch those classes
-
-        # Recommended future enhancement:
-        # The execute_agno_tool function in packages/aiqtoolkit_agno/src/aiq/plugins/agno/tool_wrapper.py
-        # should be updated to directly push IntermediateStepPayload events to the step manager
-        # at the beginning and end of tool execution, similar to what this handler does for LLM calls.
-
-        logger.debug("AgnoProfilerHandler instrumentation completed.")
-
-    def _tool_execute_monkey_patch(self) -> Callable[..., Any]:
-        """
-        Returns a function that wraps tool execution calls with usage-logging.
-
-        Note: This method is currently not used in the instrument() function since
-        Agno doesn't have a class-based tool structure to patch. It's kept for
-        reference or future use if Agno changes its architecture.
-        """
-        original_func = self._original_tool_execute
-
-        def wrapped_tool_execute(*args, **kwargs) -> Any:
-            """
-            Collects usage stats for tool execution, calls the original, and captures output stats.
-            """
-            now = time.time()
-            tool_name = kwargs.get("tool_name", "")
-            uuid = str(uuid4())
-
-            try:
-                # Pre-call usage event
-                stats = IntermediateStepPayload(event_type=IntermediateStepType.TOOL_START,
-                                                framework=LLMFrameworkEnum.AGNO,
-                                                name=tool_name,
-                                                UUID=uuid,
-                                                data=StreamEventData(),
-                                                metadata=TraceMetadata(tool_inputs={
-                                                    "args": args, "kwargs": dict(kwargs)
-                                                }),
-                                                usage_info=UsageInfo(token_usage=TokenUsageBaseModel()))
-
-                self.step_manager.push_intermediate_step(stats)
-                self.last_call_ts = now
-
-                # Call the original execute
-                result = original_func(*args, **kwargs)
-                now = time.time()
-
-                # Post-call usage stats
-                usage_stat = IntermediateStepPayload(
-                    event_type=IntermediateStepType.TOOL_END,
-                    span_event_timestamp=now,
-                    framework=LLMFrameworkEnum.AGNO,
-                    name=tool_name,
-                    UUID=uuid,
-                    data=StreamEventData(input={
-                        "args": args, "kwargs": dict(kwargs)
-                    }, output=str(result)),
-                    metadata=TraceMetadata(tool_outputs={"result": str(result)}),
-                    usage_info=UsageInfo(token_usage=TokenUsageBaseModel()),
-                )
-
-                self.step_manager.push_intermediate_step(usage_stat)
-                return result
-
-            except Exception as e:
-                logger.exception("Tool execution error: %s", e)
-                raise
-
-        return wrapped_tool_execute
-
-    def _llm_call_monkey_patch(self) -> Callable[..., Any]:
-        """
-        Returns a function that wraps calls to litellm.completion(...) with usage-logging.
-        """
-        original_func = self._original_llm_call
-
-        def wrapped_llm_call(*args, **kwargs) -> Any:
-            """
-            Collects usage stats for LLM calls, calls the original, and captures output stats.
-            """
-            now = time.time()
-            seconds_between_calls = int(now - self.last_call_ts)
-            model_name = kwargs.get('model', "")
-
-            model_input = ""
-            try:
-                for message in kwargs.get('messages', []):
-                    model_input += message.get('content', "")
-            except Exception as e:
-                logger.exception("Error getting model input: %s", e)
-
-            uuid = str(uuid4())
-
-            # Record the start event
-            input_stats = IntermediateStepPayload(
-                event_type=IntermediateStepType.LLM_START,
-                framework=LLMFrameworkEnum.AGNO,
-                name=model_name,
-                UUID=uuid,
-                data=StreamEventData(input=model_input),
-                metadata=TraceMetadata(chat_inputs=copy.deepcopy(kwargs.get('messages', []))),
-                usage_info=UsageInfo(token_usage=TokenUsageBaseModel(),
-                                     num_llm_calls=1,
-                                     seconds_between_calls=seconds_between_calls))
-
-            self.step_manager.push_intermediate_step(input_stats)
-
-            # Verify we have a valid original function before calling it
-            if original_func is None:
-                logger.error("Original litellm.completion function is None - cannot call it")
-                output = None
-            else:
-                # Call the original litellm.completion(...)
-                logger.debug(
-                    f"Calling litellm.completion for {model_name} with {len(args)} args and {len(kwargs)} kwargs")
-                try:
-                    output = original_func(*args, **kwargs)
-                    logger.debug(f"Original litellm.completion returned: {type(output)}")
-                except Exception as e:
-                    logger.exception(f"Error calling original litellm.completion: {e}")
-                    output = None
-
-            # Initialize default values
-            model_output = ""
-            chat_responses = None
-            token_usage = TokenUsageBaseModel()
-
-            # Log what we received to help with debugging
-            logger.debug(f"LLM call to {model_name} received output type: {type(output)}")
-
-            # Safely process the output if it's not None
-            if output is not None:
-                try:
-                    # Extract model output text from choices
-                    if hasattr(output, 'choices') and output.choices:
-                        logger.debug(f"Output has {len(output.choices)} choices")
-                        for i, choice in enumerate(output.choices):
-                            logger.debug(f"Processing choice {i} of type {type(choice)}")
-                            if hasattr(choice, 'model_extra') and 'message' in choice.model_extra:
-                                msg = choice.model_extra["message"]
-                                content = msg.get('content', "")
-                                logger.debug(f"Got content from model_extra.message: {content[:50]}...")
-                                model_output += content
-                            elif hasattr(choice, 'message') and hasattr(choice.message, 'content'):
-                                content = choice.message.content or ""
-                                logger.debug(f"Got content from message.content: {content[:50]}...")
-                                model_output += content
-                            else:
-                                logger.debug(f"Could not extract content from choice: {choice}")
-
-                    # Try to get chat responses
-                    if hasattr(output, 'choices') and len(output.choices) > 0:
-                        choice = output.choices[0]
-                        if hasattr(choice, 'model_dump'):
-                            logger.debug("Using model_dump to extract chat responses")
-                            chat_responses = choice.model_dump()
-                        else:
-                            # Fall back to a simpler representation
-                            logger.debug("Falling back to simple representation for chat responses")
-                            chat_responses = {"content": model_output}
-
-                    # Try to get token usage
-                    if hasattr(output, 'model_extra') and 'usage' in output.model_extra:
-                        usage_data = output.model_extra['usage']
-                        logger.debug(f"Found usage data of type {type(usage_data)}")
-
-                        # Special debug for the test case
-                        if hasattr(usage_data, 'prompt_tokens'
-                                   ) and usage_data.prompt_tokens == 20 and usage_data.completion_tokens == 15:
-                            logger.debug("Found test case token usage object with 20/15/35 tokens")
-
-                        if hasattr(usage_data, 'model_dump'):
-                            logger.debug("Using model_dump to extract token usage")
-                            token_usage = TokenUsageBaseModel(**usage_data.model_dump())
-                        elif isinstance(usage_data, dict):
-                            logger.debug("Extracting token usage from dictionary")
-                            token_usage = TokenUsageBaseModel(prompt_tokens=usage_data.get('prompt_tokens', 0),
-                                                              completion_tokens=usage_data.get('completion_tokens', 0),
-                                                              total_tokens=usage_data.get('total_tokens', 0))
-                        elif isinstance(usage_data, TokenUsageBaseModel):
-                            # If it's already a TokenUsageBaseModel instance, use it directly
-                            logger.debug("Using TokenUsageBaseModel directly")
-                            token_usage = usage_data
-                        elif hasattr(usage_data, 'prompt_tokens') and hasattr(
-                                usage_data, 'completion_tokens') and hasattr(usage_data, 'total_tokens'):
-                            # For objects that have the needed properties but aren't TokenUsageBaseModel
-                            logger.debug("Using object with token properties")
-                            token_usage = TokenUsageBaseModel(prompt_tokens=usage_data.prompt_tokens,
-                                                              completion_tokens=usage_data.completion_tokens,
-                                                              total_tokens=usage_data.total_tokens)
-
-                        logger.debug(f"Final token usage: prompt={token_usage.prompt_tokens}, "
-                                     f"completion={token_usage.completion_tokens}, "
-                                     f"total={token_usage.total_tokens}")
-                except Exception as e:
-                    logger.exception("Error getting model output: %s", e)
-
-            now = time.time()
-            # Record the end event
-            output_stats = IntermediateStepPayload(event_type=IntermediateStepType.LLM_END,
-                                                   span_event_timestamp=now,
-                                                   framework=LLMFrameworkEnum.AGNO,
-                                                   name=model_name,
-                                                   UUID=uuid,
-                                                   data=StreamEventData(input=model_input, output=model_output),
-                                                   metadata=TraceMetadata(chat_responses=chat_responses),
-                                                   usage_info=UsageInfo(token_usage=token_usage,
-                                                                        num_llm_calls=1,
-                                                                        seconds_between_calls=seconds_between_calls))
-
-            self.step_manager.push_intermediate_step(output_stats)
-            return output
-
-        return wrapped_llm_call

aiq/profiler/callbacks/base_callback_class.py

@@ -1,20 +0,0 @@
-# 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 abc import ABC
-
-
-class BaseProfilerCallback(ABC):
-    pass