athena-intelligence 0.1.82__py3-none-any.whl → 0.1.84__py3-none-any.whl

athena/base_client.py

@@ -22,21 +22,32 @@ 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.
 
-    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.
 
-        - 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.
+        Defaults to AthenaEnvironment.DEFAULT
 
-        - 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.
 
-        - 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 : 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.
+
+    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.
+
+    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(
@@ -81,21 +92,32 @@ 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.
 
-    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
+
+
+
+        Defaults to AthenaEnvironment.DEFAULT
 
-        - environment: AthenaEnvironment. The environment to use for requests from the client. from .environment import AthenaEnvironment
 
-                                          Defaults to AthenaEnvironment.DEFAULT
 
-        - api_key: str.
+    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.
 
-        - 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.
+    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.
 
-        - 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.
+    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.
 
-        - 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(

athena/chain/client.py

@@ -8,6 +8,7 @@ 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
@@ -34,15 +35,26 @@ class ChainClient:
         request_options: typing.Optional[RequestOptions] = None,
     ) -> StructuredParseResult:
         """
-        Parameters:
-            - text_input: str. The text input to be parsed.
+        Parameters
+        ----------
+        text_input : str
+            The text input to be parsed.
 
-            - custom_type_dict: typing.Dict[str, typing.Any]. A dictionary of field names and their default values.
+        custom_type_dict : typing.Dict[str, typing.Any]
+            A dictionary of field names and their default values.
 
-            - model: LlmModel.
+        model : LlmModel
 
-            - request_options: typing.Optional[RequestOptions]. Request-specific configuration.
-        ---
+        request_options : typing.Optional[RequestOptions]
+            Request-specific configuration.
+
+        Returns
+        -------
+        StructuredParseResult
+            Successful Response
+
+        Examples
+        --------
         from athena import LlmModel
         from athena.client import Athena
 
@@ -58,8 +70,10 @@ class ChainClient:
         _response = self._client_wrapper.httpx_client.request(
             method="POST",
             url=urllib.parse.urljoin(f"{self._client_wrapper.get_base_url()}/", "api/v0/structured-parse"),
-            params=jsonable_encoder(
-                request_options.get("additional_query_parameters") if request_options is not None else None
+            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
@@ -104,19 +118,28 @@ class ChainClient:
         request_options: typing.Optional[RequestOptions] = None,
     ) -> MapReduceChainOut:
         """
-        Parameters:
-            - documents: typing.Sequence[Document].
+        Parameters
+        ----------
+        documents : typing.Sequence[Document]
 
-            - model: LlmModel.
+        model : LlmModel
 
-            - operator_prompt: str.
+        operator_prompt : str
 
-            - reducer_prompt: str.
+        reducer_prompt : str
 
-            - input: str.
+        input : str
 
-            - request_options: typing.Optional[RequestOptions]. Request-specific configuration.
-        ---
+        request_options : typing.Optional[RequestOptions]
+            Request-specific configuration.
+
+        Returns
+        -------
+        MapReduceChainOut
+            Successful Response
+
+        Examples
+        --------
         from athena import Document, LlmModel
         from athena.client import Athena
 
@@ -139,8 +162,10 @@ class ChainClient:
         _response = self._client_wrapper.httpx_client.request(
             method="POST",
             url=urllib.parse.urljoin(f"{self._client_wrapper.get_base_url()}/", "api/v0/tools/map-reduce"),
-            params=jsonable_encoder(
-                request_options.get("additional_query_parameters") if request_options is not None else None
+            params=encode_query(
+                jsonable_encoder(
+                    request_options.get("additional_query_parameters") if request_options is not None else None
+                )
             ),
             json=jsonable_encoder(
                 {
@@ -204,15 +229,26 @@ class AsyncChainClient:
         request_options: typing.Optional[RequestOptions] = None,
     ) -> StructuredParseResult:
         """
-        Parameters:
-            - text_input: str. The text input to be parsed.
+        Parameters
+        ----------
+        text_input : str
+            The text input to be parsed.
 
-            - custom_type_dict: typing.Dict[str, typing.Any]. A dictionary of field names and their default values.
+        custom_type_dict : typing.Dict[str, typing.Any]
+            A dictionary of field names and their default values.
 
-            - model: LlmModel.
+        model : LlmModel
 
-            - request_options: typing.Optional[RequestOptions]. Request-specific configuration.
-        ---
+        request_options : typing.Optional[RequestOptions]
+            Request-specific configuration.
+
+        Returns
+        -------
+        StructuredParseResult
+            Successful Response
+
+        Examples
+        --------
         from athena import LlmModel
         from athena.client import AsyncAthena
 
@@ -228,8 +264,10 @@ class AsyncChainClient:
         _response = await self._client_wrapper.httpx_client.request(
             method="POST",
             url=urllib.parse.urljoin(f"{self._client_wrapper.get_base_url()}/", "api/v0/structured-parse"),
-            params=jsonable_encoder(
-                request_options.get("additional_query_parameters") if request_options is not None else None
+            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
@@ -274,19 +312,28 @@ class AsyncChainClient:
         request_options: typing.Optional[RequestOptions] = None,
     ) -> MapReduceChainOut:
         """
-        Parameters:
-            - documents: typing.Sequence[Document].
+        Parameters
+        ----------
+        documents : typing.Sequence[Document]
 
-            - model: LlmModel.
+        model : LlmModel
 
-            - operator_prompt: str.
+        operator_prompt : str
 
-            - reducer_prompt: str.
+        reducer_prompt : str
 
-            - input: str.
+        input : str
 
-            - request_options: typing.Optional[RequestOptions]. Request-specific configuration.
-        ---
+        request_options : typing.Optional[RequestOptions]
+            Request-specific configuration.
+
+        Returns
+        -------
+        MapReduceChainOut
+            Successful Response
+
+        Examples
+        --------
         from athena import Document, LlmModel
         from athena.client import AsyncAthena
 
@@ -309,8 +356,10 @@ class AsyncChainClient:
         _response = await self._client_wrapper.httpx_client.request(
             method="POST",
             url=urllib.parse.urljoin(f"{self._client_wrapper.get_base_url()}/", "api/v0/tools/map-reduce"),
-            params=jsonable_encoder(
-                request_options.get("additional_query_parameters") if request_options is not None else None
+            params=encode_query(
+                jsonable_encoder(
+                    request_options.get("additional_query_parameters") if request_options is not None else None
+                )
             ),
             json=jsonable_encoder(
                 {

athena/client.py

@@ -1,13 +1,65 @@
-# This file was auto-generated by Fern from our API Definition.
-
+import io
 import typing
+import warnings
 
 import httpx
+from typing_extensions import TypeVar, ParamSpec
 
 from .base_client import BaseAthena, AsyncBaseAthena
 from .environment import AthenaEnvironment
-from .polling_message_client import \
-        MessagePollingClient, AsyncMessagePollingClient
+from .polling_message_client import MessagePollingClient, AsyncMessagePollingClient
+from .tools.client import AsyncToolsClient, ToolsClient
+from .types.data_frame_request_out import DataFrameRequestOut
+
+if typing.TYPE_CHECKING:
+    import pandas as pd
+
+
+P = ParamSpec("P")
+T = TypeVar("T")
+U = TypeVar('U')
+
+
+def _inherit_signature(f: typing.Callable[P, T]) -> typing.Callable[..., typing.Callable[P, U]]:
+    return lambda x: x
+
+
+class WrappedToolsClient(ToolsClient):
+
+    def get_file(self, document_id: str) -> io.BytesIO:
+        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)
+    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':
+        _check_pandas_installed()
+        file_bytes = self.get_file(document_id)
+        return _to_pandas_df(file_bytes, *args, **kwargs)
+
+
+class WrappedAsyncToolsClient(AsyncToolsClient):
+
+    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)
+    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)
+
+    async def read_data_frame(self, document_id: str, *args, **kwargs) -> 'pd.DataFrame':
+        _check_pandas_installed()
+        file_bytes = await self.get_file(document_id)
+        return _to_pandas_df(file_bytes, *args, **kwargs)
 
 
 class Athena(BaseAthena):
@@ -52,6 +104,7 @@ class Athena(BaseAthena):
         )
         self.message = MessagePollingClient(
             client_wrapper=self._client_wrapper)
+        self.tools = WrappedToolsClient(client_wrapper=self._client_wrapper)
 
 
 class AsyncAthena(AsyncBaseAthena):
@@ -96,4 +149,38 @@ class AsyncAthena(AsyncBaseAthena):
         )
         self.message = AsyncMessagePollingClient(
             client_wrapper=self._client_wrapper)
+        self.tools = WrappedAsyncToolsClient(client_wrapper=self._client_wrapper)
+
+
+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')
+
+
+def _check_pandas_installed():
+    import pandas
+    assert pandas
+
+
+def _to_pandas_df(bytes_io: io.BytesIO, *args, **kwargs):
+    import pandas as pd
+    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)
 
+    if media_type == 'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet':
+        return pd.read_excel(bytes_io, *args, engine='openpyxl', **kwargs)
+    elif media_type == 'application/vnd.ms-excel':
+        return pd.read_excel(bytes_io, *args, **kwargs)
+    elif media_type == 'text/csv':
+        return pd.read_csv(bytes_io, *args, **kwargs)
+    else:
+        raise Exception("Unknown media type")

athena/core/__init__.py

@@ -6,7 +6,8 @@ from .datetime_utils import serialize_datetime
 from .file import File, convert_file_dict_to_httpx_tuples
 from .http_client import AsyncHttpClient, HttpClient
 from .jsonable_encoder import jsonable_encoder
-from .pydantic_utilities import pydantic_v1
+from .pydantic_utilities import deep_union_pydantic_dicts, pydantic_v1
+from .query_encoder import encode_query
 from .remove_none_from_dict import remove_none_from_dict
 from .request_options import RequestOptions
 
@@ -20,6 +21,8 @@ __all__ = [
     "RequestOptions",
     "SyncClientWrapper",
     "convert_file_dict_to_httpx_tuples",
+    "deep_union_pydantic_dicts",
+    "encode_query",
     "jsonable_encoder",
     "pydantic_v1",
     "remove_none_from_dict",

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

athena/core/pydantic_utilities.py

@@ -1,5 +1,7 @@
 # This file was auto-generated by Fern from our API Definition.
 
+import typing
+
 import pydantic
 
 IS_PYDANTIC_V2 = pydantic.VERSION.startswith("2.")
@@ -9,4 +11,18 @@ if IS_PYDANTIC_V2:
 else:
     import pydantic as pydantic_v1  # type: ignore  # nopycln: import
 
+
+def deep_union_pydantic_dicts(
+    source: typing.Dict[str, typing.Any], destination: typing.Dict[str, typing.Any]
+) -> typing.Dict[str, typing.Any]:
+    for key, value in source.items():
+        if isinstance(value, dict):
+            node = destination.setdefault(key, {})
+            deep_union_pydantic_dicts(value, node)
+        else:
+            destination[key] = value
+
+    return destination
+
+
 __all__ = ["pydantic_v1"]

athena/core/query_encoder.py

@@ -0,0 +1,33 @@
+# This file was auto-generated by Fern from our API Definition.
+
+from collections import ChainMap
+from typing import Any, Dict, Optional
+
+from .pydantic_utilities import pydantic_v1
+
+
+# Flattens dicts to be of the form {"key[subkey][subkey2]": value} where value is not a dict
+def traverse_query_dict(dict_flat: Dict[str, Any], key_prefix: Optional[str] = None) -> Dict[str, Any]:
+    result = {}
+    for k, v in dict_flat.items():
+        key = f"{key_prefix}[{k}]" if key_prefix is not None else k
+        if isinstance(v, dict):
+            result.update(traverse_query_dict(v, key))
+        else:
+            result[key] = v
+    return result
+
+
+def single_query_encoder(query_key: str, query_value: Any) -> Dict[str, Any]:
+    if isinstance(query_value, pydantic_v1.BaseModel) or isinstance(query_value, dict):
+        if isinstance(query_value, pydantic_v1.BaseModel):
+            obj_dict = query_value.dict(by_alias=True)
+        else:
+            obj_dict = query_value
+        return traverse_query_dict(obj_dict, query_key)
+
+    return {query_key: query_value}
+
+
+def encode_query(query: Optional[Dict[str, Any]]) -> Optional[Dict[str, Any]]:
+    return dict(ChainMap(*[single_query_encoder(k, v) for k, v in query.items()])) if query is not None else None

athena/dataset/client.py

@@ -8,6 +8,7 @@ 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
@@ -27,13 +28,24 @@ class DatasetClient:
         request_options: typing.Optional[RequestOptions] = None,
     ) -> GetDatasetsResponse:
         """
-        Parameters:
-            - page: typing.Optional[int]. Page number starting from 1
+        Parameters
+        ----------
+        page : typing.Optional[int]
+            Page number starting from 1
 
-            - page_size: typing.Optional[int]. Number of items per page
+        page_size : typing.Optional[int]
+            Number of items per page
 
-            - request_options: typing.Optional[RequestOptions]. Request-specific configuration.
-        ---
+        request_options : typing.Optional[RequestOptions]
+            Request-specific configuration.
+
+        Returns
+        -------
+        GetDatasetsResponse
+            Successful Response
+
+        Examples
+        --------
         from athena.client import Athena
 
         client = Athena(
@@ -44,17 +56,19 @@ class DatasetClient:
         _response = self._client_wrapper.httpx_client.request(
             method="GET",
             url=urllib.parse.urljoin(f"{self._client_wrapper.get_base_url()}/", "api/v0/datasets"),
-            params=jsonable_encoder(
-                remove_none_from_dict(
-                    {
-                        "page": page,
-                        "page_size": page_size,
-                        **(
-                            request_options.get("additional_query_parameters", {})
-                            if request_options is not None
-                            else {}
-                        ),
-                    }
+            params=encode_query(
+                jsonable_encoder(
+                    remove_none_from_dict(
+                        {
+                            "page": page,
+                            "page_size": page_size,
+                            **(
+                                request_options.get("additional_query_parameters", {})
+                                if request_options is not None
+                                else {}
+                            ),
+                        }
+                    )
                 )
             ),
             headers=jsonable_encoder(
@@ -96,13 +110,24 @@ class AsyncDatasetClient:
         request_options: typing.Optional[RequestOptions] = None,
     ) -> GetDatasetsResponse:
         """
-        Parameters:
-            - page: typing.Optional[int]. Page number starting from 1
+        Parameters
+        ----------
+        page : typing.Optional[int]
+            Page number starting from 1
 
-            - page_size: typing.Optional[int]. Number of items per page
+        page_size : typing.Optional[int]
+            Number of items per page
 
-            - request_options: typing.Optional[RequestOptions]. Request-specific configuration.
-        ---
+        request_options : typing.Optional[RequestOptions]
+            Request-specific configuration.
+
+        Returns
+        -------
+        GetDatasetsResponse
+            Successful Response
+
+        Examples
+        --------
         from athena.client import AsyncAthena
 
         client = AsyncAthena(
@@ -113,17 +138,19 @@ class AsyncDatasetClient:
         _response = await self._client_wrapper.httpx_client.request(
             method="GET",
             url=urllib.parse.urljoin(f"{self._client_wrapper.get_base_url()}/", "api/v0/datasets"),
-            params=jsonable_encoder(
-                remove_none_from_dict(
-                    {
-                        "page": page,
-                        "page_size": page_size,
-                        **(
-                            request_options.get("additional_query_parameters", {})
-                            if request_options is not None
-                            else {}
-                        ),
-                    }
+            params=encode_query(
+                jsonable_encoder(
+                    remove_none_from_dict(
+                        {
+                            "page": page,
+                            "page_size": page_size,
+                            **(
+                                request_options.get("additional_query_parameters", {})
+                                if request_options is not None
+                                else {}
+                            ),
+                        }
+                    )
                 )
             ),
             headers=jsonable_encoder(