fix two errors for title test

Author: taichan03

server/api/views/uploadFile/test_title.py

@@ -1,67 +1,69 @@
 import unittest
 from unittest.mock import MagicMock, patch
 
-import title
+from . 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)
+    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("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

@@ -2,11 +2,14 @@
 
 import fitz
 
-from server.api.services.openai_services import openAIServices
+from 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)
+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"]
@@ -23,7 +26,7 @@ def generate_title(pdf: fitz.Document) -> str | None:
     text_blocks = [
         block[4].strip().replace("\n", " ")
         for block in first_page_blocks
-        if block[6] == 0 # only include text blocks.
+        if block[6] == 0  # only include text blocks.
     ]
 
     # For some reason, extracted PDF text has extra spaces. Collapse them here.
@@ -35,7 +38,8 @@ def generate_title(pdf: fitz.Document) -> str | None:
             if title_regex.match(text):
                 return text
 
-    print("no suitable title found in first page text. Using GPT-4 to summarize the PDF")
+    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
 
@@ -52,5 +56,6 @@ def summarize_pdf(pdf: fitz.Document) -> str:
 
     # 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)
+    response = openAIServices.openAI(
+        first_page_content, prompt, model='gpt-4o', temp=0.0)
     return response.choices[0].message.content