athena-intelligence 0.1.124__py3-none-any.whl → 0.1.125__py3-none-any.whl

athena/client.py

@@ -7,7 +7,7 @@ import typing
 import warnings
 
 import httpx
-from typing import cast, List, Union
+from typing import cast, List, Tuple, Union
 from typing_extensions import TypeVar, ParamSpec
 
 from . import core
@@ -116,8 +116,8 @@ class WrappedToolsClient(ToolsClient):
         client.tools.read_data_frame(asset_id="asset_id")
         """
         _check_pandas_installed()
-        file_bytes = self.get_file(asset_id)
-        return _to_pandas_df(file_bytes, *args, **kwargs)
+        file_bytes, media_type = self._get_file_and_media_type(asset_id=asset_id)
+        return _to_pandas_df(file_bytes, *args, media_type=media_type, **kwargs)
 
 
     def save_asset(    # type: ignore[override]
@@ -131,7 +131,7 @@ class WrappedToolsClient(ToolsClient):
         """
         Parameters
         ----------
-        asset_object : pd.DataFrame | pd.Series | core.File
+        asset_object : pd.DataFrame | pd.Series | matplotlib.figure.Figure | core.File
             A pandas data frame, series, matplotlib figure, or core.File
 
         parent_folder_id : typing.Optional[str]
@@ -143,6 +143,8 @@ class WrappedToolsClient(ToolsClient):
         request_options : typing.Optional[RequestOptions]
             Request-specific configuration.
 
+        **kwargs : passed down to conversion methods
+
         Returns
         -------
         SaveAssetRequestOut
@@ -155,27 +157,14 @@ class WrappedToolsClient(ToolsClient):
         client = Athena(api_key="YOUR_API_KEY")
         client.tools.save_asset(df)
         """
-        asset_object = _convert_asset_object(asset_object=asset_object, name=name)
+        asset_object = _convert_asset_object(asset_object=asset_object, name=name, **kwargs)
         return super().save_asset(
-            file=asset_object, parent_folder_id=parent_folder_id, **kwargs
+            file=asset_object, parent_folder_id=parent_folder_id
         )
 
-    def get_asset(self, asset_id: str) -> Union["pd.DataFrame", AthenaAsset]:
+    def _get_file_and_media_type(self, asset_id: str) -> Tuple[io.BytesIO, str]:
         """
-        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")
+        Gets the file togehter with media type returned by server
         """
         # while we wait for https://github.com/fern-api/fern/issues/4316
         result = self._client_wrapper.httpx_client.request(
@@ -196,6 +185,27 @@ class WrappedToolsClient(ToolsClient):
             # fallback to `libmagic` inference
             media_type = _infer_media_type(bytes_io=file_bytes)
 
+        return file_bytes, media_type
+
+    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")
+        """
+        file_bytes, media_type = self._get_file_and_media_type(asset_id=asset_id)
+
         media_type_aliases = {"image/jpg": "image/jpeg"}
         media_type = media_type_aliases.get(media_type, media_type)
 
@@ -289,9 +299,9 @@ class WrappedAsyncToolsClient(AsyncToolsClient):
         name: Union[str, None] = None,
         **kwargs,
     ) -> SaveAssetRequestOut:
-        asset_object = _convert_asset_object(asset_object=asset_object, name=name)
+        asset_object = _convert_asset_object(asset_object=asset_object, name=name, **kwargs)
         return await super().save_asset(
-            file=asset_object, parent_folder_id=parent_folder_id, **kwargs
+            file=asset_object, parent_folder_id=parent_folder_id
         )
 
 
@@ -514,6 +524,7 @@ def _to_pandas_df(
 def _convert_asset_object(
     asset_object: Union["pd.DataFrame", "pd.Series", core.File],
     name: Union[str, None] = None,
+    **kwargs
 ) -> core.File:
     import pandas as pd
     try:
@@ -526,7 +537,7 @@ def _convert_asset_object(
     if isinstance(asset_object, pd.DataFrame):
         return (
             name or "Uploaded data frame",
-            asset_object.to_parquet(path=None),
+            asset_object.to_parquet(path=None, **kwargs),
             "application/vnd.apache.parquet",
         )
     if format_display_data:

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.124",
+            "X-Fern-SDK-Version": "0.1.125",
         }
         headers["X-API-KEY"] = self.api_key
         return headers

{athena_intelligence-0.1.124.dist-info → athena_intelligence-0.1.125.dist-info}/METADATA

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

{athena_intelligence-0.1.124.dist-info → athena_intelligence-0.1.125.dist-info}/RECORD

@@ -10,10 +10,10 @@ athena/agents/research/client.py,sha256=mo63cUvLw8AC9I29WI4Ym_jOsvRsco1wRqsNp0rt
 athena/agents/sql/__init__.py,sha256=FTtvy8EDg9nNNg9WCatVgKTRYV8-_v1roeGPAKoa_pw,65
 athena/agents/sql/client.py,sha256=1r2aT4JLxVCeL0FzEBeHmttwerdRmJOxVHim2_CHkvg,4767
 athena/base_client.py,sha256=-kVdOlIibBz48lxWratdQAzT7fTvZsORvOMF3KoPDPw,5647
-athena/client.py,sha256=4PUPrBPCMTFpHR1yuKVR5eC1AYBl_25SMf6ZH82JHB0,19039
+athena/client.py,sha256=5n5pOXAcMljMR2-MnI8fGvHY-TgqLzWmihDur0XrpM8,19503
 athena/core/__init__.py,sha256=UFXpYzcGxWQUucU1TkjOQ9mGWN3A5JohluOIWVYKU4I,973
 athena/core/api_error.py,sha256=RE8LELok2QCjABadECTvtDp7qejA1VmINCh6TbqPwSE,426
-athena/core/client_wrapper.py,sha256=PsnlYk4Sr5_x1d3Z8hUp0g3i5pD1dqfmKXH-mHvU04Q,1806
+athena/core/client_wrapper.py,sha256=nziI1y3xFUHOGDnMU6SL8HlIHNCXDaJtpfCCe4o2tJ4,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
@@ -81,6 +81,6 @@ athena/types/text_content.py,sha256=uG2poNIkM6o7tFgf-eKzZk9kZHYImY3JdI-NkYiqWgU,
 athena/types/tool.py,sha256=6H2BFZiBgQOtYUAwSYBeGZKhwev17IEwnIjgmno6dZw,436
 athena/types/type.py,sha256=JaUIt4ogmO4XxCQ9c56fqKN5qANKkrnpuZGmdqOCIow,581
 athena/version.py,sha256=8aYAOJtVLaJLpRp6mTiEIhnl8gXA7yE0aDtZ-3mKQ4k,87
-athena_intelligence-0.1.124.dist-info/METADATA,sha256=lwtK__HdkgNj6iBeRlDN4JVY5z7XoZPv1ABXIwaBRT0,5274
-athena_intelligence-0.1.124.dist-info/WHEEL,sha256=Zb28QaM1gQi8f4VCBhsUklF61CTlNYfs9YAZn-TOGFk,88
-athena_intelligence-0.1.124.dist-info/RECORD,,
+athena_intelligence-0.1.125.dist-info/METADATA,sha256=YZtDUfJx8ZHA6L-it-tdB-28N1MR0T3LELbTW0bWDBs,5274
+athena_intelligence-0.1.125.dist-info/WHEEL,sha256=Zb28QaM1gQi8f4VCBhsUklF61CTlNYfs9YAZn-TOGFk,88
+athena_intelligence-0.1.125.dist-info/RECORD,,