Skip to content

Commit 290610b

Browse files
intitniintitni
committed
Update system prompt to be contextual
1 parent e9dacf8 commit 290610b

5 files changed

Lines changed: 106 additions & 56 deletions

File tree

‎Core/Sources/ChatService/ChatService.swift‎

Lines changed: 26 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,10 +3,22 @@ import Combine
33
import Foundation
44
import OpenAIService
55

6+
let defaultSystemPrompt = """
7+
You are an AI programming assistant.
8+
You reply should be concise, clear, informative and logical.
9+
You MUST reply in the format of markdown.
10+
You MUST embed every code you provide in a markdown code block.
11+
You MUST add the programming language name at the start of the markdown code block.
12+
If you are asked to help perform a task, think step-by-step.
13+
"""
14+
615
public final class ChatService: ObservableObject {
716
public let chatGPTService: any ChatGPTServiceType
817
let pluginController: ChatPluginController
18+
let contextController: DynamicContextController
919
var cancellable = Set<AnyCancellable>()
20+
var systemPrompt = defaultSystemPrompt
21+
var extraSystemPrompt = ""
1022

1123
public init<T: ChatGPTServiceType>(chatGPTService: T) {
1224
self.chatGPTService = chatGPTService
@@ -16,6 +28,7 @@ public final class ChatService: ObservableObject {
1628
TerminalChatPlugin.self,
1729
AITerminalChatPlugin.self
1830
)
31+
contextController = DynamicContextController(chatGPTService: chatGPTService)
1932

2033
chatGPTService.objectWillChange.sink { [weak self] _ in
2134
self?.objectWillChange.send()
@@ -25,7 +38,11 @@ public final class ChatService: ObservableObject {
2538
public func send(content: String) async throws {
2639
let handledInPlugin = try await pluginController.handleContent(content)
2740
if handledInPlugin { return }
28-
41+
try await contextController.updatePromptToMatchContent(systemPrompt: """
42+
\(systemPrompt)
43+
\(extraSystemPrompt)
44+
""")
45+
2946
_ = try await chatGPTService.send(content: content, summary: nil)
3047
}
3148

@@ -51,11 +68,17 @@ public final class ChatService: ObservableObject {
5168
}
5269
}
5370

54-
public func mutateSystemPrompt(_ newPrompt: String) async {
55-
await chatGPTService.mutateSystemPrompt(newPrompt)
71+
/// Setting it to `nil` to reset the system prompt
72+
public func mutateSystemPrompt(_ newPrompt: String?) {
73+
systemPrompt = newPrompt ?? defaultSystemPrompt
5674
}
5775

76+
public func mutateExtraSystemPrompt(_ newPrompt: String) {
77+
extraSystemPrompt = newPrompt
78+
}
79+
5880
public func mutateHistory(_ mutator: @escaping (inout [ChatMessage]) -> Void) async {
5981
await chatGPTService.mutateHistory(mutator)
6082
}
6183
}
84+

‎Core/Sources/ChatService/DynamicContextController.swift‎

Lines changed: 73 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,24 +1,89 @@
11
import Foundation
22
import OpenAIService
3+
import SuggestionModel
34
import XcodeInspector
5+
import Preferences
46

57
final class DynamicContextController {
68
let chatGPTService: any ChatGPTServiceType
7-
9+
810
init(chatGPTService: any ChatGPTServiceType) {
911
self.chatGPTService = chatGPTService
1012
}
11-
13+
1214
func updatePromptToMatchContent(systemPrompt: String) async throws {
13-
15+
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
20+
let contextualSystemPrompt = """
21+
\(language.isEmpty ? "" : "You must always reply in \(language)")
22+
\(systemPrompt)
23+
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: Same as selection range end
33+
Selected Code (start from line \(selectionRange.end.line)):```\(content.language.rawValue)
34+
\(content.selectedContent)
35+
```
36+
37+
Line Annotations:
38+
\(content.editorContent?.lineAnnotations.map { "- \($0)" }.joined(separator: "\n") ?? "N/A")
39+
###
40+
"""
41+
print(contextualSystemPrompt)
42+
await chatGPTService.mutateSystemPrompt(contextualSystemPrompt)
1443
}
1544
}
1645

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

‎Core/Sources/Service/GUI/WidgetDataSource.swift‎

Lines changed: 1 addition & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -42,13 +42,7 @@ final class WidgetDataSource {
4242
func createChatIfNeeded(for url: URL) -> ChatService {
4343
let build = {
4444
let language = UserDefaults.shared.value(for: \.chatGPTLanguage)
45-
let systemPrompt = """
46-
\(language.isEmpty ? "" : "You must always reply in \(language)")
47-
You are a senior programmer, you will answer my questions concisely. If you are replying with code, embed the code in a code block in markdown.
48-
49-
You don't have any code in advance, ask me to provide it when needed.
50-
"""
51-
let service = ChatService(chatGPTService: ChatGPTService(systemPrompt: systemPrompt))
45+
let service = ChatService(chatGPTService: ChatGPTService())
5246
let provider = ChatProvider(
5347
service: service,
5448
fileURL: url,

‎Core/Sources/Service/SuggestionCommandHandler/WindowBaseCommandHandler.swift‎

Lines changed: 5 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -397,40 +397,17 @@ extension WindowBaseCommandHandler {
397397
defer { presenter.markAsProcessing(false) }
398398

399399
let fileURL = try await Environment.fetchCurrentFileURL()
400-
let language = UserDefaults.shared.value(for: \.chatGPTLanguage)
401-
let codeLanguage = languageIdentifierFromFileURL(fileURL)
402400

403401
let code = {
404402
guard let selection = editor.selections.last,
405403
selection.start != selection.end else { return "" }
406404
return editor.selectedCode(in: selection)
407405
}()
408406

409-
var systemPrompt = specifiedSystemPrompt ?? {
410-
if code.isEmpty {
411-
return """
412-
\(language.isEmpty ? "" : "You must always reply in \(language)")
413-
You are a senior programmer, you will answer my questions concisely. If you are replying with code, embed the code in a code block in markdown.
414-
415-
You don't have any code in advance, ask me to provide it when needed.
416-
"""
417-
}
418-
return """
419-
\(language.isEmpty ? "" : "You must always reply in \(language)")
420-
You are a senior programmer, you will answer my questions concisely about the code below, or modify it according to my requests. When you receive a modification request, reply with the modified code in a code block.
421-
```\(codeLanguage.rawValue)
422-
\(code)
423-
```
424-
"""
425-
}()
426-
427-
if let extraSystemPrompt {
428-
systemPrompt += "\n\(extraSystemPrompt)"
429-
}
430-
431407
let chat = WidgetDataSource.shared.createChatIfNeeded(for: fileURL)
432-
433-
await chat.mutateSystemPrompt(systemPrompt)
408+
409+
chat.mutateSystemPrompt(specifiedSystemPrompt)
410+
chat.mutateExtraSystemPrompt(extraSystemPrompt ?? "")
434411

435412
Task {
436413
let customCommandPrefix = {
@@ -484,20 +461,10 @@ extension WindowBaseCommandHandler {
484461
let focusedElementURI = try await Environment.fetchFocusedElementURI()
485462
let language = UserDefaults.shared.value(for: \.chatGPTLanguage)
486463

487-
var systemPrompt = specifiedSystemPrompt ?? """
488-
\(language.isEmpty ? "" : "You must always reply in \(language)")
489-
You are a senior programmer, you will answer my questions concisely. If you are replying with code, embed the code in a code block in markdown.
490-
491-
You don't have any code in advance, ask me to provide it when needed.
492-
"""
493-
494-
if let extraSystemPrompt {
495-
systemPrompt += "\n\(extraSystemPrompt)"
496-
}
497-
498464
let chat = WidgetDataSource.shared.createChatIfNeeded(for: focusedElementURI)
499465

500-
await chat.mutateSystemPrompt(systemPrompt)
466+
chat.mutateSystemPrompt(specifiedSystemPrompt)
467+
chat.mutateExtraSystemPrompt(extraSystemPrompt ?? "")
501468

502469
Task {
503470
let customCommandPrefix = {

‎ExtensionService/ServiceDelegate.swift‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,3 +17,4 @@ class ServiceDelegate: NSObject, NSXPCListenerDelegate {
1717
return true
1818
}
1919
}
20+

0 commit comments

Comments
 (0)