mlrun 1.8.0rc1__py3-none-any.whl → 1.8.0rc3__py3-none-any.whl

mlrun/common/schemas/project.py

@@ -15,7 +15,7 @@
 import datetime
 import typing
 
-import pydantic
+import pydantic.v1
 from deprecated import deprecated
 
 import mlrun.common.types
@@ -40,14 +40,14 @@ class ProjectsFormat(mlrun.common.types.StrEnum):
     leader = "leader"
 
 
-class ProjectMetadata(pydantic.BaseModel):
+class ProjectMetadata(pydantic.v1.BaseModel):
     name: str
     created: typing.Optional[datetime.datetime] = None
     labels: typing.Optional[dict] = {}
     annotations: typing.Optional[dict] = {}
 
     class Config:
-        extra = pydantic.Extra.allow
+        extra = pydantic.v1.Extra.allow
 
 
 class ProjectDesiredState(mlrun.common.types.StrEnum):
@@ -77,7 +77,7 @@ class ProjectStatus(ObjectStatus):
     state: typing.Optional[ProjectState]
 
 
-class ProjectSpec(pydantic.BaseModel):
+class ProjectSpec(pydantic.v1.BaseModel):
     description: typing.Optional[str] = None
     owner: typing.Optional[str] = None
     goals: typing.Optional[str] = None
@@ -97,10 +97,10 @@ class ProjectSpec(pydantic.BaseModel):
     default_function_node_selector: typing.Optional[dict] = {}
 
     class Config:
-        extra = pydantic.Extra.allow
+        extra = pydantic.v1.Extra.allow
 
 
-class ProjectSpecOut(pydantic.BaseModel):
+class ProjectSpecOut(pydantic.v1.BaseModel):
     description: typing.Optional[str] = None
     owner: typing.Optional[str] = None
     goals: typing.Optional[str] = None
@@ -120,11 +120,11 @@ class ProjectSpecOut(pydantic.BaseModel):
     default_function_node_selector: typing.Optional[dict] = {}
 
     class Config:
-        extra = pydantic.Extra.allow
+        extra = pydantic.v1.Extra.allow
 
 
-class Project(pydantic.BaseModel):
-    kind: ObjectKind = pydantic.Field(ObjectKind.project, const=True)
+class Project(pydantic.v1.BaseModel):
+    kind: ObjectKind = pydantic.v1.Field(ObjectKind.project, const=True)
     metadata: ProjectMetadata
     spec: ProjectSpec = ProjectSpec()
     status: ObjectStatus = ObjectStatus()
@@ -132,19 +132,19 @@ class Project(pydantic.BaseModel):
 
 # The reason we have a different schema for the response model is that we don't want to validate project.spec.build in
 # the response as the validation was added late and there may be corrupted values in the DB.
-class ProjectOut(pydantic.BaseModel):
-    kind: ObjectKind = pydantic.Field(ObjectKind.project, const=True)
+class ProjectOut(pydantic.v1.BaseModel):
+    kind: ObjectKind = pydantic.v1.Field(ObjectKind.project, const=True)
     metadata: ProjectMetadata
     spec: ProjectSpecOut = ProjectSpecOut()
     status: ObjectStatus = ObjectStatus()
 
 
-class ProjectOwner(pydantic.BaseModel):
+class ProjectOwner(pydantic.v1.BaseModel):
     username: str
     access_key: str
 
 
-class ProjectSummary(pydantic.BaseModel):
+class ProjectSummary(pydantic.v1.BaseModel):
     name: str
     files_count: int = 0
     feature_sets_count: int = 0
@@ -161,7 +161,7 @@ class ProjectSummary(pydantic.BaseModel):
     updated: typing.Optional[datetime.datetime] = None
 
 
-class IguazioProject(pydantic.BaseModel):
+class IguazioProject(pydantic.v1.BaseModel):
     data: dict
 
 
@@ -175,13 +175,18 @@ class IguazioProject(pydantic.BaseModel):
 # to add a specific classes for them. it's frustrating but couldn't find other workaround, see:
 # https://github.com/samuelcolvin/pydantic/issues/1423, https://github.com/samuelcolvin/pydantic/issues/619
 ProjectOutput = typing.TypeVar(
-    "ProjectOutput", ProjectOut, str, ProjectSummary, IguazioProject
+    "ProjectOutput",
+    ProjectOut,
+    str,
+    ProjectSummary,
+    IguazioProject,
+    tuple[str, datetime.datetime],
 )
 
 
-class ProjectsOutput(pydantic.BaseModel):
+class ProjectsOutput(pydantic.v1.BaseModel):
     projects: list[ProjectOutput]
 
 
-class ProjectSummariesOutput(pydantic.BaseModel):
+class ProjectSummariesOutput(pydantic.v1.BaseModel):
     project_summaries: list[ProjectSummary]

mlrun/common/schemas/runs.py

@@ -14,13 +14,13 @@
 
 import typing
 
-import pydantic
+import pydantic.v1
 from deprecated import deprecated
 
 import mlrun.common.types
 
 
-class RunIdentifier(pydantic.BaseModel):
+class RunIdentifier(pydantic.v1.BaseModel):
     kind: typing.Literal["run"] = "run"
     uid: typing.Optional[str]
     iter: typing.Optional[int]

mlrun/common/schemas/runtime_resource.py

@@ -14,7 +14,7 @@
 #
 import typing
 
-import pydantic
+import pydantic.v1
 
 import mlrun.common.types
 
@@ -24,23 +24,23 @@ class ListRuntimeResourcesGroupByField(mlrun.common.types.StrEnum):
     project = "project"
 
 
-class RuntimeResource(pydantic.BaseModel):
+class RuntimeResource(pydantic.v1.BaseModel):
     name: str
     labels: dict[str, str] = {}
     status: typing.Optional[dict]
 
 
-class RuntimeResources(pydantic.BaseModel):
+class RuntimeResources(pydantic.v1.BaseModel):
     crd_resources: list[RuntimeResource] = []
     pod_resources: list[RuntimeResource] = []
     # only for dask runtime
     service_resources: typing.Optional[list[RuntimeResource]] = None
 
     class Config:
-        extra = pydantic.Extra.allow
+        extra = pydantic.v1.Extra.allow
 
 
-class KindRuntimeResources(pydantic.BaseModel):
+class KindRuntimeResources(pydantic.v1.BaseModel):
     kind: str
     resources: RuntimeResources
 

mlrun/common/schemas/schedule.py

@@ -15,7 +15,7 @@
 from datetime import datetime
 from typing import Any, Literal, Optional, Union
 
-from pydantic import BaseModel
+from pydantic.v1 import BaseModel
 
 import mlrun.common.types
 from mlrun.common.schemas.auth import Credentials

mlrun/common/schemas/secret.py

@@ -14,7 +14,7 @@
 #
 from typing import Optional
 
-from pydantic import BaseModel, Field
+from pydantic.v1 import BaseModel, Field
 
 import mlrun.common.types
 

mlrun/common/schemas/tag.py

@@ -13,17 +13,17 @@
 # limitations under the License.
 #
 
-import pydantic
+import pydantic.v1
 
 from .artifact import ArtifactIdentifier
 
 
-class Tag(pydantic.BaseModel):
+class Tag(pydantic.v1.BaseModel):
     name: str
     project: str
 
 
-class TagObjects(pydantic.BaseModel):
+class TagObjects(pydantic.v1.BaseModel):
     """Tag object"""
 
     kind: str

mlrun/common/schemas/workflow.py

@@ -14,14 +14,14 @@
 #
 import typing
 
-import pydantic
+import pydantic.v1
 
 from mlrun.common.schemas.notification import Notification
 from mlrun.common.schemas.schedule import ScheduleCronTrigger
 from mlrun.common.types import StrEnum
 
 
-class WorkflowSpec(pydantic.BaseModel):
+class WorkflowSpec(pydantic.v1.BaseModel):
     name: str
     engine: typing.Optional[str] = None
     code: typing.Optional[str] = None
@@ -36,7 +36,7 @@ class WorkflowSpec(pydantic.BaseModel):
     workflow_runner_node_selector: typing.Optional[dict[str, str]] = None
 
 
-class WorkflowRequest(pydantic.BaseModel):
+class WorkflowRequest(pydantic.v1.BaseModel):
     spec: typing.Optional[WorkflowSpec] = None
     arguments: typing.Optional[dict] = None
     artifact_path: typing.Optional[str] = None
@@ -46,7 +46,7 @@ class WorkflowRequest(pydantic.BaseModel):
     notifications: typing.Optional[list[Notification]] = None
 
 
-class WorkflowResponse(pydantic.BaseModel):
+class WorkflowResponse(pydantic.v1.BaseModel):
     project: str = None
     name: str = None
     status: str = None
@@ -54,7 +54,7 @@ class WorkflowResponse(pydantic.BaseModel):
     schedule: typing.Union[str, ScheduleCronTrigger] = None
 
 
-class GetWorkflowResponse(pydantic.BaseModel):
+class GetWorkflowResponse(pydantic.v1.BaseModel):
     workflow_id: str = None
 
 

mlrun/config.py

@@ -138,6 +138,9 @@ default_config = {
     "object_retentions": {
         "alert_activation": 14 * 7,  # days
     },
+    # A safety margin to account for delays
+    # This ensures that extra partitions are available beyond the specified retention period
+    "partitions_buffer_multiplier": 3,
     # the grace period (in seconds) that will be given to runtime resources (after they're in terminal state)
     # before deleting them (4 hours)
     "runtime_resources_deletion_grace_period": "14400",
@@ -530,7 +533,7 @@ default_config = {
             "verbose": True,
         },
         "pagination": {
-            "default_page_size": 20,
+            "default_page_size": 200,
             "pagination_cache": {
                 "interval": 60,
                 "ttl": 3600,
@@ -798,11 +801,30 @@ default_config = {
         "max_allowed": 10000,
         # maximum allowed value for count in criteria field inside AlertConfig
         "max_criteria_count": 100,
+        # interval for periodic events generation job
+        "events_generation_interval": "30",
     },
     "auth_with_client_id": {
         "enabled": False,
         "request_timeout": 5,
     },
+    "services": {
+        # The running service name. One of: "api", "alerts"
+        "service_name": "api",
+        "hydra": {
+            # Comma separated list of services to run on the instance.
+            # Currently, this is only considered when the service_name is "api".
+            # "*" starts all services on the same instance,
+            # other options are considered as running only the api service.
+            "services": "*",
+        },
+    },
+    "notifications": {
+        "smtp": {
+            "config_secret_name": "mlrun-smtp-config",
+            "refresh_interval": "30",
+        }
+    },
 }
 _is_running_as_api = None
 

mlrun/datastore/datastore_profile.py

@@ -19,7 +19,7 @@ import typing
 import warnings
 from urllib.parse import ParseResult, urlparse, urlunparse
 
-import pydantic
+import pydantic.v1
 from mergedeep import merge
 
 import mlrun
@@ -28,15 +28,15 @@ import mlrun.errors
 from ..secrets import get_secret_or_env
 
 
-class DatastoreProfile(pydantic.BaseModel):
+class DatastoreProfile(pydantic.v1.BaseModel):
     type: str
     name: str
     _private_attributes: list = ()
 
     class Config:
-        extra = pydantic.Extra.forbid
+        extra = pydantic.v1.Extra.forbid
 
-    @pydantic.validator("name")
+    @pydantic.v1.validator("name")
     @classmethod
     def lower_case(cls, v):
         return v.lower()
@@ -75,14 +75,32 @@ class TemporaryClientDatastoreProfiles(metaclass=mlrun.utils.singleton.Singleton
 
 
 class DatastoreProfileBasic(DatastoreProfile):
-    type: str = pydantic.Field("basic")
+    type: str = pydantic.v1.Field("basic")
     _private_attributes = "private"
     public: str
     private: typing.Optional[str] = None
 
 
+class VectorStoreProfile(DatastoreProfile):
+    type: str = pydantic.Field("vector")
+    _private_attributes = ("kwargs_private",)
+    vector_store_class: str
+    kwargs_public: typing.Optional[dict] = None
+    kwargs_private: typing.Optional[dict] = None
+
+    def attributes(self, kwargs=None):
+        attributes = {}
+        if self.kwargs_public:
+            attributes = merge(attributes, self.kwargs_public)
+        if self.kwargs_private:
+            attributes = merge(attributes, self.kwargs_private)
+        if kwargs:
+            attributes = merge(attributes, kwargs)
+        return attributes
+
+
 class DatastoreProfileKafkaTarget(DatastoreProfile):
-    type: str = pydantic.Field("kafka_target")
+    type: str = pydantic.v1.Field("kafka_target")
     _private_attributes = "kwargs_private"
     bootstrap_servers: typing.Optional[str] = None
     brokers: typing.Optional[str] = None
@@ -123,7 +141,7 @@ class DatastoreProfileKafkaTarget(DatastoreProfile):
 
 
 class DatastoreProfileKafkaSource(DatastoreProfile):
-    type: str = pydantic.Field("kafka_source")
+    type: str = pydantic.v1.Field("kafka_source")
     _private_attributes = ("kwargs_private", "sasl_user", "sasl_pass")
     brokers: typing.Union[str, list[str]]
     topics: typing.Union[str, list[str]]
@@ -162,7 +180,7 @@ class DatastoreProfileKafkaSource(DatastoreProfile):
 
 
 class DatastoreProfileV3io(DatastoreProfile):
-    type: str = pydantic.Field("v3io")
+    type: str = pydantic.v1.Field("v3io")
     v3io_access_key: typing.Optional[str] = None
     _private_attributes = "v3io_access_key"
 
@@ -178,7 +196,7 @@ class DatastoreProfileV3io(DatastoreProfile):
 
 
 class DatastoreProfileS3(DatastoreProfile):
-    type: str = pydantic.Field("s3")
+    type: str = pydantic.v1.Field("s3")
     _private_attributes = ("access_key_id", "secret_key")
     endpoint_url: typing.Optional[str] = None
     force_non_anonymous: typing.Optional[str] = None
@@ -188,7 +206,7 @@ class DatastoreProfileS3(DatastoreProfile):
     secret_key: typing.Optional[str] = None
     bucket: typing.Optional[str] = None
 
-    @pydantic.validator("bucket")
+    @pydantic.v1.validator("bucket")
     @classmethod
     def check_bucket(cls, v):
         if not v:
@@ -226,7 +244,7 @@ class DatastoreProfileS3(DatastoreProfile):
 
 
 class DatastoreProfileRedis(DatastoreProfile):
-    type: str = pydantic.Field("redis")
+    type: str = pydantic.v1.Field("redis")
     _private_attributes = ("username", "password")
     endpoint_url: str
     username: typing.Optional[str] = None
@@ -269,7 +287,7 @@ class DatastoreProfileRedis(DatastoreProfile):
 
 
 class DatastoreProfileDBFS(DatastoreProfile):
-    type: str = pydantic.Field("dbfs")
+    type: str = pydantic.v1.Field("dbfs")
     _private_attributes = ("token",)
     endpoint_url: typing.Optional[str] = None  # host
     token: typing.Optional[str] = None
@@ -287,13 +305,13 @@ class DatastoreProfileDBFS(DatastoreProfile):
 
 
 class DatastoreProfileGCS(DatastoreProfile):
-    type: str = pydantic.Field("gcs")
+    type: str = pydantic.v1.Field("gcs")
     _private_attributes = ("gcp_credentials",)
     credentials_path: typing.Optional[str] = None  # path to file.
     gcp_credentials: typing.Optional[typing.Union[str, dict]] = None
     bucket: typing.Optional[str] = None
 
-    @pydantic.validator("bucket")
+    @pydantic.v1.validator("bucket")
     @classmethod
     def check_bucket(cls, v):
         if not v:
@@ -304,7 +322,7 @@ class DatastoreProfileGCS(DatastoreProfile):
             )
         return v
 
-    @pydantic.validator("gcp_credentials", pre=True, always=True)
+    @pydantic.v1.validator("gcp_credentials", pre=True, always=True)
     @classmethod
     def convert_dict_to_json(cls, v):
         if isinstance(v, dict):
@@ -332,7 +350,7 @@ class DatastoreProfileGCS(DatastoreProfile):
 
 
 class DatastoreProfileAzureBlob(DatastoreProfile):
-    type: str = pydantic.Field("az")
+    type: str = pydantic.v1.Field("az")
     _private_attributes = (
         "connection_string",
         "account_key",
@@ -350,7 +368,7 @@ class DatastoreProfileAzureBlob(DatastoreProfile):
     credential: typing.Optional[str] = None
     container: typing.Optional[str] = None
 
-    @pydantic.validator("container")
+    @pydantic.v1.validator("container")
     @classmethod
     def check_container(cls, v):
         if not v:
@@ -392,7 +410,7 @@ class DatastoreProfileAzureBlob(DatastoreProfile):
 
 
 class DatastoreProfileHdfs(DatastoreProfile):
-    type: str = pydantic.Field("hdfs")
+    type: str = pydantic.v1.Field("hdfs")
     _private_attributes = "token"
     host: typing.Optional[str] = None
     port: typing.Optional[int] = None
@@ -415,7 +433,7 @@ class DatastoreProfileHdfs(DatastoreProfile):
         return f"webhdfs://{self.host}:{self.http_port}{subpath}"
 
 
-class DatastoreProfile2Json(pydantic.BaseModel):
+class DatastoreProfile2Json(pydantic.v1.BaseModel):
     @staticmethod
     def _to_json(attributes):
         # First, base64 encode the values
@@ -476,6 +494,7 @@ class DatastoreProfile2Json(pydantic.BaseModel):
             "gcs": DatastoreProfileGCS,
             "az": DatastoreProfileAzureBlob,
             "hdfs": DatastoreProfileHdfs,
+            "vector": VectorStoreProfile,
         }
         if datastore_type in ds_profile_factory:
             return ds_profile_factory[datastore_type].parse_obj(decoded_dict)

mlrun/datastore/vectorstore.py

@@ -0,0 +1,186 @@
+# Copyright 2024 Iguazio
+#
+# 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
+from importlib import import_module
+from typing import Union
+
+from mlrun.artifacts import DocumentArtifact
+
+
+class VectorStoreCollection:
+    """
+    VectorStoreCollection is a class that manages a collection of vector stores, providing methods to add and delete
+    documents and artifacts, and to interact with an MLRun context.
+
+    Attributes:
+        _collection_impl (object): The underlying collection implementation.
+        _mlrun_context (Union[MlrunProject, MLClientCtx]): The MLRun context associated with the collection.
+        collection_name (str): The name of the collection.
+        id (str): The unique identifier of the collection, composed of the datastore profile and collection name.
+
+    Methods:
+        add_documents(documents: list["Document"], **kwargs):
+            Adds a list of documents to the collection and updates the MLRun artifacts associated with the documents
+            if an MLRun context is present.
+
+        add_artifacts(artifacts: list[DocumentArtifact], splitter=None, **kwargs):
+            Adds a list of DocumentArtifact objects to the collection, optionally using a splitter to convert
+            artifacts to documents.
+
+        remove_itself_from_artifact(artifact: DocumentArtifact):
+            Removes the current object from the given artifact's collection and updates the artifact.
+
+        delete_artifacts(artifacts: list[DocumentArtifact]):
+            Deletes a list of DocumentArtifact objects from the collection and updates the MLRun context.
+            Raises NotImplementedError if the delete operation is not supported for the collection implementation.
+    """
+
+    def __init__(
+        self,
+        vector_store_class: str,
+        mlrun_context: Union["MlrunProject", "MLClientCtx"],  # noqa: F821
+        datastore_profile: str,
+        collection_name: str,
+        **kwargs,
+    ):
+        # Import the vector store class dynamically
+        module_name, class_name = vector_store_class.rsplit(".", 1)
+        module = import_module(module_name)
+        vector_store_class = getattr(module, class_name)
+
+        signature = inspect.signature(vector_store_class)
+
+        # Create the vector store instance
+        if "collection_name" in signature.parameters.keys():
+            vector_store = vector_store_class(collection_name=collection_name, **kwargs)
+        else:
+            vector_store = vector_store_class(**kwargs)
+
+        self._collection_impl = vector_store
+        self._mlrun_context = mlrun_context
+        self.collection_name = collection_name
+        self.id = datastore_profile + "/" + collection_name
+
+    def __getattr__(self, name):
+        # This method is called when an attribute is not found in the usual places
+        # Forward the attribute access to _collection_impl
+        return getattr(self._collection_impl, name)
+
+    def __setattr__(self, name, value):
+        if name in ["_collection_impl", "_mlrun_context"] or name in self.__dict__:
+            # Use the base class method to avoid recursion
+            super().__setattr__(name, value)
+        else:
+            # Forward the attribute setting to _collection_impl
+            setattr(self._collection_impl, name, value)
+
+    def add_documents(
+        self,
+        documents: list["Document"],  # noqa: F821
+        **kwargs,
+    ):
+        """
+        Add a list of documents to the collection.
+
+        If the instance has an MLRun context, it will update the MLRun artifacts
+        associated with the documents.
+
+        Args:
+            documents (list[Document]): A list of Document objects to be added.
+            **kwargs: Additional keyword arguments to be passed to the underlying
+                      collection implementation.
+
+        Returns:
+            The result of the underlying collection implementation's add_documents method.
+        """
+        if self._mlrun_context:
+            for document in documents:
+                mlrun_uri = document.metadata.get(
+                    DocumentArtifact.METADATA_ARTIFACT_URI_KEY
+                )
+                if mlrun_uri:
+                    artifact = self._mlrun_context.get_store_resource(mlrun_uri)
+                    artifact.collection_add(self.id)
+                    self._mlrun_context.update_artifact(artifact)
+        return self._collection_impl.add_documents(documents, **kwargs)
+
+    def add_artifacts(self, artifacts: list[DocumentArtifact], splitter=None, **kwargs):
+        """
+        Add a list of DocumentArtifact objects to the collection.
+
+        Args:
+            artifacts (list[DocumentArtifact]): A list of DocumentArtifact objects to be added.
+            splitter (optional): An optional splitter to be used when converting artifacts to documents.
+            **kwargs: Additional keyword arguments to be passed to the collection's add_documents method.
+
+        Returns:
+            list: A list of IDs of the added documents.
+        """
+        all_ids = []
+        for artifact in artifacts:
+            documents = artifact.to_langchain_documents(splitter)
+            artifact.collection_add(self.id)
+            self._mlrun_context.update_artifact(artifact)
+            ids = self._collection_impl.add_documents(documents, **kwargs)
+            all_ids.extend(ids)
+        return all_ids
+
+    def remove_itself_from_artifact(self, artifact: DocumentArtifact):
+        """
+        Remove the current object from the given artifact's collection and update the artifact.
+
+        Args:
+            artifact (DocumentArtifact): The artifact from which the current object should be removed.
+        """
+        artifact.collection_remove(self.id)
+        self._mlrun_context.update_artifact(artifact)
+
+    def delete_artifacts(self, artifacts: list[DocumentArtifact]):
+        """
+        Delete a list of DocumentArtifact objects from the collection.
+
+        This method removes the specified artifacts from the collection and updates the MLRun context.
+        The deletion process varies depending on the type of the underlying collection implementation.
+
+        Args:
+            artifacts (list[DocumentArtifact]): A list of DocumentArtifact objects to be deleted.
+
+        Raises:
+            NotImplementedError: If the delete operation is not supported for the collection implementation.
+        """
+        store_class = self._collection_impl.__class__.__name__.lower()
+        for artifact in artifacts:
+            artifact.collection_remove(self.id)
+            self._mlrun_context.update_artifact(artifact)
+            if store_class == "milvus":
+                expr = f"{DocumentArtifact.METADATA_SOURCE_KEY} == '{artifact.source}'"
+                return self._collection_impl.delete(expr=expr)
+            elif store_class == "chroma":
+                where = {DocumentArtifact.METADATA_SOURCE_KEY: artifact.source}
+                return self._collection_impl.delete(where=where)
+
+            elif (
+                hasattr(self._collection_impl, "delete")
+                and "filter"
+                in inspect.signature(self._collection_impl.delete).parameters
+            ):
+                filter = {
+                    "metadata": {DocumentArtifact.METADATA_SOURCE_KEY: artifact.source}
+                }
+                return self._collection_impl.delete(filter=filter)
+            else:
+                raise NotImplementedError(
+                    f"delete_artifacts() operation not supported for {store_class}"
+                )