nvidia-nat 1.3.0a20250910__py3-none-any.whl → 1.3.0a20250922__py3-none-any.whl

nat/control_flow/sequential_executor.py

@@ -0,0 +1,167 @@
+# 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 logging
+import typing
+
+from langchain_core.tools.base import BaseTool
+from pydantic import BaseModel
+from pydantic import Field
+
+from nat.builder.builder import Builder
+from nat.builder.framework_enum import LLMFrameworkEnum
+from nat.builder.function import Function
+from nat.builder.function_info import FunctionInfo
+from nat.cli.register_workflow import register_function
+from nat.data_models.component_ref import FunctionRef
+from nat.data_models.function import FunctionBaseConfig
+from nat.utils.type_utils import DecomposedType
+
+logger = logging.getLogger(__name__)
+
+
+class ToolExecutionConfig(BaseModel):
+    """Configuration for individual tool execution within sequential execution."""
+
+    use_streaming: bool = Field(default=False, description="Whether to use streaming output for the tool.")
+
+
+class SequentialExecutorConfig(FunctionBaseConfig, name="sequential_executor"):
+    """Configuration for sequential execution of a list of functions."""
+
+    tool_list: list[FunctionRef] = Field(default_factory=list,
+                                         description="A list of functions to execute sequentially.")
+    tool_execution_config: dict[str, ToolExecutionConfig] = Field(default_factory=dict,
+                                                                  description="Optional configuration for each"
+                                                                  "tool in the sequential execution tool list."
+                                                                  "Keys must match the tool names from the"
+                                                                  "tool_list.")
+    raise_type_incompatibility: bool = Field(
+        default=False,
+        description="Default to False. Check if the adjacent tools are type compatible,"
+        "which means the output type of the previous function is compatible with the input type of the next function."
+        "If set to True, any incompatibility will raise an exception. If set to false, the incompatibility will only"
+        "generate a warning message and the sequential execution will continue.")
+
+
+def _get_function_output_type(function: Function, tool_execution_config: dict[str, ToolExecutionConfig]) -> type:
+    function_config = tool_execution_config.get(function.instance_name, None)
+    if function_config:
+        return function.streaming_output_type if function_config.use_streaming else function.single_output_type
+    else:
+        return function.single_output_type
+
+
+def _validate_function_type_compatibility(src_fn: Function,
+                                          target_fn: Function,
+                                          tool_execution_config: dict[str, ToolExecutionConfig]) -> None:
+    src_output_type = _get_function_output_type(src_fn, tool_execution_config)
+    target_input_type = target_fn.input_type
+    logger.debug(
+        f"The output type of the {src_fn.instance_name} function is {str(src_output_type)}, is not compatible with"
+        f"the input type of the {target_fn.instance_name} function, which is {str(target_input_type)}")
+
+    is_compatible = DecomposedType.is_type_compatible(src_output_type, target_input_type)
+    if not is_compatible:
+        raise ValueError(
+            f"The output type of the {src_fn.instance_name} function is {str(src_output_type)}, is not compatible with"
+            f"the input type of the {target_fn.instance_name} function, which is {str(target_input_type)}")
+
+
+def _validate_tool_list_type_compatibility(sequential_executor_config: SequentialExecutorConfig,
+                                           builder: Builder) -> tuple[type, type]:
+    tool_list = sequential_executor_config.tool_list
+    tool_execution_config = sequential_executor_config.tool_execution_config
+
+    function_list: list[Function] = []
+    for function_ref in tool_list:
+        function_list.append(builder.get_function(function_ref))
+    if not function_list:
+        raise RuntimeError("The function list is empty")
+    input_type = function_list[0].input_type
+
+    if len(function_list) > 1:
+        for src_fn, target_fn in zip(function_list[0:-1], function_list[1:]):
+            try:
+                _validate_function_type_compatibility(src_fn, target_fn, tool_execution_config)
+            except ValueError as e:
+                raise ValueError(f"The sequential tool list has incompatible types: {e}")
+
+    output_type = _get_function_output_type(function_list[-1], tool_execution_config)
+    logger.debug(f"The input type of the sequential executor tool list is {str(input_type)},"
+                 f"the output type is {str(output_type)}")
+
+    return (input_type, output_type)
+
+
+@register_function(config_type=SequentialExecutorConfig, framework_wrappers=[LLMFrameworkEnum.LANGCHAIN])
+async def sequential_execution(config: SequentialExecutorConfig, builder: Builder):
+    logger.debug(f"Initializing sequential executor with tool list: {config.tool_list}")
+
+    tools: list[BaseTool] = builder.get_tools(tool_names=config.tool_list, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
+    tools_dict: dict[str, BaseTool] = {tool.name: tool for tool in tools}
+
+    try:
+        input_type, output_type = _validate_tool_list_type_compatibility(config, builder)
+    except ValueError as e:
+        if config.raise_type_incompatibility:
+            logger.error(f"The sequential executor tool list has incompatible types: {e}")
+            raise
+        else:
+            logger.warning(f"The sequential executor tool list has incompatible types: {e}")
+            input_type = typing.Any
+            output_type = typing.Any
+    except Exception as e:
+        raise ValueError(f"Error with the sequential executor tool list: {e}")
+
+    # The type annotation of _sequential_function_execution is dynamically set according to the tool list
+    async def _sequential_function_execution(initial_tool_input):
+        logger.debug(f"Executing sequential executor with tool list: {config.tool_list}")
+
+        tool_list: list[FunctionRef] = config.tool_list
+        tool_input = initial_tool_input
+        tool_response = None
+
+        for tool_name in tool_list:
+            tool = tools_dict[tool_name]
+            tool_execution_config = config.tool_execution_config.get(tool_name, None)
+            logger.debug(f"Executing tool {tool_name} with input: {tool_input}")
+            try:
+                if tool_execution_config:
+                    if tool_execution_config.use_streaming:
+                        output = ""
+                        async for chunk in tool.astream(tool_input):
+                            output += chunk.content
+                        tool_response = output
+                    else:
+                        tool_response = await tool.ainvoke(tool_input)
+                else:
+                    tool_response = await tool.ainvoke(tool_input)
+            except Exception as e:
+                logger.error(f"Error with tool {tool_name}: {e}")
+                raise
+
+            # The input of the next tool is the response of the previous tool
+            tool_input = tool_response
+
+        return tool_response
+
+    # Dynamically set the annotations for the function
+    _sequential_function_execution.__annotations__ = {"initial_tool_input": input_type, "return": output_type}
+    logger.debug(f"Sequential executor function annotations: {_sequential_function_execution.__annotations__}")
+
+    yield FunctionInfo.from_fn(_sequential_function_execution,
+                               description="Executes a list of functions sequentially."
+                               "The input of the next tool is the response of the previous tool.")

nat/data_models/agent.py

@@ -0,0 +1,34 @@
+# 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 Field
+from pydantic import PositiveInt
+
+from nat.data_models.component_ref import LLMRef
+from nat.data_models.function import FunctionBaseConfig
+
+
+class AgentBaseConfig(FunctionBaseConfig):
+    """Base configuration class for all NAT agents with common fields."""
+
+    workflow_alias: str | None = Field(
+        default=None,
+        description=("The alias of the workflow. Useful when the agent is configured as a workflow "
+                     "and needs to expose a customized name as a tool."))
+    llm_name: LLMRef = Field(description="The LLM model to use with the agent.")
+    verbose: bool = Field(default=False, description="Set the verbosity of the agent's logging.")
+    description: str = Field(description="The description of this function's use.")
+    log_response_max_chars: PositiveInt = Field(
+        default=1000, description="Maximum number of characters to display in logs when logging responses.")

nat/data_models/authentication.py

@@ -177,6 +177,26 @@ Credential = typing.Annotated[
 ]
 
 
+class TokenValidationResult(BaseModel):
+    """
+    Standard result for Bearer Token Validation.
+    """
+    model_config = ConfigDict(extra="forbid")
+
+    client_id: str | None = Field(description="OAuth2 client identifier")
+    scopes: list[str] | None = Field(default=None, description="List of granted scopes (introspection only)")
+    expires_at: int | None = Field(default=None, description="Token expiration time (Unix timestamp)")
+    audience: list[str] | None = Field(default=None, description="Token audiences (aud claim)")
+    subject: str | None = Field(default=None, description="Token subject (sub claim)")
+    issuer: str | None = Field(default=None, description="Token issuer (iss claim)")
+    token_type: str = Field(description="Token type")
+    active: bool | None = Field(default=True, description="Token active status")
+    nbf: int | None = Field(default=None, description="Not before time (Unix timestamp)")
+    iat: int | None = Field(default=None, description="Issued at time (Unix timestamp)")
+    jti: str | None = Field(default=None, description="JWT ID")
+    username: str | None = Field(default=None, description="Username (introspection only)")
+
+
 class AuthResult(BaseModel):
     """
     Represents the result of an authentication process.
@@ -229,3 +249,21 @@ class AuthResult(BaseModel):
                 target_kwargs.setdefault(k, {}).update(v)
             else:
                 target_kwargs[k] = v
+
+
+class AuthReason(str, Enum):
+    """
+    Why the caller is asking for auth now.
+    """
+    NORMAL = "normal"
+    RETRY_AFTER_401 = "retry_after_401"
+
+
+class AuthRequest(BaseModel):
+    """
+    Authentication request payload for provider.authenticate(...).
+    """
+    model_config = ConfigDict(extra="forbid")
+
+    reason: AuthReason = Field(default=AuthReason.NORMAL, description="Purpose of this auth attempt.")
+    www_authenticate: str | None = Field(default=None, description="Raw WWW-Authenticate header from a 401 response.")

nat/data_models/component.py

@@ -27,6 +27,7 @@ class ComponentEnum(StrEnum):
     EVALUATOR = "evaluator"
     FRONT_END = "front_end"
     FUNCTION = "function"
+    FUNCTION_GROUP = "function_group"
     TTC_STRATEGY = "ttc_strategy"
     LLM_CLIENT = "llm_client"
     LLM_PROVIDER = "llm_provider"
@@ -47,6 +48,7 @@ class ComponentGroup(StrEnum):
     AUTHENTICATION = "authentication"
     EMBEDDERS = "embedders"
     FUNCTIONS = "functions"
+    FUNCTION_GROUPS = "function_groups"
     TTC_STRATEGIES = "ttc_strategies"
     LLMS = "llms"
     MEMORY = "memory"

nat/data_models/component_ref.py

@@ -102,6 +102,17 @@ class FunctionRef(ComponentRef):
         return ComponentGroup.FUNCTIONS
 
 
+class FunctionGroupRef(ComponentRef):
+    """
+    A reference to a function group in a NAT configuration object.
+    """
+
+    @property
+    @override
+    def component_group(self):
+        return ComponentGroup.FUNCTION_GROUPS
+
+
 class LLMRef(ComponentRef):
     """
     A reference to an LLM in a NAT configuration object.

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
@@ -185,10 +189,14 @@ class GeneralConfig(BaseModel):
 
     model_config = ConfigDict(protected_namespaces=())
 
-    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/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/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,
+    )