Add MVP of one-shot implementation

compass/_cli/finalize.py

@@ -1,15 +1,14 @@
 """COMPASS CLI finalize subcommand"""
 
-import json
 from datetime import datetime
 
 import click
 from rich.theme import Theme
 from rich.console import Console
 
 from compass.utilities import Directories
+from compass.utilities.io import load_config
 from compass.utilities.jurisdictions import Jurisdiction
-from compass.utilities.parsing import load_config
 from compass.utilities.finalize import save_run_meta, doc_infos_to_db, save_db
 from compass.scripts.process import _initialize_model_params
 
@@ -62,8 +61,7 @@ def finalize(config):
     start_datetime = datetime.fromtimestamp(dirs.out.stat().st_ctime)
     end_datetime = datetime.fromtimestamp(jurisdictions_fp.stat().st_mtime)
 
-    with jurisdictions_fp.open("r", encoding="utf-8") as fh:
-        jurisdictions = json.load(fh)
+    jurisdictions = load_config(jurisdictions_fp)
 
     console.print("Compiling databases...")
     jurisdictions = jurisdictions.get("jurisdictions", [])

compass/_cli/process.py

@@ -12,9 +12,10 @@
 from rich.console import Console
 
 from compass.pb import COMPASS_PB
+from compass.plugin import create_schema_based_one_shot_extraction_plugin
 from compass.scripts.process import process_jurisdictions_with_openai
 from compass.utilities.logs import AddLocationFilter
-from compass.utilities.parsing import load_config
+from compass.utilities.io import load_config
 
 
 @click.command
@@ -41,10 +42,22 @@
     is_flag=True,
     help="Flag to hide progress bars during processing.",
 )
-def process(config, verbose, no_progress):
+@click.option(
+    "--plugin",
+    "-p",
+    required=False,
+    default=None,
+    help="One-shot plugin configuration to add to COMPASS before processing",
+)
+def process(config, verbose, no_progress, plugin):
     """Download and extract ordinances for a list of jurisdictions"""
     config = load_config(config)
 
+    if plugin is not None:
+        create_schema_based_one_shot_extraction_plugin(
+            config=plugin, tech=config["tech"]
+        )
+
     custom_theme = Theme({"logging.level.trace": "rgb(94,79,162)"})
     console = Console(theme=custom_theme)
 

compass/exceptions.py

@@ -17,6 +17,10 @@ def __init__(self, *args, **kwargs):
             )
 
 
+class COMPASSFileNotFoundError(COMPASSError, FileNotFoundError):
+    """COMPASS FileNotFoundError"""
+
+
 class COMPASSNotInitializedError(COMPASSError):
     """COMPASS not initialized error"""
 

compass/extraction/apply.py

@@ -3,7 +3,7 @@
 import logging
 from warnings import warn
 
-from compass.llm import StructuredLLMCaller
+from compass.llm import JSONFromTextLLMCaller
 from compass.extraction.date import DateExtractor
 from compass.validation import (
     ParseChunksWithMemory,
@@ -170,7 +170,7 @@ async def extract_date(doc, model_config, usage_tracker=None):
         )
         return doc
 
-    date_llm_caller = StructuredLLMCaller(
+    date_llm_caller = JSONFromTextLLMCaller(
         llm_service=model_config.llm_service,
         usage_tracker=usage_tracker,
         **model_config.llm_call_kwargs,

compass/extraction/date.py

@@ -37,20 +37,20 @@ class DateExtractor:
     )
     """System message for date extraction LLM calls"""
 
-    def __init__(self, structured_llm_caller, text_splitter=None):
+    def __init__(self, json_llm_caller, text_splitter=None):
         """
 
         Parameters
         ----------
-        structured_llm_caller : StructuredLLMCaller
+        json_llm_caller : JSONFromTextLLMCaller
             Instance used for structured validation queries.
         text_splitter : LCTextSplitter, optional
             Optional text splitter (or subclass instance, or any object
             that implements a `split_text` method) to attach to doc
             (used for splitting out pages in an HTML document).
             By default, ``None``.
         """
-        self.slc = structured_llm_caller
+        self.jlc = json_llm_caller
         self.text_splitter = text_splitter
 
     async def parse(self, doc):
@@ -84,7 +84,7 @@ async def _parse(self, doc):
         )
         if can_check_url_for_date:
             logger.debug("Checking URL for date: %s", url)
-            response = await self.slc.call(
+            response = await self.jlc.call(
                 sys_msg=self.SYSTEM_MESSAGE,
                 content=(
                     "Please extract the date from the URL for this "
@@ -105,7 +105,7 @@ async def _parse(self, doc):
             if not text:
                 continue
 
-            response = await self.slc.call(
+            response = await self.jlc.call(
                 sys_msg=self.SYSTEM_MESSAGE,
                 content=f"Please extract the date for this ordinance:\n{text}",
                 usage_sub_label=LLMUsageCategory.DATE_EXTRACTION,

compass/extraction/small_wind/ordinance.py

@@ -369,12 +369,12 @@ class SmallWindOrdinanceTextExtractor(PromptBasedTextExtractor):
     PROMPTS = [
         {
             "key": "wind_energy_systems_text",
-            "out_fn": "{jurisdiction} Wind Ordinance Text.txt",
+            "out_fn": "{jurisdiction} Wind Ordinance.txt",
             "prompt": _WECS_TEXT_EXTRACTION_PROMPT,
         },
         {
             "key": "cleaned_text_for_extraction",
-            "out_fn": "{jurisdiction} Cleaned Text.txt",
+            "out_fn": "{jurisdiction} Small Wind Ordinance.txt",
             "prompt": _SMALL_WECS_TEXT_EXTRACTION_PROMPT,
         },
     ]
@@ -398,12 +398,12 @@ class SmallWindPermittedUseDistrictsTextExtractor(PromptBasedTextExtractor):
     PROMPTS = [
         {
             "key": "permitted_use_only_text",
-            "out_fn": "{jurisdiction} Permitted Use Only.txt",
+            "out_fn": "{jurisdiction} Permitted Use.txt",
             "prompt": _PERMITTED_USES_TEXT_EXTRACTION_PROMPT,
         },
         {
             "key": "districts_text",
-            "out_fn": "{jurisdiction} Districts.txt",
+            "out_fn": "{jurisdiction} Permitted Use Districts.txt",
             "prompt": _WECS_PERMITTED_USES_TEXT_EXTRACTION_PROMPT,
         },
     ]

compass/extraction/small_wind/parse.py

@@ -242,7 +242,7 @@ class StructuredSmallWindOrdinanceParser(StructuredSmallWindParser):
            a decision-tree-based chain-of-thought prompt on the text for
            each value to be extracted.
     Key Relationships:
-        Uses a StructuredLLMCaller for LLM queries and multiple
+        Uses a JSONFromTextLLMCaller for LLM queries and multiple
         AsyncDecisionTree instances to guide the extraction of
         individual values.
     """
@@ -558,7 +558,7 @@ class StructuredSmallWindPermittedUseDistrictsParser(
            a decision-tree-based chain-of-thought prompt on the text for
            each value to be extracted.
     Key Relationships:
-        Uses a StructuredLLMCaller for LLM queries and multiple
+        Uses a JSONFromTextLLMCaller for LLM queries and multiple
         AsyncDecisionTree instances to guide the extraction of
         individual values.
     """

compass/extraction/small_wind/plugin.py

@@ -20,7 +20,7 @@
     SmallWindPermittedUseDistrictsTextExtractor.OUT_LABEL
 )
 
-SMALL_WIND_QUESTION_TEMPLATES = [
+SMALL_WIND_QUERY_TEMPLATES = [
     "filetype:pdf {jurisdiction} wind energy conversion system ordinances",
     "wind energy conversion system ordinances {jurisdiction}",
     "{jurisdiction} wind WECS ordinance",
@@ -66,8 +66,8 @@ class COMPASSSmallWindExtractor(OrdinanceExtractionPlugin):
     IDENTIFIER = "small wind"
     """str: Identifier for extraction task """
 
-    QUESTION_TEMPLATES = SMALL_WIND_QUESTION_TEMPLATES
-    """list: List of search engine question templates for extraction"""
+    QUERY_TEMPLATES = SMALL_WIND_QUERY_TEMPLATES
+    """list: List of search engine query templates for extraction"""
 
     WEBSITE_KEYWORDS = BEST_SMALL_WIND_ORDINANCE_WEBSITE_URL_KEYWORDS
     """list: List of keywords
@@ -76,8 +76,8 @@ class COMPASSSmallWindExtractor(OrdinanceExtractionPlugin):
     a website scrape for a wind ordinance document.
     """
 
-    heuristic = SmallWindHeuristic()
-    """BaseHeuristic: Object with a ``check()`` method"""
+    HEURISTIC = SmallWindHeuristic
+    """BaseHeuristic: Class with a ``check()`` method"""
 
     TEXT_COLLECTORS = [
         SmallWindOrdinanceTextCollector,

compass/extraction/solar/ordinance.py

@@ -285,7 +285,7 @@ class SolarOrdinanceTextExtractor(PromptBasedTextExtractor):
     PROMPTS = [
         {
             "key": "cleaned_text_for_extraction",
-            "out_fn": "{jurisdiction} Cleaned Text.txt",
+            "out_fn": "{jurisdiction} Utility Scale Solar Ordinance.txt",
             "prompt": _SEF_TEXT_EXTRACTION_PROMPT,
         },
     ]
@@ -309,12 +309,12 @@ class SolarPermittedUseDistrictsTextExtractor(PromptBasedTextExtractor):
     PROMPTS = [
         {
             "key": "permitted_use_only_text",
-            "out_fn": "{jurisdiction} Permitted Use Only.txt",
+            "out_fn": "{jurisdiction} Permitted Use.txt",
             "prompt": _PERMITTED_USES_TEXT_EXTRACTION_PROMPT,
         },
         {
             "key": "districts_text",
-            "out_fn": "{jurisdiction} Districts.txt",
+            "out_fn": "{jurisdiction} Permitted Use Districts.txt",
             "prompt": _SEF_PERMITTED_USES_TEXT_EXTRACTION_PROMPT,
         },
     ]

compass/extraction/solar/parse.py

@@ -189,7 +189,7 @@ class StructuredSolarOrdinanceParser(StructuredSolarParser):
            a decision-tree-based chain-of-thought prompt on the text for
            each value to be extracted.
     Key Relationships:
-        Uses a StructuredLLMCaller for LLM queries and multiple
+        Uses a JSONFromTextLLMCaller for LLM queries and multiple
         AsyncDecisionTree instances to guide the extraction of
         individual values.
     """
@@ -494,7 +494,7 @@ class StructuredSolarPermittedUseDistrictsParser(StructuredSolarParser):
            a decision-tree-based chain-of-thought prompt on the text for
            each value to be extracted.
     Key Relationships:
-        Uses a StructuredLLMCaller for LLM queries and multiple
+        Uses a JSONFromTextLLMCaller for LLM queries and multiple
         AsyncDecisionTree instances to guide the extraction of
         individual values.
     """

compass/extraction/solar/plugin.py

@@ -18,7 +18,7 @@
     SolarPermittedUseDistrictsTextExtractor.OUT_LABEL
 )
 
-SOLAR_QUESTION_TEMPLATES = [
+SOLAR_QUERY_TEMPLATES = [
     "filetype:pdf {jurisdiction} solar energy conversion system ordinances",
     "solar energy conversion system ordinances {jurisdiction}",
     "{jurisdiction} solar energy farm ordinance",
@@ -67,8 +67,8 @@ class COMPASSSolarExtractor(OrdinanceExtractionPlugin):
     IDENTIFIER = "solar"
     """str: Identifier for extraction task """
 
-    QUESTION_TEMPLATES = SOLAR_QUESTION_TEMPLATES
-    """list: List of search engine question templates for extraction"""
+    QUERY_TEMPLATES = SOLAR_QUERY_TEMPLATES
+    """list: List of search engine query templates for extraction"""
 
     WEBSITE_KEYWORDS = BEST_SOLAR_ORDINANCE_WEBSITE_URL_KEYWORDS
     """list: List of keywords
@@ -77,8 +77,8 @@ class COMPASSSolarExtractor(OrdinanceExtractionPlugin):
     a website scrape for a wind ordinance document.
     """
 
-    heuristic = SolarHeuristic()
-    """BaseHeuristic: Object with a ``check()`` method"""
+    HEURISTIC = SolarHeuristic
+    """BaseHeuristic: Class with a ``check()`` method"""
 
     TEXT_COLLECTORS = [
         SolarOrdinanceTextCollector,

compass/extraction/water/plugin.py

@@ -19,7 +19,7 @@
 logger = logging.getLogger(__name__)
 
 
-WATER_RIGHTS_QUESTION_TEMPLATES = [
+WATER_RIGHTS_QUERY_TEMPLATES = [
     "{jurisdiction} rules",
     "{jurisdiction} management plan",
     "{jurisdiction} well permits",
@@ -66,26 +66,40 @@ class TexasWaterRightsExtractor(BaseExtractionPlugin):
     IDENTIFIER = "tx water rights"
     """str: Identifier for extraction task """
 
-    QUESTION_TEMPLATES = WATER_RIGHTS_QUESTION_TEMPLATES
-    """list: List of search engine question templates for extraction"""
-
-    WEBSITE_KEYWORDS = BEST_WATER_RIGHTS_ORDINANCE_WEBSITE_URL_KEYWORDS
-    """list: List of keywords
-
-    Keywords indicate links which should be prioritized when performing
-    a website scrape for a wind ordinance document.
-    """
-
-    heuristic = WaterRightsHeuristic()
-    """BaseHeuristic: Object with a ``check()`` method"""
-
     JURISDICTION_DATA_FP = (
         importlib.resources.files("compass")
         / "data"
         / "tx_water_districts.csv"
     )
     """:term:`path-like <path-like object>`: Path to Texas GCW names"""
 
+    async def get_query_templates(self):  # noqa: PLR6301
+        """Get a list of search engine query templates for extraction
+
+        Query templates can contain the placeholder ``{jurisdiction}``
+        which will be replaced with the full jurisdiction name during
+        the search engine query.
+        """
+        return WATER_RIGHTS_QUERY_TEMPLATES
+
+    async def get_website_keywords(self):  # noqa: PLR6301
+        """Get a dict of website search keyword scores
+
+        Dictionary mapping keywords to scores that indicate links which
+        should be prioritized when performing a website scrape for a
+        document.
+        """
+        return BEST_WATER_RIGHTS_ORDINANCE_WEBSITE_URL_KEYWORDS
+
+    async def get_heuristic(self):  # noqa: PLR6301
+        """Get a `BaseHeuristic` instance with a `check()` method
+
+        The ``check()`` method should accept a string of text and return
+        ``True`` if the text passes the heuristic check and ``False``
+        otherwise.
+        """
+        return WaterRightsHeuristic()
+
     async def filter_docs(
         self,
         extraction_context,

compass/extraction/wind/ordinance.py

@@ -343,12 +343,12 @@ class WindOrdinanceTextExtractor(PromptBasedTextExtractor):
     PROMPTS = [
         {
             "key": "wind_energy_systems_text",
-            "out_fn": "{jurisdiction} Wind Ordinance Text.txt",
+            "out_fn": "{jurisdiction} Wind Ordinance.txt",
             "prompt": _WECS_TEXT_EXTRACTION_PROMPT,
         },
         {
             "key": "cleaned_text_for_extraction",
-            "out_fn": "{jurisdiction} Cleaned Text.txt",
+            "out_fn": "{jurisdiction} Utility Scale Wind Ordinance.txt",
             "prompt": _LARGE_WECS_TEXT_EXTRACTION_PROMPT,
         },
     ]
@@ -372,12 +372,12 @@ class WindPermittedUseDistrictsTextExtractor(PromptBasedTextExtractor):
     PROMPTS = [
         {
             "key": "permitted_use_only_text",
-            "out_fn": "{jurisdiction} Permitted Use Only.txt",
+            "out_fn": "{jurisdiction} Permitted Use.txt",
             "prompt": _PERMITTED_USES_TEXT_EXTRACTION_PROMPT,
         },
         {
             "key": "districts_text",
-            "out_fn": "{jurisdiction} Districts.txt",
+            "out_fn": "{jurisdiction} Permitted Use Districts.txt",
             "prompt": _WECS_PERMITTED_USES_TEXT_EXTRACTION_PROMPT,
         },
     ]

compass/extraction/wind/parse.py

@@ -179,7 +179,7 @@ class StructuredWindOrdinanceParser(StructuredWindParser):
            a decision-tree-based chain-of-thought prompt on the text for
            each value to be extracted.
     Key Relationships:
-        Uses a StructuredLLMCaller for LLM queries and multiple
+        Uses a JSONFromTextLLMCaller for LLM queries and multiple
         AsyncDecisionTree instances to guide the extraction of
         individual values.
     """
@@ -497,7 +497,7 @@ class StructuredWindPermittedUseDistrictsParser(StructuredWindParser):
            a decision-tree-based chain-of-thought prompt on the text for
            each value to be extracted.
     Key Relationships:
-        Uses a StructuredLLMCaller for LLM queries and multiple
+        Uses a JSONFromTextLLMCaller for LLM queries and multiple
         AsyncDecisionTree instances to guide the extraction of
         individual values.
     """