aiqtoolkit 1.2.0.dev0__py3-none-any.whl → 1.2.0rc2__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

@@ -21,6 +21,8 @@ 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")
 
@@ -67,8 +69,8 @@ def subclass_depth(cls: type) -> int:
     Compute a class' subclass depth.
     """
     depth = 0
-    while (cls is not object):
-        cls = cls.__base__
+    while (cls is not object and cls.__base__ is not None):
+        cls = cls.__base__  # type: ignore
         depth += 1
     return depth
 
@@ -93,7 +95,8 @@ class TypedBaseModel(BaseModel):
     Subclass of Pydantic BaseModel that allows for specifying the object type. Use in Pydantic discriminated unions.
     """
 
-    type: str = Field(init=False,
+    type: str = Field(default="unknown",
+                      init=False,
                       serialization_alias="_type",
                       validation_alias=AliasChoices('type', '_type'),
                       description="The type of the object",
@@ -101,6 +104,7 @@ class TypedBaseModel(BaseModel):
                       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__()
@@ -117,14 +121,38 @@ class TypedBaseModel(BaseModel):
 
             full_name = f"{package_name}/{name}"
 
-            type_field = cls.model_fields.get("type")
-            if type_field is not None:
-                type_field.default = 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 cls.model_fields.get("type").default
+        return getattr(cls, '_typed_model_name')
 
     @classmethod
     def static_full_type(cls):

aiq/data_models/component.py

@@ -20,27 +20,35 @@ 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"
-    TOOL_WRAPPER = "tool_wrapper"
-    LLM_PROVIDER = "llm_provider"
+    ITS_STRATEGY = "its_strategy"
     LLM_CLIENT = "llm_client"
-    EMBEDDER_PROVIDER = "embedder_provider"
-    EMBEDDER_CLIENT = "embedder_client"
-    EVALUATOR = "evaluator"
+    LLM_PROVIDER = "llm_provider"
+    LOGGING = "logging"
     MEMORY = "memory"
-    RETRIEVER_PROVIDER = "retriever_provider"
-    RETRIEVER_CLIENT = "retriever_client"
+    OBJECT_STORE = "object_store"
+    PACKAGE = "package"
     REGISTRY_HANDLER = "registry_handler"
-    LOGGING = "logging"
+    RETRIEVER_CLIENT = "retriever_client"
+    RETRIEVER_PROVIDER = "retriever_provider"
+    TOOL_WRAPPER = "tool_wrapper"
     TRACING = "tracing"
-    PACKAGE = "package"
     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

@@ -124,6 +124,17 @@ class MemoryRef(ComponentRef):
         return ComponentGroup.MEMORY
 
 
+class ObjectStoreRef(ComponentRef):
+    """
+    A reference to an object store in an AIQ toolkit configuration object.
+    """
+
+    @property
+    @typing.override
+    def component_group(self):
+        return ComponentGroup.OBJECT_STORES
+
+
 class RetrieverRef(ComponentRef):
     """
     A reference to a retriever in an AIQ Toolkit configuration object.
@@ -133,3 +144,25 @@ class RetrieverRef(ComponentRef):
     @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

aiq/data_models/config.py

@@ -29,15 +29,18 @@ 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.its_strategy import ITSStrategyBaseConfig
 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 .authentication import AuthProviderBaseConfig
 from .common import HashableBaseModel
 from .common import TypedBaseModel
 from .embedder import EmbedderBaseConfig
 from .llm import LLMBaseConfig
 from .memory import MemoryBaseConfig
+from .object_store import ObjectStoreBaseConfig
 from .retriever import RetrieverBaseConfig
 
 logger = logging.getLogger(__name__)
@@ -57,12 +60,16 @@ def _process_validation_error(err: ValidationError, handler: ValidatorFunctionWr
 
             if (info.field_name in ('workflow', 'functions')):
                 registered_keys = GlobalTypeRegistry.get().get_registered_functions()
+            elif (info.field_name == "authentication"):
+                registered_keys = GlobalTypeRegistry.get().get_registered_auth_providers()
             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 == "object_stores"):
+                registered_keys = GlobalTypeRegistry.get().get_registered_object_stores()
             elif (info.field_name == "retrievers"):
                 registered_keys = GlobalTypeRegistry.get().get_registered_retriever_providers()
             elif (info.field_name == "tracing"):
@@ -73,6 +80,8 @@ def _process_validation_error(err: ValidationError, handler: ValidatorFunctionWr
                 registered_keys = GlobalTypeRegistry.get().get_registered_evaluators()
             elif (info.field_name == "front_ends"):
                 registered_keys = GlobalTypeRegistry.get().get_registered_front_ends()
+            elif (info.field_name == "its_strategies"):
+                registered_keys = GlobalTypeRegistry.get().get_registered_its_strategies()
 
             else:
                 assert False, f"Unknown field name {info.field_name} in validator"
@@ -242,12 +251,21 @@ class AIQConfig(HashableBaseModel):
     # Memory Configuration
     memory: dict[str, MemoryBaseConfig] = {}
 
+    # Object Stores Configuration
+    object_stores: dict[str, ObjectStoreBaseConfig] = {}
+
     # Retriever Configuration
     retrievers: dict[str, RetrieverBaseConfig] = {}
 
+    # ITS Strategies
+    its_strategies: dict[str, ITSStrategyBaseConfig] = {}
+
     # Workflow Configuration
     workflow: FunctionBaseConfig = EmptyFunctionConfig()
 
+    # Authentication Configuration
+    authentication: dict[str, AuthProviderBaseConfig] = {}
+
     # Evaluation Options
     eval: EvalConfig = EvalConfig()
 
@@ -263,9 +281,20 @@ class AIQConfig(HashableBaseModel):
         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 Object Stores: {len(self.object_stores)}\n")
         stream.write(f"Number of Retrievers: {len(self.retrievers)}\n")
-
-    @field_validator("functions", "llms", "embedders", "memory", "retrievers", "workflow", mode="wrap")
+        stream.write(f"Number of ITS Strategies: {len(self.its_strategies)}\n")
+        stream.write(f"Number of Authentication Providers: {len(self.authentication)}\n")
+
+    @field_validator("functions",
+                     "llms",
+                     "embedders",
+                     "memory",
+                     "retrievers",
+                     "workflow",
+                     "its_strategies",
+                     "authentication",
+                     mode="wrap")
     @classmethod
     def validate_components(cls, value: typing.Any, handler: ValidatorFunctionWrapHandler, info: ValidationInfo):
 
@@ -286,27 +315,45 @@ class AIQConfig(HashableBaseModel):
                               typing.Annotated[type_registry.compute_annotation(LLMBaseConfig),
                                                Discriminator(TypedBaseModel.discriminator)]]
 
+        AuthenticationProviderAnnotation = dict[str,
+                                                typing.Annotated[
+                                                    type_registry.compute_annotation(AuthProviderBaseConfig),
+                                                    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, ),
+                                   typing.Annotated[type_registry.compute_annotation(FunctionBaseConfig),
                                                     Discriminator(TypedBaseModel.discriminator)]]
 
         MemoryAnnotation = dict[str,
                                 typing.Annotated[type_registry.compute_annotation(MemoryBaseConfig),
                                                  Discriminator(TypedBaseModel.discriminator)]]
 
+        ObjectStoreAnnotation = dict[str,
+                                     typing.Annotated[type_registry.compute_annotation(ObjectStoreBaseConfig),
+                                                      Discriminator(TypedBaseModel.discriminator)]]
+
         RetrieverAnnotation = dict[str,
                                    typing.Annotated[type_registry.compute_annotation(RetrieverBaseConfig),
                                                     Discriminator(TypedBaseModel.discriminator)]]
 
+        ITSStrategyAnnotation = dict[str,
+                                     typing.Annotated[type_registry.compute_annotation(ITSStrategyBaseConfig),
+                                                      Discriminator(TypedBaseModel.discriminator)]]
+
         WorkflowAnnotation = typing.Annotated[type_registry.compute_annotation(FunctionBaseConfig),
                                               Discriminator(TypedBaseModel.discriminator)]
 
         should_rebuild = False
 
+        auth_providers_field = cls.model_fields.get("authentication")
+        if auth_providers_field is not None and auth_providers_field.annotation != AuthenticationProviderAnnotation:
+            auth_providers_field.annotation = AuthenticationProviderAnnotation
+            should_rebuild = True
+
         llms_field = cls.model_fields.get("llms")
         if llms_field is not None and llms_field.annotation != LLMsAnnotation:
             llms_field.annotation = LLMsAnnotation
@@ -327,11 +374,21 @@ class AIQConfig(HashableBaseModel):
             memory_field.annotation = MemoryAnnotation
             should_rebuild = True
 
+        object_stores_field = cls.model_fields.get("object_stores")
+        if object_stores_field is not None and object_stores_field.annotation != ObjectStoreAnnotation:
+            object_stores_field.annotation = ObjectStoreAnnotation
+            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
 
+        its_strategies_field = cls.model_fields.get("its_strategies")
+        if its_strategies_field is not None and its_strategies_field.annotation != ITSStrategyAnnotation:
+            its_strategies_field.annotation = ITSStrategyAnnotation
+            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

aiq/data_models/dataset_handler.py

@@ -30,7 +30,8 @@ from aiq.data_models.common import TypedBaseModel
 
 class EvalS3Config(BaseModel):
 
-    endpoint_url: str
+    endpoint_url: str | None = None
+    region_name: str | None = None
     bucket: str
     access_key: str
     secret_key: str

aiq/data_models/embedder.py

@@ -20,6 +20,7 @@ from .common import TypedBaseModel
 
 
 class EmbedderBaseConfig(TypedBaseModel, BaseModelRegistryTag):
+    """ Base configuration for embedding model providers. """
     pass
 
 

aiq/data_models/evaluate.py

@@ -14,6 +14,7 @@
 # limitations under the License.
 
 import typing
+from enum import Enum
 from pathlib import Path
 
 from pydantic import BaseModel
@@ -28,6 +29,12 @@ from aiq.data_models.intermediate_step import IntermediateStepType
 from aiq.data_models.profiler import ProfilerConfig
 
 
+class JobEvictionPolicy(str, Enum):
+    """Policy for evicting old jobs when max_jobs is exceeded."""
+    TIME_CREATED = "time_created"
+    TIME_MODIFIED = "time_modified"
+
+
 class EvalCustomScriptConfig(BaseModel):
     # Path to the script to run
     script: Path
@@ -35,6 +42,16 @@ class EvalCustomScriptConfig(BaseModel):
     kwargs: dict[str, str] = {}
 
 
+class JobManagementConfig(BaseModel):
+    # Whether to append a unique job ID to the output directory for each run
+    append_job_id_to_output_dir: bool = False
+    # Maximum number of jobs to keep in the output directory. Oldest jobs will be evicted.
+    # A value of 0 means no limit.
+    max_jobs: int = 0
+    # Policy for evicting old jobs. Defaults to using time_created.
+    eviction_policy: JobEvictionPolicy = JobEvictionPolicy.TIME_CREATED
+
+
 class EvalOutputConfig(BaseModel):
     # Output directory for the workflow and evaluation results
     dir: Path = Path("/tmp/aiq/examples/default/")
@@ -46,6 +63,8 @@ class EvalOutputConfig(BaseModel):
     s3: EvalS3Config | None = None
     # Whether to cleanup the output directory before running the workflow
     cleanup: bool = True
+    # Job management configuration (job id, eviction, etc.)
+    job_management: JobManagementConfig = JobManagementConfig()
     # Filter for the workflow output steps
     workflow_output_step_filter: list[IntermediateStepType] | None = None
 
@@ -53,6 +72,10 @@ class EvalOutputConfig(BaseModel):
 class EvalGeneralConfig(BaseModel):
     max_concurrency: int = 8
 
+    # Workflow alias for displaying in evaluation UI, if not provided,
+    # the workflow type will be used
+    workflow_alias: str | None = None
+
     # Output directory for the workflow and evaluation results
     output_dir: Path = Path("/tmp/aiq/examples/default/")
 

aiq/data_models/function_dependencies.py

@@ -26,6 +26,7 @@ class FunctionDependencies(BaseModel):
     llms: set[str] = Field(default_factory=set)
     embedders: set[str] = Field(default_factory=set)
     memory_clients: set[str] = Field(default_factory=set)
+    object_stores: set[str] = Field(default_factory=set)
     retrievers: set[str] = Field(default_factory=set)
 
     @field_serializer("functions", when_used="json")
@@ -44,6 +45,10 @@ class FunctionDependencies(BaseModel):
     def serialize_memory_clients(self, v: set[str]) -> list[str]:
         return list(v)
 
+    @field_serializer("object_stores", when_used="json")
+    def serialize_object_stores(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)
@@ -60,5 +65,8 @@ class FunctionDependencies(BaseModel):
     def add_memory_client(self, memory_client: str):
         self.memory_clients.add(memory_client)  # pylint: disable=no-member
 
+    def add_object_store(self, object_store: str):
+        self.object_stores.add(object_store)  # pylint: disable=no-member
+
     def add_retriever(self, retriever: str):
         self.retrievers.add(retriever)  # pylint: disable=no-member