mlrun 1.8.0rc12__py3-none-any.whl → 1.8.0rc13__py3-none-any.whl

mlrun/artifacts/document.py

@@ -89,12 +89,17 @@ class MLRunLoader:
     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.
+        artifact_key (str, optional): The key for the artifact to be logged. Special characters and symbols
+            not valid in artifact names will be encoded as their hexadecimal representation. The '%%' pattern
+            in the key will be replaced by the hex-encoded version of the source path. Defaults to "doc%%".
         local_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
+        producer (Optional[Union[MlrunProject, str, MLClientCtx]], optional): The producer of the document.
+                                If not specified, will try to get the current MLRun context or project.
+                                Defaults to None.
         upload (bool, optional): Flag indicating whether to upload the document.
+        labels (Optional[Dict[str, str]], optional): Key-value labels to attach to the artifact. Defaults to None.
+        tag (str, optional): Version tag for the artifact. Defaults to "".
 
     Returns:
         DynamicDocumentLoader: An instance of a dynamically defined subclass of BaseLoader.
@@ -146,6 +151,8 @@ class MLRunLoader:
         artifact_key="doc%%",
         producer: Optional[Union["MlrunProject", str, "MLClientCtx"]] = None,  # noqa: F821
         upload: bool = False,
+        tag: str = "",
+        labels: Optional[dict[str, str]] = None,
     ):
         # Dynamically import BaseLoader
         from langchain_community.document_loaders.base import BaseLoader
@@ -158,6 +165,8 @@ class MLRunLoader:
                 artifact_key,
                 producer,
                 upload,
+                tag,
+                labels,
             ):
                 self.producer = producer
                 self.artifact_key = (
@@ -168,6 +177,8 @@ class MLRunLoader:
                 self.loader_spec = loader_spec
                 self.local_path = local_path
                 self.upload = upload
+                self.tag = tag
+                self.labels = labels
 
                 # Resolve the producer
                 if not self.producer:
@@ -181,9 +192,11 @@ class MLRunLoader:
                     document_loader_spec=self.loader_spec,
                     local_path=self.local_path,
                     upload=self.upload,
+                    labels=self.labels,
+                    tag=self.tag,
                 )
                 res = artifact.to_langchain_documents()
-                yield res[0]
+                return res
 
         # Return an instance of the dynamically defined subclass
         instance = DynamicDocumentLoader(
@@ -192,6 +205,8 @@ class MLRunLoader:
             loader_spec=loader_spec,
             producer=producer,
             upload=upload,
+            tag=tag,
+            labels=labels,
         )
         return instance
 
@@ -257,6 +272,9 @@ class DocumentArtifact(Artifact):
     METADATA_CHUNK_KEY = "mlrun_chunk"
     METADATA_ARTIFACT_URI_KEY = "mlrun_object_uri"
     METADATA_ARTIFACT_TARGET_PATH_KEY = "mlrun_target_path"
+    METADATA_ARTIFACT_TAG = "mlrun_tag"
+    METADATA_ARTIFACT_KEY = "mlrun_key"
+    METADATA_ARTIFACT_PROJECT = "mlrun_project"
 
     def __init__(
         self,
@@ -331,6 +349,10 @@ class DocumentArtifact(Artifact):
             metadata[self.METADATA_ORIGINAL_SOURCE_KEY] = self.spec.original_source
             metadata[self.METADATA_SOURCE_KEY] = self.get_source()
             metadata[self.METADATA_ARTIFACT_URI_KEY] = self.uri
+            metadata[self.METADATA_ARTIFACT_TAG] = self.tag or "latest"
+            metadata[self.METADATA_ARTIFACT_KEY] = self.key
+            metadata[self.METADATA_ARTIFACT_PROJECT] = self.metadata.project
+
             if self.get_target_path():
                 metadata[self.METADATA_ARTIFACT_TARGET_PATH_KEY] = (
                     self.get_target_path()
@@ -346,7 +368,7 @@ class DocumentArtifact(Artifact):
                 idx = idx + 1
         return results
 
-    def collection_add(self, collection_id: str) -> None:
+    def collection_add(self, collection_id: str) -> bool:
         """
         Add a collection ID to the artifact's collection list.
 
@@ -361,8 +383,10 @@ class DocumentArtifact(Artifact):
         """
         if collection_id not in self.spec.collections:
             self.spec.collections[collection_id] = "1"
+            return True
+        return False
 
-    def collection_remove(self, collection_id: str) -> None:
+    def collection_remove(self, collection_id: str) -> bool:
         """
         Remove a collection ID from the artifact's collection list.
 
@@ -376,3 +400,5 @@ class DocumentArtifact(Artifact):
         """
         if collection_id in self.spec.collections:
             self.spec.collections.pop(collection_id)
+            return True
+        return False

mlrun/common/formatters/artifact.py

@@ -32,7 +32,7 @@ class ArtifactFormat(ObjectFormat, mlrun.common.types.StrEnum):
                 [
                     "kind",
                     "metadata",
-                    "status",
+                    "status.state",
                     "project",
                     "spec.producer",
                     "spec.db_key",

mlrun/common/schemas/partition.py

@@ -46,24 +46,23 @@ class PartitionInterval(StrEnum):
             return timedelta(weeks=1)
 
     @classmethod
-    def from_function(cls, partition_function: str):
+    def from_expression(cls, partition_expression: str):
         """
-        Returns the corresponding PartitionInterval for a given partition function,
+        Returns the corresponding PartitionInterval for a given partition expression,
         or None if the function is not mapped.
 
-        :param partition_function: The partition function to map to an interval.
-        :return: PartitionInterval corresponding to the function, or None if no match is found.
+        :param partition_expression: The partition expression to map to an interval.
+        :return: PartitionInterval corresponding to the expression, or `month` if no match is found.
         """
-        partition_function_to_partitions_interval = {
-            "DAY": "DAY",
-            "DAYOFMONTH": "DAY",
-            "MONTH": "MONTH",
-            "YEARWEEK": "YEARWEEK",
-        }
-        interval = partition_function_to_partitions_interval.get(partition_function)
-        if interval and cls.is_valid(interval):
-            return cls[interval]
-        raise KeyError(f"Partition function: {partition_function} isn't supported")
+
+        # Match the provided function string to the correct interval
+        partition_expression = partition_expression.upper()
+        if "YEARWEEK" in partition_expression:
+            return cls.YEARWEEK
+        elif "DAYOFMONTH" in partition_expression:
+            return cls.DAY
+        else:
+            return cls.MONTH
 
     def get_partition_info(
         self,
@@ -120,11 +119,17 @@ class PartitionInterval(StrEnum):
             year, week, _ = current_datetime.isocalendar()
             return f"{year}{week:02d}"
 
-    def get_partition_expression(self):
+    def get_partition_expression(self, column_name: str):
         if self == PartitionInterval.YEARWEEK:
-            return "YEARWEEK(activation_time, 1)"
-        else:
-            return f"{self}(activation_time)"
+            return f"YEARWEEK({column_name}, 1)"
+        elif self == PartitionInterval.DAY:
+            # generates value in format %Y%m%d in mysql
+            # mysql query example: `select YEAR(NOW())*10000 + MONTH(NOW())*100 + DAY(NOW());`
+            return f"YEAR({column_name}) * 10000 + MONTH({column_name}) * 100 + DAY({column_name})"
+        elif self == PartitionInterval.MONTH:
+            # generates value in format %Y%m in mysql
+            # mysql query example: `select YEAR(NOW())*100 + MONTH(NOW());`
+            return f"YEAR({column_name}) * 100 + MONTH({column_name})"
 
     def get_number_of_partitions(self, days: int) -> int:
         # Calculate the number partitions based on given number of days

mlrun/datastore/vectorstore.py

@@ -19,19 +19,42 @@ from typing import Optional, Union
 from mlrun.artifacts import DocumentArtifact
 
 
-def _extract_collection_name(vectorstore: "VectorStore") -> str:  # noqa: F821
-    # List of possible attribute names for collection name
-    possible_attributes = ["collection_name", "_collection_name"]
+def find_existing_attribute(obj, base_name="name", parent_name="collection"):
+    # Define all possible patterns
+
+    return None
+
 
-    for attr in possible_attributes:
-        if hasattr(vectorstore, attr):
-            collection_name = getattr(vectorstore, attr)
-            if collection_name:
-                return collection_name
+def _extract_collection_name(vectorstore: "VectorStore") -> str:  # noqa: F821
+    patterns = [
+        "collection.name",
+        "collection._name",
+        "_collection.name",
+        "_collection._name",
+        "collection_name",
+        "_collection_name",
+    ]
+
+    def resolve_attribute(obj, pattern):
+        if "." in pattern:
+            parts = pattern.split(".")
+            current = vectorstore
+            for part in parts:
+                if hasattr(current, part):
+                    current = getattr(current, part)
+                else:
+                    return None
+            return current
+        else:
+            return getattr(obj, pattern, None)
 
-    store_class = vectorstore.__class__.__name__.lower()
-    if store_class == "mongodbatlasvectorsearch":
-        return vectorstore.collection.name
+    for pattern in patterns:
+        try:
+            value = resolve_attribute(vectorstore, pattern)
+            if value is not None:
+                return value
+        except (AttributeError, TypeError):
+            continue
 
     # If we get here, we couldn't find a valid collection name
     raise ValueError(
@@ -82,6 +105,19 @@ class VectorStoreCollection:
             # Forward the attribute setting to _collection_impl
             setattr(self._collection_impl, name, value)
 
+    def _get_mlrun_project_name(self):
+        import mlrun
+
+        if self._mlrun_context and isinstance(
+            self._mlrun_context, mlrun.projects.MlrunProject
+        ):
+            return self._mlrun_context.name
+        if self._mlrun_context and isinstance(
+            self._mlrun_context, mlrun.execution.MLClientCtx
+        ):
+            return self._mlrun_context.get_project_object().name
+        return None
+
     def delete(self, *args, **kwargs):
         self._collection_impl.delete(*args, **kwargs)
 
@@ -106,13 +142,22 @@ class VectorStoreCollection:
         """
         if self._mlrun_context:
             for document in documents:
-                mlrun_uri = document.metadata.get(
-                    DocumentArtifact.METADATA_ARTIFACT_URI_KEY
+                mlrun_key = document.metadata.get(
+                    DocumentArtifact.METADATA_ARTIFACT_KEY, None
+                )
+                mlrun_project = document.metadata.get(
+                    DocumentArtifact.METADATA_ARTIFACT_PROJECT, None
                 )
-                if mlrun_uri:
-                    artifact = self._mlrun_context.get_store_resource(mlrun_uri)
-                    artifact.collection_add(self.collection_name)
-                    self._mlrun_context.update_artifact(artifact)
+
+                if mlrun_key and mlrun_project == self._get_mlrun_project_name():
+                    mlrun_tag = document.metadata.get(
+                        DocumentArtifact.METADATA_ARTIFACT_TAG, None
+                    )
+                    artifact = self._mlrun_context.get_artifact(
+                        key=mlrun_key, tag=mlrun_tag
+                    )
+                    if artifact.collection_add(self.collection_name):
+                        self._mlrun_context.update_artifact(artifact)
 
         return self._collection_impl.add_documents(documents, **kwargs)
 
@@ -159,8 +204,7 @@ class VectorStoreCollection:
                 )
         for index, artifact in enumerate(artifacts):
             documents = artifact.to_langchain_documents(splitter)
-            artifact.collection_add(self.collection_name)
-            if self._mlrun_context:
+            if artifact.collection_add(self.collection_name) and self._mlrun_context:
                 self._mlrun_context.update_artifact(artifact)
             if user_ids:
                 num_of_documents = len(documents)
@@ -182,8 +226,8 @@ class VectorStoreCollection:
         Args:
             artifact (DocumentArtifact): The artifact from which the current object should be removed.
         """
-        artifact.collection_remove(self.collection_name)
-        if self._mlrun_context:
+
+        if artifact.collection_remove(self.collection_name) and self._mlrun_context:
             self._mlrun_context.update_artifact(artifact)
 
     def delete_artifacts(self, artifacts: list[DocumentArtifact]):
@@ -201,16 +245,15 @@ class VectorStoreCollection:
         """
         store_class = self._collection_impl.__class__.__name__.lower()
         for artifact in artifacts:
-            artifact.collection_remove(self.collection_name)
-            if self._mlrun_context:
+            if artifact.collection_remove(self.collection_name) and self._mlrun_context:
                 self._mlrun_context.update_artifact(artifact)
 
             if store_class == "milvus":
                 expr = f"{DocumentArtifact.METADATA_SOURCE_KEY} == '{artifact.get_source()}'"
-                return self._collection_impl.delete(expr=expr)
+                self._collection_impl.delete(expr=expr)
             elif store_class == "chroma":
                 where = {DocumentArtifact.METADATA_SOURCE_KEY: artifact.get_source()}
-                return self._collection_impl.delete(where=where)
+                self._collection_impl.delete(where=where)
 
             elif (
                 hasattr(self._collection_impl, "delete")
@@ -222,7 +265,7 @@ class VectorStoreCollection:
                         DocumentArtifact.METADATA_SOURCE_KEY: artifact.get_source()
                     }
                 }
-                return self._collection_impl.delete(filter=filter)
+                self._collection_impl.delete(filter=filter)
             else:
                 raise NotImplementedError(
                     f"delete_artifacts() operation not supported for {store_class}"

mlrun/db/base.py

@@ -23,6 +23,7 @@ import mlrun.common
 import mlrun.common.formatters
 import mlrun.common.runtimes.constants
 import mlrun.common.schemas
+import mlrun.common.schemas.model_monitoring.constants as mm_constants
 import mlrun.common.schemas.model_monitoring.model_endpoints as mm_endpoints
 import mlrun.model_monitoring
 
@@ -58,6 +59,15 @@ class RunDBInterface(ABC):
     def abort_run(self, uid, project="", iter=0, timeout=45, status_text=""):
         pass
 
+    @abstractmethod
+    def push_run_notifications(
+        self,
+        uid,
+        project="",
+        timeout=45,
+    ):
+        pass
+
     @abstractmethod
     def read_run(
         self,
@@ -666,7 +676,9 @@ class RunDBInterface(ABC):
     def create_model_endpoint(
         self,
         model_endpoint: mlrun.common.schemas.ModelEndpoint,
-        creation_strategy: mlrun.common.schemas.ModelEndpointCreationStrategy = "inplace",
+        creation_strategy: Optional[
+            mm_constants.ModelEndpointCreationStrategy
+        ] = mm_constants.ModelEndpointCreationStrategy.INPLACE,
     ) -> mlrun.common.schemas.ModelEndpoint:
         pass
 

mlrun/db/httpdb.py

@@ -35,6 +35,7 @@ import mlrun.common.constants
 import mlrun.common.formatters
 import mlrun.common.runtimes
 import mlrun.common.schemas
+import mlrun.common.schemas.model_monitoring.constants as mm_constants
 import mlrun.common.schemas.model_monitoring.model_endpoints as mm_endpoints
 import mlrun.common.types
 import mlrun.platforms
@@ -755,6 +756,34 @@ class HTTPRunDB(RunDBInterface):
             )
         return None
 
+    def push_run_notifications(
+        self,
+        uid,
+        project="",
+        timeout=45,
+    ):
+        """
+        Push notifications for a run.
+
+        :param uid: Unique ID of the run.
+        :param project: Project that the run belongs to.
+        :returns: :py:class:`~mlrun.common.schemas.BackgroundTask`.
+        """
+        project = project or config.default_project
+
+        response = self.api_call(
+            "POST",
+            path=f"projects/{project}/runs/{uid}/push_notifications",
+            error="Failed push notifications",
+            timeout=timeout,
+        )
+        if response.status_code == http.HTTPStatus.ACCEPTED:
+            background_task = mlrun.common.schemas.BackgroundTask(**response.json())
+            return self._wait_for_background_task_to_reach_terminal_state(
+                background_task.metadata.name, project=project
+            )
+        return None
+
     def read_run(
         self,
         uid,
@@ -3582,18 +3611,24 @@ class HTTPRunDB(RunDBInterface):
     def create_model_endpoint(
         self,
         model_endpoint: mlrun.common.schemas.ModelEndpoint,
-        creation_strategy: mlrun.common.schemas.ModelEndpointCreationStrategy = "inplace",
+        creation_strategy: Optional[
+            mm_constants.ModelEndpointCreationStrategy
+        ] = mm_constants.ModelEndpointCreationStrategy.INPLACE,
     ) -> mlrun.common.schemas.ModelEndpoint:
         """
         Creates a DB record with the given model_endpoint record.
 
         :param model_endpoint: An object representing the model endpoint.
-        :param creation_strategy: model endpoint creation strategy :
-                            * overwrite - Create a new model endpoint and delete the last old one if it exists.
-                            * inplace - Use the existing model endpoint if it already exists (default).
-                            * archive - Preserve the old model endpoint and create a new one,
-                            tagging it as the latest.
-
+        :param creation_strategy: Strategy for creating or updating the model endpoint:
+            * **overwrite**:
+            1. If model endpoints with the same name exist, delete the `latest` one.
+            2. Create a new model endpoint entry and set it as `latest`.
+            * **inplace** (default):
+            1. If model endpoints with the same name exist, update the `latest` entry.
+            2. Otherwise, create a new entry.
+            * **archive**:
+            1. If model endpoints with the same name exist, preserve them.
+            2. Create a new model endpoint with the same name and set it to `latest`.
         :return: The created model endpoint object.
         """
 

mlrun/db/nopdb.py

@@ -20,6 +20,7 @@ import mlrun.alerts
 import mlrun.common.formatters
 import mlrun.common.runtimes.constants
 import mlrun.common.schemas
+import mlrun.common.schemas.model_monitoring.constants as mm_constants
 import mlrun.errors
 import mlrun.lists
 import mlrun.model_monitoring
@@ -75,6 +76,14 @@ class NopDB(RunDBInterface):
     def abort_run(self, uid, project="", iter=0, timeout=45, status_text=""):
         pass
 
+    def push_run_notifications(
+        self,
+        uid,
+        project="",
+        timeout=45,
+    ):
+        pass
+
     def list_runtime_resources(
         self,
         project: Optional[str] = None,
@@ -575,7 +584,9 @@ class NopDB(RunDBInterface):
     def create_model_endpoint(
         self,
         model_endpoint: mlrun.common.schemas.ModelEndpoint,
-        creation_strategy: mlrun.common.schemas.ModelEndpointCreationStrategy = "inplace",
+        creation_strategy: Optional[
+            mm_constants.ModelEndpointCreationStrategy
+        ] = mm_constants.ModelEndpointCreationStrategy.INPLACE,
     ) -> mlrun.common.schemas.ModelEndpoint:
         pass
 

mlrun/execution.py

@@ -880,7 +880,7 @@ class MLClientCtx:
         tag: str = "",
         local_path: str = "",
         artifact_path: Optional[str] = None,
-        document_loader: DocumentLoaderSpec = DocumentLoaderSpec(),
+        document_loader_spec: DocumentLoaderSpec = DocumentLoaderSpec(),
         upload: Optional[bool] = False,
         labels: Optional[dict[str, str]] = None,
         target_path: Optional[str] = None,
@@ -891,22 +891,48 @@ class MLClientCtx:
 
         :param key: Artifact key
         :param tag: Version tag
-        :param local_path:    path to the local file we upload, will also be use
-                              as the destination subpath (under "artifact_path")
-        :param artifact_path:   Target artifact path (when not using the default)
-                                to define a subpath under the default location use:
-                                `artifact_path=context.artifact_subpath('data')`
-        :param document_loader: Spec to use to load the artifact as langchain document
+        :param local_path: path to the local file we upload, will also be use
+                        as the destination subpath (under "artifact_path")
+        :param artifact_path: Target artifact path (when not using the default)
+                            to define a subpath under the default location use:
+                            `artifact_path=context.artifact_subpath('data')`
+        :param document_loader_spec: Spec to use to load the artifact as langchain document.
+
+            By default, uses DocumentLoaderSpec() which initializes with:
+
+            * loader_class_name="langchain_community.document_loaders.TextLoader"
+            * src_name="file_path"
+            * kwargs=None
+
+            Can be customized for different document types, e.g.::
+
+                DocumentLoaderSpec(
+                    loader_class_name="langchain_community.document_loaders.PDFLoader",
+                    src_name="file_path",
+                    kwargs={"extract_images": True}
+                )
         :param upload: Whether to upload the artifact
         :param labels: Key-value labels
         :param target_path: Path to the local file
         :param kwargs: Additional keyword arguments
         :return: DocumentArtifact object
+
+        Example:
+            >>> # Log a PDF document with custom loader
+            >>> project.log_document(
+            ...     key="my_doc",
+            ...     local_path="path/to/doc.pdf",
+            ...     document_loader_spec=DocumentLoaderSpec(
+            ...         loader_class_name="langchain_community.document_loaders.PDFLoader",
+            ...         src_name="file_path",
+            ...         kwargs={"extract_images": True},
+            ...     ),
+            ... )
         """
         doc_artifact = DocumentArtifact(
             key=key,
             original_source=local_path or target_path,
-            document_loader=document_loader,
+            document_loader_spec=document_loader_spec,
             **kwargs,
         )
 
@@ -929,9 +955,15 @@ class MLClientCtx:
         )
         return self.get_artifact(key)
 
-    def get_artifact(self, key: str) -> Artifact:
-        artifact_uri = self._artifacts_manager.artifact_uris[key]
-        return self.get_store_resource(artifact_uri)
+    def get_artifact(
+        self, key, tag=None, iter=None, tree=None, uid=None
+    ) -> Optional[Artifact]:
+        if tag or iter or tree or uid:
+            project = self.get_project_object()
+            return project.get_artifact(key=key, tag=tag, iter=iter, tree=tree, uid=uid)
+        else:
+            artifact_uri = self._artifacts_manager.artifact_uris[key]
+            return self.get_store_resource(artifact_uri)
 
     def update_artifact(self, artifact_object: Artifact):
         """Update an artifact object in the DB and the cached uri"""

mlrun/model_monitoring/applications/_application_steps.py

@@ -136,7 +136,7 @@ class _PrepareMonitoringEvent(StepToDict):
         :param event: Application event.
         :return: Application context.
         """
-        application_context = MonitoringApplicationContext(
+        application_context = MonitoringApplicationContext._from_graph_ctx(
             application_name=self.application_name,
             event=event,
             model_endpoint_dict=self.model_endpoints,

mlrun/model_monitoring/applications/base.py

@@ -112,11 +112,10 @@ class ModelMonitoringApplicationBase(MonitoringApplicationToDict, ABC):
         def call_do_tracking(event: Optional[dict] = None):
             if event is None:
                 event = {}
-            monitoring_context = mm_context.MonitoringApplicationContext(
+            monitoring_context = mm_context.MonitoringApplicationContext._from_ml_ctx(
                 event=event,
                 application_name=self.__class__.__name__,
-                logger=context.logger,
-                artifacts_logger=context,
+                context=context,
                 sample_df=sample_data,
                 feature_stats=feature_stats,
             )