athena-intelligence 0.1.119__py3-none-any.whl → 0.1.121__py3-none-any.whl

athena/__init__.py

@@ -2,20 +2,24 @@
 
 from .types import (
     AssetNotFoundError,
-    AthenaAssistantConfig,
-    AthenaAssistantConfigEnabledToolsItem,
-    AthenaAssistantReponse,
-    AthenaAssistantRequest,
+    CustomAgentResponse,
     DataFrameRequestOut,
     DataFrameRequestOutColumnsItem,
     DataFrameRequestOutDataItemItem,
     DataFrameRequestOutIndexItem,
     DataFrameUnknownFormatError,
     DocumentChunk,
+    DriveAgentResponse,
     FileChunkRequestOut,
     FileTooLargeError,
+    GeneralAgentConfig,
+    GeneralAgentConfigEnabledToolsItem,
+    GeneralAgentRequest,
+    GeneralAgentResponse,
     ParentFolderError,
+    ResearchAgentResponse,
     SaveAssetRequestOut,
+    SqlAgentResponse,
     StructuredDataExtractorReponse,
     Tool,
 )
@@ -36,26 +40,30 @@ from .version import __version__
 
 __all__ = [
     "AssetNotFoundError",
-    "AthenaAssistantConfig",
-    "AthenaAssistantConfigEnabledToolsItem",
-    "AthenaAssistantReponse",
-    "AthenaAssistantRequest",
     "AthenaEnvironment",
     "BadRequestError",
     "ContentTooLargeError",
+    "CustomAgentResponse",
     "DataFrameRequestOut",
     "DataFrameRequestOutColumnsItem",
     "DataFrameRequestOutDataItemItem",
     "DataFrameRequestOutIndexItem",
     "DataFrameUnknownFormatError",
     "DocumentChunk",
+    "DriveAgentResponse",
     "FileChunkRequestOut",
     "FileTooLargeError",
+    "GeneralAgentConfig",
+    "GeneralAgentConfigEnabledToolsItem",
+    "GeneralAgentRequest",
+    "GeneralAgentResponse",
     "InternalServerError",
     "NotFoundError",
     "ParentFolderError",
     "QueryExecuteRequestDatabaseAssetIds",
+    "ResearchAgentResponse",
     "SaveAssetRequestOut",
+    "SqlAgentResponse",
     "StructuredDataExtractorReponse",
     "Tool",
     "ToolsDataFrameRequestColumnsItem",

athena/agents/__init__.py

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

athena/agents/client.py

@@ -1,19 +1,169 @@
 # 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 .athena_assistant.client import AsyncAthenaAssistantClient, AthenaAssistantClient
-from .structured_data_extractor.client import AsyncStructuredDataExtractorClient, StructuredDataExtractorClient
+from ..core.jsonable_encoder import jsonable_encoder
+from ..core.pydantic_utilities import pydantic_v1
+from ..core.request_options import RequestOptions
+from ..errors.unprocessable_entity_error import UnprocessableEntityError
+from ..types.custom_agent_response import CustomAgentResponse
+from .drive.client import AsyncDriveClient, DriveClient
+from .general.client import AsyncGeneralClient, GeneralClient
+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, ...)
 
 
 class AgentsClient:
     def __init__(self, *, client_wrapper: SyncClientWrapper):
         self._client_wrapper = client_wrapper
-        self.athena_assistant = AthenaAssistantClient(client_wrapper=self._client_wrapper)
-        self.structured_data_extractor = StructuredDataExtractorClient(client_wrapper=self._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)
+
+    def invoke_by_id(
+        self,
+        agent_id: str,
+        *,
+        config: typing.Dict[str, typing.Any],
+        messages: typing.Sequence[typing.Dict[str, typing.Any]],
+        request_options: typing.Optional[RequestOptions] = None,
+    ) -> CustomAgentResponse:
+        """
+        Coming soon!
+
+        Invoke a custom agent created in [spaces](https://resources.athenaintel.com/docs/agents/create-your-agent).
+
+        Custom agents can be created and configured in spaces to perform specialized tasks.
+        Refer to the specific agent's documentation for details on configuration options
+        and expected responses.
+
+        Parameters
+        ----------
+        agent_id : str
+            The ID of the custom agent to invoke. Create custom agents in [spaces](https://resources.athenaintel.com/docs/agents/create-your-agent).
+
+        config : typing.Dict[str, typing.Any]
+            Configuration for the custom agent. See the agent's documentation for specific configuration options.
+
+        messages : typing.Sequence[typing.Dict[str, typing.Any]]
+            The messages to send to the custom agent
+
+        request_options : typing.Optional[RequestOptions]
+            Request-specific configuration.
+
+        Returns
+        -------
+        CustomAgentResponse
+            Successful Response
+
+        Examples
+        --------
+        from athena.client import Athena
+
+        client = Athena(
+            api_key="YOUR_API_KEY",
+        )
+        client.agents.invoke_by_id(
+            agent_id="agent_id",
+            config={"key": "value"},
+            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},
+            request_options=request_options,
+            omit=OMIT,
+        )
+        if 200 <= _response.status_code < 300:
+            return pydantic_v1.parse_obj_as(CustomAgentResponse, _response.json())  # type: ignore
+        if _response.status_code == 422:
+            raise UnprocessableEntityError(pydantic_v1.parse_obj_as(typing.Any, _response.json()))  # type: ignore
+        try:
+            _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)
 
 
 class AsyncAgentsClient:
     def __init__(self, *, client_wrapper: AsyncClientWrapper):
         self._client_wrapper = client_wrapper
-        self.athena_assistant = AsyncAthenaAssistantClient(client_wrapper=self._client_wrapper)
-        self.structured_data_extractor = AsyncStructuredDataExtractorClient(client_wrapper=self._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)
+
+    async def invoke_by_id(
+        self,
+        agent_id: str,
+        *,
+        config: typing.Dict[str, typing.Any],
+        messages: typing.Sequence[typing.Dict[str, typing.Any]],
+        request_options: typing.Optional[RequestOptions] = None,
+    ) -> CustomAgentResponse:
+        """
+        Coming soon!
+
+        Invoke a custom agent created in [spaces](https://resources.athenaintel.com/docs/agents/create-your-agent).
+
+        Custom agents can be created and configured in spaces to perform specialized tasks.
+        Refer to the specific agent's documentation for details on configuration options
+        and expected responses.
+
+        Parameters
+        ----------
+        agent_id : str
+            The ID of the custom agent to invoke. Create custom agents in [spaces](https://resources.athenaintel.com/docs/agents/create-your-agent).
+
+        config : typing.Dict[str, typing.Any]
+            Configuration for the custom agent. See the agent's documentation for specific configuration options.
+
+        messages : typing.Sequence[typing.Dict[str, typing.Any]]
+            The messages to send to the custom agent
+
+        request_options : typing.Optional[RequestOptions]
+            Request-specific configuration.
+
+        Returns
+        -------
+        CustomAgentResponse
+            Successful Response
+
+        Examples
+        --------
+        from athena.client import AsyncAthena
+
+        client = AsyncAthena(
+            api_key="YOUR_API_KEY",
+        )
+        await client.agents.invoke_by_id(
+            agent_id="agent_id",
+            config={"key": "value"},
+            messages=[{"key": "value"}],
+        )
+        """
+        _response = await self._client_wrapper.httpx_client.request(
+            f"api/v0/agents/{jsonable_encoder(agent_id)}/invoke",
+            method="POST",
+            json={"config": config, "messages": messages},
+            request_options=request_options,
+            omit=OMIT,
+        )
+        if 200 <= _response.status_code < 300:
+            return pydantic_v1.parse_obj_as(CustomAgentResponse, _response.json())  # type: ignore
+        if _response.status_code == 422:
+            raise UnprocessableEntityError(pydantic_v1.parse_obj_as(typing.Any, _response.json()))  # type: ignore
+        try:
+            _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)

athena/agents/drive/client.py

@@ -0,0 +1,134 @@
+# 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.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 DriveClient:
+    def __init__(self, *, client_wrapper: SyncClientWrapper):
+        self._client_wrapper = client_wrapper
+
+    def invoke(
+        self,
+        *,
+        config: typing.Dict[str, typing.Any],
+        messages: typing.Sequence[typing.Dict[str, typing.Any]],
+        request_options: typing.Optional[RequestOptions] = None
+    ) -> DriveAgentResponse:
+        """
+        Coming soon! Manage folders and search for files in the internal drive.
+
+        Parameters
+        ----------
+        config : typing.Dict[str, typing.Any]
+            Configuration for the drive agent including folder paths and search parameters
+
+        messages : typing.Sequence[typing.Dict[str, typing.Any]]
+            The messages to send to the drive agent
+
+        request_options : typing.Optional[RequestOptions]
+            Request-specific configuration.
+
+        Returns
+        -------
+        DriveAgentResponse
+            Successful Response
+
+        Examples
+        --------
+        from athena.client import Athena
+
+        client = Athena(
+            api_key="YOUR_API_KEY",
+        )
+        client.agents.drive.invoke(
+            config={"key": "value"},
+            messages=[{"key": "value"}],
+        )
+        """
+        _response = self._client_wrapper.httpx_client.request(
+            "api/v0/agents/drive/invoke",
+            method="POST",
+            json={"config": config, "messages": messages},
+            request_options=request_options,
+            omit=OMIT,
+        )
+        if 200 <= _response.status_code < 300:
+            return pydantic_v1.parse_obj_as(DriveAgentResponse, _response.json())  # type: ignore
+        if _response.status_code == 422:
+            raise UnprocessableEntityError(pydantic_v1.parse_obj_as(typing.Any, _response.json()))  # type: ignore
+        try:
+            _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)
+
+
+class AsyncDriveClient:
+    def __init__(self, *, client_wrapper: AsyncClientWrapper):
+        self._client_wrapper = client_wrapper
+
+    async def invoke(
+        self,
+        *,
+        config: typing.Dict[str, typing.Any],
+        messages: typing.Sequence[typing.Dict[str, typing.Any]],
+        request_options: typing.Optional[RequestOptions] = None
+    ) -> DriveAgentResponse:
+        """
+        Coming soon! Manage folders and search for files in the internal drive.
+
+        Parameters
+        ----------
+        config : typing.Dict[str, typing.Any]
+            Configuration for the drive agent including folder paths and search parameters
+
+        messages : typing.Sequence[typing.Dict[str, typing.Any]]
+            The messages to send to the drive agent
+
+        request_options : typing.Optional[RequestOptions]
+            Request-specific configuration.
+
+        Returns
+        -------
+        DriveAgentResponse
+            Successful Response
+
+        Examples
+        --------
+        from athena.client import AsyncAthena
+
+        client = AsyncAthena(
+            api_key="YOUR_API_KEY",
+        )
+        await client.agents.drive.invoke(
+            config={"key": "value"},
+            messages=[{"key": "value"}],
+        )
+        """
+        _response = await self._client_wrapper.httpx_client.request(
+            "api/v0/agents/drive/invoke",
+            method="POST",
+            json={"config": config, "messages": messages},
+            request_options=request_options,
+            omit=OMIT,
+        )
+        if 200 <= _response.status_code < 300:
+            return pydantic_v1.parse_obj_as(DriveAgentResponse, _response.json())  # type: ignore
+        if _response.status_code == 422:
+            raise UnprocessableEntityError(pydantic_v1.parse_obj_as(typing.Any, _response.json()))  # type: ignore
+        try:
+            _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)