Skip to content

Commit 68575c2

Browse files
intitniintitni
committed
ChatContextCollector
1 parent d631234 commit 68575c2

File tree

7 files changed

+130
-46
lines changed

7 files changed

+130
-46
lines changed

‎Core/Sources/ChatContextCollectors/ActiveDocumentChatContextCollector/ActiveDocumentChatContextCollector.swift‎

Lines changed: 12 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -17,8 +17,8 @@ public final class ActiveDocumentChatContextCollector: ChatContextCollector {
1717
scopes: Set<String>,
1818
content: String,
1919
configuration: ChatGPTConfiguration
20-
) -> ChatContext? {
21-
guard let info = getEditorInformation() else { return nil }
20+
) -> ChatContext {
21+
guard let info = getEditorInformation() else { return .empty }
2222
let context = getActiveDocumentContext(info)
2323
activeDocumentContext = context
2424

@@ -27,11 +27,16 @@ public final class ActiveDocumentChatContextCollector: ChatContextCollector {
2727
var removedCode = context
2828
removedCode.focusedContext = nil
2929
return .init(
30-
systemPrompt: extractSystemPrompt(removedCode),
30+
systemPrompt: [
31+
.init(
32+
content: extractSystemPrompt(removedCode),
33+
priority: .high
34+
),
35+
],
3136
functions: []
3237
)
3338
}
34-
return nil
39+
return .empty
3540
}
3641

3742
var functions = [any ChatGPTFunction]()
@@ -65,7 +70,9 @@ public final class ActiveDocumentChatContextCollector: ChatContextCollector {
6570
}
6671

6772
return .init(
68-
systemPrompt: extractSystemPrompt(context),
73+
systemPrompt: [
74+
.init(content: extractSystemPrompt(context), priority: .high)
75+
],
6976
functions: functions
7077
)
7178
}

‎Core/Sources/ChatContextCollectors/ActiveDocumentChatContextCollector/LegacyActiveDocumentChatContextCollector.swift‎

Lines changed: 26 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,9 @@
1+
import ChatContextCollector
12
import Foundation
23
import OpenAIService
34
import Preferences
45
import SuggestionModel
56
import XcodeInspector
6-
import ChatContextCollector
77

88
public struct LegacyActiveDocumentChatContextCollector: ChatContextCollector {
99
public init() {}
@@ -13,8 +13,8 @@ public struct LegacyActiveDocumentChatContextCollector: ChatContextCollector {
1313
scopes: Set<String>,
1414
content: String,
1515
configuration: ChatGPTConfiguration
16-
) -> ChatContext? {
17-
guard let content = getEditorInformation() else { return nil }
16+
) -> ChatContext {
17+
guard let content = getEditorInformation() else { return .empty }
1818
let relativePath = content.relativePath
1919
let selectionRange = content.editorContent?.selections.first ?? .outOfScope
2020
let editorContent = {
@@ -78,30 +78,31 @@ public struct LegacyActiveDocumentChatContextCollector: ChatContextCollector {
7878
}()
7979

8080
return .init(
81-
systemPrompt: """
82-
Active Document Context:###
83-
Document Relative Path: \(relativePath)
84-
Selection Range Start: \
85-
Line \(selectionRange.start.line) \
86-
Character \(selectionRange.start.character)
87-
Selection Range End: \
88-
Line \(selectionRange.end.line) \
89-
Character \(selectionRange.end.character)
90-
Cursor Position: \
91-
Line \(selectionRange.end.line) \
92-
Character \(selectionRange.end.character)
93-
\(editorContent)
94-
Line Annotations:
95-
\(
96-
content.editorContent?.lineAnnotations
97-
.map { " - \($0)" }
98-
.joined(separator: "\n") ?? "N/A"
99-
)
100-
###
101-
""",
81+
systemPrompt: [
82+
.init(content: """
83+
Active Document Context:###
84+
Document Relative Path: \(relativePath)
85+
Selection Range Start: \
86+
Line \(selectionRange.start.line) \
87+
Character \(selectionRange.start.character)
88+
Selection Range End: \
89+
Line \(selectionRange.end.line) \
90+
Character \(selectionRange.end.character)
91+
Cursor Position: \
92+
Line \(selectionRange.end.line) \
93+
Character \(selectionRange.end.character)
94+
\(editorContent)
95+
Line Annotations:
96+
\(
97+
content.editorContent?.lineAnnotations
98+
.map { " - \($0)" }
99+
.joined(separator: "\n") ?? "N/A"
100+
)
101+
###
102+
""", priority: .high),
103+
],
102104
functions: []
103105
)
104106
}
105107
}
106108

107-

‎Core/Sources/ChatContextCollectors/SystemInfoChatContextCollector/SystemInfoChatContextCollector.swift‎

Lines changed: 11 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -16,11 +16,18 @@ public final class SystemInfoChatContextCollector: ChatContextCollector {
1616
scopes: Set<String>,
1717
content: String,
1818
configuration: ChatGPTConfiguration
19-
) -> ChatContext? {
19+
) -> ChatContext {
2020
return .init(
21-
systemPrompt: """
22-
Current Time: \(Self.dateFormatter.string(from: Date())) (You can use it to calculate time in another time zone)
23-
""",
21+
systemPrompt: [
22+
.init(
23+
content: """
24+
Current Time: \(
25+
Self.dateFormatter.string(from: Date())
26+
) (You can use it to calculate time in another time zone)
27+
""",
28+
priority: .custom(999999)
29+
),
30+
],
2431
functions: []
2532
)
2633
}

‎Core/Sources/ChatContextCollectors/WebChatContextCollector/WebChatContextCollector.swift‎

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -12,16 +12,21 @@ public final class WebChatContextCollector: ChatContextCollector {
1212
scopes: Set<String>,
1313
content: String,
1414
configuration: ChatGPTConfiguration
15-
) -> ChatContext? {
16-
guard scopes.contains("web") || scopes.contains("w") else { return nil }
15+
) -> ChatContext {
16+
guard scopes.contains("web") || scopes.contains("w") else { return .empty }
1717
let links = Self.detectLinks(from: history) + Self.detectLinks(from: content)
1818
let functions: [(any ChatGPTFunction)?] = [
1919
SearchFunction(maxTokens: configuration.maxTokens),
2020
// allow this function only when there is a link in the memory.
2121
links.isEmpty ? nil : QueryWebsiteFunction(),
2222
]
2323
return .init(
24-
systemPrompt: "You prefer to answer questions with latest content on the internet.",
24+
systemPrompt: [
25+
.init(
26+
content: "You prefer to answer questions with latest content on the internet.",
27+
priority: .low
28+
),
29+
],
2530
functions: functions.compactMap { $0 }
2631
)
2732
}

‎Core/Sources/ChatService/DynamicContextController.swift‎

Lines changed: 13 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,7 @@ final class DynamicContextController {
4747
let language = UserDefaults.shared.value(for: \.chatGPTLanguage)
4848
let oldMessages = await memory.history
4949
let contexts = await withTaskGroup(
50-
of: ChatContext?.self
50+
of: ChatContext.self
5151
) { [scopes, content, configuration] group in
5252
for collector in contextCollectors {
5353
group.addTask {
@@ -61,17 +61,25 @@ final class DynamicContextController {
6161
}
6262
var contexts = [ChatContext]()
6363
for await context in group {
64-
if let context = context {
65-
contexts.append(context)
66-
}
64+
contexts.append(context)
6765
}
6866
return contexts
6967
}
68+
69+
let separator = String(repeating: "=", count: 32) // only 1 token
70+
71+
let contextPrompts = contexts
72+
.flatMap(\.systemPrompt)
73+
.filter { !$0.content.isEmpty }
74+
.sorted { $0.priority > $1.priority }
75+
7076
let contextualSystemPrompt = """
7177
\(language.isEmpty ? "" : "You must always reply in \(language)")
7278
\(systemPrompt)
7379
74-
\(contexts.map(\.systemPrompt).filter { !$0.isEmpty }.joined(separator: "\n\n"))
80+
Below are information related to the conversation, separated by \(separator)
81+
82+
\(contextPrompts.map(\.content).joined(separator: "\n\(separator)\n"))
7583
"""
7684
await memory.mutateSystemPrompt(contextualSystemPrompt)
7785
functionProvider.append(functions: contexts.flatMap(\.functions))

‎Pro‎

Submodule Pro updated from ba9117b to 9b5986f

‎Tool/Sources/ChatContextCollector/ChatContextCollector.swift‎

Lines changed: 59 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,12 +2,68 @@ import Foundation
22
import OpenAIService
33

44
public struct ChatContext {
5-
public var systemPrompt: String
5+
public struct RetrievedPrompt {
6+
public enum Priority: Equatable, Comparable {
7+
case low
8+
case medium
9+
case high
10+
case custom(Int)
11+
12+
public var rawValue: Int {
13+
switch self {
14+
case .low:
15+
return 20
16+
case .medium:
17+
return 60
18+
case .high:
19+
return 80
20+
case let .custom(value):
21+
return value
22+
}
23+
}
24+
25+
public static func < (lhs: Self, rhs: Self) -> Bool {
26+
lhs.rawValue < rhs.rawValue
27+
}
28+
29+
public static func == (lhs: Self, rhs: Self) -> Bool {
30+
lhs.rawValue == rhs.rawValue
31+
}
32+
}
33+
34+
public var content: String
35+
public var priority: Priority
36+
37+
public init(content: String, priority: Priority) {
38+
self.content = content
39+
self.priority = priority
40+
}
41+
}
42+
43+
public var systemPrompt: [RetrievedPrompt]
644
public var functions: [any ChatGPTFunction]
7-
public init(systemPrompt: String, functions: [any ChatGPTFunction]) {
45+
public init(systemPrompt: [RetrievedPrompt], functions: [any ChatGPTFunction]) {
846
self.systemPrompt = systemPrompt
947
self.functions = functions
1048
}
49+
50+
public static var empty: Self {
51+
.init(systemPrompt: [], functions: [])
52+
}
53+
}
54+
55+
public func + (
56+
lhs: ChatContext.RetrievedPrompt.Priority,
57+
rhs: Int
58+
) -> ChatContext.RetrievedPrompt.Priority {
59+
.custom(lhs.rawValue + rhs)
60+
}
61+
62+
public func - (
63+
lhs: ChatContext.RetrievedPrompt.Priority,
64+
rhs: Int
65+
) -> ChatContext.RetrievedPrompt.Priority {
66+
.custom(lhs.rawValue - rhs)
1167
}
1268

1369
public protocol ChatContextCollector {
@@ -16,6 +72,6 @@ public protocol ChatContextCollector {
1672
scopes: Set<String>,
1773
content: String,
1874
configuration: ChatGPTConfiguration
19-
) async -> ChatContext?
75+
) async -> ChatContext
2076
}
2177

0 commit comments

Comments
 (0)