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

athena/agents/sql/client.py

@@ -1,14 +1,14 @@
 # This file was auto-generated by Fern from our API Definition.
 
 import typing
-from json.decoder import JSONDecodeError
-
-from ...core.api_error import ApiError
-from ...core.client_wrapper import AsyncClientWrapper, SyncClientWrapper
-from ...core.pydantic_utilities import pydantic_v1
+from ...core.client_wrapper import SyncClientWrapper
 from ...core.request_options import RequestOptions
-from ...errors.unprocessable_entity_error import UnprocessableEntityError
 from ...types.sql_agent_response import SqlAgentResponse
+from ...core.pydantic_utilities import parse_obj_as
+from ...errors.unprocessable_entity_error import UnprocessableEntityError
+from json.decoder import JSONDecodeError
+from ...core.api_error import ApiError
+from ...core.client_wrapper import AsyncClientWrapper
 
 # this is used as the default value for optional parameters
 OMIT = typing.cast(typing.Any, ...)
@@ -21,19 +21,19 @@ class SqlClient:
     def invoke(
         self,
         *,
-        config: typing.Dict[str, typing.Any],
-        messages: typing.Sequence[typing.Dict[str, typing.Any]],
-        request_options: typing.Optional[RequestOptions] = None
+        config: typing.Dict[str, typing.Optional[typing.Any]],
+        messages: typing.Sequence[typing.Dict[str, typing.Optional[typing.Any]]],
+        request_options: typing.Optional[RequestOptions] = None,
     ) -> SqlAgentResponse:
         """
         Coming soon! Generate, execute, and test SQL queries. Returns an asset ID for the query object.
 
         Parameters
         ----------
-        config : typing.Dict[str, typing.Any]
+        config : typing.Dict[str, typing.Optional[typing.Any]]
             Configuration for the SQL agent including database connection details and query parameters
 
-        messages : typing.Sequence[typing.Dict[str, typing.Any]]
+        messages : typing.Sequence[typing.Dict[str, typing.Optional[typing.Any]]]
             The messages to send to the SQL agent
 
         request_options : typing.Optional[RequestOptions]
@@ -46,7 +46,7 @@ class SqlClient:
 
         Examples
         --------
-        from athena.client import Athena
+        from athena import Athena
 
         client = Athena(
             api_key="YOUR_API_KEY",
@@ -59,15 +59,35 @@ class SqlClient:
         _response = self._client_wrapper.httpx_client.request(
             "api/v0/agents/sql/invoke",
             method="POST",
-            json={"config": config, "messages": messages},
+            json={
+                "config": config,
+                "messages": messages,
+            },
+            headers={
+                "content-type": "application/json",
+            },
             request_options=request_options,
             omit=OMIT,
         )
-        if 200 <= _response.status_code < 300:
-            return pydantic_v1.parse_obj_as(SqlAgentResponse, _response.json())  # type: ignore
-        if _response.status_code == 422:
-            raise UnprocessableEntityError(pydantic_v1.parse_obj_as(typing.Any, _response.json()))  # type: ignore
         try:
+            if 200 <= _response.status_code < 300:
+                return typing.cast(
+                    SqlAgentResponse,
+                    parse_obj_as(
+                        type_=SqlAgentResponse,  # type: ignore
+                        object_=_response.json(),
+                    ),
+                )
+            if _response.status_code == 422:
+                raise UnprocessableEntityError(
+                    typing.cast(
+                        typing.Optional[typing.Any],
+                        parse_obj_as(
+                            type_=typing.Optional[typing.Any],  # type: ignore
+                            object_=_response.json(),
+                        ),
+                    )
+                )
             _response_json = _response.json()
         except JSONDecodeError:
             raise ApiError(status_code=_response.status_code, body=_response.text)
@@ -81,19 +101,19 @@ class AsyncSqlClient:
     async def invoke(
         self,
         *,
-        config: typing.Dict[str, typing.Any],
-        messages: typing.Sequence[typing.Dict[str, typing.Any]],
-        request_options: typing.Optional[RequestOptions] = None
+        config: typing.Dict[str, typing.Optional[typing.Any]],
+        messages: typing.Sequence[typing.Dict[str, typing.Optional[typing.Any]]],
+        request_options: typing.Optional[RequestOptions] = None,
     ) -> SqlAgentResponse:
         """
         Coming soon! Generate, execute, and test SQL queries. Returns an asset ID for the query object.
 
         Parameters
         ----------
-        config : typing.Dict[str, typing.Any]
+        config : typing.Dict[str, typing.Optional[typing.Any]]
             Configuration for the SQL agent including database connection details and query parameters
 
-        messages : typing.Sequence[typing.Dict[str, typing.Any]]
+        messages : typing.Sequence[typing.Dict[str, typing.Optional[typing.Any]]]
             The messages to send to the SQL agent
 
         request_options : typing.Optional[RequestOptions]
@@ -106,28 +126,56 @@ class AsyncSqlClient:
 
         Examples
         --------
-        from athena.client import AsyncAthena
+        import asyncio
+
+        from athena import AsyncAthena
 
         client = AsyncAthena(
             api_key="YOUR_API_KEY",
         )
-        await client.agents.sql.invoke(
-            config={"key": "value"},
-            messages=[{"key": "value"}],
-        )
+
+
+        async def main() -> None:
+            await client.agents.sql.invoke(
+                config={"key": "value"},
+                messages=[{"key": "value"}],
+            )
+
+
+        asyncio.run(main())
         """
         _response = await self._client_wrapper.httpx_client.request(
             "api/v0/agents/sql/invoke",
             method="POST",
-            json={"config": config, "messages": messages},
+            json={
+                "config": config,
+                "messages": messages,
+            },
+            headers={
+                "content-type": "application/json",
+            },
             request_options=request_options,
             omit=OMIT,
         )
-        if 200 <= _response.status_code < 300:
-            return pydantic_v1.parse_obj_as(SqlAgentResponse, _response.json())  # type: ignore
-        if _response.status_code == 422:
-            raise UnprocessableEntityError(pydantic_v1.parse_obj_as(typing.Any, _response.json()))  # type: ignore
         try:
+            if 200 <= _response.status_code < 300:
+                return typing.cast(
+                    SqlAgentResponse,
+                    parse_obj_as(
+                        type_=SqlAgentResponse,  # type: ignore
+                        object_=_response.json(),
+                    ),
+                )
+            if _response.status_code == 422:
+                raise UnprocessableEntityError(
+                    typing.cast(
+                        typing.Optional[typing.Any],
+                        parse_obj_as(
+                            type_=typing.Optional[typing.Any],  # type: ignore
+                            object_=_response.json(),
+                        ),
+                    )
+                )
             _response_json = _response.json()
         except JSONDecodeError:
             raise ApiError(status_code=_response.status_code, body=_response.text)

athena/base_client.py

@@ -1,14 +1,16 @@
 # This file was auto-generated by Fern from our API Definition.
 
 import typing
-
-import httpx
-
-from .agents.client import AgentsClient, AsyncAgentsClient
-from .core.client_wrapper import AsyncClientWrapper, SyncClientWrapper
 from .environment import AthenaEnvironment
-from .query.client import AsyncQueryClient, QueryClient
-from .tools.client import AsyncToolsClient, ToolsClient
+import httpx
+from .core.client_wrapper import SyncClientWrapper
+from .agents.client import AgentsClient
+from .query.client import QueryClient
+from .tools.client import ToolsClient
+from .core.client_wrapper import AsyncClientWrapper
+from .agents.client import AsyncAgentsClient
+from .query.client import AsyncQueryClient
+from .tools.client import AsyncToolsClient
 
 
 class BaseAthena:
@@ -41,7 +43,7 @@ class BaseAthena:
 
     Examples
     --------
-    from athena.client import Athena
+    from athena import Athena
 
     client = Athena(
         api_key="YOUR_API_KEY",
@@ -56,7 +58,7 @@ class BaseAthena:
         api_key: str,
         timeout: typing.Optional[float] = None,
         follow_redirects: typing.Optional[bool] = True,
-        httpx_client: typing.Optional[httpx.Client] = None
+        httpx_client: typing.Optional[httpx.Client] = None,
     ):
         _defaulted_timeout = timeout if timeout is not None else 60 if httpx_client is None else None
         self._client_wrapper = SyncClientWrapper(
@@ -104,7 +106,7 @@ class AsyncBaseAthena:
 
     Examples
     --------
-    from athena.client import AsyncAthena
+    from athena import AsyncAthena
 
     client = AsyncAthena(
         api_key="YOUR_API_KEY",
@@ -119,7 +121,7 @@ class AsyncBaseAthena:
         api_key: str,
         timeout: typing.Optional[float] = None,
         follow_redirects: typing.Optional[bool] = True,
-        httpx_client: typing.Optional[httpx.AsyncClient] = None
+        httpx_client: typing.Optional[httpx.AsyncClient] = None,
     ):
         _defaulted_timeout = timeout if timeout is not None else 60 if httpx_client is None else None
         self._client_wrapper = AsyncClientWrapper(

athena/client.py

@@ -7,8 +7,9 @@ import typing
 import warnings
 
 import httpx
-from typing import cast, List, Tuple, Union
+from typing import cast, List, Tuple, Union, Optional
 from typing_extensions import TypeVar, ParamSpec
+from langserve import RemoteRunnable
 
 from . import core
 from .base_client import BaseAthena, AsyncBaseAthena
@@ -24,23 +25,23 @@ if typing.TYPE_CHECKING:
 
 P = ParamSpec("P")
 T = TypeVar("T")
-U = TypeVar('U')
+U = TypeVar("U")
 
 
 def _inherit_signature_and_doc(
-    f: typing.Callable[P, T],
-    replace_in_doc: typing.Dict[str, str]
+    f: typing.Callable[P, T], replace_in_doc: typing.Dict[str, str]
 ) -> typing.Callable[..., typing.Callable[P, U]]:
     def decorator(decorated):
         for old, new in replace_in_doc.items():
             assert old in f.__doc__
             decorated.__doc__ = f.__doc__.replace(old, new)
         return decorated
+
     return decorator
 
 
 class SpecialEnvironments(enum.Enum):
-    AUTODETECT_ENVIRONMENT = 'AUTO'
+    AUTODETECT_ENVIRONMENT = "AUTO"
 
 
 @dataclasses.dataclass
@@ -53,7 +54,7 @@ class AthenaAsset:
         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]
+            from IPython import display  # type: ignore[import]
 
             code = display.Code(
                 data=self.data.decode(),
@@ -84,17 +85,19 @@ class WrappedToolsClient(ToolsClient):
         file_io = client.tools.get_file(asset_id="asset_id")
         pl.read_csv(file_io)
         """
-        file_bytes = b''.join(self.raw_data(asset_id=asset_id))
+        file_bytes = b"".join(self.raw_data(asset_id=asset_id))
         bytes_io = io.BytesIO(file_bytes)
         return bytes_io
 
-    @_inherit_signature_and_doc(ToolsClient.data_frame, {'DataFrameRequestOut': 'pd.DataFrame'})
-    def data_frame(self, *, asset_id: str, **kwargs) -> 'pd.DataFrame':
+    @_inherit_signature_and_doc(
+        ToolsClient.data_frame, {"DataFrameRequestOut": "pd.DataFrame"}
+    )
+    def data_frame(self, *, asset_id: str, **kwargs) -> "pd.DataFrame":
         _check_pandas_installed()
         model = super().data_frame(asset_id=asset_id, **kwargs)
         return _read_json_frame(model)
 
-    def read_data_frame(self, asset_id: str, *args, **kwargs) -> 'pd.DataFrame':
+    def read_data_frame(self, asset_id: str, *args, **kwargs) -> "pd.DataFrame":
         """
         Parameters
         ----------
@@ -119,8 +122,7 @@ class WrappedToolsClient(ToolsClient):
         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]
+    def save_asset(  # type: ignore[override]
         self,
         asset_object: Union["pd.DataFrame", "pd.Series", core.File],
         *,
@@ -157,10 +159,10 @@ 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, **kwargs)
-        return super().save_asset(
-            file=asset_object, parent_folder_id=parent_folder_id
+        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)
 
     def _get_file_and_media_type(self, asset_id: str) -> Tuple[io.BytesIO, str]:
         """
@@ -236,22 +238,26 @@ class WrappedToolsClient(ToolsClient):
             return _to_pandas_df(file_bytes, media_type=media_type)
 
         raise NotImplementedError("Assets of `{media_type}` type are not yet supported")
-    
+
 
 class WrappedQueryClient(QueryClient):
 
-    @_inherit_signature_and_doc(QueryClient.execute, {'DataFrameRequestOut': 'pd.DataFrame'})
-    def execute(self, *, sql_command: str, database_asset_ids: Union[str, List[str]], **kwargs) -> 'pd.DataFrame':
+    @_inherit_signature_and_doc(
+        QueryClient.execute, {"DataFrameRequestOut": "pd.DataFrame"}
+    )
+    def execute(
+        self, *, sql_command: str, database_asset_ids: Union[str, List[str]], **kwargs
+    ) -> "pd.DataFrame":
         _check_pandas_installed()
         model = super().execute(
-            sql_command=sql_command,
-            database_asset_ids=database_asset_ids,
-            **kwargs
+            sql_command=sql_command, database_asset_ids=database_asset_ids, **kwargs
         )
         return _read_json_frame(model)
 
-    @_inherit_signature_and_doc(QueryClient.execute_snippet, {'DataFrameRequestOut': 'pd.DataFrame'})
-    def execute_snippet(self, *, snippet_asset_id: str, **kwargs) -> 'pd.DataFrame':
+    @_inherit_signature_and_doc(
+        QueryClient.execute_snippet, {"DataFrameRequestOut": "pd.DataFrame"}
+    )
+    def execute_snippet(self, *, snippet_asset_id: str, **kwargs) -> "pd.DataFrame":
         _check_pandas_installed()
         model = super().execute_snippet(snippet_asset_id=snippet_asset_id, **kwargs)
         return _read_json_frame(model)
@@ -261,12 +267,11 @@ def _add_docs_for_async_variant(obj):
     def decorator(decorated):
         doc = obj.__doc__
         name = obj.__name__
-        decorated.__doc__ = (
-            doc
-            .replace('client = Athena', 'client = AsyncAthena')
-            .replace(f'client.tools.{name}', f'await client.tools.{name}')
-        )
+        decorated.__doc__ = doc.replace(
+            "client = Athena", "client = AsyncAthena"
+        ).replace(f"client.tools.{name}", f"await client.tools.{name}")
         return decorated
+
     return decorator
 
 
@@ -274,18 +279,20 @@ class WrappedAsyncToolsClient(AsyncToolsClient):
 
     @_add_docs_for_async_variant(WrappedToolsClient.get_file)
     async def get_file(self, asset_id: str) -> io.BytesIO:
-        file_bytes = b''.join([gen async for gen in self.raw_data(asset_id=asset_id)])
+        file_bytes = b"".join([gen async for gen in self.raw_data(asset_id=asset_id)])
         bytes_io = io.BytesIO(file_bytes)
         return bytes_io
 
-    @_inherit_signature_and_doc(AsyncToolsClient.data_frame, {'DataFrameRequestOut': 'pd.DataFrame'})
-    async def data_frame(self, *, asset_id: str, **kwargs) -> 'pd.DataFrame':
+    @_inherit_signature_and_doc(
+        AsyncToolsClient.data_frame, {"DataFrameRequestOut": "pd.DataFrame"}
+    )
+    async def data_frame(self, *, asset_id: str, **kwargs) -> "pd.DataFrame":
         _check_pandas_installed()
         model = await super().data_frame(asset_id=asset_id, **kwargs)
         return _read_json_frame(model)
 
     @_add_docs_for_async_variant(WrappedToolsClient.read_data_frame)
-    async def read_data_frame(self, asset_id: str, *args, **kwargs) -> 'pd.DataFrame':
+    async def read_data_frame(self, asset_id: str, *args, **kwargs) -> "pd.DataFrame":
         _check_pandas_installed()
         file_bytes = await self.get_file(asset_id)
         return _to_pandas_df(file_bytes, *args, **kwargs)
@@ -299,7 +306,9 @@ class WrappedAsyncToolsClient(AsyncToolsClient):
         name: Union[str, None] = None,
         **kwargs,
     ) -> SaveAssetRequestOut:
-        asset_object = _convert_asset_object(asset_object=asset_object, name=name, **kwargs)
+        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
         )
@@ -307,23 +316,109 @@ class WrappedAsyncToolsClient(AsyncToolsClient):
 
 class WrappedAsyncQueryClient(AsyncQueryClient):
 
-    @_inherit_signature_and_doc(AsyncQueryClient.execute, {'DataFrameRequestOut': 'pd.DataFrame'})
-    async def execute(self, *, sql_command: str, database_asset_ids: Union[str, List[str]], **kwargs) -> 'pd.DataFrame':
+    @_inherit_signature_and_doc(
+        AsyncQueryClient.execute, {"DataFrameRequestOut": "pd.DataFrame"}
+    )
+    async def execute(
+        self, *, sql_command: str, database_asset_ids: Union[str, List[str]], **kwargs
+    ) -> "pd.DataFrame":
         _check_pandas_installed()
         model = await super().execute(
-            sql_command=sql_command,
-            database_asset_ids=database_asset_ids,
-            **kwargs
+            sql_command=sql_command, database_asset_ids=database_asset_ids, **kwargs
         )
         return _read_json_frame(model)
 
-    @_inherit_signature_and_doc(AsyncQueryClient.execute_snippet, {'DataFrameRequestOut': 'pd.DataFrame'})
-    async def execute_snippet(self, *, snippet_asset_id: str, **kwargs) -> 'pd.DataFrame':
+    @_inherit_signature_and_doc(
+        AsyncQueryClient.execute_snippet, {"DataFrameRequestOut": "pd.DataFrame"}
+    )
+    async def execute_snippet(
+        self, *, snippet_asset_id: str, **kwargs
+    ) -> "pd.DataFrame":
         _check_pandas_installed()
-        model = await super().execute_snippet(snippet_asset_id=snippet_asset_id, **kwargs)
+        model = await super().execute_snippet(
+            snippet_asset_id=snippet_asset_id, **kwargs
+        )
         return _read_json_frame(model)
 
 
+class AthenaModel(RemoteRunnable):
+    """Use Athena's models directly with a Langchain-compatible client.
+
+    The Langchain Runnable interface is supported:
+    - `invoke`: Invoke the model with a string or a Langchain message.
+    - `batch`: Batch invoke the model on multiple inputs.
+    - `astream_events`: Streaming for real-time applications.
+
+    See Langchain documentation for more details.
+
+    Examples
+    --------
+    from src.athena.client import Athena
+    import asyncio
+    from langchain_core.messages import HumanMessage
+
+    client = Athena()
+    llm = client.llm
+
+    # sync invoke -- use strings or langchain messages
+    result = llm.invoke("Hello")
+    print(result.content)
+
+    result = llm.invoke("Hello")
+    print(result.content)
+
+    # choose the model explicitly
+    claude = llm.with_config(configurable={"model": "claude_3_7_sonnet"})
+    print(claude.invoke("Who are you?").content)
+
+    # batch (for multiple parallel requests)
+    results = llm.batch(["Hello", "World"] * 5)
+    print([r.content for r in results])
+
+    # handle stream events
+    async def stream_response():
+        separator = "---------------------------------------"
+        print(separator)
+        print("Starting stream")
+        print(separator)
+        async for event in llm.astream_events("Hello.  Please respond in 5 sentences."):
+            data = event["data"]
+            if "chunk" in data:
+                print(data["chunk"].content)
+            elif "output" in data:
+                print(separator)
+                print("Final response")
+                print(separator)
+                print(data["output"].content)
+            else:
+                print("(other event) ", event.keys())
+
+
+    asyncio.run(stream_response())
+    from athena.client import Athena
+    client = Athena()
+    llm = client.llm
+
+    # sync invoke -- use strings or langchain messages
+    result = llm.invoke("Hello")
+    print(result.content)
+
+    # batch
+    results = llm.batch(["Hello", "World"])
+    print(r.content for r in results)
+    """
+
+    def __init__(self, base_url: str, api_key: str, timeout: Optional[float] = None):
+        self.base_url = base_url
+        self.api_key = api_key
+        self.timeout = timeout
+        super().__init__(
+            base_url + "/api/models/general",
+            headers={"X-API-KEY": api_key},
+            timeout=timeout,
+        )
+
+
 class Athena(BaseAthena):
     """
     Use this class to access the different functions within the SDK. You can instantiate any number of clients with different configuration that will propogate to these functions.
@@ -360,37 +455,40 @@ class Athena(BaseAthena):
         self,
         *,
         base_url: typing.Optional[str] = None,
-        environment: Union[AthenaEnvironment, SpecialEnvironments] = SpecialEnvironments.AUTODETECT_ENVIRONMENT,    # type: ignore[arg-type]
+        environment: Union[AthenaEnvironment, SpecialEnvironments] = SpecialEnvironments.AUTODETECT_ENVIRONMENT,  # type: ignore[arg-type]
         api_key: typing.Optional[str] = None,
         timeout: typing.Optional[float] = 60,
-        httpx_client: typing.Optional[httpx.Client] = None
+        httpx_client: typing.Optional[httpx.Client] = None,
     ):
         if api_key is None:
             try:
-                api_key = os.environ['ATHENA_API_KEY']
+                api_key = os.environ["ATHENA_API_KEY"]
             except KeyError:
                 raise TypeError(
                     "Athena() missing 1 required keyword-only argument: 'api_key'"
-                    ' (ATHENA_API_KEY environment variable not found)'
+                    " (ATHENA_API_KEY environment variable not found)"
                 )
         if environment == SpecialEnvironments.AUTODETECT_ENVIRONMENT:
-            if 'ATHENA_API_URL' in os.environ:
+            if "ATHENA_API_URL" in os.environ:
 
                 class AutodetectedEnvironments(enum.Enum):
-                    CURRENT = os.environ['ATHENA_API_URL']
+                    CURRENT = os.environ["ATHENA_API_URL"]
 
                 environment = cast(AthenaEnvironment, AutodetectedEnvironments.CURRENT)
             else:
                 environment = AthenaEnvironment.PRODUCTION
         super().__init__(
             base_url=base_url,
-            environment=environment,    # type: ignore[arg-type]
+            environment=environment,  # type: ignore[arg-type]
             api_key=api_key,
             timeout=timeout,
             httpx_client=httpx_client,
         )
         self.tools = WrappedToolsClient(client_wrapper=self._client_wrapper)
         self.query = WrappedQueryClient(client_wrapper=self._client_wrapper)
+        self.llm = AthenaModel(
+            base_url=base_url or environment.value, api_key=api_key, timeout=timeout
+        )
 
 
 class AsyncAthena(AsyncBaseAthena):
@@ -430,31 +528,31 @@ class AsyncAthena(AsyncBaseAthena):
         self,
         *,
         base_url: typing.Optional[str] = None,
-        environment: Union[AthenaEnvironment, SpecialEnvironments] = SpecialEnvironments.AUTODETECT_ENVIRONMENT,    # type: ignore[arg-type]
+        environment: Union[AthenaEnvironment, SpecialEnvironments] = SpecialEnvironments.AUTODETECT_ENVIRONMENT,  # type: ignore[arg-type]
         api_key: typing.Optional[str] = None,
         timeout: typing.Optional[float] = 60,
-        httpx_client: typing.Optional[httpx.AsyncClient] = None
+        httpx_client: typing.Optional[httpx.AsyncClient] = None,
     ):
         if api_key is None:
             try:
-                api_key = os.environ['ATHENA_API_KEY']
+                api_key = os.environ["ATHENA_API_KEY"]
             except KeyError:
                 raise TypeError(
                     "AsyncAthena() missing 1 required keyword-only argument: 'api_key'"
-                    ' (ATHENA_API_KEY environment variable not found)'
+                    " (ATHENA_API_KEY environment variable not found)"
                 )
         if environment == SpecialEnvironments.AUTODETECT_ENVIRONMENT:
-            if 'ATHENA_API_URL' in os.environ:
+            if "ATHENA_API_URL" in os.environ:
 
                 class AutodetectedEnvironments(enum.Enum):
-                    CURRENT = os.environ['ATHENA_API_URL']
+                    CURRENT = os.environ["ATHENA_API_URL"]
 
                 environment = cast(AthenaEnvironment, AutodetectedEnvironments.CURRENT)
             else:
                 environment = AthenaEnvironment.PRODUCTION
         super().__init__(
             base_url=base_url,
-            environment=environment,    # type: ignore[arg-type]
+            environment=environment,  # type: ignore[arg-type]
             api_key=api_key,
             timeout=timeout,
             httpx_client=httpx_client,
@@ -463,19 +561,20 @@ class AsyncAthena(AsyncBaseAthena):
         self.query = WrappedAsyncQueryClient(client_wrapper=self._client_wrapper)
 
 
-def _read_json_frame(model: DataFrameRequestOut) -> 'pd.DataFrame':
+def _read_json_frame(model: DataFrameRequestOut) -> "pd.DataFrame":
     import pandas as pd
 
     string_io = io.StringIO(model.json())
 
     with warnings.catch_warnings():
         # Filter warnings due to https://github.com/pandas-dev/pandas/issues/59511
-        warnings.simplefilter(action='ignore', category=FutureWarning)
-        return pd.read_json(string_io, orient='split')
+        warnings.simplefilter(action="ignore", category=FutureWarning)
+        return pd.read_json(string_io, orient="split")
 
 
 def _check_pandas_installed():
     import pandas
+
     assert pandas
 
 
@@ -524,11 +623,12 @@ def _to_pandas_df(
 def _convert_asset_object(
     asset_object: Union["pd.DataFrame", "pd.Series", core.File],
     name: Union[str, None] = None,
-    **kwargs
+    **kwargs,
 ) -> core.File:
     import pandas as pd
+
     try:
-        from IPython.core.formatters import format_display_data    # type: ignore[import]
+        from IPython.core.formatters import format_display_data  # type: ignore[import]
     except ImportError:
         format_display_data = None
 
@@ -559,4 +659,4 @@ def _convert_asset_object(
                     image_bytes,
                     media_type,
                 )
-    return asset_object    # type: ignore[return-value]
+    return asset_object  # type: ignore[return-value]