remote: refactor HuggingFaceApiService; implement download feature in HuggingFaceRemoteDataSource

Author: hanyin-arm

examples/llama.android/app/src/main/java/com/example/llama/data/remote/HuggingFaceApiService.kt

@@ -1,7 +1,9 @@
 package com.example.llama.data.remote
 
 import okhttp3.ResponseBody
+import retrofit2.Response
 import retrofit2.http.GET
+import retrofit2.http.HEAD
 import retrofit2.http.Path
 import retrofit2.http.Query
 import retrofit2.http.Streaming
@@ -21,10 +23,17 @@ interface HuggingFaceApiService {
     @GET("api/models/{modelId}")
     suspend fun getModelDetails(@Path("modelId") modelId: String): HuggingFaceModelDetails
 
-    @GET("{modelId}/resolve/main/{filePath}")
+    @HEAD("{modelId}/resolve/main/{filename}")
+    suspend fun getModelFileHeader(
+        @Path("modelId", encoded = true) modelId: String,
+        @Path("filename", encoded = true) filename: String
+    ): Response<Void>
+
+    @Deprecated("Use DownloadManager instead!")
+    @GET("{modelId}/resolve/main/{filename}")
     @Streaming
     suspend fun downloadModelFile(
         @Path("modelId") modelId: String,
-        @Path("filePath") filePath: String
+        @Path("filename") filename: String
     ): ResponseBody
 }

examples/llama.android/app/src/main/java/com/example/llama/data/remote/HuggingFaceModel.kt

@@ -1,31 +1,50 @@
 package com.example.llama.data.remote
 
+import android.net.Uri
+import androidx.core.net.toUri
+import com.example.llama.di.HUGGINGFACE_HOST
 import java.util.Date
 
+private const val FILE_EXTENSION_GGUF = ".gguf"
+
 data class HuggingFaceModel(
     val _id: String,
     val id: String,
     val modelId: String,
 
     val author: String,
-    val createdAt: Date?,
-    val lastModified: Date?,
+    val createdAt: Date,
+    val lastModified: Date,
 
-    val library_name: String?,
-    val pipeline_tag: String?,
-    val tags: List<String>?,
+    val pipeline_tag: String,
+    val tags: List<String>,
 
-    val private: Boolean?,
-    val gated: Boolean?,
+    val private: Boolean,
+    val gated: Boolean,
 
-    val likes: Int?,
-    val downloads: Int?,
+    val likes: Int,
+    val downloads: Int,
 
-    val sha: String?,
+    val sha: String,
+    val siblings: List<Sibling>,
 
-    val siblings: List<Sibling>?,
+    val library_name: String?,
 ) {
     data class Sibling(
         val rfilename: String,
     )
+
+    fun getGgufFilename(): String? =
+        siblings.map { it.rfilename }.first { it.endsWith(FILE_EXTENSION_GGUF) }
+
+    fun toDownloadInfo() = getGgufFilename()?.let { HuggingFaceDownloadInfo(_id, modelId, it) }
+}
+
+data class HuggingFaceDownloadInfo(
+    val _id: String,
+    val modelId: String,
+    val filename: String,
+) {
+    val uri: Uri
+        get() = "$HUGGINGFACE_HOST${modelId}/resolve/main/$filename".toUri()
 }

examples/llama.android/app/src/main/java/com/example/llama/data/remote/HuggingFaceRemoteDataSource.kt

@@ -1,25 +1,47 @@
 package com.example.llama.data.remote
 
+import android.app.DownloadManager
+import android.content.Context
+import android.os.Environment
 import android.util.Log
 import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.delay
 import kotlinx.coroutines.withContext
-import java.io.File
 import javax.inject.Inject
 import javax.inject.Singleton
 
+private const val QUERY_Q4_0_GGUF = "gguf q4_0"
+private const val FILTER_TEXT_GENERATION = "text-generation"
+private const val SORT_BY_DOWNLOADS = "downloads"
+private const val SEARCH_RESULT_LIMIT = 20
+
 interface HuggingFaceRemoteDataSource {
+    /**
+     * Query openly available Q4_0 GGUF models on HuggingFace
+     */
     suspend fun searchModels(
-        query: String? = "gguf q4_0",
-        filter: String? = "text-generation", // Only generative models,
-        sort: String? = "downloads",
+        query: String? = QUERY_Q4_0_GGUF,
+        filter: String? = FILTER_TEXT_GENERATION,
+        sort: String? = SORT_BY_DOWNLOADS,
         direction: String? = "-1",
-        limit: Int? = 20,
+        limit: Int? = SEARCH_RESULT_LIMIT,
         full: Boolean = true,
     ): List<HuggingFaceModel>
 
     suspend fun getModelDetails(modelId: String): HuggingFaceModelDetails
 
-    suspend fun downloadModelFile(modelId: String, filePath: String, outputFile: File): Result<File>
+    /**
+     * Obtain selected HuggingFace model's GGUF file size from HTTP header
+     */
+    suspend fun getFileSize(modelId: String, filePath: String): Long?
+
+    /**
+     * Download selected HuggingFace model's GGUF file via DownloadManager
+     */
+    suspend fun downloadModelFile(
+        context: Context,
+        downloadInfo: HuggingFaceDownloadInfo,
+    ): Result<Unit>
 }
 
 @Singleton
@@ -42,7 +64,7 @@ class HuggingFaceRemoteDataSourceImpl @Inject constructor(
             direction = direction,
             limit = limit,
             full = full,
-        )
+        ).filter { it.gated != true && it.private != true }
     }
 
     override suspend fun getModelDetails(
@@ -51,32 +73,98 @@ class HuggingFaceRemoteDataSourceImpl @Inject constructor(
         apiService.getModelDetails(modelId)
     }
 
-    override suspend fun downloadModelFile(
+    override suspend fun getFileSize(
         modelId: String,
-        filePath: String,
-        outputFile: File
-    ): Result<File> = withContext(Dispatchers.IO) {
+        filePath: String
+    ): Long? = withContext(Dispatchers.IO) {
         try {
-            val response = apiService.downloadModelFile(modelId, filePath)
+            apiService.getModelFileHeader(modelId, filePath).let {
+                if (it.isSuccessful) {
+                    it.headers()[HTTP_HEADER_CONTENT_LENGTH]?.toLongOrNull()
+                } else {
+                    null
+                }
+            }
+        } catch (e: Exception) {
+            Log.e(TAG, "Error getting file size for $modelId: ${e.message}")
+            null
+        }
+    }
 
-            // Create parent directories if needed
-            outputFile.parentFile?.mkdirs()
+    override suspend fun downloadModelFile(
+        context: Context,
+        downloadInfo: HuggingFaceDownloadInfo,
+    ): Result<Unit> = withContext(Dispatchers.IO) {
+        try {
+            val downloadManager =
+                context.getSystemService(Context.DOWNLOAD_SERVICE) as DownloadManager
+            val request = DownloadManager.Request(downloadInfo.uri).apply {
+                setTitle("HuggingFace model download")
+                setDescription("Downloading ${downloadInfo.filename}")
+                setNotificationVisibility(DownloadManager.Request.VISIBILITY_VISIBLE_NOTIFY_COMPLETED)
+                setDestinationInExternalPublicDir(
+                    Environment.DIRECTORY_DOWNLOADS,
+                    downloadInfo.filename
+                )
+                setAllowedNetworkTypes(
+                    DownloadManager.Request.NETWORK_WIFI or DownloadManager.Request.NETWORK_MOBILE
+                )
+                setAllowedOverMetered(true)
+                setAllowedOverRoaming(false)
+            }
+            Log.d(TAG, "Enqueuing download request for: ${downloadInfo.modelId}")
+            val downloadId = downloadManager.enqueue(request)
+
+            delay(DOWNLOAD_MANAGER_DOUBLE_CHECK_DELAY)
 
-            // Save the file
-            response.byteStream().use { input ->
-                outputFile.outputStream().use { output ->
-                    input.copyTo(output)
+            val cursor = downloadManager.query(DownloadManager.Query().setFilterById(downloadId))
+            if (cursor != null && cursor.moveToFirst()) {
+                val statusIndex = cursor.getColumnIndex(DownloadManager.COLUMN_STATUS)
+                if (statusIndex >= 0) {
+                    val status = cursor.getInt(statusIndex)
+                    cursor.close()
+
+                    when (status) {
+                        DownloadManager.STATUS_FAILED -> {
+                            // Get failure reason if available
+                            val reasonIndex = cursor.getColumnIndex(DownloadManager.COLUMN_REASON)
+                            val reason = if (reasonIndex >= 0) cursor.getInt(reasonIndex) else -1
+                            val errorMessage = when (reason) {
+                                DownloadManager.ERROR_HTTP_DATA_ERROR -> "HTTP error"
+                                DownloadManager.ERROR_INSUFFICIENT_SPACE -> "Insufficient storage"
+                                DownloadManager.ERROR_TOO_MANY_REDIRECTS -> "Too many redirects"
+                                DownloadManager.ERROR_UNHANDLED_HTTP_CODE -> "Unhandled HTTP code"
+                                DownloadManager.ERROR_CANNOT_RESUME -> "Cannot resume download"
+                                DownloadManager.ERROR_FILE_ERROR -> "File error"
+                                else -> "Unknown error"
+                            }
+                            Result.failure(Exception(errorMessage))
+                        }
+                        else -> {
+                            // Download is pending, paused, or running
+                            Result.success(Unit)
+                        }
+                    }
+                } else {
+                    // Assume success if we can't check status
+                    cursor.close()
+                    Result.success(Unit)
                 }
+            } else {
+                // Assume success if cursor is empty
+                cursor?.close()
+                Result.success(Unit)
             }
-
-            Result.success(outputFile)
         } catch (e: Exception) {
-            Log.e(TAG, "Error downloading file $filePath: ${e.message}")
+            Log.e(TAG, "Failed to enqueue download: ${e.message}")
             Result.failure(e)
         }
     }
 
     companion object {
         private val TAG = HuggingFaceRemoteDataSourceImpl::class.java.simpleName
+
+        private const val HTTP_HEADER_CONTENT_LENGTH = "content-length"
+        private const val DOWNLOAD_MANAGER_DOUBLE_CHECK_DELAY = 500L
     }
 }