refactor: ♻️ general refactoring

Author: marcvergees

src/file_manipulator.py

@@ -1,7 +1,8 @@
 import os
 from src.filler import Filler
 from src.llm import LLM
-from commonforms import prepare_form 
+from commonforms import prepare_form
+
 
 class FileManipulator:
     def __init__(self):
@@ -13,8 +14,8 @@ def create_template(self, pdf_path: str):
         By using commonforms, we create an editable .pdf template and we store it.
         """
         template_path = pdf_path[:-4] + "_template.pdf"
-        prepare_form(pdf_path, template_path)   
-        return template_path 
+        prepare_form(pdf_path, template_path)
+        return template_path
 
     def fill_form(self, user_input: str, fields: list, pdf_form_path: str):
         """
@@ -23,28 +24,24 @@ def fill_form(self, user_input: str, fields: list, pdf_form_path: str):
         """
         print("[1] Received request from frontend.")
         print(f"[2] PDF template path: {pdf_form_path}")
-        
+
         if not os.path.exists(pdf_form_path):
             print(f"Error: PDF template not found at {pdf_form_path}")
-            return None # Or raise an exception
+            return None  # Or raise an exception
 
         print("[3] Starting extraction and PDF filling process...")
         try:
             self.llm._target_fields = fields
             self.llm._transcript_text = user_input
-            output_name = self.filler.fill_form(
-                pdf_form=pdf_form_path,
-                llm=self.llm
-            )
-            
+            output_name = self.filler.fill_form(pdf_form=pdf_form_path, llm=self.llm)
+
             print("\n----------------------------------")
-            print(f"✅ Process Complete.")
+            print("✅ Process Complete.")
             print(f"Output saved to: {output_name}")
-            
+
             return output_name
-            
+
         except Exception as e:
             print(f"An error occurred during PDF generation: {e}")
             # Re-raise the exception so the frontend can handle it
             raise e
-

src/filler.py

@@ -3,48 +3,50 @@
 from datetime import datetime
 
 
-class Filler():
+class Filler:
     def __init__(self):
         pass
 
-    def fill_form(self, pdf_form: str, llm : LLM):
+    def fill_form(self, pdf_form: str, llm: LLM):
         """
         Fill a PDF form with values from user_input using LLM.
         Fields are filled in the visual order (top-to-bottom, left-to-right).
         """
-        output_pdf = pdf_form[:-4] + "_" + datetime.now().strftime('%Y%m%d_%H%M%S') + "_filled.pdf"
-
-        # Generate dictionary of answers from your original function 
+        output_pdf = (
+            pdf_form[:-4]
+            + "_"
+            + datetime.now().strftime("%Y%m%d_%H%M%S")
+            + "_filled.pdf"
+        )
+
+        # Generate dictionary of answers from your original function
         t2j = llm.main_loop()
         textbox_answers = t2j.get_data()  # This is a dictionary
 
         answers_list = list(textbox_answers.values())
 
-        # Read PDF 
+        # Read PDF
         pdf = PdfReader(pdf_form)
 
-        # Loop through pages 
+        # Loop through pages
         for page in pdf.pages:
             if page.Annots:
                 sorted_annots = sorted(
-                    page.Annots,
-                    key=lambda a: (-float(a.Rect[1]), float(a.Rect[0]))
+                    page.Annots, key=lambda a: (-float(a.Rect[1]), float(a.Rect[0]))
                 )
 
                 i = 0
                 for annot in sorted_annots:
-                    if annot.Subtype == '/Widget' and annot.T:
-                        field_name = annot.T[1:-1]
-                        
+                    if annot.Subtype == "/Widget" and annot.T:
                         if i < len(answers_list):
-                            annot.V = f'{answers_list[i]}'
+                            annot.V = f"{answers_list[i]}"
                             annot.AP = None
                             i += 1
                         else:
                             # Stop if we run out of answers
-                            break 
+                            break
 
         PdfWriter().write(output_pdf, pdf)
-        
+
         # Your main.py expects this function to return the path
         return output_pdf

src/llm.py

@@ -2,27 +2,31 @@
 import os
 import requests
 
-class LLM():
+
+class LLM:
     def __init__(self, transcript_text=None, target_fields=None, json=None):
         if json is None:
             json = {}
-        self._transcript_text = transcript_text # str
-        self._target_fields = target_fields # List, contains the template field.
-        self._json = json # dictionary
-    
+        self._transcript_text = transcript_text  # str
+        self._target_fields = target_fields  # List, contains the template field.
+        self._json = json  # dictionary
+
     def type_check_all(self):
-        if type(self._transcript_text) != str:
-            raise TypeError(f"ERROR in LLM() attributes ->\
-                Transcript must be text. Input:\n\ttranscript_text: {self._transcript_text}")
-        elif type(self._target_fields) != list:  
-            raise TypeError(f"ERROR in LLM() attributes ->\
-                Target fields must be a list. Input:\n\ttarget_fields: {self._target_fields}")
-
-   
+        if type(self._transcript_text) is not str:
+            raise TypeError(
+                f"ERROR in LLM() attributes ->\
+                Transcript must be text. Input:\n\ttranscript_text: {self._transcript_text}"
+            )
+        elif type(self._target_fields) is not list:
+            raise TypeError(
+                f"ERROR in LLM() attributes ->\
+                Target fields must be a list. Input:\n\ttarget_fields: {self._target_fields}"
+            )
+
     def build_prompt(self, current_field):
-        """ 
-            This method is in charge of the prompt engineering. It creates a specific prompt for each target field. 
-            @params: current_field -> represents the current element of the json that is being prompted.
+        """
+        This method is in charge of the prompt engineering. It creates a specific prompt for each target field.
+        @params: current_field -> represents the current element of the json that is being prompted.
         """
         prompt = f""" 
             SYSTEM PROMPT:
@@ -52,7 +56,7 @@ def main_loop(self):
             payload = {
                 "model": "mistral",
                 "prompt": prompt,
-                "stream": False # don't really know why --> look into this later.
+                "stream": False,  # don't really know why --> look into this later.
             }
 
             try:
@@ -68,11 +72,10 @@ def main_loop(self):
 
             # parse response
             json_data = response.json()
-            parsed_response = json_data['response']
+            parsed_response = json_data["response"]
             # print(parsed_response)
             self.add_response_to_json(field, parsed_response)
 
-            
         print("----------------------------------")
         print("\t[LOG] Resulting JSON created from the input text:")
         print(json.dumps(self._json, indent=2))
@@ -81,52 +84,52 @@ def main_loop(self):
         return self
 
     def add_response_to_json(self, field, value):
-        """ 
-            this method adds the following value under the specified field, 
-            or under a new field if the field doesn't exist, to the json dict 
         """
-        value = value.strip().replace('"', '')
+        this method adds the following value under the specified field,
+        or under a new field if the field doesn't exist, to the json dict
+        """
+        value = value.strip().replace('"', "")
         parsed_value = None
-        plural = False
- 
+
         if value != "-1":
-            parsed_value = value       
-        
+            parsed_value = value
+
         if ";" in value:
             parsed_value = self.handle_plural_values(value)
-            plural = True
-
 
         if field in self._json.keys():
             self._json[field].append(parsed_value)
-        else: 
+        else:
             self._json[field] = parsed_value
-                
+
         return
 
     def handle_plural_values(self, plural_value):
-        """ 
-            This method handles plural values.
-            Takes in strings of the form 'value1; value2; value3; ...; valueN' 
-            returns a list with the respective values -> [value1, value2, value3, ..., valueN]
+        """
+        This method handles plural values.
+        Takes in strings of the form 'value1; value2; value3; ...; valueN'
+        returns a list with the respective values -> [value1, value2, value3, ..., valueN]
         """
         if ";" not in plural_value:
-            raise ValueError(f"Value is not plural, doesn't have ; separator, Value: {plural_value}")
-        
-        print(f"\t[LOG]: Formating plural values for JSON, [For input {plural_value}]...")
+            raise ValueError(
+                f"Value is not plural, doesn't have ; separator, Value: {plural_value}"
+            )
+
+        print(
+            f"\t[LOG]: Formating plural values for JSON, [For input {plural_value}]..."
+        )
         values = plural_value.split(";")
-        
+
         # Remove trailing leading whitespace
         for i in range(len(values)):
-            current = i+1 
+            current = i + 1
             if current < len(values):
                 clean_value = values[current].lstrip()
                 values[current] = clean_value
 
         print(f"\t[LOG]: Resulting formatted list of values: {values}")
-        
+
         return values
-        
 
     def get_data(self):
         return self._json