Skip to content

Commit 78ba023

Browse files
intitniintitni
committed
Make ChatPlugin support string only response
1 parent 71bdcf5 commit 78ba023

4 files changed

Lines changed: 137 additions & 54 deletions

File tree

‎ChatPlugins/Sources/ShortcutChatPlugin/ShortcutChatPlugin.swift‎

Lines changed: 30 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,36 @@ public final class ShortcutChatPlugin: ChatPlugin {
2020
terminal = Terminal()
2121
}
2222

23-
public func send(_ request: Request) async -> AsyncThrowingStream<Response, any Error> {
23+
public func sendForTextResponse(_ request: Request) async
24+
-> AsyncThrowingStream<String, any Error>
25+
{
26+
let stream = await sendForComplicatedResponse(request)
27+
return .init { continuation in
28+
let task = Task {
29+
do {
30+
for try await response in stream {
31+
switch response {
32+
case let .content(.text(content)):
33+
continuation.yield(content)
34+
default:
35+
break
36+
}
37+
}
38+
continuation.finish()
39+
} catch {
40+
continuation.finish(throwing: error)
41+
}
42+
}
43+
44+
continuation.onTermination = { _ in
45+
task.cancel()
46+
}
47+
}
48+
}
49+
50+
public func sendForComplicatedResponse(_ request: Request) async
51+
-> AsyncThrowingStream<Response, any Error>
52+
{
2453
return .init { continuation in
2554
let task = Task {
2655
let id = "\(Self.command)-\(UUID().uuidString)"

‎ChatPlugins/Sources/TerminalChatPlugin/TerminalChatPlugin.swift‎

Lines changed: 77 additions & 50 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@ public final class TerminalChatPlugin: ChatPlugin {
99
public static var name: String { "Terminal" }
1010
public static var description: String { """
1111
Run the command in the message from terminal.
12-
12+
1313
You can use environment variable `$FILE_PATH` and `$PROJECT_ROOT` to access the current file path and project root.
1414
""" }
1515

@@ -23,41 +23,11 @@ public final class TerminalChatPlugin: ChatPlugin {
2323
terminal = Terminal()
2424
}
2525

26-
public func formatContent(_ content: Response.Content) -> Response.Content {
27-
switch content {
28-
case let .text(content):
29-
return .text("""
30-
```sh
31-
\(content)
32-
```
33-
""")
34-
}
35-
}
36-
37-
public func send(_ request: Request) async -> AsyncThrowingStream<Response, any Error> {
26+
public func sendForTextResponse(_ request: Request) async
27+
-> AsyncThrowingStream<String, any Error>
28+
{
3829
return .init { continuation in
3930
let task = Task {
40-
var updateTime = Date()
41-
42-
func streamOutput(_ content: String) {
43-
defer { updateTime = Date() }
44-
if Date().timeIntervalSince(updateTime) > 60 * 2 {
45-
continuation.yield(.startNewMessage)
46-
continuation.yield(.startAction(
47-
id: "run",
48-
task: "Continue `\(request.text)`"
49-
))
50-
continuation.yield(.finishAction(
51-
id: "run",
52-
result: .success("Executed.")
53-
))
54-
continuation.yield(.content(.text("[continue]\n")))
55-
continuation.yield(.content(.text(content)))
56-
} else {
57-
continuation.yield(.content(.text(content)))
58-
}
59-
}
60-
6131
do {
6232
let fileURL = XcodeInspector.shared.realtimeActiveDocumentURL
6333
let projectURL = XcodeInspector.shared.realtimeActiveProjectURL
@@ -75,34 +45,25 @@ public final class TerminalChatPlugin: ChatPlugin {
7545
let env = ProcessInfo.processInfo.environment
7646
let shell = env["SHELL"] ?? "/bin/bash"
7747

78-
continuation.yield(.startAction(id: "run", task: "Run `\(request.text)`"))
79-
8048
let output = terminal.streamCommand(
8149
shell,
8250
arguments: ["-i", "-l", "-c", request.text],
8351
currentDirectoryURL: projectURL,
8452
environment: environment
8553
)
8654

87-
continuation.yield(.finishAction(
88-
id: "run",
89-
result: .success("Executed.")
90-
))
91-
55+
var accumulatedOutput = ""
9256
for try await content in output {
9357
try Task.checkCancellation()
94-
streamOutput(content)
58+
accumulatedOutput += content
59+
continuation.yield(accumulatedOutput)
9560
}
9661
} catch let error as Terminal.TerminationError {
97-
continuation.yield(.content(.text("""
98-
99-
[error: \(error.reason)]
100-
""")))
62+
let errorMessage = "\n\n[error: \(error.reason)]"
63+
continuation.yield(errorMessage)
10164
} catch {
102-
continuation.yield(.content(.text("""
103-
104-
[error: \(error.localizedDescription)]
105-
""")))
65+
let errorMessage = "\n\n[error: \(error.localizedDescription)]"
66+
continuation.yield(errorMessage)
10667
}
10768

10869
continuation.finish()
@@ -116,5 +77,71 @@ public final class TerminalChatPlugin: ChatPlugin {
11677
}
11778
}
11879
}
80+
81+
public func formatContent(_ content: Response.Content) -> Response.Content {
82+
switch content {
83+
case let .text(content):
84+
return .text("""
85+
```sh
86+
\(content)
87+
```
88+
""")
89+
}
90+
}
91+
92+
public func sendForComplicatedResponse(_ request: Request) async
93+
-> AsyncThrowingStream<Response, any Error>
94+
{
95+
return .init { continuation in
96+
let task = Task {
97+
var updateTime = Date()
98+
99+
continuation.yield(.startAction(id: "run", task: "Run `\(request.text)`"))
100+
101+
let textStream = await sendForTextResponse(request)
102+
var previousOutput = ""
103+
104+
continuation.yield(.finishAction(
105+
id: "run",
106+
result: .success("Executed.")
107+
))
108+
109+
do {
110+
for try await accumulatedOutput in textStream {
111+
try Task.checkCancellation()
112+
113+
let newContent = accumulatedOutput.dropFirst(previousOutput.count)
114+
previousOutput = accumulatedOutput
115+
116+
if !newContent.isEmpty {
117+
if Date().timeIntervalSince(updateTime) > 60 * 2 {
118+
continuation.yield(.startNewMessage)
119+
continuation.yield(.startAction(
120+
id: "run",
121+
task: "Continue `\(request.text)`"
122+
))
123+
continuation.yield(.finishAction(
124+
id: "run",
125+
result: .success("Executed.")
126+
))
127+
continuation.yield(.content(.text("[continue]\n")))
128+
updateTime = Date()
129+
}
130+
131+
continuation.yield(.content(.text(String(newContent))))
132+
}
133+
}
134+
135+
continuation.finish()
136+
} catch {
137+
continuation.finish(throwing: error)
138+
}
139+
}
140+
141+
continuation.onTermination = { _ in
142+
task.cancel()
143+
}
144+
}
145+
}
119146
}
120147

‎Core/Sources/ChatService/AllPlugins.swift‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -56,7 +56,7 @@ final class LegacyChatPluginWrapper<Plugin: ChatPlugin>: LegacyChatPlugin {
5656

5757
let plugin = Plugin()
5858

59-
let stream = await plugin.send(.init(
59+
let stream = await plugin.sendForComplicatedResponse(.init(
6060
text: content,
6161
arguments: [],
6262
history: chatGPTService.memory.history

‎Tool/Sources/ChatBasic/ChatPlugin.swift‎

Lines changed: 29 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,13 @@ public protocol ChatPlugin {
1919
static var command: String { get }
2020
static var name: String { get }
2121
static var description: String { get }
22-
func send(_ request: Request) async -> AsyncThrowingStream<Response, any Error>
22+
// In this method, the plugin is able to send more complicated response. It also enables it to
23+
// perform special tasks like starting a new message or reporting progress.
24+
func sendForComplicatedResponse(
25+
_ request: Request
26+
) async -> AsyncThrowingStream<Response, any Error>
27+
// This method allows the plugin to respond a stream of text content only.
28+
func sendForTextResponse(_ request: Request) async -> AsyncThrowingStream<String, any Error>
2329
func formatContent(_ content: Response.Content) -> Response.Content
2430
init()
2531
}
@@ -28,5 +34,26 @@ public extension ChatPlugin {
2834
func formatContent(_ content: Response.Content) -> Response.Content {
2935
return content
3036
}
37+
38+
func sendForComplicatedResponse(
39+
_ request: Request
40+
) async -> AsyncThrowingStream<Response, any Error> {
41+
let textStream = await sendForTextResponse(request)
42+
return AsyncThrowingStream<Response, any Error> { continuation in
43+
let task = Task {
44+
do {
45+
for try await text in textStream {
46+
continuation.yield(Response.content(.text(text)))
47+
}
48+
continuation.finish()
49+
} catch {
50+
continuation.finish(throwing: error)
51+
}
52+
}
53+
54+
continuation.onTermination = { _ in
55+
task.cancel()
56+
}
57+
}
58+
}
3159
}
32-

0 commit comments

Comments
 (0)