diff --git a/Sources/Hub/HubApi.swift b/Sources/Hub/HubApi.swift index fb54ce1e..86679b04 100644 --- a/Sources/Hub/HubApi.swift +++ b/Sources/Hub/HubApi.swift @@ -10,6 +10,53 @@ import Foundation import Network import os +/// https://datatracker.ietf.org/doc/html/rfc7540#section-8.1.2 +/// `requests` in Python leaves headers as their original casing, +/// where as Swift strictly adheres to RFC 7540 and can force lower case. +/// This is relevant for Xet +enum HFHttpHeaders { + static let location = "Location" + static let etag = "Etag" + static let contentLength = "Content-Length" + static let repoCommit = "X-Repo-Commit" + static let linkedEtag = "X-Linked-Etag" + static let linkedSize = "X-Linked-Size" + static let xetHash = "x-xet-hash" + static let xetRefreshRoute = "X-Xet-Refresh-Route" + static let linkXetAuthKey = "xet-auth" +} + +public struct XetFileData { + let fileHash: String + let refreshRoute: String +} + +/// `requests` automatically parses Link headers into `response.links`, +/// we implement similar functionality here. +extension HTTPURLResponse { + func getLinkURL(for rel: String) -> String? { + guard let linkHeader = allHeaderFields["Link"] as? String else { + return nil + } + + for link in linkHeader.split(separator: ",") { + let trimmed = link.trimmingCharacters(in: .whitespaces) + + if trimmed.contains("rel=\"\(rel)\"") || trimmed.contains("rel=\(rel)") { + if let start = trimmed.firstIndex(of: "<"), + let end = trimmed.firstIndex(of: ">"), + start < end + { + let startIndex = trimmed.index(after: start) + return String(trimmed[startIndex.. String { + ProcessInfo.processInfo.environment["HF_ENDPOINT"] ?? "https://huggingface.co" + } + static func hfTokenFromEnv() -> String? { let possibleTokens = [ { ProcessInfo.processInfo.environment["HF_TOKEN"] }, @@ -573,6 +624,9 @@ public extension HubApi { /// Size of the file. In case of an LFS file, contains the size of the actual LFS file, not the pointer. public let size: Int? + + /// Xet file data, if available. Contains the file hash and the refresh route. + public let xetFileData: XetFileData? } /// Metadata about a file in the local directory related to a download process @@ -601,12 +655,49 @@ public extension HubApi { let location = response.statusCode == 302 ? response.value(forHTTPHeaderField: "Location") : response.url?.absoluteString return FileMetadata( - commitHash: response.value(forHTTPHeaderField: "X-Repo-Commit"), + commitHash: response.value(forHTTPHeaderField: HFHttpHeaders.repoCommit), etag: normalizeEtag( - (response.value(forHTTPHeaderField: "X-Linked-Etag")) ?? (response.value(forHTTPHeaderField: "Etag")) + (response.value(forHTTPHeaderField: HFHttpHeaders.linkedEtag)) ?? (response.value(forHTTPHeaderField: HFHttpHeaders.etag)) ), location: location ?? url.absoluteString, - size: Int(response.value(forHTTPHeaderField: "X-Linked-Size") ?? response.value(forHTTPHeaderField: "Content-Length") ?? "") + size: Int(response.value(forHTTPHeaderField: HFHttpHeaders.linkedSize) ?? response.value(forHTTPHeaderField: HFHttpHeaders.contentLength) ?? ""), + xetFileData: parseXetFileDataFromResponse(response: response, endpoint: endpoint) + ) + } + + /// https://github.com/huggingface/huggingface_hub/blob/b698915d6b582c72806ac3e91c43bfd8dde35856/src/huggingface_hub/utils/_xet.py#L29 + private func parseXetFileDataFromResponse( + response: HTTPURLResponse?, + endpoint: String? = nil + ) -> XetFileData? { + guard let response else { + return nil + } + + guard let fileHash = response.allHeaderFields[HFHttpHeaders.xetHash] as? String else { + return nil + } + + guard var refreshRoute = response.getLinkURL(for: HFHttpHeaders.linkXetAuthKey) + ?? response.allHeaderFields[HFHttpHeaders.xetRefreshRoute] as? String + else { + return nil + } + + let endpoint = endpoint ?? "https://huggingface.co" + + let defaultEndpoint = "https://huggingface.co" + + if refreshRoute.hasPrefix(defaultEndpoint) { + refreshRoute = refreshRoute.replacingOccurrences( + of: defaultEndpoint.trimmingCharacters(in: CharacterSet(charactersIn: "/")), + with: endpoint.trimmingCharacters(in: CharacterSet(charactersIn: "/")) + ) + } + + return XetFileData( + fileHash: fileHash, + refreshRoute: refreshRoute ) } diff --git a/Tests/HubTests/HubApiTests.swift b/Tests/HubTests/HubApiTests.swift index 40e75453..cbb70e3a 100644 --- a/Tests/HubTests/HubApiTests.swift +++ b/Tests/HubTests/HubApiTests.swift @@ -103,6 +103,18 @@ class HubApiTests: XCTestCase { } } + func testGetXetMetadata() async throws { + do { + let url = URL(string: "https://huggingface.co/FL33TW00D-HF/xet-test/resolve/main/tokenizer.json") + let metadata = try await Hub.getFileMetadata(fileURL: url!) + + XCTAssertNotNil(metadata.xetFileData) + XCTAssertEqual(metadata.xetFileData?.fileHash, "6aec39639a0a2d1ca966356b8c2b8426a484f80ff80731f44fa8482040713bdf") + } catch { + XCTFail("\(error)") + } + } + func testGetFileMetadataBlobPath() async throws { do { let url = URL(string: "https://huggingface.co/enterprise-explorers/Llama-2-7b-chat-coreml/resolve/main/config.json")