chore: start Xetification (#215)

* chore: start Xetification

* chore: need to fuzz Xet headers

* chore: comment

* chore: fmt

* chore: move to repo I control

* chore: simplify

---------

Co-authored-by: FL33TW00D <[email protected]>

Author: FL33TW00D

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..<end])
+                }
+            }
+        }
+
+        return nil
+    }
+}
+
 public struct HubApi: Sendable {
     var downloadBase: URL
     var hfToken: String?
@@ -24,7 +71,7 @@ public struct HubApi: Sendable {
     public init(
         downloadBase: URL? = nil,
         hfToken: String? = nil,
-        endpoint: String = "https://huggingface.co",
+        endpoint: String? = nil,
         useBackgroundSession: Bool = false,
         useOfflineMode: Bool? = nil
     ) {
@@ -35,7 +82,7 @@ public struct HubApi: Sendable {
             let documents = FileManager.default.urls(for: .documentDirectory, in: .userDomainMask).first!
             self.downloadBase = documents.appending(component: "huggingface")
         }
-        self.endpoint = endpoint
+        self.endpoint = endpoint ?? Self.hfEndpointfromEnv()
         self.useBackgroundSession = useBackgroundSession
         self.useOfflineMode = useOfflineMode
         NetworkMonitor.shared.startMonitoring()
@@ -50,6 +97,10 @@ public struct HubApi: Sendable {
 }
 
 private extension HubApi {
+    static func hfEndpointfromEnv() -> 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
         )
     }
 

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")