Skip to content

Commit 41472e3

Browse files
intitniintitni
committed
Add chat context collector
1 parent b80af0d commit 41472e3

5 files changed

Lines changed: 122 additions & 72 deletions

File tree

‎Core/Package.swift‎

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -173,12 +173,28 @@ let package = Package(
173173

174174
.target(
175175
name: "ChatService",
176-
dependencies: ["OpenAIService", "ChatPlugins", "Environment", "XcodeInspector"]
176+
dependencies: [
177+
"ChatPlugins",
178+
"ChatContextCollector",
179+
"OpenAIService",
180+
"Environment",
181+
"XcodeInspector",
182+
"Preferences",
183+
]
177184
),
178185
.target(
179186
name: "ChatPlugins",
180187
dependencies: ["OpenAIService", "Environment", "Terminal"]
181188
),
189+
.target(
190+
name: "ChatContextCollector",
191+
dependencies: [
192+
"OpenAIService",
193+
"Environment",
194+
"Preferences",
195+
"SuggestionModel"
196+
]
197+
),
182198

183199
// MARK: - UI
184200

‎Core/Sources/ChatContextCollector/ActiveDocumentChatContextCollector.swift‎

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,83 @@
1+
import Foundation
2+
import Preferences
3+
import SuggestionModel
4+
import XcodeInspector
5+
6+
public struct ActiveDocumentChatContextCollector: ChatContextCollector {
7+
public init() {}
8+
9+
public func generateSystemPrompt(oldMessages: [String]) -> String {
10+
let content = getEditorInformation()
11+
let relativePath = content.documentURL.path
12+
.replacingOccurrences(of: content.projectURL.path, with: "")
13+
let selectionRange = content.editorContent?.selections.first ?? .outOfScope
14+
return """
15+
Active Document Context:###
16+
Document Relative Path: \(relativePath)
17+
Selection Range Start: \
18+
Line \(selectionRange.start.line) \
19+
Character \(selectionRange.start.character)
20+
Selection Range End: \
21+
Line \(selectionRange.end.line) \
22+
Character \(selectionRange.end.character)
23+
Cursor Position: \
24+
Line \(selectionRange.end.line) \
25+
Character \(selectionRange.end.character)
26+
Selected Code (start from line \(selectionRange.start.line)):```\(content.language.rawValue)
27+
\(content.selectedContent)
28+
```
29+
Line Annotations:
30+
\(
31+
content.editorContent?.lineAnnotations
32+
.map { " - \($0)" }
33+
.joined(separator: "\n") ?? "N/A"
34+
)
35+
###
36+
"""
37+
}
38+
}
39+
40+
extension ActiveDocumentChatContextCollector {
41+
struct Information {
42+
let editorContent: SourceEditor.Content?
43+
let selectedContent: String
44+
let documentURL: URL
45+
let projectURL: URL
46+
let language: CodeLanguage
47+
}
48+
49+
func getEditorInformation() -> Information {
50+
let editorContent = XcodeInspector.shared.focusedEditor?.content
51+
let documentURL = XcodeInspector.shared.activeDocumentURL
52+
let projectURL = XcodeInspector.shared.activeProjectURL
53+
let language = languageIdentifierFromFileURL(documentURL)
54+
55+
if let editorContent, let range = editorContent.selections.first {
56+
let startIndex = min(
57+
max(0, range.start.line),
58+
editorContent.lines.endIndex - 1
59+
)
60+
let endIndex = min(
61+
max(startIndex, range.end.line),
62+
editorContent.lines.endIndex - 1
63+
)
64+
let selectedContent = editorContent.lines[startIndex...endIndex]
65+
return .init(
66+
editorContent: editorContent,
67+
selectedContent: selectedContent.joined(),
68+
documentURL: documentURL,
69+
projectURL: projectURL,
70+
language: language
71+
)
72+
}
73+
74+
return .init(
75+
editorContent: editorContent,
76+
selectedContent: "",
77+
documentURL: documentURL,
78+
projectURL: projectURL,
79+
language: language
80+
)
81+
}
82+
}
83+

‎Core/Sources/ChatContextCollector/ChatContextCollector.swift‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
import Foundation
2+
3+
public protocol ChatContextCollector {
4+
func generateSystemPrompt(oldMessages: [String]) -> String
5+
}

‎Core/Sources/ChatService/ChatService.swift‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import ChatContextCollector
12
import ChatPlugins
23
import Combine
34
import Foundation
@@ -30,7 +31,10 @@ public final class ChatService: ObservableObject {
3031
TerminalChatPlugin.self,
3132
AITerminalChatPlugin.self
3233
)
33-
contextController = DynamicContextController(chatGPTService: chatGPTService)
34+
contextController = DynamicContextController(
35+
chatGPTService: chatGPTService,
36+
contextCollectors: ActiveDocumentChatContextCollector()
37+
)
3438

3539
chatGPTService.objectWillChange.sink { [weak self] _ in
3640
self?.objectWillChange.send()
@@ -74,7 +78,7 @@ public final class ChatService: ObservableObject {
7478
public func mutateSystemPrompt(_ newPrompt: String?) {
7579
systemPrompt = newPrompt ?? defaultSystemPrompt
7680
}
77-
81+
7882
public func mutateExtraSystemPrompt(_ newPrompt: String) {
7983
extraSystemPrompt = newPrompt
8084
}

‎Core/Sources/ChatService/DynamicContextController.swift‎

Lines changed: 11 additions & 69 deletions
Original file line numberDiff line numberDiff line change
@@ -1,90 +1,32 @@
1+
import ChatContextCollector
12
import Foundation
23
import OpenAIService
3-
import SuggestionModel
4-
import XcodeInspector
54
import Preferences
5+
import XcodeInspector
66

77
final class DynamicContextController {
88
let chatGPTService: any ChatGPTServiceType
9+
let contextCollectors: [ChatContextCollector]
910

10-
init(chatGPTService: any ChatGPTServiceType) {
11+
init(chatGPTService: any ChatGPTServiceType, contextCollectors: ChatContextCollector...) {
1112
self.chatGPTService = chatGPTService
13+
self.contextCollectors = contextCollectors
1214
}
1315

1416
func updatePromptToMatchContent(systemPrompt: String) async throws {
1517
let language = UserDefaults.shared.value(for: \.chatGPTLanguage)
16-
let content = getEditorInformation()
17-
let relativePath = content.documentURL.path
18-
.replacingOccurrences(of: content.projectURL.path, with: "")
19-
let selectionRange = content.editorContent?.selections.first ?? .outOfScope
18+
let oldMessages = (await chatGPTService.history).map(\.content)
2019
let contextualSystemPrompt = """
2120
\(language.isEmpty ? "" : "You must always reply in \(language)")
2221
\(systemPrompt)
2322
24-
Active Document Context:###
25-
Document Relative Path: \(relativePath)
26-
Selection Range Start: \
27-
Line \(selectionRange.start.line) \
28-
Character \(selectionRange.start.character)
29-
Selection Range End: \
30-
Line \(selectionRange.end.line) \
31-
Character \(selectionRange.end.character)
32-
Cursor Position: \
33-
Line \(selectionRange.end.line) \
34-
Character \(selectionRange.end.character)
35-
Selected Code (start from line \(selectionRange.end.line)):```\(content.language.rawValue)
36-
\(content.selectedContent)
37-
```
38-
39-
Line Annotations:
40-
\(content.editorContent?.lineAnnotations.map { "- \($0)" }.joined(separator: "\n") ?? "N/A")
41-
###
23+
\(
24+
contextCollectors
25+
.map { $0.generateSystemPrompt(oldMessages: oldMessages) }
26+
.joined(separator: "\n")
27+
)
4228
"""
4329
await chatGPTService.mutateSystemPrompt(contextualSystemPrompt)
4430
}
4531
}
4632

47-
extension DynamicContextController {
48-
struct Information {
49-
let editorContent: SourceEditor.Content?
50-
let selectedContent: String
51-
let documentURL: URL
52-
let projectURL: URL
53-
let language: CodeLanguage
54-
}
55-
56-
func getEditorInformation() -> Information {
57-
let editorContent = XcodeInspector.shared.focusedEditor?.content
58-
let documentURL = XcodeInspector.shared.activeDocumentURL
59-
let projectURL = XcodeInspector.shared.activeProjectURL
60-
let language = languageIdentifierFromFileURL(documentURL)
61-
62-
if let editorContent, let range = editorContent.selections.first {
63-
let startIndex = min(
64-
max(0, range.start.line),
65-
editorContent.lines.endIndex - 1
66-
)
67-
let endIndex = min(
68-
max(startIndex, range.end.line),
69-
editorContent.lines.endIndex - 1
70-
)
71-
let selectedContent = editorContent.lines[startIndex...endIndex]
72-
return .init(
73-
editorContent: editorContent,
74-
selectedContent: selectedContent.joined(),
75-
documentURL: documentURL,
76-
projectURL: projectURL,
77-
language: language
78-
)
79-
}
80-
81-
return .init(
82-
editorContent: editorContent,
83-
selectedContent: "",
84-
documentURL: documentURL,
85-
projectURL: projectURL,
86-
language: language
87-
)
88-
}
89-
}
90-

0 commit comments

Comments
 (0)