rune/oai-swift
1
0
Fork 0
Code Issues Pull Requests Releases 8 Activity

Fix EmbeddingService

Browse Source
This commit is contained in:
Rune Olsen
2026-02-27 14:50:47 +01:00
parent e9d0ad3c66
commit 98d9ee2b51
3 changed files with 56 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
oAI/Services/EmbeddingService.swift
View File

@@ -75,6 +75,16 @@ final class EmbeddingService {
private let settings = SettingsService.shared
/// Dedicated session for embedding requests — keeps embedding traffic isolated
/// from the chat API sessions and self-limits concurrent connections.
private let session: URLSession = {
let config = URLSessionConfiguration.default
config.httpMaximumConnectionsPerHost = 2 // max 2 concurrent embedding requests
config.timeoutIntervalForRequest = 60
config.timeoutIntervalForResource = 120
return URLSession(configuration: config)
}()
private init() {}
// MARK: - Provider Detection
@@ -164,7 +174,7 @@ final class EmbeddingService {
]
request.httpBody = try JSONSerialization.data(withJSONObject: body)
let (data, response) = try await URLSession.shared.data(for: request)
let (data, response) = try await session.data(for: request)
guard let httpResponse = response as? HTTPURLResponse else {
throw EmbeddingError.invalidResponse
@@ -205,7 +215,7 @@ final class EmbeddingService {
]
request.httpBody = try JSONSerialization.data(withJSONObject: body)
let (data, response) = try await URLSession.shared.data(for: request)
let (data, response) = try await session.data(for: request)
guard let httpResponse = response as? HTTPURLResponse else {
throw EmbeddingError.invalidResponse
@@ -247,7 +257,7 @@ final class EmbeddingService {
]
request.httpBody = try JSONSerialization.data(withJSONObject: body)
let (data, response) = try await URLSession.shared.data(for: request)
let (data, response) = try await session.data(for: request)
guard let httpResponse = response as? HTTPURLResponse else {
throw EmbeddingError.invalidResponse

77
oAI/ViewModels/ChatViewModel.swift
View File

@@ -1793,16 +1793,15 @@ Don't narrate future actions ("Let me...") - just use the tools.
await checkAndSummarizeOldMessages(conversationId: conversation.id)
}
// Generate embeddings for messages that don't have them yet
// Generate embeddings for messages that don't have them yet.
// Run sequentially at background priority so this never blocks the chat.
if settings.embeddingsEnabled {
Task {
Task(priority: .background) {
guard let provider = EmbeddingService.shared.getSelectedProvider() else { return }
for message in chatMessages {
// Skip if already embedded
if let _ = try? EmbeddingService.shared.getMessageEmbedding(messageId: message.id) {
continue
}
// Generate embedding (this will now succeed since message is in DB)
generateEmbeddingForMessage(message)
await embedMessage(message, provider: provider)
// Yield briefly between requests to avoid bursting the API
try? await Task.sleep(for: .milliseconds(150))
}
}
}
@@ -1927,38 +1926,44 @@ Don't narrate future actions ("Let me...") - just use the tools.
// MARK: - Embedding Generation
/// Generate embedding for a message in the background
/// Generate embedding for a single message (awaitable, no Task spawned).
/// Call this directly when you want to sequence or rate-limit requests.
private func embedMessage(_ message: Message, provider: EmbeddingProvider) async {
guard message.content.count > 20 else { return }
// Skip if already embedded
if let _ = try? EmbeddingService.shared.getMessageEmbedding(messageId: message.id) { return }
do {
let embedding = try await EmbeddingService.shared.generateEmbedding(
text: message.content,
provider: provider
)
try EmbeddingService.shared.saveMessageEmbedding(
messageId: message.id,
embedding: embedding,
model: provider.defaultModel
)
Log.api.info("Generated embedding for message \(message.id) using \(provider.displayName)")
} catch {
let errorString = String(describing: error)
if errorString.contains("FOREIGN KEY constraint failed") {
Log.api.debug("Message \(message.id) not in database yet - will embed later during save or batch operation")
} else {
Log.api.error("Failed to generate embedding for message \(message.id): \(error)")
}
}
}
/// Fire-and-forget wrapper — runs at background priority so it never blocks the chat.
func generateEmbeddingForMessage(_ message: Message) {
guard settings.embeddingsEnabled else { return }
guard message.content.count > 20 else { return } // Skip very short messages
guard message.content.count > 20 else { return }
Task {
do {
// Use user's selected provider, or fall back to best available
guard let provider = EmbeddingService.shared.getSelectedProvider() else {
Log.api.warning("No embedding providers available - skipping embedding generation")
return
}
let embedding = try await EmbeddingService.shared.generateEmbedding(
text: message.content,
provider: provider
)
try EmbeddingService.shared.saveMessageEmbedding(
messageId: message.id,
embedding: embedding,
model: provider.defaultModel
)
Log.api.info("Generated embedding for message \(message.id) using \(provider.displayName)")
} catch {
// Check if it's a foreign key constraint error (message not saved to DB yet)
let errorString = String(describing: error)
if errorString.contains("FOREIGN KEY constraint failed") {
Log.api.debug("Message \(message.id) not in database yet - will embed later during save or batch operation")
} else {
Log.api.error("Failed to generate embedding for message \(message.id): \(error)")
}
Task(priority: .background) {
guard let provider = EmbeddingService.shared.getSelectedProvider() else {
Log.api.warning("No embedding providers available - skipping embedding generation")
return
}
await embedMessage(message, provider: provider)
}
}