nvidia-nat 1.3.0a20250822__py3-none-any.whl → 1.3.0a20250824__py3-none-any.whl

nat/data_models/temperature_mixin.py

@@ -0,0 +1,36 @@
+# 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 re
+
+from pydantic import BaseModel
+from pydantic import Field
+
+from nat.data_models.model_gated_field_mixin import ModelGatedFieldMixin
+
+_UNSUPPORTED_TEMPERATURE_MODELS = (re.compile(r"gpt-?5", re.IGNORECASE), )
+
+
+class TemperatureMixin(
+        BaseModel,
+        ModelGatedFieldMixin[float],
+        field_name="temperature",
+        default_if_supported=0.0,
+        unsupported_models=_UNSUPPORTED_TEMPERATURE_MODELS,
+):
+    """
+    Mixin class for temperature configuration.
+    """
+    temperature: float | None = Field(default=None, ge=0.0, le=1.0, description="Sampling temperature in [0, 1].")

nat/data_models/top_p_mixin.py

@@ -0,0 +1,36 @@
+# 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 re
+
+from pydantic import BaseModel
+from pydantic import Field
+
+from nat.data_models.model_gated_field_mixin import ModelGatedFieldMixin
+
+_UNSUPPORTED_TOP_P_MODELS = (re.compile(r"gpt-?5", re.IGNORECASE), )
+
+
+class TopPMixin(
+        BaseModel,
+        ModelGatedFieldMixin[float],
+        field_name="top_p",
+        default_if_supported=1.0,
+        unsupported_models=_UNSUPPORTED_TOP_P_MODELS,
+):
+    """
+    Mixin class for top-p configuration.
+    """
+    top_p: float | None = Field(default=None, ge=0.0, le=1.0, description="Top-p for distribution sampling.")

nat/embedder/register.py

@@ -13,7 +13,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# pylint: disable=unused-import
 # flake8: noqa
 # isort:skip_file
 

nat/eval/dataset_handler/dataset_handler.py

@@ -146,13 +146,12 @@ class DatasetHandler:
             # When num_passes is specified, always use concurrency * num_passes
             # This respects the user's intent for exact number of passes
             target_size = self.concurrency * self.num_passes
+        # When num_passes = 0, use the largest multiple of concurrency <= original_size
+        # If original_size < concurrency, we need at least concurrency rows
+        elif original_size >= self.concurrency:
+            target_size = (original_size // self.concurrency) * self.concurrency
         else:
-            # When num_passes = 0, use the largest multiple of concurrency <= original_size
-            # If original_size < concurrency, we need at least concurrency rows
-            if original_size >= self.concurrency:
-                target_size = (original_size // self.concurrency) * self.concurrency
-            else:
-                target_size = self.concurrency
+            target_size = self.concurrency
 
         if target_size == 0:
             raise ValueError("Input dataset too small for even one batch at given concurrency.")

nat/eval/evaluate.py

@@ -42,7 +42,7 @@ from nat.runtime.session import SessionManager
 logger = logging.getLogger(__name__)
 
 
-class EvaluationRun:  # pylint: disable=too-many-public-methods
+class EvaluationRun:
     """
     Instantiated for each evaluation run and used to store data for that single run.
 
@@ -319,7 +319,7 @@ class EvaluationRun:  # pylint: disable=too-many-public-methods
             except Exception as e:
                 logger.exception("Failed to delete old job directory: %s: %s", dir_to_delete, e, exc_info=True)
 
-    def write_output(self, dataset_handler: DatasetHandler, profiler_results: ProfilerResults):  # pylint: disable=unused-argument  # noqa: E501
+    def write_output(self, dataset_handler: DatasetHandler, profiler_results: ProfilerResults):
         workflow_output_file = self.eval_config.general.output_dir / "workflow_output.json"
         workflow_output_file.parent.mkdir(parents=True, exist_ok=True)
 
@@ -511,12 +511,11 @@ class EvaluationRun:  # pylint: disable=too-many-public-methods
             with self.eval_trace_context.evaluation_context():
                 if self.config.endpoint:
                     await self.run_workflow_remote()
-                else:
-                    if not self.config.skip_workflow:
-                        if session_manager is None:
-                            session_manager = SessionManager(eval_workflow.build(),
-                                                             max_concurrency=self.eval_config.general.max_concurrency)
-                        await self.run_workflow_local(session_manager)
+                elif not self.config.skip_workflow:
+                    if session_manager is None:
+                        session_manager = SessionManager(eval_workflow.build(),
+                                                         max_concurrency=self.eval_config.general.max_concurrency)
+                    await self.run_workflow_local(session_manager)
 
                 # Evaluate
                 evaluators = {name: eval_workflow.get_evaluator(name) for name in self.eval_config.evaluators}

nat/eval/rag_evaluator/register.py

@@ -73,7 +73,7 @@ class RagasEvaluatorConfig(EvaluatorBaseConfig, name="ragas"):
         if isinstance(self.metric, str):
             return self.metric
         if isinstance(self.metric, dict) and self.metric:
-            return next(iter(self.metric.keys()))  # pylint: disable=no-member
+            return next(iter(self.metric.keys()))
         return ""
 
     @property
@@ -82,7 +82,7 @@ class RagasEvaluatorConfig(EvaluatorBaseConfig, name="ragas"):
         if isinstance(self.metric, str):
             return RagasMetricConfig()  # Default config when only a metric name is given
         if isinstance(self.metric, dict) and self.metric:
-            return next(iter(self.metric.values()))  # pylint: disable=no-member
+            return next(iter(self.metric.values()))
         return RagasMetricConfig()  # Default config when an invalid type is provided
 
 

nat/eval/register.py

@@ -14,7 +14,6 @@
 # limitations under the License.
 
 # flake8: noqa
-# pylint: disable=unused-import
 
 # Import evaluators which need to be automatically registered here
 from .rag_evaluator.register import register_ragas_evaluator

nat/eval/tunable_rag_evaluator/evaluate.py

@@ -13,7 +13,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import asyncio
 import logging
 from typing import Callable
 
@@ -23,7 +22,6 @@ from langchain.schema import HumanMessage
 from langchain.schema import SystemMessage
 from langchain_core.language_models import BaseChatModel
 from langchain_core.runnables import RunnableLambda
-from tqdm import tqdm
 
 from nat.eval.evaluator.base_evaluator import BaseEvaluator
 from nat.eval.evaluator.evaluator_model import EvalInputItem
@@ -31,7 +29,6 @@ from nat.eval.evaluator.evaluator_model import EvalOutputItem
 
 logger = logging.getLogger(__name__)
 
-# pylint: disable=line-too-long
 # flake8: noqa: E501
 
 

nat/eval/utils/weave_eval.py

@@ -31,7 +31,7 @@ if TYPE_CHECKING:
 logger = logging.getLogger(__name__)
 
 
-class WeaveEvaluationIntegration:  # pylint: disable=too-many-public-methods
+class WeaveEvaluationIntegration:
     """
     Class to handle all Weave integration functionality.
     """
@@ -47,8 +47,8 @@ class WeaveEvaluationIntegration:  # pylint: disable=too-many-public-methods
             from weave.flow.eval_imperative import EvaluationLogger
             from weave.flow.eval_imperative import ScoreLogger
             from weave.trace.context import weave_client_context
-            self.evaluation_logger_cls = EvaluationLogger  # pylint: disable=invalid-name
-            self.score_logger_cls = ScoreLogger  # pylint: disable=invalid-name
+            self.evaluation_logger_cls = EvaluationLogger
+            self.score_logger_cls = ScoreLogger
             self.weave_client_context = weave_client_context
             self.available = True
         except ImportError:

nat/experimental/test_time_compute/models/strategy_base.py

@@ -17,9 +17,10 @@ from abc import ABC
 from abc import abstractmethod
 
 from nat.builder.builder import Builder
-from nat.experimental.test_time_compute.models.ttc_item import TTCItem
-from nat.experimental.test_time_compute.models.stage_enums import StageTypeEnum, PipelineTypeEnum
 from nat.data_models.ttc_strategy import TTCStrategyBaseConfig
+from nat.experimental.test_time_compute.models.stage_enums import PipelineTypeEnum
+from nat.experimental.test_time_compute.models.stage_enums import StageTypeEnum
+from nat.experimental.test_time_compute.models.ttc_item import TTCItem
 
 
 class StrategyBase(ABC):

nat/experimental/test_time_compute/register.py

@@ -13,7 +13,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# pylint: disable=unused-import
 # flake8: noqa
 
 from .editing import iterative_plan_refinement_editor

nat/experimental/test_time_compute/selection/llm_based_output_merging_selector.py

@@ -135,8 +135,6 @@ class LLMBasedOutputMergingSelector(StrategyBase):
             except Exception as e:
                 logger.error(f"Error parsing merged output: {e}")
                 raise ValueError("Failed to parse merged output.")
-        else:
-            merged_output = merged_output
 
         logger.info("Merged output: %s", str(merged_output))
 

nat/front_ends/fastapi/fastapi_front_end_plugin_worker.py

@@ -307,7 +307,7 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
         async def start_evaluation(request: EvaluateRequest, background_tasks: BackgroundTasks, http_request: Request):
             """Handle evaluation requests."""
 
-            async with session_manager.session(request=http_request):
+            async with session_manager.session(http_connection=http_request):
 
                 # if job_id is present and already exists return the job info
                 if request.job_id:
@@ -336,7 +336,7 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
             """Get the status of an evaluation job."""
             logger.info("Getting status for job %s", job_id)
 
-            async with session_manager.session(request=http_request):
+            async with session_manager.session(http_connection=http_request):
 
                 job = job_store.get_job(job_id)
                 if not job:
@@ -349,7 +349,7 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
             """Get the status of the last created evaluation job."""
             logger.info("Getting last job status")
 
-            async with session_manager.session(request=http_request):
+            async with session_manager.session(http_connection=http_request):
 
                 job = job_store.get_last_job()
                 if not job:
@@ -361,7 +361,7 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
         async def get_jobs(http_request: Request, status: str | None = None) -> list[EvaluateStatusResponse]:
             """Get all jobs, optionally filtered by status."""
 
-            async with session_manager.session(request=http_request):
+            async with session_manager.session(http_connection=http_request):
 
                 if status is None:
                     logger.info("Getting all jobs")
@@ -522,9 +522,9 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
 
         workflow = session_manager.workflow
 
-        GenerateBodyType = workflow.input_schema  # pylint: disable=invalid-name
-        GenerateStreamResponseType = workflow.streaming_output_schema  # pylint: disable=invalid-name
-        GenerateSingleResponseType = workflow.single_output_schema  # pylint: disable=invalid-name
+        GenerateBodyType = workflow.input_schema
+        GenerateStreamResponseType = workflow.streaming_output_schema
+        GenerateSingleResponseType = workflow.single_output_schema
 
         # Append job_id and expiry_seconds to the input schema, this effectively makes these reserved keywords
         # Consider prefixing these with "nat_" to avoid conflicts
@@ -572,7 +572,7 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
 
                 response.headers["Content-Type"] = "application/json"
 
-                async with session_manager.session(request=request,
+                async with session_manager.session(http_connection=request,
                                                    user_authentication_callback=self._http_flow_handler.authenticate):
 
                     return await generate_single_response(None, session_manager, result_type=result_type)
@@ -583,7 +583,7 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
 
             async def get_stream(request: Request):
 
-                async with session_manager.session(request=request,
+                async with session_manager.session(http_connection=request,
                                                    user_authentication_callback=self._http_flow_handler.authenticate):
 
                     return StreamingResponse(headers={"Content-Type": "text/event-stream; charset=utf-8"},
@@ -618,7 +618,7 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
 
                 response.headers["Content-Type"] = "application/json"
 
-                async with session_manager.session(request=request,
+                async with session_manager.session(http_connection=request,
                                                    user_authentication_callback=self._http_flow_handler.authenticate):
 
                     return await generate_single_response(payload, session_manager, result_type=result_type)
@@ -632,7 +632,7 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
 
             async def post_stream(request: Request, payload: request_type):
 
-                async with session_manager.session(request=request,
+                async with session_manager.session(http_connection=request,
                                                    user_authentication_callback=self._http_flow_handler.authenticate):
 
                     return StreamingResponse(headers={"Content-Type": "text/event-stream; charset=utf-8"},
@@ -677,7 +677,7 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
                 # Check if streaming is requested
                 stream_requested = getattr(payload, 'stream', False)
 
-                async with session_manager.session(request=request):
+                async with session_manager.session(http_connection=request):
                     if stream_requested:
                         # Return streaming response
                         return StreamingResponse(headers={"Content-Type": "text/event-stream; charset=utf-8"},
@@ -688,42 +688,41 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
                                                      step_adaptor=self.get_step_adaptor(),
                                                      result_type=ChatResponseChunk,
                                                      output_type=ChatResponseChunk))
-                    else:
-                        # Return single response - check if workflow supports non-streaming
-                        try:
+
+                    # Return single response - check if workflow supports non-streaming
+                    try:
+                        response.headers["Content-Type"] = "application/json"
+                        return await generate_single_response(payload, session_manager, result_type=ChatResponse)
+                    except ValueError as e:
+                        if "Cannot get a single output value for streaming workflows" in str(e):
+                            # Workflow only supports streaming, but client requested non-streaming
+                            # Fall back to streaming and collect the result
+                            chunks = []
+                            async for chunk_str in generate_streaming_response_as_str(
+                                    payload,
+                                    session_manager=session_manager,
+                                    streaming=True,
+                                    step_adaptor=self.get_step_adaptor(),
+                                    result_type=ChatResponseChunk,
+                                    output_type=ChatResponseChunk):
+                                if chunk_str.startswith("data: ") and not chunk_str.startswith("data: [DONE]"):
+                                    chunk_data = chunk_str[6:].strip()  # Remove "data: " prefix
+                                    if chunk_data:
+                                        try:
+                                            chunk_json = ChatResponseChunk.model_validate_json(chunk_data)
+                                            if (chunk_json.choices and len(chunk_json.choices) > 0
+                                                    and chunk_json.choices[0].delta
+                                                    and chunk_json.choices[0].delta.content is not None):
+                                                chunks.append(chunk_json.choices[0].delta.content)
+                                        except Exception:
+                                            continue
+
+                            # Create a single response from collected chunks
+                            content = "".join(chunks)
+                            single_response = ChatResponse.from_string(content)
                             response.headers["Content-Type"] = "application/json"
-                            return await generate_single_response(payload, session_manager, result_type=ChatResponse)
-                        except ValueError as e:
-                            if "Cannot get a single output value for streaming workflows" in str(e):
-                                # Workflow only supports streaming, but client requested non-streaming
-                                # Fall back to streaming and collect the result
-                                chunks = []
-                                async for chunk_str in generate_streaming_response_as_str(
-                                        payload,
-                                        session_manager=session_manager,
-                                        streaming=True,
-                                        step_adaptor=self.get_step_adaptor(),
-                                        result_type=ChatResponseChunk,
-                                        output_type=ChatResponseChunk):
-                                    if chunk_str.startswith("data: ") and not chunk_str.startswith("data: [DONE]"):
-                                        chunk_data = chunk_str[6:].strip()  # Remove "data: " prefix
-                                        if chunk_data:
-                                            try:
-                                                chunk_json = ChatResponseChunk.model_validate_json(chunk_data)
-                                                if (chunk_json.choices and len(chunk_json.choices) > 0
-                                                        and chunk_json.choices[0].delta
-                                                        and chunk_json.choices[0].delta.content is not None):
-                                                    chunks.append(chunk_json.choices[0].delta.content)
-                                            except Exception:
-                                                continue
-
-                                # Create a single response from collected chunks
-                                content = "".join(chunks)
-                                single_response = ChatResponse.from_string(content)
-                                response.headers["Content-Type"] = "application/json"
-                                return single_response
-                            else:
-                                raise
+                            return single_response
+                        raise
 
             return post_openai_api_compatible
 
@@ -758,7 +757,7 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
                     http_request: Request) -> AsyncGenerateResponse | AsyncGenerationStatusResponse:
                 """Handle async generation requests."""
 
-                async with session_manager.session(request=http_request):
+                async with session_manager.session(http_connection=http_request):
 
                     # if job_id is present and already exists return the job info
                     if request.job_id:
@@ -804,7 +803,7 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
             """Get the status of an async job."""
             logger.info("Getting status for job %s", job_id)
 
-            async with session_manager.session(request=http_request):
+            async with session_manager.session(http_connection=http_request):
 
                 job = job_store.get_job(job_id)
                 if not job:

nat/front_ends/fastapi/message_handler.py

@@ -86,7 +86,7 @@ class WebSocketMessageHandler:
 
     async def __aexit__(self, exc_type, exc_value, traceback) -> None:
 
-        # TODO: Handle the exit  # pylint: disable=fixme
+        # TODO: Handle the exit
         pass
 
     async def run(self) -> None:
@@ -105,12 +105,10 @@ class WebSocketMessageHandler:
                 if (isinstance(validated_message, WebSocketUserMessage)):
                     await self.process_workflow_request(validated_message)
 
-                elif isinstance(
-                        validated_message,
-                    (  # noqa: E131
-                        WebSocketSystemResponseTokenMessage,
-                        WebSocketSystemIntermediateStepMessage,
-                        WebSocketSystemInteractionMessage)):
+                elif isinstance(validated_message,
+                                (WebSocketSystemResponseTokenMessage,
+                                 WebSocketSystemIntermediateStepMessage,
+                                 WebSocketSystemInteractionMessage)):
                     # These messages are already handled by self.create_websocket_message(data_model=value, …)
                     # No further processing is needed here.
                     pass
@@ -119,11 +117,9 @@ class WebSocketMessageHandler:
                     user_content = await self.process_user_message_content(validated_message)
                     self._user_interaction_response.set_result(user_content)
             except (asyncio.CancelledError, WebSocketDisconnect):
-                # TODO: Handle the disconnect  # pylint: disable=fixme
+                # TODO: Handle the disconnect
                 break
 
-        return None
-
     async def process_user_message_content(
             self, user_content: WebSocketUserMessage | WebSocketUserInteractionResponseMessage) -> BaseModel | None:
         """
@@ -162,12 +158,13 @@ class WebSocketMessageHandler:
 
             if isinstance(content, TextContent) and (self._running_workflow_task is None):
 
-                def _done_callback(task: asyncio.Task):  # pylint: disable=unused-argument
+                def _done_callback(task: asyncio.Task):
                     self._running_workflow_task = None
 
                 self._running_workflow_task = asyncio.create_task(
-                    self._run_workflow(content.text,
-                                       self._conversation_id,
+                    self._run_workflow(payload=content.text,
+                                       user_message_id=self._message_parent_id,
+                                       conversation_id=self._conversation_id,
                                        result_type=self._schema_output_mapping[self._workflow_schema_type],
                                        output_type=self._schema_output_mapping[
                                            self._workflow_schema_type])).add_done_callback(_done_callback)
@@ -290,14 +287,16 @@ class WebSocketMessageHandler:
 
     async def _run_workflow(self,
                             payload: typing.Any,
+                            user_message_id: str | None = None,
                             conversation_id: str | None = None,
                             result_type: type | None = None,
                             output_type: type | None = None) -> None:
 
         try:
             async with self._session_manager.session(
+                    user_message_id=user_message_id,
                     conversation_id=conversation_id,
-                    request=self._socket,
+                    http_connection=self._socket,
                     user_input_callback=self.human_interaction_callback,
                     user_authentication_callback=(self._flow_handler.authenticate
                                                   if self._flow_handler else None)) as session:

nat/front_ends/fastapi/message_validator.py

@@ -232,7 +232,7 @@ class MessageValidator:
         """
         return data_model.parent_id or "root"
 
-    async def create_system_response_token_message(  # pylint: disable=R0917:too-many-positional-arguments
+    async def create_system_response_token_message(
         self,
         message_type: Literal[WebSocketMessageType.RESPONSE_MESSAGE,
                               WebSocketMessageType.ERROR_MESSAGE] = WebSocketMessageType.RESPONSE_MESSAGE,
@@ -272,7 +272,7 @@ class MessageValidator:
             logger.error("Error creating system response token message: %s", str(e), exc_info=True)
             return None
 
-    async def create_system_intermediate_step_message(  # pylint: disable=R0917:too-many-positional-arguments
+    async def create_system_intermediate_step_message(
         self,
         message_type: Literal[WebSocketMessageType.INTERMEDIATE_STEP_MESSAGE] = (
             WebSocketMessageType.INTERMEDIATE_STEP_MESSAGE),
@@ -311,7 +311,7 @@ class MessageValidator:
             logger.error("Error creating system intermediate step message: %s", str(e), exc_info=True)
             return None
 
-    async def create_system_interaction_message(  # pylint: disable=R0917:too-many-positional-arguments
+    async def create_system_interaction_message(
         self,
         *,
         message_type: Literal[WebSocketMessageType.SYSTEM_INTERACTION_MESSAGE] = (
@@ -323,7 +323,7 @@ class MessageValidator:
         content: HumanPrompt,
         status: WebSocketMessageStatus = WebSocketMessageStatus.IN_PROGRESS,
         timestamp: str = str(datetime.datetime.now(datetime.timezone.utc))
-    ) -> WebSocketSystemInteractionMessage | None:  # noqa: E125 continuation line with same indent as next logical line
+    ) -> WebSocketSystemInteractionMessage | None:
         """
         Creates a system interaction message with default values.
 

nat/front_ends/fastapi/step_adaptor.py

@@ -289,7 +289,7 @@ class StepAdaptor:
 
         return event
 
-    def process(self, step: IntermediateStep) -> ResponseSerializable | None:  # pylint: disable=R1710
+    def process(self, step: IntermediateStep) -> ResponseSerializable | None:
 
         # Track the chunk
         self._history.append(step)

nat/front_ends/register.py

@@ -13,7 +13,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# pylint: disable=unused-import
 # flake8: noqa
 # isort:skip_file
 

nat/llm/aws_bedrock_llm.py

@@ -22,9 +22,10 @@ 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.retry_mixin import RetryMixin
+from nat.data_models.temperature_mixin import TemperatureMixin
 
 
-class AWSBedrockModelConfig(LLMBaseConfig, RetryMixin, name="aws_bedrock"):
+class AWSBedrockModelConfig(LLMBaseConfig, RetryMixin, TemperatureMixin, name="aws_bedrock"):
     """An AWS Bedrock llm provider to be used with an LLM client."""
 
     model_config = ConfigDict(protected_namespaces=())
@@ -33,7 +34,6 @@ class AWSBedrockModelConfig(LLMBaseConfig, RetryMixin, name="aws_bedrock"):
     model_name: str = Field(validation_alias=AliasChoices("model_name", "model"),
                             serialization_alias="model",
                             description="The model name for the hosted AWS Bedrock.")
-    temperature: float = Field(default=0.0, ge=0.0, le=1.0, description="Sampling temperature in [0, 1].")
     max_tokens: int | None = Field(default=1024,
                                    gt=0,
                                    description="Maximum number of tokens to generate."
@@ -52,6 +52,6 @@ class AWSBedrockModelConfig(LLMBaseConfig, RetryMixin, name="aws_bedrock"):
 
 
 @register_llm_provider(config_type=AWSBedrockModelConfig)
-async def aws_bedrock_model(llm_config: AWSBedrockModelConfig, builder: Builder):
+async def aws_bedrock_model(llm_config: AWSBedrockModelConfig, _builder: Builder):
 
     yield LLMProviderInfo(config=llm_config, description="A AWS Bedrock model for use with an LLM client.")

nat/llm/azure_openai_llm.py

@@ -22,9 +22,11 @@ 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.retry_mixin import RetryMixin
+from nat.data_models.temperature_mixin import TemperatureMixin
+from nat.data_models.top_p_mixin import TopPMixin
 
 
-class AzureOpenAIModelConfig(LLMBaseConfig, RetryMixin, name="azure_openai"):
+class AzureOpenAIModelConfig(LLMBaseConfig, RetryMixin, TemperatureMixin, TopPMixin, name="azure_openai"):
     """An Azure OpenAI LLM provider to be used with an LLM client."""
 
     model_config = ConfigDict(protected_namespaces=(), extra="allow")
@@ -38,10 +40,7 @@ class AzureOpenAIModelConfig(LLMBaseConfig, RetryMixin, name="azure_openai"):
     azure_deployment: str = Field(validation_alias=AliasChoices("azure_deployment", "model_name", "model"),
                                   serialization_alias="azure_deployment",
                                   description="The Azure OpenAI hosted model/deployment name.")
-    temperature: float = Field(default=0.0, description="Sampling temperature in [0, 1].")
-    top_p: float = Field(default=1.0, description="Top-p for distribution sampling.")
     seed: int | None = Field(default=None, description="Random seed to set for generation.")
-    max_retries: int = Field(default=10, description="The max number of retries for the request.")
 
 
 @register_llm_provider(config_type=AzureOpenAIModelConfig)

nat/llm/nim_llm.py

@@ -23,9 +23,11 @@ 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.retry_mixin import RetryMixin
+from nat.data_models.temperature_mixin import TemperatureMixin
+from nat.data_models.top_p_mixin import TopPMixin
 
 
-class NIMModelConfig(LLMBaseConfig, RetryMixin, name="nim"):
+class NIMModelConfig(LLMBaseConfig, RetryMixin, TemperatureMixin, TopPMixin, name="nim"):
     """An NVIDIA Inference Microservice (NIM) llm provider to be used with an LLM client."""
 
     model_config = ConfigDict(protected_namespaces=())
@@ -35,12 +37,10 @@ class NIMModelConfig(LLMBaseConfig, RetryMixin, name="nim"):
     model_name: str = Field(validation_alias=AliasChoices("model_name", "model"),
                             serialization_alias="model",
                             description="The model name for the hosted NIM.")
-    temperature: float = Field(default=0.0, description="Sampling temperature in [0, 1].")
-    top_p: float = Field(default=1.0, description="Top-p for distribution sampling.")
     max_tokens: PositiveInt = Field(default=300, description="Maximum number of tokens to generate.")
 
 
 @register_llm_provider(config_type=NIMModelConfig)
-async def nim_model(llm_config: NIMModelConfig, builder: Builder):
+async def nim_model(llm_config: NIMModelConfig, _builder: Builder):
 
     yield LLMProviderInfo(config=llm_config, description="A NIM model for use with an LLM client.")