Track usage and estimate cost for calls

README.md

@@ -96,6 +96,38 @@ The specific type of this object depends entirely on the pipeline configuration.
 Most commonly, results from the `inspect` modules will be stored here.
 For example, if you use the `inspect:quality` module, `context.quality` will contain those results.
 
+#### Usage reporting
+
+Set `report_usage` in the runtime configuration to collect one usage record for
+each OpenAI, Azure OpenAI, embedding, or Azure Document Intelligence call:
+
+```py
+context = pipe.run({
+    "report_usage": True,
+    # ... input and output runtime configuration
+})
+print(context.usage)
+```
+
+Set `estimate_cost` to enable usage reporting and add best-effort Azure retail
+cost estimates to each call. Azure pricing is fetched from the Azure Retail
+Prices API and cached in memory for 24 hours. The Azure region is required
+because retail prices can vary by region:
+
+```py
+context = pipe.run({
+    "estimate_cost": True,
+    "azure_region": "eastus",
+    # Optional; these are the defaults:
+    "azure_deployment_type": "global",  # global, data_zone, or regional
+    "azure_context_tier": "short",      # short or long
+})
+```
+
+Cost estimates currently support Azure only. If a price cannot be fetched or
+matched unambiguously, the pipeline continues and the call's `cost_estimate`
+contains a null estimate and an error message.
+
 #### Validation
 
 The pipeline will validate correctness before it runs.

bc2/core/analyze/azuredi.py

@@ -11,6 +11,7 @@
 from ..common.file import MemoryFile
 from ..common.json import date_aware_json_dumps
 from ..common.preprocess import register_preprocessor
+from ..common.usage import record_usage
 from .base import BaseAnalyzeDriver
 
 logger = logging.getLogger(__name__)
@@ -81,13 +82,27 @@ def _analyze_document(
         # object, which would duplicate the entire document in memory.
         doc.seek(0)
 
+        features = self._get_features()
         poller = self.di_client.begin_analyze_document(
             self.config.document_model,
             body=doc,
             locale=self.config.locale,
-            features=self._get_features(),
+            features=features,
         )
-        return poller.result()
+        result = poller.result()
+        record_usage(
+            {
+                "provider": "azure",
+                "service": "document_intelligence",
+                "model": self.config.document_model,
+                "api_version": self.config.api_version,
+                "features": [
+                    getattr(feature, "value", str(feature)) for feature in features
+                ],
+                "usage": {"pages": len(result.pages or [])},
+            }
+        )
+        return result
 
     def _get_features(self) -> list[DocumentAnalysisFeature]:
         features = list[DocumentAnalysisFeature]()

bc2/core/analyze/test_azuredi.py

@@ -0,0 +1,36 @@
+from io import BytesIO
+from unittest.mock import MagicMock
+
+from ..common.usage import (
+    create_usage_tracker,
+    usage_operation,
+    usage_tracking,
+)
+from .azuredi import AzureDIAnalyze, AzureDIAnalyzeConfig
+
+
+def test_document_intelligence_records_page_usage():
+    driver = AzureDIAnalyze.__new__(AzureDIAnalyze)
+    driver.config = AzureDIAnalyzeConfig(
+        endpoint="https://example.cognitiveservices.azure.com",
+        api_key="test",
+    )
+    result = MagicMock()
+    result.pages = [MagicMock(), MagicMock(), MagicMock()]
+    poller = MagicMock()
+    poller.result.return_value = result
+    driver.di_client = MagicMock()
+    driver.di_client.begin_analyze_document.return_value = poller
+    created = create_usage_tracker({"report_usage": True})
+    assert created is not None
+    report, tracker = created
+
+    with usage_tracking(tracker), usage_operation("analyze:azuredi"):
+        driver._analyze_document(BytesIO(b"document"))
+
+    call = report["calls"][0]
+    assert call["provider"] == "azure"
+    assert call["service"] == "document_intelligence"
+    assert call["operation"] == "analyze:azuredi"
+    assert call["model"] == "prebuilt-read"
+    assert call["usage"] == {"pages": 3}