aiqtoolkit 1.1.0__py3-none-any.whl

aiq/data_models/dataset_handler.py

@@ -0,0 +1,122 @@
+# 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 json
+import typing
+from collections.abc import Callable
+from pathlib import Path
+
+import pandas as pd
+from pydantic import BaseModel
+from pydantic import Discriminator
+from pydantic import FilePath
+from pydantic import Tag
+
+from aiq.data_models.common import BaseModelRegistryTag
+from aiq.data_models.common import TypedBaseModel
+
+
+class EvalS3Config(BaseModel):
+
+    endpoint_url: str
+    bucket: str
+    access_key: str
+    secret_key: str
+
+
+class EvalFilterEntryConfig(BaseModel):
+    # values are lists of allowed/blocked values
+    field: dict[str, list[str | int | float]] = {}
+
+
+class EvalFilterConfig(BaseModel):
+    allowlist: EvalFilterEntryConfig | None = None
+    denylist: EvalFilterEntryConfig | None = None
+
+
+class EvalDatasetStructureConfig(BaseModel):
+    disable: bool = False
+    question_key: str = "question"
+    answer_key: str = "answer"
+    generated_answer_key: str = "generated_answer"
+    trajectory_key: str = "intermediate_steps"
+    expected_trajectory_key: str = "expected_intermediate_steps"
+
+
+# Base model
+class EvalDatasetBaseConfig(TypedBaseModel, BaseModelRegistryTag):
+
+    id_key: str = "id"
+    structure: EvalDatasetStructureConfig = EvalDatasetStructureConfig()
+
+    # Filters
+    filter: EvalFilterConfig | None = EvalFilterConfig()
+
+    s3: EvalS3Config | None = None
+
+    remote_file_path: str | None = None  # only for s3
+    file_path: Path | str = Path(".tmp/aiq/examples/default/default.json")
+
+
+class EvalDatasetJsonConfig(EvalDatasetBaseConfig, name="json"):
+
+    @staticmethod
+    def parser() -> tuple[Callable, dict]:
+        return pd.read_json, {}
+
+
+def read_jsonl(file_path: FilePath, **kwargs):
+    with open(file_path, 'r', encoding='utf-8') as f:
+        data = [json.loads(line) for line in f]
+    return pd.DataFrame(data)
+
+
+class EvalDatasetJsonlConfig(EvalDatasetBaseConfig, name="jsonl"):
+
+    @staticmethod
+    def parser() -> tuple[Callable, dict]:
+        return read_jsonl, {}
+
+
+class EvalDatasetCsvConfig(EvalDatasetBaseConfig, name="csv"):
+
+    @staticmethod
+    def parser() -> tuple[Callable, dict]:
+        return pd.read_csv, {}
+
+
+class EvalDatasetParquetConfig(EvalDatasetBaseConfig, name="parquet"):
+
+    @staticmethod
+    def parser() -> tuple[Callable, dict]:
+        return pd.read_parquet, {}
+
+
+class EvalDatasetXlsConfig(EvalDatasetBaseConfig, name="xls"):
+
+    @staticmethod
+    def parser() -> tuple[Callable, dict]:
+        return pd.read_excel, {"engine": "openpyxl"}
+
+
+# Union model with discriminator
+EvalDatasetConfig = typing.Annotated[typing.Annotated[EvalDatasetJsonConfig, Tag(EvalDatasetJsonConfig.static_type())]
+                                     | typing.Annotated[EvalDatasetCsvConfig, Tag(EvalDatasetCsvConfig.static_type())]
+                                     | typing.Annotated[EvalDatasetXlsConfig, Tag(EvalDatasetXlsConfig.static_type())]
+                                     | typing.Annotated[EvalDatasetParquetConfig,
+                                                        Tag(EvalDatasetParquetConfig.static_type())]
+                                     | typing.Annotated[EvalDatasetJsonlConfig,
+                                                        Tag(EvalDatasetJsonlConfig.static_type())],
+                                     Discriminator(TypedBaseModel.discriminator)]

aiq/data_models/discovery_metadata.py

@@ -0,0 +1,286 @@
+# 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 importlib.metadata
+import inspect
+import json
+import logging
+import typing
+from enum import Enum
+from functools import lru_cache
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+from pydantic import BaseModel
+from pydantic import field_validator
+
+from aiq.builder.framework_enum import LLMFrameworkEnum
+from aiq.data_models.component import AIQComponentEnum
+from aiq.utils.metadata_utils import generate_config_type_docs
+
+if TYPE_CHECKING:
+    from aiq.cli.type_registry import ToolWrapperBuildCallableT
+    from aiq.data_models.common import TypedBaseModelT
+
+logger = logging.getLogger(__name__)
+
+
+class DiscoveryStatusEnum(str, Enum):
+    SUCCESS = "success"
+    FAILURE = "failure"
+
+
+class DiscoveryContractFieldsEnum(str, Enum):
+    PACKAGE = "package"
+    VERSION = "version"
+    COMPONENT_TYPE = "component_type"
+    COMPONENT_NAME = "component_name"
+    DESCRIPTION = "description"
+    DEVELOPER_NOTES = "developer_notes"
+
+
+class DiscoveryMetadata(BaseModel):
+    """A data model representing metadata about each registered component to faciliate its discovery.
+
+    Args:
+        package (str): The name of the package containing the AIQ Toolkit component.
+        version (str): The version number of the package containing the AIQ Toolkit component.
+        component_type (AIQComponentEnum): The type of AIQ Toolkit component this metadata represents.
+        component_name (str): The registered name of the AIQ Toolkit component.
+        description (str): Description of the AIQ Toolkit component pulled from its config objects docstrings.
+        developer_notes (str): Other notes to a developers to aid in the use of the component.
+        status (DiscoveryStatusEnum): Provides the status of the metadata discovery process.
+    """
+
+    package: str = ""
+    version: str = ""
+    component_type: AIQComponentEnum = AIQComponentEnum.UNDEFINED
+    component_name: str = ""
+    description: str = ""
+    developer_notes: str = ""
+    status: DiscoveryStatusEnum = DiscoveryStatusEnum.SUCCESS
+
+    @field_validator("description", mode="before")
+    @classmethod
+    def ensure_description_string(cls, v: typing.Any):
+        if not isinstance(v, str):
+            return ""
+        return v
+
+    @staticmethod
+    def get_preferred_item(items: list, preferred: str) -> str:
+        return preferred if preferred in items else items[0]
+
+    @staticmethod
+    @lru_cache
+    def get_distribution_name_from_metadata(root_package_name: str) -> str | None:
+        """
+        This is not performant and is only present to be used (not used
+        currently) as a fallback when the distro name doesn't match the
+        module name and private_data is not available to map it.
+        """
+        mapping = importlib.metadata.packages_distributions()
+        try:
+            distro_names = mapping.get(root_package_name, [None])
+            distro_name = DiscoveryMetadata.get_preferred_item(distro_names, "aiqtoolkit")
+        except KeyError:
+            return root_package_name
+
+        return distro_name if distro_name else root_package_name
+
+    @staticmethod
+    @lru_cache
+    def get_distribution_name_from_private_data(root_package: str) -> str | None:
+        # Locate distibution mapping stored in the packages private data
+        module = __import__(root_package)
+        for path in module.__path__:
+            package_dir = Path(path).resolve()
+            distinfo_path = package_dir / "meta" / "module_to_distro.json"
+
+            if distinfo_path.exists():
+                with distinfo_path.open("r") as f:
+                    data = json.load(f)
+                    return data.get(root_package, None)
+        return None
+
+    @staticmethod
+    @lru_cache
+    def get_distribution_name(root_package: str) -> str:
+        """
+        The aiq library packages use a distro name 'aiqtoolkit[]' and
+        root package name 'aiq'. They provide mapping in a metadata file
+        for optimized installation.
+        """
+        distro_name = DiscoveryMetadata.get_distribution_name_from_private_data(root_package)
+        return distro_name if distro_name else root_package
+
+    @staticmethod
+    def from_config_type(config_type: type["TypedBaseModelT"],
+                         component_type: AIQComponentEnum = AIQComponentEnum.UNDEFINED) -> "DiscoveryMetadata":
+        """Generates discovery metadata from an AIQ Toolkit config object.
+
+        Args:
+            config_type (type[TypedBaseModelT]): A registered component's configuration object.
+            component_type (AIQComponentEnum, optional): The type of the registered component. Defaults to
+            AIQComponentEnum.UNDEFINED.
+
+        Returns:
+            DiscoveryMetadata: A an object containing component metadata to facilitate discovery and reuse.
+        """
+
+        try:
+            module = inspect.getmodule(config_type)
+            root_package: str = module.__package__.split(".")[0]
+            distro_name = DiscoveryMetadata.get_distribution_name(root_package)
+
+            if not distro_name:
+                # raise an exception
+                logger.error("Encountered issue getting distro_name for module %s", module.__name__)
+                return DiscoveryMetadata(status=DiscoveryStatusEnum.FAILURE)
+
+            try:
+                version = importlib.metadata.version(distro_name) if distro_name != "" else ""
+            except importlib.metadata.PackageNotFoundError:
+                logger.warning("Package metadata not found for %s", distro_name)
+                version = ""
+        except Exception as e:
+            logger.exception("Encountered issue extracting module metadata for %s: %s", config_type, e, exc_info=True)
+            return DiscoveryMetadata(status=DiscoveryStatusEnum.FAILURE)
+
+        description = generate_config_type_docs(config_type=config_type)
+
+        return DiscoveryMetadata(package=distro_name,
+                                 version=version,
+                                 component_type=component_type,
+                                 component_name=config_type.static_type(),
+                                 description=description)
+
+    @staticmethod
+    def from_fn_wrapper(fn: "ToolWrapperBuildCallableT",
+                        wrapper_type: LLMFrameworkEnum | str,
+                        component_type: AIQComponentEnum = AIQComponentEnum.TOOL_WRAPPER) -> "DiscoveryMetadata":
+        """Generates discovery metadata from function with specified wrapper type.
+
+        Args:
+            fn (ToolWrapperBuildCallableT): A tool wrapper callable to source component metadata.
+            wrapper_type (LLMFrameworkEnum): The wrapper to apply to the callable to faciliate inter-framwork
+            interoperability.
+
+            component_type (AIQComponentEnum, optional): The type of the registered component. Defaults to
+            AIQComponentEnum.TOOL_WRAPPER.
+
+        Returns:
+            DiscoveryMetadata: A an object containing component metadata to facilitate discovery and reuse.
+        """
+
+        try:
+            module = inspect.getmodule(fn)
+            root_package: str = module.__package__.split(".")[0]
+            root_package = DiscoveryMetadata.get_distribution_name(root_package)
+            try:
+                version = importlib.metadata.version(root_package) if root_package != "" else ""
+            except importlib.metadata.PackageNotFoundError:
+                logger.warning("Package metadata not found for %s", root_package)
+                version = ""
+        except Exception as e:
+            logger.exception("Encountered issue extracting module metadata for %s: %s", fn, e, exc_info=True)
+            return DiscoveryMetadata(status=DiscoveryStatusEnum.FAILURE)
+
+        if isinstance(wrapper_type, LLMFrameworkEnum):
+            wrapper_type = wrapper_type.value
+
+        return DiscoveryMetadata(package=root_package,
+                                 version=version,
+                                 component_type=component_type,
+                                 component_name=wrapper_type,
+                                 description=fn.__doc__ or "")
+
+    @staticmethod
+    def from_package_name(package_name: str, package_version: str | None) -> "DiscoveryMetadata":
+        """Generates discovery metadata from an installed package name.
+
+        Args:
+            package_name (str): The name of the AIQ Toolkit plugin package containing registered components.
+            package_version (str, optional): The version of the package, Defaults to None.
+
+        Returns:
+            DiscoveryMetadata: A an object containing component metadata to facilitate discovery and reuse.
+        """
+
+        try:
+            package_name = DiscoveryMetadata.get_distribution_name(package_name)
+            try:
+                metadata = importlib.metadata.metadata(package_name)
+                description = metadata.get("Summary", "")
+                if (package_version is None):
+                    package_version = importlib.metadata.version(package_name)
+            except importlib.metadata.PackageNotFoundError:
+                logger.warning("Package metadata not found for %s", package_name)
+                description = ""
+                package_version = package_version or ""
+        except Exception as e:
+            logger.exception("Encountered issue extracting module metadata for %s: %s", package_name, e, exc_info=True)
+            return DiscoveryMetadata(status=DiscoveryStatusEnum.FAILURE)
+
+        return DiscoveryMetadata(package=package_name,
+                                 version=package_version,
+                                 component_type=AIQComponentEnum.PACKAGE,
+                                 component_name=package_name,
+                                 description=description)
+
+    @staticmethod
+    def from_provider_framework_map(
+            config_type: type["TypedBaseModelT"],
+            wrapper_type: LLMFrameworkEnum | str | None,
+            provider_type: AIQComponentEnum,
+            component_type: AIQComponentEnum = AIQComponentEnum.UNDEFINED) -> "DiscoveryMetadata":
+        """Generates discovery metadata from provider and framework mapping information.
+
+        Args:
+            config_type (type[TypedBaseModelT]): A registered component's configuration object.
+            wrapper_type (LLMFrameworkEnum | str): The wrapper to apply to the callable to faciliate inter-framwork
+            interoperability.
+
+            provider_type (AIQComponentEnum): The type of provider the registered component supports.
+            component_type (AIQComponentEnum, optional): The type of the registered component. Defaults to
+            AIQComponentEnum.UNDEFINED.
+
+        Returns:
+            DiscoveryMetadata: A an object containing component metadata to facilitate discovery and reuse.
+        """
+
+        try:
+            module = inspect.getmodule(config_type)
+            root_package: str = module.__package__.split(".")[0]
+            root_package = DiscoveryMetadata.get_distribution_name(root_package)
+            try:
+                version = importlib.metadata.version(root_package) if root_package != "" else ""
+            except importlib.metadata.PackageNotFoundError:
+                logger.warning("Package metadata not found for %s", root_package)
+                version = ""
+        except Exception as e:
+            logger.exception("Encountered issue extracting module metadata for %s: %s", config_type, e, exc_info=True)
+            return DiscoveryMetadata(status=DiscoveryStatusEnum.FAILURE)
+
+        wrapper_type = wrapper_type.value if isinstance(wrapper_type, LLMFrameworkEnum) else wrapper_type
+        component_name = f"{config_type.static_type()} ({provider_type.value}) - {wrapper_type}"
+
+        description = generate_config_type_docs(config_type=config_type)
+
+        return DiscoveryMetadata(package=root_package,
+                                 version=version,
+                                 component_type=component_type,
+                                 component_name=component_name,
+                                 description=description)

aiq/data_models/embedder.py

@@ -0,0 +1,26 @@
+# 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 .common import BaseModelRegistryTag
+from .common import TypedBaseModel
+
+
+class EmbedderBaseConfig(TypedBaseModel, BaseModelRegistryTag):
+    pass
+
+
+EmbedderBaseConfigT = typing.TypeVar("EmbedderBaseConfigT", bound=EmbedderBaseConfig)

aiq/data_models/evaluate.py

@@ -0,0 +1,104 @@
+# 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 pathlib import Path
+
+from pydantic import BaseModel
+from pydantic import Discriminator
+from pydantic import model_validator
+
+from aiq.data_models.common import TypedBaseModel
+from aiq.data_models.dataset_handler import EvalDatasetConfig
+from aiq.data_models.dataset_handler import EvalS3Config
+from aiq.data_models.evaluator import EvaluatorBaseConfig
+from aiq.data_models.intermediate_step import IntermediateStepType
+from aiq.data_models.profiler import ProfilerConfig
+
+
+class EvalCustomScriptConfig(BaseModel):
+    # Path to the script to run
+    script: Path
+    # Keyword arguments to pass to the script
+    kwargs: dict[str, str] = {}
+
+
+class EvalOutputConfig(BaseModel):
+    # Output directory for the workflow and evaluation results
+    dir: Path = Path("/tmp/aiq/examples/default/")
+    # S3 prefix for the workflow and evaluation results
+    remote_dir: str | None = None
+    # Custom scripts to run after the workflow and evaluation results are saved
+    custom_scripts: dict[str, EvalCustomScriptConfig] = {}
+    # S3 config for uploading the contents of the output directory
+    s3: EvalS3Config | None = None
+    # Whether to cleanup the output directory before running the workflow
+    cleanup: bool = True
+    # Filter for the workflow output steps
+    workflow_output_step_filter: list[IntermediateStepType] | None = None
+
+
+class EvalGeneralConfig(BaseModel):
+    max_concurrency: int = 8
+
+    # Output directory for the workflow and evaluation results
+    output_dir: Path = Path("/tmp/aiq/examples/default/")
+
+    # If present overrides output_dir
+    output: EvalOutputConfig | None = None
+
+    # Dataset for running the workflow and evaluating
+    dataset: EvalDatasetConfig | None = None
+
+    # Inference profiler
+    profiler: ProfilerConfig | None = None
+
+    # overwrite the output_dir with the output config if present
+    @model_validator(mode="before")
+    @classmethod
+    def override_output_dir(cls, values):
+        if values.get("output") and values["output"].get("dir"):
+            values["output_dir"] = values["output"]["dir"]
+        return values
+
+
+class EvalConfig(BaseModel):
+
+    # General Evaluation Options
+    general: EvalGeneralConfig = EvalGeneralConfig()
+
+    # Evaluators
+    evaluators: dict[str, EvaluatorBaseConfig] = {}
+
+    @classmethod
+    def rebuild_annotations(cls):
+
+        from aiq.cli.type_registry import GlobalTypeRegistry  # pylint: disable=cyclic-import
+
+        type_registry = GlobalTypeRegistry.get()
+
+        EvaluatorsAnnotation = dict[str,
+                                    typing.Annotated[type_registry.compute_annotation(EvaluatorBaseConfig),
+                                                     Discriminator(TypedBaseModel.discriminator)]]
+
+        should_rebuild = False
+
+        evaluators_field = cls.model_fields.get("evaluators")
+        if evaluators_field is not None and evaluators_field.annotation != EvaluatorsAnnotation:
+            evaluators_field.annotation = EvaluatorsAnnotation
+            should_rebuild = True
+
+        if (should_rebuild):
+            cls.model_rebuild(force=True)

aiq/data_models/evaluator.py

@@ -0,0 +1,26 @@
+# 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 .common import BaseModelRegistryTag
+from .common import TypedBaseModel
+
+
+class EvaluatorBaseConfig(TypedBaseModel, BaseModelRegistryTag):
+    pass
+
+
+EvaluatorBaseConfigT = typing.TypeVar("EvaluatorBaseConfigT", bound=EvaluatorBaseConfig)

aiq/data_models/front_end.py

@@ -0,0 +1,26 @@
+# 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 .common import BaseModelRegistryTag
+from .common import TypedBaseModel
+
+
+class FrontEndBaseConfig(TypedBaseModel, BaseModelRegistryTag):
+    pass
+
+
+FrontEndConfigT = typing.TypeVar("FrontEndConfigT", bound=FrontEndBaseConfig)

aiq/data_models/function.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.
+
+import typing
+
+from .common import BaseModelRegistryTag
+from .common import TypedBaseModel
+
+
+class FunctionBaseConfig(TypedBaseModel, BaseModelRegistryTag):
+    pass
+
+
+class EmptyFunctionConfig(FunctionBaseConfig, name="EmptyFunctionConfig"):
+    pass
+
+
+FunctionConfigT = typing.TypeVar("FunctionConfigT", bound=FunctionBaseConfig)

aiq/data_models/function_dependencies.py

@@ -0,0 +1,64 @@
+# 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 BaseModel
+from pydantic import Field
+from pydantic import field_serializer
+
+
+class FunctionDependencies(BaseModel):
+    """
+    A class to represent the dependencies of a function.
+    """
+    functions: set[str] = Field(default_factory=set)
+    llms: set[str] = Field(default_factory=set)
+    embedders: set[str] = Field(default_factory=set)
+    memory_clients: set[str] = Field(default_factory=set)
+    retrievers: set[str] = Field(default_factory=set)
+
+    @field_serializer("functions", when_used="json")
+    def serialize_functions(self, v: set[str]) -> list[str]:
+        return list(v)
+
+    @field_serializer("llms", when_used="json")
+    def serialize_llms(self, v: set[str]) -> list[str]:
+        return list(v)
+
+    @field_serializer("embedders", when_used="json")
+    def serialize_embedders(self, v: set[str]) -> list[str]:
+        return list(v)
+
+    @field_serializer("memory_clients", when_used="json")
+    def serialize_memory_clients(self, v: set[str]) -> list[str]:
+        return list(v)
+
+    @field_serializer("retrievers", when_used="json")
+    def serialize_retrievers(self, v: set[str]) -> list[str]:
+        return list(v)
+
+    def add_function(self, function: str):
+        self.functions.add(function)  # pylint: disable=no-member
+
+    def add_llm(self, llm: str):
+        self.llms.add(llm)  # pylint: disable=no-member
+
+    def add_embedder(self, embedder: str):
+        self.embedders.add(embedder)  # pylint: disable=no-member
+
+    def add_memory_client(self, memory_client: str):
+        self.memory_clients.add(memory_client)  # pylint: disable=no-member
+
+    def add_retriever(self, retriever: str):
+        self.retrievers.add(retriever)  # pylint: disable=no-member