athena-intelligence 0.1.127__py3-none-any.whl → 0.1.184__py3-none-any.whl

athena/__init__.py

@@ -1,5 +1,7 @@
 # This file was auto-generated by Fern from our API Definition.
 
+# isort: skip_file
+
 from .types import (
     AssetContentRequestOut,
     AssetNode,
@@ -11,6 +13,7 @@ from .types import (
     ChunkContentItem_Text,
     ChunkResult,
     ChunkResultChunkId,
+    Content,
     CustomAgentResponse,
     DataFrameRequestOut,
     DataFrameRequestOutColumnsItem,
@@ -26,15 +29,23 @@ from .types import (
     GeneralAgentConfigEnabledToolsItem,
     GeneralAgentRequest,
     GeneralAgentResponse,
+    GeneralAgentResponseMessage,
+    GeneralAgentResponseMessageKwargs,
+    Id,
     ImageUrlContent,
+    InputMessage,
+    InputMessageContentItem,
+    InputMessageContentItem_ImageUrl,
+    InputMessageContentItem_Text,
+    PaginatedAssetsOut,
     ParentFolderError,
     PromptMessage,
+    PublicAssetOut,
     ResearchAgentResponse,
     SaveAssetRequestOut,
     SqlAgentResponse,
     StructuredDataExtractorResponse,
     TextContent,
-    Tool,
     Type,
 )
 from .errors import (
@@ -46,7 +57,7 @@ from .errors import (
     UnprocessableEntityError,
     UnsupportedMediaTypeError,
 )
-from . import agents, query, tools
+from . import agents, assets, query, tools
 from .client import AsyncAthena, Athena
 from .environment import AthenaEnvironment
 from .query import QueryExecuteRequestDatabaseAssetIds
@@ -68,6 +79,7 @@ __all__ = [
     "ChunkContentItem_Text",
     "ChunkResult",
     "ChunkResultChunkId",
+    "Content",
     "ContentTooLargeError",
     "CustomAgentResponse",
     "DataFrameRequestOut",
@@ -84,18 +96,26 @@ __all__ = [
     "GeneralAgentConfigEnabledToolsItem",
     "GeneralAgentRequest",
     "GeneralAgentResponse",
+    "GeneralAgentResponseMessage",
+    "GeneralAgentResponseMessageKwargs",
+    "Id",
     "ImageUrlContent",
+    "InputMessage",
+    "InputMessageContentItem",
+    "InputMessageContentItem_ImageUrl",
+    "InputMessageContentItem_Text",
     "InternalServerError",
     "NotFoundError",
+    "PaginatedAssetsOut",
     "ParentFolderError",
     "PromptMessage",
+    "PublicAssetOut",
     "QueryExecuteRequestDatabaseAssetIds",
     "ResearchAgentResponse",
     "SaveAssetRequestOut",
     "SqlAgentResponse",
     "StructuredDataExtractorResponse",
     "TextContent",
-    "Tool",
     "ToolsDataFrameRequestColumnsItem",
     "Type",
     "UnauthorizedError",
@@ -103,6 +123,7 @@ __all__ = [
     "UnsupportedMediaTypeError",
     "__version__",
     "agents",
+    "assets",
     "query",
     "tools",
 ]

athena/agents/__init__.py

@@ -1,5 +1,7 @@
 # This file was auto-generated by Fern from our API Definition.
 
+# isort: skip_file
+
 from . import drive, general, research, sql
 
 __all__ = ["drive", "general", "research", "sql"]

athena/agents/client.py

@@ -1,23 +1,15 @@
 # This file was auto-generated by Fern from our API Definition.
 
 import typing
-from ..core.client_wrapper import SyncClientWrapper
-from .drive.client import DriveClient
-from .general.client import GeneralClient
-from .research.client import ResearchClient
-from .sql.client import SqlClient
+
+from ..core.client_wrapper import AsyncClientWrapper, SyncClientWrapper
 from ..core.request_options import RequestOptions
 from ..types.custom_agent_response import CustomAgentResponse
-from ..core.jsonable_encoder import jsonable_encoder
-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
-from .drive.client import AsyncDriveClient
-from .general.client import AsyncGeneralClient
-from .research.client import AsyncResearchClient
-from .sql.client import AsyncSqlClient
+from .drive.client import AsyncDriveClient, DriveClient
+from .general.client import AsyncGeneralClient, GeneralClient
+from .raw_client import AsyncRawAgentsClient, RawAgentsClient
+from .research.client import AsyncResearchClient, ResearchClient
+from .sql.client import AsyncSqlClient, SqlClient
 
 # this is used as the default value for optional parameters
 OMIT = typing.cast(typing.Any, ...)
@@ -25,11 +17,25 @@ OMIT = typing.cast(typing.Any, ...)
 
 class AgentsClient:
     def __init__(self, *, client_wrapper: SyncClientWrapper):
-        self._client_wrapper = client_wrapper
-        self.drive = DriveClient(client_wrapper=self._client_wrapper)
-        self.general = GeneralClient(client_wrapper=self._client_wrapper)
-        self.research = ResearchClient(client_wrapper=self._client_wrapper)
-        self.sql = SqlClient(client_wrapper=self._client_wrapper)
+        self._raw_client = RawAgentsClient(client_wrapper=client_wrapper)
+        self.drive = DriveClient(client_wrapper=client_wrapper)
+
+        self.general = GeneralClient(client_wrapper=client_wrapper)
+
+        self.research = ResearchClient(client_wrapper=client_wrapper)
+
+        self.sql = SqlClient(client_wrapper=client_wrapper)
+
+    @property
+    def with_raw_response(self) -> RawAgentsClient:
+        """
+        Retrieves a raw implementation of this client that returns raw responses.
+
+        Returns
+        -------
+        RawAgentsClient
+        """
+        return self._raw_client
 
     def invoke_by_id(
         self,
@@ -80,51 +86,33 @@ class AgentsClient:
             messages=[{"key": "value"}],
         )
         """
-        _response = self._client_wrapper.httpx_client.request(
-            f"api/v0/agents/{jsonable_encoder(agent_id)}/invoke",
-            method="POST",
-            json={
-                "config": config,
-                "messages": messages,
-            },
-            headers={
-                "content-type": "application/json",
-            },
-            request_options=request_options,
-            omit=OMIT,
+        _response = self._raw_client.invoke_by_id(
+            agent_id, config=config, messages=messages, request_options=request_options
         )
-        try:
-            if 200 <= _response.status_code < 300:
-                return typing.cast(
-                    CustomAgentResponse,
-                    parse_obj_as(
-                        type_=CustomAgentResponse,  # 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)
-        raise ApiError(status_code=_response.status_code, body=_response_json)
+        return _response.data
 
 
 class AsyncAgentsClient:
     def __init__(self, *, client_wrapper: AsyncClientWrapper):
-        self._client_wrapper = client_wrapper
-        self.drive = AsyncDriveClient(client_wrapper=self._client_wrapper)
-        self.general = AsyncGeneralClient(client_wrapper=self._client_wrapper)
-        self.research = AsyncResearchClient(client_wrapper=self._client_wrapper)
-        self.sql = AsyncSqlClient(client_wrapper=self._client_wrapper)
+        self._raw_client = AsyncRawAgentsClient(client_wrapper=client_wrapper)
+        self.drive = AsyncDriveClient(client_wrapper=client_wrapper)
+
+        self.general = AsyncGeneralClient(client_wrapper=client_wrapper)
+
+        self.research = AsyncResearchClient(client_wrapper=client_wrapper)
+
+        self.sql = AsyncSqlClient(client_wrapper=client_wrapper)
+
+    @property
+    def with_raw_response(self) -> AsyncRawAgentsClient:
+        """
+        Retrieves a raw implementation of this client that returns raw responses.
+
+        Returns
+        -------
+        AsyncRawAgentsClient
+        """
+        return self._raw_client
 
     async def invoke_by_id(
         self,
@@ -183,39 +171,7 @@ class AsyncAgentsClient:
 
         asyncio.run(main())
         """
-        _response = await self._client_wrapper.httpx_client.request(
-            f"api/v0/agents/{jsonable_encoder(agent_id)}/invoke",
-            method="POST",
-            json={
-                "config": config,
-                "messages": messages,
-            },
-            headers={
-                "content-type": "application/json",
-            },
-            request_options=request_options,
-            omit=OMIT,
+        _response = await self._raw_client.invoke_by_id(
+            agent_id, config=config, messages=messages, request_options=request_options
         )
-        try:
-            if 200 <= _response.status_code < 300:
-                return typing.cast(
-                    CustomAgentResponse,
-                    parse_obj_as(
-                        type_=CustomAgentResponse,  # 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)
-        raise ApiError(status_code=_response.status_code, body=_response_json)
+        return _response.data

athena/agents/drive/__init__.py

@@ -1,2 +1,4 @@
 # This file was auto-generated by Fern from our API Definition.
 
+# isort: skip_file
+

athena/agents/drive/client.py

@@ -1,14 +1,11 @@
 # This file was auto-generated by Fern from our API Definition.
 
 import typing
-from ...core.client_wrapper import SyncClientWrapper
+
+from ...core.client_wrapper import AsyncClientWrapper, SyncClientWrapper
 from ...core.request_options import RequestOptions
 from ...types.drive_agent_response import DriveAgentResponse
-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
+from .raw_client import AsyncRawDriveClient, RawDriveClient
 
 # this is used as the default value for optional parameters
 OMIT = typing.cast(typing.Any, ...)
@@ -16,7 +13,18 @@ OMIT = typing.cast(typing.Any, ...)
 
 class DriveClient:
     def __init__(self, *, client_wrapper: SyncClientWrapper):
-        self._client_wrapper = client_wrapper
+        self._raw_client = RawDriveClient(client_wrapper=client_wrapper)
+
+    @property
+    def with_raw_response(self) -> RawDriveClient:
+        """
+        Retrieves a raw implementation of this client that returns raw responses.
+
+        Returns
+        -------
+        RawDriveClient
+        """
+        return self._raw_client
 
     def invoke(
         self,
@@ -56,47 +64,24 @@ class DriveClient:
             messages=[{"key": "value"}],
         )
         """
-        _response = self._client_wrapper.httpx_client.request(
-            "api/v0/agents/drive/invoke",
-            method="POST",
-            json={
-                "config": config,
-                "messages": messages,
-            },
-            headers={
-                "content-type": "application/json",
-            },
-            request_options=request_options,
-            omit=OMIT,
-        )
-        try:
-            if 200 <= _response.status_code < 300:
-                return typing.cast(
-                    DriveAgentResponse,
-                    parse_obj_as(
-                        type_=DriveAgentResponse,  # 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)
-        raise ApiError(status_code=_response.status_code, body=_response_json)
+        _response = self._raw_client.invoke(config=config, messages=messages, request_options=request_options)
+        return _response.data
 
 
 class AsyncDriveClient:
     def __init__(self, *, client_wrapper: AsyncClientWrapper):
-        self._client_wrapper = client_wrapper
+        self._raw_client = AsyncRawDriveClient(client_wrapper=client_wrapper)
+
+    @property
+    def with_raw_response(self) -> AsyncRawDriveClient:
+        """
+        Retrieves a raw implementation of this client that returns raw responses.
+
+        Returns
+        -------
+        AsyncRawDriveClient
+        """
+        return self._raw_client
 
     async def invoke(
         self,
@@ -144,39 +129,5 @@ class AsyncDriveClient:
 
         asyncio.run(main())
         """
-        _response = await self._client_wrapper.httpx_client.request(
-            "api/v0/agents/drive/invoke",
-            method="POST",
-            json={
-                "config": config,
-                "messages": messages,
-            },
-            headers={
-                "content-type": "application/json",
-            },
-            request_options=request_options,
-            omit=OMIT,
-        )
-        try:
-            if 200 <= _response.status_code < 300:
-                return typing.cast(
-                    DriveAgentResponse,
-                    parse_obj_as(
-                        type_=DriveAgentResponse,  # 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)
-        raise ApiError(status_code=_response.status_code, body=_response_json)
+        _response = await self._raw_client.invoke(config=config, messages=messages, request_options=request_options)
+        return _response.data

athena/agents/drive/raw_client.py

@@ -0,0 +1,155 @@
+# 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.http_response import AsyncHttpResponse, HttpResponse
+from ...core.pydantic_utilities import parse_obj_as
+from ...core.request_options import RequestOptions
+from ...errors.unprocessable_entity_error import UnprocessableEntityError
+from ...types.drive_agent_response import DriveAgentResponse
+
+# this is used as the default value for optional parameters
+OMIT = typing.cast(typing.Any, ...)
+
+
+class RawDriveClient:
+    def __init__(self, *, client_wrapper: SyncClientWrapper):
+        self._client_wrapper = client_wrapper
+
+    def invoke(
+        self,
+        *,
+        config: typing.Dict[str, typing.Optional[typing.Any]],
+        messages: typing.Sequence[typing.Dict[str, typing.Optional[typing.Any]]],
+        request_options: typing.Optional[RequestOptions] = None,
+    ) -> HttpResponse[DriveAgentResponse]:
+        """
+        Coming soon! Manage folders and search for files in the internal drive.
+
+        Parameters
+        ----------
+        config : typing.Dict[str, typing.Optional[typing.Any]]
+            Configuration for the drive agent including folder paths and search parameters
+
+        messages : typing.Sequence[typing.Dict[str, typing.Optional[typing.Any]]]
+            The messages to send to the drive agent
+
+        request_options : typing.Optional[RequestOptions]
+            Request-specific configuration.
+
+        Returns
+        -------
+        HttpResponse[DriveAgentResponse]
+            Successful Response
+        """
+        _response = self._client_wrapper.httpx_client.request(
+            "api/v0/agents/drive/invoke",
+            method="POST",
+            json={
+                "config": config,
+                "messages": messages,
+            },
+            headers={
+                "content-type": "application/json",
+            },
+            request_options=request_options,
+            omit=OMIT,
+        )
+        try:
+            if 200 <= _response.status_code < 300:
+                _data = typing.cast(
+                    DriveAgentResponse,
+                    parse_obj_as(
+                        type_=DriveAgentResponse,  # type: ignore
+                        object_=_response.json(),
+                    ),
+                )
+                return HttpResponse(response=_response, data=_data)
+            if _response.status_code == 422:
+                raise UnprocessableEntityError(
+                    headers=dict(_response.headers),
+                    body=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, headers=dict(_response.headers), body=_response.text)
+        raise ApiError(status_code=_response.status_code, headers=dict(_response.headers), body=_response_json)
+
+
+class AsyncRawDriveClient:
+    def __init__(self, *, client_wrapper: AsyncClientWrapper):
+        self._client_wrapper = client_wrapper
+
+    async def invoke(
+        self,
+        *,
+        config: typing.Dict[str, typing.Optional[typing.Any]],
+        messages: typing.Sequence[typing.Dict[str, typing.Optional[typing.Any]]],
+        request_options: typing.Optional[RequestOptions] = None,
+    ) -> AsyncHttpResponse[DriveAgentResponse]:
+        """
+        Coming soon! Manage folders and search for files in the internal drive.
+
+        Parameters
+        ----------
+        config : typing.Dict[str, typing.Optional[typing.Any]]
+            Configuration for the drive agent including folder paths and search parameters
+
+        messages : typing.Sequence[typing.Dict[str, typing.Optional[typing.Any]]]
+            The messages to send to the drive agent
+
+        request_options : typing.Optional[RequestOptions]
+            Request-specific configuration.
+
+        Returns
+        -------
+        AsyncHttpResponse[DriveAgentResponse]
+            Successful Response
+        """
+        _response = await self._client_wrapper.httpx_client.request(
+            "api/v0/agents/drive/invoke",
+            method="POST",
+            json={
+                "config": config,
+                "messages": messages,
+            },
+            headers={
+                "content-type": "application/json",
+            },
+            request_options=request_options,
+            omit=OMIT,
+        )
+        try:
+            if 200 <= _response.status_code < 300:
+                _data = typing.cast(
+                    DriveAgentResponse,
+                    parse_obj_as(
+                        type_=DriveAgentResponse,  # type: ignore
+                        object_=_response.json(),
+                    ),
+                )
+                return AsyncHttpResponse(response=_response, data=_data)
+            if _response.status_code == 422:
+                raise UnprocessableEntityError(
+                    headers=dict(_response.headers),
+                    body=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, headers=dict(_response.headers), body=_response.text)
+        raise ApiError(status_code=_response.status_code, headers=dict(_response.headers), body=_response_json)

athena/agents/general/__init__.py

@@ -1,2 +1,4 @@
 # This file was auto-generated by Fern from our API Definition.
 
+# isort: skip_file
+