nvidia-nat 1.1.0a20251020__py3-none-any.whl

nat/front_ends/cron/__init__.py

@@ -0,0 +1,14 @@
+# 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.

nat/front_ends/fastapi/__init__.py

@@ -0,0 +1,14 @@
+# 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.

nat/front_ends/fastapi/auth_flow_handlers/http_flow_handler.py

@@ -0,0 +1,27 @@
+# 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.
+
+from nat.authentication.interfaces import FlowHandlerBase
+from nat.data_models.authentication import AuthenticatedContext
+from nat.data_models.authentication import AuthFlowType
+from nat.data_models.authentication import AuthProviderBaseConfig
+
+
+class HTTPAuthenticationFlowHandler(FlowHandlerBase):
+
+    async def authenticate(self, config: AuthProviderBaseConfig, method: AuthFlowType) -> AuthenticatedContext:
+
+        raise NotImplementedError(f"Authentication method '{method}' is not supported by the HTTP frontend."
+                                  f" Do you have WebSockets enabled?")

nat/front_ends/fastapi/auth_flow_handlers/websocket_flow_handler.py

@@ -0,0 +1,142 @@
+# 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 asyncio
+import logging
+import secrets
+from collections.abc import Awaitable
+from collections.abc import Callable
+from dataclasses import dataclass
+from dataclasses import field
+
+import pkce
+from authlib.common.errors import AuthlibBaseError as OAuthError
+from authlib.integrations.httpx_client import AsyncOAuth2Client
+
+from nat.authentication.interfaces import FlowHandlerBase
+from nat.authentication.oauth2.oauth2_auth_code_flow_provider_config import OAuth2AuthCodeFlowProviderConfig
+from nat.data_models.authentication import AuthenticatedContext
+from nat.data_models.authentication import AuthFlowType
+from nat.data_models.interactive import _HumanPromptOAuthConsent
+from nat.front_ends.fastapi.message_handler import WebSocketMessageHandler
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class FlowState:
+    future: asyncio.Future = field(default_factory=asyncio.Future, init=False)
+    challenge: str | None = None
+    verifier: str | None = None
+    client: AsyncOAuth2Client | None = None
+    config: OAuth2AuthCodeFlowProviderConfig | None = None
+
+
+class WebSocketAuthenticationFlowHandler(FlowHandlerBase):
+
+    def __init__(self,
+                 add_flow_cb: Callable[[str, FlowState], Awaitable[None]],
+                 remove_flow_cb: Callable[[str], Awaitable[None]],
+                 web_socket_message_handler: WebSocketMessageHandler):
+
+        self._add_flow_cb: Callable[[str, FlowState], Awaitable[None]] = add_flow_cb
+        self._remove_flow_cb: Callable[[str], Awaitable[None]] = remove_flow_cb
+        self._web_socket_message_handler: WebSocketMessageHandler = web_socket_message_handler
+
+    async def authenticate(self, config: OAuth2AuthCodeFlowProviderConfig,
+                           method: AuthFlowType) -> AuthenticatedContext:
+        if method == AuthFlowType.OAUTH2_AUTHORIZATION_CODE:
+            return await self._handle_oauth2_auth_code_flow(config)
+
+        raise NotImplementedError(f"Authentication method '{method}' is not supported by the websocket frontend.")
+
+    def create_oauth_client(self, config: OAuth2AuthCodeFlowProviderConfig) -> AsyncOAuth2Client:
+        try:
+            return AsyncOAuth2Client(client_id=config.client_id,
+                                     client_secret=config.client_secret,
+                                     redirect_uri=config.redirect_uri,
+                                     scope=" ".join(config.scopes) if config.scopes else None,
+                                     token_endpoint=config.token_url,
+                                     code_challenge_method='S256' if config.use_pkce else None,
+                                     token_endpoint_auth_method=config.token_endpoint_auth_method)
+        except (OAuthError, ValueError, TypeError) as e:
+            raise RuntimeError(f"Invalid OAuth2 configuration: {e}") from e
+        except Exception as e:
+            raise RuntimeError(f"Failed to create OAuth2 client: {e}") from e
+
+    def _create_authorization_url(self,
+                                  client: AsyncOAuth2Client,
+                                  config: OAuth2AuthCodeFlowProviderConfig,
+                                  state: str,
+                                  verifier: str = None,
+                                  challenge: str = None) -> str:
+        """
+        Create OAuth authorization URL with proper error handling.
+
+        Args:
+            client: The OAuth2 client instance
+            config: OAuth2 configuration
+            state: OAuth state parameter
+            verifier: PKCE verifier (if using PKCE)
+            challenge: PKCE challenge (if using PKCE)
+
+        Returns:
+            The authorization URL
+        """
+        try:
+            authorization_url, _ = client.create_authorization_url(
+                config.authorization_url,
+                state=state,
+                code_verifier=verifier if config.use_pkce else None,
+                code_challenge=challenge if config.use_pkce else None,
+                **(config.authorization_kwargs or {})
+            )
+            return authorization_url
+        except (OAuthError, ValueError, TypeError) as e:
+            raise RuntimeError(f"Error creating OAuth authorization URL: {e}") from e
+
+    async def _handle_oauth2_auth_code_flow(self, config: OAuth2AuthCodeFlowProviderConfig) -> AuthenticatedContext:
+
+        state = secrets.token_urlsafe(16)
+        flow_state = FlowState(config=config)
+
+        flow_state.client = self.create_oauth_client(config)
+
+        if config.use_pkce:
+            verifier, challenge = pkce.generate_pkce_pair()
+            flow_state.verifier = verifier
+            flow_state.challenge = challenge
+
+        authorization_url = self._create_authorization_url(client=flow_state.client,
+                                                           config=config,
+                                                           state=state,
+                                                           verifier=flow_state.verifier,
+                                                           challenge=flow_state.challenge)
+
+        await self._add_flow_cb(state, flow_state)
+        await self._web_socket_message_handler.create_websocket_message(_HumanPromptOAuthConsent(text=authorization_url)
+                                                                        )
+        try:
+            token = await asyncio.wait_for(flow_state.future, timeout=300)
+        except TimeoutError as exc:
+            raise RuntimeError("Authentication flow timed out after 5 minutes.") from exc
+        finally:
+
+            await self._remove_flow_cb(state)
+
+        return AuthenticatedContext(headers={"Authorization": f"Bearer {token['access_token']}"},
+                                    metadata={
+                                        "expires_at": token.get("expires_at"), "raw_token": token
+                                    })

nat/front_ends/fastapi/dask_client_mixin.py

@@ -0,0 +1,65 @@
+# 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 typing
+from abc import ABC
+from collections.abc import AsyncGenerator
+from collections.abc import Generator
+from contextlib import asynccontextmanager
+from contextlib import contextmanager
+
+if typing.TYPE_CHECKING:
+    from dask.distributed import Client
+
+
+class DaskClientMixin(ABC):
+
+    @asynccontextmanager
+    async def client(self, address: str) -> AsyncGenerator["Client"]:
+        """
+        Async context manager for obtaining a Dask client.
+
+        Yields
+        ------
+        Client
+            An async Dask client connected to the scheduler. The client is automatically closed when exiting the
+            context manager.
+        """
+        from dask.distributed import Client
+        client = await Client(address=address, asynchronous=True)
+
+        try:
+            yield client
+        finally:
+            await client.close()
+
+    @contextmanager
+    def blocking_client(self, address: str) -> Generator["Client"]:
+        """
+        context manager for obtaining a blocking Dask client.
+
+        Yields
+        ------
+        Client
+            A blocking Dask client connected to the scheduler. The client is automatically closed when exiting the
+            context manager.
+        """
+        from dask.distributed import Client
+        client = Client(address=address)
+
+        try:
+            yield client
+        finally:
+            client.close()

nat/front_ends/fastapi/fastapi_front_end_config.py

@@ -0,0 +1,272 @@
+# 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
+import os
+import sys
+import typing
+from datetime import datetime
+from pathlib import Path
+
+from pydantic import BaseModel
+from pydantic import Field
+from pydantic import field_validator
+
+from nat.data_models.component_ref import ObjectStoreRef
+from nat.data_models.front_end import FrontEndBaseConfig
+from nat.data_models.step_adaptor import StepAdaptorConfig
+
+logger = logging.getLogger(__name__)
+
+YAML_EXTENSIONS = (".yaml", ".yml")
+
+
+def _is_reserved(path: Path) -> bool:
+    """
+    Check if a path is reserved in the current Python version and platform.
+
+    On Windows, this function checks if the path is reserved in the current Python version.
+    On other platforms, returns False
+    """
+    if sys.platform != "win32":
+        return False
+    if sys.version_info >= (3, 13):
+        return os.path.isreserved(path)
+    return path.is_reserved()
+
+
+class EvaluateRequest(BaseModel):
+    """Request model for the evaluate endpoint."""
+    config_file: str = Field(description="Path to the configuration file for evaluation")
+    job_id: str | None = Field(default=None, description="Unique identifier for the evaluation job")
+    reps: int = Field(default=1, gt=0, description="Number of repetitions for the evaluation, defaults to 1")
+    expiry_seconds: int = Field(
+        default=3600,
+        gt=0,
+        description="Optional time (in seconds) before the job expires. Clamped between 600 (10 min) and 86400 (24h).")
+
+    @field_validator('job_id', mode='after')
+    @classmethod
+    def validate_job_id(cls, job_id: str):
+        job_id = job_id.strip()
+        job_id_path = Path(job_id)
+        if len(job_id_path.parts) > 1 or job_id_path.resolve().name != job_id:
+            raise ValueError(
+                f"Job ID '{job_id}' contains invalid characters. Only alphanumeric characters and underscores are"
+                " allowed.")
+
+        if _is_reserved(job_id_path):
+            # reserved names is Windows specific
+            raise ValueError(f"Job ID '{job_id}' is a reserved name. Please choose a different name.")
+
+        return job_id
+
+    @field_validator('config_file', mode='after')
+    @classmethod
+    def validate_config_file(cls, config_file: str):
+        config_file = config_file.strip()
+        config_file_path = Path(config_file).resolve()
+
+        # Ensure the config file is a YAML file
+        if config_file_path.suffix.lower() not in YAML_EXTENSIONS:
+            raise ValueError(f"Config file '{config_file}' must be a YAML file with one of the following extensions: "
+                             f"{', '.join(YAML_EXTENSIONS)}")
+
+        if _is_reserved(config_file_path):
+            # reserved names is Windows specific
+            raise ValueError(f"Config file '{config_file}' is a reserved name. Please choose a different name.")
+
+        if not config_file_path.exists():
+            raise ValueError(f"Config file '{config_file}' does not exist. Please provide a valid path.")
+
+        return config_file
+
+
+class BaseAsyncResponse(BaseModel):
+    """Base model for async responses."""
+    job_id: str = Field(description="Unique identifier for the job")
+    status: str = Field(description="Current status of the job")
+
+
+class EvaluateResponse(BaseAsyncResponse):
+    """Response model for the evaluate endpoint."""
+    pass
+
+
+class AsyncGenerateResponse(BaseAsyncResponse):
+    """Response model for the async generation endpoint."""
+    pass
+
+
+class BaseAsyncStatusResponse(BaseModel):
+    """Base model for async status responses."""
+    job_id: str = Field(description="Unique identifier for the evaluation job")
+    status: str = Field(description="Current status of the evaluation job")
+    error: str | None = Field(default=None, description="Error message if the job failed")
+    created_at: datetime = Field(description="Timestamp when the job was created")
+    updated_at: datetime = Field(description="Timestamp when the job was last updated")
+    expires_at: datetime | None = Field(default=None, description="Timestamp when the job will expire")
+
+
+class EvaluateStatusResponse(BaseAsyncStatusResponse):
+    """Response model for the evaluate status endpoint."""
+    config_file: str = Field(description="Path to the configuration file used for evaluation")
+    output_path: str | None = Field(default=None,
+                                    description="Path to the output file if the job completed successfully")
+
+
+class AsyncGenerationStatusResponse(BaseAsyncStatusResponse):
+    output: dict | None = Field(
+        default=None,
+        description="Output of the generate request, this is only available if the job completed successfully.")
+
+
+class FastApiFrontEndConfig(FrontEndBaseConfig, name="fastapi"):
+    """
+    A FastAPI based front end that allows a NAT workflow to be served as a microservice.
+    """
+
+    class EndpointBase(BaseModel):
+
+        method: typing.Literal["GET", "POST", "PUT", "DELETE"]
+        description: str
+        path: str | None = Field(
+            default=None,
+            description=("Path for the default workflow. If None, no workflow endpoint is created."),
+        )
+        websocket_path: str | None = Field(
+            default=None,
+            description=("Path for the websocket. If None, no websocket is created."),
+        )
+        openai_api_path: str | None = Field(
+            default=None,
+            description=("Path for the default workflow using the OpenAI API Specification. "
+                         "If None, no workflow endpoint with the OpenAI API Specification is created."),
+        )
+        openai_api_v1_path: str | None = Field(
+            default=None,
+            description=("Path for the OpenAI v1 Chat Completions API compatible endpoint. "
+                         "If provided, creates a single endpoint that handles both streaming and "
+                         "non-streaming requests based on the 'stream' parameter, following the "
+                         "OpenAI Chat Completions API specification exactly."),
+        )
+
+    class Endpoint(EndpointBase):
+        function_name: str = Field(description="The name of the function to call for this endpoint")
+
+    class CrossOriginResourceSharing(BaseModel):
+        allow_origins: list[str] | None = Field(
+            default=None, description=" A list of origins that should be permitted to make cross-origin requests.")
+        allow_origin_regex: str | None = Field(
+            default=None,
+            description="A permitted regex string to match against origins to make cross-origin requests",
+        )
+        allow_methods: list[str] | None = Field(
+            default_factory=lambda: ['GET'],
+            description="A list of HTTP methods that should be allowed for cross-origin requests.")
+        allow_headers: list[str] | None = Field(
+            default_factory=list,
+            description="A list of HTTP request headers that should be supported for cross-origin requests.")
+        allow_credentials: bool | None = Field(
+            default=False,
+            description="Indicate that cookies should be supported for cross-origin requests.",
+        )
+        expose_headers: list[str] | None = Field(
+            default_factory=list,
+            description="Indicate any response headers that should be made accessible to the browser.",
+        )
+        max_age: int | None = Field(
+            default=600,
+            description="Sets a maximum time in seconds for browsers to cache CORS responses.",
+        )
+
+    root_path: str = Field(default="", description="The root path for the API")
+    host: str = Field(default="localhost", description="Host to bind the server to")
+    port: int = Field(default=8000, description="Port to bind the server to", ge=0, le=65535)
+    reload: bool = Field(default=False, description="Enable auto-reload for development")
+    workers: int = Field(default=1, description="Number of workers to run", ge=1)
+    scheduler_address: str | None = Field(
+        default=None,
+        description=("Address of the Dask scheduler to use for async jobs. If None, a Dask local cluster is created. "
+                     "Note: This requires the optional dask dependency to be installed."))
+    db_url: str | None = Field(
+        default=None,
+        description=
+        "SQLAlchemy database URL for storing async job metadata, if unset a temporary SQLite database is used.")
+    max_running_async_jobs: int = Field(
+        default=10,
+        description=(
+            "Maximum number of async jobs to run concurrently, this controls the number of dask workers created. "
+            "This parameter is only used when scheduler_address is `None` and a Dask local cluster is created."),
+        ge=1)
+    dask_log_level: str = Field(
+        default="WARNING",
+        description="Logging level for Dask.",
+    )
+    step_adaptor: StepAdaptorConfig = StepAdaptorConfig()
+
+    workflow: typing.Annotated[EndpointBase, Field(description="Endpoint for the default workflow.")] = EndpointBase(
+        method="POST",
+        path="/generate",
+        websocket_path="/websocket",
+        openai_api_path="/chat",
+        openai_api_v1_path="/v1/chat/completions",
+        description="Executes the default NAT workflow from the loaded configuration ",
+    )
+
+    evaluate: typing.Annotated[EndpointBase, Field(description="Endpoint for evaluating workflows.")] = EndpointBase(
+        method="POST",
+        path="/evaluate",
+        description="Evaluates the performance and accuracy of the workflow on a dataset",
+    )
+
+    oauth2_callback_path: str | None = Field(
+        default="/auth/redirect",
+        description="OAuth2.0 authentication callback endpoint. If None, no OAuth2 callback endpoint is created.")
+
+    endpoints: list[Endpoint] = Field(
+        default_factory=list,
+        description=("Additional endpoints to add to the FastAPI app which run functions within the NAT configuration. "
+                     "Each endpoint must have a unique path."))
+
+    cors: CrossOriginResourceSharing = Field(
+        default_factory=CrossOriginResourceSharing,
+        description="Cross origin resource sharing configuration for the FastAPI app")
+
+    use_gunicorn: bool = Field(
+        default=False,
+        description="Use Gunicorn to run the FastAPI app",
+    )
+    runner_class: str | None = Field(
+        default=None,
+        description=("The NAT runner class to use when launching the FastAPI app from multiple processes. "
+                     "Each runner is responsible for loading and running the NAT workflow. "
+                     "Note: This is different from the worker class used by Gunicorn."),
+    )
+
+    object_store: ObjectStoreRef | None = Field(
+        default=None,
+        description=(
+            "Object store reference for the FastAPI app. If present, static files can be uploaded via a POST "
+            "request to '/static' and files will be served from the object store. The files will be served from the "
+            "object store at '/static/{file_name}'."))
+
+
+# Compatibility aliases with previous releases
+AIQEvaluateRequest = EvaluateRequest
+AIQEvaluateResponse = EvaluateResponse
+AIQAsyncGenerateResponse = AsyncGenerateResponse
+AIQEvaluateStatusResponse = EvaluateStatusResponse
+AIQAsyncGenerationStatusResponse = AsyncGenerationStatusResponse

nat/front_ends/fastapi/fastapi_front_end_controller.py

@@ -0,0 +1,68 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import asyncio
+import logging
+
+from fastapi import FastAPI
+from uvicorn import Config
+from uvicorn import Server
+
+logger = logging.getLogger(__name__)
+
+
+class _FastApiFrontEndController:
+    """
+    _FastApiFrontEndController class controls the spawing and tear down of the API server in environments where
+    the server is needed and not already running.
+    """
+
+    def __init__(self, app: FastAPI):
+        self._app: FastAPI = app
+        self._server: Server | None = None
+        self._server_background_task: asyncio.Task | None = None
+
+    async def start_server(self, host: str, port: int) -> None:
+        """Starts the API server."""
+
+        server_host = host
+        server_port = port
+
+        config = Config(app=self._app, host=server_host, port=server_port, log_level="warning")
+        self._server = Server(config=config)
+
+        try:
+            self._server_background_task = asyncio.create_task(self._server.serve())
+        except asyncio.CancelledError as e:
+            error_message = f"Task error occurred while starting API server: {str(e)}"
+            logger.error(error_message)
+            raise RuntimeError(error_message) from e
+        except Exception as e:
+            error_message = f"Unexpected error occurred while starting API server: {str(e)}"
+            logger.exception(error_message)
+            raise RuntimeError(error_message) from e
+
+    async def stop_server(self) -> None:
+        """Stops the API server."""
+        if not self._server or not self._server_background_task:
+            return
+
+        try:
+            self._server.should_exit = True
+            await self._server_background_task
+        except asyncio.CancelledError as e:
+            logger.exception("Server shutdown failed: %s", str(e))
+        except Exception as e:
+            logger.exception("Unexpected error occurred: %s", str(e))