CopilotForXcode/Tool/Sources/SuggestionProvider/SuggestionServiceMiddleware.swift at main · github/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
import Foundation
import Logger
import SuggestionBasic

public protocol SuggestionServiceMiddleware {
    typealias Next = (SuggestionRequest) async throws -> [CodeSuggestion]

    func getSuggestion(
        _ request: SuggestionRequest,
        configuration: SuggestionServiceConfiguration,
        next: Next
    ) async throws -> [CodeSuggestion]

    func getNESSuggestion(
        _ request: SuggestionRequest,
        configuration: SuggestionServiceConfiguration,
        next: Next
    ) async throws -> [CodeSuggestion]
}

public enum SuggestionServiceMiddlewareContainer {
    static var builtInMiddlewares: [SuggestionServiceMiddleware] = [
        DisabledLanguageSuggestionServiceMiddleware(),
        PostProcessingSuggestionServiceMiddleware()
    ]

    static var customMiddlewares: [SuggestionServiceMiddleware] = []

    public static var middlewares: [SuggestionServiceMiddleware] {
        builtInMiddlewares + customMiddlewares
    }

    public static func addMiddleware(_ middleware: SuggestionServiceMiddleware) {
        customMiddlewares.append(middleware)
    }
}

public struct DisabledLanguageSuggestionServiceMiddleware: SuggestionServiceMiddleware {
    public init() {}

    public func getSuggestion(
        _ request: SuggestionRequest,
        configuration: SuggestionServiceConfiguration,
        next: Next
    ) async throws -> [CodeSuggestion] {
        let language = languageIdentifierFromFileURL(request.fileURL)
        if UserDefaults.shared.value(for: \.suggestionFeatureDisabledLanguageList)
            .contains(where: { $0 == language.rawValue })
        {
            #if DEBUG
            Logger.service.info("Suggestion service is disabled for \(language).")
            #endif
            return []
        }

        return try await next(request)
    }

    public func getNESSuggestion(
        _ request: SuggestionRequest,
        configuration: SuggestionServiceConfiguration,
        next: Next
    ) async throws -> [CodeSuggestion] {
        let language = languageIdentifierFromFileURL(request.fileURL)
        if UserDefaults.shared.value(for: \.suggestionFeatureDisabledLanguageList)
            .contains(where: { $0 == language.rawValue })
        {
            #if DEBUG
            Logger.service.info("Suggestion service is disabled for \(language).")
            #endif
            return []
        }

        return try await next(request)
    }
}

public struct DebugSuggestionServiceMiddleware: SuggestionServiceMiddleware {
    public init() {}

    public func getSuggestion(
        _ request: SuggestionRequest,
        configuration: SuggestionServiceConfiguration,
        next: Next
    ) async throws -> [CodeSuggestion] {
        Logger.service.info("""
        Get suggestion for \(request.fileURL) at \(request.cursorPosition)
        """)
        do {
            let suggestions = try await next(request)
            Logger.service.info("""
            Receive \(suggestions.count) suggestions for \(request.fileURL) \
            at \(request.cursorPosition)
            """)
            return suggestions
        } catch {
            Logger.service.info("""
            Error: \(error.localizedDescription)
            """)
            throw error
        }
    }

    public func getNESSuggestion(
        _ request: SuggestionRequest,
        configuration: SuggestionServiceConfiguration,
        next: Next
    ) async throws -> [CodeSuggestion] {
        Logger.service.info("""
        Get suggestion for \(request.fileURL) at \(request.cursorPosition)
        """)
        do {
            let suggestions = try await next(request)
            Logger.service.info("""
            Receive \(suggestions.count) suggestions for \(request.fileURL) \
            at \(request.cursorPosition)
            """)
            return suggestions
        } catch {
            Logger.service.info("""
            Error: \(error.localizedDescription)
            """)
            throw error
        }
    }
}