Local Agentic AI on Mac with MLX

import SwiftUI
import Foundation

// MARK: - OpenAI-compatible request/response models

struct ChatMessage: Codable {
    let role: String
    let content: String
}

struct ChatRequest: Codable {
    let model: String
    let messages: [ChatMessage]
    let stream: Bool
}

struct ChatChoice: Codable {
    struct Delta: Codable {
        let content: String?
    }
    let delta: Delta
}

struct ChatStreamChunk: Codable {
    let choices: [ChatChoice]
}

// MARK: - MLX Local LLM client

actor MLXLocalClient {
    /// Base URL of the running mlx_lm.server (default port 8080)
    private let baseURL = URL(string: "http://127.0.0.1:8080/v1/chat/completions")!
    private let modelName = "mlx-community/Qwen3-4B-4bit"

    func streamChat(prompt: String) -> AsyncThrowingStream<String, Error> {
        AsyncThrowingStream { continuation in
            Task {
                let body = ChatRequest(
                    model: modelName,
                    messages: [ChatMessage(role: "user", content: prompt)],
                    stream: true
                )
                var request = URLRequest(url: baseURL)
                request.httpMethod = "POST"
                request.setValue("application/json", forHTTPHeaderField: "Content-Type")
                request.httpBody = try? JSONEncoder().encode(body)

                let (bytes, _) = try await URLSession.shared.bytes(for: request)
                for try await line in bytes.lines {
                    // SSE lines are prefixed with "data: "
                    guard line.hasPrefix("data: "),
                          let jsonData = line.dropFirst(6).data(using: .utf8),
                          let chunk = try? JSONDecoder().decode(ChatStreamChunk.self, from: jsonData),
                          let token = chunk.choices.first?.delta.content
                    else { continue }
                    continuation.yield(token)
                }
                continuation.finish()
            }
        }
    }
}

// MARK: - SwiftUI View

struct LocalAgentChatView: View {
    @State private var prompt: String = "Summarize what MLX is in one sentence."
    @State private var reply: String = ""
    @State private var isLoading = false

    private let client = MLXLocalClient()

    var body: some View {
        VStack(alignment: .leading, spacing: 16) {
            Text("Local MLX Agent")
                .font(.title2.bold())

            TextField("Enter prompt", text: $prompt, axis: .vertical)
                .textFieldStyle(.roundedBorder)
                .lineLimit(3...6)

            Button(action: sendPrompt) {
                Label(isLoading ? "Thinking…" : "Ask Local Model", systemImage: "brain")
            }
            .buttonStyle(.borderedProminent)
            .disabled(isLoading || prompt.isEmpty)

            if !reply.isEmpty {
                ScrollView {
                    Text(reply)
                        .font(.body)
                        .frame(maxWidth: .infinity, alignment: .leading)
                        .padding()
                        .background(.secondarySystemBackground, in: RoundedRectangle(cornerRadius: 12))
                }
            }
            Spacer()
        }
        .padding()
    }

    private func sendPrompt() {
        reply = ""
        isLoading = true
        Task {
            defer { isLoading = false }
            do {
                for try await token in await client.streamChat(prompt: prompt) {
                    reply += token
                }
            } catch {
                reply = "Error: \(error.localizedDescription)\n\nMake sure mlx_lm.server is running on port 8080."
            }
        }
    }
}

#Preview {
    LocalAgentChatView()
}