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

nat/front_ends/mcp/memory_profiler.py

@@ -0,0 +1,320 @@
+# 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.
+"""Memory profiling utilities for MCP frontend."""
+
+import gc
+import logging
+import tracemalloc
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+
+class MemoryProfiler:
+    """Memory profiler for tracking memory usage and potential leaks."""
+
+    def __init__(self, enabled: bool = False, log_interval: int = 50, top_n: int = 10, log_level: str = "DEBUG"):
+        """Initialize the memory profiler.
+
+        Args:
+            enabled: Whether memory profiling is enabled
+            log_interval: Log stats every N requests
+            top_n: Number of top allocations to log
+            log_level: Log level for memory profiling output (e.g., "DEBUG", "INFO")
+        """
+        self.enabled = enabled
+        # normalize interval to avoid modulo-by-zero
+        self.log_interval = max(1, int(log_interval))
+        self.top_n = top_n
+        self.log_level = getattr(logging, log_level.upper(), logging.DEBUG)
+        self.request_count = 0
+        self.baseline_snapshot = None
+
+        # Track whether this instance started tracemalloc (to avoid resetting external tracing)
+        self._we_started_tracemalloc = False
+
+        if self.enabled:
+            logger.info("Memory profiling ENABLED (interval=%d, top_n=%d, log_level=%s)",
+                        self.log_interval,
+                        top_n,
+                        log_level)
+            try:
+                if not tracemalloc.is_tracing():
+                    tracemalloc.start()
+                    self._we_started_tracemalloc = True
+                # Take baseline snapshot
+                gc.collect()
+                self.baseline_snapshot = tracemalloc.take_snapshot()
+            except RuntimeError as e:
+                logger.warning("tracemalloc unavailable or not tracing: %s", e)
+        else:
+            logger.info("Memory profiling DISABLED")
+
+    def _log(self, message: str, *args: Any) -> None:
+        """Log a message at the configured log level.
+
+        Args:
+            message: Log message format string
+            args: Arguments for the format string
+        """
+        logger.log(self.log_level, message, *args)
+
+    def on_request_complete(self) -> None:
+        """Called after each request completes."""
+        if not self.enabled:
+            return
+        self.request_count += 1
+        if self.request_count % self.log_interval == 0:
+            self.log_memory_stats()
+
+    def _ensure_tracing(self) -> bool:
+        """Ensure tracemalloc is running if we started it originally.
+
+        Returns:
+            True if tracemalloc is active, False otherwise
+        """
+        if tracemalloc.is_tracing():
+            return True
+
+        # Only restart if we started it originally (respect external control)
+        if not self._we_started_tracemalloc:
+            return False
+
+        # Attempt to restart
+        try:
+            logger.warning("tracemalloc was stopped externally; restarting (we started it originally)")
+            tracemalloc.start()
+
+            # Reset baseline since old tracking data is lost
+            gc.collect()
+            self.baseline_snapshot = tracemalloc.take_snapshot()
+            logger.info("Baseline snapshot reset after tracemalloc restart")
+
+            return True
+        except RuntimeError as e:
+            logger.error("Failed to restart tracemalloc: %s", e)
+            return False
+
+    def _safe_traced_memory(self) -> tuple[float, float] | None:
+        """Return (current, peak usage in MB) if tracemalloc is active, else None."""
+        if not self._ensure_tracing():
+            return None
+
+        try:
+            current, peak = tracemalloc.get_traced_memory()
+            megabyte = (1 << 20)
+            return (current / megabyte, peak / megabyte)
+        except RuntimeError:
+            return None
+
+    def _safe_snapshot(self) -> tracemalloc.Snapshot | None:
+        """Return a tracemalloc Snapshot if available, else None."""
+        if not self._ensure_tracing():
+            return None
+
+        try:
+            return tracemalloc.take_snapshot()
+        except RuntimeError:
+            return None
+
+    def log_memory_stats(self) -> dict[str, Any]:
+        """Log current memory statistics and return them."""
+        if not self.enabled:
+            return {}
+
+        # Force garbage collection first
+        gc.collect()
+
+        # Get current memory usage
+        mem = self._safe_traced_memory()
+        if mem is None:
+            logger.info("tracemalloc is not active; cannot collect memory stats.")
+            # still return structural fields
+            stats = {
+                "request_count": self.request_count,
+                "current_memory_mb": None,
+                "peak_memory_mb": None,
+                "active_intermediate_managers": self._safe_intermediate_step_manager_count(),
+                "outstanding_steps": self._safe_outstanding_step_count(),
+                "active_exporters": self._safe_exporter_count(),
+                "isolated_exporters": self._safe_isolated_exporter_count(),
+                "subject_instances": self._count_instances_of_type("Subject"),
+            }
+            return stats
+
+        current_mb, peak_mb = mem
+
+        # Take snapshot and compare to baseline
+        snapshot = self._safe_snapshot()
+
+        # Track BaseExporter instances (observability layer)
+        exporter_count = self._safe_exporter_count()
+        isolated_exporter_count = self._safe_isolated_exporter_count()
+
+        # Track Subject instances (event streams)
+        subject_count = self._count_instances_of_type("Subject")
+
+        stats = {
+            "request_count": self.request_count,
+            "current_memory_mb": round(current_mb, 2),
+            "peak_memory_mb": round(peak_mb, 2),
+            "active_intermediate_managers": self._safe_intermediate_step_manager_count(),
+            "outstanding_steps": self._safe_outstanding_step_count(),
+            "active_exporters": exporter_count,
+            "isolated_exporters": isolated_exporter_count,
+            "subject_instances": subject_count,
+        }
+
+        self._log("=" * 80)
+        self._log("MEMORY PROFILE AFTER %d REQUESTS:", self.request_count)
+        self._log("  Current Memory: %.2f MB", current_mb)
+        self._log("  Peak Memory: %.2f MB", peak_mb)
+        self._log("")
+        self._log("NAT COMPONENT INSTANCES:")
+        self._log("  IntermediateStepManagers: %d active (%d outstanding steps)",
+                  stats["active_intermediate_managers"],
+                  stats["outstanding_steps"])
+        self._log("  BaseExporters: %d active (%d isolated)", stats["active_exporters"], stats["isolated_exporters"])
+        self._log("  Subject (event streams): %d instances", stats["subject_instances"])
+
+        # Show top allocations
+        if snapshot is None:
+            self._log("tracemalloc snapshot unavailable.")
+        else:
+            if self.baseline_snapshot:
+                self._log("TOP %d MEMORY GROWTH SINCE BASELINE:", self.top_n)
+                top_stats = snapshot.compare_to(self.baseline_snapshot, 'lineno')
+            else:
+                self._log("TOP %d MEMORY ALLOCATIONS:", self.top_n)
+                top_stats = snapshot.statistics('lineno')
+
+            for i, stat in enumerate(top_stats[:self.top_n], 1):
+                self._log("  #%d: %s", i, stat)
+
+        self._log("=" * 80)
+
+        return stats
+
+    def _count_instances_of_type(self, type_name: str) -> int:
+        """Count instances of a specific type in memory."""
+        count = 0
+        for obj in gc.get_objects():
+            try:
+                if type(obj).__name__ == type_name:
+                    count += 1
+            except Exception:
+                pass
+        return count
+
+    def _safe_exporter_count(self) -> int:
+        try:
+            from nat.observability.exporter.base_exporter import BaseExporter
+            return BaseExporter.get_active_instance_count()
+        except Exception as e:
+            logger.debug("Could not get BaseExporter stats: %s", e)
+            return 0
+
+    def _safe_isolated_exporter_count(self) -> int:
+        try:
+            from nat.observability.exporter.base_exporter import BaseExporter
+            return BaseExporter.get_isolated_instance_count()
+        except Exception:
+            return 0
+
+    def _safe_intermediate_step_manager_count(self) -> int:
+        try:
+            from nat.builder.intermediate_step_manager import IntermediateStepManager
+            # len() is atomic in CPython, but catch RuntimeError just in case
+            try:
+                return IntermediateStepManager.get_active_instance_count()
+            except RuntimeError:
+                # Set was modified during len() - very rare
+                logger.debug("Set changed during count, returning 0")
+                return 0
+        except Exception as e:
+            logger.debug("Could not get IntermediateStepManager stats: %s", e)
+            return 0
+
+    def _safe_outstanding_step_count(self) -> int:
+        """Get total outstanding steps across all active IntermediateStepManager instances."""
+        try:
+            from nat.builder.intermediate_step_manager import IntermediateStepManager
+
+            # Make a snapshot to avoid "Set changed size during iteration" if GC runs
+            try:
+                instances_snapshot = list(IntermediateStepManager._active_instances)
+            except RuntimeError:
+                # Set changed during list() call - rare but possible
+                logger.debug("Set changed during snapshot, returning 0 for outstanding steps")
+                return 0
+
+            total_outstanding = 0
+            # Iterate through snapshot safely
+            for ref in instances_snapshot:
+                try:
+                    manager = ref()
+                    if manager is not None:
+                        total_outstanding += manager.get_outstanding_step_count()
+                except (ReferenceError, AttributeError):
+                    # Manager was GC'd or in invalid state - skip it
+                    continue
+            return total_outstanding
+        except Exception as e:
+            logger.debug("Could not get outstanding step count: %s", e)
+            return 0
+
+    def get_stats(self) -> dict[str, Any]:
+        """Get current memory statistics without logging."""
+        if not self.enabled:
+            return {"enabled": False}
+
+        mem = self._safe_traced_memory()
+        if mem is None:
+            return {
+                "enabled": True,
+                "request_count": self.request_count,
+                "current_memory_mb": None,
+                "peak_memory_mb": None,
+                "active_intermediate_managers": self._safe_intermediate_step_manager_count(),
+                "outstanding_steps": self._safe_outstanding_step_count(),
+                "active_exporters": self._safe_exporter_count(),
+                "isolated_exporters": self._safe_isolated_exporter_count(),
+                "subject_instances": self._count_instances_of_type("Subject"),
+            }
+
+        current_mb, peak_mb = mem
+        return {
+            "enabled": True,
+            "request_count": self.request_count,
+            "current_memory_mb": round(current_mb, 2),
+            "peak_memory_mb": round(peak_mb, 2),
+            "active_intermediate_managers": self._safe_intermediate_step_manager_count(),
+            "outstanding_steps": self._safe_outstanding_step_count(),
+            "active_exporters": self._safe_exporter_count(),
+            "isolated_exporters": self._safe_isolated_exporter_count(),
+            "subject_instances": self._count_instances_of_type("Subject"),
+        }
+
+    def reset_baseline(self) -> None:
+        """Reset the baseline snapshot to current state."""
+        if not self.enabled:
+            return
+        gc.collect()
+        snap = self._safe_snapshot()
+        if snap is None:
+            logger.info("Cannot reset baseline: tracemalloc is not active.")
+            return
+        self.baseline_snapshot = snap
+        logger.info("Memory profiling baseline reset at request %d", self.request_count)

nat/front_ends/mcp/tool_converter.py

@@ -18,9 +18,12 @@ import logging
 from inspect import Parameter
 from inspect import Signature
 from typing import TYPE_CHECKING
+from typing import Any
 
 from mcp.server.fastmcp import FastMCP
 from pydantic import BaseModel
+from pydantic.fields import FieldInfo
+from pydantic_core import PydanticUndefined
 
 from nat.builder.context import ContextState
 from nat.builder.function import Function
@@ -28,9 +31,45 @@ from nat.builder.function_base import FunctionBase
 
 if TYPE_CHECKING:
     from nat.builder.workflow import Workflow
+    from nat.front_ends.mcp.memory_profiler import MemoryProfiler
 
 logger = logging.getLogger(__name__)
 
+# Sentinel: marks "optional; let Pydantic supply default/factory"
+_USE_PYDANTIC_DEFAULT = object()
+
+
+def is_field_optional(field: FieldInfo) -> tuple[bool, Any]:
+    """Determine if a Pydantic field is optional and extract its default value for MCP signatures.
+
+    For MCP tool signatures, we need to distinguish:
+    - Required fields: marked with Parameter.empty
+    - Optional with concrete default: use that default
+    - Optional with factory: use sentinel so Pydantic can apply the factory later
+
+    Args:
+        field: The Pydantic FieldInfo to check
+
+    Returns:
+        A tuple of (is_optional, default_value):
+        - (False, Parameter.empty) for required fields
+        - (True, actual_default) for optional fields with explicit defaults
+        - (True, _USE_PYDANTIC_DEFAULT) for optional fields with default_factory
+    """
+    if field.is_required():
+        return False, Parameter.empty
+
+    # Field is optional - has either default or factory
+    if field.default is not PydanticUndefined:
+        return True, field.default
+
+    # Factory case: mark optional in signature but don't fabricate a value
+    if field.default_factory is not None:
+        return True, _USE_PYDANTIC_DEFAULT
+
+    # Rare corner case: non-required yet no default surfaced
+    return True, _USE_PYDANTIC_DEFAULT
+
 
 def create_function_wrapper(
     function_name: str,
@@ -38,6 +77,7 @@ def create_function_wrapper(
     schema: type[BaseModel],
     is_workflow: bool = False,
     workflow: 'Workflow | None' = None,
+    memory_profiler: 'MemoryProfiler | None' = None,
 ):
     """Create a wrapper function that exposes the actual parameters of a NAT Function as an MCP tool.
 
@@ -47,6 +87,7 @@ def create_function_wrapper(
         schema (type[BaseModel]): The input schema of the function
         is_workflow (bool): Whether the function is a Workflow
         workflow (Workflow | None): The parent workflow for observability context
+        memory_profiler: Optional memory profiler to track requests
 
     Returns:
         A wrapper function suitable for registration with MCP
@@ -76,12 +117,15 @@ def create_function_wrapper(
             # Get the field type and convert to appropriate Python type
             field_type = field.annotation
 
+            # Check if field is optional and get its default value
+            _is_optional, param_default = is_field_optional(field)
+
             # Add the parameter to our list
             parameters.append(
                 Parameter(
                     name=name,
                     kind=Parameter.KEYWORD_ONLY,
-                    default=Parameter.empty if field.is_required else None,
+                    default=param_default,
                     annotation=field_type,
                 ))
 
@@ -140,47 +184,46 @@ def create_function_wrapper(
                         result = await call_with_observability(lambda: function.ainvoke(chat_request, to_type=str))
                 else:
                     # Regular handling
-                    # Handle complex input schema - if we extracted fields from a nested schema,
-                    # we need to reconstruct the input
-                    if len(schema.model_fields) == 1 and len(parameters) > 1:
-                        # Get the field name from the original schema
-                        field_name = next(iter(schema.model_fields.keys()))
-                        field_type = schema.model_fields[field_name].annotation
-
-                        # If it's a pydantic model, we need to create an instance
-                        if field_type and hasattr(field_type, "model_validate"):
-                            # Create the nested object
-                            nested_obj = field_type.model_validate(kwargs)
-                            # Call with the nested object
-                            kwargs = {field_name: nested_obj}
+                    # Strip sentinel values so Pydantic can apply defaults/factories
+                    cleaned_kwargs = {k: v for k, v in kwargs.items() if v is not _USE_PYDANTIC_DEFAULT}
+
+                    # Always validate with the declared schema
+                    # This handles defaults, factories, nested models, validators, etc.
+                    model_input = schema.model_validate(cleaned_kwargs)
 
                     # Call the NAT function with the parameters - special handling for Workflow
                     if is_workflow:
-                        # For workflow with regular input, we'll assume the first parameter is the input
-                        input_value = list(kwargs.values())[0] if kwargs else ""
-
-                        # Workflows have a run method that is an async context manager
-                        # that returns a Runner
-                        async with function.run(input_value) as runner:
+                        # Workflows expect the model instance directly
+                        async with function.run(model_input) as runner:
                             # Get the result from the runner
                             result = await runner.result(to_type=str)
                     else:
-                        # Regular function call
-                        result = await call_with_observability(lambda: function.acall_invoke(**kwargs))
+                        # Regular function call - unpack the validated model
+                        result = await call_with_observability(lambda: function.acall_invoke(**model_input.model_dump())
+                                                               )
 
                 # Report completion
                 if ctx:
                     await ctx.report_progress(100, 100)
 
+                # Track request completion for memory profiling
+                if memory_profiler:
+                    memory_profiler.on_request_complete()
+
                 # Handle different result types for proper formatting
                 if isinstance(result, str):
                     return result
-                if isinstance(result, (dict, list)):
+                if isinstance(result, dict | list):
                     return json.dumps(result, default=str)
                 return str(result)
             except Exception as e:
                 if ctx:
                     ctx.error("Error calling function %s: %s", function_name, str(e))
+
+                # Track request completion even on error
+                if memory_profiler:
+                    memory_profiler.on_request_complete()
+
                 raise
 
         return wrapper_with_ctx
@@ -229,6 +272,9 @@ def get_function_description(function: FunctionBase) -> str:
         # Try to get anything that might be a description
         elif hasattr(config, "topic") and config.topic:
             function_description = config.topic
+        # Try to get description from the workflow config
+        elif hasattr(config, "workflow") and hasattr(config.workflow, "description") and config.workflow.description:
+            function_description = config.workflow.description
 
     elif isinstance(function, Function):
         function_description = function.description
@@ -239,7 +285,8 @@ def get_function_description(function: FunctionBase) -> str:
 def register_function_with_mcp(mcp: FastMCP,
                                function_name: str,
                                function: FunctionBase,
-                               workflow: 'Workflow | None' = None) -> None:
+                               workflow: 'Workflow | None' = None,
+                               memory_profiler: 'MemoryProfiler | None' = None) -> None:
     """Register a NAT Function as an MCP tool.
 
     Args:
@@ -247,6 +294,7 @@ def register_function_with_mcp(mcp: FastMCP,
         function_name: The name to register the function under
         function: The NAT Function to register
         workflow: The parent workflow for observability context (if available)
+        memory_profiler: Optional memory profiler to track requests
     """
     logger.info("Registering function %s with MCP", function_name)
 
@@ -264,5 +312,10 @@ def register_function_with_mcp(mcp: FastMCP,
     function_description = get_function_description(function)
 
     # Create and register the wrapper function with MCP
-    wrapper_func = create_function_wrapper(function_name, function, input_schema, is_workflow, workflow)
+    wrapper_func = create_function_wrapper(function_name,
+                                           function,
+                                           input_schema,
+                                           is_workflow,
+                                           workflow,
+                                           memory_profiler)
     mcp.tool(name=function_name, description=function_description)(wrapper_func)

nat/front_ends/simple_base/simple_front_end_plugin_base.py

@@ -35,6 +35,8 @@ class SimpleFrontEndPluginBase(FrontEndBase[FrontEndConfigT], ABC):
 
     async def run(self):
 
+        await self.pre_run()
+
         # Must yield the workflow function otherwise it cleans up
         async with WorkflowBuilder.from_config(config=self.full_config) as builder:
 
@@ -45,7 +47,7 @@ class SimpleFrontEndPluginBase(FrontEndBase[FrontEndConfigT], ABC):
 
                 click.echo(stream.getvalue())
 
-            workflow = builder.build()
+            workflow = await builder.build()
             session_manager = SessionManager(workflow)
             await self.run_workflow(session_manager)
 

nat/llm/aws_bedrock_llm.py

@@ -21,22 +21,25 @@ from nat.builder.builder import Builder
 from nat.builder.llm import LLMProviderInfo
 from nat.cli.register_workflow import register_llm_provider
 from nat.data_models.llm import LLMBaseConfig
+from nat.data_models.optimizable import OptimizableField
+from nat.data_models.optimizable import OptimizableMixin
+from nat.data_models.optimizable import SearchSpace
 from nat.data_models.retry_mixin import RetryMixin
-from nat.data_models.temperature_mixin import TemperatureMixin
 from nat.data_models.thinking_mixin import ThinkingMixin
-from nat.data_models.top_p_mixin import TopPMixin
 
 
-class AWSBedrockModelConfig(LLMBaseConfig, RetryMixin, TemperatureMixin, TopPMixin, ThinkingMixin, name="aws_bedrock"):
+class AWSBedrockModelConfig(LLMBaseConfig, RetryMixin, OptimizableMixin, ThinkingMixin, name="aws_bedrock"):
     """An AWS Bedrock llm provider to be used with an LLM client."""
 
-    model_config = ConfigDict(protected_namespaces=())
+    model_config = ConfigDict(protected_namespaces=(), extra="allow")
 
     # Completion parameters
-    model_name: str = Field(validation_alias=AliasChoices("model_name", "model"),
-                            serialization_alias="model",
-                            description="The model name for the hosted AWS Bedrock.")
-    max_tokens: int | None = Field(default=1024, gt=0, description="Maximum number of tokens to generate.")
+    model_name: str = OptimizableField(validation_alias=AliasChoices("model_name", "model"),
+                                       serialization_alias="model",
+                                       description="The model name for the hosted AWS Bedrock.")
+    max_tokens: int = OptimizableField(default=300,
+                                       description="Maximum number of tokens to generate.",
+                                       space=SearchSpace(high=2176, low=128, step=512))
     context_size: int | None = Field(
         default=1024,
         gt=0,
@@ -50,6 +53,16 @@ class AWSBedrockModelConfig(LLMBaseConfig, RetryMixin, TemperatureMixin, TopPMix
         default=None, description="Bedrock endpoint to use. Needed if you don't want to default to us-east-1 endpoint.")
     credentials_profile_name: str | None = Field(
         default=None, description="The name of the profile in the ~/.aws/credentials or ~/.aws/config files.")
+    temperature: float | None = OptimizableField(
+        default=None,
+        ge=0.0,
+        description="Sampling temperature to control randomness in the output.",
+        space=SearchSpace(high=0.9, low=0.1, step=0.2))
+    top_p: float | None = OptimizableField(default=None,
+                                           ge=0.0,
+                                           le=1.0,
+                                           description="Top-p for distribution sampling.",
+                                           space=SearchSpace(high=1.0, low=0.5, step=0.1))
 
 
 @register_llm_provider(config_type=AWSBedrockModelConfig)

nat/llm/azure_openai_llm.py

@@ -20,18 +20,17 @@ from pydantic import Field
 from nat.builder.builder import Builder
 from nat.builder.llm import LLMProviderInfo
 from nat.cli.register_workflow import register_llm_provider
+from nat.data_models.common import OptionalSecretStr
 from nat.data_models.llm import LLMBaseConfig
+from nat.data_models.optimizable import OptimizableField
+from nat.data_models.optimizable import SearchSpace
 from nat.data_models.retry_mixin import RetryMixin
-from nat.data_models.temperature_mixin import TemperatureMixin
 from nat.data_models.thinking_mixin import ThinkingMixin
-from nat.data_models.top_p_mixin import TopPMixin
 
 
 class AzureOpenAIModelConfig(
         LLMBaseConfig,
         RetryMixin,
-        TemperatureMixin,
-        TopPMixin,
         ThinkingMixin,
         name="azure_openai",
 ):
@@ -39,7 +38,7 @@ class AzureOpenAIModelConfig(
 
     model_config = ConfigDict(protected_namespaces=(), extra="allow")
 
-    api_key: str | None = Field(default=None, description="Azure OpenAI API key to interact with hosted model.")
+    api_key: OptionalSecretStr = Field(default=None, description="Azure OpenAI API key to interact with hosted model.")
     api_version: str = Field(default="2025-04-01-preview", description="Azure OpenAI API version.")
     azure_endpoint: str | None = Field(validation_alias=AliasChoices("azure_endpoint", "base_url"),
                                        serialization_alias="azure_endpoint",
@@ -49,6 +48,16 @@ class AzureOpenAIModelConfig(
                                   serialization_alias="azure_deployment",
                                   description="The Azure OpenAI hosted model/deployment name.")
     seed: int | None = Field(default=None, description="Random seed to set for generation.")
+    temperature: float | None = OptimizableField(
+        default=None,
+        ge=0.0,
+        description="Sampling temperature to control randomness in the output.",
+        space=SearchSpace(high=0.9, low=0.1, step=0.2))
+    top_p: float | None = OptimizableField(default=None,
+                                           ge=0.0,
+                                           le=1.0,
+                                           description="Top-p for distribution sampling.",
+                                           space=SearchSpace(high=1.0, low=0.5, step=0.1))
 
 
 @register_llm_provider(config_type=AzureOpenAIModelConfig)