From 2bb53681d93e07d414928da25f062e55e503eeab Mon Sep 17 00:00:00 2001 From: Mike Gromer <385630+gromer@users.noreply.github.com> Date: Tue, 25 Feb 2025 16:59:04 -0700 Subject: [PATCH] Add bearer token support to API requests --- Sources/OllamaKit/OllamaKit+Chat.swift | 4 +- Sources/OllamaKit/OllamaKit+CopyModel.swift | 4 +- Sources/OllamaKit/OllamaKit+DeleteModel.swift | 4 +- Sources/OllamaKit/OllamaKit+Embeddings.swift | 4 +- Sources/OllamaKit/OllamaKit+Generate.swift | 4 +- Sources/OllamaKit/OllamaKit+ModelInfo.swift | 4 +- Sources/OllamaKit/OllamaKit+Models.swift | 4 +- Sources/OllamaKit/OllamaKit+PullModel.swift | 4 +- Sources/OllamaKit/OllamaKit+Reachable.swift | 4 +- Sources/OllamaKit/OllamaKit.swift | 34 ++++- Sources/OllamaKit/Utils/OKRouter.swift | 117 ++++++++++-------- 11 files changed, 111 insertions(+), 76 deletions(-) diff --git a/Sources/OllamaKit/OllamaKit+Chat.swift b/Sources/OllamaKit/OllamaKit+Chat.swift index 71dd0bd..06f5adf 100644 --- a/Sources/OllamaKit/OllamaKit+Chat.swift +++ b/Sources/OllamaKit/OllamaKit+Chat.swift @@ -100,7 +100,7 @@ extension OllamaKit { /// - Returns: An `AsyncThrowingStream` emitting the live stream of chat responses from the Ollama API. public func chat(data: OKChatRequestData) -> AsyncThrowingStream { do { - let request = try OKRouter.chat(data: data).asURLRequest(with: baseURL) + let request = try OKRouter.chat(data: data).asURLRequest(with: baseURL, with: bearerToken) return OKHTTPClient.shared.stream(request: request, with: OKChatResponse.self) } catch { @@ -197,7 +197,7 @@ extension OllamaKit { /// - Returns: An `AnyPublisher` emitting the live stream of chat responses from the Ollama API. public func chat(data: OKChatRequestData) -> AnyPublisher { do { - let request = try OKRouter.chat(data: data).asURLRequest(with: baseURL) + let request = try OKRouter.chat(data: data).asURLRequest(with: baseURL, with: bearerToken) return OKHTTPClient.shared.stream(request: request, with: OKChatResponse.self) } catch { diff --git a/Sources/OllamaKit/OllamaKit+CopyModel.swift b/Sources/OllamaKit/OllamaKit+CopyModel.swift index af63a6d..fb69f4f 100644 --- a/Sources/OllamaKit/OllamaKit+CopyModel.swift +++ b/Sources/OllamaKit/OllamaKit+CopyModel.swift @@ -22,7 +22,7 @@ extension OllamaKit { /// - Parameter data: The ``OKCopyModelRequestData`` containing the details needed to copy the model. /// - Throws: An error if the request to copy the model fails. public func copyModel(data: OKCopyModelRequestData) async throws -> Void { - let request = try OKRouter.copyModel(data: data).asURLRequest(with: baseURL) + let request = try OKRouter.copyModel(data: data).asURLRequest(with: baseURL, with: bearerToken) try await OKHTTPClient.shared.send(request: request) } @@ -48,7 +48,7 @@ extension OllamaKit { /// - Returns: A `AnyPublisher` that completes when the copy operation is done. public func copyModel(data: OKCopyModelRequestData) -> AnyPublisher { do { - let request = try OKRouter.copyModel(data: data).asURLRequest(with: baseURL) + let request = try OKRouter.copyModel(data: data).asURLRequest(with: baseURL, with: bearerToken) return OKHTTPClient.shared.send(request: request) } catch { diff --git a/Sources/OllamaKit/OllamaKit+DeleteModel.swift b/Sources/OllamaKit/OllamaKit+DeleteModel.swift index 82a6916..041b5da 100644 --- a/Sources/OllamaKit/OllamaKit+DeleteModel.swift +++ b/Sources/OllamaKit/OllamaKit+DeleteModel.swift @@ -22,7 +22,7 @@ extension OllamaKit { /// - Parameter data: The ``OKDeleteModelRequestData`` containing the details needed to delete the model. /// - Throws: An error if the request to delete the model fails. public func deleteModel(data: OKDeleteModelRequestData) async throws -> Void { - let request = try OKRouter.deleteModel(data: data).asURLRequest(with: baseURL) + let request = try OKRouter.deleteModel(data: data).asURLRequest(with: baseURL, with: bearerToken) try await OKHTTPClient.shared.send(request: request) } @@ -48,7 +48,7 @@ extension OllamaKit { /// - Returns: A `AnyPublisher` that completes when the deletion operation is done. public func deleteModel(data: OKDeleteModelRequestData) -> AnyPublisher { do { - let request = try OKRouter.deleteModel(data: data).asURLRequest(with: baseURL) + let request = try OKRouter.deleteModel(data: data).asURLRequest(with: baseURL, with: bearerToken) return OKHTTPClient.shared.send(request: request) } catch { diff --git a/Sources/OllamaKit/OllamaKit+Embeddings.swift b/Sources/OllamaKit/OllamaKit+Embeddings.swift index 85d457f..d0c8209 100644 --- a/Sources/OllamaKit/OllamaKit+Embeddings.swift +++ b/Sources/OllamaKit/OllamaKit+Embeddings.swift @@ -23,7 +23,7 @@ extension OllamaKit { /// - Returns: An ``OKEmbeddingsResponse`` containing the embeddings from the model. /// - Throws: An error if the request fails or the response can't be decoded. public func embeddings(data: OKEmbeddingsRequestData) async throws -> OKEmbeddingsResponse { - let request = try OKRouter.embeddings(data: data).asURLRequest(with: baseURL) + let request = try OKRouter.embeddings(data: data).asURLRequest(with: baseURL, with: bearerToken) return try await OKHTTPClient.shared.send(request: request, with: OKEmbeddingsResponse.self) } @@ -49,7 +49,7 @@ extension OllamaKit { /// - Returns: A `AnyPublisher` that emits embeddings. public func embeddings(data: OKEmbeddingsRequestData) -> AnyPublisher { do { - let request = try OKRouter.embeddings(data: data).asURLRequest(with: baseURL) + let request = try OKRouter.embeddings(data: data).asURLRequest(with: baseURL, with: bearerToken) return OKHTTPClient.shared.send(request: request, with: OKEmbeddingsResponse.self) } catch { diff --git a/Sources/OllamaKit/OllamaKit+Generate.swift b/Sources/OllamaKit/OllamaKit+Generate.swift index 6c8d9b8..ec897c2 100644 --- a/Sources/OllamaKit/OllamaKit+Generate.swift +++ b/Sources/OllamaKit/OllamaKit+Generate.swift @@ -32,7 +32,7 @@ extension OllamaKit { /// - Returns: An `AsyncThrowingStream` emitting the live stream of responses from the Ollama API. public func generate(data: OKGenerateRequestData) -> AsyncThrowingStream { do { - let request = try OKRouter.generate(data: data).asURLRequest(with: baseURL) + let request = try OKRouter.generate(data: data).asURLRequest(with: baseURL, with: bearerToken) return OKHTTPClient.shared.stream(request: request, with: OKGenerateResponse.self) } catch { @@ -63,7 +63,7 @@ extension OllamaKit { /// - Returns: An `AnyPublisher` emitting the live stream of responses from the Ollama API. public func generate(data: OKGenerateRequestData) -> AnyPublisher { do { - let request = try OKRouter.generate(data: data).asURLRequest(with: baseURL) + let request = try OKRouter.generate(data: data).asURLRequest(with: baseURL, with: bearerToken) return OKHTTPClient.shared.stream(request: request, with: OKGenerateResponse.self) } catch { diff --git a/Sources/OllamaKit/OllamaKit+ModelInfo.swift b/Sources/OllamaKit/OllamaKit+ModelInfo.swift index cabb515..3bdd9e9 100644 --- a/Sources/OllamaKit/OllamaKit+ModelInfo.swift +++ b/Sources/OllamaKit/OllamaKit+ModelInfo.swift @@ -23,7 +23,7 @@ extension OllamaKit { /// - Returns: An ``OKModelInfoResponse`` containing detailed information about the model. /// - Throws: An error if the request fails or the response can't be decoded. public func modelInfo(data: OKModelInfoRequestData) async throws -> OKModelInfoResponse { - let request = try OKRouter.modelInfo(data: data).asURLRequest(with: baseURL) + let request = try OKRouter.modelInfo(data: data).asURLRequest(with: baseURL, with: bearerToken) return try await OKHTTPClient.shared.send(request: request, with: OKModelInfoResponse.self) } @@ -49,7 +49,7 @@ extension OllamaKit { /// - Returns: A `AnyPublisher` that emits detailed information about the model. public func modelInfo(data: OKModelInfoRequestData) -> AnyPublisher { do { - let request = try OKRouter.modelInfo(data: data).asURLRequest(with: baseURL) + let request = try OKRouter.modelInfo(data: data).asURLRequest(with: baseURL, with: bearerToken) return OKHTTPClient.shared.send(request: request, with: OKModelInfoResponse.self) } catch { diff --git a/Sources/OllamaKit/OllamaKit+Models.swift b/Sources/OllamaKit/OllamaKit+Models.swift index 6447301..74e0808 100644 --- a/Sources/OllamaKit/OllamaKit+Models.swift +++ b/Sources/OllamaKit/OllamaKit+Models.swift @@ -21,7 +21,7 @@ extension OllamaKit { /// - Returns: An ``OKModelResponse`` object listing the available models. /// - Throws: An error if the request fails or the response can't be decoded. public func models() async throws -> OKModelResponse { - let request = try OKRouter.models.asURLRequest(with: baseURL) + let request = try OKRouter.models.asURLRequest(with: baseURL, with: bearerToken) return try await OKHTTPClient.shared.send(request: request, with: OKModelResponse.self) } @@ -45,7 +45,7 @@ extension OllamaKit { /// - Returns: A `AnyPublisher` that emits the list of available models. public func models() -> AnyPublisher { do { - let request = try OKRouter.models.asURLRequest(with: baseURL) + let request = try OKRouter.models.asURLRequest(with: baseURL, with: bearerToken) return OKHTTPClient.shared.send(request: request, with: OKModelResponse.self) } catch { diff --git a/Sources/OllamaKit/OllamaKit+PullModel.swift b/Sources/OllamaKit/OllamaKit+PullModel.swift index 3657421..e31b027 100644 --- a/Sources/OllamaKit/OllamaKit+PullModel.swift +++ b/Sources/OllamaKit/OllamaKit+PullModel.swift @@ -31,7 +31,7 @@ extension OllamaKit { /// - Returns: An asynchronous stream that emits ``OKPullModelResponse``. public func pullModel(data: OKPullModelRequestData) -> AsyncThrowingStream { do { - let request = try OKRouter.pullModel(data: data).asURLRequest(with: baseURL) + let request = try OKRouter.pullModel(data: data).asURLRequest(with: baseURL, with: bearerToken) return OKHTTPClient.shared.stream(request: request, with: OKPullModelResponse.self) } catch { @@ -67,7 +67,7 @@ extension OllamaKit { /// - Returns: A combine publisher that emits ``OKPullModelResponse``. public func pullModel(data: OKPullModelRequestData) -> AnyPublisher { do { - let request = try OKRouter.pullModel(data: data).asURLRequest(with: baseURL) + let request = try OKRouter.pullModel(data: data).asURLRequest(with: baseURL, with: bearerToken) return OKHTTPClient.shared.stream(request: request, with: OKPullModelResponse.self) } catch { diff --git a/Sources/OllamaKit/OllamaKit+Reachable.swift b/Sources/OllamaKit/OllamaKit+Reachable.swift index 20b9b4d..9e299ef 100644 --- a/Sources/OllamaKit/OllamaKit+Reachable.swift +++ b/Sources/OllamaKit/OllamaKit+Reachable.swift @@ -21,7 +21,7 @@ extension OllamaKit { /// - Returns: `true` if the Ollama API is reachable, `false` otherwise. public func reachable() async -> Bool { do { - let request = try OKRouter.root.asURLRequest(with: baseURL) + let request = try OKRouter.root.asURLRequest(with: baseURL, with: bearerToken) try await OKHTTPClient.shared.send(request: request) return true @@ -47,7 +47,7 @@ extension OllamaKit { /// - Returns: A `AnyPublisher` that emits `true` if the API is reachable, `false` otherwise. public func reachable() -> AnyPublisher { do { - let request = try OKRouter.root.asURLRequest(with: baseURL) + let request = try OKRouter.root.asURLRequest(with: baseURL, with: bearerToken) return OKHTTPClient.shared.send(request: request) .map { _ in true } diff --git a/Sources/OllamaKit/OllamaKit.swift b/Sources/OllamaKit/OllamaKit.swift index 59a2bbc..913986c 100644 --- a/Sources/OllamaKit/OllamaKit.swift +++ b/Sources/OllamaKit/OllamaKit.swift @@ -12,16 +12,15 @@ public struct OllamaKit: Sendable { var router: OKRouter.Type var decoder: JSONDecoder = .default var baseURL: URL + var bearerToken: String? - /// Initializes a new instance of `OllamaKit` with the default base URL for the Ollama API. + /// Initializes a new instance of `OllamaKit` with the default base URL and no bearer token for the Ollama API. /// /// ```swift /// let ollamaKit = OllamaKit() /// ``` public init() { - let router = OKRouter.self - self.router = router - self.baseURL = URL(string: "http://localhost:11434")! + self.init(baseURL: URL(string: "http://localhost:11434")!, bearerToken: nil) } /// Initializes a new instance of `OllamaKit` with a custom base URL for the Ollama API. @@ -33,8 +32,35 @@ public struct OllamaKit: Sendable { /// /// - Parameter baseURL: The base URL to use for API requests. public init(baseURL: URL) { + self.init(baseURL: baseURL, bearerToken: nil) + } + + /// Initializes a new instance of `OllamaKit` with a bearer token for the Ollama API. + /// + /// ```swift + /// let bearerToken = "super-sected-token" + /// let ollamaKit = OllamaKit(bearerToken: bearerToken) + /// ``` + /// + /// - Parameter bearerToken: The bearer token to use for API requests. + public init(bearerToken: String) { + self.init(baseURL: URL(string: "http://localhost:11434")!, bearerToken: bearerToken) + } + + /// Initializes a new instance of `OllamaKit` with a custom base URL and bearer token for the Ollama API. + /// + /// ```swift + /// let bearerToken = "super-sected-token" + /// let customBaseURL = URL(string: "https://api.customollama.com")! + /// let ollamaKit = OllamaKit(baseURL: customBaseURL, bearerToken: bearerToken) + /// ``` + /// + /// - Parameter baseURL: The base URL to use for API requests. + /// - Parameter bearerToken: The bearer token to use for API requests. + public init(baseURL: URL, bearerToken: String?) { let router = OKRouter.self self.router = router self.baseURL = baseURL + self.bearerToken = bearerToken } } diff --git a/Sources/OllamaKit/Utils/OKRouter.swift b/Sources/OllamaKit/Utils/OKRouter.swift index 0e3d280..382cf53 100644 --- a/Sources/OllamaKit/Utils/OKRouter.swift +++ b/Sources/OllamaKit/Utils/OKRouter.swift @@ -20,47 +20,47 @@ internal enum OKRouter { internal var path: String { switch self { - case .root: - return "/" - case .models: - return "/api/tags" - case .modelInfo: - return "/api/show" - case .generate: - return "/api/generate" - case .chat: - return "/api/chat" - case .copyModel: - return "/api/copy" - case .deleteModel: - return "/api/delete" - case .pullModel: - return "/api/pull" - case .embeddings: - return "/api/embeddings" + case .root: + return "/" + case .models: + return "/api/tags" + case .modelInfo: + return "/api/show" + case .generate: + return "/api/generate" + case .chat: + return "/api/chat" + case .copyModel: + return "/api/copy" + case .deleteModel: + return "/api/delete" + case .pullModel: + return "/api/pull" + case .embeddings: + return "/api/embeddings" } } internal var method: String { switch self { - case .root: - return "HEAD" - case .models: - return "GET" - case .modelInfo: - return "POST" - case .generate: - return "POST" - case .chat: - return "POST" - case .copyModel: - return "POST" - case .deleteModel: - return "DELETE" - case .pullModel: - return "POST" - case .embeddings: - return "POST" + case .root: + return "HEAD" + case .models: + return "GET" + case .modelInfo: + return "POST" + case .generate: + return "POST" + case .chat: + return "POST" + case .copyModel: + return "POST" + case .deleteModel: + return "DELETE" + case .pullModel: + return "POST" + case .embeddings: + return "POST" } } @@ -70,30 +70,39 @@ internal enum OKRouter { } extension OKRouter { - func asURLRequest(with baseURL: URL) throws -> URLRequest { + func asURLRequest(with baseURL: URL, with bearerToken: String?) throws -> URLRequest { + let hasBearerToken = bearerToken != nil && bearerToken!.isEmpty == false let url = baseURL.appendingPathComponent(path) var request = URLRequest(url: url) request.httpMethod = method - request.allHTTPHeaderFields = headers + + if (hasBearerToken) { + request.allHTTPHeaderFields = [ + "Authorization": "Bearer " + bearerToken!, + "Content-Type": "application/json" + ] + } else { + request.allHTTPHeaderFields = headers + } switch self { - case .modelInfo(let data): - request.httpBody = try JSONEncoder.default.encode(data) - case .generate(let data): - request.httpBody = try JSONEncoder.default.encode(data) - case .chat(let data): - request.httpBody = try JSONEncoder.default.encode(data) - case .copyModel(let data): - request.httpBody = try JSONEncoder.default.encode(data) - case .deleteModel(let data): - request.httpBody = try JSONEncoder.default.encode(data) - case .pullModel(let data): - request.httpBody = try JSONEncoder.default.encode(data) - case .embeddings(let data): - request.httpBody = try JSONEncoder.default.encode(data) - default: - break + case .modelInfo(let data): + request.httpBody = try JSONEncoder.default.encode(data) + case .generate(let data): + request.httpBody = try JSONEncoder.default.encode(data) + case .chat(let data): + request.httpBody = try JSONEncoder.default.encode(data) + case .copyModel(let data): + request.httpBody = try JSONEncoder.default.encode(data) + case .deleteModel(let data): + request.httpBody = try JSONEncoder.default.encode(data) + case .pullModel(let data): + request.httpBody = try JSONEncoder.default.encode(data) + case .embeddings(let data): + request.httpBody = try JSONEncoder.default.encode(data) + default: + break } return request