nvidia-nat 1.3.0a20250910__py3-none-any.whl → 1.4.0a20251112__py3-none-any.whl

nat/data_models/config.py

@@ -20,6 +20,7 @@ import typing
 from pydantic import BaseModel
 from pydantic import ConfigDict
 from pydantic import Discriminator
+from pydantic import Field
 from pydantic import ValidationError
 from pydantic import ValidationInfo
 from pydantic import ValidatorFunctionWrapHandler
@@ -29,7 +30,9 @@ from nat.data_models.evaluate import EvalConfig
 from nat.data_models.front_end import FrontEndBaseConfig
 from nat.data_models.function import EmptyFunctionConfig
 from nat.data_models.function import FunctionBaseConfig
+from nat.data_models.function import FunctionGroupBaseConfig
 from nat.data_models.logging import LoggingBaseConfig
+from nat.data_models.optimizer import OptimizerConfig
 from nat.data_models.telemetry_exporter import TelemetryExporterBaseConfig
 from nat.data_models.ttc_strategy import TTCStrategyBaseConfig
 from nat.front_ends.fastapi.fastapi_front_end_config import FastApiFrontEndConfig
@@ -57,9 +60,10 @@ def _process_validation_error(err: ValidationError, handler: ValidatorFunctionWr
         error_type = e['type']
         if error_type == 'union_tag_invalid' and "ctx" in e and not logged_once:
             requested_type = e["ctx"]["tag"]
-
             if (info.field_name in ('workflow', 'functions')):
                 registered_keys = GlobalTypeRegistry.get().get_registered_functions()
+            elif (info.field_name == "function_groups"):
+                registered_keys = GlobalTypeRegistry.get().get_registered_function_groups()
             elif (info.field_name == "authentication"):
                 registered_keys = GlobalTypeRegistry.get().get_registered_auth_providers()
             elif (info.field_name == "llms"):
@@ -135,8 +139,8 @@ def _process_validation_error(err: ValidationError, handler: ValidatorFunctionWr
 
 class TelemetryConfig(BaseModel):
 
-    logging: dict[str, LoggingBaseConfig] = {}
-    tracing: dict[str, TelemetryExporterBaseConfig] = {}
+    logging: dict[str, LoggingBaseConfig] = Field(default_factory=dict)
+    tracing: dict[str, TelemetryExporterBaseConfig] = Field(default_factory=dict)
 
     @field_validator("logging", "tracing", mode="wrap")
     @classmethod
@@ -183,12 +187,16 @@ class TelemetryConfig(BaseModel):
 
 class GeneralConfig(BaseModel):
 
-    model_config = ConfigDict(protected_namespaces=())
+    model_config = ConfigDict(protected_namespaces=(), extra="forbid")
 
-    use_uvloop: bool = True
+    use_uvloop: bool | None = Field(
+        default=None,
+        deprecated=
+        "`use_uvloop` field is deprecated and will be removed in a future release. The use of `uv_loop` is now" +
+        "automatically determined based on platform")
     """
-    Whether to use uvloop for the event loop. This can provide a significant speedup in some cases. Disable to provide
-    better error messages when debugging.
+    This field is deprecated and ignored. It previously controlled whether to use uvloop as the event loop. uvloop
+    usage is now determined automatically based on the platform.
     """
 
     telemetry: TelemetryConfig = TelemetryConfig()
@@ -240,31 +248,37 @@ class Config(HashableBaseModel):
     general: GeneralConfig = GeneralConfig()
 
     # Functions Configuration
-    functions: dict[str, FunctionBaseConfig] = {}
+    functions: dict[str, FunctionBaseConfig] = Field(default_factory=dict)
+
+    # Function Groups Configuration
+    function_groups: dict[str, FunctionGroupBaseConfig] = Field(default_factory=dict)
 
     # LLMs Configuration
-    llms: dict[str, LLMBaseConfig] = {}
+    llms: dict[str, LLMBaseConfig] = Field(default_factory=dict)
 
     # Embedders Configuration
-    embedders: dict[str, EmbedderBaseConfig] = {}
+    embedders: dict[str, EmbedderBaseConfig] = Field(default_factory=dict)
 
     # Memory Configuration
-    memory: dict[str, MemoryBaseConfig] = {}
+    memory: dict[str, MemoryBaseConfig] = Field(default_factory=dict)
 
     # Object Stores Configuration
-    object_stores: dict[str, ObjectStoreBaseConfig] = {}
+    object_stores: dict[str, ObjectStoreBaseConfig] = Field(default_factory=dict)
+
+    # Optimizer Configuration
+    optimizer: OptimizerConfig = OptimizerConfig()
 
     # Retriever Configuration
-    retrievers: dict[str, RetrieverBaseConfig] = {}
+    retrievers: dict[str, RetrieverBaseConfig] = Field(default_factory=dict)
 
     # TTC Strategies
-    ttc_strategies: dict[str, TTCStrategyBaseConfig] = {}
+    ttc_strategies: dict[str, TTCStrategyBaseConfig] = Field(default_factory=dict)
 
     # Workflow Configuration
     workflow: FunctionBaseConfig = EmptyFunctionConfig()
 
     # Authentication Configuration
-    authentication: dict[str, AuthProviderBaseConfig] = {}
+    authentication: dict[str, AuthProviderBaseConfig] = Field(default_factory=dict)
 
     # Evaluation Options
     eval: EvalConfig = EvalConfig()
@@ -278,6 +292,7 @@ class Config(HashableBaseModel):
             stream.write(f"Workflow Type: {self.workflow.type}\n")
 
         stream.write(f"Number of Functions: {len(self.functions)}\n")
+        stream.write(f"Number of Function Groups: {len(self.function_groups)}\n")
         stream.write(f"Number of LLMs: {len(self.llms)}\n")
         stream.write(f"Number of Embedders: {len(self.embedders)}\n")
         stream.write(f"Number of Memory: {len(self.memory)}\n")
@@ -287,6 +302,7 @@ class Config(HashableBaseModel):
         stream.write(f"Number of Authentication Providers: {len(self.authentication)}\n")
 
     @field_validator("functions",
+                     "function_groups",
                      "llms",
                      "embedders",
                      "memory",
@@ -328,6 +344,10 @@ class Config(HashableBaseModel):
                                    typing.Annotated[type_registry.compute_annotation(FunctionBaseConfig),
                                                     Discriminator(TypedBaseModel.discriminator)]]
 
+        FunctionGroupsAnnotation = dict[str,
+                                        typing.Annotated[type_registry.compute_annotation(FunctionGroupBaseConfig),
+                                                         Discriminator(TypedBaseModel.discriminator)]]
+
         MemoryAnnotation = dict[str,
                                 typing.Annotated[type_registry.compute_annotation(MemoryBaseConfig),
                                                  Discriminator(TypedBaseModel.discriminator)]]
@@ -335,7 +355,6 @@ class Config(HashableBaseModel):
         ObjectStoreAnnotation = dict[str,
                                      typing.Annotated[type_registry.compute_annotation(ObjectStoreBaseConfig),
                                                       Discriminator(TypedBaseModel.discriminator)]]
-
         RetrieverAnnotation = dict[str,
                                    typing.Annotated[type_registry.compute_annotation(RetrieverBaseConfig),
                                                     Discriminator(TypedBaseModel.discriminator)]]
@@ -344,7 +363,7 @@ class Config(HashableBaseModel):
                                      typing.Annotated[type_registry.compute_annotation(TTCStrategyBaseConfig),
                                                       Discriminator(TypedBaseModel.discriminator)]]
 
-        WorkflowAnnotation = typing.Annotated[type_registry.compute_annotation(FunctionBaseConfig),
+        WorkflowAnnotation = typing.Annotated[(type_registry.compute_annotation(FunctionBaseConfig)),
                                               Discriminator(TypedBaseModel.discriminator)]
 
         should_rebuild = False
@@ -369,6 +388,11 @@ class Config(HashableBaseModel):
             functions_field.annotation = FunctionsAnnotation
             should_rebuild = True
 
+        function_groups_field = cls.model_fields.get("function_groups")
+        if function_groups_field is not None and function_groups_field.annotation != FunctionGroupsAnnotation:
+            function_groups_field.annotation = FunctionGroupsAnnotation
+            should_rebuild = True
+
         memory_field = cls.model_fields.get("memory")
         if memory_field is not None and memory_field.annotation != MemoryAnnotation:
             memory_field.annotation = MemoryAnnotation

nat/data_models/dataset_handler.py

@@ -26,6 +26,7 @@ from pydantic import FilePath
 from pydantic import Tag
 
 from nat.data_models.common import BaseModelRegistryTag
+from nat.data_models.common import SerializableSecretStr
 from nat.data_models.common import TypedBaseModel
 
 
@@ -34,8 +35,8 @@ class EvalS3Config(BaseModel):
     endpoint_url: str | None = None
     region_name: str | None = None
     bucket: str
-    access_key: str
-    secret_key: str
+    access_key: SerializableSecretStr
+    secret_key: SerializableSecretStr
 
 
 class EvalFilterEntryConfig(BaseModel):
@@ -80,7 +81,7 @@ class EvalDatasetJsonConfig(EvalDatasetBaseConfig, name="json"):
 
 
 def read_jsonl(file_path: FilePath):
-    with open(file_path, 'r', encoding='utf-8') as f:
+    with open(file_path, encoding='utf-8') as f:
         data = [json.loads(line) for line in f]
     return pd.DataFrame(data)
 

nat/data_models/function.py

@@ -15,6 +15,10 @@
 
 import typing
 
+from pydantic import Field
+from pydantic import field_validator
+from pydantic import model_validator
+
 from .common import BaseModelRegistryTag
 from .common import TypedBaseModel
 
@@ -23,8 +27,38 @@ class FunctionBaseConfig(TypedBaseModel, BaseModelRegistryTag):
     pass
 
 
+class FunctionGroupBaseConfig(TypedBaseModel, BaseModelRegistryTag):
+    """Base configuration for function groups.
+
+    Function groups enable sharing of configurations and resources across multiple functions.
+    """
+    include: list[str] = Field(
+        default_factory=list,
+        description="The list of function names which should be added to the global Function registry",
+    )
+    exclude: list[str] = Field(
+        default_factory=list,
+        description="The list of function names which should be excluded from default access to the group",
+    )
+
+    @field_validator("include", "exclude")
+    @classmethod
+    def _validate_fields_include_exclude(cls, value: list[str]) -> list[str]:
+        if len(set(value)) != len(value):
+            raise ValueError("Function names must be unique")
+        return sorted(value)
+
+    @model_validator(mode="after")
+    def _validate_include_exclude(self):
+        if self.include and self.exclude:
+            raise ValueError("include and exclude cannot be used together")
+        return self
+
+
 class EmptyFunctionConfig(FunctionBaseConfig, name="EmptyFunctionConfig"):
     pass
 
 
 FunctionConfigT = typing.TypeVar("FunctionConfigT", bound=FunctionBaseConfig)
+
+FunctionGroupConfigT = typing.TypeVar("FunctionGroupConfigT", bound=FunctionGroupBaseConfig)

nat/data_models/function_dependencies.py

@@ -23,6 +23,7 @@ class FunctionDependencies(BaseModel):
     A class to represent the dependencies of a function.
     """
     functions: set[str] = Field(default_factory=set)
+    function_groups: set[str] = Field(default_factory=set)
     llms: set[str] = Field(default_factory=set)
     embedders: set[str] = Field(default_factory=set)
     memory_clients: set[str] = Field(default_factory=set)
@@ -33,6 +34,10 @@ class FunctionDependencies(BaseModel):
     def serialize_functions(self, v: set[str]) -> list[str]:
         return list(v)
 
+    @field_serializer("function_groups", when_used="json")
+    def serialize_function_groups(self, v: set[str]) -> list[str]:
+        return list(v)
+
     @field_serializer("llms", when_used="json")
     def serialize_llms(self, v: set[str]) -> list[str]:
         return list(v)
@@ -56,6 +61,9 @@ class FunctionDependencies(BaseModel):
     def add_function(self, function: str):
         self.functions.add(function)
 
+    def add_function_group(self, function_group: str):
+        self.function_groups.add(function_group)  # pylint: disable=no-member
+
     def add_llm(self, llm: str):
         self.llms.add(llm)
 

nat/data_models/intermediate_step.py

@@ -103,11 +103,19 @@ class ToolSchema(BaseModel):
     function: ToolDetails = Field(..., description="The function details.")
 
 
+class ServerToolUseSchema(BaseModel):
+    name: str
+    arguments: str | dict[str, typing.Any] | typing.Any
+    output: typing.Any
+
+    model_config = ConfigDict(extra="ignore")
+
+
 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_outputs: list[ServerToolUseSchema] | typing.Any | None = None
     tool_info: typing.Any | None = None
     span_inputs: typing.Any | None = None
     span_outputs: typing.Any | None = None

nat/data_models/llm.py

@@ -14,14 +14,28 @@
 # limitations under the License.
 
 import typing
+from enum import Enum
+
+from pydantic import Field
 
 from .common import BaseModelRegistryTag
 from .common import TypedBaseModel
 
 
+class APITypeEnum(str, Enum):
+    CHAT_COMPLETION = "chat_completion"
+    RESPONSES = "responses"
+
+
 class LLMBaseConfig(TypedBaseModel, BaseModelRegistryTag):
     """Base configuration for LLM providers."""
-    pass
+
+    api_type: APITypeEnum = Field(default=APITypeEnum.CHAT_COMPLETION,
+                                  description="The type of API to use for the LLM provider.",
+                                  json_schema_extra={
+                                      "enum": [e.value for e in APITypeEnum],
+                                      "examples": [e.value for e in APITypeEnum],
+                                  })
 
 
 LLMBaseConfigT = typing.TypeVar("LLMBaseConfigT", bound=LLMBaseConfig)

nat/data_models/openai_mcp.py

@@ -0,0 +1,46 @@
+# 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 enum import Enum
+
+from pydantic import BaseModel
+from pydantic import ConfigDict
+from pydantic import Field
+
+
+class MCPApprovalRequiredEnum(str, Enum):
+    """
+    Enum to specify if approval is required for tool usage in the OpenAI MCP schema.
+    """
+    NEVER = "never"
+    ALWAYS = "always"
+    AUTO = "auto"
+
+
+class OpenAIMCPSchemaTool(BaseModel):
+    """
+    Represents a tool in the OpenAI MCP schema.
+    """
+    type: str = "mcp"
+    server_label: str = Field(description="Label for the server where the tool is hosted.")
+    server_url: str = Field(description="URL of the server hosting the tool.")
+    allowed_tools: list[str] | None = Field(default=None,
+                                            description="List of allowed tool names that can be used by the agent.")
+    require_approval: MCPApprovalRequiredEnum = Field(default=MCPApprovalRequiredEnum.NEVER,
+                                                      description="Specifies if approval is required for tool usage.")
+    headers: dict[str, str] | None = Field(default=None,
+                                           description="Optional headers to include in requests to the tool server.")
+
+    model_config = ConfigDict(use_enum_values=True)

nat/data_models/optimizable.py

@@ -0,0 +1,208 @@
+# 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
+
+import numpy as np
+from optuna import Trial
+from pydantic import BaseModel
+from pydantic import ConfigDict
+from pydantic import Field
+from pydantic import model_validator
+from pydantic_core import PydanticUndefined
+
+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 SearchSpace configuration."""
+        # 1. Prompt-specific validation
+        if self.is_prompt:
+            # When optimizing prompts, numeric parameters don't make sense
+            if self.low is not None or self.high is not None:
+                raise ValueError("SearchSpace with 'is_prompt=True' cannot have 'low' or 'high' parameters")
+            if self.log:
+                raise ValueError("SearchSpace with 'is_prompt=True' cannot have 'log=True'")
+            if self.step is not None:
+                raise ValueError("SearchSpace with 'is_prompt=True' cannot have 'step' parameter")
+            return self
+
+        # 2. Values-based validation
+        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'")
+            # Ensure values is not empty
+            if len(self.values) == 0:
+                raise ValueError("SearchSpace 'values' must not be empty")
+            return self
+
+        # 3. Range-based validation
+        if (self.low is None) != (self.high is None):  # XOR using !=
+            raise ValueError(f"SearchSpace range requires both 'low' and 'high'; got low={self.low}, high={self.high}")
+        if self.low is not None and self.high is not None and self.low >= self.high:
+            raise ValueError(f"SearchSpace 'low' must be less than 'high'; got low={self.low}, high={self.high}")
+
+        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 to_grid_values(self) -> list[Any]:
+        """
+        Convert SearchSpace to a list of values for GridSampler.
+
+        Grid search requires explicit values. This can be provided in two ways:
+        1. Explicit values: SearchSpace(values=[0.1, 0.5, 0.9])
+        2. Range with step: SearchSpace(low=0.1, high=0.9, step=0.2)
+
+        For ranges, step is required (no default will be applied) to avoid
+        unintentional combinatorial explosion.
+        """
+
+        if self.is_prompt:
+            raise ValueError("Prompt optimization not currently supported using Optuna. "
+                             "Use the genetic algorithm implementation instead.")
+
+        # Option 1: Explicit values provided
+        if self.values is not None:
+            return list(self.values)
+
+        # Option 2: Range with required step
+        if self.low is None or self.high is None:
+            raise ValueError("Grid search requires either 'values' or both 'low' and 'high' to be defined")
+
+        if self.step is None:
+            raise ValueError(
+                f"Grid search with range (low={self.low}, high={self.high}) requires 'step' to be specified. "
+                "Please define the step size to discretize the range, for example: step=0.1")
+
+        # Validate step is positive
+        step_float = float(self.step)
+        if step_float <= 0:
+            raise ValueError(f"Grid search step must be positive; got step={self.step}")
+
+        # Generate grid values from range with step
+        # Use integer range only if low, high, and step are all integral
+        if (isinstance(self.low, int) and isinstance(self.high, int) and step_float.is_integer()):
+            step = int(step_float)
+
+            if self.log:
+                raise ValueError("Log scale is not supported for integer ranges in grid search. "
+                                 "Please use linear scale or provide explicit values.")
+            values = list(range(self.low, self.high + 1, step))
+            if values and values[-1] != self.high:
+                values.append(self.high)
+            return values
+
+        # Float range (including integer low/high with float step)
+        low_val = float(self.low)
+        high_val = float(self.high)
+        step_val = step_float
+
+        if self.log:
+            raise ValueError("Log scale is not yet supported for grid search with ranges. "
+                             "Please provide explicit values using the 'values' field.")
+
+        # Use arange to respect step size
+        values = np.arange(low_val, high_val, step_val).tolist()
+
+        # Always include the high endpoint if not already present (within tolerance)
+        # This ensures the full range is explored in grid search
+        if not values or abs(values[-1] - high_val) > 1e-9:
+            values.append(high_val)
+
+        return values
+
+
+def OptimizableField(
+    default: Any = PydanticUndefined,
+    *,
+    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,
+    )