Rmp plugin

compass/plugin/one_shot/base.py

@@ -17,6 +17,8 @@
     OrdinanceExtractionPlugin,
     KeywordBasedHeuristic,
 )
+import compass.utilities.finalize as _finalize_default
+import compass.utilities.finalize_rmp as _finalize_rmp
 from compass.plugin.one_shot.generators import (
     generate_query_templates,
     generate_website_keywords,
@@ -223,6 +225,18 @@ class SchemaBasedExtractionPlugin(OrdinanceExtractionPlugin):
         WEBSITE_KEYWORDS = {}  # set by user or LLM-generated
         """dict: Keyword weight mapping for link crawl prioritization"""
 
+        @classmethod
+        def save_structured_data(cls, doc_infos, out_dir):
+            """Write extracted data using finalize module from config"""
+            fin = (
+                _finalize_rmp
+                if config.get("finalize") == "rmp"
+                else _finalize_default
+            )
+            db, num_docs_found = fin.doc_infos_to_db(doc_infos)
+            fin.save_db(db, out_dir)
+            return num_docs_found
+
         async def get_heuristic(self):
             """Get a `BaseHeuristic` instance with a `check()` method
 

compass/plugin/one_shot/components.py

@@ -404,8 +404,12 @@ def _to_dataframe(self, data):
         full_df = full_df.merge(df, on="feature", how="left")
 
         possible_out_cols = [
+            "location",
+            "restriction_type",
+            "geothermal_applicability",
             "value",
             "units",
+            "ammendment",
             "summary",
             "year",
             "section",

compass/services/cpu.py

@@ -273,6 +273,8 @@
 
 def _read_pdf_file(pdf_fp, **kwargs):
     """Utility func so that pdftotext.PDF doesn't have to be pickled"""
+    kwargs.pop("image_to_string_kwargs", None)
+    kwargs.pop("convert_from_bytes_kwargs", None)
     pdf_bytes = Path(pdf_fp).read_bytes()
     pages = read_pdf(pdf_bytes, verbose=False)
     return PDFDocument(pages, **kwargs), pdf_bytes
@@ -283,8 +285,16 @@
     if tesseract_cmd:
         _configure_pytesseract(tesseract_cmd)
 
+    image_to_string_kwargs = kwargs.pop("image_to_string_kwargs", None)
+    convert_from_bytes_kwargs = kwargs.pop("convert_from_bytes_kwargs", None)
+
     pdf_bytes = Path(pdf_fp).read_bytes()
-    pages = read_pdf_ocr(pdf_bytes, verbose=False)
+    pages = read_pdf_ocr(
+        pdf_bytes,
+        verbose=True,
+        image_to_string_kwargs=image_to_string_kwargs,
+        convert_from_bytes_kwargs=convert_from_bytes_kwargs,
+    )
     doc = PDFDocument(_try_decode_ocr_pages(pages), **kwargs)
     doc.attrs["from_ocr"] = True
     return doc, pdf_bytes
@@ -366,9 +376,21 @@
 def _configure_pytesseract(tesseract_cmd):
     """Set the tesseract_cmd"""
     import pytesseract  # noqa: PLC0415
+    from glob import iglob  # noqa: PLC0415
+    from os import remove  # noqa: PLC0415
 
     pytesseract.pytesseract.tesseract_cmd = tesseract_cmd
 
+    # On Windows, Tesseract may still hold the temp PPM file open when
+    # pytesseract's cleanup runs, causing WinError 32. Patch cleanup to
+    # suppress all OSErrors so the OCR result is not lost.
+    def _cleanup_win(temp_name):
+        for filename in iglob(f'{temp_name}*' if temp_name else temp_name):
+            with contextlib.suppress(OSError):
+                remove(filename)
+
+    pytesseract.pytesseract.cleanup = _cleanup_win
+
 
 def _try_decode_ocr_pages(pages):
     """Try to decode pages into strings"""