copilot-sdk/python/e2e/test_system_message_transform_e2e.py at main · Blixtleod/copilot-sdk · GitHub

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 ""