From 227ca64e13229a56c6b17e09173403cf5aaf107a Mon Sep 17 00:00:00 2001 From: -LAN- Date: Thu, 6 Nov 2025 16:26:09 +0800 Subject: [PATCH] refactor: enforce non-null message annotation questions --- ...0_make_message_annotation_question_not_.py | 60 +++++++++++++++++++ api/models/model.py | 2 +- api/services/annotation_service.py | 14 ++++- .../services/test_annotation_service.py | 17 ++++++ 4 files changed, 89 insertions(+), 4 deletions(-) create mode 100644 api/migrations/versions/2025_11_06_1603-9e6fa5cbcd80_make_message_annotation_question_not_.py diff --git a/api/migrations/versions/2025_11_06_1603-9e6fa5cbcd80_make_message_annotation_question_not_.py b/api/migrations/versions/2025_11_06_1603-9e6fa5cbcd80_make_message_annotation_question_not_.py new file mode 100644 index 0000000000..57c8d71fcb --- /dev/null +++ b/api/migrations/versions/2025_11_06_1603-9e6fa5cbcd80_make_message_annotation_question_not_.py @@ -0,0 +1,60 @@ +"""make message annotation question not nullable + +Revision ID: 9e6fa5cbcd80 +Revises: 03f8dcbc611e +Create Date: 2025-11-06 16:03:54.549378 + +""" +from alembic import op +import sqlalchemy as sa + + +# revision identifiers, used by Alembic. +revision = '9e6fa5cbcd80' +down_revision = '03f8dcbc611e' +branch_labels = None +depends_on = None + + +def upgrade(): + bind = op.get_bind() + message_annotations = sa.table( + "message_annotations", + sa.column("id", sa.String), + sa.column("message_id", sa.String), + sa.column("question", sa.Text), + ) + messages = sa.table( + "messages", + sa.column("id", sa.String), + sa.column("query", sa.Text), + ) + update_question_from_message = ( + sa.update(message_annotations) + .where( + sa.and_( + message_annotations.c.question.is_(None), + message_annotations.c.message_id.isnot(None), + ) + ) + .values( + question=sa.select(sa.func.coalesce(messages.c.query, "")) + .where(messages.c.id == message_annotations.c.message_id) + .scalar_subquery() + ) + ) + bind.execute(update_question_from_message) + + fill_remaining_questions = ( + sa.update(message_annotations) + .where(message_annotations.c.question.is_(None)) + .values(question="") + ) + bind.execute(fill_remaining_questions) + with op.batch_alter_table('message_annotations', schema=None) as batch_op: + batch_op.alter_column('question', existing_type=sa.TEXT(), nullable=False) + + +def downgrade(): + with op.batch_alter_table('message_annotations', schema=None) as batch_op: + batch_op.alter_column('question', existing_type=sa.TEXT(), nullable=True) diff --git a/api/models/model.py b/api/models/model.py index 9b7db259dd..bc96ac089c 100644 --- a/api/models/model.py +++ b/api/models/model.py @@ -1373,7 +1373,7 @@ class MessageAnnotation(Base): app_id: Mapped[str] = mapped_column(StringUUID) conversation_id: Mapped[str | None] = mapped_column(StringUUID, sa.ForeignKey("conversations.id")) message_id: Mapped[str | None] = mapped_column(StringUUID) - question = mapped_column(sa.Text, nullable=True) + question: Mapped[str] = mapped_column(sa.Text, nullable=False) content = mapped_column(sa.Text, nullable=False) hit_count: Mapped[int] = mapped_column(sa.Integer, nullable=False, server_default=sa.text("0")) account_id = mapped_column(StringUUID, nullable=False) diff --git a/api/services/annotation_service.py b/api/services/annotation_service.py index 9258def907..95f1e4d879 100644 --- a/api/services/annotation_service.py +++ b/api/services/annotation_service.py @@ -186,8 +186,12 @@ class AppAnnotationService: if not app: raise NotFound("App not found") + question = args.get("question") + if question is None: + raise ValueError("'question' is required") + annotation = MessageAnnotation( - app_id=app.id, content=args["answer"], question=args["question"], account_id=current_user.id + app_id=app.id, content=args["answer"], question=question, account_id=current_user.id ) db.session.add(annotation) db.session.commit() @@ -196,7 +200,7 @@ class AppAnnotationService: if annotation_setting: add_annotation_to_index_task.delay( annotation.id, - args["question"], + question, current_tenant_id, app_id, annotation_setting.collection_binding_id, @@ -221,8 +225,12 @@ class AppAnnotationService: if not annotation: raise NotFound("Annotation not found") + question = args.get("question") + if question is None: + raise ValueError("'question' is required") + annotation.content = args["answer"] - annotation.question = args["question"] + annotation.question = question db.session.commit() # if annotation reply is enabled , add annotation to index diff --git a/api/tests/test_containers_integration_tests/services/test_annotation_service.py b/api/tests/test_containers_integration_tests/services/test_annotation_service.py index 2b03ec1c26..c2144c1f77 100644 --- a/api/tests/test_containers_integration_tests/services/test_annotation_service.py +++ b/api/tests/test_containers_integration_tests/services/test_annotation_service.py @@ -220,6 +220,23 @@ class TestAnnotationService: # Note: In this test, no annotation setting exists, so task should not be called mock_external_service_dependencies["add_task"].delay.assert_not_called() + def test_insert_app_annotation_directly_requires_question( + self, db_session_with_containers, mock_external_service_dependencies + ): + """ + Question must be provided when inserting annotations directly. + """ + fake = Faker() + app, _ = self._create_test_app_and_account(db_session_with_containers, mock_external_service_dependencies) + + annotation_args = { + "question": None, + "answer": fake.text(max_nb_chars=200), + } + + with pytest.raises(ValueError): + AppAnnotationService.insert_app_annotation_directly(annotation_args, app.id) + def test_insert_app_annotation_directly_app_not_found( self, db_session_with_containers, mock_external_service_dependencies ):