Fixed _info bug, added async streaming functionality (#113)

Author: dialectician4

Cargo.lock

@@ -1,6 +1,6 @@
 [package]
 name = "pyfusion"
-version = "2.0.10"
+version = "2.0.11-dev0"
 edition = "2021"
 
 
@@ -27,6 +27,17 @@ jsonwebtoken = {version = "9.3.0"}
 state = { version = "0.6"}
 
 
+# Set ceiling on Windows crates due to auth issue on some Windows environments
+[target.'cfg(target_arch = "windows")'.dependencies]
+windows-targets = { version = "<=0.48.5" }
+windows-registry = { version = "<=0.2.0" }
+windows-strings = { version = "<=0.1.0" }
+windows-link = { version = "<=0.1.0" }
+# wasi = { version = "<= 0.13.3+wasi-0.2.2" }
+# tokio-util = { version = "<= 0.7.13" }
+# End of WINDOWS limit
+
+
 # Struggling to get OPEN_SSL env vars to point to the installed libs on these archs
 # Remove these once we figure out the magic
 [target.'cfg(target_arch = "aarch64")'.dependencies]
@@ -40,6 +51,7 @@ openssl = { version = "0.10.64", features = ["vendored"] }
 [package.metadata.cargo-llvm-cov]
 # End OPEN_SSL 
 
+
 # Ignore directories and files that are not part of the Rust source code
 ignore-filename-regex = "(.cargo|py_integ|py_src|py_tests|dist|docs|downloads|runnable|site|target|\\.git|\\.github)/"
 ignore-glob = ["py_integ/*", "py_src/*", "py_tests/*", "dist/*", "docs/*", "downloads/*", "runnable/*", "site/*", "target/*"]

Cargo.toml

@@ -2,7 +2,7 @@
 
 __author__ = """Fusion Devs"""
 __email__ = "fusion_developers@jpmorgan.com"
-__version__ = "2.0.10"
+__version__ = "2.0.11-dev0"
 
 from fusion._fusion import FusionCredentials
 from fusion.fs_sync import fsync

docs/fusion_filesystem.ipynb

@@ -91,7 +91,7 @@ class FusionEmbeddingsConnection(Connection):  # type: ignore
         credentials (FusionCredentials or str, optional): Credentials for the Fusion API. Can be a `FusionCredentials`
             object or a path to a credentials file. Defaults to "config/client_credentials.json".
         catalog (str, optional): Catalog name. Defaults to "common".
-        knowledge_base (str | list[str], optional): Knowledge base name. A dataset identifier. If multiple identifiers 
+        knowledge_base (str | list[str], optional): Knowledge base name. A dataset identifier. If multiple identifiers
             are provided, the connection will perform searches across all provided knowledge bases. Any other operations
             will not be supported, as a single knowledge base is required for those operations.
         index (str, optional): Index name. Defaults to `None`. Used to determine index when the _bulk operation
@@ -200,7 +200,6 @@ def _tidy_url(self, url: str) -> str:
     def _remap_endpoints(url: str) -> str:
         return url.replace("_bulk", "embeddings").replace("_search", "search")
 
-
     def _make_url_valid(self, url: str, body: bytes | None = None) -> str:
         if url == "/_bulk":
             index_name = self.index_name if self.index_name else _retrieve_index_name_from_bulk_body(body)
@@ -400,7 +399,7 @@ def __init__(  # noqa: PLR0912, PLR0913, PLR0915
         credentials (FusionCredentials or str, optional): Credentials for the Fusion API. Can be a `FusionCredentials`
             object or a path to a credentials file. Defaults to "config/client_credentials.json".
         catalog (str, optional): Catalog name. Defaults to "common".
-        knowledge_base (str, optional): Knowledge base name. A dataset identifier. If multiple identifiers 
+        knowledge_base (str, optional): Knowledge base name. A dataset identifier. If multiple identifiers
             are provided, the connection will perform searches across all provided knowledge bases. Any other operations
             will not be supported, as a single knowledge base is required for those operations.
         index (str, optional): Index name. Defaults to `None`. Used to determine index when the _bulk operation

py_src/fusion/__init__.py

@@ -13,6 +13,7 @@
 
 logger = logging.getLogger(__name__)
 
+
 def _format_full_index_response(response: requests.Response) -> pd.DataFrame:
     """Format get index response.
 
@@ -77,7 +78,7 @@ def _retrieve_index_name_from_bulk_body(body: bytes | None) -> str:
             if "index" in json_obj and "_index" in json_obj["index"]:
                 index_name = str(json_obj["index"]["_index"])
                 return index_name
-    
+
     raise ValueError("Index name not found in bulk body")
 
 
@@ -115,6 +116,7 @@ def _modify_post_response_langchain(raw_data: str | bytes | bytearray) -> str |
 
     return raw_data
 
+
 def _modify_post_haystack(knowledge_base: str | list[str] | None, body: bytes | None, method: str) -> bytes | None:
     """Method to modify haystack POST body to match the embeddings API, which expects the embedding field to be
         named "vector".

py_src/fusion/embeddings.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import copy
 import json as js
 import logging
 import re
@@ -53,6 +54,8 @@
 )
 
 if TYPE_CHECKING:
+    from collections.abc import AsyncGenerator
+
     import fsspec
     import requests
     from opensearchpy import AsyncOpenSearch, OpenSearch
@@ -770,6 +773,46 @@ def download(  # noqa: PLR0912, PLR0913
                     warnings.warn(f"The download of {r[1]} was not successful", stacklevel=2)
         return res if return_paths else None
 
+    async def _async_stream_file(self, url: str, chunk_size: int = 100) -> AsyncGenerator[bytes, None]:
+        """Return async stream of a file fro mthe given url.
+
+        Args:
+            url (str): File url. Appends Fusion.root_url if http prefix not present.
+            chunk_size (int, optional): Size for each chunk in async stream. Defaults to 100.
+
+        Returns:
+            AsyncGenerator[bytes, None]: Async generator object.
+
+        Yields:
+            Iterator[AsyncGenerator[bytes, None]]: Next set of bytes read from the file at given url.
+        """
+        dup_credentials = copy.deepcopy(self.credentials)
+        async_fs = FusionHTTPFileSystem(
+            client_kwargs={"root_url": self.root_url, "credentials": dup_credentials}, asynchronous=True
+        )
+        session = await async_fs.set_session()
+        async with session:
+            async for chunk in async_fs._stream_file(url, chunk_size):
+                yield chunk
+
+    async def _async_get_file(self, url: str, chunk_size: int = 1000) -> bytes:
+        """Return a file from url as a bytes object, asynchronously.
+
+        Under the hood, opens up an async stream downloading file in chunk_size bytes per chunk.
+        Larger chunk sizes results in shorter execution time for this function.
+
+        Args:
+            url (str): File url. Appends Fusion.root_url if http prefix not present.
+            chunk_size (int, optional): Size of chunks to get from async stream. Defaults to 1000.
+
+        Returns:
+            bytes: File from url as a bytes object.
+        """
+        async_generator = self._async_stream_file(url, chunk_size)
+        bytes_list: list[bytes] = [chunk async for chunk in async_generator]
+        final_bytes: bytes = b"".join(bytes_list)
+        return final_bytes
+
     def to_df(  # noqa: PLR0913
         self,
         dataset: str,
@@ -2517,7 +2560,7 @@ def get_async_fusion_vector_store_client(self, knowledge_base: str, catalog: str
             root_url=self.root_url,
             credentials=self.credentials,
         )
-    
+
     def list_datasetmembers_distributions(
         self,
         dataset: str,

py_src/fusion/embeddings_utils.py

@@ -207,6 +207,7 @@ def info(self, path: str, **kwargs: Any) -> Any:
         kwargs["keep_protocol"] = True
         res = super().ls(path, detail=True, **kwargs)
         if res[0]["type"] != "file":
+            kwargs.pop("keep_protocol", None)
             res = super().info(path, **kwargs)
             if path.split("/")[-2] == "datasets":
                 target = path.split("/")[-1]
@@ -220,15 +221,20 @@ def info(self, path: str, **kwargs: Any) -> Any:
 
     async def _info(self, path: str, **kwargs: Any) -> Any:
         await self._async_startup()
-        path = await self._decorate_url_a(path)
+        path = self._decorate_url(path)
         kwargs["keep_protocol"] = True
         res = await super()._ls(path, detail=True, **kwargs)
         if res[0]["type"] != "file":
+            kwargs.pop("keep_protocol", None)
             res = await super()._info(path, **kwargs)
             if path.split("/")[-2] == "datasets":
                 target = path.split("/")[-1]
                 args = ["/".join(path.split("/")[:-1]) + f"/changes?datasets={quote(target)}"]
                 res["changes"] = await self._changes(*args)
+            if res["size"] is None and (res["mimetype"] == "application/json") and (res["type"] == "file"):
+                res["size"] = 0
+                res.pop("mimetype", None)
+                res["type"] = "directory"
         split_path = path.split("/")
         if len(split_path) > 1 and split_path[-2] == "distributions":
             res = res[0]
@@ -325,6 +331,29 @@ async def _cat(self, url: str, start: Optional[int] = None, end: Optional[int] =
         out = await super()._cat(url, start=start, end=end, **kwargs)
         return out
 
+    async def _stream_file(self, url: str, chunk_size: int = 100) -> AsyncGenerator[bytes, None]:
+        """Return an async stream to file at the given url.
+
+        Args:
+            url (str): File url. Appends Fusion.root_url if http prefix not present.
+            chunk_size (int, optional): Size for each chunk in async stream. Defaults to 100.
+
+        Returns:
+            AsyncGenerator[bytes, None]: Async generator object.
+
+        Yields:
+            Iterator[AsyncGenerator[bytes, None]]: Next set of bytes read from the file at given url.
+        """
+        await self._async_startup()
+        url = self._decorate_url(url)
+        f = await self.open_async(url, "rb")
+        async with f:
+            while True:
+                chunk = await f.read(chunk_size)
+                if not chunk:
+                    break
+                yield chunk
+
     async def _fetch_range(
         self,
         session: aiohttp.ClientSession,