CopilotForXcode/Core/Tests/OpenAIServiceTests/LimitMessagesTests.swift at 2985b4e65043273c1f6e8a8db3a9990fbedb4f54 · UXLabPro/CopilotForXcode · 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
124
125
126
127
128
129
130
import Foundation
import XCTest

@testable import OpenAIService

final class LimitMessagesTests: XCTestCase {
    func test_send_all_messages_if_not_reached_token_limit() async {
        let service = await createService(systemPrompt: "system", messages: [
            "hi",
            "hello",
            "world",
        ])

        let (messages, remainingTokens) = await runService(
            service,
            minimumReplyTokens: 200,
            maxNumberOfMessages: 0, // smaller than 1 means no limit
            maxTokens: 10000
        )
        XCTAssertEqual(messages, [
            "system",
            "hi",
            "hello",
            "world",
        ])

        XCTAssertEqual(remainingTokens, 10000 - 12 - 6)
        let history = await service.history
        XCTAssertEqual(history.map(\.tokensCount), [
            2,
            5,
            5,
        ])
    }

    func test_send_max_message_if_not_reached_token_limit() async {
        let service = await createService(systemPrompt: "system", messages: [
            "hi",
            "hello",
            "world",
        ])

        let (messages, remainingTokens) = await runService(
            service,
            minimumReplyTokens: 200,
            maxNumberOfMessages: 2,
            maxTokens: 10000
        )
        XCTAssertEqual(messages, [
            "system",
            "hello",
            "world",
        ], "Count from end to start.")

        XCTAssertEqual(remainingTokens, 10000 - 10 - 6)
    }

    func test_reached_token_limit() async {
        let service = await createService(systemPrompt: "system", messages: [
            "hi",
            "hello",
            "world",
        ])

        let (messages, remainingTokens) = await runService(
            service,
            minimumReplyTokens: 200,
            maxNumberOfMessages: 100,
            maxTokens: 212
        )
        XCTAssertEqual(messages, [
            "system",
            "world",
        ])

        XCTAssertEqual(remainingTokens, 201)
    }

    func test_minimum_reply_tokens_count() async {
        let service = await createService(systemPrompt: "system", messages: [
            "hi",
            "hello",
            "world",
        ])

        let (messages, remainingTokens) = await runService(
            service,
            minimumReplyTokens: 200,
            maxNumberOfMessages: 100,
            maxTokens: 200
        )
        XCTAssertEqual(messages, [
            "system",
        ])

        XCTAssertEqual(remainingTokens, 200)
    }
}

class MockEncoder: TokenEncoder {
    func encode(text: String) -> [Int] {
        return .init(repeating: 0, count: text.count)
    }
}

private func createService(systemPrompt: String, messages: [String]) async -> ChatGPTService {
    let service = ChatGPTService(systemPrompt: systemPrompt)
    await service.mutateHistory { history in
        messages.forEach { message in
            history.append(.init(role: .user, content: message))
        }
    }
    return service
}

private func runService(
    _ service: ChatGPTService,
    minimumReplyTokens: Int,
    maxNumberOfMessages: Int,
    maxTokens: Int
) async -> (messages: [String], remainingTokens: Int) {
    let (messages, remainingTokens) = await service.combineHistoryWithSystemPrompt(
        minimumReplyTokens: minimumReplyTokens,
        maxNumberOfMessages: maxNumberOfMessages,
        maxTokens: maxTokens,
        encoder: MockEncoder()
    )

    return (messages.map(\.content), remainingTokens)
}