athena-intelligence 0.1.87__py3-none-any.whl → 0.1.89__py3-none-any.whl

athena/base_client.py

@@ -20,7 +20,7 @@ from .workflow.client import AsyncWorkflowClient, WorkflowClient
 
 class 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.
+    Use this class to access the different functions within the SDK. You can instantiate any number of clients with different configuration that will propagate to these functions.
 
     Parameters
     ----------
@@ -38,7 +38,7 @@ class BaseAthena:
 
     api_key : str
     timeout : typing.Optional[float]
-        The timeout to be used, in seconds, for requests by default the timeout is 60 seconds, unless a custom httpx client is used, in which case a default is not set.
+        The timeout to be used, in seconds, for requests. By default the timeout is 60 seconds, unless a custom httpx client is used, in which case this default is not enforced.
 
     follow_redirects : typing.Optional[bool]
         Whether the default httpx client follows redirects or not, this is irrelevant if a custom httpx client is passed in.
@@ -90,7 +90,7 @@ class BaseAthena:
 
 class AsyncBaseAthena:
     """
-    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.
+    Use this class to access the different functions within the SDK. You can instantiate any number of clients with different configuration that will propagate to these functions.
 
     Parameters
     ----------
@@ -108,7 +108,7 @@ class AsyncBaseAthena:
 
     api_key : str
     timeout : typing.Optional[float]
-        The timeout to be used, in seconds, for requests by default the timeout is 60 seconds, unless a custom httpx client is used, in which case a default is not set.
+        The timeout to be used, in seconds, for requests. By default the timeout is 60 seconds, unless a custom httpx client is used, in which case this default is not enforced.
 
     follow_redirects : typing.Optional[bool]
         Whether the default httpx client follows redirects or not, this is irrelevant if a custom httpx client is passed in.

athena/chain/client.py

@@ -1,15 +1,11 @@
 # This file was auto-generated by Fern from our API Definition.
 
 import typing
-import urllib.parse
 from json.decoder import JSONDecodeError
 
 from ..core.api_error import ApiError
 from ..core.client_wrapper import AsyncClientWrapper, SyncClientWrapper
-from ..core.jsonable_encoder import jsonable_encoder
 from ..core.pydantic_utilities import pydantic_v1
-from ..core.query_encoder import encode_query
-from ..core.remove_none_from_dict import remove_none_from_dict
 from ..core.request_options import RequestOptions
 from ..errors.unprocessable_entity_error import UnprocessableEntityError
 from ..types.document import Document
@@ -32,7 +28,7 @@ class ChainClient:
         text_input: str,
         custom_type_dict: typing.Dict[str, typing.Any],
         model: LlmModel,
-        request_options: typing.Optional[RequestOptions] = None,
+        request_options: typing.Optional[RequestOptions] = None
     ) -> StructuredParseResult:
         """
         Parameters
@@ -68,32 +64,11 @@ class ChainClient:
         )
         """
         _response = self._client_wrapper.httpx_client.request(
+            "api/v0/structured-parse",
             method="POST",
-            url=urllib.parse.urljoin(f"{self._client_wrapper.get_base_url()}/", "api/v0/structured-parse"),
-            params=encode_query(
-                jsonable_encoder(
-                    request_options.get("additional_query_parameters") if request_options is not None else None
-                )
-            ),
-            json=jsonable_encoder({"text_input": text_input, "custom_type_dict": custom_type_dict, "model": model})
-            if request_options is None or request_options.get("additional_body_parameters") is None
-            else {
-                **jsonable_encoder({"text_input": text_input, "custom_type_dict": custom_type_dict, "model": model}),
-                **(jsonable_encoder(remove_none_from_dict(request_options.get("additional_body_parameters", {})))),
-            },
-            headers=jsonable_encoder(
-                remove_none_from_dict(
-                    {
-                        **self._client_wrapper.get_headers(),
-                        **(request_options.get("additional_headers", {}) if request_options is not None else {}),
-                    }
-                )
-            ),
-            timeout=request_options.get("timeout_in_seconds")
-            if request_options is not None and request_options.get("timeout_in_seconds") is not None
-            else self._client_wrapper.get_timeout(),
-            retries=0,
-            max_retries=request_options.get("max_retries") if request_options is not None else 0,  # type: ignore
+            json={"text_input": text_input, "custom_type_dict": custom_type_dict, "model": model},
+            request_options=request_options,
+            omit=OMIT,
         )
         if 200 <= _response.status_code < 300:
             return pydantic_v1.parse_obj_as(StructuredParseResult, _response.json())  # type: ignore
@@ -115,7 +90,7 @@ class ChainClient:
         operator_prompt: str,
         reducer_prompt: str,
         input: str,
-        request_options: typing.Optional[RequestOptions] = None,
+        request_options: typing.Optional[RequestOptions] = None
     ) -> MapReduceChainOut:
         """
         Parameters
@@ -160,48 +135,17 @@ class ChainClient:
         )
         """
         _response = self._client_wrapper.httpx_client.request(
+            "api/v0/tools/map-reduce",
             method="POST",
-            url=urllib.parse.urljoin(f"{self._client_wrapper.get_base_url()}/", "api/v0/tools/map-reduce"),
-            params=encode_query(
-                jsonable_encoder(
-                    request_options.get("additional_query_parameters") if request_options is not None else None
-                )
-            ),
-            json=jsonable_encoder(
-                {
-                    "documents": documents,
-                    "model": model,
-                    "operator_prompt": operator_prompt,
-                    "reducer_prompt": reducer_prompt,
-                    "input": input,
-                }
-            )
-            if request_options is None or request_options.get("additional_body_parameters") is None
-            else {
-                **jsonable_encoder(
-                    {
-                        "documents": documents,
-                        "model": model,
-                        "operator_prompt": operator_prompt,
-                        "reducer_prompt": reducer_prompt,
-                        "input": input,
-                    }
-                ),
-                **(jsonable_encoder(remove_none_from_dict(request_options.get("additional_body_parameters", {})))),
+            json={
+                "documents": documents,
+                "model": model,
+                "operator_prompt": operator_prompt,
+                "reducer_prompt": reducer_prompt,
+                "input": input,
             },
-            headers=jsonable_encoder(
-                remove_none_from_dict(
-                    {
-                        **self._client_wrapper.get_headers(),
-                        **(request_options.get("additional_headers", {}) if request_options is not None else {}),
-                    }
-                )
-            ),
-            timeout=request_options.get("timeout_in_seconds")
-            if request_options is not None and request_options.get("timeout_in_seconds") is not None
-            else self._client_wrapper.get_timeout(),
-            retries=0,
-            max_retries=request_options.get("max_retries") if request_options is not None else 0,  # type: ignore
+            request_options=request_options,
+            omit=OMIT,
         )
         if 200 <= _response.status_code < 300:
             return pydantic_v1.parse_obj_as(MapReduceChainOut, _response.json())  # type: ignore
@@ -226,7 +170,7 @@ class AsyncChainClient:
         text_input: str,
         custom_type_dict: typing.Dict[str, typing.Any],
         model: LlmModel,
-        request_options: typing.Optional[RequestOptions] = None,
+        request_options: typing.Optional[RequestOptions] = None
     ) -> StructuredParseResult:
         """
         Parameters
@@ -262,32 +206,11 @@ class AsyncChainClient:
         )
         """
         _response = await self._client_wrapper.httpx_client.request(
+            "api/v0/structured-parse",
             method="POST",
-            url=urllib.parse.urljoin(f"{self._client_wrapper.get_base_url()}/", "api/v0/structured-parse"),
-            params=encode_query(
-                jsonable_encoder(
-                    request_options.get("additional_query_parameters") if request_options is not None else None
-                )
-            ),
-            json=jsonable_encoder({"text_input": text_input, "custom_type_dict": custom_type_dict, "model": model})
-            if request_options is None or request_options.get("additional_body_parameters") is None
-            else {
-                **jsonable_encoder({"text_input": text_input, "custom_type_dict": custom_type_dict, "model": model}),
-                **(jsonable_encoder(remove_none_from_dict(request_options.get("additional_body_parameters", {})))),
-            },
-            headers=jsonable_encoder(
-                remove_none_from_dict(
-                    {
-                        **self._client_wrapper.get_headers(),
-                        **(request_options.get("additional_headers", {}) if request_options is not None else {}),
-                    }
-                )
-            ),
-            timeout=request_options.get("timeout_in_seconds")
-            if request_options is not None and request_options.get("timeout_in_seconds") is not None
-            else self._client_wrapper.get_timeout(),
-            retries=0,
-            max_retries=request_options.get("max_retries") if request_options is not None else 0,  # type: ignore
+            json={"text_input": text_input, "custom_type_dict": custom_type_dict, "model": model},
+            request_options=request_options,
+            omit=OMIT,
         )
         if 200 <= _response.status_code < 300:
             return pydantic_v1.parse_obj_as(StructuredParseResult, _response.json())  # type: ignore
@@ -309,7 +232,7 @@ class AsyncChainClient:
         operator_prompt: str,
         reducer_prompt: str,
         input: str,
-        request_options: typing.Optional[RequestOptions] = None,
+        request_options: typing.Optional[RequestOptions] = None
     ) -> MapReduceChainOut:
         """
         Parameters
@@ -354,48 +277,17 @@ class AsyncChainClient:
         )
         """
         _response = await self._client_wrapper.httpx_client.request(
+            "api/v0/tools/map-reduce",
             method="POST",
-            url=urllib.parse.urljoin(f"{self._client_wrapper.get_base_url()}/", "api/v0/tools/map-reduce"),
-            params=encode_query(
-                jsonable_encoder(
-                    request_options.get("additional_query_parameters") if request_options is not None else None
-                )
-            ),
-            json=jsonable_encoder(
-                {
-                    "documents": documents,
-                    "model": model,
-                    "operator_prompt": operator_prompt,
-                    "reducer_prompt": reducer_prompt,
-                    "input": input,
-                }
-            )
-            if request_options is None or request_options.get("additional_body_parameters") is None
-            else {
-                **jsonable_encoder(
-                    {
-                        "documents": documents,
-                        "model": model,
-                        "operator_prompt": operator_prompt,
-                        "reducer_prompt": reducer_prompt,
-                        "input": input,
-                    }
-                ),
-                **(jsonable_encoder(remove_none_from_dict(request_options.get("additional_body_parameters", {})))),
+            json={
+                "documents": documents,
+                "model": model,
+                "operator_prompt": operator_prompt,
+                "reducer_prompt": reducer_prompt,
+                "input": input,
             },
-            headers=jsonable_encoder(
-                remove_none_from_dict(
-                    {
-                        **self._client_wrapper.get_headers(),
-                        **(request_options.get("additional_headers", {}) if request_options is not None else {}),
-                    }
-                )
-            ),
-            timeout=request_options.get("timeout_in_seconds")
-            if request_options is not None and request_options.get("timeout_in_seconds") is not None
-            else self._client_wrapper.get_timeout(),
-            retries=0,
-            max_retries=request_options.get("max_retries") if request_options is not None else 0,  # type: ignore
+            request_options=request_options,
+            omit=OMIT,
         )
         if 200 <= _response.status_code < 300:
             return pydantic_v1.parse_obj_as(MapReduceChainOut, _response.json())  # type: ignore

athena/client.py

@@ -1,4 +1,5 @@
 import io
+import os
 import typing
 import warnings
 
@@ -20,42 +21,103 @@ T = TypeVar("T")
 U = TypeVar('U')
 
 
-def _inherit_signature(f: typing.Callable[P, T]) -> typing.Callable[..., typing.Callable[P, U]]:
-    return lambda x: x
+def _inherit_signature_and_doc(
+    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 WrappedToolsClient(ToolsClient):
 
     def get_file(self, document_id: str) -> io.BytesIO:
+        """
+        Parameters
+        ----------
+        document_id : str
+
+        Returns
+        -------
+        io.BytesIO
+
+        Examples
+        --------
+        import polars as pl
+        from athena.client import Athena
+
+        client = Athena(api_key="YOUR_API_KEY")
+        file_io = client.tools.get_file(document_id="document_id")
+        pl.read_csv(file_io)
+        """
         file_bytes = b''.join(self.raw_data(document_id=document_id))
         bytes_io = io.BytesIO(file_bytes)
         return bytes_io
 
-    @_inherit_signature(ToolsClient.data_frame)
+    @_inherit_signature_and_doc(ToolsClient.data_frame, {'DataFrameRequestOut': 'pd.DataFrame'})
     def data_frame(self, *, document_id: str, **kwargs) -> 'pd.DataFrame':
         _check_pandas_installed()
         model = super().data_frame(document_id=document_id, **kwargs)
         return _read_json_frame(model)
 
     def read_data_frame(self, document_id: str, *args, **kwargs) -> 'pd.DataFrame':
+        """
+        Parameters
+        ----------
+        document_id : str
+
+        **kwargs : dict
+          keyword arguments passed to pandas `read_csv` or `read_excel` function,
+          depending on the file type of the document identified by `document_id`.
+
+        Returns
+        -------
+        pd.DataFrame
+
+        Examples
+        --------
+        from athena.client import Athena
+
+        client = Athena(api_key="YOUR_API_KEY")
+        client.tools.read_data_frame(document_id="document_id")
+        """
         _check_pandas_installed()
         file_bytes = self.get_file(document_id)
         return _to_pandas_df(file_bytes, *args, **kwargs)
 
 
+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}')
+        )
+        return decorated
+    return decorator
+
+
 class WrappedAsyncToolsClient(AsyncToolsClient):
 
+    @_add_docs_for_async_variant(WrappedToolsClient.get_file)
     async def get_file(self, document_id: str) -> io.BytesIO:
         file_bytes = b''.join([gen async for gen in self.raw_data(document_id=document_id)])
         bytes_io = io.BytesIO(file_bytes)
         return bytes_io
 
-    @_inherit_signature(ToolsClient.data_frame)
+    @_inherit_signature_and_doc(AsyncToolsClient.data_frame, {'DataFrameRequestOut': 'pd.DataFrame'})
     async def data_frame(self, *, document_id: str, **kwargs) -> 'pd.DataFrame':
         _check_pandas_installed()
         model = await super().data_frame(document_id=document_id, **kwargs)
         return _read_json_frame(model)
 
+    @_add_docs_for_async_variant(WrappedToolsClient.read_data_frame)
     async def read_data_frame(self, document_id: str, *args, **kwargs) -> 'pd.DataFrame':
         _check_pandas_installed()
         file_bytes = await self.get_file(document_id)
@@ -66,19 +128,29 @@ 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.
 
-    Parameters:
-        - base_url: typing.Optional[str]. The base url to use for requests from the client.
+    Parameters
+    ----------
+    base_url : typing.Optional[str]
+        The base url to use for requests from the client.
+
+    environment : AthenaEnvironment
+        The environment to use for requests from the client. from .environment import AthenaEnvironment
 
-        - environment: AthenaEnvironment. The environment to use for requests from the client. from .environment import AthenaEnvironment
 
-                                          Defaults to AthenaEnvironment.DEFAULT
 
-        - api_key: str.
+        Defaults to AthenaEnvironment.DEFAULT
 
-        - timeout: typing.Optional[float]. The timeout to be used, in seconds, for requests by default the timeout is 60 seconds.
 
-        - httpx_client: typing.Optional[httpx.Client]. The httpx client to use for making requests, a preconfigured client is used by default, however this is useful should you want to pass in any custom httpx configuration.
-    ---
+
+    api_key: typing.Optional[str].  The API key. Required when used outside of the athena notebook environment.
+    timeout : typing.Optional[float]
+        The timeout to be used, in seconds, for requests. By default the timeout is 60 seconds, unless a custom httpx client is used, in which case this default is not enforced.
+
+    httpx_client : typing.Optional[httpx.Client]
+        The httpx client to use for making requests, a preconfigured client is used by default, however this is useful should you want to pass in any custom httpx configuration.
+
+    Examples
+    --------
     from athena.client import Athena
 
     client = Athena(
@@ -86,15 +158,25 @@ class Athena(BaseAthena):
     )
     """
 
+    tools: WrappedToolsClient
+
     def __init__(
         self,
         *,
         base_url: typing.Optional[str] = None,
         environment: AthenaEnvironment = AthenaEnvironment.DEFAULT,
-        api_key: str,
+        api_key: typing.Optional[str] = None,
         timeout: typing.Optional[float] = 60,
         httpx_client: typing.Optional[httpx.Client] = None
     ):
+        if api_key is None:
+            try:
+                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)'
+                )
         super().__init__(
             base_url=base_url,
             environment=environment,
@@ -109,21 +191,32 @@ class Athena(BaseAthena):
 
 class AsyncAthena(AsyncBaseAthena):
     """
-    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.
+    Use this class to access the different functions within the SDK. You can instantiate any number of clients with different configuration that will propagate to these functions.
+
+
+    Parameters
+    ----------
+    base_url : typing.Optional[str]
+        The base url to use for requests from the client.
+
+    environment : AthenaEnvironment
+        The environment to use for requests from the client. from .environment import AthenaEnvironment
 
-    Parameters:
-        - base_url: typing.Optional[str]. The base url to use for requests from the client.
 
-        - environment: AthenaEnvironment. The environment to use for requests from the client. from .environment import AthenaEnvironment
 
-                                          Defaults to AthenaEnvironment.DEFAULT
+        Defaults to AthenaEnvironment.DEFAULT
 
-        - api_key: str.
 
-        - timeout: typing.Optional[float]. The timeout to be used, in seconds, for requests by default the timeout is 60 seconds.
 
-        - httpx_client: typing.Optional[httpx.AsyncClient]. The httpx client to use for making requests, a preconfigured client is used by default, however this is useful should you want to pass in any custom httpx configuration.
-    ---
+    api_key: typing.Optional[str].  The API key. Required when used outside of the athena notebook environment.
+    timeout : typing.Optional[float]
+        The timeout to be used, in seconds, for requests. By default the timeout is 60 seconds, unless a custom httpx client is used, in which case this default is not enforced.
+
+    httpx_client : typing.Optional[httpx.AsyncClient]
+        The httpx client to use for making requests, a preconfigured client is used by default, however this is useful should you want to pass in any custom httpx configuration.
+
+    Examples
+    --------
     from athena.client import AsyncAthena
 
     client = AsyncAthena(
@@ -131,15 +224,25 @@ class AsyncAthena(AsyncBaseAthena):
     )
     """
 
+    tools: WrappedAsyncToolsClient
+
     def __init__(
         self,
         *,
         base_url: typing.Optional[str] = None,
         environment: AthenaEnvironment = AthenaEnvironment.DEFAULT,
-        api_key: str,
+        api_key: typing.Optional[str] = None,
         timeout: typing.Optional[float] = 60,
         httpx_client: typing.Optional[httpx.AsyncClient] = None
     ):
+        if api_key is None:
+            try:
+                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)'
+                )
         super().__init__(
             base_url=base_url,
             environment=environment,

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.87",
+            "X-Fern-SDK-Version": "0.1.89",
         }
         headers["X-API-KEY"] = self.api_key
         return headers
@@ -34,7 +34,12 @@ class SyncClientWrapper(BaseClientWrapper):
         self, *, api_key: str, base_url: str, timeout: typing.Optional[float] = None, httpx_client: httpx.Client
     ):
         super().__init__(api_key=api_key, base_url=base_url, timeout=timeout)
-        self.httpx_client = HttpClient(httpx_client=httpx_client)
+        self.httpx_client = HttpClient(
+            httpx_client=httpx_client,
+            base_headers=self.get_headers(),
+            base_timeout=self.get_timeout(),
+            base_url=self.get_base_url(),
+        )
 
 
 class AsyncClientWrapper(BaseClientWrapper):
@@ -42,4 +47,9 @@ class AsyncClientWrapper(BaseClientWrapper):
         self, *, api_key: str, base_url: str, timeout: typing.Optional[float] = None, httpx_client: httpx.AsyncClient
     ):
         super().__init__(api_key=api_key, base_url=base_url, timeout=timeout)
-        self.httpx_client = AsyncHttpClient(httpx_client=httpx_client)
+        self.httpx_client = AsyncHttpClient(
+            httpx_client=httpx_client,
+            base_headers=self.get_headers(),
+            base_timeout=self.get_timeout(),
+            base_url=self.get_base_url(),
+        )