aiqtoolkit 1.2.0rc4__py3-none-any.whl → 1.2rc9__py3-none-any.whl

aiq/data_models/authentication.py

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

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

@@ -1,54 +0,0 @@
-# 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"
-    ITS_STRATEGY = "its_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"
-    ITS_STRATEGIES = "its_strategies"
-    LLMS = "llms"
-    MEMORY = "memory"
-    OBJECT_STORES = "object_stores"
-    RETRIEVERS = "retrievers"

aiq/data_models/component_ref.py

@@ -1,168 +0,0 @@
-# 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 ITSStrategyRef(ComponentRef):
-    """
-    A reference to an ITS strategy in an AgentIQ configuration object.
-    """
-
-    @property
-    @override
-    def component_group(self):
-        return ComponentGroup.ITS_STRATEGIES