Skip to content

Commit 4817eac

Browse files
intitniintitni
committed
Update prompt and support command modification
1 parent ed20d8c commit 4817eac

File tree

3 files changed

+123
-33
lines changed

3 files changed

+123
-33
lines changed

‎Core/Sources/ChatPlugins/AITerminalChatPlugin.swift‎

Lines changed: 92 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -32,19 +32,40 @@ public actor AITerminalChatPlugin: ChatPlugin {
3232
}
3333
delegate?.pluginDidStartResponding(self)
3434
if isCancelled { return }
35-
if try await checkConfirmation(content: content) {
35+
switch try await checkConfirmation(content: content) {
36+
case .confirmation:
3637
delegate?.pluginDidEndResponding(self)
3738
delegate?.pluginDidEnd(self)
3839
delegate?.shouldStartAnotherPlugin(
3940
TerminalChatPlugin.self,
4041
withContent: command
4142
)
42-
} else {
43+
case .cancellation:
4344
delegate?.pluginDidEndResponding(self)
4445
delegate?.pluginDidEnd(self)
4546
await chatGPTService.mutateHistory { history in
4647
history.append(.init(role: .assistant, content: "Cancelled"))
4748
}
49+
case .modification:
50+
let result = try await modifyCommand(command: command, requirement: content)
51+
self.command = result
52+
delegate?.pluginDidEndResponding(self)
53+
await chatGPTService.mutateHistory { history in
54+
history.append(.init(role: .assistant, content: """
55+
Confirm to run?
56+
```
57+
\(result)
58+
```
59+
"""))
60+
}
61+
case .other:
62+
delegate?.pluginDidEndResponding(self)
63+
await chatGPTService.mutateHistory { history in
64+
history.append(.init(
65+
role: .assistant,
66+
content: "Should I run it? Or should I modify it?"
67+
))
68+
}
4869
}
4970
} else {
5071
await chatGPTService.mutateHistory { history in
@@ -81,53 +102,91 @@ public actor AITerminalChatPlugin: ChatPlugin {
81102

82103
public func stopResponding() async {}
83104

84-
func callAIFunction(
85-
function: String,
86-
args: [Any?],
87-
description: String
88-
) async throws -> String {
89-
let args = args.map { arg -> String in
90-
if let arg = arg {
91-
return String(describing: arg)
92-
} else {
93-
return "None"
94-
}
95-
}
96-
let argsString = args.joined(separator: ", ")
97-
let service = ChatGPTService(
98-
systemPrompt: "You are now the following python function: ```# \(description)\n\(function)```\n\nOnly respond with your `return` value."
99-
)
100-
return try await service.sendAndWait(content: argsString)
105+
func generateCommand(task: String) async throws -> String {
106+
let p = """
107+
Available environment variables:
108+
- $PROJECT_ROOT: the root path of the project
109+
- $FILE_PATH: the currently editing file
110+
111+
Current directory path is the project root.
112+
113+
Generate a terminal command to solve the given task on macOS. If one command is not enough, you can use && to concatenate multiple commands.
114+
115+
The reply should contains only the command and nothing else.
116+
"""
117+
118+
return extractCodeFromMarkdown(try await askChatGPT(
119+
systemPrompt: p,
120+
question: "the task is: \"\(task)\""
121+
))
101122
}
102123

103-
func generateCommand(task: String) async throws -> String {
104-
let f = "def generate_terminal_command(task: str) -> string:"
105-
let d = """
124+
func modifyCommand(command: String, requirement: String) async throws -> String {
125+
let p = """
106126
Available environment variables:
107127
- $PROJECT_ROOT: the root path of the project
108128
- $FILE_PATH: the currently editing file
109129
110130
Current directory path is the project root.
111131
112-
The return value should not be embedded in a markdown code block.
132+
Modify the terminal command `\(
133+
command
134+
)` in macOS with the given requirement. If one command is not enough, you can use && to concatenate multiple commands.
113135
114-
Generate a terminal command to solve the given task on macOS. If one command is not enough, you can use && to concatenate multiple commands.
136+
The reply should contains only the command and nothing else.
115137
"""
116138

117-
return try await callAIFunction(function: f, args: [task], description: d)
118-
.replacingOccurrences(of: "`", with: "")
119-
.replacingOccurrences(of: "\n", with: "")
139+
return extractCodeFromMarkdown(try await askChatGPT(
140+
systemPrompt: p,
141+
question: "The requirement is: \"\(requirement)\""
142+
))
120143
}
121144

122-
func checkConfirmation(content: String) async throws -> Bool {
123-
let f = "def check_confirmation(content: str) -> bool:"
124-
let d = """
125-
Check if the given content is a phrase or sentence that considered a confirmation to run a command.
145+
func checkConfirmation(content: String) async throws -> Tone {
146+
let p = """
147+
Check the tone of the content, reply with only the number representing the tone.
148+
149+
1: If the given content is a phrase or sentence that considered a confirmation to run a command.
126150
127151
For example: "Yes", "Confirm", "True", "Run it". It can be in any language.
152+
153+
2: If the given content is a phrase or sentence that considered a cancellation to run a command.
154+
155+
For example: "No", "Cancel", "False", "Don't run it", "Stop". It can be in any language.
156+
157+
3: If the given content is a modification request.
158+
159+
For example: "Use echo instead", "Remove the argument", "Change to path".
160+
161+
4: Everything else.
128162
"""
129163

130-
let result = try await callAIFunction(function: f, args: [content], description: d)
131-
return result.lowercased().contains("true")
164+
let result = try await askChatGPT(
165+
systemPrompt: p,
166+
question: "The content is: \"\(content)\""
167+
)
168+
return Tone(rawValue: Int(result) ?? 2) ?? .cancellation
169+
}
170+
171+
enum Tone: Int {
172+
case confirmation = 1
173+
case cancellation = 2
174+
case modification = 3
175+
case other = 4
176+
}
177+
178+
func extractCodeFromMarkdown(_ markdown: String) -> String {
179+
let codeBlockRegex = try! NSRegularExpression(
180+
pattern: "```[\n](.*?)[\n]```",
181+
options: .dotMatchesLineSeparators
182+
)
183+
let range = NSRange(markdown.startIndex..<markdown.endIndex, in: markdown)
184+
guard let match = codeBlockRegex.firstMatch(in: markdown, options: [], range: range) else {
185+
return markdown
186+
.replacingOccurrences(of: "`", with: "")
187+
.replacingOccurrences(of: "\n", with: "")
188+
}
189+
let codeBlockRange = Range(match.range(at: 1), in: markdown)!
190+
return String(markdown[codeBlockRange])
132191
}
133192
}

‎Core/Sources/ChatPlugins/AskChatGPT.swift‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
import Foundation
2+
import OpenAIService
3+
4+
/// Quickly ask a question to ChatGPT.
5+
func askChatGPT(systemPrompt: String, question: String) async throws -> String {
6+
let service = ChatGPTService(systemPrompt: systemPrompt)
7+
return try await service.sendAndWait(content: question)
8+
}

‎Core/Sources/ChatPlugins/CallAIFunction.swift‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,23 @@
1+
import Foundation
2+
import OpenAIService
3+
4+
/// This is a magic function that can do anything with no-code. See
5+
/// https://github.com/Torantulino/AI-Functions for more info.
6+
func callAIFunction(
7+
function: String,
8+
args: [Any?],
9+
description: String
10+
) async throws -> String {
11+
let args = args.map { arg -> String in
12+
if let arg = arg {
13+
return String(describing: arg)
14+
} else {
15+
return "None"
16+
}
17+
}
18+
let argsString = args.joined(separator: ", ")
19+
let service = ChatGPTService(
20+
systemPrompt: "You are now the following python function: ```# \(description)\n\(function)```\n\nOnly respond with your `return` value."
21+
)
22+
return try await service.sendAndWait(content: argsString)
23+
}

0 commit comments

Comments
 (0)