-
Notifications
You must be signed in to change notification settings - Fork 1.1k
Expand file tree
/
Copy pathtest_system_message_transform.py
More file actions
123 lines (92 loc) · 3.92 KB
/
test_system_message_transform.py
File metadata and controls
123 lines (92 loc) · 3.92 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
"""
Copyright (c) Microsoft Corporation.
Tests for system message transform functionality
"""
import pytest
from copilot.session import PermissionHandler
from .testharness import E2ETestContext
from .testharness.helper import write_file
pytestmark = pytest.mark.asyncio(loop_scope="module")
class TestSystemMessageTransform:
async def test_should_invoke_transform_callbacks_with_section_content(
self, ctx: E2ETestContext
):
"""Test that transform callbacks are invoked with the section content"""
identity_contents = []
tone_contents = []
async def identity_transform(content: str) -> str:
identity_contents.append(content)
return content
async def tone_transform(content: str) -> str:
tone_contents.append(content)
return content
session = await ctx.client.create_session(
system_message={
"mode": "customize",
"sections": {
"identity": {"action": identity_transform},
"tone": {"action": tone_transform},
},
},
on_permission_request=PermissionHandler.approve_all,
)
write_file(ctx.work_dir, "test.txt", "Hello transform!")
await session.send_and_wait("Read the contents of test.txt and tell me what it says")
# Both transform callbacks should have been invoked
assert len(identity_contents) > 0
assert len(tone_contents) > 0
# Callbacks should have received non-empty content
assert all(len(c) > 0 for c in identity_contents)
assert all(len(c) > 0 for c in tone_contents)
await session.disconnect()
async def test_should_apply_transform_modifications_to_section_content(
self, ctx: E2ETestContext
):
"""Test that transform modifications are applied to the section content"""
async def identity_transform(content: str) -> str:
return content + "\nTRANSFORM_MARKER"
session = await ctx.client.create_session(
system_message={
"mode": "customize",
"sections": {
"identity": {"action": identity_transform},
},
},
on_permission_request=PermissionHandler.approve_all,
)
write_file(ctx.work_dir, "hello.txt", "Hello!")
await session.send_and_wait("Read the contents of hello.txt")
# Verify the transform result was actually applied to the system message
traffic = await ctx.get_exchanges()
system_message = _get_system_message(traffic[0])
assert "TRANSFORM_MARKER" in system_message
await session.disconnect()
async def test_should_work_with_static_overrides_and_transforms_together(
self, ctx: E2ETestContext
):
"""Test that static overrides and transforms work together"""
identity_contents = []
async def identity_transform(content: str) -> str:
identity_contents.append(content)
return content
session = await ctx.client.create_session(
system_message={
"mode": "customize",
"sections": {
"safety": {"action": "remove"},
"identity": {"action": identity_transform},
},
},
on_permission_request=PermissionHandler.approve_all,
)
write_file(ctx.work_dir, "combo.txt", "Combo test!")
await session.send_and_wait("Read the contents of combo.txt and tell me what it says")
# The transform callback should have been invoked
assert len(identity_contents) > 0
await session.disconnect()
def _get_system_message(exchange: dict) -> str:
messages = exchange.get("request", {}).get("messages", [])
for msg in messages:
if msg.get("role") == "system":
return msg.get("content", "")
return ""