swift-apple-intelligence-grpc/Sources/AppleIntelligenceCore/Providers/AppleIntelligenceProvider.swift

import Foundation
import GRPCCore
import GRPCProtobuf
import GRPCNIOTransportHTTP2

/// gRPC service provider for Apple Intelligence
public struct AppleIntelligenceProvider: RegistrableRPCService {
    /// Service descriptor
    public static let serviceDescriptor = ServiceDescriptor(
        fullyQualifiedService: "appleintelligence.AppleIntelligence"
    )

    /// Method descriptors
    enum Methods {
        static let complete = MethodDescriptor(
            service: AppleIntelligenceProvider.serviceDescriptor,
            method: "Complete"
        )
        static let streamComplete = MethodDescriptor(
            service: AppleIntelligenceProvider.serviceDescriptor,
            method: "StreamComplete"
        )
        static let health = MethodDescriptor(
            service: AppleIntelligenceProvider.serviceDescriptor,
            method: "Health"
        )
    }

    /// The underlying AI service
    private let service: AppleIntelligenceService

    /// Optional API key for authentication
    private let apiKey: String?

    public init(service: AppleIntelligenceService, apiKey: String? = nil) {
        self.service = service
        self.apiKey = apiKey
    }

    public func registerMethods<Transport: ServerTransport>(with router: inout RPCRouter<Transport>) {
        // Register Complete method (unary)
        router.registerHandler(
            forMethod: Methods.complete,
            deserializer: ProtobufDeserializer<Appleintelligence_CompletionRequest>(),
            serializer: ProtobufSerializer<Appleintelligence_CompletionResponse>()
        ) { request, context in
            try self.validateApiKey(metadata: request.metadata)

            // Collect the single message from the request stream
            var requestMessage: Appleintelligence_CompletionRequest?
            for try await message in request.messages {
                requestMessage = message
                break
            }

            guard let message = requestMessage else {
                throw RPCError(code: .invalidArgument, message: "No request message received")
            }

            let text = try await self.service.complete(
                prompt: message.prompt,
                temperature: message.hasTemperature ? message.temperature : nil,
                maxTokens: message.hasMaxTokens ? Int(message.maxTokens) : nil
            )

            var response = Appleintelligence_CompletionResponse()
            response.id = UUID().uuidString
            response.text = text
            response.finishReason = "stop"

            return StreamingServerResponse(single: ServerResponse(message: response))
        }

        // Register StreamComplete method (server streaming)
        router.registerHandler(
            forMethod: Methods.streamComplete,
            deserializer: ProtobufDeserializer<Appleintelligence_CompletionRequest>(),
            serializer: ProtobufSerializer<Appleintelligence_CompletionChunk>()
        ) { request, context in
            try self.validateApiKey(metadata: request.metadata)

            // Collect the single message from the request stream
            var requestMessage: Appleintelligence_CompletionRequest?
            for try await message in request.messages {
                requestMessage = message
                break
            }

            guard let message = requestMessage else {
                throw RPCError(code: .invalidArgument, message: "No request message received")
            }

            let completionId = UUID().uuidString
            let prompt = message.prompt
            let temperature = message.hasTemperature ? message.temperature : nil
            let maxTokens = message.hasMaxTokens ? Int(message.maxTokens) : nil

            return StreamingServerResponse { writer in
                let stream = await self.service.streamComplete(
                    prompt: prompt,
                    temperature: temperature,
                    maxTokens: maxTokens
                )

                var lastContent = ""
                for try await partialResponse in stream {
                    // Calculate the delta (new text since last response)
                    let delta: String
                    if partialResponse.hasPrefix(lastContent) {
                        delta = String(partialResponse.dropFirst(lastContent.count))
                    } else {
                        delta = partialResponse
                    }
                    lastContent = partialResponse

                    if !delta.isEmpty {
                        var chunk = Appleintelligence_CompletionChunk()
                        chunk.id = completionId
                        chunk.delta = delta
                        chunk.isFinal = false
                        try await writer.write(chunk)
                    }
                }

                // Send final chunk
                var finalChunk = Appleintelligence_CompletionChunk()
                finalChunk.id = completionId
                finalChunk.delta = ""
                finalChunk.isFinal = true
                finalChunk.finishReason = "stop"
                try await writer.write(finalChunk)

                return [:]
            }
        }

        // Register Health method (unary)
        router.registerHandler(
            forMethod: Methods.health,
            deserializer: ProtobufDeserializer<Appleintelligence_HealthRequest>(),
            serializer: ProtobufSerializer<Appleintelligence_HealthResponse>()
        ) { request, context in
            // Consume request messages (empty for health check)
            for try await _ in request.messages {}

            let isHealthy = await self.service.isAvailable
            let modelStatus = await self.service.getModelStatus()

            var response = Appleintelligence_HealthResponse()
            response.healthy = isHealthy
            response.modelStatus = modelStatus

            return StreamingServerResponse(single: ServerResponse(message: response))
        }
    }

    /// Validate API key if configured
    private func validateApiKey(metadata: Metadata) throws {
        guard let expectedKey = apiKey else {
            return // No API key required
        }

        // Look for Authorization header in metadata
        let authValues = metadata["authorization"]
        guard let authHeader = authValues.first(where: { _ in true }),
              case .string(let authString) = authHeader,
              authString.hasPrefix("Bearer ") else {
            throw RPCError(code: .unauthenticated, message: "Missing or invalid Authorization header")
        }

        let providedKey = String(authString.dropFirst("Bearer ".count))
        guard providedKey == expectedKey else {
            throw RPCError(code: .unauthenticated, message: "Invalid API key")
        }
    }
}