forked from intitni/CopilotForXcode
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathOpenAIChat.swift
More file actions
52 lines (46 loc) · 1.7 KB
/
OpenAIChat.swift
File metadata and controls
52 lines (46 loc) · 1.7 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
import Foundation
import OpenAIService
public struct OpenAIChat: ChatModel {
public var configuration: ChatGPTConfiguration
public var memory: ChatGPTMemory?
public var functionProvider: ChatGPTFunctionProvider
public var stream: Bool
public init(
configuration: ChatGPTConfiguration = UserPreferenceChatGPTConfiguration(),
memory: ChatGPTMemory? = ConversationChatGPTMemory(systemPrompt: ""),
functionProvider: ChatGPTFunctionProvider = NoChatGPTFunctionProvider(),
stream: Bool
) {
self.configuration = configuration
self.memory = memory
self.functionProvider = functionProvider
self.stream = stream
}
public func generate(
prompt: [ChatMessage],
stops: [String],
callbackManagers: [CallbackManager]
) async throws -> ChatMessage {
let memory = memory ?? EmptyChatGPTMemory()
let service = ChatGPTService(
memory: memory,
configuration: configuration,
functionProvider: functionProvider
)
for message in prompt {
await memory.appendMessage(message)
}
if stream {
let stream = try await service.send(content: "")
var message = ""
for try await chunk in stream {
message.append(chunk)
callbackManagers.send(CallbackEvents.LLMDidProduceNewToken(info: chunk))
}
return await memory.history.last ?? .init(role: .assistant, content: "")
} else {
let _ = try await service.sendAndWait(content: "")
return await memory.history.last ?? .init(role: .assistant, content: "")
}
}
}