Initial commit
This commit is contained in:
947
oAI/ViewModels/ChatViewModel.swift
Normal file
947
oAI/ViewModels/ChatViewModel.swift
Normal file
@@ -0,0 +1,947 @@
|
||||
//
|
||||
// ChatViewModel.swift
|
||||
// oAI
|
||||
//
|
||||
// Main chat view model
|
||||
//
|
||||
|
||||
import Foundation
|
||||
import os
|
||||
import SwiftUI
|
||||
|
||||
@Observable
|
||||
@MainActor
|
||||
class ChatViewModel {
|
||||
// MARK: - Observable State
|
||||
|
||||
var messages: [Message] = []
|
||||
var inputText: String = ""
|
||||
var isGenerating: Bool = false
|
||||
var sessionStats = SessionStats()
|
||||
var selectedModel: ModelInfo?
|
||||
var currentProvider: Settings.Provider = .openrouter
|
||||
var onlineMode: Bool = false
|
||||
var memoryEnabled: Bool = true
|
||||
var mcpEnabled: Bool = false
|
||||
var mcpStatus: String? = nil
|
||||
var availableModels: [ModelInfo] = []
|
||||
var isLoadingModels: Bool = false
|
||||
var showConversations: Bool = false
|
||||
var showModelSelector: Bool = false
|
||||
var showSettings: Bool = false
|
||||
var showStats: Bool = false
|
||||
var showHelp: Bool = false
|
||||
var showCredits: Bool = false
|
||||
var modelInfoTarget: ModelInfo? = nil
|
||||
|
||||
// MARK: - Private State
|
||||
|
||||
private var commandHistory: [String] = []
|
||||
private var historyIndex: Int = -1
|
||||
private var streamingTask: Task<Void, Never>?
|
||||
private let settings = SettingsService.shared
|
||||
private let providerRegistry = ProviderRegistry.shared
|
||||
|
||||
// MARK: - Initialization
|
||||
|
||||
init() {
|
||||
// Load settings
|
||||
self.currentProvider = settings.defaultProvider
|
||||
self.onlineMode = settings.onlineMode
|
||||
self.memoryEnabled = settings.memoryEnabled
|
||||
self.mcpEnabled = settings.mcpEnabled
|
||||
}
|
||||
|
||||
// MARK: - Public Methods
|
||||
|
||||
/// Switch to a different provider (from header dropdown)
|
||||
func changeProvider(_ newProvider: Settings.Provider) {
|
||||
guard newProvider != currentProvider else { return }
|
||||
Log.ui.info("Switching provider to \(newProvider.rawValue)")
|
||||
settings.defaultProvider = newProvider
|
||||
currentProvider = newProvider
|
||||
selectedModel = nil
|
||||
availableModels = []
|
||||
Task { await loadAvailableModels() }
|
||||
}
|
||||
|
||||
/// Start a new conversation
|
||||
func newConversation() {
|
||||
messages = []
|
||||
sessionStats = SessionStats()
|
||||
inputText = ""
|
||||
}
|
||||
|
||||
/// Re-sync local state from SettingsService (called when Settings sheet dismisses)
|
||||
func syncFromSettings() {
|
||||
let newProvider = settings.defaultProvider
|
||||
let providerChanged = currentProvider != newProvider
|
||||
currentProvider = newProvider
|
||||
onlineMode = settings.onlineMode
|
||||
memoryEnabled = settings.memoryEnabled
|
||||
mcpEnabled = settings.mcpEnabled
|
||||
mcpStatus = mcpEnabled ? "MCP" : nil
|
||||
|
||||
if providerChanged {
|
||||
selectedModel = nil
|
||||
availableModels = []
|
||||
Task { await loadAvailableModels() }
|
||||
}
|
||||
}
|
||||
|
||||
func loadAvailableModels() async {
|
||||
isLoadingModels = true
|
||||
|
||||
do {
|
||||
guard let provider = providerRegistry.getCurrentProvider() else {
|
||||
Log.ui.warning("No API key configured for current provider")
|
||||
isLoadingModels = false
|
||||
showSystemMessage("⚠️ No API key configured. Add your API key in Settings to load models.")
|
||||
return
|
||||
}
|
||||
|
||||
let models = try await provider.listModels()
|
||||
availableModels = models
|
||||
if selectedModel == nil, let firstModel = models.first {
|
||||
selectedModel = firstModel
|
||||
}
|
||||
isLoadingModels = false
|
||||
|
||||
} catch {
|
||||
Log.api.error("Failed to load models: \(error.localizedDescription)")
|
||||
isLoadingModels = false
|
||||
showSystemMessage("⚠️ Could not load models: \(error.localizedDescription)")
|
||||
}
|
||||
}
|
||||
|
||||
func sendMessage() {
|
||||
guard !inputText.trimmingCharacters(in: .whitespaces).isEmpty else { return }
|
||||
|
||||
let trimmedInput = inputText.trimmingCharacters(in: .whitespaces)
|
||||
|
||||
// Check if it's a slash command
|
||||
if trimmedInput.hasPrefix("/") {
|
||||
handleCommand(trimmedInput)
|
||||
inputText = ""
|
||||
return
|
||||
}
|
||||
|
||||
// Parse file attachments
|
||||
let (cleanText, filePaths) = trimmedInput.parseFileAttachments()
|
||||
|
||||
// Read file attachments from disk
|
||||
let attachments: [FileAttachment]? = filePaths.isEmpty ? nil : readFileAttachments(filePaths)
|
||||
|
||||
// Create user message
|
||||
let userMessage = Message(
|
||||
role: .user,
|
||||
content: cleanText,
|
||||
tokens: cleanText.estimateTokens(),
|
||||
cost: nil,
|
||||
timestamp: Date(),
|
||||
attachments: attachments
|
||||
)
|
||||
|
||||
messages.append(userMessage)
|
||||
sessionStats.addMessage(inputTokens: userMessage.tokens, outputTokens: nil, cost: nil)
|
||||
|
||||
// Clear input
|
||||
inputText = ""
|
||||
|
||||
// Add to command history
|
||||
commandHistory.append(trimmedInput)
|
||||
historyIndex = commandHistory.count
|
||||
|
||||
// Generate real AI response
|
||||
generateAIResponse(to: cleanText, attachments: userMessage.attachments)
|
||||
}
|
||||
|
||||
func cancelGeneration() {
|
||||
streamingTask?.cancel()
|
||||
streamingTask = nil
|
||||
isGenerating = false
|
||||
}
|
||||
|
||||
func clearChat() {
|
||||
messages.removeAll()
|
||||
sessionStats.reset()
|
||||
showSystemMessage("Chat cleared")
|
||||
}
|
||||
|
||||
func loadConversation(_ conversation: Conversation) {
|
||||
do {
|
||||
guard let (_, loadedMessages) = try DatabaseService.shared.loadConversation(id: conversation.id) else {
|
||||
showSystemMessage("Could not load conversation '\(conversation.name)'")
|
||||
return
|
||||
}
|
||||
|
||||
messages.removeAll()
|
||||
sessionStats.reset()
|
||||
messages = loadedMessages
|
||||
|
||||
// Rebuild session stats from loaded messages
|
||||
for msg in loadedMessages {
|
||||
sessionStats.addMessage(
|
||||
inputTokens: msg.role == .user ? msg.tokens : nil,
|
||||
outputTokens: msg.role == .assistant ? msg.tokens : nil,
|
||||
cost: msg.cost
|
||||
)
|
||||
}
|
||||
|
||||
showSystemMessage("Loaded conversation '\(conversation.name)'")
|
||||
} catch {
|
||||
showSystemMessage("Failed to load: \(error.localizedDescription)")
|
||||
}
|
||||
}
|
||||
|
||||
func retryLastMessage() {
|
||||
guard let lastUserMessage = messages.last(where: { $0.role == .user }) else {
|
||||
showSystemMessage("No previous message to retry")
|
||||
return
|
||||
}
|
||||
|
||||
// Remove last assistant response if exists
|
||||
if let lastMessage = messages.last, lastMessage.role == .assistant {
|
||||
messages.removeLast()
|
||||
}
|
||||
|
||||
generateAIResponse(to: lastUserMessage.content, attachments: lastUserMessage.attachments)
|
||||
}
|
||||
|
||||
// MARK: - Command Handling
|
||||
|
||||
private func handleCommand(_ command: String) {
|
||||
guard let (cmd, args) = command.parseCommand() else {
|
||||
showSystemMessage("Invalid command")
|
||||
return
|
||||
}
|
||||
|
||||
switch cmd.lowercased() {
|
||||
case "/help":
|
||||
showHelp = true
|
||||
|
||||
case "/model":
|
||||
showModelSelector = true
|
||||
|
||||
case "/clear":
|
||||
clearChat()
|
||||
|
||||
case "/retry":
|
||||
retryLastMessage()
|
||||
|
||||
case "/memory":
|
||||
if let arg = args.first?.lowercased() {
|
||||
memoryEnabled = arg == "on"
|
||||
showSystemMessage("Memory \(memoryEnabled ? "enabled" : "disabled")")
|
||||
} else {
|
||||
showSystemMessage("Usage: /memory on|off")
|
||||
}
|
||||
|
||||
case "/online":
|
||||
if let arg = args.first?.lowercased() {
|
||||
onlineMode = arg == "on"
|
||||
showSystemMessage("Online mode \(onlineMode ? "enabled" : "disabled")")
|
||||
} else {
|
||||
showSystemMessage("Usage: /online on|off")
|
||||
}
|
||||
|
||||
case "/stats":
|
||||
showStats = true
|
||||
|
||||
case "/config", "/settings":
|
||||
showSettings = true
|
||||
|
||||
case "/provider":
|
||||
if let providerName = args.first?.lowercased() {
|
||||
if let provider = Settings.Provider.allCases.first(where: { $0.rawValue == providerName }) {
|
||||
currentProvider = provider
|
||||
showSystemMessage("Switched to \(provider.displayName) provider")
|
||||
} else {
|
||||
showSystemMessage("Unknown provider: \(providerName)")
|
||||
}
|
||||
} else {
|
||||
showSystemMessage("Current provider: \(currentProvider.displayName)")
|
||||
}
|
||||
|
||||
case "/save":
|
||||
if let name = args.first {
|
||||
let chatMessages = messages.filter { $0.role != .system }
|
||||
guard !chatMessages.isEmpty else {
|
||||
showSystemMessage("Nothing to save — no messages in this conversation")
|
||||
return
|
||||
}
|
||||
do {
|
||||
let _ = try DatabaseService.shared.saveConversation(name: name, messages: chatMessages)
|
||||
showSystemMessage("Conversation saved as '\(name)'")
|
||||
} catch {
|
||||
showSystemMessage("Failed to save: \(error.localizedDescription)")
|
||||
}
|
||||
} else {
|
||||
showSystemMessage("Usage: /save <name>")
|
||||
}
|
||||
|
||||
case "/load", "/list":
|
||||
showConversations = true
|
||||
|
||||
case "/delete":
|
||||
if let name = args.first {
|
||||
do {
|
||||
let deleted = try DatabaseService.shared.deleteConversation(name: name)
|
||||
if deleted {
|
||||
showSystemMessage("Deleted conversation '\(name)'")
|
||||
} else {
|
||||
showSystemMessage("No conversation found with name '\(name)'")
|
||||
}
|
||||
} catch {
|
||||
showSystemMessage("Failed to delete: \(error.localizedDescription)")
|
||||
}
|
||||
} else {
|
||||
showSystemMessage("Usage: /delete <name>")
|
||||
}
|
||||
|
||||
case "/export":
|
||||
if args.count >= 1 {
|
||||
let format = args[0].lowercased()
|
||||
let filename = args.count >= 2 ? args[1] : "conversation.\(format)"
|
||||
exportConversation(format: format, filename: filename)
|
||||
} else {
|
||||
showSystemMessage("Usage: /export md|json <filename>")
|
||||
}
|
||||
|
||||
case "/info":
|
||||
if let modelId = args.first {
|
||||
if let model = availableModels.first(where: { $0.id == modelId || $0.name.lowercased() == modelId.lowercased() }) {
|
||||
showModelInfo(model)
|
||||
} else {
|
||||
showSystemMessage("Model not found: \(modelId)")
|
||||
}
|
||||
} else if let model = selectedModel {
|
||||
showModelInfo(model)
|
||||
} else {
|
||||
showSystemMessage("No model selected")
|
||||
}
|
||||
|
||||
case "/credits":
|
||||
showCredits = true
|
||||
|
||||
case "/mcp":
|
||||
handleMCPCommand(args: args)
|
||||
|
||||
default:
|
||||
showSystemMessage("Unknown command: \(cmd)\nType /help for available commands")
|
||||
}
|
||||
}
|
||||
|
||||
// MARK: - AI Response Generation
|
||||
|
||||
private func generateAIResponse(to prompt: String, attachments: [FileAttachment]?) {
|
||||
// Get provider
|
||||
guard let provider = providerRegistry.getCurrentProvider() else {
|
||||
Log.ui.warning("Cannot generate: no API key configured")
|
||||
showSystemMessage("❌ No API key configured. Please add your API key in Settings.")
|
||||
return
|
||||
}
|
||||
|
||||
guard let modelId = selectedModel?.id else {
|
||||
Log.ui.warning("Cannot generate: no model selected")
|
||||
showSystemMessage("❌ No model selected. Please select a model first.")
|
||||
return
|
||||
}
|
||||
|
||||
Log.ui.info("Sending message: model=\(modelId), messages=\(self.messages.count)")
|
||||
|
||||
// Dispatch to tool-aware path when MCP is enabled with folders
|
||||
// Skip for image generation models — they don't support tool calling
|
||||
let mcp = MCPService.shared
|
||||
let mcpActive = mcpEnabled || settings.mcpEnabled
|
||||
let modelSupportTools = selectedModel?.capabilities.tools ?? false
|
||||
if mcpActive && !mcp.allowedFolders.isEmpty && modelSupportTools {
|
||||
generateAIResponseWithTools(provider: provider, modelId: modelId)
|
||||
return
|
||||
}
|
||||
|
||||
isGenerating = true
|
||||
|
||||
// Cancel any existing task
|
||||
streamingTask?.cancel()
|
||||
|
||||
// Start streaming
|
||||
streamingTask = Task {
|
||||
do {
|
||||
// Create empty assistant message for streaming
|
||||
let assistantMessage = Message(
|
||||
role: .assistant,
|
||||
content: "",
|
||||
tokens: nil,
|
||||
cost: nil,
|
||||
timestamp: Date(),
|
||||
attachments: nil,
|
||||
isStreaming: true
|
||||
)
|
||||
|
||||
// Already on MainActor
|
||||
messages.append(assistantMessage)
|
||||
|
||||
// Build chat request AFTER adding the assistant message
|
||||
// Only include messages up to (but not including) the streaming assistant message
|
||||
var messagesToSend = Array(messages.dropLast()) // Remove the empty assistant message
|
||||
|
||||
// Web search via our WebSearchService (skip Anthropic — uses native search tool)
|
||||
// Append results to last user message content (matching Python oAI approach)
|
||||
if onlineMode && currentProvider != .anthropic && currentProvider != .openrouter && !messagesToSend.isEmpty {
|
||||
if let lastUserIdx = messagesToSend.lastIndex(where: { $0.role == .user }) {
|
||||
Log.search.info("Running web search for \(currentProvider.displayName)")
|
||||
let results = await WebSearchService.shared.search(query: messagesToSend[lastUserIdx].content)
|
||||
if !results.isEmpty {
|
||||
let searchContext = "\n\n\(WebSearchService.shared.formatResults(results))\n\nPlease use the above web search results to help answer the user's question."
|
||||
messagesToSend[lastUserIdx].content += searchContext
|
||||
Log.search.info("Injected \(results.count) search results into user message")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let isImageGen = selectedModel?.capabilities.imageGeneration ?? false
|
||||
if isImageGen {
|
||||
Log.ui.info("Image generation mode for model \(modelId)")
|
||||
}
|
||||
let chatRequest = ChatRequest(
|
||||
messages: Array(memoryEnabled ? messagesToSend : [messagesToSend.last!]),
|
||||
model: modelId,
|
||||
stream: settings.streamEnabled,
|
||||
maxTokens: settings.maxTokens > 0 ? settings.maxTokens : nil,
|
||||
temperature: settings.temperature > 0 ? settings.temperature : nil,
|
||||
topP: nil,
|
||||
systemPrompt: nil,
|
||||
tools: nil,
|
||||
onlineMode: onlineMode,
|
||||
imageGeneration: isImageGen
|
||||
)
|
||||
|
||||
let messageId = assistantMessage.id
|
||||
|
||||
if isImageGen {
|
||||
// Image generation: use non-streaming request
|
||||
// Image models don't reliably support streaming
|
||||
if let index = messages.firstIndex(where: { $0.id == messageId }) {
|
||||
messages[index].content = "Generating image..."
|
||||
}
|
||||
|
||||
let nonStreamRequest = ChatRequest(
|
||||
messages: chatRequest.messages,
|
||||
model: chatRequest.model,
|
||||
stream: false,
|
||||
maxTokens: chatRequest.maxTokens,
|
||||
temperature: chatRequest.temperature,
|
||||
imageGeneration: true
|
||||
)
|
||||
let response = try await provider.chat(request: nonStreamRequest)
|
||||
|
||||
if let index = messages.firstIndex(where: { $0.id == messageId }) {
|
||||
messages[index].content = response.content
|
||||
messages[index].isStreaming = false
|
||||
messages[index].generatedImages = response.generatedImages
|
||||
|
||||
if let usage = response.usage {
|
||||
messages[index].tokens = usage.completionTokens
|
||||
if let model = selectedModel {
|
||||
let cost = (Double(usage.promptTokens) * model.pricing.prompt / 1_000_000) +
|
||||
(Double(usage.completionTokens) * model.pricing.completion / 1_000_000)
|
||||
messages[index].cost = cost
|
||||
sessionStats.addMessage(inputTokens: usage.promptTokens, outputTokens: usage.completionTokens, cost: cost)
|
||||
}
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// Regular text: stream response
|
||||
var fullContent = ""
|
||||
var totalTokens: ChatResponse.Usage? = nil
|
||||
|
||||
for try await chunk in provider.streamChat(request: chatRequest) {
|
||||
if Task.isCancelled { break }
|
||||
|
||||
if let content = chunk.deltaContent {
|
||||
fullContent += content
|
||||
if let index = messages.firstIndex(where: { $0.id == messageId }) {
|
||||
messages[index].content = fullContent
|
||||
}
|
||||
}
|
||||
|
||||
if let usage = chunk.usage {
|
||||
totalTokens = usage
|
||||
}
|
||||
}
|
||||
|
||||
if let index = messages.firstIndex(where: { $0.id == messageId }) {
|
||||
messages[index].content = fullContent
|
||||
messages[index].isStreaming = false
|
||||
|
||||
if let usage = totalTokens {
|
||||
messages[index].tokens = usage.completionTokens
|
||||
if let model = selectedModel {
|
||||
let cost = (Double(usage.promptTokens) * model.pricing.prompt / 1_000_000) +
|
||||
(Double(usage.completionTokens) * model.pricing.completion / 1_000_000)
|
||||
messages[index].cost = cost
|
||||
sessionStats.addMessage(inputTokens: usage.promptTokens, outputTokens: usage.completionTokens, cost: cost)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
isGenerating = false
|
||||
streamingTask = nil
|
||||
|
||||
} catch {
|
||||
// Remove the empty streaming message
|
||||
if let index = messages.lastIndex(where: { $0.role == .assistant && $0.content.isEmpty }) {
|
||||
messages.remove(at: index)
|
||||
}
|
||||
|
||||
Log.api.error("Generation failed: \(error.localizedDescription)")
|
||||
showSystemMessage("❌ \(friendlyErrorMessage(from: error))")
|
||||
isGenerating = false
|
||||
streamingTask = nil
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// MARK: - File Attachment Reading
|
||||
|
||||
private let maxFileSize: Int = 10 * 1024 * 1024 // 10 MB
|
||||
private let maxTextSize: Int = 50 * 1024 // 50 KB before truncation
|
||||
|
||||
private func readFileAttachments(_ paths: [String]) -> [FileAttachment] {
|
||||
var attachments: [FileAttachment] = []
|
||||
let fm = FileManager.default
|
||||
|
||||
for rawPath in paths {
|
||||
// Expand ~ and resolve path
|
||||
let expanded = (rawPath as NSString).expandingTildeInPath
|
||||
let resolvedPath = expanded.hasPrefix("/") ? expanded : (fm.currentDirectoryPath as NSString).appendingPathComponent(expanded)
|
||||
|
||||
// Check file exists
|
||||
guard fm.fileExists(atPath: resolvedPath) else {
|
||||
showSystemMessage("⚠️ File not found: \(rawPath)")
|
||||
continue
|
||||
}
|
||||
|
||||
// Check file size
|
||||
guard let attrs = try? fm.attributesOfItem(atPath: resolvedPath),
|
||||
let fileSize = attrs[.size] as? Int else {
|
||||
showSystemMessage("⚠️ Cannot read file: \(rawPath)")
|
||||
continue
|
||||
}
|
||||
|
||||
if fileSize > maxFileSize {
|
||||
let sizeMB = String(format: "%.1f", Double(fileSize) / 1_000_000)
|
||||
showSystemMessage("⚠️ File too large (\(sizeMB) MB, max 10 MB): \(rawPath)")
|
||||
continue
|
||||
}
|
||||
|
||||
let type = FileAttachment.typeFromExtension(resolvedPath)
|
||||
|
||||
switch type {
|
||||
case .image, .pdf:
|
||||
// Read as raw data
|
||||
guard let data = fm.contents(atPath: resolvedPath) else {
|
||||
showSystemMessage("⚠️ Could not read file: \(rawPath)")
|
||||
continue
|
||||
}
|
||||
attachments.append(FileAttachment(path: rawPath, type: type, data: data))
|
||||
|
||||
case .text:
|
||||
// Read as string
|
||||
guard let content = try? String(contentsOfFile: resolvedPath, encoding: .utf8) else {
|
||||
showSystemMessage("⚠️ Could not read file as text: \(rawPath)")
|
||||
continue
|
||||
}
|
||||
|
||||
var finalContent = content
|
||||
// Truncate large text files
|
||||
if content.utf8.count > maxTextSize {
|
||||
let lines = content.components(separatedBy: "\n")
|
||||
if lines.count > 600 {
|
||||
let head = lines.prefix(500).joined(separator: "\n")
|
||||
let tail = lines.suffix(100).joined(separator: "\n")
|
||||
let omitted = lines.count - 600
|
||||
finalContent = head + "\n\n... [\(omitted) lines omitted] ...\n\n" + tail
|
||||
}
|
||||
}
|
||||
|
||||
attachments.append(FileAttachment(path: rawPath, type: .text, data: finalContent.data(using: .utf8)))
|
||||
}
|
||||
}
|
||||
|
||||
return attachments
|
||||
}
|
||||
|
||||
// MARK: - MCP Command Handling
|
||||
|
||||
private func handleMCPCommand(args: [String]) {
|
||||
let mcp = MCPService.shared
|
||||
guard let sub = args.first?.lowercased() else {
|
||||
showSystemMessage("Usage: /mcp on|off|status|add|remove|list")
|
||||
return
|
||||
}
|
||||
|
||||
switch sub {
|
||||
case "on":
|
||||
mcpEnabled = true
|
||||
settings.mcpEnabled = true
|
||||
mcpStatus = "MCP"
|
||||
showSystemMessage("MCP enabled (\(mcp.allowedFolders.count) folder\(mcp.allowedFolders.count == 1 ? "" : "s") registered)")
|
||||
|
||||
case "off":
|
||||
mcpEnabled = false
|
||||
settings.mcpEnabled = false
|
||||
mcpStatus = nil
|
||||
showSystemMessage("MCP disabled")
|
||||
|
||||
case "add":
|
||||
if args.count >= 2 {
|
||||
let path = args.dropFirst().joined(separator: " ")
|
||||
if let error = mcp.addFolder(path) {
|
||||
showSystemMessage("MCP: \(error)")
|
||||
} else {
|
||||
showSystemMessage("MCP: Added folder — \(mcp.allowedFolders.count) folder\(mcp.allowedFolders.count == 1 ? "" : "s") registered")
|
||||
}
|
||||
} else {
|
||||
showSystemMessage("Usage: /mcp add <path>")
|
||||
}
|
||||
|
||||
case "remove":
|
||||
if args.count >= 2 {
|
||||
let ref = args.dropFirst().joined(separator: " ")
|
||||
if let index = Int(ref) {
|
||||
if mcp.removeFolder(at: index) {
|
||||
showSystemMessage("MCP: Removed folder at index \(index)")
|
||||
} else {
|
||||
showSystemMessage("MCP: Invalid index \(index)")
|
||||
}
|
||||
} else {
|
||||
if mcp.removeFolder(path: ref) {
|
||||
showSystemMessage("MCP: Removed folder")
|
||||
} else {
|
||||
showSystemMessage("MCP: Folder not found: \(ref)")
|
||||
}
|
||||
}
|
||||
} else {
|
||||
showSystemMessage("Usage: /mcp remove <index|path>")
|
||||
}
|
||||
|
||||
case "list":
|
||||
if mcp.allowedFolders.isEmpty {
|
||||
showSystemMessage("MCP: No folders registered. Use /mcp add <path>")
|
||||
} else {
|
||||
let list = mcp.allowedFolders.enumerated().map { "\($0): \($1)" }.joined(separator: "\n")
|
||||
showSystemMessage("MCP folders:\n\(list)")
|
||||
}
|
||||
|
||||
case "write":
|
||||
guard args.count >= 2 else {
|
||||
showSystemMessage("Usage: /mcp write on|off")
|
||||
return
|
||||
}
|
||||
let toggle = args[1].lowercased()
|
||||
if toggle == "on" {
|
||||
settings.mcpCanWriteFiles = true
|
||||
settings.mcpCanDeleteFiles = true
|
||||
settings.mcpCanCreateDirectories = true
|
||||
settings.mcpCanMoveFiles = true
|
||||
showSystemMessage("MCP: All write permissions enabled (write, edit, delete, create dirs, move, copy)")
|
||||
} else if toggle == "off" {
|
||||
settings.mcpCanWriteFiles = false
|
||||
settings.mcpCanDeleteFiles = false
|
||||
settings.mcpCanCreateDirectories = false
|
||||
settings.mcpCanMoveFiles = false
|
||||
showSystemMessage("MCP: All write permissions disabled")
|
||||
} else {
|
||||
showSystemMessage("Usage: /mcp write on|off")
|
||||
}
|
||||
|
||||
case "status":
|
||||
let enabled = mcpEnabled ? "enabled" : "disabled"
|
||||
let folders = mcp.allowedFolders.count
|
||||
var perms: [String] = []
|
||||
if settings.mcpCanWriteFiles { perms.append("write") }
|
||||
if settings.mcpCanDeleteFiles { perms.append("delete") }
|
||||
if settings.mcpCanCreateDirectories { perms.append("mkdir") }
|
||||
if settings.mcpCanMoveFiles { perms.append("move/copy") }
|
||||
let permStr = perms.isEmpty ? "read-only" : "read + \(perms.joined(separator: ", "))"
|
||||
showSystemMessage("MCP: \(enabled), \(folders) folder\(folders == 1 ? "" : "s"), \(permStr), gitignore: \(settings.mcpRespectGitignore ? "on" : "off")")
|
||||
|
||||
default:
|
||||
showSystemMessage("MCP subcommands: on, off, status, add, remove, list, write")
|
||||
}
|
||||
}
|
||||
|
||||
// MARK: - AI Response with Tool Calls
|
||||
|
||||
private func generateAIResponseWithTools(provider: AIProvider, modelId: String) {
|
||||
let mcp = MCPService.shared
|
||||
isGenerating = true
|
||||
streamingTask?.cancel()
|
||||
|
||||
streamingTask = Task {
|
||||
do {
|
||||
let tools = mcp.getToolSchemas()
|
||||
|
||||
// Apply :online suffix for OpenRouter when online mode is active
|
||||
var effectiveModelId = modelId
|
||||
if onlineMode && currentProvider == .openrouter && !modelId.hasSuffix(":online") {
|
||||
effectiveModelId = modelId + ":online"
|
||||
}
|
||||
|
||||
// Build initial messages as raw dictionaries for the tool loop
|
||||
let folderList = mcp.allowedFolders.joined(separator: "\n - ")
|
||||
var capabilities = "You can read files, list directories, and search for files."
|
||||
var writeCapabilities: [String] = []
|
||||
if mcp.canWriteFiles { writeCapabilities.append("write and edit files") }
|
||||
if mcp.canDeleteFiles { writeCapabilities.append("delete files") }
|
||||
if mcp.canCreateDirectories { writeCapabilities.append("create directories") }
|
||||
if mcp.canMoveFiles { writeCapabilities.append("move and copy files") }
|
||||
if !writeCapabilities.isEmpty {
|
||||
capabilities += " You can also \(writeCapabilities.joined(separator: ", "))."
|
||||
}
|
||||
let systemContent = "You have access to the user's filesystem through tool calls. \(capabilities) The user has granted you access to these folders:\n - \(folderList)\n\nWhen the user asks about their files, use the tools proactively with the allowed paths. Always use absolute paths."
|
||||
|
||||
var messagesToSend: [Message] = memoryEnabled
|
||||
? messages.filter { $0.role != .system }
|
||||
: [messages.last(where: { $0.role == .user })].compactMap { $0 }
|
||||
|
||||
// Web search via our WebSearchService (skip Anthropic — uses native search tool)
|
||||
// Append results to last user message content (matching Python oAI approach)
|
||||
if onlineMode && currentProvider != .anthropic && currentProvider != .openrouter {
|
||||
if let lastUserIdx = messagesToSend.lastIndex(where: { $0.role == .user }) {
|
||||
Log.search.info("Running web search for tool-aware path (\(currentProvider.displayName))")
|
||||
let results = await WebSearchService.shared.search(query: messagesToSend[lastUserIdx].content)
|
||||
if !results.isEmpty {
|
||||
let searchContext = "\n\n\(WebSearchService.shared.formatResults(results))\n\nPlease use the above web search results to help answer the user's question."
|
||||
messagesToSend[lastUserIdx].content += searchContext
|
||||
Log.search.info("Injected \(results.count) search results into user message")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let systemPrompt: [String: Any] = [
|
||||
"role": "system",
|
||||
"content": systemContent
|
||||
]
|
||||
|
||||
var apiMessages: [[String: Any]] = [systemPrompt] + messagesToSend.map { msg in
|
||||
["role": msg.role.rawValue, "content": msg.content]
|
||||
}
|
||||
|
||||
let maxIterations = 5
|
||||
var finalContent = ""
|
||||
var totalUsage: ChatResponse.Usage?
|
||||
|
||||
for iteration in 0..<maxIterations {
|
||||
if Task.isCancelled { break }
|
||||
|
||||
let response = try await provider.chatWithToolMessages(
|
||||
model: effectiveModelId,
|
||||
messages: apiMessages,
|
||||
tools: tools,
|
||||
maxTokens: settings.maxTokens > 0 ? settings.maxTokens : nil,
|
||||
temperature: settings.temperature > 0 ? settings.temperature : nil
|
||||
)
|
||||
|
||||
if let usage = response.usage { totalUsage = usage }
|
||||
|
||||
// Check if the model wants to call tools
|
||||
guard let toolCalls = response.toolCalls, !toolCalls.isEmpty else {
|
||||
// No tool calls — this is the final text response
|
||||
finalContent = response.content
|
||||
break
|
||||
}
|
||||
|
||||
// Show what tools the model is calling
|
||||
let toolNames = toolCalls.map { $0.functionName }.joined(separator: ", ")
|
||||
showSystemMessage("🔧 Calling: \(toolNames)")
|
||||
|
||||
// Append assistant message with tool_calls to conversation
|
||||
var assistantMsg: [String: Any] = ["role": "assistant"]
|
||||
if !response.content.isEmpty {
|
||||
assistantMsg["content"] = response.content
|
||||
}
|
||||
let toolCallDicts: [[String: Any]] = toolCalls.map { tc in
|
||||
[
|
||||
"id": tc.id,
|
||||
"type": tc.type,
|
||||
"function": [
|
||||
"name": tc.functionName,
|
||||
"arguments": tc.arguments
|
||||
]
|
||||
]
|
||||
}
|
||||
assistantMsg["tool_calls"] = toolCallDicts
|
||||
apiMessages.append(assistantMsg)
|
||||
|
||||
// Execute each tool and append results
|
||||
for tc in toolCalls {
|
||||
if Task.isCancelled { break }
|
||||
|
||||
let result = mcp.executeTool(name: tc.functionName, arguments: tc.arguments)
|
||||
let resultJSON: String
|
||||
if let data = try? JSONSerialization.data(withJSONObject: result),
|
||||
let str = String(data: data, encoding: .utf8) {
|
||||
resultJSON = str
|
||||
} else {
|
||||
resultJSON = "{\"error\": \"Failed to serialize result\"}"
|
||||
}
|
||||
|
||||
apiMessages.append([
|
||||
"role": "tool",
|
||||
"tool_call_id": tc.id,
|
||||
"name": tc.functionName,
|
||||
"content": resultJSON
|
||||
])
|
||||
}
|
||||
|
||||
// If this was the last iteration, note it
|
||||
if iteration == maxIterations - 1 {
|
||||
finalContent = response.content.isEmpty
|
||||
? "[Tool loop reached maximum iterations]"
|
||||
: response.content
|
||||
}
|
||||
}
|
||||
|
||||
// Display the final response as an assistant message
|
||||
let assistantMessage = Message(
|
||||
role: .assistant,
|
||||
content: finalContent,
|
||||
tokens: totalUsage?.completionTokens,
|
||||
cost: nil,
|
||||
timestamp: Date()
|
||||
)
|
||||
messages.append(assistantMessage)
|
||||
|
||||
// Calculate cost
|
||||
if let usage = totalUsage, let model = selectedModel {
|
||||
let cost = (Double(usage.promptTokens) * model.pricing.prompt / 1_000_000) +
|
||||
(Double(usage.completionTokens) * model.pricing.completion / 1_000_000)
|
||||
if let index = messages.lastIndex(where: { $0.id == assistantMessage.id }) {
|
||||
messages[index].cost = cost
|
||||
}
|
||||
sessionStats.addMessage(
|
||||
inputTokens: usage.promptTokens,
|
||||
outputTokens: usage.completionTokens,
|
||||
cost: cost
|
||||
)
|
||||
}
|
||||
|
||||
isGenerating = false
|
||||
streamingTask = nil
|
||||
|
||||
} catch {
|
||||
Log.api.error("Tool generation failed: \(error.localizedDescription)")
|
||||
showSystemMessage("❌ \(friendlyErrorMessage(from: error))")
|
||||
isGenerating = false
|
||||
streamingTask = nil
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private func showSystemMessage(_ text: String) {
|
||||
let message = Message(
|
||||
role: .system,
|
||||
content: text,
|
||||
tokens: nil,
|
||||
cost: nil,
|
||||
timestamp: Date(),
|
||||
attachments: nil
|
||||
)
|
||||
messages.append(message)
|
||||
}
|
||||
|
||||
// MARK: - Error Helpers
|
||||
|
||||
private func friendlyErrorMessage(from error: Error) -> String {
|
||||
let desc = error.localizedDescription
|
||||
|
||||
// Network connectivity
|
||||
if let urlError = error as? URLError {
|
||||
switch urlError.code {
|
||||
case .notConnectedToInternet, .networkConnectionLost:
|
||||
return "Unable to reach the server. Check your internet connection."
|
||||
case .timedOut:
|
||||
return "Request timed out. Try a shorter message or different model."
|
||||
case .cannotFindHost, .cannotConnectToHost:
|
||||
return "Cannot connect to the server. Check your network or provider URL."
|
||||
default:
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
// HTTP status codes in error messages
|
||||
if desc.contains("401") || desc.contains("403") || desc.lowercased().contains("unauthorized") || desc.lowercased().contains("invalid.*key") {
|
||||
return "Invalid API key. Update it in Settings (\u{2318},)."
|
||||
}
|
||||
if desc.contains("429") || desc.lowercased().contains("rate limit") {
|
||||
return "Rate limited. Wait a moment and try again."
|
||||
}
|
||||
if desc.contains("404") || desc.lowercased().contains("model not found") || desc.lowercased().contains("not available") {
|
||||
return "Model not available. Select a different model (\u{2318}M)."
|
||||
}
|
||||
if desc.contains("500") || desc.contains("502") || desc.contains("503") {
|
||||
return "Server error. The provider may be experiencing issues. Try again shortly."
|
||||
}
|
||||
|
||||
// Timeout patterns
|
||||
if desc.lowercased().contains("timed out") || desc.lowercased().contains("timeout") {
|
||||
return "Request timed out. Try a shorter message or different model."
|
||||
}
|
||||
|
||||
// Fallback
|
||||
return desc
|
||||
}
|
||||
|
||||
// MARK: - Helpers
|
||||
|
||||
private func showModelInfo(_ model: ModelInfo) {
|
||||
modelInfoTarget = model
|
||||
}
|
||||
|
||||
private func exportConversation(format: String, filename: String) {
|
||||
let chatMessages = messages.filter { $0.role != .system }
|
||||
guard !chatMessages.isEmpty else {
|
||||
showSystemMessage("Nothing to export — no messages")
|
||||
return
|
||||
}
|
||||
|
||||
let content: String
|
||||
switch format {
|
||||
case "md", "markdown":
|
||||
content = chatMessages.map { msg in
|
||||
let header = msg.role == .user ? "**User**" : "**Assistant**"
|
||||
return "\(header)\n\n\(msg.content)"
|
||||
}.joined(separator: "\n\n---\n\n")
|
||||
case "json":
|
||||
let dicts = chatMessages.map { msg -> [String: String] in
|
||||
["role": msg.role.rawValue, "content": msg.content]
|
||||
}
|
||||
if let data = try? JSONSerialization.data(withJSONObject: dicts, options: .prettyPrinted),
|
||||
let json = String(data: data, encoding: .utf8) {
|
||||
content = json
|
||||
} else {
|
||||
showSystemMessage("Failed to encode JSON")
|
||||
return
|
||||
}
|
||||
default:
|
||||
showSystemMessage("Unsupported format: \(format). Use md or json.")
|
||||
return
|
||||
}
|
||||
|
||||
// Write to Downloads folder
|
||||
let downloads = FileManager.default.urls(for: .downloadsDirectory, in: .userDomainMask).first
|
||||
?? FileManager.default.temporaryDirectory
|
||||
let fileURL = downloads.appendingPathComponent(filename)
|
||||
|
||||
do {
|
||||
try content.write(to: fileURL, atomically: true, encoding: .utf8)
|
||||
showSystemMessage("Exported to \(fileURL.path)")
|
||||
} catch {
|
||||
showSystemMessage("Export failed: \(error.localizedDescription)")
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user