Refactor message handling

Author: domonkosadam

libs/canvas-api-2/src/main/java/com/instructure/canvasapi2/models/journey/JourneyAssistChatMessage.kt

@@ -1,13 +1,7 @@
 package com.instructure.canvasapi2.models.journey
 
 data class JourneyAssistChatMessage(
-    val id: String,
-    val prompt: String = "",
     val text: String = "",
     val role: JourneyAssistRole,
-    val chipOptions: List<JourneyAssistChipOption> = emptyList(),
-    val flashCards: List<JourneyAssistFlashCard> = emptyList(),
-    val quizItems: List<JourneyAssistQuizItem> = emptyList(),
-    val citations: List<JourneyAssistCitation> = emptyList(),
 )
 

libs/horizon/src/main/java/com/instructure/horizon/features/aiassistant/chat/AiAssistChatScreen.kt

@@ -31,11 +31,11 @@ import androidx.compose.ui.text.input.TextFieldValue
 import androidx.compose.ui.tooling.preview.Preview
 import androidx.compose.ui.unit.dp
 import androidx.navigation.NavHostController
-import com.instructure.canvasapi2.models.journey.JourneyAssistChatMessage
 import com.instructure.canvasapi2.models.journey.JourneyAssistRole
 import com.instructure.canvasapi2.utils.ContextKeeper
 import com.instructure.horizon.features.aiassistant.common.composable.AiAssistMessage
 import com.instructure.horizon.features.aiassistant.common.composable.AiAssistScaffold
+import com.instructure.horizon.features.aiassistant.common.model.AiAssistMessage
 import com.instructure.horizon.features.aiassistant.navigation.AiAssistRoute
 import com.instructure.horizon.horizonui.foundation.HorizonColors
 import com.instructure.horizon.horizonui.molecules.Spinner
@@ -107,15 +107,11 @@ private fun AssistChatScreenPreview() {
     ContextKeeper.appContext = LocalContext.current
     val state = AiAssistChatUiState(
         messages = listOf(
-            JourneyAssistChatMessage(
-                id = "1",
-                prompt = "Hello",
+            AiAssistMessage(
                 text = "Hello",
                 role = JourneyAssistRole.User,
             ),
-            JourneyAssistChatMessage(
-                id = "2",
-                prompt = "Hi there! How can I assist you today?",
+            AiAssistMessage(
                 text = "Hi there! How can I assist you today?",
                 role = JourneyAssistRole.Assistant
             )

libs/horizon/src/main/java/com/instructure/horizon/features/aiassistant/chat/AiAssistChatUiState.kt

@@ -1,10 +1,10 @@
 package com.instructure.horizon.features.aiassistant.chat
 
 import androidx.compose.ui.text.input.TextFieldValue
-import com.instructure.canvasapi2.models.journey.JourneyAssistChatMessage
+import com.instructure.horizon.features.aiassistant.common.model.AiAssistMessage
 
 data class AiAssistChatUiState(
-    val messages: List<JourneyAssistChatMessage> = emptyList(),
+    val messages: List<AiAssistMessage> = emptyList(),
     val isLoading: Boolean = false,
     val error: String? = null,
     val isFeedbackEnabled: Boolean = false,

libs/horizon/src/main/java/com/instructure/horizon/features/aiassistant/chat/AiAssistChatViewModel.kt

@@ -19,18 +19,18 @@ package com.instructure.horizon.features.aiassistant.chat
 import androidx.compose.ui.text.input.TextFieldValue
 import androidx.lifecycle.ViewModel
 import androidx.lifecycle.viewModelScope
-import com.instructure.canvasapi2.models.journey.JourneyAssistChatMessage
 import com.instructure.canvasapi2.models.journey.JourneyAssistRole
 import com.instructure.canvasapi2.utils.weave.catch
 import com.instructure.canvasapi2.utils.weave.tryLaunch
 import com.instructure.horizon.features.aiassistant.common.AiAssistContextProvider
 import com.instructure.horizon.features.aiassistant.common.AiAssistRepository
+import com.instructure.horizon.features.aiassistant.common.model.AiAssistMessage
 import com.instructure.horizon.features.aiassistant.common.model.toContextSourceList
+import com.instructure.horizon.features.aiassistant.common.model.toJourneyAssistChatMessages
 import dagger.hilt.android.lifecycle.HiltViewModel
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.asStateFlow
 import kotlinx.coroutines.flow.update
-import java.util.UUID
 import javax.inject.Inject
 
 @HiltViewModel
@@ -72,13 +72,13 @@ class AiAssistChatViewModel @Inject constructor(
         evaluatePrompt(message)
     }
 
-    private fun evaluatePrompt(message: JourneyAssistChatMessage) {
+    private fun evaluatePrompt(message: AiAssistMessage) {
         viewModelScope.tryLaunch {
             _uiState.update {
                 it.copy(isLoading = true)
             }
 
-            val response = answerPrompt(message.prompt)
+            val response = answerPrompt(message.text)
             aiAssistMessages.add(response)
 
             _uiState.update {
@@ -96,21 +96,19 @@ class AiAssistChatViewModel @Inject constructor(
         }
     }
 
-    private suspend fun answerPrompt(prompt: String): JourneyAssistChatMessage {
+    private suspend fun answerPrompt(prompt: String): AiAssistMessage {
         val response = repository.answerPrompt(
             prompt,
-            aiAssistMessages,
+            aiAssistMessages.toJourneyAssistChatMessages(),
             aiAssistContextState
         )
         aiAssistContextState = response.state ?: aiAssistContextState
         return response.message
     }
 
-    private fun addMessage(prompt: String): JourneyAssistChatMessage {
-        val message = JourneyAssistChatMessage(
-            id = UUID.randomUUID().toString(),
+    private fun addMessage(prompt: String): AiAssistMessage {
+        val message = AiAssistMessage(
             text = prompt,
-            prompt = prompt,
             role = JourneyAssistRole.User,
         )
         aiAssistMessages.add(message)

libs/horizon/src/main/java/com/instructure/horizon/features/aiassistant/common/AiAssistContextProvider.kt

@@ -16,18 +16,18 @@
  */
 package com.instructure.horizon.features.aiassistant.common
 
-import com.instructure.canvasapi2.models.journey.JourneyAssistChatMessage
 import com.instructure.canvasapi2.models.journey.JourneyAssistState
 import com.instructure.horizon.features.aiassistant.common.model.AiAssistContext
 import com.instructure.horizon.features.aiassistant.common.model.AiAssistContextSource
+import com.instructure.horizon.features.aiassistant.common.model.AiAssistMessage
 import javax.inject.Inject
 import javax.inject.Singleton
 
 @Singleton
 class AiAssistContextProvider @Inject constructor() {
     var aiAssistContext = AiAssistContext()
 
-    fun addMessageToChatHistory(message: JourneyAssistChatMessage) {
+    fun addMessageToChatHistory(message: AiAssistMessage) {
         aiAssistContext = aiAssistContext.copy(
             chatHistory = aiAssistContext.chatHistory + message
         )

libs/horizon/src/main/java/com/instructure/horizon/features/aiassistant/common/AiAssistRepository.kt

@@ -5,11 +5,11 @@ import com.instructure.canvasapi2.models.journey.JourneyAssistChatMessage
 import com.instructure.canvasapi2.models.journey.JourneyAssistRequestBody
 import com.instructure.canvasapi2.models.journey.JourneyAssistRole
 import com.instructure.canvasapi2.models.journey.JourneyAssistState
-import java.util.UUID
+import com.instructure.horizon.features.aiassistant.common.model.AiAssistMessage
 import javax.inject.Inject
 
 data class AiAssistResponse(
-    val message: JourneyAssistChatMessage,
+    val message: AiAssistMessage,
     val state: JourneyAssistState?
 )
 
@@ -23,15 +23,14 @@ class AiAssistRepository @Inject constructor(
     ): AiAssistResponse {
         val requestBody = JourneyAssistRequestBody(prompt, history, state)
         val response = journeyAssistAPI.answerPrompt(requestBody)
-        val message = JourneyAssistChatMessage(
-            id = UUID.randomUUID().toString(),
-            prompt = response.response.orEmpty(),
+        val message = AiAssistMessage(
             text = response.response.orEmpty(),
             role = JourneyAssistRole.Assistant,
             chipOptions = response.chips,
             flashCards = response.flashCards,
             quizItems = response.quizItems,
-            citations = response.citations
+            citations = response.citations,
+            errorMessage = response.error,
         )
         return AiAssistResponse(message, response.state)
     }

libs/horizon/src/main/java/com/instructure/horizon/features/aiassistant/common/composable/AiAssistMessage.kt

@@ -22,15 +22,19 @@ import androidx.compose.foundation.layout.fillMaxWidth
 import androidx.compose.runtime.Composable
 import androidx.compose.ui.Alignment
 import androidx.compose.ui.Modifier
-import com.instructure.canvasapi2.models.journey.JourneyAssistChatMessage
 import com.instructure.canvasapi2.models.journey.JourneyAssistRole
+import com.instructure.horizon.features.aiassistant.common.model.AiAssistMessage
 
 @Composable
 fun AiAssistMessage(
-    message: JourneyAssistChatMessage,
+    message: AiAssistMessage,
     onSendPrompt: (String) -> Unit
 ) {
-    if (message.role == JourneyAssistRole.Assistant) {
+    if (!message.errorMessage.isNullOrBlank()) {
+        AiAssistResponseTextBlock(
+            text = message.errorMessage
+        )
+    } else if (message.role == JourneyAssistRole.Assistant) {
         AiAssistResponseTextBlock(
             text = message.text,
             footerState = AiAssistResponseTextBlockFooterState(

libs/horizon/src/main/java/com/instructure/horizon/features/aiassistant/common/model/AiAssistContext.kt

@@ -1,12 +1,11 @@
 package com.instructure.horizon.features.aiassistant.common.model
 
-import com.instructure.canvasapi2.models.journey.JourneyAssistChatMessage
 import com.instructure.canvasapi2.models.journey.JourneyAssistState
 
 data class AiAssistContext(
     val contextString: String? = null,
     val contextSources: List<AiAssistContextSource> = emptyList(),
-    val chatHistory: List<JourneyAssistChatMessage> = emptyList(),
+    val chatHistory: List<AiAssistMessage> = emptyList(),
 ) {
     val state: JourneyAssistState
         get() {

libs/horizon/src/main/java/com/instructure/horizon/features/aiassistant/common/model/AiAssistMessage.kt

@@ -0,0 +1,52 @@
+/*
+ * Copyright (C) 2026 - present Instructure, Inc.
+ *
+ *     This program is free software: you can redistribute it and/or modify
+ *     it under the terms of the GNU General Public License as published by
+ *     the Free Software Foundation, version 3 of the License.
+ *
+ *     This program is distributed in the hope that it will be useful,
+ *     but WITHOUT ANY WARRANTY; without even the implied warranty of
+ *     MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ *     GNU General Public License for more details.
+ *
+ *     You should have received a copy of the GNU General Public License
+ *     along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ *
+ */
+package com.instructure.horizon.features.aiassistant.common.model
+
+import com.instructure.canvasapi2.models.journey.JourneyAssistChatMessage
+import com.instructure.canvasapi2.models.journey.JourneyAssistChipOption
+import com.instructure.canvasapi2.models.journey.JourneyAssistCitation
+import com.instructure.canvasapi2.models.journey.JourneyAssistFlashCard
+import com.instructure.canvasapi2.models.journey.JourneyAssistQuizItem
+import com.instructure.canvasapi2.models.journey.JourneyAssistRole
+
+data class AiAssistMessage(
+    val text: String = "",
+    val role: JourneyAssistRole,
+    val chipOptions: List<JourneyAssistChipOption> = emptyList(),
+    val flashCards: List<JourneyAssistFlashCard> = emptyList(),
+    val quizItems: List<JourneyAssistQuizItem> = emptyList(),
+    val citations: List<JourneyAssistCitation> = emptyList(),
+    val errorMessage: String? = null
+)
+
+fun List<JourneyAssistChatMessage>.toAiAssistMessages(): List<AiAssistMessage> {
+    return this.map {
+        AiAssistMessage(
+            text = it.text,
+            role = it.role,
+        )
+    }
+}
+
+fun List<AiAssistMessage>.toJourneyAssistChatMessages(): List<JourneyAssistChatMessage> {
+    return this.map {
+        JourneyAssistChatMessage(
+            text = it.text,
+            role = it.role,
+        )
+    }
+}

libs/horizon/src/main/java/com/instructure/horizon/features/aiassistant/flashcard/AiAssistFlashcardViewModel.kt

@@ -23,6 +23,7 @@ import com.instructure.canvasapi2.utils.weave.catch
 import com.instructure.canvasapi2.utils.weave.tryLaunch
 import com.instructure.horizon.features.aiassistant.common.AiAssistContextProvider
 import com.instructure.horizon.features.aiassistant.common.AiAssistRepository
+import com.instructure.horizon.features.aiassistant.common.model.toJourneyAssistChatMessages
 import dagger.hilt.android.lifecycle.HiltViewModel
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.asStateFlow
@@ -73,8 +74,8 @@ class AiAssistFlashcardViewModel @Inject constructor(
 
             val response = aiAssistRepository.answerPrompt(
                 prompt = aiAssistContextProvider.aiAssistContext.chatHistory
-                    .lastOrNull { it.role == JourneyAssistRole.User }?.prompt.orEmpty(),
-                history = aiAssistContextProvider.aiAssistContext.chatHistory,
+                    .lastOrNull { it.role == JourneyAssistRole.User }?.text.orEmpty(),
+                history = aiAssistContextProvider.aiAssistContext.chatHistory.toJourneyAssistChatMessages(),
                 state = aiAssistContextProvider.aiAssistContext.state
             )