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

aiq/authentication/exceptions/request_exceptions.py

@@ -0,0 +1,54 @@
+# 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.
+
+
+class BaseUrlValidationError(Exception):
+    """Raised when HTTP Base URL validation fails unexpectedly."""
+
+    def __init__(self, error_code: str, message: str, *args):
+        self.error_code = error_code
+        super().__init__(f"[{error_code}] {message}", *args)
+
+
+class HTTPMethodValidationError(Exception):
+    """Raised when HTTP Method validation fails unexpectedly."""
+
+    def __init__(self, error_code: str, message: str, *args):
+        self.error_code = error_code
+        super().__init__(f"[{error_code}] {message}", *args)
+
+
+class QueryParameterValidationError(Exception):
+    """Raised when HTTP Query Parameter validation fails unexpectedly."""
+
+    def __init__(self, error_code: str, message: str, *args):
+        self.error_code = error_code
+        super().__init__(f"[{error_code}] {message}", *args)
+
+
+class HTTPHeaderValidationError(Exception):
+    """Raised when HTTP Header validation fails unexpectedly."""
+
+    def __init__(self, error_code: str, message: str, *args):
+        self.error_code = error_code
+        super().__init__(f"[{error_code}] {message}", *args)
+
+
+class BodyValidationError(Exception):
+    """Raised when HTTP Body validation fails unexpectedly."""
+
+    def __init__(self, error_code: str, message: str, *args):
+        self.error_code = error_code
+        super().__init__(f"[{error_code}] {message}", *args)

aiq/authentication/http_basic_auth/http_basic_auth_provider.py

@@ -0,0 +1,81 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from pydantic import SecretStr
+
+from aiq.authentication.interfaces import AuthProviderBase
+from aiq.builder.context import AIQContext
+from aiq.data_models.authentication import AuthenticatedContext
+from aiq.data_models.authentication import AuthFlowType
+from aiq.data_models.authentication import AuthProviderBaseConfig
+from aiq.data_models.authentication import AuthResult
+from aiq.data_models.authentication import BasicAuthCred
+from aiq.data_models.authentication import BearerTokenCred
+
+
+class HTTPBasicAuthProvider(AuthProviderBase):
+    """
+    Abstract base class for HTTP Basic Authentication exchangers.
+    """
+
+    def __init__(self, config: AuthProviderBaseConfig):
+        """
+        Initialize the HTTP Basic Auth Exchanger with the given configuration.
+        """
+        super().__init__(config)
+
+        self._authenticated_tokens: dict[str, AuthResult] = {}
+
+    async def authenticate(self, user_id: str | None = None) -> AuthResult:
+        """
+        Performs simple HTTP Authentication using the provided user ID.
+        """
+
+        context = AIQContext.get()
+
+        if user_id is None and hasattr(context, "metadata") and hasattr(
+                context.metadata, "cookies") and context.metadata.cookies is not None:
+            session_id = context.metadata.cookies.get("aiqtoolkit-session", None)
+            if not session_id:
+                raise RuntimeError("Authentication failed. No session ID found. Cannot identify user.")
+
+            user_id = session_id
+
+        if user_id and user_id in self._authenticated_tokens:
+            return self._authenticated_tokens[user_id]
+
+        auth_callback = context.user_auth_callback
+
+        try:
+            auth_context: AuthenticatedContext = await auth_callback(self.config, AuthFlowType.HTTP_BASIC)
+        except RuntimeError as e:
+            raise RuntimeError(f"Authentication callback failed: {str(e)}. Did you forget to set a "
+                               f"callback handler for your frontend?") from e
+
+        basic_auth_credentials = BasicAuthCred(username=SecretStr(auth_context.metadata.get("username", "")),
+                                               password=SecretStr(auth_context.metadata.get("password", "")))
+
+        # Get the auth token from the headers of auth context
+        bearer_token = auth_context.headers.get("Authorization", "").split(" ")[-1]
+        if not bearer_token:
+            raise RuntimeError("Authentication failed: No Authorization header found in the response.")
+
+        bearer_token_cred = BearerTokenCred(token=SecretStr(bearer_token), scheme="Basic")
+
+        auth_result = AuthResult(credentials=[basic_auth_credentials, bearer_token_cred])
+
+        self._authenticated_tokens[user_id] = auth_result
+
+        return auth_result

aiq/authentication/http_basic_auth/register.py

@@ -0,0 +1,30 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from aiq.builder.builder import Builder
+from aiq.cli.register_workflow import register_auth_provider
+from aiq.data_models.authentication import AuthProviderBaseConfig
+
+
+class HTTPBasicAuthProviderConfig(AuthProviderBaseConfig, name="http_basic"):
+    pass
+
+
+@register_auth_provider(config_type=HTTPBasicAuthProviderConfig)
+async def http_basic_auth_provider(config: HTTPBasicAuthProviderConfig, builder: Builder):
+
+    from aiq.authentication.http_basic_auth.http_basic_auth_provider import HTTPBasicAuthProvider
+
+    yield HTTPBasicAuthProvider(config)

aiq/authentication/interfaces.py

@@ -0,0 +1,93 @@
+# 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 typing
+from abc import ABC
+from abc import abstractmethod
+
+from aiq.data_models.authentication import AuthenticatedContext
+from aiq.data_models.authentication import AuthFlowType
+from aiq.data_models.authentication import AuthProviderBaseConfig
+from aiq.data_models.authentication import AuthProviderBaseConfigT
+from aiq.data_models.authentication import AuthResult
+
+AUTHORIZATION_HEADER = "Authorization"
+
+
+class AuthProviderBase(typing.Generic[AuthProviderBaseConfigT], ABC):
+    """
+    Base class for authenticating to API services.
+    This class provides an interface for authenticating to API services.
+    """
+
+    def __init__(self, config: AuthProviderBaseConfigT):
+        """
+        Initialize the AuthProviderBase with the given configuration.
+
+        Args:
+            config (AuthProviderBaseConfig): Configuration items for authentication.
+        """
+        self._config = config
+
+    @property
+    def config(self) -> AuthProviderBaseConfigT:
+        """
+        Returns the auth provider configuration object.
+
+        Returns
+        -------
+        AuthProviderBaseConfigT
+            The auth provider configuration object.
+        """
+        return self._config
+
+    @abstractmethod
+    async def authenticate(self, user_id: str | None = None) -> AuthResult:
+        """
+        Perform the authentication process for the client.
+
+        This method handles the necessary steps to authenticate the client with the
+        target API service, which may include obtaining tokens, refreshing credentials,
+        or completing multi-step authentication flows.
+
+        Raises:
+            NotImplementedError: Must be implemented by subclasses.
+        """
+        # This method will call the frontend FlowHandlerBase `authenticate` method
+        pass
+
+
+class FlowHandlerBase(ABC):
+    """
+    Handles front-end specifc flows for authentication clients.
+
+    Each front end will define a FlowHandler that will implement the authenticate method.
+
+    The `authenticate` method will be stored as the callback in the AIQContextState.user_auth_callback
+    """
+
+    @abstractmethod
+    async def authenticate(self, config: AuthProviderBaseConfig, method: AuthFlowType) -> AuthenticatedContext:
+        """
+        Perform the authentication process for the client.
+
+        This method handles the necessary steps to authenticate the client with the
+        target API service, which may include obtaining tokens, refreshing credentials,
+        or completing multistep authentication flows.
+
+        Raises:
+            NotImplementedError: Must be implemented by subclasses.
+        """
+        pass

aiq/authentication/oauth2/__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.

aiq/authentication/oauth2/oauth2_auth_code_flow_provider.py

@@ -0,0 +1,107 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from datetime import datetime
+from datetime import timezone
+
+from authlib.integrations.httpx_client import OAuth2Client as AuthlibOAuth2Client
+from pydantic import SecretStr
+
+from aiq.authentication.interfaces import AuthProviderBase
+from aiq.authentication.oauth2.oauth2_auth_code_flow_provider_config import OAuth2AuthCodeFlowProviderConfig
+from aiq.builder.context import AIQContext
+from aiq.data_models.authentication import AuthFlowType
+from aiq.data_models.authentication import AuthResult
+from aiq.data_models.authentication import BearerTokenCred
+
+
+class OAuth2AuthCodeFlowProvider(AuthProviderBase[OAuth2AuthCodeFlowProviderConfig]):
+
+    def __init__(self, config: OAuth2AuthCodeFlowProviderConfig):
+        super().__init__(config)
+        self._authenticated_tokens: dict[str, AuthResult] = {}
+        self._context = AIQContext.get()
+
+    async def _attempt_token_refresh(self, user_id: str, auth_result: AuthResult) -> AuthResult | None:
+        refresh_token = auth_result.raw.get("refresh_token")
+        if not isinstance(refresh_token, str):
+            return None
+
+        with AuthlibOAuth2Client(
+                client_id=self.config.client_id,
+                client_secret=self.config.client_secret,
+        ) as client:
+            try:
+                new_token_data = client.refresh_token(self.config.token_url, refresh_token=refresh_token)
+            except Exception:
+                # On any failure, we'll fall back to the full auth flow.
+                return None
+
+        expires_at_ts = new_token_data.get("expires_at")
+        new_expires_at = datetime.fromtimestamp(expires_at_ts, tz=timezone.utc) if expires_at_ts else None
+
+        new_auth_result = AuthResult(
+            credentials=[BearerTokenCred(token=SecretStr(new_token_data["access_token"]))],
+            token_expires_at=new_expires_at,
+            raw=new_token_data,
+        )
+
+        self._authenticated_tokens[user_id] = new_auth_result
+
+        return new_auth_result
+
+    async def authenticate(self, user_id: str | None = None) -> AuthResult:
+        if user_id is None and hasattr(AIQContext.get(), "metadata") and hasattr(
+                AIQContext.get().metadata, "cookies") and AIQContext.get().metadata.cookies is not None:
+            session_id = AIQContext.get().metadata.cookies.get("aiqtoolkit-session", None)
+            if not session_id:
+                raise RuntimeError("Authentication failed. No session ID found. Cannot identify user.")
+
+            user_id = session_id
+
+        if user_id and user_id in self._authenticated_tokens:
+            auth_result = self._authenticated_tokens[user_id]
+            if not auth_result.is_expired():
+                return auth_result
+
+            refreshed_auth_result = await self._attempt_token_refresh(user_id, auth_result)
+            if refreshed_auth_result:
+                return refreshed_auth_result
+
+        auth_callback = self._context.user_auth_callback
+        if not auth_callback:
+            raise RuntimeError("Authentication callback not set on AIQContext.")
+
+        try:
+            authenticated_context = await auth_callback(self.config, AuthFlowType.OAUTH2_AUTHORIZATION_CODE)
+        except Exception as e:
+            raise RuntimeError(f"Authentication callback failed: {e}") from e
+
+        auth_header = authenticated_context.headers.get("Authorization", "")
+        if not auth_header.startswith("Bearer "):
+            raise RuntimeError("Invalid Authorization header")
+
+        token = auth_header.split(" ")[1]
+
+        auth_result = AuthResult(
+            credentials=[BearerTokenCred(token=SecretStr(token))],
+            token_expires_at=authenticated_context.metadata.get("expires_at"),
+            raw=authenticated_context.metadata.get("raw_token"),
+        )
+
+        if user_id:
+            self._authenticated_tokens[user_id] = auth_result
+
+        return auth_result

aiq/authentication/oauth2/oauth2_auth_code_flow_provider_config.py

@@ -0,0 +1,39 @@
+# 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 pydantic import Field
+
+from aiq.data_models.authentication import AuthProviderBaseConfig
+
+
+class OAuth2AuthCodeFlowProviderConfig(AuthProviderBaseConfig, name="oauth2_auth_code_flow"):
+
+    client_id: str = Field(description="The client ID for OAuth 2.0 authentication.")
+    client_secret: str = Field(description="The secret associated with the client_id.")
+    authorization_url: str = Field(description="The authorization URL for OAuth 2.0 authentication.")
+    token_url: str = Field(description="The token URL for OAuth 2.0 authentication.")
+    token_endpoint_auth_method: str | None = Field(
+        description=("The authentication method for the token endpoint. "
+                     "Usually one of `client_secret_post` or `client_secret_basic`."),
+        default=None)
+    redirect_uri: str = Field(description="The redirect URI for OAuth 2.0 authentication. Must match the registered "
+                              "redirect URI with the OAuth provider.")
+    scopes: list[str] = Field(description="The scopes for OAuth 2.0 authentication.", default_factory=list)
+    use_pkce: bool = Field(default=False,
+                           description="Whether to use PKCE (Proof Key for Code Exchange) in the OAuth 2.0 flow.")
+
+    authorization_kwargs: dict[str, str] | None = Field(description=("Additional keyword arguments for the "
+                                                                     "authorization request."),
+                                                        default=None)

aiq/authentication/oauth2/register.py

@@ -0,0 +1,25 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from aiq.authentication.oauth2.oauth2_auth_code_flow_provider_config import OAuth2AuthCodeFlowProviderConfig
+from aiq.builder.builder import Builder
+from aiq.cli.register_workflow import register_auth_provider
+
+
+@register_auth_provider(config_type=OAuth2AuthCodeFlowProviderConfig)
+async def oauth2_client(authentication_provider: OAuth2AuthCodeFlowProviderConfig, builder: Builder):
+    from aiq.authentication.oauth2.oauth2_auth_code_flow_provider import OAuth2AuthCodeFlowProvider
+
+    yield OAuth2AuthCodeFlowProvider(authentication_provider)

aiq/authentication/register.py

@@ -0,0 +1,21 @@
+# 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.
+
+# pylint: disable=unused-import
+# flake8: noqa
+
+from aiq.authentication.api_key import register as register_api_key
+from aiq.authentication.http_basic_auth import register as register_http_basic_auth
+from aiq.authentication.oauth2 import register as register_oauth2

aiq/builder/builder.py

@@ -20,22 +20,32 @@ from abc import abstractmethod
 from collections.abc import Sequence
 from pathlib import Path
 
+from aiq.authentication.interfaces import AuthProviderBase
 from aiq.builder.context import AIQContext
 from aiq.builder.framework_enum import LLMFrameworkEnum
 from aiq.builder.function import Function
+from aiq.data_models.authentication import AuthProviderBaseConfig
+from aiq.data_models.component_ref import AuthenticationRef
 from aiq.data_models.component_ref import EmbedderRef
 from aiq.data_models.component_ref import FunctionRef
+from aiq.data_models.component_ref import ITSStrategyRef
 from aiq.data_models.component_ref import LLMRef
 from aiq.data_models.component_ref import MemoryRef
+from aiq.data_models.component_ref import ObjectStoreRef
 from aiq.data_models.component_ref import RetrieverRef
 from aiq.data_models.embedder import EmbedderBaseConfig
 from aiq.data_models.evaluator import EvaluatorBaseConfig
 from aiq.data_models.function import FunctionBaseConfig
 from aiq.data_models.function_dependencies import FunctionDependencies
+from aiq.data_models.its_strategy import ITSStrategyBaseConfig
 from aiq.data_models.llm import LLMBaseConfig
 from aiq.data_models.memory import MemoryBaseConfig
+from aiq.data_models.object_store import ObjectStoreBaseConfig
 from aiq.data_models.retriever import RetrieverBaseConfig
+from aiq.experimental.inference_time_scaling.models.stage_enums import PipelineTypeEnum
+from aiq.experimental.inference_time_scaling.models.stage_enums import StageTypeEnum
 from aiq.memory.interfaces import MemoryEditor
+from aiq.object_store.interfaces import ObjectStore
 from aiq.retriever.interface import AIQRetriever
 
 
@@ -91,6 +101,10 @@ class Builder(ABC):  # pylint: disable=too-many-public-methods
     async def add_llm(self, name: str | LLMRef, config: LLMBaseConfig):
         pass
 
+    @abstractmethod
+    async def get_llm(self, llm_name: str | LLMRef, wrapper_type: LLMFrameworkEnum | str) -> typing.Any:
+        pass
+
     async def get_llms(self, llm_names: Sequence[str | LLMRef],
                        wrapper_type: LLMFrameworkEnum | str) -> list[typing.Any]:
 
@@ -101,11 +115,41 @@ class Builder(ABC):  # pylint: disable=too-many-public-methods
         return list(llms)
 
     @abstractmethod
-    async def get_llm(self, llm_name: str | LLMRef, wrapper_type: LLMFrameworkEnum | str) -> typing.Any:
+    def get_llm_config(self, llm_name: str | LLMRef) -> LLMBaseConfig:
         pass
 
     @abstractmethod
-    def get_llm_config(self, llm_name: str | LLMRef) -> LLMBaseConfig:
+    async def add_auth_provider(self, name: str | AuthenticationRef, config: AuthProviderBaseConfig):
+        pass
+
+    @abstractmethod
+    async def get_auth_provider(self, auth_provider_name: str | AuthenticationRef) -> AuthProviderBase:
+        pass
+
+    async def get_auth_providers(self, auth_provider_names: list[str | AuthenticationRef]):
+
+        coros = [self.get_auth_provider(auth_provider_name=n) for n in auth_provider_names]
+
+        auth_providers = await asyncio.gather(*coros, return_exceptions=False)
+
+        return list(auth_providers)
+
+    @abstractmethod
+    async def add_object_store(self, name: str | ObjectStoreRef, config: ObjectStoreBaseConfig):
+        pass
+
+    async def get_object_store_clients(self, object_store_names: Sequence[str | ObjectStoreRef]) -> list[ObjectStore]:
+        """
+        Return a list of all object store clients.
+        """
+        return list(await asyncio.gather(*[self.get_object_store_client(name) for name in object_store_names]))
+
+    @abstractmethod
+    async def get_object_store_client(self, object_store_name: str | ObjectStoreRef) -> ObjectStore:
+        pass
+
+    @abstractmethod
+    def get_object_store_config(self, object_store_name: str | ObjectStoreRef) -> ObjectStoreBaseConfig:
         pass
 
     @abstractmethod
@@ -187,6 +231,24 @@ class Builder(ABC):  # pylint: disable=too-many-public-methods
     async def get_retriever_config(self, retriever_name: str | RetrieverRef) -> RetrieverBaseConfig:
         pass
 
+    @abstractmethod
+    async def add_its_strategy(self, name: str | str, config: ITSStrategyBaseConfig):
+        pass
+
+    @abstractmethod
+    async def get_its_strategy(self,
+                               strategy_name: str | ITSStrategyRef,
+                               pipeline_type: PipelineTypeEnum,
+                               stage_type: StageTypeEnum):
+        pass
+
+    @abstractmethod
+    async def get_its_strategy_config(self,
+                                      strategy_name: str | ITSStrategyRef,
+                                      pipeline_type: PipelineTypeEnum,
+                                      stage_type: StageTypeEnum) -> ITSStrategyBaseConfig:
+        pass
+
     @abstractmethod
     def get_user_manager(self) -> UserManagerHolder:
         pass

aiq/builder/component_utils.py

@@ -21,6 +21,7 @@ from collections.abc import Iterable
 import networkx as nx
 from pydantic import BaseModel
 
+from aiq.data_models.authentication import AuthProviderBaseConfig
 from aiq.data_models.common import TypedBaseModel
 from aiq.data_models.component import ComponentGroup
 from aiq.data_models.component_ref import ComponentRef
@@ -29,8 +30,10 @@ from aiq.data_models.component_ref import generate_instance_id
 from aiq.data_models.config import AIQConfig
 from aiq.data_models.embedder import EmbedderBaseConfig
 from aiq.data_models.function import FunctionBaseConfig
+from aiq.data_models.its_strategy import ITSStrategyBaseConfig
 from aiq.data_models.llm import LLMBaseConfig
 from aiq.data_models.memory import MemoryBaseConfig
+from aiq.data_models.object_store import ObjectStoreBaseConfig
 from aiq.data_models.retriever import RetrieverBaseConfig
 from aiq.utils.type_utils import DecomposedType
 
@@ -38,11 +41,14 @@ logger = logging.getLogger(__name__)
 
 # Order in which we want to process the component groups
 _component_group_order = [
+    ComponentGroup.AUTHENTICATION,
     ComponentGroup.EMBEDDERS,
     ComponentGroup.LLMS,
     ComponentGroup.MEMORY,
+    ComponentGroup.OBJECT_STORES,
     ComponentGroup.RETRIEVERS,
-    ComponentGroup.FUNCTIONS
+    ComponentGroup.ITS_STRATEGIES,
+    ComponentGroup.FUNCTIONS,
 ]
 
 
@@ -95,6 +101,8 @@ def group_from_component(component: TypedBaseModel) -> ComponentGroup | None:
             component is not a valid runtime instance, None is returned.
     """
 
+    if (isinstance(component, AuthProviderBaseConfig)):
+        return ComponentGroup.AUTHENTICATION
     if (isinstance(component, EmbedderBaseConfig)):
         return ComponentGroup.EMBEDDERS
     if (isinstance(component, FunctionBaseConfig)):
@@ -103,8 +111,12 @@ def group_from_component(component: TypedBaseModel) -> ComponentGroup | None:
         return ComponentGroup.LLMS
     if (isinstance(component, MemoryBaseConfig)):
         return ComponentGroup.MEMORY
+    if (isinstance(component, ObjectStoreBaseConfig)):
+        return ComponentGroup.OBJECT_STORES
     if (isinstance(component, RetrieverBaseConfig)):
         return ComponentGroup.RETRIEVERS
+    if (isinstance(component, ITSStrategyBaseConfig)):
+        return ComponentGroup.ITS_STRATEGIES
 
     return None
 
@@ -142,7 +154,7 @@ def recursive_componentref_discovery(cls: TypedBaseModel, value: typing.Any,
             yield from recursive_componentref_discovery(cls, field_data, field_info.annotation)
     if (decomposed_type.is_union):
         for arg in decomposed_type.args:
-            if (isinstance(value, DecomposedType(arg).root)):
+            if arg is typing.Any or (isinstance(value, DecomposedType(arg).root)):
                 yield from recursive_componentref_discovery(cls, value, arg)
     else:
         for arg in decomposed_type.args:
@@ -243,7 +255,8 @@ def build_dependency_sequence(config: "AIQConfig") -> list[ComponentInstanceData
     """
 
     total_node_count = len(config.embedders) + len(config.functions) + len(config.llms) + len(config.memory) + len(
-        config.retrievers) + 1  # +1 for the workflow
+        config.object_stores) + len(config.retrievers) + len(config.its_strategies) + len(
+            config.authentication) + 1  # +1 for the workflow
 
     dependency_map: dict
     dependency_graph: nx.DiGraph