aiqtoolkit 1.2.0rc4__py3-none-any.whl → 1.2.0rc5__py3-none-any.whl

aiq/profiler/callbacks/semantic_kernel_callback_handler.py

@@ -1,238 +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 semantic_kernel
-from pydantic import BaseModel
-from pydantic import Field
-from semantic_kernel.connectors.ai.open_ai.services.open_ai_chat_completion_base import OpenAIChatCompletionBase
-
-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 SemanticKernelPatchMethod(BaseModel):
-    """
-    Stores the module and function to patch in Semantic Kernel.
-    """
-    module: Any = Field(..., description="The module to patch")
-    function: str = Field(..., description="The function to patch")
-
-
-class SemanticKernelProfilerHandler(BaseProfilerCallback):
-    """
-    A callback manager/handler for Msft Semantic Kernel that intercepts calls to:
-
-    - Chat Completions Endpoints
-    - Tool 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, workflow_llms: dict) -> None:
-        super().__init__()
-        self._lock = threading.Lock()
-        self.last_call_ts = time.time()
-        self.step_manager = AIQContext.get().intermediate_step_manager
-        self._builder_llms = workflow_llms
-
-        # Original references to SK methods
-        self._original_tool_call = None
-
-        # Store a mapping of current SK methods we support patching for
-        self._patch_methods = {
-            "openai_streaming":
-                SemanticKernelPatchMethod(module=OpenAIChatCompletionBase,
-                                          function="_inner_get_streaming_chat_message_contents"),
-            "openai_non_streaming":
-                SemanticKernelPatchMethod(module=OpenAIChatCompletionBase, function="_inner_get_chat_message_contents")
-        }
-
-    def instrument(self) -> None:
-        """
-        Monkey-patch the relevant Semantic Kernel methods with usage-stat collection logic.
-        """
-
-        functions_to_patch = []
-
-        # Gather the appropriate modules/functions based on your builder config
-        for llm in self._builder_llms:
-            if self._builder_llms[llm].provider_type == 'openai':  # pylint: disable=consider-using-in
-                functions_to_patch.extend(["openai_non_streaming", "openai_streaming"])
-
-        # Grab original reference for the tool call
-        self._original_tool_call = getattr(semantic_kernel.Kernel, "invoke_function_call", None)
-
-        # Now do direct monkey-patching: replace each function with a closure
-        for method in functions_to_patch:
-            patch_method = self._patch_methods[method]
-            setattr(patch_method.module,
-                    patch_method.function,
-                    self._build_llm_call_patch(getattr(patch_method.module, patch_method.function)))
-
-        if self._original_tool_call:
-            patched_tool_call = self._build_tool_call_patch(self._original_tool_call)
-            setattr(semantic_kernel.Kernel, "invoke_function_call", patched_tool_call)
-
-        logger.debug("SemanticKernelProfilerHandler instrumentation applied successfully.")
-
-    def _build_llm_call_patch(self, original_func: Callable[..., Any]) -> Callable[..., Any]:
-        """
-        Returns an async monkey-patch that wraps the original chat-completion method.
-        Replicates the usage collection from _llm_call_wrapper.
-        """
-
-        async def patched_llm_call(sk_self, *args, **kwargs) -> Any:
-            now = time.time()
-            seconds_between_calls = int(now - self.last_call_ts)
-            uuid = str(uuid4())
-
-            # Build the input stats
-            if args:
-                chat_input = [copy.deepcopy(args[0].model_dump())]
-            else:
-                # if no args, fallback on kwargs["chat_history"]
-                chat_input = [kwargs["chat_history"].model_dump()]
-
-            model_name = sk_self.ai_model_id
-
-            model_input = ""
-            try:
-                for message in chat_input[0]["messages"]:
-                    for item in message["items"]:
-                        if "text" in item:
-                            model_input += item["text"]
-            except Exception as e:
-                logger.exception("Error in getting model input: %s", e, exc_info=True)
-
-            input_stats = IntermediateStepPayload(event_type=IntermediateStepType.LLM_START,
-                                                  framework=LLMFrameworkEnum.SEMANTIC_KERNEL,
-                                                  name=model_name,
-                                                  UUID=uuid,
-                                                  data=StreamEventData(input=model_input),
-                                                  metadata=TraceMetadata(chat_inputs=copy.deepcopy(chat_input)),
-                                                  usage_info=UsageInfo(token_usage=TokenUsageBaseModel(),
-                                                                       num_llm_calls=1,
-                                                                       seconds_between_calls=seconds_between_calls))
-
-            self.step_manager.push_intermediate_step(input_stats)
-
-            # Call the original method
-            output = await original_func(sk_self, *args, **kwargs)
-
-            model_output = output[0].content
-            now = time.time()
-            # Build the output stats
-            output_stats = IntermediateStepPayload(
-                event_type=IntermediateStepType.LLM_END,
-                span_event_timestamp=now,
-                framework=LLMFrameworkEnum.SEMANTIC_KERNEL,
-                name=model_name,
-                UUID=uuid,
-                data=StreamEventData(input=model_input, output=model_output),
-                metadata=TraceMetadata(chat_responses=output[0].model_dump()),
-                usage_info=UsageInfo(token_usage=TokenUsageBaseModel(**output[0].metadata["usage"].model_dump())))
-
-            self.step_manager.push_intermediate_step(output_stats)
-
-            # Update last_call_ts
-            self.last_call_ts = time.time()
-
-            return output
-
-        return patched_llm_call
-
-    def _build_tool_call_patch(self, original_func: Callable[..., Any]) -> Callable[..., Any]:
-        """
-        Returns an async monkey-patch that wraps the original tool call (invoke_function_call).
-        Replicates usage collection from _tool_use_wrapper.
-        """
-
-        async def patched_tool_call(kernel_self, *args, **kwargs) -> Any:
-
-            uuid = str(uuid4())
-            now = time.time()
-            # Extract the tool input
-            if kwargs:
-                tool_input = kwargs["function_call"].model_dump(exclude="content_type")
-            else:
-                tool_input = args[0].model_dump(exclude="content_type")
-
-            try:
-                # Pre-call usage event
-                input_stat = IntermediateStepPayload(event_type=IntermediateStepType.TOOL_START,
-                                                     framework=LLMFrameworkEnum.SEMANTIC_KERNEL,
-                                                     name=tool_input["name"],
-                                                     UUID=uuid,
-                                                     data=StreamEventData(input=tool_input),
-                                                     metadata=TraceMetadata(tool_inputs=copy.deepcopy(tool_input),
-                                                                            tool_info=copy.deepcopy(tool_input)),
-                                                     usage_info=UsageInfo(token_usage=TokenUsageBaseModel()))
-
-                self.step_manager.push_intermediate_step(input_stat)
-                now = time.time()
-                # Call the original invoke_function_call
-                result = await original_func(kernel_self, *args, **kwargs)
-
-                # Try to get the chat history from kwargs or args
-                if kwargs:
-                    chat_history = copy.deepcopy(kwargs["chat_history"])
-                else:
-                    chat_history = copy.deepcopy(args[1])
-
-                # Post-call usage event
-                output_stat = IntermediateStepPayload(event_type=IntermediateStepType.TOOL_END,
-                                                      span_event_timestamp=now,
-                                                      framework=LLMFrameworkEnum.SEMANTIC_KERNEL,
-                                                      name=tool_input["name"],
-                                                      UUID=uuid,
-                                                      data=StreamEventData(input=tool_input,
-                                                                           output=[
-                                                                               item.model_dump(exclude="content_type")
-                                                                               for item in chat_history[-1].items
-                                                                           ]),
-                                                      metadata=TraceMetadata(tool_outputs=[
-                                                          item.model_dump(exclude="content_type")
-                                                          for item in chat_history[-1].items
-                                                      ],
-                                                                             tool_info=copy.deepcopy(tool_input)),
-                                                      usage_info=UsageInfo(token_usage=TokenUsageBaseModel()))
-
-                self.step_manager.push_intermediate_step(output_stat)
-
-                return result
-
-            except Exception as e:
-                logger.exception("ToolUsage._use error: %s", e)
-                raise
-
-        return patched_tool_call

aiq/profiler/callbacks/token_usage_base_model.py

@@ -1,27 +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 pydantic import BaseModel
-from pydantic import Field
-
-
-class TokenUsageBaseModel(BaseModel):
-    """
-    Base model for token usage callbacks.
-    """
-
-    prompt_tokens: int = Field(default=0, description="Number of tokens in the prompt.")
-    completion_tokens: int = Field(default=0, description="Number of tokens in the completion.")
-    total_tokens: int = Field(default=0, description="Number of tokens total.")

aiq/profiler/data_frame_row.py

@@ -1,51 +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 typing import Any
-
-from pydantic import BaseModel
-from pydantic import ConfigDict
-from pydantic import field_validator
-
-
-class DataFrameRow(BaseModel):
-    model_config = ConfigDict(extra="allow", validate_assignment=True)
-
-    event_type: Any
-    event_timestamp: float | None
-    example_number: int | None
-    prompt_tokens: int | None
-    completion_tokens: int | None
-    total_tokens: int | None
-    llm_text_input: str | None
-    llm_text_output: str | None
-    llm_new_token: str | None
-    llm_name: str | None
-    tool_name: str | None
-    function_name: str | None
-    function_id: str | None
-    parent_function_name: str | None
-    parent_function_id: str | None
-    UUID: str | None
-    framework: str | None
-
-    @field_validator('llm_text_input', 'llm_text_output', 'llm_new_token', mode='before')
-    def cast_to_str(cls, v):  # pylint: disable=no-self-argument
-        if v is None:
-            return v
-        try:
-            return str(v)
-        except Exception as e:
-            raise ValueError(f"Value {v} cannot be cast to str: {e}") from e

aiq/profiler/data_models.py

@@ -1,24 +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 pydantic import BaseModel
-
-from aiq.profiler.inference_metrics_model import InferenceMetricsModel
-from aiq.profiler.inference_optimization.data_models import WorkflowRuntimeMetrics
-
-
-class ProfilerResults(BaseModel):
-    workflow_runtime_metrics: WorkflowRuntimeMetrics | None = None
-    llm_latency_ci: InferenceMetricsModel | None = None

aiq/profiler/decorators/framework_wrapper.py

@@ -1,131 +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.
-
-# pylint disable=ungrouped-imports
-
-from __future__ import annotations
-
-import functools
-import logging
-from collections.abc import Callable
-from contextlib import AbstractAsyncContextManager as AsyncContextManager
-from contextlib import asynccontextmanager
-from contextvars import ContextVar
-from typing import Any
-
-from aiq.builder.framework_enum import LLMFrameworkEnum
-
-logger = logging.getLogger(__name__)
-
-_library_instrumented = {
-    "langchain": False,
-    "crewai": False,
-    "semantic_kernel": False,
-    "agno": False,
-}
-
-callback_handler_var: ContextVar[Any | None] = ContextVar("callback_handler_var", default=None)
-
-
-def set_framework_profiler_handler(
-    workflow_llms: dict = None,
-    frameworks: list[LLMFrameworkEnum] = None,
-) -> Callable[[Callable[..., AsyncContextManager[Any]]], Callable[..., AsyncContextManager[Any]]]:
-    """
-    Decorator that wraps an async context manager function to set up framework-specific profiling.
-    """
-
-    def decorator(func: Callable[..., AsyncContextManager[Any]]) -> Callable[..., AsyncContextManager[Any]]:
-
-        @functools.wraps(func)
-        @asynccontextmanager
-        async def wrapper(workflow_config, builder):
-
-            if LLMFrameworkEnum.LANGCHAIN in frameworks and not _library_instrumented["langchain"]:
-                from langchain_core.tracers.context import register_configure_hook
-
-                from aiq.profiler.callbacks.langchain_callback_handler import LangchainProfilerHandler
-
-                handler = LangchainProfilerHandler()
-                callback_handler_var.set(handler)
-                register_configure_hook(callback_handler_var, inheritable=True)
-                _library_instrumented["langchain"] = True
-                logger.debug("Langchain callback handler registered")
-
-            if LLMFrameworkEnum.LLAMA_INDEX in frameworks:
-                from llama_index.core import Settings
-                from llama_index.core.callbacks import CallbackManager
-
-                from aiq.profiler.callbacks.llama_index_callback_handler import LlamaIndexProfilerHandler
-
-                handler = LlamaIndexProfilerHandler()
-                Settings.callback_manager = CallbackManager([handler])
-                logger.debug("LlamaIndex callback handler registered")
-
-            if LLMFrameworkEnum.CREWAI in frameworks and not _library_instrumented["crewai"]:
-                from aiq.plugins.crewai.crewai_callback_handler import \
-                    CrewAIProfilerHandler  # pylint: disable=ungrouped-imports,line-too-long  # noqa E501
-
-                handler = CrewAIProfilerHandler()
-                handler.instrument()
-                _library_instrumented["crewai"] = True
-                logger.debug("CrewAI callback handler registered")
-
-            if LLMFrameworkEnum.SEMANTIC_KERNEL in frameworks and not _library_instrumented["semantic_kernel"]:
-                from aiq.profiler.callbacks.semantic_kernel_callback_handler import SemanticKernelProfilerHandler
-
-                handler = SemanticKernelProfilerHandler(workflow_llms=workflow_llms)
-                handler.instrument()
-                _library_instrumented["semantic_kernel"] = True
-                logger.debug("SemanticKernel callback handler registered")
-
-            if LLMFrameworkEnum.AGNO in frameworks and not _library_instrumented["agno"]:
-                from aiq.profiler.callbacks.agno_callback_handler import AgnoProfilerHandler
-
-                handler = AgnoProfilerHandler()
-                handler.instrument()
-                _library_instrumented["agno"] = True
-                logger.info("Agno callback handler registered")
-
-            # IMPORTANT: actually call the wrapped function as an async context manager
-            async with func(workflow_config, builder) as result:
-                yield result
-
-        return wrapper
-
-    return decorator
-
-
-def chain_wrapped_build_fn(
-    original_build_fn: Callable[..., AsyncContextManager],
-    workflow_llms: dict,
-    function_frameworks: list[LLMFrameworkEnum],
-) -> Callable[..., AsyncContextManager]:
-    """
-    Convert an original build function into an async context manager that
-    wraps it with a single call to set_framework_profiler_handler, passing
-    all frameworks at once.
-    """
-
-    # Define a base async context manager that simply calls the original build function.
-    @asynccontextmanager
-    async def base_fn(*args, **kwargs):
-        async with original_build_fn(*args, **kwargs) as w:
-            yield w
-
-    # Instead of wrapping iteratively, we now call the decorator once,
-    # passing the entire list of frameworks along with the workflow_llms.
-    wrapped_fn = set_framework_profiler_handler(workflow_llms, function_frameworks)(base_fn)
-    return wrapped_fn