diff --git a/.github/workflows/ci-python.yml b/.github/workflows/ci-python.yml index 446925384..a683bfdd1 100644 --- a/.github/workflows/ci-python.yml +++ b/.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 diff --git a/compass/common/base.py b/compass/common/base.py index 74c3356ca..e0b752451 100644 --- a/compass/common/base.py +++ b/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", diff --git a/compass/llm/config.py b/compass/llm/config.py index bce90750b..467f9aa70 100644 --- a/compass/llm/config.py +++ b/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")) diff --git a/compass/pipeline/collection/base.py b/compass/pipeline/collection/base.py index e271147b0..5b236ab26 100644 --- a/compass/pipeline/collection/base.py +++ b/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 diff --git a/compass/pipeline/collection/persistence.py b/compass/pipeline/collection/persistence.py index 420428429..36cc8ef6f 100644 --- a/compass/pipeline/collection/persistence.py +++ b/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): diff --git a/compass/pipeline/data_classes.py b/compass/pipeline/data_classes.py index 278703ce5..20127c1e1 100644 --- a/compass/pipeline/data_classes.py +++ b/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 diff --git a/compass/plugin/one_shot/base.py b/compass/plugin/one_shot/base.py index 4c7c8d011..f235b90bf 100644 --- a/compass/plugin/one_shot/base.py +++ b/compass/plugin/one_shot/base.py @@ -52,7 +52,9 @@ class _CacheKey(StrEnum): HEURISTIC_KEYWORDS = auto() -def create_schema_based_one_shot_extraction_plugin(config, tech): # ruff:ignore[complex-structure] +# ruff:ignore[complex-structure] +# complexipy: ignore +def create_schema_based_one_shot_extraction_plugin(config, tech): """Create a one-shot extraction plugin based on a configuration Parameters @@ -641,6 +643,18 @@ def _normalize_heuristic_keywords(raw): "GOOD_TECH_PHRASES", } + normalized = _normalize_input_kw(raw, expected_keys) + + _verify_expected_kw_are_not_missing(normalized, expected_keys) + _verify_kw_list_not_empty(normalized) + _verify_min_number_of_kw_provided(normalized) + _warn_if_not_enough_kw_provided(normalized) + + return normalized + + +def _normalize_input_kw(raw, expected_keys): + """Normalize the input keyword dictionary""" normalized = {} for raw_key, value in raw.items(): if not isinstance(raw_key, str): @@ -656,6 +670,11 @@ def _normalize_heuristic_keywords(raw): normalized[target_key] = _normalize_keyword_list(value) + return normalized + + +def _verify_expected_kw_are_not_missing(normalized, expected_keys): + """Verify that all expected keyword lists are present""" missing = expected_keys - set(normalized) if missing: msg = ( @@ -663,11 +682,17 @@ def _normalize_heuristic_keywords(raw): ) raise COMPASSPluginConfigurationError(msg) + +def _verify_kw_list_not_empty(normalized): + """Verify that no keyword list is empty""" empty = [key for key, value in normalized.items() if not value] if empty: msg = f"Heuristic keyword lists must not be empty: {sorted(empty)}" raise COMPASSPluginConfigurationError(msg) + +def _verify_min_number_of_kw_provided(normalized): + """Verify that the minimum number of "Good" keywords is provided""" num_good_kw = sum( len(kw_list) for key, kw_list in normalized.items() @@ -683,6 +708,14 @@ def _normalize_heuristic_keywords(raw): ) raise COMPASSPluginConfigurationError(msg) + +def _warn_if_not_enough_kw_provided(normalized): + """Warn if the number of "Good" keywords is too low""" + num_good_kw = sum( + len(kw_list) + for key, kw_list in normalized.items() + if key != "NOT_TECH_WORDS" + ) if num_good_kw < 10: # ruff:ignore[magic-value-comparison] msg = ( 'It is recommended to provide at least 10 total "Good" ' @@ -692,8 +725,6 @@ def _normalize_heuristic_keywords(raw): ) warn(msg, COMPASSPluginConfigurationWarning) - return normalized - def _normalize_keyword_list(items): """Normalize keyword list entries""" diff --git a/compass/services/cpu.py b/compass/services/cpu.py index 2080f8935..98fafc932 100644 --- a/compass/services/cpu.py +++ b/compass/services/cpu.py @@ -543,6 +543,7 @@ def _run_docling_subprocess(sender, fn, args, kwargs): sender.close() +# complexipy: ignore def _receive_docling_result(receiver, process, timeout): """Receive a child conversion result before its deadline expires""" deadline = time.monotonic() + timeout diff --git a/compass/utilities/parsing.py b/compass/utilities/parsing.py index 901285796..77e49d71b 100644 --- a/compass/utilities/parsing.py +++ b/compass/utilities/parsing.py @@ -236,31 +236,13 @@ def raw_pages_from_doc( # Do NOT use `is_pdf_doc` here because MDDocuments could have # "doc_type" == "pdf" and be treated as a single page doc if isinstance(doc, PDFDocument) and hasattr(doc, "raw_pages"): - raw_pages = doc.raw_pages - # failsafe check - if text_splitter is not None and len(raw_pages) == 1: - raw_pages = text_splitter.split_text(raw_pages[0]) - raw_pages = _down_select_pages( - raw_pages, - percent_raw_pages_to_keep, - max_raw_pages, - num_end_pages_to_keep, - ) - logger.debug( - "PDF Document from %s had 1 raw page; " - "has %d raw %s after splitting", - doc.attrs.get("source", "unknown source"), - len(raw_pages), - "page" if len(raw_pages) == 1 else "pages", - ) - else: - logger.debug( - "PDF Document from %s has %d raw %s", - doc.attrs.get("source", "unknown source"), - len(raw_pages), - "page" if len(raw_pages) == 1 else "pages", - ) - return raw_pages + return _raw_pages_from_pdf_doc( + doc, + text_splitter, + percent_raw_pages_to_keep, + max_raw_pages, + num_end_pages_to_keep, + ) if text_splitter is None: logger.debug( @@ -288,6 +270,42 @@ def raw_pages_from_doc( return raw_pages +def _raw_pages_from_pdf_doc( + doc, + text_splitter, + percent_raw_pages_to_keep, + max_raw_pages, + num_end_pages_to_keep, +): + """Get raw pages from an input PDF doc""" + raw_pages = doc.raw_pages + # failsafe check + if text_splitter is not None and len(raw_pages) == 1: + raw_pages = text_splitter.split_text(raw_pages[0]) + raw_pages = _down_select_pages( + raw_pages, + percent_raw_pages_to_keep, + max_raw_pages, + num_end_pages_to_keep, + ) + logger.debug( + "PDF Document from %s had 1 raw page; " + "has %d raw %s after splitting", + doc.attrs.get("source", "unknown source"), + len(raw_pages), + "page" if len(raw_pages) == 1 else "pages", + ) + return raw_pages + + logger.debug( + "PDF Document from %s has %d raw %s", + doc.attrs.get("source", "unknown source"), + len(raw_pages), + "page" if len(raw_pages) == 1 else "pages", + ) + return raw_pages + + def _down_select_pages( pages, percent_raw_pages_to_keep, max_raw_pages, num_end_pages_to_keep ): @@ -306,6 +324,7 @@ def _down_select_pages( return raw_pages +# complexipy: ignore def convert_paths_to_strings(obj): """[NOT PUBLIC API] Convert all Path instances to strings""" logger.trace("Converting paths to strings in object: %s", obj) diff --git a/compass/validation/content.py b/compass/validation/content.py index 39025840c..822ae15ac 100644 --- a/compass/validation/content.py +++ b/compass/validation/content.py @@ -471,20 +471,45 @@ async def parse_by_chunks( callbacks = callbacks or [] outer_task_name = asyncio.current_task().get_name() + async for ind in _chunks_to_check( + chunk_parser, + heuristic, + text_kind_validator, + min_chunks_to_process, + passed_heuristic_mem, + ): + if not callbacks: + continue + + cb_futures = [ + asyncio.create_task(cb(chunk_parser, ind), name=outer_task_name) + for cb in callbacks + ] + cb_results = await asyncio.gather(*cb_futures) + + # mask this chunk if we got a good result - this avoids forcing + # the following chunk to be checked (it will only be checked if + # it itself passes the heuristic) + passed_heuristic_mem[-1] = not any(cb_results) + + +async def _chunks_to_check( + chunk_parser, + heuristic, + text_kind_validator, + min_chunks_to_process, + passed_heuristic_mem, +): + """Yield indices of chunks that should be checked""" for ind, text in enumerate(chunk_parser.text_chunks): passed_heuristic_mem.append(heuristic.check(text)) if ind < min_chunks_to_process: - if text_kind_validator is not None: - is_correct_text_kind = await text_kind_validator.check_chunk( - chunk_parser, ind - ) - if not is_correct_text_kind: - continue # don't bother checking this chunk + if not await _chunk_is_correct_text_kind( + chunk_parser, ind, text_kind_validator + ): + continue - elif ( - text_kind_validator is not None - and not text_kind_validator.is_correct_kind_of_text - ): + elif not _document_is_correct_text_kind(text_kind_validator): return # don't bother checking this document # hasn't passed heuristic, so don't pass it to callbacks @@ -493,17 +518,20 @@ async def parse_by_chunks( logger.debug("Processing text at ind %d", ind) logger.debug_to_file("Text:\n%s", text) + yield ind - if not callbacks: - continue - cb_futures = [ - asyncio.create_task(cb(chunk_parser, ind), name=outer_task_name) - for cb in callbacks - ] - cb_results = await asyncio.gather(*cb_futures) +async def _chunk_is_correct_text_kind(chunk_parser, ind, text_kind_validator): + """Check if a specific chunk of text is of the correct kind""" + if text_kind_validator is None: + return True # assume chunk is good if no validator given - # mask this chunk if we got a good result - this avoids forcing - # the following chunk to be checked (it will only be checked if - # it itself passes the heuristic) - passed_heuristic_mem[-1] = not any(cb_results) + return await text_kind_validator.check_chunk(chunk_parser, ind) + + +def _document_is_correct_text_kind(text_kind_validator): + """Check if the entire document is of the correct kind""" + if text_kind_validator is None: + return True # assume doc is good if no validator given + + return text_kind_validator.is_correct_kind_of_text diff --git a/compass/web/file_loader.py b/compass/web/file_loader.py index 3a4f01747..1bbe85835 100644 --- a/compass/web/file_loader.py +++ b/compass/web/file_loader.py @@ -288,27 +288,9 @@ async def _maybe_fetch_failed_docs_with_elm(self, docs, sources): if self.failed_fetcher is None: return docs - out_docs = [] - partial_fail_docs = {} - failed_searches = [] - for source in sources: - source_docs = [ - doc for doc in docs if doc.attrs["source"] == source - ] - if not source_docs: - failed_searches.append(source) - continue - - if len(source_docs) > 1: - out_docs.extend(source_docs) - continue - - doc = source_docs[0] - if doc.attrs.get("conversion_status") != "success": - failed_searches.append(source) - partial_fail_docs[source] = doc - else: - out_docs.append(doc) + out_docs, partial_fail_docs, failed_searches = ( + _collect_failed_searches(docs, sources) + ) if not failed_searches: return out_docs @@ -319,16 +301,12 @@ async def _maybe_fetch_failed_docs_with_elm(self, docs, sources): failed_searches, ) elm_docs = await self.failed_fetcher.fetch_all(*failed_searches) + elm_docs = [ + _select_elm_or_partial_doc(elm_doc, partial_fail_docs) + for elm_doc in elm_docs + ] - for elm_doc in elm_docs: - docling_doc = partial_fail_docs.get(elm_doc.attrs["source"]) - elm_doc_failed = elm_doc.empty or "cache_fn" not in elm_doc.attrs - if elm_doc_failed and docling_doc is not None: - out_docs.append(docling_doc) - else: - out_docs.append(elm_doc) - - return out_docs + return out_docs + elm_docs async def _fetch_doc(self, url): """Fetch a doc using Docling""" @@ -473,6 +451,42 @@ async def _fetch_doc_with_url_in_metadata(self, source): return doc, raw_content +# complexipy: ignore +def _collect_failed_searches(docs, sources): + """Collect failed searches and categorize documents""" + out_docs = [] + partial_fail_docs = {} + failed_searches = [] + for source in sources: + source_docs = [doc for doc in docs if doc.attrs["source"] == source] + if not source_docs: + failed_searches.append(source) + continue + + if len(source_docs) > 1: + out_docs.extend(source_docs) + continue + + doc = source_docs[0] + if doc.attrs.get("conversion_status") != "success": + failed_searches.append(source) + partial_fail_docs[source] = doc + else: + out_docs.append(doc) + + return out_docs, partial_fail_docs, failed_searches + + +def _select_elm_or_partial_doc(elm_doc, partial_fail_docs): + """Select elm doc if it is valid; otherwise use partial doc""" + docling_doc = partial_fail_docs.get(elm_doc.attrs["source"]) + elm_doc_failed = elm_doc.empty or "cache_fn" not in elm_doc.attrs + + return ( + docling_doc if elm_doc_failed and docling_doc is not None else elm_doc + ) + + if os.environ.get("COMPASS_FILE_LOAD_BACKEND", "elm") == "docling": COMPASSWebFileLoader = AsyncDoclingWebFileLoader COMPASSLocalFileLoader = AsyncLocalDoclingFileLoader diff --git a/compass/web/search.py b/compass/web/search.py index c8600b1ac..5db0098d4 100644 --- a/compass/web/search.py +++ b/compass/web/search.py @@ -185,26 +185,38 @@ def _flatten_results(results): def _apply_blacklist_filters(results, url_blacklist, url_whitelist): """Mark rows that match any blacklist substring""" - blacklist_terms = [sub.casefold() for sub in url_blacklist or [] if sub] - whitelist_terms = [sub.casefold() for sub in url_whitelist or [] if sub] + blacklist_terms = _parsed_list(url_blacklist) + whitelist_terms = _parsed_list(url_whitelist) + for entry in results: url_cf = entry["url"].casefold() - if any(sub in url_cf for sub in whitelist_terms): + if _url_is_whitelisted(url_cf, whitelist_terms): continue - match_index = next( - ( - i - for i, sub_cf in enumerate(blacklist_terms) - if sub_cf in url_cf - ), - None, - ) + match_index = _blacklist_match_index(url_cf, blacklist_terms) if match_index is None: continue entry["filtered_reason"] = f"blacklist:{blacklist_terms[match_index]}" +def _parsed_list(url_list): + """Parse a list of URL substrings; normalize each non-empty entry""" + return [sub.casefold() for sub in url_list or [] if sub] + + +def _url_is_whitelisted(url_cf, whitelist_terms): + """Check if the URL matches any whitelist substring""" + return any(sub in url_cf for sub in whitelist_terms) + + +def _blacklist_match_index(url_cf, blacklist_terms): + """Return the index of the first matching blacklist substring""" + return next( + (i for i, sub_cf in enumerate(blacklist_terms) if sub_cf in url_cf), + None, + ) + + def _apply_duplicate_filters(results): """Mark duplicate rows by URL, across all search engines diff --git a/compass/web/website_crawl.py b/compass/web/website_crawl.py index 2871d49c1..0a64de851 100644 --- a/compass/web/website_crawl.py +++ b/compass/web/website_crawl.py @@ -242,6 +242,7 @@ async def run( return self._out_docs + # complexipy: ignore async def _run( self, base_url, @@ -394,7 +395,10 @@ async def _website_link_is_pdf(self, link, depth, score): return False logger.debug("Loading Link: %s", link) + return await self._loaded_link_is_pdf(link, depth, score, parsed) + async def _loaded_link_is_pdf(self, link, depth, score, parsed): + """Check if the loaded link is a PDF document""" try: doc = await self.fast_afl.fetch(link.href) except KeyboardInterrupt: @@ -571,31 +575,30 @@ def _extract_links_from_html(text, base_url): (a.get_text().strip(), a["href"]) for a in soup.find_all("a", href=True) ] + return set(_sanitized_links(links, base_url)) - out_links = set() + +def _sanitized_links(links, base_url): + """Sanitized links from the given list of (title, path) tuples""" for title, path in links: if not title or not path: continue - if any(substr in title.lower() for substr in _BLACKLIST_SUBSTRINGS): - continue - - if any(substr in path.lower() for substr in _BLACKLIST_SUBSTRINGS): + if _is_blacklisted(title, path): continue href = sanitize_url(urljoin(base_url, path)) if urlsplit(href).scheme not in {"http", "https"}: continue - out_links.add( - _Link( - title=title, - href=href, - base_domain=base_url, - ) - ) + yield _Link(title=title, href=href, base_domain=base_url) + - return out_links +def _is_blacklisted(title, path): + """Check if a link is blacklisted based on title or path""" + if any(substr in title.lower() for substr in _BLACKLIST_SUBSTRINGS): + return True + return any(substr in path.lower() for substr in _BLACKLIST_SUBSTRINGS) async def _get_text_from_all_locators(page): diff --git a/pixi.lock b/pixi.lock index 6b422f60d..2742ccaad 100644 --- a/pixi.lock +++ b/pixi.lock @@ -4416,6 +4416,7 @@ environments: - pypi: https://files.pythonhosted.org/packages/d9/16/7409957243cd7413eda85f7caf81090e0b1849db1a3e673f74c69dafe2e8/scrapling-0.2.99-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/d9/4f/00be2196329ebbff56ce564aa94efb0fbc828d00de250b1980de1a34ab49/python_pptx-1.0.2-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/db/48/dd301d62c1529efdd721b47b9e5fb52120fcdac5f4d3405cfc0d2f391414/pyclipper-1.4.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl + - pypi: https://files.pythonhosted.org/packages/db/85/bd15c109459a9b34050a5965c5b3fa5da40f010572c556796441eef99112/complexipy-7.0.1-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl - pypi: https://files.pythonhosted.org/packages/dc/61/e24b560ab2e2eaeb3c839129175fb330dfcfc29e5203196e5541a4c44682/nvidia_cublas_cu12-12.8.4.1-py3-none-manylinux_2_27_x86_64.whl - pypi: https://files.pythonhosted.org/packages/dd/34/b6f19941adcdaf415b5e8a8d577499f5b6a76b59cbae37f9b125a9ffe9f2/polyfactory-3.3.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/dd/35/ce962f738ae28ffce6293e7607b129075633e6bb185a5ab87e49246eedc2/browserforge-1.2.4-py3-none-any.whl @@ -4975,6 +4976,7 @@ environments: - pypi: https://files.pythonhosted.org/packages/e3/42/178db21aab1815583fcdb8ae465fc006b384fbe679412b11ddf8aae90f38/ua_parser_builtins-202605-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/e3/94/1843518e420fa3ed6919835845df698c7e27e183cb997394e4a670973a65/omegaconf-2.3.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/e4/ad/77fad9d6f974ec58d837cb49fb9b483d6227a420c4f908c3578633de1d47/alphashape-1.3.1-py2.py3-none-any.whl + - pypi: https://files.pythonhosted.org/packages/e5/e3/c86048d3fba3b3d6c888de45080fda26261f2c0599b9e323ea3255e79ddb/complexipy-7.0.1-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl - pypi: https://files.pythonhosted.org/packages/e8/5f/32de4d99220eb559b7b1cd1c529a1856efa8097f7a3e10b6c207aa95e36c/ddgs-9.14.4-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/e8/b1/c488b530994c4f68e46efa99a4d6ca6741aaf158e35779fe6c4d8a9a427d/latex2mathml-3.81.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/ea/3e/47088ffc9c33eeaee3e013c838bb1be8e307619faed12e245e16d74cefef/rebrowser_playwright-1.49.1-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl @@ -5518,6 +5520,7 @@ environments: - pypi: https://files.pythonhosted.org/packages/05/98/716a473cfb24750858ddd5d14e6527539dd206583a46408d08eeb2844a75/trimesh-4.12.2-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/09/bd/bea0f66f3c8e8fc7a7e8af54c3b6d9139c195e7138db2cc9a5ce025a495a/nlr_elm-0.0.44-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/0a/21/b48b9e8408b3faec9fb7cb2f68352c7724add35a980c948eaceb29ac41e4/language_tags-1.3.1-py3-none-any.whl + - pypi: https://files.pythonhosted.org/packages/10/9c/49ccd33634d155bfb4c98a6c8be392d2beb9aa6e10e0ce16ba66864ce993/complexipy-7.0.1-cp313-cp313-macosx_10_12_x86_64.whl - pypi: https://files.pythonhosted.org/packages/1a/55/d62c85fff36e9e9e515ee92407b02acb556e6832d4fbcc8624b638cf70bb/patchright-1.51.3-py3-none-macosx_11_0_universal2.whl - pypi: https://files.pythonhosted.org/packages/20/0c/7bb51e3acfafd16c48875bf3db03607674df16f5b6ef8d056586af7e2b8b/cssselect-1.4.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/2a/21/f691fb2613100a62b3fa91e9988c991e9ca5b89ea31c0d3152a3210344f9/rank_bm25-0.2.2-py3-none-any.whl @@ -6101,6 +6104,7 @@ environments: - pypi: https://files.pythonhosted.org/packages/55/20/4df3f817c28938e22ee7c7c4b28d8b3a212e5a111c3bd9633bc410267daa/patchright-1.51.3-py3-none-macosx_11_0_arm64.whl - pypi: https://files.pythonhosted.org/packages/55/5f/25bdec773905bff0ff6cf35ca73b17bd05593b4f87bd8c5fa43705f7167d/chardet-7.4.3-cp313-cp313-macosx_11_0_arm64.whl - pypi: https://files.pythonhosted.org/packages/58/c1/c965cc23f96a364803d44b4331f33e4465bb6f269add37e39d0ad77ffe33/primp-1.3.1-cp310-abi3-macosx_11_0_arm64.whl + - pypi: https://files.pythonhosted.org/packages/64/d8/e1e3f0b088d6efedee48246cf016a073ccde8535f7b5b8d732f6244a02d7/complexipy-7.0.1-cp313-cp313-macosx_11_0_arm64.whl - pypi: https://files.pythonhosted.org/packages/66/c3/f8b216cbd742e5b84c40f045204c764ccb7524d2aeab021054ec69446b0a/w3lib-2.4.1-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/6d/42/0e49d6d0aac449ca71952ec5bae764af009754fcb2e76a5cc097543747b3/tldextract-5.3.1-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/6e/bf/c5205d480307bef660e56544b9e3d7ff687da776abb30c9cb3f330887570/screeninfo-0.8.1-py3-none-any.whl @@ -6598,6 +6602,7 @@ environments: - pypi: https://files.pythonhosted.org/packages/20/0f/098488de02e3d52fc77e8d55c1467f6703701b6ea6788f40409bb8c00dd4/playwright-1.51.0-py3-none-win_amd64.whl - pypi: https://files.pythonhosted.org/packages/20/14/1db1729ad6db4999c3a16c47937d601fcb909aaa4224f5eca5a2f145a605/mpire-2.10.2-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/22/1d/5573ef3624ceb7abf4a46073d3554e37191c868abc3aecd5289a72f9810a/fastuuid-0.14.0-cp313-cp313-win_amd64.whl + - pypi: https://files.pythonhosted.org/packages/29/9e/e8d107a796b91690b2b763e8d2d49152cda02c583376942f97bba1569e7b/complexipy-7.0.1-cp313-cp313-win_amd64.whl - pypi: https://files.pythonhosted.org/packages/2a/21/f691fb2613100a62b3fa91e9988c991e9ca5b89ea31c0d3152a3210344f9/rank_bm25-0.2.2-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/2e/45/6adc5efeb19fd5fed4027e520b5c668ce58236a2b271ade5533c4c116276/sphinx_tabs-3.5.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/2f/4a/ef8bb2b86988e7e45b8dfa75a9387fe346f5f52792bfdd0d530ef36c9afe/rebrowser_playwright-1.49.1-py3-none-win_amd64.whl @@ -34235,6 +34240,7 @@ packages: - rich>=13.9.4,<14 - toml>=0.10.2,<0.11 - pytesseract>=0.3.13,<0.4 ; extra == 'ocr' + - complexipy>=7.0.1,<8 ; extra == 'dev' - jupyter>=1.0.0,<1.1 ; extra == 'dev' - pipreqs>=0.4.13,<0.5 ; extra == 'dev' - ruff>=0.16,<0.17 ; extra == 'dev' @@ -34950,6 +34956,14 @@ packages: - skia-pathops>=0.5.0 ; extra == 'all' - uharfbuzz>=0.45.0 ; extra == 'all' requires_python: '>=3.10' +- pypi: https://files.pythonhosted.org/packages/10/9c/49ccd33634d155bfb4c98a6c8be392d2beb9aa6e10e0ce16ba66864ce993/complexipy-7.0.1-cp313-cp313-macosx_10_12_x86_64.whl + name: complexipy + version: 7.0.1 + sha256: b960ababd486b62f18e548c252f5b5b37931c1a236923f04b626f283f1df590b + requires_dist: + - tomli>=2.2.1 + - typer>=0.12.5 + requires_python: '>=3.8' - pypi: https://files.pythonhosted.org/packages/11/c8/26f2e4aae92f11522a96043892ba39a90eac610d5242523aa863212bc1c7/cuda_pathfinder-1.5.5-py3-none-any.whl name: cuda-pathfinder version: 1.5.5 @@ -35069,6 +35083,14 @@ packages: version: 0.14.0 sha256: 0df14e92e7ad3276327631c9e7cec09e32572ce82089c55cb1bb8df71cf394ed requires_python: '>=3.8' +- pypi: https://files.pythonhosted.org/packages/29/9e/e8d107a796b91690b2b763e8d2d49152cda02c583376942f97bba1569e7b/complexipy-7.0.1-cp313-cp313-win_amd64.whl + name: complexipy + version: 7.0.1 + sha256: 54650d957684333852005e8cc26b61eb3979f83428a1051e872545733387ad9b + requires_dist: + - tomli>=2.2.1 + - typer>=0.12.5 + requires_python: '>=3.8' - pypi: https://files.pythonhosted.org/packages/2a/21/f691fb2613100a62b3fa91e9988c991e9ca5b89ea31c0d3152a3210344f9/rank_bm25-0.2.2-py3-none-any.whl name: rank-bm25 version: 0.2.2 @@ -35581,6 +35603,14 @@ packages: - setuptools-scm>=7,<10 ; extra == 'dev' - setuptools>=64 ; extra == 'dev' requires_python: '>=3.10' +- pypi: https://files.pythonhosted.org/packages/64/d8/e1e3f0b088d6efedee48246cf016a073ccde8535f7b5b8d732f6244a02d7/complexipy-7.0.1-cp313-cp313-macosx_11_0_arm64.whl + name: complexipy + version: 7.0.1 + sha256: 4582c730c592bbd0a4c248523d8573305af9cc02d7677757088d5e5ff9695eba + requires_dist: + - tomli>=2.2.1 + - typer>=0.12.5 + requires_python: '>=3.8' - pypi: https://files.pythonhosted.org/packages/66/c3/f8b216cbd742e5b84c40f045204c764ccb7524d2aeab021054ec69446b0a/w3lib-2.4.1-py3-none-any.whl name: w3lib version: 2.4.1 @@ -36986,6 +37016,14 @@ packages: version: 1.4.0 sha256: 14c8bdb5a72004b721c4e6f448d2c2262d74a7f0c9e3076aeff41e564a92389f requires_python: '>=3.10' +- pypi: https://files.pythonhosted.org/packages/db/85/bd15c109459a9b34050a5965c5b3fa5da40f010572c556796441eef99112/complexipy-7.0.1-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl + name: complexipy + version: 7.0.1 + sha256: cc12b194a36299bfd9f72a2aafda718c6038b9ee747f316d9564fa58ec0f9827 + requires_dist: + - tomli>=2.2.1 + - typer>=0.12.5 + requires_python: '>=3.8' - pypi: https://files.pythonhosted.org/packages/dc/61/e24b560ab2e2eaeb3c839129175fb330dfcfc29e5203196e5541a4c44682/nvidia_cublas_cu12-12.8.4.1-py3-none-manylinux_2_27_x86_64.whl name: nvidia-cublas-cu12 version: 12.8.4.1 @@ -37137,6 +37175,14 @@ packages: - networkx>=2.5 - rtree>=0.9.7 - scipy>=1.0.0 +- pypi: https://files.pythonhosted.org/packages/e5/e3/c86048d3fba3b3d6c888de45080fda26261f2c0599b9e323ea3255e79ddb/complexipy-7.0.1-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl + name: complexipy + version: 7.0.1 + sha256: 84012b74389b0bf0d482a52a9e98f5e3086907c0d962eba9935df709759439f3 + requires_dist: + - tomli>=2.2.1 + - typer>=0.12.5 + requires_python: '>=3.8' - pypi: https://files.pythonhosted.org/packages/e7/05/c19819d5e3d95294a6f5947fb9b9629efb316b96de511b418c53d245aae6/cycler-0.12.1-py3-none-any.whl name: cycler version: 0.12.1 diff --git a/pixi.toml b/pixi.toml index c3144b5b9..dac31d779 100644 --- a/pixi.toml +++ b/pixi.toml @@ -148,6 +148,9 @@ pipreqs = ">=0.4.13,<0.5" ruff = ">=0.16,<0.17" seaborn = ">=0.13.2,<0.14" +[feature.python-dev.pypi-dependencies] +complexipy = ">=7.0.1,<8" + [feature.python-test.dependencies] codespell = ">=2.4,<3" flaky = ">=3.8.1,<4" diff --git a/pyproject.toml b/pyproject.toml index 890b54426..bd77e572a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -61,6 +61,7 @@ ocr = [ "pytesseract>=0.3.13,<0.4" ] dev = [ + "complexipy>=7.0.1,<8", "jupyter>=1.0.0,<1.1", "pipreqs>=0.4.13,<0.5", "ruff>=0.16,<0.17", @@ -106,6 +107,7 @@ homepage = "https://github.com/NatLabRockies/COMPASS" documentation = "https://natlabrockies.github.io/COMPASS/" repository = "https://github.com/NatLabRockies/COMPASS" + [tool.setuptools.packages] find = { include = ["compass*"] } @@ -116,6 +118,7 @@ fallback_version = "9999" git_describe_command = [ "git", "describe", "--dirty", "--tags", "--long", "--match", "v*[0-9]*" ] version_file = "compass/_version.py" + [tool.ruff] line-length = 79 @@ -217,6 +220,7 @@ convention = "numpy" max-args = 10 max-positional-args = 10 + [tool.coverage.run] branch = true @@ -251,7 +255,6 @@ exclude_also = [ "def print_logging_info*", "def __cls_name", ] - omit = [ # omit test files "tests/*", @@ -276,6 +279,7 @@ omit = [ skip = "*.lock,*.min.js,*.svg,*.pdf,*.png,*.jpg,*.jpeg,*.gif,*.ico" # ignore-words-list = "word1,word2" + [tool.pytest.ini_options] addopts = '--durations=10 --disable-warnings -m "not evals"' asyncio_mode="auto" @@ -287,3 +291,12 @@ testpaths = [ "tests/python/unit", "tests/python/integration", ] + + +[tool.complexipy] +paths = ["compass"] +max-complexity-allowed = 10 +exclude = ["tests/**"] +failed = true +sort = "desc" +check-script = true diff --git a/tests/python/unit/validation/test_validation_content.py b/tests/python/unit/validation/test_validation_content.py index 9a62bfa6a..403171223 100644 --- a/tests/python/unit/validation/test_validation_content.py +++ b/tests/python/unit/validation/test_validation_content.py @@ -1,5 +1,6 @@ """COMPASS Ordinance content validation tests""" +import asyncio import os from pathlib import Path @@ -76,6 +77,89 @@ async def call(self, key, text_chunk): ] +@pytest.mark.asyncio +async def test_parse_by_chunks_masks_chunk_after_successful_callback(): + """Test callback dispatch respects heuristic and callback results""" + + class MatchingHeuristic: + """Recognize chunks explicitly marked as matching""" + + def check(self, text): + """Return whether a chunk matches""" + return text == "match" + + processed_indices = [] + + async def callback(chunk_parser, ind): + """Record processed chunk indices""" + processed_indices.append(ind) + await asyncio.sleep(0) + return True + + chunk_parser = ParseChunksWithMemory( + ["match", "skip", "match"], num_to_recall=2 + ) + + await parse_by_chunks( + chunk_parser, + heuristic=MatchingHeuristic(), + callbacks=[callback], + min_chunks_to_process=0, + ) + + assert processed_indices == [0, 2] + + +@pytest.mark.asyncio +async def test_parse_by_chunks_stops_after_initial_invalid_chunks(): + """Test invalid initial chunks stop later callback processing""" + + class AlwaysMatchingHeuristic: + """Recognize every chunk""" + + def check(self, text): + """Return a matching result""" + return True + + class InvalidTextValidator: + """Reject every chunk and document""" + + def __init__(self): + self.checked_indices = [] + + async def check_chunk(self, chunk_parser, ind): + """Record and reject a chunk""" + self.checked_indices.append(ind) + return False + + @property + def is_correct_kind_of_text(self): + """bool: Always reject the document""" + return False + + processed_indices = [] + validator = InvalidTextValidator() + + async def callback(chunk_parser, ind): + """Record processed chunk indices""" + processed_indices.append(ind) + await asyncio.sleep(0) + return True + + chunk_parser = ParseChunksWithMemory(["one", "two", "three"]) + + await parse_by_chunks( + chunk_parser, + heuristic=AlwaysMatchingHeuristic(), + text_kind_validator=validator, + callbacks=[callback], + min_chunks_to_process=2, + ) + + assert validator.checked_indices == [0, 1] + assert processed_indices == [] + + @flaky(max_runs=3, min_passes=1) @pytest.mark.skipif(SHOULD_SKIP, reason="requires Azure OpenAI key") @pytest.mark.asyncio