"""Collection workflow for the COMPASS pipeline"""
import logging
from functools import cached_property
from elm.web.document import BaseDocument
from compass.pipeline.collection.dedupe import DocumentDeDuplicator
from compass.pipeline.collection.persistence import validate_collection_targets
from compass.pipeline.targets import shard_failures
from compass.pipeline.collection.steps import (
CompassWebsiteCrawlStep,
ElmWebsiteCrawlStep,
KnownLocalDocumentsStep,
KnownUrlDocumentsStep,
SearchEngineDocumentsStep,
SearchResultsCrawlStep,
)
logger = logging.getLogger(__name__)
class _PersistedDocument(BaseDocument):
"""Document subclass used to hold collection artifacts"""
WRITE_KWARGS = None
FILE_EXTENSION = None
def __init__(self, attrs):
super().__init__(pages=[], attrs=attrs)
def _raw_pages(self):
"""Get raw pages from document"""
def _cleaned_text(self):
"""Compute cleaned text from document"""
[docs]
class DocumentCollection:
"""Workflow object that applies a fixed pipeline of steps"""
def __init__(self, workflow):
"""
Parameters
----------
workflow : compass.pipeline.jurisdiction.SingleJurisdictionRun
The workflow for the jurisdiction being processed, which may
or may not have website search enabled. The workflow is
passed to each collection step, which may use it to access
jurisdiction information and other relevant data, and to
determine whether website search is enabled.
"""
self.workflow = workflow
self.de_duplicator = DocumentDeDuplicator()
self._collection_info = {}
self._completed_steps = set()
[docs]
@cached_property
def steps(self):
"""Collection steps in the order they should be executed"""
steps = []
if self.workflow.known_local_docs:
steps.append(KnownLocalDocumentsStep())
else:
logger.debug(
"%r processing has no known local docs configured",
self.workflow.jurisdiction.full_name,
)
if self.workflow.known_doc_urls:
steps.append(KnownUrlDocumentsStep())
else:
logger.debug(
"%r processing has no known URLs configured",
self.workflow.jurisdiction.full_name,
)
if self.workflow.perform_se_search:
steps.append(SearchEngineDocumentsStep())
else:
logger.debug(
"%r processing doesn't have SE search enabled",
self.workflow.jurisdiction.full_name,
)
if self.workflow.perform_search_based_crawl:
steps.append(SearchResultsCrawlStep())
else:
logger.debug(
"%r processing doesn't have SE-based crawl enabled",
self.workflow.jurisdiction.full_name,
)
if self.workflow.perform_website_search:
steps.extend([CompassWebsiteCrawlStep(), ElmWebsiteCrawlStep()])
else:
logger.debug(
"%r processing doesn't have website search enabled",
self.workflow.jurisdiction.full_name,
)
return steps
[docs]
async def execute(self, *, eager_extract=False):
"""Run the fixed collection sequence
The document collection has a well-defined order:
1. Process any/all known local documents
2. Process any/all known document URLs
3. Search engine-based search for ordinance documents
4. Crawl HTML pages among the ranked search results, when
configured
5. Jurisdiction website crawl-based search for ordinance
documents
Users can disable any of these steps via the workflow
configuration.
Parameters
----------
eager_extract : bool, optional
Option to apply extraction as soon as any documents are
found. If the extraction returns any structured data,
subsequent steps are skipped for that jurisdiction.
By default, ``False``.
Returns
-------
dict or None
If ``eager_extract`` is ``False``, a dictionary containing
collection information and metadata. If ``eager_extract`` is
``True``, the result of the extraction workflow if any
structured data was extracted, or ``None`` if no structured
data was extracted from any of the collected documents.
"""
await self._determine_unfinished_steps()
for step in self._unfinished_steps():
docs = await self._run_collection_step(step)
if eager_extract:
context = await self.workflow.extraction.extract_from_docs(
docs
)
if context is not None:
return self._context_with_documented_steps(context)
else:
self._collection_info = (
await self.workflow.write_collection_shard_no_fail(
self.de_duplicator, self._completed_steps
)
)
if eager_extract:
return None
self._log_execute_results()
return self._collection_info
async def _determine_unfinished_steps(self):
"""Get any previously persisted documents and completed steps"""
# Save these before loading persisted state info,
# as loading may overwrite them
user_configured_budget = self.workflow.num_search_results_to_crawl
user_configured_website = self.workflow.jurisdiction_website
await self._load_persisted_state_info()
await self._check_targets_for_completion(
user_configured_budget, user_configured_website
)
self._load_persisted_docs()
async def _load_persisted_state_info(self):
"""Load persisted state information for the collection shard"""
self._collection_info = (
await self.workflow.load_existing_collection_shard()
) or {}
self._set_num_search_results_to_crawl()
self._completed_steps |= set(
self._collection_info.get("completed_step_document_counts", {})
)
def _set_num_search_results_to_crawl(self):
"""Set the number of results to crawl based on persisted info"""
num_to_crawl = self._collection_info.get("num_search_results_to_crawl")
if num_to_crawl is not None:
self.workflow.num_search_results_to_crawl = max(
0, int(num_to_crawl)
)
async def _check_targets_for_completion(
self, user_configured_budget, user_configured_website
):
"""Check if the collection targets are met"""
re_run_targets = self.workflow.runtime.re_run_targets
if not re_run_targets or not self._collection_info:
return
unfinished = any(
step.STEP_NAME not in self._completed_steps for step in self.steps
)
if unfinished:
logger.info(
"Resuming unfinished collection steps for %s",
self.workflow.jurisdiction.full_name,
)
return
metrics = validate_collection_targets(re_run_targets)
failures = shard_failures(
self._collection_info, re_run_targets, metrics
)
if not failures:
logger.info(
"Reusing collection shard for %s",
self.workflow.jurisdiction.full_name,
)
return
logger.info(
"Resetting collection shard for %s: %s",
self.workflow.jurisdiction.full_name,
failures,
)
await self._reset_collection_state(
user_configured_budget, user_configured_website
)
async def _reset_collection_state(
self, user_configured_budget, user_configured_website
):
"""Reset the collection state to start fresh"""
await self.workflow.reset_collection_shard()
self.workflow.num_search_results_to_crawl = user_configured_budget
self.workflow.jurisdiction_website = user_configured_website
self.__dict__.pop("steps", None) # reset @cached_property
self._collection_info = {}
self._completed_steps.clear()
self.de_duplicator = DocumentDeDuplicator()
def _load_persisted_docs(self):
"""Populate the de-duplicator with persisted documents"""
docs = [
_PersistedDocument(doc_info)
for doc_info in self._collection_info.get("results", [])
]
self.de_duplicator.add_docs(docs)
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 _context_with_documented_steps(self, context):
"""Attach collection steps to each document in the context"""
for doc in context.data_docs:
doc.attrs["from_steps"] = list(
self.de_duplicator.info(doc).from_steps
)
return context
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",
self.workflow.jurisdiction.full_name,
"\n\n".join(
[f"{info.doc!r}" for info in self.de_duplicator.values()]
),
)
else:
logger.debug(
"No documents were collected for %s",
self.workflow.jurisdiction.full_name,
)
