"""COMPASS utilities for finalizing a run directory"""
import json
import getpass
import logging
from pathlib import Path
import pandas as pd
from elm.version import __version__ as elm_version
from compass import __version__ as compass_version
from compass.utilities.parsing import (
num_ordinances_dataframe,
ordinances_bool_index,
)
logger = logging.getLogger(__name__)
QUALITATIVE_UNITS = "str"
"""str: ``units`` value used to mark a qualitative ordinance row"""
MAX_ORDINANCE_TEXT_CHARS = 5000
"""int: Max characters kept in the ``ordinance_text`` column"""
[docs]
def doc_infos_to_db(doc_infos, output_columns):
"""Aggregate parsed ordinance CSV files into a normalized database
Parameters
----------
doc_infos : Iterable
Iterable of dictionaries describing ordinance extraction
results. Each dictionary must contain ``"ord_db_fp"`` (path to a
parsed CSV), ``"source"`` (document URL), ``"date"`` (tuple of
year, month, day, with ``None`` allowed), and ``"jurisdiction"``
(a :class:`~compass.utilities.jurisdictions.Jurisdiction`
instance).
output_columns : list
List of expected output columns (as
:class:`compass.plugin.interface.OutputColumn` instances)
for the consolidated database. The function enforces the
presence of these columns in the output
:class:`pandas.DataFrame`, filling in missing columns with null
values as needed.
Returns
-------
pandas.DataFrame
Consolidated ordinance dataset.
int
Number of jurisdictions contributing at least one ordinance to
the consolidated dataset.
Notes
-----
Empty or ``None`` entries in ``doc_infos`` are skipped. Ordinance
CSVs that lack parsed values (``num_ordinances_dataframe`` equals
zero) are ignored. The returned DataFrame enforces an ordered column
layout and casts the ``quantitative`` flag to nullable boolean.
"""
db = []
parsed_cols = [col.name for col in output_columns]
for doc_info in doc_infos:
if doc_info is None:
continue
ord_db_fp = doc_info.get("ord_db_fp")
if ord_db_fp is None:
continue
ord_db = pd.read_csv(ord_db_fp, encoding="utf-8-sig")
if num_ordinances_dataframe(ord_db) == 0:
continue
results = _db_results(ord_db, doc_info["jurisdiction"])
results = _formatted_db(results, parsed_cols=parsed_cols)
db.append(results)
if not db:
return pd.DataFrame(columns=parsed_cols), 0
logger.info("Compiling final database for %d jurisdiction(s)", len(db))
num_jurisdictions_found = len(db)
db = pd.concat([df.dropna(axis=1, how="all") for df in db], axis=0)
db = _empirical_adjustments(db)
return _formatted_db(db, parsed_cols=parsed_cols), num_jurisdictions_found
[docs]
def save_db(db, out_dir, output_columns):
"""Write the combined ordinance output to disk
Qualitative and quantitative ordinances are written to a single
``ordinances.csv`` file. The ``quantitative`` flag is retained as a
column so the two kinds of row remain distinguishable.
Parameters
----------
db : pandas.DataFrame
Ordinance dataset containing the full set of output columns,
plus the ``quantitative`` boolean flag.
out_dir : path-like
Directory where ``ordinances.csv`` should be written. The
directory is created by :class:`pathlib.Path` if necessary.
output_columns : list
List of expected output columns (as
:class:`compass.plugin.interface.OutputColumn` instances)
for the consolidated database. The function enforces the
presence of these columns in the output
:class:`pandas.DataFrame`, filling in missing columns with null
values as needed.
Notes
-----
Empty DataFrames short-circuit without creating output files. The
function assumes the boolean ``quantitative`` column has already
been sanitized by :func:`doc_infos_to_db`.
"""
if db.empty:
return
out_cols = [col.name for col in output_columns]
db = trim_ordinance_text(db)
db = _normalize_qualitative_rows(db)
out_dir = Path(out_dir)
db[out_cols].to_csv(
out_dir / "ordinances.csv",
index=False,
encoding="utf-8-sig",
)
[docs]
def trim_ordinance_text(db):
"""Trim the ``ordinance_text`` column to a maximum character count
LLMs occasionally return a very long excerpt for ``ordinance_text``,
which makes the output CSV unwieldy. Any excerpt longer than
:data:`MAX_ORDINANCE_TEXT_CHARS` is cut back to the last whole word
that fits and marked with a trailing ellipsis, matching the ellipsis
convention already used for elided text within an excerpt.
Parameters
----------
db : pandas.DataFrame
The database containing extraction results, which may include
an ``"ordinance_text"`` column.
Returns
-------
pandas.DataFrame
The updated database, with any over-long ``ordinance_text``
entries trimmed. Databases without that column are returned
unchanged.
"""
if db.empty or "ordinance_text" not in db.columns:
return db
db["ordinance_text"] = db["ordinance_text"].apply(_trim_excerpt)
return db
def _trim_excerpt(text):
"""Cut an excerpt to the last whole word within the char limit"""
if not isinstance(text, str) or len(text) <= MAX_ORDINANCE_TEXT_CHARS:
return text
kept = text[:MAX_ORDINANCE_TEXT_CHARS].rsplit(" ", 1)[0].rstrip()
return f"{kept} ..."
def _db_results(results, jurisdiction):
"""Extract results from doc attrs to DataFrame"""
results["FIPS"] = jurisdiction.code
results["county"] = jurisdiction.county
results["state"] = jurisdiction.state
results["subdivision"] = jurisdiction.subdivision_name
results["jurisdiction_type"] = jurisdiction.type
return results
def _empirical_adjustments(db):
"""Post-processing adjustments based on empirical observations
Current adjustments include:
- Limit adder to max of 250 ft.
- Chat GPT likes to report large values here, but in
practice all values manually observed in ordinance documents
are below 250 ft. If large value is detected, assume it's an
error on Chat GPT's part and remove it.
"""
if "adder" in db.columns:
db.loc[db["adder"] > 250, "adder"] = None # ruff:ignore[magic-value-comparison]
return db
def _formatted_db(db, parsed_cols):
"""Format DataFrame for output
``db`` must carry a ``quantitative`` column; every parser sets one
on each row it emits. It is not an output column, but it drives the
qualitative-row conventions applied here.
"""
for col in parsed_cols:
if col not in db.columns:
db[col] = None
db["quantitative"] = db["quantitative"].astype("boolean").fillna(True)
db = _normalize_qualitative_rows(db)
ord_rows = ordinances_bool_index(db)
# "quantitative" is not an output column, but it is needed
# downstream to label qualitative rows, so it rides along here
keep = [*parsed_cols, "quantitative"]
return db[ord_rows][keep].reset_index(drop=True)
def _normalize_qualitative_rows(db):
"""Apply the qualitative-row conventions to a database
Qualitative features have no measurable units, so ``units`` is set
to :data:`QUALITATIVE_UNITS` instead. That lets a reader tell the
two kinds of feature apart with a plain column comparison, since
they share a single output file and the ``quantitative`` flag is
not published.
Their ``summary`` is also copied into ``value``. The schemas ask the
LLM for a summary and a null value on these rows, so filling the
value here means the model only ever writes the requirement once,
while readers still find it in the column they expect.
"""
qualitative = ~db["quantitative"]
db.loc[qualitative, "units"] = QUALITATIVE_UNITS
# a numeric-only "value" column is float-typed, so it has to be
# widened before qualitative text can be written into it
db["value"] = db["value"].astype(object)
db.loc[qualitative, "value"] = db.loc[qualitative, "summary"]
return db
def _extract_model_info_from_all_models(models):
"""Group model info together"""
models_to_tasks = {}
for task, caller_args in models.items():
models_to_tasks.setdefault(caller_args, []).append(task)
return [
{
"name": caller_args.name,
"llm_call_kwargs": {
k: v
for k, v in (caller_args.llm_call_kwargs or {}).items()
if k != "rate_tracker"
},
"llm_service_rate_limit": caller_args.llm_service_rate_limit,
"text_splitter_chunk_size": caller_args.text_splitter_chunk_size,
"text_splitter_chunk_overlap": (
caller_args.text_splitter_chunk_overlap
),
"client_type": caller_args.client_type,
"tasks": tasks,
}
for caller_args, tasks in models_to_tasks.items()
]
[docs]
def compile_run_summary_message(
total_seconds, total_cost, out_dir, document_count
):
"""Create a human-readable summary of a completed run
Parameters
----------
total_seconds : float or int
Duration of the run in seconds.
total_cost : float or int or None
Monetary cost incurred by the run. ``None`` or zero suppresses
the cost line in the summary.
out_dir : path-like
Location of the run output directory. The value is embedded in
the summary text.
document_count : int
Number of documents discovered across all jurisdictions.
Returns
-------
str
Summary string formatted for CLI presentation with ``rich``
markup.
Notes
-----
The function does not perform I/O; callers may log or display the
returned string as needed.
"""
runtime = _elapsed_time_as_str(total_seconds)
total_cost = (
f"\nTotal cost: [#71906e]${total_cost:,.2f}[/#71906e]"
if total_cost
else ""
)
return (
f"✅ Scraping complete!\nOutput Directory: {out_dir}\n"
f"Total runtime: {runtime} {total_cost}\n"
f"Number of jurisdictions with extracted data: {document_count}"
)
[docs]
def compile_collection_summary_message(manifest_fp, collection_manifest):
"""Compile a short collection summary message
Parameters
----------
manifest_fp : path-like
File path where the collection manifest was written. The value
is embedded in the summary text.
collection_manifest : dict
Dictionary describing the collection results, including a list
of jurisdictions and their associated documents. The function
extracts jurisdiction and document counts from the manifest for
inclusion in the summary.
total_seconds : float or int
Duration of the collection phase in seconds, used to report
total runtime in the summary.
Returns
-------
str
Summary string formatted for CLI presentation with ``rich``
markup.
"""
num_jurisdictions = len(collection_manifest.get("jurisdictions", []))
num_documents = sum(
len((info or {}).get("results") or [])
for info in collection_manifest.get("jurisdictions", [])
)
runtime = _elapsed_time_as_str(collection_manifest["total_time"])
locs = "jurisdiction" if num_jurisdictions == 1 else "jurisdictions"
docs = "document" if num_documents == 1 else "documents"
return (
f"✅ Collection complete!\nCollection manifest: {manifest_fp}\n"
f"Total runtime: {runtime}\n"
f"{num_documents:,d} {docs} collected for "
f"{num_jurisdictions:,d} {locs}"
)
def _elapsed_time_as_str(seconds_elapsed):
"""Format elapsed time into human readable string"""
days, seconds = divmod(int(seconds_elapsed), 24 * 3600)
minutes, seconds = divmod(seconds, 60)
hours, minutes = divmod(minutes, 60)
time_str = f"{hours:d}:{minutes:02d}:{seconds:02d}"
if days:
time_str = f"{days:,d} day{'s' if abs(days) != 1 else ''}, {time_str}"
return time_str
