aiqtoolkit 1.1.0a20250429__py3-none-any.whl

aiq/observability/register.py

@@ -0,0 +1,97 @@
+# 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 pydantic import Field
+
+from aiq.builder.builder import Builder
+from aiq.cli.register_workflow import register_logging_method
+from aiq.cli.register_workflow import register_telemetry_exporter
+from aiq.data_models.logging import LoggingBaseConfig
+from aiq.data_models.telemetry_exporter import TelemetryExporterBaseConfig
+
+logger = logging.getLogger(__name__)
+
+
+class PhoenixTelemetryExporter(TelemetryExporterBaseConfig, name="phoenix"):
+    """A telemetry exporter to transmit traces to externally hosted phoenix service."""
+
+    endpoint: str = Field(description="The phoenix endpoint to export telemetry traces.")
+    project: str = Field(description="The project name to group the telemetry traces.")
+
+
+@register_telemetry_exporter(config_type=PhoenixTelemetryExporter)
+async def phoenix_telemetry_exporter(config: PhoenixTelemetryExporter, builder: Builder):
+
+    from phoenix.otel import HTTPSpanExporter
+    try:
+        yield HTTPSpanExporter(endpoint=config.endpoint)
+    except ConnectionError as ex:
+        logger.warning("Unable to connect to Phoenix at port 6006. Are you sure Phoenix is running?\n %s",
+                       ex,
+                       exc_info=True)
+    except Exception as ex:
+        logger.error("Error in Phoenix telemetry Exporter\n %s", ex, exc_info=True)
+
+
+class OtelCollectorTelemetryExporter(TelemetryExporterBaseConfig, name="otelcollector"):
+    """A telemetry exporter to transmit traces to externally hosted otel collector service."""
+
+    endpoint: str = Field(description="The otel endpoint to export telemetry traces.")
+    project: str = Field(description="The project name to group the telemetry traces.")
+
+
+@register_telemetry_exporter(config_type=OtelCollectorTelemetryExporter)
+async def otel_telemetry_exporter(config: OtelCollectorTelemetryExporter, builder: Builder):
+
+    from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
+
+    yield OTLPSpanExporter(endpoint=config.endpoint)
+
+
+class ConsoleLoggingMethod(LoggingBaseConfig, name="console"):
+    """A logger to write runtime logs to the console."""
+
+    level: str = Field(description="The logging level of console logger.")
+
+
+@register_logging_method(config_type=ConsoleLoggingMethod)
+async def console_logging_method(config: ConsoleLoggingMethod, builder: Builder):
+    """
+        Build and return a StreamHandler for console-based logging.
+        """
+    level = getattr(logging, config.level.upper(), logging.INFO)
+    handler = logging.StreamHandler()
+    handler.setLevel(level)
+    yield handler
+
+
+class FileLoggingMethod(LoggingBaseConfig, name="file"):
+    """A logger to write runtime logs to a file."""
+
+    path: str = Field(description="The file path to save the logging output.")
+    level: str = Field(description="The logging level of file logger.")
+
+
+@register_logging_method(config_type=FileLoggingMethod)
+async def file_logging_method(config: FileLoggingMethod, builder: Builder):
+    """
+        Build and return a FileHandler for file-based logging.
+        """
+    level = getattr(logging, config.level.upper(), logging.INFO)
+    handler = logging.FileHandler(filename=config.path, mode="a", encoding="utf-8")
+    handler.setLevel(level)
+    yield handler

aiq/plugins/.namespace

@@ -0,0 +1 @@
+Note: This is a python namespace package and this directory should remain empty. Do NOT add a `__init__.py` file or any other files to this directory. This file is also needed to ensure the directory exists in git.

aiq/profiler/callbacks/agno_callback_handler.py

@@ -0,0 +1,295 @@
+# 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 AgentIQ'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 AgentIQ, 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

@@ -0,0 +1,20 @@
+# 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