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

nat/control_flow/router_agent/register.py

@@ -0,0 +1,91 @@
+# 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 logging
+
+from pydantic import Field
+
+from nat.builder.builder import Builder
+from nat.builder.framework_enum import LLMFrameworkEnum
+from nat.builder.function_info import FunctionInfo
+from nat.cli.register_workflow import register_function
+from nat.data_models.agent import AgentBaseConfig
+from nat.data_models.component_ref import FunctionRef
+
+logger = logging.getLogger(__name__)
+
+
+class RouterAgentWorkflowConfig(AgentBaseConfig, name="router_agent"):
+    """
+    A router agent takes in the incoming message, combines it with a prompt and the list of branches,
+    and ask a LLM about which branch to take.
+    """
+    description: str = Field(default="Router Agent Workflow", description="Description of this functions use.")
+    branches: list[FunctionRef] = Field(default_factory=list,
+                                        description="The list of branches to provide to the router agent.")
+    system_prompt: str | None = Field(default=None, description="Provides the system prompt to use with the agent.")
+    user_prompt: str | None = Field(default=None, description="Provides the prompt to use with the agent.")
+    max_router_retries: int = Field(
+        default=3, description="Maximum number of retries if the router agent fails to choose a branch.")
+
+
+@register_function(config_type=RouterAgentWorkflowConfig, framework_wrappers=[LLMFrameworkEnum.LANGCHAIN])
+async def router_agent_workflow(config: RouterAgentWorkflowConfig, builder: Builder):
+    from langchain_core.messages.human import HumanMessage
+    from langgraph.graph.state import CompiledStateGraph
+
+    from nat.agent.base import AGENT_LOG_PREFIX
+    from nat.control_flow.router_agent.agent import RouterAgentGraph
+    from nat.control_flow.router_agent.agent import RouterAgentGraphState
+    from nat.control_flow.router_agent.agent import create_router_agent_prompt
+
+    prompt = create_router_agent_prompt(config)
+    llm = await builder.get_llm(config.llm_name, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
+    branches = await builder.get_tools(tool_names=config.branches, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
+    if not branches:
+        raise ValueError(f"No branches specified for Router Agent '{config.llm_name}'")
+
+    graph: CompiledStateGraph = await RouterAgentGraph(
+        llm=llm,
+        branches=branches,
+        prompt=prompt,
+        max_router_retries=config.max_router_retries,
+        detailed_logs=config.verbose,
+        log_response_max_chars=config.log_response_max_chars,
+    ).build_graph()
+
+    async def _response_fn(input_message: str) -> str:
+        try:
+            message = HumanMessage(content=input_message)
+            state = RouterAgentGraphState(forward_message=message)
+
+            result_dict = await graph.ainvoke(state)
+            result_state = RouterAgentGraphState(**result_dict)
+
+            output_message = result_state.messages[-1]
+            return str(output_message.content)
+
+        except Exception as ex:
+            logger.exception("%s Router Agent failed with exception: %s", AGENT_LOG_PREFIX, ex)
+            if config.verbose:
+                return str(ex)
+            return f"Router agent failed with exception: {ex}"
+
+    try:
+        yield FunctionInfo.from_fn(_response_fn, description=config.description)
+    except GeneratorExit:
+        logger.exception("%s Workflow exited early!", AGENT_LOG_PREFIX)
+    finally:
+        logger.debug("%s Cleaning up router_agent workflow.", AGENT_LOG_PREFIX)

nat/control_flow/sequential_executor.py

@@ -0,0 +1,166 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-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 logging
+import typing
+
+from langchain_core.tools.base import BaseTool
+from pydantic import BaseModel
+from pydantic import Field
+
+from nat.builder.builder import Builder
+from nat.builder.framework_enum import LLMFrameworkEnum
+from nat.builder.function import Function
+from nat.builder.function_info import FunctionInfo
+from nat.cli.register_workflow import register_function
+from nat.data_models.component_ref import FunctionRef
+from nat.data_models.function import FunctionBaseConfig
+from nat.utils.type_utils import DecomposedType
+
+logger = logging.getLogger(__name__)
+
+
+class ToolExecutionConfig(BaseModel):
+    """Configuration for individual tool execution within sequential execution."""
+
+    use_streaming: bool = Field(default=False, description="Whether to use streaming output for the tool.")
+
+
+class SequentialExecutorConfig(FunctionBaseConfig, name="sequential_executor"):
+    """Configuration for sequential execution of a list of functions."""
+
+    tool_list: list[FunctionRef] = Field(default_factory=list,
+                                         description="A list of functions to execute sequentially.")
+    tool_execution_config: dict[str, ToolExecutionConfig] = Field(default_factory=dict,
+                                                                  description="Optional configuration for each"
+                                                                  "tool in the sequential execution tool list."
+                                                                  "Keys must match the tool names from the"
+                                                                  "tool_list.")
+    raise_type_incompatibility: bool = Field(
+        default=False,
+        description="Default to False. Check if the adjacent tools are type compatible,"
+        "which means the output type of the previous function is compatible with the input type of the next function."
+        "If set to True, any incompatibility will raise an exception. If set to false, the incompatibility will only"
+        "generate a warning message and the sequential execution will continue.")
+
+
+def _get_function_output_type(function: Function, tool_execution_config: dict[str, ToolExecutionConfig]) -> type:
+    function_config = tool_execution_config.get(function.instance_name, None)
+    if function_config:
+        return function.streaming_output_type if function_config.use_streaming else function.single_output_type
+    else:
+        return function.single_output_type
+
+
+def _validate_function_type_compatibility(src_fn: Function,
+                                          target_fn: Function,
+                                          tool_execution_config: dict[str, ToolExecutionConfig]) -> None:
+    src_output_type = _get_function_output_type(src_fn, tool_execution_config)
+    target_input_type = target_fn.input_type
+    logger.debug(
+        f"The output type of the {src_fn.instance_name} function is {str(src_output_type)}, is not compatible with"
+        f"the input type of the {target_fn.instance_name} function, which is {str(target_input_type)}")
+
+    is_compatible = DecomposedType.is_type_compatible(src_output_type, target_input_type)
+    if not is_compatible:
+        raise ValueError(
+            f"The output type of the {src_fn.instance_name} function is {str(src_output_type)}, is not compatible with"
+            f"the input type of the {target_fn.instance_name} function, which is {str(target_input_type)}")
+
+
+async def _validate_tool_list_type_compatibility(sequential_executor_config: SequentialExecutorConfig,
+                                                 builder: Builder) -> tuple[type, type]:
+    tool_list = sequential_executor_config.tool_list
+    tool_execution_config = sequential_executor_config.tool_execution_config
+
+    function_list = await builder.get_functions(tool_list)
+    if not function_list:
+        raise RuntimeError("The function list is empty")
+    input_type = function_list[0].input_type
+
+    if len(function_list) > 1:
+        for src_fn, target_fn in zip(function_list[0:-1], function_list[1:]):
+            try:
+                _validate_function_type_compatibility(src_fn, target_fn, tool_execution_config)
+            except ValueError as e:
+                raise ValueError(f"The sequential tool list has incompatible types: {e}")
+
+    output_type = _get_function_output_type(function_list[-1], tool_execution_config)
+    logger.debug(f"The input type of the sequential executor tool list is {str(input_type)},"
+                 f"the output type is {str(output_type)}")
+
+    return (input_type, output_type)
+
+
+@register_function(config_type=SequentialExecutorConfig, framework_wrappers=[LLMFrameworkEnum.LANGCHAIN])
+async def sequential_execution(config: SequentialExecutorConfig, builder: Builder):
+    logger.debug(f"Initializing sequential executor with tool list: {config.tool_list}")
+
+    tools: list[BaseTool] = await builder.get_tools(tool_names=config.tool_list,
+                                                    wrapper_type=LLMFrameworkEnum.LANGCHAIN)
+    tools_dict: dict[str, BaseTool] = {tool.name: tool for tool in tools}
+
+    try:
+        input_type, output_type = await _validate_tool_list_type_compatibility(config, builder)
+    except ValueError as e:
+        if config.raise_type_incompatibility:
+            logger.error(f"The sequential executor tool list has incompatible types: {e}")
+            raise
+        else:
+            logger.warning(f"The sequential executor tool list has incompatible types: {e}")
+            input_type = typing.Any
+            output_type = typing.Any
+    except Exception as e:
+        raise ValueError(f"Error with the sequential executor tool list: {e}")
+
+    # The type annotation of _sequential_function_execution is dynamically set according to the tool list
+    async def _sequential_function_execution(initial_tool_input):
+        logger.debug(f"Executing sequential executor with tool list: {config.tool_list}")
+
+        tool_list: list[FunctionRef] = config.tool_list
+        tool_input = initial_tool_input
+        tool_response = None
+
+        for tool_name in tool_list:
+            tool = tools_dict[tool_name]
+            tool_execution_config = config.tool_execution_config.get(tool_name, None)
+            logger.debug(f"Executing tool {tool_name} with input: {tool_input}")
+            try:
+                if tool_execution_config:
+                    if tool_execution_config.use_streaming:
+                        output = ""
+                        async for chunk in tool.astream(tool_input):
+                            output += chunk.content
+                        tool_response = output
+                    else:
+                        tool_response = await tool.ainvoke(tool_input)
+                else:
+                    tool_response = await tool.ainvoke(tool_input)
+            except Exception as e:
+                logger.error(f"Error with tool {tool_name}: {e}")
+                raise
+
+            # The input of the next tool is the response of the previous tool
+            tool_input = tool_response
+
+        return tool_response
+
+    # Dynamically set the annotations for the function
+    _sequential_function_execution.__annotations__ = {"initial_tool_input": input_type, "return": output_type}
+    logger.debug(f"Sequential executor function annotations: {_sequential_function_execution.__annotations__}")
+
+    yield FunctionInfo.from_fn(_sequential_function_execution,
+                               description="Executes a list of functions sequentially."
+                               "The input of the next tool is the response of the previous tool.")

nat/data_models/agent.py

@@ -0,0 +1,34 @@
+# 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 Field
+from pydantic import PositiveInt
+
+from nat.data_models.component_ref import LLMRef
+from nat.data_models.function import FunctionBaseConfig
+
+
+class AgentBaseConfig(FunctionBaseConfig):
+    """Base configuration class for all NAT agents with common fields."""
+
+    workflow_alias: str | None = Field(
+        default=None,
+        description=("The alias of the workflow. Useful when the agent is configured as a workflow "
+                     "and needs to expose a customized name as a tool."))
+    llm_name: LLMRef = Field(description="The LLM model to use with the agent.")
+    verbose: bool = Field(default=False, description="Set the verbosity of the agent's logging.")
+    description: str = Field(description="The description of this function's use.")
+    log_response_max_chars: PositiveInt = Field(
+        default=1000, description="Maximum number of characters to display in logs when logging responses.")

nat/data_models/api_server.py

@@ -273,7 +273,7 @@ class ChatResponse(ResponseBaseModelOutput):
         if model is None:
             model = ""
         if created is None:
-            created = datetime.datetime.now(datetime.timezone.utc)
+            created = datetime.datetime.now(datetime.UTC)
 
         return ChatResponse(id=id_,
                             object=object_,
@@ -317,7 +317,7 @@ 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 = ""
         if object_ is None:
@@ -343,7 +343,7 @@ 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 = ""
 
@@ -485,7 +485,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 +501,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 +527,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 +551,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 +560,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 +582,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 ========
@@ -688,7 +688,7 @@ 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'''
+    '''Converts LangChain/LangGraph AINodeMessageChunk to ChatResponseChunk'''
     content = ""
     if hasattr(data, 'content') and data.content is not None:
         content = str(data.content)

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)