MathApp/ai_tutor/tests_grade_grounding.py

"""
«از پی بپرس» بدونِ lessonId هم به کتاب مستند می‌شود.

آن صفحه از نوارِ پایینِ اپ باز می‌شود و درسی در کار نیست؛ قبلاً در این حالت
retrieval کلاً رد می‌شد و همه‌ی جواب‌ها عمومی بودند. حالا پایه‌ی خودِ کاربر
کتاب را معلوم می‌کند.
"""
from unittest import mock

from django.core.cache import cache
from django.test import TestCase

from accounts.authentication import issue_access_token
from accounts.models import User
from ai_tutor.models import Conversation
from content.models import Grade


class AskWithoutLessonUsesTheStudentsBookTests(TestCase):
    def setUp(self):
        cache.clear()

    def _ask(self, user, **extra):
        return self.client.post(
            '/api/ai/ask', {'userQuestion': 'توان چیست؟', **extra},
            content_type='application/json',
            HTTP_AUTHORIZATION=f'Bearer {issue_access_token(user.id)}',
        )

    def _patched(self):
        retrieve = mock.patch(
            'rag.services.retrieval.retrieve_context',
            return_value=([{'title': 'فصل', 'summary': 'متن'}], {'gradeTitle': 'نهم'}),
        )
        generate = mock.patch(
            'rag.services.generation.generate_answer', return_value=('جواب', [], {}),
        )
        configured = mock.patch('rag.services.config.rag_is_configured', return_value=True)
        return retrieve, generate, configured

    def test_math_grade_picks_that_grades_book(self):
        grade = Grade.objects.create(title='پایه نهم', icon='', color='')
        user = User.objects.create(phone='09120000011', grade_id=grade.id)
        retrieve, generate, configured = self._patched()

        with retrieve as retrieve_mock, generate, configured:
            response = self._ask(user)

        self.assertEqual(response.status_code, 200)
        kwargs = retrieve_mock.call_args.kwargs
        self.assertEqual(kwargs['subject'], 'ریاضی')
        self.assertEqual(kwargs['grade_title'], 'نهم')
        self.assertEqual(response.json()['usedChunks'], [{'title': 'فصل', 'summary': 'متن'}])

    def test_calculus_grades_use_the_calculus_book(self):
        grade = Grade.objects.create(title='پایه دوازدهم', icon='', color='')
        user = User.objects.create(phone='09120000012', grade_id=grade.id)
        retrieve, generate, configured = self._patched()

        with retrieve as retrieve_mock, generate, configured:
            self._ask(user)

        self.assertEqual(retrieve_mock.call_args.kwargs['subject'], 'حسابان')

    def test_the_follow_up_message_keeps_the_book(self):
        grade = Grade.objects.create(title='پایه نهم', icon='', color='')
        user = User.objects.create(phone='09120000013', grade_id=grade.id)
        retrieve, generate, configured = self._patched()

        with retrieve as retrieve_mock, generate, configured:
            first = self._ask(user).json()
            self._ask(user, conversationId=first['conversationId'])

        self.assertEqual(retrieve_mock.call_args.kwargs['subject'], 'ریاضی')
        conversation = Conversation.objects.get(pk=first['conversationId'])
        self.assertEqual(conversation.subject, 'ریاضی')

    def test_a_user_without_a_grade_still_gets_an_answer(self):
        user = User.objects.create(phone='09120000014')
        retrieve, generate, configured = self._patched()

        with retrieve as retrieve_mock, generate, configured:
            response = self._ask(user)

        self.assertEqual(response.status_code, 200)
        self.assertEqual(retrieve_mock.call_args.kwargs['subject'], '')


class AskFollowsTheBookTheStudentPickedTests(TestCase):
    """۲۰۲۶-۱۰-۱۱: یک پایه چند کتاب دارد؛ دستیار از کتابی جواب می‌دهد که دانش‌آموز انتخاب کرده یا درس به آن تعلق دارد."""

    def setUp(self):
        cache.clear()
        self.grade = Grade.objects.create(title='پایه دوازدهم', icon='', color='')
        self.user = User.objects.create(phone='09120000021', grade_id=self.grade.id)

    def _ask(self, **extra):
        retrieve = mock.patch('rag.services.retrieval.retrieve_context',
                              return_value=([], {'gradeTitle': 'دوازدهم'}))
        generate = mock.patch('rag.services.generation.generate_answer', return_value=('جواب', [], {}))
        configured = mock.patch('rag.services.config.rag_is_configured', return_value=True)
        with retrieve as retrieve_mock, generate, configured:
            response = self.client.post(
                '/api/ai/ask', {'userQuestion': 'ماتریس چیست؟', **extra}, content_type='application/json',
                HTTP_AUTHORIZATION=f'Bearer {issue_access_token(self.user.id)}',
            )
        self.assertEqual(response.status_code, 200)
        return retrieve_mock.call_args.kwargs

    def test_explicit_subject_without_lesson_gets_the_grade_title(self):
        kwargs = self._ask(subject='هندسه')
        self.assertEqual((kwargs['subject'], kwargs['grade_title']), ('هندسه', 'دوازدهم'))

    def test_lesson_beats_the_subject_stored_on_the_conversation(self):
        from content.models import Lesson, Node, Subject, Unit
        geo = Unit.objects.create(grade=self.grade, subject=Subject.objects.get(key='geometry'),
                                  title='فصل ۱: ماتریس و کاربردها', order=1)
        calc = Unit.objects.create(grade=self.grade, subject=Subject.objects.get(key='calculus'),
                                   title='فصل ۴: مشتق', order=4)
        geo_lesson = Lesson.objects.create(node=Node.objects.create(unit=geo, title='n'), title='l')
        calc_lesson = Lesson.objects.create(node=Node.objects.create(unit=calc, title='n'), title='l')

        first = self.client.post(
            '/api/ai/ask', {'userQuestion': 'سلام', 'lessonId': str(geo_lesson.id)},
            content_type='application/json', HTTP_AUTHORIZATION=f'Bearer {issue_access_token(self.user.id)}',
        )
        self.assertEqual(first.status_code, 200)
        conversation_id = first.json()['conversationId']

        kwargs = self._ask(conversationId=conversation_id, lessonId=str(calc_lesson.id))
        self.assertEqual((kwargs['subject'], kwargs['lesson_number']), ('حسابان', 4))