diff --git a/Makefile b/Makefile index 53eb56a..a0d5555 100644 --- a/Makefile +++ b/Makefile @@ -55,7 +55,7 @@ pull-model: docker compose exec ollama ollama pull mistral test: - docker compose exec app python3 -m pytest src/test/ + docker compose exec app python3 -m pytest tests/ clean: docker compose down -v diff --git a/src/filler.py b/src/filler.py index e31e535..4073799 100644 --- a/src/filler.py +++ b/src/filler.py @@ -39,8 +39,32 @@ def fill_form(self, pdf_form: str, llm: LLM): for annot in sorted_annots: if annot.Subtype == "/Widget" and annot.T: if i < len(answers_list): - annot.V = f"{answers_list[i]}" - annot.AP = None + answer = answers_list[i] + + # Check if the field type is a Button (Checkbox/Radio) + field_type = annot.FT if annot.FT else (annot.Parent.FT if annot.Parent else None) + if str(field_type) == "/Btn": + is_truthy = str(answer).lower() in ["yes", "true", "1", "x", "on"] + + # Find the 'ON' state from the appearance dictionary + on_state = "/Yes" # Default assumption + if annot.AP and annot.AP.N: + keys = [k for k in annot.AP.N.keys() if k != "/Off"] + if keys: + on_state = keys[0] + + if is_truthy: + from pdfrw import PdfName + annot.V = PdfName(on_state.strip("/")) + annot.AS = PdfName(on_state.strip("/")) + else: + from pdfrw import PdfName + annot.V = PdfName("Off") + annot.AS = PdfName("Off") + else: + annot.V = f"{answer}" + annot.AP = None + i += 1 else: # Stop if we run out of answers diff --git a/src/main.py b/src/main.py index 5bb632b..8ed7e9e 100644 --- a/src/main.py +++ b/src/main.py @@ -1,4 +1,5 @@ import os +from typing import Union # from backend import Fill from commonforms import prepare_form from pypdf import PdfReader @@ -30,10 +31,11 @@ def run_pdf_fill_process(user_input: str, definitions: list, pdf_form_path: Unio print("[3] Starting extraction and PDF filling process...") try: - output_name = Fill.fill_form( + controller = Controller() + output_name = controller.fill_form( user_input=user_input, - definitions=definitions, - pdf_form=pdf_form_path + fields=definitions, + pdf_form_path=pdf_form_path ) print("\n----------------------------------")