aiqtoolkit 1.1.0rc5__py3-none-any.whl → 1.2.0__py3-none-any.whl

aiq/profiler/callbacks/llama_index_callback_handler.py

@@ -1,205 +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 __future__ import annotations
-
-import copy
-import logging
-import threading
-import time
-from typing import Any
-
-from llama_index.core.callbacks import CBEventType
-from llama_index.core.callbacks import EventPayload
-from llama_index.core.callbacks.base_handler import BaseCallbackHandler
-from llama_index.core.llms import ChatResponse
-
-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 LlamaIndexProfilerHandler(BaseCallbackHandler, BaseProfilerCallback):
-    """
-    A callback handler for LlamaIndex that tracks usage stats similarly to NIMCallbackHandler.
-    Collects:
-
-    - Prompts
-    - Token usage
-    - Response data
-    - Time intervals between calls
-
-    and appends them to AIQContextState.usage_stats.
-    """
-
-    def __init__(self) -> None:
-        BaseCallbackHandler.__init__(self, event_starts_to_ignore=[], event_ends_to_ignore=[])
-        BaseProfilerCallback.__init__(self)
-        self._lock = threading.Lock()
-        self.last_call_ts = time.time()
-        self._last_tool_map: dict[str, str] = {}
-        self.step_manager = AIQContext.get().intermediate_step_manager
-
-        self._run_id_to_llm_input = {}
-        self._run_id_to_tool_input = {}
-        self._run_id_to_timestamp = {}
-
-    def on_event_start(
-        self,
-        event_type: CBEventType,
-        payload: dict[str, Any] | None = None,
-        event_id: str = "",
-        parent_id: str = "",
-        **kwargs: Any,
-    ) -> str:
-        """
-        Called at the *start* of a LlamaIndex "event" (LLM call, Embedding, etc.).
-        We capture the prompts or query strings here, if any.
-        """
-
-        prompts_or_messages = None
-        now = time.time()
-        seconds_between_calls = int(now - self.last_call_ts)
-
-        # For LLM or chat calls, look in `payload` for messages/prompts
-        if event_type == CBEventType.LLM and payload:
-            # For example, "PROMPT" or "MESSAGES" might be in the payload.
-            # If found, store them in usage stats (just like your NIMCallbackHandler).
-            if EventPayload.PROMPT in payload:
-                prompts_or_messages = [payload[EventPayload.PROMPT]]
-            elif EventPayload.MESSAGES in payload:
-                prompts_or_messages = [str(msg) for msg in payload[EventPayload.MESSAGES]]
-
-            model_name = ""
-            try:
-                model_name = payload.get(EventPayload.SERIALIZED)['model']
-            except Exception as e:
-                logger.exception("Error getting model name: %s", e, exc_info=True)
-
-            llm_text_input = " ".join(prompts_or_messages) if prompts_or_messages else ""
-
-            if prompts_or_messages:
-                stats = IntermediateStepPayload(event_type=IntermediateStepType.LLM_START,
-                                                framework=LLMFrameworkEnum.LLAMA_INDEX,
-                                                name=model_name,
-                                                UUID=event_id,
-                                                data=StreamEventData(input=llm_text_input),
-                                                metadata=TraceMetadata(chat_inputs=copy.deepcopy(prompts_or_messages)),
-                                                usage_info=UsageInfo(token_usage=TokenUsageBaseModel(),
-                                                                     num_llm_calls=1,
-                                                                     seconds_between_calls=seconds_between_calls))
-
-                self.step_manager.push_intermediate_step(stats)
-                self._run_id_to_llm_input[event_id] = llm_text_input
-                self.last_call_ts = now
-                self._run_id_to_timestamp[event_id] = time.time()
-
-        elif event_type == CBEventType.FUNCTION_CALL and payload:
-            tool_metadata = payload.get(EventPayload.TOOL)
-            tool_metadata = {
-                "description": tool_metadata.description if hasattr(tool_metadata, "description") else "",
-                "fn_schema_str": tool_metadata.fn_schema_str if hasattr(tool_metadata, "fn_schema_str") else "",
-                "name": tool_metadata.name if hasattr(tool_metadata, "name") else "",
-            }
-            stats = IntermediateStepPayload(
-                event_type=IntermediateStepType.TOOL_START,
-                framework=LLMFrameworkEnum.LLAMA_INDEX,
-                name=payload.get(EventPayload.TOOL).name,
-                UUID=event_id,
-                data=StreamEventData(input=copy.deepcopy(payload.get(EventPayload.FUNCTION_CALL))),
-                metadata=TraceMetadata(tool_inputs=copy.deepcopy(payload.get(EventPayload.FUNCTION_CALL)),
-                                       tool_info=tool_metadata),
-                usage_info=UsageInfo(token_usage=TokenUsageBaseModel()))
-
-            self._run_id_to_tool_input[event_id] = copy.deepcopy(payload.get(EventPayload.FUNCTION_CALL))
-            self._last_tool_map[event_id] = payload.get(EventPayload.TOOL).name
-            self.step_manager.push_intermediate_step(stats)
-            self._run_id_to_timestamp[event_id] = time.time()
-        return event_id  # must return the event_id
-
-    def on_event_end(
-        self,
-        event_type: CBEventType,
-        payload: dict[str, Any] | None = None,
-        event_id: str = "",
-        **kwargs: Any,
-    ) -> None:
-        """
-        Called at the *end* of a LlamaIndex "event".
-        We collect token usage (if available) and the returned response text.
-        """
-
-        if payload and event_type == CBEventType.LLM:
-            # Often, token usage is embedded in e.g. payload["RESPONSE"].raw["usage"] for OpenAI-based calls
-            response = payload.get(EventPayload.RESPONSE)
-            if isinstance(response, ChatResponse):
-                llm_text_output = ""
-
-                try:
-                    for block in response.message.blocks:
-                        llm_text_output += block.text
-                except Exception as e:
-                    logger.exception("Error getting LLM text output: %s", e, exc_info=True)
-
-                model_name = ""
-                try:
-                    model_name = response.raw.model
-                except Exception as e:
-                    logger.exception("Error getting model name: %s", e, exc_info=True)
-
-                # Append usage data to AIQ Toolkit usage stats
-                with self._lock:
-                    stats = IntermediateStepPayload(
-                        event_type=IntermediateStepType.LLM_END,
-                        span_event_timestamp=self._run_id_to_timestamp.get(event_id),
-                        framework=LLMFrameworkEnum.LLAMA_INDEX,
-                        name=model_name,
-                        UUID=event_id,
-                        data=StreamEventData(input=self._run_id_to_llm_input.get(event_id), output=llm_text_output),
-                        metadata=TraceMetadata(chat_responses=response.message if response.message else None),
-                        usage_info=UsageInfo(token_usage=TokenUsageBaseModel(**response.additional_kwargs)))
-                    self.step_manager.push_intermediate_step(stats)
-
-        elif event_type == CBEventType.FUNCTION_CALL and payload:
-            stats = IntermediateStepPayload(
-                event_type=IntermediateStepType.TOOL_END,
-                span_event_timestamp=self._run_id_to_timestamp.get(event_id),
-                framework=LLMFrameworkEnum.LLAMA_INDEX,
-                name=self._last_tool_map.get(event_id),
-                UUID=event_id,
-                data=StreamEventData(output=copy.deepcopy(payload.get(EventPayload.FUNCTION_OUTPUT))),
-                usage_info=UsageInfo(token_usage=TokenUsageBaseModel()))
-
-            self.step_manager.push_intermediate_step(stats)
-
-    def start_trace(self, trace_id: str | None = None) -> None:
-        """Run when an overall trace is launched."""
-        pass
-
-    def end_trace(
-        self,
-        trace_id: str | None = None,
-        trace_map: dict[str, list[str]] | None = None,
-    ) -> None:
-        """Run when an overall trace is exited."""
-        pass

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/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