Add complexipy

.github/workflows/ci-python.yml

@@ -28,20 +28,30 @@ on:
 
 jobs:
   lint:
-    name: Lint Python Code Base with Ruff
+    name: Lint Python Code Base
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
-      - uses: astral-sh/ruff-action@278981a28ce3188b1e39527901f38254bf3aac89 # v4.1.0
+      - name: Lint Python Code (Ruff)
+        uses: astral-sh/ruff-action@278981a28ce3188b1e39527901f38254bf3aac89 # v4.1.0
         with:
           version: "latest"
           args: "check"
           src: "./compass"
-      - uses: astral-sh/ruff-action@278981a28ce3188b1e39527901f38254bf3aac89 # v4.1.0
+      - name: Check Python Code Format (Ruff)
+        uses: astral-sh/ruff-action@278981a28ce3188b1e39527901f38254bf3aac89 # v4.1.0
         with:
           version: "latest"
           args: "format --check"
           src: "./compass"
+      - name: Check Python Code Complexity (Complexipy)
+        uses: rohaquinlop/complexipy-action@e2b05bcc06d899a24e2b6bb8b1354ac42800a95e # v7.0.1
+        with:
+          paths: "./compass"
+          max_complexity_allowed: 10
+          failed: false  # true
+          sort: desc
+          ignore_complexity: false  # Set to true to ignore complexity checks
 
   locked-tests:
     needs: lint

compass/common/base.py

@@ -417,6 +417,7 @@ def setup_participating_owner(**kwargs):
     return G
 
 
+# complexipy: ignore
 def setup_graph_extra_restriction(is_numerical=True, **kwargs):
     """Setup Graph to extract non-setback ordinance values from text
 
@@ -434,9 +435,8 @@ def setup_graph_extra_restriction(is_numerical=True, **kwargs):
     kwargs.setdefault("unit_clarification", "")
     kwargs.setdefault("feature_clarifications", "")
     feature_id = kwargs.get("feature_id", "")
-    G = setup_graph_no_nodes(  # ruff:ignore[non-lowercase-variable-in-function]
-        d_tree_name="Extra restriction", **kwargs
-    )
+    # ruff:ignore[non-lowercase-variable-in-function]
+    G = setup_graph_no_nodes(d_tree_name="Extra restriction", **kwargs)
 
     G.add_node(
         "init",

compass/llm/config.py

@@ -1,6 +1,7 @@
 """Ordinances LLM Configurations"""
 
 import os
+import logging
 from collections import Counter
 from abc import ABC, abstractmethod
 from functools import partial, cached_property
@@ -14,6 +15,9 @@
 from compass.exceptions import COMPASSValueError
 
 
+logger = logging.getLogger(__name__)
+
+
 class _PrintableRecursiveCharacterTextSplitter(RecursiveCharacterTextSplitter):
     """RecursiveCharacterTextSplitter with __str__ method"""
 
@@ -86,7 +90,7 @@ def text_splitter(self):
             RTS_SEPARATORS,
             chunk_size=self.text_splitter_chunk_size,
             chunk_overlap=self.text_splitter_chunk_overlap,
-            length_function=partial(ApiBase.count_tokens, model=self.name),
+            length_function=partial(_count_tokens_safely, model=self.name),
             is_separator_regex=True,
         )
 
@@ -204,23 +208,23 @@ def _validate_tag(self):
     @cached_property
     def client_kwargs(self):
         """dict: Parameters to pass to client initializer"""
+
+        arg_env_pairs = []
         if self.client_type == "azure":
             arg_env_pairs = [
                 ("api_key", "AZURE_OPENAI_API_KEY"),
                 ("api_version", "AZURE_OPENAI_VERSION"),
                 ("azure_endpoint", "AZURE_OPENAI_ENDPOINT"),
             ]
-            for key, env_var in arg_env_pairs:
-                if self._client_kwargs.get(key) is None:
-                    self._client_kwargs[key] = os.environ.get(env_var)
         elif self.client_type == "openai":
             arg_env_pairs = [
                 ("api_key", "OPENAI_API_KEY"),
                 ("base_url", "OPENAI_BASE_URL"),
             ]
-            for key, env_var in arg_env_pairs:
-                if self._client_kwargs.get(key) is None:
-                    self._client_kwargs[key] = os.environ.get(env_var)
+
+        for key, env_var in arg_env_pairs:
+            val = self._client_kwargs.get(key)
+            self._client_kwargs[key] = val or os.environ.get(env_var)
 
         return self._client_kwargs
 
@@ -234,3 +238,23 @@ def llm_service(self):
             rate_limit=self.llm_service_rate_limit,
             service_tag=self._tag,
         )
+
+
+def _count_tokens_safely(text, model):
+    """Count tokens with a conservative fallback
+
+    ``len(text.encode("utf-8"))`` is a conservative upper bound on BPE
+    token count, so it will not permit oversized chunks
+    """
+    try:
+        return ApiBase.count_tokens(text, model=model)
+    except ValueError as err:
+        if "Max stack size exceeded for backtracking" not in str(err):
+            raise
+
+        logger.warning(
+            "Using byte length after tokenizer backtracking failure "
+            "for %d characters",
+            len(text),
+        )
+        return len(text.encode("utf-8"))

compass/pipeline/collection/base.py

@@ -50,6 +50,8 @@ def __init__(self, workflow):
         """
         self.workflow = workflow
         self.de_duplicator = DocumentDeDuplicator()
+        self._collection_info = {}
+        self._completed_steps = set()
 
     @cached_property
     def steps(self):
@@ -121,22 +123,9 @@ async def execute(self, *, eager_extract=False):
             structured data was extracted, or ``None`` if no structured
             data was extracted from any of the collected documents.
         """
-        collection_info = await self._load_persisted_docs()
-        completed_steps = set(
-            collection_info.get("completed_step_document_counts", {})
-        )
-        for step in self.steps:
-            if step.STEP_NAME in completed_steps:
-                logger.info(
-                    "Skipping completed collection step %s for %s",
-                    step.STEP_NAME,
-                    self.workflow.jurisdiction.full_name,
-                )
-                continue
-
-            docs = await step.collect(self.workflow)
-            self.de_duplicator.add_docs(docs, step_name=str(step.STEP_NAME))
-            completed_steps.add(step.STEP_NAME)
+        await self._load_persisted_docs()
+        for step in self._unfinished_steps():
+            docs = await self._run_collection_step(step)
             if eager_extract:
                 context = (
                     await self.workflow.extraction_workflow.extract_from_docs(
@@ -146,15 +135,55 @@ async def execute(self, *, eager_extract=False):
                 if context is not None:
                     return context
             else:
-                collection_info = (
+                self._collection_info = (
                     await self.workflow.write_collection_shard_no_fail(
-                        self.de_duplicator, completed_steps
+                        self.de_duplicator, self._completed_steps
                     )
                 )
 
         if eager_extract:
             return None
 
+        self._log_execute_results()
+        return self._collection_info
+
+    async def _load_persisted_docs(self):
+        """Get any previously persisted documents and completed steps"""
+        self._collection_info = (
+            await self.workflow.load_existing_collection_shard()
+        ) or {}
+
+        docs = [
+            _PersistedDocument(doc_info)
+            for doc_info in self._collection_info.get("documents", [])
+        ]
+        self.de_duplicator.add_docs(docs)
+
+        self._completed_steps |= set(
+            self._collection_info.get("completed_step_document_counts", {})
+        )
+
+    def _unfinished_steps(self):
+        """Yield unfinished collection steps"""
+        for step in self.steps:
+            if step.STEP_NAME in self._completed_steps:
+                logger.info(
+                    "Skipping completed collection step %s for %s",
+                    step.STEP_NAME,
+                    self.workflow.jurisdiction.full_name,
+                )
+                continue
+            yield step
+
+    async def _run_collection_step(self, step):
+        """Run collection step and record results"""
+        docs = await step.collect(self.workflow)
+        self.de_duplicator.add_docs(docs, step_name=str(step.STEP_NAME))
+        self._completed_steps.add(step.STEP_NAME)
+        return docs
+
+    def _log_execute_results(self):
+        """Log the results of the collection execution"""
         if self.de_duplicator:
             logger.debug(
                 "Collected the following documents for %s:\n\n%s",
@@ -168,17 +197,3 @@ async def execute(self, *, eager_extract=False):
                 "No documents were collected for %s",
                 self.workflow.jurisdiction.full_name,
             )
-
-        return collection_info
-
-    async def _load_persisted_docs(self):
-        """Get any previously persisted documents and completed steps"""
-        existing_collection_info = (
-            await self.workflow.load_existing_collection_shard()
-        ) or {}
-        docs = [
-            _PersistedDocument(doc_info)
-            for doc_info in existing_collection_info.get("documents", [])
-        ]
-        self.de_duplicator.add_docs(docs)
-        return existing_collection_info

compass/pipeline/collection/persistence.py

@@ -166,6 +166,24 @@ async def load_collection_manifest_jurisdictions(manifest_fp, expected_tech):
     if isinstance(manifest_fp, (str, os.PathLike)):
         manifest_fp = [str(manifest_fp)]
 
+    manifests = await _load_jur_manifests(manifest_fp, expected_tech)
+
+    jurisdictions_by_fips = {}
+    for jurisdiction in chain.from_iterable(
+        manifest.get("jurisdictions", []) for manifest in manifests
+    ):
+        if jurisdiction is None:
+            continue
+
+        fips = jurisdiction.get("FIPS")
+        _validate_not_duplicate_jurisdiction(fips, jurisdictions_by_fips)
+        jurisdictions_by_fips[fips] = jurisdiction
+
+    return jurisdictions_by_fips
+
+
+async def _load_jur_manifests(manifest_fp, expected_tech):
+    """Load one or more collection manifest(s) for jurisdictions"""
     task_fps = []
     for maybe_glob in manifest_fp:
         # ruff: ignore[glob]
@@ -178,22 +196,14 @@ async def load_collection_manifest_jurisdictions(manifest_fp, expected_tech):
         GenericFuncRunner.call(_load_collection_manifest, fp, expected_tech)
         for fp in task_fps
     ]
-    manifests = await asyncio.gather(*tasks)
-
-    jurisdictions_by_fips = {}
-    for jurisdiction in chain.from_iterable(
-        manifest.get("jurisdictions", []) for manifest in manifests
-    ):
-        if jurisdiction is None:
-            continue
+    return await asyncio.gather(*tasks)
 
-        fips = jurisdiction.get("FIPS")
-        if fips in jurisdictions_by_fips:
-            msg = f"Duplicate collection manifest entry for FIPS '{fips}'"
-            raise COMPASSValueError(msg)
-        jurisdictions_by_fips[fips] = jurisdiction
 
-    return jurisdictions_by_fips
+def _validate_not_duplicate_jurisdiction(fips, jurisdictions_by_fips):
+    """Validate that a jurisdiction is not duplicated in the manifest"""
+    if fips in jurisdictions_by_fips:
+        msg = f"Duplicate collection manifest entry for FIPS '{fips}'"
+        raise COMPASSValueError(msg)
 
 
 async def load_specific_collection_manifest_shard(shard_dir, jurisdiction):

compass/pipeline/data_classes.py

@@ -1290,22 +1290,39 @@ def build_models(user_input, *, allow_empty=False):
 
     caller_instances = {}
     for raw_kwargs in user_input:
-        kwargs = dict(raw_kwargs)
-        tasks = kwargs.pop("tasks", LLMTasks.DEFAULT)
-        if isinstance(tasks, str):
-            tasks = [tasks]
-
-        model_config = OpenAIConfig(**kwargs)
-        for task in tasks:
-            if task in caller_instances:
-                msg = (
-                    f"Found duplicated task: {task!r}. Please ensure "
-                    "each LLM caller definition has uniquely-assigned "
-                    "tasks."
-                )
-                raise COMPASSValueError(msg)
+        for task, model_config in _config_for_tasks(raw_kwargs):
+            _verify_task_not_duplicate(task, caller_instances)
             caller_instances[task] = model_config
 
+    _verify_default_case_handled(caller_instances, allow_empty)
+    return caller_instances
+
+
+def _config_for_tasks(kwargs):
+    """Yield (task, model_config) pairs for the given raw kwargs"""
+    kwargs = dict(kwargs)
+    tasks = kwargs.pop("tasks", LLMTasks.DEFAULT)
+    if isinstance(tasks, str):
+        tasks = [tasks]
+
+    model_config = OpenAIConfig(**kwargs)
+    for task in tasks:
+        yield task, model_config
+
+
+def _verify_task_not_duplicate(task, caller_instances):
+    """Verify that the given task has not already been defined"""
+    if task in caller_instances:
+        msg = (
+            f"Found duplicated task: {task!r}. Please ensure "
+            "each LLM caller definition has uniquely-assigned "
+            "tasks."
+        )
+        raise COMPASSValueError(msg)
+
+
+def _verify_default_case_handled(caller_instances, allow_empty):
+    """Verify that the default LLM task is handled correctly"""
     if not allow_empty and LLMTasks.DEFAULT not in caller_instances:
         msg = (
             "No 'default' LLM caller defined in the `model` portion "
@@ -1314,5 +1331,3 @@ def build_models(user_input, *, allow_empty=False):
             f"unspecified. Found tasks: {list(caller_instances)}"
         )
         raise COMPASSValueError(msg)
-
-    return caller_instances