17
17
# The Bedrock agent is hosted on Sage Bionetworks AWS infrastructure.
18
18
# CFN Template:
19
19
# https://raw.githubusercontent.com/Sage-Bionetworks-Workflows/dpe-agents/refs/heads/main/client_integration_test/template.json
20
+
20
21
AGENT_AWS_ID = "QOTV3KQM1X"
21
- AGENT_REGISTRATION_ID = "29"
22
22
23
23
24
24
class TestAgentPrompt :
@@ -27,6 +27,10 @@ class TestAgentPrompt:
27
27
@pytest .fixture (autouse = True , scope = "function" )
28
28
def init (self , syn : Synapse ) -> None :
29
29
self .syn = syn
30
+ if syn .repoEndpoint == "https://repo-dev.dev.sagebase.org/repo/v1" :
31
+ self .AGENT_REGISTRATION_ID = "7"
32
+ else :
33
+ self .AGENT_REGISTRATION_ID = "29"
30
34
31
35
async def test_send_job_and_wait_async_with_post_exchange_args (self ) -> None :
32
36
# GIVEN an AgentPrompt with a valid concrete type, prompt, and enable_trace
@@ -37,7 +41,7 @@ async def test_send_job_and_wait_async_with_post_exchange_args(self) -> None:
37
41
)
38
42
# AND the ID of an existing agent session
39
43
test_session = await AgentSession (
40
- agent_registration_id = AGENT_REGISTRATION_ID
44
+ agent_registration_id = self . AGENT_REGISTRATION_ID
41
45
).start_async (synapse_client = self .syn )
42
46
test_prompt .session_id = test_session .id
43
47
# WHEN I send the job and wait for it to complete
@@ -59,7 +63,7 @@ def init(self, syn: Synapse) -> None:
59
63
60
64
async def test_start (self ) -> None :
61
65
# GIVEN an agent session with a valid agent registration id
62
- agent_session = AgentSession (agent_registration_id = AGENT_REGISTRATION_ID )
66
+ agent_session = AgentSession (agent_registration_id = self . AGENT_REGISTRATION_ID )
63
67
64
68
# WHEN the start method is called
65
69
result_session = await agent_session .start_async (synapse_client = self .syn )
@@ -73,13 +77,13 @@ async def test_start(self) -> None:
73
77
assert result_session .started_on is not None
74
78
assert result_session .started_by is not None
75
79
assert result_session .modified_on is not None
76
- assert result_session .agent_registration_id == AGENT_REGISTRATION_ID
80
+ assert result_session .agent_registration_id == self . AGENT_REGISTRATION_ID
77
81
assert result_session .etag is not None
78
82
assert result_session .chat_history == []
79
83
80
84
async def test_get (self ) -> None :
81
85
# GIVEN an agent session with a valid agent registration id
82
- agent_session = AgentSession (agent_registration_id = AGENT_REGISTRATION_ID )
86
+ agent_session = AgentSession (agent_registration_id = self . AGENT_REGISTRATION_ID )
83
87
# WHEN I start a session
84
88
await agent_session .start_async (synapse_client = self .syn )
85
89
# THEN I expect to be able to get the session with its id
@@ -92,7 +96,7 @@ async def test_update(self) -> None:
92
96
# GIVEN an agent session with a valid agent
93
97
# registration id and access level set
94
98
agent_session = AgentSession (
95
- agent_registration_id = AGENT_REGISTRATION_ID ,
99
+ agent_registration_id = self . AGENT_REGISTRATION_ID ,
96
100
access_level = AgentSessionAccessLevel .PUBLICLY_ACCESSIBLE ,
97
101
)
98
102
# WHEN I start a session
@@ -111,7 +115,7 @@ async def test_update(self) -> None:
111
115
112
116
async def test_prompt (self ) -> None :
113
117
# GIVEN an agent session with a valid agent registration id
114
- agent_session = AgentSession (agent_registration_id = AGENT_REGISTRATION_ID )
118
+ agent_session = AgentSession (agent_registration_id = self . AGENT_REGISTRATION_ID )
115
119
# WHEN I start a session
116
120
await agent_session .start_async (synapse_client = self .syn )
117
121
# THEN I expect to be able to prompt the agent
@@ -133,7 +137,7 @@ def get_test_agent(self) -> Agent:
133
137
return Agent (
134
138
cloud_agent_id = AGENT_AWS_ID ,
135
139
cloud_alias_id = "TSTALIASID" ,
136
- registration_id = AGENT_REGISTRATION_ID ,
140
+ registration_id = self . AGENT_REGISTRATION_ID ,
137
141
registered_on = "2025-01-16T18:57:35.680Z" ,
138
142
type = "CUSTOM" ,
139
143
sessions = {},
@@ -155,7 +159,7 @@ async def test_register(self) -> None:
155
159
156
160
async def test_get (self ) -> None :
157
161
# GIVEN an Agent with a valid agent registration id
158
- agent = Agent (registration_id = AGENT_REGISTRATION_ID )
162
+ agent = Agent (registration_id = self . AGENT_REGISTRATION_ID )
159
163
# WHEN I get the agent
160
164
await agent .get_async (synapse_client = self .syn )
161
165
# THEN I expect the agent to be returned
@@ -171,7 +175,7 @@ async def test_get_no_registration_id(self) -> None:
171
175
172
176
async def test_start_session (self ) -> None :
173
177
# GIVEN an Agent with a valid agent registration id
174
- agent = Agent (registration_id = AGENT_REGISTRATION_ID )
178
+ agent = Agent (registration_id = self . AGENT_REGISTRATION_ID )
175
179
# WHEN I start a session
176
180
await agent .start_session_async (synapse_client = self .syn )
177
181
# THEN I expect a current session to be set
@@ -181,7 +185,7 @@ async def test_start_session(self) -> None:
181
185
182
186
async def test_get_session (self ) -> None :
183
187
# GIVEN an Agent with a valid agent registration id
184
- agent = Agent (registration_id = AGENT_REGISTRATION_ID )
188
+ agent = Agent (registration_id = self . AGENT_REGISTRATION_ID )
185
189
# WHEN I start a session
186
190
await agent .start_session_async (synapse_client = self .syn )
187
191
# THEN I expect to be able to get the session with its id
@@ -193,12 +197,12 @@ async def test_get_session(self) -> None:
193
197
194
198
async def test_prompt_with_session (self ) -> None :
195
199
# GIVEN an Agent with a valid agent registration id
196
- agent = await Agent (registration_id = AGENT_REGISTRATION_ID ).get_async (
200
+ agent = await Agent (registration_id = self . AGENT_REGISTRATION_ID ).get_async (
197
201
synapse_client = self .syn
198
202
)
199
203
# AND a session started separately
200
204
session = await AgentSession (
201
- agent_registration_id = AGENT_REGISTRATION_ID
205
+ agent_registration_id = self . AGENT_REGISTRATION_ID
202
206
).start_async (synapse_client = self .syn )
203
207
# WHEN I prompt the agent with a session
204
208
await agent .prompt_async (prompt = "hello" , enable_trace = True , session = session )
@@ -213,7 +217,7 @@ async def test_prompt_with_session(self) -> None:
213
217
214
218
async def test_prompt_no_session (self ) -> None :
215
219
# GIVEN an Agent with a valid agent registration id
216
- agent = await Agent (registration_id = AGENT_REGISTRATION_ID ).get_async (
220
+ agent = await Agent (registration_id = self . AGENT_REGISTRATION_ID ).get_async (
217
221
synapse_client = self .syn
218
222
)
219
223
# WHEN I prompt the agent without a current session set
0 commit comments