athena-intelligence 0.1.104__tar.gz → 0.1.105__tar.gz

{athena_intelligence-0.1.104 → athena_intelligence-0.1.105}/PKG-INFO

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

{athena_intelligence-0.1.104 → athena_intelligence-0.1.105}/pyproject.toml

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

{athena_intelligence-0.1.104 → athena_intelligence-0.1.105}/src/athena/client.py

@@ -8,11 +8,13 @@ import httpx
 from typing import cast, List, Union
 from typing_extensions import TypeVar, ParamSpec
 
+from . import core
 from .base_client import BaseAthena, AsyncBaseAthena
 from .environment import AthenaEnvironment
 from .tools.client import AsyncToolsClient, ToolsClient
 from .query.client import AsyncQueryClient, QueryClient
 from .types.data_frame_request_out import DataFrameRequestOut
+from .types.save_asset_request_out import SaveAssetRequestOut
 
 if typing.TYPE_CHECKING:
     import pandas as pd
@@ -96,6 +98,43 @@ class WrappedToolsClient(ToolsClient):
         return _to_pandas_df(file_bytes, *args, **kwargs)
 
 
+    def save_asset(    # type: ignore[override]
+        self,
+        asset_object: Union["pd.DataFrame", "pd.Series", core.File],
+        *,
+        parent_folder_id: Union[str, None] = None,
+        name: Union[str, None] = None,
+        **kwargs,
+    ) -> SaveAssetRequestOut:
+        """
+        Parameters
+        ----------
+        asset_object : pd.DataFrame | pd.Series | core.File
+            A pandas data frame, series, or core.File
+
+        parent_folder_id : typing.Optional[str]
+
+        request_options : typing.Optional[RequestOptions]
+            Request-specific configuration.
+
+        Returns
+        -------
+        SaveAssetRequestOut
+            Successful Response
+
+        Examples
+        --------
+        from athena.client import Athena
+
+        client = Athena(api_key="YOUR_API_KEY")
+        client.tools.save_asset(df)
+        """
+        asset_object = _convert_asset_object(asset_object=asset_object, name=name)
+        return super().save_asset(
+            file=asset_object, parent_folder_id=parent_folder_id, **kwargs
+        )
+    
+
 class WrappedQueryClient(QueryClient):
 
     @_inherit_signature_and_doc(QueryClient.execute, {'DataFrameRequestOut': 'pd.DataFrame'})
@@ -148,6 +187,20 @@ class WrappedAsyncToolsClient(AsyncToolsClient):
         file_bytes = await self.get_file(asset_id)
         return _to_pandas_df(file_bytes, *args, **kwargs)
 
+    @_add_docs_for_async_variant(WrappedToolsClient.save_asset)
+    async def save_asset(
+        self,
+        asset_object: Union["pd.DataFrame", "pd.Series", core.File],
+        *,
+        parent_folder_id: Union[str, None] = None,
+        name: Union[str, None] = None,
+        **kwargs,
+    ) -> SaveAssetRequestOut:
+        asset_object = _convert_asset_object(asset_object=asset_object, name=name)
+        return await super().save_asset(
+            file=asset_object, parent_folder_id=parent_folder_id, **kwargs
+        )
+
 
 class WrappedAsyncQueryClient(AsyncQueryClient):
 
@@ -333,9 +386,28 @@ def _to_pandas_df(bytes_io: io.BytesIO, *args, **kwargs):
 
     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':
+        return pd.read_parquet(bytes_io, *args, **kwargs)
     elif media_type == 'application/vnd.ms-excel':
         return pd.read_excel(bytes_io, *args, **kwargs)
-    elif media_type == 'text/csv':
+    elif media_type in {'text/csv', 'application/csv'}:
         return pd.read_csv(bytes_io, *args, **kwargs)
     else:
         raise Exception("Unknown media type")
+
+
+def _convert_asset_object(
+    asset_object: Union["pd.DataFrame", "pd.Series", core.File],
+    name: Union[str, None] = None,
+) -> core.File:
+    import pandas as pd
+
+    if isinstance(asset_object, pd.Series):
+        asset_object = asset_object.to_frame()
+    if isinstance(asset_object, pd.DataFrame):
+        return (
+            name or "Uploaded data frame",
+            asset_object.to_parquet(path=None),
+            "application/vnd.apache.parquet",
+        )
+    return asset_object    # type: ignore[return-value]

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

{athena_intelligence-0.1.104 → athena_intelligence-0.1.105}/src/athena/tools/client.py

@@ -179,6 +179,7 @@ class ToolsClient:
             See core.File for more documentation
 
         parent_folder_id : typing.Optional[str]
+            Identifier of the folder into which the asset should be saved
 
         request_options : typing.Optional[RequestOptions]
             Request-specific configuration.
@@ -373,6 +374,7 @@ class AsyncToolsClient:
             See core.File for more documentation
 
         parent_folder_id : typing.Optional[str]
+            Identifier of the folder into which the asset should be saved
 
         request_options : typing.Optional[RequestOptions]
             Request-specific configuration.

{athena_intelligence-0.1.104 → athena_intelligence-0.1.105}/src/athena/tools/types/tools_data_frame_request_columns_item.py

@@ -2,4 +2,4 @@
 
 import typing
 
-ToolsDataFrameRequestColumnsItem = typing.Union[str, int]
+ToolsDataFrameRequestColumnsItem = typing.Union[int, str]