athena-intelligence 0.1.124__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(