aiqtoolkit 1.2.0a20250706__py3-none-any.whl → 1.2.0a20250730__py3-none-any.whl

aiq/front_ends/fastapi/fastapi_front_end_plugin_worker.py

@@ -20,8 +20,9 @@ import time
 import typing
 from abc import ABC
 from abc import abstractmethod
+from collections.abc import Awaitable
+from collections.abc import Callable
 from contextlib import asynccontextmanager
-from functools import partial
 from pathlib import Path
 
 from fastapi import BackgroundTasks
@@ -29,11 +30,13 @@ from fastapi import Body
 from fastapi import FastAPI
 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 aiq.builder.workflow_builder import WorkflowBuilder
 from aiq.data_models.api_server import AIQChatRequest
@@ -41,9 +44,14 @@ from aiq.data_models.api_server import AIQChatResponse
 from aiq.data_models.api_server import AIQChatResponseChunk
 from aiq.data_models.api_server import AIQResponseIntermediateStep
 from aiq.data_models.config import AIQConfig
+from aiq.data_models.object_store import KeyAlreadyExistsError
+from aiq.data_models.object_store import NoSuchKeyError
 from aiq.eval.config import EvaluationRunOutput
 from aiq.eval.evaluate import EvaluationRun
 from aiq.eval.evaluate import EvaluationRunConfig
+from aiq.front_ends.fastapi.auth_flow_handlers.http_flow_handler import HTTPAuthenticationFlowHandler
+from aiq.front_ends.fastapi.auth_flow_handlers.websocket_flow_handler import FlowState
+from aiq.front_ends.fastapi.auth_flow_handlers.websocket_flow_handler import WebSocketAuthenticationFlowHandler
 from aiq.front_ends.fastapi.fastapi_front_end_config import AIQAsyncGenerateResponse
 from aiq.front_ends.fastapi.fastapi_front_end_config import AIQAsyncGenerationStatusResponse
 from aiq.front_ends.fastapi.fastapi_front_end_config import AIQEvaluateRequest
@@ -52,11 +60,12 @@ from aiq.front_ends.fastapi.fastapi_front_end_config import AIQEvaluateStatusRes
 from aiq.front_ends.fastapi.fastapi_front_end_config import FastApiFrontEndConfig
 from aiq.front_ends.fastapi.job_store import JobInfo
 from aiq.front_ends.fastapi.job_store import JobStore
+from aiq.front_ends.fastapi.message_handler import WebSocketMessageHandler
 from aiq.front_ends.fastapi.response_helpers import generate_single_response
 from aiq.front_ends.fastapi.response_helpers import generate_streaming_response_as_str
 from aiq.front_ends.fastapi.response_helpers import generate_streaming_response_full_as_str
 from aiq.front_ends.fastapi.step_adaptor import StepAdaptor
-from aiq.front_ends.fastapi.websocket import AIQWebSocket
+from aiq.object_store.models import ObjectStoreItem
 from aiq.runtime.session import AIQSessionManager
 
 logger = logging.getLogger(__name__)
@@ -74,6 +83,7 @@ class FastApiFrontEndPluginWorkerBase(ABC):
 
         self._cleanup_tasks: list[str] = []
         self._cleanup_tasks_lock = asyncio.Lock()
+        self._http_flow_handler: HTTPAuthenticationFlowHandler | None = HTTPAuthenticationFlowHandler()
 
     @property
     def config(self) -> AIQConfig:
@@ -81,7 +91,6 @@ class FastApiFrontEndPluginWorkerBase(ABC):
 
     @property
     def front_end_config(self) -> FastApiFrontEndConfig:
-
         return self._front_end_config
 
     def build_app(self) -> FastAPI:
@@ -116,8 +125,13 @@ class FastApiFrontEndPluginWorkerBase(ABC):
 
         aiq_app = FastAPI(lifespan=lifespan)
 
+        # Configure app CORS.
         self.set_cors_config(aiq_app)
 
+        @aiq_app.middleware("http")
+        async def authentication_log_filter(request: Request, call_next: Callable[[Request], Awaitable[Response]]):
+            return await self._suppress_authentication_logs(request, call_next)
+
         return aiq_app
 
     def set_cors_config(self, aiq_app: FastAPI) -> None:
@@ -152,6 +166,26 @@ class FastApiFrontEndPluginWorkerBase(ABC):
             **cors_kwargs,
         )
 
+    async def _suppress_authentication_logs(self, request: Request,
+                                            call_next: Callable[[Request], Awaitable[Response]]) -> Response:
+        """
+        Intercepts authentication request and supreses logs that contain sensitive data.
+        """
+        from aiq.utils.log_utils import LogFilter
+
+        logs_to_suppress: list[str] = []
+
+        if (self.front_end_config.oauth2_callback_path):
+            logs_to_suppress.append(self.front_end_config.oauth2_callback_path)
+
+        logging.getLogger("uvicorn.access").addFilter(LogFilter(logs_to_suppress))
+        try:
+            response = await call_next(request)
+        finally:
+            logging.getLogger("uvicorn.access").removeFilter(LogFilter(logs_to_suppress))
+
+        return response
+
     @abstractmethod
     async def configure(self, app: FastAPI, builder: WorkflowBuilder):
         pass
@@ -168,6 +202,12 @@ class RouteInfo(BaseModel):
 
 class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
 
+    def __init__(self, config: AIQConfig):
+        super().__init__(config)
+
+        self._outstanding_flows: dict[str, FlowState] = {}
+        self._outstanding_flows_lock = asyncio.Lock()
+
     @staticmethod
     async def _periodic_cleanup(name: str, job_store: JobStore, sleep_time_sec: int = 300):
         while True:
@@ -209,6 +249,8 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
 
         await self.add_default_route(app, AIQSessionManager(builder.build()))
         await self.add_evaluate_route(app, AIQSessionManager(builder.build()))
+        await self.add_static_files_route(app, builder)
+        await self.add_authorization_route(app)
 
         for ep in self.front_end_config.endpoints:
 
@@ -381,6 +423,100 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
                 responses={500: response_500},
             )
 
+    async def add_static_files_route(self, app: FastAPI, builder: WorkflowBuilder):
+
+        if not self.front_end_config.object_store:
+            logger.debug("No object store configured, skipping static files route")
+            return
+
+        object_store_client = await builder.get_object_store_client(self.front_end_config.object_store)
+
+        def sanitize_path(path: str) -> str:
+            sanitized_path = os.path.normpath(path.strip("/"))
+            if sanitized_path == ".":
+                raise HTTPException(status_code=400, detail="Invalid file path.")
+            filename = os.path.basename(sanitized_path)
+            if not filename:
+                raise HTTPException(status_code=400, detail="Filename cannot be empty.")
+            return sanitized_path
+
+        # Upload static files to the object store; if key is present, it will fail with 409 Conflict
+        async def add_static_file(file_path: str, file: UploadFile):
+            sanitized_file_path = sanitize_path(file_path)
+            file_data = await file.read()
+
+            try:
+                await object_store_client.put_object(sanitized_file_path,
+                                                     ObjectStoreItem(data=file_data, content_type=file.content_type))
+            except KeyAlreadyExistsError as e:
+                raise HTTPException(status_code=409, detail=str(e)) from e
+
+            return {"filename": sanitized_file_path}
+
+        # Upsert static files to the object store; if key is present, it will overwrite the file
+        async def upsert_static_file(file_path: str, file: UploadFile):
+            sanitized_file_path = sanitize_path(file_path)
+            file_data = await file.read()
+
+            await object_store_client.upsert_object(sanitized_file_path,
+                                                    ObjectStoreItem(data=file_data, content_type=file.content_type))
+
+            return {"filename": sanitized_file_path}
+
+        # Get static files from the object store
+        async def get_static_file(file_path: str):
+
+            try:
+                file_data = await object_store_client.get_object(file_path)
+            except NoSuchKeyError as e:
+                raise HTTPException(status_code=404, detail=str(e)) from e
+
+            filename = file_path.split("/")[-1]
+
+            async def reader():
+                yield file_data.data
+
+            return StreamingResponse(reader(),
+                                     media_type=file_data.content_type,
+                                     headers={"Content-Disposition": f"attachment; filename={filename}"})
+
+        async def delete_static_file(file_path: str):
+            try:
+                await object_store_client.delete_object(file_path)
+            except NoSuchKeyError as e:
+                raise HTTPException(status_code=404, detail=str(e)) from e
+
+            return Response(status_code=204)
+
+        # Add the static files route to the FastAPI app
+        app.add_api_route(
+            path="/static/{file_path:path}",
+            endpoint=add_static_file,
+            methods=["POST"],
+            description="Upload a static file to the object store",
+        )
+
+        app.add_api_route(
+            path="/static/{file_path:path}",
+            endpoint=upsert_static_file,
+            methods=["PUT"],
+            description="Upsert a static file to the object store",
+        )
+
+        app.add_api_route(
+            path="/static/{file_path:path}",
+            endpoint=get_static_file,
+            methods=["GET"],
+            description="Get a static file from the object store",
+        )
+
+        app.add_api_route(
+            path="/static/{file_path:path}",
+            endpoint=delete_static_file,
+            methods=["DELETE"],
+            description="Delete a static file from the object store",
+        )
+
     async def add_route(self,
                         app: FastAPI,
                         endpoint: FastApiFrontEndConfig.EndpointBase,
@@ -388,10 +524,6 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
 
         workflow = session_manager.workflow
 
-        if (endpoint.websocket_path):
-            app.add_websocket_route(endpoint.websocket_path,
-                                    partial(AIQWebSocket, session_manager, self.get_step_adaptor()))
-
         GenerateBodyType = workflow.input_schema  # pylint: disable=invalid-name
         GenerateStreamResponseType = workflow.streaming_output_schema  # pylint: disable=invalid-name
         GenerateSingleResponseType = workflow.single_output_schema  # pylint: disable=invalid-name
@@ -442,7 +574,8 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
 
                 response.headers["Content-Type"] = "application/json"
 
-                async with session_manager.session(request=request):
+                async with session_manager.session(request=request,
+                                                   user_authentication_callback=self._http_flow_handler.authenticate):
 
                     return await generate_single_response(None, session_manager, result_type=result_type)
 
@@ -452,7 +585,8 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
 
             async def get_stream(request: Request):
 
-                async with session_manager.session(request=request):
+                async with session_manager.session(request=request,
+                                                   user_authentication_callback=self._http_flow_handler.authenticate):
 
                     return StreamingResponse(headers={"Content-Type": "text/event-stream; charset=utf-8"},
                                              content=generate_streaming_response_as_str(
@@ -486,7 +620,8 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
 
                 response.headers["Content-Type"] = "application/json"
 
-                async with session_manager.session(request=request):
+                async with session_manager.session(request=request,
+                                                   user_authentication_callback=self._http_flow_handler.authenticate):
 
                     return await generate_single_response(payload, session_manager, result_type=result_type)
 
@@ -499,7 +634,8 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
 
             async def post_stream(request: Request, payload: request_type):
 
-                async with session_manager.session(request=request):
+                async with session_manager.session(request=request,
+                                                   user_authentication_callback=self._http_flow_handler.authenticate):
 
                     return StreamingResponse(headers={"Content-Type": "text/event-stream; charset=utf-8"},
                                              content=generate_streaming_response_as_str(
@@ -533,6 +669,66 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
 
             return post_stream
 
+        def post_openai_api_compatible_endpoint(request_type: type):
+            """
+            OpenAI-compatible endpoint that handles both streaming and non-streaming
+            based on the 'stream' parameter in the request.
+            """
+
+            async def post_openai_api_compatible(response: Response, request: Request, payload: request_type):
+                # Check if streaming is requested
+                stream_requested = getattr(payload, 'stream', False)
+
+                async with session_manager.session(request=request):
+                    if stream_requested:
+                        # Return streaming response
+                        return StreamingResponse(headers={"Content-Type": "text/event-stream; charset=utf-8"},
+                                                 content=generate_streaming_response_as_str(
+                                                     payload,
+                                                     session_manager=session_manager,
+                                                     streaming=True,
+                                                     step_adaptor=self.get_step_adaptor(),
+                                                     result_type=AIQChatResponseChunk,
+                                                     output_type=AIQChatResponseChunk))
+                    else:
+                        # Return single response - check if workflow supports non-streaming
+                        try:
+                            response.headers["Content-Type"] = "application/json"
+                            return await generate_single_response(payload, session_manager, result_type=AIQChatResponse)
+                        except ValueError as e:
+                            if "Cannot get a single output value for streaming workflows" in str(e):
+                                # Workflow only supports streaming, but client requested non-streaming
+                                # Fall back to streaming and collect the result
+                                chunks = []
+                                async for chunk_str in generate_streaming_response_as_str(
+                                        payload,
+                                        session_manager=session_manager,
+                                        streaming=True,
+                                        step_adaptor=self.get_step_adaptor(),
+                                        result_type=AIQChatResponseChunk,
+                                        output_type=AIQChatResponseChunk):
+                                    if chunk_str.startswith("data: ") and not chunk_str.startswith("data: [DONE]"):
+                                        chunk_data = chunk_str[6:].strip()  # Remove "data: " prefix
+                                        if chunk_data:
+                                            try:
+                                                chunk_json = AIQChatResponseChunk.model_validate_json(chunk_data)
+                                                if (chunk_json.choices and len(chunk_json.choices) > 0
+                                                        and chunk_json.choices[0].delta
+                                                        and chunk_json.choices[0].delta.content is not None):
+                                                    chunks.append(chunk_json.choices[0].delta.content)
+                                            except Exception:
+                                                continue
+
+                                # Create a single response from collected chunks
+                                content = "".join(chunks)
+                                single_response = AIQChatResponse.from_string(content)
+                                response.headers["Content-Type"] = "application/json"
+                                return single_response
+                            else:
+                                raise
+
+            return post_openai_api_compatible
+
         async def run_generation(job_id: str,
                                  payload: typing.Any,
                                  session_manager: AIQSessionManager,
@@ -623,7 +819,56 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
                 logger.info("Found job %s with status %s", job_id, job.status)
                 return _job_status_to_response(job)
 
+        async def websocket_endpoint(websocket: WebSocket):
+
+            # Universal cookie handling: works for both cross-origin and same-origin connections
+            session_id = websocket.query_params.get("session")
+            if session_id:
+                headers = list(websocket.scope.get("headers", []))
+                cookie_header = f"aiqtoolkit-session={session_id}"
+
+                # Check if the session cookie already exists to avoid duplicates
+                cookie_exists = False
+                existing_session_cookie = False
+
+                for i, (name, value) in enumerate(headers):
+                    if name == b"cookie":
+                        cookie_exists = True
+                        cookie_str = value.decode()
+
+                        # Check if aiqtoolkit-session already exists in cookies
+                        if "aiqtoolkit-session=" in cookie_str:
+                            existing_session_cookie = True
+                            logger.info("WebSocket: Session cookie already present in headers (same-origin)")
+                        else:
+                            # Append to existing cookie header (cross-origin case)
+                            headers[i] = (name, f"{cookie_str}; {cookie_header}".encode())
+                            logger.info("WebSocket: Added session cookie to existing cookie header: %s",
+                                        session_id[:10] + "...")
+                        break
+
+                # Add new cookie header only if no cookies exist and no session cookie found
+                if not cookie_exists and not existing_session_cookie:
+                    headers.append((b"cookie", cookie_header.encode()))
+                    logger.info("WebSocket: Added new session cookie header: %s", session_id[:10] + "...")
+
+                # Update the websocket scope with the modified headers
+                websocket.scope["headers"] = headers
+
+            async with WebSocketMessageHandler(websocket, session_manager, self.get_step_adaptor()) as handler:
+
+                flow_handler = WebSocketAuthenticationFlowHandler(self._add_flow, self._remove_flow, handler)
+
+                # Ugly hack to set the flow handler on the message handler. Both need eachother to be set.
+                handler.set_flow_handler(flow_handler)
+
+                await handler.run()
+
+        if (endpoint.websocket_path):
+            app.add_websocket_route(endpoint.websocket_path, websocket_endpoint)
+
         if (endpoint.path):
+
             if (endpoint.method == "GET"):
 
                 app.add_api_route(
@@ -745,26 +990,103 @@ class FastApiFrontEndPluginWorker(FastApiFrontEndPluginWorkerBase):
 
             elif (endpoint.method == "POST"):
 
-                app.add_api_route(
-                    path=endpoint.openai_api_path,
-                    endpoint=post_single_endpoint(request_type=AIQChatRequest, result_type=AIQChatResponse),
-                    methods=[endpoint.method],
-                    response_model=AIQChatResponse,
-                    description=endpoint.description,
-                    responses={500: response_500},
-                )
-
-                app.add_api_route(
-                    path=f"{endpoint.openai_api_path}/stream",
-                    endpoint=post_streaming_endpoint(request_type=AIQChatRequest,
-                                                     streaming=True,
-                                                     result_type=AIQChatResponseChunk,
-                                                     output_type=AIQChatResponseChunk),
-                    methods=[endpoint.method],
-                    response_model=AIQChatResponseChunk | AIQResponseIntermediateStep,
-                    description=endpoint.description,
-                    responses={500: response_500},
-                )
+                # Check if OpenAI v1 compatible endpoint is configured
+                openai_v1_path = getattr(endpoint, 'openai_api_v1_path', None)
+
+                # Always create legacy endpoints for backward compatibility (unless they conflict with v1 path)
+                if not openai_v1_path or openai_v1_path != endpoint.openai_api_path:
+                    # <openai_api_path> = non-streaming (legacy behavior)
+                    app.add_api_route(
+                        path=endpoint.openai_api_path,
+                        endpoint=post_single_endpoint(request_type=AIQChatRequest, result_type=AIQChatResponse),
+                        methods=[endpoint.method],
+                        response_model=AIQChatResponse,
+                        description=endpoint.description,
+                        responses={500: response_500},
+                    )
+
+                    # <openai_api_path>/stream = streaming (legacy behavior)
+                    app.add_api_route(
+                        path=f"{endpoint.openai_api_path}/stream",
+                        endpoint=post_streaming_endpoint(request_type=AIQChatRequest,
+                                                         streaming=True,
+                                                         result_type=AIQChatResponseChunk,
+                                                         output_type=AIQChatResponseChunk),
+                        methods=[endpoint.method],
+                        response_model=AIQChatResponseChunk | AIQResponseIntermediateStep,
+                        description=endpoint.description,
+                        responses={500: response_500},
+                    )
+
+                # Create OpenAI v1 compatible endpoint if configured
+                if openai_v1_path:
+                    # OpenAI v1 Compatible Mode: Create single endpoint that handles both streaming and non-streaming
+                    app.add_api_route(
+                        path=openai_v1_path,
+                        endpoint=post_openai_api_compatible_endpoint(request_type=AIQChatRequest),
+                        methods=[endpoint.method],
+                        response_model=AIQChatResponse | AIQChatResponseChunk,
+                        description=f"{endpoint.description} (OpenAI Chat Completions API compatible)",
+                        responses={500: response_500},
+                    )
 
             else:
                 raise ValueError(f"Unsupported method {endpoint.method}")
+
+    async def add_authorization_route(self, app: FastAPI):
+
+        from fastapi.responses import HTMLResponse
+
+        from aiq.front_ends.fastapi.html_snippets.auth_code_grant_success import AUTH_REDIRECT_SUCCESS_HTML
+
+        async def redirect_uri(request: Request):
+            """
+            Handle the redirect URI for OAuth2 authentication.
+            Args:
+                request: The FastAPI request object containing query parameters.
+
+            Returns:
+                HTMLResponse: A response indicating the success of the authentication flow.
+            """
+            state = request.query_params.get("state")
+
+            async with self._outstanding_flows_lock:
+                if not state or state not in self._outstanding_flows:
+                    return "Invalid state. Please restart the authentication process."
+
+                flow_state = self._outstanding_flows[state]
+
+            config = flow_state.config
+            verifier = flow_state.verifier
+            client = flow_state.client
+
+            try:
+                res = await client.fetch_token(url=config.token_url,
+                                               authorization_response=str(request.url),
+                                               code_verifier=verifier,
+                                               state=state)
+                flow_state.future.set_result(res)
+            except Exception as e:
+                flow_state.future.set_exception(e)
+
+            return HTMLResponse(content=AUTH_REDIRECT_SUCCESS_HTML,
+                                status_code=200,
+                                headers={
+                                    "Content-Type": "text/html; charset=utf-8", "Cache-Control": "no-cache"
+                                })
+
+        if (self.front_end_config.oauth2_callback_path):
+            # Add the redirect URI route
+            app.add_api_route(
+                path=self.front_end_config.oauth2_callback_path,
+                endpoint=redirect_uri,
+                methods=["GET"],
+                description="Handles the authorization code and state returned from the Authorization Code Grant Flow.")
+
+    async def _add_flow(self, state: str, flow_state: FlowState):
+        async with self._outstanding_flows_lock:
+            self._outstanding_flows[state] = flow_state
+
+    async def _remove_flow(self, state: str):
+        async with self._outstanding_flows_lock:
+            del self._outstanding_flows[state]

aiq/front_ends/fastapi/html_snippets/__init__.py

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

aiq/front_ends/fastapi/html_snippets/auth_code_grant_success.py

@@ -0,0 +1,35 @@
+# 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.
+
+AUTH_REDIRECT_SUCCESS_HTML = """
+<!DOCTYPE html>
+<html>
+<head>
+    <title>Authentication Complete</title>
+    <script>
+        (function () {
+            window.history.replaceState(null, "", window.location.pathname);
+
+            window.opener?.postMessage({ type: 'AUTH_SUCCESS' }, '*');
+
+            window.close();
+        })();
+    </script>
+</head>
+<body>
+    <p>Authentication complete. You may now close this window.</p>
+</body>
+</html>
+"""

aiq/front_ends/fastapi/main.py

@@ -68,3 +68,5 @@ def get_app():
 
     except ImportError as e:
         raise ValueError(f"Front end worker {front_end_worker_full_name} not found.") from e
+    except Exception as e:
+        raise ValueError(f"Error loading front end worker {front_end_worker_full_name}: {e}") from e