mlrun 1.8.0rc5__py3-none-any.whl → 1.8.0rc6__py3-none-any.whl

mlrun/artifacts/__init__.py

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

mlrun/artifacts/base.py

@@ -36,7 +36,16 @@ from ..utils import (
 
 
 class ArtifactMetadata(ModelObj):
-    _dict_fields = ["key", "project", "iter", "tree", "description", "hash", "tag"]
+    _dict_fields = [
+        "key",
+        "project",
+        "iter",
+        "tree",
+        "description",
+        "hash",
+        "tag",
+        "uid",
+    ]
     _extra_fields = ["updated", "labels"]
 
     def __init__(
@@ -48,6 +57,7 @@ class ArtifactMetadata(ModelObj):
         description=None,
         hash=None,
         tag=None,
+        uid=None,
     ):
         self.key = key
         self.project = project
@@ -58,6 +68,7 @@ class ArtifactMetadata(ModelObj):
         self.labels = {}
         self.updated = None
         self.tag = tag  # temp store of the tag
+        self.uid = uid
 
     def base_dict(self):
         return super().to_dict()

mlrun/artifacts/document.py

@@ -12,7 +12,6 @@
 # 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
@@ -74,14 +73,14 @@ class DocumentLoaderSpec(ModelObj):
         return loader
 
 
-class DocumentLoader:
+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.
-        source_path (str): The source path of the document to be loaded.
+        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
         upload (bool, optional): Flag indicating whether to upload the document.
@@ -104,7 +103,7 @@ class DocumentLoader:
         class DynamicDocumentLoader(BaseLoader):
             def __init__(
                 self,
-                source_path,
+                local_path,
                 loader_spec,
                 artifact_key,
                 producer,
@@ -112,12 +111,12 @@ class DocumentLoader:
             ):
                 self.producer = producer
                 self.artifact_key = (
-                    DocumentLoader.artifact_key_instance(artifact_key, source_path)
+                    MLRunLoader.artifact_key_instance(artifact_key, local_path)
                     if "%%" in artifact_key
                     else artifact_key
                 )
                 self.loader_spec = loader_spec
-                self.source_path = source_path
+                self.local_path = local_path
                 self.upload = upload
 
                 # Resolve the producer
@@ -129,16 +128,17 @@ class DocumentLoader:
             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,
+                    document_loader_spec=self.loader_spec,
+                    local_path=self.local_path,
                     upload=self.upload,
                 )
-                yield artifact.to_langchain_documents()
+                res = artifact.to_langchain_documents()
+                yield res[0]
 
         # Return an instance of the dynamically defined subclass
         instance = DynamicDocumentLoader(
             artifact_key=artifact_key,
-            source_path=source_path,
+            local_path=source_path,
             loader_spec=loader_spec,
             producer=producer,
             upload=upload,
@@ -195,12 +195,15 @@ class DocumentArtifact(Artifact):
         def __init__(
             self,
             *args,
+            document_loader: Optional[DocumentLoaderSpec] = None,
+            collections: Optional[dict] = None,
+            original_source: Optional[str] = None,
             **kwargs,
         ):
             super().__init__(*args, **kwargs)
-            self.document_loader = None
-            self.collections = set()
-            self.original_source = None
+            self.document_loader = document_loader
+            self.collections = collections if collections is not None else {}
+            self.original_source = original_source
 
     """
     A specific artifact class inheriting from generic artifact, used to maintain Document meta-data.
@@ -216,14 +219,17 @@ class DocumentArtifact(Artifact):
 
     def __init__(
         self,
-        key=None,
-        document_loader: DocumentLoaderSpec = DocumentLoaderSpec(),
+        original_source: Optional[str] = None,
+        document_loader_spec: Optional[DocumentLoaderSpec] = None,
         **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"]
+        super().__init__(**kwargs)
+        self.spec.document_loader = (
+            document_loader_spec.to_dict()
+            if document_loader_spec
+            else self.spec.document_loader
+        )
+        self.spec.original_source = original_source or self.spec.original_source
 
     @property
     def spec(self) -> DocumentArtifactSpec:
@@ -234,17 +240,8 @@ class DocumentArtifact(Artifact):
         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):
+    def get_source(self):
         return generate_artifact_uri(self.metadata.project, self.spec.db_key)
 
     def to_langchain_documents(
@@ -262,9 +259,8 @@ class DocumentArtifact(Artifact):
         Returns:
             list[Document]: A list of LangChain Document objects.
         """
-        dictionary = ast.literal_eval(self.spec.document_loader)
-        loader_spec = DocumentLoaderSpec.from_dict(dictionary)
 
+        loader_spec = DocumentLoaderSpec.from_dict(self.spec.document_loader)
         if self.get_target_path():
             with tempfile.NamedTemporaryFile() as tmp_file:
                 mlrun.datastore.store_manager.object(
@@ -272,8 +268,8 @@ class DocumentArtifact(Artifact):
                 ).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)
+        elif self.spec.original_source:
+            loader = loader_spec.make_loader(self.spec.original_source)
             documents = loader.load()
         else:
             raise ValueError(
@@ -289,8 +285,8 @@ class DocumentArtifact(Artifact):
 
             metadata = document.metadata
 
-            metadata[self.METADATA_ORIGINAL_SOURCE_KEY] = self.src_path
-            metadata[self.METADATA_SOURCE_KEY] = self.source
+            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
             if self.get_target_path():
                 metadata[self.METADATA_ARTIFACT_TARGET_PATH_KEY] = (
@@ -301,13 +297,38 @@ class DocumentArtifact(Artifact):
                 metadata[self.METADATA_CHUNK_KEY] = str(idx)
                 doc = Document(
                     page_content=text,
-                    metadata=metadata,
+                    metadata=metadata.copy(),
                 )
                 results.append(doc)
         return results
 
     def collection_add(self, collection_id: str) -> None:
-        self.spec.collections.add(collection_id)
+        """
+        Add a collection ID to the artifact's collection list.
+
+        Adds the specified collection ID to the artifact's collection mapping if it
+        doesn't already exist.
+        This method only modifies the client-side artifact object and does not persist
+        the changes to the MLRun DB. To save the changes permanently, you must call
+        project.update_artifact() after this method.
+
+        Args:
+            collection_id (str): The ID of the collection to add
+        """
+        if collection_id not in self.spec.collections:
+            self.spec.collections[collection_id] = "1"
 
     def collection_remove(self, collection_id: str) -> None:
-        return self.spec.collections.discard(collection_id)
+        """
+        Remove a collection ID from the artifact's collection list.
+
+        Removes the specified collection ID from the artifact's local collection mapping.
+        This method only modifies the client-side artifact object and does not persist
+        the changes to the MLRun DB. To save the changes permanently, you must call
+        project.update_artifact() or context.update_artifact() after this method.
+
+        Args:
+            collection_id (str): The ID of the collection to remove
+        """
+        if collection_id in self.spec.collections:
+            self.spec.collections.pop(collection_id)

mlrun/common/model_monitoring/__init__.py

@@ -11,5 +11,3 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
-from .helpers import create_model_endpoint_uid

mlrun/common/model_monitoring/helpers.py

@@ -17,11 +17,6 @@ import typing
 
 import mlrun.common
 import mlrun.common.schemas.model_monitoring.constants as mm_constants
-from mlrun.common.schemas.model_monitoring import (
-    EndpointUID,
-    FunctionURI,
-    VersionedModel,
-)
 
 FeatureStats = typing.NewType("FeatureStats", dict[str, dict[str, typing.Any]])
 Histogram = typing.NewType("Histogram", list[list])
@@ -31,29 +26,6 @@ BinEdges = typing.NewType("BinEdges", list[float])
 _MAX_FLOAT = sys.float_info.max
 
 
-def create_model_endpoint_uid(function_uri: str, versioned_model: str):
-    function_uri = FunctionURI.from_string(function_uri)
-    versioned_model = VersionedModel.from_string(versioned_model)
-
-    if (
-        not function_uri.project
-        or not function_uri.function
-        or not versioned_model.model
-    ):
-        raise ValueError("Both function_uri and versioned_model have to be initialized")
-
-    uid = EndpointUID(
-        function_uri.project,
-        function_uri.function,
-        function_uri.tag,
-        function_uri.hash_key,
-        versioned_model.model,
-        versioned_model.version,
-    )
-
-    return uid
-
-
 def parse_model_endpoint_project_prefix(path: str, project_name: str):
     return path.split(project_name, 1)[0] + project_name
 

mlrun/common/schemas/__init__.py

@@ -149,10 +149,6 @@ from .model_monitoring import (
     ModelEndpointMetadata,
     ModelEndpointSpec,
     ModelEndpointStatus,
-    ModelEndpointV2,
-    ModelEndpointV2Metadata,
-    ModelEndpointV2Spec,
-    ModelEndpointV2Status,
     ModelMonitoringMode,
     ModelMonitoringStoreKinds,
     MonitoringFunctionNames,
@@ -165,6 +161,7 @@ from .notification import (
     NotificationSeverity,
     NotificationState,
     NotificationStatus,
+    NotificationSummary,
     SetNotificationRequest,
 )
 from .object import ObjectKind, ObjectMetadata, ObjectSpec, ObjectStatus

mlrun/common/schemas/client_spec.py

@@ -57,7 +57,6 @@ class ClientSpec(pydantic.v1.BaseModel):
     redis_url: typing.Optional[str]
     redis_type: typing.Optional[str]
     sql_url: typing.Optional[str]
-    model_endpoint_monitoring_endpoint_store_connection: typing.Optional[str]
     model_monitoring_tsdb_connection: typing.Optional[str]
     ce: typing.Optional[dict]
     # not passing them as one object as it possible client user would like to override only one of the params

mlrun/common/schemas/model_monitoring/__init__.py

@@ -55,12 +55,6 @@ from .grafana import (
     GrafanaTable,
     GrafanaTimeSeriesTarget,
 )
-from .model_endpoint_v2 import (
-    ModelEndpointV2,
-    ModelEndpointV2Metadata,
-    ModelEndpointV2Spec,
-    ModelEndpointV2Status,
-)
 from .model_endpoints import (
     Features,
     FeatureValues,

mlrun/common/schemas/model_monitoring/constants.py

@@ -41,6 +41,7 @@ class ModelEndpointSchema(MonitoringStrEnum):
 
     # spec
     FUNCTION_NAME = "function_name"
+    FUNCTION_TAG = "function_tag"
     FUNCTION_UID = "function_uid"
     MODEL_NAME = "model_name"
     MODEL_TAG = "model_tag"
@@ -48,23 +49,23 @@ class ModelEndpointSchema(MonitoringStrEnum):
     MODEL_UID = "model_uid"
     FEATURE_NAMES = "feature_names"
     LABEL_NAMES = "label_names"
-
-    # status
-    STATE = "state"
-    MONITORING_MODE = "monitoring_mode"
+    FEATURE_STATS = "feature_stats"
     MONITORING_FEATURE_SET_URI = "monitoring_feature_set_uri"
     CHILDREN = "children"
     CHILDREN_UIDS = "children_uids"
-    FIRST_REQUEST = "first_request"
     FUNCTION_URI = "function_uri"
     MODEL_URI = "model_uri"
 
+    # status
+    STATE = "state"
+    MONITORING_MODE = "monitoring_mode"
+    FIRST_REQUEST = "first_request"
+
     # status - operative
     LAST_REQUEST = "last_request"
-    DRIFT_STATUS = "drift_status"
+    RESULT_STATUS = "result_status"
     AVG_LATENCY = "avg_latency"
     ERROR_COUNT = "error_count"
-    FEATURE_STATS = "feature_stats"
     CURRENT_STATS = "current_stats"
     DRIFT_MEASURES = "drift_measures"
 
@@ -80,6 +81,7 @@ class EventFieldType:
     TIMESTAMP = "timestamp"
     # `endpoint_id` is deprecated as a field in the model endpoint schema since 1.3.1, replaced by `uid`.
     ENDPOINT_ID = "endpoint_id"
+    ENDPOINT_NAME = "endpoint_name"
     UID = "uid"
     ENDPOINT_TYPE = "endpoint_type"
     REQUEST_ID = "request_id"
@@ -148,10 +150,12 @@ class ApplicationEvent:
     START_INFER_TIME = "start_infer_time"
     END_INFER_TIME = "end_infer_time"
     ENDPOINT_ID = "endpoint_id"
+    ENDPOINT_NAME = "endpoint_name"
     OUTPUT_STREAM_URI = "output_stream_uri"
 
 
 class WriterEvent(MonitoringStrEnum):
+    ENDPOINT_NAME = "endpoint_name"
     APPLICATION_NAME = "application_name"
     ENDPOINT_ID = "endpoint_id"
     START_INFER_TIME = "start_infer_time"
@@ -222,7 +226,6 @@ class TSDBTarget(MonitoringStrEnum):
 
 
 class ProjectSecretKeys:
-    ENDPOINT_STORE_CONNECTION = "MODEL_MONITORING_ENDPOINT_STORE_CONNECTION"
     ACCESS_KEY = "MODEL_MONITORING_ACCESS_KEY"
     STREAM_PATH = "STREAM_PATH"
     TSDB_CONNECTION = "TSDB_CONNECTION"
@@ -230,7 +233,6 @@ class ProjectSecretKeys:
     @classmethod
     def mandatory_secrets(cls):
         return [
-            cls.ENDPOINT_STORE_CONNECTION,
             cls.STREAM_PATH,
             cls.TSDB_CONNECTION,
         ]