Skip to content

Commit 4147cb4

Browse files
andrewklatzkeandrewklatzke
committed
[REL-11697] update tests to match server response
1 parent 206e42e commit 4147cb4

File tree

1 file changed

+53
-75
lines changed

1 file changed

+53
-75
lines changed

‎packages/sdk/server-ai/tests/test_agent_graph.py‎

Lines changed: 53 additions & 75 deletions
Original file line numberDiff line numberDiff line change
@@ -19,27 +19,17 @@ def td() -> TestData:
1919
td.flag("test-agent-graph")
2020
.variations(
2121
{
22-
"key": "test-agent-graph",
23-
"name": "Test Agent Graph",
24-
"rootConfigKey": "customer-support-agent",
25-
"edges": [
26-
{
27-
"key": "edge-customer-support-agent-personalized-agent",
28-
"sourceConfig": "customer-support-agent",
29-
"targetConfig": "personalized-agent",
30-
},
31-
{
32-
"key": "edge-customer-support-agent-multi-context-agent",
33-
"sourceConfig": "customer-support-agent",
34-
"targetConfig": "multi-context-agent",
35-
},
36-
{
37-
"key": "edge-customer-support-agent-minimal-agent",
38-
"sourceConfig": "customer-support-agent",
39-
"targetConfig": "minimal-agent",
40-
},
41-
],
42-
"description": "Test agent graph",
22+
"root": "customer-support-agent",
23+
"edges": {
24+
"customer-support-agent": [
25+
{
26+
"key": "personalized-agent",
27+
"handoff": {"state": "from-root-to-personalized"},
28+
},
29+
{"key": "multi-context-agent", "handoff": {}},
30+
{"key": "minimal-agent", "handoff": {}},
31+
]
32+
},
4333
"_ldMeta": {
4434
"enabled": True,
4535
"variationKey": "test-agent-graph",
@@ -54,35 +44,28 @@ def td() -> TestData:
5444
td.flag("test-agent-graph-depth-3")
5545
.variations(
5646
{
57-
"key": "test-agent-graph-depth-3",
58-
"name": "Test Agent Graph with Depth of 3",
59-
"rootConfigKey": "customer-support-agent",
60-
"edges": [
61-
{
62-
"key": "edge-customer-support-agent-personalized-agent",
63-
"sourceConfig": "customer-support-agent",
64-
"targetConfig": "personalized-agent",
65-
"handoff": { "state": "from-root-to-personalized" }
66-
},
67-
{
68-
"key": "edge-personalized-agent-multi-context-agent",
69-
"sourceConfig": "personalized-agent",
70-
"targetConfig": "multi-context-agent",
71-
},
72-
{
73-
"key": "edge-multi-context-agent-minimal-agent",
74-
"sourceConfig": "multi-context-agent",
75-
"targetConfig": "minimal-agent",
76-
"handoff": {"state": "from-multi-context-to-minimal"},
77-
},
78-
{
79-
"key": "edge-customer-support-agent-minimal-agent",
80-
"sourceConfig": "customer-support-agent",
81-
"targetConfig": "minimal-agent",
82-
"handoff": { "state": "from-root-to-minimal" }
83-
},
84-
],
85-
"description": "Test agent graph with depth of 3",
47+
"root": "customer-support-agent",
48+
"edges": {
49+
"customer-support-agent": [
50+
{
51+
"key": "personalized-agent",
52+
"handoff": {"state": "from-root-to-personalized"},
53+
},
54+
{
55+
"key": "minimal-agent",
56+
"handoff": {"state": "from-root-to-minimal"},
57+
},
58+
],
59+
"personalized-agent": [
60+
{"key": "multi-context-agent", "handoff": {}}
61+
],
62+
"multi-context-agent": [
63+
{
64+
"key": "minimal-agent",
65+
"handoff": {"state": "from-multi-context-to-minimal"},
66+
}
67+
],
68+
},
8669
"_ldMeta": {
8770
"enabled": True,
8871
"variationKey": "test-agent-graph-depth-3",
@@ -98,17 +81,10 @@ def td() -> TestData:
9881
td.flag("test-agent-graph-disabled-agent")
9982
.variations(
10083
{
101-
"key": "test-agent-graph-disabled-agent",
102-
"name": "Test Agent Graph with Disabled Agent",
103-
"rootConfigKey": "customer-support-agent",
104-
"edges": [
105-
{
106-
"key": "edge-customer-support-agent-personalized-agent",
107-
"sourceConfig": "customer-support-agent",
108-
"targetConfig": "disabled-agent",
109-
},
110-
],
111-
"description": "Test agent graph with disabled agent",
84+
"root": "customer-support-agent",
85+
"edges": {
86+
"customer-support-agent": [{"key": "disabled-agent", "handoff": {}}]
87+
},
11288
"_ldMeta": {
11389
"enabled": True,
11490
"variationKey": "test-agent-graph-disabled-agent",
@@ -124,9 +100,12 @@ def td() -> TestData:
124100
td.flag("test-agent-graph-no-root-key")
125101
.variations(
126102
{
127-
"name": "Test Agent Graph with No Root Key",
128-
"key": "test-agent-graph-no-root-key",
129-
"edges": [],
103+
"edges": {},
104+
"_ldMeta": {
105+
"enabled": True,
106+
"variationKey": "test-agent-graph-no-root-key",
107+
"version": 1,
108+
},
130109
}
131110
)
132111
.variation_for_all(0)
@@ -271,19 +250,18 @@ def test_agent_graph_build_nodes(ldai_client: LDAIClient):
271250
)
272251

273252
ai_graph_config = AIAgentGraphConfig(
274-
key=graph_config["key"],
275-
name=graph_config["name"],
276-
root_config_key=graph_config["rootConfigKey"],
253+
key="test-agent-graph",
254+
root_config_key=graph_config["root"],
277255
edges=[
278256
Edge(
279-
key=edge.get("key", ""),
280-
source_config=edge.get("sourceConfig", ""),
281-
target_config=edge.get("targetConfig", ""),
257+
key=edge_key + "-" + edge.get("key", ""),
258+
source_config=edge_key,
259+
target_config=edge.get("key", ""),
282260
handoff=edge.get("handoff", {}),
283261
)
284-
for edge in graph_config["edges"]
262+
for edge_key, edges in graph_config["edges"].items()
263+
for edge in edges
285264
],
286-
description=graph_config["description"],
287265
)
288266

289267
nodes = AgentGraphDefinition.build_nodes(
@@ -422,18 +400,18 @@ def handle_traverse(node, context):
422400
assert first_edge.handoff == {"state": "from-multi-context-to-minimal"}
423401
assert first_edge.source_config == "multi-context-agent"
424402
assert first_edge.target_config == "minimal-agent"
425-
assert first_edge.key == "edge-multi-context-agent-minimal-agent"
403+
assert first_edge.key == "multi-context-agent-minimal-agent"
426404
if node.get_key() == "customer-support-agent":
427405
first_edge = node.get_edges()[0]
428406
second_edge = node.get_edges()[1]
429407
assert first_edge.handoff == {"state": "from-root-to-personalized"}
430408
assert second_edge.handoff == {"state": "from-root-to-minimal"}
431409
assert first_edge.source_config == "customer-support-agent"
432410
assert first_edge.target_config == "personalized-agent"
433-
assert first_edge.key == "edge-customer-support-agent-personalized-agent"
411+
assert first_edge.key == "customer-support-agent-personalized-agent"
434412
assert second_edge.source_config == "customer-support-agent"
435413
assert second_edge.target_config == "minimal-agent"
436-
assert second_edge.key == "edge-customer-support-agent-minimal-agent"
414+
assert second_edge.key == "customer-support-agent-minimal-agent"
437415
return None
438416

439417
graph.traverse(handle_traverse, context)

0 commit comments

Comments
 (0)