aiqtoolkit 1.1.0rc6__py3-none-any.whl → 1.2.0__py3-none-any.whl

aiq/data_models/common.py

@@ -1,143 +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
-
-_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):
-        cls = cls.__base__
-        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(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]
-
-    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}"
-
-            type_field = cls.model_fields.get("type")
-            if type_field is not None:
-                type_field.default = name
-            cls.full_type = full_name
-
-    @classmethod
-    def static_type(cls):
-        return cls.model_fields.get("type").default
-
-    @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,46 +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):
-    FRONT_END = "front_end"
-    FUNCTION = "function"
-    TOOL_WRAPPER = "tool_wrapper"
-    LLM_PROVIDER = "llm_provider"
-    LLM_CLIENT = "llm_client"
-    EMBEDDER_PROVIDER = "embedder_provider"
-    EMBEDDER_CLIENT = "embedder_client"
-    EVALUATOR = "evaluator"
-    MEMORY = "memory"
-    RETRIEVER_PROVIDER = "retriever_provider"
-    RETRIEVER_CLIENT = "retriever_client"
-    REGISTRY_HANDLER = "registry_handler"
-    LOGGING = "logging"
-    TRACING = "tracing"
-    PACKAGE = "package"
-    UNDEFINED = "undefined"
-
-
-class ComponentGroup(StrEnum):
-    EMBEDDERS = "embedders"
-    FUNCTIONS = "functions"
-    LLMS = "llms"
-    MEMORY = "memory"
-    RETRIEVERS = "retrievers"

aiq/data_models/component_ref.py

@@ -1,135 +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 RetrieverRef(ComponentRef):
-    """
-    A reference to a retriever in an AIQ Toolkit configuration object.
-    """
-
-    @property
-    @override
-    def component_group(self):
-        return ComponentGroup.RETRIEVERS

aiq/data_models/config.py

@@ -1,349 +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 logging
-import sys
-import typing
-
-from pydantic import BaseModel
-from pydantic import ConfigDict
-from pydantic import Discriminator
-from pydantic import ValidationError
-from pydantic import ValidationInfo
-from pydantic import ValidatorFunctionWrapHandler
-from pydantic import field_validator
-
-from aiq.data_models.evaluate import EvalConfig
-from aiq.data_models.front_end import FrontEndBaseConfig
-from aiq.data_models.function import EmptyFunctionConfig
-from aiq.data_models.function import FunctionBaseConfig
-from aiq.data_models.logging import LoggingBaseConfig
-from aiq.data_models.telemetry_exporter import TelemetryExporterBaseConfig
-from aiq.front_ends.fastapi.fastapi_front_end_config import FastApiFrontEndConfig
-
-from .common import HashableBaseModel
-from .common import TypedBaseModel
-from .embedder import EmbedderBaseConfig
-from .llm import LLMBaseConfig
-from .memory import MemoryBaseConfig
-from .retriever import RetrieverBaseConfig
-
-logger = logging.getLogger(__name__)
-
-
-def _process_validation_error(err: ValidationError, handler: ValidatorFunctionWrapHandler, info: ValidationInfo):
-    from aiq.cli.type_registry import GlobalTypeRegistry  # pylint: disable=cyclic-import
-
-    new_errors = []
-    logged_once = False
-    needs_reraise = False
-    for e in err.errors():
-
-        error_type = e['type']
-        if error_type == 'union_tag_invalid' and "ctx" in e and not logged_once:
-            requested_type = e["ctx"]["tag"]
-
-            if (info.field_name in ('workflow', 'functions')):
-                registered_keys = GlobalTypeRegistry.get().get_registered_functions()
-            elif (info.field_name == "llms"):
-                registered_keys = GlobalTypeRegistry.get().get_registered_llm_providers()
-            elif (info.field_name == "embedders"):
-                registered_keys = GlobalTypeRegistry.get().get_registered_embedder_providers()
-            elif (info.field_name == "memory"):
-                registered_keys = GlobalTypeRegistry.get().get_registered_memorys()
-            elif (info.field_name == "retrievers"):
-                registered_keys = GlobalTypeRegistry.get().get_registered_retriever_providers()
-            elif (info.field_name == "tracing"):
-                registered_keys = GlobalTypeRegistry.get().get_registered_telemetry_exporters()
-            elif (info.field_name == "logging"):
-                registered_keys = GlobalTypeRegistry.get().get_registered_logging_method()
-            elif (info.field_name == "evaluators"):
-                registered_keys = GlobalTypeRegistry.get().get_registered_evaluators()
-            elif (info.field_name == "front_ends"):
-                registered_keys = GlobalTypeRegistry.get().get_registered_front_ends()
-
-            else:
-                assert False, f"Unknown field name {info.field_name} in validator"
-
-            # Check and see if the there are multiple full types which match this short type
-            matching_keys = [k for k in registered_keys if k.local_name == requested_type]
-
-            assert len(matching_keys) != 1, "Exact match should have been found. Contact developers"
-
-            matching_key_names = [x.full_type for x in matching_keys]
-            registered_key_names = [x.full_type for x in registered_keys]
-
-            if (len(matching_keys) == 0):
-                # This is a case where the requested type is not found. Show a helpful message about what is
-                # available
-                logger.error(("Requested %s type `%s` not found. "
-                              "Have you ensured the necessary package has been installed with `uv pip install`?"
-                              "\nAvailable %s names:\n - %s\n"),
-                             info.field_name,
-                             requested_type,
-                             info.field_name,
-                             '\n - '.join(registered_key_names))
-            else:
-                # This is a case where the requested type is ambiguous.
-                logger.error(("Requested %s type `%s` is ambiguous. "
-                              "Matched multiple %s by their local name: %s. "
-                              "Please use the fully qualified %s name."
-                              "\nAvailable %s names:\n - %s\n"),
-                             info.field_name,
-                             requested_type,
-                             info.field_name,
-                             matching_key_names,
-                             info.field_name,
-                             info.field_name,
-                             '\n - '.join(registered_key_names))
-
-            # Only show one error
-            logged_once = True
-
-        elif error_type == 'missing':
-            location = e["loc"]
-            if len(location) > 1:  # remove the _type field from the location
-                e['loc'] = (location[0], ) + location[2:]
-                needs_reraise = True
-
-        new_errors.append(e)
-
-    if needs_reraise:
-        raise ValidationError.from_exception_data(title=err.title, line_errors=new_errors)
-
-
-class TelemetryConfig(BaseModel):
-
-    logging: dict[str, LoggingBaseConfig] = {}
-    tracing: dict[str, TelemetryExporterBaseConfig] = {}
-
-    @field_validator("logging", "tracing", mode="wrap")
-    @classmethod
-    def validate_components(cls, value: typing.Any, handler: ValidatorFunctionWrapHandler, info: ValidationInfo):
-
-        try:
-            return handler(value)
-        except ValidationError as err:
-            _process_validation_error(err, handler, info)
-            raise
-
-    @classmethod
-    def rebuild_annotations(cls):
-
-        from aiq.cli.type_registry import GlobalTypeRegistry
-
-        type_registry = GlobalTypeRegistry.get()
-
-        TracingAnnotation = dict[str,
-                                 typing.Annotated[type_registry.compute_annotation(TelemetryExporterBaseConfig),
-                                                  Discriminator(TypedBaseModel.discriminator)]]
-
-        LoggingAnnotation = dict[str,
-                                 typing.Annotated[type_registry.compute_annotation(LoggingBaseConfig),
-                                                  Discriminator(TypedBaseModel.discriminator)]]
-
-        should_rebuild = False
-
-        tracing_field = cls.model_fields.get("tracing")
-        if tracing_field is not None and tracing_field.annotation != TracingAnnotation:
-            tracing_field.annotation = TracingAnnotation
-            should_rebuild = True
-
-        logging_field = cls.model_fields.get("logging")
-        if logging_field is not None and logging_field.annotation != LoggingAnnotation:
-            logging_field.annotation = LoggingAnnotation
-            should_rebuild = True
-
-        if (should_rebuild):
-            return cls.model_rebuild(force=True)
-
-        return False
-
-
-class GeneralConfig(BaseModel):
-
-    model_config = ConfigDict(protected_namespaces=())
-
-    use_uvloop: bool = True
-    """
-    Whether to use uvloop for the event loop. This can provide a significant speedup in some cases. Disable to provide
-    better error messages when debugging.
-    """
-
-    telemetry: TelemetryConfig = TelemetryConfig()
-
-    # FrontEnd Configuration
-    front_end: FrontEndBaseConfig = FastApiFrontEndConfig()
-
-    @field_validator("front_end", mode="wrap")
-    @classmethod
-    def validate_components(cls, value: typing.Any, handler: ValidatorFunctionWrapHandler, info: ValidationInfo):
-
-        try:
-            return handler(value)
-        except ValidationError as err:
-            _process_validation_error(err, handler, info)
-            raise
-
-    @classmethod
-    def rebuild_annotations(cls):
-
-        from aiq.cli.type_registry import GlobalTypeRegistry
-
-        type_registry = GlobalTypeRegistry.get()
-
-        FrontEndAnnotation = typing.Annotated[type_registry.compute_annotation(FrontEndBaseConfig),
-                                              Discriminator(TypedBaseModel.discriminator)]
-
-        should_rebuild = False
-
-        front_end_field = cls.model_fields.get("front_end")
-        if front_end_field is not None and front_end_field.annotation != FrontEndAnnotation:
-            front_end_field.annotation = FrontEndAnnotation
-            should_rebuild = True
-
-        if (TelemetryConfig.rebuild_annotations()):
-            should_rebuild = True
-
-        if (should_rebuild):
-            return cls.model_rebuild(force=True)
-
-        return False
-
-
-class AIQConfig(HashableBaseModel):
-
-    model_config = ConfigDict(extra="forbid")
-
-    # Global Options
-    general: GeneralConfig = GeneralConfig()
-
-    # Functions Configuration
-    functions: dict[str, FunctionBaseConfig] = {}
-
-    # LLMs Configuration
-    llms: dict[str, LLMBaseConfig] = {}
-
-    # Embedders Configuration
-    embedders: dict[str, EmbedderBaseConfig] = {}
-
-    # Memory Configuration
-    memory: dict[str, MemoryBaseConfig] = {}
-
-    # Retriever Configuration
-    retrievers: dict[str, RetrieverBaseConfig] = {}
-
-    # Workflow Configuration
-    workflow: FunctionBaseConfig = EmptyFunctionConfig()
-
-    # Evaluation Options
-    eval: EvalConfig = EvalConfig()
-
-    def print_summary(self, stream: typing.TextIO = sys.stdout):
-        """Print a summary of the configuration"""
-
-        stream.write("\nConfiguration Summary:\n")
-        stream.write("-" * 20 + "\n")
-        if self.workflow:
-            stream.write(f"Workflow Type: {self.workflow.type}\n")
-
-        stream.write(f"Number of Functions: {len(self.functions)}\n")
-        stream.write(f"Number of LLMs: {len(self.llms)}\n")
-        stream.write(f"Number of Embedders: {len(self.embedders)}\n")
-        stream.write(f"Number of Memory: {len(self.memory)}\n")
-        stream.write(f"Number of Retrievers: {len(self.retrievers)}\n")
-
-    @field_validator("functions", "llms", "embedders", "memory", "retrievers", "workflow", mode="wrap")
-    @classmethod
-    def validate_components(cls, value: typing.Any, handler: ValidatorFunctionWrapHandler, info: ValidationInfo):
-
-        try:
-            return handler(value)
-        except ValidationError as err:
-            _process_validation_error(err, handler, info)
-            raise
-
-    @classmethod
-    def rebuild_annotations(cls):
-
-        from aiq.cli.type_registry import GlobalTypeRegistry
-
-        type_registry = GlobalTypeRegistry.get()
-
-        LLMsAnnotation = dict[str,
-                              typing.Annotated[type_registry.compute_annotation(LLMBaseConfig),
-                                               Discriminator(TypedBaseModel.discriminator)]]
-
-        EmbeddersAnnotation = dict[str,
-                                   typing.Annotated[type_registry.compute_annotation(EmbedderBaseConfig),
-                                                    Discriminator(TypedBaseModel.discriminator)]]
-
-        FunctionsAnnotation = dict[str,
-                                   typing.Annotated[type_registry.compute_annotation(FunctionBaseConfig, ),
-                                                    Discriminator(TypedBaseModel.discriminator)]]
-
-        MemoryAnnotation = dict[str,
-                                typing.Annotated[type_registry.compute_annotation(MemoryBaseConfig),
-                                                 Discriminator(TypedBaseModel.discriminator)]]
-
-        RetrieverAnnotation = dict[str,
-                                   typing.Annotated[type_registry.compute_annotation(RetrieverBaseConfig),
-                                                    Discriminator(TypedBaseModel.discriminator)]]
-
-        WorkflowAnnotation = typing.Annotated[type_registry.compute_annotation(FunctionBaseConfig),
-                                              Discriminator(TypedBaseModel.discriminator)]
-
-        should_rebuild = False
-
-        llms_field = cls.model_fields.get("llms")
-        if llms_field is not None and llms_field.annotation != LLMsAnnotation:
-            llms_field.annotation = LLMsAnnotation
-            should_rebuild = True
-
-        embedders_field = cls.model_fields.get("embedders")
-        if embedders_field is not None and embedders_field.annotation != EmbeddersAnnotation:
-            embedders_field.annotation = EmbeddersAnnotation
-            should_rebuild = True
-
-        functions_field = cls.model_fields.get("functions")
-        if functions_field is not None and functions_field.annotation != FunctionsAnnotation:
-            functions_field.annotation = FunctionsAnnotation
-            should_rebuild = True
-
-        memory_field = cls.model_fields.get("memory")
-        if memory_field is not None and memory_field.annotation != MemoryAnnotation:
-            memory_field.annotation = MemoryAnnotation
-            should_rebuild = True
-
-        retrievers_field = cls.model_fields.get("retrievers")
-        if retrievers_field is not None and retrievers_field.annotation != RetrieverAnnotation:
-            retrievers_field.annotation = RetrieverAnnotation
-            should_rebuild = True
-
-        workflow_field = cls.model_fields.get("workflow")
-        if workflow_field is not None and workflow_field.annotation != WorkflowAnnotation:
-            workflow_field.annotation = WorkflowAnnotation
-            should_rebuild = True
-
-        if (GeneralConfig.rebuild_annotations()):
-            should_rebuild = True
-
-        if (EvalConfig.rebuild_annotations()):
-            should_rebuild = True
-
-        if (should_rebuild):
-            return cls.model_rebuild(force=True)
-
-        return False