aiqtoolkit 1.2.0a20250609__py3-none-any.whl → 1.2.0a20250610__py3-none-any.whl

aiq/data_models/intermediate_step.py

@@ -17,6 +17,7 @@ import time
 import typing
 import uuid
 from enum import Enum
+from typing import Literal
 
 from pydantic import BaseModel
 from pydantic import ConfigDict
@@ -82,6 +83,26 @@ class UsageInfo(BaseModel):
     seconds_between_calls: int = 0
 
 
+class ToolParameters(BaseModel):
+    properties: dict[str, typing.Any] = Field(..., description="The properties of the function parameters.")
+    required: list[str] = Field(default_factory=list, description="The required properties of the function parameters.")
+    type_: Literal["object"] = Field(default="object", description="The type of the function parameters.", alias="type")
+    additionalProperties: bool = Field(default=False,
+                                       description="Enable function parameters allow additional properties.")
+    strict: bool = Field(default=True, description="Ensure function calls reliably adhere to the function schema.")
+
+
+class ToolDetails(BaseModel):
+    name: str = Field(..., description="The name of the function.")
+    description: str = Field(..., description="The description of the function.")
+    parameters: ToolParameters = Field(..., description="The parameters of the function.")
+
+
+class ToolSchema(BaseModel):
+    type: Literal["function"] = Field(..., description="The type of the tool.")
+    function: ToolDetails = Field(..., description="The function details.")
+
+
 class TraceMetadata(BaseModel):
     chat_responses: typing.Any | None = None
     chat_inputs: typing.Any | None = None
@@ -91,6 +112,8 @@ class TraceMetadata(BaseModel):
     span_inputs: typing.Any | None = None
     span_outputs: typing.Any | None = None
     provided_metadata: typing.Any | None = None
+    tools_schema: list[ToolSchema] = Field(default_factory=list,
+                                           description="The schema of tools used in a tool calling request.")
 
     # Allow extra fields in the model_config to support derived models
     model_config = ConfigDict(extra="allow")

aiq/observability/async_otel_listener.py

@@ -13,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import json
 import logging
 import re
 import warnings
@@ -75,6 +76,24 @@ except TelemetryOptionalImportError:
     set_span_in_context = dummy_set_span_in_context
 
 
+def merge_dicts(dict1: dict, dict2: dict) -> dict:
+    """
+    Merge two dictionaries, prioritizing non-null values from the first dictionary.
+
+    Args:
+        dict1 (dict): First dictionary (higher priority)
+        dict2 (dict): Second dictionary (lower priority)
+
+    Returns:
+        dict: Merged dictionary with non-null values from dict1 taking precedence
+    """
+    result = dict2.copy()  # Start with a copy of the second dictionary
+    for key, value in dict1.items():
+        if value is not None:  # Only update if value is not None
+            result[key] = value
+    return result
+
+
 def _ns_timestamp(seconds_float: float) -> int:
     """
     Convert AIQ Toolkit's float `event_timestamp` (in seconds) into an integer number
@@ -285,6 +304,10 @@ class AsyncOtelSpanListener:
                 sub_span.set_attribute(SpanAttributes.INPUT_VALUE, serialized_input)
                 sub_span.set_attribute(SpanAttributes.INPUT_MIME_TYPE, "application/json" if is_json else "text/plain")
 
+        # Optional: add metadata to the span from TraceMetadata
+        if step.payload.metadata:
+            sub_span.set_attribute("aiq.metadata", step.payload.metadata.model_dump_json())
+
         self._span_stack[step.UUID] = sub_span
 
         self._outstanding_spans[step.UUID] = sub_span
@@ -323,6 +346,13 @@ class AsyncOtelSpanListener:
             sub_span.set_attribute(SpanAttributes.OUTPUT_VALUE, serialized_output)
             sub_span.set_attribute(SpanAttributes.OUTPUT_MIME_TYPE, "application/json" if is_json else "text/plain")
 
+        # # Optional: add metadata to the span from TraceMetadata
+        if step.payload.metadata:
+            start_event_metadata = json.loads(sub_span.attributes.get("aiq.metadata", {}))
+            end_event_metadata = json.loads(step.payload.metadata.model_dump_json())
+            merged_event_metadata = merge_dicts(start_event_metadata, end_event_metadata)
+            sub_span.set_attribute("aiq.metadata", json.dumps(merged_event_metadata))
+
         end_ns = _ns_timestamp(step.payload.event_timestamp)
 
         # End the subspan

aiq/profiler/callbacks/langchain_callback_handler.py

@@ -34,6 +34,7 @@ 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
@@ -42,6 +43,16 @@ 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."""
 
@@ -138,16 +149,17 @@ class LangchainProfilerHandler(AsyncCallbackHandler, BaseProfilerCallback):  # p
         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])),
-                                        usage_info=UsageInfo(token_usage=TokenUsageBaseModel(),
-                                                             num_llm_calls=1,
-                                                             seconds_between_calls=int(time.time() -
-                                                                                       self.last_call_ts)))
+        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

{aiqtoolkit-1.2.0a20250609.dist-info → aiqtoolkit-1.2.0a20250610.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: aiqtoolkit
-Version: 1.2.0a20250609
+Version: 1.2.0a20250610
 Summary: NVIDIA Agent Intelligence toolkit
 Author: NVIDIA Corporation
 Maintainer: NVIDIA Corporation

{aiqtoolkit-1.2.0a20250609.dist-info → aiqtoolkit-1.2.0a20250610.dist-info}/RECORD

@@ -88,7 +88,7 @@ aiq/data_models/front_end.py,sha256=z8k6lSWjt1vMOYFbjWQxodpwAqPeuGS0hRBjsriDW2s,
 aiq/data_models/function.py,sha256=M_duXVXL5MvYe0WVLvqEgEzXs0UAYNSMfy9ZTpxuKPA,1013
 aiq/data_models/function_dependencies.py,sha256=bIdfrwDT_rR5eEER2WhYtl43QaOAnTL1p8IkyGJtB_Y,2403
 aiq/data_models/interactive.py,sha256=cYBIp3N3AaP0pxMRJFMwVoMrAWpWVT2turidhr__36A,8418
-aiq/data_models/intermediate_step.py,sha256=F0QW-uZODehEV0NKy55zl9fmbZrF2qggi8R4iW_XKHk,10057
+aiq/data_models/intermediate_step.py,sha256=XVn8MdAJAk5Yy3JfWfgUQHKqjsr5E-XrC_sluFAU8OA,11385
 aiq/data_models/invocation_node.py,sha256=nDRylgzBfJduGA-lme9xN4P6BdOYj0L6ytLHnTuetMI,1618
 aiq/data_models/llm.py,sha256=McbDdUUtWfp9WCdMMJA2xh7mvlmyNdGDCH8P_7l2iKU,920
 aiq/data_models/logging.py,sha256=1QtVjIQ99PgMYUuzw4h1FAoPRteZY7uf3oFTqV3ONgA,940
@@ -174,7 +174,7 @@ aiq/memory/models.py,sha256=c5dA7nKHQ4AS1_ptQZcfC_oXO495-ehocnf_qXTE6c8,4319
 aiq/meta/module_to_distro.json,sha256=1XV7edobFrdDKvsSoynfodXg_hczUWpDrQzGkW9qqEs,28
 aiq/meta/pypi.md,sha256=Ukj1J--Q6GF7Zh1tKonygttN3aiq_Lf2-445nz0sE-o,4362
 aiq/observability/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-aiq/observability/async_otel_listener.py,sha256=gR8fEdpZ9L8vGZiXLI7FwbwtSiGgQHMDDj5d-vKjZGc,18407
+aiq/observability/async_otel_listener.py,sha256=2Ye9bkHfAssuxFS_ECyRyl-bTa73yYvsPyO4BaK5Beg,19662
 aiq/observability/register.py,sha256=WLPGY1b_RmVv48od-hIytvJrCqNoJupriHCzyxvTbUw,6723
 aiq/plugins/.namespace,sha256=Gace0pOC3ETEJf-TBVuNw0TQV6J_KtOPpEiSzMH-odo,215
 aiq/profiler/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -186,7 +186,7 @@ aiq/profiler/utils.py,sha256=hNh_JfxXDrACIp4usXtlriTfVuYUkk3Pv-x74K34MQg,8180
 aiq/profiler/callbacks/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 aiq/profiler/callbacks/agno_callback_handler.py,sha256=aDAUY6GDIUtly6KowXXKUqLc7NbE6khg1aXT1AritaA,14930
 aiq/profiler/callbacks/base_callback_class.py,sha256=BFZMkb-dPoHAnM_4jVXX08hD8Vi2n8thyftVlUxfE24,745
-aiq/profiler/callbacks/langchain_callback_handler.py,sha256=sgxnbEMGgqHx7_IY6uAAFURHdRvDOzG8G_brr_3KJAU,12388
+aiq/profiler/callbacks/langchain_callback_handler.py,sha256=tG9OX2C3Rsqpew2Jaj2LLkZkwEztbeQ-oLVIBUDPL48,12511
 aiq/profiler/callbacks/llama_index_callback_handler.py,sha256=AXUtfsUwif-rUFoVaRnMeSe29TcpT9Z8nxWcJWD2PiQ,9334
 aiq/profiler/callbacks/semantic_kernel_callback_handler.py,sha256=uYlq0Lku6U1AMg5GoPpJoZm6lvgIjLQN5_uIbswL4bA,11157
 aiq/profiler/callbacks/token_usage_base_model.py,sha256=X0b_AbBgVQAAbgbDMim-3S3XdQ7PaPs9qMUACvMAe5o,1104
@@ -308,10 +308,10 @@ aiq/utils/reactive/base/observer_base.py,sha256=UAlyAY_ky4q2t0P81RVFo2Bs_R7z5Nde
 aiq/utils/reactive/base/subject_base.py,sha256=Ed-AC6P7cT3qkW1EXjzbd5M9WpVoeN_9KCe3OM3FLU4,2521
 aiq/utils/settings/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 aiq/utils/settings/global_settings.py,sha256=U9TCLdoZsKq5qOVGjREipGVv9e-FlStzqy5zv82_VYk,7454
-aiqtoolkit-1.2.0a20250609.dist-info/licenses/LICENSE-3rd-party.txt,sha256=8o7aySJa9CBvFshPcsRdJbczzdNyDGJ8b0J67WRUQ2k,183936
-aiqtoolkit-1.2.0a20250609.dist-info/licenses/LICENSE.md,sha256=QwcOLU5TJoTeUhuIXzhdCEEDDvorGiC6-3YTOl4TecE,11356
-aiqtoolkit-1.2.0a20250609.dist-info/METADATA,sha256=QOm_3b4TK8JrHBnrAtNBolujyy6DYJtTyAiicNN46Ng,20250
-aiqtoolkit-1.2.0a20250609.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-aiqtoolkit-1.2.0a20250609.dist-info/entry_points.txt,sha256=gRlPfR5g21t328WNEQ4CcEz80S1sJNS8A7rMDYnzl4A,452
-aiqtoolkit-1.2.0a20250609.dist-info/top_level.txt,sha256=fo7AzYcNhZ_tRWrhGumtxwnxMew4xrT1iwouDy_f0Kc,4
-aiqtoolkit-1.2.0a20250609.dist-info/RECORD,,
+aiqtoolkit-1.2.0a20250610.dist-info/licenses/LICENSE-3rd-party.txt,sha256=8o7aySJa9CBvFshPcsRdJbczzdNyDGJ8b0J67WRUQ2k,183936
+aiqtoolkit-1.2.0a20250610.dist-info/licenses/LICENSE.md,sha256=QwcOLU5TJoTeUhuIXzhdCEEDDvorGiC6-3YTOl4TecE,11356
+aiqtoolkit-1.2.0a20250610.dist-info/METADATA,sha256=IU2IX0jWhO68QeQ8nZMzpCNnlGHCShFaijLFkWxKajo,20250
+aiqtoolkit-1.2.0a20250610.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+aiqtoolkit-1.2.0a20250610.dist-info/entry_points.txt,sha256=gRlPfR5g21t328WNEQ4CcEz80S1sJNS8A7rMDYnzl4A,452
+aiqtoolkit-1.2.0a20250610.dist-info/top_level.txt,sha256=fo7AzYcNhZ_tRWrhGumtxwnxMew4xrT1iwouDy_f0Kc,4
+aiqtoolkit-1.2.0a20250610.dist-info/RECORD,,