CopilotForXcode/Core/Sources/ChatService/DynamicContextController.swift at 94aaadedd7fa05077da28e934caa72d31939a39a · Colacn/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
131
import ChatContextCollector
import Foundation
import OpenAIService
import Preferences
import XcodeInspector

final class DynamicContextController {
    let contextCollectors: [ChatContextCollector]
    let memory: AutoManagedChatGPTMemory
    let functionProvider: ChatFunctionProvider
    let configuration: OverridingChatGPTConfiguration
    var defaultScopes = [] as Set<ChatContext.Scope>

    convenience init(
        memory: AutoManagedChatGPTMemory,
        functionProvider: ChatFunctionProvider,
        configuration: OverridingChatGPTConfiguration,
        contextCollectors: ChatContextCollector...
    ) {
        self.init(
            memory: memory,
            functionProvider: functionProvider,
            configuration: configuration,
            contextCollectors: contextCollectors
        )
    }

    init(
        memory: AutoManagedChatGPTMemory,
        functionProvider: ChatFunctionProvider,
        configuration: OverridingChatGPTConfiguration,
        contextCollectors: [ChatContextCollector]
    ) {
        self.memory = memory
        self.functionProvider = functionProvider
        self.configuration = configuration
        self.contextCollectors = contextCollectors
    }

    func collectContextInformation(systemPrompt: String, content: String) async throws {
        var content = content
        var scopes = Self.parseScopes(&content)
        scopes.formUnion(defaultScopes)

        let overridingChatModelId = {
            var ids = [String]()
            if scopes.contains(.sense) {
                ids.append(UserDefaults.shared.value(for: \.preferredChatModelIdForSenseScope))
            }

            if scopes.contains(.project) {
                ids.append(UserDefaults.shared.value(for: \.preferredChatModelIdForProjectScope))
            }

            if scopes.contains(.web) {
                ids.append(UserDefaults.shared.value(for: \.preferredChatModelIdForWebScope))
            }

            let chatModels = UserDefaults.shared.value(for: \.chatModels)
            let idIndexMap = chatModels.enumerated().reduce(into: [String: Int]()) {
                $0[$1.element.id] = $1.offset
            }
            return ids.sorted(by: {
                let lhs = idIndexMap[$0] ?? Int.max
                let rhs = idIndexMap[$1] ?? Int.max
                return lhs < rhs
            }).first
        }()

        configuration.overriding.modelId = overridingChatModelId

        functionProvider.removeAll()
        let language = UserDefaults.shared.value(for: \.chatGPTLanguage)
        let oldMessages = await memory.history
        let contexts = await withTaskGroup(
            of: ChatContext.self
        ) { [scopes, content, configuration] group in
            for collector in contextCollectors {
                group.addTask {
                    await collector.generateContext(
                        history: oldMessages,
                        scopes: scopes,
                        content: content,
                        configuration: configuration
                    )
                }
            }
            var contexts = [ChatContext]()
            for await context in group {
                contexts.append(context)
            }
            return contexts
        }

        let contextSystemPrompt = contexts
            .map(\.systemPrompt)
            .filter { !$0.isEmpty }
            .joined(separator: "\n\n")

        let contextPrompts = contexts
            .flatMap(\.retrievedContent)
            .filter { !$0.content.isEmpty }
            .sorted { $0.priority > $1.priority }

        let contextualSystemPrompt = """
        \(language.isEmpty ? "" : "You must always reply in \(language)")
        \(systemPrompt)
        """
        await memory.mutateSystemPrompt(contextualSystemPrompt)
        await memory.mutateContextSystemPrompt(contextSystemPrompt)
        await memory.mutateRetrievedContent(contextPrompts.map {
            .init(
                title: "",
                subTitle: "",
                uri: "",
                content: $0.content,
                startLine: nil,
                endLine: nil,
                metadata: [:]
            )
        })
        functionProvider.append(functions: contexts.flatMap(\.functions))
    }
}

extension DynamicContextController {
    static func parseScopes(_ prompt: inout String) -> Set<ChatContext.Scope> {
        let parser = MessageScopeParser()
        return parser(&prompt)
    }
}