From bc93c384a7fd392ea897dd80a3487210c8cd05ca Mon Sep 17 00:00:00 2001
From: Han Yin <han.yin@arm.com>
Date: Mon, 14 Apr 2025 15:19:09 -0700
Subject: [PATCH] data: introduce Model entity and DAO; update DI module

---
 .../llama/revamp/data/local/AppDatabase.kt    |  8 ++++-
 .../llama/revamp/data/local/ModelDao.kt       | 32 +++++++++++++++++++
 .../llama/revamp/data/local/ModelEntity.kt    | 32 +++++++++++++++++++
 .../com/example/llama/revamp/di/AppModule.kt  |  3 ++
 4 files changed, 74 insertions(+), 1 deletion(-)
 create mode 100644 examples/llama.android/app/src/main/java/com/example/llama/revamp/data/local/ModelDao.kt
 create mode 100644 examples/llama.android/app/src/main/java/com/example/llama/revamp/data/local/ModelEntity.kt

diff --git a/examples/llama.android/app/src/main/java/com/example/llama/revamp/data/local/AppDatabase.kt b/examples/llama.android/app/src/main/java/com/example/llama/revamp/data/local/AppDatabase.kt
index 711db848ee..5eb4b4aa99 100644
--- a/examples/llama.android/app/src/main/java/com/example/llama/revamp/data/local/AppDatabase.kt
+++ b/examples/llama.android/app/src/main/java/com/example/llama/revamp/data/local/AppDatabase.kt
@@ -10,9 +10,15 @@ import javax.inject.Singleton
  * Main database for the application.
  */
 @Singleton
-@Database(entities = [SystemPromptEntity::class], version = 1, exportSchema = false)
+@Database(
+    entities = [ModelEntity::class, SystemPromptEntity::class],
+    version = 1,
+    exportSchema = false
+)
 abstract class AppDatabase : RoomDatabase() {
 
+    abstract fun modelDao(): ModelDao
+
     abstract fun systemPromptDao(): SystemPromptDao
 
     companion object {
diff --git a/examples/llama.android/app/src/main/java/com/example/llama/revamp/data/local/ModelDao.kt b/examples/llama.android/app/src/main/java/com/example/llama/revamp/data/local/ModelDao.kt
new file mode 100644
index 0000000000..c59ad1e751
--- /dev/null
+++ b/examples/llama.android/app/src/main/java/com/example/llama/revamp/data/local/ModelDao.kt
@@ -0,0 +1,32 @@
+package com.example.llama.revamp.data.local
+
+import androidx.room.Dao
+import androidx.room.Delete
+import androidx.room.Insert
+import androidx.room.OnConflictStrategy
+import androidx.room.Query
+import kotlinx.coroutines.flow.Flow
+
+@Dao
+interface ModelDao {
+    @Query("SELECT * FROM models ORDER BY dateAdded DESC")
+    fun getAllModels(): Flow<List<ModelEntity>>  // Changed to Flow
+
+    @Query("SELECT * FROM models WHERE id = :id")
+    suspend fun getModelById(id: String): ModelEntity?
+
+    @Query("SELECT * FROM models WHERE id IN (:ids)")
+    suspend fun getModelsByIds(ids: Collection<String>): List<ModelEntity>
+
+    @Insert(onConflict = OnConflictStrategy.REPLACE)
+    suspend fun insertModel(model: ModelEntity)
+
+    @Delete
+    suspend fun deleteModel(model: ModelEntity)
+
+    @Delete
+    suspend fun deleteModels(models: List<ModelEntity>)
+
+    @Query("UPDATE models SET lastUsed = :timestamp WHERE id = :id")
+    suspend fun updateLastUsed(id: String, timestamp: Long)
+}
diff --git a/examples/llama.android/app/src/main/java/com/example/llama/revamp/data/local/ModelEntity.kt b/examples/llama.android/app/src/main/java/com/example/llama/revamp/data/local/ModelEntity.kt
new file mode 100644
index 0000000000..c5a67babf1
--- /dev/null
+++ b/examples/llama.android/app/src/main/java/com/example/llama/revamp/data/local/ModelEntity.kt
@@ -0,0 +1,32 @@
+package com.example.llama.revamp.data.local
+
+import androidx.room.Entity
+import androidx.room.PrimaryKey
+import com.example.llama.revamp.data.model.ModelInfo
+
+@Entity(tableName = "models")
+data class ModelEntity(
+    @PrimaryKey
+    val id: String,
+    val name: String,
+    val path: String,
+    val sizeInBytes: Long,
+    val parameters: String,
+    val quantization: String,
+    val type: String,
+    val contextLength: Int,
+    val lastUsed: Long?,
+    val dateAdded: Long
+) {
+    fun toModelInfo() = ModelInfo(
+        id = id,
+        name = name,
+        path = path,
+        sizeInBytes = sizeInBytes,
+        parameters = parameters,
+        quantization = quantization,
+        type = type,
+        contextLength = contextLength,
+        lastUsed = lastUsed
+    )
+}
diff --git a/examples/llama.android/app/src/main/java/com/example/llama/revamp/di/AppModule.kt b/examples/llama.android/app/src/main/java/com/example/llama/revamp/di/AppModule.kt
index 5045c89383..70b20e3cdf 100644
--- a/examples/llama.android/app/src/main/java/com/example/llama/revamp/di/AppModule.kt
+++ b/examples/llama.android/app/src/main/java/com/example/llama/revamp/di/AppModule.kt
@@ -38,6 +38,9 @@ abstract class AppModule {
         @Provides
         fun provideAppDatabase(@ApplicationContext context: Context) = AppDatabase.getDatabase(context)
 
+        @Provides
+        fun providesModelDao(appDatabase: AppDatabase) = appDatabase.modelDao()
+
         @Provides
         fun providesSystemPromptDao(appDatabase: AppDatabase) = appDatabase.systemPromptDao()
     }