Real Inference — MLXClient and the Local Model

//  Sources/gateway/MLX/MLXClient.swift

import Foundation
import Logging

enum MLXError: Error, LocalizedError {
    case invalidURL(String)
    case backendUnavailable(String)
    case inferenceError(Int, String)
    case decodingError(String)

    var errorDescription: String? {
        switch self {
        case .invalidURL(let url):
            return "Invalid MLX backend URL: \(url)"
        case .backendUnavailable(let reason):
            return "MLX backend unavailable: \(reason)"
        case .inferenceError(let status, let body):
            return "MLX backend returned HTTP \(status): \(body)"
        case .decodingError(let detail):
            return "Failed to decode MLX response: \(detail)"
        }
    }
}

actor MLXClient {
    private let baseURL: String
    private let session: URLSession
    private let logger: Logger

    init(baseURL: String, logger: Logger) {
        self.baseURL = baseURL.hasSuffix("/") ? String(baseURL.dropLast()) : baseURL
        self.logger = logger

        let config = URLSessionConfiguration.default
        config.timeoutIntervalForRequest = 120
        config.timeoutIntervalForResource = 300
        self.session = URLSession(configuration: config)
    }

    func complete(
        messages: [ChatMessage],
        model: String,
        maxTokens: Int = 1024,
        temperature: Double = 1.0
    ) async throws -> (text: String, inputTokens: Int, outputTokens: Int) {
        guard let url = URL(string: "\(baseURL)/v1/chat/completions") else {
            throw MLXError.invalidURL(baseURL)
        }

        let body = ChatCompletionRequest(
            model: model,
            messages: messages,
            maxTokens: maxTokens,
            temperature: temperature,
            topP: nil,
            stream: false,
            stop: nil,
            presencePenalty: nil,
            frequencyPenalty: nil,
            user: nil
        )

        var request = URLRequest(url: url)
        request.httpMethod = "POST"
        request.setValue("application/json", forHTTPHeaderField: "Content-Type")
        request.httpBody = try JSONEncoder().encode(body)

        logger.debug("MLX request", metadata: [
            "model": .string(model),
            "messages": .string("\(messages.count)"),
        ])

        let (data, response): (Data, URLResponse)
        do {
            (data, response) = try await session.data(for: request)
        } catch let urlError as URLError {
            throw MLXError.backendUnavailable(urlError.localizedDescription)
        }

        guard let http = response as? HTTPURLResponse else {
            throw MLXError.backendUnavailable("Non-HTTP response")
        }
        guard http.statusCode == 200 else {
            let errorBody = String(data: data, encoding: .utf8) ?? "(no body)"
            throw MLXError.inferenceError(http.statusCode, errorBody)
        }

        let decoded: ChatCompletionResponse
        do {
            decoded = try JSONDecoder().decode(ChatCompletionResponse.self, from: data)
        } catch {
            throw MLXError.decodingError(error.localizedDescription)
        }

        let text = decoded.choices.first?.message.content ?? ""
        let inputTokens = decoded.usage.promptTokens
        let outputTokens = decoded.usage.completionTokens

        logger.debug("MLX response", metadata: [
            "input_tokens": .string("\(inputTokens)"),
            "output_tokens": .string("\(outputTokens)"),
        ])

        return (text: text, inputTokens: inputTokens, outputTokens: outputTokens)
    }
}

private func toOpenAIMessages(from request: MessageRequest) -> [ChatMessage] {
    var result: [ChatMessage] = []
    if let system = request.system {
        result.append(ChatMessage(role: .system, content: system.asText))
    }
    result.append(contentsOf: request.messages.map {
        ChatMessage(role: $0.role == .user ? .user : .assistant, content: $0.content.asText)
    })
    return result
}

private func withMLXError<T>(_ operation: () async throws -> T) async throws -> T {
    do {
        return try await operation()
    } catch let error as MLXError {
        switch error {
        case .backendUnavailable:
            throw HTTPError(.serviceUnavailable, message: error.localizedDescription)
        case .inferenceError:
            throw HTTPError(.badGateway, message: error.localizedDescription)
        case .invalidURL, .decodingError:
            throw HTTPError(.internalServerError, message: error.localizedDescription)
        }
    }
}

// Application+build.swift
protocol AppArguments {
    var host: String { get }
    var port: Int { get }
    var logLevel: Logger.Level { get }
    var mlxURL: String { get }
    var mlxModel: String { get }
}

func buildApplication(_ args: some AppArguments) async throws -> some ApplicationProtocol {
    var logger = Logger(label: "swift-mlx-gateway")
    logger.logLevel = args.logLevel

    let mlxClient = MLXClient(baseURL: args.mlxURL, logger: logger)
    let router = buildRouter(mlxClient: mlxClient, modelID: args.mlxModel)
    // ...
}

// Anthropic Messages API route (from Router+build.swift)
router.post("v1/messages") { request, context -> MessageResponse in
    let payload = try await request.decode(as: MessageRequest.self, context: context)
    try validate(payload)
    let messages = toOpenAIMessages(from: payload)
    let result = try await withMLXError {
        try await mlxClient.complete(
            messages: messages,
            model: modelID,
            maxTokens: payload.maxTokens,
            temperature: payload.temperature ?? 1.0
        )
    }
    return buildAnthropicResponse(
        text: result.text,
        model: modelID,
        inputTokens: result.inputTokens,
        outputTokens: result.outputTokens
    )
}

@Option(help: "Base URL of the local MLX inference server (mlx_lm.server)")
var mlxURL: String = "http://localhost:8081"

@Option(help: "MLX model ID (must match mlx_lm.server --model)")
var mlxModel: String = "mlx-community/Qwen3-8B-4bit"

# Terminal 1: MLX backend
mlx_lm.server --model mlx-community/Qwen3-8B-4bit --port 8081

# Terminal 2: Gateway
swift run gateway \
  --mlx-url http://localhost:8081 \
  --mlx-model mlx-community/Qwen3-8B-4bit

curl -s -X POST localhost:8080/v1/messages \
  -H "Content-Type: application/json" \
  -d '{
    "model": "mlx-community/Qwen3-8B-4bit",
    "max_tokens": 1024,
    "messages": [{"role": "user", "content": "Write a haiku about Swift."}]
  }' | jq '.content[0].text'

{"role": "user", "content": "Write a haiku about Swift. /no_think"}

curl -s -X POST localhost:8080/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{
    "model": "mlx-community/Qwen3-8B-4bit",
    "messages": [{"role": "user", "content": "Explain Structured Concurrency in two sentences."}]
  }' | jq '.choices[0].message.content'

export ANTHROPIC_BASE_URL=http://localhost:8080
claude

curl -s -i -X POST localhost:8080/v1/messages \
  -H "Content-Type: application/json" \
  -d '{"model":"x","max_tokens":256,"messages":[{"role":"user","content":"test"}]}'

git add .
git commit -m "article-03: MLXClient — real inference via mlx_lm.server"
git tag article-03
git push origin main --tags