Skip to content

Commit 4201df1

Browse files
intitniintitni
committed
Merge branch 'feature/ai-terminal' into develop
2 parents 0bb28e3 + 01260c7 commit 4201df1

9 files changed

Lines changed: 281 additions & 9 deletions

File tree

‎Core/Sources/ChatPlugins/AITerminalChatPlugin.swift‎

Lines changed: 192 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,192 @@
1+
import Environment
2+
import Foundation
3+
import OpenAIService
4+
import Terminal
5+
6+
public actor AITerminalChatPlugin: ChatPlugin {
7+
public static var command: String { "airun" }
8+
public nonisolated var name: String { "AI Terminal" }
9+
10+
let chatGPTService: any ChatGPTServiceType
11+
var terminal: TerminalType = Terminal()
12+
var isCancelled = false
13+
weak var delegate: ChatPluginDelegate?
14+
var isStarted = false
15+
var command: String?
16+
17+
public init(inside chatGPTService: any ChatGPTServiceType, delegate: ChatPluginDelegate) {
18+
self.chatGPTService = chatGPTService
19+
self.delegate = delegate
20+
}
21+
22+
public func send(content: String) async {
23+
if !isStarted {
24+
isStarted = true
25+
delegate?.pluginDidStart(self)
26+
}
27+
28+
do {
29+
if let command {
30+
await chatGPTService.mutateHistory { history in
31+
history.append(.init(role: .user, content: content))
32+
}
33+
delegate?.pluginDidStartResponding(self)
34+
if isCancelled { return }
35+
switch try await checkConfirmation(content: content) {
36+
case .confirmation:
37+
delegate?.pluginDidEndResponding(self)
38+
delegate?.pluginDidEnd(self)
39+
delegate?.shouldStartAnotherPlugin(
40+
TerminalChatPlugin.self,
41+
withContent: command
42+
)
43+
case .cancellation:
44+
delegate?.pluginDidEndResponding(self)
45+
delegate?.pluginDidEnd(self)
46+
await chatGPTService.mutateHistory { history in
47+
history.append(.init(role: .assistant, content: "Cancelled"))
48+
}
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+
}
69+
}
70+
} else {
71+
await chatGPTService.mutateHistory { history in
72+
history.append(.init(role: .user, content: "Run a command to \(content)"))
73+
}
74+
delegate?.pluginDidStartResponding(self)
75+
let result = try await generateCommand(task: content)
76+
command = result
77+
if isCancelled { return }
78+
await chatGPTService.mutateHistory { history in
79+
history.append(.init(role: .assistant, content: """
80+
Confirm to run?
81+
```
82+
\(result)
83+
```
84+
"""))
85+
}
86+
delegate?.pluginDidEndResponding(self)
87+
}
88+
} catch {
89+
await chatGPTService.mutateHistory { history in
90+
history.append(.init(role: .assistant, content: error.localizedDescription))
91+
}
92+
delegate?.pluginDidEndResponding(self)
93+
delegate?.pluginDidEnd(self)
94+
}
95+
}
96+
97+
public func cancel() async {
98+
isCancelled = true
99+
delegate?.pluginDidEndResponding(self)
100+
delegate?.pluginDidEnd(self)
101+
}
102+
103+
public func stopResponding() async {}
104+
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+
))
122+
}
123+
124+
func modifyCommand(command: String, requirement: String) async throws -> String {
125+
let p = """
126+
Available environment variables:
127+
- $PROJECT_ROOT: the root path of the project
128+
- $FILE_PATH: the currently editing file
129+
130+
Current directory path is the project root.
131+
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.
135+
136+
The reply should contains only the command and nothing else.
137+
"""
138+
139+
return extractCodeFromMarkdown(try await askChatGPT(
140+
systemPrompt: p,
141+
question: "The requirement is: \"\(requirement)\""
142+
))
143+
}
144+
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.
150+
151+
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.
162+
"""
163+
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])
191+
}
192+
}

‎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+
}

‎Core/Sources/ChatPlugins/ChatPlugin.swift‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
import Foundation
22
import OpenAIService
33

4-
public protocol ChatPlugin {
4+
public protocol ChatPlugin: AnyObject {
55
/// Should be [a-zA-Z0-9]+
66
static var command: String { get }
77
var name: String { get }
@@ -17,4 +17,5 @@ public protocol ChatPluginDelegate: AnyObject {
1717
func pluginDidEnd(_ plugin: ChatPlugin)
1818
func pluginDidStartResponding(_ plugin: ChatPlugin)
1919
func pluginDidEndResponding(_ plugin: ChatPlugin)
20+
func shouldStartAnotherPlugin(_ type: ChatPlugin.Type, withContent: String)
2021
}

‎Core/Sources/ChatPlugins/TerminalChatPlugin.swift‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -43,9 +43,12 @@ public actor TerminalChatPlugin: ChatPlugin {
4343

4444
if isCancelled { throw CancellationError() }
4545

46+
let env = ProcessInfo.processInfo.environment
47+
let shell = env["SHELL"] ?? "/bin/bash"
48+
4649
let output = terminal.streamCommand(
47-
"/bin/bash",
48-
arguments: ["-c", content],
50+
shell,
51+
arguments: ["-l", "-c", content],
4952
currentDirectoryPath: projectURL?.path ?? fileURL.path,
5053
environment: [
5154
"PROJECT_ROOT": projectURL?.path ?? fileURL.path,

‎Core/Sources/ChatService/ChatService.swift‎

Lines changed: 32 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,8 @@ import OpenAIService
66
public final class ChatService: ObservableObject {
77
public let chatGPTService: any ChatGPTServiceType
88
let plugins = registerPlugins(
9-
TerminalChatPlugin.self
9+
TerminalChatPlugin.self,
10+
AITerminalChatPlugin.self
1011
)
1112
var runningPlugin: ChatPlugin?
1213
var cancellable = Set<AnyCancellable>()
@@ -27,8 +28,9 @@ public final class ChatService: ObservableObject {
2728
let matches = regex.matches(in: content, range: NSRange(content.startIndex..., in: content))
2829
if let match = matches.first {
2930
let command = String(content[Range(match.range(at: 1), in: content)!])
30-
if command == "/exit" {
31+
if command == "exit" {
3132
if let plugin = runningPlugin {
33+
runningPlugin = nil
3234
_ = await chatGPTService.mutateHistory { history in
3335
history.append(.init(
3436
role: .user,
@@ -50,18 +52,33 @@ public final class ChatService: ObservableObject {
5052
))
5153
}
5254
}
55+
} else if let runningPlugin {
56+
await runningPlugin.send(content: content)
5357
} else if let pluginType = plugins[command] {
5458
let plugin = pluginType.init(inside: chatGPTService, delegate: self)
55-
await plugin.send(content: String(content.dropFirst(command.count + 1)))
59+
if #available(macOS 13.0, *) {
60+
await plugin.send(
61+
content: String(
62+
content.dropFirst(command.count + 1)
63+
.trimmingPrefix(while: { $0 == " " })
64+
)
65+
)
66+
} else {
67+
await plugin.send(content: String(content.dropFirst(command.count + 1)))
68+
}
69+
} else {
70+
_ = try await chatGPTService.send(content: content, summary: nil)
5671
}
72+
} else if let runningPlugin {
73+
await runningPlugin.send(content: content)
5774
} else {
5875
_ = try await chatGPTService.send(content: content, summary: nil)
5976
}
6077
}
6178

6279
public func stopReceivingMessage() async {
6380
if let runningPlugin {
64-
await runningPlugin.cancel()
81+
await runningPlugin.stopResponding()
6582
}
6683
await chatGPTService.stopReceivingMessage()
6784
}
@@ -95,8 +112,17 @@ extension ChatService: ChatPluginDelegate {
95112
runningPlugin = plugin
96113
}
97114

98-
public func pluginDidEnd(_: ChatPlugin) {
99-
runningPlugin = nil
115+
public func pluginDidEnd(_ plugin: ChatPlugin) {
116+
if runningPlugin === plugin {
117+
runningPlugin = nil
118+
}
119+
}
120+
121+
public func shouldStartAnotherPlugin(_ type: ChatPlugin.Type, withContent content: String) {
122+
let plugin = type.init(inside: chatGPTService, delegate: self)
123+
Task {
124+
await plugin.send(content: content)
125+
}
100126
}
101127
}
102128

‎Core/Sources/OpenAIService/ChatGPTService.swift‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -172,6 +172,18 @@ public actor ChatGPTService: ChatGPTServiceType {
172172
}
173173
}
174174
}
175+
176+
public func sendAndWait(
177+
content: String,
178+
summary: String? = nil
179+
) async throws -> String {
180+
let stream = try await send(content: content, summary: summary)
181+
var content = ""
182+
for try await fragment in stream {
183+
content.append(fragment)
184+
}
185+
return content
186+
}
175187

176188
public func stopReceivingMessage() {
177189
cancelTask?()

‎Core/Sources/SuggestionWidget/SuggestionPanelContent/ChatPanel.swift‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -227,6 +227,7 @@ struct ChatPanelInputArea: View {
227227
.textFieldStyle(.plain)
228228
.padding(8)
229229
.onSubmit {
230+
if chat.isReceivingMessage { return }
230231
if typedMessage.isEmpty { return }
231232
chat.send(typedMessage)
232233
typedMessage = ""

‎README.md‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -143,9 +143,15 @@ The chat panel supports chat plugins that may not require an OpenAI API key. For
143143
/run echo hello
144144
```
145145

146+
If you need to end a plugin, you can just type
147+
```
148+
/exit
149+
```
150+
146151
| Command | Description |
147152
|:---:|---|
148153
| `/run` | Runs the command under the project root. You can also use environment variable `PROJECT_ROOT` to get the project root and `FILE_PATH` to get the editing file path.|
154+
| `/airun` | Create a command with natural language. You can ask to modify the command if it is not what you want. After confirming, the command will be executed by calling the `/run` plugin. |
149155

150156
## Key Bindings
151157

0 commit comments

Comments
 (0)