nvidia-nat 1.3.0a20250827__py3-none-any.whl → 1.3.0a20250829__py3-none-any.whl

nat/data_models/thinking_mixin.py

@@ -0,0 +1,68 @@
+# 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.gated_field_mixin import GatedFieldMixin
+
+# The system prompt format for thinking is different for these, so we need to distinguish them here with two separate
+# regex patterns
+_NVIDIA_NEMOTRON_REGEX = re.compile(r"^nvidia/nvidia.*nemotron", re.IGNORECASE)
+_LLAMA_NEMOTRON_REGEX = re.compile(r"^nvidia/llama.*nemotron", re.IGNORECASE)
+_MODEL_KEYS = ("model_name", "model", "azure_deployment")
+
+
+class ThinkingMixin(
+        BaseModel,
+        GatedFieldMixin,
+        field_name="thinking",
+        default_if_supported=None,
+        keys=_MODEL_KEYS,
+        supported=(_NVIDIA_NEMOTRON_REGEX, _LLAMA_NEMOTRON_REGEX),
+):
+    """
+    Mixin class for thinking configuration. Only supported on Nemotron models.
+
+    Attributes:
+        thinking: Whether to enable thinking. Defaults to None when supported on the model.
+    """
+    thinking: bool | None = Field(
+        default=None,
+        description="Whether to enable thinking. Defaults to None when supported on the model.",
+        exclude=True,
+    )
+
+    @property
+    def thinking_system_prompt(self) -> str | None:
+        """
+        Returns the system prompt to use for thinking.
+        For NVIDIA Nemotron, returns "/think" if enabled, else "/no_think".
+        For Llama Nemotron, returns "detailed thinking on" if enabled, else "detailed thinking off".
+        If thinking is not supported on the model, returns None.
+
+        Returns:
+            str | None: The system prompt to use for thinking.
+        """
+        if self.thinking is None:
+            return None
+        for key in _MODEL_KEYS:
+            if hasattr(self, key):
+                if _NVIDIA_NEMOTRON_REGEX.match(getattr(self, key)):
+                    return "/think" if self.thinking else "/no_think"
+                elif _LLAMA_NEMOTRON_REGEX.match(getattr(self, key)):
+                    return f"detailed thinking {'on' if self.thinking else 'off'}"

nat/data_models/top_p_mixin.py

@@ -23,7 +23,7 @@ from nat.data_models.gated_field_mixin import GatedFieldMixin
 
 class TopPMixin(
         BaseModel,
-        GatedFieldMixin[float],
+        GatedFieldMixin,
         field_name="top_p",
         default_if_supported=1.0,
         keys=("model_name", "model", "azure_deployment"),

nat/eval/evaluate.py

@@ -168,17 +168,17 @@ class EvaluationRun:
                 intermediate_future = None
 
                 try:
-
                     # Start usage stats and intermediate steps collection in parallel
                     intermediate_future = pull_intermediate()
                     runner_result = runner.result()
                     base_output = await runner_result
                     intermediate_steps = await intermediate_future
                 except NotImplementedError as e:
+                    logger.error("Failed to run the workflow: %s", e)
                     # raise original error
-                    raise e
+                    raise
                 except Exception as e:
-                    logger.exception("Failed to run the workflow: %s", e, exc_info=True)
+                    logger.exception("Failed to run the workflow: %s", e)
                     # stop processing if a workflow error occurs
                     self.workflow_interrupted = True
 
@@ -317,7 +317,7 @@ class EvaluationRun:
                 logger.info("Deleting old job directory: %s", dir_to_delete)
                 shutil.rmtree(dir_to_delete)
             except Exception as e:
-                logger.exception("Failed to delete old job directory: %s: %s", dir_to_delete, e, exc_info=True)
+                logger.exception("Failed to delete old job directory: %s: %s", dir_to_delete, e)
 
     def write_output(self, dataset_handler: DatasetHandler, profiler_results: ProfilerResults):
         workflow_output_file = self.eval_config.general.output_dir / "workflow_output.json"
@@ -367,7 +367,7 @@ class EvaluationRun:
 
             await self.weave_eval.alog_score(eval_output, evaluator_name)
         except Exception as e:
-            logger.exception("An error occurred while running evaluator %s: %s", evaluator_name, e, exc_info=True)
+            logger.exception("An error occurred while running evaluator %s: %s", evaluator_name, e)
 
     async def run_evaluators(self, evaluators: dict[str, Any]):
         """Run all configured evaluators asynchronously."""
@@ -380,7 +380,7 @@ class EvaluationRun:
         try:
             await asyncio.gather(*tasks)
         except Exception as e:
-            logger.exception("An error occurred while running evaluators: %s", e, exc_info=True)
+            logger.error("An error occurred while running evaluators: %s", e)
             raise
         finally:
             # Finish prediction loggers in Weave

nat/eval/intermediate_step_adapter.py

@@ -40,7 +40,7 @@ class IntermediateStepAdapter:
             try:
                 validated_steps.append(IntermediateStep.model_validate(step_data))
             except Exception as e:
-                logger.exception("Validation failed for step: %r, Error: %s", step_data, e, exc_info=True)
+                logger.exception("Validation failed for step: %r, Error: %s", step_data, e)
         return validated_steps
 
     def serialize_intermediate_steps(self, intermediate_steps: list[IntermediateStep]) -> list[dict]:

nat/eval/rag_evaluator/evaluate.py

@@ -102,7 +102,7 @@ class RAGEvaluator:
         """Converts the ragas EvaluationResult to nat EvalOutput"""
 
         if not results_dataset:
-            logger.error("Ragas evaluation failed with no results")
+            logger.error("Ragas evaluation failed with no results", exc_info=True)
             return EvalOutput(average_score=0.0, eval_output_items=[])
 
         scores: list[dict[str, float]] = results_dataset.scores
@@ -169,7 +169,7 @@ class RAGEvaluator:
                                              _pbar=pbar)
         except Exception as e:
             # On exception we still continue with other evaluators. Log and return an avg_score of 0.0
-            logger.exception("Error evaluating ragas metric, Error: %s", e, exc_info=True)
+            logger.exception("Error evaluating ragas metric, Error: %s", e)
             results_dataset = None
         finally:
             pbar.close()

nat/eval/rag_evaluator/register.py

@@ -104,7 +104,7 @@ async def register_ragas_evaluator(config: RagasEvaluatorConfig, builder: EvalBu
             raise ValueError(message) from e
         except AttributeError as e:
             message = f"Ragas metric {metric_name} not found {e}."
-            logger.error(message)
+            logger.exception(message)
             return None
 
     async def evaluate_fn(eval_input: EvalInput) -> EvalOutput:

nat/eval/remote_workflow.py

@@ -74,7 +74,7 @@ class EvaluationRemoteWorkflowHandler:
                             if chunk_data.get("value"):
                                 final_response = chunk_data.get("value")
                         except json.JSONDecodeError as e:
-                            logger.error("Failed to parse generate response chunk: %s", e)
+                            logger.exception("Failed to parse generate response chunk: %s", e)
                             continue
                     elif line.startswith(INTERMEDIATE_DATA_PREFIX):
                         # This is an intermediate step
@@ -90,12 +90,12 @@ class EvaluationRemoteWorkflowHandler:
                                                                  payload=payload)
                             intermediate_steps.append(intermediate_step)
                         except (json.JSONDecodeError, ValidationError) as e:
-                            logger.error("Failed to parse intermediate step: %s", e)
+                            logger.exception("Failed to parse intermediate step: %s", e)
                             continue
 
         except aiohttp.ClientError as e:
             # Handle connection or HTTP-related errors
-            logger.error("Request failed for question %s: %s", question, e)
+            logger.exception("Request failed for question %s: %s", question, e)
             item.output_obj = None
             item.trajectory = []
             return

nat/eval/swe_bench_evaluator/evaluate.py

@@ -69,13 +69,13 @@ class SweBenchEvaluator:
         try:
             shutil.move(swe_bench_report_file, report_dir)
         except Exception as e:
-            logger.exception("Error moving report file: %s", e, exc_info=True)
+            logger.exception("Error moving report file: %s", e)
 
         try:
             dest_logs_dir = os.path.join(report_dir, 'logs')
             shutil.move(logs_dir, dest_logs_dir)
         except Exception as e:
-            logger.exception("Error moving logs directory: %s", e, exc_info=True)
+            logger.exception("Error moving logs directory: %s", e)
 
     def is_repo_supported(self, repo: str, version: str) -> bool:
         """Check if the repo is supported by swebench"""
@@ -106,7 +106,7 @@ class SweBenchEvaluator:
                     self._model_name_or_path = swebench_output.model_name_or_path
 
             except Exception as e:
-                logger.exception("Failed to parse EvalInputItem %s: %s", item.id, e, exc_info=True)
+                logger.exception("Failed to parse EvalInputItem %s: %s", item.id, e)
 
         # Filter out repos/version not supported by SWEBench
         supported_inputs = [
@@ -114,7 +114,7 @@ class SweBenchEvaluator:
         ]
 
         if not supported_inputs:
-            logger.error("No supported instances; nothing to evaluate")
+            logger.exception("No supported instances; nothing to evaluate")
             return None, None
 
         if len(supported_inputs) < len(swebench_inputs):
@@ -135,7 +135,7 @@ class SweBenchEvaluator:
         filtered_outputs = [output for output in swebench_outputs if output.instance_id in valid_instance_ids]
 
         if not filtered_outputs:
-            logger.error("No supported outputs; nothing to evaluate")
+            logger.error("No supported outputs; nothing to evaluate", exc_info=True)
             return None, None
 
         # Write SWEBenchOutput to file

nat/eval/trajectory_evaluator/evaluate.py

@@ -65,7 +65,7 @@ class TrajectoryEvaluator(BaseEvaluator):
                 prediction=generated_answer,
             )
         except Exception as e:
-            logger.exception("Error evaluating trajectory for question: %s, Error: %s", question, e, exc_info=True)
+            logger.exception("Error evaluating trajectory for question: %s, Error: %s", question, e)
             return EvalOutputItem(id=item.id, score=0.0, reasoning=f"Error evaluating trajectory: {e}")
 
         reasoning = {

nat/eval/tunable_rag_evaluator/evaluate.py

@@ -182,8 +182,8 @@ class TunableRagEvaluator(BaseEvaluator):
                     relevance_score = parsed_response["relevance_score"]
                     reasoning = parsed_response["reasoning"]
                 except KeyError as e:
-                    logger.error("Missing required keys in default scoring response: %s",
-                                 ", ".join(str(arg) for arg in e.args))
+                    logger.exception("Missing required keys in default scoring response: %s",
+                                     ", ".join(str(arg) for arg in e.args))
                     reasoning = f"Error in evaluator from parsing judge LLM response. Missing required key(s): {', '.join(str(arg) for arg in e.args)}"
 
                 coverage_weight = self.default_score_weights.get("coverage", 1 / 3)
@@ -215,7 +215,7 @@ class TunableRagEvaluator(BaseEvaluator):
                     reasoning = f"Error in evaluator from parsing judge LLM response. Missing required key(s): {', '.join(str(arg) for arg in e.args)}"
                     raise
         except (KeyError, ValueError) as e:
-            logger.error("Error parsing judge LLM response: %s", e)
+            logger.exception("Error parsing judge LLM response: %s", e)
             score = 0.0
             reasoning = "Error in evaluator from parsing judge LLM response."
 

nat/experimental/test_time_compute/functions/ttc_tool_orchestration_function.py

@@ -148,13 +148,13 @@ async def register_ttc_tool_orchestration_function(
                 result = await fn.acall_invoke(item.output)
                 return item, result, None
             except Exception as e:
-                logger.error(f"Error invoking function '{item.name}': {e}")
+                logger.exception(f"Error invoking function '{item.name}': {e}")
                 return item, None, str(e)
 
         tasks = []
         for item in ttc_items:
             if item.name not in function_map:
-                logger.error(f"Function '{item.name}' not found in function map.")
+                logger.error(f"Function '{item.name}' not found in function map.", exc_info=True)
                 item.output = f"Error: Function '{item.name}' not found in function map. Check your input"
             else:
                 fn = function_map[item.name]

nat/front_ends/fastapi/fastapi_front_end_controller.py

@@ -47,11 +47,11 @@ class _FastApiFrontEndController:
             self._server_background_task = asyncio.create_task(self._server.serve())
         except asyncio.CancelledError as e:
             error_message = f"Task error occurred while starting API server: {str(e)}"
-            logger.error(error_message, exc_info=True)
+            logger.error(error_message)
             raise RuntimeError(error_message) from e
         except Exception as e:
             error_message = f"Unexpected error occurred while starting API server: {str(e)}"
-            logger.error(error_message, exc_info=True)
+            logger.exception(error_message)
             raise RuntimeError(error_message) from e
 
     async def stop_server(self) -> None:
@@ -63,6 +63,6 @@ class _FastApiFrontEndController:
             self._server.should_exit = True
             await self._server_background_task
         except asyncio.CancelledError as e:
-            logger.error("Server shutdown failed: %s", str(e), exc_info=True)
+            logger.exception("Server shutdown failed: %s", str(e))
         except Exception as e:
-            logger.error("Unexpected error occurred: %s", str(e), exc_info=True)
+            logger.exception("Unexpected error occurred: %s", str(e))

nat/front_ends/fastapi/fastapi_front_end_plugin.py

@@ -113,4 +113,4 @@ class FastApiFrontEndPlugin(FrontEndBase[FastApiFrontEndConfig]):
             try:
                 os.remove(config_file_name)
             except OSError as e:
-                logger.error(f"Warning: Failed to delete temp file {config_file_name}: {e}")
+                logger.exception(f"Warning: Failed to delete temp file {config_file_name}: {e}")

nat/front_ends/fastapi/fastapi_front_end_plugin_worker.py

@@ -215,7 +215,7 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
                 job_store.cleanup_expired_jobs()
                 logger.debug("Expired %s jobs cleaned up", name)
             except Exception as e:
-                logger.error("Error during %s job cleanup: %s", name, e)
+                logger.exception("Error during %s job cleanup: %s", name, e)
             await asyncio.sleep(sleep_time_sec)
 
     async def create_cleanup_task(self, app: FastAPI, name: str, job_store: JobStore, sleep_time_sec: int = 300):
@@ -301,7 +301,7 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
 
                         job_store.update_status(job_id, "success", output_path=str(parent_dir))
                 except Exception as e:
-                    logger.error("Error in evaluation job %s: %s", job_id, str(e))
+                    logger.exception("Error in evaluation job %s: %s", job_id, str(e))
                     job_store.update_status(job_id, "failure", error=str(e))
 
         async def start_evaluation(request: EvaluateRequest, background_tasks: BackgroundTasks, http_request: Request):
@@ -735,7 +735,7 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
                                                             result_type=result_type)
                     job_store.update_status(job_id, "success", output=result)
                 except Exception as e:
-                    logger.error("Error in evaluation job %s: %s", job_id, e)
+                    logger.exception("Error in evaluation job %s: %s", job_id, e)
                     job_store.update_status(job_id, "failure", error=str(e))
 
         def _job_status_to_response(job: JobInfo) -> AsyncGenerationStatusResponse:

nat/front_ends/fastapi/message_handler.py

@@ -170,7 +170,7 @@ class WebSocketMessageHandler:
                                            self._workflow_schema_type])).add_done_callback(_done_callback)
 
         except ValueError as e:
-            logger.error("User message content not found: %s", str(e), exc_info=True)
+            logger.exception("User message content not found: %s", str(e))
             await self.create_websocket_message(data_model=Error(code=ErrorTypes.INVALID_USER_MESSAGE_CONTENT,
                                                                  message="User message content could not be found",
                                                                  details=str(e)),
@@ -238,7 +238,7 @@ class WebSocketMessageHandler:
                     f"Message type could not be resolved by input data model: {data_model.model_dump_json()}")
 
         except (ValidationError, TypeError, ValueError) as e:
-            logger.error("A data vaidation error ocurred creating websocket message: %s", str(e), exc_info=True)
+            logger.exception("A data vaidation error ocurred creating websocket message: %s", str(e))
             message = await self._message_validator.create_system_response_token_message(
                 message_type=WebSocketMessageType.ERROR_MESSAGE,
                 conversation_id=self._conversation_id,

nat/front_ends/fastapi/message_validator.py

@@ -97,7 +97,7 @@ class MessageValidator:
             return validated_message
 
         except (ValidationError, TypeError, ValueError) as e:
-            logger.error("A data validation error %s occurred for message: %s", str(e), str(message), exc_info=True)
+            logger.exception("A data validation error %s occurred for message: %s", str(e), str(message))
             return await self.create_system_response_token_message(message_type=WebSocketMessageType.ERROR_MESSAGE,
                                                                    content=Error(code=ErrorTypes.INVALID_MESSAGE,
                                                                                  message="Error validating message.",
@@ -119,7 +119,7 @@ class MessageValidator:
             return schema
 
         except (TypeError, ValueError) as e:
-            logger.error("Error retrieving schema for message type '%s': %s", message_type, str(e), exc_info=True)
+            logger.exception("Error retrieving schema for message type '%s': %s", message_type, str(e))
             return Error
 
     async def convert_data_to_message_content(self, data_model: BaseModel) -> BaseModel:
@@ -156,7 +156,7 @@ class MessageValidator:
             return validated_message_content
 
         except ValueError as e:
-            logger.error("Input data could not be converted to validated message content: %s", str(e), exc_info=True)
+            logger.exception("Input data could not be converted to validated message content: %s", str(e))
             return Error(code=ErrorTypes.INVALID_DATA_CONTENT, message="Input data not supported.", details=str(e))
 
     async def convert_text_content_to_human_response(self, text_content: TextContent,
@@ -191,7 +191,7 @@ class MessageValidator:
             return human_response
 
         except ValueError as e:
-            logger.error("Error human response content not found: %s", str(e), exc_info=True)
+            logger.exception("Error human response content not found: %s", str(e))
             return HumanResponseText(text=str(e))
 
     async def resolve_message_type_by_data(self, data_model: BaseModel) -> str:
@@ -218,9 +218,7 @@ class MessageValidator:
             return validated_message_type
 
         except ValueError as e:
-            logger.error("Error type not found converting data to validated websocket message content: %s",
-                         str(e),
-                         exc_info=True)
+            logger.exception("Error type not found converting data to validated websocket message content: %s", str(e))
             return WebSocketMessageType.ERROR_MESSAGE
 
     async def get_intermediate_step_parent_id(self, data_model: ResponseIntermediateStep) -> str:
@@ -269,7 +267,7 @@ class MessageValidator:
                                                        timestamp=timestamp)
 
         except Exception as e:
-            logger.error("Error creating system response token message: %s", str(e), exc_info=True)
+            logger.exception("Error creating system response token message: %s", str(e))
             return None
 
     async def create_system_intermediate_step_message(
@@ -308,7 +306,7 @@ class MessageValidator:
                                                           timestamp=timestamp)
 
         except Exception as e:
-            logger.error("Error creating system intermediate step message: %s", str(e), exc_info=True)
+            logger.exception("Error creating system intermediate step message: %s", str(e))
             return None
 
     async def create_system_interaction_message(
@@ -348,5 +346,5 @@ class MessageValidator:
                                                      timestamp=timestamp)
 
         except Exception as e:
-            logger.error("Error creating system interaction message: %s", str(e), exc_info=True)
+            logger.exception("Error creating system interaction message: %s", str(e))
             return None

nat/front_ends/fastapi/response_helpers.py

@@ -98,9 +98,9 @@ async def generate_streaming_response(payload: typing.Any,
                     yield item
                 else:
                     yield ResponsePayloadOutput(payload=item)
-        except Exception as e:
+        except Exception:
             # Handle exceptions here
-            raise e
+            raise
         finally:
             await q.close()
 
@@ -165,9 +165,9 @@ async def generate_streaming_response_full(payload: typing.Any,
                         yield item
                 else:
                     yield ResponsePayloadOutput(payload=item)
-        except Exception as e:
+        except Exception:
             # Handle exceptions here
-            raise e
+            raise
         finally:
             await q.close()
 

nat/front_ends/fastapi/step_adaptor.py

@@ -314,6 +314,6 @@ class StepAdaptor:
                 return self._handle_custom(payload, ancestry)
 
         except Exception as e:
-            logger.error("Error processing intermediate step: %s", e, exc_info=True)
+            logger.exception("Error processing intermediate step: %s", e)
 
         return None

nat/llm/aws_bedrock_llm.py

@@ -23,9 +23,11 @@ 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.thinking_mixin import ThinkingMixin
+from nat.data_models.top_p_mixin import TopPMixin
 
 
-class AWSBedrockModelConfig(LLMBaseConfig, RetryMixin, TemperatureMixin, name="aws_bedrock"):
+class AWSBedrockModelConfig(LLMBaseConfig, RetryMixin, TemperatureMixin, TopPMixin, ThinkingMixin, name="aws_bedrock"):
     """An AWS Bedrock llm provider to be used with an LLM client."""
 
     model_config = ConfigDict(protected_namespaces=())
@@ -34,14 +36,13 @@ class AWSBedrockModelConfig(LLMBaseConfig, RetryMixin, TemperatureMixin, name="a
     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."
-                                   "This field is ONLY required when using AWS Bedrock with Langchain.")
-    context_size: int | None = Field(default=1024,
-                                     gt=0,
-                                     description="Maximum number of tokens to generate."
-                                     "This field is ONLY required when using AWS Bedrock with LlamaIndex.")
+    max_tokens: int | None = Field(default=1024, gt=0, description="Maximum number of tokens to generate.")
+    context_size: int | None = Field(
+        default=1024,
+        gt=0,
+        description="The maximum number of tokens available for input. This is only required for LlamaIndex. "
+        "This field is ignored for Langchain.",
+    )
 
     # Client parameters
     region_name: str | None = Field(default="None", description="AWS region to use.")

nat/llm/azure_openai_llm.py

@@ -23,10 +23,18 @@ 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.thinking_mixin import ThinkingMixin
 from nat.data_models.top_p_mixin import TopPMixin
 
 
-class AzureOpenAIModelConfig(LLMBaseConfig, RetryMixin, TemperatureMixin, TopPMixin, name="azure_openai"):
+class AzureOpenAIModelConfig(
+        LLMBaseConfig,
+        RetryMixin,
+        TemperatureMixin,
+        TopPMixin,
+        ThinkingMixin,
+        name="azure_openai",
+):
     """An Azure OpenAI LLM provider to be used with an LLM client."""
 
     model_config = ConfigDict(protected_namespaces=(), extra="allow")

nat/llm/nim_llm.py

@@ -24,10 +24,11 @@ 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.thinking_mixin import ThinkingMixin
 from nat.data_models.top_p_mixin import TopPMixin
 
 
-class NIMModelConfig(LLMBaseConfig, RetryMixin, TemperatureMixin, TopPMixin, name="nim"):
+class NIMModelConfig(LLMBaseConfig, RetryMixin, TemperatureMixin, TopPMixin, ThinkingMixin, name="nim"):
     """An NVIDIA Inference Microservice (NIM) llm provider to be used with an LLM client."""
 
     model_config = ConfigDict(protected_namespaces=())

nat/llm/openai_llm.py

@@ -23,10 +23,11 @@ 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.thinking_mixin import ThinkingMixin
 from nat.data_models.top_p_mixin import TopPMixin
 
 
-class OpenAIModelConfig(LLMBaseConfig, RetryMixin, TemperatureMixin, TopPMixin, name="openai"):
+class OpenAIModelConfig(LLMBaseConfig, RetryMixin, TemperatureMixin, TopPMixin, ThinkingMixin, name="openai"):
     """An OpenAI LLM provider to be used with an LLM client."""
 
     model_config = ConfigDict(protected_namespaces=(), extra="allow")