diff --git a/review_agent/regulatory_info_package/workflow.py b/review_agent/regulatory_info_package/workflow.py index 6a9f05b..37250ba 100644 --- a/review_agent/regulatory_info_package/workflow.py +++ b/review_agent/regulatory_info_package/workflow.py @@ -128,6 +128,7 @@ class RegulatoryInfoPackageWorkflowExecutor: self.batch.status = RegulatoryInfoPackageBatch.Status.SUCCESS self.batch.finished_at = timezone.now() self.batch.save(update_fields=["status", "finished_at"]) + self._append_completion_message() record_event(self.batch, "workflow_completed", {"batch_id": self.batch.pk}) def _nodes(self): @@ -309,6 +310,42 @@ class RegulatoryInfoPackageWorkflowExecutor: ) return + def _append_completion_message(self) -> None: + if ( + Message.objects.filter( + conversation=self.batch.conversation, + role=Message.Role.ASSISTANT, + content__contains=self.batch.batch_no, + ) + .filter(content__contains=self.batch.output_zip_name) + .exists() + ): + return + exports = list( + ExportedSummaryFile.objects.filter( + workflow_type=WORKFLOW_TYPE, + workflow_batch_id=self.batch.pk, + ) + ) + exports = sorted(exports, key=lambda export: 0 if export.export_type == ExportedSummaryFile.ExportType.ZIP else 1) + content = build_assistant_summary( + batch_no=self.batch.batch_no, + exports=[ + { + "file_name": export.file_name, + "download_url": f"/api/review-agent/file-summary/exports/{export.pk}/download/", + "export_type": export.export_type, + } + for export in exports + ], + failed_files=[item for item in self.batch.generated_files if item.get("status") == "failed"], + ) + Message.objects.create( + conversation=self.batch.conversation, + role=Message.Role.ASSISTANT, + content=content, + ) + def _create_export(self, *, path: str, export_type: str, export_category: str) -> ExportedSummaryFile: from pathlib import Path diff --git a/tests/test_regulatory_info_package_workflow.py b/tests/test_regulatory_info_package_workflow.py index fc1331f..4f2b699 100644 --- a/tests/test_regulatory_info_package_workflow.py +++ b/tests/test_regulatory_info_package_workflow.py @@ -1,6 +1,8 @@ +from pathlib import Path + import pytest -from review_agent.models import Conversation, RegulatoryInfoPackageBatch, WorkflowNodeRun +from review_agent.models import Conversation, FileAttachment, Message, RegulatoryInfoPackageBatch, WorkflowNodeRun from review_agent.regulatory_info_package.constants import ( REGULATORY_INFO_PACKAGE_NODE_DEFINITIONS, WORKFLOW_TYPE, @@ -60,3 +62,31 @@ def test_empty_workflow_skeleton_completes(django_user_model, settings): status=WorkflowNodeRun.Status.SUCCESS, ).count() == len(REGULATORY_INFO_PACKAGE_NODE_DEFINITIONS) + +def test_completed_workflow_appends_download_summary_message(django_user_model, settings): + settings.REGULATORY_INFO_PACKAGE_ASYNC = False + user = django_user_model.objects.create_user(username="owner", password="pass") + conversation = Conversation.objects.create(user=user, title="会话") + trigger = Message.objects.create(conversation=conversation, role=Message.Role.USER, content="根据说明书生成第1章监管信息") + source = Path("docs/0.原始材料/目标产品说明书.docx").resolve() + attachment = FileAttachment.objects.create( + conversation=conversation, + user=user, + original_name="目标产品说明书.docx", + storage_path=str(source), + file_size=source.stat().st_size, + ) + batch = create_regulatory_info_package_batch( + conversation=conversation, + user=user, + trigger_message=trigger, + source_attachment=attachment, + source_file_name=attachment.original_name, + source_storage_path=attachment.storage_path, + ) + + start_regulatory_info_package_workflow(batch, async_run=False) + + message = conversation.messages.filter(role=Message.Role.ASSISTANT, content__contains=batch.batch_no).latest("id") + assert "第1章 监管信息(预生成版).zip" in message.content + assert "/api/review-agent/file-summary/exports/" in message.content