Custom LLM Provider for Foundation Models

import FoundationModels

// MARK: - Configuration (Hashable lookup key for Executor reuse)
struct EchoModelConfiguration: LanguageModelConfiguration {
    let serverURL: URL

    func hash(into hasher: inout Hasher) {
        hasher.combine(serverURL)
    }
    static func == (lhs: EchoModelConfiguration, rhs: EchoModelConfiguration) -> Bool {
        lhs.serverURL == rhs.serverURL
    }
}

// MARK: - Executor (does the real work)
final class EchoModelExecutor: LanguageModelExecutor {
    typealias Configuration = EchoModelConfiguration

    private let config: Configuration

    required init(configuration: Configuration) {
        self.config = configuration
    }

    func prewarm() async {
        // Open connection or pre-load weights here.
        // No-op for this echo example.
    }

    func respond(
        to request: LanguageModelRequest,
        model: some LanguageModel,
        channel: LanguageModelResponseChannel
    ) async throws {
        // 1. Emit metadata so callers can log request IDs immediately.
        await channel.send(.metadataUpdate(.init(modelID: "echo-1.0", requestID: UUID().uuidString)))

        // 2. Extract the last user prompt from the transcript.
        let userText = request.transcript.entries.compactMap { entry -> String? in
            if case .prompt(let content) = entry { return content.text }
            return nil
        }.last ?? "(no prompt)"

        // 3. Emit token-count usage upfront.
        let tokenCount = userText.split(separator: " ").count
        await channel.send(.usageUpdate(.init(promptTokenCount: tokenCount)))

        // 4. Stream echoed response word-by-word.
        for word in userText.split(separator: " ") {
            await channel.send(.textDelta(String(word) + " "))
            try await Task.sleep(for: .milliseconds(80))
        }
    }
}

// MARK: - LanguageModel (trivially constructable)
struct EchoLanguageModel: LanguageModel {
    typealias Executor = EchoModelExecutor

    var capabilities: LanguageModelCapabilities { .init(textGeneration: true) }

    func makeConfiguration() -> EchoModelConfiguration {
        EchoModelConfiguration(serverURL: URL(string: "https://echo.example.com")!)
    }
}

// MARK: - Usage (identical to using the system model)
async func runEchoDemo() async throws {
    let model = EchoLanguageModel()
    let session = LanguageModelSession(model: model)
    let response = try await session.respond(to: "Hello, Foundation Models!")
    print(response.content) // "Hello, Foundation Models! "
}