nvidia-nat 1.2.0rc5__py3-none-any.whl

aiq/profiler/callbacks/langchain_callback_handler.py

@@ -0,0 +1,290 @@
+# 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 uuid import UUID
+from uuid import uuid4
+
+from langchain_core.callbacks import AsyncCallbackHandler
+from langchain_core.messages import AIMessage
+from langchain_core.messages import BaseMessage
+from langchain_core.outputs import ChatGeneration
+from langchain_core.outputs import LLMResult
+
+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 ToolSchema
+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__)
+
+
+def _extract_tools_schema(invocation_params: dict) -> list:
+
+    tools_schema = []
+    if invocation_params is not None:
+        for tool in invocation_params.get("tools", []):
+            tools_schema.append(ToolSchema(**tool))
+
+    return tools_schema
+
+
+class LangchainProfilerHandler(AsyncCallbackHandler, BaseProfilerCallback):  # pylint: disable=R0901
+    """Callback Handler that tracks NIM info."""
+
+    total_tokens: int = 0
+    prompt_tokens: int = 0
+    completion_tokens: int = 0
+    successful_requests: int = 0
+    raise_error = True  # Override to raise error and run inline
+    run_inline = True
+
+    def __init__(self) -> None:
+        super().__init__()
+        self._lock = threading.Lock()
+        self.last_call_ts = time.time()
+
+        self.step_manager = AIQContext.get().intermediate_step_manager
+        self._state = IntermediateStepType.LLM_END
+
+        self._run_id_to_model_name = {}
+        self._run_id_to_llm_input = {}
+        self._run_id_to_tool_input = {}
+        self._run_id_to_start_time = {}
+
+    def __repr__(self) -> str:
+        return (f"Tokens Used: {self.total_tokens}\n"
+                f"\tPrompt Tokens: {self.prompt_tokens}\n"
+                f"\tCompletion Tokens: {self.completion_tokens}\n"
+                f"Successful Requests: {self.successful_requests}\n")
+
+    @property
+    def always_verbose(self) -> bool:
+        """Whether to call verbose callbacks even if verbose is False."""
+        return True
+
+    def _extract_token_base_model(self, usage_metadata: dict[str, Any]) -> TokenUsageBaseModel:
+        if usage_metadata:
+            prompt_tokens = usage_metadata.get("input_tokens", 0)
+            completion_tokens = usage_metadata.get("output_tokens", 0)
+            total_tokens = usage_metadata.get("total_tokens", 0)
+
+            return TokenUsageBaseModel(
+                prompt_tokens=prompt_tokens,
+                completion_tokens=completion_tokens,
+                total_tokens=total_tokens,
+            )
+        return TokenUsageBaseModel()
+
+    async def on_llm_start(self, serialized: dict[str, Any], prompts: list[str], **kwargs: Any) -> None:
+
+        model_name = ""
+        try:
+            model_name = kwargs.get("metadata")["ls_model_name"]
+        except Exception as e:
+            logger.exception("Error getting model name: %s", e, exc_info=True)
+
+        run_id = str(kwargs.get("run_id", str(uuid4())))
+        self._run_id_to_model_name[run_id] = model_name
+
+        stats = IntermediateStepPayload(event_type=IntermediateStepType.LLM_START,
+                                        framework=LLMFrameworkEnum.LANGCHAIN,
+                                        name=model_name,
+                                        UUID=run_id,
+                                        data=StreamEventData(input=prompts[-1]),
+                                        metadata=TraceMetadata(chat_inputs=copy.deepcopy(prompts)),
+                                        usage_info=UsageInfo(token_usage=TokenUsageBaseModel(),
+                                                             num_llm_calls=1,
+                                                             seconds_between_calls=int(time.time() -
+                                                                                       self.last_call_ts)))
+
+        self.step_manager.push_intermediate_step(stats)
+        self._run_id_to_llm_input[run_id] = prompts[-1]
+        self._state = IntermediateStepType.LLM_START
+        self.last_call_ts = time.time()
+        self._run_id_to_start_time[run_id] = time.time()
+
+    async def on_chat_model_start(
+        self,
+        serialized: dict[str, Any],
+        messages: list[list[BaseMessage]],
+        *,
+        run_id: UUID,
+        parent_run_id: UUID | None = None,
+        tags: list[str] | None = None,
+        metadata: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> Any:
+
+        model_name = ""
+        try:
+            model_name = metadata["ls_model_name"] if metadata else kwargs.get("metadata")["ls_model_name"]
+        except Exception as e:
+            logger.exception("Error getting model name: %s", e, exc_info=True)
+
+        run_id = str(run_id)
+        self._run_id_to_model_name[run_id] = model_name
+
+        stats = IntermediateStepPayload(
+            event_type=IntermediateStepType.LLM_START,
+            framework=LLMFrameworkEnum.LANGCHAIN,
+            name=model_name,
+            UUID=run_id,
+            data=StreamEventData(input=copy.deepcopy(messages[0])),
+            metadata=TraceMetadata(chat_inputs=copy.deepcopy(messages[0]),
+                                   tools_schema=_extract_tools_schema(kwargs.get("invocation_params", {}))),
+            usage_info=UsageInfo(token_usage=TokenUsageBaseModel(),
+                                 num_llm_calls=1,
+                                 seconds_between_calls=int(time.time() - self.last_call_ts)))
+
+        self.step_manager.push_intermediate_step(stats)
+        self._run_id_to_llm_input[run_id] = messages[0][-1].content
+        self._state = IntermediateStepType.LLM_START
+        self.last_call_ts = time.time()
+        self._run_id_to_start_time[run_id] = time.time()
+
+    async def on_llm_new_token(self, token: str, **kwargs: Any) -> None:
+        """Collect stats for just the token"""
+        model_name = ""
+        try:
+            model_name = self._run_id_to_model_name.get(str(kwargs.get("run_id", "")), "")
+        except Exception as e:
+            logger.exception("Error getting model name: %s", e, exc_info=True)
+
+        usage_metadata = {}
+        try:
+            usage_metadata = kwargs.get("chunk").message.usage_metadata if kwargs.get("chunk") else {}
+        except Exception as e:
+            logger.exception("Error getting usage metadata: %s", e, exc_info=True)
+
+        stats = IntermediateStepPayload(
+            event_type=IntermediateStepType.LLM_NEW_TOKEN,
+            framework=LLMFrameworkEnum.LANGCHAIN,
+            name=model_name,
+            UUID=str(kwargs.get("run_id", str(uuid4()))),
+            data=StreamEventData(input=self._run_id_to_llm_input.get(str(kwargs.get("run_id", "")), ""), chunk=token),
+            usage_info=UsageInfo(token_usage=self._extract_token_base_model(usage_metadata),
+                                 num_llm_calls=1,
+                                 seconds_between_calls=int(time.time() - self.last_call_ts)),
+            metadata=TraceMetadata(chat_responses=[kwargs.get("chunk")] if kwargs.get("chunk") else []))
+
+        self.step_manager.push_intermediate_step(stats)
+
+    async def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
+        """Collect token usage."""
+
+        usage_metadata = {}
+
+        model_name = ""
+        try:
+            model_name = response.llm_output["model_name"]
+        except Exception as e:
+            try:
+                model_name = self._run_id_to_model_name.get(str(kwargs.get("run_id", "")), "")
+            except Exception as e_inner:
+                logger.exception("Error getting model name: %s from outer error %s", e_inner, e, exc_info=True)
+
+        try:
+            generation = response.generations[0][0]
+        except IndexError:
+            generation = None
+
+        if isinstance(generation, ChatGeneration):
+            try:
+                message = generation.message
+                if isinstance(message, AIMessage):
+                    usage_metadata = message.usage_metadata
+                else:
+                    usage_metadata = {}
+            except AttributeError:
+                usage_metadata = {}
+
+        llm_text_output = generation.message.content if generation else ""
+
+        # update shared state behind lock
+        with self._lock:
+            usage_stat = IntermediateStepPayload(
+                span_event_timestamp=self._run_id_to_start_time.get(str(kwargs.get("run_id", "")), time.time()),
+                event_type=IntermediateStepType.LLM_END,
+                framework=LLMFrameworkEnum.LANGCHAIN,
+                name=model_name,
+                UUID=str(kwargs.get("run_id", str(uuid4()))),
+                data=StreamEventData(input=self._run_id_to_llm_input.get(str(kwargs.get("run_id", "")), ""),
+                                     output=llm_text_output),
+                usage_info=UsageInfo(token_usage=self._extract_token_base_model(usage_metadata)),
+                metadata=TraceMetadata(chat_responses=[generation] if generation else []))
+
+            self.step_manager.push_intermediate_step(usage_stat)
+
+        self._state = IntermediateStepType.LLM_END
+
+    async def on_tool_start(
+        self,
+        serialized: dict[str, Any],
+        input_str: str,
+        *,
+        run_id: UUID,
+        parent_run_id: UUID | None = None,
+        tags: list[str] | None = None,
+        metadata: dict[str, Any] | None = None,
+        inputs: dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> Any:
+
+        stats = IntermediateStepPayload(event_type=IntermediateStepType.TOOL_START,
+                                        framework=LLMFrameworkEnum.LANGCHAIN,
+                                        name=serialized.get("name", ""),
+                                        UUID=str(run_id),
+                                        data=StreamEventData(input=input_str),
+                                        metadata=TraceMetadata(tool_inputs=copy.deepcopy(inputs),
+                                                               tool_info=copy.deepcopy(serialized)),
+                                        usage_info=UsageInfo(token_usage=TokenUsageBaseModel()))
+
+        self.step_manager.push_intermediate_step(stats)
+        self._run_id_to_tool_input[str(run_id)] = input_str
+        self._run_id_to_start_time[str(run_id)] = time.time()
+
+    async def on_tool_end(
+        self,
+        output: Any,
+        *,
+        run_id: UUID,
+        parent_run_id: UUID | None = None,
+        **kwargs: Any,
+    ) -> Any:
+
+        stats = IntermediateStepPayload(event_type=IntermediateStepType.TOOL_END,
+                                        span_event_timestamp=self._run_id_to_start_time.get(str(run_id), time.time()),
+                                        framework=LLMFrameworkEnum.LANGCHAIN,
+                                        name=kwargs.get("name", ""),
+                                        UUID=str(run_id),
+                                        metadata=TraceMetadata(tool_outputs=output),
+                                        usage_info=UsageInfo(token_usage=TokenUsageBaseModel()),
+                                        data=StreamEventData(input=self._run_id_to_tool_input.get(str(run_id), ""),
+                                                             output=output))
+
+        self.step_manager.push_intermediate_step(stats)

aiq/profiler/callbacks/llama_index_callback_handler.py

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