3
3
from beartype import beartype
4
4
from temporalio import activity
5
5
6
+ from ..app import app
6
7
from ..autogen .openapi_model import BaseIntegrationDef
7
8
from ..clients import integrations
8
9
from ..common .exceptions .tools import IntegrationExecutionException
9
10
from ..common .protocol .tasks import ExecutionInput , StepContext
10
- from ..common .storage_handler import auto_blob_store
11
11
from ..env import testing
12
- from ..models . tools import get_tool_args_from_metadata
12
+ from ..queries import tools
13
13
14
14
15
- @auto_blob_store (deep = True )
16
15
@beartype
17
16
async def execute_integration (
18
17
context : StepContext ,
@@ -22,23 +21,35 @@ async def execute_integration(
22
21
setup : dict [str , Any ] = {},
23
22
) -> Any :
24
23
if not isinstance (context .execution_input , ExecutionInput ):
25
- raise TypeError ("Expected ExecutionInput type for context.execution_input" )
24
+ msg = "Expected ExecutionInput type for context.execution_input"
25
+ raise TypeError (msg )
26
26
27
27
developer_id = context .execution_input .developer_id
28
28
agent_id = context .execution_input .agent .id
29
+
30
+ if context .execution_input .task is None :
31
+ msg = "Task cannot be None in execution_input"
32
+ raise ValueError (msg )
33
+
29
34
task_id = context .execution_input .task .id
30
35
31
- merged_tool_args = get_tool_args_from_metadata (
32
- developer_id = developer_id , agent_id = agent_id , task_id = task_id , arg_type = "args"
36
+ merged_tool_args = await tools .get_tool_args_from_metadata (
37
+ developer_id = developer_id ,
38
+ agent_id = agent_id ,
39
+ task_id = task_id ,
40
+ arg_type = "args" ,
41
+ connection_pool = app .state .postgres_pool ,
33
42
)
34
43
35
- merged_tool_setup = get_tool_args_from_metadata (
36
- developer_id = developer_id , agent_id = agent_id , task_id = task_id , arg_type = "setup"
44
+ merged_tool_setup = await tools .get_tool_args_from_metadata (
45
+ developer_id = developer_id ,
46
+ agent_id = agent_id ,
47
+ task_id = task_id ,
48
+ arg_type = "setup" ,
49
+ connection_pool = app .state .postgres_pool ,
37
50
)
38
51
39
- arguments = (
40
- merged_tool_args .get (tool_name , {}) | (integration .arguments or {}) | arguments
41
- )
52
+ arguments = merged_tool_args .get (tool_name , {}) | (integration .arguments or {}) | arguments
42
53
43
54
setup = merged_tool_setup .get (tool_name , {}) | (integration .setup or {}) | setup
44
55
@@ -53,10 +64,7 @@ async def execute_integration(
53
64
arguments = arguments ,
54
65
)
55
66
56
- if (
57
- "error" in integration_service_response
58
- and integration_service_response ["error" ]
59
- ):
67
+ if integration_service_response .get ("error" ):
60
68
raise IntegrationExecutionException (
61
69
integration = integration ,
62
70
error = integration_service_response ["error" ],
@@ -69,9 +77,7 @@ async def execute_integration(
69
77
integration_str = integration .provider + (
70
78
"." + integration .method if integration .method else ""
71
79
)
72
- activity .logger .error (
73
- f"Error in execute_integration { integration_str } : { e } "
74
- )
80
+ activity .logger .error (f"Error in execute_integration { integration_str } : { e } " )
75
81
76
82
raise
77
83
0 commit comments