CopilotForXcode/Tool/Tests/OpenAIServiceTests/LimitMessagesTests.swift at main · codehobby2022/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
import Foundation
import TokenEncoder
import XCTest

@testable import OpenAIService

final class AutoManagedChatGPTMemoryTests: XCTestCase {
    func test_send_all_messages_if_not_reached_token_limit() async {
        let (messages, _, memory) = await runService(
            systemPrompt: "system", messages: [
                "hi",
                "hello",
                "world",
            ], maxTokens: 10000, minimumReplyTokens: 200,
            maxNumberOfMessages: 0 // smaller than 1 means no limit
        )
        XCTAssertEqual(messages, [
            "system",
            "hi",
            "hello",
            "world",
        ])

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

    func test_send_max_message_if_not_reached_token_limit() async {
        let (messages, _, _) = await runService(
            systemPrompt: "system", messages: [
                "hi",
                "hello",
                "world",
            ], maxTokens: 10000, minimumReplyTokens: 200,
            maxNumberOfMessages: 2
        )
        XCTAssertEqual(messages, [
            "system",
            "hello",
            "world",
        ], "Count from end to start.")

//        XCTAssertEqual(remainingTokens, 10000 - 10 - 6)
    }

    func test_reached_token_limit() async {
        let (messages, _, _) = await runService(
            systemPrompt: "system", messages: [
                "hi",
                "hello",
                "world",
            ], maxTokens: 212, minimumReplyTokens: 200,
            maxNumberOfMessages: 100
        )
        XCTAssertEqual(messages, [
            "system",
        ])

//        XCTAssertEqual(remainingTokens, 201)
    }

    func test_minimum_reply_tokens_count() async {
        let (messages, _, _) = await runService(
            systemPrompt: "system", messages: [
                "hi",
                "hello",
                "world",
            ],
            maxTokens: 200, minimumReplyTokens: 200,
            maxNumberOfMessages: 100
        )
        XCTAssertEqual(messages, [
            "system",
        ])

//        XCTAssertEqual(remainingTokens, 200)
    }
}

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

private func runService(
    systemPrompt: String,
    messages: [String],
    maxTokens: Int,
    minimumReplyTokens: Int,
    maxNumberOfMessages: Int
) async -> (messages: [String], remainingTokens: Int?, memory: AutoManagedChatGPTMemory) {
    let configuration = UserPreferenceChatGPTConfiguration().overriding(.init(
        maxTokens: maxTokens,
        minimumReplyTokens: minimumReplyTokens
    ))
    let memory = AutoManagedChatGPTMemory(
        systemPrompt: systemPrompt,
        configuration: configuration,
        functionProvider: NoChatGPTFunctionProvider()
    )

    for message in messages {
        await memory.appendMessage(.init(role: .user, content: message))
    }

    let messages = await memory.generateSendingHistory(
        maxNumberOfMessages: maxNumberOfMessages,
        encoder: MockEncoder()
    )
    let remainingTokens = await memory.generateRemainingTokens(
        maxNumberOfMessages: maxNumberOfMessages,
        encoder: MockEncoder()
    )

    let contents = messages.map { $0.content ?? "" }
    return (contents, remainingTokens, memory)
}