diff --git a/src-ui/src/app/components/document-detail/document-detail.component.spec.ts b/src-ui/src/app/components/document-detail/document-detail.component.spec.ts index 504ec79ee..bbfc53b29 100644 --- a/src-ui/src/app/components/document-detail/document-detail.component.spec.ts +++ b/src-ui/src/app/components/document-detail/document-detail.component.spec.ts @@ -1381,6 +1381,7 @@ describe('DocumentDetailComponent', () => { it('should get suggestions', () => { const suggestionsSpy = jest.spyOn(documentService, 'getSuggestions') + const aiSuggestionsSpy = jest.spyOn(documentService, 'getAiSuggestions') suggestionsSpy.mockReturnValue( of({ tags: [42, 43], @@ -1391,6 +1392,35 @@ describe('DocumentDetailComponent', () => { ) initNormally() expect(suggestionsSpy).toHaveBeenCalled() + expect(aiSuggestionsSpy).not.toHaveBeenCalled() + expect(component.suggestions).toEqual({ + tags: [42, 43], + suggested_tags: [], + suggested_document_types: [], + suggested_correspondents: [], + }) + }) + + it('should get AI suggestions when AI is enabled', () => { + const getSetting = settingsService.get.bind(settingsService) + jest + .spyOn(settingsService, 'get') + .mockImplementation((key) => + key === SETTINGS_KEYS.AI_ENABLED ? true : getSetting(key) + ) + const suggestionsSpy = jest.spyOn(documentService, 'getSuggestions') + const aiSuggestionsSpy = jest.spyOn(documentService, 'getAiSuggestions') + aiSuggestionsSpy.mockReturnValue( + of({ + tags: [42, 43], + suggested_tags: [], + suggested_document_types: [], + suggested_correspondents: [], + }) + ) + initNormally() + expect(suggestionsSpy).not.toHaveBeenCalled() + expect(aiSuggestionsSpy).toHaveBeenCalled() expect(component.suggestions).toEqual({ tags: [42, 43], suggested_tags: [], diff --git a/src-ui/src/app/components/document-detail/document-detail.component.ts b/src-ui/src/app/components/document-detail/document-detail.component.ts index 91f448056..9c7fd2cad 100644 --- a/src-ui/src/app/components/document-detail/document-detail.component.ts +++ b/src-ui/src/app/components/document-detail/document-detail.component.ts @@ -981,8 +981,10 @@ export class DocumentDetailComponent getSuggestions() { this.suggestionsLoading = true - this.documentsService - .getSuggestions(this.documentId) + const suggestionsObservable = this.aiEnabled + ? this.documentsService.getAiSuggestions(this.documentId) + : this.documentsService.getSuggestions(this.documentId) + suggestionsObservable .pipe( first(), takeUntil(this.unsubscribeNotifier), diff --git a/src-ui/src/app/services/rest/document.service.spec.ts b/src-ui/src/app/services/rest/document.service.spec.ts index 03375e367..d6f5799b1 100644 --- a/src-ui/src/app/services/rest/document.service.spec.ts +++ b/src-ui/src/app/services/rest/document.service.spec.ts @@ -193,6 +193,14 @@ describe(`DocumentService`, () => { expect(req.request.method).toEqual('GET') }) + it('should call appropriate api endpoint for getting AI suggestions', () => { + subscription = service.getAiSuggestions(documents[0].id).subscribe() + const req = httpTestingController.expectOne( + `${environment.apiBaseUrl}${endpoint}/${documents[0].id}/ai_suggestions/` + ) + expect(req.request.method).toEqual('GET') + }) + it('should call appropriate api endpoint for bulk download', () => { const ids = [1, 2, 3] const content = 'both' diff --git a/src-ui/src/app/services/rest/document.service.ts b/src-ui/src/app/services/rest/document.service.ts index cfee4c405..bc87cb1fb 100644 --- a/src-ui/src/app/services/rest/document.service.ts +++ b/src-ui/src/app/services/rest/document.service.ts @@ -404,6 +404,12 @@ export class DocumentService extends AbstractPaperlessService { ) } + getAiSuggestions(id: number): Observable { + return this.http.get( + this.getResourceUrl(id, 'ai_suggestions') + ) + } + getHistory(id: number): Observable { return this.http.get(this.getResourceUrl(id, 'history')) } diff --git a/src/documents/tests/test_api_documents.py b/src/documents/tests/test_api_documents.py index 9b640753f..e7892072a 100644 --- a/src/documents/tests/test_api_documents.py +++ b/src/documents/tests/test_api_documents.py @@ -2145,6 +2145,29 @@ class TestDocumentApi(DirectoriesMixin, ConsumeTaskMixin, APITestCase): response = self.client.get("/api/documents/34676/suggestions/") self.assertEqual(response.status_code, status.HTTP_404_NOT_FOUND) + @mock.patch("documents.views.get_ai_document_classification") + @override_settings(AI_ENABLED=True) + def test_suggestions_still_uses_classifier_when_ai_enabled( + self, + mock_get_ai_classification, + ) -> None: + doc = Document.objects.create(title="test", mime_type="application/pdf") + + response = self.client.get(f"/api/documents/{doc.pk}/suggestions/") + + self.assertEqual(response.status_code, status.HTTP_200_OK) + self.assertEqual( + response.data, + { + "correspondents": [], + "tags": [], + "document_types": [], + "storage_paths": [], + "dates": [], + }, + ) + mock_get_ai_classification.assert_not_called() + @mock.patch("documents.views.match_storage_paths") @mock.patch("documents.views.match_document_types") @mock.patch("documents.views.match_tags") diff --git a/src/documents/tests/test_views.py b/src/documents/tests/test_views.py index 3fe171fd6..4876ec6ab 100644 --- a/src/documents/tests/test_views.py +++ b/src/documents/tests/test_views.py @@ -306,7 +306,7 @@ class TestAISuggestions(DirectoriesMixin, TestCase): AI_ENABLED=True, LLM_BACKEND="mock_backend", ) - def test_suggestions_with_cached_llm( + def test_ai_suggestions_with_cached_llm( self, mock_refresh_cache, mock_get_cache, @@ -314,7 +314,9 @@ class TestAISuggestions(DirectoriesMixin, TestCase): mock_get_cache.return_value = MagicMock(suggestions={"tags": ["tag1", "tag2"]}) self.client.force_login(user=self.user) - response = self.client.get(f"/api/documents/{self.document.pk}/suggestions/") + response = self.client.get( + f"/api/documents/{self.document.pk}/ai_suggestions/", + ) self.assertEqual(response.status_code, status.HTTP_200_OK) self.assertEqual(response.json(), {"tags": ["tag1", "tag2"]}) mock_refresh_cache.assert_called_once_with(self.document.pk) @@ -324,7 +326,7 @@ class TestAISuggestions(DirectoriesMixin, TestCase): AI_ENABLED=True, LLM_BACKEND="mock_backend", ) - def test_suggestions_with_ai_enabled( + def test_ai_suggestions_with_ai_enabled( self, mock_get_ai_classification, ) -> None: @@ -338,7 +340,9 @@ class TestAISuggestions(DirectoriesMixin, TestCase): } self.client.force_login(user=self.user) - response = self.client.get(f"/api/documents/{self.document.pk}/suggestions/") + response = self.client.get( + f"/api/documents/{self.document.pk}/ai_suggestions/", + ) self.assertEqual(response.status_code, status.HTTP_200_OK) self.assertEqual( response.json(), @@ -361,7 +365,7 @@ class TestAISuggestions(DirectoriesMixin, TestCase): AI_ENABLED=True, LLM_BACKEND="openai-like", ) - def test_suggestions_with_invalid_ai_configuration( + def test_ai_suggestions_with_invalid_ai_configuration( self, mock_get_ai_classification, ) -> None: @@ -370,7 +374,9 @@ class TestAISuggestions(DirectoriesMixin, TestCase): ) self.client.force_login(user=self.user) - response = self.client.get(f"/api/documents/{self.document.pk}/suggestions/") + response = self.client.get( + f"/api/documents/{self.document.pk}/ai_suggestions/", + ) self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST) self.assertEqual( diff --git a/src/documents/views.py b/src/documents/views.py index 700780ace..eb0f8b8f0 100644 --- a/src/documents/views.py +++ b/src/documents/views.py @@ -782,6 +782,43 @@ class EmailDocumentDetailSchema(EmailSerializer): 404: None, }, ), + ai_suggestions=extend_schema( + description="View AI suggestions for the document", + responses={ + 200: inline_serializer( + name="AISuggestions", + fields={ + "title": serializers.CharField(allow_null=True), + "correspondents": serializers.ListField( + child=serializers.IntegerField(), + ), + "suggested_correspondents": serializers.ListField( + child=serializers.CharField(), + ), + "tags": serializers.ListField(child=serializers.IntegerField()), + "suggested_tags": serializers.ListField( + child=serializers.CharField(), + ), + "document_types": serializers.ListField( + child=serializers.IntegerField(), + ), + "suggested_document_types": serializers.ListField( + child=serializers.CharField(), + ), + "storage_paths": serializers.ListField( + child=serializers.IntegerField(), + ), + "suggested_storage_paths": serializers.ListField( + child=serializers.CharField(), + ), + "dates": serializers.ListField(child=serializers.CharField()), + }, + ), + 400: None, + 403: None, + 404: None, + }, + ), thumb=extend_schema( description="View the document thumbnail", responses={200: OpenApiTypes.BINARY}, @@ -1308,114 +1345,134 @@ class DocumentViewSet( ): return HttpResponseForbidden("Insufficient permissions") - ai_config = AIConfig() + document_suggestions = get_suggestion_cache(doc.pk) - if ai_config.ai_enabled: - cached_llm_suggestions = get_llm_suggestion_cache( - doc.pk, - backend=ai_config.llm_backend, - ) + if document_suggestions is not None: + refresh_suggestions_cache(doc.pk) + return Response(document_suggestions.suggestions) - if cached_llm_suggestions: - refresh_suggestions_cache(doc.pk) - return Response(cached_llm_suggestions.suggestions) + classifier = load_classifier() - try: - llm_suggestions = get_ai_document_classification(doc, request.user) - except ValueError as exc: - logger.exception( - "Invalid AI configuration while generating suggestions for " - "document %s: %s", - doc.pk, - exc, - exc_info=True, + dates = [] + if settings.NUMBER_OF_SUGGESTED_DATES > 0: + with get_date_parser() as date_parser: + gen = date_parser.parse(doc.filename, doc.content) + dates = sorted( + { + i + for i in itertools.islice( + gen, + settings.NUMBER_OF_SUGGESTED_DATES, + ) + }, ) - raise ValidationError({"ai": [_("Invalid AI configuration.")]}) from exc - matched_tags = match_tags_by_name( + resp_data = { + "correspondents": [ + c.id for c in match_correspondents(doc, classifier, request.user) + ], + "tags": [t.id for t in match_tags(doc, classifier, request.user)], + "document_types": [ + dt.id for dt in match_document_types(doc, classifier, request.user) + ], + "storage_paths": [ + dt.id for dt in match_storage_paths(doc, classifier, request.user) + ], + "dates": [date.strftime("%Y-%m-%d") for date in dates if date is not None], + } + + # Cache the suggestions and the classifier hash for later + set_suggestions_cache(doc.pk, resp_data, classifier) + + return Response(resp_data) + + @action( + methods=["get"], + detail=True, + filter_backends=[], + url_path="ai_suggestions", + ) + @method_decorator(cache_control(no_cache=True)) + def ai_suggestions(self, request, pk=None): + doc = get_object_or_404( + Document.objects.select_related("owner").prefetch_related("versions"), + pk=pk, + ) + if request.user is not None and not has_perms_owner_aware( + request.user, + "view_document", + doc, + ): + return HttpResponseForbidden("Insufficient permissions") + + ai_config = AIConfig() + if not ai_config.ai_enabled: + return HttpResponseBadRequest("AI is required for this feature") + + cached_llm_suggestions = get_llm_suggestion_cache( + doc.pk, + backend=ai_config.llm_backend, + ) + + if cached_llm_suggestions: + refresh_suggestions_cache(doc.pk) + return Response(cached_llm_suggestions.suggestions) + + try: + llm_suggestions = get_ai_document_classification(doc, request.user) + except ValueError as exc: + logger.exception( + "Invalid AI configuration while generating suggestions for " + "document %s: %s", + doc.pk, + exc, + exc_info=True, + ) + raise ValidationError({"ai": [_("Invalid AI configuration.")]}) from exc + + matched_tags = match_tags_by_name( + llm_suggestions.get("tags", []), + request.user, + ) + matched_correspondents = match_correspondents_by_name( + llm_suggestions.get("correspondents", []), + request.user, + ) + matched_types = match_document_types_by_name( + llm_suggestions.get("document_types", []), + request.user, + ) + matched_paths = match_storage_paths_by_name( + llm_suggestions.get("storage_paths", []), + request.user, + ) + + resp_data = { + "title": llm_suggestions.get("title"), + "tags": [t.id for t in matched_tags], + "suggested_tags": extract_unmatched_names( llm_suggestions.get("tags", []), - request.user, - ) - matched_correspondents = match_correspondents_by_name( + matched_tags, + ), + "correspondents": [c.id for c in matched_correspondents], + "suggested_correspondents": extract_unmatched_names( llm_suggestions.get("correspondents", []), - request.user, - ) - matched_types = match_document_types_by_name( + matched_correspondents, + ), + "document_types": [d.id for d in matched_types], + "suggested_document_types": extract_unmatched_names( llm_suggestions.get("document_types", []), - request.user, - ) - matched_paths = match_storage_paths_by_name( + matched_types, + ), + "storage_paths": [s.id for s in matched_paths], + "suggested_storage_paths": extract_unmatched_names( llm_suggestions.get("storage_paths", []), - request.user, - ) + matched_paths, + ), + "dates": llm_suggestions.get("dates", []), + } - resp_data = { - "title": llm_suggestions.get("title"), - "tags": [t.id for t in matched_tags], - "suggested_tags": extract_unmatched_names( - llm_suggestions.get("tags", []), - matched_tags, - ), - "correspondents": [c.id for c in matched_correspondents], - "suggested_correspondents": extract_unmatched_names( - llm_suggestions.get("correspondents", []), - matched_correspondents, - ), - "document_types": [d.id for d in matched_types], - "suggested_document_types": extract_unmatched_names( - llm_suggestions.get("document_types", []), - matched_types, - ), - "storage_paths": [s.id for s in matched_paths], - "suggested_storage_paths": extract_unmatched_names( - llm_suggestions.get("storage_paths", []), - matched_paths, - ), - "dates": llm_suggestions.get("dates", []), - } - - set_llm_suggestions_cache(doc.pk, resp_data, backend=ai_config.llm_backend) - else: - document_suggestions = get_suggestion_cache(doc.pk) - - if document_suggestions is not None: - refresh_suggestions_cache(doc.pk) - return Response(document_suggestions.suggestions) - - classifier = load_classifier() - - dates = [] - if settings.NUMBER_OF_SUGGESTED_DATES > 0: - with get_date_parser() as date_parser: - gen = date_parser.parse(doc.filename, doc.content) - dates = sorted( - { - i - for i in itertools.islice( - gen, - settings.NUMBER_OF_SUGGESTED_DATES, - ) - }, - ) - - resp_data = { - "correspondents": [ - c.id for c in match_correspondents(doc, classifier, request.user) - ], - "tags": [t.id for t in match_tags(doc, classifier, request.user)], - "document_types": [ - dt.id for dt in match_document_types(doc, classifier, request.user) - ], - "storage_paths": [ - dt.id for dt in match_storage_paths(doc, classifier, request.user) - ], - "dates": [ - date.strftime("%Y-%m-%d") for date in dates if date is not None - ], - } - - # Cache the suggestions and the classifier hash for later - set_suggestions_cache(doc.pk, resp_data, classifier) + set_llm_suggestions_cache(doc.pk, resp_data, backend=ai_config.llm_backend) return Response(resp_data)