nvidia-nat 1.1.0a20251020__py3-none-any.whl

nat/data_models/intermediate_step.py

@@ -0,0 +1,302 @@
+# 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 time
+import typing
+import uuid
+from enum import Enum
+from typing import Literal
+
+from pydantic import BaseModel
+from pydantic import ConfigDict
+from pydantic import Field
+from pydantic import model_validator
+
+from nat.builder.framework_enum import LLMFrameworkEnum
+from nat.data_models.invocation_node import InvocationNode
+from nat.profiler.callbacks.token_usage_base_model import TokenUsageBaseModel
+
+
+class IntermediateStepCategory(str, Enum):
+    LLM = "LLM"
+    TOOL = "TOOL"
+    WORKFLOW = "WORKFLOW"
+    TASK = "TASK"
+    FUNCTION = "FUNCTION"
+    CUSTOM = "CUSTOM"
+    SPAN = "SPAN"
+
+
+class IntermediateStepType(str, Enum):
+    LLM_START = "LLM_START"
+    LLM_END = "LLM_END"
+    LLM_NEW_TOKEN = "LLM_NEW_TOKEN"
+    TOOL_START = "TOOL_START"
+    TOOL_END = "TOOL_END"
+    WORKFLOW_START = "WORKFLOW_START"
+    WORKFLOW_END = "WORKFLOW_END"
+    TASK_START = "TASK_START"
+    TASK_END = "TASK_END"
+    FUNCTION_START = "FUNCTION_START"
+    FUNCTION_END = "FUNCTION_END"
+    CUSTOM_START = "CUSTOM_START"
+    CUSTOM_END = "CUSTOM_END"
+    SPAN_START = "SPAN_START"
+    SPAN_CHUNK = "SPAN_CHUNK"
+    SPAN_END = "SPAN_END"
+
+
+class IntermediateStepState(str, Enum):
+    START = "START"
+    CHUNK = "CHUNK"
+    END = "END"
+
+
+class StreamEventData(BaseModel):
+    """
+    StreamEventData is a data model that represents the data field in an streaming event.
+    """
+
+    # Allow extra fields in the model_config to support derived models
+    model_config = ConfigDict(extra="allow")
+
+    input: typing.Any | None = None
+    output: typing.Any | None = None
+    chunk: typing.Any | None = None
+
+
+class UsageInfo(BaseModel):
+    token_usage: TokenUsageBaseModel = TokenUsageBaseModel()
+    num_llm_calls: int = 0
+    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
+    tool_inputs: typing.Any | None = None
+    tool_outputs: typing.Any | None = None
+    tool_info: typing.Any | None = None
+    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")
+
+
+class IntermediateStepPayload(BaseModel):
+    """
+    IntermediateStep is a data model that represents an intermediate step in the NAT. Intermediate steps are
+    captured while a request is running and can be used to show progress or to evaluate the path a workflow took to get
+    a response.
+    """
+
+    # Allow extra fields in the model_config to support derived models
+    model_config = ConfigDict(extra="allow")
+
+    event_type: IntermediateStepType
+    # Create an event timestamp field with the default being a lambda that returns the current time
+    event_timestamp: float = Field(default_factory=lambda: time.time())
+    span_event_timestamp: float | None = None  # Used for tracking the start time of a task if this is end
+    framework: LLMFrameworkEnum | None = None
+    name: str | None = None
+    tags: list[str] | None = None
+    metadata: dict[str, typing.Any] | TraceMetadata | None = None
+    data: StreamEventData | None = None
+    usage_info: UsageInfo | None = None
+    UUID: str = Field(default_factory=lambda: str(uuid.uuid4()))
+
+    @property
+    def event_category(self) -> IntermediateStepCategory:
+        match self.event_type:
+            case IntermediateStepType.LLM_START:
+                return IntermediateStepCategory.LLM
+            case IntermediateStepType.LLM_END:
+                return IntermediateStepCategory.LLM
+            case IntermediateStepType.LLM_NEW_TOKEN:
+                return IntermediateStepCategory.LLM
+            case IntermediateStepType.TOOL_START:
+                return IntermediateStepCategory.TOOL
+            case IntermediateStepType.TOOL_END:
+                return IntermediateStepCategory.TOOL
+            case IntermediateStepType.WORKFLOW_START:
+                return IntermediateStepCategory.WORKFLOW
+            case IntermediateStepType.WORKFLOW_END:
+                return IntermediateStepCategory.WORKFLOW
+            case IntermediateStepType.TASK_START:
+                return IntermediateStepCategory.TASK
+            case IntermediateStepType.TASK_END:
+                return IntermediateStepCategory.TASK
+            case IntermediateStepType.FUNCTION_START:
+                return IntermediateStepCategory.FUNCTION
+            case IntermediateStepType.FUNCTION_END:
+                return IntermediateStepCategory.FUNCTION
+            case IntermediateStepType.CUSTOM_START:
+                return IntermediateStepCategory.CUSTOM
+            case IntermediateStepType.CUSTOM_END:
+                return IntermediateStepCategory.CUSTOM
+            case IntermediateStepType.SPAN_START:
+                return IntermediateStepCategory.SPAN
+            case IntermediateStepType.SPAN_CHUNK:
+                return IntermediateStepCategory.SPAN
+            case IntermediateStepType.SPAN_END:
+                return IntermediateStepCategory.SPAN
+            case _:
+                raise ValueError(f"Unknown event type: {self.event_type}")
+
+    @property
+    def event_state(self) -> IntermediateStepState:
+        match self.event_type:
+            case IntermediateStepType.LLM_START:
+                return IntermediateStepState.START
+            case IntermediateStepType.LLM_END:
+                return IntermediateStepState.END
+            case IntermediateStepType.LLM_NEW_TOKEN:
+                return IntermediateStepState.CHUNK
+            case IntermediateStepType.TOOL_START:
+                return IntermediateStepState.START
+            case IntermediateStepType.TOOL_END:
+                return IntermediateStepState.END
+            case IntermediateStepType.WORKFLOW_START:
+                return IntermediateStepState.START
+            case IntermediateStepType.WORKFLOW_END:
+                return IntermediateStepState.END
+            case IntermediateStepType.TASK_START:
+                return IntermediateStepState.START
+            case IntermediateStepType.TASK_END:
+                return IntermediateStepState.END
+            case IntermediateStepType.FUNCTION_START:
+                return IntermediateStepState.START
+            case IntermediateStepType.FUNCTION_END:
+                return IntermediateStepState.END
+            case IntermediateStepType.CUSTOM_START:
+                return IntermediateStepState.START
+            case IntermediateStepType.CUSTOM_END:
+                return IntermediateStepState.END
+            case IntermediateStepType.SPAN_START:
+                return IntermediateStepState.START
+            case IntermediateStepType.SPAN_CHUNK:
+                return IntermediateStepState.CHUNK
+            case IntermediateStepType.SPAN_END:
+                return IntermediateStepState.END
+            case _:
+                raise ValueError(f"Unknown event type: {self.event_type}")
+
+    @model_validator(mode="after")
+    def check_span_event_timestamp(self) -> "IntermediateStepPayload":
+        if self.event_state != IntermediateStepState.END and self.span_event_timestamp is not None:
+            raise ValueError("span_event_timestamp can only be provided for events with an END state")
+        return self
+
+
+class IntermediateStep(BaseModel):
+    """
+    IntermediateStep is a data model that represents an intermediate step in the NAT. Intermediate steps are
+    captured while a request is running and can be used to show progress or to evaluate the path a workflow took to get
+    a response.
+    """
+
+    # Allow extra fields in the model_config to support derived models
+    model_config = ConfigDict(extra="forbid")
+
+    parent_id: str
+    """
+    The parent step ID for the current step. The parent ID is the ID of the last START step which has a different UUID
+    than the current step. This value is different from the function_ancestry.parent_id value which tracks the last
+    parent FUNCTION step. For the first START step, the parent_id is 'root'.
+    """
+
+    function_ancestry: InvocationNode
+    """
+    The function ancestry for the current step showing the current NAT function that was being executed when the step
+    was created.
+    """
+
+    payload: IntermediateStepPayload
+    """
+    The payload for the current step.
+    """
+
+    # ===== Payload Properties =====
+    @property
+    def event_type(self) -> IntermediateStepType:
+        return self.payload.event_type
+
+    @property
+    def event_timestamp(self) -> float:
+        return self.payload.event_timestamp
+
+    @property
+    def span_event_timestamp(self) -> float | None:
+        return self.payload.span_event_timestamp
+
+    @property
+    def framework(self) -> LLMFrameworkEnum | None:
+        return self.payload.framework
+
+    @property
+    def name(self) -> str | None:
+        return self.payload.name
+
+    @property
+    def tags(self) -> list[str] | None:
+        return self.payload.tags
+
+    @property
+    def metadata(self) -> dict[str, typing.Any] | TraceMetadata | None:
+        return self.payload.metadata
+
+    @property
+    def data(self) -> StreamEventData | None:
+        return self.payload.data
+
+    @property
+    def usage_info(self) -> UsageInfo | None:
+        return self.payload.usage_info
+
+    @property
+    def UUID(self) -> str:
+        return self.payload.UUID
+
+    @property
+    def event_category(self) -> IntermediateStepCategory:
+        return self.payload.event_category
+
+    @property
+    def event_state(self) -> IntermediateStepState:
+        return self.payload.event_state

nat/data_models/invocation_node.py

@@ -0,0 +1,38 @@
+# 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
+
+
+class InvocationNode(BaseModel):
+    """
+    Represents a node in an invocation call graph.
+
+    The InvocationNode class encapsulates the details of a specific function
+    invocation within a call graph. It stores the unique identifier of the
+    invocation, the function name, and optional details about the parent
+    node (if any). This class is useful for tracing the execution flow
+    in a system or application.
+
+    Attributes:
+        function_id (str): Unique identifier for the function invocation.
+        function_name (str): Name of the function invoked.
+        parent_id (str | None): Unique identifier of the parent invocation, if applicable. Defaults to None.
+        parent_name (str | None): Name of the parent function invoked, if applicable. Defaults to None.
+    """
+    function_id: str
+    function_name: str
+    parent_id: str | None = None
+    parent_name: str | None = None

nat/data_models/llm.py

@@ -0,0 +1,27 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-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 typing
+
+from .common import BaseModelRegistryTag
+from .common import TypedBaseModel
+
+
+class LLMBaseConfig(TypedBaseModel, BaseModelRegistryTag):
+    """Base configuration for LLM providers."""
+    pass
+
+
+LLMBaseConfigT = typing.TypeVar("LLMBaseConfigT", bound=LLMBaseConfig)

nat/data_models/logging.py

@@ -0,0 +1,26 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-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 typing
+
+from .common import BaseModelRegistryTag
+from .common import TypedBaseModel
+
+
+class LoggingBaseConfig(TypedBaseModel, BaseModelRegistryTag):
+    pass
+
+
+LoggingMethodConfigT = typing.TypeVar("LoggingMethodConfigT", bound=LoggingBaseConfig)

nat/data_models/memory.py

@@ -0,0 +1,27 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-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 typing
+
+from .common import BaseModelRegistryTag
+from .common import TypedBaseModel
+
+
+class MemoryBaseConfig(TypedBaseModel, BaseModelRegistryTag):
+    """ The base level config object for a memory object. Memories provide an interface for storing and retrieving. """
+    pass
+
+
+MemoryBaseConfigT = typing.TypeVar("MemoryBaseConfigT", bound=MemoryBaseConfig)

nat/data_models/object_store.py

@@ -0,0 +1,44 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-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 typing
+
+from .common import BaseModelRegistryTag
+from .common import TypedBaseModel
+
+
+class ObjectStoreBaseConfig(TypedBaseModel, BaseModelRegistryTag):
+    pass
+
+
+ObjectStoreBaseConfigT = typing.TypeVar("ObjectStoreBaseConfigT", bound=ObjectStoreBaseConfig)
+
+
+class KeyAlreadyExistsError(Exception):
+
+    def __init__(self, key: str, additional_message: str | None = None):
+        parts = [f"Key already exists: {key}."]
+        if additional_message:
+            parts.append(additional_message)
+        super().__init__(" ".join(parts))
+
+
+class NoSuchKeyError(Exception):
+
+    def __init__(self, key: str, additional_message: str | None = None):
+        parts = [f"No object found with key: {key}."]
+        if additional_message:
+            parts.append(additional_message)
+        super().__init__(" ".join(parts))

nat/data_models/optimizable.py

@@ -0,0 +1,119 @@
+# 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 collections.abc import Sequence
+from typing import Any
+from typing import Generic
+from typing import TypeVar
+
+from optuna import Trial
+from pydantic import BaseModel
+from pydantic import ConfigDict
+from pydantic import Field
+from pydantic import model_validator
+
+T = TypeVar("T", int, float, bool, str)
+
+
+# --------------------------------------------------------------------- #
+# 1.  Hyper‑parameter metadata container                                #
+# --------------------------------------------------------------------- #
+class SearchSpace(BaseModel, Generic[T]):
+    values: Sequence[T] | None = None
+    low: T | None = None
+    high: T | None = None
+    log: bool = False  # log scale
+    step: float | None = None
+    is_prompt: bool = False
+    prompt: str | None = None  # prompt to optimize
+    prompt_purpose: str | None = None  # purpose of the prompt
+
+    model_config = ConfigDict(protected_namespaces=(), extra="forbid")
+
+    @model_validator(mode="after")
+    def validate_search_space_parameters(self):
+        """Validate that either values is provided, or both high and low."""
+        if self.values is not None:
+            # If values is provided, we don't need high/low
+            if self.high is not None or self.low is not None:
+                raise ValueError("SearchSpace 'values' is mutually exclusive with 'high' and 'low'")
+            return self
+
+        return self
+
+    # Helper for Optuna Trials
+    def suggest(self, trial: Trial, name: str):
+        if self.is_prompt:
+            raise ValueError("Prompt optimization not currently supported using Optuna. "
+                             "Use the genetic algorithm implementation instead.")
+        if self.values is not None:
+            return trial.suggest_categorical(name, self.values)
+        if isinstance(self.low, int):
+            return trial.suggest_int(name, self.low, self.high, log=self.log, step=self.step)
+        return trial.suggest_float(name, self.low, self.high, log=self.log, step=self.step)
+
+
+def OptimizableField(
+    default: Any,
+    *,
+    space: SearchSpace | None = None,
+    merge_conflict: str = "overwrite",
+    **fld_kw,
+):
+    # 1. Pull out any user‑supplied extras (must be a dict)
+    user_extra = fld_kw.pop("json_schema_extra", None) or {}
+    if not isinstance(user_extra, dict):
+        raise TypeError("`json_schema_extra` must be a mapping.")
+
+    # 2. If the space is a prompt, ensure a concrete base prompt exists
+    if space is not None and getattr(space, "is_prompt", False):
+        if getattr(space, "prompt", None) is None:
+            if default is None:
+                raise ValueError("Prompt-optimized fields require a base prompt: provide a "
+                                 "non-None field default or set space.prompt.")
+            # Default prompt not provided in space; fall back to the field's default
+            space.prompt = default
+
+    # 3. Prepare our own metadata
+    ours = {"optimizable": True}
+    if space is not None:
+        ours["search_space"] = space
+
+    # 4. Merge with user extras according to merge_conflict policy
+    intersect = ours.keys() & user_extra.keys()
+    if intersect:
+        if merge_conflict == "error":
+            raise ValueError("`json_schema_extra` already contains reserved key(s): "
+                             f"{', '.join(intersect)}")
+        if merge_conflict == "keep":
+            # remove the ones the user already set so we don't overwrite them
+            ours = {k: v for k, v in ours.items() if k not in intersect}
+
+    merged_extra = {**user_extra, **ours}  # ours wins if 'overwrite'
+
+    # 5. Return a normal Pydantic Field with merged extras
+    return Field(default, json_schema_extra=merged_extra, **fld_kw)
+
+
+class OptimizableMixin(BaseModel):
+    optimizable_params: list[str] = Field(default_factory=list,
+                                          description="List of parameters that can be optimized.",
+                                          exclude=True)
+
+    search_space: dict[str, SearchSpace] = Field(
+        default_factory=dict,
+        description="Optional search space overrides for optimizable parameters.",
+        exclude=True,
+    )