nvidia-nat 1.2.1rc1__py3-none-any.whl → 1.3.0__py3-none-any.whl

nat/data_models/config.py

@@ -20,6 +20,7 @@ import typing
 from pydantic import BaseModel
 from pydantic import ConfigDict
 from pydantic import Discriminator
+from pydantic import Field
 from pydantic import ValidationError
 from pydantic import ValidationInfo
 from pydantic import ValidatorFunctionWrapHandler
@@ -29,7 +30,9 @@ from nat.data_models.evaluate import EvalConfig
 from nat.data_models.front_end import FrontEndBaseConfig
 from nat.data_models.function import EmptyFunctionConfig
 from nat.data_models.function import FunctionBaseConfig
+from nat.data_models.function import FunctionGroupBaseConfig
 from nat.data_models.logging import LoggingBaseConfig
+from nat.data_models.optimizer import OptimizerConfig
 from nat.data_models.telemetry_exporter import TelemetryExporterBaseConfig
 from nat.data_models.ttc_strategy import TTCStrategyBaseConfig
 from nat.front_ends.fastapi.fastapi_front_end_config import FastApiFrontEndConfig
@@ -47,7 +50,7 @@ logger = logging.getLogger(__name__)
 
 
 def _process_validation_error(err: ValidationError, handler: ValidatorFunctionWrapHandler, info: ValidationInfo):
-    from nat.cli.type_registry import GlobalTypeRegistry  # pylint: disable=cyclic-import
+    from nat.cli.type_registry import GlobalTypeRegistry
 
     new_errors = []
     logged_once = False
@@ -57,9 +60,10 @@ def _process_validation_error(err: ValidationError, handler: ValidatorFunctionWr
         error_type = e['type']
         if error_type == 'union_tag_invalid' and "ctx" in e and not logged_once:
             requested_type = e["ctx"]["tag"]
-
             if (info.field_name in ('workflow', 'functions')):
                 registered_keys = GlobalTypeRegistry.get().get_registered_functions()
+            elif (info.field_name == "function_groups"):
+                registered_keys = GlobalTypeRegistry.get().get_registered_function_groups()
             elif (info.field_name == "authentication"):
                 registered_keys = GlobalTypeRegistry.get().get_registered_auth_providers()
             elif (info.field_name == "llms"):
@@ -135,8 +139,8 @@ def _process_validation_error(err: ValidationError, handler: ValidatorFunctionWr
 
 class TelemetryConfig(BaseModel):
 
-    logging: dict[str, LoggingBaseConfig] = {}
-    tracing: dict[str, TelemetryExporterBaseConfig] = {}
+    logging: dict[str, LoggingBaseConfig] = Field(default_factory=dict)
+    tracing: dict[str, TelemetryExporterBaseConfig] = Field(default_factory=dict)
 
     @field_validator("logging", "tracing", mode="wrap")
     @classmethod
@@ -183,12 +187,16 @@ class TelemetryConfig(BaseModel):
 
 class GeneralConfig(BaseModel):
 
-    model_config = ConfigDict(protected_namespaces=())
+    model_config = ConfigDict(protected_namespaces=(), extra="forbid")
 
-    use_uvloop: bool = True
+    use_uvloop: bool | None = Field(
+        default=None,
+        deprecated=
+        "`use_uvloop` field is deprecated and will be removed in a future release. The use of `uv_loop` is now" +
+        "automatically determined based on platform")
     """
-    Whether to use uvloop for the event loop. This can provide a significant speedup in some cases. Disable to provide
-    better error messages when debugging.
+    This field is deprecated and ignored. It previously controlled whether to use uvloop as the event loop. uvloop
+    usage is now determined automatically based on the platform.
     """
 
     telemetry: TelemetryConfig = TelemetryConfig()
@@ -240,31 +248,37 @@ class Config(HashableBaseModel):
     general: GeneralConfig = GeneralConfig()
 
     # Functions Configuration
-    functions: dict[str, FunctionBaseConfig] = {}
+    functions: dict[str, FunctionBaseConfig] = Field(default_factory=dict)
+
+    # Function Groups Configuration
+    function_groups: dict[str, FunctionGroupBaseConfig] = Field(default_factory=dict)
 
     # LLMs Configuration
-    llms: dict[str, LLMBaseConfig] = {}
+    llms: dict[str, LLMBaseConfig] = Field(default_factory=dict)
 
     # Embedders Configuration
-    embedders: dict[str, EmbedderBaseConfig] = {}
+    embedders: dict[str, EmbedderBaseConfig] = Field(default_factory=dict)
 
     # Memory Configuration
-    memory: dict[str, MemoryBaseConfig] = {}
+    memory: dict[str, MemoryBaseConfig] = Field(default_factory=dict)
 
     # Object Stores Configuration
-    object_stores: dict[str, ObjectStoreBaseConfig] = {}
+    object_stores: dict[str, ObjectStoreBaseConfig] = Field(default_factory=dict)
+
+    # Optimizer Configuration
+    optimizer: OptimizerConfig = OptimizerConfig()
 
     # Retriever Configuration
-    retrievers: dict[str, RetrieverBaseConfig] = {}
+    retrievers: dict[str, RetrieverBaseConfig] = Field(default_factory=dict)
 
     # TTC Strategies
-    ttc_strategies: dict[str, TTCStrategyBaseConfig] = {}
+    ttc_strategies: dict[str, TTCStrategyBaseConfig] = Field(default_factory=dict)
 
     # Workflow Configuration
     workflow: FunctionBaseConfig = EmptyFunctionConfig()
 
     # Authentication Configuration
-    authentication: dict[str, AuthProviderBaseConfig] = {}
+    authentication: dict[str, AuthProviderBaseConfig] = Field(default_factory=dict)
 
     # Evaluation Options
     eval: EvalConfig = EvalConfig()
@@ -278,6 +292,7 @@ class Config(HashableBaseModel):
             stream.write(f"Workflow Type: {self.workflow.type}\n")
 
         stream.write(f"Number of Functions: {len(self.functions)}\n")
+        stream.write(f"Number of Function Groups: {len(self.function_groups)}\n")
         stream.write(f"Number of LLMs: {len(self.llms)}\n")
         stream.write(f"Number of Embedders: {len(self.embedders)}\n")
         stream.write(f"Number of Memory: {len(self.memory)}\n")
@@ -287,6 +302,7 @@ class Config(HashableBaseModel):
         stream.write(f"Number of Authentication Providers: {len(self.authentication)}\n")
 
     @field_validator("functions",
+                     "function_groups",
                      "llms",
                      "embedders",
                      "memory",
@@ -328,6 +344,10 @@ class Config(HashableBaseModel):
                                    typing.Annotated[type_registry.compute_annotation(FunctionBaseConfig),
                                                     Discriminator(TypedBaseModel.discriminator)]]
 
+        FunctionGroupsAnnotation = dict[str,
+                                        typing.Annotated[type_registry.compute_annotation(FunctionGroupBaseConfig),
+                                                         Discriminator(TypedBaseModel.discriminator)]]
+
         MemoryAnnotation = dict[str,
                                 typing.Annotated[type_registry.compute_annotation(MemoryBaseConfig),
                                                  Discriminator(TypedBaseModel.discriminator)]]
@@ -335,7 +355,6 @@ class Config(HashableBaseModel):
         ObjectStoreAnnotation = dict[str,
                                      typing.Annotated[type_registry.compute_annotation(ObjectStoreBaseConfig),
                                                       Discriminator(TypedBaseModel.discriminator)]]
-
         RetrieverAnnotation = dict[str,
                                    typing.Annotated[type_registry.compute_annotation(RetrieverBaseConfig),
                                                     Discriminator(TypedBaseModel.discriminator)]]
@@ -344,7 +363,7 @@ class Config(HashableBaseModel):
                                      typing.Annotated[type_registry.compute_annotation(TTCStrategyBaseConfig),
                                                       Discriminator(TypedBaseModel.discriminator)]]
 
-        WorkflowAnnotation = typing.Annotated[type_registry.compute_annotation(FunctionBaseConfig),
+        WorkflowAnnotation = typing.Annotated[(type_registry.compute_annotation(FunctionBaseConfig)),
                                               Discriminator(TypedBaseModel.discriminator)]
 
         should_rebuild = False
@@ -369,6 +388,11 @@ class Config(HashableBaseModel):
             functions_field.annotation = FunctionsAnnotation
             should_rebuild = True
 
+        function_groups_field = cls.model_fields.get("function_groups")
+        if function_groups_field is not None and function_groups_field.annotation != FunctionGroupsAnnotation:
+            function_groups_field.annotation = FunctionGroupsAnnotation
+            should_rebuild = True
+
         memory_field = cls.model_fields.get("memory")
         if memory_field is not None and memory_field.annotation != MemoryAnnotation:
             memory_field.annotation = MemoryAnnotation

nat/data_models/dataset_handler.py

@@ -80,7 +80,7 @@ class EvalDatasetJsonConfig(EvalDatasetBaseConfig, name="json"):
 
 
 def read_jsonl(file_path: FilePath):
-    with open(file_path, 'r', encoding='utf-8') as f:
+    with open(file_path, encoding='utf-8') as f:
         data = [json.loads(line) for line in f]
     return pd.DataFrame(data)
 

nat/data_models/discovery_metadata.py

@@ -177,7 +177,7 @@ class DiscoveryMetadata(BaseModel):
                 logger.warning("Package metadata not found for %s", distro_name)
                 version = ""
         except Exception as e:
-            logger.exception("Encountered issue extracting module metadata for %s: %s", config_type, e, exc_info=True)
+            logger.exception("Encountered issue extracting module metadata for %s: %s", config_type, e)
             return DiscoveryMetadata(status=DiscoveryStatusEnum.FAILURE)
 
         description = generate_config_type_docs(config_type=config_type)
@@ -217,7 +217,7 @@ class DiscoveryMetadata(BaseModel):
                 logger.warning("Package metadata not found for %s", distro_name)
                 version = ""
         except Exception as e:
-            logger.exception("Encountered issue extracting module metadata for %s: %s", fn, e, exc_info=True)
+            logger.exception("Encountered issue extracting module metadata for %s: %s", fn, e)
             return DiscoveryMetadata(status=DiscoveryStatusEnum.FAILURE)
 
         if isinstance(wrapper_type, LLMFrameworkEnum):
@@ -252,7 +252,7 @@ class DiscoveryMetadata(BaseModel):
                 description = ""
                 package_version = package_version or ""
         except Exception as e:
-            logger.exception("Encountered issue extracting module metadata for %s: %s", package_name, e, exc_info=True)
+            logger.exception("Encountered issue extracting module metadata for %s: %s", package_name, e)
             return DiscoveryMetadata(status=DiscoveryStatusEnum.FAILURE)
 
         return DiscoveryMetadata(package=package_name,
@@ -290,7 +290,7 @@ class DiscoveryMetadata(BaseModel):
                 logger.warning("Package metadata not found for %s", distro_name)
                 version = ""
         except Exception as e:
-            logger.exception("Encountered issue extracting module metadata for %s: %s", config_type, e, exc_info=True)
+            logger.exception("Encountered issue extracting module metadata for %s: %s", config_type, e)
             return DiscoveryMetadata(status=DiscoveryStatusEnum.FAILURE)
 
         wrapper_type = wrapper_type.value if isinstance(wrapper_type, LLMFrameworkEnum) else wrapper_type

nat/data_models/evaluate.py

@@ -57,6 +57,9 @@ class EvalOutputConfig(BaseModel):
     dir: Path = Path("./.tmp/nat/examples/default/")
     # S3 prefix for the workflow and evaluation results
     remote_dir: str | None = None
+    # Custom function to pre-evaluation process the eval input
+    # Format: "module.path.function_name"
+    custom_pre_eval_process_function: str | None = None
     # Custom scripts to run after the workflow and evaluation results are saved
     custom_scripts: dict[str, EvalCustomScriptConfig] = {}
     # S3 config for uploading the contents of the output directory
@@ -108,7 +111,7 @@ class EvalConfig(BaseModel):
     @classmethod
     def rebuild_annotations(cls):
 
-        from nat.cli.type_registry import GlobalTypeRegistry  # pylint: disable=cyclic-import
+        from nat.cli.type_registry import GlobalTypeRegistry
 
         type_registry = GlobalTypeRegistry.get()
 

nat/data_models/function.py

@@ -15,6 +15,10 @@
 
 import typing
 
+from pydantic import Field
+from pydantic import field_validator
+from pydantic import model_validator
+
 from .common import BaseModelRegistryTag
 from .common import TypedBaseModel
 
@@ -23,8 +27,38 @@ class FunctionBaseConfig(TypedBaseModel, BaseModelRegistryTag):
     pass
 
 
+class FunctionGroupBaseConfig(TypedBaseModel, BaseModelRegistryTag):
+    """Base configuration for function groups.
+
+    Function groups enable sharing of configurations and resources across multiple functions.
+    """
+    include: list[str] = Field(
+        default_factory=list,
+        description="The list of function names which should be added to the global Function registry",
+    )
+    exclude: list[str] = Field(
+        default_factory=list,
+        description="The list of function names which should be excluded from default access to the group",
+    )
+
+    @field_validator("include", "exclude")
+    @classmethod
+    def _validate_fields_include_exclude(cls, value: list[str]) -> list[str]:
+        if len(set(value)) != len(value):
+            raise ValueError("Function names must be unique")
+        return sorted(value)
+
+    @model_validator(mode="after")
+    def _validate_include_exclude(self):
+        if self.include and self.exclude:
+            raise ValueError("include and exclude cannot be used together")
+        return self
+
+
 class EmptyFunctionConfig(FunctionBaseConfig, name="EmptyFunctionConfig"):
     pass
 
 
 FunctionConfigT = typing.TypeVar("FunctionConfigT", bound=FunctionBaseConfig)
+
+FunctionGroupConfigT = typing.TypeVar("FunctionGroupConfigT", bound=FunctionGroupBaseConfig)

nat/data_models/function_dependencies.py

@@ -23,6 +23,7 @@ class FunctionDependencies(BaseModel):
     A class to represent the dependencies of a function.
     """
     functions: set[str] = Field(default_factory=set)
+    function_groups: set[str] = Field(default_factory=set)
     llms: set[str] = Field(default_factory=set)
     embedders: set[str] = Field(default_factory=set)
     memory_clients: set[str] = Field(default_factory=set)
@@ -33,6 +34,10 @@ class FunctionDependencies(BaseModel):
     def serialize_functions(self, v: set[str]) -> list[str]:
         return list(v)
 
+    @field_serializer("function_groups", when_used="json")
+    def serialize_function_groups(self, v: set[str]) -> list[str]:
+        return list(v)
+
     @field_serializer("llms", when_used="json")
     def serialize_llms(self, v: set[str]) -> list[str]:
         return list(v)
@@ -54,19 +59,22 @@ class FunctionDependencies(BaseModel):
         return list(v)
 
     def add_function(self, function: str):
-        self.functions.add(function)  # pylint: disable=no-member
+        self.functions.add(function)
+
+    def add_function_group(self, function_group: str):
+        self.function_groups.add(function_group)  # pylint: disable=no-member
 
     def add_llm(self, llm: str):
-        self.llms.add(llm)  # pylint: disable=no-member
+        self.llms.add(llm)
 
     def add_embedder(self, embedder: str):
-        self.embedders.add(embedder)  # pylint: disable=no-member
+        self.embedders.add(embedder)
 
     def add_memory_client(self, memory_client: str):
-        self.memory_clients.add(memory_client)  # pylint: disable=no-member
+        self.memory_clients.add(memory_client)
 
     def add_object_store(self, object_store: str):
-        self.object_stores.add(object_store)  # pylint: disable=no-member
+        self.object_stores.add(object_store)
 
     def add_retriever(self, retriever: str):
-        self.retrievers.add(retriever)  # pylint: disable=no-member
+        self.retrievers.add(retriever)

nat/data_models/gated_field_mixin.py

@@ -0,0 +1,242 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from collections.abc import Sequence
+from dataclasses import dataclass
+from re import Pattern
+
+from pydantic import model_validator
+
+
+@dataclass
+class GatedFieldMixinConfig:
+    """Configuration for a gated field mixin."""
+
+    field_name: str
+    default_if_supported: object | None
+    unsupported: Sequence[Pattern[str]] | None
+    supported: Sequence[Pattern[str]] | None
+    keys: Sequence[str]
+
+
+class GatedFieldMixin:
+    """
+    A mixin that gates a field based on specified keys.
+
+    This should be used to automatically validate a field based on a given key.
+
+    Parameters
+    ----------
+    field_name: `str`
+                The name of the field.
+    default_if_supported: `object | None`
+                          The default value of the field if it is supported for the key.
+    keys: `Sequence[str]`
+          A sequence of keys that are used to validate the field.
+    unsupported: `Sequence[Pattern[str]] | None`
+                 A sequence of regex patterns that match the key names NOT supported for the field.
+                 Defaults to None.
+    supported: `Sequence[Pattern[str]] | None`
+               A sequence of regex patterns that match the key names supported for the field.
+               Defaults to None.
+    """
+
+    def __init_subclass__(
+        cls,
+        field_name: str | None = None,
+        default_if_supported: object | None = None,
+        keys: Sequence[str] | None = None,
+        unsupported: Sequence[Pattern[str]] | None = None,
+        supported: Sequence[Pattern[str]] | None = None,
+    ) -> None:
+        """Store the class variables for the field and define the gated field validator."""
+        super().__init_subclass__()
+
+        # Check if this class directly inherits from GatedFieldMixin
+        has_gated_field_mixin = GatedFieldMixin in cls.__bases__
+
+        if has_gated_field_mixin:
+            if keys is None:
+                raise ValueError("keys must be provided when subclassing GatedFieldMixin")
+            if field_name is None:
+                raise ValueError("field_name must be provided when subclassing GatedFieldMixin")
+
+            cls._setup_direct_mixin(field_name, default_if_supported, unsupported, supported, keys)
+
+        # Always try to collect mixins and create validators for multiple inheritance
+        # This handles both direct inheritance and deep inheritance chains
+        all_mixins = cls._collect_all_mixin_configs()
+        if all_mixins:
+            cls._create_combined_validator(all_mixins)
+
+    @classmethod
+    def _setup_direct_mixin(
+        cls,
+        field_name: str,
+        default_if_supported: object | None,
+        unsupported: Sequence[Pattern[str]] | None,
+        supported: Sequence[Pattern[str]] | None,
+        keys: Sequence[str],
+    ) -> None:
+        """Set up a class that directly inherits from GatedFieldMixin."""
+        cls._validate_mixin_parameters(unsupported, supported, keys)
+
+        # Create and store validator
+        validator = cls._create_gated_field_validator(field_name, default_if_supported, unsupported, supported, keys)
+        validator_name = f"_gated_field_validator_{field_name}"
+        setattr(cls, validator_name, validator)
+
+        # Store mixin info for multiple inheritance
+        if not hasattr(cls, "_gated_field_mixins"):
+            cls._gated_field_mixins = []
+
+        cls._gated_field_mixins.append(
+            GatedFieldMixinConfig(
+                field_name,
+                default_if_supported,
+                unsupported,
+                supported,
+                keys,
+            ))
+
+    @classmethod
+    def _validate_mixin_parameters(
+        cls,
+        unsupported: Sequence[Pattern[str]] | None,
+        supported: Sequence[Pattern[str]] | None,
+        keys: Sequence[str],
+    ) -> None:
+        """Validate that all required parameters are provided."""
+        if unsupported is None and supported is None:
+            raise ValueError("Either unsupported or supported must be provided")
+        if unsupported is not None and supported is not None:
+            raise ValueError("Only one of unsupported or supported must be provided")
+        if len(keys) == 0:
+            raise ValueError("keys must be provided and non-empty when subclassing GatedFieldMixin")
+
+    @classmethod
+    def _create_gated_field_validator(
+        cls,
+        field_name: str,
+        default_if_supported: object | None,
+        unsupported: Sequence[Pattern[str]] | None,
+        supported: Sequence[Pattern[str]] | None,
+        keys: Sequence[str],
+    ):
+        """Create the model validator function."""
+
+        @model_validator(mode="after")
+        def gated_field_validator(self):
+            """Validate the gated field."""
+            current_value = getattr(self, field_name, None)
+            is_supported = cls._check_field_support(self, unsupported, supported, keys)
+            if not is_supported:
+                if current_value is not None:
+                    blocking_key = cls._find_blocking_key(self, unsupported, supported, keys)
+                    value = getattr(self, blocking_key, "<unknown>")
+                    raise ValueError(f"{field_name} is not supported for {blocking_key}: {value}")
+            elif current_value is None:
+                setattr(self, field_name, default_if_supported)
+            return self
+
+        return gated_field_validator
+
+    @classmethod
+    def _check_field_support(
+        cls,
+        instance: object,
+        unsupported: Sequence[Pattern[str]] | None,
+        supported: Sequence[Pattern[str]] | None,
+        keys: Sequence[str],
+    ) -> bool:
+        """Check if a specific field is supported based on its configuration and keys."""
+        seen = False
+        for key in keys:
+            if not hasattr(instance, key):
+                continue
+            seen = True
+            value = str(getattr(instance, key))
+            if supported is not None:
+                if any(p.search(value) for p in supported):
+                    return True
+            elif unsupported is not None:
+                if any(p.search(value) for p in unsupported):
+                    return False
+        return True if not seen else (unsupported is not None)
+
+    @classmethod
+    def _find_blocking_key(
+        cls,
+        instance: object,
+        unsupported: Sequence[Pattern[str]] | None,
+        supported: Sequence[Pattern[str]] | None,
+        keys: Sequence[str],
+    ) -> str:
+        """Find which key is blocking the field."""
+        for key in keys:
+            if not hasattr(instance, key):
+                continue
+            value = str(getattr(instance, key))
+            if supported is not None:
+                if not any(p.search(value) for p in supported):
+                    return key
+            elif unsupported is not None:
+                if any(p.search(value) for p in unsupported):
+                    return key
+
+        return "<unknown>"
+
+    @classmethod
+    def _collect_all_mixin_configs(cls) -> list[GatedFieldMixinConfig]:
+        """Collect all mixin configurations from base classes."""
+        all_mixins = []
+        for base in cls.__bases__:
+            if hasattr(base, "_gated_field_mixins"):
+                all_mixins.extend(base._gated_field_mixins)
+        return all_mixins
+
+    @classmethod
+    def _create_combined_validator(cls, all_mixins: list[GatedFieldMixinConfig]) -> None:
+        """Create a combined validator that handles all fields."""
+
+        @model_validator(mode="after")
+        def combined_gated_field_validator(self):
+            """Validate all gated fields."""
+            for mixin_config in all_mixins:
+                field_name_local = mixin_config.field_name
+                current_value = getattr(self, field_name_local, None)
+                if not self._check_field_support_instance(mixin_config):
+                    if current_value is not None:
+                        blocking_key = self._find_blocking_key_instance(mixin_config)
+                        value = getattr(self, blocking_key, "<unknown>")
+                        raise ValueError(f"{field_name_local} is not supported for {blocking_key}: {value}")
+                elif current_value is None:
+                    setattr(self, field_name_local, mixin_config.default_if_supported)
+
+            return self
+
+        cls._combined_gated_field_validator = combined_gated_field_validator
+
+        # Add helper methods
+        def _check_field_support_instance(self, mixin_config: GatedFieldMixinConfig) -> bool:
+            """Check if a specific field is supported based on its configuration and keys."""
+            return cls._check_field_support(self, mixin_config.unsupported, mixin_config.supported, mixin_config.keys)
+
+        def _find_blocking_key_instance(self, mixin_config: GatedFieldMixinConfig) -> str:
+            """Find which key is blocking the field."""
+            return cls._find_blocking_key(self, mixin_config.unsupported, mixin_config.supported, mixin_config.keys)
+
+        cls._check_field_support_instance = _check_field_support_instance
+        cls._find_blocking_key_instance = _find_blocking_key_instance

nat/data_models/intermediate_step.py

@@ -142,7 +142,7 @@ class IntermediateStepPayload(BaseModel):
     UUID: str = Field(default_factory=lambda: str(uuid.uuid4()))
 
     @property
-    def event_category(self) -> IntermediateStepCategory:  # pylint: disable=too-many-return-statements
+    def event_category(self) -> IntermediateStepCategory:
         match self.event_type:
             case IntermediateStepType.LLM_START:
                 return IntermediateStepCategory.LLM
@@ -180,7 +180,7 @@ class IntermediateStepPayload(BaseModel):
                 raise ValueError(f"Unknown event type: {self.event_type}")
 
     @property
-    def event_state(self) -> IntermediateStepState:  # pylint: disable=too-many-return-statements
+    def event_state(self) -> IntermediateStepState:
         match self.event_type:
             case IntermediateStepType.LLM_START:
                 return IntermediateStepState.START
@@ -290,7 +290,7 @@ class IntermediateStep(BaseModel):
         return self.payload.usage_info
 
     @property
-    def UUID(self) -> str:  # pylint: disable=invalid-name
+    def UUID(self) -> str:
         return self.payload.UUID
 
     @property