nvidia-nat 1.3.0.dev2__py3-none-any.whl → 1.3.0rc2__py3-none-any.whl

nat/data_models/api_server.py

@@ -36,6 +36,15 @@ from nat.utils.type_converter import GlobalTypeConverter
 FINISH_REASONS = frozenset({'stop', 'length', 'tool_calls', 'content_filter', 'function_call'})
 
 
+class UserMessageContentRoleType(str, Enum):
+    """
+    Enum representing chat message roles in API requests and responses.
+    """
+    USER = "user"
+    ASSISTANT = "assistant"
+    SYSTEM = "system"
+
+
 class Request(BaseModel):
     """
     Request is a data model that represents HTTP request attributes.
@@ -108,7 +117,7 @@ UserContent = typing.Annotated[TextContent | ImageContent | AudioContent, Discri
 
 class Message(BaseModel):
     content: str | list[UserContent]
-    role: str
+    role: UserMessageContentRoleType
 
 
 class ChatRequest(BaseModel):
@@ -164,7 +173,7 @@ class ChatRequest(BaseModel):
                     max_tokens: int | None = None,
                     top_p: float | None = None) -> "ChatRequest":
 
-        return ChatRequest(messages=[Message(content=data, role="user")],
+        return ChatRequest(messages=[Message(content=data, role=UserMessageContentRoleType.USER)],
                            model=model,
                            temperature=temperature,
                            max_tokens=max_tokens,
@@ -178,7 +187,7 @@ class ChatRequest(BaseModel):
                      max_tokens: int | None = None,
                      top_p: float | None = None) -> "ChatRequest":
 
-        return ChatRequest(messages=[Message(content=content, role="user")],
+        return ChatRequest(messages=[Message(content=content, role=UserMessageContentRoleType.USER)],
                            model=model,
                            temperature=temperature,
                            max_tokens=max_tokens,
@@ -187,29 +196,40 @@ class ChatRequest(BaseModel):
 
 class ChoiceMessage(BaseModel):
     content: str | None = None
-    role: str | None = None
+    role: UserMessageContentRoleType | None = None
 
 
 class ChoiceDelta(BaseModel):
     """Delta object for streaming responses (OpenAI-compatible)"""
     content: str | None = None
-    role: str | None = None
+    role: UserMessageContentRoleType | None = None
 
 
-class Choice(BaseModel):
+class ChoiceBase(BaseModel):
+    """Base choice model with common fields for both streaming and non-streaming responses"""
     model_config = ConfigDict(extra="allow")
-
-    message: ChoiceMessage | None = None
-    delta: ChoiceDelta | None = None
     finish_reason: typing.Literal['stop', 'length', 'tool_calls', 'content_filter', 'function_call'] | None = None
     index: int
-    # logprobs: ChoiceLogprobs | None = None
+
+
+class ChatResponseChoice(ChoiceBase):
+    """Choice model for non-streaming responses - contains message field"""
+    message: ChoiceMessage
+
+
+class ChatResponseChunkChoice(ChoiceBase):
+    """Choice model for streaming responses - contains delta field"""
+    delta: ChoiceDelta
+
+
+# Backward compatibility alias
+Choice = ChatResponseChoice
 
 
 class Usage(BaseModel):
-    prompt_tokens: int
-    completion_tokens: int
-    total_tokens: int
+    prompt_tokens: int | None = None
+    completion_tokens: int | None = None
+    total_tokens: int | None = None
 
 
 class ResponseSerializable(abc.ABC):
@@ -245,10 +265,10 @@ class ChatResponse(ResponseBaseModelOutput):
     model_config = ConfigDict(extra="allow")
     id: str
     object: str = "chat.completion"
-    model: str = ""
+    model: str = "unknown-model"
     created: datetime.datetime
-    choices: list[Choice]
-    usage: Usage | None = None
+    choices: list[ChatResponseChoice]
+    usage: Usage
     system_fingerprint: str | None = None
     service_tier: typing.Literal["scale", "default"] | None = None
 
@@ -264,22 +284,27 @@ class ChatResponse(ResponseBaseModelOutput):
                     object_: str | None = None,
                     model: str | None = None,
                     created: datetime.datetime | None = None,
-                    usage: Usage | None = None) -> "ChatResponse":
+                    usage: Usage) -> "ChatResponse":
 
         if id_ is None:
             id_ = str(uuid.uuid4())
         if object_ is None:
             object_ = "chat.completion"
         if model is None:
-            model = ""
+            model = "unknown-model"
         if created is None:
-            created = datetime.datetime.now(datetime.timezone.utc)
+            created = datetime.datetime.now(datetime.UTC)
 
         return ChatResponse(id=id_,
                             object=object_,
                             model=model,
                             created=created,
-                            choices=[Choice(index=0, message=ChoiceMessage(content=data), finish_reason="stop")],
+                            choices=[
+                                ChatResponseChoice(index=0,
+                                                   message=ChoiceMessage(content=data,
+                                                                         role=UserMessageContentRoleType.ASSISTANT),
+                                                   finish_reason="stop")
+                            ],
                             usage=usage)
 
 
@@ -293,9 +318,9 @@ class ChatResponseChunk(ResponseBaseModelOutput):
     model_config = ConfigDict(extra="allow")
 
     id: str
-    choices: list[Choice]
+    choices: list[ChatResponseChunkChoice]
     created: datetime.datetime
-    model: str = ""
+    model: str = "unknown-model"
     object: str = "chat.completion.chunk"
     system_fingerprint: str | None = None
     service_tier: typing.Literal["scale", "default"] | None = None
@@ -317,14 +342,20 @@ class ChatResponseChunk(ResponseBaseModelOutput):
         if id_ is None:
             id_ = str(uuid.uuid4())
         if created is None:
-            created = datetime.datetime.now(datetime.timezone.utc)
+            created = datetime.datetime.now(datetime.UTC)
         if model is None:
-            model = ""
+            model = "unknown-model"
         if object_ is None:
             object_ = "chat.completion.chunk"
 
         return ChatResponseChunk(id=id_,
-                                 choices=[Choice(index=0, message=ChoiceMessage(content=data), finish_reason="stop")],
+                                 choices=[
+                                     ChatResponseChunkChoice(index=0,
+                                                             delta=ChoiceDelta(
+                                                                 content=data,
+                                                                 role=UserMessageContentRoleType.ASSISTANT),
+                                                             finish_reason="stop")
+                                 ],
                                  created=created,
                                  model=model,
                                  object=object_)
@@ -335,7 +366,7 @@ class ChatResponseChunk(ResponseBaseModelOutput):
                                id_: str | None = None,
                                created: datetime.datetime | None = None,
                                model: str | None = None,
-                               role: str | None = None,
+                               role: UserMessageContentRoleType | None = None,
                                finish_reason: str | None = None,
                                usage: Usage | None = None,
                                system_fingerprint: str | None = None) -> "ChatResponseChunk":
@@ -343,9 +374,9 @@ class ChatResponseChunk(ResponseBaseModelOutput):
         if id_ is None:
             id_ = str(uuid.uuid4())
         if created is None:
-            created = datetime.datetime.now(datetime.timezone.utc)
+            created = datetime.datetime.now(datetime.UTC)
         if model is None:
-            model = ""
+            model = "unknown-model"
 
         delta = ChoiceDelta(content=content, role=role) if content is not None or role is not None else ChoiceDelta()
 
@@ -353,7 +384,14 @@ class ChatResponseChunk(ResponseBaseModelOutput):
 
         return ChatResponseChunk(
             id=id_,
-            choices=[Choice(index=0, message=None, delta=delta, finish_reason=final_finish_reason)],
+            choices=[
+                ChatResponseChunkChoice(
+                    index=0,
+                    delta=delta,
+                    finish_reason=typing.cast(
+                        typing.Literal['stop', 'length', 'tool_calls', 'content_filter', 'function_call'] | None,
+                        final_finish_reason))
+            ],
             created=created,
             model=model,
             object="chat.completion.chunk",
@@ -398,11 +436,6 @@ class GenerateResponse(BaseModel):
     value: str | None = "default"
 
 
-class UserMessageContentRoleType(str, Enum):
-    USER = "user"
-    ASSISTANT = "assistant"
-
-
 class WebSocketMessageType(str, Enum):
     """
     WebSocketMessageType is an Enum that represents WebSocket Message types.
@@ -485,7 +518,7 @@ class WebSocketUserMessage(BaseModel):
     security: Security = Security()
     error: Error = Error()
     schema_version: str = "1.0.0"
-    timestamp: str = str(datetime.datetime.now(datetime.timezone.utc))
+    timestamp: str = str(datetime.datetime.now(datetime.UTC))
 
 
 class WebSocketUserInteractionResponseMessage(BaseModel):
@@ -501,7 +534,7 @@ class WebSocketUserInteractionResponseMessage(BaseModel):
     security: Security = Security()
     error: Error = Error()
     schema_version: str = "1.0.0"
-    timestamp: str = str(datetime.datetime.now(datetime.timezone.utc))
+    timestamp: str = str(datetime.datetime.now(datetime.UTC))
 
 
 class SystemIntermediateStepContent(BaseModel):
@@ -527,7 +560,7 @@ class WebSocketSystemIntermediateStepMessage(BaseModel):
     conversation_id: str | None = None
     content: SystemIntermediateStepContent
     status: WebSocketMessageStatus
-    timestamp: str = str(datetime.datetime.now(datetime.timezone.utc))
+    timestamp: str = str(datetime.datetime.now(datetime.UTC))
 
 
 class SystemResponseContent(BaseModel):
@@ -551,7 +584,7 @@ class WebSocketSystemResponseTokenMessage(BaseModel):
     conversation_id: str | None = None
     content: SystemResponseContent | Error | GenerateResponse
     status: WebSocketMessageStatus
-    timestamp: str = str(datetime.datetime.now(datetime.timezone.utc))
+    timestamp: str = str(datetime.datetime.now(datetime.UTC))
 
     @field_validator("content")
     @classmethod
@@ -560,7 +593,7 @@ class WebSocketSystemResponseTokenMessage(BaseModel):
             raise ValueError(f"Field: content must be 'Error' when type is {WebSocketMessageType.ERROR_MESSAGE}")
 
         if info.data.get("type") == WebSocketMessageType.RESPONSE_MESSAGE and not isinstance(
-                value, (SystemResponseContent, GenerateResponse)):
+                value, SystemResponseContent | GenerateResponse):
             raise ValueError(
                 f"Field: content must be 'SystemResponseContent' when type is {WebSocketMessageType.RESPONSE_MESSAGE}")
         return value
@@ -582,7 +615,7 @@ class WebSocketSystemInteractionMessage(BaseModel):
     conversation_id: str | None = None
     content: HumanPrompt
     status: WebSocketMessageStatus
-    timestamp: str = str(datetime.datetime.now(datetime.timezone.utc))
+    timestamp: str = str(datetime.datetime.now(datetime.UTC))
 
 
 # ======== GenerateResponse Converters ========
@@ -622,7 +655,7 @@ GlobalTypeConverter.register_converter(_nat_chat_request_to_string)
 
 
 def _string_to_nat_chat_request(data: str) -> ChatRequest:
-    return ChatRequest.from_string(data, model="")
+    return ChatRequest.from_string(data, model="unknown-model")
 
 
 GlobalTypeConverter.register_converter(_string_to_nat_chat_request)
@@ -654,22 +687,12 @@ def _string_to_nat_chat_response(data: str) -> ChatResponse:
 GlobalTypeConverter.register_converter(_string_to_nat_chat_response)
 
 
-def _chat_response_to_chat_response_chunk(data: ChatResponse) -> ChatResponseChunk:
-    # Preserve original message structure for backward compatibility
-    return ChatResponseChunk(id=data.id, choices=data.choices, created=data.created, model=data.model)
-
-
-GlobalTypeConverter.register_converter(_chat_response_to_chat_response_chunk)
-
-
 # ======== ChatResponseChunk Converters ========
 def _chat_response_chunk_to_string(data: ChatResponseChunk) -> str:
     if data.choices and len(data.choices) > 0:
         choice = data.choices[0]
         if choice.delta and choice.delta.content:
             return choice.delta.content
-        if choice.message and choice.message.content:
-            return choice.message.content
     return ""
 
 
@@ -685,21 +708,6 @@ def _string_to_nat_chat_response_chunk(data: str) -> ChatResponseChunk:
 
 GlobalTypeConverter.register_converter(_string_to_nat_chat_response_chunk)
 
-
-# ======== AINodeMessageChunk Converters ========
-def _ai_message_chunk_to_nat_chat_response_chunk(data) -> ChatResponseChunk:
-    '''Converts LangChain AINodeMessageChunk to ChatResponseChunk'''
-    content = ""
-    if hasattr(data, 'content') and data.content is not None:
-        content = str(data.content)
-    elif hasattr(data, 'text') and data.text is not None:
-        content = str(data.text)
-    elif hasattr(data, 'message') and data.message is not None:
-        content = str(data.message)
-
-    return ChatResponseChunk.create_streaming_chunk(content=content, role="assistant", finish_reason=None)
-
-
 # Compatibility aliases with previous releases
 AIQChatRequest = ChatRequest
 AIQChoiceMessage = ChoiceMessage

nat/data_models/authentication.py

@@ -14,8 +14,8 @@
 # limitations under the License.
 
 import typing
+from datetime import UTC
 from datetime import datetime
-from datetime import timezone
 from enum import Enum
 
 import httpx
@@ -166,17 +166,31 @@ class BearerTokenCred(_CredBase):
 
 
 Credential = typing.Annotated[
-    typing.Union[
-        HeaderCred,
-        QueryCred,
-        CookieCred,
-        BasicAuthCred,
-        BearerTokenCred,
-    ],
+    HeaderCred | QueryCred | CookieCred | BasicAuthCred | BearerTokenCred,
     Field(discriminator="kind"),
 ]
 
 
+class TokenValidationResult(BaseModel):
+    """
+    Standard result for Bearer Token Validation.
+    """
+    model_config = ConfigDict(extra="forbid")
+
+    client_id: str | None = Field(description="OAuth2 client identifier")
+    scopes: list[str] | None = Field(default=None, description="List of granted scopes (introspection only)")
+    expires_at: int | None = Field(default=None, description="Token expiration time (Unix timestamp)")
+    audience: list[str] | None = Field(default=None, description="Token audiences (aud claim)")
+    subject: str | None = Field(default=None, description="Token subject (sub claim)")
+    issuer: str | None = Field(default=None, description="Token issuer (iss claim)")
+    token_type: str = Field(description="Token type")
+    active: bool | None = Field(default=True, description="Token active status")
+    nbf: int | None = Field(default=None, description="Not before time (Unix timestamp)")
+    iat: int | None = Field(default=None, description="Issued at time (Unix timestamp)")
+    jti: str | None = Field(default=None, description="JWT ID")
+    username: str | None = Field(default=None, description="Username (introspection only)")
+
+
 class AuthResult(BaseModel):
     """
     Represents the result of an authentication process.
@@ -193,7 +207,7 @@ class AuthResult(BaseModel):
         """
         Checks if the authentication token has expired.
         """
-        return bool(self.token_expires_at and datetime.now(timezone.utc) >= self.token_expires_at)
+        return bool(self.token_expires_at and datetime.now(UTC) >= self.token_expires_at)
 
     def as_requests_kwargs(self) -> dict[str, typing.Any]:
         """

nat/data_models/common.py

@@ -160,7 +160,7 @@ class TypedBaseModel(BaseModel):
 
     @staticmethod
     def discriminator(v: typing.Any) -> str | None:
-        # If its serialized, then we use the alias
+        # If it's serialized, then we use the alias
         if isinstance(v, dict):
             return v.get("_type", v.get("type"))
 

nat/data_models/component.py

@@ -27,6 +27,7 @@ class ComponentEnum(StrEnum):
     EVALUATOR = "evaluator"
     FRONT_END = "front_end"
     FUNCTION = "function"
+    FUNCTION_GROUP = "function_group"
     TTC_STRATEGY = "ttc_strategy"
     LLM_CLIENT = "llm_client"
     LLM_PROVIDER = "llm_provider"
@@ -47,6 +48,7 @@ class ComponentGroup(StrEnum):
     AUTHENTICATION = "authentication"
     EMBEDDERS = "embedders"
     FUNCTIONS = "functions"
+    FUNCTION_GROUPS = "function_groups"
     TTC_STRATEGIES = "ttc_strategies"
     LLMS = "llms"
     MEMORY = "memory"

nat/data_models/component_ref.py

@@ -102,6 +102,17 @@ class FunctionRef(ComponentRef):
         return ComponentGroup.FUNCTIONS
 
 
+class FunctionGroupRef(ComponentRef):
+    """
+    A reference to a function group in a NAT configuration object.
+    """
+
+    @property
+    @override
+    def component_group(self):
+        return ComponentGroup.FUNCTION_GROUPS
+
+
 class LLMRef(ComponentRef):
     """
     A reference to an LLM in a NAT configuration object.

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
@@ -185,10 +189,14 @@ class GeneralConfig(BaseModel):
 
     model_config = ConfigDict(protected_namespaces=())
 
-    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()