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

mlrun/datastore/targets.py

@@ -396,7 +396,7 @@ class BaseStoreTarget(DataTargetBase):
         self,
         name: str = "",
         path=None,
-        attributes: dict[str, str] = None,
+        attributes: Optional[dict[str, str]] = None,
         after_step=None,
         columns=None,
         partitioned: bool = False,
@@ -405,8 +405,8 @@ class BaseStoreTarget(DataTargetBase):
         time_partitioning_granularity: Optional[str] = None,
         max_events: Optional[int] = None,
         flush_after_seconds: Optional[int] = None,
-        storage_options: dict[str, str] = None,
-        schema: dict[str, Any] = None,
+        storage_options: Optional[dict[str, str]] = None,
+        schema: Optional[dict[str, Any]] = None,
         credentials_prefix=None,
     ):
         super().__init__(
@@ -834,16 +834,16 @@ class ParquetTarget(BaseStoreTarget):
         self,
         name: str = "",
         path=None,
-        attributes: dict[str, str] = None,
+        attributes: Optional[dict[str, str]] = None,
         after_step=None,
         columns=None,
-        partitioned: bool = None,
+        partitioned: Optional[bool] = None,
         key_bucketing_number: Optional[int] = None,
         partition_cols: Optional[list[str]] = None,
         time_partitioning_granularity: Optional[str] = None,
         max_events: Optional[int] = 10000,
         flush_after_seconds: Optional[int] = 900,
-        storage_options: dict[str, str] = None,
+        storage_options: Optional[dict[str, str]] = None,
     ):
         self.path = path
         if partitioned is None:
@@ -1136,7 +1136,8 @@ class CSVTarget(BaseStoreTarget):
         import pyspark.sql.functions as funcs
 
         for col_name, col_type in df.dtypes:
-            if col_type == "timestamp":
+            # covers TimestampType and TimestampNTZType, which was added in PySpark 3.4.0
+            if col_type.startswith("timestamp"):
                 # df.write.csv saves timestamps with millisecond precision, but we want microsecond precision
                 # for compatibility with storey.
                 df = df.withColumn(
@@ -1199,7 +1200,7 @@ class SnowflakeTarget(BaseStoreTarget):
         self,
         name: str = "",
         path=None,
-        attributes: dict[str, str] = None,
+        attributes: Optional[dict[str, str]] = None,
         after_step=None,
         columns=None,
         partitioned: bool = False,
@@ -1208,15 +1209,15 @@ class SnowflakeTarget(BaseStoreTarget):
         time_partitioning_granularity: Optional[str] = None,
         max_events: Optional[int] = None,
         flush_after_seconds: Optional[int] = None,
-        storage_options: dict[str, str] = None,
-        schema: dict[str, Any] = None,
+        storage_options: Optional[dict[str, str]] = None,
+        schema: Optional[dict[str, Any]] = None,
         credentials_prefix=None,
-        url: str = None,
-        user: str = None,
-        db_schema: str = None,
-        database: str = None,
-        warehouse: str = None,
-        table_name: str = None,
+        url: Optional[str] = None,
+        user: Optional[str] = None,
+        db_schema: Optional[str] = None,
+        database: Optional[str] = None,
+        warehouse: Optional[str] = None,
+        table_name: Optional[str] = None,
     ):
         attributes = attributes or {}
         if url:
@@ -1903,7 +1904,7 @@ class SQLTarget(BaseStoreTarget):
         self,
         name: str = "",
         path=None,
-        attributes: dict[str, str] = None,
+        attributes: Optional[dict[str, str]] = None,
         after_step=None,
         partitioned: bool = False,
         key_bucketing_number: Optional[int] = None,
@@ -1911,16 +1912,16 @@ class SQLTarget(BaseStoreTarget):
         time_partitioning_granularity: Optional[str] = None,
         max_events: Optional[int] = None,
         flush_after_seconds: Optional[int] = None,
-        storage_options: dict[str, str] = None,
-        db_url: str = None,
-        table_name: str = None,
-        schema: dict[str, Any] = None,
+        storage_options: Optional[dict[str, str]] = None,
+        db_url: Optional[str] = None,
+        table_name: Optional[str] = None,
+        schema: Optional[dict[str, Any]] = None,
         primary_key_column: str = "",
         if_exists: str = "append",
         create_table: bool = False,
         # create_according_to_data: bool = False,
         varchar_len: int = 50,
-        parse_dates: list[str] = None,
+        parse_dates: Optional[list[str]] = None,
     ):
         """
         Write to SqlDB as output target for a flow.

mlrun/datastore/utils.py

@@ -26,7 +26,7 @@ import mlrun.datastore
 
 
 def parse_kafka_url(
-    url: str, brokers: typing.Union[list, str] = None
+    url: str, brokers: typing.Optional[typing.Union[list, str]] = None
 ) -> tuple[str, list]:
     """Generating Kafka topic and adjusting a list of bootstrap servers.
 
@@ -71,7 +71,7 @@ def upload_tarball(source_dir, target, secrets=None):
 
 def filter_df_start_end_time(
     df: typing.Union[pd.DataFrame, typing.Iterator[pd.DataFrame]],
-    time_column: str = None,
+    time_column: typing.Optional[str] = None,
     start_time: pd.Timestamp = None,
     end_time: pd.Timestamp = None,
 ) -> typing.Union[pd.DataFrame, typing.Iterator[pd.DataFrame]]:

mlrun/datastore/v3io.py

@@ -14,6 +14,7 @@
 
 import time
 from datetime import datetime
+from typing import Optional
 
 import fsspec
 import v3io
@@ -33,7 +34,9 @@ V3IO_DEFAULT_UPLOAD_CHUNK_SIZE = 1024 * 1024 * 10
 
 
 class V3ioStore(DataStore):
-    def __init__(self, parent, schema, name, endpoint="", secrets: dict = None):
+    def __init__(
+        self, parent, schema, name, endpoint="", secrets: Optional[dict] = None
+    ):
         super().__init__(parent, name, schema, endpoint, secrets=secrets)
         self.endpoint = self.endpoint or mlrun.mlconf.v3io_api
 

mlrun/datastore/vectorstore.py

@@ -0,0 +1,229 @@
+# 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 collections.abc import Iterable
+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"]
+
+    for attr in possible_attributes:
+        if hasattr(vectorstore, attr):
+            collection_name = getattr(vectorstore, attr)
+            if collection_name:
+                return collection_name
+
+    store_class = vectorstore.__class__.__name__.lower()
+    if store_class == "mongodbatlasvectorsearch":
+        return vectorstore.collection.name
+
+    # If we get here, we couldn't find a valid collection name
+    raise ValueError(
+        "Failed to extract collection name from the vector store. "
+        "Please provide the collection name explicitly. "
+    )
+
+
+class VectorStoreCollection:
+    """
+    A wrapper class for vector store collections with MLRun integration.
+
+    This class wraps a vector store implementation (like Milvus, Chroma) and provides
+    integration with MLRun context for document and artifact management. It delegates
+    most operations to the underlying vector store while handling MLRun-specific
+    functionality.
+
+    The class implements attribute delegation through __getattr__ and __setattr__,
+    allowing direct access to the underlying vector store's methods and attributes
+    while maintaining MLRun integration.
+    """
+
+    def __init__(
+        self,
+        mlrun_context: Union["MlrunProject", "MLClientCtx"],  # noqa: F821
+        vector_store: "VectorStore",  # noqa: F821
+        collection_name: Optional[str] = None,
+    ):
+        self._collection_impl = vector_store
+        self._mlrun_context = mlrun_context
+        self.collection_name = collection_name or _extract_collection_name(vector_store)
+
+    @property
+    def __class__(self):
+        # Make isinstance() check the wrapped object's class
+        return self._collection_impl.__class__
+
+    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 delete(self, *args, **kwargs):
+        self._collection_impl.delete(*args, **kwargs)
+
+    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.collection_name)
+                    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 vector store collection.
+
+        Converts artifacts to LangChain documents, adds them to the vector store, and
+        updates the MLRun context. If documents are split, the IDs are handled appropriately.
+
+        :param artifacts: List of DocumentArtifact objects to add
+        :type artifacts: list[DocumentArtifact]
+        :param splitter: Document splitter to break artifacts into smaller chunks.
+                        If None, each artifact becomes a single document.
+        :type splitter: TextSplitter, optional
+        :param kwargs: Additional arguments passed to the underlying add_documents method.
+                    Special handling for 'ids' kwarg:
+
+                    * If provided and document is split, IDs are generated as "{original_id}_{i}"
+                        where i starts from 1 (e.g., "doc1_1", "doc1_2", etc.)
+                    * If provided and document isn't split, original IDs are used as-is
+
+        :return: List of IDs for all added documents. When no custom IDs are provided:
+
+                * Without splitting: Vector store generates IDs automatically
+                * With splitting: Vector store generates separate IDs for each chunk
+
+                When custom IDs are provided:
+
+                * Without splitting: Uses provided IDs directly
+                * With splitting: Generates sequential IDs as "{original_id}_{i}" for each chunk
+        :rtype: list
+
+        """
+        all_ids = []
+        user_ids = kwargs.pop("ids", None)
+
+        if user_ids:
+            if not isinstance(user_ids, Iterable):
+                raise ValueError("IDs must be an iterable collection")
+            if len(user_ids) != len(artifacts):
+                raise ValueError(
+                    "The number of IDs should match the number of artifacts"
+                )
+        for index, artifact in enumerate(artifacts):
+            documents = artifact.to_langchain_documents(splitter)
+            artifact.collection_add(self.collection_name)
+            if self._mlrun_context:
+                self._mlrun_context.update_artifact(artifact)
+            if user_ids:
+                num_of_documents = len(documents)
+                if num_of_documents > 1:
+                    ids_to_pass = [
+                        f"{user_ids[index]}_{i}" for i in range(1, num_of_documents + 1)
+                    ]
+                else:
+                    ids_to_pass = [user_ids[index]]
+                kwargs["ids"] = ids_to_pass
+            ids = self._collection_impl.add_documents(documents, **kwargs)
+            all_ids.extend(ids)
+        return all_ids
+
+    def remove_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.collection_name)
+        if self._mlrun_context:
+            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.collection_name)
+            if 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)
+            elif store_class == "chroma":
+                where = {DocumentArtifact.METADATA_SOURCE_KEY: artifact.get_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.get_source()
+                    }
+                }
+                return self._collection_impl.delete(filter=filter)
+            else:
+                raise NotImplementedError(
+                    f"delete_artifacts() operation not supported for {store_class}"
+                )

mlrun/datastore/wasbfs/fs.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import Optional
 from urllib.parse import urlparse
 
 from fsspec import AbstractFileSystem
@@ -22,23 +23,23 @@ class WasbFS(AbstractFileSystem):
 
     def __init__(
         self,
-        account_name: str = None,
-        account_key: str = None,
-        connection_string: str = None,
-        credential: str = None,
-        sas_token: str = None,
+        account_name: Optional[str] = None,
+        account_key: Optional[str] = None,
+        connection_string: Optional[str] = None,
+        credential: Optional[str] = None,
+        sas_token: Optional[str] = None,
         request_session=None,
-        socket_timeout: int = None,
-        blocksize: int = None,
-        client_id: str = None,
-        client_secret: str = None,
-        tenant_id: str = None,
+        socket_timeout: Optional[int] = None,
+        blocksize: Optional[int] = None,
+        client_id: Optional[str] = None,
+        client_secret: Optional[str] = None,
+        tenant_id: Optional[str] = None,
         anon: bool = True,
-        location_mode: str = None,
+        location_mode: Optional[str] = None,
         loop=None,
         asynchronous: bool = False,
         default_fill_cache: bool = True,
-        default_cache_type: str = None,
+        default_cache_type: Optional[str] = None,
         **kwargs,
     ):
         from adlfs import AzureBlobFileSystem