Track search engine in doc attrs

compass/_cli/finalize.py

@@ -15,6 +15,7 @@
 from compass.utilities.io import load_config
 from compass.utilities.jurisdictions import Jurisdiction
 from compass.utilities.finalize import save_run_meta, doc_infos_to_db, save_db
+from compass.services.usage import LLM_USAGE_RATES_KEY
 from compass.pipeline import build_models
 
 
@@ -94,10 +95,19 @@ def finalize(ctx, config):
         num_jurisdictions_found=num_jurisdictions_found,
         total_cost=total_cost,
         models=models,
+        llm_usage_rates=_load_llm_usage_rates(dirs.out),
     )
     console.print(f"✅ Finalized COMPASS run in {dirs.out!s}!")
 
 
+def _load_llm_usage_rates(out_dir):
+    """Load persisted LLM rate statistics when available"""
+    usage_fp = out_dir / "usage.json"
+    if not usage_fp.exists():
+        return None
+    return load_config(usage_fp).get(LLM_USAGE_RATES_KEY)
+
+
 def _compile_db(jurisdictions, dirs, tech):
     """Merge all jurisdiction dbs into one"""
     all_doc_infos = []

compass/extraction/apply.py

@@ -67,7 +67,7 @@ async def check_for_relevant_text(
         If the document already contains text collected by a given
         collector (i.e. the collector's ``OUT_LABEL`` is found in
         ``doc.attrs``), that collector will be skipped.
-    usage_tracker : UsageTracker, optional
+    usage_tracker : LLMUsageTracker, optional
         Optional tracker instance to monitor token usage during
         LLM calls. By default, ``None``.
     min_chunks_to_process : int, optional
@@ -154,7 +154,7 @@ async def extract_date(doc, model_config, usage_tracker=None):
     model_config : compass.llm.config.LLMConfig
         Configuration describing which LLM service, splitter, and call
         parameters should be used for date extraction.
-    usage_tracker : UsageTracker, optional
+    usage_tracker : LLMUsageTracker, optional
         Optional tracker instance to monitor token usage during
         LLM calls. By default, ``None``.
 

compass/llm/calling.py

@@ -40,7 +40,7 @@ def __init__(self, llm_service, usage_tracker=None, **kwargs):
         ----------
         llm_service : Service
             LLM service used for queries.
-        usage_tracker : UsageTracker, optional
+        usage_tracker : LLMUsageTracker, optional
             Optional tracker instance to monitor token usage during
             LLM calls. By default, ``None``.
         **kwargs
@@ -130,7 +130,7 @@ def __init__(
             LLM service used for queries.
         system_message : str
             System message to use for chat with LLM.
-        usage_tracker : UsageTracker, optional
+        usage_tracker : LLMUsageTracker, optional
             Optional tracker instance to monitor token usage during
             LLM calls. By default, ``None``.
         **kwargs

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/coordinator.py

@@ -7,7 +7,7 @@
 from abc import ABC, abstractmethod
 
 from compass.services.openai import usage_from_response
-from compass.services.usage import UsageTracker
+from compass.services.usage import LLMUsageTracker
 from compass.exceptions import COMPASSError, COMPASSValueError
 from compass.utilities import (
     compile_collection_summary_message,
@@ -128,6 +128,7 @@ def _create(self, jurisdiction, *, usage_tracker=None):
             jurisdiction=jurisdiction,
             model_configs=self.runtime.models,
             usage_tracker=usage_tracker,
+            rate_tracker=self.runtime.rate_tracker,
         )
         return SingleJurisdictionRun(
             self.runtime,
@@ -185,7 +186,7 @@ async def run(self, jurisdictions_df):
         )
         tasks = []
         for jurisdiction in jurisdictions_from_df(jurisdictions_df):
-            usage_tracker = UsageTracker(
+            usage_tracker = LLMUsageTracker(
                 jurisdiction.full_name, usage_from_response
             )
             workflow = self._create(jurisdiction, usage_tracker=usage_tracker)
@@ -315,7 +316,7 @@ async def run(self, jurisdictions_df):
                 )
                 continue
 
-            usage_tracker = UsageTracker(
+            usage_tracker = LLMUsageTracker(
                 jurisdiction.full_name, usage_from_response
             )
             workflow = self._create(jurisdiction, usage_tracker=usage_tracker)
@@ -419,7 +420,9 @@ async def _finalize_extraction(
     runtime, results, start_date, num_jurisdictions
 ):
     """Finalize process or extraction mode outputs"""
-    total_cost = await _compute_total_cost()
+    total_cost, llm_usage_rates = await _compute_total_cost(
+        runtime.rate_tracker
+    )
     doc_infos = [
         {
             "jurisdiction": result.jurisdiction,
@@ -445,6 +448,7 @@ async def _finalize_extraction(
         num_jurisdictions_found=num_docs_found,
         total_cost=total_cost,
         models=runtime.models,
+        llm_usage_rates=llm_usage_rates,
     )
     run_msg = compile_run_summary_message(
         total_seconds=total_time,
@@ -457,9 +461,13 @@ async def _finalize_extraction(
     return run_msg
 
 
-async def _compute_total_cost():
-    """Compute total cost from tracked usage"""
-    total_usage = await UsageUpdater.call(None)
+async def _compute_total_cost(rate_tracker):
+    """Compute total cost and load LLM rate statistics from usage"""
+    total_usage = await UsageUpdater.call(None, rate_tracker)
     if not total_usage:
-        return 0
-    return compute_total_cost_from_usage(total_usage)
+        return 0, None
+
+    return (
+        compute_total_cost_from_usage(total_usage),
+        total_usage.get(rate_tracker.label),
+    )

compass/pipeline/data_classes.py

@@ -7,6 +7,7 @@
 from elm.web.search.run import SEARCH_ENGINE_OPTIONS
 
 from compass.llm import OpenAIConfig
+from compass.services.usage import LLMRateTracker
 from compass.utilities.enums import COMPASSRunMode, LLMTasks
 from compass.utilities.io import load_config
 from compass.exceptions import COMPASSValueError
@@ -691,13 +692,22 @@ def __init__(  # ruff:ignore[too-many-arguments]
         )
         self.user_model_input = model
         self.llm_costs = llm_costs
+        self._rate_tracker = LLMRateTracker()
 
     @cached_property
     def models(self):
         """dict: Mapping of LLM task to OpenAIConfig for this request"""
         if not self.user_model_input:
             return {}
-        return build_models(self.user_model_input)
+
+        return build_models(
+            self.user_model_input, rate_tracker=self._rate_tracker
+        )
+
+    @property
+    def rate_tracker(self):
+        """LLMRateTracker: Rate tracker for LLM calls"""
+        return self._rate_tracker
 
 
 class ProcessRequest(BaseRequest):
@@ -1287,7 +1297,7 @@ def __bool__(self):
         return self.ord_db_fp is not None
 
 
-def build_models(user_input, *, allow_empty=False):
+def build_models(user_input, *, allow_empty=False, rate_tracker=None):
     """[NOT PUBLIC API] Build configured model registry"""
     if user_input is None:
         return {} if allow_empty else {LLMTasks.DEFAULT: OpenAIConfig()}
@@ -1297,22 +1307,23 @@ def build_models(user_input, *, allow_empty=False):
 
     caller_instances = {}
     for raw_kwargs in user_input:
-        for task, model_config in _config_for_tasks(raw_kwargs):
+        for task, model_config in _config_for_tasks(raw_kwargs, rate_tracker):
             _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):
+def _config_for_tasks(kwargs, rate_tracker=None):
     """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)
+    model_config.llm_call_kwargs.update({"rate_tracker": rate_tracker})
     for task in tasks:
         yield task, model_config
 

compass/pipeline/jurisdiction.py

@@ -52,7 +52,7 @@ def __init__(
             Configured extraction plugin instance responsible for
             parsing collected documents and persisting structured
             output for this jurisdiction.
-        usage_tracker : UsageTracker, optional
+        usage_tracker : LLMUsageTracker, optional
             Optional tracker instance used to accumulate token usage
             and cost information for LLM calls made during the
             jurisdiction workflow. By default, ``None``.
@@ -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/pipeline/runtime.py

@@ -57,6 +57,7 @@ def __init__(self, request):
         self.mode = request.MODE
         self.tech = request.tech
         self.models = request.models
+        self.rate_tracker = request.rate_tracker
         self.search_params = request.search_settings
         self.log_level = _normalize_log_level(
             request.runtime_settings.log_level

compass/plugin/base.py

@@ -14,7 +14,13 @@ class BaseExtractionPlugin(ABC):
     implementer must define most functionality on their own.
     """
 
-    def __init__(self, jurisdiction, model_configs, usage_tracker=None):
+    def __init__(
+        self,
+        jurisdiction,
+        model_configs,
+        usage_tracker=None,
+        rate_tracker=None,
+    ):
         """
 
         Parameters
@@ -26,13 +32,17 @@ def __init__(self, jurisdiction, model_configs, usage_tracker=None):
             :class:`~compass.utilities.enums.LLMTasks` and values are
             :class:`~compass.llm.config.LLMConfig` instances to be used
             for those tasks.
-        usage_tracker : UsageTracker, optional
+        usage_tracker : LLMUsageTracker, optional
             Usage tracker instance that can be used to record the LLM
             call cost. By default, ``None``.
+        rate_tracker : LLMRateTracker, optional
+            Run-level LLM rate tracker used to persist LLM request and
+            token rate statistics. By default, ``None``.
         """
         self.jurisdiction = jurisdiction
         self.model_configs = model_configs
         self.usage_tracker = usage_tracker
+        self.rate_tracker = rate_tracker
 
     JURISDICTION_DATA_FP = None
     """:term:`path-like <path-like object>`: Path to jurisdiction CSV
@@ -169,12 +179,15 @@ def save_structured_data(cls, doc_infos, out_dir):
 
     async def record_usage(self):
         """Persist usage tracking data when a tracker is available"""
-        if self.usage_tracker is None:
+        if self.usage_tracker is None and self.rate_tracker is None:
             return
 
-        total_usage = await UsageUpdater.call(self.usage_tracker)
+        total_usage = await UsageUpdater.call(
+            self.usage_tracker, self.rate_tracker
+        )
         total_cost = compute_total_cost_from_usage(total_usage)
         COMPASS_PB.update_total_cost(total_cost, replace=True)
 
-    def validate_plugin_configuration(self):  # ruff:ignore[empty-method-without-abstract-decorator]
+    # ruff:ignore[empty-method-without-abstract-decorator]
+    def validate_plugin_configuration(self):
         """[NOT PUBLIC API] Validate plugin is properly configured"""