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

mlrun/__init__.py

@@ -31,8 +31,6 @@ from typing import Optional
 
 import dotenv
 
-import mlrun_pipelines
-
 from .config import config as mlconf
 from .datastore import DataItem, store_manager
 from .db import get_run_db
@@ -63,16 +61,16 @@ from .run import (
     new_function,
     wait_for_pipeline_completion,
 )
-from .runtimes import new_model_server
+from .runtimes import mounts, new_model_server
 from .secrets import get_secret_or_env
 from .utils.version import Version
 
 __version__ = Version().get()["version"]
 
-VolumeMount = mlrun_pipelines.common.mounts.VolumeMount
-mount_v3io = mlrun_pipelines.mounts.mount_v3io
-v3io_cred = mlrun_pipelines.mounts.v3io_cred
-auto_mount = mlrun_pipelines.mounts.auto_mount
+VolumeMount = mounts.VolumeMount
+mount_v3io = mounts.mount_v3io
+v3io_cred = mounts.v3io_cred
+auto_mount = mounts.auto_mount
 
 
 def get_version():

mlrun/__main__.py

@@ -33,7 +33,7 @@ import mlrun
 import mlrun.common.constants as mlrun_constants
 import mlrun.common.schemas
 from mlrun.common.helpers import parse_versioned_object_uri
-from mlrun_pipelines.mounts import auto_mount as auto_mount_modifier
+from mlrun.runtimes.mounts import auto_mount as auto_mount_modifier
 
 from .config import config as mlconf
 from .db import get_run_db

mlrun/artifacts/__init__.py

@@ -23,6 +23,7 @@ from .base import (
     get_artifact_meta,
 )
 from .dataset import DatasetArtifact, TableArtifact, update_dataset_meta
+from .document import DocumentArtifact, DocumentLoader, DocumentLoaderSpec
 from .manager import (
     ArtifactManager,
     ArtifactProducer,

mlrun/artifacts/document.py

@@ -0,0 +1,313 @@
+# 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 ast
+import re
+import tempfile
+from collections.abc import Iterator
+from copy import deepcopy
+from importlib import import_module
+from typing import Optional, Union
+
+import mlrun
+from mlrun.artifacts import Artifact, ArtifactSpec
+from mlrun.model import ModelObj
+
+from ..utils import generate_artifact_uri
+
+
+class DocumentLoaderSpec(ModelObj):
+    """
+    A class to load a document from a file path using a specified loader class.
+
+    This class is responsible for loading documents from a given source path using a specified loader class.
+    The loader class is dynamically imported and instantiated with the provided arguments. The loaded documents
+    can be optionally uploaded as artifacts.
+
+    Attributes:
+        loader_class_name (str): The name of the loader class to use for loading documents.
+        src_name (str): The name of the source attribute to pass to the loader class.
+        kwargs (Optional[dict]): Additional keyword arguments to pass to the loader class.
+
+    Methods:
+        make_loader(src_path): Creates an instance of the loader class with the specified source path.
+    """
+
+    _dict_fields = ["loader_class_name", "src_name", "kwargs"]
+
+    def __init__(
+        self,
+        loader_class_name: str = "langchain_community.document_loaders.TextLoader",
+        src_name: str = "file_path",
+        kwargs: Optional[dict] = None,
+    ):
+        """
+        Initialize the document loader.
+
+        Args:
+            loader_class_name (str): The name of the loader class to use.
+            src_name (str): The source name for the document.
+            kwargs (Optional[dict]): Additional keyword arguments to pass to the loader class.
+        """
+        self.loader_class_name = loader_class_name
+        self.src_name = src_name
+        self.kwargs = kwargs
+
+    def make_loader(self, src_path):
+        module_name, class_name = self.loader_class_name.rsplit(".", 1)
+        module = import_module(module_name)
+        loader_class = getattr(module, class_name)
+        kwargs = deepcopy(self.kwargs or {})
+        kwargs[self.src_name] = src_path
+        loader = loader_class(**kwargs)
+        return loader
+
+
+class DocumentLoader:
+    """
+    A factory class for creating instances of a dynamically defined document loader.
+
+    Args:
+        artifact_key (str): The key for the artifact to be logged.It can include '%%' which will be replaced
+        by a hex-encoded version of the source path.
+        source_path (str): The source path of the document to be loaded.
+        loader_spec (DocumentLoaderSpec): Specification for the document loader.
+        producer (Optional[Union[MlrunProject, str, MLClientCtx]], optional): The producer of the document
+        upload (bool, optional): Flag indicating whether to upload the document.
+
+    Returns:
+        DynamicDocumentLoader: An instance of a dynamically defined subclass of BaseLoader.
+    """
+
+    def __new__(
+        cls,
+        source_path: str,
+        loader_spec: "DocumentLoaderSpec",
+        artifact_key="doc%%",
+        producer: Optional[Union["MlrunProject", str, "MLClientCtx"]] = None,  # noqa: F821
+        upload: bool = False,
+    ):
+        # Dynamically import BaseLoader
+        from langchain_community.document_loaders.base import BaseLoader
+
+        class DynamicDocumentLoader(BaseLoader):
+            def __init__(
+                self,
+                source_path,
+                loader_spec,
+                artifact_key,
+                producer,
+                upload,
+            ):
+                self.producer = producer
+                self.artifact_key = (
+                    DocumentLoader.artifact_key_instance(artifact_key, source_path)
+                    if "%%" in artifact_key
+                    else artifact_key
+                )
+                self.loader_spec = loader_spec
+                self.source_path = source_path
+                self.upload = upload
+
+                # Resolve the producer
+                if not self.producer:
+                    self.producer = mlrun.mlconf.default_project
+                if isinstance(self.producer, str):
+                    self.producer = mlrun.get_or_create_project(self.producer)
+
+            def lazy_load(self) -> Iterator["Document"]:  # noqa: F821
+                artifact = self.producer.log_document(
+                    key=self.artifact_key,
+                    document_loader=self.loader_spec,
+                    src_path=self.source_path,
+                    upload=self.upload,
+                )
+                yield artifact.to_langchain_documents()
+
+        # Return an instance of the dynamically defined subclass
+        instance = DynamicDocumentLoader(
+            artifact_key=artifact_key,
+            source_path=source_path,
+            loader_spec=loader_spec,
+            producer=producer,
+            upload=upload,
+        )
+        return instance
+
+    @staticmethod
+    def artifact_key_instance(artifact_key: str, src_path: str) -> str:
+        if "%%" in artifact_key:
+            pattern = mlrun.utils.regex.artifact_key[0]
+            # Convert anchored pattern (^...$) to non-anchored version for finditer
+            search_pattern = pattern.strip("^$")
+            result = []
+            current_pos = 0
+
+            # Find all valid sequences
+            for match in re.finditer(search_pattern, src_path):
+                # Add hex values for characters between matches
+                for char in src_path[current_pos : match.start()]:
+                    result.append(hex(ord(char))[2:].zfill(2))
+
+                # Add the valid sequence
+                result.append(match.group())
+                current_pos = match.end()
+
+            # Handle any remaining characters after the last match
+            for char in src_path[current_pos:]:
+                result.append(hex(ord(char))[2:].zfill(2))
+
+            resolved_path = "".join(result)
+
+            artifact_key = artifact_key.replace("%%", resolved_path)
+
+        return artifact_key
+
+
+class DocumentArtifact(Artifact):
+    """
+    A specific artifact class inheriting from generic artifact, used to maintain Document meta-data.
+
+    Methods:
+        to_langchain_documents(splitter): Create LC documents from the artifact.
+        collection_add(collection_id): Add a collection ID to the artifact.
+        collection_remove(collection_id): Remove a collection ID from the artifact.
+    """
+
+    class DocumentArtifactSpec(ArtifactSpec):
+        _dict_fields = ArtifactSpec._dict_fields + [
+            "document_loader",
+            "collections",
+            "original_source",
+        ]
+
+        def __init__(
+            self,
+            *args,
+            **kwargs,
+        ):
+            super().__init__(*args, **kwargs)
+            self.document_loader = None
+            self.collections = set()
+            self.original_source = None
+
+    """
+    A specific artifact class inheriting from generic artifact, used to maintain Document meta-data.
+    """
+
+    kind = "document"
+
+    METADATA_SOURCE_KEY = "source"
+    METADATA_ORIGINAL_SOURCE_KEY = "original_source"
+    METADATA_CHUNK_KEY = "mlrun_chunk"
+    METADATA_ARTIFACT_URI_KEY = "mlrun_object_uri"
+    METADATA_ARTIFACT_TARGET_PATH_KEY = "mlrun_target_path"
+
+    def __init__(
+        self,
+        key=None,
+        document_loader: DocumentLoaderSpec = DocumentLoaderSpec(),
+        **kwargs,
+    ):
+        super().__init__(key, **kwargs)
+        self.spec.document_loader = document_loader.to_str()
+        if "src_path" in kwargs:
+            self.spec.original_source = kwargs["src_path"]
+
+    @property
+    def spec(self) -> DocumentArtifactSpec:
+        return self._spec
+
+    @spec.setter
+    def spec(self, spec):
+        self._spec = self._verify_dict(
+            spec, "spec", DocumentArtifact.DocumentArtifactSpec
+        )
+        # _verify_dict doesn't handle set, so we need to convert it back
+        if isinstance(self._spec.collections, str):
+            self._spec.collections = ast.literal_eval(self._spec.collections)
+
+    @property
+    def inputs(self):
+        # To keep the interface consistent with the project.update_artifact() when we update the artifact
+        return None
+
+    @property
+    def source(self):
+        return generate_artifact_uri(self.metadata.project, self.spec.db_key)
+
+    def to_langchain_documents(
+        self,
+        splitter: Optional["TextSplitter"] = None,  # noqa: F821
+    ) -> list["Document"]:  # noqa: F821
+        from langchain.schema import Document
+
+        """
+        Create LC documents from the artifact
+
+        Args:
+            splitter (Optional[TextSplitter]): A LangChain TextSplitter to split the document into chunks.
+
+        Returns:
+            list[Document]: A list of LangChain Document objects.
+        """
+        dictionary = ast.literal_eval(self.spec.document_loader)
+        loader_spec = DocumentLoaderSpec.from_dict(dictionary)
+
+        if self.get_target_path():
+            with tempfile.NamedTemporaryFile() as tmp_file:
+                mlrun.datastore.store_manager.object(
+                    url=self.get_target_path()
+                ).download(tmp_file.name)
+                loader = loader_spec.make_loader(tmp_file.name)
+                documents = loader.load()
+        elif self.src_path:
+            loader = loader_spec.make_loader(self.src_path)
+            documents = loader.load()
+        else:
+            raise ValueError(
+                "No src_path or target_path provided. Cannot load document."
+            )
+
+        results = []
+        for document in documents:
+            if splitter:
+                texts = splitter.split_text(document.page_content)
+            else:
+                texts = [document.page_content]
+
+            metadata = document.metadata
+
+            metadata[self.METADATA_ORIGINAL_SOURCE_KEY] = self.src_path
+            metadata[self.METADATA_SOURCE_KEY] = self.source
+            metadata[self.METADATA_ARTIFACT_URI_KEY] = self.uri
+            if self.get_target_path():
+                metadata[self.METADATA_ARTIFACT_TARGET_PATH_KEY] = (
+                    self.get_target_path()
+                )
+
+            for idx, text in enumerate(texts):
+                metadata[self.METADATA_CHUNK_KEY] = str(idx)
+                doc = Document(
+                    page_content=text,
+                    metadata=metadata,
+                )
+                results.append(doc)
+        return results
+
+    def collection_add(self, collection_id: str) -> None:
+        self.spec.collections.add(collection_id)
+
+    def collection_remove(self, collection_id: str) -> None:
+        return self.spec.collections.discard(collection_id)

mlrun/artifacts/manager.py

@@ -41,6 +41,7 @@ from .dataset import (
     DatasetArtifact,
     TableArtifact,
 )
+from .document import DocumentArtifact
 from .model import ModelArtifact
 from .plots import (
     PlotArtifact,
@@ -57,6 +58,7 @@ artifact_types = {
     "model": ModelArtifact,
     "dataset": DatasetArtifact,
     "plotly": PlotlyArtifact,
+    "document": DocumentArtifact,
 }
 
 

mlrun/common/formatters/project.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 #
 
+import datetime
 import typing
 
 import mlrun.common.schemas
@@ -30,11 +31,18 @@ class ProjectFormat(ObjectFormat, mlrun.common.types.StrEnum):
     # internal - allowed only in follower mode, only for the leader for upgrade purposes
     leader = "leader"
 
+    name_and_creation_time = "name_and_creation_time"
+
     @staticmethod
     def format_method(_format: str) -> typing.Optional[typing.Callable]:
         def _name_only(project: mlrun.common.schemas.Project) -> str:
             return project.metadata.name
 
+        def _name_and_creation_time(
+            project: mlrun.common.schemas.Project,
+        ) -> tuple[str, datetime.datetime]:
+            return project.metadata.name, project.metadata.created
+
         def _minimal(
             project: mlrun.common.schemas.Project,
         ) -> mlrun.common.schemas.Project:
@@ -48,4 +56,5 @@ class ProjectFormat(ObjectFormat, mlrun.common.types.StrEnum):
             ProjectFormat.name_only: _name_only,
             ProjectFormat.minimal: _minimal,
             ProjectFormat.leader: None,
+            ProjectFormat.name_and_creation_time: _name_and_creation_time,
         }[_format]

mlrun/common/schemas/__init__.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 from .alert import (
+    AlertActivation,
     AlertActiveState,
     AlertConfig,
     AlertNotification,
@@ -61,6 +62,7 @@ from .clusterization_spec import (
 from .common import ImageBuilder
 from .constants import (
     APIStates,
+    ArtifactPartitionByField,
     ClusterizationRole,
     DeletionStrategy,
     FeatureStorePartitionByField,
@@ -161,11 +163,13 @@ from .notification import (
     Notification,
     NotificationKind,
     NotificationSeverity,
+    NotificationState,
     NotificationStatus,
     SetNotificationRequest,
 )
 from .object import ObjectKind, ObjectMetadata, ObjectSpec, ObjectStatus
 from .pagination import PaginationInfo
+from .partition import PartitionInterval
 from .pipeline import PipelinesOutput, PipelinesPagination
 from .project import (
     IguazioProject,

mlrun/common/schemas/alert.py

@@ -15,9 +15,9 @@
 from datetime import datetime
 from typing import Annotated, Optional, Union
 
-import pydantic
+import pydantic.v1
 
-from mlrun.common.schemas.notification import Notification
+import mlrun.common.schemas.notification as notification_objects
 from mlrun.common.types import StrEnum
 
 
@@ -27,10 +27,10 @@ class EventEntityKind(StrEnum):
     JOB = "job"
 
 
-class EventEntities(pydantic.BaseModel):
+class EventEntities(pydantic.v1.BaseModel):
     kind: EventEntityKind
     project: str
-    ids: pydantic.conlist(str, min_items=1, max_items=1)
+    ids: pydantic.v1.conlist(str, min_items=1, max_items=1)
 
 
 class EventKind(StrEnum):
@@ -64,11 +64,11 @@ _event_kind_entity_map = {
 }
 
 
-class Event(pydantic.BaseModel):
+class Event(pydantic.v1.BaseModel):
     kind: EventKind
     timestamp: Union[str, datetime] = None  # occurrence time
     entity: EventEntities
-    value_dict: Optional[dict] = pydantic.Field(default_factory=dict)
+    value_dict: Optional[dict] = pydantic.v1.Field(default_factory=dict)
 
     def is_valid(self):
         return self.entity.kind in _event_kind_entity_map[self.kind]
@@ -86,7 +86,7 @@ class AlertSeverity(StrEnum):
 
 
 # what should trigger the alert. must be either event (at least 1), or prometheus query
-class AlertTrigger(pydantic.BaseModel):
+class AlertTrigger(pydantic.v1.BaseModel):
     events: list[EventKind] = []
     prometheus_alert: str = None
 
@@ -97,16 +97,16 @@ class AlertTrigger(pydantic.BaseModel):
         )
 
 
-class AlertCriteria(pydantic.BaseModel):
+class AlertCriteria(pydantic.v1.BaseModel):
     count: Annotated[
         int,
-        pydantic.Field(
+        pydantic.v1.Field(
             description="Number of events to wait until notification is sent"
         ),
     ] = 1
     period: Annotated[
         str,
-        pydantic.Field(
+        pydantic.v1.Field(
             description="Time period during which event occurred. e.g. 1d, 3h, 5m, 15s"
         ),
     ] = None
@@ -120,11 +120,11 @@ class ResetPolicy(StrEnum):
     AUTO = "auto"
 
 
-class AlertNotification(pydantic.BaseModel):
-    notification: Notification
+class AlertNotification(pydantic.v1.BaseModel):
+    notification: notification_objects.Notification
     cooldown_period: Annotated[
         str,
-        pydantic.Field(
+        pydantic.v1.Field(
             description="Period during which notifications "
             "will not be sent after initial send. The format of this would be in time."
             " e.g. 1d, 3h, 5m, 15s"
@@ -132,14 +132,14 @@ class AlertNotification(pydantic.BaseModel):
     ] = None
 
 
-class AlertConfig(pydantic.BaseModel):
+class AlertConfig(pydantic.v1.BaseModel):
     project: str
     id: int = None
     name: str
     description: Optional[str] = ""
     summary: Annotated[
         str,
-        pydantic.Field(
+        pydantic.v1.Field(
             description=(
                 "String to be sent in the notifications generated."
                 "e.g. 'Model {{project}}/{{entity}} is drifting.'"
@@ -153,11 +153,11 @@ class AlertConfig(pydantic.BaseModel):
     trigger: AlertTrigger
     criteria: Optional[AlertCriteria]
     reset_policy: ResetPolicy = ResetPolicy.AUTO
-    notifications: pydantic.conlist(AlertNotification, min_items=1)
+    notifications: pydantic.v1.conlist(AlertNotification, min_items=1)
     state: AlertActiveState = AlertActiveState.INACTIVE
     count: Optional[int] = 0
 
-    def get_raw_notifications(self) -> list[Notification]:
+    def get_raw_notifications(self) -> list[notification_objects.Notification]:
         return [
             alert_notification.notification for alert_notification in self.notifications
         ]
@@ -169,7 +169,7 @@ class AlertsModes(StrEnum):
 
 
 class AlertTemplate(
-    pydantic.BaseModel
+    pydantic.v1.BaseModel
 ):  # Template fields that are not shared with created configs
     template_id: int = None
     template_name: str
@@ -200,3 +200,16 @@ class AlertTemplate(
             or self.reset_policy != other.reset_policy
             or self.criteria != other.criteria
         )
+
+
+class AlertActivation(pydantic.v1.BaseModel):
+    name: str
+    project: str
+    severity: AlertSeverity
+    activation_time: datetime
+    entity_id: str
+    entity_kind: EventEntityKind
+    criteria: AlertCriteria
+    event_kind: EventKind
+    number_of_events: int
+    notifications: list[notification_objects.NotificationState]

mlrun/common/schemas/api_gateway.py

@@ -15,7 +15,7 @@
 import typing
 from typing import Optional
 
-import pydantic
+import pydantic.v1
 
 import mlrun.common.constants as mlrun_constants
 import mlrun.common.types
@@ -49,9 +49,9 @@ class APIGatewayState(mlrun.common.types.StrEnum):
     waiting_for_provisioning = "waitingForProvisioning"
 
 
-class _APIGatewayBaseModel(pydantic.BaseModel):
+class _APIGatewayBaseModel(pydantic.v1.BaseModel):
     class Config:
-        extra = pydantic.Extra.allow
+        extra = pydantic.v1.Extra.allow
 
 
 class APIGatewayMetadata(_APIGatewayBaseModel):

mlrun/common/schemas/artifact.py

@@ -14,7 +14,7 @@
 #
 import typing
 
-import pydantic
+import pydantic.v1
 from deprecated import deprecated
 
 import mlrun.common.types
@@ -48,7 +48,7 @@ class ArtifactCategories(mlrun.common.types.StrEnum):
             )
 
 
-class ArtifactIdentifier(pydantic.BaseModel):
+class ArtifactIdentifier(pydantic.v1.BaseModel):
     # artifact kind
     kind: typing.Optional[str]
     key: typing.Optional[str]
@@ -69,7 +69,7 @@ class ArtifactsFormat(mlrun.common.types.StrEnum):
     full = "full"
 
 
-class ArtifactMetadata(pydantic.BaseModel):
+class ArtifactMetadata(pydantic.v1.BaseModel):
     key: str
     project: str
     iter: typing.Optional[int]
@@ -77,10 +77,10 @@ class ArtifactMetadata(pydantic.BaseModel):
     tag: typing.Optional[str]
 
     class Config:
-        extra = pydantic.Extra.allow
+        extra = pydantic.v1.Extra.allow
 
 
-class ArtifactSpec(pydantic.BaseModel):
+class ArtifactSpec(pydantic.v1.BaseModel):
     src_path: typing.Optional[str]
     target_path: typing.Optional[str]
     viewer: typing.Optional[str]
@@ -91,10 +91,10 @@ class ArtifactSpec(pydantic.BaseModel):
     unpackaging_instructions: typing.Optional[dict[str, typing.Any]]
 
     class Config:
-        extra = pydantic.Extra.allow
+        extra = pydantic.v1.Extra.allow
 
 
-class Artifact(pydantic.BaseModel):
+class Artifact(pydantic.v1.BaseModel):
     kind: str
     metadata: ArtifactMetadata
     spec: ArtifactSpec