diff --git a/src/llm.py b/src/llm.py index 70937f9..3ed6761 100644 --- a/src/llm.py +++ b/src/llm.py @@ -46,7 +46,7 @@ def build_prompt(self, current_field): def main_loop(self): # self.type_check_all() - for field in self._target_fields.keys(): + for field in self._target_fields: prompt = self.build_prompt(field) # print(prompt) # ollama_url = "http://localhost:11434/api/generate" diff --git a/src/main.py b/src/main.py index 5bb632b..a1bd4f1 100644 --- a/src/main.py +++ b/src/main.py @@ -3,6 +3,9 @@ from commonforms import prepare_form from pypdf import PdfReader from controller import Controller +from typing import Union +import os +import tempfile def input_fields(num_fields: int): fields = [] @@ -68,16 +71,26 @@ def run_pdf_fill_process(user_input: str, definitions: list, pdf_form_path: Unio if __name__ == "__main__": file = "./src/inputs/file.pdf" user_input = "Hi. The employee's name is John Doe. His job title is managing director. His department supervisor is Jane Doe. His phone number is 123456. His email is jdoe@ucsc.edu. The signature is , and the date is 01/02/2005" - fields = ["Employee's name", "Employee's job title", "Employee's department supervisor", "Employee's phone number", "Employee's email", "Signature", "Date"] - prepared_pdf = "temp_outfile.pdf" - prepare_form(file, prepared_pdf) + descriptive_fields = ["Employee's name", "Employee's job title", "Employee's department supervisor", "Employee's phone number", "Employee's email", "Signature", "Date"] + fd, temp_path = tempfile.mkstemp(suffix=".pdf") + os.close(fd) # Close the file descriptor so prepare_form can safely write to it - reader = PdfReader(prepared_pdf) - fields = reader.get_fields() - if(fields): - num_fields = len(fields) - else: - num_fields = 0 + try: + # 2. Use our secure, collision-proof path + prepare_form(file, temp_path) + + reader = PdfReader(temp_path) + fields = reader.get_fields() + + if fields: + num_fields = len(fields) + else: + num_fields = 0 + + controller = Controller() + controller.fill_form(user_input, descriptive_fields, file) - controller = Controller() - controller.fill_form(user_input, fields, file) + finally: + # 3. Guarantee the OS deletes this specific file when we are done + if os.path.exists(temp_path): + os.remove(temp_path)