athena-intelligence 0.1.106__tar.gz → 0.1.108__tar.gz

{athena_intelligence-0.1.106 → athena_intelligence-0.1.108}/PKG-INFO

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

{athena_intelligence-0.1.106 → athena_intelligence-0.1.108}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "athena-intelligence"
-version = "0.1.106"
+version = "0.1.108"
 description = "Athena Intelligence Python Library"
 readme = "README.md"
 authors = []

{athena_intelligence-0.1.106 → athena_intelligence-0.1.108}/src/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, include=None, exclude=None):
+        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, media_type=media_type)
+
+        raise NotImplementedError("Assets of `{media_type}` type are not yet supported")
     
 
 class WrappedQueryClient(QueryClient):
@@ -376,24 +469,46 @@ 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
 
-    if media_type == 'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet':
-        return pd.read_excel(bytes_io, *args, engine='openpyxl', **kwargs)
-    elif media_type in {'application/vnd.apache.parquet', 'application/octet-stream'}:
+
+_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, media_type: Union[str, None] = None, **kwargs
+):
+    import pandas as pd
+
+    if media_type is None:
+        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 in {"application/vnd.apache.parquet", "application/octet-stream"}:
         return pd.read_parquet(bytes_io, *args, **kwargs)
-    elif media_type == 'application/vnd.ms-excel':
+    elif media_type == "application/vnd.ms-excel":
         return pd.read_excel(bytes_io, *args, **kwargs)
-    elif media_type in {'text/csv', 'application/csv'}:
+    elif media_type in {"text/csv", "application/csv"}:
         return pd.read_csv(bytes_io, *args, **kwargs)
     else:
-        raise Exception("Unknown media type")
+        raise Exception(f"Unknown media type: {media_type}")
 
 
 def _convert_asset_object(

{athena_intelligence-0.1.106 → athena_intelligence-0.1.108}/src/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.106",
+            "X-Fern-SDK-Version": "0.1.108",
         }
         headers["X-API-KEY"] = self.api_key
         return headers