upstream-zulip/zerver/tests/test_message_summary.py

from datetime import datetime, timezone
from unittest import mock

import orjson
import time_machine
from django.conf import settings
from openai.resources.chat.completions import Completions
from openai.types.chat import ChatCompletion
from typing_extensions import override

from analytics.models import UserCount
from zerver.actions.realm_settings import do_change_realm_permission_group_setting
from zerver.lib.test_classes import ZulipTestCase
from zerver.models import NamedUserGroup
from zerver.models.groups import SystemGroups
from zerver.models.realms import get_realm

# Fixture file to store recorded responses
LLM_FIXTURES_FILE = "zerver/tests/fixtures/llm/summary.json"


class MessagesSummaryTestCase(ZulipTestCase):
    @override
    def setUp(self) -> None:
        super().setUp()
        self.user = self.example_user("iago")
        self.topic_name = "New feature launch"
        self.channel_name = "Zulip features"

        self.login_user(self.user)
        self.subscribe(self.user, self.channel_name)
        content = "Zulip just launched a feature to generate summary of messages."
        self.send_stream_message(
            self.user, self.channel_name, content=content, topic_name=self.topic_name
        )

        content = "Sounds awesome! This will **greatly** help me when catching up."
        self.send_stream_message(
            self.user, self.channel_name, content=content, topic_name=self.topic_name
        )

        # Tests fail on the last day of the month due to us capturing the credit usage for that day
        # on the first of the next month, so we need to set the date to a different day.
        not_last_day_of_any_month = datetime(2025, 2, 18, 1, tzinfo=timezone.utc)

        self.mocked_time_patcher = time_machine.travel(not_last_day_of_any_month, tick=False)
        self.mocked_time_patcher.start()
        if settings.GENERATE_LLM_FIXTURES:  # nocoverage
            # Forward calls to the real `create` method via side_effect.
            # autospec=True makes the patched mock honor the descriptor
            # protocol so the action's `client.chat.completions.create(...)`
            # call passes `self` (the action's own Completions instance).
            self.patcher = mock.patch.object(
                Completions, "create", autospec=True, side_effect=Completions.create
            )
            self.mocked_completion = self.patcher.start()

    @override
    def tearDown(self) -> None:
        self.mocked_time_patcher.stop()
        if settings.GENERATE_LLM_FIXTURES:  # nocoverage
            self.patcher.stop()
        super().tearDown()

    def test_summarize_messages_in_topic(self) -> None:
        narrow = orjson.dumps([["channel", self.channel_name], ["topic", self.topic_name]]).decode()

        if settings.GENERATE_LLM_FIXTURES:  # nocoverage
            # NOTE: You need have proper credentials in zproject/dev-secrets.conf
            # to generate the fixtures.
            # Trigger the API call to extract the arguments.
            self.client_get("/json/messages/summary", dict(narrow=narrow))
            call_args = self.mocked_completion.call_args

            # Once we have the arguments, call the original method and save its response.
            response = self.mocked_completion(*call_args.args, **call_args.kwargs)
            with open(LLM_FIXTURES_FILE, "wb") as f:
                fixture_data = {
                    # Only store model and messages.
                    # We don't want to store any secrets.
                    "model": call_args.kwargs["model"],
                    "messages": call_args.kwargs["messages"],
                    "response": response.model_dump(mode="json"),
                }
                f.write(orjson.dumps(fixture_data, option=orjson.OPT_INDENT_2) + b"\n")
            return

        # In this code path, we test using the fixtures.
        with open(LLM_FIXTURES_FILE, "rb") as f:
            fixture_data = orjson.loads(f.read())

        fake_response = ChatCompletion.model_validate(fixture_data["response"])

        # Block summary requests if budget set to 0.
        with self.settings(
            TOPIC_SUMMARIZATION_MODEL="llama-3.3-70b-versatile",
            MAX_PER_USER_MONTHLY_AI_COST=0,
        ):
            response = self.client_get("/json/messages/summary")
            self.assert_json_error_contains(response, "Reached monthly limit for AI credits.")

        # Fake credentials to ensure we crash if actual network
        # requests occur, which would reflect a problem with how the
        # fixtures were set up.
        with self.settings(
            TOPIC_SUMMARIZATION_MODEL="llama-3.3-70b-versatile",
            TOPIC_SUMMARIZATION_API_KEY="test",
        ):
            input_tokens = fixture_data["response"]["usage"]["prompt_tokens"]
            output_tokens = fixture_data["response"]["usage"]["completion_tokens"]
            credits_used = (output_tokens * settings.OUTPUT_COST_PER_GIGATOKEN) + (
                input_tokens * settings.INPUT_COST_PER_GIGATOKEN
            )
            self.assertFalse(
                UserCount.objects.filter(
                    property="ai_credit_usage::day", value=credits_used, user_id=self.user.id
                ).exists()
            )
            with mock.patch.object(Completions, "create", return_value=fake_response):
                payload = self.client_get("/json/messages/summary", dict(narrow=narrow))
                self.assertEqual(payload.status_code, 200)
            # Check that we recorded this usage.
            self.assertTrue(
                UserCount.objects.filter(
                    property="ai_credit_usage::day", value=credits_used, user_id=self.user.id
                ).exists()
            )

        # If we reached the credit usage limit, block summary requests.
        with self.settings(
            TOPIC_SUMMARIZATION_MODEL="llama-3.3-70b-versatile",
            MAX_PER_USER_MONTHLY_AI_COST=credits_used / 1000000000,
        ):
            response = self.client_get("/json/messages/summary")
            self.assert_json_error_contains(response, "Reached monthly limit for AI credits.")

    def test_permission_to_summarize_message_in_topics(self) -> None:
        narrow = orjson.dumps([["channel", self.channel_name], ["topic", self.topic_name]]).decode()

        realm = get_realm("zulip")
        moderators_group = NamedUserGroup.objects.get(
            name=SystemGroups.MODERATORS, realm_for_sharding=realm, is_system_group=True
        )

        do_change_realm_permission_group_setting(
            realm,
            "can_summarize_topics_group",
            moderators_group,
            acting_user=None,
        )

        # In this code path, we test using the fixtures.
        with open(LLM_FIXTURES_FILE, "rb") as f:
            fixture_data = orjson.loads(f.read())

        fake_response = ChatCompletion.model_validate(fixture_data["response"])

        def check_message_summary_permission(user: str, expect_fail: bool = False) -> None:
            self.login(user)
            with (
                self.settings(
                    TOPIC_SUMMARIZATION_MODEL="llama-3.3-70b-versatile",
                    TOPIC_SUMMARIZATION_API_KEY="test",
                ),
                mock.patch.object(Completions, "create", return_value=fake_response),
            ):
                result = self.client_get("/json/messages/summary", dict(narrow=narrow))

            if expect_fail:
                self.assert_json_error(result, "Insufficient permission")
            else:
                self.assert_json_success(result)

        check_message_summary_permission("hamlet", expect_fail=True)
        check_message_summary_permission("shiva")

        nobody_group = NamedUserGroup.objects.get(
            name=SystemGroups.NOBODY, realm_for_sharding=realm, is_system_group=True
        )
        do_change_realm_permission_group_setting(
            realm,
            "can_summarize_topics_group",
            nobody_group,
            acting_user=None,
        )

        check_message_summary_permission("desdemona", expect_fail=True)

        hamletcharacters_group = NamedUserGroup.objects.get(
            name="hamletcharacters", realm_for_sharding=realm
        )
        do_change_realm_permission_group_setting(
            realm,
            "can_summarize_topics_group",
            hamletcharacters_group,
            acting_user=None,
        )

        check_message_summary_permission("desdemona", expect_fail=True)
        check_message_summary_permission("othello", expect_fail=True)
        check_message_summary_permission("hamlet")
        check_message_summary_permission("cordelia")

        setting_group = self.create_or_update_anonymous_group_for_setting(
            [self.example_user("othello")], [moderators_group]
        )
        do_change_realm_permission_group_setting(
            realm,
            "can_summarize_topics_group",
            setting_group,
            acting_user=None,
        )

        check_message_summary_permission("cordelia", expect_fail=True)
        check_message_summary_permission("hamlet", expect_fail=True)
        check_message_summary_permission("othello")
        check_message_summary_permission("shiva")
        check_message_summary_permission("desdemona")