Files
oai-swift/oAI/Services/DatabaseService.swift

886 lines
31 KiB
Swift

//
// DatabaseService.swift
// oAI
//
// SQLite persistence layer for conversations using GRDB
//
import Foundation
import GRDB
import os
// MARK: - Database Record Types
struct ConversationRecord: Codable, FetchableRecord, PersistableRecord, Sendable {
static let databaseTableName = "conversations"
var id: String
var name: String
var createdAt: String
var updatedAt: String
var primaryModel: String?
}
struct MessageRecord: Codable, FetchableRecord, PersistableRecord, Sendable {
static let databaseTableName = "messages"
var id: String
var conversationId: String
var role: String
var content: String
var tokens: Int?
var cost: Double?
var timestamp: String
var sortOrder: Int
var modelId: String?
}
struct SettingRecord: Codable, FetchableRecord, PersistableRecord, Sendable {
static let databaseTableName = "settings"
var key: String
var value: String
}
struct HistoryRecord: Codable, FetchableRecord, PersistableRecord, Sendable {
static let databaseTableName = "command_history"
var id: String
var input: String
var timestamp: String
}
struct EmailLogRecord: Codable, FetchableRecord, PersistableRecord, Sendable {
static let databaseTableName = "email_logs"
var id: String
var timestamp: String
var sender: String
var subject: String
var emailContent: String
var aiResponse: String?
var status: String // "success" or "error"
var errorMessage: String?
var tokens: Int?
var cost: Double?
var responseTime: Double?
var modelId: String?
}
struct MessageMetadataRecord: Codable, FetchableRecord, PersistableRecord, Sendable {
static let databaseTableName = "message_metadata"
var message_id: String
var importance_score: Double
var user_starred: Int
var summary: String?
var chunk_index: Int
}
struct MessageEmbeddingRecord: Codable, FetchableRecord, PersistableRecord, Sendable {
static let databaseTableName = "message_embeddings"
var message_id: String
var embedding: Data
var embedding_model: String
var embedding_dimension: Int
var created_at: String
}
struct ConversationEmbeddingRecord: Codable, FetchableRecord, PersistableRecord, Sendable {
static let databaseTableName = "conversation_embeddings"
var conversation_id: String
var embedding: Data
var embedding_model: String
var embedding_dimension: Int
var created_at: String
}
struct ConversationSummaryRecord: Codable, FetchableRecord, PersistableRecord, Sendable {
static let databaseTableName = "conversation_summaries"
var id: String
var conversation_id: String
var start_message_index: Int
var end_message_index: Int
var summary: String
var token_count: Int?
var created_at: String
var summary_model: String?
}
// MARK: - DatabaseService
final class DatabaseService: Sendable {
nonisolated static let shared = DatabaseService()
private let dbQueue: DatabaseQueue
private let isoFormatter: ISO8601DateFormatter
// Command history limit - keep most recent 5000 entries
private static let maxHistoryEntries = 5000
nonisolated private init() {
isoFormatter = ISO8601DateFormatter()
isoFormatter.formatOptions = [.withInternetDateTime, .withFractionalSeconds]
let fileManager = FileManager.default
let appSupport = fileManager.urls(for: .applicationSupportDirectory, in: .userDomainMask).first!
let dbDirectory = appSupport.appendingPathComponent("oAI", isDirectory: true)
try! fileManager.createDirectory(at: dbDirectory, withIntermediateDirectories: true)
let dbPath = dbDirectory.appendingPathComponent("oai_conversations.db").path
Log.db.info("Opening database at \(dbPath)")
dbQueue = try! DatabaseQueue(path: dbPath)
try! migrator.migrate(dbQueue)
}
private var migrator: DatabaseMigrator {
var migrator = DatabaseMigrator()
migrator.registerMigration("v1") { db in
try db.create(table: "conversations") { t in
t.primaryKey("id", .text)
t.column("name", .text).notNull()
t.column("createdAt", .text).notNull()
t.column("updatedAt", .text).notNull()
}
try db.create(table: "messages") { t in
t.primaryKey("id", .text)
t.column("conversationId", .text).notNull()
.references("conversations", onDelete: .cascade)
t.column("role", .text).notNull()
t.column("content", .text).notNull()
t.column("tokens", .integer)
t.column("cost", .double)
t.column("timestamp", .text).notNull()
t.column("sortOrder", .integer).notNull()
}
try db.create(
index: "messages_on_conversationId",
on: "messages",
columns: ["conversationId"]
)
}
migrator.registerMigration("v2") { db in
try db.create(table: "settings") { t in
t.primaryKey("key", .text)
t.column("value", .text).notNull()
}
}
migrator.registerMigration("v3") { db in
try db.create(table: "command_history") { t in
t.primaryKey("id", .text)
t.column("input", .text).notNull()
t.column("timestamp", .text).notNull()
}
try db.create(
index: "command_history_on_timestamp",
on: "command_history",
columns: ["timestamp"]
)
}
migrator.registerMigration("v4") { db in
// Add modelId to messages table (nullable for existing messages)
try db.alter(table: "messages") { t in
t.add(column: "modelId", .text)
}
// Add primaryModel to conversations table (nullable)
try db.alter(table: "conversations") { t in
t.add(column: "primaryModel", .text)
}
}
migrator.registerMigration("v5") { db in
// Email handler logs
try db.create(table: "email_logs") { t in
t.primaryKey("id", .text)
t.column("timestamp", .text).notNull()
t.column("sender", .text).notNull()
t.column("subject", .text).notNull()
t.column("emailContent", .text).notNull()
t.column("aiResponse", .text)
t.column("status", .text).notNull() // "success" or "error"
t.column("errorMessage", .text)
t.column("tokens", .integer)
t.column("cost", .double)
t.column("responseTime", .double)
t.column("modelId", .text)
}
try db.create(
index: "email_logs_on_timestamp",
on: "email_logs",
columns: ["timestamp"]
)
}
migrator.registerMigration("v6") { db in
// Message metadata for smart context selection
try db.create(table: "message_metadata") { t in
t.primaryKey("message_id", .text)
.references("messages", onDelete: .cascade)
t.column("importance_score", .double).notNull().defaults(to: 0.0)
t.column("user_starred", .integer).notNull().defaults(to: 0)
t.column("summary", .text)
t.column("chunk_index", .integer).notNull().defaults(to: 0)
}
try db.create(
index: "idx_message_metadata_importance",
on: "message_metadata",
columns: ["importance_score"]
)
try db.create(
index: "idx_message_metadata_starred",
on: "message_metadata",
columns: ["user_starred"]
)
}
migrator.registerMigration("v7") { db in
// Message embeddings for semantic search
try db.create(table: "message_embeddings") { t in
t.primaryKey("message_id", .text)
.references("messages", onDelete: .cascade)
t.column("embedding", .blob).notNull()
t.column("embedding_model", .text).notNull()
t.column("embedding_dimension", .integer).notNull()
t.column("created_at", .text).notNull()
}
// Conversation embeddings (aggregate of all messages)
try db.create(table: "conversation_embeddings") { t in
t.primaryKey("conversation_id", .text)
.references("conversations", onDelete: .cascade)
t.column("embedding", .blob).notNull()
t.column("embedding_model", .text).notNull()
t.column("embedding_dimension", .integer).notNull()
t.column("created_at", .text).notNull()
}
}
migrator.registerMigration("v8") { db in
// Conversation summaries for progressive summarization
try db.create(table: "conversation_summaries") { t in
t.primaryKey("id", .text)
t.column("conversation_id", .text).notNull()
.references("conversations", onDelete: .cascade)
t.column("start_message_index", .integer).notNull()
t.column("end_message_index", .integer).notNull()
t.column("summary", .text).notNull()
t.column("token_count", .integer)
t.column("created_at", .text).notNull()
t.column("summary_model", .text)
}
try db.create(
index: "idx_conversation_summaries_conv",
on: "conversation_summaries",
columns: ["conversation_id"]
)
}
return migrator
}
// MARK: - Settings Operations
nonisolated func loadAllSettings() throws -> [String: String] {
try dbQueue.read { db in
let records = try SettingRecord.fetchAll(db)
return Dictionary(uniqueKeysWithValues: records.map { ($0.key, $0.value) })
}
}
nonisolated func setSetting(key: String, value: String) {
try? dbQueue.write { db in
let record = SettingRecord(key: key, value: value)
try record.save(db)
}
}
nonisolated func deleteSetting(key: String) {
try? dbQueue.write { db in
_ = try SettingRecord.deleteOne(db, key: key)
}
}
// MARK: - Encrypted Settings Operations
/// Store an encrypted setting (for sensitive data like API keys)
nonisolated func setEncryptedSetting(key: String, value: String) throws {
let encryptedValue = try EncryptionService.shared.encrypt(value)
try dbQueue.write { db in
let record = SettingRecord(key: "encrypted_\(key)", value: encryptedValue)
try record.save(db)
}
}
/// Retrieve and decrypt an encrypted setting
nonisolated func getEncryptedSetting(key: String) throws -> String? {
let encryptedValue = try dbQueue.read { db in
try SettingRecord.fetchOne(db, key: "encrypted_\(key)")?.value
}
guard let encryptedValue = encryptedValue else {
return nil
}
return try EncryptionService.shared.decrypt(encryptedValue)
}
/// Delete an encrypted setting
nonisolated func deleteEncryptedSetting(key: String) {
try? dbQueue.write { db in
_ = try SettingRecord.deleteOne(db, key: "encrypted_\(key)")
}
}
// MARK: - Conversation Operations
nonisolated func saveConversation(name: String, messages: [Message]) throws -> Conversation {
return try saveConversation(id: UUID(), name: name, messages: messages, primaryModel: nil)
}
nonisolated func saveConversation(id: UUID, name: String, messages: [Message], primaryModel: String?) throws -> Conversation {
Log.db.info("Saving conversation '\(name)' with \(messages.count) messages (primaryModel: \(primaryModel ?? "none"))")
let now = Date()
let nowString = isoFormatter.string(from: now)
let convRecord = ConversationRecord(
id: id.uuidString,
name: name,
createdAt: nowString,
updatedAt: nowString,
primaryModel: primaryModel
)
let messageRecords = messages.enumerated().compactMap { index, msg -> MessageRecord? in
guard msg.role != .system else { return nil }
return MessageRecord(
id: UUID().uuidString,
conversationId: id.uuidString,
role: msg.role.rawValue,
content: msg.content,
tokens: msg.tokens,
cost: msg.cost,
timestamp: isoFormatter.string(from: msg.timestamp),
sortOrder: index,
modelId: msg.modelId
)
}
try dbQueue.write { db in
try convRecord.insert(db)
for record in messageRecords {
try record.insert(db)
}
}
let savedMessages = messages.filter { $0.role != .system }
return Conversation(
id: id,
name: name,
messages: savedMessages,
createdAt: now,
updatedAt: now,
primaryModel: primaryModel
)
}
nonisolated func loadConversation(id: UUID) throws -> (Conversation, [Message])? {
try dbQueue.read { db in
guard let convRecord = try ConversationRecord.fetchOne(db, key: id.uuidString) else {
return nil
}
let messageRecords = try MessageRecord
.filter(Column("conversationId") == id.uuidString)
.order(Column("sortOrder"))
.fetchAll(db)
let messages = messageRecords.compactMap { record -> Message? in
guard let msgId = UUID(uuidString: record.id),
let role = MessageRole(rawValue: record.role),
let timestamp = self.isoFormatter.date(from: record.timestamp)
else { return nil }
return Message(
id: msgId,
role: role,
content: record.content,
tokens: record.tokens,
cost: record.cost,
timestamp: timestamp,
modelId: record.modelId
)
}
guard let convId = UUID(uuidString: convRecord.id),
let createdAt = self.isoFormatter.date(from: convRecord.createdAt),
let updatedAt = self.isoFormatter.date(from: convRecord.updatedAt)
else { return nil }
let conversation = Conversation(
id: convId,
name: convRecord.name,
messages: messages,
createdAt: createdAt,
updatedAt: updatedAt,
primaryModel: convRecord.primaryModel
)
return (conversation, messages)
}
}
nonisolated func listConversations() throws -> [Conversation] {
try dbQueue.read { db in
let records = try ConversationRecord
.order(Column("updatedAt").desc)
.fetchAll(db)
return records.compactMap { record -> Conversation? in
guard let id = UUID(uuidString: record.id),
let createdAt = self.isoFormatter.date(from: record.createdAt),
let updatedAt = self.isoFormatter.date(from: record.updatedAt)
else { return nil }
// Fetch message count without loading all messages
let messageCount = (try? MessageRecord
.filter(Column("conversationId") == record.id)
.fetchCount(db)) ?? 0
// Get last message date
let lastMsg = try? MessageRecord
.filter(Column("conversationId") == record.id)
.order(Column("sortOrder").desc)
.fetchOne(db)
let lastDate = lastMsg.flatMap { self.isoFormatter.date(from: $0.timestamp) } ?? updatedAt
// Create conversation with empty messages array but correct metadata
var conv = Conversation(
id: id,
name: record.name,
messages: Array(repeating: Message(role: .user, content: ""), count: messageCount),
createdAt: createdAt,
updatedAt: lastDate
)
// We store placeholder messages just for the count; lastMessageDate uses updatedAt
conv.updatedAt = lastDate
return conv
}
}
}
nonisolated func deleteConversation(id: UUID) throws -> Bool {
Log.db.info("Deleting conversation \(id.uuidString)")
return try dbQueue.write { db in
try MessageRecord.filter(Column("conversationId") == id.uuidString).deleteAll(db)
return try ConversationRecord.deleteOne(db, key: id.uuidString)
}
}
nonisolated func deleteConversation(name: String) throws -> Bool {
try dbQueue.write { db in
guard let record = try ConversationRecord
.filter(Column("name") == name)
.fetchOne(db)
else { return false }
try MessageRecord.filter(Column("conversationId") == record.id).deleteAll(db)
return try ConversationRecord.deleteOne(db, key: record.id)
}
}
nonisolated func updateConversation(id: UUID, name: String?, messages: [Message]?) throws -> Bool {
try dbQueue.write { db in
guard var convRecord = try ConversationRecord.fetchOne(db, key: id.uuidString) else {
return false
}
if let name = name {
convRecord.name = name
}
convRecord.updatedAt = self.isoFormatter.string(from: Date())
try convRecord.update(db)
if let messages = messages {
try MessageRecord.filter(Column("conversationId") == id.uuidString).deleteAll(db)
let messageRecords = messages.enumerated().compactMap { index, msg -> MessageRecord? in
guard msg.role != .system else { return nil }
return MessageRecord(
id: UUID().uuidString,
conversationId: id.uuidString,
role: msg.role.rawValue,
content: msg.content,
tokens: msg.tokens,
cost: msg.cost,
timestamp: self.isoFormatter.string(from: msg.timestamp),
sortOrder: index
)
}
for record in messageRecords {
try record.insert(db)
}
}
return true
}
}
// MARK: - Command History Operations
nonisolated func saveCommandHistory(input: String) {
let now = Date()
let record = HistoryRecord(
id: UUID().uuidString,
input: input,
timestamp: isoFormatter.string(from: now)
)
try? dbQueue.write { db in
try record.insert(db)
// Clean up old entries if we exceed the limit
let count = try HistoryRecord.fetchCount(db)
if count > Self.maxHistoryEntries {
// Delete oldest entries to get back to limit
let excess = count - Self.maxHistoryEntries
try db.execute(
sql: """
DELETE FROM command_history
WHERE id IN (
SELECT id FROM command_history
ORDER BY timestamp ASC
LIMIT ?
)
""",
arguments: [excess]
)
}
}
}
nonisolated func loadCommandHistory() throws -> [(input: String, timestamp: Date)] {
try dbQueue.read { db in
let records = try HistoryRecord
.order(Column("timestamp").desc)
.fetchAll(db)
return records.compactMap { record in
guard let date = isoFormatter.date(from: record.timestamp) else {
return nil
}
return (input: record.input, timestamp: date)
}
}
}
nonisolated func searchCommandHistory(query: String) throws -> [(input: String, timestamp: Date)] {
try dbQueue.read { db in
let records = try HistoryRecord
.filter(Column("input").like("%\(query)%"))
.order(Column("timestamp").desc)
.fetchAll(db)
return records.compactMap { record in
guard let date = isoFormatter.date(from: record.timestamp) else {
return nil
}
return (input: record.input, timestamp: date)
}
}
}
// MARK: - Email Log Operations
nonisolated func saveEmailLog(_ log: EmailLog) {
let record = EmailLogRecord(
id: log.id.uuidString,
timestamp: isoFormatter.string(from: log.timestamp),
sender: log.sender,
subject: log.subject,
emailContent: log.emailContent,
aiResponse: log.aiResponse,
status: log.status.rawValue,
errorMessage: log.errorMessage,
tokens: log.tokens,
cost: log.cost,
responseTime: log.responseTime,
modelId: log.modelId
)
try? dbQueue.write { db in
try record.insert(db)
}
}
nonisolated func loadEmailLogs(limit: Int = 100) throws -> [EmailLog] {
try dbQueue.read { db in
let records = try EmailLogRecord
.order(Column("timestamp").desc)
.limit(limit)
.fetchAll(db)
return records.compactMap { record in
guard let timestamp = isoFormatter.date(from: record.timestamp),
let status = EmailLogStatus(rawValue: record.status),
let id = UUID(uuidString: record.id) else {
return nil
}
return EmailLog(
id: id,
timestamp: timestamp,
sender: record.sender,
subject: record.subject,
emailContent: record.emailContent,
aiResponse: record.aiResponse,
status: status,
errorMessage: record.errorMessage,
tokens: record.tokens,
cost: record.cost,
responseTime: record.responseTime,
modelId: record.modelId
)
}
}
}
nonisolated func deleteEmailLog(id: UUID) {
try? dbQueue.write { db in
try db.execute(
sql: "DELETE FROM email_logs WHERE id = ?",
arguments: [id.uuidString]
)
}
}
nonisolated func clearEmailLogs() {
try? dbQueue.write { db in
try db.execute(sql: "DELETE FROM email_logs")
}
}
nonisolated func getEmailLogCount() throws -> Int {
try dbQueue.read { db in
try EmailLogRecord.fetchCount(db)
}
}
// MARK: - Message Metadata Operations
nonisolated func getMessageMetadata(messageId: UUID) throws -> MessageMetadataRecord? {
try dbQueue.read { db in
try MessageMetadataRecord.fetchOne(db, key: messageId.uuidString)
}
}
nonisolated func setMessageStarred(messageId: UUID, starred: Bool) throws {
try dbQueue.write { db in
// Check if metadata exists
if var existing = try MessageMetadataRecord.fetchOne(db, key: messageId.uuidString) {
existing.user_starred = starred ? 1 : 0
try existing.update(db)
} else {
// Create new metadata record
let record = MessageMetadataRecord(
message_id: messageId.uuidString,
importance_score: 0.0,
user_starred: starred ? 1 : 0,
summary: nil,
chunk_index: 0
)
try record.insert(db)
}
}
}
nonisolated func setMessageImportance(messageId: UUID, score: Double) throws {
try dbQueue.write { db in
if var existing = try MessageMetadataRecord.fetchOne(db, key: messageId.uuidString) {
existing.importance_score = score
try existing.update(db)
} else {
let record = MessageMetadataRecord(
message_id: messageId.uuidString,
importance_score: score,
user_starred: 0,
summary: nil,
chunk_index: 0
)
try record.insert(db)
}
}
}
nonisolated func getStarredMessages(conversationId: UUID) throws -> [String] {
try dbQueue.read { db in
let sql = """
SELECT mm.message_id
FROM message_metadata mm
JOIN messages m ON m.id = mm.message_id
WHERE m.conversationId = ? AND mm.user_starred = 1
ORDER BY m.sortOrder
"""
return try String.fetchAll(db, sql: sql, arguments: [conversationId.uuidString])
}
}
// MARK: - Embedding Operations
nonisolated func saveMessageEmbedding(messageId: UUID, embedding: Data, model: String, dimension: Int) throws {
let now = isoFormatter.string(from: Date())
let record = MessageEmbeddingRecord(
message_id: messageId.uuidString,
embedding: embedding,
embedding_model: model,
embedding_dimension: dimension,
created_at: now
)
try dbQueue.write { db in
try record.save(db)
}
}
nonisolated func getMessageEmbedding(messageId: UUID) throws -> Data? {
try dbQueue.read { db in
try MessageEmbeddingRecord.fetchOne(db, key: messageId.uuidString)?.embedding
}
}
nonisolated func saveConversationEmbedding(conversationId: UUID, embedding: Data, model: String, dimension: Int) throws {
let now = isoFormatter.string(from: Date())
let record = ConversationEmbeddingRecord(
conversation_id: conversationId.uuidString,
embedding: embedding,
embedding_model: model,
embedding_dimension: dimension,
created_at: now
)
try dbQueue.write { db in
try record.save(db)
}
}
nonisolated func getConversationEmbedding(conversationId: UUID) throws -> Data? {
try dbQueue.read { db in
try ConversationEmbeddingRecord.fetchOne(db, key: conversationId.uuidString)?.embedding
}
}
nonisolated func getAllConversationEmbeddings() throws -> [(UUID, Data)] {
try dbQueue.read { db in
let records = try ConversationEmbeddingRecord.fetchAll(db)
return records.compactMap { record in
guard let id = UUID(uuidString: record.conversation_id) else { return nil }
return (id, record.embedding)
}
}
}
/// Search conversations by semantic similarity
nonisolated func searchConversationsBySemantic(queryEmbedding: [Float], limit: Int = 10) throws -> [(Conversation, Float)] {
// Get all conversations
let allConversations = try listConversations()
// Get all conversation embeddings
let embeddingData = try getAllConversationEmbeddings()
let embeddingMap = Dictionary(uniqueKeysWithValues: embeddingData)
// Calculate similarity scores
var results: [(Conversation, Float)] = []
for conv in allConversations {
guard let embeddingData = embeddingMap[conv.id] else { continue }
// Deserialize embedding
let embedding = deserializeEmbedding(embeddingData)
// Calculate cosine similarity
let similarity = EmbeddingService.shared.cosineSimilarity(queryEmbedding, embedding)
results.append((conv, similarity))
}
// Sort by similarity (highest first) and take top N
results.sort { $0.1 > $1.1 }
return Array(results.prefix(limit))
}
private func deserializeEmbedding(_ data: Data) -> [Float] {
var embedding: [Float] = []
embedding.reserveCapacity(data.count / 4)
for offset in stride(from: 0, to: data.count, by: 4) {
let bytes = data.subdata(in: offset..<(offset + 4))
let bitPattern = bytes.withUnsafeBytes { $0.load(as: UInt32.self) }
let value = Float(bitPattern: UInt32(littleEndian: bitPattern))
embedding.append(value)
}
return embedding
}
// MARK: - Conversation Summary Operations
nonisolated func saveConversationSummary(
conversationId: UUID,
startIndex: Int,
endIndex: Int,
summary: String,
model: String?,
tokenCount: Int?
) throws {
let now = isoFormatter.string(from: Date())
let record = ConversationSummaryRecord(
id: UUID().uuidString,
conversation_id: conversationId.uuidString,
start_message_index: startIndex,
end_message_index: endIndex,
summary: summary,
token_count: tokenCount,
created_at: now,
summary_model: model
)
try dbQueue.write { db in
try record.insert(db)
}
}
nonisolated func getConversationSummaries(conversationId: UUID) throws -> [ConversationSummaryRecord] {
try dbQueue.read { db in
try ConversationSummaryRecord
.filter(Column("conversation_id") == conversationId.uuidString)
.order(Column("start_message_index"))
.fetchAll(db)
}
}
nonisolated func hasSummaryForRange(conversationId: UUID, startIndex: Int, endIndex: Int) throws -> Bool {
try dbQueue.read { db in
let count = try ConversationSummaryRecord
.filter(Column("conversation_id") == conversationId.uuidString)
.filter(Column("start_message_index") == startIndex)
.filter(Column("end_message_index") == endIndex)
.fetchCount(db)
return count > 0
}
}
}