Track search engine in doc attrs

compass/pipeline/collection/base.py

@@ -127,13 +127,11 @@ async def execute(self, *, eager_extract=False):
         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(
-                        docs
-                    )
+                context = await self.workflow.extraction.extract_from_docs(
+                    docs
                 )
                 if context is not None:
-                    return context
+                    return self._context_with_documented_steps(context)
             else:
                 self._collection_info = (
                     await self.workflow.write_collection_shard_no_fail(
@@ -182,6 +180,14 @@ async def _run_collection_step(self, step):
         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:

compass/pipeline/collection/dedupe.py

@@ -55,6 +55,23 @@ def add_docs(self, docs, *, step_name=None):
             doc_info = self.data.setdefault(key, _DocInfo.from_doc(doc))
             doc_info.add_step(step_name)
 
+    def info(self, doc):
+        """Get the info for a given document
+
+        Parameters
+        ----------
+        doc : BaseDocument
+            The document for which to retrieve the deduplication info.
+
+        Returns
+        -------
+        object
+            The deduplication info for the given document.
+        """
+        return self.data.get(
+            _collection_doc_key(doc.attrs), _DocInfo(doc=doc, from_steps=[])
+        )
+
 
 def _collection_doc_key(doc_info):
     """Build the deduplication key for a collected document"""

compass/pipeline/jurisdiction.py

@@ -81,8 +81,8 @@ def __init__(
         self.perform_website_search = perform_website_search
         self.jurisdiction_website = jurisdiction.website_url
         self.last_scrape_results = []
-        self.extraction_workflow = DocumentExtraction(self)
-        self.collection_workflow = DocumentCollection(self)
+        self.collection = DocumentCollection(self)
+        self.extraction = DocumentExtraction(self)
 
     async def process(self):
         """Run process mode for one jurisdiction
@@ -101,8 +101,8 @@ async def process(self):
             self.jurisdiction.code,
         )
         try:
-            extraction_context = await self.collection_workflow.execute(
-                eager_extract=True,
+            extraction_context = await self.collection.execute(
+                eager_extract=True
             )
         finally:
             await self.extractor.record_usage()
@@ -143,9 +143,7 @@ async def collect(self):
             self.jurisdiction.full_name,
         )
 
-        collection_info = await self.collection_workflow.execute(
-            eager_extract=False
-        )
+        collection_info = await self.collection.execute(eager_extract=False)
 
         logger.info(
             "Completed collection for jurisdiction: %s",
@@ -187,9 +185,7 @@ async def extract_from_collection_info(self, collection_info):
                 collection_info, task_name=self.jurisdiction.full_name
             )
             docs = [doc for doc in docs if doc is not None]
-            extraction_context = (
-                await self.extraction_workflow.extract_from_docs(docs)
-            )
+            extraction_context = await self.extraction.extract_from_docs(docs)
         finally:
             await self.extractor.record_usage()
             await _record_jurisdiction_info(

compass/scripts/download.py

@@ -767,21 +767,29 @@ async def _docs_from_web_search(
         **kwargs,
     )
     ranked_results = {
-        res.get("url"): res.get("overall_rank") or 1
+        res.get("url"): res
         for res in out["results"]
         if res.get("filtered_reason") is None and res.get("url") is not None
     }
-    urls = sorted(ranked_results, key=ranked_results.get)
+    urls = sorted(
+        ranked_results,
+        key=lambda url: ranked_results[url].get("overall_rank") or 1,
+    )
     if not urls:
         return []
 
     docs = await _docs_from_urls(
         urls, jurisdiction.full_name, browser_semaphore, **kwargs
     )
     for doc in docs:
-        doc.attrs[_COLLECTION_SCORE_KEY] = ranked_results.get(
-            doc.attrs.get("source")
-        )
+        result = ranked_results.get(doc.attrs.get("source"))
+        if result is None:
+            doc.attrs[_COLLECTION_SCORE_KEY] = None
+            continue
+
+        doc.attrs[_COLLECTION_SCORE_KEY] = result.get("overall_rank") or 1
+        if "search_engines" in result:
+            doc.attrs["search_engines"] = list(result["search_engines"])
     return docs
 
 

compass/services/openai.py

@@ -117,7 +117,7 @@ def __init__(
             Name of model being used.
         rate_limit : int or float, optional
             Token rate limit (typically per minute, but the time
-            interval is ultimately controlled by the `rate_tracker`
+            interval is ultimately controlled by the `timed_tracker`
             instance). By default, ``1e3``.
         timed_tracker : TimeBoundedUsageTracker, optional
             Instance used to track usage per time interval and compare
@@ -186,7 +186,7 @@ async def process(
         return _get_response_message(response)
 
     def _record_prompt_tokens(self, kwargs):
-        """Add prompt token count to rate tracker"""
+        """Add prompt token count to timed tracker"""
         num_tokens = count_tokens(kwargs.get("messages", []), self.model_name)
         return self.timed_tracker.add(num_tokens)
 
@@ -197,7 +197,7 @@ def _record_request_rate(self, timestamp, rate_tracker):
         rate_tracker.record_request(self.model_name, timestamp)
 
     def _record_completion_tokens(self, response):
-        """Add completion token count to rate tracker"""
+        """Add completion token count to timed tracker"""
         if response is None:
             return None
         return self.timed_tracker.add(response.usage.completion_tokens)

compass/services/threaded.py

@@ -683,6 +683,7 @@ def _compile_doc_info(doc):
             "permitted_use_text_ngram_score"
         ),
         "from_steps": doc.attrs.get("from_steps"),
+        "search_engines": doc.attrs.get("search_engines"),
     }
 
 

compass/utilities/costs.py

@@ -86,7 +86,7 @@ def compute_cost_from_totals(totals):
         should contain "prompt_tokens" and "response_tokens" keys
         indicating the number of tokens used for prompts and responses,
         respectively. This dictionary is typically obtained from the
-        `tracker_totals` property of a LLMUsageTracker instance.
+        `tracker_totals` entry of a LLMUsageTracker instance.
 
     Returns
     -------

compass/web/search.py

@@ -3,7 +3,7 @@
 import logging
 from urllib.parse import urlsplit, urlunsplit
 
-from elm.web.search.run import search_with_fallback, search_all_se
+from elm.web.search.run import search_all_se, search_with_fallback_with_attrs
 
 from compass.utilities.url import URLPartFilter
 
@@ -150,7 +150,7 @@ async def _run_simple_sort_search(
     **se_kwargs,
 ):
     """Run search with fallback search engines, applying simple sort"""
-    urls = await search_with_fallback(
+    return await search_with_fallback_with_attrs(
         queries,
         num_urls=num_urls,
         url_ignore_substrings=url_ignore_substrings,
@@ -159,7 +159,6 @@ async def _run_simple_sort_search(
         task_name=jurisdiction_full_name,
         **se_kwargs,
     )
-    return [{"url": url} for url in urls]
 
 
 async def _run_holistic_sort_search(
@@ -236,6 +235,9 @@ def _apply_duplicate_filters(results):
     config only wins among entries that are otherwise tied (same
     ``query_rank`` and ``query_index``).
     """
+    for entry in results:
+        entry["search_engines"] = [entry["search_engine"]]
+
     winners = {}
     for entry in _active_results_sorted(results):
         key = entry["url"]
@@ -252,6 +254,8 @@ def _apply_duplicate_filters(results):
                 "query_rank": entry["query_rank"],
             }
         )
+        if entry["search_engine"] not in winner["search_engines"]:
+            winner["search_engines"].append(entry["search_engine"])
 
         entry["filtered_reason"] = "duplicate"
 

examples/execution_basics/config_kitchen_sink.json5

@@ -83,7 +83,7 @@
     // futures to sit idly, awaiting their turn to query the LLM
     "max_num_concurrent_jurisdictions": 25,
     "url_ignore_substrings": [
-        // ignore any URLs with these strings
+        // ignore search results and COMPASS crawl URLs with these strings
         // useful to filter out known noise
         "wiki",
         "nrel.gov",
@@ -94,6 +94,10 @@
         // e.g. a giant file with generic ordinances that should be skipped
         "www.co.delaware.in.us/egov/documents/1649699794_0382.pdf",
     ],
+    "url_keep_substrings": [
+        // override default and configured URL exclusions in both phases
+        "my_ordinance_collection.edu",
+    ],
     "known_doc_urls": {
         "18017": "https://www.in.gov/counties/cass/files/Cass-County-Zoning-Ordinance-2024.pdf",
         // Can include more FIPS -> URL mappings here