mlrun 1.7.1rc10__py3-none-any.whl → 1.8.0rc11__py3-none-any.whl

mlrun/artifacts/document.py

@@ -0,0 +1,378 @@
+# 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 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.
+
+    """
+
+    _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.
+
+        Example:
+            >>> # Create a loader specification for PDF documents
+            >>> loader_spec = DocumentLoaderSpec(
+            ...     loader_class_name="langchain_community.document_loaders.PDFLoader",
+            ...     src_name="file_path",
+            ...     kwargs={"extract_images": True},
+            ... )
+            >>> # Create a loader instance for a specific PDF file
+            >>> pdf_loader = loader_spec.make_loader("/path/to/document.pdf")
+            >>> # Load the documents
+            >>> documents = pdf_loader.load()
+
+        """
+        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 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.
+        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.
+
+    Returns:
+        DynamicDocumentLoader: An instance of a dynamically defined subclass of BaseLoader.
+
+    Example:
+        >>> # Create a document loader specification
+        >>> loader_spec = DocumentLoaderSpec(
+        ...     loader_class_name="langchain_community.document_loaders.TextLoader",
+        ...     src_name="file_path",
+        ... )
+        >>> # Create a basic loader for a single file
+        >>> loader = MLRunLoader(
+        ...     source_path="/path/to/document.txt",
+        ...     loader_spec=loader_spec,
+        ...     artifact_key="my_doc",
+        ...     producer=project,
+        ...     upload=True,
+        ... )
+        >>> documents = loader.load()
+        >>> # Create a loader with auto-generated keys
+        >>> loader = MLRunLoader(
+        ...     source_path="/path/to/document.txt",
+        ...     loader_spec=loader_spec,
+        ...     artifact_key="doc%%",  # %% will be replaced with encoded path
+        ...     producer=project,
+        ... )
+        >>> documents = loader.load()
+        >>> # Use with DirectoryLoader
+        >>> from langchain_community.document_loaders import DirectoryLoader
+        >>> dir_loader = DirectoryLoader(
+        ...     "/path/to/directory",
+        ...     glob="**/*.txt",
+        ...     loader_cls=MLRunLoader,
+        ...     loader_kwargs={
+        ...         "loader_spec": loader_spec,
+        ...         "artifact_key": "doc%%",
+        ...         "producer": project,
+        ...         "upload": True,
+        ...     },
+        ... )
+        >>> documents = dir_loader.load()
+
+    """
+
+    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,
+                local_path,
+                loader_spec,
+                artifact_key,
+                producer,
+                upload,
+            ):
+                self.producer = producer
+                self.artifact_key = (
+                    MLRunLoader.artifact_key_instance(artifact_key, local_path)
+                    if "%%" in artifact_key
+                    else artifact_key
+                )
+                self.loader_spec = loader_spec
+                self.local_path = local_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_spec=self.loader_spec,
+                    local_path=self.local_path,
+                    upload=self.upload,
+                )
+                res = artifact.to_langchain_documents()
+                yield res[0]
+
+        # Return an instance of the dynamically defined subclass
+        instance = DynamicDocumentLoader(
+            artifact_key=artifact_key,
+            local_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.
+    """
+
+    class DocumentArtifactSpec(ArtifactSpec):
+        _dict_fields = ArtifactSpec._dict_fields + [
+            "document_loader",
+            "collections",
+            "original_source",
+        ]
+
+        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 = document_loader
+            self.collections = collections if collections is not None else {}
+            self.original_source = original_source
+
+    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,
+        original_source: Optional[str] = None,
+        document_loader_spec: Optional[DocumentLoaderSpec] = None,
+        **kwargs,
+    ):
+        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:
+        return self._spec
+
+    @spec.setter
+    def spec(self, spec):
+        self._spec = self._verify_dict(
+            spec, "spec", DocumentArtifact.DocumentArtifactSpec
+        )
+
+    def get_source(self):
+        """Get the source URI for this artifact."""
+        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.
+        """
+
+        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(
+                    url=self.get_target_path()
+                ).download(tmp_file.name)
+                loader = loader_spec.make_loader(tmp_file.name)
+                documents = loader.load()
+        elif self.spec.original_source:
+            loader = loader_spec.make_loader(self.spec.original_source)
+            documents = loader.load()
+        else:
+            raise ValueError(
+                "No src_path or target_path provided. Cannot load document."
+            )
+
+        results = []
+        idx = 0
+        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.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] = (
+                    self.get_target_path()
+                )
+
+            for text in texts:
+                metadata[self.METADATA_CHUNK_KEY] = str(idx)
+                doc = Document(
+                    page_content=text,
+                    metadata=metadata.copy(),
+                )
+                results.append(doc)
+                idx = idx + 1
+        return results
+
+    def collection_add(self, collection_id: str) -> None:
+        """
+        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:
+        """
+        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/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,
 }
 
 
@@ -124,7 +126,7 @@ class ArtifactManager:
 
         self.artifact_db = db
         self.input_artifacts = {}
-        self.artifacts = {}
+        self.artifact_uris = {}
 
     @staticmethod
     def ensure_artifact_source_file_exists(item, path, body):
@@ -156,14 +158,12 @@ class ArtifactManager:
 
     def artifact_list(self, full=False):
         artifacts = []
-        for artifact in self.artifacts.values():
-            if isinstance(artifact, dict):
-                artifacts.append(artifact)
+        for artifacts_uri in self.artifact_uris.values():
+            artifact: Artifact = mlrun.datastore.get_store_resource(artifacts_uri)
+            if full:
+                artifacts.append(artifact.to_dict())
             else:
-                if full:
-                    artifacts.append(artifact.to_dict())
-                else:
-                    artifacts.append(artifact.base_dict())
+                artifacts.append(artifact.base_dict())
         return artifacts
 
     def log_artifact(
@@ -246,6 +246,8 @@ class ArtifactManager:
                 # otherwise, we do not want to override it.
                 # this is mainly relevant for imported artifacts that have an explicit db_key value already set
                 db_key = item.db_key or key
+        if db_key != key:
+            validate_artifact_key_name(db_key, "artifact.db_key")
         item.db_key = db_key or ""
         item.viewer = viewer or item.viewer
         item.tree = producer.tag
@@ -304,7 +306,6 @@ class ArtifactManager:
         item.target_path = target_path
 
         item.before_log()
-        self.artifacts[key] = item
 
         if ((upload is None and item.kind != "dir") or upload) and not item.is_inline():
             # before uploading the item, we want to ensure that its tags are valid,
@@ -313,32 +314,38 @@ class ArtifactManager:
             item.upload(artifact_path=artifact_path)
 
         if db_key:
-            self._log_to_db(db_key, project, producer.inputs, item)
+            artifact_uid = self._log_to_db(db_key, project, producer.inputs, item)
+            if artifact_uid is not None:
+                item.uid = artifact_uid
+        # Generate the artifact URI after logging to the database and retrieving the artifact UID, if available.
+        self.artifact_uris[key] = item.uri
+
         size = str(item.size) or "?"
         db_str = "Y" if (self.artifact_db and db_key) else "N"
         logger.debug(
-            f"log artifact {key} at {item.target_path}, size: {size}, db: {db_str}"
+            f"Log artifact {key} at {item.target_path}, size: {size}, db: {db_str}"
         )
         return item
 
-    def update_artifact(self, producer, item):
-        self.artifacts[item.key] = item
+    def update_artifact(self, producer, item: Artifact):
+        self.artifact_uris[item.key] = item.uri
         self._log_to_db(item.db_key, producer.project, producer.inputs, item)
 
-    def _log_to_db(self, key, project, sources, item, tag=None):
+    def _log_to_db(self, key, project, sources, item, tag=None) -> typing.Optional[str]:
         """
         log artifact to db
         :param key: Identifying key of the artifact.
         :param project: Project that the artifact belongs to.
-        :param sources: List of artifact sources ( Mainly passed from the producer.items ).
+        :param sources: List of artifact sources ( Mainly passed from the `producer.items` ).
         :param item: The actual artifact to store.
         :param tag: The name of the Tag of the artifact.
+        :return: The logged artifact uid.
         """
         if self.artifact_db:
             item.updated = None
             if sources:
                 item.sources = [{"name": k, "path": str(v)} for k, v in sources.items()]
-            self.artifact_db.store_artifact(
+            artifact_item = self.artifact_db.store_artifact(
                 key,
                 item.to_dict(),
                 iter=item.iter,
@@ -346,6 +353,8 @@ class ArtifactManager:
                 project=project,
                 tree=item.tree,
             )
+            if artifact_item:
+                return artifact_item.get("metadata", {}).get("uid")
 
     def link_artifact(
         self,
@@ -387,7 +396,7 @@ class ArtifactManager:
         deletion_strategy: mlrun.common.schemas.artifact.ArtifactsDeletionStrategies = (
             mlrun.common.schemas.artifact.ArtifactsDeletionStrategies.metadata_only
         ),
-        secrets: dict = None,
+        secrets: typing.Optional[dict] = None,
     ):
         self.artifact_db.del_artifact(
             key=item.db_key,

mlrun/artifacts/model.py

@@ -303,7 +303,7 @@ class ModelArtifact(Artifact):
             self.metadata.labels = self.metadata.labels or {}
             self.metadata.labels["framework"] = self.spec.framework
 
-    def upload(self, artifact_path: str = None):
+    def upload(self, artifact_path: Optional[str] = None):
         """
         internal, upload to target store
         :param artifact_path: required only for when generating target_path from artifact hash
@@ -324,9 +324,7 @@ class ModelArtifact(Artifact):
             artifact=self, extra_data=self.spec.extra_data, artifact_path=artifact_path
         )
 
-        # the model spec yaml should not include the tag, as the same model can be used with different tags,
-        # and the tag is not part of the model spec but the metadata of the model artifact
-        spec_body = _remove_tag_from_spec_yaml(self)
+        spec_body = _sanitize_and_serialize_model_spec_yaml(self)
         spec_target_path = None
 
         if mlrun.mlconf.artifacts.generate_target_path_from_artifact_hash:
@@ -355,7 +353,7 @@ class ModelArtifact(Artifact):
     def _upload_body_or_file(
         self,
         artifact_path: str,
-        target_model_path: str = None,
+        target_model_path: Optional[str] = None,
     ):
         body = self.spec.get_body()
         if body:
@@ -403,12 +401,6 @@ class ModelArtifact(Artifact):
         return mlrun.get_dataitem(target_model_path).get()
 
 
-def _get_src_path(model_spec: ModelArtifact, filename):
-    if model_spec.src_path:
-        return path.join(model_spec.src_path, filename)
-    return filename
-
-
 def get_model(model_dir, suffix=""):
     """return model file, model spec object, and list of extra data items
 
@@ -483,49 +475,20 @@ def get_model(model_dir, suffix=""):
     return temp_path, model_spec, extra_dataitems
 
 
-def _load_model_spec(spec_path):
-    data = mlrun.datastore.store_manager.object(url=spec_path).get()
-    spec = yaml.load(data, Loader=yaml.FullLoader)
-    return ModelArtifact.from_dict(spec)
-
-
-def _get_file_path(base_path: str, name: str, isdir=False):
-    if not is_relative_path(name):
-        return name
-    if not isdir:
-        base_path = path.dirname(base_path)
-    return path.join(base_path, name).replace("\\", "/")
-
-
-def _get_extra(target, extra_data, is_dir=False):
-    extra_dataitems = {}
-    for k, v in extra_data.items():
-        extra_dataitems[k] = mlrun.datastore.store_manager.object(
-            url=_get_file_path(target, v, isdir=is_dir), key=k
-        )
-    return extra_dataitems
-
-
-def _remove_tag_from_spec_yaml(model_spec):
-    spec_dict = model_spec.to_dict()
-    spec_dict["metadata"].pop("tag", None)
-    return yaml.safe_dump(spec_dict)
-
-
 def update_model(
     model_artifact,
-    parameters: dict = None,
-    metrics: dict = None,
-    extra_data: dict = None,
-    inputs: list[Feature] = None,
-    outputs: list[Feature] = None,
-    feature_vector: str = None,
-    feature_weights: list = None,
+    parameters: Optional[dict] = None,
+    metrics: Optional[dict] = None,
+    extra_data: Optional[dict] = None,
+    inputs: Optional[list[Feature]] = None,
+    outputs: Optional[list[Feature]] = None,
+    feature_vector: Optional[str] = None,
+    feature_weights: Optional[list] = None,
     key_prefix: str = "",
-    labels: dict = None,
+    labels: Optional[dict] = None,
     write_spec_copy=True,
     store_object: bool = True,
-):
+) -> ModelArtifact:
     """Update model object attributes
 
     this method will edit or add attributes to a model object
@@ -593,10 +556,7 @@ def update_model(
 
     if write_spec_copy:
         spec_path = path.join(model_spec.target_path, model_spec_filename)
-
-        # the model spec yaml should not include the tag, as the same model can be used with different tags,
-        # and the tag is not part of the model spec but the metadata of the model artifact
-        model_spec_yaml = _remove_tag_from_spec_yaml(model_spec)
+        model_spec_yaml = _sanitize_and_serialize_model_spec_yaml(model_spec)
         mlrun.datastore.store_manager.object(url=spec_path).put(model_spec_yaml)
 
     model_spec.db_key = model_spec.db_key or model_spec.key
@@ -609,3 +569,56 @@ def update_model(
             project=model_spec.project,
         )
     return model_spec
+
+
+def _get_src_path(model_spec: ModelArtifact, filename: str) -> str:
+    return path.join(model_spec.src_path, filename) if model_spec.src_path else filename
+
+
+def _load_model_spec(spec_path) -> ModelArtifact:
+    data = mlrun.datastore.store_manager.object(url=spec_path).get()
+    spec = yaml.load(data, Loader=yaml.FullLoader)
+    return ModelArtifact.from_dict(spec)
+
+
+def _get_file_path(base_path: str, name: str, isdir: bool = False) -> str:
+    if not is_relative_path(name):
+        return name
+    if not isdir:
+        base_path = path.dirname(base_path)
+    return path.join(base_path, name).replace("\\", "/")
+
+
+def _get_extra(target: str, extra_data: dict, is_dir: bool = False) -> dict:
+    extra_dataitems = {}
+    for k, v in extra_data.items():
+        extra_dataitems[k] = mlrun.datastore.store_manager.object(
+            url=_get_file_path(target, v, isdir=is_dir), key=k
+        )
+    return extra_dataitems
+
+
+def _sanitize_and_serialize_model_spec_yaml(model: ModelArtifact) -> str:
+    model_dict = _sanitize_model_spec(model)
+    return _serialize_model_spec_yaml(model_dict)
+
+
+def _sanitize_model_spec(model: ModelArtifact) -> dict:
+    model_dict = model.to_dict()
+
+    # The model spec yaml should not include the tag, as the same model can be used with different tags,
+    # and the tag is not part of the model spec but the metadata of the model artifact
+    model_dict["metadata"].pop("tag", None)
+
+    # Remove future packaging links
+    if model_dict["spec"].get("extra_data"):
+        model_dict["spec"]["extra_data"] = {
+            key: item
+            for key, item in model_dict["spec"]["extra_data"].items()
+            if item is not ...
+        }
+    return model_dict
+
+
+def _serialize_model_spec_yaml(model_dict: dict) -> str:
+    return yaml.safe_dump(model_dict)