nvidia-nat 1.3.0a20250929__py3-none-any.whl → 1.3.0a20250930__py3-none-any.whl

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,13 +166,7 @@ class BearerTokenCred(_CredBase):
 
 
 Credential = typing.Annotated[
-    typing.Union[
-        HeaderCred,
-        QueryCred,
-        CookieCred,
-        BasicAuthCred,
-        BearerTokenCred,
-    ],
+    HeaderCred | QueryCred | CookieCred | BasicAuthCred | BearerTokenCred,
     Field(discriminator="kind"),
 ]
 
@@ -213,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/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/eval/evaluator/base_evaluator.py

@@ -71,7 +71,7 @@ class BaseEvaluator(ABC):
             TqdmPositionRegistry.release(tqdm_position)
 
         # Compute average if possible
-        numeric_scores = [item.score for item in output_items if isinstance(item.score, (int, float))]
+        numeric_scores = [item.score for item in output_items if isinstance(item.score, int | float)]
         avg_score = round(sum(numeric_scores) / len(numeric_scores), 2) if numeric_scores else None
 
         return EvalOutput(average_score=avg_score, eval_output_items=output_items)

nat/eval/swe_bench_evaluator/evaluate.py

@@ -204,7 +204,7 @@ class SweBenchEvaluator:
         # if report file is not present, return empty EvalOutput
         avg_score = 0.0
         if report_file.exists():
-            with open(report_file, "r", encoding="utf-8") as f:
+            with open(report_file, encoding="utf-8") as f:
                 report = json.load(f)
                 resolved_instances = report.get("resolved_instances", 0)
                 total_instances = report.get("total_instances", 0)

nat/eval/tunable_rag_evaluator/evaluate.py

@@ -14,7 +14,7 @@
 # limitations under the License.
 
 import logging
-from typing import Callable
+from collections.abc import Callable
 
 from langchain.output_parsers import ResponseSchema
 from langchain.output_parsers import StructuredOutputParser

nat/experimental/decorators/experimental_warning_decorator.py

@@ -137,8 +137,7 @@ def experimental(func: Any = None, *, feature_name: str | None = None, metadata:
         @functools.wraps(func)
         def sync_gen_wrapper(*args, **kwargs) -> Generator[Any, Any, Any]:
             issue_experimental_warning(function_name, feature_name, metadata)
-            for item in func(*args, **kwargs):
-                yield item  # yield the original item
+            yield from func(*args, **kwargs)  # yield the original item
 
         return sync_gen_wrapper
 

nat/experimental/test_time_compute/selection/llm_based_output_merging_selector.py

@@ -71,7 +71,7 @@ class LLMBasedOutputMergingSelector(StrategyBase):
             raise ImportError("langchain-core is not installed. Please install it to use SingleShotMultiPlanPlanner.\n"
                               "This error can be resolved by installing nvidia-nat-langchain.")
 
-        from typing import Callable
+        from collections.abc import Callable
 
         from pydantic import BaseModel
 

nat/front_ends/console/authentication_flow_handler.py

@@ -14,13 +14,16 @@
 # limitations under the License.
 
 import asyncio
+import logging
 import secrets
 import webbrowser
 from dataclasses import dataclass
 from dataclasses import field
 
 import click
+import httpx
 import pkce
+from authlib.common.errors import AuthlibBaseError as OAuthError
 from authlib.integrations.httpx_client import AsyncOAuth2Client
 from fastapi import FastAPI
 from fastapi import Request
@@ -32,6 +35,8 @@ from nat.data_models.authentication import AuthFlowType
 from nat.data_models.authentication import AuthProviderBaseConfig
 from nat.front_ends.fastapi.fastapi_front_end_controller import _FastApiFrontEndController
 
+logger = logging.getLogger(__name__)
+
 
 # --------------------------------------------------------------------------- #
 # Helpers                                                                     #
@@ -87,17 +92,53 @@ class ConsoleAuthenticationFlowHandler(FlowHandlerBase):
         """
         Separated for easy overriding in tests (to inject ASGITransport).
         """
-        client = AsyncOAuth2Client(
-            client_id=cfg.client_id,
-            client_secret=cfg.client_secret,
-            redirect_uri=cfg.redirect_uri,
-            scope=" ".join(cfg.scopes) if cfg.scopes else None,
-            token_endpoint=cfg.token_url,
-            token_endpoint_auth_method=cfg.token_endpoint_auth_method,
-            code_challenge_method="S256" if cfg.use_pkce else None,
-        )
-        self._oauth_client = client
-        return client
+        try:
+            client = AsyncOAuth2Client(
+                client_id=cfg.client_id,
+                client_secret=cfg.client_secret,
+                redirect_uri=cfg.redirect_uri,
+                scope=" ".join(cfg.scopes) if cfg.scopes else None,
+                token_endpoint=cfg.token_url,
+                token_endpoint_auth_method=cfg.token_endpoint_auth_method,
+                code_challenge_method="S256" if cfg.use_pkce else None,
+            )
+            self._oauth_client = client
+            return client
+        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 = None,
+                                  challenge: str | None = 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:
+            auth_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 auth_url
+        except (OAuthError, ValueError, TypeError) as e:
+            raise RuntimeError(f"Error creating OAuth authorization URL: {e}") from e
 
     # --------------------------- HTTP Basic ------------------------------ #
     @staticmethod
@@ -131,13 +172,12 @@ class ConsoleAuthenticationFlowHandler(FlowHandlerBase):
             flow_state.verifier = verifier
             flow_state.challenge = challenge
 
-        auth_url, _ = client.create_authorization_url(
-            cfg.authorization_url,
-            state=state,
-            code_verifier=flow_state.verifier if cfg.use_pkce else None,
-            code_challenge=flow_state.challenge if cfg.use_pkce else None,
-            **(cfg.authorization_kwargs or {})
-        )
+        # Create authorization URL using helper function
+        auth_url = self._create_authorization_url(client=client,
+                                                  config=cfg,
+                                                  state=state,
+                                                  verifier=flow_state.verifier,
+                                                  challenge=flow_state.challenge)
 
         # Register flow + maybe spin up redirect handler
         async with self._server_lock:
@@ -149,14 +189,18 @@ class ConsoleAuthenticationFlowHandler(FlowHandlerBase):
             self._flows[state] = flow_state
             self._active_flows += 1
 
-        click.echo("Your browser has been opened for authentication.")
-        webbrowser.open(auth_url)
+        try:
+            webbrowser.open(auth_url)
+            click.echo("Your browser has been opened for authentication.")
+        except Exception as e:
+            logger.error("Browser open failed: %s", e)
+            raise RuntimeError(f"Browser open failed: {e}") from e
 
         # Wait for the redirect to land
         try:
             token = await asyncio.wait_for(flow_state.future, timeout=300)
-        except asyncio.TimeoutError:
-            raise RuntimeError("Authentication timed out (5 min).")
+        except TimeoutError as exc:
+            raise RuntimeError("Authentication timed out (5 min).") from exc
         finally:
             async with self._server_lock:
                 self._flows.pop(state, None)
@@ -175,9 +219,9 @@ class ConsoleAuthenticationFlowHandler(FlowHandlerBase):
     # --------------- redirect server / in‑process app -------------------- #
     async def _build_redirect_app(self) -> FastAPI:
         """
-        * If cfg.run_redirect_local_server == True → start a uvicorn server (old behaviour).
-        * Else → only build the FastAPI app and save it to `self._redirect_app`
-                 for in‑process testing with ASGITransport.
+        * If cfg.run_redirect_local_server == True → start a local server.
+        * Else → only build the redirect app and save it to `self._redirect_app`
+                 for in‑process testing.
         """
         app = FastAPI()
 
@@ -195,8 +239,16 @@ class ConsoleAuthenticationFlowHandler(FlowHandlerBase):
                     state=state,
                 )
                 flow_state.future.set_result(token)
-            except Exception as exc:  # noqa: BLE001
-                flow_state.future.set_exception(exc)
+            except OAuthError as e:
+                flow_state.future.set_exception(
+                    RuntimeError(f"Authorization server rejected request: {e.error} ({e.description})"))
+                return "Authentication failed: Authorization server rejected the request. You may close this tab."
+            except httpx.HTTPError as e:
+                flow_state.future.set_exception(RuntimeError(f"Network error during token fetch: {e}"))
+                return "Authentication failed: Network error occurred. You may close this tab."
+            except Exception as e:
+                flow_state.future.set_exception(RuntimeError(f"Authentication failed: {e}"))
+                return "Authentication failed: An unexpected error occurred. You may close this tab."
             return "Authentication successful – you may close this tab."
 
         return app
@@ -213,7 +265,7 @@ class ConsoleAuthenticationFlowHandler(FlowHandlerBase):
 
             asyncio.create_task(self._server_controller.start_server(host="localhost", port=8000))
 
-            # Give uvicorn a moment to bind sockets before we return
+            # Give the server a moment to bind sockets before we return
             await asyncio.sleep(0.3)
         except Exception as exc:  # noqa: BLE001
             raise RuntimeError(f"Failed to start redirect server: {exc}") from exc
@@ -227,7 +279,7 @@ class ConsoleAuthenticationFlowHandler(FlowHandlerBase):
     @property
     def redirect_app(self) -> FastAPI | None:
         """
-        In “test‑mode” (run_redirect_local_server=False) the in‑memory FastAPI
-        app is exposed so you can mount it on `httpx.ASGITransport`.
+        In test mode (run_redirect_local_server=False) the in‑memory redirect
+        app is exposed for testing purposes.
         """
         return self._redirect_app

nat/front_ends/console/console_front_end_plugin.py

@@ -88,7 +88,7 @@ class ConsoleFrontEndPlugin(SimpleFrontEndPluginBase[ConsoleFrontEndConfig]):
         elif (self.front_end_config.input_file):
 
             # Run the workflow
-            with open(self.front_end_config.input_file, "r", encoding="utf-8") as f:
+            with open(self.front_end_config.input_file, encoding="utf-8") as f:
 
                 async with session_manager.workflow.run(f) as runner:
                     runner_outputs = await runner.result(to_type=str)

nat/front_ends/fastapi/auth_flow_handlers/websocket_flow_handler.py

@@ -22,6 +22,7 @@ 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
@@ -61,14 +62,50 @@ class WebSocketAuthenticationFlowHandler(FlowHandlerBase):
 
         raise NotImplementedError(f"Authentication method '{method}' is not supported by the websocket frontend.")
 
-    def create_oauth_client(self, config: OAuth2AuthCodeFlowProviderConfig):
-        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)
+    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:
 
@@ -82,21 +119,19 @@ class WebSocketAuthenticationFlowHandler(FlowHandlerBase):
             flow_state.verifier = verifier
             flow_state.challenge = challenge
 
-        authorization_url, _ = flow_state.client.create_authorization_url(
-            config.authorization_url,
-            state=state,
-            code_verifier=flow_state.verifier if config.use_pkce else None,
-            code_challenge=flow_state.challenge if config.use_pkce else None,
-            **(config.authorization_kwargs or {})
-        )
+        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 asyncio.TimeoutError:
-            raise RuntimeError("Authentication flow timed out after 5 minutes.")
+        except TimeoutError as exc:
+            raise RuntimeError("Authentication flow timed out after 5 minutes.") from exc
         finally:
 
             await self._remove_flow_cb(state)

nat/front_ends/fastapi/fastapi_front_end_plugin_worker.py

@@ -25,18 +25,21 @@ from collections.abc import Callable
 from contextlib import asynccontextmanager
 from pathlib import Path
 
+import httpx
+from authlib.common.errors import AuthlibBaseError as OAuthError
 from fastapi import Body
 from fastapi import FastAPI
+from fastapi import HTTPException
 from fastapi import Request
 from fastapi import Response
 from fastapi import UploadFile
-from fastapi.exceptions import HTTPException
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import StreamingResponse
 from pydantic import BaseModel
 from pydantic import Field
 from starlette.websockets import WebSocket
 
+from nat.builder.function import Function
 from nat.builder.workflow_builder import WorkflowBuilder
 from nat.data_models.api_server import ChatRequest
 from nat.data_models.api_server import ChatResponse
@@ -241,6 +244,7 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
         await self.add_evaluate_route(app, SessionManager(await builder.build()))
         await self.add_static_files_route(app, builder)
         await self.add_authorization_route(app)
+        await self.add_mcp_client_tool_list_route(app, builder)
 
         for ep in self.front_end_config.endpoints:
 
@@ -1071,8 +1075,13 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
                                                code_verifier=verifier,
                                                state=state)
                 flow_state.future.set_result(res)
+            except OAuthError as e:
+                flow_state.future.set_exception(
+                    RuntimeError(f"Authorization server rejected request: {e.error} ({e.description})"))
+            except httpx.HTTPError as e:
+                flow_state.future.set_exception(RuntimeError(f"Network error during token fetch: {e}"))
             except Exception as e:
-                flow_state.future.set_exception(e)
+                flow_state.future.set_exception(RuntimeError(f"Authentication failed: {e}"))
 
             return HTMLResponse(content=AUTH_REDIRECT_SUCCESS_HTML,
                                 status_code=200,
@@ -1088,6 +1097,183 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
                 methods=["GET"],
                 description="Handles the authorization code and state returned from the Authorization Code Grant Flow.")
 
+    async def add_mcp_client_tool_list_route(self, app: FastAPI, builder: WorkflowBuilder):
+        """Add the MCP client tool list endpoint to the FastAPI app."""
+        from typing import Any
+
+        from pydantic import BaseModel
+
+        class MCPToolInfo(BaseModel):
+            name: str
+            description: str
+            server: str
+            available: bool
+
+        class MCPClientToolListResponse(BaseModel):
+            mcp_clients: list[dict[str, Any]]
+
+        async def get_mcp_client_tool_list() -> MCPClientToolListResponse:
+            """
+            Get the list of MCP tools from all MCP clients in the workflow configuration.
+            Checks session health and compares with workflow function group configuration.
+            """
+            mcp_clients_info = []
+
+            try:
+                # Get all function groups from the builder
+                function_groups = builder._function_groups
+
+                # Find MCP client function groups
+                for group_name, configured_group in function_groups.items():
+                    if configured_group.config.type != "mcp_client":
+                        continue
+
+                    from nat.plugins.mcp.client_impl import MCPClientConfig
+
+                    config = configured_group.config
+                    assert isinstance(config, MCPClientConfig)
+
+                    # Reuse the existing MCP client session stored on the function group instance
+                    group_instance = configured_group.instance
+
+                    client = group_instance.mcp_client
+                    if client is None:
+                        raise RuntimeError(f"MCP client not found for group {group_name}")
+
+                    try:
+                        session_healthy = False
+                        server_tools: dict[str, Any] = {}
+
+                        try:
+                            server_tools = await client.get_tools()
+                            session_healthy = True
+                        except Exception as e:
+                            logger.exception(f"Failed to connect to MCP server {client.server_name}: {e}")
+                            session_healthy = False
+
+                        # Get workflow function group configuration (configured client-side tools)
+                        configured_short_names: set[str] = set()
+                        configured_full_to_fn: dict[str, Function] = {}
+                        try:
+                            # Pass a no-op filter function to bypass any default filtering that might check
+                            # health status, preventing potential infinite recursion during health status checks.
+                            async def pass_through_filter(fn):
+                                return fn
+
+                            accessible_functions = await group_instance.get_accessible_functions(
+                                filter_fn=pass_through_filter)
+                            configured_full_to_fn = accessible_functions
+                            configured_short_names = {name.split('.', 1)[1] for name in accessible_functions.keys()}
+                        except Exception as e:
+                            logger.exception(f"Failed to get accessible functions for group {group_name}: {e}")
+
+                        # Build alias->original mapping and override configs from overrides
+                        alias_to_original: dict[str, str] = {}
+                        override_configs: dict[str, Any] = {}
+                        try:
+                            if config.tool_overrides is not None:
+                                for orig_name, override in config.tool_overrides.items():
+                                    if override.alias is not None:
+                                        alias_to_original[override.alias] = orig_name
+                                        override_configs[override.alias] = override
+                                    else:
+                                        override_configs[orig_name] = override
+                        except Exception:
+                            pass
+
+                        # Create tool info list (always return configured tools; mark availability)
+                        tools_info: list[dict[str, Any]] = []
+                        available_count = 0
+                        for wf_fn, fn_short in zip(configured_full_to_fn.values(), configured_short_names):
+                            orig_name = alias_to_original.get(fn_short, fn_short)
+                            available = session_healthy and (orig_name in server_tools)
+                            if available:
+                                available_count += 1
+
+                            # Prefer tool override description, then workflow function description,
+                            # then server description
+                            description = ""
+                            if fn_short in override_configs and override_configs[fn_short].description:
+                                description = override_configs[fn_short].description
+                            elif wf_fn.description:
+                                description = wf_fn.description
+                            elif available and orig_name in server_tools:
+                                description = server_tools[orig_name].description or ""
+
+                            tools_info.append(
+                                MCPToolInfo(name=fn_short,
+                                            description=description or "",
+                                            server=client.server_name,
+                                            available=available).model_dump())
+
+                        # Sort tools_info by name to maintain consistent ordering
+                        tools_info.sort(key=lambda x: x['name'])
+
+                        mcp_clients_info.append({
+                            "function_group": group_name,
+                            "server": client.server_name,
+                            "transport": config.server.transport,
+                            "session_healthy": session_healthy,
+                            "tools": tools_info,
+                            "total_tools": len(configured_short_names),
+                            "available_tools": available_count
+                        })
+
+                    except Exception as e:
+                        logger.error(f"Error processing MCP client {group_name}: {e}")
+                        mcp_clients_info.append({
+                            "function_group": group_name,
+                            "server": "unknown",
+                            "transport": config.server.transport if config.server else "unknown",
+                            "session_healthy": False,
+                            "error": str(e),
+                            "tools": [],
+                            "total_tools": 0,
+                            "workflow_tools": 0
+                        })
+
+                return MCPClientToolListResponse(mcp_clients=mcp_clients_info)
+
+            except Exception as e:
+                logger.error(f"Error in MCP client tool list endpoint: {e}")
+                raise HTTPException(status_code=500, detail=f"Failed to retrieve MCP client information: {str(e)}")
+
+        # Add the route to the FastAPI app
+        app.add_api_route(
+            path="/mcp/client/tool/list",
+            endpoint=get_mcp_client_tool_list,
+            methods=["GET"],
+            response_model=MCPClientToolListResponse,
+            description="Get list of MCP client tools with session health and workflow configuration comparison",
+            responses={
+                200: {
+                    "description": "Successfully retrieved MCP client tool information",
+                    "content": {
+                        "application/json": {
+                            "example": {
+                                "mcp_clients": [{
+                                    "function_group": "mcp_tools",
+                                    "server": "streamable-http:http://localhost:9901/mcp",
+                                    "transport": "streamable-http",
+                                    "session_healthy": True,
+                                    "tools": [{
+                                        "name": "tool_a",
+                                        "description": "Tool A description",
+                                        "server": "streamable-http:http://localhost:9901/mcp",
+                                        "available": True
+                                    }],
+                                    "total_tools": 1,
+                                    "available_tools": 1
+                                }]
+                            }
+                        }
+                    }
+                },
+                500: {
+                    "description": "Internal Server Error"
+                }
+            })
+
     async def _add_flow(self, state: str, flow_state: FlowState):
         async with self._outstanding_flows_lock:
             self._outstanding_flows[state] = flow_state

nat/front_ends/fastapi/job_store.py

@@ -370,7 +370,7 @@ class JobStore(DaskClientMixin):
                 # Convert BaseModel to JSON string for storage
                 output = output.model_dump_json(round_trip=True)
 
-            if isinstance(output, (dict, list)):
+            if isinstance(output, dict | list):
                 # Convert dict or list to JSON string for storage
                 output = json.dumps(output)
 
@@ -555,7 +555,7 @@ class JobStore(DaskClientMixin):
                         logger.exception("Failed to expire %s", job_id)
 
                 await session.execute(
-                    (update(JobInfo).where(JobInfo.job_id.in_(successfully_expired)).values(is_expired=True)))
+                    update(JobInfo).where(JobInfo.job_id.in_(successfully_expired)).values(is_expired=True))
 
 
 def get_db_engine(db_url: str | None = None, echo: bool = False, use_async: bool = True) -> "Engine | AsyncEngine":

nat/front_ends/fastapi/message_handler.py

@@ -105,10 +105,10 @@ class WebSocketMessageHandler:
                 if (isinstance(validated_message, WebSocketUserMessage)):
                     await self.process_workflow_request(validated_message)
 
-                elif isinstance(validated_message,
-                                (WebSocketSystemResponseTokenMessage,
-                                 WebSocketSystemIntermediateStepMessage,
-                                 WebSocketSystemInteractionMessage)):
+                elif isinstance(
+                        validated_message,
+                        WebSocketSystemResponseTokenMessage | WebSocketSystemIntermediateStepMessage
+                        | WebSocketSystemInteractionMessage):
                     # These messages are already handled by self.create_websocket_message(data_model=value, …)
                     # No further processing is needed here.
                     pass