From 85adf95b3aaa4d0aa82337ab4ce452554548b832 Mon Sep 17 00:00:00 2001 From: Tushar Sadhwani Date: Mon, 23 Dec 2024 18:16:52 +0530 Subject: [PATCH 1/2] feat: Support requesting `connected_account` in custom action --- python/composio/tools/base/runtime.py | 32 ++++++++++++++++++++++----- 1 file changed, 26 insertions(+), 6 deletions(-) diff --git a/python/composio/tools/base/runtime.py b/python/composio/tools/base/runtime.py index daaea03c7db..c4eedf13088 100644 --- a/python/composio/tools/base/runtime.py +++ b/python/composio/tools/base/runtime.py @@ -9,7 +9,7 @@ from pydantic import BaseModel, Field from composio import Composio -from composio.client.collections import CustomAuthParameter +from composio.client.collections import ConnectedAccountModel, CustomAuthParameter from composio.client.enums.base import ActionData, SentinalObject, add_runtime_action from composio.client.exceptions import ComposioClientError from composio.exceptions import ComposioSDKError @@ -253,12 +253,23 @@ def _parse_docstring( return header, params, returns +def _get_connected_account(app: str, entity_id: str) -> t.Optional[ConnectedAccountModel]: + try: + client = Composio.get_latest() + connected_account = client.connected_accounts.get( + connection_id=client.get_entity(entity_id).get_connection(app=app).id + ) + return connected_account + except ComposioClientError: + return None + def _get_auth_params(app: str, entity_id: str) -> t.Optional[t.Dict]: try: client = Composio.get_latest() - connection_params = client.connected_accounts.get( + connected_account = client.connected_accounts.get( connection_id=client.get_entity(entity_id).get_connection(app=app).id - ).connectionParams + ) + connection_params = connected_account.connectionParams return { "headers": connection_params.headers, "base_url": connection_params.base_url, @@ -303,7 +314,8 @@ def _build_executable_from_args( # pylint: disable=too-many-statements } shell_argument = None - auth_params = False + auth_param = False + connected_account_param = False request_executor = False if "return" not in argspec.annotations: raise InvalidRuntimeAction( @@ -316,7 +328,11 @@ def _build_executable_from_args( # pylint: disable=too-many-statements continue if arg == "auth": - auth_params = True + auth_param = True + continue + + if arg == "connected_account": + connected_account_param = True continue if arg == "execute_request": @@ -376,7 +392,11 @@ def execute(request: BaseModel, metadata: t.Dict) -> BaseModel: if shell_argument is not None: kwargs[shell_argument] = metadata["workspace"].shells.recent - if auth_params > 0: + if connected_account_param: + kwargs["connected_account"] = ( + _get_connected_account(app=app, entity_id=metadata["entity_id"]) or {} + ) + if auth_param: kwargs["auth"] = ( _get_auth_params(app=app, entity_id=metadata["entity_id"]) or {} ) From b1626af896d5904f7c9aa238546b3a0483cd74b8 Mon Sep 17 00:00:00 2001 From: Tushar Sadhwani Date: Wed, 25 Dec 2024 15:17:00 +0530 Subject: [PATCH 2/2] format --- python/composio/tools/base/runtime.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/python/composio/tools/base/runtime.py b/python/composio/tools/base/runtime.py index c4eedf13088..49e07f777ce 100644 --- a/python/composio/tools/base/runtime.py +++ b/python/composio/tools/base/runtime.py @@ -253,7 +253,9 @@ def _parse_docstring( return header, params, returns -def _get_connected_account(app: str, entity_id: str) -> t.Optional[ConnectedAccountModel]: +def _get_connected_account( + app: str, entity_id: str +) -> t.Optional[ConnectedAccountModel]: try: client = Composio.get_latest() connected_account = client.connected_accounts.get( @@ -263,6 +265,7 @@ def _get_connected_account(app: str, entity_id: str) -> t.Optional[ConnectedAcco except ComposioClientError: return None + def _get_auth_params(app: str, entity_id: str) -> t.Optional[t.Dict]: try: client = Composio.get_latest()