aiqtoolkit 1.2.0a20250707__py3-none-any.whl → 1.2.0rc1__py3-none-any.whl

aiq/front_ends/console/authentication_flow_handler.py

@@ -0,0 +1,233 @@
+# 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 secrets
+import webbrowser
+from dataclasses import dataclass
+from dataclasses import field
+
+import click
+import pkce
+from authlib.integrations.httpx_client import AsyncOAuth2Client
+from fastapi import FastAPI
+from fastapi import Request
+
+from aiq.authentication.interfaces import FlowHandlerBase
+from aiq.authentication.oauth2.oauth2_auth_code_flow_provider_config import OAuth2AuthCodeFlowProviderConfig
+from aiq.data_models.authentication import AuthenticatedContext
+from aiq.data_models.authentication import AuthFlowType
+from aiq.data_models.authentication import AuthProviderBaseConfig
+from aiq.front_ends.fastapi.fastapi_front_end_controller import _FastApiFrontEndController
+
+
+# --------------------------------------------------------------------------- #
+# Helpers                                                                     #
+# --------------------------------------------------------------------------- #
+@dataclass
+class _FlowState:
+    future: asyncio.Future = field(default_factory=asyncio.Future, init=False)
+    challenge: str | None = None
+    verifier: str | None = None
+    token_url: str | None = None
+    use_pkce: bool | None = None
+
+
+# --------------------------------------------------------------------------- #
+# Main handler                                                                #
+# --------------------------------------------------------------------------- #
+class ConsoleAuthenticationFlowHandler(FlowHandlerBase):
+    """
+    Authentication helper for CLI / console environments.  Supports:
+
+      • HTTP Basic (username/password)
+      • OAuth 2 Authorization‑Code with optional PKCE
+    """
+
+    # ----------------------------- lifecycle ----------------------------- #
+    def __init__(self) -> None:
+        super().__init__()
+        self._server_controller: _FastApiFrontEndController | None = None
+        self._redirect_app: FastAPI | None = None  # ★ NEW
+        self._flows: dict[str, _FlowState] = {}
+        self._active_flows = 0
+        self._server_lock = asyncio.Lock()
+        self._oauth_client: AsyncOAuth2Client | None = None
+
+    # ----------------------------- public API ---------------------------- #
+    async def authenticate(
+        self,
+        config: AuthProviderBaseConfig,
+        method: AuthFlowType,
+    ) -> AuthenticatedContext:
+        if method == AuthFlowType.HTTP_BASIC:
+            return self._handle_http_basic()
+        if method == AuthFlowType.OAUTH2_AUTHORIZATION_CODE:
+            if (not isinstance(config, OAuth2AuthCodeFlowProviderConfig)):
+                raise ValueError("Requested OAuth2 Authorization Code Flow but passed invalid config")
+
+            return await self._handle_oauth2_auth_code_flow(config)
+
+        raise NotImplementedError(f"Auth method “{method}” not supported.")
+
+    # --------------------- OAuth2 helper factories ----------------------- #
+    def construct_oauth_client(self, cfg: OAuth2AuthCodeFlowProviderConfig) -> AsyncOAuth2Client:
+        """
+        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
+
+    # --------------------------- HTTP Basic ------------------------------ #
+    @staticmethod
+    def _handle_http_basic() -> AuthenticatedContext:
+        username = click.prompt("Username", type=str)
+        password = click.prompt("Password", type=str, hide_input=True)
+
+        import base64
+        credentials = f"{username}:{password}"
+        encoded_credentials = base64.b64encode(credentials.encode("utf-8")).decode("ascii")
+
+        return AuthenticatedContext(
+            headers={"Authorization": f"Bearer {encoded_credentials}"},
+            metadata={
+                "username": username, "password": password
+            },
+        )
+
+    # --------------------- OAuth2 Authorization‑Code --------------------- #
+    async def _handle_oauth2_auth_code_flow(self, cfg: OAuth2AuthCodeFlowProviderConfig) -> AuthenticatedContext:
+        state = secrets.token_urlsafe(16)
+        flow_state = _FlowState()
+        client = self.construct_oauth_client(cfg)
+
+        flow_state.token_url = cfg.token_url
+        flow_state.use_pkce = cfg.use_pkce
+
+        # PKCE bits
+        if cfg.use_pkce:
+            verifier, challenge = pkce.generate_pkce_pair()
+            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 {})
+        )
+
+        # Register flow + maybe spin up redirect handler
+        async with self._server_lock:
+            if (not self._redirect_app):
+                self._redirect_app = await self._build_redirect_app()
+
+            await self._start_redirect_server()
+
+            self._flows[state] = flow_state
+            self._active_flows += 1
+
+        click.echo("Your browser has been opened for authentication.")
+        webbrowser.open(auth_url)
+
+        # 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).")
+        finally:
+            async with self._server_lock:
+                self._flows.pop(state, None)
+                self._active_flows -= 1
+
+                if self._active_flows == 0:
+                    await self._stop_redirect_server()
+
+        return AuthenticatedContext(
+            headers={"Authorization": f"Bearer {token['access_token']}"},
+            metadata={
+                "expires_at": token.get("expires_at"), "raw_token": token
+            },
+        )
+
+    # --------------- 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.
+        """
+        app = FastAPI()
+
+        @app.get("/auth/redirect")
+        async def handle_redirect(request: Request):
+            state = request.query_params.get("state")
+            if not state or state not in self._flows:
+                return "Invalid state; restart authentication."
+            flow_state = self._flows[state]
+            try:
+                token = await self._oauth_client.fetch_token(  # type: ignore[arg-type]
+                    url=flow_state.token_url,
+                    authorization_response=str(request.url),
+                    code_verifier=flow_state.verifier if flow_state.use_pkce else None,
+                    state=state,
+                )
+                flow_state.future.set_result(token)
+            except Exception as exc:  # noqa: BLE001
+                flow_state.future.set_exception(exc)
+            return "Authentication successful – you may close this tab."
+
+        return app
+
+    async def _start_redirect_server(self) -> None:
+        # If the server is already running, do nothing
+        if self._server_controller:
+            return
+        try:
+            if not self._redirect_app:
+                raise RuntimeError("Redirect app not built.")
+
+            self._server_controller = _FastApiFrontEndController(self._redirect_app)
+
+            asyncio.create_task(self._server_controller.start_server(host="localhost", port=8000))
+
+            # Give uvicorn 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
+
+    async def _stop_redirect_server(self) -> None:
+        if self._server_controller:
+            await self._server_controller.stop_server()
+            self._server_controller = None
+
+    # ------------------------- test helpers ------------------------------ #
+    @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`.
+        """
+        return self._redirect_app

aiq/front_ends/console/console_front_end_plugin.py

@@ -25,6 +25,7 @@ from aiq.data_models.interactive import HumanPromptModelType
 from aiq.data_models.interactive import HumanResponse
 from aiq.data_models.interactive import HumanResponseText
 from aiq.data_models.interactive import InteractionPrompt
+from aiq.front_ends.console.authentication_flow_handler import ConsoleAuthenticationFlowHandler
 from aiq.front_ends.console.console_front_end_config import ConsoleFrontEndConfig
 from aiq.front_ends.simple_base.simple_front_end_plugin_base import SimpleFrontEndPluginBase
 from aiq.runtime.session import AIQSessionManager
@@ -43,12 +44,18 @@ async def prompt_for_input_cli(question: InteractionPrompt) -> HumanResponse:
 
         return HumanResponseText(text=user_response)
 
-    raise ValueError("Unsupported human propmt input type. The run command only supports the 'HumanPromptText' "
+    raise ValueError("Unsupported human prompt input type. The run command only supports the 'HumanPromptText' "
                      "input type. Please use the 'serve' command to ensure full support for all input types.")
 
 
 class ConsoleFrontEndPlugin(SimpleFrontEndPluginBase[ConsoleFrontEndConfig]):
 
+    def __init__(self, full_config):
+        super().__init__(full_config=full_config)
+
+        # Set the authentication flow handler
+        self.auth_flow_handler = ConsoleAuthenticationFlowHandler()
+
     async def pre_run(self):
 
         if (not self.front_end_config.input_query and not self.front_end_config.input_file):
@@ -81,7 +88,9 @@ class ConsoleFrontEndPlugin(SimpleFrontEndPluginBase[ConsoleFrontEndConfig]):
 
             async def run_single_query(query):
 
-                async with session_manager.session(user_input_callback=prompt_for_input_cli) as session:
+                async with session_manager.session(
+                        user_input_callback=prompt_for_input_cli,
+                        user_authentication_callback=self.auth_flow_handler.authenticate) as session:
                     async with session.run(query) as runner:
                         base_output = await runner.result(to_type=str)
 

aiq/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 aiq.authentication.interfaces import FlowHandlerBase
+from aiq.data_models.authentication import AuthenticatedContext
+from aiq.data_models.authentication import AuthFlowType
+from aiq.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?")

aiq/front_ends/fastapi/auth_flow_handlers/websocket_flow_handler.py

@@ -0,0 +1,107 @@
+# 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.integrations.httpx_client import AsyncOAuth2Client
+
+from aiq.authentication.interfaces import FlowHandlerBase
+from aiq.authentication.oauth2.oauth2_auth_code_flow_provider_config import OAuth2AuthCodeFlowProviderConfig
+from aiq.data_models.authentication import AuthenticatedContext
+from aiq.data_models.authentication import AuthFlowType
+from aiq.data_models.interactive import _HumanPromptOAuthConsent
+from aiq.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):
+        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)
+
+    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, _ = 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 {})
+        )
+
+        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.")
+        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
+                                    })

aiq/front_ends/fastapi/fastapi_front_end_config.py

@@ -22,6 +22,7 @@ from pydantic import BaseModel
 from pydantic import Field
 from pydantic import field_validator
 
+from aiq.data_models.component_ref import ObjectStoreRef
 from aiq.data_models.front_end import FrontEndBaseConfig
 from aiq.data_models.step_adaptor import StepAdaptorConfig
 
@@ -138,6 +139,13 @@ class FastApiFrontEndConfig(FrontEndBaseConfig, name="fastapi"):
             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")
@@ -183,6 +191,7 @@ class FastApiFrontEndConfig(FrontEndBaseConfig, name="fastapi"):
         path="/generate",
         websocket_path="/websocket",
         openai_api_path="/chat",
+        openai_api_v1_path="/v1/chat/completions",
         description="Executes the default AIQ Toolkit workflow from the loaded configuration ",
     )
 
@@ -192,6 +201,10 @@ class FastApiFrontEndConfig(FrontEndBaseConfig, name="fastapi"):
         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=(
@@ -212,3 +225,10 @@ class FastApiFrontEndConfig(FrontEndBaseConfig, name="fastapi"):
                      "Each runner is responsible for loading and running the AIQ Toolkit 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}'."))

aiq/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, exc_info=True)
+            raise RuntimeError(error_message) from e
+        except Exception as e:
+            error_message = f"Unexpected error occurred while starting API server: {str(e)}"
+            logger.error(error_message, exc_info=True)
+            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.error("Server shutdown failed: %s", str(e), exc_info=True)
+        except Exception as e:
+            logger.error("Unexpected error occurred: %s", str(e), exc_info=True)

aiq/front_ends/fastapi/fastapi_front_end_plugin.py

@@ -13,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import logging
 import os
 import tempfile
 import typing
@@ -23,6 +24,8 @@ from aiq.front_ends.fastapi.fastapi_front_end_plugin_worker import FastApiFrontE
 from aiq.front_ends.fastapi.main import get_app
 from aiq.utils.io.yaml_tools import yaml_dump
 
+logger = logging.getLogger(__name__)
+
 
 class FastApiFrontEndPlugin(FrontEndBase[FastApiFrontEndConfig]):
 
@@ -44,7 +47,7 @@ class FastApiFrontEndPlugin(FrontEndBase[FastApiFrontEndConfig]):
     async def run(self):
 
         # Write the entire config to a temporary file
-        with tempfile.NamedTemporaryFile(mode="w", prefix="aiq_config", suffix=".yml", delete=True) as config_file:
+        with tempfile.NamedTemporaryFile(mode="w", prefix="aiq_config", suffix=".yml", delete=False) as config_file:
 
             # Get as dict
             config_dict = self.full_config.model_dump(mode="json", by_alias=True, round_trip=True)
@@ -52,12 +55,16 @@ class FastApiFrontEndPlugin(FrontEndBase[FastApiFrontEndConfig]):
             # Write to YAML file
             yaml_dump(config_dict, config_file)
 
+            # Save the config file path for cleanup (required on Windows due to delete=False workaround)
+            config_file_name = config_file.name
+
             # Set the config file in the environment
             os.environ["AIQ_CONFIG_FILE"] = str(config_file.name)
 
             # Set the worker class in the environment
             os.environ["AIQ_FRONT_END_WORKER"] = self.get_worker_class_name()
 
+        try:
             if not self.front_end_config.use_gunicorn:
                 import uvicorn
 
@@ -101,3 +108,9 @@ class FastApiFrontEndPlugin(FrontEndBase[FastApiFrontEndConfig]):
                 }
 
                 StandaloneApplication(app, options=options).run()
+
+        finally:
+            try:
+                os.remove(config_file_name)
+            except OSError as e:
+                logger.error(f"Warning: Failed to delete temp file {config_file_name}: {e}")