1083 lines
44 KiB
Swift
1083 lines
44 KiB
Swift
//
|
|
// 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 showHistory: Bool = false
|
|
var modelInfoTarget: ModelInfo? = nil
|
|
var commandHistory: [String] = []
|
|
var historyIndex: Int = 0
|
|
|
|
// MARK: - Private State
|
|
|
|
private var streamingTask: Task<Void, Never>?
|
|
private let settings = SettingsService.shared
|
|
private let providerRegistry = ProviderRegistry.shared
|
|
|
|
// Default system prompt
|
|
private let defaultSystemPrompt = """
|
|
You are a helpful AI assistant. Follow these guidelines:
|
|
|
|
1. **Accuracy First**: Never invent information or make assumptions. If you're unsure about something, say so clearly.
|
|
|
|
2. **Ask for Clarification**: When a request is ambiguous or lacks necessary details, ask clarifying questions before proceeding.
|
|
|
|
3. **Be Honest About Limitations**: If you cannot or should not help with a request, explain why clearly and politely. Don't attempt tasks outside your capabilities.
|
|
|
|
4. **Stay Grounded**: Base responses on facts and known information. Clearly distinguish between facts, informed opinions, and speculation.
|
|
|
|
5. **Be Direct**: Provide concise, relevant answers. Avoid unnecessary preambles or apologies.
|
|
|
|
6. **Use Markdown Formatting**: Always format your responses using standard Markdown syntax:
|
|
- Use **bold** for emphasis
|
|
- Use bullet points and numbered lists for organization
|
|
- Use code blocks with language tags for code (e.g., ```python)
|
|
- Use proper headings (##, ###) to structure long responses
|
|
- If the user requests output in other formats (HTML, JSON, XML, etc.), wrap those in appropriate code blocks
|
|
|
|
7. **Break Down Complex Tasks**: When working with tools (file access, search, etc.), break complex tasks into smaller, manageable steps. If a task requires many operations:
|
|
- Complete one logical step at a time
|
|
- Present findings or progress after each step
|
|
- Ask the user if you should continue to the next step
|
|
- Be mindful of tool usage limits (typically 25-30 tool calls per request)
|
|
|
|
8. **Incremental Progress**: For large codebases or complex analyses, work incrementally. Don't try to explore everything at once. Focus on what's immediately relevant to the user's question.
|
|
|
|
Remember: It's better to ask questions or admit uncertainty than to provide incorrect or fabricated information.
|
|
"""
|
|
|
|
/// Builds the complete system prompt by combining default + custom
|
|
private var effectiveSystemPrompt: String {
|
|
var prompt = defaultSystemPrompt
|
|
if let customPrompt = settings.systemPrompt, !customPrompt.isEmpty {
|
|
prompt += "\n\n---\n\nAdditional Instructions:\n" + customPrompt
|
|
}
|
|
return prompt
|
|
}
|
|
|
|
// MARK: - Initialization
|
|
|
|
init() {
|
|
// Load settings
|
|
self.currentProvider = settings.defaultProvider
|
|
self.onlineMode = settings.onlineMode
|
|
self.memoryEnabled = settings.memoryEnabled
|
|
self.mcpEnabled = settings.mcpEnabled
|
|
|
|
// Load command history from database
|
|
if let history = try? DatabaseService.shared.loadCommandHistory() {
|
|
self.commandHistory = history.map { $0.input }
|
|
self.historyIndex = self.commandHistory.count
|
|
}
|
|
|
|
// Load models on startup
|
|
Task {
|
|
await loadAvailableModels()
|
|
}
|
|
}
|
|
|
|
// 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
|
|
|
|
// Select model priority: saved default > current selection > first available
|
|
if let defaultModelId = settings.defaultModel,
|
|
let defaultModel = models.first(where: { $0.id == defaultModelId }) {
|
|
selectedModel = defaultModel
|
|
} else 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)
|
|
|
|
// Handle slash escape: "//" becomes "/"
|
|
var effectiveInput = trimmedInput
|
|
if effectiveInput.hasPrefix("//") {
|
|
effectiveInput = String(effectiveInput.dropFirst())
|
|
} else if effectiveInput.hasPrefix("/") {
|
|
// Check if it's a slash command
|
|
handleCommand(effectiveInput)
|
|
inputText = ""
|
|
return
|
|
}
|
|
|
|
// Parse file attachments
|
|
let (cleanText, filePaths) = effectiveInput.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)
|
|
|
|
// Add to command history (in-memory and database)
|
|
commandHistory.append(trimmedInput)
|
|
historyIndex = commandHistory.count
|
|
DatabaseService.shared.saveCommandHistory(input: trimmedInput)
|
|
|
|
// Clear input
|
|
inputText = ""
|
|
|
|
// 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 "/history":
|
|
showHistory = 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 {
|
|
let startTime = Date()
|
|
var messageId: UUID?
|
|
do {
|
|
// Create empty assistant message for streaming
|
|
let assistantMessage = Message(
|
|
role: .assistant,
|
|
content: "",
|
|
tokens: nil,
|
|
cost: nil,
|
|
timestamp: Date(),
|
|
attachments: nil,
|
|
isStreaming: true
|
|
)
|
|
messageId = assistantMessage.id
|
|
|
|
// 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: effectiveSystemPrompt,
|
|
tools: nil,
|
|
onlineMode: onlineMode,
|
|
imageGeneration: isImageGen
|
|
)
|
|
|
|
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)
|
|
let responseTime = Date().timeIntervalSince(startTime)
|
|
|
|
if let index = messages.firstIndex(where: { $0.id == messageId }) {
|
|
messages[index].content = response.content
|
|
messages[index].isStreaming = false
|
|
messages[index].generatedImages = response.generatedImages
|
|
messages[index].responseTime = responseTime
|
|
|
|
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
|
|
var wasCancelled = false
|
|
|
|
for try await chunk in provider.streamChat(request: chatRequest) {
|
|
if Task.isCancelled {
|
|
wasCancelled = true
|
|
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
|
|
}
|
|
}
|
|
|
|
// Check for cancellation one more time after loop exits
|
|
// (in case it was cancelled after the last chunk)
|
|
if Task.isCancelled {
|
|
wasCancelled = true
|
|
}
|
|
|
|
let responseTime = Date().timeIntervalSince(startTime)
|
|
|
|
if let index = messages.firstIndex(where: { $0.id == messageId }) {
|
|
messages[index].content = fullContent
|
|
messages[index].isStreaming = false
|
|
messages[index].responseTime = responseTime
|
|
messages[index].wasInterrupted = wasCancelled
|
|
|
|
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 {
|
|
let responseTime = Date().timeIntervalSince(startTime)
|
|
|
|
// Check if this was a cancellation (either by checking Task state or error type)
|
|
let isCancellation = Task.isCancelled || error is CancellationError
|
|
|
|
if isCancellation, let msgId = messageId {
|
|
// Mark the message as interrupted instead of removing it
|
|
if let index = messages.firstIndex(where: { $0.id == msgId }) {
|
|
messages[index].isStreaming = false
|
|
messages[index].wasInterrupted = true
|
|
messages[index].responseTime = responseTime
|
|
}
|
|
} else if let msgId = messageId {
|
|
// For real errors, remove the empty streaming message
|
|
if let index = messages.firstIndex(where: { $0.id == msgId && $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 {
|
|
let startTime = Date()
|
|
var wasCancelled = false
|
|
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: ", "))."
|
|
}
|
|
var 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."
|
|
|
|
// Append the complete system prompt (default + custom)
|
|
systemContent += "\n\n---\n\n" + effectiveSystemPrompt
|
|
|
|
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 {
|
|
wasCancelled = true
|
|
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 {
|
|
wasCancelled = true
|
|
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
|
|
}
|
|
}
|
|
|
|
// Check for cancellation one more time after loop exits
|
|
if Task.isCancelled {
|
|
wasCancelled = true
|
|
}
|
|
|
|
// Display the final response as an assistant message
|
|
let responseTime = Date().timeIntervalSince(startTime)
|
|
let assistantMessage = Message(
|
|
role: .assistant,
|
|
content: finalContent,
|
|
tokens: totalUsage?.completionTokens,
|
|
cost: nil,
|
|
timestamp: Date(),
|
|
attachments: nil,
|
|
responseTime: responseTime,
|
|
wasInterrupted: wasCancelled
|
|
)
|
|
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 {
|
|
let responseTime = Date().timeIntervalSince(startTime)
|
|
|
|
// Check if this was a cancellation
|
|
let isCancellation = Task.isCancelled || wasCancelled || error is CancellationError
|
|
|
|
if isCancellation {
|
|
// Create an interrupted message
|
|
let assistantMessage = Message(
|
|
role: .assistant,
|
|
content: "",
|
|
timestamp: Date(),
|
|
responseTime: responseTime,
|
|
wasInterrupted: true
|
|
)
|
|
messages.append(assistantMessage)
|
|
} else {
|
|
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)")
|
|
}
|
|
}
|
|
}
|