"""Persistence for search results documents"""
import json
from datetime import UTC, datetime
from functools import partial
from os import PathLike
from pathlib import Path
from statistics import median
from warnings import warn
from elm.version import __version__ as elm_version
from elm.web.search.run import SEARCH_ENGINE_OPTIONS
from compass import __version__ as compass_version
from compass.services.threaded import GenericFuncRunner
from compass.pipeline.targets import normalize_targets
from compass.utilities.io import (
load_config,
normalize_output_stem,
write_text_atomic,
)
from compass.utilities.parsing import convert_paths_to_strings
from compass.exceptions import COMPASSFileNotFoundError, COMPASSValueError
from compass.warn import COMPASSWarning
SEARCH_RESULT_MANIFEST_FILENAME = "search_result_manifest.json"
[docs]
def validate_search_targets(targets, num_urls=5):
"""Validate shard count targets and return search metric accessors
Parameters
----------
targets : dict
Shard metric names and nonnegative integer minimums.
num_urls : int, default=5
Maximum retained URLs per jurisdiction. By default, ``5``.
Returns
-------
dict
Supported metric names mapped to shard value accessors.
"""
metrics = {
"num_results": lambda record: record.get("num_results", 0),
"num_kept_results": _kept_result_count,
}
for option in SEARCH_ENGINE_OPTIONS.values():
label = option.se_class._SE_NAME # ruff:ignore[private-member-access]
metrics[f"search_engine_counts.{label}"] = partial(
_engine_result_count, engine=label
)
for metric, minimum in normalize_targets(targets).items():
if metric not in metrics:
msg = (
f"Unknown search shard target '{metric}'. "
f"Available targets: {', '.join(sorted(metrics))}"
)
raise COMPASSValueError(msg)
if not isinstance(minimum, int) or minimum < 0:
msg = f"Search target '{metric}' requires a nonnegative integer"
raise COMPASSValueError(msg)
if metric == "num_kept_results" and minimum > num_urls:
msg = (
f"Target num_kept_results={minimum} exceeds the configured "
f"retained URL limit ({num_urls})"
)
raise COMPASSValueError(msg)
return metrics
[docs]
async def build_search_manifest(
runtime, jur_results, time_start_utc, failed_targets=None
):
"""Aggregate jurisdiction search results into a report
Parameters
----------
runtime : compass.pipeline.runtime.PipelineRuntime
Runtime containing the search settings and technology.
jur_results : list of dict or None
Results for each requested jurisdiction. Failed workflows may
return ``None``.
time_start_utc : datetime.datetime
UTC timestamp when the search run started.
failed_targets : list of str, optional
List of failed target descriptions/specifications, if any.
By default, ``None``.
Returns
-------
dict
JSON-serializable report containing per-jurisdiction ranked
URLs and filtering reasons.
"""
qt = await runtime.extractor_class(None, None).get_query_templates()
se_kwargs = runtime.search_params.se_kwargs
num_urls = runtime.search_params.num_urls_to_check_per_jurisdiction
num_jurisdictions_searched = len(jur_results)
time_end_utc = datetime.now(UTC)
time_elapsed = time_end_utc - time_start_utc
result_counts = []
filtered_counts = []
se_counts = {}
out_results = []
for results in jur_results:
if results is None:
result_counts.append(0)
filtered_counts.append(0)
continue
out_results.append(results)
result_counts.append(results.get("num_results", 0))
filtered_counts.append(
sum(
row.get("filtered_reason") is None
for row in results["results"]
)
)
for se, count in results["search_engine_counts"].items():
se_counts[se] = se_counts.get(se, 0) + count
return {
"tech": runtime.tech,
"versions": {"compass": compass_version, "elm": elm_version},
"num_urls_requested": num_urls,
"search_engines": list(se_kwargs.get("search_engines", [])),
"query_templates": list(qt),
"time_start_utc": time_start_utc.isoformat(),
"time_end_utc": time_end_utc.isoformat(),
"total_time": time_elapsed.total_seconds(),
"total_time_string": str(time_elapsed),
"num_jurisdictions_searched": num_jurisdictions_searched,
"num_jurisdictions_found": sum(
results.get("num_results", 0) > 0 for results in out_results
),
"search_engine_totals": dict(se_counts),
"result_stats": {
"min": min(result_counts, default=0),
"max": max(result_counts, default=0),
"median": median(result_counts) if result_counts else 0,
"total": sum(result_counts),
},
"filtered_result_stats": {
"min": min(filtered_counts, default=0),
"max": max(filtered_counts, default=0),
"median": median(filtered_counts) if filtered_counts else 0,
"total": sum(filtered_counts),
},
"failed_targets": failed_targets,
"jurisdictions": out_results,
}
[docs]
async def write_search_manifest(manifest_dir, collection_manifest):
"""Write a collection manifest to disk
Parameters
----------
manifest_dir : path-like
Path to the directory where the manifest should be written.
collection_manifest : dict
Dictionary containing collection manifest information to be
serialized and written to disk.
Returns
-------
pathlib.Path
Path to the written manifest file.
"""
return await GenericFuncRunner.call(
_write_search_manifest, manifest_dir, collection_manifest
)
[docs]
async def write_search_manifest_shard(manifest_dir, se_results, jurisdiction):
"""Write one jurisdiction's existing search result shard format
Parameters
----------
out_dir : path-like
Directory for search result shards.
se_results : dict
Search results for one jurisdiction.
jurisdiction : compass.utilities.jurisdictions.Jurisdiction
Jurisdiction represented by the results.
Returns
-------
pathlib.Path
Written shard path.
"""
return await GenericFuncRunner.call(
_write_search_result_shard, manifest_dir, se_results, jurisdiction
)
[docs]
async def reset_search_manifest_shards(manifest_dir, paths):
"""Remove one collection checkpoint and its stale aggregate manifest
Parameters
----------
manifest_dir : path-like
Directory containing the aggregate collection manifest.
paths : iterable of path-like
Search shards to purge.
"""
await GenericFuncRunner.call(_purge_search_shards, manifest_dir, paths)
[docs]
def load_search_result_jurisdictions(manifest_fp, expected_tech):
"""Load saved search results indexed by jurisdiction code
Parameters
----------
manifest_fp : path-like or list of path-like
Manifest, shard, run directory, shard directory, or glob paths.
expected_tech : str
Technology required when the input declares its technology.
Returns
-------
dict
Saved jurisdiction results indexed by their FIPS codes.
"""
if isinstance(manifest_fp, (str, PathLike)):
manifest_fp = [manifest_fp]
jurisdictions = {}
for pattern in manifest_fp:
jurisdictions.update(
_records_from_files(pattern, expected_tech, jurisdictions)
)
return jurisdictions
[docs]
def load_search_shards(out_dir, jurisdictions, expected_tech):
"""Inventory individual requested shards
Parameters
----------
out_dir : path-like
Existing search output directory.
jurisdictions : iterable of Jurisdiction
Jurisdictions requested in the current configuration.
expected_tech : str
Required technology when saved outputs declare one.
Returns
-------
dict
FIPS codes mapped to dictionaries with ``record`` and ``path``.
Missing shards have no entry.
"""
out_dir = Path(out_dir)
requested = {
jurisdiction.code: jurisdiction for jurisdiction in jurisdictions
}
shard_dir = out_dir / "se_results"
expected_paths = _expected_shard_paths(shard_dir, requested)
shards = {}
for path in sorted(shard_dir.glob("*.json")):
record = load_config(path, resolve_paths=False)
_validate_search_result_record(record, path)
code = record["FIPS"]
_validate_code_path_combination(code, path, expected_paths)
if code not in requested:
continue
_validate_search_result_tech(record, expected_tech, path)
_ensure_code_not_duplicate(code, shards)
_validate_search_shard_counts(record, path)
shards[code] = {"record": record, "path": path}
return shards
def _search_result_shard_path(out_dir, jurisdiction):
"""Build the normalized search result shard path for a jur"""
filename = normalize_output_stem(
f"{jurisdiction.full_name} search results"
)
return Path(out_dir) / f"{filename}.json"
def _write_search_manifest(manifest_dir, report):
"""Write a search-only report as JSON"""
manifest_fp = Path(manifest_dir) / SEARCH_RESULT_MANIFEST_FILENAME
payload = json.dumps(report, indent=4, ensure_ascii=False)
write_text_atomic(manifest_fp, payload)
return manifest_fp
def _write_search_result_shard(out_dir, se_results, jurisdiction):
"""Write one jurisdiction's existing search result shard format"""
out_fp = _search_result_shard_path(out_dir, jurisdiction)
write_text_atomic(
out_fp,
json.dumps(convert_paths_to_strings(se_results), indent=4),
)
return out_fp
def _purge_search_shards(manifest_dir, paths):
"""Purge aggregate and selected search shards"""
(Path(manifest_dir) / SEARCH_RESULT_MANIFEST_FILENAME).unlink(
missing_ok=True
)
for path in paths:
path.unlink(missing_ok=True)
def _records_from_files(pattern, expected_tech, jurisdictions):
"""Yield search records from files matching the given pattern"""
for shard_fp in _search_result_paths(Path(pattern).expanduser()):
yield from _records_from_single_file(
shard_fp, expected_tech, jurisdictions
)
def _search_result_paths(path):
"""Select an aggregate manifest or the existing result shards"""
if not path.exists():
return _try_find_from_relative(path)
if not path.is_dir():
return [path]
manifest = path / SEARCH_RESULT_MANIFEST_FILENAME
if manifest.is_file():
return [manifest]
shard_dir = path / "se_results"
if not shard_dir.is_dir():
shard_dir = path
paths = sorted(shard_dir.glob("*.json"))
if not paths:
msg = f"No search result shards found in {shard_dir}"
raise COMPASSFileNotFoundError(msg)
return paths
def _try_find_from_relative(path):
"""Try to find search result files from a relative path"""
base = Path(path.anchor or ".")
pattern = path.relative_to(base) if path.is_absolute() else path
matches = sorted(base.glob(str(pattern)))
if not matches:
msg = f"Search result input not found: {path}"
raise COMPASSFileNotFoundError(msg)
return [
shard_fp
for match in matches
for shard_fp in _search_result_paths(match)
]
def _records_from_single_file(shard_fp, expected_tech, jurisdictions):
"""Extract search result records from a single file"""
payload = load_config(shard_fp, resolve_paths=False)
_validate_search_result_tech(payload, expected_tech, shard_fp)
if payload.get("tech") is None:
msg = (
f"Search result input {shard_fp} has no technology "
"metadata; its technology cannot be verified"
)
warn(msg, COMPASSWarning)
records = payload.get("jurisdictions", [payload])
if not isinstance(records, list):
msg = f"Invalid search result jurisdictions: {shard_fp}"
raise COMPASSValueError(msg)
for record in records:
_validate_search_result_record(record, shard_fp)
_validate_search_result_tech(record, expected_tech, shard_fp)
code = str(record["FIPS"])
if code in jurisdictions:
msg = f"Duplicate search result entry for FIPS '{code}'"
raise COMPASSValueError(msg)
yield code, record
def _validate_search_result_tech(payload, expected_tech, path):
"""Validate declared search result technology"""
if not isinstance(payload, dict):
msg = f"Invalid search result payload: {path}"
raise COMPASSValueError(msg)
if (tech := payload.get("tech")) is not None and tech != expected_tech:
msg = (
f"Search result technology '{tech}' does not match "
f"'{expected_tech}': {path}"
)
raise COMPASSValueError(msg)
def _validate_search_result_record(record, path):
"""Validate replay identity and result structure"""
if _invalid_record(record):
msg = f"Invalid search result FIPS: {path}"
raise COMPASSValueError(msg)
if _invalid_results_list(record.get("results")):
msg = f"Invalid search results: {path}"
raise COMPASSValueError(msg)
def _invalid_record(record):
"""Validate a SE record
Checks:
- record is a dict
- record contains a non-empty "FIPS" string
"""
if not isinstance(record, dict):
return True
if not isinstance(record.get("FIPS"), str):
return True
return not record["FIPS"].strip()
def _invalid_results_list(results):
"""Validate a SE results list
Checks:
- results is a list
- all result dicts within are valid
"""
if not isinstance(results, list):
return True
return any(_invalid_result(result) for result in results)
def _invalid_result(result):
"""Validate a SE result dictionary
Checks:
- result is a dict
- URL exists and is str
- URL is non-empty
- overall_rank is either None or a positive integer
"""
if not isinstance(result, dict):
return True
if not isinstance(result.get("url"), str):
return True
if not result["url"].strip():
return True
return result.get("overall_rank") is not None and (
type(result["overall_rank"]) is not int or result["overall_rank"] < 1
)
def _kept_result_count(record):
"""Read retained counts or derive them for older search shards"""
return record.get(
"num_kept_results",
sum(row.get("filtered_reason") is None for row in record["results"]),
)
def _engine_result_count(record, engine):
"""Read one engine count, treating an absent engine as zero"""
return record.get("search_engine_counts", {}).get(engine, 0)
def _expected_shard_paths(shard_dir, requested):
"""Compile expected shard paths for requested jurisdictions"""
expected_paths = {
_search_result_shard_path(shard_dir, jurisdiction): code
for code, jurisdiction in requested.items()
}
if len(expected_paths) != len(requested):
msg = "Requested jurisdiction shard filenames collide"
raise COMPASSValueError(msg)
return expected_paths
def _validate_code_path_combination(code, path, expected_paths):
"""Ensure the FIPS code matches the expected path"""
if path in expected_paths and code != expected_paths[path]:
msg = f"Search shard FIPS does not match {path}"
raise COMPASSValueError(msg)
def _ensure_code_not_duplicate(code, shards):
"""Ensure the FIPS code is not already present in the shards"""
if code in shards:
msg = f"Duplicate search result entry for FIPS '{code}'"
raise COMPASSValueError(msg)
def _validate_search_shard_counts(record, path):
"""Reject malformed shard counts before any continuation purge"""
counts = record.get("search_engine_counts")
if not isinstance(counts, dict):
msg = f"Invalid search engine counts: {path}"
raise COMPASSValueError(msg)
values = [
record.get("num_results"),
_kept_result_count(record),
*counts.values(),
]
if any(type(value) is not int or value < 0 for value in values):
msg = f"Invalid search shard counts: {path}"
raise COMPASSValueError(msg)
