nvidia-nat 1.2.0rc5__py3-none-any.whl

aiq/data_models/authentication.py

@@ -0,0 +1,231 @@
+# 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 datetime import datetime
+from datetime import timezone
+from enum import Enum
+
+import httpx
+from pydantic import BaseModel
+from pydantic import ConfigDict
+from pydantic import Field
+from pydantic import SecretStr
+
+from aiq.data_models.common import BaseModelRegistryTag
+from aiq.data_models.common import TypedBaseModel
+
+
+class AuthProviderBaseConfig(TypedBaseModel, BaseModelRegistryTag):
+    """
+    Base configuration for authentication providers.
+    """
+
+    # Default, forbid extra fields to prevent unexpected behavior or miss typed options
+    model_config = ConfigDict(extra="forbid")
+
+
+AuthProviderBaseConfigT = typing.TypeVar("AuthProviderBaseConfigT", bound=AuthProviderBaseConfig)
+
+
+class CredentialLocation(str, Enum):
+    """
+    Enum representing the location of credentials in an HTTP request.
+    """
+    HEADER = "header"
+    QUERY = "query"
+    COOKIE = "cookie"
+    BODY = "body"
+
+
+class AuthFlowType(str, Enum):
+    """
+    Enum representing different types of authentication flows.
+    """
+    API_KEY = "api_key"
+    OAUTH2_CLIENT_CREDENTIALS = "oauth2_client_credentials"
+    OAUTH2_AUTHORIZATION_CODE = "oauth2_auth_code_flow"
+    OAUTH2_PASSWORD = "oauth2_password"
+    OAUTH2_DEVICE_CODE = "oauth2_device_code"
+    HTTP_BASIC = "http_basic"
+    NONE = "none"
+
+
+class AuthenticatedContext(BaseModel):
+    """
+    Represents an authenticated context for making requests.
+    """
+    model_config = ConfigDict(extra="forbid", arbitrary_types_allowed=True)
+    headers: dict[str, str] | httpx.Headers | None = Field(default=None,
+                                                           description="HTTP headers used for authentication.")
+    query_params: dict[str, str] | httpx.QueryParams | None = Field(
+        default=None, description="Query parameters used for authentication.")
+    cookies: dict[str, str] | httpx.Cookies | None = Field(default=None, description="Cookies used for authentication.")
+    body: dict[str, str] | None = Field(default=None, description="Authenticated Body value, if applicable.")
+    metadata: dict[str, typing.Any] | None = Field(default=None, description="Additional metadata for the request.")
+
+
+class HeaderAuthScheme(str, Enum):
+    """
+    Enum representing different header authentication schemes.
+    """
+    BEARER = "Bearer"
+    X_API_KEY = "X-API-Key"
+    BASIC = "Basic"
+    CUSTOM = "Custom"
+
+
+class HTTPMethod(str, Enum):
+    """
+    Enum representing HTTP methods used in requests.
+    """
+    GET = "GET"
+    POST = "POST"
+    PUT = "PUT"
+    DELETE = "DELETE"
+    PATCH = "PATCH"
+    HEAD = "HEAD"
+    OPTIONS = "OPTIONS"
+
+
+class CredentialKind(str, Enum):
+    """
+    Enum representing different kinds of credentials used for authentication.
+    """
+    HEADER = "header"
+    QUERY = "query"
+    COOKIE = "cookie"
+    BASIC = "basic_auth"
+    BEARER = "bearer_token"
+
+
+class _CredBase(BaseModel):
+    """
+    Base class for credentials used in authentication.
+    """
+    kind: CredentialKind
+    model_config = ConfigDict(extra="forbid")
+
+
+class HeaderCred(_CredBase):
+    """
+    Represents a credential that is sent in the HTTP header.
+    """
+    kind: typing.Literal[CredentialKind.HEADER] = CredentialKind.HEADER
+    name: str
+    value: SecretStr
+
+
+class QueryCred(_CredBase):
+    """
+    Represents a credential that is sent as a query parameter in the URL.
+    """
+    kind: typing.Literal[CredentialKind.QUERY] = CredentialKind.QUERY
+    name: str
+    value: SecretStr
+
+
+class CookieCred(_CredBase):
+    """
+    Represents a credential that is sent as a cookie in the HTTP request.
+    """
+    kind: typing.Literal[CredentialKind.COOKIE] = CredentialKind.COOKIE
+    name: str
+    value: SecretStr
+
+
+class BasicAuthCred(_CredBase):
+    """
+    Represents credentials for HTTP Basic Authentication.
+    """
+    kind: typing.Literal[CredentialKind.BASIC] = CredentialKind.BASIC
+    username: SecretStr
+    password: SecretStr
+
+
+class BearerTokenCred(_CredBase):
+    """
+    Represents a credential for Bearer Token Authentication.
+    """
+    kind: typing.Literal[CredentialKind.BEARER] = CredentialKind.BEARER
+    token: SecretStr
+    scheme: str = "Bearer"
+    header_name: str = "Authorization"
+
+
+Credential = typing.Annotated[
+    typing.Union[
+        HeaderCred,
+        QueryCred,
+        CookieCred,
+        BasicAuthCred,
+        BearerTokenCred,
+    ],
+    Field(discriminator="kind"),
+]
+
+
+class AuthResult(BaseModel):
+    """
+    Represents the result of an authentication process.
+    """
+    credentials: list[Credential] = Field(default_factory=list,
+                                          description="List of credentials used for authentication.")
+    token_expires_at: datetime | None = Field(default=None, description="Expiration time of the token, if applicable.")
+    raw: dict[str, typing.Any] = Field(default_factory=dict,
+                                       description="Raw response data from the authentication process.")
+
+    model_config = ConfigDict(extra="forbid")
+
+    def is_expired(self) -> bool:
+        """
+        Checks if the authentication token has expired.
+        """
+        return bool(self.token_expires_at and datetime.now(timezone.utc) >= self.token_expires_at)
+
+    def as_requests_kwargs(self) -> dict[str, typing.Any]:
+        """
+        Converts the authentication credentials into a format suitable for use with the `httpx` library.
+        """
+        kw: dict[str, typing.Any] = {"headers": {}, "params": {}, "cookies": {}}
+
+        for cred in self.credentials:
+            match cred:
+                case HeaderCred():
+                    kw["headers"][cred.name] = cred.value.get_secret_value()
+                case QueryCred():
+                    kw["params"][cred.name] = cred.value.get_secret_value()
+                case CookieCred():
+                    kw["cookies"][cred.name] = cred.value.get_secret_value()
+                case BearerTokenCred():
+                    kw["headers"][cred.header_name] = (f"{cred.scheme} {cred.token.get_secret_value()}")
+                case BasicAuthCred():
+                    kw["auth"] = (
+                        cred.username.get_secret_value(),
+                        cred.password.get_secret_value(),
+                    )
+
+        return kw
+
+    def attach(self, target_kwargs: dict[str, typing.Any]) -> None:
+        """
+        Attaches the authentication credentials to the target request kwargs.
+        """
+        merged = self.as_requests_kwargs()
+        for k, v in merged.items():
+            if isinstance(v, dict):
+                target_kwargs.setdefault(k, {}).update(v)
+            else:
+                target_kwargs[k] = v

aiq/data_models/common.py

@@ -0,0 +1,171 @@
+# 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 inspect
+import sys
+import typing
+from hashlib import sha512
+
+from pydantic import AliasChoices
+from pydantic import BaseModel
+from pydantic import Field
+from pydantic.json_schema import GenerateJsonSchema
+from pydantic.json_schema import JsonSchemaMode
+
+_LT = typing.TypeVar("_LT")
+
+
+class HashableBaseModel(BaseModel):
+    """
+    Subclass of a Pydantic BaseModel that is hashable. Use in objects that need to be hashed for caching purposes.
+    """
+
+    def __hash__(self):
+        return int.from_bytes(bytes=sha512(f"{self.__class__.__qualname__}::{self.model_dump_json()}".encode(
+            'utf-8', errors='ignore')).digest(),
+                              byteorder=sys.byteorder)
+
+    def __lt__(self, other):
+        return self.__hash__() < other.__hash__()
+
+    def __eq__(self, other):
+        return self.__hash__() == other.__hash__()
+
+    def __ne__(self, other):
+        return self.__hash__() != other.__hash__()
+
+    def __gt__(self, other):
+        return self.__hash__() > other.__hash__()
+
+    @classmethod
+    def generate_json_schema(cls) -> dict[str, typing.Any]:
+        return cls.model_json_schema()
+
+    @classmethod
+    def write_json_schema(cls, schema_path: str) -> None:
+
+        import json
+
+        schema = cls.generate_json_schema()
+
+        with open(schema_path, "w", encoding="utf-8") as f:
+            json.dump(schema, f, indent=2)
+
+
+def subclass_depth(cls: type) -> int:
+    """
+    Compute a class' subclass depth.
+    """
+    depth = 0
+    while (cls is not object and cls.__base__ is not None):
+        cls = cls.__base__  # type: ignore
+        depth += 1
+    return depth
+
+
+def _get_origin_or_base(cls: type) -> type:
+    """
+    Get the origin of a type or the base class if it is not a generic.
+    """
+    origin = typing.get_origin(cls)
+    if origin is None:
+        return cls
+    return origin
+
+
+class BaseModelRegistryTag:
+
+    pass
+
+
+class TypedBaseModel(BaseModel):
+    """
+    Subclass of Pydantic BaseModel that allows for specifying the object type. Use in Pydantic discriminated unions.
+    """
+
+    type: str = Field(default="unknown",
+                      init=False,
+                      serialization_alias="_type",
+                      validation_alias=AliasChoices('type', '_type'),
+                      description="The type of the object",
+                      title="Type",
+                      repr=False)
+
+    full_type: typing.ClassVar[str]
+    _typed_model_name: typing.ClassVar[str | None] = None
+
+    def __init_subclass__(cls, name: str | None = None):
+        super().__init_subclass__()
+
+        if (name is not None):
+            module = inspect.getmodule(cls)
+
+            assert module is not None, f"Module not found for class {cls} when registering {name}"
+            package_name: str | None = module.__package__
+
+            # If the package name is not set, then we use the module name. Must have some namespace which will be unique
+            if (not package_name):
+                package_name = module.__name__
+
+            full_name = f"{package_name}/{name}"
+
+            # Store the type name as a class attribute - no field manipulation needed!
+            cls._typed_model_name = name  # type: ignore
+            cls.full_type = full_name
+
+    def model_post_init(self, __context):
+        """Set the type field to the correct value after instance creation."""
+        if hasattr(self.__class__, '_typed_model_name') and self.__class__._typed_model_name is not None:
+            object.__setattr__(self, 'type', self.__class__._typed_model_name)
+        # If no type name is set, the field retains its default "unknown" value
+
+    @classmethod
+    def model_json_schema(cls,
+                          by_alias: bool = True,
+                          ref_template: str = '#/$defs/{model}',
+                          schema_generator: "type[GenerateJsonSchema]" = GenerateJsonSchema,
+                          mode: JsonSchemaMode = 'validation') -> dict:
+        """Override to provide correct default for type field in schema."""
+        schema = super().model_json_schema(by_alias=by_alias,
+                                           ref_template=ref_template,
+                                           schema_generator=schema_generator,
+                                           mode=mode)
+
+        # Fix the type field default to show the actual component type instead of "unknown"
+        if ('properties' in schema and 'type' in schema['properties'] and hasattr(cls, '_typed_model_name')
+                and cls._typed_model_name is not None):
+            schema['properties']['type']['default'] = cls._typed_model_name
+
+        return schema
+
+    @classmethod
+    def static_type(cls):
+        return getattr(cls, '_typed_model_name')
+
+    @classmethod
+    def static_full_type(cls):
+        return cls.full_type
+
+    @staticmethod
+    def discriminator(v: typing.Any) -> str | None:
+        # If its serialized, then we use the alias
+        if isinstance(v, dict):
+            return v.get("_type", v.get("type"))
+
+        # Otherwise we use the property
+        return getattr(v, "type")
+
+
+TypedBaseModelT = typing.TypeVar("TypedBaseModelT", bound=TypedBaseModel)

aiq/data_models/component.py

@@ -0,0 +1,54 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+from enum import StrEnum
+
+logger = logging.getLogger(__name__)
+
+
+class AIQComponentEnum(StrEnum):
+    # Keep sorted!!!
+    AUTHENTICATION_PROVIDER = "auth_provider"
+    EMBEDDER_CLIENT = "embedder_client"
+    EMBEDDER_PROVIDER = "embedder_provider"
+    EVALUATOR = "evaluator"
+    FRONT_END = "front_end"
+    FUNCTION = "function"
+    TTC_STRATEGY = "ttc_strategy"
+    LLM_CLIENT = "llm_client"
+    LLM_PROVIDER = "llm_provider"
+    LOGGING = "logging"
+    MEMORY = "memory"
+    OBJECT_STORE = "object_store"
+    PACKAGE = "package"
+    REGISTRY_HANDLER = "registry_handler"
+    RETRIEVER_CLIENT = "retriever_client"
+    RETRIEVER_PROVIDER = "retriever_provider"
+    TOOL_WRAPPER = "tool_wrapper"
+    TRACING = "tracing"
+    UNDEFINED = "undefined"
+
+
+class ComponentGroup(StrEnum):
+    # Keep sorted!!!
+    AUTHENTICATION = "authentication"
+    EMBEDDERS = "embedders"
+    FUNCTIONS = "functions"
+    TTC_STRATEGIES = "ttc_strategies"
+    LLMS = "llms"
+    MEMORY = "memory"
+    OBJECT_STORES = "object_stores"
+    RETRIEVERS = "retrievers"

aiq/data_models/component_ref.py

@@ -0,0 +1,168 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import typing
+from abc import ABC
+from abc import abstractmethod
+
+from pydantic_core import CoreSchema
+from pydantic_core import core_schema
+
+from aiq.data_models.common import HashableBaseModel
+from aiq.data_models.component import ComponentGroup
+from aiq.utils.type_utils import override
+
+
+def generate_instance_id(input_object: typing.Any) -> str:
+    """Generates a unique identifier for a python object derived from its python unique id.
+
+    Args:
+        input_object (typing.Any): The input object to receive a unique identifier.
+
+    Returns:
+        str: Unique identifier.
+    """
+
+    return str(id(input_object))
+
+
+class ComponentRefNode(HashableBaseModel):
+    """A node type for component runtime instances reference names in a networkx digraph.
+
+    Args:
+        ref_name (ComponentRef): The name of the component runtime instance.
+        component_group (ComponentGroup): The component group in an AIQ Toolkit configuration object.
+    """
+
+    ref_name: "ComponentRef"
+    component_group: ComponentGroup
+
+
+class ComponentRef(str, ABC):
+    """
+    Abstract class used for the interface to derive ComponentRef objects.
+    """
+
+    def __new__(cls, value: "ComponentRef | str"):
+        # Sublcassing str skips abstractmethod enforcement.
+        if len(cls.__abstractmethods__ - set(cls.__dict__)):
+            abstract_methods = ", ".join([f"'{method}'" for method in cls.__abstractmethods__])
+            raise TypeError(f"Can't instantiate abstract class {cls.__name__} "
+                            f"without an implementation for abstract method(s) {abstract_methods}")
+
+        return super().__new__(cls, value)
+
+    @property
+    @abstractmethod
+    def component_group(self) -> ComponentGroup:
+        """Provides the component group this ComponentRef object represents.
+
+        Returns:
+            ComponentGroup: A component group of the AIQ Toolkit configuration object
+        """
+
+        pass
+
+    @classmethod
+    def __get_pydantic_core_schema__(cls, source_type, handler, **kwargs) -> CoreSchema:
+        return core_schema.no_info_plain_validator_function(cls)
+
+
+class EmbedderRef(ComponentRef):
+    """
+    A reference to an embedder in an AIQ Toolkit configuration object.
+    """
+
+    @property
+    @override
+    def component_group(self):
+        return ComponentGroup.EMBEDDERS
+
+
+class FunctionRef(ComponentRef):
+    """
+    A reference to a function in an AIQ Toolkit configuration object.
+    """
+
+    @property
+    @override
+    def component_group(self):
+        return ComponentGroup.FUNCTIONS
+
+
+class LLMRef(ComponentRef):
+    """
+    A reference to an LLM in an AIQ Toolkit configuration object.
+    """
+
+    @property
+    @override
+    def component_group(self):
+        return ComponentGroup.LLMS
+
+
+class MemoryRef(ComponentRef):
+    """
+    A reference to a memory in an AIQ Toolkit configuration object.
+    """
+
+    @property
+    @override
+    def component_group(self):
+        return ComponentGroup.MEMORY
+
+
+class ObjectStoreRef(ComponentRef):
+    """
+    A reference to an object store in an AIQ toolkit configuration object.
+    """
+
+    @property
+    @override
+    def component_group(self):
+        return ComponentGroup.OBJECT_STORES
+
+
+class RetrieverRef(ComponentRef):
+    """
+    A reference to a retriever in an AIQ Toolkit configuration object.
+    """
+
+    @property
+    @override
+    def component_group(self):
+        return ComponentGroup.RETRIEVERS
+
+
+class AuthenticationRef(ComponentRef):
+    """
+    A reference to an API Authentication Provider in an AIQ Toolkit configuration object.
+    """
+
+    @property
+    @override
+    def component_group(self):
+        return ComponentGroup.AUTHENTICATION
+
+
+class TTCStrategyRef(ComponentRef):
+    """
+    A reference to an TTC strategy in an NeMo Agent Toolkit configuration object.
+    """
+
+    @property
+    @override
+    def component_group(self):
+        return ComponentGroup.TTC_STRATEGIES