feat: Added save_chat_message() to API

learning_assistant/api.py

@@ -4,14 +4,15 @@
 import logging
 
 from django.conf import settings
+from django.contrib.auth import get_user_model
 from django.core.cache import cache
 from edx_django_utils.cache import get_cache_key
 from jinja2 import BaseLoader, Environment
 from opaque_keys import InvalidKeyError
 
 from learning_assistant.constants import ACCEPTED_CATEGORY_TYPES, CATEGORY_TYPE_MAP
 from learning_assistant.data import LearningAssistantCourseEnabledData
-from learning_assistant.models import LearningAssistantCourseEnabled
+from learning_assistant.models import LearningAssistantCourseEnabled, LearningAssistantMessage
 from learning_assistant.platform_imports import (
     block_get_children,
     block_leaf_filter,
@@ -24,6 +25,7 @@
 from learning_assistant.text_utils import html_to_text
 
 log = logging.getLogger(__name__)
+User = get_user_model()
 
 
 def _extract_block_contents(child, category):
@@ -187,3 +189,21 @@ def get_course_id(course_run_id):
     course_data = get_cache_course_run_data(course_run_id, ['course'])
     course_key = course_data['course']
     return course_key
+
+def save_chat_message(user_id, chat_role, message):
+    """
+    Saves the chat message to the database.
+    """
+
+    user = None
+    try:
+        user = User.objects.get(id=user_id)
+    except User.DoesNotExist:
+        raise Exception("User does not exists.")
+
+    # Save the user message to the database.
+    LearningAssistantMessage.objects.create(
+        user=user,
+        role=chat_role,
+        content=message,
+    )

tests/test_api.py

@@ -6,6 +6,7 @@
 
 import ddt
 from django.conf import settings
+from django.contrib.auth import get_user_model
 from django.core.cache import cache
 from django.test import TestCase, override_settings
 from opaque_keys import InvalidKeyError
@@ -19,13 +20,14 @@
     learning_assistant_available,
     learning_assistant_enabled,
     render_prompt_template,
+    save_chat_message,
     set_learning_assistant_enabled,
 )
 from learning_assistant.data import LearningAssistantCourseEnabledData
-from learning_assistant.models import LearningAssistantCourseEnabled
+from learning_assistant.models import LearningAssistantCourseEnabled, LearningAssistantMessage
 
 fake_transcript = 'This is the text version from the transcript'
-
+User = get_user_model()
 
 class FakeChild:
     """Fake child block for testing"""
@@ -230,6 +232,30 @@ def test_render_prompt_template_invalid_unit_key(self, mock_get_content):
 
         self.assertNotIn('The following text is useful.', prompt_text)
 
+@ddt.ddt
+class TestLearningAssistantCourseEnabledApi(TestCase):
+    """
+    Test suite for save_chat_message.
+    """
+    def setUp(self):
+        super().setUp()
+
+        self.test_user = User.objects.create(username='username', password='password')
+
+    @ddt.data(
+        (LearningAssistantMessage.USER_ROLE, 'What is the meaning of life, the universe and everything?'),
+        (LearningAssistantMessage.ASSISTANT_ROLE, '42'),
+    )
+    @ddt.unpack
+    def test_save_chat_message(self, chat_role, message):
+        save_chat_message(self.test_user.id, chat_role, message)
+
+        row = LearningAssistantMessage.objects.all().last()
+
+        self.assertEqual(row.role, chat_role)
+        self.assertEqual(row.content, message)
+
+
 
 @ddt.ddt
 class LearningAssistantCourseEnabledApiTests(TestCase):