6
6
//
7
7
8
8
import Foundation
9
- import GoogleGenerativeAI
10
9
import SwiftAnthropic
11
10
import SwiftOpenAI
12
11
@@ -16,20 +15,12 @@ enum PolyAIError: Error {
16
15
case missingLLMConfiguration( String )
17
16
}
18
17
19
- // MARK: Gemini specific
20
-
21
- extension GenerativeModel {
22
- public struct Configuration {
23
- let apiKey : String
24
- }
25
- }
26
-
27
18
struct DefaultPolyAIService : PolyAIService {
28
19
29
20
private var openAIService : OpenAIService ?
30
21
private var ollamaOpenAIServiceCompatible : OpenAIService ?
31
22
private var anthropicService : AnthropicService ?
32
- private var gemini : GenerativeModel . Configuration ?
23
+ private var gemini : OpenAIService ?
33
24
34
25
init ( configurations: [ LLMConfiguration ] )
35
26
{
@@ -51,7 +42,14 @@ struct DefaultPolyAIService: PolyAIService {
51
42
anthropicService = AnthropicServiceFactory . service ( apiKey: apiKey, betaHeaders: betaHeaders, configuration: configuration)
52
43
53
44
case . gemini( let apiKey) :
54
- gemini = . init( apiKey: apiKey)
45
+ let baseURL = " https://generativelanguage.googleapis.com "
46
+ let version = " v1beta "
47
+
48
+ let service = OpenAIServiceFactory . service (
49
+ apiKey: apiKey,
50
+ overrideBaseURL: baseURL,
51
+ overrideVersion: version)
52
+ gemini = service
55
53
56
54
case . ollama( let url) :
57
55
ollamaOpenAIServiceCompatible = OpenAIServiceFactory . service ( baseURL: url)
@@ -93,37 +91,13 @@ struct DefaultPolyAIService: PolyAIService {
93
91
return try await anthropicService. createMessage ( messageParameter)
94
92
95
93
case . gemini( let model, let messages, let maxTokens) :
96
- guard let gemini else {
97
- throw PolyAIError . missingLLMConfiguration ( " You Must provide a valid configuration for the \( parameter. llmService) API " )
98
- }
99
-
100
- // Extract system message if present
101
- let systemInstruction : ModelContent ?
102
- if let systemMessage = messages. first ( where: { $0. role == " system " } ) ? . content {
103
- systemInstruction = ModelContent ( parts: [ . text( systemMessage) ] )
104
- } else {
105
- systemInstruction = nil
106
- }
107
-
108
- // Create the model with system instruction
109
- let generativeModel = GenerativeModel (
110
- name: model,
111
- apiKey: gemini. apiKey,
112
- generationConfig: . init( GenerationConfig ( maxOutputTokens: maxTokens) ) ,
113
- systemInstruction: systemInstruction
114
- )
115
-
116
- // Convert messages to ModelContent array for chat history
117
- let chatHistory = messages. filter { $0. role != " system " } . map { message in
118
- ModelContent (
119
- role: message. role,
120
- parts: [ . text( message. content) ]
121
- )
122
- }
123
-
124
- // Create chat with history and send the last message
125
- let chat = generativeModel. startChat ( history: chatHistory)
126
- return try await chat. sendMessage ( " " ) // Empty message since history contains everything
94
+ guard let gemini else {
95
+ throw PolyAIError . missingLLMConfiguration ( " You Must provide a valid configuration for the \( parameter. llmService) API " )
96
+ }
97
+ let messageParams : [ SwiftOpenAI . ChatCompletionParameters . Message ] = messages. map { . init( role: . init( rawValue: $0. role) ?? . user, content: . text( $0. content) ) }
98
+ let messageParameter = ChatCompletionParameters ( messages: messageParams, model: . custom( model) , maxTokens: maxTokens)
99
+ return try await gemini. startChat ( parameters: messageParameter)
100
+
127
101
case . ollama( let model, let messages, let maxTokens) :
128
102
guard let ollamaOpenAIServiceCompatible else {
129
103
throw PolyAIError . missingLLMConfiguration ( " You Must provide a valid configuration for the \( parameter. llmService) API " )
@@ -146,7 +120,6 @@ struct DefaultPolyAIService: PolyAIService {
146
120
let messageParams : [ SwiftOpenAI . ChatCompletionParameters . Message ] = messages. map { . init( role: . init( rawValue: $0. role) ?? . user, content: . text( $0. content) ) }
147
121
let messageParameter = ChatCompletionParameters ( messages: messageParams, model: model, maxTokens: maxTokens)
148
122
149
-
150
123
let stream = try await openAIService. startStreamedChat ( parameters: messageParameter)
151
124
return try mapToLLMMessageStreamResponse ( stream: stream)
152
125
@@ -173,20 +146,12 @@ struct DefaultPolyAIService: PolyAIService {
173
146
guard let gemini else {
174
147
throw PolyAIError . missingLLMConfiguration ( " You Must provide a valid configuration for the \( parameter. llmService) API " )
175
148
}
176
- let systemInstruction : ModelContent ?
177
- if let systemMessage = messages. first ( where: { message in
178
- message. role == " system "
179
- } ) ? . content {
180
- systemInstruction = ModelContent ( parts: [ . text( systemMessage) ] )
181
- } else {
182
- systemInstruction = nil
183
- }
184
- let generativeModel = GenerativeModel ( name: model, apiKey: gemini. apiKey, generationConfig: . init( GenerationConfig ( maxOutputTokens: maxTokens) ) , systemInstruction: systemInstruction)
185
- let userMessage = messages. first { message in
186
- message. role == " user "
187
- } ? . content ?? " "
188
- let stream = generativeModel. generateContentStream ( userMessage)
149
+ let messageParams : [ SwiftOpenAI . ChatCompletionParameters . Message ] = messages. map { . init( role: . init( rawValue: $0. role) ?? . user, content: . text( $0. content) ) }
150
+ let messageParameter = ChatCompletionParameters ( messages: messageParams, model: . custom( model) , maxTokens: maxTokens)
151
+
152
+ let stream = try await gemini. startStreamedChat ( parameters: messageParameter)
189
153
return try mapToLLMMessageStreamResponse ( stream: stream)
154
+
190
155
case . ollama( let model, let messages, let maxTokens) :
191
156
guard let ollamaOpenAIServiceCompatible else {
192
157
throw PolyAIError . missingLLMConfiguration ( " You Must provide a valid configuration for the \( parameter. llmService) API " )
0 commit comments