athena-intelligence 0.1.105__py3-none-any.whl → 0.1.107__py3-none-any.whl

athena/client.py

@@ -1,3 +1,4 @@
+import dataclasses
 import enum
 import io
 import os
@@ -41,6 +42,26 @@ class SpecialEnvironments(enum.Enum):
     AUTODETECT_ENVIRONMENT = 'AUTO'
 
 
+@dataclasses.dataclass
+class AthenaAsset:
+    asset_id: str
+    data: bytes
+    media_type: str
+
+    def _repr_mimebundle_(self):
+        if self.media_type == "application/sql":
+            # it is safe to import IPython in `_repr_mimebundle_`
+            # as this is only intended to be invoked by IPython.
+            from IPython import display    # type: ignore[import]
+
+            code = display.Code(
+                data=self.data,
+                language="sql",
+            )
+            return {"text/html": code._repr_html_()}
+        return {self.media_type: self.data}
+
+
 class WrappedToolsClient(ToolsClient):
 
     def get_file(self, asset_id: str) -> io.BytesIO:
@@ -113,6 +134,10 @@ class WrappedToolsClient(ToolsClient):
             A pandas data frame, series, or core.File
 
         parent_folder_id : typing.Optional[str]
+            Identifier of the folder into which the asset should be saved
+
+        name : typing.Optional[str]
+            The name for the asset
 
         request_options : typing.Optional[RequestOptions]
             Request-specific configuration.
@@ -133,6 +158,74 @@ class WrappedToolsClient(ToolsClient):
         return super().save_asset(
             file=asset_object, parent_folder_id=parent_folder_id, **kwargs
         )
+
+    def get_asset(self, asset_id: str) -> Union["pd.DataFrame", AthenaAsset]:
+        """
+        Parameters
+        ----------
+        asset_id : str
+
+        Returns
+        -------
+        pd.DataFrame or AthenaAsset
+
+        Examples
+        --------
+        from athena.client import Athena
+
+        client = Athena(api_key="YOUR_API_KEY")
+        client.tools.get_asset(asset_id="asset_id")
+        """
+        # while we wait for https://github.com/fern-api/fern/issues/4316
+        result = self._client_wrapper.httpx_client.request(
+            "api/v0/tools/file/raw-data", method="GET", params={"asset_id": asset_id}
+        )
+        if result.status_code != 200:
+            # let fern handle errors codes
+            self.raw_data(asset_id=asset_id)
+            raise Exception(
+                f"Could not get assset - unhandled error code: {result.status_code}"
+            )
+
+        file_bytes = io.BytesIO(result.read())
+
+        media_type = result.headers.get("content-type", "").split(";")[0]
+
+        if media_type == "":
+            # fallback to `libmagic` inference
+            media_type = _infer_media_type(bytes_io=file_bytes)
+
+        media_type_aliases = {"image/jpg": "image/jpeg"}
+        media_type = media_type_aliases.get(media_type, media_type)
+
+        supported_media_types = {
+            "application/json",
+            "application/pdf",
+            "image/jpeg",
+            "image/gif",
+            "image/png",
+            "image/svg+xml",
+            "image/webp",
+            "text/html",
+            "text/latex",
+            "text/markdown",
+            "text/plain",
+            "application/sql",
+        }
+
+        data = file_bytes.read()
+
+        if media_type in supported_media_types:
+            return AthenaAsset(
+                asset_id=asset_id,
+                data=data,
+                media_type=media_type,
+            )
+
+        if media_type in _pandas_media_types:
+            return _to_pandas_df(file_bytes)
+
+        raise NotImplementedError("Assets of `{media_type}` type are not yet supported")
     
 
 class WrappedQueryClient(QueryClient):
@@ -376,17 +469,33 @@ def _check_pandas_installed():
     assert pandas
 
 
-def _to_pandas_df(bytes_io: io.BytesIO, *args, **kwargs):
-    import pandas as pd
+def _infer_media_type(bytes_io: io.BytesIO) -> str:
     import magic
 
     # ideally this would be read from response header, but fern SDK for Python hides this info from us
     media_type = magic.from_buffer(bytes_io.read(2048), mime=True)
     bytes_io.seek(0)
+    return media_type
+
+
+_pandas_media_types = {
+    "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
+    "application/vnd.apache.parquet",
+    "application/octet-stream",
+    "application/vnd.ms-excel",
+    "text/csv",
+    "application/csv",
+}
+
+
+def _to_pandas_df(bytes_io: io.BytesIO, *args, **kwargs):
+    import pandas as pd
+
+    media_type = _infer_media_type(bytes_io)
 
     if media_type == 'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet':
         return pd.read_excel(bytes_io, *args, engine='openpyxl', **kwargs)
-    elif media_type == 'application/vnd.apache.parquet':
+    elif media_type in {'application/vnd.apache.parquet', 'application/octet-stream'}:
         return pd.read_parquet(bytes_io, *args, **kwargs)
     elif media_type == 'application/vnd.ms-excel':
         return pd.read_excel(bytes_io, *args, **kwargs)

athena/core/client_wrapper.py

@@ -17,7 +17,7 @@ class BaseClientWrapper:
         headers: typing.Dict[str, str] = {
             "X-Fern-Language": "Python",
             "X-Fern-SDK-Name": "athena-intelligence",
-            "X-Fern-SDK-Version": "0.1.105",
+            "X-Fern-SDK-Version": "0.1.107",
         }
         headers["X-API-KEY"] = self.api_key
         return headers

{athena_intelligence-0.1.105.dist-info → athena_intelligence-0.1.107.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: athena-intelligence
-Version: 0.1.105
+Version: 0.1.107
 Summary: Athena Intelligence Python Library
 Requires-Python: >=3.8,<4.0
 Classifier: Intended Audience :: Developers

{athena_intelligence-0.1.105.dist-info → athena_intelligence-0.1.107.dist-info}/RECORD

@@ -1,9 +1,9 @@
 athena/__init__.py,sha256=EtDa2cZbHfqJnBxCrjErWF1oJWH0cbJbFB7bkxeOFnE,1392
 athena/base_client.py,sha256=TlP599mOBvj7Tk8IpFe5dgrDM98GJu3lEQh_zwl4vtA,5439
-athena/client.py,sha256=9hj9SzW7OhAIY8yKHA0WObOtx8-1W4aLSHduqkBHJT0,14722
+athena/client.py,sha256=OhZA99pgQ2kEy18GoyUt25AZdqjQ49YdXvaewz3llLU,17939
 athena/core/__init__.py,sha256=UFXpYzcGxWQUucU1TkjOQ9mGWN3A5JohluOIWVYKU4I,973
 athena/core/api_error.py,sha256=RE8LELok2QCjABadECTvtDp7qejA1VmINCh6TbqPwSE,426
-athena/core/client_wrapper.py,sha256=ijXhms74o95_1sms0mC87d_xv0uBTKXzxH7pmz9YEKE,1806
+athena/core/client_wrapper.py,sha256=A07xONnQhSHb3pTSjJC_rKQUFDhHJycf4wvFEh5XkfI,1806
 athena/core/datetime_utils.py,sha256=nBys2IsYrhPdszxGKCNRPSOCwa-5DWOHG95FB8G9PKo,1047
 athena/core/file.py,sha256=sy1RUGZ3aJYuw998bZytxxo6QdgKmlnlgBaMvwEKCGg,1480
 athena/core/http_client.py,sha256=Z4NuAsJD-51yqmoME17O5sxwx5orSp1wsnd6bPyKcgA,17768
@@ -41,6 +41,6 @@ athena/types/file_too_large_error.py,sha256=AinkrcgR7lcTILAD8RX0x48P3GlSoAh1Oihx
 athena/types/parent_folder_error.py,sha256=ZMF-i3mZY6Mu1n5uQ60Q3mIIfehlWuXtgFUkSYspkx8,1120
 athena/types/save_asset_request_out.py,sha256=5bpBaUV3oeuL_hz4s07c-6MQHkn4cBsyxgT_SD5oi6I,1193
 athena/version.py,sha256=8aYAOJtVLaJLpRp6mTiEIhnl8gXA7yE0aDtZ-3mKQ4k,87
-athena_intelligence-0.1.105.dist-info/METADATA,sha256=0-dSyqtK9SICscX0vlve617c46HjDvYdj6NdWuZcn9k,5274
-athena_intelligence-0.1.105.dist-info/WHEEL,sha256=Zb28QaM1gQi8f4VCBhsUklF61CTlNYfs9YAZn-TOGFk,88
-athena_intelligence-0.1.105.dist-info/RECORD,,
+athena_intelligence-0.1.107.dist-info/METADATA,sha256=wuYxYlnBqbPOgg6rVwFafRSoIWxRjvn8QnbhenZExgg,5274
+athena_intelligence-0.1.107.dist-info/WHEEL,sha256=Zb28QaM1gQi8f4VCBhsUklF61CTlNYfs9YAZn-TOGFk,88
+athena_intelligence-0.1.107.dist-info/RECORD,,