Merge pull request #301 from ricanontherun/generate-upload-ttle

Generate upload ttle

Author: taichan03

server/api/views/uploadFile/test_title.py

@@ -0,0 +1,67 @@
+import unittest
+from unittest.mock import MagicMock, patch
+
+import title
+
+class TestGenerateTitle(unittest.TestCase):
+  def test_prefers_metadata_title_if_valid(self):
+    doc = MagicMock()
+    doc.metadata = {"title": "A Study Regarding The Efficacy of Drugs"}
+    self.assertEqual("A Study Regarding The Efficacy of Drugs", title.generate_title(doc))
+
+  def test_falls_back_to_first_page_text_if_metadata_title_is_empty(self):
+    doc = MagicMock()
+    doc.metadata = {"title": ""}
+    doc[0].get_text = MagicMock()
+
+    foo_block = [None] * 7
+    foo_block[4] = "foo"
+    foo_block[6] = 0
+
+    title_block = [None] * 7
+    title_block[4] = "Advances in Mood Disorder Pharmacotherapy: Evaluating New Antipsychotics and Mood Stabilizers for Bipolar Disorder and Schizophrenia"
+    title_block[6] = 0
+
+    bar_block = [None] * 7
+    bar_block[4] = "bar"
+    bar_block[6] = 0
+    doc[0].get_text.return_value = [foo_block, title_block, bar_block]
+
+    expected_title = "Advances in Mood Disorder Pharmacotherapy: Evaluating New Antipsychotics and Mood Stabilizers for Bipolar Disorder and Schizophrenia"
+    self.assertEqual(expected_title, title.generate_title(doc))
+
+  def test_falls_back_to_first_page_text_if_metadata_title_does_not_match_regex(self):
+    doc = MagicMock()
+    doc.metadata = {"title": "abcd1234"}
+    doc[0].get_text = MagicMock()
+
+    foo_block = [None] * 7
+    foo_block[4] = "foo"
+    foo_block[6] = 0
+
+    title_block = [None] * 7
+    title_block[4] = "Advances in Mood Disorder Pharmacotherapy: Evaluating New Antipsychotics and Mood Stabilizers for Bipolar Disorder and Schizophrenia"
+    title_block[6] = 0
+
+    bar_block = [None] * 7
+    bar_block[4] = "bar"
+    bar_block[6] = 0
+    doc[0].get_text.return_value = [foo_block, title_block, bar_block]
+
+    expected_title = "Advances in Mood Disorder Pharmacotherapy: Evaluating New Antipsychotics and Mood Stabilizers for Bipolar Disorder and Schizophrenia"
+    self.assertEqual(expected_title, title.generate_title(doc))
+
+  @patch("server.api.services.openai_services.openAIServices.openAI")
+  def test_falls_back_to_chatgpt_if_no_title_found(self, mock_openAI):
+    doc = MagicMock()
+    doc.metadata = {"title": None}
+    doc.get_text.return_value = []
+
+    mock_response = MagicMock()
+    mock_response.choices = [MagicMock()]
+    mock_response.choices[0].message.content = "A Study Regarding The Efficacy of Drugs"
+    mock_openAI.return_value = mock_response
+
+    title.generate_title(doc)
+
+    self.assertTrue(mock_openAI.called)

server/api/views/uploadFile/title.py

@@ -0,0 +1,56 @@
+import re
+
+import fitz
+
+from server.api.services.openai_services import openAIServices
+
+# regular expression to match common research white paper titles. Created by Chat-gpt
+# requires at least 3 words, no dates, no version numbers.
+title_regex = re.compile(r'^(?=(?:\b\w+\b[\s:,\-\(\)]*){3,})(?!.*\b(?:19|20)\d{2}\b)(?!.*\bv\d+\b)[A-Za-z0-9][\w\s:,\-\(\)]*[A-Za-z\)]$', re.IGNORECASE)
+
+def generate_title(pdf: fitz.Document) -> str | None:
+    document_metadata_title = pdf.metadata["title"]
+    if document_metadata_title is not None and document_metadata_title != "":
+        if title_regex.match(document_metadata_title):
+            print("suitable title was found in metadata")
+            return document_metadata_title.strip()
+        else:
+            print("metadata title did not match regex")
+
+    print("Looking for title in first page text")
+    first_page = pdf[0]
+    first_page_blocks = first_page.get_text("blocks")
+    text_blocks = [
+        block[4].strip().replace("\n", " ")
+        for block in first_page_blocks
+        if block[6] == 0 # only include text blocks.
+    ]
+
+    # For some reason, extracted PDF text has extra spaces. Collapse them here.
+    regex = r"\s{2,}"
+    text_blocks = [re.sub(regex, " ", text) for text in text_blocks]
+
+    if len(text_blocks) != 0:
+        for text in text_blocks:
+            if title_regex.match(text):
+                return text
+
+    print("no suitable title found in first page text. Using GPT-4 to summarize the PDF")
+    gpt_title = summarize_pdf(pdf)
+    return gpt_title or None
+
+
+def summarize_pdf(pdf: fitz.Document) -> str:
+    """
+    Summarize a PDF document using OpenAI's GPT-4 model.
+    """
+    first_page = pdf[0]
+    first_page_content = first_page.get_text()
+
+    if first_page_content is None:
+        raise Exception("Failed to read the first page of the PDF file")
+
+    # UploadFile model title is limited to 255 chars.
+    prompt = "Please provide a title for this document. The title should be less than 256 characters and will be displayed on a webpage."
+    response = openAIServices.openAI(first_page_content, prompt, model='gpt-4o', temp=0.0)
+    return response.choices[0].message.content

server/api/views/uploadFile/views.py

@@ -15,6 +15,7 @@
 from ...models.model_embeddings import Embeddings
 import fitz
 from django.db import transaction
+from .title import generate_title
 
 
 @method_decorator(csrf_exempt, name='dispatch')
@@ -77,22 +78,27 @@ def post(self, request, format=None):
                     uploaded_by=request.user,  # Set to the user instance
                     uploaded_by_email=request.user.email  # Also store the email separately
                 )
-                new_file.save()
-
-                if new_file.id is None:
-                    return Response({"message": "Failed to save the upload file."}, status=status.HTTP_500_INTERNAL_SERVER_ERROR)
 
                 with fitz.open(stream=pdf_binary, filetype="pdf") as doc:
                     text = ""
                     page_number = 1  # Initialize page_number
                     page_texts = []  # List to hold text for each page with page number
 
+                    title = generate_title(doc)
+                    if title is not None:
+                        new_file.title = title
+
                     for page in doc:
                         page_text = page.get_text()
                         text += page_text
                         page_texts.append((page_number, page_text))
+
                         page_number += 1
 
+                new_file.save()
+                if new_file.id is None:
+                    return Response({"message": "Failed to save the upload file."}, status=status.HTTP_500_INTERNAL_SERVER_ERROR)
+
                 chunks_with_page = []
 
                 # Create chunks along with their corresponding page number