nvidia-nat 1.3a20250819__py3-none-any.whl → 1.3.0a20250823__py3-none-any.whl

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

@@ -0,0 +1,49 @@
+# 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 AliasChoices
+from pydantic import ConfigDict
+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.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, TemperatureMixin, TopPMixin, name="azure_openai"):
+    """An Azure OpenAI LLM provider to be used with an LLM client."""
+
+    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_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",
+                                       default=None,
+                                       description="Base URL for the hosted model.")
+    azure_deployment: str = Field(validation_alias=AliasChoices("azure_deployment", "model_name", "model"),
+                                  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.")
+
+
+@register_llm_provider(config_type=AzureOpenAIModelConfig)
+async def azure_openai_llm(config: AzureOpenAIModelConfig, _builder: Builder):
+
+    yield LLMProviderInfo(config=config, description="An Azure OpenAI model for use with an LLM client.")

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.")

nat/llm/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 OpenAIModelConfig(LLMBaseConfig, RetryMixin, name="openai"):
+class OpenAIModelConfig(LLMBaseConfig, RetryMixin, TemperatureMixin, TopPMixin, name="openai"):
     """An OpenAI LLM provider to be used with an LLM client."""
 
     model_config = ConfigDict(protected_namespaces=(), extra="allow")
@@ -34,13 +36,11 @@ class OpenAIModelConfig(LLMBaseConfig, RetryMixin, name="openai"):
     model_name: str = Field(validation_alias=AliasChoices("model_name", "model"),
                             serialization_alias="model",
                             description="The OpenAI hosted model 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=OpenAIModelConfig)
-async def openai_llm(config: OpenAIModelConfig, builder: Builder):
+async def openai_llm(config: OpenAIModelConfig, _builder: Builder):
 
     yield LLMProviderInfo(config=config, description="An OpenAI model for use with an LLM client.")

nat/llm/register.py

@@ -13,11 +13,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# pylint: disable=unused-import
 # flake8: noqa
 # isort:skip_file
 
 # Import any providers which need to be automatically registered here
 from . import aws_bedrock_llm
+from . import azure_openai_llm
 from . import nim_llm
 from . import openai_llm

nat/llm/utils/env_config_value.py

@@ -72,9 +72,8 @@ class EnvConfigValue(ABC):
                     f"{message} Try passing a value to the constructor, or setting the `{self.__class__._ENV_KEY}` "
                     "environment variable.")
 
-        else:
-            if not self.__class__._ALLOW_NONE and value is None:
-                raise ValueError("value must not be none")
+        elif not self.__class__._ALLOW_NONE and value is None:
+            raise ValueError("value must not be none")
 
         assert isinstance(value, str) or value is None
 

nat/meta/pypi.md

@@ -23,19 +23,19 @@ NeMo Agent toolkit is a flexible library designed to seamlessly integrate your e
 
 ## Key Features
 
-- [**Framework Agnostic:**](https://docs.nvidia.com/nemo/agent-toolkit/1.2.0/extend/plugins.html) Works with any agentic framework, so you can use your current technology stack without replatforming.
-- [**Reusability:**](https://docs.nvidia.com/nemo/agent-toolkit/1.2.0/extend/sharing-components.html) Every agent, tool, or workflow can be combined and repurposed, allowing developers to leverage existing work in new scenarios.
-- [**Rapid Development:**](https://docs.nvidia.com/nemo/agent-toolkit/1.2.0/tutorials/index.html) Start with a pre-built agent, tool, or workflow, and customize it to your needs.
-- [**Profiling:**](https://docs.nvidia.com/nemo/agent-toolkit/1.2.0/workflows/profiler.html) Profile entire workflows down to the tool and agent level, track input/output tokens and timings, and identify bottlenecks.
-- [**Observability:**](https://docs.nvidia.com/nemo/agent-toolkit/1.2.0/workflows/observe/observe-workflow-with-phoenix.html) Monitor and debug your workflows with any OpenTelemetry-compatible observability tool, with examples using [Phoenix](https://docs.nvidia.com/nemo/agent-toolkit/1.2.0/workflows/observe/observe-workflow-with-phoenix.html) and [W&B Weave](https://docs.nvidia.com/nemo/agent-toolkit/1.2.0/workflows/observe/observe-workflow-with-weave.html).
-- [**Evaluation System:**](https://docs.nvidia.com/nemo/agent-toolkit/1.2.0/workflows/evaluate.html) Validate and maintain accuracy of agentic workflows with built-in evaluation tools.
-- [**User Interface:**](https://docs.nvidia.com/nemo/agent-toolkit/1.2.0/quick-start/launching-ui.html) Use the NeMo Agent toolkit UI chat interface to interact with your agents, visualize output, and debug workflows.
-- [**MCP Compatibility**](https://docs.nvidia.com/nemo/agent-toolkit/1.2.0/workflows/mcp/mcp-client.html) Compatible with Model Context Protocol (MCP), allowing tools served by MCP Servers to be used as NeMo Agent toolkit functions.
+- [**Framework Agnostic:**](https://docs.nvidia.com/nemo/agent-toolkit/1.3/extend/plugins.html) Works with any agentic framework, so you can use your current technology stack without replatforming.
+- [**Reusability:**](https://docs.nvidia.com/nemo/agent-toolkit/1.3/extend/sharing-components.html) Every agent, tool, or workflow can be combined and repurposed, allowing developers to leverage existing work in new scenarios.
+- [**Rapid Development:**](https://docs.nvidia.com/nemo/agent-toolkit/1.3/tutorials/index.html) Start with a pre-built agent, tool, or workflow, and customize it to your needs.
+- [**Profiling:**](https://docs.nvidia.com/nemo/agent-toolkit/1.3/workflows/profiler.html) Profile entire workflows down to the tool and agent level, track input/output tokens and timings, and identify bottlenecks.
+- [**Observability:**](https://docs.nvidia.com/nemo/agent-toolkit/1.3/workflows/observe/observe-workflow-with-phoenix.html) Monitor and debug your workflows with any OpenTelemetry-compatible observability tool, with examples using [Phoenix](https://docs.nvidia.com/nemo/agent-toolkit/1.3/workflows/observe/observe-workflow-with-phoenix.html) and [W&B Weave](https://docs.nvidia.com/nemo/agent-toolkit/1.3/workflows/observe/observe-workflow-with-weave.html).
+- [**Evaluation System:**](https://docs.nvidia.com/nemo/agent-toolkit/1.3/workflows/evaluate.html) Validate and maintain accuracy of agentic workflows with built-in evaluation tools.
+- [**User Interface:**](https://docs.nvidia.com/nemo/agent-toolkit/1.3/quick-start/launching-ui.html) Use the NeMo Agent toolkit UI chat interface to interact with your agents, visualize output, and debug workflows.
+- [**MCP Compatibility**](https://docs.nvidia.com/nemo/agent-toolkit/1.3/workflows/mcp/mcp-client.html) Compatible with Model Context Protocol (MCP), allowing tools served by MCP Servers to be used as NeMo Agent toolkit functions.
 
 With NeMo Agent toolkit, you can move quickly, experiment freely, and ensure reliability across all your agent-driven projects.
 
 ## Links
- * [Documentation](https://docs.nvidia.com/nemo/agent-toolkit/1.2.0/index.html): Explore the full documentation for NeMo Agent toolkit.
+ * [Documentation](https://docs.nvidia.com/nemo/agent-toolkit/1.3/index.html): Explore the full documentation for NeMo Agent toolkit.
 
 ## First time user?
  If this is your first time using NeMo Agent toolkit, it is recommended to install the latest version from the [source repository](https://github.com/NVIDIA/NeMo-Agent-Toolkit?tab=readme-ov-file#quick-start) on GitHub. This package is intended for users who are familiar with NeMo Agent toolkit applications and need to add NeMo Agent toolkit as a dependency to their project.

nat/object_store/models.py

@@ -14,6 +14,7 @@
 # limitations under the License.
 
 from pydantic import BaseModel
+from pydantic import ConfigDict
 from pydantic import Field
 
 
@@ -30,6 +31,7 @@ class ObjectStoreItem(BaseModel):
     metadata : dict[str, str] | None
         Metadata providing context and utility for management operations.
     """
+    model_config = ConfigDict(ser_json_bytes="base64", val_json_bytes="base64")
 
     data: bytes = Field(description="The data to store in the object store.")
     content_type: str | None = Field(description="The content type of the data.", default=None)

nat/object_store/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/observability/exporter/base_exporter.py

@@ -357,7 +357,7 @@ class BaseExporter(Exporter):
                 except Exception as e:
                     logger.warning("Error while canceling task %s: %s", task.get_name(), e)
 
-    async def _wait_for_tasks(self, timeout: float = 5.0):
+    async def wait_for_tasks(self, timeout: float = 5.0):
         """Wait for all tracked tasks to complete with a timeout.
 
         Note: This method is NOT called during normal stop() operation for performance.

nat/observability/exporter/file_exporter.py

@@ -24,7 +24,7 @@ from nat.observability.processor.intermediate_step_serializer import Intermediat
 logger = logging.getLogger(__name__)
 
 
-class FileExporter(FileExportMixin, RawExporter[IntermediateStep, str]):  # pylint: disable=R0901
+class FileExporter(FileExportMixin, RawExporter[IntermediateStep, str]):
     """A File exporter that exports telemetry traces to a local file."""
 
     def __init__(self, context_state: ContextState | None = None, **file_kwargs):

nat/observability/register.py

@@ -45,7 +45,7 @@ class FileTelemetryExporterConfig(TelemetryExporterBaseConfig, name="file"):
 
 
 @register_telemetry_exporter(config_type=FileTelemetryExporterConfig)
-async def file_telemetry_exporter(config: FileTelemetryExporterConfig, builder: Builder):  # pylint: disable=W0613
+async def file_telemetry_exporter(config: FileTelemetryExporterConfig, builder: Builder):
     """
     Build and return a FileExporter for file-based telemetry export with optional rolling.
     """
@@ -68,7 +68,7 @@ class ConsoleLoggingMethodConfig(LoggingBaseConfig, name="console"):
 
 
 @register_logging_method(config_type=ConsoleLoggingMethodConfig)
-async def console_logging_method(config: ConsoleLoggingMethodConfig, builder: Builder):  # pylint: disable=W0613
+async def console_logging_method(config: ConsoleLoggingMethodConfig, builder: Builder):
     """
     Build and return a StreamHandler for console-based logging.
     """
@@ -86,7 +86,7 @@ class FileLoggingMethod(LoggingBaseConfig, name="file"):
 
 
 @register_logging_method(config_type=FileLoggingMethod)
-async def file_logging_method(config: FileLoggingMethod, builder: Builder):  # pylint: disable=W0613
+async def file_logging_method(config: FileLoggingMethod, builder: Builder):
     """
     Build and return a FileHandler for file-based logging.
     """

nat/profiler/callbacks/langchain_callback_handler.py

@@ -53,7 +53,7 @@ def _extract_tools_schema(invocation_params: dict) -> list:
     return tools_schema
 
 
-class LangchainProfilerHandler(AsyncCallbackHandler, BaseProfilerCallback):  # pylint: disable=R0901
+class LangchainProfilerHandler(AsyncCallbackHandler, BaseProfilerCallback):
     """Callback Handler that tracks NIM info."""
 
     total_tokens: int = 0
@@ -223,7 +223,14 @@ class LangchainProfilerHandler(AsyncCallbackHandler, BaseProfilerCallback):  # p
             except AttributeError:
                 usage_metadata = {}
 
-        llm_text_output = generation.message.content if generation else ""
+        if generation:
+            llm_text_output = generation.message.content
+            if "tool_calls" in generation.message.additional_kwargs:
+                # add tool calls if included in the output
+                tool_calls = generation.message.additional_kwargs['tool_calls']
+                llm_text_output = f"{llm_text_output}\n\nTool calls: {tool_calls}"
+        else:
+            llm_text_output = ""
 
         # update shared state behind lock
         with self._lock:

nat/profiler/callbacks/semantic_kernel_callback_handler.py

@@ -86,7 +86,7 @@ class SemanticKernelProfilerHandler(BaseProfilerCallback):
 
         # Gather the appropriate modules/functions based on your builder config
         for llm in self._builder_llms:
-            if self._builder_llms[llm].provider_type == 'openai':  # pylint: disable=consider-using-in
+            if self._builder_llms[llm].provider_type == 'openai':
                 functions_to_patch.extend(["openai_non_streaming", "openai_streaming"])
 
         # Grab original reference for the tool call

nat/profiler/data_frame_row.py

@@ -42,7 +42,7 @@ class DataFrameRow(BaseModel):
     framework: str | None
 
     @field_validator('llm_text_input', 'llm_text_output', 'llm_new_token', mode='before')
-    def cast_to_str(cls, v):  # pylint: disable=no-self-argument
+    def cast_to_str(cls, v):
         if v is None:
             return v
         try:

nat/profiler/decorators/framework_wrapper.py

@@ -13,8 +13,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# pylint disable=ungrouped-imports
-
 from __future__ import annotations
 
 import functools
@@ -75,8 +73,7 @@ def set_framework_profiler_handler(
                 logger.debug("LlamaIndex callback handler registered")
 
             if LLMFrameworkEnum.CREWAI in frameworks and not _library_instrumented["crewai"]:
-                from nat.plugins.crewai.crewai_callback_handler import \
-                    CrewAIProfilerHandler  # pylint: disable=ungrouped-imports,line-too-long  # noqa E501
+                from nat.plugins.crewai.crewai_callback_handler import CrewAIProfilerHandler
 
                 handler = CrewAIProfilerHandler()
                 handler.instrument()

nat/profiler/forecasting/models/forecasting_base_model.py

@@ -15,7 +15,9 @@
 
 # forecasting/models/base_model.py
 
-from abc import ABC, abstractmethod
+from abc import ABC
+from abc import abstractmethod
+
 import numpy as np
 
 

nat/profiler/inference_optimization/bottleneck_analysis/simple_stack_analysis.py

@@ -195,7 +195,7 @@ def profile_workflow_bottlenecks(all_steps: list[list[IntermediateStep]]) -> Sim
         c_max = 0
         for ts, delta in events_sub:
             c_curr += delta
-            if c_curr > c_max:  # pylint: disable=consider-using-max-builtin
+            if c_curr > c_max:  # noqa: PLR1730 - don't use max built-in
                 c_max = c_curr
         max_concurrency_by_name[op_name] = c_max
 

nat/profiler/inference_optimization/data_models.py

@@ -172,7 +172,7 @@ class CallNode(BaseModel):
         if not self.children:
             return self.duration
 
-        intervals = [(c.start_time, c.end_time) for c in self.children]  # pylint: disable=not-an-iterable
+        intervals = [(c.start_time, c.end_time) for c in self.children]
         # Sort by start time
         intervals.sort(key=lambda x: x[0])
 
@@ -204,7 +204,7 @@ class CallNode(BaseModel):
         This ensures no overlap double-counting among children.
         """
         total = self.compute_self_time()
-        for c in self.children:  # pylint: disable=not-an-iterable
+        for c in self.children:
             total += c.compute_subtree_time()
         return total
 
@@ -216,7 +216,7 @@ class CallNode(BaseModel):
         info = (f"{indent}- {self.operation_type} '{self.operation_name}' "
                 f"(uuid={self.uuid}, start={self.start_time:.2f}, "
                 f"end={self.end_time:.2f}, dur={self.duration:.2f})")
-        child_strs = [child._repr(level + 1) for child in self.children]  # pylint: disable=not-an-iterable
+        child_strs = [child._repr(level + 1) for child in self.children]
         return "\n".join([info] + child_strs)
 
 

nat/profiler/inference_optimization/experimental/prefix_span_analysis.py

@@ -228,7 +228,7 @@ def run_prefixspan(sequences_map: dict[int, list[PrefixCallNode]],
     else:
         abs_min_support = min_support
 
-    freq_patterns = ps.frequent(abs_min_support)  # pylint: disable=not-callable
+    freq_patterns = ps.frequent(abs_min_support)
     # freq_patterns => [(count, [item1, item2, ...])]
 
     results = []
@@ -321,13 +321,12 @@ def compute_coverage_and_duration(sequences_map: dict[int, list[PrefixCallNode]]
 # --------------------------------------------------------------------------------
 
 
-def prefixspan_subworkflow_with_text(  # pylint: disable=too-many-positional-arguments
-        all_steps: list[list[IntermediateStep]],
-        min_support: int | float = 2,
-        top_k: int = 10,
-        min_coverage: float = 0.0,
-        max_text_len: int = 700,
-        prefix_list: list[str] = None) -> PrefixSpanSubworkflowResult:
+def prefixspan_subworkflow_with_text(all_steps: list[list[IntermediateStep]],
+                                     min_support: int | float = 2,
+                                     top_k: int = 10,
+                                     min_coverage: float = 0.0,
+                                     max_text_len: int = 700,
+                                     prefix_list: list[str] = None) -> PrefixSpanSubworkflowResult:
     """
     1) Build sequences of calls for each example (with llm_text_input).
     2) Convert to token lists, run PrefixSpan with min_support.

nat/profiler/inference_optimization/token_uniqueness.py

@@ -66,7 +66,7 @@ def compute_inter_query_token_uniqueness_by_llm(all_steps: list[list[Intermediat
     # 2) Group by (llm_name, example_number), then sort each group
     grouped = cdf.groupby(['llm_name', 'example_number'], as_index=False, group_keys=True)
 
-    for (llm, ex_num), group_df in grouped:  # pylint: disable=unused-variable
+    for (llm, ex_num), group_df in grouped:
         # Sort by event_timestamp
         group_df = group_df.sort_values('event_timestamp', ascending=True)
 

nat/profiler/profile_runner.py

@@ -88,14 +88,19 @@ class ProfilerRunner:
         writes out combined requests JSON, then computes and saves additional metrics,
         and optionally fits a forecasting model.
         """
-        from nat.profiler.inference_optimization.bottleneck_analysis.nested_stack_analysis import \
-            multi_example_call_profiling
-        from nat.profiler.inference_optimization.bottleneck_analysis.simple_stack_analysis import \
-            profile_workflow_bottlenecks
-        from nat.profiler.inference_optimization.experimental.concurrency_spike_analysis import \
-            concurrency_spike_analysis
-        from nat.profiler.inference_optimization.experimental.prefix_span_analysis import \
-            prefixspan_subworkflow_with_text
+        # Yapf and ruff disagree on how to format long imports, disable yapf go with ruff
+        from nat.profiler.inference_optimization.bottleneck_analysis.nested_stack_analysis import (
+            multi_example_call_profiling,
+        )  # yapf: disable
+        from nat.profiler.inference_optimization.bottleneck_analysis.simple_stack_analysis import (
+            profile_workflow_bottlenecks,
+        )  # yapf: disable
+        from nat.profiler.inference_optimization.experimental.concurrency_spike_analysis import (
+            concurrency_spike_analysis,
+        )  # yapf: disable
+        from nat.profiler.inference_optimization.experimental.prefix_span_analysis import (
+            prefixspan_subworkflow_with_text,
+        )  # yapf: disable
         from nat.profiler.inference_optimization.llm_metrics import LLMMetrics
         from nat.profiler.inference_optimization.prompt_caching import get_common_prefixes
         from nat.profiler.inference_optimization.token_uniqueness import compute_inter_query_token_uniqueness_by_llm

nat/registry_handlers/package_utils.py

@@ -29,7 +29,6 @@ from nat.registry_handlers.schemas.publish import Artifact
 from nat.runtime.loader import PluginTypes
 from nat.runtime.loader import discover_entrypoints
 
-# pylint: disable=redefined-outer-name
 logger = logging.getLogger(__name__)
 
 

nat/registry_handlers/pypi/pypi_handler.py

@@ -44,13 +44,12 @@ class PypiRegistryHandler(AbstractRegistryHandler):
     https://github.com/pypiserver/pypiserver
     """
 
-    def __init__(  # pylint: disable=R0917
-            self,
-            endpoint: str,
-            token: str | None = None,
-            publish_route: str = "",
-            pull_route: str = "",
-            search_route: str = ""):
+    def __init__(self,
+                 endpoint: str,
+                 token: str | None = None,
+                 publish_route: str = "",
+                 pull_route: str = "",
+                 search_route: str = ""):
         super().__init__()
         self._endpoint = endpoint.rstrip("/")
         self._token = token
@@ -126,17 +125,16 @@ class PypiRegistryHandler(AbstractRegistryHandler):
 
             versioned_packages_str = " ".join(versioned_packages)
 
-            result = subprocess.run(
-                [
-                    "uv",
-                    "pip",
-                    "install",
-                    "--prerelease=allow",
-                    "--index-url",
-                    f"{self._endpoint}/{self._pull_route}/",
-                    versioned_packages_str
-                ],  # pylint: disable=W0631
-                check=True)
+            result = subprocess.run([
+                "uv",
+                "pip",
+                "install",
+                "--prerelease=allow",
+                "--index-url",
+                f"{self._endpoint}/{self._pull_route}/",
+                versioned_packages_str
+            ],
+                                    check=True)
 
             result.check_returncode()
 
@@ -171,11 +169,10 @@ class PypiRegistryHandler(AbstractRegistryHandler):
         """
 
         try:
-            completed_process = subprocess.run(
-                ["pip", "search", "--index", f"{self._endpoint}", query.query],  # pylint: disable=W0631
-                text=True,
-                capture_output=True,
-                check=True)
+            completed_process = subprocess.run(["pip", "search", "--index", f"{self._endpoint}", query.query],
+                                               text=True,
+                                               capture_output=True,
+                                               check=True)
             search_response_list = []
             search_results = completed_process.stdout
             package_results = search_results.split("\n")

nat/registry_handlers/register.py

@@ -13,9 +13,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# pylint: disable=unused-import
 # flake8: noqa
 
-from .local import register_local  # pylint: disable=E0611
-from .pypi import register_pypi  # pylint: disable=E0611
-from .rest import register_rest  # pylint: disable=E0611
+from .local import register_local
+from .pypi import register_pypi
+from .rest import register_rest

nat/registry_handlers/rest/rest_handler.py

@@ -42,15 +42,14 @@ logger = logging.getLogger(__name__)
 class RestRegistryHandler(AbstractRegistryHandler):
     """A registry handler for interactions with a remote REST registry."""
 
-    def __init__(  # pylint: disable=R0917
-            self,
-            endpoint: str,
-            token: str,
-            timeout: int = 30,
-            publish_route: str = "",
-            pull_route: str = "",
-            search_route: str = "",
-            remove_route: str = ""):
+    def __init__(self,
+                 endpoint: str,
+                 token: str,
+                 timeout: int = 30,
+                 publish_route: str = "",
+                 pull_route: str = "",
+                 search_route: str = "",
+                 remove_route: str = ""):
         super().__init__()
         self._endpoint = endpoint.rstrip("/")
         self._timeout = timeout

nat/retriever/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